processor

class anemoi.inference.processor.Processor(context: Context, metadata: Metadata)

Bases: ABC

Abstract base class for processors.

Parameters:

context (Context) – The context in which the processor operates.
metadata (Metadata) – Metadata corresponding to the dataset this processor is handling.

abstractmethod process(state: dict[str, Any]) → dict[str, Any]

Process the given state.

Parameters:: state (State) – The state to be processed.
Returns:: The processed state.
Return type:: State

patch_data_request(data_request: dict[str, Any]) → dict[str, Any]

Override if a processor needs to patch the data request (e.g. mars or cds).

Parameters:: data_request (DataRequest) – The data request to be patched.
Returns:: The patched data request.
Return type:: DataRequest

pre_processors

extract

class anemoi.inference.pre_processors.extract.ExtractBase(context: Context, metadata: Metadata)

Bases: Processor

Base class for pre-processors that extract data from the state.

property indexer: ndarray[tuple[Any, ...], dtype[Any]] | slice

process(state: dict[str, Any]) → dict[str, Any]: Process the state to extract a subset of points based on the indexer.

class anemoi.inference.pre_processors.extract.ExtractSlice(context: Context, metadata: Metadata, *slice_args: int)

Bases: ExtractBase

Extract a subset of points from the state based on a slice.

Parameters:

context (Context) – The context in which the processor is running.
metadata (Metadata) – Metadata corresponding to the dataset this processor is handling.
slice_args (int) – Arguments to create a slice object. This can be a single integer or a tuple of integers representing the start, stop, and step of the slice.

forward_transform_filter

mask

class anemoi.inference.pre_processors.mask.MaskValues(context: Context, *, mask: str, param: str | list[str], **kwargs)

Bases: Processor

Replace values in a field with nans from a specified mask in the supporting arrays.

process(state: dict[str, Any]) → dict[str, Any]

Apply the mask to the specified parameters in the state.

Parameters:: state (State) – The state containing the fields to be masked.
Returns:: The state with the specified parameters masked.
Return type:: State

no_missing_values

class anemoi.inference.pre_processors.no_missing_values.NoMissingValues(context: Context, metadata: Metadata, **kwargs: Any)

Bases: Processor

Replace NaNs with mean.

process(state: dict[str, Any]) → dict[str, Any]

Process the fields to replace NaNs with the mean value.

Parameters:: state (State) – The state to process.
Returns:: List of processed state with NaNs replaced by the mean value.
Return type:: list

post_processors

accumulate

class anemoi.inference.post_processors.accumulate.Accumulate(context: Context, metadata: Metadata, *, accumulations: list[str] | None = None, allow_negative: bool = False)

Bases: Processor

Accumulate fields from zero and return the accumulated fields.

Parameters:

context (Any) – The context in which the processor is running.
metadata (Metadata) – Metadata corresponding to the dataset this processor is handling.
accumulations (Optional[List[str]], optional) – List of fields to accumulate, by default None. If None, the fields are taken from the context’s checkpoint.
allow_negative (bool, optional) – Whether to allow negative values in the accumulation, by default False.

process(state: dict[str, Any]) → dict[str, Any]

Process the state to accumulate specified fields.

Parameters:: state (State) – The state containing fields to be accumulated.
Returns:: The updated state with accumulated fields.
Return type:: State

assign

backward_transform_filter

class anemoi.inference.post_processors.backward_transform_filter.ForwardTransformFilter(*args: Any, use_forward: bool = False, **kwargs: Any)

Bases: BackwardTransformFilter

Apply a transform forward or reversed backward as a post-processor.

earthkit_state

Code to wrap and unwrap state dictionaries into earthkit.data field lists. So that we can pass them through anemoi-transorm filters.

class anemoi.inference.post_processors.earthkit_state.StateFieldGeography(field: Any)

Bases: object

Geographical information of a state field.

Parameters:: field (Any) – The field containing geographical data.

property shape: tuple[int, ...]

Shape of the geographical field.

Type:: Tuple

class anemoi.inference.post_processors.earthkit_state.StateFieldMetadata(field: Any)

Bases: RawMetadata

Metadata for a state field.

Parameters:: field (Any) – The field containing metadata.

as_namespace(ns: str) → dict[str, Any]

Convert metadata to a specific namespace.

Parameters:: ns (str) – The namespace to convert the metadata to.
Returns:: The metadata in the specified namespace.
Return type:: Dict[str, Any]

property geography: StateFieldGeography

Geographical information of the field.

Type:: StateFieldGeography

class anemoi.inference.post_processors.earthkit_state.StateField(*args, **kwargs)

Bases: Field

State field containing name, values, and state information.

Parameters:

name (str) – The name of the field.
values (FloatArray) – The values of the field.
state (Dict[str, Any]) – The state information associated with the field.

property shape: tuple[int, ...]

Shape of the field.

Type:: Tuple

anemoi.inference.post_processors.earthkit_state.wrap_state(state: dict[str, Any]) → FieldList

Transform a state dictionary into an earthkit.data field list.

Parameters:: state (Dict[str, Any]) – The state dictionary to be transformed.
Returns:: The transformed field list.
Return type:: ekd.FieldList

anemoi.inference.post_processors.earthkit_state.unwrap_state(fields: FieldList, state: dict[str, Any], namer: Callable) → dict[str, Any]

Transform a earthkit.data field list into a state dictionary.

Parameters:

fields (ekd.FieldList) – The field list to be transformed.
state (State) – The original state dictionary.
namer (Callable) – A function to generate new field names.

Returns:

The transformed state dictionary.

Return type:

Dict[str, Any]

extract

class anemoi.inference.post_processors.extract.ExtractBase(context: Context, metadata: Metadata)

Bases: Processor

Base class for processors that extract data from the state.

property indexer: ndarray[tuple[Any, ...], dtype[Any]] | slice

process(state: dict[str, Any]) → dict[str, Any]

Process the state to extract a subset of points based on the indexer.

Parameters:: state (State) – The state containing fields to be extracted.
Returns:: The updated state with extracted fields.
Return type:: State

class anemoi.inference.post_processors.extract.ExtractSlice(context: Context, metadata: Metadata, *slice_args: int)

Bases: ExtractBase

Extract a subset of points from the state based on a slice.

Parameters:

context (Context) – The context in which the processor is running.
metadata (Metadata) – Metadata corresponding to the dataset this processor is handling.
slice_args (int) – Arguments to create a slice object. This can be a single integer or a tuple of integers representing the start, stop, and step of the slice.