Skip to content

Models

BaseModel

Bases: LightningModule, ABC

A base class for all models used in the CryoCast project.

ignored_hparams = frozenset(('latitudes_fn', 'longitudes_fn', 'mask_dir', 'metrics')) class-attribute

name = name instance-attribute

hemisphere = Hemisphere(hemisphere) instance-attribute

latitudes_fn = latitudes_fn instance-attribute

longitudes_fn = longitudes_fn instance-attribute

channel_names = list(channel_names) if channel_names else [] instance-attribute

checkpoint_epoch = None instance-attribute

training_metadata = None instance-attribute

n_forecast_steps = n_forecast_steps instance-attribute

n_history_steps = n_history_steps instance-attribute

input_spaces = [DataSpace.from_dict(space) for space in input_spaces] instance-attribute

output_space = DataSpace.from_dict(output_space) instance-attribute

target_variable_indices = list(target_variable_indices or []) instance-attribute

optimizer_cfg = optimizer instance-attribute

scheduler_cfg = scheduler instance-attribute

lr_scheduler_cfg = lr_scheduler instance-attribute

loss_cfg property writable

Get the loss configuration.

metric_cfgs = {} instance-attribute

test_metrics = self.build_metrics(land_mask) instance-attribute

train_metrics = self.build_metrics(land_mask) instance-attribute

validation_metrics = self.build_metrics(land_mask) instance-attribute

climatology_metrics = self.build_metrics(land_mask) instance-attribute

latitudes cached property

longitudes cached property

multistage_only property

build_metrics(land_mask)

Build a metric collection from the configured metrics.

Each configured metric is a mapping with a name (its key in logs), a Hydra _target_ and any other constructor arguments, e.g. {"name": "my_metric", "_target_": "my_package.MyMetric", "k": 3}. Targets that subclass LandMaskMixin are also given the land mask.

This should include only metrics that are compatible with the output space. We therefore filter out single-channel metrics from the metric collection if the model will predict multiple channels.

Parameters:

Name Type Description Default
land_mask Tensor | None

Optional boolean tensor of shape (H, W), True for ocean cells and False for land. This is passed to every LandMaskMixin metric, and is used to exclude land cells from the metric.

 required

Returns:

Type Description
MetricCollection

A MetricCollection containing the requested metrics.

configure_optimizers()

Construct the optimizer and optional scheduler from the config.

forward(inputs) abstractmethod

Forward step of the model.

  • start with multiple NTCHW inputs, one for each input dataset
  • return a single NTCHW output representing the predicted output

Parameters:

Name Type Description Default
inputs dict[str, TensorNTCHW]

Dictionary of dataset name to TensorNTCHW with shape (batch, n_history_steps, C_input_k, H_input_k, W_input_k)

 required

Returns:

Type Description
TensorNTCHW

Predicted TensorNTCHW with shape (batch, n_forecast_steps, C_output, H_output, W_output)

load_state_dict(state_dict, *args, **kwargs)

Load a state dict, ignoring any metric collections it contains.

on_fit_start()

Record the data that this model is being trained on.

on_load_checkpoint(checkpoint)

Restore the training metadata, if the checkpoint has any.

on_save_checkpoint(checkpoint)

Save the training metadata as plain types, so it loads with weights_only.

loss(prediction, target)

Calculate the loss given a prediction and target.

process_batch(batch)

Process a batch before the forward pass and loss computation.

Subclasses can override this to extract or transform inputs before the standard training/validation steps. The returned dict must include a "target" key.

test_step(batch, _batch_idx)

Run the test step, in PyTorch eval model (i.e. no gradients).

  • Separate the batch into inputs and target
  • Run inputs through the model
  • Update the test metrics (and climatology metrics if appropriate)
  • Return the prediction, target and loss

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary mapping dataset name to its contents. There is one entry for each input dataset and one for the target. Each of these is a TensorNTCHW with (batch_size, n_history_steps, C, H, W).

 required

Returns:

Type Description
ModelStepOutput

A ModelStepOutput containing the prediction, target and loss for the batch.

training_step(batch, _batch_idx)

Run the training step.

  • Separate the batch into inputs and target
  • Run inputs and target through the model
  • Calculate the loss wrt. the target

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary mapping dataset name to its contents. There is one entry for each input dataset and one for the target. Each of these is a TensorNTCHW with (batch_size, n_history_steps, C, H, W).

 required

Returns:

Type Description
ModelStepOutput

A ModelStepOutput containing the prediction, target and loss for the batch.

validation_step(batch, _batch_idx)

Run the validation step.

A batch contains one tensor for each input dataset and one for the target These are [NTCHW] tensors with (batch_size, n_history_steps, C, H, W)

  • Separate the batch into inputs and target
  • Run inputs through the model
  • Calculate and log the loss wrt. the target

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary mapping dataset name to its contents. There is one entry for each input dataset and one for the target. Each of these is a TensorNTCHW with (batch_size, n_history_steps, C, H, W).

 required

Returns:

Type Description
ModelStepOutput

A ModelStepOutput containing the prediction, target and loss for the batch.

Climatology

Bases: BaseModel

automatic_optimization = False instance-attribute

model = nn.Identity() instance-attribute

configure_optimizers()

Climatology model does not need an optimizer.

forward(inputs)

Forward step of the model.

Return the climatology (calendar-day mean) field for each forecast step, supplied by the data loader under the climatology key with shape [batch, n_forecast_steps, C_target, H_target, W_target].

DDPM

Bases: BaseModel

Denoising Diffusion Probabilistic Model (DDPM).

Input space

TensorNTCHW with shape (batch_size, n_history_steps + n_history_steps * n_era5_channels, height, width) - OSISAF input: T historical steps, singleton channel squeezed - ERA5 input: T historical steps times number of channels, resized to OSISAF resolution

Output space

TensorNTCHW with shape (batch_size, n_forecast_steps * n_output_channels, height, width) - Forecasted outputs per timestep and channel, flattened along the channel dimension

ignored_hparams = BaseModel.ignored_hparams | {'mask_dir'} class-attribute

use_autoregressive = use_autoregressive instance-attribute

osisaf_key = self.output_space.name instance-attribute

restrict = RestrictRange(RangeRestriction(restrict_range), min_val=0, max_val=1) instance-attribute

mask = Mask(mask_type=mask_type, output_shape=self.output_space.shape, mask_dir=mask_dir) instance-attribute

era5_space = next(space for space in self.input_spaces if space.name == 'era5').channels instance-attribute

osisaf_channels = next(space for space in self.input_spaces if space.name == self.osisaf_key).channels instance-attribute

base_output_channels = self.output_space.channels instance-attribute

output_channels = self.base_output_channels instance-attribute

timesteps = timesteps instance-attribute

cond_channels = 64 instance-attribute

input_channels = self.cond_channels instance-attribute

era5_norm = torch.nn.InstanceNorm3d(self.era5_space, affine=True) instance-attribute

era5_compressed_channels = 32 instance-attribute

era5_projector = torch.nn.Sequential(torch.nn.Conv3d(self.era5_space, self.era5_compressed_channels, kernel_size=1), torch.nn.SiLU()) instance-attribute

osisaf_encoder = SimpleEncoder2D(in_channels=self.n_history_steps * self.osisaf_channels, out_channels=self.cond_channels // 2) instance-attribute

era5_encoder = torch.nn.Sequential(torch.nn.Conv2d(in_channels=self.era5_compressed_channels * self.n_history_steps, out_channels=self.cond_channels // 2, kernel_size=3, padding=1), torch.nn.GroupNorm(4, self.cond_channels // 2), torch.nn.SiLU()) instance-attribute

model = UNetDiffusion(input_channels=self.input_channels, output_channels=self.output_channels, timesteps=self.timesteps, kernel_size=kernel_size, start_out_channels=start_out_channels, time_embed_dim=time_embed_dim, normalization=normalization, activation=activation, dropout_rate=dropout_rate) instance-attribute

diffusion = GaussianDiffusion(timesteps=timesteps) instance-attribute

learning_rate = learning_rate instance-attribute

forward(*args, **kwargs)

sample(batch)

Generate forecasts using a reverse diffusion process.

This method selects between two diffusion sampling strategies:

  1. Non-autoregressive (parallel) sampling:
  2. The model generates the entire future sequence in a single diffusion process.

  3. No temporal dependency exists between forecast steps.

  4. Autoregressive sampling:

  5. Forecast steps are generated sequentially.
  6. Each step is produced via an independent diffusion process.
  7. The conditioning tensor is updated after each step to incorporate previously generated outputs.

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary containing the input data.

 required

Returns:

Type Description
TensorNCHW

torch.Tensor: Forecast tensor of shape: [B, n_forecast_steps * base_output_channels, H, W]

The output format is identical in both modes.

  • Parallel mode: Produced in a single reverse diffusion process.

  • Autoregressive mode: Constructed by concatenating step-wise diffusion outputs.
Notes
  • The diffusion process follows v-parameterization.
  • Sampling begins from standard Gaussian noise.

prepare_inputs(batch)

Encode OSISAF and ERA5 separately, then concatenate.

ERA5 -> Norm -> Project -> Resize -> Flatten Time -> Encode

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary with osisaf key (e.g. 'osisaf-south') [B, T, C, H, W] 'era5' [B, T, C, H2, W2]

 required

Returns:

Type Description
TensorNCHW

Conditioning tensor [B, cond_channels, H, W]

training_step(batch, _batch_idx)

One training step using DDPM v-prediction loss.

During training, the clean target (SIC) is corrupted using the forward diffusion process by adding noise at a randomly sampled timestep. The model is trained to predict the corresponding v-target.

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary containing: - input tensors (used to prepare conditioning inputs) - "target": groundtruth SIC tensor

 required

Returns:

Name Type Description
ModelStepOutput ModelStepOutput
ModelStepOutput
  • prediction: reconstructed noisy SIC (pred_v)
ModelStepOutput
  • target: generated noisy SIC (target_v)
ModelStepOutput
  • loss: training loss value

validation_step(batch, _batch_idx)

One validation step using full diffusion sampling.

During validation, samples are generated by starting from noise and iteratively denoising conditioned on the inputs. The final prediction is compared to the groundtruth SIC using the configured evaluation loss.

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary containing: - input tensors (used to prepare conditioning inputs) - "target": groundtruth SIC tensor

 required

Returns:

Name Type Description
ModelStepOutput ModelStepOutput
ModelStepOutput
  • prediction: reconstructed SIC (y_hat)
ModelStepOutput
  • target: groundtruth SIC (y)
ModelStepOutput
  • loss: validation loss value

test_step(batch, _batch_idx)

One test step using full diffusion sampling and metric evaluation.

During testing, predictions are generated by starting from noise and running the reverse diffusion process conditioned on the inputs. The final reconstructed SIC is compared to the groundtruth target using the configured loss and test metrics.

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary containing: - input tensors (used to prepare conditioning inputs) - "target": groundtruth SIC tensor

 required

Returns:

Name Type Description
ModelStepOutput ModelStepOutput
ModelStepOutput
  • prediction: reconstructed SIC (y_hat)
ModelStepOutput
  • target: groundtruth SIC (y)
ModelStepOutput
  • loss: test loss value

EncodeProcessDecode

Bases: BaseModel

Model that encodes to latent space, processes, then decodes back.

ignored_hparams = BaseModel.ignored_hparams | {'decoder', 'encoders', 'mask_dir', 'processor'} class-attribute

encoders = [hydra.utils.instantiate(encoders[input_space.name], data_space_in=input_space, latent_space=encoders['latent_space'], latitudes_fn=self.latitudes_fn, longitudes_fn=self.longitudes_fn) for input_space in self.input_spaces] if isinstance(encoders, Mapping) else [encoder for input_space in self.input_spaces for encoder in encoders if encoder.name == input_space.name] instance-attribute

target_encoder = hydra.utils.instantiate(encoders[self.output_space.name], data_space_in=DataSpace(name='target', channels=self.output_space.channels, shape=self.output_space.shape), latent_space=encoders['latent_space'], latitudes_fn=self.latitudes_fn, longitudes_fn=self.longitudes_fn) if isinstance(encoders, Mapping) else encoders.pop(next(idx for idx, encoder in enumerate(encoders) if encoder.name == 'target')) instance-attribute

processor = processor if isinstance(processor, BaseProcessor) else hydra.utils.instantiate(processor, _recursive_=False, data_space=combined_latent_space, data_space_target=(self.target_encoder if target_input_encoder is None else target_input_encoder).data_space_out, n_forecast_steps=self.n_forecast_steps, n_history_steps=self.n_history_steps, target_channel_offset=self._find_target_channel_offset()) instance-attribute

decoder = decoder if isinstance(decoder, BaseDecoder) else hydra.utils.instantiate(decoder, data_space_in=combined_latent_space, data_space_out=self.output_space, mask_dir=mask_dir) instance-attribute

rollout_space = self._validate_rollout_options(rollout_space) instance-attribute

multistage_only property

target_input_encoder property

Return the input encoder whose dataset contains the forecast target.

forward(inputs)

Forward step of the model (used for inference).

Delegates to either the latent-space or physical-space path depending on the rollout_space setting.

train(mode=True)

Set training mode, with decoder frozen if computing loss in latent space.

training_step(batch, _batch_idx)

Run the training step.

The standard path is to call the forward step with self(batch) and compute the loss by comparing the decoded prediction to the target.

Processors that set the flag computes_loss_in_latent_space=True are a special case. We encode the inputs and target into latent space, call rollout() on the processor, which returns both a latent prediction and a loss. We pass this prediction through a frozen decoder in order to calculate metrics, but we use the loss returned by the processor for backpropagation.

Parameters:

Name Type Description Default
batch dict[str, TensorNTCHW]

Dictionary with one NTCHW entry per input dataset (n_history_steps) and a "target" entry (n_forecast_steps).

 required

Returns:

Type Description
ModelStepOutput

A ModelStepOutput containing the prediction, target and loss.

Persistence

Bases: BaseModel

automatic_optimization = False instance-attribute

model = nn.Identity() instance-attribute

configure_optimizers()

Persistence model does not need an optimizer.

forward(inputs)

Forward step of the model.

  • start with multiple NTCHW inputs each with shape (batch, n_history_steps, C_input_k, H_input_k, W_input_k)
  • find the input with the same name as the output space
  • select the channels corresponding to the target variables
  • take the last time step and repeat it n_forecast_steps times