---
file_format: mystnb
kernelspec:
  name: python3
---

# Limited-area HEALPix

The HEALPix grid is designed to cover an entire sphere, be it Space or the Earth.
This raises the question of whether HEALPix can also be used to store regional cutouts, as is often the case in weather forecasting or regional climate modelling.

In this article we will try to outline some technical issues and suggest ways around them.

```{code-cell} ipython3
import cartopy.crs as ccrs
import easygems.healpix as egh
import healpix as hp
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
```

Define two helper functions to create a cartopy map with regional extent and plot boxes on it.

```{code-cell} ipython3
def get_axis_europe():
    """Create cartyop.GeoAxis with an extent centered around Europe."""
    _, ax = plt.subplots(subplot_kw={"projection": ccrs.PlateCarree()})
    ax.coastlines()
    ax.set_extent([-40, 60, 20, 65])

    return ax


def plot_extent_box(ax, extent, edgecolor="m", linewidth=2):
    """Add a rectangular patch around an extent to a cartopy.GeoAxis."""
    return ax.add_patch(
        plt.Rectangle(
            xy=[extent[0], extent[2]],
            width=extent[1] - extent[0],
            height=extent[3] - extent[2],
            facecolor="none",
            edgecolor=edgecolor,
            linewidth=linewidth,
            transform=ccrs.PlateCarree(),
        )
    )
```

## Querying grid points

In a first step, we will define a simple lat/lon box over Europe as our region of interest.

```{code-cell} ipython3
extent = [-20, 40, 30, 60]  # Europe

ax = get_axis_europe()
plot_extent_box(ax, extent);
```

Next, we can use existing HEALPix libraries to compute lat/lon coordinates at a given order (zoom level).
These coordinates can then be compared again our boundaries to query a regional cutout.

```{code-cell} ipython3
# Compute HEALPix lat/lon coordinates for a given order (zoom level)
order = zoom = 8
nside = hp.order2nside(order)
npix = hp.nside2npix(nside)

hp_lon, hp_lat = hp.pix2ang(nside, np.arange(npix), nest=True, lonlat=True)
hp_lon = (hp_lon + 180) % 360 - 180

# Find all grid points within the defined extent
icell, = np.where(
    (hp_lon > extent[0]) &
    (hp_lon < extent[1]) &
    (hp_lat > extent[2]) &
    (hp_lat < extent[3])
)

# Plot the selected points
ax = get_axis_europe()
plot_extent_box(ax, extent)

region = np.full_like(hp_lon, fill_value=np.nan)  # Full HEALPix with NaNs
region[icell] = 1  # Set selected values to 1
egh.healpix_show(region);
```

It worked! However, a technical problem with these slices is that they distort the grid index ordering[^1].

## Maintaining the grid index ordering

[^1]: You can have a look at [this article about ICON's grid](../../Processing/playing_with_triangles/grid_index_ordering) if you are not familiar with the topic.

The nested HEALPix grid ordering ensures that geographical proximity of data is mapped to proximity in the array.
However, this relationship is broken when we cut along simple lines on the map.

Fortunately, we can maintain the proximity by keeping full chunks that follow the $4^n$ relation.
We will use the `get_full_chunks()` convenience function provided by `easygems`.
This function takes an array of indices and a chunk size and returns all indices of grid points in chunks that are touched.

```{code-cell} ipython3
# Get indices of **full chunks**
ichunk = egh.get_full_chunks(icell, chunksize=4**5)

# Plot the selected points
ax = get_axis_europe()
plot_extent_box(ax, extent)

region = np.full_like(hp_lon, fill_value=np.nan)  # Full HEALPix with NaNs
region[ichunk] = 1  # Set selected values to 1
egh.healpix_show(region);
```

## Setting the Coordinate References System

An [open question](https://github.com/cf-convention/cf-conventions/issues/433) is what to store in the Coordinate Reference System (CRS).
Our suggestion is to store the information about the full global HEALPix grid.
This way downstream libraries will be able to compute the full grid.

```{code-cell} ipython3
crs = xr.DataArray(
    name="crs",
    attrs={
        "grid_mapping_name": "healpix",
        "healpix_nside": 2**zoom,
        "healpix_order": "nest",
    },
)
crs
```

In addition, the arrays in the dataset **must** contain the global indices of the original HEALPix grid.
This allows the original location of any selected grid point to be diagnosed.
You can set the values of indices outside the defined extent (not inside `icell`) to `NaN`, but the index must remain to keep the array order in tact.

Here is a complete example of a limited-area HEALPix dataset:

```{code-cell} ipython3
# Construct global HEALPix dataset with proper `cell` **coordinate** and CRS information.
ds = xr.Dataset(
    data_vars={
        "crs": crs,
    },
    coords={
        "lat": (("cell",), hp_lat),
        "lon": (("cell",), hp_lon),
        "cell": (("cell",), np.arange(npix)),
    }
)

# Select the limited-area while maintaining global index information (and chunks)
ds.isel(cell=ichunk)
```