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: str, n_forward_steps: int, checkpoint_path: str, logging: ~fme.core.logging_utils.LoggingConfig, loader: ~fme.core.data_loading.inference.InferenceDataLoaderConfig, prediction_loader: ~fme.core.data_loading.inference.InferenceDataLoaderConfig | None = None, forward_steps_in_memory: int = 1, data_writer: ~fme.ace.inference.data_writer.main.DataWriterConfig = <factory>, aggregator: ~fme.core.aggregator.inference.main.InferenceEvaluatorAggregatorConfig = <factory>, ocean: ~fme.core.ocean.OceanConfig | None = None)[source]

Bases: object

Configuration for running inference including comparison to reference data.

experiment_dir

Directory to save results to.

Type:: str

n_forward_steps

Number of steps to run the model forward for.

Type:: int

checkpoint_path

Path to stepper checkpoint to load.

Type:: str

logging

configuration for logging.

Type:: fme.core.logging_utils.LoggingConfig

loader

Configuration for data to be used as initial conditions, forcing, and target in inference.

Type:: fme.core.data_loading.inference.InferenceDataLoaderConfig

prediction_loader

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.

Type:: fme.core.data_loading.inference.InferenceDataLoaderConfig | None

forward_steps_in_memory

Number of forward steps to complete in memory at a time, will load one more step for initial condition.

Type:: int

data_writer

Configuration for data writers.

Type:: fme.ace.inference.data_writer.main.DataWriterConfig

aggregator

Configuration for inference evaluator aggregator.

Type:: fme.core.aggregator.inference.main.InferenceEvaluatorAggregatorConfig

ocean

Ocean configuration for running inference with a different one than what is used in training.

Type:: fme.core.ocean.OceanConfig | None

The sub-configurations are:

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

Bases: object

Configuration for logging.

project

name of the project in Weights & Biases

Type:: str

entity

name of the entity in Weights & Biases

Type:: str

log_to_screen

whether to log to the screen

Type:: bool

log_to_file

whether to log to a file

Type:: bool

log_to_wandb

whether to log to Weights & Biases

Type:: bool

log_format

format of the log messages

Type:: str

class fme.ace.InferenceDataLoaderConfig(dataset: XarrayDataConfig, start_indices: InferenceInitialConditionIndices | ExplicitIndices | TimestampList, num_data_workers: int = 0)[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.

dataset

Configuration to define the dataset.

Type:: fme.core.data_loading.config.XarrayDataConfig

start_indices

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.

Type:: fme.core.data_loading.inference.InferenceInitialConditionIndices | fme.core.data_loading.inference.ExplicitIndices | fme.core.data_loading.inference.TimestampList

num_data_workers

Number of parallel workers to use for data loading.

Type:: int

class fme.ace.InferenceInitialConditionIndices(n_initial_conditions: int, first: int = 0, interval: int = 1)[source]

Bases: object

Configuration of the indices for initial conditions during inference.

n_initial_conditions

Number of initial conditions to use.

Type:: int

first

Index of the first initial condition.

Type:: int

interval

Interval between initial conditions.

Type:: int

class fme.ace.ExplicitIndices(list: Sequence[int])[source]

Bases: object

Configure indices providing them explicitly.

list

List of integer indices.

Type:: Sequence[int]

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

Bases: object

Configuration for a list of timestamps.

times

List of timestamps.

Type:: Sequence[str]

timestamp_format

Format of the timestamps.

Type:: str

class fme.ace.XarrayDataConfig(data_path: str, file_pattern: str = '*.nc', n_repeats: int = 1, engine: ~typing.Literal['netcdf4', 'h5netcdf', 'zarr'] | None = None, spatial_dimensions: ~typing.Literal['healpix', 'latlon'] = 'latlon', subset: ~fme.core.data_loading.config.Slice | ~fme.core.data_loading.config.TimeSlice = <factory>, infer_timestep: bool = True)[source]

Bases: object

data_path

Path to the data.

Type:: str

file_pattern

Glob pattern to match files in the data_path.

Type:: str

n_repeats

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.

Type:: int

engine

Backend for xarray.open_dataset. Currently supported options are “netcdf4” (the default) and “h5netcdf”. Only valid when using XarrayDataset.

Type:: Literal[‘netcdf4’, ‘h5netcdf’, ‘zarr’] | None

spatial_dimensions

Specifies the spatial dimensions for the grid, default is lat/lon.

Type:: Literal[‘healpix’, ‘latlon’]

subset

Slice defining a subset of the XarrayDataset to load. This can either be a Slice of integer indices or a TimeSlice of timestamps.

Type:: fme.core.data_loading.config.Slice | fme.core.data_loading.config.TimeSlice

infer_timestep

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.

Type:: bool

class fme.ace.DataWriterConfig(log_extended_video_netcdfs: bool = False, save_prediction_files: bool = True, save_monthly_files: bool = True, names: Sequence[str] | None = None, save_histogram_files: bool = False, time_coarsen: TimeCoarsenConfig | None = None)[source]

Bases: object

Configuration for inference data writers.

log_extended_video_netcdfs

Whether to enable writing of netCDF files containing video metrics.

Type:: bool

save_prediction_files

Whether to enable writing of netCDF files containing the predictions and target values.

Type:: bool

save_monthly_files

Whether to enable writing of netCDF files containing the monthly predictions and target values.

Type:: bool

names

Names of variables to save in the prediction, histogram, and monthly netCDF files.

Type:: Sequence[str] | None

save_histogram_files

Enable writing of netCDF files containing histograms.

Type:: bool

time_coarsen

Configuration for time coarsening of written outputs.

Type:: fme.ace.inference.data_writer.time_coarsen.TimeCoarsenConfig | None

class fme.ace.InferenceEvaluatorAggregatorConfig(log_histograms: bool = False, log_video: bool = False, log_extended_video: bool = False, log_zonal_mean_images: bool = True, log_seasonal_means: bool = False, monthly_reference_data: str | None = None, time_mean_reference_data: str | None = None)[source]

Bases: object

Configuration for inference evaluator aggregator.

log_histograms

Whether to log histograms of the targets and predictions.

Type:: bool

log_video

Whether to log videos of the state evolution.

Type:: bool

log_extended_video

Whether to log wandb videos of the predictions with statistical metrics, only done if log_video is True.

Type:: bool

log_zonal_mean_images

Whether to log zonal-mean images (hovmollers) with a time dimension.

Type:: bool

log_seasonal_means

Whether to log seasonal mean metrics and images.

Type:: bool

monthly_reference_data

Path to monthly reference data to compare against.

Type:: str | None

time_mean_reference_data

Path to reference time means to compare against.

Type:: str | None

class fme.ace.OceanConfig(surface_temperature_name: str, ocean_fraction_name: str, interpolate: bool = False, slab: SlabOceanConfig | None = None)[source]

Bases: object

Configuration for determining sea surface temperature from an ocean model.

surface_temperature_name

Name of the sea surface temperature field.

Type:: str

ocean_fraction_name

Name of the ocean fraction field.

Type:: str

interpolate

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.

Type:: bool

slab

If provided, use a slab ocean model to predict surface temperature.

Type:: fme.core.ocean.SlabOceanConfig | None