Inference Config

The following is an example configuration for running inference. While you can use absolute paths in the config yamls (we encourage it!), the example uses paths relative to the directory you run the command. The example assumes you are running in a directory structure like:

.
├── ace2_era5_ckpt.tar
├── initial_conditions
│   ├── ic_1940.nc
│   ├── ic_1950.nc
│   ├── ...
│   └── ic_2020.nc
├── forcing_data
│   ├── forcing_1940.nc
│   ├── forcing_1941.nc
│   ├── ...
│   └── forcing_1989.nc
└── inference-config.yaml

that includes a model checkpoint (ace2_era5_ckpt.tar), forcing data (e.g., forcing_1940.nc), and initial conditions (e.g., ic_1940.nc). You can find the checkpoint and forcing and initial condition data in the ACE2-ERA5 Hugging Face page.

The specified initial condition file should contain a time dimension of at least length 1, but can also contain multiple times. If multiple times are present and start_indices is not specified in the fme.ace.InitialConditionConfig configuration, the inference will run an ensemble using all times in the initial condition file. The ic_1940.nc file is an example of a file with multiple times, containing initial conditions for each month of 1940. For examples of selecting specific initial conditions, see fme.ace.InitialConditionConfig Examples.

While netCDFs files are specified in the example, zarr stores are also compatible, e.g., specifying the parent folder containing the zarr store directory as the path, setting engine to “zarr”, and setting file_pattern to “<zarr_store_name>.zarr” in the dataset configuration. See fme.ace.XarrayDataConfig for more information.

Example YAML Configuration

Example YAML Configuration
experiment_dir: inference_output
n_forward_steps: 400 # 100 days
forward_steps_in_memory: 50
checkpoint_path: ace2_era5_ckpt.tar
logging:
  log_to_screen: true
  log_to_wandb: false
  log_to_file: true
  project: ace
  entity: your_wandb_entity
initial_condition:
  path: initial_conditions/ic_1940.nc
  start_indices:
    n_initial_conditions: 2
    first: 0
    interval: 3
forcing_loader:
  dataset:
    data_path: forcing_data
  num_data_workers: 2
data_writer:
  save_prediction_files: false

Configuration structure

We use the Builder pattern to load this configuration into a multi-level dataclass structure. The configuration is divided into several sub-configurations, each with its own dataclass. The top-level configuration is the fme.ace.InferenceConfig class.

class fme.ace.InferenceConfig(experiment_dir, n_forward_steps, checkpoint_path, logging, initial_condition, forcing_loader, forward_steps_in_memory=10, data_writer=<factory>, aggregator=<factory>, stepper_override=None, allow_incompatible_dataset=False, labels=<factory>)[source]

Bases: object

Configuration for running inference.

Parameters:

  • experiment_dir (str) – Directory to save results to.

  • n_forward_steps (int) – Number of steps to run the model forward for.

  • checkpoint_path (str) – Path to stepper checkpoint to load.

  • logging (LoggingConfig) – Configuration for logging.

  • initial_condition (InitialConditionConfig) – Configuration for initial condition data.

  • forcing_loader (ForcingDataLoaderConfig) – Configuration for forcing data.

  • forward_steps_in_memory (int, default: 10) – Number of forward steps to complete in memory at a time.

  • data_writer (DataWriterConfig, default: <factory>) – Configuration for data writers.

  • aggregator (InferenceAggregatorConfig, default: <factory>) – Configuration for inference aggregator.

  • stepper_override (Optional[StepperOverrideConfig], default: None) – Configuration for overriding select stepper configuration options at inference time (optional).

  • allow_incompatible_dataset (bool, default: False) – If True, allow the dataset used for inference to be incompatible with the dataset used for stepper training. This should be used with caution, as it may allow the stepper to make scientifically invalid predictions, but it can allow running inference with incorrectly formatted or missing grid information.

  • labels (list[str], default: <factory>) – Dataset labels to use for inference. If provided, these labels will be provided to the stepper for every initial condition.

The sub-configurations are:

class fme.ace.LoggingConfig(project='ace', entity='ai2cm', log_to_screen=True, log_to_file=True, log_to_wandb=True, log_format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=20, wandb_dir_in_experiment_dir=False)[source]

Bases: object

Configuration for logging.

Parameters:

  • project (str, default: 'ace') – Name of the project in Weights & Biases.

  • entity (str, default: 'ai2cm') – Name of the entity in Weights & Biases.

  • log_to_screen (bool, default: True) – Whether to log to the screen.

  • log_to_file (bool, default: True) – Whether to log to a file.

  • log_to_wandb (bool, default: True) – Whether to log to Weights & Biases.

  • log_format (str, default: '%(asctime)s - %(name)s - %(levelname)s - %(message)s') – Format of the log messages.

  • level (str | int, default: 20) – Sets the logging level.

  • wandb_dir_in_experiment_dir (bool, default: False) – Whether to create the wandb_dir in the experiment_dir or in local /tmp (default False).

class fme.ace.InitialConditionConfig(path, engine='netcdf4', start_indices=None)[source]

Bases: object

Configuration for initial conditions.

Note

The data specified under path should contain a time dimension of at least length 1. If multiple times are present in the dataset specified by path, the inference will start an ensemble simulation using each IC along a leading sample dimension. Specific times can be selected from the dataset by using start_indices.

Parameters:
class fme.ace.ForcingDataLoaderConfig(dataset, num_data_workers=0, perturbations=None, persistence_names=None)[source]

Bases: object

Configuration for the forcing data.

Parameters:

  • dataset (XarrayDataConfig | MergeNoConcatDatasetConfig) – Configuration to define the dataset.

  • num_data_workers (int, default: 0) – Number of parallel workers to use for data loading.

  • perturbations (Optional[SSTPerturbation], default: None) – Configuration for SST perturbations used in forcing data.

  • persistence_names (Optional[Sequence[str]], default: None) – Names of variables for which all returned values will be the same as the initial condition. When evaluating initial condition predictability, set this to forcing variables that should not be updated during inference (e.g. surface temperature).

class fme.ace.XarrayDataConfig(data_path, file_pattern='*.nc', n_repeats=1, engine='netcdf4', spatial_dimensions='latlon', subset=<factory>, infer_timestep=True, dtype='float32', overwrite=<factory>, fill_nans=None, isel=<factory>, labels=<factory>)[source]

Bases: DatasetConfigABC

Parameters:

  • data_path (str) – Path to the data.

  • file_pattern (str, default: '*.nc') – Glob pattern to match files in the data_path.

  • n_repeats (int, default: 1) – Number of times to repeat the dataset (in time). It is up to the user to ensure that the input dataset to repeat results in data that is reasonably continuous across repetitions.

  • engine (Literal['netcdf4', 'h5netcdf', 'zarr'], default: 'netcdf4') – Backend used in xarray.open_dataset call.

  • spatial_dimensions (Literal['healpix', 'latlon'], default: 'latlon') – Specifies the spatial dimensions for the grid, default is lat/lon. If ‘latlon’, it is assumed that the last two dimensions are latitude and longitude, respectively. If ‘healpix’, it is assumed that the last three dimensions are face, height, and width, respectively.

  • subset (Slice | TimeSlice | RepeatedInterval, default: <factory>) – Slice defining a subset of the XarrayDataset to load. This can either be a Slice of integer indices or a TimeSlice of timestamps. This feature is applied directly to the dataset samples. For example, if the file(s) have the time coordinate (t0, t1, t2, t3) and requirements.n_timesteps=2, then subset=Slice(stop=2) will provide two samples: (t0, t1), (t1, t2).

  • infer_timestep (bool, default: True) – Whether to infer the timestep from the provided data. This should be set to True (the default) for ACE training. It may be useful to toggle this to False for applications like downscaling, which do not depend on the timestep of the data and therefore lack the additional requirement that the data be ordered and evenly spaced in time. It must be set to True if n_repeats > 1 in order to be able to infer the full time coordinate.

  • dtype (Optional[str], default: 'float32') – Data type to cast the data to. If None, no casting is done. It is required that ‘torch.{dtype}’ is a valid dtype.

  • overwrite (OverwriteConfig, default: <factory>) – Optional OverwriteConfig to overwrite loaded field values.

  • fill_nans (Optional[FillNaNsConfig], default: None) – Optional FillNaNsConfig to fill NaNs with a constant value.

  • isel (Mapping[str, Slice | int], default: <factory>) – Optional xarray isel arguments to be passed to the dataset. Will raise ValueError if time is included here, since the subset argument is used specifically for selecting times. Horizontal dimensions are also not currently supported.

  • labels (list[str], default: <factory>) – Optional list of labels to be returned with the data.

Examples

If data is stored in a directory with multiple netCDF files which can be concatenated along the time dimension, use:

>>> fme.ace.XarrayDataConfig(data_path="/some/directory", file_pattern="*.nc")

If data is stored in a single zarr store at /some/directory/dataset.zarr, use:

>>> fme.ace.XarrayDataConfig(
...     data_path="/some/directory",
...     file_pattern="dataset.zarr",
...     engine="zarr"
... )
class fme.ace.DataWriterConfig(save_prediction_files=True, save_monthly_files=True, names=None, time_coarsen=None, files=None)[source]

Bases: object

Configuration for inference data writers.

Parameters:

  • save_prediction_files (bool, default: True) – Whether to enable writing of netCDF files containing the predictions and target values.

  • save_monthly_files (bool, default: True) – Whether to enable writing of netCDF files containing the monthly predictions and target values.

  • names (Optional[Sequence[str]], default: None) – Names of variables to save in the prediction and monthly netCDF files.

  • time_coarsen (Optional[TimeCoarsenConfig], default: None) – Configuration for time coarsening of written outputs to the raw data writer.

  • files (Optional[list[FileWriterConfig]], default: None) – Configuration for a sequence of individual data writers.

class fme.ace.InferenceAggregatorConfig(time_mean_reference_data=None, log_global_mean_time_series=True)[source]

Bases: object

Configuration for inference aggregator.

Parameters:

  • time_mean_reference_data (Optional[str], default: None) – Path to reference time means to compare against.

  • log_global_mean_time_series (bool, default: True) – Whether to log global mean time series metrics.

class fme.ace.StepperOverrideConfig(ocean='keep', multi_call='keep')[source]

Bases: object

Configuration for overriding stepper configuration options.

The default value for each parameter is "keep", which denotes that the serialized stepper’s configuration will not be modified when loaded. Passing other values will override the configuration of the loaded stepper.

Parameters:

  • ocean (Union[Literal['keep'], OceanConfig, None], default: 'keep') – Ocean configuration to override that used in producing a serialized stepper.

  • multi_call (Union[Literal['keep'], MultiCallConfig, None], default: 'keep') – MultiCall configuration to override that used in producing a serialized stepper.

fme.ace.InitialConditionConfig Examples

The start_indices attribute can be used to specify which initial conditions to use when multiple are present in the dataset (instead of using all available). The following examples show example selections using the yaml builder pattern for an InitialConditionConfig.

fme.ace.InferenceInitialConditionIndices

Select a number of regularly spaced initial conditions.

path: initial_conditions.nc
start_indices:
  n_initial_conditions: 3
  first: 1
  interval: 2

fme.ace.TimestampList

Selecting two timestamps from the initial conditions.

path: initial_conditions.nc
start_indices:
   times:
      - "2021-01-01T00:00:00"
      - "2021-02-01T00:00:00"

fme.ace.ExplicitIndices

Selecting specific indices from the initial conditions.

path: initial_conditions.nc
start_indices:
   list: [0, 3, 7]