Skip to content

Models

BaseModel

Bases: LightningModule, ABC

A base class for all models used in the IceNet-MP project.

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

name = name instance-attribute

hemisphere = hemisphere instance-attribute

latitudes_fn = latitudes_fn instance-attribute

longitudes_fn = longitudes_fn 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

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.

test_metrics = MetricCollection(deepcopy(_common_metrics)) instance-attribute

train_metrics = MetricCollection(deepcopy(_common_metrics)) instance-attribute

validation_metrics = MetricCollection(deepcopy(_common_metrics)) instance-attribute

latitudes cached property

longitudes cached property

multistage_only property

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]

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

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 = era5_space['channels'] instance-attribute

osisaf_channels = osisaf_space['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 | {'mask_dir'} class-attribute

target_variable_indices = target_variable_indices instance-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] 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) instance-attribute

processor = hydra.utils.instantiate(processor, data_space=combined_latent_space, data_space_target=self.target_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 = hydra.utils.instantiate(decoder, data_space_in=combined_latent_space, data_space_out=self.output_space, mask_dir=mask_dir) instance-attribute

multistage_only property

encode_inputs(inputs)

Encode all input datasets and concatenate along the channel dimension.

Parameters:

Name Type Description Default
inputs dict[str, TensorNTCHW]

Dictionary with one TensorNTCHW entry per input dataset with shape (batch, n_history_steps, n_input_channels_k, H_input_k, W_input_k)

 required

Returns:

Type Description
TensorNTCHW

TensorNTCHW with shape (batch_size, n_history_steps, n_latent_channels_total, latent_height, latent_width)

forward(inputs)

Forward step of the model (used for inference).

  • start with multiple [NTCHW] inputs each with shape [batch, n_history_steps, n_input_channels_k, H_input_k, W_input_k]
  • encode inputs to [NTCHW] latent space [batch, n_history_steps, n_latent_channels, H_latent, W_latent]
  • concatenate inputs in [NTCHW] latent space [batch, n_history_steps, n_latent_channels_total, H_latent, W_latent]
  • process in latent space [NTCHW][batch, n_forecast_steps, n_latent_channels_total, H_latent, W_latent]
  • decode back to [NTCHW] output space [batch, n_forecast_steps, n_output_channels, H_output, W_output]
  • add a skip connection from the most recent target value to every forecast step

find_target_channel_offset()

Find the channel offset of the target dataset within the combined latent space, if present.

get_persistence(inputs)

Extract persistence if needed for a skip connection.

train(mode=True)

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

training_step(batch, _batch_idx)

Run the training step.

If the processor returns a loss in its ProcessorOutput (rather than None), this is used for backpropagation. We use no_grad to compute the decoded prediction, which allows us to calculate metrics and log outputs, but the usefulness of these will depend on what ProcessorOutput.prediction contains.

Otherwise, the standard encode-process-decode path is used and the loss is computed by comparing the decoded prediction to the target.

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

variable_indices = target_variable_indices 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