API design

[keewis]

From GRID4EARTH/healpix-geo#132, it looks like there is a need to be able to subsample and plot data using matplotlib and cartopy.

My current understanding of the process is that it has the following steps:

1. define a target grid in lon / lat space (equirectangular projection)
2. find the cells closest to each target grid point
3. index the data to perform nearest-neighbour interpolation
4. use the standard cartopy / matplotlib functionality to plot

I'll use this discussion to design the API for the library (whose name definitely needs to be improved).

cc @jmdelouis @tinaok

[keewis]

According to GRID4EARTH/healpix-geo#132 (comment), the interface of existing packages looks something like

def lgnomproject(
    cell_ids, # (N,) HEALPix pixel indices
    data, # (N,) scalar values OR (N,3) RGB values in [0,1]
    nside: int,
    rot=None,
    xsize: int = 400,
    ysize: int = 400,
    reso: float = None,
    fov_deg=None,
    nest: bool = True,
    reduce: str = "mean", # 'mean'|'median'|'sum'|'first'
    mask_outside: bool = True,
    unseen_value=None, # defaults to hp.UNSEEN (scalar) or np.nan (RGB)
    return_image_only: bool = False,
    title: str = None, cmap: str = "viridis", vmin=None, vmax=None,
    notext: bool = False,
    hold: bool = True,
    interp: bool = False,
    sub=(1,1,1),
    cbar: bool = False,
    unit: str = "Value",
    rgb_clip=(0.0, 1.0), # clip range for RGB
):
    pass

And a usage example:

ilevel=14
ndata=128
cell_ids=34**ilevel+np.arange(ndatandata)
val=np.cos(cell_ids/100.) # fake data
lgnomproject(cell_ids,val,2**(ilevel),hold=False,fov_deg=1,interp=False,sub=(1,2,1),cbar=True)

[keewis]

I believe we can group these parameters into the following (excluding cell_ids and data):

• grid parameters: nside, nest
• sampling grid parameters: xsize, ysize, fov_deg, reso, rot
• masking strategy: mask_outside, unseen_value
• value transformation: reduce, interp, rgb_clip
• matplotlib control: title, sub, hold, cbar, unit, cmap, vmin, vmax, notext
• others: return_image_only

[keewis]

Proposed new parameters:

grid parameters

As in xdggs, the new grid parameters will be:

• level / refinement_level
• indexing_scheme
• ellipsoid

sampling grid parameters

The current sampling grid parameters are shape, step sizes, and central position although the step size has two different ways of specifying that (reso, which is explicit, and fov_deg, ).

New parameter could be an affine transform, which is a lot easier to deal with (if we stick with axis-aligned bounding boxes).

value transforms

reduce deduplicates values, rgb_clip defines the clip range for rgb values, and interp controls whether to use nearest-neighbour or bilinear interpolation.

New parameters:

• reduce could be aggregation="method" (or agg or stay reduce)
• rgb_clip does clipping, but so do vmin and vmax. Does this provide anything new, then?
• interp should most likely be interpolation="nearest" or interpolation="bilinear"

masking

There are two parameters: unseen_value that defines the fill value, and mask_outside that controls whether to replace values that are outside the domain (I'll have to figure out what "outside" means in this case).

matplotlib control

These control different things, but it may be better to trim a bit:

• title sets the ax' title
• cbar could be a bool, or a dict of options for the color bar (including the label)
• hold will be replaced by ax
• cmap, vmin, and vmax will stay, but there is also going to be a norm parameter that is forwarded to the plotting function for more control
• notext will be replaced by axis_labels, a dict containing mapping for x and / or y. If empty or None, there won't be axis labels

return_image_only

This one will instead be a separate function (which is called internally by the plotting function).

new signature

With that, the new signature would become:

class SamplingGridParameters(TypedDict):
    shape: int | tuple[int, int] = 1024  # expands to 2-tuple if an int
    resolution: float | tuple[float, float] | None = None # expands to 2-tuple if an int, derives from spatial extent and shape if `None`,
    center: tuple[float, float] | None = None # inferred from the data if `None`

@dataclass
class SamplingGrid:
    x: np.ndarray
    y: np.ndarray

    @classmethod
    def from_transform(cls, transform: affine.Affine) -> Self:
        pass

    @classmethod
    def from_parameters(cls, shape: tuple[int, int], resolution: tuple[float, float], center: tuple[float, float]) -> Self:
        pass

    @classmethod
    def from_dict(self, mapping: SamplingGridParameters) -> Self:
        pass

@dataclass
class HealpixParameters:
    level: int
    indexing_scheme: Literal["nested", "ring", "zuniq"]
    ellipsoid: EllipsoidLike

def resample(
    cell_ids: np.ndarray,
    data: np.ndarray,
    sampling_grid: SamplingGrid,
    *,
    grid_info: HealpixParameters,
    interpolation: Literal["nearest", "bilinear"],
) -> np.ndarray:
    pass

def plot(
    cell_ids: np.ndarray,
    data: np.ndarray,
    *,
    healpix_parameters: dict[str, Any] | HealpixParameters,
    sampling_grid: SamplingGridParameters | SamplingGrid,
    projection: str | cartopy.crs.CRS,
    agg: Literal["mean", "median", "std", "var"] = "mean",
    interpolation: Literal["nearest", "bilinear"] = "nearest",
    rgb_clip: tuple[float, float] = (0.0, 1.0),
    ax: matplotlib.axis.Axis | None = None,
    title: str | None = None,
    colorbar: bool | dict[str, Any] = False,
    cmap: str | matplotlib.colormap.ColorMap = "viridis",
    vmin: float | None = None,
    vmax: float | None = None,
    norm: matplotlib.norm.Norm | None = None,
    axis_labels: dict[str, str] | Literal["none"] | None = None,
) -> matplotlib.axis.Axis:
    pass

# and possibly a few more functions that select the projection in advance

Usage would be something like:

fig, axes = plt.subplots(1, 3, figsize=(15, 12))
plot(
    cell_ids,
    data,
    healpix_parameters={"level": 10, "indexing_scheme": "nested", "ellipsoid": "WGS84"},
    sampling_grid={"shape": 1024, "resolution": 0.5, "center": (45, 45)},
    projection="mollweide",
    agg="mean",
    interpolation="bilinear",
    ax=axes[0],
    norm=matplotlib.norm.CenteredNorm(center=273.15),
)

[jmdelouis]

What you describe is almost what is done here https://github.com/jmdelouis/FOSCAT/blob/main/src/foscat/Plot.py

[keewis]

I just copied what you posted, so I can more easily analyze the requirements and put it in a consistent API

[keewis]

I'll probably use cartopy (and thus indirectly pyproj) to do the projections, though