Modules

ACE’s code uses a module registry system to allow different machine learning architectures to plug into the framework. This is managed by the fme.ace.ModuleSelector configuration class, which is used to select a module type and version.

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

Bases: object

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

config: Mapping[str, Any]

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.

type: str

The following module types are available:

dict_keys(['prebuilt', 'SphericalFourierNeuralOperatorNet', 'SFNO-v0.1.0'])

fme.ace.get_available_module_types()

The following module builders are available:

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

Bases: ModuleConfig

Configuration for the SFNO architecture used in FourCastNet-SFNO.

activation_function: str = 'gelu'

big_skip: bool = True

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

checkpointing: int = 0

complex_activation: str = 'real'

complex_network: bool = True

data_grid: Literal['legendre-gauss', 'equiangular', 'healpix'] = 'legendre-gauss'

embed_dim: int = 256

encoder_layers: int = 1

factorization: str | None = None

filter_type: str = 'non-linear'

hard_thresholding_fraction: float = 1.0

normalization_layer: str = 'instance_norm'

num_layers: int = 12

operator_type: str = 'diagonal'

pos_embed: bool = True

rank: float = 1.0

scale_factor: int = 16

separable: bool = False

spectral_layers: int = 1

spectral_transform: str = 'sht'

use_mlp: bool = True

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

Bases: ModuleConfig

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

activation_function: str = 'gelu'

big_skip: bool = True

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

checkpointing: int = 0

complex_activation: str = 'real'

data_grid: Literal['legendre-gauss', 'equiangular', 'healpix'] = 'legendre-gauss'

embed_dim: int = 256

encoder_layers: int = 1

factorization: str | None = None

filter_type: Literal['linear'] = 'linear'

hard_thresholding_fraction: float = 1.0

normalization_layer: str = 'instance_norm'

num_layers: int = 12

operator_type: str = 'dhconv'

pos_embed: Literal['none', 'direct', 'frequency'] = 'direct'

rank: float = 1.0

repeat_layers: int = 1

scale_factor: int = 16

separable: bool = False

spectral_layers: int = 1

spectral_transform: str = 'sht'

use_mlp: bool = True