API Reference

fme

class fme.Packer(names: List[str])

Responsible for packing tensors into a single tensor.

pack(tensors: Dict[str, Tensor], axis=0) → Tensor

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

Parameters:

• tensors – Dict from names to tensors.

• axis – index for new concatenation axis.

Raises:

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

class fme.StandardNormalizer(means: Dict[str, Tensor], stds: Dict[str, Tensor])

Responsible for normalizing tensors.

classmethod from_state(state) → StandardNormalizer

Loads state from a serializable data structure.

get_state()

Returns state as a serializable data structure.

fme.get_device() → 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.

fme.gradient_magnitude(tensor: Tensor, dim: int | Iterable[int] = ()) → Tensor

Compute the magnitude of gradient across the specified dimensions.

fme.gradient_magnitude_percent_diff(truth: Tensor, predicted: Tensor, weights: Tensor | None = None, dim: int | Iterable[int] = ()) → Tensor

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

fme.rmse_of_time_mean(truth: Tensor, predicted: Tensor, weights: Tensor | None = None, time_dim: int | Iterable[int] = 0, spatial_dims: int | Iterable[int] = (-2, -1)) → Tensor

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

Parameters:

• truth – truth tensor

• predicted – predicted tensor

• weights – weights to use for computing spatial RMSE

• time_dim – time dimension

• spatial_dims – spatial dimensions over which RMSE is calculated

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: Tensor, predicted: Tensor, weights: Tensor | None = None, dim: int | Iterable[int] = ()) → Tensor

Computes the weighted global RMSE over all variables. Namely, for each variable:

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

If you want to compute the RMSE over the time dimension, then pass in truth.mean(time_dim) and predicted.mean(time_dim) and specify dims=space_dims.

Parameters:

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

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

• weights – torch.Tensor to apply to the squared bias.

• dim – Dimensions to average over.

Returns:

a tensor of shape (variable,) of weighted RMSEs.

fme.spherical_area_weights(lats: ndarray | Tensor, num_lon: int) → Tensor

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

Parameters:

• lats – tensor of shape (num_lat,) with the latitudes of the cell centers.

• num_lon – Number of longitude points.

• device – Device to place the tensor on.

Returns:

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

fme.time_and_global_mean_bias(truth: Tensor, predicted: Tensor, weights: Tensor | None = None, time_dim: int | Iterable[int] = 0, spatial_dims: int | Iterable[int] = (-2, -1)) → Tensor

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

Parameters:

• truth – truth tensor

• predicted – predicted tensor

• weights – weights to use for computing the global mean

• time_dim – time dimension

• spatial_dims – spatial dimensions over which global mean is calculated

Returns:

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

time and spatial dims are reduced.

fme.weighted_mean(tensor: Tensor, weights: Tensor | None = None, dim: int | Iterable[int] = (), keepdim: bool = False) → Tensor

Computes the weighted mean across the specified list of dimensions.

Parameters:

• tensor – torch.Tensor

• weights – Weights to apply to the mean.

• dim – Dimensions to compute the mean over.

• keepdim – Whether the output tensor has dim retained or not.

Returns:

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

fme.weighted_mean_bias(truth: Tensor, predicted: Tensor, weights: Tensor | None = None, dim: int | Iterable[int] = ()) → Tensor

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 – torch.Tensor

• predicted – torch.Tensor

• dim – Dimensions to compute the mean over.

• weights – Weights to apply to the mean.

Returns:

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

fme.weighted_mean_gradient_magnitude(tensor: Tensor, weights: Tensor | None = None, dim: int | Iterable[int] = ()) → Tensor

Compute weighted mean of gradient magnitude across the specified dimensions.

fme.ace

class fme.ace.CopyWeightsConfig(include: ~typing.List[str] = <factory>, exclude: ~typing.List[str] = <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.

include

list of wildcard patterns to overwrite

Type:

List[str]

exclude

list of wildcard patterns to exclude from overwriting

Type:

List[str]

apply(weights: List[Mapping[str, Any]], modules: List[Module])

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 of base weights to apply

• modules – list of modules to apply the weights to

class fme.ace.CorrectorConfig(conserve_dry_air: bool = False, zero_global_mean_moisture_advection: bool = False, moisture_budget_correction: ~typing.Literal['precipitation', 'evaporation', 'advection_and_precipitation', 'advection_and_evaporation'] | None = None, force_positive_names: ~typing.List[str] = <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.

conserve_dry_air

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.

Type:

bool

zero_global_mean_moisture_advection

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.

Type:

bool

moisture_budget_correction

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

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

Type:

Literal[‘precipitation’, ‘evaporation’, ‘advection_and_precipitation’, ‘advection_and_evaporation’] | None

force_positive_names

Names of fields that should be forced to be greater than or equal to zero. This is useful for fields like precipitation.

Type:

List[str]

class fme.ace.DataLoaderConfig(dataset: Sequence[XarrayDataConfig], batch_size: int, num_data_workers: int, prefetch_factor: int | None = None, strict_ensemble: bool = True)

dataset

A sequence of configurations each defining a dataset to be loaded. This sequence of datasets will be concatenated.

Type:

Sequence[fme.core.data_loading.config.XarrayDataConfig]

batch_size

Number of samples per batch.

Type:

int

num_data_workers

Number of parallel workers to use for data loading.

Type:

int

prefetch_factor

how many batches a single data worker will attempt to hold in host memory at a given time.

Type:

int | None

strict_ensemble

Whether to enforce that the ensemble members have the same dimensions and coordinates.

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)

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.EMAConfig(decay: float = 0.9999)

Configuration for exponential moving average of model weights.

decay

decay rate for the moving average

Type:

float

class fme.ace.ExistingStepperConfig(checkpoint_path: str)

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.

checkpoint_path

The path to the serialized checkpoint.

Type:

str

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

Configure indices providing them explicitly.

list

List of integer indices.

Type:

Sequence[int]

class fme.ace.ForcingDataLoaderConfig(dataset: XarrayDataConfig, num_data_workers: int = 0)

Configuration for the forcing data.

dataset

Configuration to define the dataset.

Type:

fme.core.data_loading.config.XarrayDataConfig

num_data_workers

Number of parallel workers to use for data loading.

Type:

int

class fme.ace.FromStateNormalizer(state: Dict[str, Dict[str, float]])

An alternative to NormalizationConfig which provides a normalizer initialized from a serializable state. This is not a public configuration class, but instead allows for loading trained models that have been serialized to disk, using the pre-existing normalization state.

state

State dict of a normalizer.

Type:

Dict[str, Dict[str, float]]

class fme.ace.FrozenParameterConfig(include: ~typing.List[str] = <factory>, exclude: ~typing.List[str] = <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.

include

list of parameter names to freeze (set requires_grad = False)

Type:

List[str]

exclude

list of parameter names to ignore

Type:

List[str]

class fme.ace.InferenceAggregatorConfig(time_mean_reference_data: str | None = None)

Configuration for inference aggregator.

time_mean_reference_data

Path to reference time means to compare against.

Type:

str | None

class fme.ace.InferenceConfig(experiment_dir: str, n_forward_steps: int, checkpoint_path: str, logging: ~fme.core.logging_utils.LoggingConfig, initial_condition: ~fme.ace.inference.inference.InitialConditionConfig, forcing_loader: ~fme.core.data_loading.inference.ForcingDataLoaderConfig, forward_steps_in_memory: int = 10, data_writer: ~fme.ace.inference.data_writer.main.DataWriterConfig = <factory>, aggregator: ~fme.core.aggregator.inference.main.InferenceAggregatorConfig = <factory>, ocean: ~fme.core.ocean.OceanConfig | None = None)

Configuration for running inference.

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

initial_condition

Configuration for initial condition data.

Type:

fme.ace.inference.inference.InitialConditionConfig

forcing_loader

Configuration for forcing data.

Type:

fme.core.data_loading.inference.ForcingDataLoaderConfig

forward_steps_in_memory

Number of forward steps to complete in memory at a time.

Type:

int

data_writer

Configuration for data writers.

Type:

fme.ace.inference.data_writer.main.DataWriterConfig

aggregator

Configuration for inference aggregator.

Type:

fme.core.aggregator.inference.main.InferenceAggregatorConfig

ocean

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

Type:

fme.core.ocean.OceanConfig | None

load_stepper(area: Tensor, sigma_coordinates: SigmaCoordinates) → SingleModuleStepper

Parameters:

• area – A tensor of shape (n_lat, n_lon) containing the area of each grid cell.

• sigma_coordinates – The sigma coordinates of the model.

class fme.ace.InferenceDataLoaderConfig(dataset: XarrayDataConfig, start_indices: InferenceInitialConditionIndices | ExplicitIndices | TimestampList, num_data_workers: int = 0)

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

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

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

load_stepper(area: Tensor, sigma_coordinates: SigmaCoordinates) → SingleModuleStepper

Parameters:

• area – A tensor of shape (n_lat, n_lon) containing the area of each grid cell.

• sigma_coordinates – The sigma coordinates of the model.

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

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.InitialConditionConfig(path: str, engine: Literal['netcdf4', 'h5netcdf', 'zarr'] = 'netcdf4', start_indices: InferenceInitialConditionIndices | ExplicitIndices | TimestampList | None = 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.

path

The path to the initial conditions dataset.

Type:

str

engine

The engine used to open the dataset.

Type:

Literal[‘netcdf4’, ‘h5netcdf’, ‘zarr’]

start_indices

optional specification of the subset of initial conditions to use.

Type:

fme.core.data_loading.inference.InferenceInitialConditionIndices | fme.core.data_loading.inference.ExplicitIndices | fme.core.data_loading.inference.TimestampList | None

class fme.ace.InlineInferenceConfig(loader: ~fme.core.data_loading.inference.InferenceDataLoaderConfig, n_forward_steps: int = 2, forward_steps_in_memory: int = 2, epochs: ~fme.core.data_loading.config.Slice = Slice(start=0, stop=None, step=1), aggregator: ~fme.core.aggregator.inference.main.InferenceEvaluatorAggregatorConfig = <factory>)

loader

configuration for the data loader used during inference

Type:

fme.core.data_loading.inference.InferenceDataLoaderConfig

n_forward_steps

number of forward steps to take

Type:

int

forward_steps_in_memory

number of forward steps to take before re-reading data from disk

Type:

int

epochs

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.

Type:

fme.core.data_loading.config.Slice

aggregator

configuration of inline inference aggregator.

Type:

fme.core.aggregator.inference.main.InferenceEvaluatorAggregatorConfig

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)

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

configure_logging(experiment_dir: str, log_filename: str)

Configure the global logging module based on this LoggingConfig.

class fme.ace.ModuleSelector(type: str, config: Mapping[str, Any])

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.

type

the type of the ModuleConfig

Type:

str

config

data for a ModuleConfig instance of the indicated type

Type:

Mapping[str, Any]

build(n_in_channels: int, n_out_channels: int, img_shape: Tuple[int, int]) → Module

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

Parameters:

• n_in_channels – number of input channels

• n_out_channels – number of output channels

• img_shape – last two dimensions of data, corresponding to lat and lon when using FourCastNet conventions

Returns:

a nn.Module

classmethod from_state(state: Mapping[str, Any]) → ModuleSelector

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

get_state() → Mapping[str, Any]

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

class fme.ace.NormalizationConfig(global_means_path: str | None = None, global_stds_path: str | None = None, exclude_names: ~typing.List[str] | None = None, means: ~typing.Mapping[str, float] = <factory>, stds: ~typing.Mapping[str, float] = <factory>)

Configuration for normalizing data.

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

global_means_path

Path to a netCDF file containing global means.

Type:

str | None

global_stds_path

Path to a netCDF file containing global stds.

Type:

str | None

exclude_names

Names to exclude from normalization.

Type:

List[str] | None

means

Mapping from variable names to means.

Type:

Mapping[str, float]

stds

Mapping from variable names to stds.

Type:

Mapping[str, float]

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

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

class fme.ace.OptimizationConfig(optimizer_type: ~typing.Literal['Adam', 'FusedAdam'] = 'Adam', lr: float = 0.001, kwargs: ~typing.Mapping[str, ~typing.Any] = <factory>, enable_automatic_mixed_precision: bool = False, scheduler: ~fme.core.scheduler.SchedulerConfig = <factory>)

Configuration for optimization.

optimizer_type

The type of optimizer to use.

Type:

Literal[‘Adam’, ‘FusedAdam’]

lr

The learning rate.

Type:

float

kwargs

Additional keyword arguments to pass to the optimizer.

Type:

Mapping[str, Any]

enable_automatic_mixed_precision

Whether to use automatic mixed precision.

Type:

bool

scheduler

The type of scheduler to use. If none is given, no scheduler will be used.

Type:

fme.core.scheduler.SchedulerConfig

class fme.ace.ParameterInitializationConfig(weights_path: str | None = None, exclude_parameters: ~typing.List[str] = <factory>, frozen_parameters: ~fme.core.parameter_init.FrozenParameterConfig = <factory>, alpha: float = 0.0, beta: float = 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.

weight_path

path to a SingleModuleStepper checkpoint containing weights to load

exclude_parameters

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

Type:

List[str]

frozen_parameters

configuration for freezing parameters in the built model

Type:

fme.core.parameter_init.FrozenParameterConfig

alpha

L2 regularization coefficient keeping initialized weights close to their intiial values

Type:

float

beta

L2 regularization coefficient keeping uninitialized weights close to zero

Type:

float

apply(module: Module, init_weights: bool) → Tuple[Module, Callable[[], Tensor]]

Apply the weight initialization to a module.

Parameters:

• module – a nn.Module to initialize

• init_weights – whether to initialize the weight values

Returns:

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

get_base_weights() → Mapping[str, Any] | None

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

class fme.ace.SFNO_V0_1_0(spectral_transform: str = 'sht', filter_type: Literal['linear'] = 'linear', operator_type: str = 'dhconv', scale_factor: int = 16, embed_dim: int = 256, num_layers: int = 12, repeat_layers: int = 1, hard_thresholding_fraction: float = 1.0, normalization_layer: str = 'instance_norm', use_mlp: bool = True, activation_function: str = 'gelu', encoder_layers: int = 1, pos_embed: Literal['none', 'direct', 'frequency'] = 'direct', big_skip: bool = True, rank: float = 1.0, factorization: str | None = None, separable: bool = False, complex_activation: str = 'real', spectral_layers: int = 1, checkpointing: int = 0, data_grid: Literal['legendre-gauss', 'equiangular', 'healpix'] = 'legendre-gauss')

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

build(n_in_channels: int, n_out_channels: int, img_shape: Tuple[int, int])

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

Parameters:

• n_in_channels – number of input channels

• n_out_channels – number of output channels

• img_shape – last two dimensions of data, corresponding to lat and lon when using FourCastNet conventions

Returns:

a nn.Module

class fme.ace.SchedulerConfig(type: str | None = None, kwargs: ~typing.Mapping[str, ~typing.Any] = <factory>)

Configuration for a scheduler to use during training.

type

Name of scheduler class from torch.optim.lr_scheduler, no scheduler is used by default.

Type:

str | None

kwargs

Keyword arguments to pass to the scheduler constructor.

Type:

Mapping[str, Any]

build(optimizer, max_epochs) → _LRScheduler | None

Build the scheduler.

class fme.ace.SingleModuleStepperConfig(builder: ~fme.core.registry.ModuleSelector, in_names: ~typing.List[str], out_names: ~typing.List[str], normalization: ~fme.core.normalizer.NormalizationConfig | ~fme.core.normalizer.FromStateNormalizer, parameter_init: ~fme.core.parameter_init.ParameterInitializationConfig = <factory>, ocean: ~fme.core.ocean.OceanConfig | None = None, loss: ~fme.core.loss.WeightedMappingLossConfig = <factory>, corrector: ~fme.core.corrector.CorrectorConfig = <factory>, next_step_forcing_names: ~typing.List[str] = <factory>, loss_normalization: ~fme.core.normalizer.NormalizationConfig | ~fme.core.normalizer.FromStateNormalizer | None = None, residual_normalization: ~fme.core.normalizer.NormalizationConfig | ~fme.core.normalizer.FromStateNormalizer | None = None)

Configuration for a single module stepper.

builder

The module builder.

Type:

fme.core.registry.ModuleSelector

in_names

Names of input variables.

Type:

List[str]

out_names

Names of output variables.

Type:

List[str]

normalization

The normalization configuration.

Type:

fme.core.normalizer.NormalizationConfig | fme.core.normalizer.FromStateNormalizer

parameter_init

The parameter initialization configuration.

Type:

fme.core.parameter_init.ParameterInitializationConfig

ocean

The ocean configuration.

Type:

fme.core.ocean.OceanConfig | None

loss

The loss configuration.

Type:

fme.core.loss.WeightedMappingLossConfig

corrector

The corrector configuration.

Type:

fme.core.corrector.CorrectorConfig

next_step_forcing_names

Names of forcing variables for the next timestep.

Type:

List[str]

loss_normalization

The normalization configuration for the loss.

Type:

fme.core.normalizer.NormalizationConfig | fme.core.normalizer.FromStateNormalizer | None

residual_normalization

Optional alternative to configure loss normalization. If provided, it will be used for all prognostic variables in loss scaling.

Type:

fme.core.normalizer.NormalizationConfig | fme.core.normalizer.FromStateNormalizer | None

property all_names

Names of all variables required, including auxiliary ones.

property forcing_names: List[str]

Names of variables which are inputs only.

get_base_weights() → List[Mapping[str, Any]] | None

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.

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: str, q_flux_name: str)

Configuration for a slab ocean model.

mixed_layer_depth_name

Name of the mixed layer depth field.

Type:

str

q_flux_name

Name of the heat flux field.

Type:

str

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

Configuration of a python slice built-in.

Required because slice cannot be initialized directly by dacite.

start

Start index of the slice.

Type:

int | None

stop

Stop index of the slice.

Type:

int | None

step

Step of the slice.

Type:

int | None

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

Configuration for the SFNO architecture used in FourCastNet-SFNO.

build(n_in_channels: int, n_out_channels: int, img_shape: Tuple[int, int])

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

Parameters:

• n_in_channels – number of input channels

• n_out_channels – number of output channels

• img_shape – last two dimensions of data, corresponding to lat and lon when using FourCastNet conventions

Returns:

a nn.Module

class fme.ace.TimeCoarsenConfig(coarsen_factor: int)

Config for inference data time coarsening.

Parameters:

coarsen_factor – 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: int) → int

Assumes initial condition is NOT in n_timesteps

class fme.ace.TimeSlice(start_time: str | None = None, stop_time: str | None = None, step: int | None = 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.

start_time

Start time of the slice.

Type:

str | None

stop_time

Stop time of the slice.

Type:

str | None

step

Step of the slice.

Type:

int | None

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

Configuration for a list of timestamps.

times

List of timestamps.

Type:

Sequence[str]

timestamp_format

Format of the timestamps.

Type:

str

class fme.ace.TrainConfig(train_loader: ~fme.core.data_loading.config.DataLoaderConfig, validation_loader: ~fme.core.data_loading.config.DataLoaderConfig, stepper: ~fme.core.stepper.SingleModuleStepperConfig | ~fme.core.stepper.ExistingStepperConfig, optimization: ~fme.core.optimization.OptimizationConfig, logging: ~fme.core.logging_utils.LoggingConfig, max_epochs: int, save_checkpoint: bool, experiment_dir: str, inference: ~fme.ace.train.train_config.InlineInferenceConfig, n_forward_steps: int, copy_weights_after_batch: ~fme.core.weight_ops.CopyWeightsConfig = <factory>, ema: ~fme.core.ema.EMAConfig = <factory>, validate_using_ema: bool = False, checkpoint_save_epochs: ~fme.core.data_loading.config.Slice | None = None, ema_checkpoint_save_epochs: ~fme.core.data_loading.config.Slice | None = None, log_train_every_n_batches: int = 100, segment_epochs: int | None = None)

Configuration for training a model.

train_loader

Configuration for the training data loader.

Type:

fme.core.data_loading.config.DataLoaderConfig

validation_loader

Configuration for the validation data loader.

Type:

fme.core.data_loading.config.DataLoaderConfig

stepper

Configuration for the stepper.

Type:

fme.core.stepper.SingleModuleStepperConfig | fme.core.stepper.ExistingStepperConfig

optimization

Configuration for the optimization.

Type:

fme.core.optimization.OptimizationConfig

logging

Configuration for logging.

Type:

fme.core.logging_utils.LoggingConfig

max_epochs

Total number of epochs to train for.

Type:

int

save_checkpoint

Whether to save checkpoints.

Type:

bool

experiment_dir

Directory where checkpoints and logs are saved.

Type:

str

inference

Configuration for inline inference.

Type:

fme.ace.train.train_config.InlineInferenceConfig

n_forward_steps

Number of forward steps to take gradient over.

Type:

int

copy_weights_after_batch

Configuration for copying weights from the base model to the training model after each batch.

Type:

fme.core.weight_ops.CopyWeightsConfig

ema

Configuration for exponential moving average of model weights.

Type:

fme.core.ema.EMAConfig

validate_using_ema

Whether to validate and perform inference using the EMA model.

Type:

bool

checkpoint_save_epochs

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

Type:

fme.core.data_loading.config.Slice | None

ema_checkpoint_save_epochs

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

Type:

fme.core.data_loading.config.Slice | None

log_train_every_n_batches

How often to log batch_loss during training.

Type:

int

segment_epochs

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.

Type:

int | None

property checkpoint_dir: str

The directory where checkpoints are saved.

class fme.ace.WeightedMappingLossConfig(type: ~typing.Literal['LpLoss', 'MSE', 'AreaWeightedMSE'] = 'MSE', kwargs: ~typing.Mapping[str, ~typing.Any] = <factory>, global_mean_type: ~typing.Literal['LpLoss'] | None = None, global_mean_kwargs: ~typing.Mapping[str, ~typing.Any] = <factory>, global_mean_weight: float = 1.0, weights: ~typing.Dict[str, float] = <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 – the type of the loss function

• kwargs – data for a loss function instance of the indicated type

• global_mean_type – the type of the loss function to apply to the global mean of each sample, by default no loss is applied

• global_mean_kwargs – data for a loss function instance of the indicated type to apply to the global mean of each sample

• global_mean_weight – the weight to apply to the global mean loss relative to the main loss

• weights – A dictionary of variable names with individual weights to apply to their normalized losses

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)

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

fme.ace.register(name: str) → Callable[[Type[ModuleConfig]], Type[ModuleConfig]]

Register a new ModuleConfig type with the NET_REGISTRY.

This is useful for adding new ModuleConfig types to the registry from other modules.

Parameters:

name – name of the ModuleConfig type to register

Returns:

a decorator which registers the decorated class with the NET_REGISTRY