Coupled Modeling

Tip

We highly recommend familiarizing yourself with the standalone atmosphere only model (ACE) before using a coupled atmosphere-ocean model. See quickstart for running, evaluating, and training ACE.

Currently, coupled modeling supports coupling an atmosphere model (e.g. ACE) with an ocean model (e.g. Samudra).

Commands

The following commands are available, and can be run with --help for more information:

  • python3 -m fme.coupled.validate_config - Validate a configuration file

  • python3 -m fme.coupled.inference - Run a saved model checkpoint

Accessing coupled checkpoints and datasets

Checkpoints for running coupled inference are available in the ACE Collection on Hugging Face. You can find an example coupled model checkpoint under the name SamudrACE-CM4-piControl. The minimum requirements for running inference with SamudrACE are:

  • a coupled model checkpoint

  • initial conditions files containing all prognostic variables for the atmosphere and ocean respectively

  • forcing dataset containing all input-only variables for the atmosphere

Save a config-inference.yaml file based on the example config with updated initial conditions and forcing paths for the downloaded data. Specifically, initial_condition.atmosphere.path is the local initial condition file for the atmosphere, and initial_condition.ocean.path for the ocean. forcing_loader.atmosphere.dataset.data_path should be the local directory containing the forcing data files for the atmosphere. forcing_loader.ocean.dataset.data_path is not required because there are no ocean specific forcing data.

Then in the fme conda environment, run inference with:

python -m fme.coupled.inference config-inference.yaml

If you run into configuration issues, you can validate your configuration with

python -m fme.coupled.validate_config config-inference.yaml --config_type inference

Coupled inference config

The example uses paths relative to the directory you run the command. You should have a directory structure like:

.
├── ckpt.tar
├── initial_conditions
│   ├── atmosphere-ic.nc
│   └── ocean-ic.nc
├── forcing_data
│   ├── forcing_data_00.nc
│   └── forcing_data_....nc

including a model checkpoint (ckpt.tar), atmosphere and ocean forcing and initial condition files. You can find the checkpoint and forcing and initial condition data in the SamudrACE Hugging Face page.

Example inference YAML Configuration

Example YAML Configuration
experiment_dir: inference_output
n_coupled_steps: 40 # 200 days
coupled_steps_in_memory: 20
checkpoint_path: samudrACE_CM4_piControl_ckpt.tar
data_writer:
  ocean:
    save_prediction_files: false
    save_monthly_files: true
  atmosphere:
    save_prediction_files: false
    save_monthly_files: true
forcing_loader:
  atmosphere:
    dataset:
      data_path: forcing_data
      file_pattern: forcing_0311.nc
  num_data_workers: 1
initial_condition:
  ocean:
    path: "initial_conditions/0311-01-01T00:00:00/ocean/initial_condition.nc"
  atmosphere:
    path: "initial_conditions/0311-01-01T00:00:00/atmosphere/initial_condition.nc"
  start_indices:
    first: 0
    n_initial_conditions: 1
logging:
  log_to_screen: true
  log_to_wandb: false
  log_to_file: true
  project: project_name
  entity: your_wandb_entity

Configuration structure

We use the Builder pattern to load this configuration into a multi-level dataclass structure. The configuration is divided into several sub-configurations, each with its own dataclass. The top-level configuration is the fme.coupled.InferenceConfig class.

class fme.coupled.InferenceConfig(experiment_dir, n_coupled_steps, checkpoint_path, logging, initial_condition, forcing_loader, coupled_steps_in_memory=1, data_writer=<factory>, aggregator=<factory>, labels=<factory>)[source]

Bases: object

Configuration for running inference.

Parameters:

  • experiment_dir (str) – Directory to save results to.

  • n_coupled_steps (int) – Number of steps to run the model forward for.

  • checkpoint_path (str | StandaloneComponentCheckpointsConfig) – Path to a CoupledStepper training checkpoint to load, or a mapping to two separate Stepper training checkpoints.

  • logging (LoggingConfig) – configuration for logging.

  • initial_condition (CoupledInitialConditionConfig) – Configuration for initial condition data.

  • forcing_loader (CoupledForcingDataLoaderConfig) – Configuration for forcing data.

  • coupled_steps_in_memory (int, default: 1) – Number of coupled steps to complete in memory at a time, will load one more step for initial condition.

  • data_writer (CoupledDataWriterConfig, default: <factory>) – Configuration for data writers.

  • aggregator (InferenceAggregatorConfig, default: <factory>) – Configuration for inference aggregator.

  • labels (list[str], default: <factory>) – Dataset labels to use for inference.

Information for fme.coupled.CoupledInitialConditionConfig

Initial condition configuration is similar to standalone ACE as documtned in inference config, but in coupled configuration, start_indices correspond to the ocean initial condition file , and the atmosphere initial condition file must contain the same time stamps.