Evaluator Config

The following is an example configuration for running inference while evaluating against target data. 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:

.
├── ckpt.tar
└── validation
    ├── data1.nc  # files can have any name, but must sort into time-sequential order
    ├── data2.nc  # can have any number of netCDF files
    └── ...

The .nc files correspond to data files like training_validation_data/training_validation/1940010100.nc in the ACE2-ERA5 Hugging Face page, while ckpt.tar corresponds to a file like ace2_era5_ckpt.tar in that repository.

Example YAML Configuration
experiment_dir: evaluator_output
n_forward_steps: 400  # 100 days
forward_steps_in_memory: 50
checkpoint_path: ckpt.tar
logging:
  log_to_screen: true
  log_to_wandb: false
  log_to_file: true
  project: ace
  entity: your_wandb_entity
loader:
  dataset:
    data_path: validation
  start_indices:
    first: 0
    n_initial_conditions: 1
  num_data_workers: 4
data_writer:
  save_prediction_files: false
  save_monthly_files: false

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.InferenceEvaluatorConfig class.

class fme.ace.InferenceEvaluatorConfig(experiment_dir, n_forward_steps, checkpoint_path, logging, loader, forward_steps_in_memory, prediction_loader=None, data_writer=<factory>, aggregator=<factory>, stepper_override=None, allow_incompatible_dataset=False, validation=None, n_ensemble_per_ic=1)[source]

Bases: object

Configuration for running inference including comparison to reference data.

Parameters:

  • experiment_dir (str) –

    Directory to save results to. This can be a local directory, like /results, or a remote directory prefixed with a protocol recognized by fsspec, like gs://bucket/results.

    Note

    While most types of output can be written to a remote experiment_dir, there are some limitations:

    • To write raw or time-coarsened data, the zarr writer must be used. See the files parameter of the fme.ace.DataWriterConfig for more details on how this can be configured. Note that monthly coarsened data cannot currently be written to zarr, and hence a remote directory, since it uses a different code path than uniformly coarsened data.

    • Piping logging output to a file in the experiment_dir is not supported. To silence the warning related to this, set log_to_file to False in the fme.ace.LoggingConfig.

    There are no restrictions on the types of output that can be written to a local experiment_dir.

  • 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.

  • loader (InferenceDataLoaderConfig) – Configuration for data to be used as initial conditions, forcing, and target in inference.

  • prediction_loader (Optional[InferenceDataLoaderConfig], default: None) – Configuration for prediction data to evaluate. If given, model evaluation will not run, and instead predictions will be evaluated. Model checkpoint will still be used to determine inputs and outputs.

  • forward_steps_in_memory (int) – Number of forward steps to complete in memory at a time, will load one more step for initial condition.

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

  • aggregator (InferenceEvaluatorAggregatorConfig, default: <factory>) – Configuration for inference evaluator 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 forcing 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.

  • validation (Optional[ValidationConfig], default: None) – Optional configuration for running a one-step validation loop before inference. When provided, validation runs first and produces metrics prefixed with val/ (e.g. val/mean/weighted_rmse), mirroring the validation done at the end of each training epoch.

  • n_ensemble_per_ic (int, default: 1) – Number of ensemble members per initial condition. Useful for stochastic model weather inference. n_ensemble_per_ic = 1 is default inference behavior.

The sub-configurations are:

class fme.ace.LoggingConfig(project='ace', entity='ai2cm', log_to_screen=True, log_to_file=True, log_to_wandb=True, metrics_log_dir=None, 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.

  • metrics_log_dir (Optional[str], default: None) – Directory to write scalar metrics to disk as JSONL. If None, disk metric logging is disabled.

  • 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.InferenceDataLoaderConfig(dataset, start_indices, num_data_workers=0, perturbations=None, persistence_names=None)[source]

Bases: object

Configuration for inference data.

This is like the DataLoaderConfig class, but with some additional constraints. During inference, we have only one batch, so the number of samples directly determines the size of that batch.

Parameters:

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

  • start_indices (InferenceInitialConditionIndices | ExplicitIndices | TimestampList) – Configuration of the indices for initial conditions during inference. This can be a list of timestamps, a list of integer indices, or a slice configuration of the integer indices. Values following the initial condition will still come from the full 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.

  • 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.InferenceInitialConditionIndices(n_initial_conditions, first=0, interval=1)[source]

Bases: object

Configuration of the indices for initial conditions during inference.

Parameters:

  • n_initial_conditions (int) – Number of initial conditions to use.

  • first (int, default: 0) – Index of the first initial condition.

  • interval (int, default: 1) – Interval between initial conditions.

class fme.ace.ExplicitIndices(list)[source]

Bases: object

Configure indices providing them explicitly.

Parameters:

list (Sequence[int]) – List of integer indices.

class fme.ace.TimestampList(times, timestamp_format='%Y-%m-%dT%H:%M:%S')[source]

Bases: object

Configuration for a list of timestamps.

Parameters:

  • times (Sequence[str]) – List of timestamps.

  • timestamp_format (str, default: '%Y-%m-%dT%H:%M:%S') – Format of the timestamps.

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=None)[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 (Optional[list[str]], default: None) – 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. Each data writer must have a unique label to avoid filename collisions.

class fme.ace.InferenceEvaluatorAggregatorConfig(log_histograms=False, log_video=False, log_extended_video=False, log_zonal_mean_images=4096, log_seasonal_means=False, log_global_mean_time_series=True, log_global_mean_norm_time_series=True, monthly_reference_data=None, time_mean_reference_data=None, log_nino34_index=True, log_step_means=<factory>)[source]

Bases: object

Configuration for inference evaluator aggregator.

Parameters:

  • log_histograms (bool, default: False) – Whether to log histograms of the targets and predictions.

  • log_video (bool, default: False) – Whether to log videos of the state evolution.

  • log_extended_video (bool, default: False) – Whether to log wandb videos of the predictions with statistical metrics, only done if log_video is True.

  • log_zonal_mean_images (bool | int, default: 4096) – Whether to log zonal-mean images (hovmollers) with a time dimension. If greater than 0 zonal-mean images will be logged. The value of log_zonal_mean_images is default to 4096 (2**12) and can be set with a maximum of 32768 (2**15) (limited by matplotlib).

  • log_seasonal_means (bool, default: False) – Whether to log seasonal mean metrics and images.

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

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

  • monthly_reference_data (Optional[str], default: None) – Path to monthly reference data to compare against.

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

  • log_step_means (list[StepMeanEntry], default: <factory>) – List of StepMeanEntry objects specifying steps at which to log mean metrics.

  • log_nino34_index (bool) –

class fme.ace.StepperOverrideConfig(ocean='keep', multi_call='keep', derived_forcings='keep', prescribed_prognostic_names='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.

  • derived_forcings (Union[Literal['keep'], DerivedForcingsConfig], default: 'keep') – Derived forcings configuration to override that used in producing a serialized stepper.

  • prescribed_prognostic_names (Union[Literal['keep'], list[str]], default: 'keep') – List of prognostic variable names to overwrite from forcing at each step during inference.