Extension (rojak.atmosphere.feature_composite.extension)

Interfaces for adding a new kind of feature or pairing to rojak.atmosphere.feature_composite.

Users running composites with the existing features do not need this module.

Adding a feature. A feature has two halves, so that no dask collection is ever sent to a task:

1. A FeatureSource, the graph side, holds the lazy, time-chunked fields the feature is found in. Its chunk_inputs() returns them by name, each with dimensions (time, lat, lon) on the same grid and chunks, and its tabulator() returns the task side.

2. A FeatureTabulator, the task side, runs in each chunk task on the in-memory blocks of those fields and returns the chunk’s centres as a FeatureChunk, with the coordinates of the chunk given as a ChunkInfo. It is sent to every task, so it must hold parameters only, never a dask collection or a data-sized array.

Implement both as frozen dataclasses. Declare the FeatureGeometry with a one-line geometry property, and override has_labels to return True if the tabulator labels the region each centre owns:

@dataclass(frozen=True)
class TableTabulator(FeatureTabulator):
    centres: pd.DataFrame  # Small: every centre, with a Column.TIME_INDEX

    def __call__(self, blocks: Mapping[str, np.ndarray], chunk: ChunkInfo) -> FeatureChunk:
        stop = chunk.first_time_index + chunk.times.size
        in_chunk = self.centres[self.centres[Column.TIME_INDEX].between(chunk.first_time_index, stop - 1)]
        return FeatureChunk(centres=in_chunk.reset_index(drop=True), labels=None)


@dataclass(frozen=True, eq=False)
class TableFeature(FeatureSource):
    field: xr.DataArray  # Only gives the time-chunked grid
    centres: pd.DataFrame

    @property
    def geometry(self) -> FeatureGeometry:
        return FeatureGeometry.POINT

    def chunk_inputs(self) -> Mapping[str, xr.DataArray]:
        return {"values": self.field}

    def tabulator(self) -> FeatureTabulator:
        return TableTabulator(self.centres)

Adding a pairing. Subclass PairingStrategy, implement pair so that it returns the same pair table as the existing strategies, and override check_compatible to reject features it cannot handle.

Classes

ChunkInfo(times, first_time_index, lats, lons)	Coordinates of the time chunk that a FeatureTabulator is called on.
FeatureChunk(centres, labels)	The centres of a feature found in one time chunk, as returned by a FeatureTabulator.
FeatureGeometry(*values)	Shape of one instance of a feature.
FeatureSource()	Graph side of a feature: holds the lazy fields it is found in and creates its FeatureTabulator.
FeatureTabulator()	Task side of a feature: tabulates its centres from the in-memory blocks of one time chunk.
PairingStrategy()	How the observations of a time chunk are matched with the centres of a feature in that chunk.