API Reference

fme

class fme.Packer(names)

Responsible for packing tensors into a single tensor.

Parameters:

names (List[str]) –

pack(tensors, axis=0)

Packs tensors into a single tensor, concatenated along a new axis.

Parameters:

• tensors (Dict[str, Tensor]) – Dict from names to tensors.

• axis (default: 0) – index for new concatenation axis.

Raises:

DataShapesNotUniform – when packed tensors do not all have the same shape.

Return type:

Tensor

class fme.StandardNormalizer(means, stds)

Responsible for normalizing tensors.

Parameters:

• means (Dict[str, Tensor]) –

• stds (Dict[str, Tensor]) –

classmethod from_state(state)

Loads state from a serializable data structure.

Return type:

StandardNormalizer

get_state()

Returns state as a serializable data structure.

fme.get_device()

If CUDA is available, return a CUDA device. Otherwise, return a CPU device unless FME_USE_MPS is set, in which case return an MPS device if available.

Return type:

device

fme.gradient_magnitude(tensor, dim=())

Compute the magnitude of gradient across the specified dimensions.

Return type:

Tensor

Parameters:

• tensor (Tensor) –

• dim (int | Iterable[int]) –

fme.gradient_magnitude_percent_diff(truth, predicted, weights=None, dim=())

Compute the percent difference of the weighted mean gradient magnitude across the specified dimensions.

Return type:

Tensor

Parameters:

• truth (Tensor) –

• predicted (Tensor) –

• weights (Tensor | None) –

• dim (int | Iterable[int]) –

fme.rmse_of_time_mean(truth, predicted, weights=None, time_dim=0, spatial_dims=(-2, -1))

Compute the RMSE of the time-average given truth and predicted.

Parameters:

• truth (Tensor) – truth tensor

• predicted (Tensor) – predicted tensor

• weights (Optional[Tensor], default: None) – weights to use for computing spatial RMSE

• time_dim (Union[int, Iterable[int]], default: 0) – time dimension

• spatial_dims (Union[int, Iterable[int]], default: (-2, -1)) – spatial dimensions over which RMSE is calculated

Return type:

Tensor

Returns:

The RMSE between the time-mean of the two input tensors. The time and

spatial dims are reduced.

fme.root_mean_squared_error(truth, predicted, weights=None, dim=())

Compute a weighted root mean square error between truth and predicted.

Namely:

sqrt((weights * ((xhat - x) ** 2)).mean(dims))

Parameters:

• truth (Tensor) – torch.Tensor whose last dimensions are to be weighted

• predicted (Tensor) – torch.Tensor whose last dimensions are to be weighted

• weights (Optional[Tensor], default: None) – torch.Tensor to apply to the squared bias.

• dim (Union[int, Iterable[int]], default: ()) – Dimensions to average over.

Return type:

Tensor

Returns:

A tensor of weighted RMSEs.

fme.spherical_area_weights(lats, num_lon)

Computes area weights given the latitudes of a regular lat-lon grid.

Parameters:

• lats (Union[ndarray, Tensor]) – tensor of shape (num_lat,) with the latitudes of the cell centers.

• num_lon (int) – Number of longitude points.

• device – Device to place the tensor on.

Return type:

Tensor

Returns:

a torch.tensor of shape (num_lat, num_lon).

fme.time_and_global_mean_bias(truth, predicted, weights=None, time_dim=0, spatial_dims=(-2, -1))

Compute the global- and time-mean bias given truth and predicted.

Parameters:

• truth (Tensor) – truth tensor

• predicted (Tensor) – predicted tensor

• weights (Optional[Tensor], default: None) – weights to use for computing the global mean

• time_dim (Union[int, Iterable[int]], default: 0) – time dimension

• spatial_dims (Union[int, Iterable[int]], default: (-2, -1)) – spatial dimensions over which global mean is calculated

Return type:

Tensor

Returns:

The global- and time-mean bias between the predicted and truth tensors. The

time and spatial dims are reduced.

fme.weighted_mean(tensor, weights=None, dim=(), keepdim=False)

Computes the weighted mean across the specified list of dimensions.

Parameters:

• tensor (Tensor) – torch.Tensor

• weights (Optional[Tensor], default: None) – Weights to apply to the mean.

• dim (Union[int, Iterable[int]], default: ()) – Dimensions to compute the mean over.

• keepdim (bool, default: False) – Whether the output tensor has dim retained or not.

Return type:

Tensor

Returns:

a tensor of the weighted mean averaged over the specified dimensions dim.

fme.weighted_mean_bias(truth, predicted, weights=None, dim=())

Computes the mean bias across the specified list of dimensions assuming that the weights are applied to the last dimensions, e.g. the spatial dimensions.

Parameters:

• truth (Tensor) – torch.Tensor

• predicted (Tensor) – torch.Tensor

• dim (Union[int, Iterable[int]], default: ()) – Dimensions to compute the mean over.

• weights (Optional[Tensor], default: None) – Weights to apply to the mean.

Return type:

Tensor

Returns:

a tensor of the mean biases averaged over the specified dimensions dim.

fme.weighted_mean_gradient_magnitude(tensor, weights=None, dim=())

Compute weighted mean of gradient magnitude across the specified dimensions.

Return type:

Tensor

Parameters:

• tensor (Tensor) –

• weights (Tensor | None) –

• dim (int | Iterable[int]) –

fme.ace

class fme.ace.CappedGELUConfig(cap_value=10, enable_nhwc=False, enable_healpixpad=False)

Configuration for the CappedGELU activation function.

Parameters:

• cap_value (int, default: 10) – Cap value for the GELU function, default is 10.

• enable_nhwc (bool, default: False) – Flag to enable NHWC data format, default is False.

• enable_healpixpad (bool, default: False) – Flag to enable HEALPix padding, default is False.

build()

Builds the CappedGELU activation function.

Return type:

Module

Returns:

CappedGELU activation function.

class fme.ace.ConstantConfig(amplitude=1.0)

Configuration for a constant perturbation.

Parameters:

amplitude (float) –

class fme.ace.ConvBlockConfig(in_channels=3, out_channels=1, kernel_size=3, dilation=1, n_layers=1, upsampling=2, upscale_factor=4, latent_channels=None, activation=None, enable_nhwc=False, enable_healpixpad=False, block_type='BasicConvBlock')

Configuration for the convolutional block.

Parameters:

• in_channels (int, default: 3) – Number of input channels, default is 3.

• out_channels (int, default: 1) – Number of output channels, default is 1.

• kernel_size (int, default: 3) – Size of the kernel, default is 3.

• dilation (int, default: 1) – Dilation rate, default is 1.

• n_layers (int, default: 1) – Number of layers, default is 1.

• upsampling (int, default: 2) – Upsampling factor for TransposedConvUpsample, default is 2.

• upscale_factor (int, default: 4) – Upscale factor for ConvNeXtBlock and SymmetricConvNeXtBlock, default is 4.

• latent_channels (Optional[int], default: None) – Number of latent channels, default is None.

• activation (Optional[CappedGELUConfig], default: None) – Activation configuration, default is None.

• enable_nhwc (bool, default: False) – Flag to enable NHWC data format, default is False.

• enable_healpixpad (bool, default: False) – Flag to enable HEALPix padding, default is False.

• block_type (Literal['BasicConvBlock', 'ConvNeXtBlock', 'SymmetricConvNeXtBlock', 'TransposedConvUpsample'], default: 'BasicConvBlock') – Type of block, default is “BasicConvBlock”.

build()

Builds the convolutional block model.

Return type:

Module

Returns:

Convolutional block model.

class fme.ace.CopyWeightsConfig(include=<factory>, exclude=<factory>)

Configuration for copying weights from a base model to a target model.

Used during training to overwrite weights after every batch of data, to have the effect of “freezing” the overwritten weights. When the target parameters have longer dimensions than the base model, only the initial slice is overwritten.

This is used to achieve an effect of freezing model parameters that can freeze a subset of each weight that comes from a smaller base weight. This is less efficient than true parameter freezing, but layer freezing is all-or-nothing for each parameter.

All parameters must be covered by either the include or exclude list, but not both.

Parameters:

• include (List[str], default: <factory>) – list of wildcard patterns to overwrite

• exclude (List[str], default: <factory>) – list of wildcard patterns to exclude from overwriting

apply(weights, modules)

Apply base weights to modules according to the include/exclude lists of this instance.

In order to “freeze” the weights during training, this must be called after each time the weights are updated in the training loop.

Parameters:

• weights (List[Mapping[str, Any]]) – list of base weights to apply

• modules (List[Module]) – list of modules to apply the weights to

class fme.ace.CorrectorConfig(conserve_dry_air=False, zero_global_mean_moisture_advection=False, moisture_budget_correction=None, force_positive_names=<factory>)

Configuration for the post-step state corrector.

conserve_dry_air enforces the constraint that:

\[global\_dry\_air = global\_mean(ps - sum_k((ak\_diff + bk\_diff \* ps) \* wat_k))\]

in the generated data is equal to its value in the input data. This is done by adding a globally-constant correction to the surface pressure in each column. As per-mass values such as mixing ratios of water are unchanged, this can cause changes in total water or energy. Note all global means here are area-weighted.

zero_global_mean_moisture_advection enforces the constraint that:

\[global\_mean(tendency\_of\_total\_water\_path\_due\_to\_advection) = 0\]

in the generated data. This is done by adding a globally-constant correction to the moisture advection tendency in each column.

moisture_budget_correction enforces closure of the moisture budget equation:

\[\begin{split}tendency\_of\_total\_water\_path = (evaporation\_rate - precipitation\_rate \\\\ + tendency\_of\_total\_water\_path\_due\_to\_advection)\end{split}\]

in the generated data, where tendency_of_total_water_path is the difference between the total water path at the current timestep and the previous timestep divided by the time difference. This is done by modifying the precipitation, evaporation, and/or moisture advection tendency fields as described in the moisture_budget_correction attribute. When advection tendency is modified, this budget equation is enforced in each column, while when only precipitation or evaporation are modified, only the global mean of the budget equation is enforced.

When enforcing moisture budget closure, we assume the global mean moisture advection is zero. Therefore zero_global_mean_moisture_advection must be True if using a moisture_budget_correction option other than None.

Parameters:

• conserve_dry_air (bool, default: False) – If True, force the generated data to conserve dry air by subtracting a constant offset from the surface pressure of each column. This can cause changes in per-mass values such as total water or energy.

• zero_global_mean_moisture_advection (bool, default: False) – If True, force the generated data to have zero global mean moisture advection by subtracting a constant offset from the moisture advection tendency of each column.

• moisture_budget_correction (Optional[Literal['precipitation', 'evaporation', 'advection_and_precipitation', 'advection_and_evaporation']], default: None) –

  If not “None”, force the generated data to conserve global or column-local moisture by modifying budget fields. Options are:

  • precipitation: multiply precipitation by a scale factor to close the global moisture budget.

  • evaporation: multiply evaporation by a scale factor to close the global moisture budget.

  • advection_and_precipitation: after applying the “precipitation” global-mean correction above, recompute the column-integrated advective tendency as the budget residual, ensuring column budget closure.

  • advection_and_evaporation: after applying the “evaporation” global-mean correction above, recompute the column-integrated advective tendency as the budget residual, ensuring column budget closure.

• force_positive_names (List[str], default: <factory>) – Names of fields that should be forced to be greater than or equal to zero. This is useful for fields like precipitation.

classmethod from_state(state)

Create a ModuleSelector from a dictionary containing all the information needed to build a ModuleConfig.

Return type:

CorrectorConfig

Parameters:

state (Mapping[str, Any]) –

class fme.ace.CorrectorSelector(type, config)

A dataclass containing all the information needed to build a CorrectorConfigProtocol, including the type of the CorrectorConfigProtocol and the data needed to build it.

This is helpful as CorrectorSelector can be serialized and deserialized without any additional information, whereas to load a CorrectorConfigProtocol you would need to know the type of the CorrectorConfigProtocol being loaded.

It is also convenient because CorrectorSelector is a single class that can be used to represent any CorrectorConfigProtocol, whereas CorrectorConfigProtocol is a protocol that can be implemented by many different classes.

Parameters:

• type (str) – the type of the CorrectorConfigProtocol

• config (Mapping[str, Any]) – data for a CorrectorConfigProtocol instance of the indicated type

classmethod from_state(state)

Create a CorrectorSelector from a dictionary containing all the information needed to build a CorrectorConfigProtocol.

Return type:

CorrectorSelector

Parameters:

state (Mapping[str, Any]) –

classmethod get_available_types()

This class method is used to expose all available types of Correctors.

get_state()

Get a dictionary containing all the information needed to build a CorrectorConfigProtocol.

Return type:

Mapping[str, Any]

class fme.ace.DataLoaderConfig(dataset, batch_size, num_data_workers=0, prefetch_factor=None, strict_ensemble=True)

Parameters:

• dataset (Sequence[XarrayDataConfig]) – A sequence of configurations each defining a dataset to be loaded. This sequence of datasets will be concatenated.

• batch_size (int) – Number of samples per batch.

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

• prefetch_factor (Optional[int], default: None) – how many batches a single data worker will attempt to hold in host memory at a given time.

• strict_ensemble (bool, default: True) – Whether to enforce that the datasets to be concatened have the same dimensions and coordinates.

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)

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.DownsamplingBlockConfig(block_type, pooling=2, enable_nhwc=False, enable_healpixpad=False)

Configuration for the downsampling block. Generally, either a pooling block or a striding conv block.

Parameters:

• block_type (Literal['MaxPool', 'AvgPool']) – Type of recurrent block, either “MaxPool” or “AvgPool”

• pooling (int, default: 2) – Pooling size

• enable_nhwc (bool, default: False) – Flag to enable NHWC data format, default is False.

• enable_healpixpad (bool, default: False) – Flag to enable HEALPix padding, default is False.

build()

Builds the recurrent block model.

Return type:

Module

Returns:

Recurrent block.

class fme.ace.EMAConfig(decay=0.9999)

Configuration for exponential moving average of model weights.

Parameters:

decay (float, default: 0.9999) – decay rate for the moving average

class fme.ace.ExistingStepperConfig(checkpoint_path)

Configuration for an existing stepper. This is only designed to point to a serialized stepper checkpoint for loading, e.g., in the case of training resumption.

Parameters:

checkpoint_path (str) – The path to the serialized checkpoint.

class fme.ace.ExplicitIndices(list)

Configure indices providing them explicitly.

Parameters:

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

class fme.ace.ForcingDataLoaderConfig(dataset, num_data_workers=0, perturbations=None, persistence_names=None)

Configuration for the forcing data.

Parameters:

• dataset (XarrayDataConfig) – 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.FrozenParameterConfig(include=<factory>, exclude=<factory>)

Configuration for freezing parameters in a model.

Parameter names can include wildcards, e.g. “encoder.*” will select all parameters in the encoder, while “encoder.*.bias” will select all bias parameters in the encoder. All parameters must be specified in either the include or exclude list, or an exception will be raised.

An exception is raised if a parameter is included by both lists.

Parameters:

• include (List[str], default: <factory>) – list of parameter names to freeze (set requires_grad = False)

• exclude (List[str], default: <factory>) – list of parameter names to ignore

class fme.ace.GreensFunctionConfig(amplitude=1.0, lat_center=0.0, lon_center=0.0, lat_width=10.0, lon_width=10.0)

Configuration for a single sinusoidal patch of a Green’s function perturbation. See equation 1 in Bloch‐Johnson, J., et al. (2024).

Parameters:

• amplitude (float, default: 1.0) – The amplitude of the perturbation, maximum is reached at (lat_center, lon_center).

• lat_center (float, default: 0.0) – The latitude at the center of the patch in degrees.

• lon_center (float, default: 0.0) – The longitude at the center of the patch in degrees.

• lat_width (float, default: 10.0) – latitudinal width of the patch in degrees.

• lon_width (float, default: 10.0) – longitudinal width of the patch in degrees.

class fme.ace.GriddedOperations

classmethod from_state(state)

Given a dictionary with a “type” key and a “state” key, return the GriddedOperations it describes.

The “type” key should be the name of a subclass of GriddedOperations, and the “state” key should be a dictionary specific to that subclass.

Parameters:

state (Dict[str, Any]) – A dictionary with a “type” key and a “state” key.

Return type:

GriddedOperations

Returns:

An instance of the subclass.

abstract get_initialization_kwargs()

Get the keyword arguments needed to initialize the instance.

Return type:

Dict[str, Any]

class fme.ace.HEALPixRecUNetBuilder(encoder, decoder, presteps=1, input_time_size=0, output_time_size=0, delta_time='6h', reset_cycle='24h', n_constants=2, decoder_input_channels=1, prognostic_variables=7, enable_nhwc=False, enable_healpixpad=False)

Configuration for the HEALPixRecUNet architecture used in DLWP.

Parameters:

• presteps (int, default: 1) – Number of pre-steps, by default 1.

• input_time_size (int, default: 0) – Input time dimension, by default 0.

• output_time_size (int, default: 0) – Output time dimension, by default 0.

• delta_time (str, default: '6h') – Delta time interval, by default “6h”.

• reset_cycle (str, default: '24h') – Reset cycle interval, by default “24h”.

• input_channels – Number of input channels, by default 8.

• output_channels – Number of output channels, by default 8.

• n_constants (int, default: 2) – Number of constant input channels, by default 2.

• decoder_input_channels (int, default: 1) – Number of input channels for the decoder, by default 1.

• enable_nhwc (bool, default: False) – Flag to enable NHWC data format, by default False.

• enable_healpixpad (bool, default: False) – Flag to enable HEALPix padding, by default False.

• encoder (UNetEncoderConfig) –

• decoder (UNetDecoderConfig) –

• prognostic_variables (int) –

build(n_in_channels, n_out_channels, img_shape)

Builds the HEALPixRecUNet model.

Parameters:

• n_in_channels (int) – Number of input channels.

• n_out_channels (int) – Number of output channels.

• img_shape (Tuple[int, int]) – Shape of the input image.

Return type:

Module

Returns:

HEALPixRecUNet model.

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

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.InferenceConfig(experiment_dir, n_forward_steps, checkpoint_path, logging, initial_condition, forcing_loader, forward_steps_in_memory=10, data_writer=<factory>, aggregator=<factory>, ocean=None)

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.

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

class fme.ace.InferenceDataLoaderConfig(dataset, start_indices, num_data_workers=0, perturbations=None, persistence_names=None)

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

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

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.

class fme.ace.InferenceInitialConditionIndices(n_initial_conditions, first=0, interval=1)

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.InitialConditionConfig(path, engine='netcdf4', start_indices=None)

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:

• path (str) – The path to the initial conditions dataset.

• engine (Literal['netcdf4', 'h5netcdf', 'zarr'], default: 'netcdf4') – The engine used to open the dataset.

• start_indices (Union[InferenceInitialConditionIndices, ExplicitIndices, TimestampList, None], default: None) – optional specification of the subset of initial conditions to use.

class fme.ace.InlineInferenceConfig(loader, n_forward_steps=2, forward_steps_in_memory=2, epochs=Slice(start=0, stop=None, step=1), aggregator=<factory>)

Parameters:

• loader (InferenceDataLoaderConfig) – configuration for the data loader used during inference

• n_forward_steps (int, default: 2) – number of forward steps to take

• forward_steps_in_memory (int, default: 2) – number of forward steps to take before re-reading data from disk

• epochs (Slice, default: Slice(start=0, stop=None, step=1)) – epochs on which to run inference, where the first epoch is defined as epoch 0 (unlike in logs which show epochs as starting from 1). By default runs inference every epoch.

• aggregator (InferenceEvaluatorAggregatorConfig, default: <factory>) – configuration of inline inference aggregator.

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)

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

configure_logging(experiment_dir, log_filename)

Configure the global logging module based on this LoggingConfig.

Parameters:

• experiment_dir (str) –

• log_filename (str) –

class fme.ace.ModuleSelector(type, config)

A dataclass containing all the information needed to build a ModuleConfig, including the type of the ModuleConfig and the data needed to build it.

This is helpful as ModuleSelector can be serialized and deserialized without any additional information, whereas to load a ModuleConfig you would need to know the type of the ModuleConfig being loaded.

It is also convenient because ModuleSelector is a single class that can be used to represent any ModuleConfig, whereas ModuleConfig is a protocol that can be implemented by many different classes.

Parameters:

• type (str) – the type of the ModuleConfig

• config (Mapping[str, Any]) – data for a ModuleConfig instance of the indicated type

build(n_in_channels, n_out_channels, img_shape)

Build a nn.Module given information about the input and output channels and the image shape.

Parameters:

• n_in_channels (int) – number of input channels

• n_out_channels (int) – number of output channels

• img_shape (Tuple[int, int]) – last two dimensions of data, corresponding to lat and lon when using FourCastNet conventions

Return type:

Module

Returns:

a nn.Module

classmethod from_state(state)

Create a ModuleSelector from a dictionary containing all the information needed to build a ModuleConfig.

Return type:

ModuleSelector

Parameters:

state (Mapping[str, Any]) –

classmethod get_available_types()

This class method is used to expose all available types of Modules.

get_state()

Get a dictionary containing all the information needed to build a ModuleConfig.

Return type:

Union[Mapping[str, Any], dict]

class fme.ace.NormalizationConfig(global_means_path=None, global_stds_path=None, means=<factory>, stds=<factory>)

Configuration for normalizing data.

Either global_means_path and global_stds_path or explicit means and stds must be provided.

Parameters:

• global_means_path (Optional[str], default: None) – Path to a netCDF file containing global means.

• global_stds_path (Optional[str], default: None) – Path to a netCDF file containing global stds.

• means (Mapping[str, float], default: <factory>) – Mapping from variable names to means.

• stds (Mapping[str, float], default: <factory>) – Mapping from variable names to stds.

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

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.

class fme.ace.OceanCorrectorConfig(masking=None, force_positive_names=<factory>)

Parameters:

• masking (MaskingConfig | None) –

• force_positive_names (List[str]) –

classmethod from_state(state)

Create a ModuleSelector from a dictionary containing all the information needed to build a ModuleConfig.

Return type:

OceanCorrectorConfig

Parameters:

state (Mapping[str, Any]) –

class fme.ace.OptimizationConfig(optimizer_type='Adam', lr=0.001, kwargs=<factory>, enable_automatic_mixed_precision=False, scheduler=<factory>)

Configuration for optimization.

Parameters:

• optimizer_type (Literal['Adam', 'FusedAdam'], default: 'Adam') – The type of optimizer to use.

• lr (float, default: 0.001) – The learning rate.

• kwargs (Mapping[str, Any], default: <factory>) – Additional keyword arguments to pass to the optimizer.

• enable_automatic_mixed_precision (bool, default: False) – Whether to use automatic mixed precision.

• scheduler (SchedulerConfig, default: <factory>) – The type of scheduler to use. If none is given, no scheduler will be used.

class fme.ace.ParameterInitializationConfig(weights_path=None, exclude_parameters=<factory>, frozen_parameters=<factory>, alpha=0.0, beta=0.0)

A class which applies custom initialization to module parameters.

Assumes the module weights have already been randomly initialized.

Supports overwriting the weights of the built model with weights from a pre-trained model. If the built model has larger weights than the pre-trained model, only the initial slice of the weights is overwritten.

Parameters:

• weight_path – path to a SingleModuleStepper checkpoint containing weights to load

• exclude_parameters (List[str], default: <factory>) – list of parameter names to exclude from the loaded weights. Used for example to keep the random initialization for final layer(s) of a model, and only overwrite the weights for earlier layers. Takes values like “decoder.2.weight”.

• frozen_parameters (FrozenParameterConfig, default: <factory>) – configuration for freezing parameters in the built model

• alpha (float, default: 0.0) – L2 regularization coefficient keeping initialized weights close to their intiial values

• beta (float, default: 0.0) – L2 regularization coefficient keeping uninitialized weights close to zero

• weights_path (str | None) –

apply(module, init_weights)

Apply the weight initialization to a module.

Parameters:

• module (Module) – a nn.Module to initialize

• init_weights (bool) – whether to initialize the weight values

Return type:

Tuple[Module, Callable[[], Tensor]]

Returns:

a nn.Module with initialization applied a function which returns the regularization loss term

get_base_weights()

If a weights_path is provided, return the model base weights used for initialization.

Return type:

Optional[Mapping[str, Any]]

class fme.ace.PerturbationSelector(type, config)

Parameters:

• type (str) –

• config (Mapping[str, Any]) –

classmethod get_available_types()

This class method is used to expose all available types of Perturbations.

get_state()

Get a dictionary containing all the information needed to build a PerturbationConfig.

Return type:

Mapping[str, Any]

class fme.ace.RecurrentBlockConfig(in_channels=3, kernel_size=1, enable_nhwc=False, enable_healpixpad=False, block_type='ConvGRUBlock')

Configuration for the recurrent block.

Parameters:

• in_channels (int, default: 3) – Number of input channels, default is 3.

• kernel_size (int, default: 1) – Size of the kernel, default is 1.

• enable_nhwc (bool, default: False) – Flag to enable NHWC data format, default is False.

• enable_healpixpad (bool, default: False) – Flag to enable HEALPix padding, default is False.

• block_type (Literal['ConvGRUBlock', 'ConvLSTMBlock'], default: 'ConvGRUBlock') – Type of recurrent block, either “ConvGRUBlock” or “ConvLSTMBlock”,

• "ConvGRUBlock". (default is) –

build()

Builds the recurrent block model.

Return type:

Module

Returns:

Recurrent block.

class fme.ace.SFNO_V0_1_0(spectral_transform='sht', filter_type='linear', operator_type='dhconv', scale_factor=16, embed_dim=256, num_layers=12, repeat_layers=1, hard_thresholding_fraction=1.0, normalization_layer='instance_norm', use_mlp=True, activation_function='gelu', encoder_layers=1, pos_embed='direct', big_skip=True, rank=1.0, factorization=None, separable=False, complex_activation='real', spectral_layers=1, checkpointing=0, data_grid='legendre-gauss')

Configuration for the SFNO architecture in modulus-makani version 0.1.0.

Parameters:

• spectral_transform (str) –

• filter_type (Literal['linear']) –

• operator_type (str) –

• scale_factor (int) –

• embed_dim (int) –

• num_layers (int) –

• repeat_layers (int) –

• hard_thresholding_fraction (float) –

• normalization_layer (str) –

• use_mlp (bool) –

• activation_function (str) –

• encoder_layers (int) –

• pos_embed (Literal['none', 'direct', 'frequency']) –

• big_skip (bool) –

• rank (float) –

• factorization (str | None) –

• separable (bool) –

• complex_activation (str) –

• spectral_layers (int) –

• checkpointing (int) –

• data_grid (Literal['legendre-gauss', 'equiangular', 'healpix']) –

build(n_in_channels, n_out_channels, img_shape)

Build a nn.Module given information about the input and output channels and the image shape.

Parameters:

• n_in_channels (int) – number of input channels

• n_out_channels (int) – number of output channels

• img_shape (Tuple[int, int]) – shape of last two dimensions of data, e.g. latitude and longitude.

Returns:

a nn.Module

class fme.ace.SSTPerturbation(sst)

Configuration for sea surface temperature perturbations applied to initial condition and forcing data. Currently, this is strictly applied to both.

Parameters:

sst (list[PerturbationSelector]) – List of perturbation selectors for SST perturbations.

class fme.ace.SchedulerConfig(type=None, kwargs=<factory>)

Configuration for a scheduler to use during training.

Parameters:

• type (Optional[str], default: None) – Name of scheduler class from torch.optim.lr_scheduler, no scheduler is used by default.

• kwargs (Mapping[str, Any], default: <factory>) – Keyword arguments to pass to the scheduler constructor.

build(optimizer, max_epochs)

Build the scheduler.

Return type:

Optional[_LRScheduler]

class fme.ace.SingleModuleStepperConfig(builder, in_names, out_names, normalization, parameter_init=<factory>, ocean=None, loss=<factory>, corrector=<factory>, next_step_forcing_names=<factory>, loss_normalization=None, residual_normalization=None)

Configuration for a single module stepper.

Parameters:

• builder (ModuleSelector) – The module builder.

• in_names (List[str]) – Names of input variables.

• out_names (List[str]) – Names of output variables.

• normalization (NormalizationConfig) – The normalization configuration.

• parameter_init (ParameterInitializationConfig, default: <factory>) – The parameter initialization configuration.

• ocean (Optional[OceanConfig], default: None) – The ocean configuration.

• loss (WeightedMappingLossConfig, default: <factory>) – The loss configuration.

• corrector (Union[CorrectorConfig, CorrectorSelector], default: <factory>) – The corrector configuration.

• next_step_forcing_names (List[str], default: <factory>) – Names of forcing variables for the next timestep.

• loss_normalization (Optional[NormalizationConfig], default: None) – The normalization configuration for the loss.

• residual_normalization (Optional[NormalizationConfig], default: None) – Optional alternative to configure loss normalization. If provided, it will be used for all prognostic variables in loss scaling.

property all_names

Names of all variables required, including auxiliary ones.

property diagnostic_names: List[str]

Names of variables which both inputs and outputs.

property forcing_names: List[str]

Names of variables which are inputs only.

get_base_weights()

If the model is being initialized from another model’s weights for fine-tuning, returns those weights. Otherwise, returns None.

The list mirrors the order of modules in the SingleModuleStepper class.

Return type:

Optional[List[Mapping[str, Any]]]

property normalize_names

Names of variables which require normalization. I.e. inputs/outputs.

property prognostic_names: List[str]

Names of variables which both inputs and outputs.

class fme.ace.SlabOceanConfig(mixed_layer_depth_name, q_flux_name)

Configuration for a slab ocean model.

Parameters:

• mixed_layer_depth_name (str) – Name of the mixed layer depth field.

• q_flux_name (str) – Name of the heat flux field.

class fme.ace.Slice(start=None, stop=None, step=None)

Configuration of a python slice built-in.

Required because slice cannot be initialized directly by dacite.

Parameters:

• start (Optional[int], default: None) – Start index of the slice.

• stop (Optional[int], default: None) – Stop index of the slice.

• step (Optional[int], default: None) – Step of the slice.

class fme.ace.SphericalFourierNeuralOperatorBuilder(spectral_transform='sht', filter_type='non-linear', operator_type='diagonal', scale_factor=16, embed_dim=256, num_layers=12, hard_thresholding_fraction=1.0, normalization_layer='instance_norm', use_mlp=True, activation_function='gelu', encoder_layers=1, pos_embed=True, big_skip=True, rank=1.0, factorization=None, separable=False, complex_network=True, complex_activation='real', spectral_layers=1, checkpointing=0, data_grid='legendre-gauss')

Configuration for the SFNO architecture used in FourCastNet-SFNO.

Parameters:

• spectral_transform (str) –

• filter_type (str) –

• operator_type (str) –

• scale_factor (int) –

• embed_dim (int) –

• num_layers (int) –

• hard_thresholding_fraction (float) –

• normalization_layer (str) –

• use_mlp (bool) –

• activation_function (str) –

• encoder_layers (int) –

• pos_embed (bool) –

• big_skip (bool) –

• rank (float) –

• factorization (str | None) –

• separable (bool) –

• complex_network (bool) –

• complex_activation (str) –

• spectral_layers (int) –

• checkpointing (int) –

• data_grid (Literal['legendre-gauss', 'equiangular', 'healpix']) –

build(n_in_channels, n_out_channels, img_shape)

Build a nn.Module given information about the input and output channels and the image shape.

Parameters:

• n_in_channels (int) – number of input channels

• n_out_channels (int) – number of output channels

• img_shape (Tuple[int, int]) – shape of last two dimensions of data, e.g. latitude and longitude.

Returns:

a nn.Module

class fme.ace.TimeCoarsenConfig(coarsen_factor)

Config for inference data time coarsening.

Parameters:

coarsen_factor (int) – Factor by which to coarsen in time, an integer 1 or greater. The resulting time labels will be coarsened to the mean of the original labels.

n_coarsened_timesteps(n_timesteps)

Assumes initial condition is NOT in n_timesteps.

Return type:

int

Parameters:

n_timesteps (int) –

class fme.ace.TimeSlice(start_time=None, stop_time=None, step=None)

Configuration of a slice of times. Step is an integer-valued index step.

Note: start_time and stop_time may be provided as partial time strings and the

stop_time will be included in the slice. See more details in Xarray docs.

Parameters:

• start_time (Optional[str], default: None) – Start time of the slice.

• stop_time (Optional[str], default: None) – Stop time of the slice.

• step (Optional[int], default: None) – Step of the slice.

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

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.TrainConfig(train_loader, validation_loader, stepper, optimization, logging, max_epochs, save_checkpoint, experiment_dir, inference, n_forward_steps, copy_weights_after_batch=<factory>, ema=<factory>, validate_using_ema=False, checkpoint_save_epochs=None, ema_checkpoint_save_epochs=None, log_train_every_n_batches=100, segment_epochs=None)

Configuration for training a model.

Parameters:

• train_loader (DataLoaderConfig) – Configuration for the training data loader.

• validation_loader (DataLoaderConfig) – Configuration for the validation data loader.

• stepper (Union[SingleModuleStepperConfig, ExistingStepperConfig]) – Configuration for the stepper.

• optimization (OptimizationConfig) – Configuration for the optimization.

• logging (LoggingConfig) – Configuration for logging.

• max_epochs (int) – Total number of epochs to train for.

• save_checkpoint (bool) – Whether to save checkpoints.

• experiment_dir (str) – Directory where checkpoints and logs are saved.

• inference (InlineInferenceConfig) – Configuration for inline inference.

• n_forward_steps (int) – Number of forward steps to take gradient over.

• copy_weights_after_batch (CopyWeightsConfig, default: <factory>) – Configuration for copying weights from the base model to the training model after each batch.

• ema (EMAConfig, default: <factory>) – Configuration for exponential moving average of model weights.

• validate_using_ema (bool, default: False) – Whether to validate and perform inference using the EMA model.

• checkpoint_save_epochs (Optional[Slice], default: None) – How often to save epoch-based checkpoints, if save_checkpoint is True. If None, checkpoints are only saved for the most recent epoch (and the best epochs if validate_using_ema == False).

• ema_checkpoint_save_epochs (Optional[Slice], default: None) – How often to save epoch-based EMA checkpoints, if save_checkpoint is True. If None, EMA checkpoints are only saved for the most recent epoch (and the best epochs if validate_using_ema == True).

• log_train_every_n_batches (int, default: 100) – How often to log batch_loss during training.

• segment_epochs (Optional[int], default: None) – Exit after training for at most this many epochs in current job, without exceeding max_epochs. Use this if training must be run in segments, e.g. due to wall clock limit.

property checkpoint_dir: str

The directory where checkpoints are saved.

class fme.ace.UNetDecoderConfig(conv_block, up_sampling_block, output_layer, recurrent_block=None, n_channels=<factory>, n_layers=<factory>, output_channels=1, dilations=None, enable_nhwc=False, enable_healpixpad=False)

Configuration for the UNet Decoder.

Parameters:

• conv_block (ConvBlockConfig) – Configuration for the convolutional block.

• up_sampling_block (ConvBlockConfig) – Configuration for the up-sampling block.

• output_layer (ConvBlockConfig) – Configuration for the output layer block.

• recurrent_block (Optional[RecurrentBlockConfig], default: None) – Configuration for the recurrent block, by default None.

• n_channels (List[int], default: <factory>) – Number of channels for each layer, by default (34, 68, 136).

• n_layers (List[int], default: <factory>) – Number of layers in each block, by default (1, 2, 2).

• output_channels (int, default: 1) – Number of output channels, by default 1.

• dilations (Optional[list], default: None) – List of dilation rates for the layers, by default None.

• enable_nhwc (bool, default: False) – Flag to enable NHWC data format, by default False.

• enable_healpixpad (bool, default: False) – Flag to enable HEALPix padding, by default False.

build()

Builds the UNet Decoder model.

Return type:

Module

Returns:

UNet Decoder model.

class fme.ace.UNetEncoderConfig(conv_block, down_sampling_block, input_channels=3, n_channels=<factory>, n_layers=<factory>, dilations=None, enable_nhwc=False, enable_healpixpad=False)

Configuration for the UNet Encoder.

Parameters:

• conv_block (ConvBlockConfig) – Configuration for the convolutional block.

• down_sampling_block (DownsamplingBlockConfig) – Configuration for the down-sampling block.

• input_channels (int, default: 3) – Number of input channels, by default 3.

• n_channels (List[int], default: <factory>) – Number of channels for each layer, by default (136, 68, 34).

• n_layers (List[int], default: <factory>) – Number of layers in each block, by default (2, 2, 1).

• dilations (Optional[list], default: None) – List of dilation rates for the layers, by default None.

• enable_nhwc (bool, default: False) – Flag to enable NHWC data format, by default False.

• enable_healpixpad (bool, default: False) – Flag to enable HEALPix padding, by default False.

build()

Builds the UNet Encoder model.

Return type:

Module

Returns:

UNet Encoder model.

class fme.ace.WeightedMappingLossConfig(type='MSE', kwargs=<factory>, global_mean_type=None, global_mean_kwargs=<factory>, global_mean_weight=1.0, weights=<factory>)

Loss configuration class that has the same fields as LossConfig but also has additional weights field. The build method will apply the weights to the inputs of the loss function. The loss returned by build will be a MappingLoss, which takes Dict[str, tensor] as inputs instead of packed tensors.

Parameters:

• type (Literal['LpLoss', 'MSE', 'AreaWeightedMSE'], default: 'MSE') – the type of the loss function

• kwargs (Mapping[str, Any], default: <factory>) – data for a loss function instance of the indicated type

• global_mean_type (Optional[Literal['LpLoss']], default: None) – the type of the loss function to apply to the global mean of each sample, by default no loss is applied

• global_mean_kwargs (Mapping[str, Any], default: <factory>) – data for a loss function instance of the indicated type to apply to the global mean of each sample

• global_mean_weight (float, default: 1.0) – the weight to apply to the global mean loss relative to the main loss

• weights (Dict[str, float], default: <factory>) – A dictionary of variable names with individual weights to apply to their normalized losses

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)

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"
... )