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 2021010100.nc in the Zenodo repository, while ckpt.tar corresponds to a file like ace_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: 8
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, prediction_loader=None, forward_steps_in_memory=1, data_writer=<factory>, aggregator=<factory>, ocean=None)[source]

Bases: object

Configuration for running inference including comparison to reference data.

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.

  • 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, default: 1) – 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.

  • ocean (Optional[OceanConfig], default: None) – Ocean configuration for running inference with a different one than what is used in training.

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)[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) –

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) – Configuration to define the dataset.

  • start_indices (Union[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>, renamed_variables=None, fill_nans=None)[source]

Bases: object

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.

  • subset (Union[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.

  • 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. If this is configured for a renamed field, the key should be the final updated name.

  • renamed_variables (Optional[Mapping[str, str]], default: None) – Optional mapping of {old_name: new_name} to rename variables

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

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(log_extended_video_netcdfs=False, save_prediction_files=True, save_monthly_files=True, names=None, save_histogram_files=False, time_coarsen=None)[source]

Bases: object

Configuration for inference data writers.

Parameters:

  • log_extended_video_netcdfs (bool, default: False) – Whether to enable writing of netCDF files containing video metrics.

  • 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, histogram, and monthly netCDF files.

  • save_histogram_files (bool, default: False) – Enable writing of netCDF files containing histograms.

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

class fme.ace.InferenceEvaluatorAggregatorConfig(log_histograms=False, log_video=False, log_extended_video=False, log_zonal_mean_images=True, 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)[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, default: True) – Whether to log zonal-mean images (hovmollers) with a time dimension.

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

class fme.ace.OceanConfig(surface_temperature_name, ocean_fraction_name, interpolate=False, slab=None)[source]

Bases: object

Configuration for determining sea surface temperature from an ocean model.

Parameters:

  • surface_temperature_name (str) – Name of the sea surface temperature field.

  • ocean_fraction_name (str) – Name of the ocean fraction field.

  • interpolate (bool, default: False) – If True, interpolate between ML-predicted surface temperature and ocean-predicted surface temperature according to ocean_fraction. If False, only use ocean-predicted surface temperature where ocean_fraction>=0.5.

  • slab (Optional[SlabOceanConfig], default: None) – If provided, use a slab ocean model to predict surface temperature.