outputs

grib

class anemoi.inference.outputs.grib.HindcastOutput(reference_year: int)

Bases: object

Hindcast output class.

Initialize the HindcastOutput object.

Parameters:

reference_year (int) – The reference year.

class anemoi.inference.outputs.grib.Matching(config)

Bases: object

patch(variable, request)
class anemoi.inference.outputs.grib.PatchOutput(*patches)

Bases: object

anemoi.inference.outputs.grib.modifier_factory(modifiers: list) → list

Create a list of modifier instances.

Parameters:

modifiers (list) – A list of modifier configurations.

Returns:

A list of modifier instances.

Return type:

list

class anemoi.inference.outputs.grib.BaseGribOutput(context: dict, metadata: Metadata, *, post_processors: list[str | dict[str, Any]] | None = None, encoding: dict[str, Any] | None = None, templates: list[str] | str | None = None, grib1_keys: dict[str, Any] | None = None, grib2_keys: dict[str, Any] | None = None, modifiers: list[str] | None = None, variables: OutputVariableConfig | str | list[str] | dict[Literal['select', 'drop'], str] | None = None, output_frequency: int | None = None, write_initial_state: bool | None = None, negative_step_mode: Literal['error', 'write', 'skip'] = 'error')

Bases: Output

Handles grib.

Initialise the GribOutput object.

Parameters:

  • context (dict) – The context dictionary.

  • metadata (Metadata) – Metadata corresponding to the dataset this output is handling.

  • post_processors (Optional[List[ProcessorConfig]] = None) – Post-processors to apply to the input

  • encoding (dict, optional) – The encoding dictionary, by default None.

  • templates (list or str, optional) – The templates list or string, by default None.

  • grib1_keys (dict, optional) – The grib1 keys dictionary, by default None.

  • grib2_keys (dict, optional) – The grib2 keys dictionary, by default None.

  • modifiers (list, optional) – The list of modifiers, by default None.

  • output_frequency (int, optional) – The frequency of output, by default None.

  • write_initial_state (bool, optional) – Whether to write the initial state, by default None.

  • variables (OutputVariableConfigUnion) – Variable settings for inclusion/exclusion, by default None.

  • negative_step_mode (Literal["error", "write", "skip"], optional) – What to do when writing a variable that has a base time before the forecast base time. This can happen when the initial conditions contain an accumulated variable, or a variable period is longer than the model step time. In all cases a warning will be shown. - error: (default) raise an exception - write: write the variable as normal - skip: skip writing the variable

write_initial_state(state: dict[str, Any]) → None

Write the initial step of the state.

Parameters:

state (State) – The state object.

write_step(state: dict[str, Any]) → None

Write a step of the state.

Parameters:

state (State) – The state object.

abstractmethod write_message(message: ndarray[tuple[Any, ...], dtype[Any]], *args: Any, **kwargs: Any) → None

Write a message to the grib file.

Parameters:

  • message (FloatArray) – The message array.

  • *args (Any) – Additional arguments.

  • **kwargs (Any) – Additional keyword arguments.

template(state: dict[str, Any], name: str) → object

Get the template for a variable.

Parameters:

  • state (State) – The state object.

  • name (str) – The variable name.

Returns:

The template object.

Return type:

object

template_lookup(name: str) → dict

Lookup the template for a variable.

Parameters:

name (str) – The variable name.

Returns:

The template dictionary.

Return type:

dict

gribfile

class anemoi.inference.outputs.gribfile.ArchiveCollector

Bases: object

Collects archive requests.

UNIQUE = {'date', 'expver', 'hdate', 'referenceDate', 'stream', 'time', 'type'}
add(field: dict[str, Any]) → None

Add a field to the archive request.

Parameters:

field (Dict[str,Any]) – The field dictionary.

property request: dict[str, Any]

Get the archive request.

class anemoi.inference.outputs.gribfile.GribIoOutput(context: Context, metadata: Metadata, *, out: Path | IOBase, post_processors: list[str | dict[str, Any]] | None = None, encoding: dict[str, Any] | None = None, archive_requests: dict[str, Any] | None = None, check_encoding: bool = True, templates: list[str] | str | None = None, grib1_keys: dict[str, Any] | None = None, grib2_keys: dict[str, Any] | None = None, modifiers: list[str] | None = None, variables: OutputVariableConfig | str | list[str] | dict[Literal['select', 'drop'], str] | None = None, output_frequency: int | None = None, write_initial_state: bool | None = None, split_output: bool = True, negative_step_mode: Literal['error', 'write', 'skip'] = 'error', missing_value: int | float = -9999)

Bases: BaseGribOutput

Output class for grib io.

This class handles writing grib and collecting archive requests. It extends the BaseGribOutput class and implements the write_message method.

Initialize the GribIOOutput.

Parameters:

  • context (Context) – The context.

  • metadata (Metadata) – Metadata corresponding to the dataset this output is handling.

  • out (Union[Path, IOBase]) – Path or file-like object to write the grib data to. If a file-like object, it should be opened in binary write mode.

  • post_processors (Optional[List[ProcessorConfig]], default None) – Post-processors to apply to the input

  • encoding (dict, optional) – The encoding dictionary, by default None.

  • archive_requests (dict, optional) – The archive requests dictionary, by default None.

  • check_encoding (bool, optional) – Whether to check encoding, by default True.

  • templates (list or str, optional) – The templates list or string, by default None.

  • grib1_keys (dict, optional) – The grib1 keys dictionary, by default None.

  • missing_value (int or float, optional) – Sentinel value used to replace NaNs before GRIB bitmap encoding (requires allow_nans=True). Default -9999 avoids collision with physically valid geophysical values which are always positive for most fields (e.g. z_850 ~ 9999 m2/s2 over Switzerland).

  • grib2_keys (dict, optional) – The grib2 keys dictionary, by default None.

  • modifiers (list, optional) – The list of modifiers, by default None.

  • output_frequency (int, optional) – The frequency of output, by default None.

  • write_initial_state (bool, optional) – Whether to write the initial state, by default None.

  • variables (OutputVariableConfigUnion) – Variable settings for inclusion/exclusion, by default None.

  • split_output (bool, optional) – Whether to split the output, by default True. Cannot be True if out is a file-like object.

  • negative_step_mode (Literal["error", "write", "skip"], optional) – What to do when writing a variable that has a base time before the forecast base time. This can happen when the initial conditions contain an accumulated variable, or a variable period is longer than the model step time. In all cases a warning will be shown. - error: (default) raise an exception - write: write the variable as normal - skip: skip writing the variable

write_message(message: ndarray[tuple[Any, ...], dtype[Any]], template: Field, **keys: dict[str, Any]) → None

Write a message to the grib file.

Parameters:

  • message (FloatArray) – The message array.

  • template (ekd.Field) – A ekd.Field use as a template for GRIB encoding.

  • **keys (Dict[str, Any]) – Additional keys for the message.

collect_archive_requests(written: tuple, template: object, **keys: Any) → None

Collect archive requests.

Parameters:

  • written (tuple) – The written tuple.

  • template (object) – The template object.

  • **keys (Any) – Additional keys for the archive requests.

close() → None

Close the grib file.

gribmemory

class anemoi.inference.outputs.gribmemory.GribMemoryOutput(context: Context, metadata: Metadata, *, out: IOBase, post_processors: list[str | dict[str, Any]] | None = None, encoding: dict[str, Any] | None = None, archive_requests: dict[str, Any] | None = None, check_encoding: bool = True, templates: list[str] | str | None = None, grib1_keys: dict[str, Any] | None = None, grib2_keys: dict[str, Any] | None = None, modifiers: list[str] | None = None, variables: OutputVariableConfig | str | list[str] | dict[Literal['select', 'drop'], str] | None = None, output_frequency: int | None = None, write_initial_state: bool | None = None)

Bases: GribIoOutput

Handles grib files in memory.

Initialize the GribFileOutput.

Parameters:

  • context (Context) – The context.

  • metadata (Metadata) – Metadata corresponding to the dataset this output is handling.

  • out (IOBase) – Output stream or file-like object for writing GRIB data.

  • post_processors (Optional[List[ProcessorConfig]], default None) – Post-processors to apply to the input

  • encoding (dict, optional) – The encoding dictionary, by default None.

  • archive_requests (dict, optional) – The archive requests dictionary, by default None.

  • check_encoding (bool, optional) – Whether to check encoding, by default True.

  • templates (list or str, optional) – The templates list or string, by default None.

  • grib1_keys (dict, optional) – The grib1 keys dictionary, by default None.

  • grib2_keys (dict, optional) – The grib2 keys dictionary, by default None.

  • modifiers (list, optional) – The list of modifiers, by default None.

  • output_frequency (int, optional) – The frequency of output, by default None.

  • write_initial_state (bool, optional) – Whether to write the initial state, by default None.

  • variables (OutputVariableConfigUnion) – Variable settings for inclusion/exclusion, by default None.

netcdf

none

class anemoi.inference.outputs.none.NoneOutput(context: Context, metadata: Metadata, *, variables: OutputVariableConfig | str | list[str] | dict[Literal['select', 'drop'], str] | None = None, post_processors: list[str | dict[str, Any]] | None = None, output_frequency: int | None = None, write_initial_state: bool | None = None)

Bases: Output

None output class.

Initialize the Output object.

Parameters:

  • context (Context) – The context in which the output operates.

  • metadata (Metadata) – Metadata corresponding to the dataset this output is handling.

  • variables (OutputVariableConfigUnion) – Variable settings for inclusion/exclusion.

  • post_processors (Optional[List[ProcessorConfig]], default None) – Post-processors to apply to the output

  • output_frequency (Optional[int], optional) – The frequency at which to output states, by default None.

  • write_initial_state (Optional[bool], optional) – Whether to write the initial state, by default None.

write_step(state: dict[str, Any]) → None

Write a step of the state.

Parameters:

state (State) – The state dictionary.

parallel

class anemoi.inference.outputs.parallel.Chunker

Bases: object

Chunking strategies for dividing the state into smaller parts for parallel outputting.

Parameters:

typed_variables (dict[str, Variable]) – Dictionary of field name to its corresponding Variable metadata.

by_metadata(keys: list[str], max_groups: int = -1) → Callable[[dict[str, Any]], Generator[dict[str, Any], None, None]]

Chunk the state into smaller parts based on the specified metadata keys, ensuring that variables with the same key=value pair are kept together in the same chunk.

by_size(fields_per_chunk: int) → Callable[[dict[str, Any]], Generator[dict[str, Any], None, None]]

Chunk the state into smaller parts, each containing a specified number of fields.

by_worker() → Callable[[dict[str, Any]], Generator[dict[str, Any], None, None]]

Chunk the state into smaller parts, one for each worker.

from_config(chunk_strategy: Literal['by_worker', 'by_metadata', 'by_size'] | dict[Literal['by_worker', 'by_metadata', 'by_size'], dict[str, Any]]) → Callable[[dict[str, Any]], Generator[dict[str, Any], None, None]]

Build a chunking function from a chunk_strategy configuration.

Parameters:

chunk_strategy (CHUNK_STRATEGIES | dict[CHUNK_STRATEGIES, dict[str, Any]]) – Either a strategy name (one of VALID_CHUNK_STRATEGIES) or a dictionary with a single key being the strategy name and the value being a dictionary of keyword arguments for that strategy.

Returns:

The chunking function that splits a state into chunks.

Return type:

Callable[[State], Generator[State, None, None]]

class anemoi.inference.outputs.parallel.MessageType(value)

Bases: str, Enum

Types of messages sent from the main process to the writer processes. Used for logging and control flow in the writer loop.

TERMINATE = 'Terminate'
INITIAL_STATE = 'InitialState'
STATE = 'State'
OPEN = 'Open'
class anemoi.inference.outputs.parallel.ParallelOutput(context: Context, metadata: Metadata, *, output: Output | Any | None = None, num_writers: int = 1, chunk_strategy: Literal['by_worker', 'by_metadata', 'by_size'] | dict[Literal['by_worker', 'by_metadata', 'by_size'], dict[str, Any]] = 'by_worker', **kwargs: Any)

Bases: Output

Wraps another Output and offloads write_state calls to one or more forked writer processes. The output is split along the field dimension and each chunk is sent to a different writer process via multiprocessing queues. Each writer process writes the initial state into its own file.

When writing a file output, a suffix ‘_{writer_id}’ is appended to the file name to avoid conflicts between writers.

Usage in YAML:

output:
  parallel:
    num_writers: 2
    output:
      grib:
        path: output.grib
This yaml will result in the following outputs being written:

  • output_w0.grib

  • output_w1.grib

Initialise the ParallelOutput.

Parameters:

  • context (Context) – The inference context.

  • metadata (Metadata) – Metadata for the dataset.

  • output (Output | Any | None) – The inner output (or its config dict) that will be forked into writer processes.

  • num_writers (int) – Number of writer processes to spawn. Must be >= 1. Defaults to 1 (single output file, asynchronous writes).

  • chunk_strategy (CHUNK_STRATEGIES | dict[CHUNK_STRATEGIES, dict[str, Any]], default="by_worker") –

    The strategy for chunking the output among writer processes. Can be a string (one of “by_size”, “by_metadata”, “by_worker”) or a dictionary with a single key being the strategy name and the value being a dictionary of keyword arguments for that strategy. - by_worker (default), splits the fields into exactly num_writers contiguous

    chunks of roughly equal size, one per writer. No additional arguments are needed.

    • by_size, requires an additional argument fields_per_chunk specifying the number of fields per chunk.

    • by_metadata, requires an additional argument keys specifying a list of metadata keys to keep in a chunk.

  • **kwargs (Any) – Forwarded to the inner output.

open(state: dict[str, Any]) → None

Spawn the writer processes during open() instead of __init__() to ensure they have access to the full context.

Pass the open() message to the writers so that they can call the output-appropriate open() method.

write_step(state: dict[str, Any]) → None

Write a step of the state.

Parameters:

state (State) – The state to write.

write_state(state: dict[str, Any], message=MessageType.STATE) → None

Write the state, dispatching to writer processes when enabled.

dispatch_state_to_writers(state: dict[str, Any], message=MessageType.STATE) → None

Send the state to each writer process via multiprocessing queues.

Takes an optional ‘message’ argument to indicate the type of message being sent, which is used for control flow in the writer loop.

close() → None

Terminate writer processes, then close the wrapped output.

print_summary(depth: int = 0) → None

Print a summary of the output configuration.

Parameters:

depth (int, optional) – The indentation depth for the summary, by default 0.

write_initial_state(state: dict[str, Any]) → None

Write the initial state.

plot

anemoi.inference.outputs.plot.fix(lons: ndarray[tuple[Any, ...], dtype[Any]]) → ndarray[tuple[Any, ...], dtype[Any]]

Fix longitudes greater than 180 degrees.

Parameters:

lons (FloatArray) – Array of longitudes.

Returns:

Fixed array of longitudes.

Return type:

FloatArray

printer

anemoi.inference.outputs.printer.print_state(state: dict[str, ~typing.Any], print: ~typing.Callable[[...], None] = <built-in function print>, max_lines: int = 4, variables: ~anemoi.inference.schemas.OutputVariableConfig | str | list[str] | dict[~typing.Literal['select', 'drop'], str] | None = None) → None

Print the state.

Parameters:

  • state (State) – The state dictionary.

  • print (function, optional) – The print function to use, by default print.

  • max_lines (int, optional) – The maximum number of lines to print, by default 4. If variables is provided, this option is ignored.

  • variables (list, optional) – The list of variables to print, by default None. This should match the structure described in OutputVariableConfig.

raw

tee

class anemoi.inference.outputs.tee.TeeOutput(context: Context, metadata: Metadata, *args, outputs: Sequence | None = None, **kwargs: Any)

Bases: Output

TeeOutput class to manage multiple outputs.

Initialise the TeeOutput.

Parameters:

  • context (object) – The context object.

  • metadata (Metadata) – Metadata corresponding to the dataset this output is handling.

  • *args – Additional positional arguments.

  • outputs (Sequence, optional) – Outputs to be created.

  • **kwargs (Any) – Additional keyword arguments.

write_initial_state(state: dict[str, Any]) → None

Write the initial state to all outputs.

Parameters:

state (State) – The state dictionary.

write_state(state: dict[str, Any]) → None

Write the state to all outputs.

Parameters:

state (State) – The state dictionary.

write_step(state: dict[str, Any]) → None

Raise NotImplementedError as TeeOutput does not support write_step.

Parameters:

state (State) – The state dictionary.

open(state: dict[str, Any]) → None

Open all outputs.

Parameters:

state (State) – The state dictionary.

close() → None

Close all outputs.

print_summary(depth: int = 0) → None

Print the summary of all outputs.

Parameters:

depth (int, optional) – The depth of the summary.

truth

anemoi.inference.outputs.truth.create_diagnostic_input(context, metadata, dataset_name)

zarr

anemoi.inference.outputs.zarr.create_zarr_array(store: StoreLike, name: str, shape: tuple, dtype: str, dimensions: tuple[str, ...], chunks: tuple[int, ...] | str | bool, fill_value: float | None = None) → Any

Create a Zarr array with the given parameters.

Parses the Zarr version to handle differences in API between versions 2 and 3.