"""Lightning module handling training, validation and testing loops."""
# Standard library
import os
import warnings
from typing import Any
# Third-party
import matplotlib.pyplot as plt
import numpy as np
import pytorch_lightning as pl
import torch
import xarray as xr
from PIL import Image
# First-party
from neural_lam.utils import get_integer_time
# Local
from .. import metrics, vis
from ..config import NeuralLAMConfig
from ..datastore.base import BaseDatastore
from ..loss_weighting import get_state_feature_weighting
from ..weather_dataset import WeatherDataset
from .forecasters.base import Forecaster
# (init_states, target_states, forcing, boundary, target_times)
Batch = tuple[
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
]
[docs]
class ForecasterModule(pl.LightningModule):
"""
Lightning module handling training, validation and testing loops.
Wraps a Forecaster instance which performs the actual prediction.
"""
# pylint: disable=arguments-differ
interior_mask_bool: torch.Tensor
per_var_std: torch.Tensor | None
state_mean: torch.Tensor
state_std: torch.Tensor
forcing_mean: torch.Tensor | None
forcing_std: torch.Tensor | None
forcing_mean_tiled: torch.Tensor | None
forcing_std_tiled: torch.Tensor | None
boundary_mean: torch.Tensor | None
boundary_std: torch.Tensor | None
boundary_mean_tiled: torch.Tensor | None
boundary_std_tiled: torch.Tensor | None
def __init__(
self,
forecaster: Forecaster,
config: NeuralLAMConfig,
datastore: BaseDatastore,
datastore_boundary: BaseDatastore | None = None,
loss: str = "wmse",
lr: float = 1e-3,
restore_opt: bool = False,
n_example_pred: int = 1,
create_gif: bool = False,
val_steps_to_log: list[int] | None = None,
train_steps_to_log: list[int] | None = None,
metrics_watch: list[str] | None = None,
var_leads_metrics_watch: dict[int, list[int]] | None = None,
args: Any | None = None,
) -> None:
"""
Initialize the ForecasterModule.
Parameters
----------
forecaster : Forecaster
The forecaster model to use for predictions.
config : NeuralLAMConfig
Configuration object for the neural LAM model.
datastore : BaseDatastore
Datastore providing grid metadata and data access.
datastore_boundary : BaseDatastore, optional
Boundary forcing datastore, used for its standardization
statistics.
loss : str, default "wmse"
The loss function to use.
lr : float, default 1e-3
Learning rate for the optimizer.
restore_opt : bool, default False
Whether to restore optimizer state from checkpoint.
n_example_pred : int, default 1
Number of example predictions to plot during testing.
create_gif : bool, default False
Whether to create GIFs of example predictions.
val_steps_to_log : list of int, optional
Specific rollout steps to log during validation/testing.
metrics_watch : list of str, optional
List of metrics to watch and log specifically.
var_leads_metrics_watch : dict of {int: list of int}, optional
Mapping from variable index to a list of rollout steps to log
individually for the configured metrics.
args : argparse.Namespace, optional
Pre-refactor ``ARModel`` checkpoint hyperparameters. When
provided, attributes on ``args`` take precedence over the
corresponding explicit kwargs (``loss``, ``lr``, ``restore_opt``,
``n_example_pred``, ``create_gif``, ``val_steps_to_log``,
``metrics_watch``, ``var_leads_metrics_watch``) so legacy
checkpoints round-trip through ``load_from_checkpoint``
correctly.
"""
super().__init__()
# Pre-refactor ``ARModel`` checkpoints saved every hyperparameter nested
# inside an argparse Namespace under the single key 'args'. When
# Lightning calls __init__ during load_from_checkpoint it would
# otherwise drop 'args' (not in the new signature) and silently fall
# back to defaults for loss/lr/create_gif/etc. Unpack the namespace
# here so legacy checkpoints round-trip correctly.
if args is not None:
loss = getattr(args, "loss", loss)
lr = getattr(args, "lr", lr)
restore_opt = getattr(args, "restore_opt", restore_opt)
n_example_pred = getattr(args, "n_example_pred", n_example_pred)
create_gif = getattr(args, "create_gif", create_gif)
val_steps_to_log = getattr(
args, "val_steps_to_log", val_steps_to_log
)
train_steps_to_log = getattr(
args, "train_steps_to_log", train_steps_to_log
)
metrics_watch = getattr(args, "metrics_watch", metrics_watch)
var_leads_metrics_watch = getattr(
args, "var_leads_metrics_watch", var_leads_metrics_watch
)
# Resolve mutable defaults
if val_steps_to_log is None:
val_steps_to_log = [
1,
]
if train_steps_to_log is None:
train_steps_to_log = []
if metrics_watch is None:
metrics_watch = []
if var_leads_metrics_watch is None:
var_leads_metrics_watch = {}
# datastore and forecaster are excluded from saved hparams and must
# be provided explicitly when calling load_from_checkpoint. Saving
# args makes the checkpoint self-describing: it carries model,
# graph_name, hidden_dim, etc. so the caller can reconstruct the
# exact forecaster architecture from the checkpoint alone.
self.save_hyperparameters(
ignore=["datastore", "datastore_boundary", "forecaster"]
)
self.datastore = datastore
self.datastore_boundary = datastore_boundary
self.forecaster = forecaster
self.matched_metrics: set = set()
# Compute interior_mask_bool directly from datastore
boundary_mask = (
torch.tensor(datastore.boundary_mask.values, dtype=torch.float32)
.unsqueeze(0)
.unsqueeze(-1)
) # (1, num_grid_nodes, 1)
interior_mask = 1.0 - boundary_mask
self.register_buffer(
"interior_mask_bool",
interior_mask[0, :, 0].to(torch.bool),
persistent=False,
)
# Store per_var_std here if predictor does not output std
if not self.forecaster.predicts_std:
da_state_stats = datastore.get_standardization_dataarray(
category="state"
)
state_feature_weights = get_state_feature_weighting(
config=config, datastore=datastore
)
diff_std = torch.tensor(
da_state_stats.state_diff_std_standardized.values,
dtype=torch.float32,
)
feature_weights_t = torch.tensor(
state_feature_weights, dtype=torch.float32
)
self.register_buffer(
"per_var_std",
diff_std / torch.sqrt(feature_weights_t),
persistent=False,
)
else:
self.per_var_std = None
# Standardization statistics used to normalize each batch on-device in
# `on_after_batch_transfer`. WeatherDataset returns unstandardized
# data, so state and forcing are normalized here rather than on CPU.
eps = torch.finfo(torch.float32).eps
da_state_stats = datastore.get_standardization_dataarray(
category="state"
)
self.register_buffer(
"state_mean",
torch.tensor(da_state_stats.state_mean.values, dtype=torch.float32),
persistent=False,
)
self.register_buffer(
"state_std",
self._safe_std(da_state_stats.state_std.values, eps, "state"),
persistent=False,
)
if datastore.get_num_data_vars(category="forcing") > 0:
da_forcing_stats = datastore.get_standardization_dataarray(
category="forcing"
)
self.register_buffer(
"forcing_mean",
torch.tensor(
da_forcing_stats.forcing_mean.values, dtype=torch.float32
),
persistent=False,
)
self.register_buffer(
"forcing_std",
self._safe_std(
da_forcing_stats.forcing_std.values, eps, "forcing"
),
persistent=False,
)
# The windowed forcing width is only known once the first batch
# arrives, so the tiled mean/std are built and cached then (in
# on_after_batch_transfer) instead of recomputed every batch.
self.register_buffer("forcing_mean_tiled", None, persistent=False)
self.register_buffer("forcing_std_tiled", None, persistent=False)
else:
self.forcing_mean = None
self.forcing_std = None
# Boundary tiling matches forcing, see `on_after_batch_transfer`.
if (
datastore_boundary is not None
and datastore_boundary.get_num_data_vars(category="forcing") > 0
):
da_boundary_stats = (
datastore_boundary.get_standardization_dataarray(
category="forcing"
)
)
self.register_buffer(
"boundary_mean",
torch.tensor(
da_boundary_stats.forcing_mean.values, dtype=torch.float32
),
persistent=False,
)
self.register_buffer(
"boundary_std",
self._safe_std(
da_boundary_stats.forcing_std.values, eps, "boundary"
),
persistent=False,
)
self.register_buffer("boundary_mean_tiled", None, persistent=False)
self.register_buffer("boundary_std_tiled", None, persistent=False)
else:
self.boundary_mean = None
self.boundary_std = None
# Instantiate loss function
self.loss = metrics.get_metric(loss)
self.val_metrics: dict[str, list] = {
"mse": [],
}
self.test_metrics: dict[str, list] = {
"mse": [],
"mae": [],
}
if self.forecaster.predicts_std:
self.test_metrics["output_std"] = [] # Treat as metric
# For making restoring of optimizer state optional
self.restore_opt = restore_opt
# For example plotting
self.n_example_pred = n_example_pred
self.create_gif = create_gif
self.plotted_examples = 0
# For storing spatial loss maps during evaluation
self.spatial_loss_maps: list[Any] = []
self._test_spatial_steps: list[int] = []
# Warn once per phase if steps_to_log exceeds the actual rollout
self._steps_warn_issued = {
"val": False,
"test": False,
"train": False,
}
self.time_step_int, self.time_step_unit = get_integer_time(
self.datastore.step_length
)
def _create_dataarray_from_tensor(
self,
tensor: torch.Tensor,
time: torch.Tensor,
split: str,
category: str,
) -> xr.DataArray:
"""
Create an xarray DataArray from a torch tensor.
Parameters
----------
tensor : torch.Tensor
The tensor to convert.
time : torch.Tensor
The time coordinates for the data.
split : str
The data split (e.g., 'train', 'val', 'test').
category : str
The category of data (e.g., 'state', 'forcing').
Returns
-------
xr.DataArray
The resulting xarray DataArray.
"""
weather_dataset = WeatherDataset(datastore=self.datastore, split=split)
time_np = np.array(time.cpu(), dtype="datetime64[ns]")
da = weather_dataset.create_dataarray_from_tensor(
tensor=tensor, time=time_np, category=category
)
return da
@staticmethod
def _safe_std(
std_values: np.ndarray | torch.Tensor, eps: float, category: str
) -> torch.Tensor:
"""Build a float32 std tensor, clamping near-zero values to `eps`.
Mirrors the previous CPU-side WeatherDataset behavior: features with
near-zero std are clamped to machine epsilon to avoid division-by-zero
NaNs, and a single warning is emitted naming the affected category.
"""
std = torch.tensor(std_values, dtype=torch.float32)
if bool((std <= eps).any()):
warnings.warn(
f"Some {category} features have near-zero std and will be "
"standardized using machine epsilon to avoid NaN.",
UserWarning,
stacklevel=2,
)
return torch.clamp(std, min=eps)
[docs]
def on_after_batch_transfer(
self,
batch: Batch,
dataloader_idx: int,
) -> Batch:
"""Standardize a batch on-device after transfer to the accelerator.
Lightning calls this for every train/val/test/predict batch.
WeatherDataset returns unstandardized state, forcing and boundary;
all three are normalized here so the work runs on the accelerator.
Forcing and boundary share the same feature-major
``(feature, window)`` stacking, so per-feature mean/std are tiled
once and cached.
Parameters
----------
batch : Batch
``(init_states, target_states, forcing, boundary, target_times)``
as returned by the dataloader.
dataloader_idx : int
Index of the dataloader the batch comes from.
Returns
-------
Batch
The batch with state, forcing and boundary standardized.
"""
init_states, target_states, forcing, boundary, batch_times = batch
init_states = (init_states - self.state_mean) / self.state_std
target_states = (target_states - self.state_mean) / self.state_std
if forcing.shape[-1] > 0:
assert self.forcing_mean is not None
assert self.forcing_std is not None
if self.forcing_mean_tiled is None:
# Forcing is (..., num_forcing_vars * window_size).
# WeatherDataset stacks (forcing_feature, window)
# feature-major, so each per-feature mean/std is repeated
# window_size times. Tile once and cache as buffers rather
# than recomputing every batch.
window_size = forcing.shape[-1] // self.forcing_mean.shape[-1]
self.forcing_mean_tiled = self.forcing_mean.repeat_interleave(
window_size
)
self.forcing_std_tiled = self.forcing_std.repeat_interleave(
window_size
)
assert self.forcing_mean_tiled is not None
assert self.forcing_std_tiled is not None
forcing = (
forcing - self.forcing_mean_tiled
) / self.forcing_std_tiled
if boundary.shape[-1] > 0:
if self.boundary_mean is None:
raise ValueError(
"Batch has non-empty boundary forcing but no boundary "
"standardization statistics are registered. Pass "
"`datastore_boundary` when constructing or loading "
"ForecasterModule."
)
assert self.boundary_std is not None
if self.boundary_mean_tiled is None:
window_size = boundary.shape[-1] // self.boundary_mean.shape[-1]
self.boundary_mean_tiled = self.boundary_mean.repeat_interleave(
window_size
)
self.boundary_std_tiled = self.boundary_std.repeat_interleave(
window_size
)
assert self.boundary_mean_tiled is not None
assert self.boundary_std_tiled is not None
boundary = (
boundary - self.boundary_mean_tiled
) / self.boundary_std_tiled
return init_states, target_states, forcing, boundary, batch_times
[docs]
def common_step(
self,
batch: Batch,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]:
"""
Perform a common prediction step for training, validation, and testing.
Parameters
----------
batch : tuple
The batch of data containing initial states, target states,
forcing features, boundary forcing features, and batch times.
Returns
-------
tuple
A tuple containing prediction, target states, predicted standard
deviation, and batch times.
"""
(
init_states,
target_states,
forcing_features,
boundary_features,
batch_times,
) = batch
# boundary_features is standardized here but not yet passed to the
# forecaster.
prediction, pred_std = self.forecaster(
init_states, forcing_features, target_states
)
return prediction, target_states, pred_std, batch_times
[docs]
def training_step(
self,
batch: Batch,
) -> torch.Tensor:
"""
Perform a single training step.
Parameters
----------
batch : tuple
The batch of data.
Returns
-------
torch.Tensor
The computed loss for the training step.
"""
_, _, _, time_step_loss = self._compute_prediction_and_loss(batch)
batch_loss = torch.mean(time_step_loss)
batch_size = batch[0].shape[0]
self._log_step_loss(time_step_loss, batch_loss, "train", batch_size)
return batch_loss
[docs]
def all_gather_cat(self, tensor_to_gather: torch.Tensor) -> torch.Tensor:
"""
Gather tensors from all GPUs and concatenate them.
Parameters
----------
tensor_to_gather : torch.Tensor
The tensor to gather from all processes.
Returns
-------
torch.Tensor
The concatenated tensor from all processes.
"""
gathered = self.all_gather(tensor_to_gather)
assert isinstance(gathered, torch.Tensor)
# all_gather adds dim 0 only on multi-device; on single
# device it returns the same tensor unchanged.
if gathered.dim() > tensor_to_gather.dim():
gathered = gathered.flatten(0, 1)
return gathered
# pylint: disable-next=unused-argument
def _warn_skipped_steps(self, pred_steps: int, phase: str) -> None:
"""Warn once per phase if any log steps exceed pred_steps."""
if self._steps_warn_issued[phase]:
return
hparam_name = (
"train_steps_to_log" if phase == "train" else "val_steps_to_log"
)
steps_to_log = getattr(self.hparams, hparam_name, [])
invalid = [s for s in steps_to_log if s > pred_steps]
if invalid:
warnings.warn(
f"{hparam_name} contains steps {invalid} that exceed "
f"the {phase} rollout length ({pred_steps}). "
"These steps will be skipped from logging. "
f"Adjust --{hparam_name} or the ar_steps.",
UserWarning,
stacklevel=2,
)
self._steps_warn_issued[phase] = True
def _compute_prediction_and_loss(
self,
batch: Batch,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Compute predicted mean, standard deviation, and step-wise loss.
Also extract and return corresponding target from batch.
Parameters
----------
batch : tuple of torch.Tensor
The batch of data.
Returns
-------
prediction : torch.Tensor
Model predictions, shape
`(B, pred_steps, num_grid_nodes, num_state_vars)`.
target_states : torch.Tensor
Target states, shape
`(B, pred_steps, num_grid_nodes, num_state_vars)`.
pred_std : torch.Tensor
Predicted or pre-defined standard deviation, shape
`(B, pred_steps, num_grid_nodes, num_state_vars)` or
`(num_state_vars,)`.
time_step_loss : torch.Tensor
Loss for each unroll step, shape `(pred_steps,)`.
"""
prediction, target_states, pred_std, _ = self.common_step(batch)
if pred_std is None:
pred_std = self.per_var_std
assert pred_std is not None
time_step_loss = torch.mean(
self.loss(
prediction,
target_states,
pred_std,
mask=self.interior_mask_bool,
),
dim=0,
)
return (
prediction,
target_states,
pred_std,
time_step_loss,
)
def _log_step_loss(
self,
time_step_loss: torch.Tensor,
mean_loss: torch.Tensor,
phase: str,
batch_size: int,
) -> None:
"""Log mean and step-wise losses for a given phase."""
self._warn_skipped_steps(len(time_step_loss), phase)
loss_key = "train_loss" if phase == "train" else f"{phase}_mean_loss"
log_dict = {loss_key: mean_loss}
hparam_name = (
"train_steps_to_log" if phase == "train" else "val_steps_to_log"
)
steps_to_log = getattr(self.hparams, hparam_name, [])
for step in steps_to_log:
if step <= len(time_step_loss):
log_dict[f"{phase}_loss_unroll{step}"] = time_step_loss[
step - 1
]
is_train = phase == "train"
self.log_dict(
log_dict,
prog_bar=is_train,
on_step=is_train,
on_epoch=True,
sync_dist=True,
batch_size=batch_size,
)
[docs]
def validation_step(
self,
batch: Batch,
batch_idx: int,
) -> None:
"""
Perform a single validation step.
Parameters
----------
batch : tuple
The batch of data.
batch_idx : int
The index of the batch.
"""
prediction, target_states, pred_std, time_step_loss = (
self._compute_prediction_and_loss(batch)
)
mean_loss = torch.mean(time_step_loss)
batch_size = batch[0].shape[0]
self._log_step_loss(time_step_loss, mean_loss, "val", batch_size)
entry_mses = metrics.mse(
prediction,
target_states,
pred_std,
mask=self.interior_mask_bool,
sum_vars=False,
)
self.val_metrics["mse"].append(entry_mses)
[docs]
def on_validation_epoch_end(self) -> None:
"""
Perform actions at the end of the validation epoch.
Aggregates and plots validation metrics.
"""
self.aggregate_and_plot_metrics(self.val_metrics, prefix="val")
if (
self.trainer.is_global_zero and self.hparams.metrics_watch # ty: ignore[unresolved-attribute]
):
metrics_watch = (
self.hparams.metrics_watch # ty: ignore[unresolved-attribute]
)
unmatched = set(metrics_watch) - self.matched_metrics
if unmatched:
warnings.warn(
"The following metrics in --metrics_watch "
"were not found during validation phase: "
f"{sorted(unmatched)}. Ensure the metric prefix "
"matches the evaluation mode (expected 'val_')."
)
self.matched_metrics = set()
for metric_list in self.val_metrics.values():
metric_list.clear()
# pylint: disable-next=unused-argument
[docs]
def test_step(
self,
batch: Batch,
batch_idx: int,
) -> None:
"""
Perform a single test step.
Parameters
----------
batch : tuple
The batch of data.
batch_idx : int
The index of the batch.
"""
prediction, target_states, pred_std, time_step_loss = (
self._compute_prediction_and_loss(batch)
)
if self.forecaster.predicts_std:
mean_pred_std = torch.mean(
pred_std[..., self.interior_mask_bool, :], dim=-2
)
self.test_metrics["output_std"].append(mean_pred_std)
mean_loss = torch.mean(time_step_loss)
batch_size = batch[0].shape[0]
self._log_step_loss(time_step_loss, mean_loss, "test", batch_size)
hparams = self.hparams
val_steps_to_log = (
hparams.val_steps_to_log # ty: ignore[unresolved-attribute]
)
for metric_name in ("mse", "mae"):
metric_func = metrics.get_metric(metric_name)
batch_metric_vals = metric_func(
prediction,
target_states,
pred_std,
mask=self.interior_mask_bool,
sum_vars=False,
)
self.test_metrics[metric_name].append(batch_metric_vals)
spatial_loss = self.loss(
prediction, target_states, pred_std, average_grid=False
)
spatial_loss[..., ~self.interior_mask_bool] = float("nan")
self._test_spatial_steps = [
step
for step in val_steps_to_log
if 1 <= step <= spatial_loss.shape[1]
]
log_spatial_losses = spatial_loss[
:,
[step - 1 for step in self._test_spatial_steps],
]
self.spatial_loss_maps.append(log_spatial_losses)
if (
self.trainer.is_global_zero
and self.plotted_examples < self.n_example_pred
):
n_examples = min(
prediction.shape[0],
self.n_example_pred - self.plotted_examples,
)
self.plot_examples(
batch,
n_examples,
prediction=prediction,
split="test",
)
[docs]
def plot_examples(
self,
batch: Batch,
n_examples: int,
split: str,
prediction: torch.Tensor,
) -> None:
"""
Plot example predictions.
Parameters
----------
batch : tuple
The batch of data.
n_examples : int
Number of examples to plot.
split : str
The data split.
prediction : torch.Tensor
The model predictions.
"""
target = batch[1]
time = batch[4]
da_state_stats = self.datastore.get_standardization_dataarray("state")
state_std = torch.tensor(
da_state_stats.state_std.values,
dtype=torch.float32,
device=prediction.device,
)
state_mean = torch.tensor(
da_state_stats.state_mean.values,
dtype=torch.float32,
device=prediction.device,
)
prediction_rescaled = prediction * state_std + state_mean
target_rescaled = target * state_std + state_mean
for pred_slice, target_slice, time_slice in zip(
prediction_rescaled[:n_examples],
target_rescaled[:n_examples],
time[:n_examples],
):
self.plotted_examples += 1
da_prediction = self._create_dataarray_from_tensor(
tensor=pred_slice,
time=time_slice,
split=split,
category="state",
).unstack("grid_index")
da_target = self._create_dataarray_from_tensor(
tensor=target_slice,
time=time_slice,
split=split,
category="state",
).unstack("grid_index")
var_vmin = (
torch.minimum(
pred_slice.flatten(0, 1).min(dim=0)[0],
target_slice.flatten(0, 1).min(dim=0)[0],
)
.cpu()
.numpy()
)
var_vmax = (
torch.maximum(
pred_slice.flatten(0, 1).max(dim=0)[0],
target_slice.flatten(0, 1).max(dim=0)[0],
)
.cpu()
.numpy()
)
var_vranges = list(zip(var_vmin, var_vmax))
example_i = self.plotted_examples
if self.create_gif:
plot_dir_path = os.path.join(
self.logger.save_dir, # ty: ignore[unresolved-attribute]
f"example_plots_{example_i}",
)
os.makedirs(plot_dir_path, exist_ok=True)
png_frames: dict[str, list[str]] = {
var_name: []
for var_name in self.datastore.get_vars_names("state")
}
for t_i, _ in enumerate(zip(pred_slice, target_slice), start=1):
var_figs = [
vis.plot_prediction(
datastore=self.datastore,
title=f"{var_name} ({var_unit}), "
f"t={t_i} ({(self.time_step_int * t_i)}"
f"{self.time_step_unit})",
vrange=var_vrange,
da_prediction=da_prediction.isel(
state_feature=var_i, time=t_i - 1
).squeeze(),
da_target=da_target.isel(
state_feature=var_i, time=t_i - 1
).squeeze(),
)
for var_i, (var_name, var_unit, var_vrange) in enumerate(
zip(
self.datastore.get_vars_names("state"),
self.datastore.get_vars_units("state"),
var_vranges,
)
)
]
for var_name, fig in zip(
self.datastore.get_vars_names("state"), var_figs
):
if isinstance(self.logger, pl.loggers.WandbLogger):
key = f"{var_name}_example"
else:
key = f"{var_name}_example_{example_i}"
if hasattr(self.logger, "log_image"):
self.logger.log_image( # ty: ignore[call-non-callable]
key=key, images=[fig], step=t_i
)
else:
warnings.warn(
f"{self.logger} does not support image logging."
)
if self.create_gif:
png_path = os.path.join(
plot_dir_path,
f"{var_name}_example_{example_i}"
f"_prediction_t_{t_i:02d}.png",
)
fig.savefig(png_path, dpi=100, bbox_inches="tight")
png_frames[var_name].append(png_path)
plt.close("all")
if self.create_gif:
for var_name, frames_for_var in png_frames.items():
if frames_for_var:
gif_path = os.path.join(
plot_dir_path,
f"{var_name}_example_{example_i}_prediction.gif",
)
frames = [Image.open(f) for f in frames_for_var]
try:
frames[0].save(
gif_path,
save_all=True,
append_images=frames[1:],
loop=0,
duration=1000,
)
finally:
for frame in frames:
frame.close()
torch.save(
pred_slice.cpu(),
os.path.join(
self.logger.save_dir, # ty: ignore[unresolved-attribute]
f"example_pred_{self.plotted_examples}.pt",
),
)
torch.save(
target_slice.cpu(),
os.path.join(
self.logger.save_dir, # ty: ignore[unresolved-attribute]
f"example_target_{self.plotted_examples}.pt",
),
)
[docs]
def create_metric_log_dict(
self, metric_tensor: torch.Tensor, prefix: str, metric_name: str
) -> dict[str, Any]:
"""
Create a dictionary of metrics to log.
Parameters
----------
metric_tensor : torch.Tensor
The metric values.
prefix : str
Prefix for the log names (e.g., 'val', 'test').
metric_name : str
The name of the metric.
Returns
-------
dict
Dictionary of logged metrics and figures.
"""
log_dict: dict[str, Any] = {}
metric_fig = vis.plot_error_heatmap(
errors=metric_tensor,
datastore=self.datastore,
)
full_log_name = f"{prefix}_{metric_name}"
log_dict[full_log_name] = metric_fig
if prefix == "test":
metric_fig.savefig(
os.path.join(
self.logger.save_dir, # ty: ignore[unresolved-attribute]
f"{full_log_name}.pdf",
)
)
np.savetxt(
os.path.join(
self.logger.save_dir, # ty: ignore[unresolved-attribute]
f"{full_log_name}.csv",
),
metric_tensor.cpu().numpy(),
delimiter=",",
)
var_names = self.datastore.get_vars_names(category="state")
hparams = self.hparams
metrics_watch = (
hparams.metrics_watch # ty: ignore[unresolved-attribute]
)
var_leads_metrics_watch = (
hparams.var_leads_metrics_watch # ty: ignore[unresolved-attribute]
)
if full_log_name in metrics_watch:
self.matched_metrics.add(full_log_name)
for var_i, timesteps in var_leads_metrics_watch.items():
var_name = var_names[var_i]
for step in timesteps:
key = f"{full_log_name}_{var_name}_step_{step}"
log_dict[key] = metric_tensor[step - 1, var_i]
return log_dict
[docs]
def aggregate_and_plot_metrics(
self, metrics_dict: dict[str, list[torch.Tensor]], prefix: str
) -> None:
"""
Aggregate metrics from all GPUs and plot them.
Parameters
----------
metrics_dict : dict
Dictionary of metrics lists.
prefix : str
Prefix for the log names.
"""
log_dict = {}
for metric_name, metric_val_list in metrics_dict.items():
metric_tensor = self.all_gather_cat(
torch.cat(metric_val_list, dim=0)
)
if self.trainer.is_global_zero:
metric_tensor_averaged = torch.mean(metric_tensor, dim=0)
if "mse" in metric_name:
metric_tensor_averaged = torch.sqrt(metric_tensor_averaged)
metric_name = metric_name.replace("mse", "rmse")
da_state_stats = self.datastore.get_standardization_dataarray(
"state"
)
state_std = torch.tensor(
da_state_stats.state_std.values,
dtype=torch.float32,
device=metric_tensor_averaged.device,
)
metric_rescaled = metric_tensor_averaged * state_std
log_dict.update(
self.create_metric_log_dict(
metric_rescaled, prefix, metric_name
)
)
figure_dict = {
k: v for k, v in log_dict.items() if isinstance(v, plt.Figure)
}
scalar_dict = {
k: v for k, v in log_dict.items() if not isinstance(v, plt.Figure)
}
if self.trainer.is_global_zero and not self.trainer.sanity_checking:
# Log scalars via Lightning's built-in mechanism. No sync_dist:
# we are inside `if is_global_zero`, so only rank 0 reaches this
# call. sync_dist=True would launch a distributed collective
# from one rank and hang DDP. The values are already aggregated
# via all_gather_cat above, so no further sync is needed.
if scalar_dict:
self.log_dict(scalar_dict, rank_zero_only=True)
current_epoch = self.trainer.current_epoch
for key, figure in figure_dict.items():
if not isinstance(self.logger, pl.loggers.WandbLogger):
key = f"{key}-{current_epoch}"
if hasattr(self.logger, "log_image"):
self.logger.log_image( # ty: ignore[call-non-callable]
key=key, images=[figure]
)
plt.close("all")
[docs]
def on_test_epoch_end(self) -> None:
"""
Perform actions at the end of the test epoch.
Aggregates and plots test metrics and spatial loss maps.
"""
self.aggregate_and_plot_metrics(self.test_metrics, prefix="test")
spatial_loss_tensor = self.all_gather_cat(
torch.cat(self.spatial_loss_maps, dim=0)
)
if self.trainer.is_global_zero:
mean_spatial_loss = torch.nanmean(spatial_loss_tensor, dim=0)
hparams = self.hparams
logger_save_dir = (
self.logger.save_dir # ty: ignore[unresolved-attribute]
)
loss_map_figs = [
vis.plot_spatial_error(
error=loss_map,
datastore=self.datastore,
title=f"Test loss, t={t_i} "
f"({(self.time_step_int * t_i)} {self.time_step_unit})",
)
for t_i, loss_map in zip(
self._test_spatial_steps,
mean_spatial_loss,
)
]
for i, fig in enumerate(loss_map_figs):
key = "test_loss"
if not isinstance(self.logger, pl.loggers.WandbLogger):
key = f"{key}_{i}"
if hasattr(self.logger, "log_image"):
self.logger.log_image( # ty: ignore[call-non-callable]
key=key, images=[fig]
)
pdf_loss_map_figs = [
vis.plot_spatial_error(error=loss_map, datastore=self.datastore)
for loss_map in mean_spatial_loss
]
pdf_loss_maps_dir = os.path.join(
logger_save_dir, "spatial_loss_maps"
)
os.makedirs(pdf_loss_maps_dir, exist_ok=True)
for t_i, fig in zip(
self._test_spatial_steps,
pdf_loss_map_figs,
):
fig.savefig(os.path.join(pdf_loss_maps_dir, f"loss_t{t_i}.pdf"))
torch.save(
mean_spatial_loss.cpu(),
os.path.join(
logger_save_dir,
"mean_spatial_loss.pt",
),
)
metrics_watch = (
hparams.metrics_watch # ty: ignore[unresolved-attribute]
)
if metrics_watch:
unmatched = set(metrics_watch) - self.matched_metrics
if unmatched:
warnings.warn(
"The following metrics in --metrics_watch "
"were not found during test phase: "
f"{sorted(unmatched)}. Ensure the metric prefix "
"matches the evaluation mode (expected 'test_')."
)
self.matched_metrics = set()
self.spatial_loss_maps.clear()
self._test_spatial_steps.clear()
# Clear stored test metrics so repeated `trainer.test()` calls on
# the same model instance start from a clean slate (otherwise the
# tensors accumulate and skew the aggregated metrics).
for metric_list in self.test_metrics.values():
metric_list.clear()
# Reset the example-plot counter so example prediction plots are
# generated again on every `trainer.test()` call, not just the
# first one (the guard `plotted_examples < n_example_pred` would
# otherwise stay permanently False).
self.plotted_examples = 0
[docs]
def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
"""
Perform actions when loading a checkpoint.
Handles backward compatibility for older checkpoints.
Parameters
----------
checkpoint : dict
The loaded checkpoint dictionary.
"""
loaded_state_dict = checkpoint["state_dict"]
# 1. Broad namespace remap: for pre-refactor checkpoints
# The old ``ARModel`` was a flat LightningModule. Everything that
# belonged to the predictor needs to be moved to
# 'forecaster.predictor.'
old_keys = list(loaded_state_dict.keys())
for key in old_keys:
if not key.startswith("forecaster.") and key not in (
"interior_mask_bool",
"per_var_std",
):
new_key = f"forecaster.predictor.{key}"
loaded_state_dict[new_key] = loaded_state_dict.pop(key)
# 2. Specific rename from g2m_gnn.grid_mlp -> encoding_grid_mlp
# Will be under forecaster.predictor due to the remap above, or
# already there if from a recent checkpoint before this rename.
if (
"forecaster.predictor.g2m_gnn.grid_mlp.0.weight"
in loaded_state_dict
):
replace_keys = list(
filter(
lambda key: key.startswith(
"forecaster.predictor.g2m_gnn.grid_mlp"
),
loaded_state_dict.keys(),
)
)
for old_key in replace_keys:
new_key = old_key.replace(
"forecaster.predictor.g2m_gnn.grid_mlp",
"forecaster.predictor.encoding_grid_mlp",
)
loaded_state_dict[new_key] = loaded_state_dict[old_key]
del loaded_state_dict[old_key]
if not self.restore_opt:
opt = self.configure_optimizers()
checkpoint["optimizer_states"] = [opt.state_dict()]
