From 3dad68d1919845385d956e3596693b23261106bc Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Mon, 29 Jun 2026 13:13:20 -0500 Subject: [PATCH 01/28] First pass: port DLESyM fork changes into upstream layout. Bring over modulus-uw dev changes through optimized CLN (94b785f), remapping dlwp_healpix_layers to dlwp_healpix/layers and aligning HEALPix ops with physicsnemo.nn.module.hpx. Includes Zarr datapipes, hydrostatic loss, CLN/checkpointing, and related tests. --- .../healpix/base_timeseries_dataset_zarr.py | 623 +++++++++ .../healpix/coupledtimeseries_dataset.py | 19 +- .../healpix/coupledtimeseries_dataset_zarr.py | 380 +++++ physicsnemo/datapipes/healpix/couplers.py | 536 +++---- physicsnemo/datapipes/healpix/data_modules.py | 562 ++------ .../datapipes/healpix/data_modules_zarr.py | 549 ++++++++ .../datapipes/healpix/timeseries_dataset.py | 25 +- .../healpix/timeseries_dataset_zarr.py | 271 ++++ physicsnemo/metrics/climate/healpix_loss.py | 658 ++++++++- physicsnemo/metrics/climate/hydrostasy.py | 830 +++++++++++ .../models/dlwp_healpix/HEALPixRecUNet.py | 100 +- .../models/dlwp_healpix/HEALPixUNet.py | 57 +- .../models/dlwp_healpix/layers/__init__.py | 12 + .../dlwp_healpix/layers/healpix_blocks.py | 382 +++-- .../layers/healpix_constraints.py | 51 + .../dlwp_healpix/layers/healpix_decoder.py | 100 +- .../dlwp_healpix/layers/healpix_encoder.py | 87 +- .../dlwp_healpix/layers/normalization.py | 230 +++ test/datapipes/test_healpix_couple_zarr.py | 1230 +++++++++++++++++ test/datapipes/test_healpix_zarr.py | 782 +++++++++++ test/metrics/test_hydrostatic_loss.py | 167 +++ test/models/dlwp_healpix/_cln_reference.py | 73 + .../test_conditional_layer_norm.py | 265 ++++ test/pytest_utils.py | 90 ++ 24 files changed, 7205 insertions(+), 874 deletions(-) create mode 100644 physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py create mode 100644 physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py create mode 100644 physicsnemo/datapipes/healpix/data_modules_zarr.py create mode 100644 physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py create mode 100644 physicsnemo/metrics/climate/hydrostasy.py create mode 100644 physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py create mode 100644 physicsnemo/models/dlwp_healpix/layers/normalization.py create mode 100644 test/datapipes/test_healpix_couple_zarr.py create mode 100644 test/datapipes/test_healpix_zarr.py create mode 100644 test/metrics/test_hydrostatic_loss.py create mode 100644 test/models/dlwp_healpix/_cln_reference.py create mode 100644 test/models/dlwp_healpix/test_conditional_layer_norm.py diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py new file mode 100644 index 0000000000..39f55dca62 --- /dev/null +++ b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py @@ -0,0 +1,623 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import logging +import os +import warnings +from abc import ABC, abstractmethod +from typing import List, Optional, Sequence, Tuple, Union + +import numpy as np +import pandas as pd +import xarray as xr +import zarr +from physicsnemo.datapipes.datapipe import Datapipe +from physicsnemo.datapipes.meta import DatapipeMetaData +from omegaconf import DictConfig, OmegaConf +from torch.utils.data import Dataset + +logger = logging.getLogger(__name__) + + +def _is_object_store_path(path: str) -> bool: # pragma: no cover + """Check if path is an object store path (contains :// or ::). + + Parameters + ---------- + path : str + Path to check + + Returns + ------- + bool + True if path appears to be an object store path + """ + return "://" in str(path) or "::" in str(path) + + +def _check_availability(path: str) -> None: # pragma: no cover + """ + Check if path exists or fsspec is available for object store paths + + Parameters + ---------- + path : str + Path to check + + Raises + ------ + ImportError + If path is an object store path but fsspec is not available + FileNotFoundError + If the path is file and doesn't exist + """ + if _is_object_store_path(path): + if not importlib.util.find_spec("fsspec"): + raise ImportError( + f"fsspec is required to access object store paths like '{path}'. " + "Please install fsspec with: pip install fsspec" + ) + elif not os.path.exists(path): + raise FileNotFoundError(f"Dataset not found at specified location: {path}") + + +class BaseTimeSeriesDatasetZarr(Dataset, Datapipe, ABC): + """Abstract base class for time series datasets using Zarr storage. + + This class provides the core functionality for loading and processing time series data + stored in Zarr format. It handles data loading, scaling, and time management. + Subclasses must implement the __getitem__ method to define specific data retrieval logic. + """ + + def __init__( + self, + dataset_path: str, + input_variables: Sequence, + output_variables: Sequence = None, + constant_variables: Sequence = None, + scaling: DictConfig = None, + input_time_dim: int = 1, + output_time_dim: int = 1, + data_time_step: Union[int, str] = "3h", + time_step: Union[int, str] = "6h", + gap: Union[int, str, None] = None, + batch_size: int = 32, + drop_last: bool = False, + add_insolation: bool = False, + forecast_init_times: Optional[Sequence] = None, + start_date: Optional[Union[int, str]] = None, + end_date: Optional[Union[int, str]] = None, + add_train_noise: bool = False, + train_noise_params: DictConfig = None, + train_noise_seed: int = 42, + meta: DatapipeMetaData = None, + ): + """Initialize base time series dataset. + + Parameters + ---------- + dataset_path : str + Path to the Zarr dataset + input_variables : Sequence + Variables to use as model inputs + output_variables : Sequence, optional + Variables to predict as outputs. If None, uses input_variables + constant_variables : Sequence, optional + Constant fields used as additional inputs + scaling : DictConfig, optional + Configuration for data scaling/normalization + input_time_dim : int, default=1 + Number of time steps in input sequence + output_time_dim : int, default=1 + Number of time steps to predict + data_time_step : Union[int, str], default="3h" + Either integer hours or a str interpretable by pandas + Time resolution of raw data + time_step : Union[int, str], default="6h" + Either integer hours or a str interpretable by pandas + Time step between predictions + gap : Union[int, str, None], optional + Either integer hours or a str interpretable by pandas + Time gap between input and output sequences + batch_size : int, default=32 + Number of samples per batch + drop_last : bool, default=False + Whether to drop last incomplete batch + add_insolation : bool, default=False + Whether to add solar insolation as input + forecast_init_times : Sequence, optional + A Sequence of pandas Timestamps + Specific times to initialize forecasts + start_date : Union[int, str], optional + Start date/index from which to load data + end_date : Union[int, str], optional + End date/index to which to load data + add_train_noise : bool, default=False + Whether to add train noise + train_noise_params : DictConfig, optional + Standard deviation of train noise + train_noise_seed : int, default=42 + Seed for train noise + meta : DatapipeMetaData, optional + Metadata for the datapipe + """ + Datapipe.__init__(self, meta=meta) + + self.dataset_path = dataset_path + self.scaling = OmegaConf.to_object(scaling) if scaling else None + self.input_time_dim = input_time_dim + self.output_time_dim = output_time_dim + self.data_time_step = self._convert_time_step(data_time_step) + self.time_step = self._convert_time_step(time_step) + self.gap = self._convert_time_step(gap if gap is not None else time_step) + self.batch_size = batch_size + self.drop_last = drop_last + self.add_insolation = add_insolation + self.forecast_init_times = forecast_init_times + self.forecast_mode = self.forecast_init_times is not None + self.input_variables = input_variables + self.output_variables = ( + input_variables if output_variables is None else output_variables + ) + self.constant_variables = constant_variables + self.all_variables = list( + set(self.input_variables).union(self.output_variables) + ) + self.all_scaling = None + + # Check if for fsspec if necessary and make sure path exist + _check_availability(dataset_path) + + self.ds = zarr.open(dataset_path) + + if ( + start_date is None or end_date is None + ) and self.forecast_init_times is None: + raise ValueError( + "Either start and end date or forecast_init_times must be provided" + ) + + # Validate channels exist + channels = set(self.input_variables).union(self.output_variables) + missing_channels = channels - set(self.ds["channel_in"][:]) + if len(missing_channels) > 0: + raise KeyError( + f"Requested Input, coupled, or output variables not found in dataset: {missing_channels}" + ) + + self._get_time_da(self.dataset_path, start_date, end_date) + + self.all_variable_indices = [ + int(np.where(self.ds["channel_in"][:] == ch)[0][0]) + for ch in self.all_variables + ] + + # Validate constants exist + if constant_variables: + missing_constants = set(constant_variables) - set(self.ds["channel_c"][:]) + if len(missing_constants) > 0: + raise KeyError( + f"Requested constants not found in dataset: {missing_constants}" + ) + + self.constant_variable_indices = ( + [ + int(np.where(self.ds["channel_c"][:] == ch)[0][0]) + for ch in self.constant_variables + ] + if self.constant_variables + else None + ) + self.input_variable_indices = [ + self.all_variables.index(inp_ch) for inp_ch in self.input_variables + ] + self.output_variable_indices = [ + self.all_variables.index(out_ch) for out_ch in self.output_variables + ] + + # Length of the data window needed for one sample + if self.forecast_mode: + self._window_length = self.interval * (self.input_time_dim - 1) + 1 + else: + self._window_length = ( + self.interval * (self.input_time_dim - 1) + + 1 + + (self.gap // self.data_time_step) + + self.interval * (self.output_time_dim - 1) + ) + + self._batch_window_length = self.batch_size + self._window_length - 1 + self._output_delay = self.interval * (self.input_time_dim - 1) + ( + self.gap // self.data_time_step + ) + + # Indices within a batch + self._input_indices = [ + list(range(n, n + self.interval * self.input_time_dim, self.interval)) + for n in range(self.batch_size) + ] + self._output_indices = [ + list( + range( + n + self._output_delay, + n + self.interval * self.output_time_dim + self._output_delay, + self.interval, + ) + ) + for n in range(self.batch_size) + ] + + self.spatial_dims = ( + self.ds["face"].shape[0], + self.ds["height"].shape[0], + self.ds["width"].shape[0], + ) + + # Cached values + self.lat = np.asarray(self.ds["lat"]) + self.lon = np.asarray(self.ds["lon"]) + self.input_scaling = None + self.target_scaling = None + self.constant_scaling = None + self.constants = None + + if self.scaling: + self._get_scaling_da() + if self.constant_variables: + self.constants = self.get_constants() + + self.add_train_noise = add_train_noise + self.train_noise_params = train_noise_params + if self.add_train_noise: + self.rng = np.random.default_rng(train_noise_seed) + + @staticmethod + def _convert_time_step(dt: Union[int, str]) -> pd.Timedelta: + """Convert time step specification to Timedelta. + + Parameters + ---------- + dt : Union[int, str] + Either integer hours or string time to convert to Timedelta + + Returns + ------- + pd.Timedelta + Converted time delta object + """ + return pd.Timedelta(hours=dt) if isinstance(dt, int) else pd.Timedelta(dt) + + def _get_time_da( + self, + dataset_path: str, + start_date: Optional[Union[int, str]], + end_date: Optional[Union[int, str]], + ) -> None: + """Load and decode time array from dataset. + + Sets up time-related attributes including total samples, start index, + and forecast initialization indices. + + Parameters + ---------- + dataset_path : str + Path to Zarr dataset + start_date : Optional[Union[int, str]] + Start date/index for data slice + end_date : Optional[Union[int, str]] + End date/index for data slice + """ + # Check if fsspec is available for object store paths + _check_availability(dataset_path) + + ds = xr.open_zarr(dataset_path) + + if "time" not in ds: + raise KeyError(f"Dataset missing time. Dataset provided {dataset_path}") + + if np.datetime64(start_date) < ds.time[0]: + warnings.warn( + f"Start date {start_date} is before first available date {ds.time[0].values}" + ) + if ds.time[-1] < np.datetime64(end_date): + warnings.warn( + f"End date {end_date} is after last available date {ds.time[-1].values}" + ) + + # used when we need all the dates to calculate things like offset indices + self.time_da = ds.time.copy(deep=True) + # a list of dates that is available to fetch, used for things like inferencers + self.times = self.time_da.sel(time=slice(start_date, end_date)) + self.total_samples = self.times.shape[0] + + if start_date: + if isinstance(start_date, int): + self.start_index = start_date + else: + self.start_index = int( + np.where(self.time_da == np.datetime64(start_date))[0][0] + ) + else: + self.start_index = 0 + + # Validate time stepping + if (self.time_step % self.data_time_step).total_seconds() != 0: + raise ValueError( + f"'time_step' must be a multiple of 'data_time_step' " + f"(got {self.time_step} and {self.data_time_step}" + ) + if (self.gap % self.data_time_step).total_seconds() != 0: + raise ValueError( + f"'gap' must be a multiple of 'data_time_step' " + f"(got {self.gap} and {self.data_time_step}" + ) + self.interval = self.time_step // self.data_time_step + + # Verify timestep matches data + ds_dt = pd.Timedelta((self.time_da[1] - self.time_da[0]).values) + if not (ds_dt == self.data_time_step): + warnings.warn( + f"Dataset dt {ds_dt} doesn't match configuration dt {self.data_time_step}. " + "This could be a configuration error or a dataset mismatch." + ) + + # Find indices of init times for forecast mode + if self.forecast_mode: + if self.batch_size != 1: + self.batch_size = 1 + warnings.warn( + "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`; " + "setting it now" + ) + self._forecast_init_indices = np.array( + [ + int(np.where(self.time_da == s)[0][0]) + for s in self.forecast_init_times + ], + dtype="int", + ) - ((self.input_time_dim - 1) * self.interval) + else: + self._forecast_init_indices = None + + def _get_scaling_da(self) -> None: + """Setup data scaling parameters. + + Processes scaling configuration and sets up scaling parameters for: + - Input variables + - Target/output variables + - All variables combined + - Constant fields + + Raises + ------ + KeyError + If scaling parameters are missing for any variables + """ + scaling_df = pd.DataFrame.from_dict(self.scaling).T + scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} + scaling_da = scaling_df.to_xarray().astype("float32") + + try: + self.input_scaling = scaling_da.sel(index=self.input_variables).rename( + {"index": "channel_in"} + ) + self.input_scaling = { + "mean": np.expand_dims( + self.input_scaling["mean"].values.copy(), (0, 2, 3, 4) + ), + "std": np.expand_dims( + self.input_scaling["std"].values.copy(), (0, 2, 3, 4) + ), + } + except (ValueError, KeyError): + missing = [ + m for m in self.input_variables if m not in list(self.scaling.keys()) + ] + raise KeyError( + f"Input channels {missing} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" + ) + + try: + self.target_scaling = scaling_da.sel(index=self.output_variables).rename( + {"index": "channel_out"} + ) + self.target_scaling = { + "mean": np.expand_dims( + self.target_scaling["mean"].values.copy(), (0, 2, 3, 4) + ), + "std": np.expand_dims( + self.target_scaling["std"].values.copy(), (0, 2, 3, 4) + ), + } + except (ValueError, KeyError): + missing = [ + m for m in self.output_variables if m not in list(self.scaling.keys()) + ] + raise KeyError( + f"Target channels {missing} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" + ) + + self.all_scaling = scaling_da.sel(index=self.all_variables).rename( + {"index": "channel_in"} + ) + self.all_scaling = { + "mean": np.expand_dims( + self.all_scaling["mean"].values.copy(), (0, 2, 3, 4) + ), + "std": np.expand_dims(self.all_scaling["std"].values.copy(), (0, 2, 3, 4)), + } + + if self.constant_variables: + # Check that all constant variables are present in scaling data + missing_constants = [ + var + for var in self.constant_variables + if var not in list(self.scaling.keys()) + ] + if missing_constants: + raise KeyError( + f"Constant variables {missing_constants} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" + ) + + try: + + self.constant_scaling = scaling_da.sel( + index=self.constant_variables + ).rename({"index": "channel_out"}) + self.constant_scaling = { + "mean": np.expand_dims( + self.constant_scaling["mean"].values.copy(), (1, 2, 3) + ), + "std": np.expand_dims( + self.constant_scaling["std"].values.copy(), (1, 2, 3) + ), + } + except (ValueError, KeyError): + missing = [ + m + for m in self.constant_variables + if m not in list(self.scaling.keys()) + ] + raise KeyError( + f"Constant channels {missing} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" + ) + + def get_constants(self) -> np.ndarray: + """Get constant fields used in dataset. + + Returns + ------- + np.ndarray + Array of constant fields with shape [F, C, H, W] + where F=faces, C=channels, H=height, W=width + """ + if self.constants is not None: + return self.constants + + if self.constant_variables is None: + return None + + const = np.asarray(self.ds["constants"][self.constant_variable_indices]) + + if self.constant_scaling: + const = (const - self.constant_scaling["mean"]) / self.constant_scaling[ + "std" + ] + + self.constants = np.transpose(const, axes=(1, 0, 2, 3)) + return self.constants + + def _get_time_index(self, item: int) -> Tuple[Tuple[int, int], int]: + """Get time indices for specified sample. + + Parameters + ---------- + item : int + Sample index + + Returns + ------- + Tuple[Tuple[int, int], int] + ((start_index, end_index), batch_size) + Time window indices and actual batch size + """ + window_start_index = ( + self._forecast_init_indices[item] + if self.forecast_mode + else item * self.batch_size + self.start_index + ) + window_max_index = ( + window_start_index + self._window_length + if self.forecast_mode + else (item + 1) * self.batch_size + self._window_length + self.start_index + ) + if not self.drop_last and window_max_index > self.total_samples: + batch_size = self.batch_size - (window_max_index - self.total_samples) + else: + batch_size = self.batch_size + return (window_start_index, window_max_index), batch_size + + def _get_forecast_sol_times(self, item: int) -> np.ndarray: + """Get times for calculating solar insolation. + + Parameters + ---------- + item : int + Sample index + + Returns + ------- + np.ndarray + Array of timestamps for insolation calculation + """ + time_index, _ = self._get_time_index(item) + if self.forecast_mode: + timedeltas = ( + np.array(self._input_indices[0] + self._output_indices[0]) + ) * self.data_time_step + return self.time_da[time_index[0]].values + timedeltas + return self.time_da[slice(*time_index)].values + + def __len__(self) -> int: + """Get number of samples available in the dataset based on + timedeltas, gaps, start and end dates. + + Returns + ------- + int + Total number of available samples + """ + if self.forecast_mode: + return len(self._forecast_init_indices) + length = (self.total_samples - self._window_length + 1) / self.batch_size + if self.drop_last: + return int(np.floor(length)) + return int(np.ceil(length)) + + @abstractmethod + def __getitem__( + self, item: int + ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]]: + """Get requested sample - must be implemented by subclasses. + + Parameters + ---------- + item : int + Sample index + + Returns + ------- + Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]] + In forecast mode: List of input arrays + In training mode: Tuple of (input arrays, target array) + + Input arrays are in order: + - Model inputs [B, F, T, C, H, W] + - Insolation (if enabled) [B, F, T, 1, H, W] + - Constants (if provided) [F, C, H, W] + - Additional data (in subclasses) + + Target array has shape [B, F, T, C, H, W] + where: + B = batch size + F = faces + T = time steps + C = channels + H = height + W = width + """ + pass diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index a0087bed00..8acc22f126 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -14,8 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations - import gc import logging import time @@ -25,17 +23,15 @@ import numpy as np import pandas as pd import torch +import xarray as xr from omegaconf import DictConfig, OmegaConf -from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.meta import DatapipeMetaData from physicsnemo.utils.insolation import insolation from . import couplers from .timeseries_dataset import TimeSeriesDataset -xr = OptionalImport("xarray") - logger = logging.getLogger(__name__) @@ -211,9 +207,14 @@ def __getitem__(self, item): axis=2, ) - input_array = (input_array - self.input_scaling["mean"]) / self.input_scaling[ - "std" - ] + # for models with extra outputs + if len(self.ds["targets"].channel_out) != (len(self.ds["inputs"].channel_in)-len(self.couplings[0].variables)): + input_array = (input_array - self.input_scaling["mean"][:,:-len(self.couplings[0].variables)]) \ + / self.input_scaling["std"][:,:-len(self.couplings[0].variables)] + else: + input_array = (input_array - self.input_scaling["mean"]) \ + / self.input_scaling["std"] + if not self.forecast_mode: # BAD NEWS: Indexing the array as commented out below causes unexpected behavior in target creation. # leaving this in here as a warning diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py new file mode 100644 index 0000000000..65f7df4f66 --- /dev/null +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple, Union + +import numpy as np +import pandas as pd +import torch +import zarr +from omegaconf import DictConfig, OmegaConf + +from physicsnemo.datapipes.meta import DatapipeMetaData +from physicsnemo.utils.insolation import insolation + +from . import couplers +from .timeseries_dataset_zarr import TimeSeriesDatasetZarr +from .base_timeseries_dataset_zarr import _check_availability + +logger = logging.getLogger(__name__) + + +@dataclass +class MetaData(DatapipeMetaData): + """Metadata for this datapipe""" + + name: str = "CoupledTimeSeries" + # Optimization + auto_device: bool = False + cuda_graphs: bool = False + # Parallel + ddp_sharding: bool = False + + +class CoupledTimeSeriesDatasetZarr(TimeSeriesDatasetZarr): + """Dataset for coupling time series data with external earth system components. + + This class extends the base time series functionality to include coupling with + external data sources like ocean models, land models, etc. It supports: + - Integration of coupled data during training/inference + - Addition of training noise to improve generalization + - Step-by-step integration for forecasting + """ + + def __init__( + self, + dataset_path: str, + scaling: DictConfig, + input_variables: Sequence, + output_variables: Sequence = None, + constant_variables: Sequence = None, + input_time_dim: int = 1, + output_time_dim: int = 1, + data_time_step: Union[int, str] = "3h", + time_step: Union[int, str] = "6h", + gap: Union[int, str, None] = None, + batch_size: int = 32, + drop_last: bool = False, + add_insolation: bool = False, + forecast_init_times: Optional[Sequence] = None, + start_date: Optional[Union[int, str]] = None, + end_date: Optional[Union[int, str]] = None, + couplings: Sequence = [], + meta: DatapipeMetaData = MetaData(), + add_train_noise: bool = False, + train_noise_params: DictConfig = None, + train_noise_seed: int = 42, + ): + """Initialize coupled time series dataset. + + Parameters + ---------- + couplings : Sequence + List of coupling configurations, each containing: + - coupler: Name of coupler class to use + - params: Parameters for the coupler + add_train_noise : bool, default=False + Whether to add noise during training + train_noise_params : DictConfig, optional + Configuration for training noise, containing: + - inputs: Dict mapping variable names to noise std + - couplings: Dict mapping variable names to noise std + train_noise_seed : int, default=42 + Random seed for noise generation + + Other parameters are same as BaseTimeSeriesDatasetZarr. + See base class for detailed parameter descriptions. + """ + self.coupled_variables = [] + for c in couplings: + self.coupled_variables.append(c["params"]["variables"]) + + # We setup couplers first so superclass can properly initialize + # and set the scaling + _check_availability(dataset_path) + self.ds = zarr.open(dataset_path) + self.couplings = [ + getattr(couplers, c["coupler"])( + self.ds, + **OmegaConf.to_object(DictConfig(c))["params"], + ) + for c in couplings + ] + + super().__init__( + dataset_path=dataset_path, + scaling=scaling, + input_variables=input_variables, + output_variables=output_variables, + constant_variables=constant_variables, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + data_time_step=data_time_step, + time_step=time_step, + gap=gap, + batch_size=batch_size, + drop_last=drop_last, + add_insolation=add_insolation, + forecast_init_times=forecast_init_times, + start_date=start_date, + end_date=end_date, + add_train_noise=add_train_noise, + train_noise_params=train_noise_params, + train_noise_seed=train_noise_seed, + meta=meta, + ) + + # calculate static indices for coupling + for c in self.couplings: + c.compute_coupled_indices(self.interval, self.data_time_step) + # keep track of integration steps + self.integration_step = ( + 1 # starts at 1 because first step is done by __getitem__ + ) + self.last_batch = None # keeps track of batch info for coupler + self.curr_item = None # keeps track of current initialization + + def _get_scaling_da(self) -> None: + """extend scaling to include coupling-specific additions. + + Extends base scaling setup to: + - Create scaling parameters for coupled variables + - Pass scaling info to coupler objects + """ + scaling_df = pd.DataFrame.from_dict(self.scaling).T + scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} + scaling_da = scaling_df.to_xarray().astype("float32") + + for c in self.couplings: + c.set_scaling(scaling_da) + super()._get_scaling_da() + + def __getitem__( + self, item: int + ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]]: + """Get a batch of coupled time series data. + + This implementation extends the base time series data loading to include: + 1. Integration of coupled data sources + 2. Addition of training noise if enabled + 3. Tracking of batch info for step-by-step integration + + Parameters + ---------- + item : int + Sample index + + Returns + ------- + Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]] + In forecast mode: List of input arrays + In training mode: Tuple of (input arrays, target array) + + Input arrays are in order: + - Model inputs [B, F, T, C, H, W] + - Insolation (if enabled) [B, F, T, 1, H, W] + - Constants (if provided) [F, C, H, W] + - Coupled inputs (if provided) [T, B, C, F, H, W] + + Target array has shape [B, F, T, C, H, W] + where: + B = batch size + F = faces + T = time steps + C = channels + H = height + W = width + + Raises + ------ + IndexError + If item is out of range + """ + + # check for valid item and adjust for negative indices + if item < 0: + item = len(self) + item + if item < 0 or item > len(self): + raise IndexError( + f"index {item} out of range for dataset with length {len(self)}" + ) + + # start range + torch.cuda.nvtx.range_push("CoupledTimeSeriesDataset:__getitem__") + + if self.forecast_mode: + inputs_result = super().__getitem__(item) + else: + inputs_result, targets = super().__getitem__(item) + + # used by the couplers to determine what time index to load + # see method "next_integration()" for details + time_index, this_batch = self._get_time_index(item) + batch = {"time": slice(*time_index)} + self.last_batch = batch + + torch.cuda.nvtx.range_push( + "CoupledTimeSeriesDataset:__getitem__:retrieve_coupled" + ) + # retrieve coupled inputs + if len(self.couplings) > 0: + integrated_couplings = np.concatenate( + [ + c.construct_integrated_couplings(batch, this_batch) + for c in self.couplings + ], + axis=2, + ) + torch.cuda.nvtx.range_pop() # CoupledTimeSeriesDataset:__getitem__:retrieve_coupled + + # Insolation + if self.add_insolation: + # update current item and reset integration_step counter for further integrations which need + # insolation but bypass this method see method "next_integration()" for details + self.curr_item = item + self.integration_step = 1 + + if not self.forecast_mode and self.add_train_noise: + torch.cuda.nvtx.range_push( + "CoupledTimeSeriesDataset:__getitem__:add_train_noise" + ) + for c in self.couplings: + for i, v in enumerate(c.variables): + integrated_couplings[i, :, :] += self.rng.normal( + loc=0, + scale=self.train_noise_params["couplings"][v]["std"], + size=integrated_couplings[i, :, :].shape, + ) + torch.cuda.nvtx.range_pop() # CoupledTimeSeriesDataset:__getitem__:add_train_noise + + # append integrated couplings + if len(self.couplings) > 0: + inputs_result.append(integrated_couplings) + + torch.cuda.nvtx.range_pop() # CoupledTimeSeriesDataset:__getitem__ + if self.forecast_mode: + return inputs_result + + return inputs_result, targets + + def _get_next_insolation(self, time_offset: int) -> torch.Tensor: + """Calculate insolation for next integration step. + + Parameters + ---------- + time_offset : int + Time offset for insolation calculation + + Returns + ------- + torch.Tensor + Insolation tensor + """ + sol = torch.tensor( + insolation( + self._get_forecast_sol_times(self.curr_item) + time_offset, + self.lat, + self.lon, + )[:, None] + ) + decoder_inputs = np.empty( + (1, self.input_time_dim + self.output_time_dim, 1) + self.spatial_dims, + dtype="float32", + ) + decoder_inputs[0] = sol + return torch.tensor(decoder_inputs.transpose(0, 3, 1, 2, 4, 5)) + + def _get_next_couplings(self, offset: int) -> Optional[torch.Tensor]: + """Get coupled inputs for next integration step. + + Parameters + ---------- + offset : int + Integration offset + + Returns + ------- + Optional[torch.Tensor] + Coupled inputs tensor if couplings exist, None otherwise + """ + if not self.couplings: + return None + + # account for multiple integration steps without fetching data + batch = self.last_batch.copy() + batch = { + "time": slice( + self.last_batch["time"].start + offset, + self.last_batch["time"].stop + offset, + self.last_batch["time"].step, + ) + } + integrated_couplings = np.concatenate( + [ + c.construct_integrated_couplings( + batch=batch, + ) + for c in self.couplings + ], + axis=2, + ) + + return torch.tensor(integrated_couplings) + + def next_integration( + self, model_outputs: torch.Tensor, constants: torch.Tensor + ) -> List[torch.Tensor]: + """Get inputs for next integration step with coupling data. + + Parameters + ---------- + model_outputs : torch.Tensor + Model outputs from previous step [B, F, T, C, H, W] + constants : torch.Tensor + Constant fields [F, C, H, W] + + Returns + ------- + List[torch.Tensor] + List of input tensors for next step + """ + inputs_result = [] + + # Get prognostic inputs from model outputs + init_time_dim = len(self._input_indices[0]) + prognostic_inputs = model_outputs[:, :, -init_time_dim:] + inputs_result.append(prognostic_inputs) + + # Add insolation if needed + if self.add_insolation: + time_offset = self.time_step * self.output_time_dim * self.integration_step + inputs_result.append(self._get_next_insolation(time_offset)) + + # Add constants + inputs_result.append(constants) + + # Add coupled inputs if any + offset = self.interval * self.output_time_dim * self.integration_step + coupled_inputs = self._get_next_couplings(offset) + if coupled_inputs is not None: + inputs_result.append(coupled_inputs) + + # Increment integration step + self.integration_step += 1 + + return inputs_result diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index a0960b2e6d..607b14dcfc 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -14,28 +14,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations - import logging +from abc import ABC, abstractmethod from typing import Sequence +import cftime import numpy as np import pandas as pd import torch as th - -from physicsnemo.core.version_check import OptionalImport - -xr = OptionalImport("xarray") +import xarray as xr +import zarr as zr logger = logging.getLogger(__name__) -class ConstantCoupler: +class BaseCoupler(ABC): """ - coupler used to interface two component of earth system + Base class for couplers used to interface two components of earth system. - constant coupler will take the the coupled field at integration time and - force the model with this field consistently + This class contains common functionality shared by different coupler implementations. """ def __init__( @@ -47,7 +44,8 @@ def __init__( input_time_dim: int = 2, output_time_dim: int = 2, input_times: Sequence = [pd.Timedelta("24h"), pd.Timedelta("48h")], - prepared_coupled_data=True, + prepared_coupled_data: bool = True, + time_first: bool = True, ): """ Parameters @@ -59,7 +57,8 @@ def __init__( forecasting batch size should be 1 variables: Sequence sequence of strings that indicate the coupled variable - names in the dataset + names in the dataset. All names should be in the dataset with + an optional time component at the end, eg ttr-48h presteps: int, optional the number of model steps used to initialize the hidden state. If not using a GRU, prestep is 0, default 0 @@ -71,15 +70,18 @@ def __init__( sequence of pandas Timedelta objects that indicate which times are to be coupled, default [pd.Timedelta("24h"), pd.Timedelta("48h")] prepared_coupled_data: boolean, optional - If True assumes data in dataset has been prepared approiately for training: + If True assumes data in dataset has been prepared appropriately for training: averages have already been calculated so that each time step denotes the right side of a averaging_window window. - This is highly remcommended for training, default True + This is highly recommended for training, default True + time_first: boolean, optional + Whether the coupled data should be permuted to have the time dimension first + [T, B, C, F, H, W] rather than [B, F, T, C, H, W] """ # extract important meta data from ds self.ds = dataset self.batch_size = batch_size - self.spatial_dims = self.ds.inputs.shape[2:] + self.spatial_dims = self.ds["inputs"].shape[2:] self.variables = variables self.presteps = presteps self.input_time_dim = input_time_dim @@ -89,38 +91,42 @@ def __init__( self.output_channels = len(self.variables) * len(self.input_times) self.timevar_dim = self._compute_timevar_dim() self.coupled_inputs_shape = None - self.scaling_dict = None + self.coupled_scaling = None self._coupled_offsets = None self.coupled_mode = False self.integrated_couplings = None + self.ds_variable_indices = [] + self.time_first = time_first if not prepared_coupled_data: - logger.log( - logging.DEBUG, - "Assuming coupled data is not preprocessed preparing data.", - ) - self._prepare_coupled_data() + raise NotImplementedError("Data preparation not yet implemented") + + if type(self.ds) == xr.Dataset: + self.use_zarr = False + elif type(self.ds) == zr.Group: + self.use_zarr = True + self.ds_variable_indices = [ + i + for i, ic in enumerate(self.ds["channel_in"]) + for v in self.variables + if ic == v + ] else: - logger.log( - logging.DEBUG, - "**Assuming coupled data has been prepared properly, using coupled field[s] from " - 'dataset "as-is"**', + raise TypeError( + f"Coupler only supports xarray Datasets or zarr Groups, got {type(self.ds)}" ) - def _prepare_coupled_data(self): - # TODO: write function to lazily compute average as spcified in time scheme - raise NotImplementedError("Data preparation not yet implemented") - def _compute_coupled_integration_dim(self): return self.presteps + max(self.output_time_dim // self.input_time_dim, 1) def _compute_timevar_dim(self): return len(self.input_times) * len(self.variables) + @abstractmethod def compute_coupled_indices(self, interval, data_time_step): """ - Called by CoupledDataset to compute static indices for training - samples + Called by CoupledDataset to compute static indices for training samples. + Must be implemented by subclasses as the logic varies between coupler types. Parameters ---------- @@ -129,28 +135,25 @@ def compute_coupled_indices(self, interval, data_time_step): data_time_step: dataset timestep """ - # create array of static coupled offstes that accompany each batch - self._coupled_offsets = np.empty( - [self.batch_size, self.coupled_integration_dim, len(self.input_times)] - ) - for b in range(self.batch_size): - for i in range(self.coupled_integration_dim): - self._coupled_offsets[b, i, :] = b + np.array( - [ts / data_time_step for ts in self.input_times] - ) - - self._coupled_offsets = self._coupled_offsets.astype(int) + pass def set_scaling(self, scaling_da): """ - Called by CoupledDataset to compute static indices for training - samples + Called by CoupledDataset to compute static indices for training samples Parameters ---------- scaling_da: xarray.DataArray values used to scale input data, uses mean and std """ + # verify all the channels are there for scaling, this avoids an opaque + # "not all values found in index 'index'"" error that looks like its from hydra + missing_channels = set(self.variables) - set(scaling_da.index.values) + if len(missing_channels) > 0: + raise KeyError( + f"Coupled variable(s) not found in scaling values: {missing_channels}" + ) + coupled_scaling = scaling_da.sel(index=self.variables).rename( {"index": "channel_in"} ) @@ -160,22 +163,40 @@ def set_scaling(self, scaling_da): } def setup_coupling(self, coupled_module): - # To expediate the coupling process the coupled_forecast + """ + Sets up the coupling between the coupled variables and the provided module + + Parameters + ---------- + coupled_module: physicsnemo.datapipes.healpix.TimeSeriesDataset + The module which this coupler will be coupled against. + """ + # To expedite the coupling process the coupled_forecast # get proper channels from coupled component output output_channels = coupled_module.output_variables - # A bit convoluted. Prepared coupled variables - # are given a suffix for training associated with their - # trailing average increment e.g. 'z1000-48H'. To extract - # thr proper field from the coupled model output, we see if - # we check if the coupled model output var is in self.variables. - # - # for example 'z1000' is in 'z1000-48H' + # A bit convoluted. Some variable names are present in the dataset as is, + # Some prepared coupled variables are given a suffix for training associated + # with a time increment suach as a trailing average increment e.g. 'z1000-48H'. + # Some variables may have an additional suffix, e.g. 'z1000-3H-48H'. The final + # suffix (if it exists) is used to determine the coupling increment. channel_indices = [ i for i, oc in enumerate(output_channels) for v in self.variables - if oc == v.split("-")[0] + # extract everthing before the last "-" if there is one in the name + if (("-" not in v and oc == v) or (oc == "-".join(v.split("-")[:-1]))) ] + # check for missing variables + if len(self.variables) != len(channel_indices): + found_channels = [ + oc + for oc in output_channels + for v in self.variables + # extract everthing before the last - + if (("-" not in v and oc == v) or (oc == "-".join(v.split("-")[:-1]))) + ] + missing_channels = set(self.variables) - set(found_channels) + raise ValueError(f"Missing variables in coupled module: {missing_channels}") self.coupled_channel_indices = channel_indices def reset_coupler(self): @@ -183,11 +204,11 @@ def reset_coupler(self): self.integrated_couplings = None self.preset_coupled_fields = None + @abstractmethod def set_coupled_fields(self, coupled_fields: th.tensor): """ Set the data for the coupled field for the next iteration of the dataloader. - Instead of loading data from the dataset the data from coupled_fields will - be returned instead. + Must be implemented by subclasses as the processing logic varies. Parameters ---------- @@ -195,26 +216,37 @@ def set_coupled_fields(self, coupled_fields: th.tensor): The data to use when the dataloader requests coupled fields. Expected format is [B, F, T, C, H, W] """ - if coupled_fields.shape[0] != self.batch_size: - raise ValueError( - f"Batch size of coupled field {coupled_fields.shape[0]} doesn't " - f" match configured batch size {self.batch_size}" - ) - # create buffer for coupling - coupled_fields = coupled_fields[ - :, :, :, self.coupled_channel_indices, :, : - ].permute(2, 0, 3, 1, 4, 5) - self.preset_coupled_fields = th.empty( - [self.coupled_integration_dim, self.batch_size, self.timevar_dim] - + list(self.spatial_dims) + pass + + def _construct_integrated_couplings_from_dataset(self, batch, bsize): + """ + Common logic for constructing integrated couplings from dataset. + Used by both ConstantCoupler and TrailingAverageCoupler. + """ + # reset integrated couplings + self.integrated_couplings = np.empty( + (bsize, self.coupled_integration_dim, self.timevar_dim) + self.spatial_dims ) - # we use a constant set of values so we just copy time 0 - for i in range(len(self.preset_coupled_fields)): - self.preset_coupled_fields[i, :, :, :, :, :] = coupled_fields[ - 0, :, -1, :, :, : - ] - # flag for construct integrated coupling method to use this array - self.coupled_mode = True + + index_range = slice( + batch["time"].start, + batch["time"].start + self._coupled_offsets[-1, -1, -1] + 1, + ) + + # extract coupled variables + if self.use_zarr: + # Loading the contiguous time slice into memory and then pulling out the semi-random + # variable indices is quicker than trying to do this all at once. + ds_index_range = self.ds["inputs"][index_range] + ds_index_range = ds_index_range[:, self.ds_variable_indices] + else: + ds_index_range = ( + self.ds.inputs.sel(channel_in=self.variables) + .isel(time=index_range) + .compute() + ) + + return ds_index_range def construct_integrated_couplings( self, @@ -239,48 +271,160 @@ def construct_integrated_couplings( if self.coupled_mode: return self.preset_coupled_fields else: - # reset integrated couplings - self.integrated_couplings = np.empty( - (bsize, self.coupled_integration_dim, self.timevar_dim) - + self.spatial_dims - ) + if (batch is None) or (bsize is None): + raise ValueError( + "batch and bsize must be provided when not in coupled_mode" + ) - index_range = slice( - batch["time"].start, - batch["time"].start + self._coupled_offsets[-1, -1, -1] + 1, + ds_index_range = self._construct_integrated_couplings_from_dataset( + batch, bsize ) - # extract coupled variables and scale lazily - input_array = self.ds.inputs.sel(channel_in=self.variables) - ds = (input_array - self.coupled_scaling["mean"]) / self.coupled_scaling[ - "std" - ] - # load before entering loop for efficiency - ds_index_range = ds.isel(time=index_range).load() + # Apply scaling if available + if self.coupled_scaling is not None: + ds_index_range -= self.coupled_scaling["mean"] + ds_index_range /= self.coupled_scaling["std"] # use static offsets to create integrated coupling array for b in range(bsize): for i in range(self.coupled_integration_dim): - coupling_temp = ds_index_range.isel( - time=self._coupled_offsets[b, i, :] - ) - self.integrated_couplings[b, i, :, :, :] = ( - coupling_temp.to_numpy().reshape( - (self.timevar_dim,) + coupling_temp.shape[2:] - ) + if self.use_zarr: + coupling_temp = ds_index_range[ + self._coupled_offsets[b, i, :], : + ] + else: + coupling_temp = ds_index_range.isel( + time=self._coupled_offsets[b, i, :] + ).to_numpy() + self.integrated_couplings[b, i, :, :, :] = coupling_temp.reshape( + (self.timevar_dim,) + coupling_temp.shape[2:] ) + if self.time_first: + return self.integrated_couplings.transpose((1, 0, 2, 3, 4, 5)).astype( + "float32" + ) # cast to float for compatibility + else: + return self.integrated_couplings.astype("float32") + + +class ConstantCoupler(BaseCoupler): + """ + coupler used to interface two component of earth system + + constant coupler will take the the coupled field at integration time and + force the model with this field consistently + """ + + def __init__( + self, + dataset: xr.Dataset, + batch_size: int, + variables: Sequence, + presteps: int = 0, + input_time_dim: int = 2, + output_time_dim: int = 2, + input_times: Sequence = [pd.Timedelta("24h"), pd.Timedelta("48h")], + prepared_coupled_data=True, + ): + """ + Parameters + ---------- + dataset: xr.Dataset + xarray Dataset that holds coupled data + batch_size: int + number of batch size during training. + forecasting batch size should be 1 + variables: Sequence + sequence of strings that indicate the coupled variable + names in the dataset + presteps: int, optional + the number of model steps used to initialize the hidden state. + If not using a GRU, prestep is 0, default 0 + input_time_dim: int, optional + number of input times into the model, default 2 + output_time_dim: int, optional + number of output times for each model step, default 2 + input_times: Sequence, optional + sequence of pandas Timedelta objects that indicate which times are to be coupled, + default [pd.Timedelta("24h"), pd.Timedelta("48h")] + prepared_coupled_data: boolean, optional + If True assumes data in dataset has been prepared appropriately for training: + averages have already been calculated so that each time step denotes + the right side of a averaging_window window. + This is highly recommended for training, default True + """ + super().__init__( + dataset=dataset, + batch_size=batch_size, + variables=variables, + presteps=presteps, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + input_times=input_times, + prepared_coupled_data=prepared_coupled_data, + ) + + def compute_coupled_indices(self, interval, data_time_step): + """ + Called by CoupledDataset to compute static indices for training + samples + + Parameters + ---------- + interval: int + ratio of dataset timestep to model dt + data_time_step: + dataset timestep + """ + # create array of static coupled offsets that accompany each batch + self._coupled_offsets = np.empty( + [self.batch_size, self.coupled_integration_dim, len(self.input_times)] + ) + for b in range(self.batch_size): + for i in range(self.coupled_integration_dim): + self._coupled_offsets[b, i, :] = b + np.array( + [ts / data_time_step for ts in self.input_times] + ) - return self.integrated_couplings.transpose((1, 0, 2, 3, 4, 5)).astype( - "float32" - ) # cast to float for compatability + self._coupled_offsets = self._coupled_offsets.astype(int) + def set_coupled_fields(self, coupled_fields: th.tensor): + """ + Set the data for the coupled field for the next iteration of the dataloader. + Instead of loading data from the dataset the data from coupled_fields will + be returned instead. -class TrailingAverageCoupler: + Parameters + ---------- + coupled_fields: th.tensor + The data to use when the dataloader requests coupled fields. Expected + format is [B, F, T, C, H, W] + """ + # create buffer for coupling + coupled_fields = coupled_fields[ + :, :, :, self.coupled_channel_indices, :, : + ] + self.preset_coupled_fields = th.empty( + [coupled_fields.shape[0], self.spatial_dims[0], self.coupled_integration_dim, self.timevar_dim] + + list(self.spatial_dims[1:]) + ) + # we use a constant set of values so we just copy time 0 + for i in range(self.coupled_integration_dim): + self.preset_coupled_fields[:, :, i, :, :, :] = coupled_fields[ + :, :, 0, -1:, :, : + ] + if self.time_first: + self.preset_coupled_fields = self.preset_coupled_fields.permute(2, 0, 3, 1, 4, 5) + # flag for construct integrated coupling method to use this array + self.coupled_mode = True + + +class TrailingAverageCoupler(BaseCoupler): """ - coupler used to inferface two components of the earth system + coupler used to interface two components of the earth system Trailing average coupler uses coupled input times as the right side of - an averag that is taken over an "averaging_window" window size. + an average that is taken over an "averaging_window" window size. """ def __init__( @@ -319,55 +463,50 @@ def __init__( sequence of pandas Timedelta objects that indicate which times are to be coupled, default [pd.Timedelta("24h"), pd.Timedelta("48h")] prepared_coupled_data: boolean, optional - If True assumes data in dataset has been prepared approiately for training: + If True assumes data in dataset has been prepared appropriately for training: averages have already been calculated so that each time step denotes the right side of a averaging_window window. - This is highly remcommended for training, default True + This is highly recommended for training, default True """ - # extract important meta data from ds - self.ds = dataset - self.batch_size = batch_size - self.spatial_dims = self.ds.inputs.shape[2:] - self.variables = variables - self.presteps = presteps - self.input_time_dim = input_time_dim - self.output_time_dim = output_time_dim - self.input_times = [pd.Timedelta(t) for t in input_times] + super().__init__( + dataset=dataset, + batch_size=batch_size, + variables=variables, + presteps=presteps, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + input_times=input_times, + prepared_coupled_data=prepared_coupled_data, + ) + + # TrailingAverageCoupler-specific attributes self.averaging_window = pd.Timedelta(averaging_window) - self.output_channels = len(self.variables) * len(self.input_times) - self._set_time_increments() - self.coupled_integration_dim = self._compute_coupled_integration_dim() - self.timevar_dim = self._compute_timevar_dim() - self.coupled_inputs_shape = None - self.scaling_dict = None - self._coupled_offsets = None - self.integrated_couplings = None - self.coupled_mode = False # if forecasting with another coupled model - if not prepared_coupled_data: - logger.log( - logging.DEBUG, - "Assuming coupled data is not preprocessed, averaging fields in as designed in" - "TrailingAverageCoupler. See docs for specifics.", + if self.use_zarr: + cf_dates = cftime.num2pydate( + self.ds["time"][:], + units=self.ds["time"].attrs["units"], + calendar=self.ds["time"].attrs["calendar"], ) - self._prepare_coupled_data() + dates = [np.datetime64(date.isoformat()) for date in cf_dates] + self.time_da = np.asarray(dates) else: - logger.log( - logging.DEBUG, - "**Assuming coupled data has been prepared properly, using coupled field[s] from" - 'dataset "as-is"**', - ) + self.time_da = self.ds.time.values + self._set_time_increments() def compute_coupled_indices(self, interval, data_time_step): """ Called by CoupledDataset to compute static indices for training samples - :param interval: int ratio of dataset timestep to model dt - :param data_time_step: dataset timestep + Parameters + ---------- + interval: int + ratio of dataset timestep to model dt + data_time_step: + dataset timestep """ - - # create array of static coupled offstes that accompany each batch + # create array of static coupled offsets that accompany each batch self._coupled_offsets = np.empty( [self.batch_size, self.coupled_integration_dim, len(self.input_times)] ) @@ -381,24 +520,9 @@ def compute_coupled_indices(self, interval, data_time_step): self._coupled_offsets = self._coupled_offsets.astype(int) - def _prepare_coupled_data(self): - # TODO: write function to lazily compute average as spcified in time scheme - raise NotImplementedError("Data preparation not yet implemented") - - def set_scaling(self, scaling_da): - coupled_scaling = scaling_da.sel(index=self.variables).rename( - {"index": "channel_in"} - ) - self.coupled_scaling = { - "mean": np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)), - "std": np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)), - } - def _set_time_increments(self): # get the dt of the dataset - dt = pd.Timedelta( - self.ds.time[1].values - self.ds.time[0].values - ).total_seconds() + dt = pd.Timedelta(self.time_da[1] - self.time_da[0]).total_seconds() # assert that the time increments are divisible by the dt of the dataset if np.any([t.total_seconds() % dt != 0 for t in self.input_times]): raise ValueError( @@ -407,32 +531,12 @@ def _set_time_increments(self): ) self.time_increments = [t.total_seconds() / dt for t in self.input_times] - def _compute_timevar_dim(self): - return len(self.input_times) * len(self.variables) - - def _compute_coupled_integration_dim(self): - return self.presteps + max(self.output_time_dim // self.input_time_dim, 1) - def setup_coupling(self, coupled_module): - # To expediate the coupling process the coupled_forecast - # get proper channels from coupled component output - output_channels = coupled_module.output_variables - # A bit convoluted. Prepared coupled variables - # are given a suffix for training associated with their - # trailing average increment e.g. 'z1000-48H'. To extract - # thr proper field from the coupled model output, we see if - # we check if the coupled model output var is in self.variables. - # - # for example 'z1000' is in 'z1000-48H' - channel_indices = [ - i - for i, oc in enumerate(output_channels) - for v in self.variables - if oc == v.split("-")[0] - ] - self.coupled_channel_indices = channel_indices + # Call parent method first to set basic coupling + super().setup_coupling(coupled_module) - # find averaging periods from componenet output + # TrailingAverageCoupler-specific setup + # find averaging periods from component output averaging_window_max_indices = [ i // pd.Timedelta(coupled_module.time_step) for i in self.input_times ] @@ -450,12 +554,7 @@ def setup_coupling(self, coupled_module): ) self.averaging_slices = averaging_slices - def reset_coupler(self): - self.coupled_mode = False - self.integrated_couplings = None - self.preset_coupled_fields = None - - def set_coupled_fields(self, coupled_fields): + def set_coupled_fields(self, coupled_fields: th.tensor): """ Set the data for the coupled field for the next iteration of the dataloader. Instead of loading data from the dataset the data from coupled_fields will @@ -467,12 +566,6 @@ def set_coupled_fields(self, coupled_fields): The data to use when the dataloader requests coupled fields. Expected format is [B, F, T, C, H, W] """ - if coupled_fields.shape[0] != self.batch_size: - raise ValueError( - f"Batch size of coupled field {coupled_fields.shape[0]} doesn't " - f" match configured batch size {self.batch_size}" - ) - coupled_fields = coupled_fields[:, :, :, self.coupled_channel_indices, :, :] # TODO: Now support output_time_dim =/= input_time_dim, but presteps need to be 0, will add support for presteps>0 coupled_averaging_periods = [] @@ -484,65 +577,8 @@ def set_coupled_fields(self, coupled_fields): coupled_averaging_periods.append(th.concat(averaging_periods, dim=3)) self.preset_coupled_fields = th.concat( coupled_averaging_periods, dim=2 - ).permute(2, 0, 3, 1, 4, 5) + ) + if self.time_first: + self.preset_coupled_fields = self.preset_coupled_fields.permute(2, 0, 3, 1, 4, 5) # flag for construct integrated coupling method to use this array self.coupled_mode = True - - def construct_integrated_couplings( - self, - batch=None, - bsize=None, - ): - """ - Construct array of coupled inputs that includes values required for - model integration steps. - - :param batch: indices of dataset sample dimension associated with - current batch. - :param bsize: int batch size - """ - - if self.coupled_mode: - return self.preset_coupled_fields - else: - if (batch is None) or (bsize is None): - raise ValueError( - "Coupled fields must be set if no batch or batch size is provided" - ) - # reset integrated couplings - self.integrated_couplings = np.empty( - (bsize, self.coupled_integration_dim, self.timevar_dim) - + self.spatial_dims - ) - - index_range = slice( - batch["time"].start, - batch["time"].start + self._coupled_offsets[-1, -1, -1] + 1, - ) - - # extract coupled variables and scale lazily - ds = ( - self.ds.inputs.sel(channel_in=self.variables) - - self.coupled_scaling["mean"] - ) / self.coupled_scaling["std"] - # load before entering loop for efficiency - ds_index_range = ds.isel(time=index_range).load() - - # use static offsets to create integrated coupling array - for b in range(bsize): - for i in range(self.coupled_integration_dim): - # coupling_temp = \ - # ds.isel(time=batch["time"].start+self._coupled_offsets[b,i,:]) # changed i to 0 for debugging - # added to test speed, original is commented above - coupling_temp = ds_index_range.isel( - time=self._coupled_offsets[b, i, :] - ) - self.integrated_couplings[b, i, :, :, :] = ( - coupling_temp.to_numpy().reshape( - (self.timevar_dim,) + coupling_temp.shape[2:] - ) - ) - - return self.integrated_couplings.transpose((1, 0, 2, 3, 4, 5)).astype( - "float32" - ) # cast to float32 for pytroch compatibility. diff --git a/physicsnemo/datapipes/healpix/data_modules.py b/physicsnemo/datapipes/healpix/data_modules.py index 8d955bf50c..7d8fc1fa77 100644 --- a/physicsnemo/datapipes/healpix/data_modules.py +++ b/physicsnemo/datapipes/healpix/data_modules.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -16,175 +16,33 @@ # System modules import logging -import os -import time from pathlib import Path -from typing import DefaultDict, Optional, Sequence, Union +from typing import Optional, Sequence, Union # numpy import numpy as np # distributed stuff -import torch +import xarray as xr + +import warnings + +# External modules from omegaconf import DictConfig from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler -from physicsnemo.core.version_check import OptionalImport from physicsnemo.distributed import DistributedManager from .coupledtimeseries_dataset import CoupledTimeSeriesDataset from .timeseries_dataset import TimeSeriesDataset -xr = OptionalImport("xarray") - logger = logging.getLogger(__name__) -def _get_file_name(path, prefix, var, suffix): - """ - Helper that returns a fully formed path for a given variable - - Parameters - ---------- - path: str - The base path where the file is located - prefix: str - The prefix used for the filename - var: str - The variable stored in the file - suffix: str - The suffix used for the files - - Returns - ------- - str: The fully formed path - """ - return os.path.join(path, f"{prefix}{var}{suffix}.nc") - - -def open_time_series_dataset_classic_on_the_fly( - directory: str, - input_variables: Sequence, - output_variables: Optional[Sequence], - constants: Optional[DefaultDict] = None, - prefix: Optional[str] = None, - suffix: Optional[str] = None, - batch_size: int = 32, - scaling: Optional[DictConfig] = None, -) -> "xr.Dataset": - """ - Opens and merges multiple datasets that that contain individual variables - into a single dataset - - Parameters - ---------- - directory: str - The directory that contains the input datasets - input_variables: Sequence - The input variables to be merged into the new dataset - output_variables: Sequence, optional - The output variables to be merged into the new dataset - If no output variables are provided the input set is used - constants: DefaultDict, optional - A set of constants to add to the merged dataset - default None - prefix: str, optional - The prefix of the input datasets, default None - suffix: str, optional - The suffix of the input datasets, default None - batch_size: str, optional - The chunk size to use for the input datasets, default 32 - scaling: DictConfig, optional - Not used for open_time_series_dataset_classic_on_the_fly - - Returns - ------- - xr.Dataset: The merged dataset - """ - output_variables = output_variables or input_variables - all_variables = np.union1d(input_variables, output_variables) - prefix = prefix or "" - suffix = suffix or "" - - merge_time = time.time() - logger.info("merging input datasets") - - datasets = [] - remove_attrs = ["mean", "std"] if "LL" in prefix else ["varlev", "mean", "std"] - for variable in all_variables: - file_name = _get_file_name(directory, prefix, variable, suffix) - logger.debug("open nc dataset %s", file_name) - - ds = xr.open_dataset(file_name, autoclose=True) - - if "LL" in prefix: - ds = ds.rename({"lat": "height", "lon": "width"}) - ds = ds.isel({"height": slice(0, 180)}) - try: - ds = ds.isel(varlev=0) - except ValueError: - pass - - # remove unused - for attr in remove_attrs: - if attr in ds.indexes or attr in ds.variables: - ds = ds.drop(attr) - - # Rename variable - if "sample" in ds.variables or "sample" in ds.dims: - ds = ds.rename({"sample": "time"}) - - ds = ds.chunk({"time": batch_size}) - - # Change lat/lon to coordinates - try: - ds = ds.set_coords(["lat", "lon"]) - except (ValueError, KeyError): - pass - datasets.append(ds) - # Merge datasets - data = xr.merge(datasets, compat="override") - - # Convert to input/target array by merging along the variables - input_da = ( - data[list(input_variables)] - .to_array("channel_in", name="inputs") - .transpose("time", "channel_in", "face", "height", "width") - ) - target_da = ( - data[list(output_variables)] - .to_array("channel_out", name="targets") - .transpose("time", "channel_out", "face", "height", "width") - ) - - result = xr.Dataset() - result["inputs"] = input_da - result["targets"] = target_da - - # Get constants - if constants is not None: - constants_ds = [] - for name, var in constants.items(): - constants_ds.append( - xr.open_dataset( - _get_file_name(directory, prefix, name, suffix), autoclose=True - ).set_coords(["lat", "lon"])[var] - ) - constants_ds = xr.merge(constants_ds, compat="override") - constants_da = constants_ds.to_array("channel_c", name="constants").transpose( - "channel_c", "face", "height", "width" - ) - result["constants"] = constants_da - - logger.info("merged datasets in %0.1f s", time.time() - merge_time) - - return result - - def open_time_series_dataset_classic_prebuilt( directory: str, dataset_name: str, constants: bool = False, batch_size: int = 32 -) -> "xr.Dataset": +) -> xr.Dataset: """ Opens an existing dataset @@ -194,7 +52,7 @@ def open_time_series_dataset_classic_prebuilt( The directory that contains the dataset dataset_name: str The name of the dataset to open - constants: DefaultDict, optional + constants: typing.DefaultDict, optional Not used for open_time_series_dataset_classic_prebuilt, default False batch_size: str, optional The chunk size to use for the input datasets, default 32 @@ -202,164 +60,21 @@ def open_time_series_dataset_classic_prebuilt( Returns ------- xr.Dataset: The opened dataset + + Raises + ------ + FileNotFoundError: If the dataset doesn't exist at the specified path """ ds_path = Path(directory, dataset_name + ".zarr") if not ds_path.exists(): - raise FileNotFoundError(f"Dataset doesn't appear to exist at {ds_path}") + raise FileNotFoundError(f"Dataset doesn't exist at {ds_path}") result = xr.open_zarr(ds_path) return result -def create_time_series_dataset_classic( - src_directory: str, - dst_directory: str, - dataset_name: str, - input_variables: Sequence, - output_variables: Optional[Sequence] = None, - constants: Optional[DefaultDict] = None, - prefix: Optional[str] = None, - suffix: Optional[str] = None, - batch_size: int = 32, - scaling: Optional[DictConfig] = None, - overwrite: bool = False, -) -> "xr.Dataset": - """ - Opens and merges multiple datasets that that contain individual variables - into a single dataset - - Parameters - ---------- - src_directory: str - The directory that contains the input datasets - dst_directory: str - dataset_name: str - input_variables: Sequence - The input variables to be merged into the new dataset - output_variables: Sequence, optional - The output variables to be merged into the new dataset - If no output variables are provided the input set is used - constants: DefaultDict, optional - A set of constants to add to the merged dataset, default None - prefix: str, optional - The prefix of the input datasets, default None - suffix: str, optional - The suffix of the input datasets, default None - batch_size: str, optional - The chunk size to use for the input datasets, default 32 - scaling: DictConfig, optional - Scale factors applied to the listed variables, default None - overwrite: bool, optional - IF an existing dataset exists at the destination replace it, default False - - Returns - ------- - xr.Dataset: The merged dataset - """ - dst_zarr = os.path.join(dst_directory, dataset_name + ".zarr") - file_exists = os.path.exists(dst_zarr) - - if file_exists and not overwrite: - logger.info("opening input datasets") - return open_time_series_dataset_classic_prebuilt( - directory=dst_directory, - dataset_name=dataset_name, - constants=constants is not None, - ) - - output_variables = output_variables or input_variables - all_variables = np.union1d(input_variables, output_variables) - prefix = prefix or "" - suffix = suffix or "" - - merge_time = time.time() - logger.info("merging input datasets") - - datasets = [] - remove_attrs = ["varlev", "mean", "std"] - for variable in all_variables: - file_name = _get_file_name(src_directory, prefix, variable, suffix) - logger.debug("open nc dataset %s", file_name) - if "sample" in list(xr.open_dataset(file_name).sizes.keys()): - ds = xr.open_dataset(file_name).rename({"sample": "time"}) - else: - ds = xr.open_dataset(file_name) - if "varlev" in ds.dims: - ds = ds.isel(varlev=0) - - for attr in remove_attrs: - if attr in ds.indexes or attr in ds.variables: - ds = ds.drop(attr) - - # Rename variable - if "predictors" in list(ds.keys()): - ds = ds.rename({"predictors": variable}) - - # Change lat/lon to coordinates - try: - ds = ds.set_coords(["lat", "lon"]) - except (ValueError, KeyError): - pass - # Apply log scaling lazily - if ( - scaling - and variable in scaling - and scaling[variable].get("log_epsilon", None) is not None - ): - ds[variable] = np.log( - ds[variable] + scaling[variable]["log_epsilon"] - ) - np.log(scaling[variable]["log_epsilon"]) - datasets.append(ds) - # Merge datasets - data = xr.merge(datasets, compat="override") - - # Convert to input/target array by merging along the variables - input_da = ( - data[list(input_variables)] - .to_array("channel_in", name="inputs") - .transpose("time", "channel_in", "face", "height", "width") - ) - target_da = ( - data[list(output_variables)] - .to_array("channel_out", name="targets") - .transpose("time", "channel_out", "face", "height", "width") - ) - - result = xr.Dataset() - result["inputs"] = input_da - result["targets"] = target_da - - # Get constants - if constants is not None: - constants_ds = [] - for name, var in constants.items(): - constants_ds.append( - xr.open_dataset(_get_file_name(src_directory, prefix, name, suffix)) - .set_coords(["lat", "lon"])[var] - .astype(np.float32) - ) - constants_ds = xr.merge(constants_ds, compat="override") - constants_da = constants_ds.to_array("channel_c", name="constants").transpose( - "channel_c", "face", "height", "width" - ) - result["constants"] = constants_da - - logger.info("merged datasets in %0.1f s", time.time() - merge_time) - logger.info("writing unified dataset to file (takes long!)") - - # writing out - def _write_zarr(data, path): - write_job = data.to_zarr(path, compute=False, mode="w") - logger.info(f"writing dataset to {path}") - write_job.compute() - - _write_zarr(data=result, path=dst_zarr) - - return result - - class TimeSeriesDataModule: """pytorch-lightning module for complete model train, validation, and test data loading. Uses dlwp.data.data_loading.TimeSeriesDataset under-the-hood. Loaded data files follow the naming scheme @@ -368,12 +83,12 @@ class TimeSeriesDataModule: def __init__( self, - src_directory: str = ".", + src_directory: str = ".", # kept for backwards compatibility dst_directory: str = ".", dataset_name: str = "dataset", prefix: Optional[str] = None, suffix: Optional[str] = None, - data_format: str = "classic", + data_format: str = "classic", # kept for backwards compatibility batch_size: int = 32, drop_last: bool = False, input_variables: Optional[Sequence] = ["t2m"], @@ -389,17 +104,15 @@ def __init__( gap: Union[int, str, None] = None, shuffle: bool = True, add_insolation: bool = False, - cube_dim: int = 64, + cube_dim: int = 64, # kept for backwards compatibility num_workers: int = 4, pin_memory: bool = True, - prebuilt_dataset: bool = True, + prebuilt_dataset: bool = True, # kept for backwards compatibility forecast_init_times: Optional[Sequence] = None, ): """ Parameters ---------- - src_directory: str, optional - The directory containing data files per variable, default "." dst_directory: str, optional The directory containing joint data files, default "." dataset_name: str, optional @@ -408,11 +121,6 @@ def __init__( Prefix appended to all data files, default None suffix: str, optional Suffix appended to all data files, default None - data_format: str, optional - str indicating data schema. - 'classic': use classic DLWP file types. Loads .nc files, assuming - dimensions [sample, varlev, face, height, width] and - data variables 'predictors', 'lat', and 'lon'. batch_size: int, optional Size of batches to draw from data, defualt 32 drop_last: bool, optional @@ -449,14 +157,10 @@ def __init__( Option to shuffle the training data, default True add_insolation: bool, optional Option to add prescribed insolation as a decoder input feature, default True - cube_dim: int, optional - Number of points on the side of a cube face. Not currently used. num_workers: int, optional Number of parallel data loading workers, default 4 pin_memory: bool, optional Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True - prebuilt_dataset: bool, optional - Create a custom dataset for training. If False, the variables are gathered on the fly, default True forecast_init_times: Sequence, optional A Sequence of pandas Timestamps dictating the specific initialization times to produce inputs for. default None @@ -466,12 +170,17 @@ def __init__( NOT produce any target array. """ super().__init__() - self.src_directory = src_directory + + warnings.warn( + "TimeSeriesDataModule will be removed in a future release, please switch to using TimeSeriesDataModuleZarr", + DeprecationWarning, + stacklevel=2 + ) + self.dst_directory = dst_directory self.dataset_name = dataset_name self.prefix = prefix self.suffix = suffix - self.data_format = data_format self.batch_size = batch_size self.drop_last = drop_last self.input_variables = input_variables @@ -489,7 +198,6 @@ def __init__( self.cube_dim = cube_dim self.num_workers = num_workers self.pin_memory = pin_memory - self.prebuilt_dataset = prebuilt_dataset self.forecast_init_times = forecast_init_times self.train_dataset = None @@ -520,93 +228,45 @@ def get_constants(self) -> Optional[np.ndarray]: def setup(self) -> None: """Setup the datasets used for this DataModule""" - if self.data_format == "classic": - create_fn = create_time_series_dataset_classic - open_fn = ( - open_time_series_dataset_classic_prebuilt - if self.prebuilt_dataset - else open_time_series_dataset_classic_on_the_fly - ) - else: - raise ValueError("'data_format' must be one of ['classic']") - # make sure distributed manager is initalized if not DistributedManager.is_initialized(): DistributedManager.initialize() - dist = DistributedManager() - - if torch.distributed.is_initialized(): - if self.prebuilt_dataset: - if dist.rank == 0: - create_fn( - src_directory=self.src_directory, - dst_directory=self.dst_directory, - dataset_name=self.dataset_name, - input_variables=self.input_variables, - output_variables=self.output_variables, - constants=self.constants, - prefix=self.prefix, - suffix=self.suffix, - batch_size=self.dataset_batch_size, - scaling=self.scaling, - overwrite=False, - ) - # wait for rank 0 to complete, because then the files are guaranteed to exist - torch.distributed.barrier() + dataset = open_time_series_dataset_classic_prebuilt( + directory=self.dst_directory, + dataset_name=self.dataset_name, + constants=self.constants is not None, + batch_size=self.batch_size, + ) - dataset = open_fn( - directory=self.dst_directory, - dataset_name=self.dataset_name, - constants=self.constants is not None, - batch_size=self.batch_size, - ) - else: - dataset = open_fn( - input_variables=self.input_variables, - output_variables=self.output_variables, - directory=self.dst_directory, - constants=self.constants, - prefix=self.prefix, - batch_size=self.batch_size, - ) - else: - if self.prebuilt_dataset: - create_fn( - src_directory=self.src_directory, - dst_directory=self.dst_directory, - dataset_name=self.dataset_name, - input_variables=self.input_variables, - output_variables=self.output_variables, - constants=self.constants, - prefix=self.prefix, - suffix=self.suffix, - batch_size=self.dataset_batch_size, - scaling=self.scaling, - overwrite=False, - ) + # Verify that all input variables are available in the dataset + missing_input_vars = set(self.input_variables) - set(dataset.channel_in.values) + if missing_input_vars: + raise ValueError( + f"Input variables not found in dataset: {missing_input_vars}" + ) - dataset = open_fn( - directory=self.dst_directory, - dataset_name=self.dataset_name, - constants=self.constants is not None, - batch_size=self.batch_size, - ) - else: - dataset = open_fn( - input_variables=self.input_variables, - output_variables=self.output_variables, - directory=self.dst_directory, - constants=self.constants, - prefix=self.prefix, - batch_size=self.batch_size, - ) + # Verify that all output variables are available in the dataset + missing_output_vars = set(self.output_variables) - set( + dataset.channel_out.values + ) + if missing_output_vars: + raise ValueError( + f"Output variables not found in dataset: {missing_output_vars}" + ) dataset = dataset.sel( channel_in=self.input_variables, channel_out=self.output_variables, ) if self.constants is not None: + # Verify that all constants are available in the dataset + missing_constants = set(list(self.constants.values())) - set( + dataset.channel_c.values + ) + if missing_constants: + raise ValueError(f"Constants not found in dataset: {missing_constants}") + dataset = dataset.sel(channel_c=list(self.constants.values())) if self.splits is not None and self.forecast_init_times is None: @@ -801,12 +461,12 @@ class CoupledTimeSeriesDataModule(TimeSeriesDataModule): def __init__( self, - src_directory: str = ".", + src_directory: str = ".", # kept for backwards compatibility dst_directory: str = ".", dataset_name: str = "dataset", prefix: Optional[str] = None, suffix: Optional[str] = None, - data_format: str = "classic", + data_format: str = "classic", # kept for backwards compatibility batch_size: int = 32, drop_last: bool = False, input_variables: Optional[Sequence] = None, @@ -822,10 +482,10 @@ def __init__( gap: Union[int, str, None] = None, shuffle: bool = True, add_insolation: bool = False, - cube_dim: int = 64, + cube_dim: int = 64, # kept for backwards compatibility num_workers: int = 4, pin_memory: bool = True, - prebuilt_dataset: bool = True, + prebuilt_dataset: bool = True, # kept for backwards compatibility forecast_init_times: Optional[Sequence] = None, couplings: Sequence = None, add_train_noise: Optional[bool] = False, @@ -835,8 +495,6 @@ def __init__( """ Parameters ---------- - src_directory: str, optional - The directory containing data files per variable, default "." dst_directory: str, optional The directory containing joint data files, default "." dataset_name: str, optional @@ -845,11 +503,6 @@ def __init__( Prefix appended to all data files, default None suffix: str, optional Suffix appended to all data files, default None - data_format: str, optional - str indicating data schema. - 'classic': use classic DLWP file types. Loads .nc files, assuming - dimensions [sample, varlev, face, height, width] and - data variables 'predictors', 'lat', and 'lon'. batch_size: int, optional Size of batches to draw from data, defualt 32 drop_last: bool, optional @@ -892,8 +545,6 @@ def __init__( Number of parallel data loading workers, default 4 pin_memory: bool, optional Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True - prebuilt_dataset: bool, optional - Create a custom dataset for training. If False, the variables are gathered on the fly, default True forecast_init_times: Sequence, optional A Sequence of pandas Timestamps dictating the specific initialization times to produce inputs for. default None @@ -946,6 +597,7 @@ def __init__( ) def _get_coupled_vars(self): + coupled_variables = [] for d in self.couplings: coupled_variables = coupled_variables + d["params"]["variables"] @@ -953,94 +605,46 @@ def _get_coupled_vars(self): def setup(self) -> None: """Setup the datasets used for this DataModule""" - if self.data_format == "classic": - create_fn = create_time_series_dataset_classic - open_fn = ( - open_time_series_dataset_classic_prebuilt - if self.prebuilt_dataset - else open_time_series_dataset_classic_on_the_fly - ) - else: - raise ValueError("'data_format' must be one of ['classic', 'zarr']") - coupled_variables = self._get_coupled_vars() # make sure distributed manager is initalized if not DistributedManager.is_initialized(): DistributedManager.initialize() dist = DistributedManager() - if torch.distributed.is_initialized(): - if self.prebuilt_dataset: - if dist.rank == 0: - create_fn( - src_directory=self.src_directory, - dst_directory=self.dst_directory, - dataset_name=self.dataset_name, - input_variables=self.input_variables + coupled_variables, - output_variables=self.output_variables, - constants=self.constants, - prefix=self.prefix, - suffix=self.suffix, - batch_size=self.dataset_batch_size, - scaling=self.scaling, - overwrite=False, - ) - - # wait for rank 0 to complete, because then the files are guaranteed to exist - torch.distributed.barrier() + dataset = open_time_series_dataset_classic_prebuilt( + directory=self.dst_directory, + dataset_name=self.dataset_name, + constants=self.constants is not None, + batch_size=self.batch_size, + ) - dataset = open_fn( - directory=self.dst_directory, - dataset_name=self.dataset_name, - constants=self.constants is not None, - batch_size=self.batch_size, - ) - else: - dataset = open_fn( - input_variables=self.input_variables + coupled_variables, - output_variables=self.output_variables, - directory=self.dst_directory, - constants=self.constants, - prefix=self.prefix, - batch_size=self.batch_size, - ) - else: - if self.prebuilt_dataset: - create_fn( - src_directory=self.src_directory, - dst_directory=self.dst_directory, - dataset_name=self.dataset_name, - input_variables=self.input_variables + coupled_variables, - output_variables=self.output_variables, - constants=self.constants, - prefix=self.prefix, - suffix=self.suffix, - batch_size=self.dataset_batch_size, - scaling=self.scaling, - overwrite=False, - ) + # Verify that all input variables are available in the dataset + missing_input_vars = set(self.input_variables) - set(dataset.channel_in.values) + if missing_input_vars: + raise ValueError( + f"Input variables not found in dataset: {missing_input_vars}" + ) - dataset = open_fn( - directory=self.dst_directory, - dataset_name=self.dataset_name, - constants=self.constants is not None, - batch_size=self.batch_size, - ) - else: - dataset = open_fn( - input_variables=self.input_variables + coupled_variables, - output_variables=self.output_variables, - directory=self.dst_directory, - constants=self.constants, - prefix=self.prefix, - batch_size=self.batch_size, - ) + # Verify that all output variables are available in the dataset + missing_output_vars = set(self.output_variables) - set( + dataset.channel_out.values + ) + if missing_output_vars: + raise ValueError( + f"Output variables not found in dataset: {missing_output_vars}" + ) dataset = dataset.sel( channel_in=self.input_variables + coupled_variables, channel_out=self.output_variables, ) if self.constants is not None: + # Verify that all constants are available in the dataset + missing_constants = set(list(self.constants.values())) - set( + dataset.channel_c.values + ) + if missing_constants: + raise ValueError(f"Constants not found in dataset: {missing_constants}") dataset = dataset.sel(channel_c=list(self.constants.values())) if self.splits is not None and self.forecast_init_times is None: diff --git a/physicsnemo/datapipes/healpix/data_modules_zarr.py b/physicsnemo/datapipes/healpix/data_modules_zarr.py new file mode 100644 index 0000000000..ad830188e0 --- /dev/null +++ b/physicsnemo/datapipes/healpix/data_modules_zarr.py @@ -0,0 +1,549 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# System modules +import logging +import warnings +from pathlib import Path +from typing import Optional, Sequence, Union + +# numpy +import numpy as np + +# External modules +from omegaconf import DictConfig +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from physicsnemo.distributed import DistributedManager + +from .coupledtimeseries_dataset_zarr import CoupledTimeSeriesDatasetZarr +from .timeseries_dataset_zarr import TimeSeriesDatasetZarr + +logger = logging.getLogger(__name__) + + +class TimeSeriesDataModuleZarr: + """Module for complete model train, validation, and test data loading. Uses + dlwp.data.data_loading.TimeSeriesDataset under-the-hood. + """ + + def __init__( + self, + dataset_path: str, + batch_size: int = 32, + dataloader_batch_size: Optional[int] = None, + drop_last: bool = True, + input_variables: Optional[Sequence] = ["t2m"], + output_variables: Optional[Sequence] = None, + constant_variables: Optional[Sequence] = None, + scaling: Optional[DictConfig] = None, + splits: Optional[DictConfig] = None, + presteps: int = 0, + input_time_dim: int = 1, + output_time_dim: int = 1, + data_time_step: Union[int, str] = "3h", + time_step: Union[int, str] = "6h", + gap: Union[int, str, None] = None, + shuffle: bool = True, + add_insolation: bool = False, + num_workers: int = 4, + pin_memory: bool = True, + forecast_init_times: Optional[Sequence] = None, + add_train_noise: Optional[bool] = False, + train_noise_params: Optional[DictConfig] = None, + train_noise_seed: Optional[int] = 42, + ): + """ + Parameters + ---------- + dataset_path: str, optional + The path to the dataset, default "." + batch_size: int, optional + The number of sequential samples to load from the dataset to load, default 32 + dataloader_batch_size: int, optional + Passed to nn.DataLoader as batch_size. Used to assemble batches of samples from the dataloader + The total number of samples will be batch_size*dataloader_batch_size, default None + drop_last: bool, optional + Whether to drop the last batch if it is smaller than batch_size, it is + recommended to set this to true to avoid issues with mismatched sizes, default True + input_variables: Sequence, optional + List of input variable names, to be found in data file name, default "t2m" + output_variables: Sequence, optional + List of output variables names. If None, defaults to `input_variables`. default None + constant_variables: Sequence, optional + List of constant variables names. default None + scaling: DictConfig, optional + Dictionary containing scaling parameters for data variables, default None + splits: DictConfig + Dictionary with train/validation/test set start/end dates. + presteps: int, optional + Number of time steps to initialize recurrent hidden states. default 0 + input_time_dim: int, optional + Number of time steps in the input array, default 1 + output_time_dim: int, optional + Number of time steps in the target/ground truth array, default 1 + data_time_step: Union[int, str], optional + Either integer hours or a str interpretable by pandas: time between steps in the + original data time series, default "3h" + time_step: Union[int, str], optional + Either integer hours or a str interpretable by pandas: desired time between effective model + time steps, default "6h" + gap: Union[int, str], optional + either integer hours or a str interpretable by pandas: time step between the last input time and + the first output time. Defaults to `time_step`. + shuffle: bool, optional + Whether to shuffle the training data, default True + add_insolation: bool, optional + Whether to add prescribed insolation as a decoder input feature, default False + num_workers: int, optional + Number of parallel data loading workers, default 4 + pin_memory: bool, optional + Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True + forecast_init_times: Sequence, optional + A Sequence of pandas Timestamps dictating the specific initialization times + to produce inputs for. default None + Note: + - this is only applied to the test dataloader + - providing this parameter configures the data loader to only produce this number of samples, and + NOT produce any target array. + add_train_noise: bool, optional + Wether to add noise to the training data to inputs and integrated couplings to improve generalization, default False + train_noise_params: DictConfig, optional + Dictionary containing parameters for adding noise to the training data + train_noise_seed: int, optional + Seed for the random number generator for adding noise to the training data, default 42 + """ + super().__init__() + self.dataset_path = Path(dataset_path) + self.dataset_batch_size = batch_size + self.dataloader_batch_size = dataloader_batch_size + self.drop_last = drop_last + self.input_variables = input_variables + self.output_variables = output_variables or input_variables + self.constant_variables = constant_variables + self.constants = constant_variables + self.scaling = scaling + self.splits = splits + self.input_time_dim = input_time_dim + (presteps * input_time_dim) + self.output_time_dim = output_time_dim + self.data_time_step = data_time_step + self.time_step = time_step + self.gap = gap + self.shuffle = shuffle + self.add_insolation = add_insolation + self.num_workers = num_workers + self.pin_memory = pin_memory + self.forecast_init_times = forecast_init_times + self.add_train_noise = add_train_noise + self.train_noise_params = train_noise_params + self.train_noise_seed = train_noise_seed + + self.train_dataset = None + self.val_dataset = None + self.test_dataset = None + + self.collate_fn = None + + self.setup() + + def get_constants(self) -> Optional[np.ndarray]: + """Returns the constants used in this dataset + + Returns + ------- + np.ndarray: The list of constants, None if there are no constants + """ + if self.constant_variables is None: + return None + + return ( + self.train_dataset.get_constants() + if self.train_dataset is not None + else self.test_dataset.get_constants() + ) + + def _get_common_dataset_kwargs(self) -> dict: + """Get common keyword arguments for dataset creation""" + return { + "dataset_path": self.dataset_path, + "input_variables": self.input_variables, + "output_variables": self.output_variables, + "constant_variables": self.constant_variables, + "batch_size": self.dataset_batch_size, + "add_insolation": self.add_insolation, + "input_time_dim": self.input_time_dim, + "output_time_dim": self.output_time_dim, + "data_time_step": self.data_time_step, + "time_step": self.time_step, + "gap": self.gap, + "scaling": self.scaling, + } + + def _validate_setup_requirements(self) -> None: + """Validate that required setup conditions are met""" + if not self.dataset_path.exists(): + raise FileNotFoundError(f"Dataset path not found: {self.dataset_path}") + + if self.splits is None and self.forecast_init_times is None: + raise ValueError("Either splits or forecast_init_times must be provided") + + # sanity check if dates overlap + if self.splits: + train_test_overlap = ( + np.datetime64(self.splits["train_date_end"]) + >= np.datetime64(self.splits["test_date_start"]) + ) and ( + np.datetime64(self.splits["train_date_start"]) + <= np.datetime64(self.splits["test_date_end"]) + ) + if train_test_overlap: + warnings.warn("Training and test date ranges overlap") + + train_val_overlap = ( + np.datetime64(self.splits["train_date_end"]) + >= np.datetime64(self.splits["val_date_start"]) + ) and ( + np.datetime64(self.splits["train_date_start"]) + <= np.datetime64(self.splits["val_date_end"]) + ) + if train_val_overlap: + warnings.warn("Training and validation date ranges overlap") + + test_val_overlap = ( + np.datetime64(self.splits["test_date_end"]) + >= np.datetime64(self.splits["val_date_start"]) + ) and ( + np.datetime64(self.splits["test_date_start"]) + <= np.datetime64(self.splits["val_date_end"]) + ) + if test_val_overlap: + warnings.warn("Test and validation date ranges overlap") + + def _initialize_distributed_manager(self): + """Initialize distributed manager if not already initialized""" + if not DistributedManager.is_initialized(): + DistributedManager.initialize() + return DistributedManager() + + def _get_dataset_class(self): + """Get the dataset class to use for creating datasets""" + return TimeSeriesDatasetZarr + + def _create_datasets(self, common_kwargs: dict, dataset_class, dist) -> None: + """Create train, validation, and test datasets based on configuration""" + # forecast_init_times are provided just create the test dataset + if self.forecast_init_times is not None: + self.test_dataset = dataset_class( + drop_last=False, + **common_kwargs, + forecast_init_times=self.forecast_init_times, + ) + # splits are provided, use them to split the dataset + else: + self.train_dataset = dataset_class( + start_date=self.splits["train_date_start"], + end_date=self.splits["train_date_end"], + drop_last=self.drop_last, + add_train_noise=self.add_train_noise, + train_noise_params=self.train_noise_params, + train_noise_seed=self.train_noise_seed + int(dist.rank), + **common_kwargs, + ) + self.val_dataset = dataset_class( + start_date=self.splits["val_date_start"], + end_date=self.splits["val_date_end"], + drop_last=self.drop_last, + **common_kwargs, + ) + self.test_dataset = dataset_class( + start_date=self.splits["test_date_start"], + end_date=self.splits["test_date_end"], + drop_last=False, + **common_kwargs, + ) + + def setup(self) -> None: + """Setup the datasets used for this DataModule""" + self._validate_setup_requirements() + dist = self._initialize_distributed_manager() + + common_kwargs = self._get_common_dataset_kwargs() + dataset_class = self._get_dataset_class() + + self._create_datasets(common_kwargs, dataset_class, dist) + + def _base_dataloader( + self, dataset, num_shards=1, shard_id=0, shuffle=False, drop_last=False + ) -> DataLoader: + """Setup a dataloader with common functionality + + Parameters + ---------- + dataset: Dataset + The dataset to create the dataloader for + num_shards: int, optional + The total total number of distributed shards + default is 1 meaning distributed training is not being used + shard_id: int, optional + The shard number of this instance of the dataloader, default 0 + shuffle: bool, optional + Whether to shuffle the data, default False + + Returns + ------- + DataLoader: The configured dataloader + """ + sampler = None + drop_last = False + if num_shards > 1: + sampler = DistributedSampler( + dataset, + num_replicas=num_shards, + rank=shard_id, + shuffle=shuffle, + drop_last=drop_last, + ) + shuffle = False + drop_last = False + + loader = DataLoader( + dataset=dataset, + pin_memory=self.pin_memory, + num_workers=self.num_workers, + shuffle=shuffle, + drop_last=drop_last, + sampler=sampler, + collate_fn=self.collate_fn, + batch_size=self.dataloader_batch_size, + ) + + return loader, sampler + + def train_dataloader(self, num_shards=1, shard_id=0) -> DataLoader: + """Setup the training dataloader + + Parameters + ---------- + num_shards: int, optional + The total total number of distributed shards + default is 1 meaning distributed training is not being used + shard_id: int, optional + The shard number of this instance of the dataloader, default 0 + + Returns + ------- + DataLoader: The training dataloader + """ + return self._base_dataloader( + dataset=self.train_dataset, + num_shards=num_shards, + shard_id=shard_id, + shuffle=self.shuffle, + drop_last=True, + ) + + def val_dataloader(self, num_shards=1, shard_id=0) -> DataLoader: + """Setup the validation dataloader + + Parameters + ---------- + num_shards: int, optional + The total total number of distributed shards + default is 1 meaning distributed validation is not being used + shard_id: int, optional + The shard number of this instance of the dataloader, default 0 + + Returns + ------- + DataLoader: The validation dataloader + """ + return self._base_dataloader( + dataset=self.val_dataset, + num_shards=num_shards, + shard_id=shard_id, + shuffle=False, + drop_last=True, + ) + + def test_dataloader(self, num_shards=1, shard_id=0) -> DataLoader: + """Setup the test dataloader + + Parameters + ---------- + num_shards: int, optional + The total total number of distributed shards + default is 1 meaning distributed test is not being used + shard_id: int, optional + The shard number of this instance of the dataloader, default 0 + + Returns + ------- + DataLoader: The test dataloader + """ + return self._base_dataloader( + dataset=self.test_dataset, + num_shards=num_shards, + shard_id=shard_id, + shuffle=False, + drop_last=True, + ) + + +class CoupledTimeSeriesDataModuleZarr(TimeSeriesDataModuleZarr): + """ + Extension of TimeSeriesDataModule, designed for coupled models that take input from other + earth system components. + """ + + def __init__( + self, + dataset_path: str, + batch_size: int = 32, + dataloader_batch_size: Optional[int] = None, + drop_last: bool = True, + input_variables: Optional[Sequence] = None, + output_variables: Optional[Sequence] = None, + constant_variables: Optional[Sequence] = None, + scaling: Optional[DictConfig] = None, + splits: Optional[DictConfig] = None, + presteps: int = 0, + input_time_dim: int = 1, + output_time_dim: int = 1, + data_time_step: Union[int, str] = "3h", + time_step: Union[int, str] = "6h", + gap: Union[int, str, None] = None, + shuffle: bool = True, + add_insolation: bool = False, + num_workers: int = 4, + pin_memory: bool = True, + forecast_init_times: Optional[Sequence] = None, + couplings: Sequence = None, + add_train_noise: Optional[bool] = False, + train_noise_params: Optional[DictConfig] = None, + train_noise_seed: Optional[int] = 42, + ): + """ + Parameters + ---------- + dataset_path: str, optional + The path to the dataset, default "." + batch_size: int, optional + The number of sequential samples to load from the dataset to load, default 32 + dataloader_batch_size: int, optional + Passed to nn.DataLoader as batch_size. Used to assemble batches of samples from the dataloader + The total number of samples will be dataloader_batch_size*dataloader_batch_size, default None + drop_last: bool, optional + Whether to drop the last batch if it is smaller than batch_size, it is + recommended to set this to true to avoid issues with mismatched sizes, default True + input_variables: Sequence, optional + List of input variable names, to be found in data file name, default None + output_variables: Sequence, optional + List of output variables names. If None, defaults to `input_variables`. default None + constant_variables: Sequence, optional + List of constant variables names. default None + scaling: DictConfig, optional + Dictionary containing scaling parameters for data variables, default None + splits: DictConfig + Dictionary with train/validation/test set start/end dates. + presteps: int, optional + Number of time steps to initialize recurrent hidden states. default 0 + input_time_dim: int, optional + Number of time steps in the input array, default 1 + output_time_dim: int, optional + Number of time steps in the target/ground truth array, default 1 + data_time_step: Union[int, str], optional + Either integer hours or a str interpretable by pandas: time between steps in the + original data time series, default "3h" + time_step: Union[int, str], optional + Either integer hours or a str interpretable by pandas: desired time between effective model + time steps, default "6h" + gap: Union[int, str, None], optional + either integer hours or a str interpretable by pandas: time step between the last input time and + the first output time. default None. + shuffle: bool, optional + Whether to shuffle the training data, default True + add_insolation: bool, optional + Whether to add prescribed insolation as a decoder input feature, default False + num_workers: int, optional + Number of parallel data loading workers, default 4 + pin_memory: bool, optional + Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True + forecast_init_times: Sequence, optional + A Sequence of pandas Timestamps dictating the specific initialization times + to produce inputs for. default None + Note: + - this is only applied to the test dataloader + - providing this parameter configures the data loader to only produce this number of samples, and + NOT produce any target array. + couplings: Sequence, optional + a Sequence of dictionaries that define the mechanics of couplings with other earth system + components. default None + add_train_noise: bool, optional + Wether to add noise to the training data to inputs and integrated couplings to improve generalization, default False + train_noise_params: DictConfig, optional + Dictionary containing parameters for adding noise to the training data + train_noise_seed: int, optional + Seed for the random number generator for adding noise to the training data, default 42 + """ + self.couplings = couplings + + super().__init__( + dataset_path, + batch_size, + dataloader_batch_size, + drop_last, + input_variables, + output_variables, + constant_variables, + scaling, + splits, + presteps, + input_time_dim, + output_time_dim, + data_time_step, + time_step, + gap, + shuffle, + add_insolation, + num_workers, + pin_memory, + forecast_init_times, + add_train_noise, + train_noise_params, + train_noise_seed, + ) + + def _get_coupled_vars(self): + """Get the coupled variables from the couplings""" + coupled_variables = [] + for d in self.couplings: + coupled_variables = coupled_variables + d["params"]["variables"] + return coupled_variables + + def _get_common_dataset_kwargs(self) -> dict: + """Get common keyword arguments for dataset creation (includes coupling-specific params)""" + base_kwargs = super()._get_common_dataset_kwargs() + base_kwargs.update( + { + "couplings": self.couplings, + } + ) + return base_kwargs + + def _get_dataset_class(self): + """Get the dataset class to use for creating datasets""" + return CoupledTimeSeriesDatasetZarr diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index 9e2fbad3ef..33d30e6e47 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -14,8 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations - import logging import time import warnings @@ -25,16 +23,14 @@ import numpy as np import pandas as pd import torch +import xarray as xr from omegaconf import DictConfig, OmegaConf from torch.utils.data import Dataset -from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.datapipe import Datapipe from physicsnemo.datapipes.meta import DatapipeMetaData from physicsnemo.utils.insolation import insolation -xr = OptionalImport("xarray") - logger = logging.getLogger(__name__) @@ -245,12 +241,17 @@ def _get_scaling_da(self): # REMARK: we remove the xarray overhead from these try: - # we use channel_out instead of channel_in because - # the list of input channels may contain data fetched outside - # the datasets such as coupled fields - self.input_scaling = scaling_da.sel( - index=self.ds.channel_out.values - ).rename({"index": "channel_in"}) + # 'if' statement used for cases where atmos model + # includes diagnostic variables like tp6 and msl. + # using 'channel_out' is still necessary for ocean models. + if len(self.ds.channel_out) != (len(self.ds.channel_in)-len(self.couplings[0].variables)): + self.input_scaling = scaling_da.sel( + index=self.ds.channel_in.values + ).rename({"index": "channel_in"}) + else: + self.input_scaling = scaling_da.sel( + index=self.ds.channel_out.values + ).rename({"index": "channel_in"}) self.input_scaling = { "mean": np.expand_dims( self.input_scaling["mean"].to_numpy(), (0, 2, 3, 4) diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py new file mode 100644 index 0000000000..adacbc9678 --- /dev/null +++ b/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple, Union + +import numpy as np +import torch +from omegaconf import DictConfig + +from physicsnemo.datapipes.meta import DatapipeMetaData +from physicsnemo.utils.insolation import insolation + +from .base_timeseries_dataset_zarr import BaseTimeSeriesDatasetZarr + +logger = logging.getLogger(__name__) + + +@dataclass +class MetaData(DatapipeMetaData): + """Metadata for this datapipe""" + + name: str = "TimeSeries" + # Optimization + auto_device: bool = False + cuda_graphs: bool = False + # Parallel + ddp_sharding: bool = False + + +class TimeSeriesDatasetZarr(BaseTimeSeriesDatasetZarr): + """Dataset for sampling from continuous time-series data stored in Zarr format. + + This class implements the basic time series dataset functionality without coupling + to external data sources. It provides data loading, scaling, and batching + capabilities for training and inference. + """ + + def __init__( + self, + dataset_path: str, + input_variables: Sequence, + output_variables: Sequence = None, + constant_variables: Sequence = None, + scaling: DictConfig = None, + input_time_dim: int = 1, + output_time_dim: int = 1, + data_time_step: Union[int, str] = "3h", + time_step: Union[int, str] = "6h", + gap: Union[int, str, None] = None, + batch_size: int = 32, + drop_last: bool = False, + add_insolation: bool = False, + forecast_init_times: Optional[Sequence] = None, + start_date: Optional[Union[int, str]] = None, + end_date: Optional[Union[int, str]] = None, + add_train_noise: bool = False, + train_noise_params: DictConfig = None, + train_noise_seed: int = 42, + meta: DatapipeMetaData = MetaData(), + ): + """Initialize time series dataset. + + Parameters are same as BaseTimeSeriesDatasetZarr. + See base class for detailed parameter descriptions. + """ + super().__init__( + dataset_path=dataset_path, + input_variables=input_variables, + output_variables=output_variables, + constant_variables=constant_variables, + scaling=scaling, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + data_time_step=data_time_step, + time_step=time_step, + gap=gap, + batch_size=batch_size, + drop_last=drop_last, + add_insolation=add_insolation, + forecast_init_times=forecast_init_times, + start_date=start_date, + end_date=end_date, + add_train_noise=add_train_noise, + train_noise_params=train_noise_params, + train_noise_seed=train_noise_seed, + meta=meta, + ) + + def __getitem__( + self, item: int + ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]]: + """Get a batch of time series data. + + This implementation provides the basic time series data loading functionality: + 1. Load data window for batch + 2. Apply scaling + 3. Extract input and target sequences + 4. Add optional insolation and constant inputs + 5. Format dimensions appropriately + + Parameters + ---------- + item : int + Sample index + + Returns + ------- + Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]] + In forecast mode: List of input arrays + In training mode: Tuple of (input arrays, target array) + + Input arrays are in order: + - Model inputs [B, F, T, C, H, W] + - Insolation (if enabled) [B, F, T, 1, H, W] + - Constants (if provided) [F, C, H, W] + + Target array has shape [B, F, T, C, H, W] + where: + B = batch size + F = faces + T = time steps + C = channels + H = height + W = width + + Raises + ------ + IndexError + If item is out of range + """ + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__") + + if item < 0: + item = len(self) + item + if item < 0 or item > len(self): + raise IndexError( + f"index {item} out of range for dataset with length {len(self)}" + ) + + # remark: load first then normalize + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_batch") + time_index, this_batch = self._get_time_index(item) + + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_staging_data") + # data from the input and target arrays overlap, to avoid 2 seperate loads + # we load both and then slice it later + staging_ds = self.ds["inputs"][slice(*time_index)] + staging_ds = staging_ds[:, self.all_variable_indices] + torch.cuda.nvtx.range_pop() + + # we scale this dataset to avoid doing twice the work when we scale as both + # input and output + # we do the scaling as in place operations to avoid creating temp arrays + # that result in a lot of data movement. This is around 4x faster than using + # standard operations + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:scale_batch") + if self.all_scaling is not None: + staging_ds -= self.all_scaling["mean"] + staging_ds /= self.all_scaling["std"] + + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_input") + input_array = staging_ds[:, self.input_variable_indices] + torch.cuda.nvtx.range_pop() + + if not self.forecast_mode: + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_target") + target_array = staging_ds[:, self.output_variable_indices] + torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:process_batch") + + # Calculate insolation if needed + if self.add_insolation: + sol = insolation( + self._get_forecast_sol_times(item), + self.lat, + self.lon, + )[:, None] + decoder_inputs = np.empty( + (this_batch, self.input_time_dim + self.output_time_dim, 1) + + self.spatial_dims, + dtype="float32", + ) + + # Get buffers for the batches, which we'll fill in iteratively. + inputs = np.empty( + (this_batch, self.input_time_dim, len(self.input_variables)) + + self.spatial_dims, + dtype="float32", + ) + if not self.forecast_mode: + targets = np.empty( + (this_batch, self.output_time_dim, len(self.output_variables)) + + self.spatial_dims, + dtype="float32", + ) + + # Iterate over valid sample windows + torch.cuda.nvtx.range_push( + "TimeSeriesDataset:__getitem__:copy_inputs_targets_insolation" + ) + for sample in range(this_batch): + inputs[sample] = input_array[self._input_indices[sample]] + if not self.forecast_mode: + targets[sample] = target_array[self._output_indices[sample]] + if self.add_insolation: + decoder_inputs[sample] = ( + sol + if self.forecast_mode + else sol[self._input_indices[sample] + self._output_indices[sample]] + ) + torch.cuda.nvtx.range_pop() + + if not self.forecast_mode and self.add_train_noise: + torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:add_train_noise") + # Iterate over C: inputs.shape = [B, T, C, F, H, W] + for i in range(inputs.shape[2]): + inputs[:, :, i] += self.rng.normal( + loc=0, + scale=self.train_noise_params["inputs"][self.input_variables[i]][ + "std" + ], + size=inputs[:, :, i].shape, + ) + torch.cuda.nvtx.range_pop() + + inputs_result = [inputs] + torch.cuda.nvtx.range_push( + "CoupledTimeSeriesDataset:__getitem__:add_insolation" + ) + if self.add_insolation: + inputs_result.append(decoder_inputs) + torch.cuda.nvtx.range_pop() + + # Transpose dimensions to match expected format + inputs_result = [ + np.transpose(x, axes=(0, 3, 1, 2, 4, 5)) for x in inputs_result + ] + + if self.constant_variables: + # Add constants + inputs_result.append(self.constants) + torch.cuda.nvtx.range_pop() + + if self.forecast_mode: + torch.cuda.nvtx.range_pop() + return inputs_result + + # Transpose targets to match input format + targets = np.transpose(targets, axes=(0, 3, 1, 2, 4, 5)) + + torch.cuda.nvtx.range_pop() + return inputs_result, targets diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index ce7f5b6d69..ef759ad924 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -22,6 +22,16 @@ from physicsnemo.core.version_check import OptionalImport xr = OptionalImport("xarray") +earth2grid = OptionalImport("earth2grid") + +try: + from cuhpx import SHTCUDA, iSHTCUDA + from earth2grid.healpix import HEALPIX_PAD_XY, PixelOrder +except ImportError: + SHTCUDA = None + iSHTCUDA = None + HEALPIX_PAD_XY = None + PixelOrder = None """ Custom dlwp compatible loss classes that allow for more sophisticated training optimization. @@ -79,6 +89,7 @@ def forward(self, prediction, target, average_channels=True): class WeightedMSE(th.nn.MSELoss): + """ Loss object that allows for user defined weighting of variables when calculating MSE """ @@ -131,6 +142,57 @@ def forward(self, prediction, target, average_channels=True): else: return d +class ConditionalWeightLoss( th.nn.MSELoss ): + """ + Conditional loss for precipitation diagnostic model. + (Total 6hr precipitation is the only output field.) + """ + + def __init__( + self, + weight=(0.01,1.0), + b=None, + w=1, + ): + """ + Parameters + ----------- + weight: tuple of floats + weight[0] is used when the target precipitation value is zero + weight[1] is used for all non-zero precipitation + b: float + Exponential scaling factor used to define weighting curve for non-zero precip. weights + w: float + Final scaling factor applied to loss. + """ + + super().__init__() + self.weight_zero = weight[0] + self.weight_nonzero = weight[1] + self.b = b + self.device = None + self.w = w + + def setup(self, trainer): + self.b = th.tensor(self.b, device=trainer.device) + self.w = th.tensor(self.w, device=trainer.device) + + def forward(self, prediction, target): + """ + Computes the MSE of model prediction and applies weights for zero and non-zero precipitation cases. + + Parameters + ----------- + prediction: torch.tensor + The prediction tensor + target: torch.Tensor + The target tensor + """ + weights_for_zero = th.ones_like(target) * self.weight_zero + weights_for_nonzero = (th.ones_like(target) * self.weight_nonzero) * th.exp(self.b*target) + weights = th.where(target > 0, weights_for_nonzero, weights_for_zero) + loss = (th.mean(weights * (prediction - target) ** 2))*self.w + return loss class OceanMSE(th.nn.MSELoss): """ @@ -177,6 +239,7 @@ def setup(self, trainer): ).to(trainer.device) def forward(self, prediction, target, average_channels=True): + if not self.lsm_sum_calculated: self.lsm_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum() self.lsm_var_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum( @@ -241,6 +304,7 @@ def setup(self, trainer): self.loss_weights = self.loss_weights.to(device=trainer.device) def forward(self, prediction, target, average_channels=True): + if not self.lsm_sum_calculated: self.lsm_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum() self.lsm_var_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum( @@ -256,3 +320,595 @@ def forward(self, prediction, target, average_channels=True): return th.sum(ocean_mean_err) / self.lsm_sum else: return ocean_mean_err / self.lsm_var_sum + +class WeightedCRPSLoss(th.nn.MSELoss): + + """ + Probabilistic loss function that allows for user defined weighting of variables when calculating CRPS. + """ + + def __init__( + self, + weights: Sequence = [], + n_members: int = 2, + alpha: float = 0.95, + mean_penalty: float = 0.0, + lsm_file: str = None, + open_dict: dict = {"engine": "zarr"}, + selection_dict: dict = {"channel_c": "land_sea_mask"}, + multiscale: float = 0.0, + masked_processing: bool = False, + ): + """ + Parameters + ---------- + weights: Sequence + list of floats that determine weighting of variable loss, assumed to be + in order consistent with order of model output channels + n_members: int + number of ensemble members in the model output + alpha: float + hyperparamter for approximating fair CRPS loss. between 0 and 1, 1 corresponds to a fair CRPS loss. + mean_penalty: float + weight for the penalty constraining the global mean of the ensemble to be close to the target mean + if 0, no penalty is applied + lsm_file: str + land-sea-mask file + open_dict: dict, optional + dictionary that store land-sea-mask file information + selection_dict: dict, optional + dictionary that store channel selection information + multiscale: float, optional + weight for the multiscale CRPS loss. Default is 0, no multiscale loss is applied. + masked_processing: bool, optional + whether masked pixels should be excluded from the processing. Default is False, no masked processing is applied. + """ + super().__init__() + self.loss_weights = th.tensor(weights) + if n_members < 2: + raise ValueError("n_members must be at least 2 for CRPS loss to be defined") + else: + self.n_members = n_members + self.device = None + self.mean_penalty = mean_penalty + self.multiscale = multiscale + self.scales = [4, 16, 32] + self.masked_processing = masked_processing + + if lsm_file is not None: + self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel(selection_dict) + self.lsm_tensor = 1 - th.tensor(np.expand_dims(self.lsm_ds.values, (0, 2, 3))) + else: + self.lsm_tensor = th.ones(1, 1, 1, 1, 1, 1) # Spoof the tensor dimensions for broadcasting + + # Parameters for "almost fair CRPS" loss. See https://arxiv.org/html/2412.15832v1 + self.coeff_eps = 1 - ((1-alpha) / (n_members)) + self.averaging_coeff = 1 / (2* n_members * (n_members - 1)) + + # For n>2, will use pairwise distance to copmute [NxN] distance matrix + # Diagonal elements of (prediciton - target) matrix are zeroed out to avoid double counting + self.pdist = th.nn.PairwiseDistance(p=1) + self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye(self.n_members) # Mask to zero out diagonal elements + + def setup(self, trainer): + """ + pushes constants to cuda device + """ + + if len(trainer.output_variables) != len(self.loss_weights): + raise ValueError("Length of outputs and loss_weights is not the same!") + + self.loss_weights = self.loss_weights.to(device=trainer.device) + self.averaging_coeff = th.tensor(self.averaging_coeff, device=trainer.device) + self.coeff_eps = th.tensor(self.coeff_eps, device=trainer.device) + self.pdist = self.pdist.to(device=trainer.device) + self.diag_mask = self.diag_mask.to(device=trainer.device) + self.lsm_tensor = self.lsm_tensor.to(device=trainer.device) + + def _2member_crps(self, prediction, target, lsm_tensor): + diff_target = th.abs(prediction - target.unsqueeze(0)).sum(dim=0) # [B, F, T, C, H, W] + diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] + crps = self.averaging_coeff*(diff_target - self.coeff_eps * diff_ensemble) # [B, F, T, C, H, W] + crps *= lsm_tensor + return crps + + def _pool(self, tensor, scale): + shape = tensor.shape + h, w = shape[-2:] + pooled = th.nn.functional.avg_pool2d(tensor.reshape(shape[0], -1, h, w), scale, scale) + return pooled.reshape(*shape[:-2], h//scale, w//scale) + + def _masked_pool(self, tensor, scale, mask): + """ + Pools a tensor while excluding masked pixels + + Parameters + ---------- + tensor: torch.Tensor + The tensor to pool + scale: int + The scale factor for the pooling + mask: torch.Tensor + The mask to apply to the tensor, masked should only contain 1s and 0s + + Returns + ------- + torch.Tensor + The pooled tensor + torch.Tensor + The pooled mask + """ + # Apply the mask to the tensor + masked_values = tensor * mask + + # Compute the non-masked values for the pooled tensor + pooled_values = self._pool(masked_values, scale) + + # pool the mask to use as a weight for the pooled tensor + pooled_mask = self._pool(mask, scale) + + pooled_tensor = pooled_values / (pooled_mask + 1e-8) # Avoid division by zero + return pooled_tensor, pooled_mask + + def forward(self, prediction, target, average_channels=True): + """ + Forward pass of the WeightedCRPSLoss + Computes the CRPS loss for the model prediction and target. + + Parameters + ---------- + prediction: torch.Tensor + The prediction tensor shape [Cond*B, F, T, C, H, W] where Cond is the number of ensemble members + target: torch.Tensor + The target tensor shape [B, F, T, C, H, W] + average_channels: bool, optional + whether the mean of the channels should be taken + """ + + # Unfold ensemble dimension from batch dimension to have shape [Cond, B, F, T, C, H, W] + b, f, t, c, h, w = target.shape + prediction = prediction.view(self.n_members, b, f, t, c, h, w) + + # checks for dimensions + if not prediction.shape[1:] == target.shape: + raise ValueError(f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}") + + if not prediction.shape[0] == self.n_members: + raise ValueError(f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}") + + n = self.n_members + + # Manual Cast + prediction = prediction.to(th.float32) + target = target.to(th.float32) + + # Apply channel weights across channel dims + prediction *= self.loss_weights[None, None, None, None, :, None, None] + target *= self.loss_weights[None, None, None, :, None, None] + + if n == 2: + # Use faster explicit implementation + crps = self._2member_crps(prediction, target, self.lsm_tensor) + + if average_channels: + loss = crps.mean() + else: + loss = crps.mean(dim=(0, 1, 2, 4, 5)) + + if self.mean_penalty > 0: + # the fraction of valid pixels in land sea mask + if self.masked_processing and self.lsm_tensor.numel() > 1: + valid_pixels = self.lsm_tensor.numel() + else: + valid_pixels = f * h * w + + prediction_count = n * b * t * valid_pixels + target_count = b * t * valid_pixels + ens_global_means = (prediction * self.lsm_tensor).sum(dim=(0, 1, 2, 3, 5, 6)) / prediction_count + target_global_means = (target * self.lsm_tensor).sum(dim=(0, 1, 2, 4, 5)) / target_count + + bias_penalty = self.mean_penalty * th.abs(ens_global_means - target_global_means) + if average_channels: + loss += bias_penalty.mean() + else: + loss += bias_penalty + + # spatial multiscale loss + if self.multiscale > 0.: + crps_scales = 0 + for scale in self.scales: + if self.masked_processing and self.lsm_tensor.numel() > 1: + masked_pred, masked_lsm = self._masked_pool(prediction, scale, self.lsm_tensor) + masked_tar, _ = self._masked_pool(target, scale, self.lsm_tensor) + crps_scale = self._2member_crps(masked_pred, masked_tar, masked_lsm) + else: + pred, tar, lsm = self._pool(prediction, scale), self._pool(target, scale), self._pool(self.lsm_tensor, scale) + crps_scale = self._2member_crps(pred, tar, lsm) + + if average_channels: + crps_scale = crps_scale.mean() + else: + crps_scale = crps_scale.mean(dim=(0, 1, 2, 4, 5)) + crps_scales += crps_scale + + crps_scales = crps_scales / len(self.scales) + loss += self.multiscale * crps_scales + + return loss + else: + # Do mean penalty + bias_penalty = 0 + if self.mean_penalty > 0: + # the fraction of valid pixels in land sea mask + if self.masked_processing and self.lsm_tensor.numel() > 1: + valid_pixels = self.lsm_tensor.numel() + else: + valid_pixels = f * h * w + + prediction_count = n * b * t * valid_pixels + target_count = b * t * valid_pixels + ens_global_means = (prediction * self.lsm_tensor).sum(dim=(0, 1, 2, 3, 5, 6)) / prediction_count + target_global_means = (target * self.lsm_tensor).sum(dim=(0, 1, 2, 4, 5)) / target_count + + bias_penalty = self.mean_penalty * th.abs(ens_global_means - target_global_means) + if average_channels: + bias_penalty = bias_penalty.mean() + + # zero out land and determine number of valid pixels + if self.masked_processing and self.lsm_tensor.numel() > 1: + prediction = prediction * self.lsm_tensor + target = target * self.lsm_tensor + valid_pixels = b * t * self.lsm_tensor.sum() + else: + valid_pixels = b * f * t * h * w + + # Use pairwise distance method + if not average_channels: + # Move channels to first dimension and exclude that dimension from the reductions + prediction = prediction.permute(4, 0, 1, 2, 3, 5, 6) # [C, Cond, B, F, T, H, W] + target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] + + prediction = prediction.reshape(c, n, -1) # [C, Cond, ...] + target = target.unsqueeze(1).reshape(c, 1, -1) # [C, 1, ...] (second dim will broadcast across ensemble) + + diff = self.pdist(prediction, target) # [C, Cond] + dist_matrix = self.pdist(prediction.unsqueeze(1), prediction.unsqueeze(2)) # [C, Cond, Cond] + + diff_terms = self.diag_mask[None, ...] * (diff.unsqueeze(1) + diff.unsqueeze(2)) # [C, Cond, Cond], diagonal elements zeroed out + crps = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1,2))/valid_pixels + else: + prediction = prediction.reshape(n, -1) + target = target.unsqueeze(0).reshape(1, -1) # [1, ...] (first dim will broadcast across ensemble) + diff = self.pdist(prediction, target) # [Cond] + dist_matrix = self.pdist(prediction.unsqueeze(1), prediction.unsqueeze(0)) # [Cond, Cond] + + diff_terms = self.diag_mask * (diff.unsqueeze(0) + diff.unsqueeze(1)) # [Cond, Cond], diagonal elements zeroed out + crps = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum()/valid_pixels + + return crps + bias_penalty + + +class WeightedCRPSLossSpectral(th.nn.MSELoss): + + """ + Probabilistic loss function that allows for user defined weighting of variables when calculating CRPS. + """ + + def __init__( + self, + weights: Sequence = [], + n_members: int = 2, + alpha: float = 0.95, + lambda_spec: float = 0.1, + nside: int = 64, + lmax: int = 3*64 - 1, + mmax: int = 3*64 - 1, + multiscale: float = 0.0, + lsm_file: str = None, + open_dict: dict = {"engine": "zarr"}, + selection_dict: dict = {"channel_c": "land_sea_mask"}, + ): + """ + Parameters + ---------- + weights: Sequence + list of floats that determine weighting of variable loss, assumed to be + in order consistent with order of model output channels + n_members: int + number of ensemble members in the model output + alpha: float + hyperparamter for approximating fair CRPS loss. between 0 and 1, 1 corresponds to a fair CRPS loss. + lambda_spec: float + weight for the spectral loss. Default is 0, no spectral loss is applied. + nside: int + nside for the HEALPix grid. Default is 64. + lmax: int + lmax for the SHT. Default is 3*nside - 1. + mmax: int + mmax for the SHT. Default is 3*nside - 1. + multiscale: float, optional + weight for the multiscale loss. Default is 0, no multiscale loss is applied. + lsm_file: str + path to the lsm file. Default is None, no lsm is applied. + open_dict: dict + dictionary of keyword arguments for xarray.open_dataset. Default is {"engine": "zarr"}. + selection_dict: dict + dictionary of keyword arguments for xarray.open_dataset. Default is {"channel_c": "land_sea_mask"}. + """ + super().__init__() + self.loss_weights = th.tensor(weights) + self.n_members = n_members + self.device = None + self.lambda_spec = lambda_spec + self.multiscale = multiscale + + # Parameters for "almost fair CRPS" loss. See https://arxiv.org/html/2412.15832v1 + self.coeff_eps = 1 - ((1-alpha) / (n_members)) + self.averaging_coeff = 1 / (2* n_members * (n_members - 1)) + + # SHT utils: transform, grid reordering, output indexing + self.lmax = lmax + self.mmax = mmax + self.nside = nside + self.sht = SHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') + src_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=HEALPIX_PAD_XY) + tar_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=PixelOrder.RING) + self.reorder_to_ring = earth2grid.get_regridder(src_grid, tar_grid).to(th.float32) + if self.multiscale > 0: + self.scales = [200, 400, 800, 1600] # in units of km + self.isht = iSHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') + self.reorder_from_ring = earth2grid.get_regridder(tar_grid, src_grid).to(th.float32) + + self.lsm_file = lsm_file + if lsm_file is not None: + self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel(selection_dict) + self.lsm_tensor = 1 - th.tensor(np.expand_dims(self.lsm_ds.values, (0, 2, 3))) + else: + self.lsm_tensor = th.ones(1, 1, 1, 1, 1, 1) # Spoof the tensor dimensions for broadcasting + + # For n>2, will use pairwise distance to copmute [NxN] distance matrix + # Diagonal elements of (prediciton - target) matrix are zeroed out to avoid double counting + self.pdist = th.nn.PairwiseDistance(p=1) + self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye(self.n_members) # Mask to zero out diagonal elements + + + def setup(self, trainer): + """ + pushes constants to cuda device + """ + + if len(trainer.output_variables) != len(self.loss_weights): + raise ValueError("Length of outputs and loss_weights is not the same!") + + self.loss_weights = self.loss_weights.to(device=trainer.device) + self.averaging_coeff = th.tensor(self.averaging_coeff, device=trainer.device) + self.coeff_eps = th.tensor(self.coeff_eps, device=trainer.device) + self.reorder_to_ring = self.reorder_to_ring.to(device=trainer.device) + self.sht = self.sht.to(device=trainer.device) + self.pdist = self.pdist.to(device=trainer.device) + self.diag_mask = self.diag_mask.to(device=trainer.device) + + if self.multiscale > 0: + self.isht = self.isht.to(device=trainer.device) + self.reorder_from_ring = self.reorder_from_ring.to(device=trainer.device) + if self.lsm_file is not None: + self.lsm_tensor = self.lsm_tensor.to(device=trainer.device) + + def _apply_sht(self, x, face_dim, return_abs=True): + """Apply SHT to a tensor + Reshape to [..., F*H*W], reorder to ring, apply SHT + If return_abs is True, return the absolute value of the SHT (real**2 + imag**2) + + Parameters + ---------- + x: torch.Tensor + The tensor to apply SHT to + face_dim: int + The dimension of the tensor corrsponding to HEALPix faces + return_abs: bool, optional + Whether to return the absolute value of the SHT (real**2 + imag**2) + """ + x = th.movedim(x, face_dim, -3) + if x.shape[-3:] != (12, self.nside, self.nside): + raise ValueError(f"Shape of input tensor should be [..., F, ..., H, W] with F in position {face_dim}, got {x.shape}") + + x = x.reshape(*x.shape[:-3], -1) + x = self.reorder_to_ring(x.contiguous()) # contiguous needed for channels first format in validation loop + x = self.sht(x) + if return_abs: + x = x.real ** 2 + x.imag ** 2 + return x + + def _apply_isht(self, x, face_dim): + """Apply inverse SHT to a tensor shape [..., l, m] + Inverse transform, reorder from ring, Reshape to [..., F, H, W], move face dim appropriately + """ + + x = self.isht(x) # [..., l, m] -> [..., F*H*W] + x = self.reorder_from_ring(x) + x = x.reshape(*x.shape[:-1], 12, self.nside, self.nside) # [..., F*H*W] -> [..., F, H, W] + x = th.movedim(x, -3, face_dim) # [..., F, H, W] -> [..., F, ..., H, W] + return x + + def _l_filter(self, scale, device="cuda"): + """Return a spherical gaussian filter of scale `scale` (in units of km) + """ + scale_radians = scale / 6371.0 + ell = th.arange(self.lmax, device=device, dtype=th.float32) + return th.exp(-0.5* ell * (ell + 1) * (scale_radians ** 2)) + + def forward(self, prediction, target, average_channels=True): + """ + Forward pass of the WeightedCRPSLoss + Computes the CRPS loss for the model prediction and target. + + Parameters + ---------- + prediction: torch.Tensor + The prediction tensor shape [Cond*B, F, T, C, H, W] where Cond is the number of ensemble members + target: torch.Tensor + The target tensor shape [B, F, T, C, H, W] + average_channels: bool, optional + whether the mean of the channels should be taken + """ + + # Unfold ensemble dimension from batch dimension to have shape [Cond, B, F, T, C, H, W] + b, f, t, c, h, w = target.shape + prediction = prediction.view(self.n_members, b, f, t, c, h, w) + + # checks for dimensions + if not prediction.shape[1:] == target.shape: + raise ValueError(f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}") + + if not prediction.shape[0] == self.n_members: + raise ValueError(f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}") + + n = self.n_members + + # Manual cast + prediction = prediction.to(th.float32) + target = target.to(th.float32) + + # Apply channel weights across channel dims + prediction *= self.loss_weights[None, None, None, None, :, None, None] + target *= self.loss_weights[None, None, None, :, None, None] + + if n == 2: + # Use faster explicit implementation + diff_target = th.abs(prediction - target.unsqueeze(0)).sum(dim=0) # [B, F, T, C, H, W] + diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] + crps = self.averaging_coeff*(diff_target - self.coeff_eps * diff_ensemble) # [B, F, T, C, H, W] + + if average_channels: + loss = crps.mean() + else: + loss = crps.mean(dim=(0, 1, 2, 4, 5)) + + if self.lambda_spec > 0: + + with th.cuda.amp.autocast(enabled=False): + # # Reorder predictions: [N, B, F, T, C, H, W] -> [N, B, T, C, F*H*W] + # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) + + # # Reorder targets: [B, F, T, C, H, W] -> [B, T, C, F*H*W] + # tar_ring = self.reorder_to_ring(target.permute(0, 2, 3, 1, 4, 5).reshape(b, t, c, f*h*w)) + + # # Compute SHT of predictions and targets + # sht_pred = self.sht(pred_ring) # [N, B, T, C, l, m] + # sht_tar = self.sht(tar_ring) # [B, T, C, l, m] + # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 + # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 + + sht_pred = self._apply_sht(prediction, face_dim=2, return_abs=True) + sht_tar = self._apply_sht(target, face_dim=1, return_abs=True) + + diff_sht_target = th.abs(sht_pred - sht_tar.unsqueeze(0)).sum(dim=(0, 4, 5)) # [B, T, C] + diff_sht_ensemble = th.abs(sht_pred[0] - sht_pred[1]).sum(dim=(-1,-2)) # [B, T, C] + crps_sht = self.averaging_coeff * (diff_sht_target - self.coeff_eps * diff_sht_ensemble) # [B, T, C] + + # Compute spectral afCRPS + if average_channels: + spec_loss = crps_sht.mean() + else: + spec_loss = crps_sht.mean(dim=(0, 1)) + + loss += self.lambda_spec * spec_loss + + if self.multiscale > 0: + for scale in self.scales: + l_filter = self._l_filter(scale, device=prediction.device) + with th.cuda.amp.autocast(enabled=False): + sht_pred= self._apply_sht(prediction, face_dim=2, return_abs=False) + sht_tar = self._apply_sht(target, face_dim=1, return_abs=False) + + l_filter_pred = l_filter[None, None, None, None, :, None] # [1, 1, 1, 1, lmax, 1] + l_filter_tar = l_filter[None, None, None, :, None] # [1, 1, 1, lmax, 1] + + pred_smooth = self._apply_isht(l_filter_pred * sht_pred, face_dim=2) + tar_smooth = self._apply_isht(l_filter_tar * sht_tar, face_dim=1) + + diff_target = th.abs(pred_smooth - tar_smooth.unsqueeze(0)).sum(dim=0) # [B, F, T, C, H, W] + diff_ensemble = th.abs(pred_smooth[0] - pred_smooth[1]) # [B, F, T, C, H, W] + crps = self.averaging_coeff*(diff_target - self.coeff_eps * diff_ensemble) # [B, F, T, C, H, W] + + crps *= self.lsm_tensor + + if average_channels: + loss += self.multiscale * crps.mean() + else: + loss += self.multiscale * crps.mean(dim=(0, 1, 2, 4, 5)) + + return loss + + else: + # Use pairwise distance method + if not average_channels: + # Move channels to first dimension and exclude that dimension from the reductions + prediction = prediction.permute(4, 0, 1, 2, 3, 5, 6) # [C, Cond, B, F, T, H, W] + target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] + + pred = prediction.reshape(c, n, -1) # [C, Cond, ...] + tar = target.unsqueeze(1).reshape(c, 1, -1) # [C, 1, ...] (second dim will broadcast across ensemble) + + diff = self.pdist(pred, tar) # [C, Cond] + dist_matrix = self.pdist(pred.unsqueeze(1), pred.unsqueeze(2)) # [C, Cond, Cond] + + diff_terms = self.diag_mask[None, ...] * (diff.unsqueeze(1) + diff.unsqueeze(2)) # [C, Cond, Cond], diagonal elements zeroed out + loss = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1,2))/(b*f*t*h*w) + + if self.lambda_spec > 0: + with th.cuda.amp.autocast(enabled=False): + # # Reorder predictions: [C, Cond, B, F, T, H, W] -> [C, Cond, B, T, F*H*W] + # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 2, 4, 3, 5, 6).reshape(c, n, b, t, f*h*w)) + + # # Reorder targets: [C, B, F, T, H, W] -> [C, B, T, F*H*W] + # tar_ring = self.reorder_to_ring(target.permute(0, 1, 3, 2, 4, 5).reshape(c, b, t, f*h*w)) + + # # Compute SHT of predictions and targets + # sht_pred = self.sht(pred_ring).reshape(c, n, -1) # [C, Cond, B, T, l, m] -> [C, Cond, ...] + # sht_tar = self.sht(tar_ring).unsqueeze(1).reshape(c, 1, -1) # [C, B, T, l, m] -> [C, 1, ...] (second dim will broadcast across ensemble) + # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 + # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 + + sht_pred = self._apply_sht(prediction, face_dim=3, return_abs=True).reshape(c, n, -1) + sht_tar = self._apply_sht(target, face_dim=2, return_abs=True).unsqueeze(1).reshape(c, 1, -1) + + diff = self.pdist(sht_pred, sht_tar) # [C, Cond] + dist_matrix = self.pdist(sht_pred.unsqueeze(1), sht_pred.unsqueeze(2)) # [C, Cond, Cond] + + diff_terms = self.diag_mask[None, ...] * (diff.unsqueeze(1) + diff.unsqueeze(2)) # [C, Cond, Cond], diagonal elements zeroed out + loss += self.lambda_spec * self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1,2))/(b*t*self.lmax*self.mmax) + + else: + pred = prediction.reshape(n, -1) + tar = target.unsqueeze(0).reshape(1, -1) # [1, ...] (first dim will broadcast across ensemble) + diff = self.pdist(pred, tar) # [Cond] + dist_matrix = self.pdist(pred.unsqueeze(1), pred.unsqueeze(0)) # [Cond, Cond] + + diff_terms = self.diag_mask * (diff.unsqueeze(0) + diff.unsqueeze(1)) # [Cond, Cond], diagonal elements zeroed out + loss = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum()/(b*f*c*t*h*w) + + if self.lambda_spec > 0: + with th.cuda.amp.autocast(enabled=False): + # # Reorder predictions: [Cond, B, F, T, C, H, W] -> [Cond, B, T, C, F*H*W] + # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) + + # # Reorder targets: [B, F, T, C, H, W] -> [B, T, C, F*H*W] + # tar_ring = self.reorder_to_ring(target.permute(0, 2, 3, 1, 4, 5).reshape(b, t, c, f*h*w)) + + # # Compute SHT of predictions and targets + # sht_pred = self.sht(pred_ring).reshape(n, -1) # [Cond, B, T, C, l, m] -> [Cond, ...] + # sht_tar = self.sht(tar_ring).unsqueeze(0).reshape(1, -1) # [B, T, C, l, m] -> [1, ...] (first dim will broadcast across ensemble) + # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 + # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 + sht_pred = self._apply_sht(prediction, face_dim=2, return_abs=True).reshape(n, -1) + sht_tar = self._apply_sht(target, face_dim=1, return_abs=True).unsqueeze(0).reshape(1, -1) + + diff = self.pdist(sht_pred, sht_tar) # [Cond] + dist_matrix = self.pdist(sht_pred.unsqueeze(1), sht_pred.unsqueeze(0)) # [Cond, Cond] + + diff_terms = self.diag_mask * (diff.unsqueeze(0) + diff.unsqueeze(1)) # [Cond, Cond], diagonal elements zeroed out + loss += self.lambda_spec * self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum()/(b*t*self.lmax*self.mmax) + + return loss + diff --git a/physicsnemo/metrics/climate/hydrostasy.py b/physicsnemo/metrics/climate/hydrostasy.py new file mode 100644 index 0000000000..91857f71f4 --- /dev/null +++ b/physicsnemo/metrics/climate/hydrostasy.py @@ -0,0 +1,830 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from typing import Dict, Sequence + +import numpy as np +import torch +import xarray as xr + +logger = logging.getLogger(__name__) + + +def _average_virtual_temperature_from_geopotential_height(z1, z2, p1, p2, R, g0): + return g0 / (R * np.log(p1 / p2)) * (z2 - z1) + + +def _virtual_temperature_from_geopotential_height(z1, z2, p1, p2, T1, R, g0): + return (2.0 * g0) / (R * np.log(p1 / p2)) * (z2 - z1) - T1 + + +class HydrostaticBalance(torch.nn.Module): + def __init__( + self, + z_pressure_levels: Dict[int, float], + anchor_z_channel: int, + anchor_T_channel: int, + R: float, + g0: float, + ) -> None: + super().__init__() + self.R = R + self.g0 = g0 + + assert ( + anchor_z_channel in z_pressure_levels.keys() + ), f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" + + # Sort channels by pressure levels for ease of use later + self.z_pressure_levels = dict( + sorted(z_pressure_levels.items(), key=lambda item: item[1]) + ) + self.anchor_z_channel = anchor_z_channel + self.anchor_T_channel = anchor_T_channel + self.anchor_pressure = z_pressure_levels[anchor_z_channel] + self.anchor_z_index = list(self.z_pressure_levels.keys()).index( + anchor_z_channel + ) + self.z_channels = list(self.z_pressure_levels.keys()) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + + Parameters + ---------- + x : torch.Tensor + Input tensor that contains the geopotential heights (Z) at the channels and pressure levels specified in the constructor + [B, F, C, H, W] is the format + + + Returns + ------- + torch.Tensor + + """ + Tv_size = [ + len(self.z_pressure_levels) if i == 2 else s for i, s in enumerate(x.size()) + ] + Tv = torch.empty(Tv_size, dtype=x.dtype, layout=x.layout, device=x.device) + Tv[:, :, self.anchor_z_index, ...] = x[:, :, self.anchor_T_channel, ...] + + # Go down in index (up in vertical level) from anchor + for i in range(self.anchor_z_index, 0, -1): + z_channel = self.z_channels[i] + z_channel_m1 = self.z_channels[i - 1] + zi = x[:, :, z_channel, ...] + zim1 = x[:, :, z_channel_m1, ...] + Tv[:, :, i - 1, ...] = _virtual_temperature_from_geopotential_height( + zi, + zim1, + self.z_pressure_levels[z_channel], + self.z_pressure_levels[z_channel_m1], + Tv[:, :, i, ...], + self.R, + self.g0, + ) + + # Go up in index (down in vertical level) from anchor + for i in range(self.anchor_z_index, len(self.z_pressure_levels) - 1): + z_channel = self.z_channels[i] + z_channel_p1 = self.z_channels[i + 1] + zi = x[:, :, z_channel, ...] + zip1 = x[:, :, z_channel_p1, ...] + Tv[:, :, i + 1, ...] = _virtual_temperature_from_geopotential_height( + zi, + zip1, + self.z_pressure_levels[z_channel], + self.z_pressure_levels[z_channel_p1], + Tv[:, :, i, ...], + self.R, + self.g0, + ) + + return Tv + + +class DifferentialHydrostaticBalanceConstraint(torch.nn.Module): + def __init__( + self, + z_pressure_levels: Dict[int, float], + Tv_pressure_levels: Dict[int, float], + anchor_z_channel: int, + anchor_T_channel: int, + R: float, + g0: float, + ) -> None: + super().__init__() + self.R = R + self.g0 = g0 + + assert ( + anchor_z_channel in z_pressure_levels.keys() + ), f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" + assert ( + anchor_T_channel in Tv_pressure_levels.keys() + ), f"anchor_T_channel ({anchor_T_channel}) not in Tv_pressure_levels ({Tv_pressure_levels.keys()})" + + # Sort channels by pressure levels for ease of use later + self.z_pressure_levels = dict( + sorted(z_pressure_levels.items(), key=lambda item: item[1]) + ) + self.Tv_pressure_levels = dict( + sorted(Tv_pressure_levels.items(), key=lambda item: item[1]) + ) + self.anchor_z_channel = anchor_z_channel + self.anchor_T_channel = anchor_T_channel + self.anchor_pressure = z_pressure_levels[anchor_z_channel] + self.anchor_z_index = list(self.z_pressure_levels.keys()).index( + anchor_z_channel + ) + self.anchor_T_index = list(self.Tv_pressure_levels.keys()).index( + anchor_T_channel + ) + self.z_channels = list(self.z_pressure_levels.keys()) + self.Tv_channels = list(self.Tv_pressure_levels.keys()) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + + Parameters + ---------- + x : torch.Tensor + Input tensor that contains the geopotential heights (Z) at the channels and pressure levels specified in the constructor + [B, F, C, H, W] is the format + + + Returns + ------- + torch.Tensor + + """ + Tv_avg_size = [ + len(self.z_pressure_levels) - 1 if i == 2 else s + for i, s in enumerate(x.size()) + ] + Tv_avg = torch.empty( + Tv_avg_size, dtype=x.dtype, layout=x.layout, device=x.device + ) + + # Go up in index (down in vertical level) from anchor + # TODO: remove constraint that first z_channel is 0 + for i in range(len(self.z_pressure_levels) - 1): + z_channel = self.z_channels[i] + z_channel_p1 = self.z_channels[i + 1] + zi = x[:, :, z_channel, ...] + zip1 = x[:, :, z_channel_p1, ...] + Tv_avg[ + :, :, i, ... + ] = _average_virtual_temperature_from_geopotential_height( + zi, + zip1, + self.z_pressure_levels[z_channel], + self.z_pressure_levels[z_channel_p1], + self.R, + self.g0, + ) + + Tv_model_avg_size = [ + len(self.Tv_pressure_levels) - 1 if i == 2 else s + for i, s in enumerate(x.size()) + ] + Tv_model_avg = torch.empty( + Tv_model_avg_size, dtype=x.dtype, layout=x.layout, device=x.device + ) + + # Go up in index (down in vertical level) from anchor + for i in range(len(self.Tv_pressure_levels) - 1): + Tv_channel = self.Tv_channels[i] + Tv_channel_p1 = self.Tv_channels[i + 1] + Tvi = x[:, :, Tv_channel, ...] + Tvip1 = x[:, :, Tv_channel_p1, ...] + Tv_model_avg[:, :, i, ...] = 0.5 * (Tvi + Tvip1) + + return Tv_avg, Tv_model_avg + + +class WeightedMSEWithHydrostasy(torch.nn.MSELoss): + + """ + Loss object that adds a differential Hydrostatic balance constraint in addition to + user defined weighting of variables when calculating MSE + """ + + def __init__( + self, + hPa_levels: Sequence[int], + channels: Sequence[str], + weights: Sequence, + alpha: Sequence[float], # K + scaling: Dict[str, Dict[str, float]], + src_directory: str, + dst_directory: str, + dataset_name: str, + data_format: str, + surface_geopotential_mean: float = -597.7115478515625, + surface_geopotential_std: float = 55658.21484375, + R: float = 287, # J K^{-1} kg^{-1} + g0: float = 9.81, # m s^{-2} + topography_masking: bool = True, + ): + """ + Parameters + ---------- + weights: Sequence + list of floats that determine weighting of variable loss, assumed to be + in order consistent with order of model output channels + """ + super().__init__() + self.loss_weights = torch.tensor(weights) + self.device = None + self.g0 = g0 + self.topography_masking = topography_masking + + # Get channel index to pressure level mapping + self.pressure_levels = sorted(hPa_levels) + self.z_pressure_levels = { + channels.index(f"z{int(pl)}"): pl + for i, pl in enumerate(self.pressure_levels) + } + self.T_pressure_levels = { + channels.index(f"t{int(pl)}"): pl + for i, pl in enumerate(self.pressure_levels) + } + self.q_pressure_levels = { + channels.index(f"q{int(pl)}"): pl + for i, pl in enumerate(self.pressure_levels) + if f"q{int(pl)}" in channels + } + # Get offset to map from q channels here to Tv channels + # Relies on all pressure levels below a threshold to have q channels + for i, pl in enumerate(self.pressure_levels): + if f"q{int(pl)}" in channels: + self.q_index_offset = i + break + + # Create mapping for new tensor that holds only the constraint variables + self.z_constraint_pressure_levels = { + i: pl for i, pl in enumerate(self.pressure_levels) + } + self.Tv_constraint_pressure_levels = { + len(self.z_constraint_pressure_levels) + i: pl + for i, pl in enumerate(self.pressure_levels) + } + + # Get scaling weights + self.z_mean = torch.Tensor( + [ + scaling[f"z{int(pl)}"]["mean"] + for i, pl in enumerate(self.pressure_levels) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + self.z_std = torch.Tensor( + [scaling[f"z{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] + ).reshape((1, 1, 1, -1, 1, 1)) + self.T_mean = torch.Tensor( + [ + scaling[f"t{int(pl)}"]["mean"] + for i, pl in enumerate(self.pressure_levels) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + self.T_std = torch.Tensor( + [scaling[f"t{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] + ).reshape((1, 1, 1, -1, 1, 1)) + self.q_mean = torch.Tensor( + [ + scaling[f"q{int(pl)}"]["mean"] + for i, pl in enumerate(self.q_pressure_levels.values()) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + self.q_std = torch.Tensor( + [ + scaling[f"q{int(pl)}"]["std"] + for i, pl in enumerate(self.q_pressure_levels.values()) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + + # Set per level alphas + if len(alpha) != len(hPa_levels) - 1: + raise AssertionError( + f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels)-1}], got {len(alpha)}" + ) + self.alpha = torch.Tensor(alpha).reshape((1, 1, -1, 1, 1)) + + # Molecular weight ratio factor of water vapor to air + self.Mw_ratio = 28.97 / 18.016 - 1.0 # 0.6078 + + # Create the constraint + # TODO: remove anchor levels since it's not needed for the + # differential constraint + self.constraint = DifferentialHydrostaticBalanceConstraint( + self.z_constraint_pressure_levels, + self.Tv_constraint_pressure_levels, + 0, + len(self.z_pressure_levels), + R, + self.g0, + ) + + self.num_z_levels = len(self.z_constraint_pressure_levels) + self.num_Tv_levels = len(self.Tv_constraint_pressure_levels) + self.z_level_mapping = torch.tensor(list(self.z_pressure_levels.keys())) + self.T_level_mapping = torch.tensor(list(self.T_pressure_levels.keys())) + self.q_level_mapping = torch.tensor(list(self.q_pressure_levels.keys())) + + # Get topography information + ds = xr.open_zarr(f"{src_directory}{dataset_name}.zarr") + self.topography = ( + surface_geopotential_std * ds.constants[1, :, :, :].values + + surface_geopotential_mean + ) / self.g0 + + self.topography = torch.tensor( + self.topography[np.newaxis, :, np.newaxis, :, :], dtype=torch.float + ) + logger.info( + f"Min/Max topography (m): {self.topography.min()}/{self.topography.max()}" + ) + + def setup(self, trainer): + """ + pushes weights to cuda device + """ + + if len(trainer.output_variables) + len(self.z_pressure_levels) - 1 != len( + self.loss_weights + ): + raise ValueError("Length of outputs and loss_weights is not the same!") + + self.loss_weights = self.loss_weights.to(device=trainer.device) + + # Move means and stds + self.z_mean = self.z_mean.to(device=trainer.device) + self.z_std = self.z_std.to(device=trainer.device) + self.T_mean = self.T_mean.to(device=trainer.device) + self.T_std = self.T_std.to(device=trainer.device) + self.q_mean = self.q_mean.to(device=trainer.device) + self.q_std = self.q_std.to(device=trainer.device) + + # Move alphas + self.alpha = self.alpha.to(device=trainer.device) + + # Move indexing arrays for CUDA graphs + self.z_level_mapping = self.z_level_mapping.to(device=trainer.device) + self.T_level_mapping = self.T_level_mapping.to(device=trainer.device) + self.q_level_mapping = self.q_level_mapping.to(device=trainer.device) + + # Move topography + self.topography = self.topography.to(device=trainer.device) + + def scale(self, x): + """ + Scale inputs to physical values and compute virtual temperature + Tensors are expected to be in the shape [N, B, F, C, H, W] + """ + # N, B, F, C, H, W = x.shape + N, F, B, C, H, W = x.shape + C_scaled = self.num_z_levels + self.num_Tv_levels + x_scaled = torch.zeros( + # (N, B, F, C_scaled, H, W), device=x.device, dtype=torch.float + (N, F, B, C_scaled, H, W), + device=x.device, + dtype=torch.float, + ) + # Get scaled geopotential heights + x_scaled[:, :, :, : self.num_z_levels, :, :] = ( + x[:, :, :, self.z_level_mapping, :, :] * self.z_std + self.z_mean + ) / self.g0 # divide by g0 for heights + # Get scaled temperatures + x_scaled[:, :, :, self.num_z_levels :, :, :] = ( + x[:, :, :, self.T_level_mapping, :, :] * self.T_std + self.T_mean + ) + # Add q correction to get virtual temperature for levels with non-zero q + x_scaled[ + :, + :, + :, + (self.num_z_levels + self.q_index_offset) :, + :, + :, + ] *= 1.0 + self.Mw_ratio * ( + x[:, :, :, self.q_level_mapping, :, :] * self.q_std + self.q_mean + ) + + # transpose B dim to before F + x_scaled = x_scaled.transpose(1, 2) + + # Combine N and B dimensions and return + return x_scaled.reshape((-1, F, C_scaled, H, W)) + + def error_histogram(self, prediction, bins, accumulator=None): + N, F, B, C, H, W = tuple(prediction.shape) + + if not (prediction.ndim == 6): + raise AssertionError("Expected predictions to have 6 dimensions") + + # Scale to physical units and compute virtual temperature + x = self.scale(prediction) + Tv_avg, Tv_model_avg = self.constraint(x) + Tv_error = Tv_avg - Tv_model_avg + # Mask out error in regions below the surface + if self.topography_masking: + Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 + + vlevels = Tv_error.shape[2] + if accumulator is None: + if isinstance(bins, int): + accumulator = torch.zeros( + (vlevels, bins), dtype=torch.float32, device=prediction.device + ) + else: + accumulator = torch.zeros( + (vlevels, bins.shape[1] - 1), + dtype=torch.float32, + device=prediction.device, + ) + if isinstance(bins, int): + bin_edges = torch.zeros( + (vlevels, bins + 1), + dtype=prediction.dtype, + device=prediction.device, + ) + else: + bin_edges = bins + + for l in range(vlevels): + hist, be = torch.histogram( + torch.absolute(Tv_error[:, :, l, :, :]), + bins=bins if isinstance(bins, int) else bin_edges[l, :], + ) + accumulator[l, :] += hist + bin_edges[l, :] = be + + return accumulator, bin_edges + + def forward(self, prediction, target, average_channels=True): + """ + Forward pass of the WeightedMSE pass + Tensors are expected to be in the shape [N, F, B, C, H, W] + + Parameters + ---------- + prediction: torch.Tensor + The prediction tensor + target: torch.Tensor + The target tensor + average_channels: bool, optional + whether the mean of the channels should be taken + """ + + # Need to scale back to physical units here so disable autocast + # and explicitly cast to float32 + with torch.cuda.amp.autocast(enabled=False): + prediction = prediction.float() + target = target.float() + + N, F, B, C, H, W = tuple(prediction.shape) + + if not (prediction.ndim == 6 and target.ndim == 6): + raise AssertionError("Expected predictions to have 6 dimensions") + + # Scale to physical units and compute virtual temperature + x = self.scale(prediction) + Tv_avg, Tv_model_avg = self.constraint(x) + Tv_error = ((Tv_avg - Tv_model_avg) / self.alpha) ** 2 + + # Mask out error in regions below the surface + if self.topography_masking: + Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 + + # Compute the error tolerant loss + Tv_loss = (Tv_error / (1 + torch.exp(1 - Tv_error))).mean(dim=(0, 1, 3, 4)) + + data_loss = ((target - prediction) ** 2).mean(dim=(0, 1, 2, 4, 5)) + d = torch.concatenate((data_loss, Tv_loss)) * self.loss_weights + + if average_channels: + return torch.mean(d) + else: + return d + + +class LossWithHydrostasy(torch.nn.MSELoss): + + """ + Loss object that adds a differential Hydrostatic balance constraint in addition to + user defined data loss + """ + + def __init__( + self, + data_loss: torch.nn.MSELoss, + hPa_levels: Sequence[int], + channels: Sequence[str], + weights: Sequence, + alpha: Sequence[float], # K + scaling: Dict[str, Dict[str, float]], + src_directory: str, + dst_directory: str, + dataset_name: str, + data_format: str, + surface_geopotential_mean: float = -597.7115478515625, + surface_geopotential_std: float = 55658.21484375, + R: float = 287, # J K^{-1} kg^{-1} + g0: float = 9.81, # m s^{-2} + topography_masking: bool = True, + ): + """ + Parameters + ---------- + weights: Sequence + list of floats that determine weighting of hydrostatic loss terms, assumed to be + in order of increasing pressure levels + """ + super().__init__() + self.data_loss = data_loss + self.loss_weights = torch.tensor(weights) + self.device = None + self.g0 = g0 + self.topography_masking = topography_masking + + # Get channel index to pressure level mapping + self.pressure_levels = sorted(hPa_levels) + self.z_pressure_levels = { + channels.index(f"z{int(pl)}"): pl + for i, pl in enumerate(self.pressure_levels) + } + self.T_pressure_levels = { + channels.index(f"t{int(pl)}"): pl + for i, pl in enumerate(self.pressure_levels) + } + self.q_pressure_levels = { + channels.index(f"q{int(pl)}"): pl + for i, pl in enumerate(self.pressure_levels) + if f"q{int(pl)}" in channels + } + # Get offset to map from q channels here to Tv channels + # Relies on all pressure levels below a threshold to have q channels + for i, pl in enumerate(self.pressure_levels): + if f"q{int(pl)}" in channels: + self.q_index_offset = i + break + + # Create mapping for new tensor that holds only the constraint variables + self.z_constraint_pressure_levels = { + i: pl for i, pl in enumerate(self.pressure_levels) + } + self.Tv_constraint_pressure_levels = { + len(self.z_constraint_pressure_levels) + i: pl + for i, pl in enumerate(self.pressure_levels) + } + + # Get scaling weights + self.z_mean = torch.Tensor( + [ + scaling[f"z{int(pl)}"]["mean"] + for i, pl in enumerate(self.pressure_levels) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + self.z_std = torch.Tensor( + [scaling[f"z{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] + ).reshape((1, 1, 1, -1, 1, 1)) + self.T_mean = torch.Tensor( + [ + scaling[f"t{int(pl)}"]["mean"] + for i, pl in enumerate(self.pressure_levels) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + self.T_std = torch.Tensor( + [scaling[f"t{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] + ).reshape((1, 1, 1, -1, 1, 1)) + self.q_mean = torch.Tensor( + [ + scaling[f"q{int(pl)}"]["mean"] + for i, pl in enumerate(self.q_pressure_levels.values()) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + self.q_std = torch.Tensor( + [ + scaling[f"q{int(pl)}"]["std"] + for i, pl in enumerate(self.q_pressure_levels.values()) + ] + ).reshape((1, 1, 1, -1, 1, 1)) + + # Set per level alphas + if len(alpha) != len(hPa_levels) - 1: + raise AssertionError( + f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels)-1}], got {len(alpha)}" + ) + self.alpha = torch.Tensor(alpha).reshape((1, 1, -1, 1, 1)) + + # Molecular weight ratio factor of water vapor to air + self.Mw_ratio = 28.97 / 18.016 - 1.0 # 0.6078 + + # Create the constraint + # TODO: remove anchor levels since it's not needed for the + # differential constraint + self.constraint = DifferentialHydrostaticBalanceConstraint( + self.z_constraint_pressure_levels, + self.Tv_constraint_pressure_levels, + 0, + len(self.z_pressure_levels), + R, + self.g0, + ) + + self.num_z_levels = len(self.z_constraint_pressure_levels) + self.num_Tv_levels = len(self.Tv_constraint_pressure_levels) + self.z_level_mapping = torch.tensor(list(self.z_pressure_levels.keys())) + self.T_level_mapping = torch.tensor(list(self.T_pressure_levels.keys())) + self.q_level_mapping = torch.tensor(list(self.q_pressure_levels.keys())) + + # Get topography information + ds = xr.open_zarr(f"{src_directory}{dataset_name}.zarr") + self.topography = ( + surface_geopotential_std * ds.constants.sel(channel_c='z').values + + surface_geopotential_mean + ) / self.g0 + + self.topography = torch.tensor( + self.topography[np.newaxis, :, np.newaxis, :, :], dtype=torch.float + ) + logger.info( + f"Min/Max topography (m): {self.topography.min()}/{self.topography.max()}" + ) + + def setup(self, trainer): + """ + pushes weights to cuda device + """ + + # Call setup for data loss first + self.data_loss.setup(trainer) + + if len(self.z_pressure_levels) - 1 != len( + self.loss_weights + ): + raise ValueError("Length of loss_weights is not one less than number of pressure levels!") + + self.loss_weights = self.loss_weights.to(device=trainer.device) + + # Move means and stds + self.z_mean = self.z_mean.to(device=trainer.device) + self.z_std = self.z_std.to(device=trainer.device) + self.T_mean = self.T_mean.to(device=trainer.device) + self.T_std = self.T_std.to(device=trainer.device) + self.q_mean = self.q_mean.to(device=trainer.device) + self.q_std = self.q_std.to(device=trainer.device) + + # Move alphas + self.alpha = self.alpha.to(device=trainer.device) + + # Move indexing arrays for CUDA graphs + self.z_level_mapping = self.z_level_mapping.to(device=trainer.device) + self.T_level_mapping = self.T_level_mapping.to(device=trainer.device) + self.q_level_mapping = self.q_level_mapping.to(device=trainer.device) + + # Move topography + self.topography = self.topography.to(device=trainer.device) + + def scale(self, x): + """ + Scale inputs to physical values and compute virtual temperature + Tensors are expected to be in the shape [N, B, F, C, H, W] + """ + # N, B, F, C, H, W = x.shape + N, F, B, C, H, W = x.shape + C_scaled = self.num_z_levels + self.num_Tv_levels + x_scaled = torch.zeros( + # (N, B, F, C_scaled, H, W), device=x.device, dtype=torch.float + (N, F, B, C_scaled, H, W), + device=x.device, + dtype=torch.float, + ) + # Get scaled geopotential heights + x_scaled[:, :, :, : self.num_z_levels, :, :] = ( + x[:, :, :, self.z_level_mapping, :, :] * self.z_std + self.z_mean + ) / self.g0 # divide by g0 for heights + # Get scaled temperatures + x_scaled[:, :, :, self.num_z_levels :, :, :] = ( + x[:, :, :, self.T_level_mapping, :, :] * self.T_std + self.T_mean + ) + # Add q correction to get virtual temperature for levels with non-zero q + x_scaled[ + :, + :, + :, + (self.num_z_levels + self.q_index_offset) :, + :, + :, + ] *= 1.0 + self.Mw_ratio * ( + x[:, :, :, self.q_level_mapping, :, :] * self.q_std + self.q_mean + ) + + # transpose B dim to before F + x_scaled = x_scaled.transpose(1, 2) + + # Combine N and B dimensions and return + return x_scaled.reshape((-1, F, C_scaled, H, W)) + + def error_histogram(self, prediction, bins, accumulator=None): + N, F, B, C, H, W = tuple(prediction.shape) + + if not (prediction.ndim == 6): + raise AssertionError("Expected predictions to have 6 dimensions") + + # Scale to physical units and compute virtual temperature + x = self.scale(prediction) + Tv_avg, Tv_model_avg = self.constraint(x) + Tv_error = Tv_avg - Tv_model_avg + # Mask out error in regions below the surface + if self.topography_masking: + Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 + + vlevels = Tv_error.shape[2] + if accumulator is None: + if isinstance(bins, int): + accumulator = torch.zeros( + (vlevels, bins), dtype=torch.float32, device=prediction.device + ) + else: + accumulator = torch.zeros( + (vlevels, bins.shape[1] - 1), + dtype=torch.float32, + device=prediction.device, + ) + if isinstance(bins, int): + bin_edges = torch.zeros( + (vlevels, bins + 1), + dtype=prediction.dtype, + device=prediction.device, + ) + else: + bin_edges = bins + + for l in range(vlevels): + hist, be = torch.histogram( + torch.absolute(Tv_error[:, :, l, :, :]), + bins=bins if isinstance(bins, int) else bin_edges[l, :], + ) + accumulator[l, :] += hist + bin_edges[l, :] = be + + return accumulator, bin_edges + + def forward(self, prediction, target, average_channels=True): + """ + Forward pass of LossWithHydrostasy + Tensors are expected to be in the shape [N, F, B, C, H, W] + + Parameters + ---------- + prediction: torch.Tensor + The prediction tensor + target: torch.Tensor + The target tensor + average_channels: bool, optional + whether the mean of the channels should be taken + """ + + # Need to scale back to physical units here so disable autocast + # and explicitly cast to float32 + with torch.cuda.amp.autocast(enabled=False): + prediction = prediction.float() + target = target.float() + + if not (prediction.ndim == 6 and target.ndim == 6): + raise AssertionError("Expected predictions to have 6 dimensions") + + # Scale to physical units and compute virtual temperature + x = self.scale(prediction) + Tv_avg, Tv_model_avg = self.constraint(x) + Tv_error = ((Tv_avg - Tv_model_avg) / self.alpha) ** 2 + + # Mask out error in regions below the surface + if self.topography_masking: + Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 + + # Compute the error tolerant loss + Tv_loss = self.loss_weights * (Tv_error / (1 + torch.exp(1 - Tv_error))).mean(dim=(0, 1, 3, 4)) + + # Compute data loss + data_loss = self.data_loss(prediction, target, average_channels=average_channels) + + if average_channels: + return data_loss + torch.mean(Tv_loss) + else: + return torch.concatenate((data_loss, Tv_loss)) diff --git a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py index 20553b5e24..70e691997e 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py @@ -145,6 +145,9 @@ def __init__( enable_nhwc: bool = False, enable_healpixpad: bool = False, couplings: list = [], + residual_prediction: bool = True, + couplings_time_first: bool = True, + constraints: list[DictConfig] = None, ): r"""Initialize the recurrent DLWP HEALPix UNet.""" super().__init__(meta=MetaData()) @@ -180,6 +183,8 @@ def __init__( self.presteps = presteps self.enable_nhwc = enable_nhwc self.enable_healpixpad = enable_healpixpad + self.residual_prediction = residual_prediction + self.couplings_time_first = couplings_time_first # Number of passes through the model, or a diagnostic model with only one output time self.is_diagnostic = self.output_time_dim == 1 and self.input_time_dim > 1 @@ -206,6 +211,9 @@ def __init__( enable_healpixpad=self.enable_healpixpad, ) + self.constraints = None + self.set_constraints(constraints) + @property def integration_steps(self): r""" @@ -297,7 +305,9 @@ def _reshape_inputs(self, inputs: Sequence, step: int = 0) -> torch.Tensor: inputs[2].expand( *tuple([inputs[0].shape[0]] + len(inputs[2].shape) * [-1]) ), # constants - inputs[3].permute(0, 2, 1, 3, 4), # coupled inputs + inputs[3].permute(0, 2, 1, 3, 4) + if self.couplings_time_first + else inputs[3], # coupled inputs ] res = torch.cat(result, dim=self.channel_dim) @@ -361,8 +371,24 @@ def _reshape_inputs(self, inputs: Sequence, step: int = 0) -> torch.Tensor: # fold faces into batch dim res = self.fold(res) + if self.enable_nhwc: + res = res.to(memory_format=torch.channels_last) return res + def set_constraints(self, constraints: list[DictConfig] = None): + r""" + Set constraints (e.g., non-negative) to be applied to model outputs. + + Parameters + ---------- + constraints : list[DictConfig], optional + Hydra instantiable constraint configurations. + """ + if constraints is not None: + self.constraints = [ + instantiate(constraints[constraint]) for constraint in constraints + ] + def _reshape_outputs(self, outputs: torch.Tensor) -> torch.Tensor: r""" Reshape decoder output back to explicit time and channel dimensions. @@ -396,7 +422,11 @@ def _reshape_outputs(self, outputs: torch.Tensor) -> torch.Tensor: return res def _initialize_hidden( - self, inputs: Sequence, outputs: Sequence, step: int + self, + inputs: Sequence, + outputs: Sequence, + step: int, + conditions_cln: Sequence = None, ) -> None: r""" Initialize the recurrent hidden states. @@ -447,19 +477,31 @@ def _initialize_hidden( s = step - self.presteps + prestep if len(self.couplings) > 0: input_tensor = self._reshape_inputs( - inputs=[outputs[s - 1]] + inputs=[outputs[s - 1][:, :, :, : self.input_channels]] + list(inputs[1:3]) + [inputs[3][step - (prestep - self.presteps)]], step=s + 1, ) else: input_tensor = self._reshape_inputs( - inputs=[outputs[s - 1]] + list(inputs[1:]), step=s + 1 + inputs=[outputs[s - 1][:, :, :, : self.input_channels]] + + list(inputs[1:]), + step=s + 1, ) - # Forward the data through the model to initialize hidden states - self.decoder(self.encoder(input_tensor)) + if conditions_cln is not None: + self.decoder( + self.encoder(input_tensor, conditions_cln=conditions_cln), + conditions_cln=conditions_cln, + ) + else: + self.decoder(self.encoder(input_tensor)) - def forward(self, inputs: Sequence, output_only_last: bool = False) -> torch.Tensor: + def forward( + self, + inputs: Sequence, + output_only_last: bool = False, + conditions_cln=None, + ) -> torch.Tensor: r""" Forward pass of the recurrent HEALPix UNet. @@ -489,7 +531,15 @@ def forward(self, inputs: Sequence, output_only_last: bool = False) -> torch.Ten for step in range(self.integration_steps): # (Re-)initialize recurrent hidden states if (step * (self.delta_t * self.input_time_dim)) % self.reset_cycle == 0: - self._initialize_hidden(inputs=inputs, outputs=outputs, step=step) + if conditions_cln is not None: + self._initialize_hidden( + inputs=inputs, + outputs=outputs, + step=step, + conditions_cln=conditions_cln[step], + ) + else: + self._initialize_hidden(inputs=inputs, outputs=outputs, step=step) # Construct concatenated input: [prognostics|TISR|constants] if step == 0: @@ -522,26 +572,40 @@ def forward(self, inputs: Sequence, output_only_last: bool = False) -> torch.Ten else: if len(self.couplings) > 0: input_tensor = self._reshape_inputs( - inputs=[outputs[-1]] + inputs=[outputs[-1][:, :, :, : self.input_channels]] + list(inputs[1:3]) + [inputs[3][self.presteps + step]], step=step + self.presteps, ) else: input_tensor = self._reshape_inputs( - inputs=[outputs[-1]] + list(inputs[1:]), + inputs=[outputs[-1][:, :, :, : self.input_channels]] + + list(inputs[1:]), step=step + self.presteps, ) - # Forward through model - encodings = self.encoder(input_tensor) - decodings = self.decoder(encodings) + if conditions_cln is not None: + kwargs = {"conditions_cln": conditions_cln[step]} + else: + kwargs = {} - # Residual prediction - reshaped = self._reshape_outputs( - input_tensor[:, : self.input_channels * self.input_time_dim] + decodings - ) - outputs.append(reshaped) + encodings = self.encoder(input_tensor, **kwargs) + decodings = self.decoder(encodings, **kwargs) + + combined = self._reshape_outputs(decodings) + prognostics = combined[:, :, :, : self.input_channels] + if self.residual_prediction: + prognostics = prognostics + self._reshape_outputs( + input_tensor[:, : self.input_channels * self.input_time_dim] + ) + diagnostics = combined[:, :, :, self.input_channels :] + out = torch.cat([prognostics, diagnostics], dim=3) + + if self.constraints is not None: + for constraint in self.constraints: + out = constraint(out) + + outputs.append(out) if output_only_last: return outputs[-1] diff --git a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py index b2d035ff1d..10ac29a6b2 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py @@ -139,6 +139,9 @@ def __init__( enable_nhwc: bool = False, enable_healpixpad: bool = False, couplings: list = [], + residual_prediction: bool = False, + couplings_time_first: bool = True, + constraints: list[DictConfig] = None, ): r"""Initialize the DLWP HEALPix UNet.""" super().__init__(meta=MetaData()) @@ -166,6 +169,8 @@ def __init__( self.channel_dim = 2 # Now 2 with [B, F, C*T, H, W]. Was 1 in old data format with [B, T*C, F, H, W] self.enable_nhwc = enable_nhwc self.enable_healpixpad = enable_healpixpad + self.residual_prediction = residual_prediction + self.couplings_time_first = couplings_time_first # Number of passes through the model, or a diagnostic model with only one output time self.is_diagnostic = self.output_time_dim == 1 and self.input_time_dim > 1 @@ -192,6 +197,9 @@ def __init__( enable_healpixpad=self.enable_healpixpad, ) + self.constraints = None + self.set_constraints(constraints) + @property def integration_steps(self): r""" @@ -283,7 +291,9 @@ def _reshape_inputs(self, inputs: Sequence, step: int = 0) -> torch.Tensor: inputs[2].expand( *tuple([inputs[0].shape[0]] + len(inputs[2].shape) * [-1]) ), # constants - inputs[3].permute(0, 2, 1, 3, 4), # coupled inputs + inputs[3].permute(0, 2, 1, 3, 4) + if self.couplings_time_first + else inputs[3], # coupled inputs ] res = torch.cat(result, dim=self.channel_dim) @@ -383,7 +393,26 @@ def _reshape_outputs(self, outputs: torch.Tensor) -> torch.Tensor: return res - def forward(self, inputs: Sequence, output_only_last: bool = False) -> torch.Tensor: + def set_constraints(self, constraints: list[DictConfig] = None): + r""" + Set constraints (e.g., non-negative) to be applied to model outputs. + + Parameters + ---------- + constraints : list[DictConfig], optional + Hydra instantiable constraint configurations. + """ + if constraints is not None: + self.constraints = [ + instantiate(constraints[constraint]) for constraint in constraints + ] + + def forward( + self, + inputs: Sequence, + output_only_last: bool = False, + conditions_cln=None, + ) -> torch.Tensor: r""" Forward pass of the HEALPix UNet. @@ -426,10 +455,28 @@ def forward(self, inputs: Sequence, output_only_last: bool = False) -> torch.Ten input_tensor = self._reshape_inputs( [outputs[-1]] + list(inputs[1:]), step ) - encodings = self.encoder(input_tensor) - decodings = self.decoder(encodings) + if conditions_cln is not None: + kwargs = {"conditions_cln": conditions_cln[step]} + else: + kwargs = {} + + encodings = self.encoder(input_tensor, **kwargs) + decodings = self.decoder(encodings, **kwargs) + + if self.residual_prediction: + prediction = ( + input_tensor[:, : self.input_channels * self.input_time_dim] + + decodings + ) + else: + prediction = decodings + + reshaped = self._reshape_outputs(prediction) + + if self.constraints is not None: + for constraint in self.constraints: + reshaped = constraint(reshaped) - reshaped = self._reshape_outputs(decodings) # Absolute prediction outputs.append(reshaped) if output_only_last: diff --git a/physicsnemo/models/dlwp_healpix/layers/__init__.py b/physicsnemo/models/dlwp_healpix/layers/__init__.py index 84f6463c3b..abfdf22788 100644 --- a/physicsnemo/models/dlwp_healpix/layers/__init__.py +++ b/physicsnemo/models/dlwp_healpix/layers/__init__.py @@ -38,8 +38,10 @@ SymmetricConvNeXtBlock, TransposedConvUpsample, ) +from .healpix_constraints import NonnegativeConstraint from .healpix_decoder import UNetDecoder from .healpix_encoder import UNetEncoder +from .normalization import ConditionalLayerNorm __all__ = [ "BasicConvBlock", @@ -52,6 +54,8 @@ "TransposedConvUpsample", "UNetDecoder", "UNetEncoder", + "ConditionalLayerNorm", + "NonnegativeConstraint", "HEALPixFoldFaces", "HEALPixLayer", "HEALPixPadding", @@ -93,6 +97,14 @@ def _remap_target(target: str) -> str: cls_name = target.split(".")[-1] return f"physicsnemo.nn.{cls_name}" + if target.startswith("physicsnemo.models.dlwp_healpix_layers.normalization."): + cls_name = target.split(".")[-1] + return f"physicsnemo.models.dlwp_healpix.layers.normalization.{cls_name}" + + if target.startswith("physicsnemo.models.dlwp_healpix_layers.healpix_constraints."): + cls_name = target.split(".")[-1] + return f"physicsnemo.models.dlwp_healpix.layers.healpix_constraints.{cls_name}" + if target.startswith("physicsnemo.models.dlwp_healpix_layers."): cls_name = target.split(".")[-1] if cls_name == "AvgPool": diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py index 32e6ffaac2..3b3ed6a243 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -14,18 +14,38 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Sequence, Tuple, Union +from typing import Sequence, Tuple, Union, Callable import torch - +import torch as th from physicsnemo.nn.module.hpx import HEALPixLayer +from .normalization import ConditionalLayerNorm + +from hydra.utils import instantiate +from omegaconf import DictConfig + + +# +# Helper: standard LayerNorm over channel dimension for (B, C, H, W) +# +class _LayerNormOverChannels(th.nn.Module): + """Applies nn.LayerNorm over the channel dimension for (B, C, H, W) tensors.""" + + def __init__(self, channel_depth: int, eps: float = 1e-5): + super().__init__() + self.norm = th.nn.LayerNorm(channel_depth, eps=eps) + + def forward(self, x): + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + return x.permute(0, 3, 1, 2) # # RECURRENT BLOCKS # -class ConvGRUBlock(torch.nn.Module): +class ConvGRUBlock(th.nn.Module): """Class that implements a Convolutional GRU Code modified from https://github.com/happyjin/ConvGRU-pytorch/blob/master/convGRU.py @@ -33,7 +53,7 @@ class ConvGRUBlock(torch.nn.Module): def __init__( self, - geometry_layer: torch.nn.Module = HEALPixLayer, + geometry_layer: th.nn.Module = HEALPixLayer, in_channels: int = 3, kernel_size: int = 1, enable_nhwc: bool = False, @@ -74,7 +94,7 @@ def __init__( enable_nhwc=enable_nhwc, enable_healpixpad=enable_healpixpad, ) - self.h = torch.zeros(1, 1, 1, 1) + self.h = th.zeros(1, 1, 1, 1) def forward(self, inputs: Sequence) -> Sequence: """Forward pass of the ConvGRUBlock @@ -90,17 +110,17 @@ def forward(self, inputs: Sequence) -> Sequence: Result of the forward pass """ if inputs.shape != self.h.shape: - self.h = torch.zeros_like(inputs) - combined = torch.cat([inputs, self.h], dim=1) + self.h = th.zeros_like(inputs) + combined = th.cat([inputs, self.h], dim=1) combined_conv = self.conv_gates(combined) - gamma, beta = torch.split(combined_conv, self.channels, dim=1) - reset_gate = torch.sigmoid(gamma) - update_gate = torch.sigmoid(beta) + gamma, beta = th.split(combined_conv, self.channels, dim=1) + reset_gate = th.sigmoid(gamma) + update_gate = th.sigmoid(beta) - combined = torch.cat([inputs, reset_gate * self.h], dim=1) + combined = th.cat([inputs, reset_gate * self.h], dim=1) cc_cnm = self.conv_can(combined) - cnm = torch.tanh(cc_cnm) + cnm = th.tanh(cc_cnm) h_next = (1 - update_gate) * self.h + update_gate * cnm self.h = h_next @@ -109,7 +129,7 @@ def forward(self, inputs: Sequence) -> Sequence: def reset(self): """Reset the update gates""" - self.h = torch.zeros_like(self.h) + self.h = th.zeros_like(self.h) # @@ -117,19 +137,19 @@ def reset(self): # -class BasicConvBlock(torch.nn.Module): +class BasicConvBlock(th.nn.Module): """Convolution block consisting of n subsequent convolutions and activations""" def __init__( self, - geometry_layer: torch.nn.Module = HEALPixLayer, + geometry_layer: th.nn.Module = HEALPixLayer, in_channels: int = 3, out_channels: int = 1, kernel_size: int = 3, dilation: int = 1, n_layers: int = 1, latent_channels: int = None, - activation: torch.nn.Module = None, + activation: th.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, ): @@ -175,7 +195,7 @@ def __init__( ) if activation is not None: convblock.append(activation) - self.convblock = torch.nn.Sequential(*convblock) + self.convblock = th.nn.Sequential(*convblock) def forward(self, x): """Forward pass of the BasicConvBlock @@ -193,14 +213,14 @@ def forward(self, x): return self.convblock(x) -class ConvNeXtBlock(torch.nn.Module): +class ConvNeXtBlock(th.nn.Module): """Class implementing a modified ConvNeXt network as described in https://arxiv.org/pdf/2201.03545.pdf and shown in figure 4 """ def __init__( self, - geometry_layer: torch.nn.Module = HEALPixLayer, + geometry_layer: th.nn.Module = HEALPixLayer, in_channels: int = 3, latent_channels: int = 1, out_channels: int = 1, @@ -208,7 +228,7 @@ def __init__( dilation: int = 1, n_layers: int = 1, # not used, but required for hydra instantiation upscale_factor: int = 4, - activation: torch.nn.Module = None, + activation: th.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, ): @@ -291,7 +311,7 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) - self.convblock = torch.nn.Sequential(*convblock) + self.convblock = th.nn.Sequential(*convblock) def forward(self, x): """Forward pass of the ConvNextBlock @@ -309,7 +329,7 @@ def forward(self, x): return self.skip_module(x) + self.convblock(x) -class DoubleConvNeXtBlock(torch.nn.Module): +class DoubleConvNeXtBlock(th.nn.Module): """Modification of ConvNeXtBlock block this time putting two sequentially in a single block with the number of channels in the middle being the number of latent channels @@ -317,7 +337,7 @@ class DoubleConvNeXtBlock(torch.nn.Module): def __init__( self, - geometry_layer: torch.nn.Module = HEALPixLayer, + geometry_layer: th.nn.Module = HEALPixLayer, in_channels: int = 3, out_channels: int = 1, kernel_size: int = 3, @@ -325,9 +345,12 @@ def __init__( n_layers: int = 1, # not used, but required for hydra instantiation upscale_factor: int = 4, latent_channels: int = 1, - activation: torch.nn.Module = None, + activation: th.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, + conditional_layer_norm: Callable = None, + conditional_layer_norm_once: bool = False, + dropout: float = 0.0, ): """ Parameters: @@ -355,6 +378,8 @@ def __init__( """ super().__init__() + self.cln_enabled = conditional_layer_norm is not None + if in_channels == int(latent_channels): self.skip_module1 = ( lambda x: x @@ -382,6 +407,25 @@ def __init__( enable_healpixpad=enable_healpixpad, ) + # check if we're applying a layer norm at the beginning + # we've got two ConvNeXt equivalent blocks in the layer, so have two entry points + # TODO: conditional layer norm once is doing two things, it's applying a norm on block entry + # and switch from conditional to non-conditional layer norm. This is not ideal and should be fixed once we determine + # what works best. + if conditional_layer_norm_once: + if conditional_layer_norm is not None: + # Conditional norm at the beginning of the block + self.entry_norm1 = conditional_layer_norm(channel_depth=in_channels) + self.entry_norm2 = conditional_layer_norm(channel_depth=latent_channels) + else: + # Regular layer normalization at the beginning of the block + self.entry_norm1 = _LayerNormOverChannels(channel_depth=in_channels) + self.entry_norm2 = _LayerNormOverChannels(channel_depth=latent_channels) + else: + # No normalization at the beginning + self.entry_norm1 = None + self.entry_norm2 = None + # 1st ConvNeXt block, the output of this one remains internal convblock1 = [] # 3x3 convolution establishing latent channels channels @@ -396,8 +440,17 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply batch norm and conditional layer norm if needed + if conditional_layer_norm is not None and not conditional_layer_norm_once: + cln = conditional_layer_norm(channel_depth=int(latent_channels)) + convblock1.append(cln) + if activation is not None: convblock1.append(activation) + if dropout > 0.0: + convblock1.append(th.nn.Dropout2d(p=dropout)) + # 1x1 convolution establishing increased channels convblock1.append( geometry_layer( @@ -410,8 +463,19 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply layer norm if needed + if conditional_layer_norm_once: + convblock1.append(_LayerNormOverChannels(channel_depth=int(latent_channels * upscale_factor))) + elif conditional_layer_norm is not None: + cln = conditional_layer_norm(channel_depth=int(latent_channels * upscale_factor)) + convblock1.append(cln) + if activation is not None: convblock1.append(activation) + if dropout > 0.0: + convblock1.append(th.nn.Dropout2d(p=dropout)) + # 1x1 convolution returning to latent channels convblock1.append( geometry_layer( @@ -426,7 +490,9 @@ def __init__( ) if activation is not None: convblock1.append(activation) - self.convblock1 = torch.nn.Sequential(*convblock1) + if dropout > 0.0: + convblock1.append(th.nn.Dropout2d(p=dropout)) + self.convblock1 = th.nn.ModuleList(convblock1) # 2nd ConNeXt block, takes the output of the first convnext block convblock2 = [] @@ -442,8 +508,16 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + # Apply batch norm and conditional layer norm if needed + if conditional_layer_norm is not None and not conditional_layer_norm_once: + cln = conditional_layer_norm(channel_depth=int(latent_channels)) + convblock2.append(cln) + if activation is not None: convblock2.append(activation) + if dropout > 0.0: + convblock2.append(th.nn.Dropout2d(p=dropout)) + # 1x1 convolution establishing increased channels convblock2.append( geometry_layer( @@ -456,8 +530,18 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + # Apply layer norm if needed + if conditional_layer_norm_once: + convblock2.append(_LayerNormOverChannels(channel_depth=int(latent_channels * upscale_factor))) + elif conditional_layer_norm is not None: + cln = conditional_layer_norm(channel_depth=int(latent_channels * upscale_factor)) + convblock2.append(cln) + if activation is not None: convblock2.append(activation) + if dropout > 0.0: + convblock2.append(th.nn.Dropout2d(p=dropout)) + # 1x1 convolution reducing to output channels convblock2.append( geometry_layer( @@ -472,9 +556,11 @@ def __init__( ) if activation is not None: convblock2.append(activation) - self.convblock2 = torch.nn.Sequential(*convblock2) + if dropout > 0.0: + convblock2.append(th.nn.Dropout2d(p=dropout)) + self.convblock2 = th.nn.ModuleList(convblock2) - def forward(self, x): + def forward(self, x, conditions_cln=None): """Forward pass of the DoubleConvNextBlock Parameters @@ -486,21 +572,56 @@ def forward(self, x): ------- torch.Tensor result of the forward pass + conditions_cln: torch.Tensor, optional + conditions for the conditional layer normalization """ - # internal convnext result - x1 = self.skip_module1(x) + self.convblock1(x) - # return second convnext result - return self.skip_module2(x1) + self.convblock2(x1) + # TODO: performance of skip connectioni hasn't been compared + # check cln(x) vs. cln(x1_residual + x) in the future + # save residual for the first block + x1_residual = self.skip_module1(x) -class Multi_SymmetricConvNeXtBlock(torch.nn.Module): + # entry norm for the first block + if self.entry_norm1 is not None: + if conditions_cln is not None: + x = self.entry_norm1(x, conditions=conditions_cln) + else: + x = self.entry_norm1(x) + + # internal convnext result + for layer in self.convblock1: + if isinstance(layer, ConditionalLayerNorm): + x = layer(x, conditions=conditions_cln) + else: + x = layer(x) + x1 = x1_residual + x + + # save residual for the second block + x2_residual = self.skip_module2(x1) + + # entry norm for the second block + if self.entry_norm2 is not None: + if conditions_cln is not None: + x1 = self.entry_norm2(x1, conditions=conditions_cln) + else: + x1 = self.entry_norm2(x1) + + # return second convnext result + for layer in self.convblock2: + if isinstance(layer, ConditionalLayerNorm): + x1 = layer(x1, conditions=conditions_cln) + else: + x1 = layer(x1) + return x2_residual + x1 + +class Multi_SymmetricConvNeXtBlock(th.nn.Module): """ - Class for creating multi-block SymmetricConvNeXtBlock. Defaults to all SymmetricConvNeXtBlocks having same parameters + Wrapper for SymmetricConvNeXtBlock that allows serial linking of blocks. """ def __init__( self, - geometry_layer: torch.nn.Module = HEALPixLayer, + geometry_layer: th.nn.Module = HEALPixLayer, in_channels: int = 3, latent_channels: int = 1, out_channels: int = 1, @@ -508,25 +629,36 @@ def __init__( dilation: int = 1, upscale_factor: int = 4, n_layers: int = 1, - activation: torch.nn.Module = None, + activation: th.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, + dropout: float = 0.0, + conditional_layer_norm: Callable = None, + conditional_layer_norm_once: bool = False, ): """ Parameters ---------- n_layers: int, optional The number of SymmetricConvNeXt Blocks + conditional_layer_norm: Callable, optional + Callable for physicsnemo.models.dlwp_healpix_layers.normalization.ConditionalLayerNorm. + Callable can be passed in by setting _partial_ to True in hydra config. If None, + conditional layer normalization is not applied. + conditional_layer_norm_once: bool, optional + Whether or not to apply conditional layer normalization only once. If True, + the conditional layer normalization is applied only once, otherwise it is applied + for each block. """ super().__init__() # Create a ModuleList to store complete blocks - self.blocks = torch.nn.ModuleList() + self.blocks = th.nn.ModuleList() + # flag for conditional layer normalization + self.cln_enabled = conditional_layer_norm is not None for i in range(n_layers): curr_in = in_channels if i == 0 else out_channels - - # Create a single block as a separate Module self.blocks.append( SymmetricConvNeXtBlock( geometry_layer=geometry_layer, @@ -539,17 +671,20 @@ def __init__( activation=activation, enable_nhwc=enable_nhwc, enable_healpixpad=enable_healpixpad, - ) + dropout=dropout, + conditional_layer_norm=conditional_layer_norm if conditional_layer_norm is not None else None, + conditional_layer_norm_once=conditional_layer_norm_once, + ), ) - def forward(self, x): + def forward(self, x, conditions_cln=None): out = x for block in self.blocks: - out = block(out) + out = block(out, conditions_cln=conditions_cln) return out -class SymmetricConvNeXtBlock(torch.nn.Module): +class SymmetricConvNeXtBlock(th.nn.Module): """Another modification of ConvNeXtBlock block this time using 4 layers and adding a layer that instead of going from in_channels to latent*upscale channesl goes to latent channels first @@ -557,7 +692,7 @@ class SymmetricConvNeXtBlock(torch.nn.Module): def __init__( self, - geometry_layer: torch.nn.Module = HEALPixLayer, + geometry_layer: th.nn.Module = HEALPixLayer, in_channels: int = 3, latent_channels: int = 1, out_channels: int = 1, @@ -565,9 +700,13 @@ def __init__( dilation: int = 1, n_layers: int = 1, # not used, but required for hydra instantiation upscale_factor: int = 4, - activation: torch.nn.Module = None, + activation: th.nn.Module = None, enable_nhwc: bool = False, + use_block_skip_connection: bool = True, enable_healpixpad: bool = False, + dropout: float = 0.0, + conditional_layer_norm: th.nn.Module = None, + conditional_layer_norm_once: bool = False, ): """ Parameters @@ -592,27 +731,61 @@ def __init__( Enable nhwc format, passed to wrapper enable_healpixpad: bool, optional If HEALPixPadding should be enabled, passed to wrapper + use_block_skip_connection: bool, optional + Whether or not to use block-level skip connection + dropout: float, optional + Dropout probability to apply after the first convolution + conditional_layer_norm: th.nn.Module, optional + conditional layer normalization. If None, + no conditional layer normalization is applied. + conditional_layer_norm_once: bool, optional + Whether or not to apply conditional layer normalization only once. If True, + the conditional layer normalization is applied only once, otherwise it is applied + for each block. """ + super().__init__() - if in_channels == int(latent_channels): - self.skip_module = lambda x: x # Identity-function required in forward pass + self.use_block_skip_connection = use_block_skip_connection + self.activation = activation + self.dropout = dropout > 0.0 + self.cln_enabled = conditional_layer_norm is not None + + if use_block_skip_connection: + if in_channels == int(out_channels): + self.skip_module = lambda x: x + else: + self.skip_module = geometry_layer( + layer=th.nn.Conv2d, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + enable_nhwc=enable_nhwc, + enable_healpixpad=enable_healpixpad, + ) + + # check if we're applying a layer norm at the beginning + # TODO: conditional layer norm once is doing two things, it's applying a norm on block entry + # and switch from conditional to non-conditional layer norm. This is not ideal and should be fixed once we determine + # what works best. + if conditional_layer_norm_once: + if conditional_layer_norm is not None: + # Conditional norm at the beginning of the block + self.entry_norm = conditional_layer_norm(channel_depth=in_channels) + else: + # Regular layer normalization at the beginning of the block + self.entry_norm = _LayerNormOverChannels(channel_depth=in_channels) else: - self.skip_module = geometry_layer( - layer=torch.nn.Conv2d, - in_channels=in_channels, - out_channels=out_channels, - kernel_size=1, - enable_nhwc=enable_nhwc, - enable_healpixpad=enable_healpixpad, - ) + # No normalization at the beginning + self.entry_norm = None - # 1st ConvNeXt block, the output of this one remains internal + # Collect conv->norm->activation->dropout operations in list for sequential execution convblock = [] - # 3x3 convolution establishing latent channels channels + + # 3x3: in → latent convblock.append( geometry_layer( - layer=torch.nn.Conv2d, + layer=th.nn.Conv2d, in_channels=in_channels, out_channels=int(latent_channels), kernel_size=kernel_size, @@ -621,12 +794,20 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + # Apply layer norm if needed + if conditional_layer_norm is not None and not conditional_layer_norm_once: + cln = conditional_layer_norm(channel_depth=int(latent_channels)) + convblock.append(cln) + if activation is not None: convblock.append(activation) - # 1x1 convolution establishing increased channels + if dropout > 0.0: + convblock.append(th.nn.Dropout2d(p=dropout)) + + # 1x1: latent → latent * upscale convblock.append( geometry_layer( - layer=torch.nn.Conv2d, + layer=th.nn.Conv2d, in_channels=int(latent_channels), out_channels=int(latent_channels * upscale_factor), kernel_size=1, @@ -635,12 +816,24 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply layer norm if needed + if conditional_layer_norm_once: + convblock.append(_LayerNormOverChannels(channel_depth=int(latent_channels * upscale_factor))) + elif conditional_layer_norm is not None: + cln = conditional_layer_norm(channel_depth=int(latent_channels * upscale_factor)) + convblock.append(cln) + if activation is not None: convblock.append(activation) - # 1x1 convolution returning to latent channels + if dropout > 0.0: + convblock.append(th.nn.Dropout2d(p=dropout)) + + + # 1x1: upscale → latent convblock.append( geometry_layer( - layer=torch.nn.Conv2d, + layer=th.nn.Conv2d, in_channels=int(latent_channels * upscale_factor), out_channels=int(latent_channels), kernel_size=1, @@ -649,14 +842,25 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply layer norm if needed + if conditional_layer_norm_once: + convblock.append(_LayerNormOverChannels(channel_depth=int(latent_channels))) + elif conditional_layer_norm is not None: + cln = conditional_layer_norm(channel_depth=int(latent_channels)) + convblock.append(cln) + if activation is not None: convblock.append(activation) - # 3x3 convolution from latent channels to latent channels + if dropout > 0.0: + convblock.append(th.nn.Dropout2d(p=dropout)) + + # 3x3: latent → out (no norm on this one, following convnext) convblock.append( geometry_layer( - layer=torch.nn.Conv2d, + layer=th.nn.Conv2d, in_channels=int(latent_channels), - out_channels=out_channels, # int(latent_channels), + out_channels=out_channels, kernel_size=kernel_size, dilation=dilation, enable_nhwc=enable_nhwc, @@ -665,23 +869,45 @@ def __init__( ) if activation is not None: convblock.append(activation) - self.convblock = torch.nn.Sequential(*convblock) + if dropout > 0.0: + convblock.append(th.nn.Dropout2d(p=dropout)) + + self.convblock = th.nn.ModuleList(convblock) - def forward(self, x): - """Forward pass of the SymmetricConvNextBlock + def forward(self, x, conditions_cln=None): + """ + Forward pass of the SymmetricConvNextBlock, broken into steps for support of conditional layer normalization Parameters ---------- x: torch.Tensor inputs to the forward pass - + conditions_cln: torch.Tensor, optional + Condition for the conditional layer normalization, if applicable Returns ------- torch.Tensor result of the forward pass """ - # residual connection with reshaped inpute and output of conv block - return self.skip_module(x) + self.convblock(x) + + # TODO: performance of skip connectioni hasn't been compared + # check cln(x) vs. cln(x1_residual + x) in the future + # Save residual + residual = self.skip_module(x) if self.use_block_skip_connection else 0 + + if self.entry_norm is not None: + if conditions_cln is not None: + x = self.entry_norm(x, conditions=conditions_cln) + else: + x = self.entry_norm(x) + + for layer in self.convblock: + if isinstance(layer, ConditionalLayerNorm): + x = layer(x, conditions=conditions_cln) + else: + x = layer(x) + + return x + residual # @@ -689,18 +915,18 @@ def forward(self, x): # -class TransposedConvUpsample(torch.nn.Module): +class TransposedConvUpsample(th.nn.Module): """This class provides a wrapper for a HEALPix (or other) tensor data around the torch.nn.ConvTranspose2d class. """ def __init__( self, - geometry_layer: torch.nn.Module = HEALPixLayer, + geometry_layer: th.nn.Module = HEALPixLayer, in_channels: int = 3, out_channels: int = 1, upsampling: int = 2, - activation: torch.nn.Module = None, + activation: th.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, ): @@ -708,7 +934,7 @@ def __init__( Parameters ---------- geometry_layer: torch.nn.Module, optional - The wrapper for the geometry of the tensor being bassed to ConvTranspose2d + The wrapper for the geometry of the tensor being bassed to MaxPool2d in_channels: int, optional The number of input channels out_channels: int, optional @@ -739,7 +965,7 @@ def __init__( ) if activation is not None: upsampler.append(activation) - self.upsampler = torch.nn.Sequential(*upsampler) + self.upsampler = th.nn.Sequential(*upsampler) def forward(self, x): """Forward pass of the TransposedConvUpsample layer @@ -762,7 +988,7 @@ def forward(self, x): # -class Interpolate(torch.nn.Module): +class Interpolate(th.nn.Module): """Helper class that handles interpolation This is done as a class so that scale and mode can be stored """ @@ -777,7 +1003,7 @@ def __init__(self, scale_factor: Union[int, Tuple], mode: str = "nearest"): Interpolation mode used for upsampling, passed to torch.nn.functional.interpolate """ super().__init__() - self.interp = torch.nn.functional.interpolate + self.interp = th.nn.functional.interpolate self.scale_factor = scale_factor self.mode = mode diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py new file mode 100644 index 0000000000..d184147d8f --- /dev/null +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py @@ -0,0 +1,51 @@ +import numpy as np +import torch +import xarray as xr + +class NonnegativeConstraint(torch.nn.Module): + def __init__( + self, + variables: list[str], + channels: list[str], + scaling: dict[str, dict[str, float]], + ): + """ + Parameters + ---------- + variables: list[str] + List of variable names to apply the constraint to. + channels: list[str] + List of all input channel names in the model. + scaling: dict[str, dict[str, float]] + Dictionary containing the mean and std for each variable. + """ + super().__init__() + self.variables = variables + self.channels = channels + self.scaling = scaling + + # Only apply constraint to variables that are used by model + self.variables = [var for var in self.variables if var in channels] + + var_indices = torch.tensor( + [channels.index(var) for var in self.variables], + dtype=torch.long + ) + self.register_buffer('var_indices', var_indices, persistent=False) + + self.var_means = torch.tensor([scaling[var]['mean'] for var in self.variables]) + self.var_stds = torch.tensor([scaling[var]['std'] for var in self.variables]) + + thresholds = (0. - self.var_means) / self.var_stds + thresholds = thresholds.view(1, 1, 1, -1, 1, 1) + self.register_buffer('thresholds', thresholds, persistent=False) + + def forward(self, x): + ''' + Tensors are expected to be in the shape [B, F, T, C, H, W] + ''' + selected_vars = torch.index_select(x, dim=3, index=self.var_indices) + clamped = torch.maximum(selected_vars, self.thresholds).to(x.dtype) + x.index_copy_(3, self.var_indices, clamped) + + return x diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py index b3b20c8e11..cadebb6689 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py @@ -19,6 +19,7 @@ import torch from hydra.utils import instantiate from omegaconf import DictConfig +from torch.utils.checkpoint import checkpoint class UNetDecoder(torch.nn.Module): @@ -36,6 +37,8 @@ def __init__( dilations: list = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, + per_level_cln: list[bool] = None, + per_level_checkpointing: list[bool] = None, ): """ Parameters @@ -61,17 +64,43 @@ def __init__( If channel last format should be used enable_healpixpad, bool, optional If the healpixpad library should be used if installed + per_level_cln: list[bool], optional + If the CLN should be applied to each level of the decoder + If None, the CLN will based on the conv_block.conditional_layer_norm attribute + per_level_checkpointing: list[bool], optional + If the checkpointing should be applied to each level of the decoder + If None, the checkpointing will not be applied """ super().__init__() - self.channel_dim = 1 # 1 in previous layout + self.channel_dim = 1 + + if per_level_cln is not None and len(per_level_cln) != len(n_channels): + raise ValueError( + "per_level_cln must be a list of booleans of the same length as n_channels" + f"Got {len(per_level_cln)} for per_level_cln and {len(n_channels)} for n_channels" + ) + per_level_cln = ( + per_level_cln if per_level_cln is not None else [True] * len(n_channels) + ) + + if per_level_checkpointing is not None and len(per_level_checkpointing) != len( + n_channels + ): + raise ValueError( + "per_level_checkpointing must be a list of booleans of the same length as n_channels" + f"Got {len(per_level_checkpointing)} for per_level_checkpointing and {len(n_channels)} for n_channels" + ) + self.per_level_checkpointing = ( + per_level_checkpointing + if per_level_checkpointing is not None + else [False] * len(n_channels) + ) if dilations is None: - # Defaults to [1, 1, 1...] in accordance with the number of unet levels dilations = [1 for _ in range(len(n_channels))] self.decoder = [] for n, curr_channel in enumerate(n_channels): - # Second half of the synoptic layer does not need an upsampling module if n == 0: up_sample_module = None else: @@ -87,11 +116,17 @@ def __init__( n_channels[n + 1] if n < len(n_channels) - 1 else n_channels[-1] ) + block_config = conv_block.copy() + if ( + "conditional_layer_norm" in block_config + and block_config.conditional_layer_norm is not None + ): + if not per_level_cln[n]: + block_config.conditional_layer_norm = None + conv_module = instantiate( - config=conv_block, - in_channels=curr_channel * 2 - if n > 0 - else curr_channel, # Considering skip connection + config=block_config, + in_channels=curr_channel * 2 if n > 0 else curr_channel, latent_channels=curr_channel, out_channels=next_channel, dilation=dilations[n], @@ -100,7 +135,6 @@ def __init__( enable_healpixpad=enable_healpixpad, ) - # Recurrent module if recurrent_block is not None: rec_module = instantiate( config=recurrent_block, @@ -121,8 +155,6 @@ def __init__( ) self.decoder = torch.nn.ModuleList(self.decoder) - - # (Linear) Output layer self.output_layer = instantiate( config=output_layer, in_channels=curr_channel, @@ -132,7 +164,31 @@ def __init__( enable_healpixpad=enable_healpixpad, ) - def forward(self, inputs: Sequence) -> torch.Tensor: + def _forward_layer_pass( + self, + layer: torch.nn.Module, + x: torch.Tensor, + skip_connection: torch.Tensor = None, + conditions_cln: torch.Tensor = None, + ) -> torch.Tensor: + if layer["upsamp"] is not None: + up = layer["upsamp"](x) + x = torch.cat([up, skip_connection], dim=self.channel_dim) + if hasattr(layer["conv"], "cln_enabled") and layer["conv"].cln_enabled: + if conditions_cln is not None: + x = layer["conv"](x, conditions_cln=conditions_cln) + else: + raise ValueError( + "Conditional inputs are required for layers with cln_enabled=True" + ) + else: + x = layer["conv"](x) + + return x + + def forward( + self, inputs: Sequence, conditions_cln: Sequence = None + ) -> torch.Tensor: """ Forward pass of the HEALPix Unet decoder @@ -140,6 +196,8 @@ def forward(self, inputs: Sequence) -> torch.Tensor: ---------- inputs: Sequence The inputs to decode + conditions_cln: Sequence, optional + The conditional inputs for the normalization layers. Returns ------- @@ -147,12 +205,24 @@ def forward(self, inputs: Sequence) -> torch.Tensor: """ x = inputs[-1] for n, layer in enumerate(self.decoder): - if layer["upsamp"] is not None: - up = layer["upsamp"](x) - x = torch.cat([up, inputs[-1 - n]], dim=self.channel_dim) - x = layer["conv"](x) + skip_connection = inputs[-1 - n] if layer["upsamp"] is not None else None + if self.per_level_checkpointing[n]: + x = checkpoint( + self._forward_layer_pass, + layer, + x, + skip_connection, + conditions_cln, + use_reentrant=False, + ) + else: + x = self._forward_layer_pass( + layer, x, skip_connection, conditions_cln + ) + if layer["recurrent"] is not None: x = layer["recurrent"](x) + return self.output_layer(x) def reset(self): diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py index cc7c9f587a..8fe0749b44 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py @@ -19,6 +19,7 @@ import torch from hydra.utils import instantiate from omegaconf import DictConfig +from torch.utils.checkpoint import checkpoint class UNetEncoder(torch.nn.Module): @@ -35,6 +36,8 @@ def __init__( dilations: list = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, + per_level_cln: Sequence[bool] = None, + per_level_checkpointing: Sequence[bool] = None, ): """ Parameters @@ -58,15 +61,41 @@ def __init__( If channel last format should be used enable_healpixpad, bool, optional If the healpixpad library should be used (if installed) + per_level_cln: list[bool] | None, optional + If the CLN should be applied to each level of the encoder + If None, the CLN will based on the conv_block.conditional_layer_norm attribute + per_level_checkpointing: list[bool] | None, optional + If the checkpointing should be applied to each level of the encoder + If None, the checkpointing will not be applied """ super().__init__() self.n_channels = n_channels + if per_level_cln is not None and len(per_level_cln) != len(n_channels): + raise ValueError( + "per_level_cln must be a list of booleans of the same length as n_channels" + f"Got {len(per_level_cln)} for per_level_cln and {len(n_channels)} for n_channels" + ) + per_level_cln = ( + per_level_cln if per_level_cln is not None else [True] * len(n_channels) + ) + + if per_level_checkpointing is not None and len(per_level_checkpointing) != len( + n_channels + ): + raise ValueError( + "per_level_checkpointing must be a list of booleans of the same length as n_channels" + f"Got {len(per_level_checkpointing)} for per_level_checkpointing and {len(n_channels)} for n_channels" + ) + self.per_level_checkpointing = ( + per_level_checkpointing + if per_level_checkpointing is not None + else [False] * len(n_channels) + ) + if dilations is None: - # Defaults to [1, 1, 1...] in accordance with the number of unet levels dilations = [1 for _ in range(len(n_channels))] - # Build encoder old_channels = input_channels self.encoder = [] for n, curr_channel in enumerate(n_channels): @@ -80,9 +109,17 @@ def __init__( ) ) + block_config = conv_block.copy() + if ( + "conditional_layer_norm" in block_config + and block_config.conditional_layer_norm is not None + ): + if not per_level_cln[n]: + block_config.conditional_layer_norm = None + modules.append( instantiate( - config=conv_block, + config=block_config, in_channels=old_channels, latent_channels=curr_channel, out_channels=curr_channel, @@ -98,22 +135,58 @@ def __init__( self.encoder = torch.nn.ModuleList(self.encoder) - def forward(self, inputs: Sequence) -> Sequence: + def _forward_layer_pass( + self, + layer_group: torch.nn.Module, + inp: torch.Tensor, + conditions_cln: torch.Tensor = None, + ) -> torch.Tensor: + interim_output = inp + for layer in layer_group: + if getattr(layer, "cln_enabled", False): + if conditions_cln is None: + raise ValueError( + "Conditional inputs are required for layers with cln_enabled=True" + ) + interim_output = layer(interim_output, conditions_cln=conditions_cln) + else: + interim_output = layer(interim_output) + + return interim_output + + def forward( + self, inputs: Sequence, conditions_cln: torch.Tensor = None + ) -> Sequence: """ Forward pass of the HEALPix Unet encoder Parameters ---------- inputs: Sequence - The inputs to enccode + The inputs to encode + conditions_cln: torch.Tensor, optional + The conditional inputs for the normalization layers. Returns ------- Sequence: The encoded values """ outputs = [] - for layer in self.encoder: - outputs.append(layer(inputs)) + for n, layer_group in enumerate(self.encoder): + interim_output = inputs + if self.per_level_checkpointing[n]: + interim_output = checkpoint( + self._forward_layer_pass, + layer_group, + interim_output, + conditions_cln, + use_reentrant=False, + ) + else: + interim_output = self._forward_layer_pass( + layer_group, interim_output, conditions_cln + ) + outputs.append(interim_output) inputs = outputs[-1] return outputs diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py new file mode 100644 index 0000000000..9461bd37d5 --- /dev/null +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Sequence, List + +import torch as th +from omegaconf import DictConfig + +try: + from apex.normalization import FusedLayerNorm + _APEX_AVAILABLE = True +except ImportError: + _APEX_AVAILABLE = False + + +@th.compile +def _cln_affine(x_norm, gamma_raw, beta, scale_center, n_faces): + """Fused affine transform: expand gamma/beta across faces and apply to normalized input.""" + C = gamma_raw.shape[-1] + gamma = (scale_center + gamma_raw).unsqueeze(1).expand(-1, n_faces, -1).reshape(-1, 1, 1, C) + beta = beta.unsqueeze(1).expand(-1, n_faces, -1).reshape(-1, 1, 1, C) + return gamma * x_norm + beta + + +class ConditionalLayerNorm(th.nn.Module): + def __init__( + self, + condition_shape: int, + channel_depth: int, + mlp_hidden_dims: List[int] = [128, 128], + activation: th.nn.Module = None, + eps: float = 1e-5, + n_faces: int = 12, + norm_op:str = "torch", + init_cln_to_zero: bool = False, + scale_center: float = 0.0, + ): + """ + Conditional LayerNorm with MLP-based conditioning. + + Parameters + ---------- + condition_shape : int + Shape of the conditioning input. + channel_depth : int + Number of channels in the input tensor. + mlp_hidden_dims : List[int] + Hidden layer sizes for MLPs predicting gamma and beta. + activation : DictConfig + Activation function configuration for the MLPs. + eps : float + Numerical stability constant. + n_faces : int + Number of faces in the Healpix grid, used for reshaping. + norm_op : str + "torch" for torch.nn.LayerNorm, "apex" for apex FusedLayerNorm. + init_cln_to_zero : bool = False + If True, initialize the last layer of the MLPs to zero. + At the start of training, the noise will be ignored + scale_center : float = 0.0 + Center of the scale parameter. Set to 1.0 and use `init_cln_to_zero=True` + to make CLN behave like standard LayerNorm at initialization. + """ + super().__init__() + self.eps = eps + self.condition_shape = condition_shape + self.channel_depth = channel_depth + self.hidden_dims = mlp_hidden_dims + self.activation = activation if activation is not None else th.nn.Identity() + self.gamma_beta_mlp = self._make_mlp( + self.condition_shape, + [2 * h for h in self.hidden_dims], + 2 * self.channel_depth, + self.activation, + ) + self.n_faces = n_faces + self.scale_center = scale_center + + if init_cln_to_zero: + self.gamma_beta_mlp[-1].weight.data.zero_() + self.gamma_beta_mlp[-1].bias.data.zero_() + + if norm_op == "torch": + self.norm = th.nn.LayerNorm(channel_depth, elementwise_affine=False) + elif norm_op == "apex": + if not _APEX_AVAILABLE: + raise ImportError("Apex FusedLayerNorm requested but apex is not available, please install it from https://github.com/NVIDIA/apex") + self.norm = FusedLayerNorm(channel_depth, elementwise_affine=False) + + def _make_mlp(self, in_dim: int, hidden_dims: List[int], out_dim: int, activation: th.nn.Module) -> th.nn.Sequential: + + layers = [] + for hdim in hidden_dims: + layers.append(th.nn.Linear(in_dim, hdim)) + if activation: + layers.append(activation) + in_dim = hdim + layers.append(th.nn.Linear(in_dim, out_dim)) + return th.nn.Sequential(*layers) + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + """Backward compatibility: merge old separate gamma_mlp/beta_mlp into fused gamma_beta_mlp. + + Old MLPs had hidden_dims [h1, h2, ...] and output dim C. + New fused MLP has hidden_dims [2*h1, 2*h2, ...] and output dim 2*C. + + For the first Linear (condition_shape → 2*h1), we vertically concatenate: + new_weight = cat([gamma_weight, beta_weight], dim=0) + + For subsequent Linear layers (2*h_i → 2*h_{i+1} or 2*h_last → 2*C), + we build a block-diagonal weight matrix: + new_weight = [[gamma_weight, 0 ], + [0, beta_weight]] + + Biases are always concatenated: cat([gamma_bias, beta_bias]). + """ + gamma_prefix = prefix + "gamma_mlp." + beta_prefix = prefix + "beta_mlp." + fused_prefix = prefix + "gamma_beta_mlp." + + has_old_keys = any(k.startswith(gamma_prefix) for k in state_dict) + + if has_old_keys: + # Collect all Linear layer indices from the old gamma MLP + gamma_layer_indices = set() + for k in state_dict: + if k.startswith(gamma_prefix): + layer_key = k[len(gamma_prefix):] + parts = layer_key.split(".") + if len(parts) == 2 and parts[1] in ("weight", "bias"): + gamma_layer_indices.add(int(parts[0])) + first_layer_idx = min(gamma_layer_indices) + + keys_to_remove = [] + keys_to_add = {} + + for k in list(state_dict.keys()): + if k.startswith(gamma_prefix): + layer_key = k[len(gamma_prefix):] # e.g. "0.weight" + parts = layer_key.split(".") + if len(parts) != 2 or parts[1] not in ("weight", "bias"): + continue + idx = int(parts[0]) + param_type = parts[1] + fused_key = fused_prefix + layer_key + beta_key = beta_prefix + layer_key + + if beta_key not in state_dict: + continue + + gamma_val = state_dict[k] + beta_val = state_dict[beta_key] + + if param_type == "bias": + # Biases are always concatenated + keys_to_add[fused_key] = th.cat([gamma_val, beta_val], dim=0) + elif idx == first_layer_idx: + # First layer: input dim is shared (condition_shape), + # just concatenate along output dim + keys_to_add[fused_key] = th.cat([gamma_val, beta_val], dim=0) + else: + # Hidden→hidden or hidden→output: block-diagonal + # gamma_val: (out_old, in_old), beta_val: (out_old, in_old) + # result: (2*out_old, 2*in_old) + out_old, in_old = gamma_val.shape + zeros = th.zeros(out_old, in_old, dtype=gamma_val.dtype, device=gamma_val.device) + keys_to_add[fused_key] = th.cat([ + th.cat([gamma_val, zeros], dim=1), + th.cat([zeros, beta_val], dim=1), + ], dim=0) + + keys_to_remove.append(k) + if beta_key not in keys_to_remove: + keys_to_remove.append(beta_key) + + for k in keys_to_remove: + if k in state_dict: + del state_dict[k] + state_dict.update(keys_to_add) + + super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) + + def forward(self, x: th.Tensor, conditions: th.Tensor) -> th.Tensor: + """ + Parameters + ---------- + x : th.Tensor + Input tensor of shape: (B, C, H, W) + conditions : th.Tensor + Conditioning tensor of shape (B*n_cond, cond_dim) + + Returns + ------- + th.Tensor + Normalized and conditioned tensor of shape: (B, C, H, W) + """ + + is_channels_last = x.is_contiguous(memory_format=th.channels_last) + + # LayerNorm on last dim: permute to (B, H, W, C) + x_nhwc = x.permute(0, 2, 3, 1) + if not is_channels_last: + x_nhwc = x_nhwc.contiguous() + x_norm = self.norm(x_nhwc) + + # Fused gamma/beta MLP: single forward pass, then split + gamma_beta = self.gamma_beta_mlp(conditions) # (B*n_cond, 2*C) + gamma_raw, beta = gamma_beta.chunk(2, dim=-1) # each (B*n_cond, C) + + # Fused affine: expand across faces + scale_center + multiply + add + result = _cln_affine(x_norm, gamma_raw, beta, self.scale_center, self.n_faces) + + # Return to NCHW logical layout, preserving channels_last memory format if input was + if is_channels_last: + return result.permute(0, 3, 1, 2).contiguous(memory_format=th.channels_last) + else: + return result.permute(0, 3, 1, 2) diff --git a/test/datapipes/test_healpix_couple_zarr.py b/test/datapipes/test_healpix_couple_zarr.py new file mode 100644 index 0000000000..61aaca854a --- /dev/null +++ b/test/datapipes/test_healpix_couple_zarr.py @@ -0,0 +1,1230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random +import warnings +from dataclasses import dataclass +from pathlib import Path + +import pytest +import torch as th +from pytest_utils import import_or_fail, nfsdata_or_fail +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from physicsnemo.distributed import DistributedManager + +omegaconf = pytest.importorskip("omegaconf") +np = pytest.importorskip("numpy") +pd = pytest.importorskip("pandas") +xr = pytest.importorskip("xarray") +zarr = pytest.importorskip("zarr") + + +@pytest.fixture +def dataset_path(): + data_dir = "/data/nfs/modulus-data/datasets/healpix/" + dataset_name = "healpix.zarr" + path = Path(data_dir, dataset_name) + return path + + +@pytest.fixture +def splits(): + split_dict = { + "train_date_start": "1979-01-01", + "train_date_end": "1979-01-01T21:00", + "val_date_start": "1979-01-02", + "val_date_end": "1979-01-02T09:00", + "test_date_start": "1979-01-02T12:00", + "test_date_end": "1979-01-02T18:00", + } + return split_dict + + +@pytest.fixture +def constant_coupler_config(): + constant_coupler = [ + { + "coupler": "ConstantCoupler", + "params": { + "batch_size": 1, + "variables": ["z250"], + "input_times": ["0h"], + "input_time_dim": 1, + "output_time_dim": 1, + "presteps": 0, + "prepared_coupled_data": True, + }, + } + ] + return constant_coupler + + +@pytest.fixture +def average_coupler_config(): + average_coupler = [ + { + "coupler": "ConstantCoupler", + "params": { + "batch_size": 1, + "variables": ["z250"], + "input_times": ["0h"], + "input_time_dim": 1, + "output_time_dim": 1, + "presteps": 0, + "prepared_coupled_data": True, + }, + } + ] + return average_coupler + + +@dataclass +class coupler_helper: + """helper class for setting up the couplers""" + + output_variables: list + time_step: str + + +@pytest.fixture +def scaling_dict(): + scaling = { + "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, + "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, + "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, + "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, + "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, + "z1000-12h": {"mean": 952.1435546875, "std": 895.7516479492188}, + "z250": {"mean": 101186.28125, "std": 5551.77978515625}, + "z500": {"mean": 55625.9609375, "std": 2681.712890625}, + "lsm": {"mean": 0, "std": 1}, + "z": {"mean": 0, "std": 1}, + "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, + "extra": {"mean": 0, "std": 0}, # doesn't appear in test dataset + } + return omegaconf.DictConfig(scaling) + + +@pytest.fixture +def scaling_double_dict(): + scaling = { + "t2m0": {"mean": 0, "std": 2}, + "t850": {"mean": 0, "std": 2}, + "tau300-700": {"mean": 0, "std": 2}, + "tcwv0": {"mean": 0, "std": 2}, + "z1000": {"mean": 0, "std": 2}, + "z1000-12h": {"mean": 0, "std": 2}, + "z250": {"mean": 0, "std": 2}, + "z500": {"mean": 0, "std": 2}, + "lsm": {"mean": 0, "std": 2}, + "z": {"mean": 0, "std": 2}, + "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, + "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset + } + return omegaconf.DictConfig(scaling) + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("pandas") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): + + from physicsnemo.datapipes.healpix.couplers import ( + ConstantCoupler, + ) + + variables = ["z500", "z1000"] + input_times = ["0h"] + input_time_dim = 1 + output_time_dim = 1 + presteps = 0 + batch_size = 2 + batch = {"time": slice(0, 2)} + + # open our test dataset + zarr_ds = zarr.open(dataset_path) + input_indices = [ + int(np.where(zarr_ds.channel_in[:] == ch)[0][0]) for ch in variables + ] + + # test fail initialization + with pytest.raises( + NotImplementedError, match=("Data preparation not yet implemented") + ): + coupler = ConstantCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + presteps=presteps, + input_times=input_times, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + prepared_coupled_data=False, + ) + + coupler = ConstantCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + presteps=presteps, + input_times=input_times, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ) + assert isinstance(coupler, ConstantCoupler) + + # check setting coupled variable indices + mock_coupled_module = coupler_helper( + output_variables=["not_coupled", "z500"], + time_step="0h", + ) + with pytest.raises(ValueError, match=("Missing variables in coupled module")): + coupler.setup_coupling(mock_coupled_module) + + mock_coupled_module.output_variables = ["z500", "z1000"] + coupler.setup_coupling(mock_coupled_module) + assert coupler.coupled_channel_indices == [0, 1] + + interval = 2 + data_time_step = "3h" + coupler.compute_coupled_indices(interval, data_time_step) + coupled_integration_dim = presteps + max(output_time_dim // input_time_dim, 1) + expected = np.empty([batch_size, coupled_integration_dim, len(input_times)]) + for b in range(batch_size): + for i in range(coupled_integration_dim): + expected[b, i, :] = b + np.array( + [pd.Timedelta(ts) / pd.Timedelta(data_time_step) for ts in input_times] + ) + expected = expected.astype(int) + assert np.array_equal(expected, coupler._coupled_offsets) + + scaling_df = pd.DataFrame.from_dict(omegaconf.OmegaConf.to_object(scaling_dict)).T + scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} + scaling_da = scaling_df.to_xarray().astype("float32") + coupler.set_scaling(scaling_da) + coupled_scaling = scaling_da.sel(index=variables).rename({"index": "channel_in"}) + expected = np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)) + assert np.array_equal(expected, coupler.coupled_scaling["mean"]) + expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) + assert np.array_equal(expected, coupler.coupled_scaling["std"]) + + coupled_fields_batch_size = batch_size + coupled_fields_timedim = 4 + expected_shape = [ + coupler.coupled_integration_dim, + coupled_fields_batch_size, + coupler.timevar_dim, + ] + list(coupler.spatial_dims) + coupled_fields = th.rand( + coupled_fields_batch_size, + coupler.spatial_dims[0], + coupled_fields_timedim, + len(coupler.coupled_channel_indices), + coupler.spatial_dims[1], + coupler.spatial_dims[2], + ) + coupler.set_coupled_fields(coupled_fields) + assert coupler.coupled_mode + assert list(coupler.construct_integrated_couplings().shape) == expected_shape + + # verify that the data is being properly transformed + expected = coupled_fields[:, :, :, coupler.coupled_channel_indices, :, :].permute( + 2, 0, 3, 1, 4, 5 + ) + expected = expected[0, :, -1:, :, :, :] + expected = expected.unsqueeze(0) + expected = expected.repeat( + coupler.coupled_integration_dim, 1, coupled_fields_batch_size, 1, 1, 1 + ) + result = coupler.construct_integrated_couplings() + assert th.equal(expected, coupler.construct_integrated_couplings()) + + # verify that dimensions aren't reordered when time_first is false + coupler.time_first = False + coupler.set_coupled_fields(coupled_fields) + # [T, B, C, F, H, W] + expected = expected.permute(1, 3, 0, 2, 4, 5) + result = coupler.construct_integrated_couplings() + assert th.equal(expected, result) + coupler.time_first = True + + # test coupler reset + coupler.reset_coupler() + assert coupler.coupled_mode is False + + # test loading from the dataset + coupled_scaling = { + "mean": np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)), + "std": np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)), + } + expected = zarr_ds.inputs[:2][:, input_indices] + expected = (expected - coupled_scaling["mean"]) / coupled_scaling["std"] + coupled_field = coupler.construct_integrated_couplings( + batch=batch, bsize=batch_size + ) + assert np.array_equal(expected, coupled_field[0]) + + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("pandas") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): + + from physicsnemo.datapipes.healpix.couplers import ( + TrailingAverageCoupler, + ) + + variables = ["z500", "z1000-12h"] + input_times = ["6h", "12h"] + input_time_dim = 2 + output_time_dim = 2 + presteps = 0 + batch_size = 2 + averaging_window = "6h" + # open our test dataset + zarr_ds = zarr.open(dataset_path) + + # test fail initialization when trying to prepare data + with pytest.raises( + NotImplementedError, match=("Data preparation not yet implemented") + ): + coupler = TrailingAverageCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + presteps=presteps, + averaging_window=averaging_window, + input_times=input_times, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + prepared_coupled_data=False, + ) + + # test fail when input times aren't evenly divisible by dataset dt + with pytest.raises(ValueError, match=("Coupled input times")): + coupler = TrailingAverageCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + presteps=presteps, + averaging_window=averaging_window, + input_times=["30m"], + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ) + + coupler = TrailingAverageCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + presteps=presteps, + averaging_window=averaging_window, + input_times=input_times, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ) + assert isinstance(coupler, TrailingAverageCoupler) + + mock_coupled_module = coupler_helper( + output_variables=["not_coupled", "z500"], + time_step="3h", + ) + with pytest.raises(ValueError, match=("Missing variables in coupled module")): + coupler.setup_coupling(mock_coupled_module) + + # veryify averaging slices computed correctly + mock_coupled_module = coupler_helper( + output_variables=["z500", "z1000"], + time_step="3h", + ) + coupler.setup_coupling(mock_coupled_module) + averaging_window_max_indices = [ + i // pd.Timedelta(mock_coupled_module.time_step) for i in input_times + ] + dt = averaging_window_max_indices[0] + # assumes only 1 integration step, otherwise would be wrong + expected_slices = [[]] + for i, window_end in enumerate(averaging_window_max_indices): + expected_slices[0].append(slice(i * dt, window_end)) + assert expected_slices == coupler.averaging_slices + + interval = 2 + data_time_step = "3h" + coupler.compute_coupled_indices(interval, data_time_step) + coupled_integration_dim = presteps + max(output_time_dim // input_time_dim, 1) + expected = np.empty([batch_size, coupled_integration_dim, len(input_times)]) + for b in range(batch_size): + for i in range(coupled_integration_dim): + expected[b, i, :] = ( + b + + (input_time_dim * i + 1) * interval + + np.array( + [ + pd.Timedelta(ts) / pd.Timedelta(data_time_step) + for ts in input_times + ] + ) + ) + expected = expected.astype(int) + assert np.array_equal(expected, coupler._coupled_offsets) + + scaling_df = pd.DataFrame.from_dict(omegaconf.OmegaConf.to_object(scaling_dict)).T + scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} + scaling_da = scaling_df.to_xarray().astype("float32") + coupler.set_scaling(scaling_da) + coupled_scaling = scaling_da.sel(index=variables).rename({"index": "channel_in"}) + expected = np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)) + assert np.array_equal(expected, coupler.coupled_scaling["mean"]) + expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) + assert np.array_equal(expected, coupler.coupled_scaling["std"]) + + averaging_window_max_indices = [ + i // pd.Timedelta(data_time_step) for i in coupler.input_times + ] + di = averaging_window_max_indices[0] + averaging_slices = [] + for j in range(coupler.coupled_integration_dim): + averaging_slices.append([]) + for i, r in enumerate(averaging_window_max_indices): + averaging_slices[j].append( + slice( + coupler.input_time_dim * j * di + i * di, + coupler.input_time_dim * j * di + r, + ) + ) + coupler.averaging_slices = averaging_slices + coupler.coupled_channel_indices = [0, 1] + + coupled_fields_batch_size = batch_size + coupled_fields_timedim = 4 + expected_shape = [ + coupler.coupled_integration_dim, + coupled_fields_batch_size, + coupler.timevar_dim, + ] + list(coupler.spatial_dims) + coupled_fields = th.rand( + coupled_fields_batch_size, + coupler.spatial_dims[0], + coupled_fields_timedim, + len(coupler.coupled_channel_indices), + coupler.spatial_dims[1], + coupler.spatial_dims[2], + ) + coupler.set_coupled_fields(coupled_fields) + assert list(coupler.preset_coupled_fields.shape) == expected_shape + + # check reset + assert coupler.coupled_mode + coupler.reset_coupler() + assert coupler.coupled_mode is False + + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_CoupledTimeSeriesDatasetZarr_initialization( + dataset_path, scaling_dict, pytestconfig +): + + from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( + CoupledTimeSeriesDatasetZarr, + ) + + # open our test dataset + time_da = xr.open_zarr(dataset_path).time.values + input_variables = ["z500", "z1000"] + valid_start_date = "1979-01-01" + valid_end_date = "1979-01-02" + + # check for failure of timestep not being a multiple of datatime step + with pytest.raises( + ValueError, match=("'time_step' must be a multiple of 'data_time_step' ") + ): + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + data_time_step="2h", + time_step="5h", + scaling=scaling_dict, + forecast_init_times=time_da[:2], + ) + + # check for failure of gap not being a multiple of datatime step + with pytest.raises( + ValueError, match=("'gap' must be a multiple of 'data_time_step' ") + ): + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + data_time_step="2h", + time_step="6h", + gap="3h", + scaling=scaling_dict, + batch_size=1, + forecast_init_times=time_da[:2], + ) + + # check for failure of invalid scaling variable on input + invalid_scaling = omegaconf.DictConfig( + { + "bogosity": {"mean": 0, "std": 42}, + } + ) + with pytest.raises(KeyError, match=("Input channels ")): + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + data_time_step="3h", + time_step="6h", + scaling=invalid_scaling, + batch_size=1, + forecast_init_times=time_da[:2], + ) + + # check for warning on batch size > 1 and forecast mode + warnings.filterwarnings("error") + with pytest.raises( + UserWarning, + match=( + "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`" + ), + ): + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + batch_size=2, + forecast_init_times=time_da[:2], + ) + warnings.resetwarnings() + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + add_train_noise=True, + batch_size=1, + forecast_init_times=time_da[:2], + ) + assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + start_date=valid_start_date, + end_date=valid_end_date, + ) + assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + batch_size=1, + forecast_init_times=time_da[:2], + ) + assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + batch_size=1, + forecast_init_times=time_da[:2], + data_time_step="3h", + time_step="6h", + ) + assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) + + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_CoupledTimeSeriesDatasetZarr_get_constants( + dataset_path, scaling_dict, constant_coupler_config, pytestconfig +): + + from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( + CoupledTimeSeriesDatasetZarr, + ) + + input_variables = ["z500", "z1000"] + constant_variables = ["lsm"] + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + constant_indices = [ + int(np.where(zarr_ds.channel_c[:] == ch)[0][0]) for ch in constant_variables + ] + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_dict, + couplings=constant_coupler_config, + forecast_init_times=zarr_ds.time[:2], + ) + assert timeseries_ds.get_constants() is None + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=constant_variables, + batch_size=1, + scaling=scaling_dict, + couplings=constant_coupler_config, + forecast_init_times=zarr_ds.time[:2], + ) + + # constants are reshaped + expected = np.transpose( + zarr_ds.constants.values[constant_indices], axes=(1, 0, 2, 3) + ) + outvar = timeseries_ds.get_constants() + assert np.array_equal( + expected, + outvar, + ) + + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_CoupledTimeSeriesDatasetZarr_len( + dataset_path, scaling_dict, constant_coupler_config, pytestconfig +): + from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( + CoupledTimeSeriesDatasetZarr, + ) + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + + variables = ["z500", "z1000"] + batch_size = 2 + + # check forecast mode + init_times = random.randint(1, zarr_ds.time.shape[0]) + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=variables, + scaling=scaling_dict, + batch_size=1, + forecast_init_times=zarr_ds.time[:init_times], + couplings=constant_coupler_config, + ) + assert len(timeseries_ds) == init_times + + batch2_coupler = constant_coupler_config.copy() + batch2_coupler[0]["params"]["batch_size"] = 2 + + # get the last index that's evenly divisible by 3 (9h / 3h) + last_index = (zarr_ds.time.shape[0] // 3) * 3 - 1 + + # check train mode + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=variables, + data_time_step="3h", + time_step="9h", + scaling=scaling_dict, + batch_size=batch_size, + couplings=batch2_coupler, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[last_index - 1].values, + ) + # Window length of 3 for one sample size + assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 3) // batch_size + + # drop incomplete last window + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=variables, + data_time_step="3h", + time_step="9h", + scaling=scaling_dict, + batch_size=batch_size, + drop_last=True, + couplings=batch2_coupler, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[last_index - 1].values, + ) + assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 4) // batch_size + + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_CoupledTimeSeriesDatasetZarr_get( + dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig +): + from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( + CoupledTimeSeriesDatasetZarr, + ) + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + + input_variables = list(zarr_ds.channel_out.values) + constant_variables = ["lsm"] + batch_size = 2 + batch2_constant_coupler = constant_coupler_config.copy() + batch2_constant_coupler[0]["params"]["batch_size"] = 2 + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=batch2_constant_coupler, + ) + + # check for invalid index + invalid_idx = len(zarr_ds.targets) + 1 + with pytest.raises( + IndexError, match=(f"index {invalid_idx} out of range for dataset with length") + ): + inputs, targets = timeseries_ds[invalid_idx] + + inputs, targets = timeseries_ds[0] + + # make sure number of targets is correct + assert len(targets) == batch_size + + # check target data + # need to transpose + targets_expected = zarr_ds.targets[batch_size].transpose( + "face", "channel_out", "height", "width" + ) + targets_expected = targets_expected.to_numpy() / 2 + assert np.array_equal(targets[0][:, 0, :, :], targets_expected) + + # check for negative index + inputs, targets = timeseries_ds[-1] + targets_expected = zarr_ds.targets[12].transpose( + "face", "channel_out", "height", "width" + ) + targets_expected = targets_expected.to_numpy() / 2 + + # we're not dropping incomplete elements by default + assert len(targets) == 0 + + # this time dropping incomplete so that we get a full sample + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=batch2_constant_coupler, + ) + + inputs, targets = timeseries_ds[-1] + targets_expected = zarr_ds.targets[-1 - batch_size].transpose( + "face", "channel_out", "height", "width" + ) + targets_expected = targets_expected.to_numpy() / 2 + assert np.array_equal(targets[0][:, 0, :, :], targets_expected) + + # without couplings + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=constant_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=[], + ) + non_perturbed_inputs = timeseries_ds + assert len(non_perturbed_inputs[0][0]) == 2 # just inputs and targets + + # without couplings but with noise + noise_params = { + "inputs": scaling_double_dict, + "couplings": scaling_double_dict, + } + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=constant_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + add_train_noise=True, + train_noise_params=noise_params, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=[], + ) + perturbed_inputs = timeseries_ds + # The first input will be the same sample, with perturbation it should have + # different values + assert non_perturbed_inputs[0][0][0].shape == perturbed_inputs[0][0][0].shape + assert not np.array_equal(non_perturbed_inputs[0][0][0], perturbed_inputs[0][0][0]) + + # With insolation we get 1 extra channel + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=None, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + add_insolation=True, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=batch2_constant_coupler, + ) + assert (len(inputs)) + 1 == len(timeseries_ds[0][0]) + + # nothing should change with forecast mode other than getting just inputs + init_times = random.randint(1, len(zarr_ds.time.values)) + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=None, + scaling=scaling_double_dict, + batch_size=1, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=constant_coupler_config, + ) + inputs = timeseries_ds[0] + + assert np.array_equal(targets[0][:, 0, :, :], targets_expected) + + # insolation adds 1 extra channel + init_times = random.randint(1, len(zarr_ds.time.values)) + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=None, + scaling=scaling_double_dict, + batch_size=1, + add_insolation=True, + forecast_init_times=zarr_ds.time[:init_times], + couplings=constant_coupler_config, + ) + assert (len(inputs)) + 1 == len(timeseries_ds[0]) + + # Constants + insolation is 2 extra channels + init_times = random.randint(1, len(zarr_ds.time.values)) + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=constant_variables, + scaling=scaling_double_dict, + batch_size=1, + add_insolation=True, + forecast_init_times=zarr_ds.time[:init_times], + couplings=constant_coupler_config, + ) + assert len(inputs) + 2 == len(timeseries_ds[0]) + + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_CoupledTimeSeriesDataModuleZarr_initialization( + dataset_path, splits, scaling_double_dict, constant_coupler_config, pytestconfig +): + + from physicsnemo.datapipes.healpix.data_modules_zarr import ( + CoupledTimeSeriesDataModuleZarr, + ) + + input_variables = ["z500", "z1000"] + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + + # test for invalid path + with pytest.raises(FileNotFoundError, match=("Dataset path not found")): + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path="DoesntExist", + input_variables=input_variables, + batch_size=1, + couplings=constant_coupler_config, + ) + + # use the prebuilt dataset + # Internally initializes DistributedManager + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + splits=omegaconf.DictConfig(splits), + couplings=constant_coupler_config, + ) + assert isinstance(timeseries_dm, CoupledTimeSeriesDataModuleZarr) + + # with init times + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + forecast_init_times=zarr_ds.time[:2], + couplings=constant_coupler_config, + ) + assert isinstance(timeseries_dm, CoupledTimeSeriesDataModuleZarr) + + # with splits + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + splits=omegaconf.DictConfig(splits), + couplings=constant_coupler_config, + ) + assert isinstance(timeseries_dm, CoupledTimeSeriesDataModuleZarr) + + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_CoupledTimeSeriesDataModuleZarr_get_constants( + dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig +): + + from physicsnemo.datapipes.healpix.data_modules_zarr import ( + CoupledTimeSeriesDataModuleZarr, + ) + + variables = ["z500", "z1000"] + constants = ["lsm"] + + # No constants + # Internally initializes DistributedManager + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=variables, + batch_size=1, + scaling=scaling_double_dict, + splits=splits, + constant_variables=None, + couplings=constant_coupler_config, + ) + + assert timeseries_dm.get_constants() is None + + # just lsm as constant + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=variables, + batch_size=1, + scaling=scaling_double_dict, + splits=splits, + constant_variables=constants, + couplings=constant_coupler_config, + ) + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + + # divide by 2 due to scaling + expected = ( + np.transpose( + zarr_ds.constants.sel(channel_c=constants).values, + axes=(1, 0, 2, 3), + ) + / 2.0 + ) + + assert np.array_equal( + timeseries_dm.get_constants(), + expected, + ) + + # with splits we're doing forecasting and get + # constants from train instead of test dataset + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=variables, + batch_size=1, + scaling=scaling_double_dict, + splits=splits, + constant_variables=constants, + couplings=constant_coupler_config, + ) + + assert np.array_equal( + timeseries_dm.get_constants(), + expected, + ) + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@nfsdata_or_fail +def test_CoupledTimeSeriesDataModuleZarr_get_dataloaders( + dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig +): + + from physicsnemo.datapipes.healpix.data_modules_zarr import ( + CoupledTimeSeriesDataModuleZarr, + ) + + input_variables = ["z500", "z1000"] + + # use the prebuilt dataset + # Internally initializes DistributedManager + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + splits=splits, + shuffle=False, + couplings=constant_coupler_config, + ) + + # with 1 shard should get no sampler + train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) + assert train_sampler is None + assert isinstance(train_dataloader, DataLoader) + + val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) + assert val_sampler is None + assert isinstance(val_dataloader, DataLoader) + + test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) + assert test_sampler is None + assert isinstance(test_dataloader, DataLoader) + + # with >1 shard should be distributed sampler + train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) + assert isinstance(train_sampler, DistributedSampler) + assert isinstance(train_dataloader, DataLoader) + + val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) + assert isinstance(val_sampler, DistributedSampler) + assert isinstance(val_dataloader, DataLoader) + + test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) + assert isinstance(test_sampler, DistributedSampler) + assert isinstance(test_dataloader, DataLoader) + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@nfsdata_or_fail +def test_CoupledTimeSeriesDataModuleZarr_get_coupled_vars( + dataset_path, + scaling_double_dict, + splits, + constant_coupler_config, + average_coupler_config, + pytestconfig, +): + from physicsnemo.datapipes.healpix.data_modules_zarr import ( + CoupledTimeSeriesDataModuleZarr, + ) + + input_variables = ["z500", "z1000"] + + # Constant coupler + # Internally initializes DistributedManager + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + splits=splits, + couplings=constant_coupler_config, + ) + + outvar = timeseries_dm._get_coupled_vars() + outvar.sort() + expected = ["z250"] + expected.sort() + + assert expected == outvar + + # Average coupler + # Internally initializes DistributedManager + timeseries_dm = CoupledTimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + splits=splits, + couplings=average_coupler_config, + ) + outvar = timeseries_dm._get_coupled_vars() + outvar.sort() + + assert expected == outvar + + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("xarray") +@nfsdata_or_fail +def test_CoupledTimeSeriesDatasetZarr_next_integration( + dataset_path, scaling_dict, pytestconfig +): + from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( + CoupledTimeSeriesDatasetZarr, + ) + + spatial_dims = [12, 32, 32] + input_variables = ["z500", "z1000"] + coupled_channel_indices = [0, 1] + coupled_variables = ["z250"] + num_variables = len(input_variables) + input_time_dim = 1 + output_time_dim = 1 + batch_size = 1 + + constant_coupler = [ + { + "coupler": "ConstantCoupler", + "params": { + "batch_size": 1, + "variables": coupled_variables, + "input_times": ["0h"], + "input_time_dim": input_time_dim, + "output_time_dim": output_time_dim, + "presteps": 0, + "prepared_coupled_data": True, + }, + } + ] + + # open our test dataset + ds = xr.open_zarr(dataset_path) + init_times = random.randint(1, len(ds.time.values)) + # channels need to be subselected before being handed over + test_ds = ds.sel( + channel_in=input_variables + coupled_variables, + channel_out=input_variables, + ) + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + batch_size=batch_size, + couplings=constant_coupler, + data_time_step="6h", + time_step="6h", + drop_last=True, + add_insolation=True, + forecast_init_times=test_ds.time[:init_times], + ) + + test_model_outputs = th.rand( + 1, + spatial_dims[0], + output_time_dim, + num_variables, + spatial_dims[1], + spatial_dims[2], + ) + constants = np.transpose(ds.constants.values, axes=(1, 0, 2, 3)) + coupled_fields = th.rand( + batch_size, + spatial_dims[0], + input_time_dim + output_time_dim, + len(input_variables), + spatial_dims[1], + spatial_dims[2], + ) + + expected_coupling = coupled_fields[:, :, :, coupled_channel_indices, :, :].permute( + 2, 0, 3, 1, 4, 5 + ) + expected_coupling = expected_coupling[0, :, -1, :, :, :] + expected_coupling = expected_coupling.unsqueeze(0).unsqueeze(0) + expected_coupling = expected_coupling.repeat(1, batch_size, 1, 1, 1, 1) + + # need to grab at least 1 sample to properly intialize everything + timeseries_ds[0] + # hacky way to setup the indices since we don't actually have any coupled fields + timeseries_ds.couplings[0].coupled_channel_indices = coupled_channel_indices + + # set the coupled fields + timeseries_ds.couplings[0].set_coupled_fields(coupled_fields) + test_integration = timeseries_ds.next_integration(test_model_outputs, constants) + # test to make sure prognostics are used, constants stay the same, and couplings + # are what we set + assert np.array_equal(test_integration[0], test_model_outputs[:, :, -1:]) + assert np.array_equal(test_integration[2], constants) + assert np.array_equal(test_integration[3], expected_coupling) + + # I have absolutely no idea why a coupled dataset has the option for 0 couplings + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + batch_size=batch_size, + couplings=[], + data_time_step="6h", + time_step="6h", + drop_last=True, + add_insolation=True, + forecast_init_times=test_ds.time[:init_times], + ) + # need to grab at least 1 sample to properly intialize everything + timeseries_ds[0] + test_integration = timeseries_ds.next_integration(test_model_outputs, constants) + assert np.array_equal(test_integration[0], test_model_outputs[:, :, -1:]) + assert np.array_equal(test_integration[2], constants) diff --git a/test/datapipes/test_healpix_zarr.py b/test/datapipes/test_healpix_zarr.py new file mode 100644 index 0000000000..fb4fe1aef8 --- /dev/null +++ b/test/datapipes/test_healpix_zarr.py @@ -0,0 +1,782 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import random +import warnings +from pathlib import Path + +import pytest +from pytest_utils import import_or_fail, nfsdata_or_fail +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from physicsnemo.distributed import DistributedManager + +omegaconf = pytest.importorskip("omegaconf") +np = pytest.importorskip("numpy") +xr = pytest.importorskip("xarray") +zarr = pytest.importorskip("zarr") + + +@pytest.fixture +def dataset_path(): + data_dir = "/data/nfs/modulus-data/datasets/healpix/" + dataset_name = "healpix.zarr" + path = Path(data_dir, dataset_name) + return path + + +@pytest.fixture +def splits(): + split_dict = { + "train_date_start": "1979-01-01", + "train_date_end": "1979-01-01T21:00", + "val_date_start": "1979-01-02", + "val_date_end": "1979-01-02T09:00", + "test_date_start": "1979-01-02T12:00", + "test_date_end": "1979-01-02T18:00", + } + return split_dict + + +@pytest.fixture +def scaling_dict(): + scaling = { + "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, + "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, + "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, + "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, + "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, + "z250": {"mean": 101186.28125, "std": 5551.77978515625}, + "z500": {"mean": 55625.9609375, "std": 2681.712890625}, + "lsm": {"mean": 0, "std": 1}, + "z": {"mean": 0, "std": 1}, + "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, + "extra": {"mean": 1, "std": 0}, # doesn't appear in test dataset + } + return omegaconf.DictConfig(scaling) + + +@pytest.fixture +def scaling_double_dict(): + scaling = { + "t2m0": {"mean": 0, "std": 2}, + "t850": {"mean": 0, "std": 2}, + "tau300-700": {"mean": 0, "std": 2}, + "tcwv0": {"mean": 0, "std": 2}, + "z1000": {"mean": 0, "std": 2}, + "z250": {"mean": 0, "std": 2}, + "z500": {"mean": 0, "std": 2}, + "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, + "lsm": {"mean": 0, "std": 2}, + "z": {"mean": 0, "std": 2}, + "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset + } + return omegaconf.DictConfig(scaling) + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@nfsdata_or_fail +def test_TimeSeriesDataset_initialization( + dataset_path, + scaling_dict, + pytestconfig, +): + from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( + TimeSeriesDatasetZarr, + ) + + input_variables = ["t2m0", "t850", "z500"] + + bad_start_date = "1900-01-01" + valid_start_date = "1979-01-01" + bad_end_date = "2000-12-31" + valid_end_date = "1979-01-02" + + zarr_ds = zarr.open(dataset_path) + time_da = xr.open_zarr(dataset_path).time + + # check for failure of invalid dataset path + # Check if fsspec is available for object store paths, optional dependency + if not importlib.util.find_spec("fsspec"): + # If fsspec is not available, expect an ImportError + with pytest.raises( + ImportError, match=("fsspec is required to access object store paths") + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path="s3://physicsnemo-data/datasets/healpix/healpix.zarr", + scaling=scaling_dict, + input_variables=input_variables, + ) + + # If path doesn't exist, expect a FileNotFoundError + with pytest.raises(FileNotFoundError, match=("Dataset not found at")): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path="/boguspath.zarr", + scaling=scaling_dict, + input_variables=input_variables, + ) + + # check for failure of timestep not being a multiple of datatime step + with pytest.raises( + ValueError, match=("'time_step' must be a multiple of 'data_time_step' ") + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + data_time_step="2h", + time_step="5h", + scaling=scaling_dict, + input_variables=input_variables, + forecast_init_times=zarr_ds.time[:2], + batch_size=1, + ) + + # check for failure of gap not being a multiple of datatime step + with pytest.raises( + ValueError, match=("'gap' must be a multiple of 'data_time_step' ") + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + data_time_step="2h", + time_step="6h", + gap="3h", + scaling=scaling_dict, + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=input_variables, + ) + + # check for failure of invalid scaling variable on input + invalid_scaling = omegaconf.DictConfig( + { + "bogosity": {"mean": 0, "std": 42}, + } + ) + with pytest.raises(KeyError, match=("Input channels ")): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + data_time_step="3h", + time_step="6h", + scaling=invalid_scaling, + forecast_init_times=time_da[:10], + input_variables=input_variables, + batch_size=1, + ) + + # check for warning on batch size > 1 and forecast mode + warnings.filterwarnings("error") + with pytest.raises( + UserWarning, + match=( + "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`" + ), + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + batch_size=2, + forecast_init_times=zarr_ds.time[:2], + input_variables=input_variables, + ) + warnings.resetwarnings() + + # check for no dates provided + with pytest.raises( + ValueError, + match=("Either start and end date or forecast_init_times must be provided"), + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + batch_size=2, + input_variables=input_variables, + ) + + # check for out of range dates + warnings.filterwarnings("error") + with pytest.raises( + UserWarning, + match=(f"Start date {bad_start_date} is before first available date"), + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + batch_size=2, + start_date=bad_start_date, + end_date=valid_end_date, + input_variables=input_variables, + ) + + with pytest.raises( + UserWarning, match=(f"End date {bad_end_date} is after last available date") + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + batch_size=2, + start_date=valid_start_date, + end_date=bad_end_date, + input_variables=input_variables, + ) + warnings.resetwarnings() + + # test no scaling + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=input_variables, + ) + assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) + + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=input_variables, + ) + assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) + + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + batch_size=1, + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=input_variables, + ) + assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) + + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + batch_size=1, + start_date=valid_start_date, + end_date=valid_end_date, + data_time_step="3h", + time_step="6h", + input_variables=input_variables, + ) + assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("numpy") +@import_or_fail("zarr") +@nfsdata_or_fail +def test_TimeSeriesDataset_get_constants(dataset_path, scaling_dict, pytestconfig): + from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( + TimeSeriesDatasetZarr, + ) + + input_variables = ["z500", "z1000"] + constant_variables = ["lsm"] + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + constant_indices = [ + int(np.where(zarr_ds.channel_c[:] == ch)[0][0]) for ch in constant_variables + ] + + # Constant that isn't in dataset + with pytest.raises(KeyError, match=("Requested constants not found in dataset")): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_dict, + constant_variables=["DoesntExist"], + forecast_init_times=zarr_ds.time[:2], + ) + + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + constant_variables=None, + forecast_init_times=zarr_ds.time[:2], + batch_size=1, + ) + + assert timeseries_ds.get_constants() is None + + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + constant_variables=constant_variables, + forecast_init_times=zarr_ds.time[:2], + batch_size=1, + ) + + # constants are reshaped + expected = np.transpose( + zarr_ds.constants.values[constant_indices], axes=(1, 0, 2, 3) + ) + outvar = timeseries_ds.get_constants() + assert np.array_equal( + expected, + outvar, + ) + zarr_ds.close() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@nfsdata_or_fail +def test_TimeSeriesDataset_len(dataset_path, scaling_dict, pytestconfig): + from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( + TimeSeriesDatasetZarr, + ) + + input_variables = ["z500", "z1000"] + batch_size = 2 + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + + # check forecast mode + init_times = random.randint(1, zarr_ds.time.shape[0]) + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_dict, + batch_size=1, + forecast_init_times=zarr_ds.time[:init_times], + ) + assert len(timeseries_ds) == init_times + + # get the last index that's evenly divisible by 3 (9h / 3h) + last_index = (zarr_ds.time.shape[0] // 3) * 3 - 1 + + # check train mode + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + data_time_step="3h", + time_step="9h", + input_variables=input_variables, + scaling=scaling_dict, + batch_size=batch_size, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[last_index - 1].values, + ) + # Window length of 3 for one sample size + assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 3) // batch_size + + # drop incomplete last window + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + data_time_step="3h", + time_step="9h", + input_variables=input_variables, + scaling=scaling_dict, + batch_size=batch_size, + drop_last=True, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[last_index - 1].values, + ) + assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 4) // batch_size + + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("numpy") +@nfsdata_or_fail +def test_TimeSeriesDataset_get(dataset_path, scaling_double_dict, splits, pytestconfig): + from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( + TimeSeriesDatasetZarr, + ) + + input_variables = ["z500", "z1000"] + constant_variables = ["lsm"] + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + time_da = zarr_ds.time.values + + batch_size = 2 + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + start_date=splits["train_date_start"], + end_date=splits["test_date_start"], + ) + + # check for invalid index + invalid_idx = len(zarr_ds.targets) + 1 + with pytest.raises( + IndexError, match=(f"index {invalid_idx} out of range for dataset with length") + ): + inputs, targets = timeseries_ds[invalid_idx] + + inputs, targets = timeseries_ds[0] + + # make sure number of targets is correct + assert len(targets) == batch_size + + # check target data + # need to transpose + targets_expected = ( + zarr_ds.targets[batch_size] + .transpose("face", "channel_out", "height", "width") + .sel(channel_out=input_variables) + ) + + # scale target data + targets_expected = targets_expected.to_numpy() / 2 + assert np.array_equal(targets[0][:, 0, :, :], targets_expected) + + # check for negative index + inputs, targets = timeseries_ds[-1] + targets_expected = zarr_ds.targets[12].transpose( + "face", "channel_out", "height", "width" + ) + targets_expected = targets_expected.to_numpy() / 2 + + # we're not dropping incomplete elements by default + assert len(targets) == 0 + + # this time dropping incomplete so that we get a full sample sample + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + start_date=time_da[0], + end_date=time_da[-1], + ) + + inputs, targets = timeseries_ds[-1] + targets_expected = ( + zarr_ds.targets[-1 - batch_size] + .transpose("face", "channel_out", "height", "width") + .sel(channel_out=input_variables) + ) + targets_expected = targets_expected.to_numpy() / 2 + assert np.array_equal(targets[0][:, 0, :, :], targets_expected) + + # With insolation we get 1 extra tensor + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + add_insolation=True, + start_date=time_da[0], + end_date=time_da[-1], + ) + + # verify underlying data doesn't change + targets_expected = ( + zarr_ds.targets[2] + .transpose("face", "channel_out", "height", "width") + .sel(channel_out=input_variables) + ) + targets_expected = targets_expected.to_numpy() / 2 + result = timeseries_ds[0] + assert (len(inputs)) + 1 == len(result[0]) + assert np.array_equal(result[1][0][:, 0, :, :], targets_expected) + + # With insolation and constants we get 2 extra tensors + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + constant_variables=constant_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + add_insolation=True, + start_date=time_da[0], + end_date=time_da[-1], + ) + assert (len(inputs)) + 2 == len(timeseries_ds[0][0]) + + # nothing should change with forecast mode other than getting just inputs + init_times = random.randint(1, len(zarr_ds.time.values)) + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=1, + forecast_init_times=zarr_ds.time[:init_times].values, + ) + inputs = timeseries_ds[0] + + assert len(inputs) == 1 + + # insolation adds 1 extra tensor, same as above but using forecast mode + init_times = random.randint(1, len(zarr_ds.time.values)) + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=1, + add_insolation=True, + forecast_init_times=zarr_ds.time[:init_times].values, + ) + assert (len(inputs) + 1) == len(timeseries_ds[0]) + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("zarr") +@nfsdata_or_fail +def test_TimeSeriesDataModule_initialization( + dataset_path, splits, scaling_double_dict, pytestconfig +): + from physicsnemo.datapipes.healpix.data_modules_zarr import ( + TimeSeriesDataModuleZarr, + ) + + input_variables = ["z500", "z1000"] + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + + # test for invalid path + with pytest.raises(FileNotFoundError, match=("Dataset path not found")): + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path="DoesntExist", + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + ) + + # test for missing times + with pytest.raises( + ValueError, match=("Either splits or forecast_init_times must be provided") + ): + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + ) + + # test for overlapping dates + warnings.filterwarnings("error") + with pytest.raises( + UserWarning, match=("Training and validation date ranges overlap") + ): + bad_splits = splits.copy() + bad_splits["val_date_start"] = splits["train_date_start"] + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + splits=bad_splits, + scaling=scaling_double_dict, + ) + warnings.resetwarnings() + + warnings.filterwarnings("error") + with pytest.raises(UserWarning, match=("Training and test date ranges overlap")): + bad_splits = splits.copy() + bad_splits["test_date_start"] = splits["train_date_start"] + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + splits=bad_splits, + scaling=scaling_double_dict, + ) + warnings.resetwarnings() + + warnings.filterwarnings("error") + with pytest.raises(UserWarning, match=("Test and validation date ranges overlap")): + bad_splits = splits.copy() + bad_splits["val_date_end"] = splits["test_date_start"] + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + splits=bad_splits, + scaling=scaling_double_dict, + ) + warnings.resetwarnings() + + # with init times + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + forecast_init_times=zarr_ds.time[:2], + ) + assert isinstance(timeseries_dm, TimeSeriesDataModuleZarr) + + # with splits + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + splits=omegaconf.DictConfig(splits), + ) + assert isinstance(timeseries_dm, TimeSeriesDataModuleZarr) + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("netCDF4") +@import_or_fail("numpy") +@import_or_fail("zarr") +@nfsdata_or_fail +def test_TimeSeriesDataModule_get_constants( + dataset_path, scaling_double_dict, splits, pytestconfig +): + from physicsnemo.datapipes.healpix.data_modules_zarr import ( + TimeSeriesDataModuleZarr, + ) + + input_variables = ["z500", "z1000"] + constants = ["lsm"] + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + forecast_times = zarr_ds.time[:2] + + # No constants + # Internally initializes DistributedManager + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + constant_variables=None, + forecast_init_times=forecast_times, + ) + + assert timeseries_dm.get_constants() is None + + # Constant that isn't in dataset + with pytest.raises(KeyError, match=("Requested constants not found in dataset")): + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + constant_variables={"missing": "missing"}, + forecast_init_times=forecast_times, + ) + + # just lsm as constant + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + constant_variables=constants, + forecast_init_times=forecast_times, + ) + + # open our test dataset + zarr_ds = xr.open_zarr(dataset_path) + + # dividing by 2 due to scaling + expected = ( + np.transpose( + zarr_ds.constants.sel(channel_c=constants).values, + axes=(1, 0, 2, 3), + ) + / 2.0 + ) + + assert np.array_equal( + timeseries_dm.get_constants(), + expected, + ) + + # with splits we're doing forecasting and get + # constants from train instead of test dataset + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + constant_variables=constants, + splits=splits, + ) + + assert np.array_equal( + timeseries_dm.get_constants(), + expected, + ) + zarr_ds.close() + DistributedManager.cleanup() + + +@import_or_fail("omegaconf") +@import_or_fail("zarr") +@nfsdata_or_fail +def test_TimeSeriesDataModule_get_dataloaders( + dataset_path, scaling_double_dict, splits, pytestconfig +): + from physicsnemo.datapipes.healpix.data_modules_zarr import ( + TimeSeriesDataModuleZarr, + ) + + input_variables = ["z500", "z1000"] + + # use the prebuilt dataset + # Internally initializes DistributedManager + timeseries_dm = TimeSeriesDataModuleZarr( + dataset_path=dataset_path, + input_variables=input_variables, + batch_size=1, + scaling=scaling_double_dict, + splits=splits, + shuffle=False, + ) + + # with 1 shard should get no sampler + train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) + assert train_sampler is None + assert isinstance(train_dataloader, DataLoader) + + val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) + assert val_sampler is None + assert isinstance(val_dataloader, DataLoader) + + test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) + assert test_sampler is None + assert isinstance(test_dataloader, DataLoader) + + # with >1 shard should be distributed sampler + train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) + assert isinstance(train_sampler, DistributedSampler) + assert isinstance(train_dataloader, DataLoader) + + val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) + assert isinstance(val_sampler, DistributedSampler) + assert isinstance(val_dataloader, DataLoader) + + test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) + assert isinstance(test_sampler, DistributedSampler) + assert isinstance(test_dataloader, DataLoader) + DistributedManager.cleanup() diff --git a/test/metrics/test_hydrostatic_loss.py b/test/metrics/test_hydrostatic_loss.py new file mode 100644 index 0000000000..4c3058814b --- /dev/null +++ b/test/metrics/test_hydrostatic_loss.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np +import pytest +import torch + +from physicsnemo.metrics.climate.hydrostasy import ( + HydrostaticBalance, +) + +xr = pytest.importorskip("xarray") + +def test_constant_temperature(rtol: float = 1e-3, atol: float = 1e-3): + R = 287 # J K^{-1} kg^{-1} + g0 = 9.81 # m s^{-2} + T = 273.15 # K + z_pressure_levels = {0: 850.0, 1: 500.0, 2: 250.0, 3: 50.0} + z_pressure_levels = { + c: p * 100 for c, p in z_pressure_levels.items() + } # convert hPa to Pa + anchor_z_channel = 0 + anchor_T_channel = 4 + + # Create HydrostaticBalance constraint object + constraint = HydrostaticBalance( + z_pressure_levels, anchor_z_channel, anchor_T_channel, R, g0 + ) + + x = torch.zeros((1, 12, 5, 1, 1)) + # Set Z + z = ( + R + * T + / g0 + * torch.log( + z_pressure_levels[anchor_z_channel] + / torch.Tensor(list(z_pressure_levels.values())) + ) + ) + x[:, :, 0:4, :, :] = z.view(1, 1, -1, 1, 1) + # Set Tv at 850 hPa + x[:, :, 4, :, :] = T + torch.zeros_like(x[:, :, 4, :, :]) + + Tv = constraint(x) + + print(Tv[0, 0, :, 0, 0]) + assert torch.allclose(T * torch.ones_like(Tv), Tv, rtol=rtol, atol=atol,) + + +@pytest.mark.parametrize("N", [10, 20]) +def test_constant_lapse_rate(N: int, rtol: float = 1e-3, atol: float = 1e-3): + R = 287 # J K^{-1} kg^{-1} + g0 = 9.81 # m s^{-2} + T0 = 273.15 # K + z0 = 0 # m + Gamma = 9.8 / 1000 # K m^{-1} + p = np.logspace(np.log10(850.0), np.log10(50.0), num=N) + # z_pressure_levels = {0: 850., 1: 500., 2: 250., 3: 50.} + # z_pressure_levels = {i: 850. - 50.*i for i in range(17)} + z_pressure_levels = {i: p[i] for i in range(p.shape[0])} + z_pressure_levels = { + c: p * 100 for c, p in z_pressure_levels.items() + } # convert hPa to Pa + anchor_z_channel = 0 + anchor_T_channel = len(z_pressure_levels) + + # Create HydrostaticBalance constraint object + constraint = HydrostaticBalance( + z_pressure_levels, anchor_z_channel, anchor_T_channel, R, g0 + ) + + x = torch.zeros((1, 12, anchor_T_channel + 1, 1, 1)) + # Set T + T = T0 * torch.pow( + torch.Tensor(list(z_pressure_levels.values())) + / z_pressure_levels[anchor_z_channel], + Gamma * R / g0, + ) + print("T: ", T) + z = z0 + (T0 - T) / Gamma + print("z: ", z) + x[:, :, 0:anchor_T_channel, :, :] = z.view(1, 1, -1, 1, 1) + # Set Tv at 850 hPa + x[:, :, anchor_T_channel, :, :] = T[0] + torch.zeros_like( + x[:, :, anchor_T_channel, :, :] + ) + + Tv = constraint(x) + + print(Tv[0, 0, :, 0, 0]) + # assert torch.allclose(T * torch.ones_like(Tv), Tv) + return p, z, T, Tv[0, 0, constraint.z_channels, 0, 0] + + +@pytest.mark.parametrize("N", [10, 20]) +def test_dual_lapse_rate(N: int, rtol: float = 1e-3, atol: float = 1e-3): + R = 287 # J K^{-1} kg^{-1} + g0 = 9.81 # m s^{-2} + T0 = 273.15 # K + z0 = 0 # m + Gamma1 = 9.8 / 1000 # K m^{-1} + Gamma2 = Gamma1 / 2.0 + zc = 10000 # m + p = np.logspace(np.log10(850.0), np.log10(50.0), num=N) + # z_pressure_levels = {0: 850., 1: 500., 2: 250., 3: 50.} + # z_pressure_levels = {i: 850. - 50.*i for i in range(17)} + z_pressure_levels = {i: p[i] for i in range(p.shape[0])} + z_pressure_levels = { + c: p * 100 for c, p in z_pressure_levels.items() + } # convert hPa to Pa + anchor_z_channel = 0 + anchor_T_channel = len(z_pressure_levels) + + # Create HydrostaticBalance constraint object + constraint = HydrostaticBalance( + z_pressure_levels, anchor_z_channel, anchor_T_channel, R, g0 + ) + + x = torch.zeros((1, 12, anchor_T_channel + 1, 1, 1)) + # Set T + T1 = T0 * torch.pow( + torch.Tensor(list(z_pressure_levels.values())) + / z_pressure_levels[anchor_z_channel], + Gamma1 * R / g0, + ) + z1 = z0 + (T0 - T1) / Gamma1 + + Tc = T0 - Gamma1 * (zc - z0) + Pc = z_pressure_levels[anchor_z_channel] * (Tc / T0) ** (g0 / (Gamma1 * R)) + T2 = Tc * torch.pow( + torch.Tensor(list(z_pressure_levels.values())) / Pc, Gamma2 * R / g0 + ) + z2 = zc + (Tc - T2) / Gamma2 + + z = torch.zeros_like(z1) + z = torch.where(z1 < zc, z1, z2) + T = torch.where(z1 < zc, T1, T2) + + x[:, :, 0:anchor_T_channel, :, :] = z.view(1, 1, -1, 1, 1) + # Set Tv at 850 hPa + x[:, :, anchor_T_channel, :, :] = T[0] + torch.zeros_like( + x[:, :, anchor_T_channel, :, :] + ) + + Tv = constraint(x) + + print(Tv[0, 0, :, 0, 0]) + # assert torch.allclose(T * torch.ones_like(Tv), Tv) + return p, z, T, Tv[0, 0, constraint.z_channels, 0, 0] diff --git a/test/models/dlwp_healpix/_cln_reference.py b/test/models/dlwp_healpix/_cln_reference.py new file mode 100644 index 0000000000..1524062627 --- /dev/null +++ b/test/models/dlwp_healpix/_cln_reference.py @@ -0,0 +1,73 @@ +# Reference (old) implementation of ConditionalLayerNorm for testing. +# This is a copy of the original code before optimization. + +from typing import List + +import torch as th + +try: + from apex.normalization import FusedLayerNorm + _APEX_AVAILABLE = True +except ImportError: + _APEX_AVAILABLE = False + + +class ConditionalLayerNormReference(th.nn.Module): + def __init__( + self, + condition_shape: int, + channel_depth: int, + mlp_hidden_dims: List[int] = [128, 128], + activation: th.nn.Module = None, + eps: float = 1e-5, + n_faces: int = 12, + norm_op: str = "torch", + init_cln_to_zero: bool = False, + scale_center: float = 0.0, + ): + super().__init__() + self.eps = eps + self.condition_shape = condition_shape + self.channel_depth = channel_depth + self.hidden_dims = mlp_hidden_dims + self.activation = activation if activation is not None else th.nn.Identity() + self.gamma_mlp = self._make_mlp(self.condition_shape, self.hidden_dims, self.channel_depth, self.activation) + self.beta_mlp = self._make_mlp(self.condition_shape, self.hidden_dims, self.channel_depth, self.activation) + self.n_faces = n_faces + self.scale_center = scale_center + + if init_cln_to_zero: + self.gamma_mlp[-1].weight.data.zero_() + self.beta_mlp[-1].weight.data.zero_() + self.gamma_mlp[-1].bias.data.zero_() + self.beta_mlp[-1].bias.data.zero_() + + if norm_op == "torch": + self.norm = th.nn.LayerNorm(channel_depth, elementwise_affine=False) + elif norm_op == "apex": + if not _APEX_AVAILABLE: + raise ImportError("Apex FusedLayerNorm requested but apex is not available") + self.norm = FusedLayerNorm(channel_depth, elementwise_affine=False) + + def _make_mlp(self, in_dim: int, hidden_dims: List[int], out_dim: int, activation: th.nn.Module) -> th.nn.Sequential: + layers = [] + for hdim in hidden_dims: + layers.append(th.nn.Linear(in_dim, hdim)) + if activation: + layers.append(activation) + in_dim = hdim + layers.append(th.nn.Linear(in_dim, out_dim)) + return th.nn.Sequential(*layers) + + def forward(self, x: th.Tensor, conditions: th.Tensor) -> th.Tensor: + x = x.permute(0, 2, 3, 1) + x_norm = self.norm(x) + + gamma = self.scale_center + self.gamma_mlp(conditions)[:, None, None, :] + beta = self.beta_mlp(conditions)[:, None, None, :] + + gamma = gamma.repeat_interleave(self.n_faces, dim=0) + beta = beta.repeat_interleave(self.n_faces, dim=0) + + x = gamma * x_norm + beta + return x.permute(0, 3, 1, 2) diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py new file mode 100644 index 0000000000..5b6f729663 --- /dev/null +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -0,0 +1,265 @@ +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import pytest +import torch +from _cln_reference import ConditionalLayerNormReference +from physicsnemo.models.dlwp_healpix.layers.normalization import ConditionalLayerNorm + + +def _make_old_cln(condition_shape, channel_depth, **kwargs): + """Instantiate the reference (old) implementation.""" + return ConditionalLayerNormReference( + condition_shape=condition_shape, channel_depth=channel_depth, **kwargs + ).cuda() + + +def _make_new_cln(condition_shape, channel_depth, **kwargs): + """Instantiate the optimized (new) implementation.""" + return ConditionalLayerNorm( + condition_shape=condition_shape, channel_depth=channel_depth, **kwargs + ).cuda() + + +def _copy_old_to_new(old_cln, new_cln): + """Copy old separate gamma/beta MLP weights into new fused MLP using block-diagonal structure. + + Old: gamma_mlp and beta_mlp each have hidden_dims [h1, h2] and output C. + New: gamma_beta_mlp has hidden_dims [2*h1, 2*h2] and output 2*C. + + Layer 0 (condition_shape → 2*h1): vertical cat of gamma/beta weights. + Layer i>0 (2*h_{i-1} → 2*h_i or 2*C): block-diagonal [[gamma, 0], [0, beta]]. + Biases: always concatenated. + """ + old_sd = old_cln.state_dict() + + # Collect Linear layer indices from the old gamma MLP + gamma_linear_indices = sorted({ + int(k.split(".")[1]) + for k in old_sd if k.startswith("gamma_mlp.") and k.endswith(".weight") + }) + first_layer_idx = gamma_linear_indices[0] + + new_sd = {} + for key in old_sd: + if key.startswith("norm."): + new_sd[key] = old_sd[key] + + for idx in gamma_linear_indices: + for param in ("weight", "bias"): + gamma_val = old_sd[f"gamma_mlp.{idx}.{param}"] + beta_val = old_sd[f"beta_mlp.{idx}.{param}"] + fused_key = f"gamma_beta_mlp.{idx}.{param}" + + if param == "bias": + new_sd[fused_key] = torch.cat([gamma_val, beta_val], dim=0) + elif idx == first_layer_idx: + # First layer: shared input dim, just cat along output dim + new_sd[fused_key] = torch.cat([gamma_val, beta_val], dim=0) + else: + # Block-diagonal: [[gamma, 0], [0, beta]] + out_old, in_old = gamma_val.shape + zeros = torch.zeros_like(gamma_val) + new_sd[fused_key] = torch.cat([ + torch.cat([gamma_val, zeros], dim=1), + torch.cat([zeros, beta_val], dim=1), + ], dim=0) + + new_cln.load_state_dict(new_sd) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("n_cond", [1, 2, 4]) +@pytest.mark.parametrize("channels_last", [False, True]) +@pytest.mark.parametrize("scale_center", [0.0, 1.0]) +def test_old_vs_new_forward(n_cond, channels_last, scale_center): + """Verify optimized CLN matches reference implementation with block-diagonal weight mapping.""" + C, H, W = 128, 16, 16 + cond_shape = 32 + B_nf = n_cond * 12 + + torch.manual_seed(42) + old_cln = _make_old_cln(cond_shape, C, scale_center=scale_center) + + new_cln = _make_new_cln(cond_shape, C, scale_center=scale_center) + _copy_old_to_new(old_cln, new_cln) + + x = torch.randn(B_nf, C, H, W, device="cuda") + cond = torch.randn(n_cond, cond_shape, device="cuda") + + if channels_last: + x = x.to(memory_format=torch.channels_last) + + with torch.no_grad(): + out_old = old_cln(x, cond) + out_new = new_cln(x, cond) + + assert out_old.shape == out_new.shape + assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), \ + f"Max diff: {(out_old - out_new).abs().max().item()}" + + if channels_last: + assert out_new.is_contiguous(memory_format=torch.channels_last), \ + "Output should preserve channels_last format" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("channels_last", [False, True]) +def test_old_vs_new_backward(channels_last): + """Verify gradients match between old and new implementations.""" + C, H, W = 64, 8, 8 + cond_shape = 16 + n_cond = 2 + B_nf = n_cond * 12 + + torch.manual_seed(42) + old_cln = _make_old_cln(cond_shape, C) + new_cln = _make_new_cln(cond_shape, C) + _copy_old_to_new(old_cln, new_cln) + + x_base = torch.randn(B_nf, C, H, W, device="cuda") + cond_base = torch.randn(n_cond, cond_shape, device="cuda") + + if channels_last: + x_base = x_base.to(memory_format=torch.channels_last) + + x_old = x_base.clone().detach().requires_grad_(True) + cond_old = cond_base.clone().detach().requires_grad_(True) + x_new = x_base.clone().detach().requires_grad_(True) + cond_new = cond_base.clone().detach().requires_grad_(True) + + out_old = old_cln(x_old, cond_old) + out_old.sum().backward() + + out_new = new_cln(x_new, cond_new) + out_new.sum().backward() + + assert torch.allclose(x_old.grad, x_new.grad, atol=1e-4, rtol=1e-3), \ + f"Input grad max diff: {(x_old.grad - x_new.grad).abs().max().item()}" + + assert torch.allclose(cond_old.grad, cond_new.grad, atol=1e-4, rtol=1e-3), \ + f"Cond grad max diff: {(cond_old.grad - cond_new.grad).abs().max().item()}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("channels_last", [False, True]) +def test_init_cln_to_zero_matches_layer_norm(channels_last): + """With scale_center=1.0 and init_cln_to_zero=True, CLN should behave like plain LayerNorm.""" + C, H, W = 64, 8, 8 + n_cond = 2 + B_nf = n_cond * 12 + + torch.manual_seed(42) + cln = _make_new_cln(32, C, scale_center=1.0, init_cln_to_zero=True) + plain_ln = torch.nn.LayerNorm(C, elementwise_affine=False).cuda() + + x = torch.randn(B_nf, C, H, W, device="cuda") + cond = torch.randn(n_cond, 32, device="cuda") + + if channels_last: + x = x.to(memory_format=torch.channels_last) + + with torch.no_grad(): + out_cln = cln(x, cond) + x_nhwc = x.permute(0, 2, 3, 1) + out_ln = plain_ln(x_nhwc).permute(0, 3, 1, 2) + + assert torch.allclose(out_cln, out_ln, atol=1e-5, rtol=1e-4), \ + f"Max diff: {(out_cln - out_ln).abs().max().item()}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("channels_last", [False, True]) +def test_backward_gradients(channels_last): + """Verify gradients flow through CLN and are finite.""" + C, H, W = 64, 8, 8 + cond_shape = 16 + n_cond = 2 + B_nf = n_cond * 12 + + torch.manual_seed(42) + cln = _make_new_cln(cond_shape, C) + + x = torch.randn(B_nf, C, H, W, device="cuda") + cond = torch.randn(n_cond, cond_shape, device="cuda") + + if channels_last: + x = x.to(memory_format=torch.channels_last) + + x = x.requires_grad_(True) + cond = cond.requires_grad_(True) + + out = cln(x, cond) + out.sum().backward() + + assert x.grad is not None, "No gradient for input x" + assert cond.grad is not None, "No gradient for conditions" + assert torch.isfinite(x.grad).all(), "Non-finite input gradients" + assert torch.isfinite(cond.grad).all(), "Non-finite condition gradients" + + for name, p in cln.named_parameters(): + assert p.grad is not None, f"No gradient for {name}" + assert torch.isfinite(p.grad).all(), f"Non-finite gradient for {name}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_backward_channels_last_matches_contiguous(): + """Verify channels_last and contiguous inputs produce the same gradients.""" + C, H, W = 64, 8, 8 + cond_shape = 16 + n_cond = 2 + B_nf = n_cond * 12 + + torch.manual_seed(42) + cln = _make_new_cln(cond_shape, C) + + x_base = torch.randn(B_nf, C, H, W, device="cuda") + cond_base = torch.randn(n_cond, cond_shape, device="cuda") + + # Contiguous path + x_cont = x_base.clone().detach().requires_grad_(True) + cond_cont = cond_base.clone().detach().requires_grad_(True) + out_cont = cln(x_cont, cond_cont) + out_cont.sum().backward() + + cln.zero_grad() + + # Channels-last path + x_cl = x_base.clone().detach().to(memory_format=torch.channels_last).requires_grad_(True) + cond_cl = cond_base.clone().detach().requires_grad_(True) + out_cl = cln(x_cl, cond_cl) + out_cl.sum().backward() + + assert torch.allclose(out_cont, out_cl, atol=1e-5, rtol=1e-4), \ + f"Output max diff: {(out_cont - out_cl).abs().max().item()}" + assert torch.allclose(x_cont.grad, x_cl.grad, atol=1e-5, rtol=1e-4), \ + f"Input grad max diff: {(x_cont.grad - x_cl.grad).abs().max().item()}" + assert torch.allclose(cond_cont.grad, cond_cl.grad, atol=1e-5, rtol=1e-4), \ + f"Cond grad max diff: {(cond_cont.grad - cond_cl.grad).abs().max().item()}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_load_old_checkpoint(): + """Verify new CLN can load old-format state dict via _load_from_state_dict.""" + C, cond_shape = 64, 16 + torch.manual_seed(42) + old_cln = _make_old_cln(cond_shape, C) + old_sd = old_cln.state_dict() + + new_cln = _make_new_cln(cond_shape, C) + new_cln.load_state_dict(old_sd, strict=False) + + # Verify outputs match after loading old checkpoint + x = torch.randn(12, C, 8, 8, device="cuda") + cond = torch.randn(1, cond_shape, device="cuda") + + with torch.no_grad(): + out_old = old_cln(x, cond) + out_new = new_cln(x, cond) + + assert out_new.shape == (12, C, 8, 8) + assert torch.isfinite(out_new).all() + assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), \ + f"Max diff after loading old checkpoint: {(out_old - out_new).abs().max().item()}" diff --git a/test/pytest_utils.py b/test/pytest_utils.py index 51c21c2188..1f00fded44 100644 --- a/test/pytest_utils.py +++ b/test/pytest_utils.py @@ -54,3 +54,93 @@ def modify_environment(*remove, **update): finally: env.update(restore_after) [env.pop(k, None) for k in purge_after] + + +import importlib +from functools import wraps + +import pytest +from packaging.version import Version + + +def import_or_fail( + module_names: str | list[str] | tuple, + min_versions: str | list[str] | tuple | None = None, +): + """Skip test if module is missing or below minimum version.""" + + def decorator(test_func): + @pytest.mark.usefixtures("pytestconfig") + @wraps(test_func) + def wrapper(*args, **kwargs): + pytestconfig = kwargs.get("pytestconfig") + if pytestconfig is None: + raise ValueError( + "pytestconfig must be passed when using import_or_fail." + ) + _import_or_fail(module_names, pytestconfig, min_versions) + return test_func(*args, **kwargs) + + return wrapper + + return decorator + + +def _import_or_fail(module_names, config, min_versions=None): + if not isinstance(module_names, (list, tuple)): + module_names = [module_names] + if min_versions is not None and not isinstance(min_versions, (list, tuple)): + min_versions = [min_versions] + + if min_versions is None: + min_versions = [None] * len(module_names) + elif len(min_versions) != len(module_names): + raise ValueError("module_names and min_versions must have the same length.") + + for module_name, min_version in zip(module_names, min_versions): + if config.getoption("--fail-on-missing-modules"): + __import__(module_name) + else: + try: + module = importlib.import_module(module_name) + if hasattr(module, "__version__"): + if isinstance(module.__version__, str) or module.__version__ is None: + pytest.importorskip(module_name, min_version) + elif isinstance(module.__version__, Version): + version_check = Version(min_version) + if module.__version__ < version_check: + pytest.skip( + f"{module_name} {module.__version__} is less than " + f"required version {version_check}" + ) + except ModuleNotFoundError: + pytest.importorskip(module_name, min_version) + + +def nfsdata_or_fail(test_func): + @pytest.mark.usefixtures("pytestconfig") + @wraps(test_func) + def wrapper(*args, **kwargs): + pytestconfig = kwargs.get("pytestconfig") + if pytestconfig is None: + raise ValueError("pytestconfig must be passed when using nfsdata_or_fail.") + _nfsdata_or_fail(pytestconfig) + return test_func(*args, **kwargs) + + return wrapper + + +def _nfsdata_or_fail(config): + nfs_data_dir_opt = config.getoption("--nfs-data-dir") + test_data_dir_env = os.environ.get("TEST_DATA_DIR") + if nfs_data_dir_opt and os.path.exists(nfs_data_dir_opt): + return + if test_data_dir_env and os.path.exists(test_data_dir_env): + return + for path in ("/data/nfs/physicsnemo-data", "/data/nfs/modulus-data"): + if os.path.exists(path): + return + pytest.skip( + "NFS volumes not set up with CI data repo. Run `make get-data` from the " + "root directory of the repo" + ) From 624ddb91c8e9f3bef12d852d45aaf96d22cc51a0 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 1 Jul 2026 10:20:39 -0500 Subject: [PATCH 02/28] Fix coupler, use OptionalImport --- physicsnemo/datapipes/healpix/couplers.py | 7 ++----- physicsnemo/metrics/climate/healpix_loss.py | 20 ++++++-------------- 2 files changed, 8 insertions(+), 19 deletions(-) diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 607b14dcfc..fa08560e23 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -408,11 +408,8 @@ def set_coupled_fields(self, coupled_fields: th.tensor): [coupled_fields.shape[0], self.spatial_dims[0], self.coupled_integration_dim, self.timevar_dim] + list(self.spatial_dims[1:]) ) - # we use a constant set of values so we just copy time 0 - for i in range(self.coupled_integration_dim): - self.preset_coupled_fields[:, :, i, :, :, :] = coupled_fields[ - :, :, 0, -1:, :, : - ] + # broadcast the first time step to all the integration steps + self.preset_coupled_fields[:, :, :, :, :, :] = coupled_fields[:, :, :1, :, :, :] if self.time_first: self.preset_coupled_fields = self.preset_coupled_fields.permute(2, 0, 3, 1, 4, 5) # flag for construct integrated coupling method to use this array diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index ef759ad924..319baf8b0e 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -22,16 +22,8 @@ from physicsnemo.core.version_check import OptionalImport xr = OptionalImport("xarray") -earth2grid = OptionalImport("earth2grid") - -try: - from cuhpx import SHTCUDA, iSHTCUDA - from earth2grid.healpix import HEALPIX_PAD_XY, PixelOrder -except ImportError: - SHTCUDA = None - iSHTCUDA = None - HEALPIX_PAD_XY = None - PixelOrder = None +earth2grid = OptionalImport("earth2grid", "Install earth2grid from https://github.com/earth2grid/earth2grid") +cuhpx = OptionalImport("cuhpx", "Install cuhpx from https://github.com/NVIDIA/cuhpx") """ Custom dlwp compatible loss classes that allow for more sophisticated training optimization. @@ -650,13 +642,13 @@ def __init__( self.lmax = lmax self.mmax = mmax self.nside = nside - self.sht = SHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') - src_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=HEALPIX_PAD_XY) - tar_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=PixelOrder.RING) + self.sht = cuhpx.SHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') + src_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=earth2grid.healpix.HEALPIX_PAD_XY) + tar_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=earth2grid.healpix.PixelOrder.RING) self.reorder_to_ring = earth2grid.get_regridder(src_grid, tar_grid).to(th.float32) if self.multiscale > 0: self.scales = [200, 400, 800, 1600] # in units of km - self.isht = iSHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') + self.isht = cuhpx.iSHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') self.reorder_from_ring = earth2grid.get_regridder(tar_grid, src_grid).to(th.float32) self.lsm_file = lsm_file From 2e746b26b56ae3dc9ba90adfe68e2ac1b9870ea6 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 1 Jul 2026 10:39:07 -0500 Subject: [PATCH 03/28] remove unnecessary import --- physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py | 3 +-- physicsnemo/datapipes/healpix/timeseries_dataset.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py index 39f55dca62..ccb379b915 100644 --- a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py @@ -28,7 +28,6 @@ from physicsnemo.datapipes.datapipe import Datapipe from physicsnemo.datapipes.meta import DatapipeMetaData from omegaconf import DictConfig, OmegaConf -from torch.utils.data import Dataset logger = logging.getLogger(__name__) @@ -75,7 +74,7 @@ def _check_availability(path: str) -> None: # pragma: no cover raise FileNotFoundError(f"Dataset not found at specified location: {path}") -class BaseTimeSeriesDatasetZarr(Dataset, Datapipe, ABC): +class BaseTimeSeriesDatasetZarr(Datapipe, ABC): """Abstract base class for time series datasets using Zarr storage. This class provides the core functionality for loading and processing time series data diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index 33d30e6e47..eca11a9ab9 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -25,7 +25,6 @@ import torch import xarray as xr from omegaconf import DictConfig, OmegaConf -from torch.utils.data import Dataset from physicsnemo.datapipes.datapipe import Datapipe from physicsnemo.datapipes.meta import DatapipeMetaData @@ -46,7 +45,7 @@ class MetaData(DatapipeMetaData): ddp_sharding: bool = False -class TimeSeriesDataset(Dataset, Datapipe): +class TimeSeriesDataset(Datapipe): """ Dataset for sampling from continuous time-series data, compatible with pytorch data loading. """ From e0717bcf735d5457c4f3bb6ab6df52e640f7505a Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 1 Jul 2026 12:40:46 -0500 Subject: [PATCH 04/28] update amp calls --- physicsnemo/metrics/climate/healpix_loss.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index 319baf8b0e..a62eb31b80 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -778,7 +778,7 @@ def forward(self, prediction, target, average_channels=True): if self.lambda_spec > 0: - with th.cuda.amp.autocast(enabled=False): + with th.amp.autocast("cuda",enabled=False): # # Reorder predictions: [N, B, F, T, C, H, W] -> [N, B, T, C, F*H*W] # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) @@ -809,7 +809,7 @@ def forward(self, prediction, target, average_channels=True): if self.multiscale > 0: for scale in self.scales: l_filter = self._l_filter(scale, device=prediction.device) - with th.cuda.amp.autocast(enabled=False): + with th.amp.autocast("cuda",enabled=False): sht_pred= self._apply_sht(prediction, face_dim=2, return_abs=False) sht_tar = self._apply_sht(target, face_dim=1, return_abs=False) @@ -849,7 +849,7 @@ def forward(self, prediction, target, average_channels=True): loss = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1,2))/(b*f*t*h*w) if self.lambda_spec > 0: - with th.cuda.amp.autocast(enabled=False): + with th.amp.autocast("cuda",enabled=False): # # Reorder predictions: [C, Cond, B, F, T, H, W] -> [C, Cond, B, T, F*H*W] # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 2, 4, 3, 5, 6).reshape(c, n, b, t, f*h*w)) @@ -881,7 +881,7 @@ def forward(self, prediction, target, average_channels=True): loss = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum()/(b*f*c*t*h*w) if self.lambda_spec > 0: - with th.cuda.amp.autocast(enabled=False): + with th.amp.autocast("cuda",enabled=False): # # Reorder predictions: [Cond, B, F, T, C, H, W] -> [Cond, B, T, C, F*H*W] # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) From af7e810b85117bdf1e09836cb1aaaeef808c6800 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Thu, 2 Jul 2026 15:22:42 -0500 Subject: [PATCH 05/28] Improve handling for checkpoints with older normalization --- .../dlwp_healpix/layers/normalization.py | 26 ++++++++++ test/models/dlwp_healpix/_cln_reference.py | 6 ++- .../test_conditional_layer_norm.py | 48 +++++++++++++++++++ 3 files changed, 79 insertions(+), 1 deletion(-) diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index 9461bd37d5..524694ca0a 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -126,6 +126,9 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss [0, beta_weight]] Biases are always concatenated: cat([gamma_bias, beta_bias]). + + Activation submodule buffers (e.g. CappedGELU ``cap``) are copied from + the gamma MLP to the fused MLP at the same sequential index. """ gamma_prefix = prefix + "gamma_mlp." beta_prefix = prefix + "beta_mlp." @@ -147,6 +150,7 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss keys_to_remove = [] keys_to_add = {} + # handle weights and biases for the gamma MLP for k in list(state_dict.keys()): if k.startswith(gamma_prefix): layer_key = k[len(gamma_prefix):] # e.g. "0.weight" @@ -186,6 +190,28 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss if beta_key not in keys_to_remove: keys_to_remove.append(beta_key) + # handle activation submodule buffers + for k in list(state_dict.keys()): + if not k.startswith(gamma_prefix): + continue + layer_key = k[len(gamma_prefix):] # e.g. "0.weight" + parts = layer_key.split(".", 1) + if len(parts) != 2: + continue + try: + int(parts[0]) + except ValueError: + continue + if parts[1] in ("weight", "bias"): + continue + + fused_key = fused_prefix + layer_key + keys_to_add[fused_key] = state_dict[k] + keys_to_remove.append(k) + beta_key = beta_prefix + layer_key + if beta_key in state_dict and beta_key not in keys_to_remove: + keys_to_remove.append(beta_key) + for k in keys_to_remove: if k in state_dict: del state_dict[k] diff --git a/test/models/dlwp_healpix/_cln_reference.py b/test/models/dlwp_healpix/_cln_reference.py index 1524062627..e27227ac54 100644 --- a/test/models/dlwp_healpix/_cln_reference.py +++ b/test/models/dlwp_healpix/_cln_reference.py @@ -3,6 +3,8 @@ from typing import List +import copy + import torch as th try: @@ -54,7 +56,9 @@ def _make_mlp(self, in_dim: int, hidden_dims: List[int], out_dim: int, activatio for hdim in hidden_dims: layers.append(th.nn.Linear(in_dim, hdim)) if activation: - layers.append(activation) + # some variations of _make_mlp may have duplicate activation submodule buffers (e.g. CappedGELU ``cap``) + # so we need to deepcopy the activation submodule buffers + layers.append(copy.deepcopy(activation)) in_dim = hdim layers.append(th.nn.Linear(in_dim, out_dim)) return th.nn.Sequential(*layers) diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py index 5b6f729663..25e45265ba 100644 --- a/test/models/dlwp_healpix/test_conditional_layer_norm.py +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -7,6 +7,7 @@ import torch from _cln_reference import ConditionalLayerNormReference from physicsnemo.models.dlwp_healpix.layers.normalization import ConditionalLayerNorm +from physicsnemo.nn import CappedGELU def _make_old_cln(condition_shape, channel_depth, **kwargs): @@ -67,6 +68,15 @@ def _copy_old_to_new(old_cln, new_cln): torch.cat([zeros, beta_val], dim=1), ], dim=0) + for key, value in old_sd.items(): + if not key.startswith("gamma_mlp."): + continue + layer_key = key[len("gamma_mlp.") :] + parts = layer_key.split(".", 1) + if len(parts) != 2 or parts[1] in ("weight", "bias"): + continue + new_sd[f"gamma_beta_mlp.{layer_key}"] = value + new_cln.load_state_dict(new_sd) @@ -263,3 +273,41 @@ def test_load_old_checkpoint(): assert torch.isfinite(out_new).all() assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), \ f"Max diff after loading old checkpoint: {(out_old - out_new).abs().max().item()}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_load_old_checkpoint_with_capped_gelu(): + """Verify strict loading of old CLN checkpoints that use CappedGELU activations.""" + C, cond_shape = 64, 16 + mlp_hidden_dims = [32, 32] + activation = CappedGELU() + + torch.manual_seed(42) + old_cln = _make_old_cln( + cond_shape, + C, + mlp_hidden_dims=mlp_hidden_dims, + activation=activation, + ) + old_sd = old_cln.state_dict() + assert any(k.endswith(".cap") for k in old_sd if k.startswith("gamma_mlp.")) + + new_cln = _make_new_cln( + cond_shape, + C, + mlp_hidden_dims=mlp_hidden_dims, + activation=CappedGELU(), + ) + missing, unexpected = new_cln.load_state_dict(old_sd, strict=True) + assert not missing + assert not unexpected + + x = torch.randn(12, C, 8, 8, device="cuda") + cond = torch.randn(1, cond_shape, device="cuda") + + with torch.no_grad(): + out_old = old_cln(x, cond) + out_new = new_cln(x, cond) + + assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), \ + f"Max diff after loading old CappedGELU checkpoint: {(out_old - out_new).abs().max().item()}" From 33f0c7b7e46da41f705aa6235b993083fcf4b278 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 8 Jul 2026 10:12:34 -0500 Subject: [PATCH 06/28] Linting and format cleanup --- .../healpix/base_timeseries_dataset_zarr.py | 4 +- .../healpix/coupledtimeseries_dataset.py | 15 +- .../healpix/coupledtimeseries_dataset_zarr.py | 2 +- physicsnemo/datapipes/healpix/couplers.py | 31 +- physicsnemo/datapipes/healpix/data_modules.py | 6 +- .../datapipes/healpix/timeseries_dataset.py | 4 +- physicsnemo/metrics/climate/healpix_loss.py | 490 ++++++++++++------ physicsnemo/metrics/climate/hydrostasy.py | 85 +-- .../models/dlwp_healpix/HEALPixRecUNet.py | 6 + .../models/dlwp_healpix/HEALPixUNet.py | 6 + .../dlwp_healpix/layers/healpix_blocks.py | 65 ++- .../layers/healpix_constraints.py | 20 +- .../dlwp_healpix/layers/healpix_decoder.py | 4 +- .../dlwp_healpix/layers/normalization.py | 71 ++- 14 files changed, 529 insertions(+), 280 deletions(-) diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py index ccb379b915..51fd9b81c1 100644 --- a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py @@ -25,9 +25,10 @@ import pandas as pd import xarray as xr import zarr +from omegaconf import DictConfig, OmegaConf + from physicsnemo.datapipes.datapipe import Datapipe from physicsnemo.datapipes.meta import DatapipeMetaData -from omegaconf import DictConfig, OmegaConf logger = logging.getLogger(__name__) @@ -473,7 +474,6 @@ def _get_scaling_da(self) -> None: ) try: - self.constant_scaling = scaling_da.sel( index=self.constant_variables ).rename({"index": "channel_out"}) diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index 8acc22f126..3cba485113 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -208,12 +208,17 @@ def __getitem__(self, item): ) # for models with extra outputs - if len(self.ds["targets"].channel_out) != (len(self.ds["inputs"].channel_in)-len(self.couplings[0].variables)): - input_array = (input_array - self.input_scaling["mean"][:,:-len(self.couplings[0].variables)]) \ - / self.input_scaling["std"][:,:-len(self.couplings[0].variables)] + if len(self.ds["targets"].channel_out) != ( + len(self.ds["inputs"].channel_in) - len(self.couplings[0].variables) + ): + input_array = ( + input_array + - self.input_scaling["mean"][:, : -len(self.couplings[0].variables)] + ) / self.input_scaling["std"][:, : -len(self.couplings[0].variables)] else: - input_array = (input_array - self.input_scaling["mean"]) \ - / self.input_scaling["std"] + input_array = ( + input_array - self.input_scaling["mean"] + ) / self.input_scaling["std"] if not self.forecast_mode: # BAD NEWS: Indexing the array as commented out below causes unexpected behavior in target creation. diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py index 65f7df4f66..4c4a71403c 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py @@ -28,8 +28,8 @@ from physicsnemo.utils.insolation import insolation from . import couplers -from .timeseries_dataset_zarr import TimeSeriesDatasetZarr from .base_timeseries_dataset_zarr import _check_availability +from .timeseries_dataset_zarr import TimeSeriesDatasetZarr logger = logging.getLogger(__name__) diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index fa08560e23..1d77761734 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -57,7 +57,7 @@ def __init__( forecasting batch size should be 1 variables: Sequence sequence of strings that indicate the coupled variable - names in the dataset. All names should be in the dataset with + names in the dataset. All names should be in the dataset with an optional time component at the end, eg ttr-48h presteps: int, optional the number of model steps used to initialize the hidden state. @@ -75,7 +75,7 @@ def __init__( the right side of a averaging_window window. This is highly recommended for training, default True time_first: boolean, optional - Whether the coupled data should be permuted to have the time dimension first + Whether the coupled data should be permuted to have the time dimension first [T, B, C, F, H, W] rather than [B, F, T, C, H, W] """ # extract important meta data from ds @@ -101,9 +101,9 @@ def __init__( if not prepared_coupled_data: raise NotImplementedError("Data preparation not yet implemented") - if type(self.ds) == xr.Dataset: + if isinstance(self.ds, xr.Dataset): self.use_zarr = False - elif type(self.ds) == zr.Group: + elif isinstance(self.ds, zr.Group): self.use_zarr = True self.ds_variable_indices = [ i @@ -401,17 +401,22 @@ def set_coupled_fields(self, coupled_fields: th.tensor): format is [B, F, T, C, H, W] """ # create buffer for coupling - coupled_fields = coupled_fields[ - :, :, :, self.coupled_channel_indices, :, : - ] + coupled_fields = coupled_fields[:, :, :, self.coupled_channel_indices, :, :] self.preset_coupled_fields = th.empty( - [coupled_fields.shape[0], self.spatial_dims[0], self.coupled_integration_dim, self.timevar_dim] + [ + coupled_fields.shape[0], + self.spatial_dims[0], + self.coupled_integration_dim, + self.timevar_dim, + ] + list(self.spatial_dims[1:]) ) # broadcast the first time step to all the integration steps self.preset_coupled_fields[:, :, :, :, :, :] = coupled_fields[:, :, :1, :, :, :] if self.time_first: - self.preset_coupled_fields = self.preset_coupled_fields.permute(2, 0, 3, 1, 4, 5) + self.preset_coupled_fields = self.preset_coupled_fields.permute( + 2, 0, 3, 1, 4, 5 + ) # flag for construct integrated coupling method to use this array self.coupled_mode = True @@ -572,10 +577,10 @@ def set_coupled_fields(self, coupled_fields: th.tensor): for s in self.averaging_slices[j] ] coupled_averaging_periods.append(th.concat(averaging_periods, dim=3)) - self.preset_coupled_fields = th.concat( - coupled_averaging_periods, dim=2 - ) + self.preset_coupled_fields = th.concat(coupled_averaging_periods, dim=2) if self.time_first: - self.preset_coupled_fields = self.preset_coupled_fields.permute(2, 0, 3, 1, 4, 5) + self.preset_coupled_fields = self.preset_coupled_fields.permute( + 2, 0, 3, 1, 4, 5 + ) # flag for construct integrated coupling method to use this array self.coupled_mode = True diff --git a/physicsnemo/datapipes/healpix/data_modules.py b/physicsnemo/datapipes/healpix/data_modules.py index 7d8fc1fa77..1e67017a39 100644 --- a/physicsnemo/datapipes/healpix/data_modules.py +++ b/physicsnemo/datapipes/healpix/data_modules.py @@ -16,6 +16,7 @@ # System modules import logging +import warnings from pathlib import Path from typing import Optional, Sequence, Union @@ -25,8 +26,6 @@ # distributed stuff import xarray as xr -import warnings - # External modules from omegaconf import DictConfig from torch.utils.data import DataLoader @@ -174,7 +173,7 @@ def __init__( warnings.warn( "TimeSeriesDataModule will be removed in a future release, please switch to using TimeSeriesDataModuleZarr", DeprecationWarning, - stacklevel=2 + stacklevel=2, ) self.dst_directory = dst_directory @@ -597,7 +596,6 @@ def __init__( ) def _get_coupled_vars(self): - coupled_variables = [] for d in self.couplings: coupled_variables = coupled_variables + d["params"]["variables"] diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index eca11a9ab9..157eb7170c 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -243,7 +243,9 @@ def _get_scaling_da(self): # 'if' statement used for cases where atmos model # includes diagnostic variables like tp6 and msl. # using 'channel_out' is still necessary for ocean models. - if len(self.ds.channel_out) != (len(self.ds.channel_in)-len(self.couplings[0].variables)): + if len(self.ds.channel_out) != ( + len(self.ds.channel_in) - len(self.couplings[0].variables) + ): self.input_scaling = scaling_da.sel( index=self.ds.channel_in.values ).rename({"index": "channel_in"}) diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index a62eb31b80..44b1e5f416 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -22,7 +22,9 @@ from physicsnemo.core.version_check import OptionalImport xr = OptionalImport("xarray") -earth2grid = OptionalImport("earth2grid", "Install earth2grid from https://github.com/earth2grid/earth2grid") +earth2grid = OptionalImport( + "earth2grid", "Install earth2grid from https://github.com/earth2grid/earth2grid" +) cuhpx = OptionalImport("cuhpx", "Install cuhpx from https://github.com/NVIDIA/cuhpx") """ @@ -81,7 +83,6 @@ def forward(self, prediction, target, average_channels=True): class WeightedMSE(th.nn.MSELoss): - """ Loss object that allows for user defined weighting of variables when calculating MSE """ @@ -134,7 +135,8 @@ def forward(self, prediction, target, average_channels=True): else: return d -class ConditionalWeightLoss( th.nn.MSELoss ): + +class ConditionalWeightLoss(th.nn.MSELoss): """ Conditional loss for precipitation diagnostic model. (Total 6hr precipitation is the only output field.) @@ -142,10 +144,10 @@ class ConditionalWeightLoss( th.nn.MSELoss ): def __init__( self, - weight=(0.01,1.0), + weight=(0.01, 1.0), b=None, w=1, - ): + ): """ Parameters ----------- @@ -181,11 +183,14 @@ def forward(self, prediction, target): The target tensor """ weights_for_zero = th.ones_like(target) * self.weight_zero - weights_for_nonzero = (th.ones_like(target) * self.weight_nonzero) * th.exp(self.b*target) + weights_for_nonzero = (th.ones_like(target) * self.weight_nonzero) * th.exp( + self.b * target + ) weights = th.where(target > 0, weights_for_nonzero, weights_for_zero) - loss = (th.mean(weights * (prediction - target) ** 2))*self.w + loss = (th.mean(weights * (prediction - target) ** 2)) * self.w return loss + class OceanMSE(th.nn.MSELoss): """ Ocean MSE class offers impementaion for MSE loss weighted by a land-sea-mask field. @@ -231,7 +236,6 @@ def setup(self, trainer): ).to(trainer.device) def forward(self, prediction, target, average_channels=True): - if not self.lsm_sum_calculated: self.lsm_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum() self.lsm_var_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum( @@ -296,7 +300,6 @@ def setup(self, trainer): self.loss_weights = self.loss_weights.to(device=trainer.device) def forward(self, prediction, target, average_channels=True): - if not self.lsm_sum_calculated: self.lsm_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum() self.lsm_var_sum = th.broadcast_to(self.lsm_tensor, target.shape).sum( @@ -313,8 +316,8 @@ def forward(self, prediction, target, average_channels=True): else: return ocean_mean_err / self.lsm_var_sum -class WeightedCRPSLoss(th.nn.MSELoss): +class WeightedCRPSLoss(th.nn.MSELoss): """ Probabilistic loss function that allows for user defined weighting of variables when calculating CRPS. """ @@ -359,7 +362,7 @@ def __init__( self.loss_weights = th.tensor(weights) if n_members < 2: raise ValueError("n_members must be at least 2 for CRPS loss to be defined") - else: + else: self.n_members = n_members self.device = None self.mean_penalty = mean_penalty @@ -368,19 +371,27 @@ def __init__( self.masked_processing = masked_processing if lsm_file is not None: - self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel(selection_dict) - self.lsm_tensor = 1 - th.tensor(np.expand_dims(self.lsm_ds.values, (0, 2, 3))) + self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel( + selection_dict + ) + self.lsm_tensor = 1 - th.tensor( + np.expand_dims(self.lsm_ds.values, (0, 2, 3)) + ) else: - self.lsm_tensor = th.ones(1, 1, 1, 1, 1, 1) # Spoof the tensor dimensions for broadcasting + self.lsm_tensor = th.ones( + 1, 1, 1, 1, 1, 1 + ) # Spoof the tensor dimensions for broadcasting # Parameters for "almost fair CRPS" loss. See https://arxiv.org/html/2412.15832v1 - self.coeff_eps = 1 - ((1-alpha) / (n_members)) - self.averaging_coeff = 1 / (2* n_members * (n_members - 1)) + self.coeff_eps = 1 - ((1 - alpha) / (n_members)) + self.averaging_coeff = 1 / (2 * n_members * (n_members - 1)) # For n>2, will use pairwise distance to copmute [NxN] distance matrix # Diagonal elements of (prediciton - target) matrix are zeroed out to avoid double counting self.pdist = th.nn.PairwiseDistance(p=1) - self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye(self.n_members) # Mask to zero out diagonal elements + self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye( + self.n_members + ) # Mask to zero out diagonal elements def setup(self, trainer): """ @@ -394,26 +405,32 @@ def setup(self, trainer): self.averaging_coeff = th.tensor(self.averaging_coeff, device=trainer.device) self.coeff_eps = th.tensor(self.coeff_eps, device=trainer.device) self.pdist = self.pdist.to(device=trainer.device) - self.diag_mask = self.diag_mask.to(device=trainer.device) + self.diag_mask = self.diag_mask.to(device=trainer.device) self.lsm_tensor = self.lsm_tensor.to(device=trainer.device) def _2member_crps(self, prediction, target, lsm_tensor): - diff_target = th.abs(prediction - target.unsqueeze(0)).sum(dim=0) # [B, F, T, C, H, W] - diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] - crps = self.averaging_coeff*(diff_target - self.coeff_eps * diff_ensemble) # [B, F, T, C, H, W] + diff_target = th.abs(prediction - target.unsqueeze(0)).sum( + dim=0 + ) # [B, F, T, C, H, W] + diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] + crps = self.averaging_coeff * ( + diff_target - self.coeff_eps * diff_ensemble + ) # [B, F, T, C, H, W] crps *= lsm_tensor return crps def _pool(self, tensor, scale): shape = tensor.shape h, w = shape[-2:] - pooled = th.nn.functional.avg_pool2d(tensor.reshape(shape[0], -1, h, w), scale, scale) - return pooled.reshape(*shape[:-2], h//scale, w//scale) + pooled = th.nn.functional.avg_pool2d( + tensor.reshape(shape[0], -1, h, w), scale, scale + ) + return pooled.reshape(*shape[:-2], h // scale, w // scale) def _masked_pool(self, tensor, scale, mask): """ Pools a tensor while excluding masked pixels - + Parameters ---------- tensor: torch.Tensor @@ -439,12 +456,12 @@ def _masked_pool(self, tensor, scale, mask): # pool the mask to use as a weight for the pooled tensor pooled_mask = self._pool(mask, scale) - pooled_tensor = pooled_values / (pooled_mask + 1e-8) # Avoid division by zero + pooled_tensor = pooled_values / (pooled_mask + 1e-8) # Avoid division by zero return pooled_tensor, pooled_mask def forward(self, prediction, target, average_channels=True): """ - Forward pass of the WeightedCRPSLoss + Forward pass of the WeightedCRPSLoss Computes the CRPS loss for the model prediction and target. Parameters @@ -456,24 +473,28 @@ def forward(self, prediction, target, average_channels=True): average_channels: bool, optional whether the mean of the channels should be taken """ - + # Unfold ensemble dimension from batch dimension to have shape [Cond, B, F, T, C, H, W] b, f, t, c, h, w = target.shape prediction = prediction.view(self.n_members, b, f, t, c, h, w) - # checks for dimensions + # checks for dimensions if not prediction.shape[1:] == target.shape: - raise ValueError(f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}") - + raise ValueError( + f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}" + ) + if not prediction.shape[0] == self.n_members: - raise ValueError(f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}") + raise ValueError( + f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}" + ) n = self.n_members - + # Manual Cast prediction = prediction.to(th.float32) target = target.to(th.float32) - + # Apply channel weights across channel dims prediction *= self.loss_weights[None, None, None, None, :, None, None] target *= self.loss_weights[None, None, None, :, None, None] @@ -493,28 +514,44 @@ def forward(self, prediction, target, average_channels=True): valid_pixels = self.lsm_tensor.numel() else: valid_pixels = f * h * w - + prediction_count = n * b * t * valid_pixels target_count = b * t * valid_pixels - ens_global_means = (prediction * self.lsm_tensor).sum(dim=(0, 1, 2, 3, 5, 6)) / prediction_count - target_global_means = (target * self.lsm_tensor).sum(dim=(0, 1, 2, 4, 5)) / target_count - - bias_penalty = self.mean_penalty * th.abs(ens_global_means - target_global_means) + ens_global_means = (prediction * self.lsm_tensor).sum( + dim=(0, 1, 2, 3, 5, 6) + ) / prediction_count + target_global_means = (target * self.lsm_tensor).sum( + dim=(0, 1, 2, 4, 5) + ) / target_count + + bias_penalty = self.mean_penalty * th.abs( + ens_global_means - target_global_means + ) if average_channels: loss += bias_penalty.mean() else: loss += bias_penalty # spatial multiscale loss - if self.multiscale > 0.: + if self.multiscale > 0.0: crps_scales = 0 for scale in self.scales: if self.masked_processing and self.lsm_tensor.numel() > 1: - masked_pred, masked_lsm = self._masked_pool(prediction, scale, self.lsm_tensor) - masked_tar, _ = self._masked_pool(target, scale, self.lsm_tensor) - crps_scale = self._2member_crps(masked_pred, masked_tar, masked_lsm) + masked_pred, masked_lsm = self._masked_pool( + prediction, scale, self.lsm_tensor + ) + masked_tar, _ = self._masked_pool( + target, scale, self.lsm_tensor + ) + crps_scale = self._2member_crps( + masked_pred, masked_tar, masked_lsm + ) else: - pred, tar, lsm = self._pool(prediction, scale), self._pool(target, scale), self._pool(self.lsm_tensor, scale) + pred, tar, lsm = ( + self._pool(prediction, scale), + self._pool(target, scale), + self._pool(self.lsm_tensor, scale), + ) crps_scale = self._2member_crps(pred, tar, lsm) if average_channels: @@ -525,7 +562,7 @@ def forward(self, prediction, target, average_channels=True): crps_scales = crps_scales / len(self.scales) loss += self.multiscale * crps_scales - + return loss else: # Do mean penalty @@ -536,13 +573,19 @@ def forward(self, prediction, target, average_channels=True): valid_pixels = self.lsm_tensor.numel() else: valid_pixels = f * h * w - + prediction_count = n * b * t * valid_pixels target_count = b * t * valid_pixels - ens_global_means = (prediction * self.lsm_tensor).sum(dim=(0, 1, 2, 3, 5, 6)) / prediction_count - target_global_means = (target * self.lsm_tensor).sum(dim=(0, 1, 2, 4, 5)) / target_count - - bias_penalty = self.mean_penalty * th.abs(ens_global_means - target_global_means) + ens_global_means = (prediction * self.lsm_tensor).sum( + dim=(0, 1, 2, 3, 5, 6) + ) / prediction_count + target_global_means = (target * self.lsm_tensor).sum( + dim=(0, 1, 2, 4, 5) + ) / target_count + + bias_penalty = self.mean_penalty * th.abs( + ens_global_means - target_global_means + ) if average_channels: bias_penalty = bias_penalty.mean() @@ -556,32 +599,53 @@ def forward(self, prediction, target, average_channels=True): # Use pairwise distance method if not average_channels: - # Move channels to first dimension and exclude that dimension from the reductions - prediction = prediction.permute(4, 0, 1, 2, 3, 5, 6) # [C, Cond, B, F, T, H, W] - target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] - - prediction = prediction.reshape(c, n, -1) # [C, Cond, ...] - target = target.unsqueeze(1).reshape(c, 1, -1) # [C, 1, ...] (second dim will broadcast across ensemble) - - diff = self.pdist(prediction, target) # [C, Cond] - dist_matrix = self.pdist(prediction.unsqueeze(1), prediction.unsqueeze(2)) # [C, Cond, Cond] - - diff_terms = self.diag_mask[None, ...] * (diff.unsqueeze(1) + diff.unsqueeze(2)) # [C, Cond, Cond], diagonal elements zeroed out - crps = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1,2))/valid_pixels + # Move channels to first dimension and exclude that dimension from the reductions + prediction = prediction.permute( + 4, 0, 1, 2, 3, 5, 6 + ) # [C, Cond, B, F, T, H, W] + target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] + + prediction = prediction.reshape(c, n, -1) # [C, Cond, ...] + target = target.unsqueeze(1).reshape( + c, 1, -1 + ) # [C, 1, ...] (second dim will broadcast across ensemble) + + diff = self.pdist(prediction, target) # [C, Cond] + dist_matrix = self.pdist( + prediction.unsqueeze(1), prediction.unsqueeze(2) + ) # [C, Cond, Cond] + + diff_terms = self.diag_mask[None, ...] * ( + diff.unsqueeze(1) + diff.unsqueeze(2) + ) # [C, Cond, Cond], diagonal elements zeroed out + crps = ( + self.averaging_coeff + * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1, 2)) + / valid_pixels + ) else: prediction = prediction.reshape(n, -1) - target = target.unsqueeze(0).reshape(1, -1) # [1, ...] (first dim will broadcast across ensemble) - diff = self.pdist(prediction, target) # [Cond] - dist_matrix = self.pdist(prediction.unsqueeze(1), prediction.unsqueeze(0)) # [Cond, Cond] - - diff_terms = self.diag_mask * (diff.unsqueeze(0) + diff.unsqueeze(1)) # [Cond, Cond], diagonal elements zeroed out - crps = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum()/valid_pixels + target = target.unsqueeze(0).reshape( + 1, -1 + ) # [1, ...] (first dim will broadcast across ensemble) + diff = self.pdist(prediction, target) # [Cond] + dist_matrix = self.pdist( + prediction.unsqueeze(1), prediction.unsqueeze(0) + ) # [Cond, Cond] + + diff_terms = self.diag_mask * ( + diff.unsqueeze(0) + diff.unsqueeze(1) + ) # [Cond, Cond], diagonal elements zeroed out + crps = ( + self.averaging_coeff + * (diff_terms - self.coeff_eps * dist_matrix).sum() + / valid_pixels + ) return crps + bias_penalty class WeightedCRPSLossSpectral(th.nn.MSELoss): - """ Probabilistic loss function that allows for user defined weighting of variables when calculating CRPS. """ @@ -593,8 +657,8 @@ def __init__( alpha: float = 0.95, lambda_spec: float = 0.1, nside: int = 64, - lmax: int = 3*64 - 1, - mmax: int = 3*64 - 1, + lmax: int = 3 * 64 - 1, + mmax: int = 3 * 64 - 1, multiscale: float = 0.0, lsm_file: str = None, open_dict: dict = {"engine": "zarr"}, @@ -635,34 +699,51 @@ def __init__( self.multiscale = multiscale # Parameters for "almost fair CRPS" loss. See https://arxiv.org/html/2412.15832v1 - self.coeff_eps = 1 - ((1-alpha) / (n_members)) - self.averaging_coeff = 1 / (2* n_members * (n_members - 1)) + self.coeff_eps = 1 - ((1 - alpha) / (n_members)) + self.averaging_coeff = 1 / (2 * n_members * (n_members - 1)) # SHT utils: transform, grid reordering, output indexing self.lmax = lmax self.mmax = mmax self.nside = nside - self.sht = cuhpx.SHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') - src_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=earth2grid.healpix.HEALPIX_PAD_XY) - tar_grid = earth2grid.healpix.Grid(level=int(np.log2(nside)), pixel_order=earth2grid.healpix.PixelOrder.RING) - self.reorder_to_ring = earth2grid.get_regridder(src_grid, tar_grid).to(th.float32) + self.sht = cuhpx.SHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights="ring") + src_grid = earth2grid.healpix.Grid( + level=int(np.log2(nside)), pixel_order=earth2grid.healpix.HEALPIX_PAD_XY + ) + tar_grid = earth2grid.healpix.Grid( + level=int(np.log2(nside)), pixel_order=earth2grid.healpix.PixelOrder.RING + ) + self.reorder_to_ring = earth2grid.get_regridder(src_grid, tar_grid).to( + th.float32 + ) if self.multiscale > 0: - self.scales = [200, 400, 800, 1600] # in units of km - self.isht = cuhpx.iSHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring') - self.reorder_from_ring = earth2grid.get_regridder(tar_grid, src_grid).to(th.float32) + self.scales = [200, 400, 800, 1600] # in units of km + self.isht = cuhpx.iSHTCUDA( + nside=nside, lmax=lmax, mmax=mmax, quad_weights="ring" + ) + self.reorder_from_ring = earth2grid.get_regridder(tar_grid, src_grid).to( + th.float32 + ) self.lsm_file = lsm_file if lsm_file is not None: - self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel(selection_dict) - self.lsm_tensor = 1 - th.tensor(np.expand_dims(self.lsm_ds.values, (0, 2, 3))) + self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel( + selection_dict + ) + self.lsm_tensor = 1 - th.tensor( + np.expand_dims(self.lsm_ds.values, (0, 2, 3)) + ) else: - self.lsm_tensor = th.ones(1, 1, 1, 1, 1, 1) # Spoof the tensor dimensions for broadcasting + self.lsm_tensor = th.ones( + 1, 1, 1, 1, 1, 1 + ) # Spoof the tensor dimensions for broadcasting # For n>2, will use pairwise distance to copmute [NxN] distance matrix # Diagonal elements of (prediciton - target) matrix are zeroed out to avoid double counting self.pdist = th.nn.PairwiseDistance(p=1) - self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye(self.n_members) # Mask to zero out diagonal elements - + self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye( + self.n_members + ) # Mask to zero out diagonal elements def setup(self, trainer): """ @@ -702,13 +783,19 @@ def _apply_sht(self, x, face_dim, return_abs=True): """ x = th.movedim(x, face_dim, -3) if x.shape[-3:] != (12, self.nside, self.nside): - raise ValueError(f"Shape of input tensor should be [..., F, ..., H, W] with F in position {face_dim}, got {x.shape}") - + face, height, width = x.shape[-3:] + raise ValueError( + f"Shape of input tensor should be [..., 12, ..., {self.nside}, {self.nside}] with F (12) in position {face_dim}," + f"got {x.shape}" + ) + x = x.reshape(*x.shape[:-3], -1) - x = self.reorder_to_ring(x.contiguous()) # contiguous needed for channels first format in validation loop + x = self.reorder_to_ring( + x.contiguous() + ) # contiguous needed for channels first format in validation loop x = self.sht(x) if return_abs: - x = x.real ** 2 + x.imag ** 2 + x = x.real**2 + x.imag**2 return x def _apply_isht(self, x, face_dim): @@ -716,22 +803,23 @@ def _apply_isht(self, x, face_dim): Inverse transform, reorder from ring, Reshape to [..., F, H, W], move face dim appropriately """ - x = self.isht(x) # [..., l, m] -> [..., F*H*W] + x = self.isht(x) # [..., l, m] -> [..., F*H*W] x = self.reorder_from_ring(x) - x = x.reshape(*x.shape[:-1], 12, self.nside, self.nside) # [..., F*H*W] -> [..., F, H, W] - x = th.movedim(x, -3, face_dim) # [..., F, H, W] -> [..., F, ..., H, W] + x = x.reshape( + *x.shape[:-1], 12, self.nside, self.nside + ) # [..., F*H*W] -> [..., F, H, W] + x = th.movedim(x, -3, face_dim) # [..., F, H, W] -> [..., F, ..., H, W] return x def _l_filter(self, scale, device="cuda"): - """Return a spherical gaussian filter of scale `scale` (in units of km) - """ - scale_radians = scale / 6371.0 + """Return a spherical gaussian filter of scale `scale` (in units of km)""" + scale_radians = scale / 6371.0 ell = th.arange(self.lmax, device=device, dtype=th.float32) - return th.exp(-0.5* ell * (ell + 1) * (scale_radians ** 2)) + return th.exp(-0.5 * ell * (ell + 1) * (scale_radians**2)) def forward(self, prediction, target, average_channels=True): """ - Forward pass of the WeightedCRPSLoss + Forward pass of the WeightedCRPSLoss Computes the CRPS loss for the model prediction and target. Parameters @@ -743,17 +831,21 @@ def forward(self, prediction, target, average_channels=True): average_channels: bool, optional whether the mean of the channels should be taken """ - + # Unfold ensemble dimension from batch dimension to have shape [Cond, B, F, T, C, H, W] b, f, t, c, h, w = target.shape prediction = prediction.view(self.n_members, b, f, t, c, h, w) - # checks for dimensions + # checks for dimensions if not prediction.shape[1:] == target.shape: - raise ValueError(f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}") - + raise ValueError( + f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}" + ) + if not prediction.shape[0] == self.n_members: - raise ValueError(f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}") + raise ValueError( + f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}" + ) n = self.n_members @@ -767,9 +859,13 @@ def forward(self, prediction, target, average_channels=True): if n == 2: # Use faster explicit implementation - diff_target = th.abs(prediction - target.unsqueeze(0)).sum(dim=0) # [B, F, T, C, H, W] - diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] - crps = self.averaging_coeff*(diff_target - self.coeff_eps * diff_ensemble) # [B, F, T, C, H, W] + diff_target = th.abs(prediction - target.unsqueeze(0)).sum( + dim=0 + ) # [B, F, T, C, H, W] + diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] + crps = self.averaging_coeff * ( + diff_target - self.coeff_eps * diff_ensemble + ) # [B, F, T, C, H, W] if average_channels: loss = crps.mean() @@ -777,8 +873,7 @@ def forward(self, prediction, target, average_channels=True): loss = crps.mean(dim=(0, 1, 2, 4, 5)) if self.lambda_spec > 0: - - with th.amp.autocast("cuda",enabled=False): + with th.amp.autocast("cuda", enabled=False): # # Reorder predictions: [N, B, F, T, C, H, W] -> [N, B, T, C, F*H*W] # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) @@ -794,9 +889,15 @@ def forward(self, prediction, target, average_channels=True): sht_pred = self._apply_sht(prediction, face_dim=2, return_abs=True) sht_tar = self._apply_sht(target, face_dim=1, return_abs=True) - diff_sht_target = th.abs(sht_pred - sht_tar.unsqueeze(0)).sum(dim=(0, 4, 5)) # [B, T, C] - diff_sht_ensemble = th.abs(sht_pred[0] - sht_pred[1]).sum(dim=(-1,-2)) # [B, T, C] - crps_sht = self.averaging_coeff * (diff_sht_target - self.coeff_eps * diff_sht_ensemble) # [B, T, C] + diff_sht_target = th.abs(sht_pred - sht_tar.unsqueeze(0)).sum( + dim=(0, 4, 5) + ) # [B, T, C] + diff_sht_ensemble = th.abs(sht_pred[0] - sht_pred[1]).sum( + dim=(-1, -2) + ) # [B, T, C] + crps_sht = self.averaging_coeff * ( + diff_sht_target - self.coeff_eps * diff_sht_ensemble + ) # [B, T, C] # Compute spectral afCRPS if average_channels: @@ -805,23 +906,39 @@ def forward(self, prediction, target, average_channels=True): spec_loss = crps_sht.mean(dim=(0, 1)) loss += self.lambda_spec * spec_loss - + if self.multiscale > 0: for scale in self.scales: l_filter = self._l_filter(scale, device=prediction.device) - with th.amp.autocast("cuda",enabled=False): - sht_pred= self._apply_sht(prediction, face_dim=2, return_abs=False) + with th.amp.autocast("cuda", enabled=False): + sht_pred = self._apply_sht( + prediction, face_dim=2, return_abs=False + ) sht_tar = self._apply_sht(target, face_dim=1, return_abs=False) - l_filter_pred = l_filter[None, None, None, None, :, None] # [1, 1, 1, 1, lmax, 1] - l_filter_tar = l_filter[None, None, None, :, None] # [1, 1, 1, lmax, 1] - - pred_smooth = self._apply_isht(l_filter_pred * sht_pred, face_dim=2) - tar_smooth = self._apply_isht(l_filter_tar * sht_tar, face_dim=1) - - diff_target = th.abs(pred_smooth - tar_smooth.unsqueeze(0)).sum(dim=0) # [B, F, T, C, H, W] - diff_ensemble = th.abs(pred_smooth[0] - pred_smooth[1]) # [B, F, T, C, H, W] - crps = self.averaging_coeff*(diff_target - self.coeff_eps * diff_ensemble) # [B, F, T, C, H, W] + l_filter_pred = l_filter[ + None, None, None, None, :, None + ] # [1, 1, 1, 1, lmax, 1] + l_filter_tar = l_filter[ + None, None, None, :, None + ] # [1, 1, 1, lmax, 1] + + pred_smooth = self._apply_isht( + l_filter_pred * sht_pred, face_dim=2 + ) + tar_smooth = self._apply_isht( + l_filter_tar * sht_tar, face_dim=1 + ) + + diff_target = th.abs(pred_smooth - tar_smooth.unsqueeze(0)).sum( + dim=0 + ) # [B, F, T, C, H, W] + diff_ensemble = th.abs( + pred_smooth[0] - pred_smooth[1] + ) # [B, F, T, C, H, W] + crps = self.averaging_coeff * ( + diff_target - self.coeff_eps * diff_ensemble + ) # [B, F, T, C, H, W] crps *= self.lsm_tensor @@ -829,27 +946,39 @@ def forward(self, prediction, target, average_channels=True): loss += self.multiscale * crps.mean() else: loss += self.multiscale * crps.mean(dim=(0, 1, 2, 4, 5)) - + return loss - + else: # Use pairwise distance method if not average_channels: - # Move channels to first dimension and exclude that dimension from the reductions - prediction = prediction.permute(4, 0, 1, 2, 3, 5, 6) # [C, Cond, B, F, T, H, W] - target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] - - pred = prediction.reshape(c, n, -1) # [C, Cond, ...] - tar = target.unsqueeze(1).reshape(c, 1, -1) # [C, 1, ...] (second dim will broadcast across ensemble) - - diff = self.pdist(pred, tar) # [C, Cond] - dist_matrix = self.pdist(pred.unsqueeze(1), pred.unsqueeze(2)) # [C, Cond, Cond] - - diff_terms = self.diag_mask[None, ...] * (diff.unsqueeze(1) + diff.unsqueeze(2)) # [C, Cond, Cond], diagonal elements zeroed out - loss = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1,2))/(b*f*t*h*w) + # Move channels to first dimension and exclude that dimension from the reductions + prediction = prediction.permute( + 4, 0, 1, 2, 3, 5, 6 + ) # [C, Cond, B, F, T, H, W] + target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] + + pred = prediction.reshape(c, n, -1) # [C, Cond, ...] + tar = target.unsqueeze(1).reshape( + c, 1, -1 + ) # [C, 1, ...] (second dim will broadcast across ensemble) + + diff = self.pdist(pred, tar) # [C, Cond] + dist_matrix = self.pdist( + pred.unsqueeze(1), pred.unsqueeze(2) + ) # [C, Cond, Cond] + + diff_terms = self.diag_mask[None, ...] * ( + diff.unsqueeze(1) + diff.unsqueeze(2) + ) # [C, Cond, Cond], diagonal elements zeroed out + loss = ( + self.averaging_coeff + * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1, 2)) + / (b * f * t * h * w) + ) if self.lambda_spec > 0: - with th.amp.autocast("cuda",enabled=False): + with th.amp.autocast("cuda", enabled=False): # # Reorder predictions: [C, Cond, B, F, T, H, W] -> [C, Cond, B, T, F*H*W] # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 2, 4, 3, 5, 6).reshape(c, n, b, t, f*h*w)) @@ -862,26 +991,53 @@ def forward(self, prediction, target, average_channels=True): # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 - sht_pred = self._apply_sht(prediction, face_dim=3, return_abs=True).reshape(c, n, -1) - sht_tar = self._apply_sht(target, face_dim=2, return_abs=True).unsqueeze(1).reshape(c, 1, -1) - - diff = self.pdist(sht_pred, sht_tar) # [C, Cond] - dist_matrix = self.pdist(sht_pred.unsqueeze(1), sht_pred.unsqueeze(2)) # [C, Cond, Cond] - - diff_terms = self.diag_mask[None, ...] * (diff.unsqueeze(1) + diff.unsqueeze(2)) # [C, Cond, Cond], diagonal elements zeroed out - loss += self.lambda_spec * self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1,2))/(b*t*self.lmax*self.mmax) + sht_pred = self._apply_sht( + prediction, face_dim=3, return_abs=True + ).reshape(c, n, -1) + sht_tar = ( + self._apply_sht(target, face_dim=2, return_abs=True) + .unsqueeze(1) + .reshape(c, 1, -1) + ) + + diff = self.pdist(sht_pred, sht_tar) # [C, Cond] + dist_matrix = self.pdist( + sht_pred.unsqueeze(1), sht_pred.unsqueeze(2) + ) # [C, Cond, Cond] + + diff_terms = self.diag_mask[None, ...] * ( + diff.unsqueeze(1) + diff.unsqueeze(2) + ) # [C, Cond, Cond], diagonal elements zeroed out + loss += ( + self.lambda_spec + * self.averaging_coeff + * (diff_terms - self.coeff_eps * dist_matrix).sum( + dim=(1, 2) + ) + / (b * t * self.lmax * self.mmax) + ) else: pred = prediction.reshape(n, -1) - tar = target.unsqueeze(0).reshape(1, -1) # [1, ...] (first dim will broadcast across ensemble) - diff = self.pdist(pred, tar) # [Cond] - dist_matrix = self.pdist(pred.unsqueeze(1), pred.unsqueeze(0)) # [Cond, Cond] - - diff_terms = self.diag_mask * (diff.unsqueeze(0) + diff.unsqueeze(1)) # [Cond, Cond], diagonal elements zeroed out - loss = self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum()/(b*f*c*t*h*w) + tar = target.unsqueeze(0).reshape( + 1, -1 + ) # [1, ...] (first dim will broadcast across ensemble) + diff = self.pdist(pred, tar) # [Cond] + dist_matrix = self.pdist( + pred.unsqueeze(1), pred.unsqueeze(0) + ) # [Cond, Cond] + + diff_terms = self.diag_mask * ( + diff.unsqueeze(0) + diff.unsqueeze(1) + ) # [Cond, Cond], diagonal elements zeroed out + loss = ( + self.averaging_coeff + * (diff_terms - self.coeff_eps * dist_matrix).sum() + / (b * f * c * t * h * w) + ) if self.lambda_spec > 0: - with th.amp.autocast("cuda",enabled=False): + with th.amp.autocast("cuda", enabled=False): # # Reorder predictions: [Cond, B, F, T, C, H, W] -> [Cond, B, T, C, F*H*W] # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) @@ -893,14 +1049,28 @@ def forward(self, prediction, target, average_channels=True): # sht_tar = self.sht(tar_ring).unsqueeze(0).reshape(1, -1) # [B, T, C, l, m] -> [1, ...] (first dim will broadcast across ensemble) # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 - sht_pred = self._apply_sht(prediction, face_dim=2, return_abs=True).reshape(n, -1) - sht_tar = self._apply_sht(target, face_dim=1, return_abs=True).unsqueeze(0).reshape(1, -1) - - diff = self.pdist(sht_pred, sht_tar) # [Cond] - dist_matrix = self.pdist(sht_pred.unsqueeze(1), sht_pred.unsqueeze(0)) # [Cond, Cond] - - diff_terms = self.diag_mask * (diff.unsqueeze(0) + diff.unsqueeze(1)) # [Cond, Cond], diagonal elements zeroed out - loss += self.lambda_spec * self.averaging_coeff * (diff_terms - self.coeff_eps * dist_matrix).sum()/(b*t*self.lmax*self.mmax) + sht_pred = self._apply_sht( + prediction, face_dim=2, return_abs=True + ).reshape(n, -1) + sht_tar = ( + self._apply_sht(target, face_dim=1, return_abs=True) + .unsqueeze(0) + .reshape(1, -1) + ) + + diff = self.pdist(sht_pred, sht_tar) # [Cond] + dist_matrix = self.pdist( + sht_pred.unsqueeze(1), sht_pred.unsqueeze(0) + ) # [Cond, Cond] + + diff_terms = self.diag_mask * ( + diff.unsqueeze(0) + diff.unsqueeze(1) + ) # [Cond, Cond], diagonal elements zeroed out + loss += ( + self.lambda_spec + * self.averaging_coeff + * (diff_terms - self.coeff_eps * dist_matrix).sum() + / (b * t * self.lmax * self.mmax) + ) return loss - diff --git a/physicsnemo/metrics/climate/hydrostasy.py b/physicsnemo/metrics/climate/hydrostasy.py index 91857f71f4..5099c7095a 100644 --- a/physicsnemo/metrics/climate/hydrostasy.py +++ b/physicsnemo/metrics/climate/hydrostasy.py @@ -45,9 +45,10 @@ def __init__( self.R = R self.g0 = g0 - assert ( - anchor_z_channel in z_pressure_levels.keys() - ), f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" + if anchor_z_channel not in z_pressure_levels: + raise ValueError( + f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" + ) # Sort channels by pressure levels for ease of use later self.z_pressure_levels = dict( @@ -131,12 +132,14 @@ def __init__( self.R = R self.g0 = g0 - assert ( - anchor_z_channel in z_pressure_levels.keys() - ), f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" - assert ( - anchor_T_channel in Tv_pressure_levels.keys() - ), f"anchor_T_channel ({anchor_T_channel}) not in Tv_pressure_levels ({Tv_pressure_levels.keys()})" + if anchor_z_channel not in z_pressure_levels.keys(): + raise ValueError( + f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" + ) + if anchor_T_channel not in Tv_pressure_levels.keys(): + raise ValueError( + f"anchor_T_channel ({anchor_T_channel}) not in Tv_pressure_levels ({Tv_pressure_levels.keys()})" + ) # Sort channels by pressure levels for ease of use later self.z_pressure_levels = dict( @@ -187,15 +190,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: z_channel_p1 = self.z_channels[i + 1] zi = x[:, :, z_channel, ...] zip1 = x[:, :, z_channel_p1, ...] - Tv_avg[ - :, :, i, ... - ] = _average_virtual_temperature_from_geopotential_height( - zi, - zip1, - self.z_pressure_levels[z_channel], - self.z_pressure_levels[z_channel_p1], - self.R, - self.g0, + Tv_avg[:, :, i, ...] = ( + _average_virtual_temperature_from_geopotential_height( + zi, + zip1, + self.z_pressure_levels[z_channel], + self.z_pressure_levels[z_channel_p1], + self.R, + self.g0, + ) ) Tv_model_avg_size = [ @@ -218,7 +221,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class WeightedMSEWithHydrostasy(torch.nn.MSELoss): - """ Loss object that adds a differential Hydrostatic balance constraint in addition to user defined weighting of variables when calculating MSE @@ -320,7 +322,7 @@ def __init__( # Set per level alphas if len(alpha) != len(hPa_levels) - 1: raise AssertionError( - f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels)-1}], got {len(alpha)}" + f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels) - 1}], got {len(alpha)}" ) self.alpha = torch.Tensor(alpha).reshape((1, 1, -1, 1, 1)) @@ -465,13 +467,13 @@ def error_histogram(self, prediction, bins, accumulator=None): else: bin_edges = bins - for l in range(vlevels): + for level_idx in range(vlevels): hist, be = torch.histogram( - torch.absolute(Tv_error[:, :, l, :, :]), - bins=bins if isinstance(bins, int) else bin_edges[l, :], + torch.absolute(Tv_error[:, :, level_idx, :, :]), + bins=bins if isinstance(bins, int) else bin_edges[level_idx, :], ) - accumulator[l, :] += hist - bin_edges[l, :] = be + accumulator[level_idx, :] += hist + bin_edges[level_idx, :] = be return accumulator, bin_edges @@ -523,10 +525,9 @@ def forward(self, prediction, target, average_channels=True): class LossWithHydrostasy(torch.nn.MSELoss): - """ Loss object that adds a differential Hydrostatic balance constraint in addition to - user defined data loss + user defined data loss function. """ def __init__( @@ -627,7 +628,7 @@ def __init__( # Set per level alphas if len(alpha) != len(hPa_levels) - 1: raise AssertionError( - f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels)-1}], got {len(alpha)}" + f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels) - 1}], got {len(alpha)}" ) self.alpha = torch.Tensor(alpha).reshape((1, 1, -1, 1, 1)) @@ -655,7 +656,7 @@ def __init__( # Get topography information ds = xr.open_zarr(f"{src_directory}{dataset_name}.zarr") self.topography = ( - surface_geopotential_std * ds.constants.sel(channel_c='z').values + surface_geopotential_std * ds.constants.sel(channel_c="z").values + surface_geopotential_mean ) / self.g0 @@ -674,10 +675,10 @@ def setup(self, trainer): # Call setup for data loss first self.data_loss.setup(trainer) - if len(self.z_pressure_levels) - 1 != len( - self.loss_weights - ): - raise ValueError("Length of loss_weights is not one less than number of pressure levels!") + if len(self.z_pressure_levels) - 1 != len(self.loss_weights): + raise ValueError( + "Length of loss_weights is not one less than number of pressure levels!" + ) self.loss_weights = self.loss_weights.to(device=trainer.device) @@ -775,13 +776,13 @@ def error_histogram(self, prediction, bins, accumulator=None): else: bin_edges = bins - for l in range(vlevels): + for level_idx in range(vlevels): hist, be = torch.histogram( - torch.absolute(Tv_error[:, :, l, :, :]), - bins=bins if isinstance(bins, int) else bin_edges[l, :], + torch.absolute(Tv_error[:, :, level_idx, :, :]), + bins=bins if isinstance(bins, int) else bin_edges[level_idx, :], ) - accumulator[l, :] += hist - bin_edges[l, :] = be + accumulator[level_idx, :] += hist + bin_edges[level_idx, :] = be return accumulator, bin_edges @@ -819,10 +820,14 @@ def forward(self, prediction, target, average_channels=True): Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 # Compute the error tolerant loss - Tv_loss = self.loss_weights * (Tv_error / (1 + torch.exp(1 - Tv_error))).mean(dim=(0, 1, 3, 4)) + Tv_loss = self.loss_weights * ( + Tv_error / (1 + torch.exp(1 - Tv_error)) + ).mean(dim=(0, 1, 3, 4)) # Compute data loss - data_loss = self.data_loss(prediction, target, average_channels=average_channels) + data_loss = self.data_loss( + prediction, target, average_channels=average_channels + ) if average_channels: return data_loss + torch.mean(Tv_loss) diff --git a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py index 70e691997e..15775b4888 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py @@ -104,6 +104,12 @@ class HEALPixRecUNet(Module): "0.1.0": _legacy_hydra_targets_warning, } + # Allows ``Module.from_checkpoint(..., override_args=...)`` to substitute + # a corrected ``encoder``/``decoder`` config. Used by checkpoint conversion + # utilities to fix up legacy checkpoints (e.g. the ``per_level_cln`` + # decoder bug) without having to hand-edit the model source. + _overridable_args: set = {"encoder", "decoder"} + @classmethod def _backward_compat_arg_mapper( cls, version: str, args: Dict[str, Any] diff --git a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py index 10ac29a6b2..1f76da6560 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py @@ -100,6 +100,12 @@ class HEALPixUNet(Module): "0.1.0": _legacy_hydra_targets_warning, } + # Allows ``Module.from_checkpoint(..., override_args=...)`` to substitute + # a corrected ``encoder``/``decoder`` config. Used by checkpoint conversion + # utilities to fix up legacy checkpoints (e.g. the ``per_level_cln`` + # decoder bug) without having to hand-edit the model source. + _overridable_args: set = {"encoder", "decoder"} + @classmethod def _backward_compat_arg_mapper( cls, version: str, args: Dict[str, Any] diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py index 3b3ed6a243..db29af0d19 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py @@ -14,15 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Sequence, Tuple, Union, Callable +from typing import Callable, Sequence, Tuple, Union import torch import torch as th + from physicsnemo.nn.module.hpx import HEALPixLayer -from .normalization import ConditionalLayerNorm -from hydra.utils import instantiate -from omegaconf import DictConfig +from .normalization import ConditionalLayerNorm # @@ -40,6 +39,7 @@ def forward(self, x): x = self.norm(x) return x.permute(0, 3, 1, 2) + # # RECURRENT BLOCKS # @@ -379,10 +379,10 @@ def __init__( super().__init__() self.cln_enabled = conditional_layer_norm is not None - + if in_channels == int(latent_channels): - self.skip_module1 = ( - lambda x: x + self.skip_module1 = lambda x: ( + x ) # Identity-function required in forward pass else: self.skip_module1 = geometry_layer( @@ -394,8 +394,8 @@ def __init__( enable_healpixpad=enable_healpixpad, ) if out_channels == int(latent_channels): - self.skip_module2 = ( - lambda x: x + self.skip_module2 = lambda x: ( + x ) # Identity-function required in forward pass else: self.skip_module2 = geometry_layer( @@ -407,7 +407,7 @@ def __init__( enable_healpixpad=enable_healpixpad, ) - # check if we're applying a layer norm at the beginning + # check if we're applying a layer norm at the beginning # we've got two ConvNeXt equivalent blocks in the layer, so have two entry points # TODO: conditional layer norm once is doing two things, it's applying a norm on block entry # and switch from conditional to non-conditional layer norm. This is not ideal and should be fixed once we determine @@ -466,9 +466,15 @@ def __init__( # Apply layer norm if needed if conditional_layer_norm_once: - convblock1.append(_LayerNormOverChannels(channel_depth=int(latent_channels * upscale_factor))) + convblock1.append( + _LayerNormOverChannels( + channel_depth=int(latent_channels * upscale_factor) + ) + ) elif conditional_layer_norm is not None: - cln = conditional_layer_norm(channel_depth=int(latent_channels * upscale_factor)) + cln = conditional_layer_norm( + channel_depth=int(latent_channels * upscale_factor) + ) convblock1.append(cln) if activation is not None: @@ -512,7 +518,7 @@ def __init__( if conditional_layer_norm is not None and not conditional_layer_norm_once: cln = conditional_layer_norm(channel_depth=int(latent_channels)) convblock2.append(cln) - + if activation is not None: convblock2.append(activation) if dropout > 0.0: @@ -532,9 +538,15 @@ def __init__( ) # Apply layer norm if needed if conditional_layer_norm_once: - convblock2.append(_LayerNormOverChannels(channel_depth=int(latent_channels * upscale_factor))) + convblock2.append( + _LayerNormOverChannels( + channel_depth=int(latent_channels * upscale_factor) + ) + ) elif conditional_layer_norm is not None: - cln = conditional_layer_norm(channel_depth=int(latent_channels * upscale_factor)) + cln = conditional_layer_norm( + channel_depth=int(latent_channels * upscale_factor) + ) convblock2.append(cln) if activation is not None: @@ -614,9 +626,10 @@ def forward(self, x, conditions_cln=None): x1 = layer(x1) return x2_residual + x1 + class Multi_SymmetricConvNeXtBlock(th.nn.Module): """ - Wrapper for SymmetricConvNeXtBlock that allows serial linking of blocks. + Wrapper for SymmetricConvNeXtBlock that allows serial linking of blocks. """ def __init__( @@ -642,7 +655,7 @@ def __init__( n_layers: int, optional The number of SymmetricConvNeXt Blocks conditional_layer_norm: Callable, optional - Callable for physicsnemo.models.dlwp_healpix_layers.normalization.ConditionalLayerNorm. + Callable for physicsnemo.models.dlwp_healpix_layers.normalization.ConditionalLayerNorm. Callable can be passed in by setting _partial_ to True in hydra config. If None, conditional layer normalization is not applied. conditional_layer_norm_once: bool, optional @@ -672,7 +685,9 @@ def __init__( enable_nhwc=enable_nhwc, enable_healpixpad=enable_healpixpad, dropout=dropout, - conditional_layer_norm=conditional_layer_norm if conditional_layer_norm is not None else None, + conditional_layer_norm=conditional_layer_norm + if conditional_layer_norm is not None + else None, conditional_layer_norm_once=conditional_layer_norm_once, ), ) @@ -819,17 +834,22 @@ def __init__( # Apply layer norm if needed if conditional_layer_norm_once: - convblock.append(_LayerNormOverChannels(channel_depth=int(latent_channels * upscale_factor))) + convblock.append( + _LayerNormOverChannels( + channel_depth=int(latent_channels * upscale_factor) + ) + ) elif conditional_layer_norm is not None: - cln = conditional_layer_norm(channel_depth=int(latent_channels * upscale_factor)) - convblock.append(cln) + cln = conditional_layer_norm( + channel_depth=int(latent_channels * upscale_factor) + ) + convblock.append(cln) if activation is not None: convblock.append(activation) if dropout > 0.0: convblock.append(th.nn.Dropout2d(p=dropout)) - # 1x1: upscale → latent convblock.append( geometry_layer( @@ -874,7 +894,6 @@ def __init__( self.convblock = th.nn.ModuleList(convblock) - def forward(self, x, conditions_cln=None): """ Forward pass of the SymmetricConvNextBlock, broken into steps for support of conditional layer normalization diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py index d184147d8f..522b514f24 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py @@ -1,6 +1,5 @@ -import numpy as np import torch -import xarray as xr + class NonnegativeConstraint(torch.nn.Module): def __init__( @@ -28,22 +27,21 @@ def __init__( self.variables = [var for var in self.variables if var in channels] var_indices = torch.tensor( - [channels.index(var) for var in self.variables], - dtype=torch.long + [channels.index(var) for var in self.variables], dtype=torch.long ) - self.register_buffer('var_indices', var_indices, persistent=False) + self.register_buffer("var_indices", var_indices, persistent=False) - self.var_means = torch.tensor([scaling[var]['mean'] for var in self.variables]) - self.var_stds = torch.tensor([scaling[var]['std'] for var in self.variables]) + self.var_means = torch.tensor([scaling[var]["mean"] for var in self.variables]) + self.var_stds = torch.tensor([scaling[var]["std"] for var in self.variables]) - thresholds = (0. - self.var_means) / self.var_stds + thresholds = (0.0 - self.var_means) / self.var_stds thresholds = thresholds.view(1, 1, 1, -1, 1, 1) - self.register_buffer('thresholds', thresholds, persistent=False) + self.register_buffer("thresholds", thresholds, persistent=False) def forward(self, x): - ''' + """ Tensors are expected to be in the shape [B, F, T, C, H, W] - ''' + """ selected_vars = torch.index_select(x, dim=3, index=self.var_indices) clamped = torch.maximum(selected_vars, self.thresholds).to(x.dtype) x.index_copy_(3, self.var_indices, clamped) diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py index cadebb6689..a719c01540 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py @@ -216,9 +216,7 @@ def forward( use_reentrant=False, ) else: - x = self._forward_layer_pass( - layer, x, skip_connection, conditions_cln - ) + x = self._forward_layer_pass(layer, x, skip_connection, conditions_cln) if layer["recurrent"] is not None: x = layer["recurrent"](x) diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index 524694ca0a..51c4b47e1f 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -14,13 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Sequence, List +from typing import List import torch as th -from omegaconf import DictConfig try: from apex.normalization import FusedLayerNorm + _APEX_AVAILABLE = True except ImportError: _APEX_AVAILABLE = False @@ -30,7 +30,12 @@ def _cln_affine(x_norm, gamma_raw, beta, scale_center, n_faces): """Fused affine transform: expand gamma/beta across faces and apply to normalized input.""" C = gamma_raw.shape[-1] - gamma = (scale_center + gamma_raw).unsqueeze(1).expand(-1, n_faces, -1).reshape(-1, 1, 1, C) + gamma = ( + (scale_center + gamma_raw) + .unsqueeze(1) + .expand(-1, n_faces, -1) + .reshape(-1, 1, 1, C) + ) beta = beta.unsqueeze(1).expand(-1, n_faces, -1).reshape(-1, 1, 1, C) return gamma * x_norm + beta @@ -44,7 +49,7 @@ def __init__( activation: th.nn.Module = None, eps: float = 1e-5, n_faces: int = 12, - norm_op:str = "torch", + norm_op: str = "torch", init_cln_to_zero: bool = False, scale_center: float = 0.0, ): @@ -97,11 +102,18 @@ def __init__( self.norm = th.nn.LayerNorm(channel_depth, elementwise_affine=False) elif norm_op == "apex": if not _APEX_AVAILABLE: - raise ImportError("Apex FusedLayerNorm requested but apex is not available, please install it from https://github.com/NVIDIA/apex") + raise ImportError( + "Apex FusedLayerNorm requested but apex is not available, please install it from https://github.com/NVIDIA/apex" + ) self.norm = FusedLayerNorm(channel_depth, elementwise_affine=False) - def _make_mlp(self, in_dim: int, hidden_dims: List[int], out_dim: int, activation: th.nn.Module) -> th.nn.Sequential: - + def _make_mlp( + self, + in_dim: int, + hidden_dims: List[int], + out_dim: int, + activation: th.nn.Module, + ) -> th.nn.Sequential: layers = [] for hdim in hidden_dims: layers.append(th.nn.Linear(in_dim, hdim)) @@ -111,7 +123,16 @@ def _make_mlp(self, in_dim: int, hidden_dims: List[int], out_dim: int, activatio layers.append(th.nn.Linear(in_dim, out_dim)) return th.nn.Sequential(*layers) - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): """Backward compatibility: merge old separate gamma_mlp/beta_mlp into fused gamma_beta_mlp. Old MLPs had hidden_dims [h1, h2, ...] and output dim C. @@ -141,7 +162,7 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss gamma_layer_indices = set() for k in state_dict: if k.startswith(gamma_prefix): - layer_key = k[len(gamma_prefix):] + layer_key = k[len(gamma_prefix) :] parts = layer_key.split(".") if len(parts) == 2 and parts[1] in ("weight", "bias"): gamma_layer_indices.add(int(parts[0])) @@ -153,7 +174,7 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss # handle weights and biases for the gamma MLP for k in list(state_dict.keys()): if k.startswith(gamma_prefix): - layer_key = k[len(gamma_prefix):] # e.g. "0.weight" + layer_key = k[len(gamma_prefix) :] # e.g. "0.weight" parts = layer_key.split(".") if len(parts) != 2 or parts[1] not in ("weight", "bias"): continue @@ -180,11 +201,19 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss # gamma_val: (out_old, in_old), beta_val: (out_old, in_old) # result: (2*out_old, 2*in_old) out_old, in_old = gamma_val.shape - zeros = th.zeros(out_old, in_old, dtype=gamma_val.dtype, device=gamma_val.device) - keys_to_add[fused_key] = th.cat([ - th.cat([gamma_val, zeros], dim=1), - th.cat([zeros, beta_val], dim=1), - ], dim=0) + zeros = th.zeros( + out_old, + in_old, + dtype=gamma_val.dtype, + device=gamma_val.device, + ) + keys_to_add[fused_key] = th.cat( + [ + th.cat([gamma_val, zeros], dim=1), + th.cat([zeros, beta_val], dim=1), + ], + dim=0, + ) keys_to_remove.append(k) if beta_key not in keys_to_remove: @@ -194,7 +223,7 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss for k in list(state_dict.keys()): if not k.startswith(gamma_prefix): continue - layer_key = k[len(gamma_prefix):] # e.g. "0.weight" + layer_key = k[len(gamma_prefix) :] # e.g. "0.weight" parts = layer_key.split(".", 1) if len(parts) != 2: continue @@ -217,7 +246,15 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, miss del state_dict[k] state_dict.update(keys_to_add) - super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) def forward(self, x: th.Tensor, conditions: th.Tensor) -> th.Tensor: """ From a0bd39d9371c0ff058fd7cc10f49aae518267f57 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 8 Jul 2026 10:37:47 -0500 Subject: [PATCH 07/28] License updates --- .gitignore | 1 + .../healpix/base_timeseries_dataset_zarr.py | 2 +- .../healpix/coupledtimeseries_dataset.py | 2 +- .../healpix/coupledtimeseries_dataset_zarr.py | 2 +- physicsnemo/datapipes/healpix/couplers.py | 2 +- physicsnemo/datapipes/healpix/data_modules.py | 2 +- .../datapipes/healpix/data_modules_zarr.py | 2 +- .../datapipes/healpix/timeseries_dataset.py | 2 +- .../datapipes/healpix/timeseries_dataset_zarr.py | 2 +- physicsnemo/metrics/climate/healpix_loss.py | 2 +- physicsnemo/metrics/climate/hydrostasy.py | 2 +- .../models/dlwp_healpix/layers/healpix_blocks.py | 2 +- .../dlwp_healpix/layers/healpix_constraints.py | 16 ++++++++++++++++ .../models/dlwp_healpix/layers/normalization.py | 2 +- test/datapipes/test_healpix_couple_zarr.py | 2 +- test/datapipes/test_healpix_zarr.py | 2 +- test/metrics/test_hydrostatic_loss.py | 2 +- test/models/dlwp_healpix/_cln_reference.py | 16 ++++++++++++++++ .../dlwp_healpix/test_conditional_layer_norm.py | 16 ++++++++++++++++ 19 files changed, 64 insertions(+), 15 deletions(-) diff --git a/.gitignore b/.gitignore index 7ec721b0e1..28570f5b9e 100644 --- a/.gitignore +++ b/.gitignore @@ -121,6 +121,7 @@ venv/ ENV/ env.bak/ venv.bak/ +.venv-dlesym/ # Spyder project settings .spyderproject diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py index 51fd9b81c1..4d347e2480 100644 --- a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index 3cba485113..92ba830390 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py index 4c4a71403c..c3df1e1412 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 1d77761734..101f916369 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/datapipes/healpix/data_modules.py b/physicsnemo/datapipes/healpix/data_modules.py index 1e67017a39..fbb45bb6f6 100644 --- a/physicsnemo/datapipes/healpix/data_modules.py +++ b/physicsnemo/datapipes/healpix/data_modules.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/datapipes/healpix/data_modules_zarr.py b/physicsnemo/datapipes/healpix/data_modules_zarr.py index ad830188e0..5ef837d1bd 100644 --- a/physicsnemo/datapipes/healpix/data_modules_zarr.py +++ b/physicsnemo/datapipes/healpix/data_modules_zarr.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index 157eb7170c..1b5debe1e8 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py index adacbc9678..f43cc7b141 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index 44b1e5f416..130670dd4b 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/metrics/climate/hydrostasy.py b/physicsnemo/metrics/climate/hydrostasy.py index 5099c7095a..8cabc6f5be 100644 --- a/physicsnemo/metrics/climate/hydrostasy.py +++ b/physicsnemo/metrics/climate/hydrostasy.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py index db29af0d19..b8b7342057 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py index 522b514f24..0eaf2460c4 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py @@ -1,3 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index 51c4b47e1f..30c541b755 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/test/datapipes/test_healpix_couple_zarr.py b/test/datapipes/test_healpix_couple_zarr.py index 61aaca854a..a5d5e32ec2 100644 --- a/test/datapipes/test_healpix_couple_zarr.py +++ b/test/datapipes/test_healpix_couple_zarr.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/test/datapipes/test_healpix_zarr.py b/test/datapipes/test_healpix_zarr.py index fb4fe1aef8..05fc8d3c34 100644 --- a/test/datapipes/test_healpix_zarr.py +++ b/test/datapipes/test_healpix_zarr.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/test/metrics/test_hydrostatic_loss.py b/test/metrics/test_hydrostatic_loss.py index 4c3058814b..4b7d660f44 100644 --- a/test/metrics/test_hydrostatic_loss.py +++ b/test/metrics/test_hydrostatic_loss.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. # SPDX-FileCopyrightText: All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/test/models/dlwp_healpix/_cln_reference.py b/test/models/dlwp_healpix/_cln_reference.py index e27227ac54..41d17716b9 100644 --- a/test/models/dlwp_healpix/_cln_reference.py +++ b/test/models/dlwp_healpix/_cln_reference.py @@ -1,3 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reference (old) implementation of ConditionalLayerNorm for testing. # This is a copy of the original code before optimization. diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py index 25e45265ba..d1f9ee5596 100644 --- a/test/models/dlwp_healpix/test_conditional_layer_norm.py +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -1,3 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys From abfbcac07d8377a39d9d9276e6aa4b42db872d17 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 8 Jul 2026 14:14:49 -0500 Subject: [PATCH 08/28] Doc updates --- physicsnemo/datapipes/healpix/__init__.py | 37 +++++++++++- .../healpix/base_timeseries_dataset_zarr.py | 8 +++ .../healpix/coupledtimeseries_dataset.py | 16 +++++ .../healpix/coupledtimeseries_dataset_zarr.py | 6 ++ physicsnemo/datapipes/healpix/couplers.py | 31 ++++++---- physicsnemo/datapipes/healpix/data_modules.py | 15 +++++ .../datapipes/healpix/data_modules_zarr.py | 8 +++ .../datapipes/healpix/timeseries_dataset.py | 7 +++ .../healpix/timeseries_dataset_zarr.py | 8 +++ physicsnemo/metrics/climate/healpix_loss.py | 1 + physicsnemo/metrics/climate/hydrostasy.py | 53 ++++++++++++++++ .../models/dlwp_healpix/HEALPixRecUNet.py | 9 +++ .../models/dlwp_healpix/HEALPixUNet.py | 8 +++ physicsnemo/models/dlwp_healpix/__init__.py | 9 +++ .../models/dlwp_healpix/layers/__init__.py | 9 ++- .../dlwp_healpix/layers/healpix_blocks.py | 49 +++++++++++++++ .../layers/healpix_constraints.py | 15 +++++ .../dlwp_healpix/layers/healpix_decoder.py | 24 ++++++++ .../dlwp_healpix/layers/healpix_encoder.py | 22 +++++++ .../dlwp_healpix/layers/normalization.py | 32 ++++++++++ test/datapipes/test_healpix_couple_zarr.py | 7 --- test/metrics/test_hydrostatic_loss.py | 12 ++-- test/models/dlwp_healpix/_cln_reference.py | 24 ++++++-- .../test_conditional_layer_norm.py | 60 +++++++++++++------ 24 files changed, 417 insertions(+), 53 deletions(-) diff --git a/physicsnemo/datapipes/healpix/__init__.py b/physicsnemo/datapipes/healpix/__init__.py index ae1a6ce394..ee5b75e6cf 100644 --- a/physicsnemo/datapipes/healpix/__init__.py +++ b/physicsnemo/datapipes/healpix/__init__.py @@ -14,4 +14,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .data_modules import CoupledTimeSeriesDataModule +""" +Healpix DataPipes. + +This package contains the Healpix DataPipes for loading and processing healpix data. +It supports loading data from xarray Datasets or zarr Groups, and coupling it with a model. + +The main classes are: +- TimeSeriesDataModule: A DataModule for loading and processing time series data. +- CoupledTimeSeriesDataModule: A DataModule for loading and processing coupled time series data. +- TimeSeriesDataset: A dataset for loading and processing time series data. +- CoupledTimeSeriesDataset: A dataset for loading and processing coupled time series data. +- Zarr versions of the above classes. +- ConstantCoupler: A coupler that uses a constant value for the coupled data. +- TrailingAverageCoupler: A coupler that uses a trailing average for the coupled data. +""" + +from .coupledtimeseries_dataset import CoupledTimeSeriesDataset +from .coupledtimeseries_dataset_zarr import CoupledTimeSeriesDatasetZarr +from .couplers import ConstantCoupler, TrailingAverageCoupler +from .data_modules import CoupledTimeSeriesDataModule, TimeSeriesDataModule +from .data_modules_zarr import CoupledTimeSeriesDataModuleZarr, TimeSeriesDataModuleZarr +from .timeseries_dataset import TimeSeriesDataset +from .timeseries_dataset_zarr import TimeSeriesDatasetZarr + +__all__ = [ + "TimeSeriesDataModule", + "CoupledTimeSeriesDataModule", + "TimeSeriesDatasetZarr", + "CoupledTimeSeriesDatasetZarr", + "TimeSeriesDataset", + "CoupledTimeSeriesDataset", + "TimeSeriesDataModuleZarr", + "CoupledTimeSeriesDataModuleZarr", + "ConstantCoupler", + "TrailingAverageCoupler", +] diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py index 4d347e2480..ada8c00d84 100644 --- a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py @@ -14,6 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +BaseTimeSeriesDatasetZarr - Abstract base class for time series healpix datasets using Zarr storage. + +This class provides the core functionality for loading and processing time series healpix data +stored in Zarr format. It handles data loading, scaling, and time management. +Subclasses must implement the __getitem__ method to define specific data retrieval logic. +""" + import importlib.util import logging import os diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index 92ba830390..d63d81c089 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -14,6 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Dataset for coupling time series healpix data with external inputs from various earth system components. + +This class extends the TimeSeriesDataset to add the core functionality for coupling time series healpix data with external inputs from various earth system components. +""" + import gc import logging import time @@ -165,6 +171,10 @@ def __init__( self.curr_item = None # keeps track of current initialization def _get_scaling_da(self): + """ + Get the scaling for the coupled values from the scaling dictionary. + This is overridden to add the scaling for the coupled values. + """ scaling_df = pd.DataFrame.from_dict(self.scaling).T scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} scaling_da = scaling_df.to_xarray().astype("float32") @@ -175,6 +185,9 @@ def _get_scaling_da(self): super()._get_scaling_da() def __getitem__(self, item): + """ + Get the item from the dataset. This is overridden to add the coupled values. + """ # start range torch.cuda.nvtx.range_push("CoupledTimeSeriesDataset:__getitem__") @@ -340,6 +353,9 @@ def __getitem__(self, item): return inputs_result, targets def next_integration(self, model_outputs, constants): + """ + Fetch the next coupled integration step + """ inputs_result = [] # grab last few model outputs for re-initialization diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py index c3df1e1412..73ade041ab 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py @@ -14,6 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Dataset for coupling time series healpix data stored in Zarr format with external inputs from various earth system components. + +This class extends the TimeSeriesDatasetZarr to add the core functionality for coupling time series healpix data stored in Zarr format with external inputs from various earth system components. +""" + import logging from dataclasses import dataclass from typing import List, Optional, Sequence, Tuple, Union diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 101f916369..704133e7f2 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -117,9 +117,18 @@ def __init__( ) def _compute_coupled_integration_dim(self): + """ + Compute the coupled integration steps across the time dimension. + This is the number of presteps plus the maximum number of output times divided by the number of input times. + """ return self.presteps + max(self.output_time_dim // self.input_time_dim, 1) def _compute_timevar_dim(self): + """ + Compute the time variable dimension. + Multiple time integrations are stacked along the time dimension. + Thus it is the number of input times multiplied by the number of variables. + """ return len(self.input_times) * len(self.variables) @abstractmethod @@ -200,6 +209,12 @@ def setup_coupling(self, coupled_module): self.coupled_channel_indices = channel_indices def reset_coupler(self): + """Clear saved coupled fields, forces recomputation on next use. + + This method is called when the dataloader is reset, and it is used to clear the + cached coupled fields. This is necessary because the coupled fields are computed + on the fly, and they are not saved to the dataset. + """ self.coupled_mode = False self.integrated_couplings = None self.preset_coupled_fields = None @@ -494,7 +509,6 @@ def __init__( self.time_da = np.asarray(dates) else: self.time_da = self.ds.time.values - self._set_time_increments() def compute_coupled_indices(self, interval, data_time_step): """ @@ -522,18 +536,11 @@ def compute_coupled_indices(self, interval, data_time_step): self._coupled_offsets = self._coupled_offsets.astype(int) - def _set_time_increments(self): - # get the dt of the dataset - dt = pd.Timedelta(self.time_da[1] - self.time_da[0]).total_seconds() - # assert that the time increments are divisible by the dt of the dataset - if np.any([t.total_seconds() % dt != 0 for t in self.input_times]): - raise ValueError( - f"Coupled input times {self.input_times} " - f"({[t.total_seconds() for t in self.input_times]} in secs) are not divisible by dataset dt: {dt}" - ) - self.time_increments = [t.total_seconds() / dt for t in self.input_times] - def setup_coupling(self, coupled_module): + """ + Setup the trailing average specific coupling between the coupled variables and the provided module. + Precomputes the averaging slices for the coupled variables. + """ # Call parent method first to set basic coupling super().setup_coupling(coupled_module) diff --git a/physicsnemo/datapipes/healpix/data_modules.py b/physicsnemo/datapipes/healpix/data_modules.py index fbb45bb6f6..931f5d20d6 100644 --- a/physicsnemo/datapipes/healpix/data_modules.py +++ b/physicsnemo/datapipes/healpix/data_modules.py @@ -14,6 +14,18 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +TimeSeriesDataModule - DataModule that sets up the dataloaders for the time series healpix data. + +This class provides the core functionality for setting up the dataloaders for the time series healpix data. +Depending on the splits and forecast_init_times, it will set up the appropriate dataloaders for the training, validation, and test sets. +It also supports coupling the time series data with external inputs from various earth system components. + +The main classes are: +- TimeSeriesDataModule: A DataModule for loading and processing time series healpix data. +- CoupledTimeSeriesDataModule: A DataModule for loading and processing coupled time series healpix data. +""" + # System modules import logging import warnings @@ -596,6 +608,9 @@ def __init__( ) def _get_coupled_vars(self): + """ " + Get the coupled variables from the couplings dictionary. + """ coupled_variables = [] for d in self.couplings: coupled_variables = coupled_variables + d["params"]["variables"] diff --git a/physicsnemo/datapipes/healpix/data_modules_zarr.py b/physicsnemo/datapipes/healpix/data_modules_zarr.py index 5ef837d1bd..988c03b738 100644 --- a/physicsnemo/datapipes/healpix/data_modules_zarr.py +++ b/physicsnemo/datapipes/healpix/data_modules_zarr.py @@ -14,6 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +DataModule for loading and processing time series healpix data stored in Zarr format. + +This class provides the core functionality for setting up the dataloaders for the time series healpix data stored in Zarr format. +Data loading, scaling, and time management are handled by the TimeSeriesDatasetZarr and CoupledTimeSeriesDatasetZarr classes. +Zarr format is used to avoid Xarray overheads when loading and processing the data. +""" + # System modules import logging import warnings diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index 1b5debe1e8..b4c06be594 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -14,6 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Dataset for sampling from continuous time-series data, compatible with pytorch data loading. + +This class provides the core functionality for sampling from continuous time-series data, compatible with pytorch data loading. +It handles data loading, scaling, and time management. It is used by the TimeSeriesDataModule to set up the dataloaders for the time series healpix data. +""" + import logging import time import warnings diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py index f43cc7b141..79f08d3a93 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py @@ -14,6 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Dataset for sampling from continuous time-series data stored in Zarr format, compatible with pytorch data loading. + +This class provides the core functionality for sampling from continuous time-series data stored in Zarr format, compatible with pytorch data loading. +It handles data loading, scaling, and time management. It is used by the TimeSeriesDataModule to set up the dataloaders for the time series healpix data stored in Zarr format. +Zarr format is used to avoid Xarray overheads when loading and processing the data. +""" + import logging from dataclasses import dataclass from typing import List, Optional, Sequence, Tuple, Union diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index 130670dd4b..169dfbc902 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -168,6 +168,7 @@ def __init__( self.w = w def setup(self, trainer): + """Setup function that moves tensors to the same device as the model""" self.b = th.tensor(self.b, device=trainer.device) self.w = th.tensor(self.w, device=trainer.device) diff --git a/physicsnemo/metrics/climate/hydrostasy.py b/physicsnemo/metrics/climate/hydrostasy.py index 8cabc6f5be..db677099e5 100644 --- a/physicsnemo/metrics/climate/hydrostasy.py +++ b/physicsnemo/metrics/climate/hydrostasy.py @@ -33,6 +33,14 @@ def _virtual_temperature_from_geopotential_height(z1, z2, p1, p2, T1, R, g0): class HydrostaticBalance(torch.nn.Module): + """Derive virtual temperature at pressure levels from geopotential height. + + Given geopotential heights (Z) at a set of pressure levels and a known + virtual temperature anchor, integrates the hypsometric equation outward + from the anchor level to recover virtual temperature at every other + level. + """ + def __init__( self, z_pressure_levels: Dict[int, float], @@ -119,6 +127,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class DifferentialHydrostaticBalanceConstraint(torch.nn.Module): + """Compute layer-average virtual temperatures implied by geopotential height and model output. + + For each pair of adjacent pressure levels, derives the layer-mean virtual + temperature implied by the hypsometric equation from geopotential height + (``Tv_avg``) alongside the layer-mean virtual temperature predicted + directly by the model (``Tv_model_avg``), so the two can be compared as a + hydrostatic-balance constraint. + """ + def __init__( self, z_pressure_levels: Dict[int, float], @@ -433,6 +450,24 @@ def scale(self, x): return x_scaled.reshape((-1, F, C_scaled, H, W)) def error_histogram(self, prediction, bins, accumulator=None): + """Accumulate a per-level histogram of absolute hydrostatic-balance error. + + Parameters + ---------- + prediction: torch.Tensor + Model prediction tensor of shape [N, F, B, C, H, W]. + bins: int or torch.Tensor + Either the number of equal-width bins to use, or an existing + tensor of per-level bin edges to reuse. + accumulator: torch.Tensor, optional + Existing per-level histogram counts to add to. If None, a new + zero-initialized accumulator is created. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + The updated ``(accumulator, bin_edges)``. + """ N, F, B, C, H, W = tuple(prediction.shape) if not (prediction.ndim == 6): @@ -742,6 +777,24 @@ def scale(self, x): return x_scaled.reshape((-1, F, C_scaled, H, W)) def error_histogram(self, prediction, bins, accumulator=None): + """Accumulate a per-level histogram of absolute hydrostatic-balance error. + + Parameters + ---------- + prediction: torch.Tensor + Model prediction tensor of shape [N, F, B, C, H, W]. + bins: int or torch.Tensor + Either the number of equal-width bins to use, or an existing + tensor of per-level bin edges to reuse. + accumulator: torch.Tensor, optional + Existing per-level histogram counts to add to. If None, a new + zero-initialized accumulator is created. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + The updated ``(accumulator, bin_edges)``. + """ N, F, B, C, H, W = tuple(prediction.shape) if not (prediction.ndim == 6): diff --git a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py index 15775b4888..9127dec57f 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py @@ -14,6 +14,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Implementation of the Deep Learning Weather Prediction (DLWP) recurrent UNet on the HEALPix mesh. + +This class provides the core functionality for the DLWP recurrent UNet on the HEALPix mesh. +It handles the forward pass of the model, the backward pass, and the initialization of the hidden states. +It also supports coupling the model with external inputs from various earth system components. + +""" + import logging from dataclasses import dataclass from typing import Any, Dict, Sequence diff --git a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py index 1f76da6560..0f46f5501f 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py @@ -14,6 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Implementation of the Deep Learning Weather Prediction (DLWP) UNet on the HEALPix mesh. + +This class provides the core functionality for the DLWP UNet on the HEALPix mesh. +It handles the forward pass of the model, the backward pass, and the initialization of the hidden states. +It also supports coupling the model with external inputs from various earth system components. +""" + import logging from dataclasses import dataclass from typing import Any, Dict, Sequence diff --git a/physicsnemo/models/dlwp_healpix/__init__.py b/physicsnemo/models/dlwp_healpix/__init__.py index b61924fea0..b1c8797e2c 100644 --- a/physicsnemo/models/dlwp_healpix/__init__.py +++ b/physicsnemo/models/dlwp_healpix/__init__.py @@ -14,6 +14,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +This package contains the implementation of the Deep Learning Weather Prediction (DLWP) recurrent UNet on the HEALPix mesh. +It handles the forward pass of the model, the backward pass, and the initialization of the hidden states. + +The main classes are: +- HEALPixRecUNet: The main class for the DLWP recurrent UNet on the HEALPix mesh. +- HEALPixUNet: The main class for the DLWP UNet on the HEALPix mesh. +""" + from .HEALPixRecUNet import HEALPixRecUNet from .HEALPixUNet import HEALPixUNet diff --git a/physicsnemo/models/dlwp_healpix/layers/__init__.py b/physicsnemo/models/dlwp_healpix/layers/__init__.py index abfdf22788..dec4636954 100644 --- a/physicsnemo/models/dlwp_healpix/layers/__init__.py +++ b/physicsnemo/models/dlwp_healpix/layers/__init__.py @@ -66,10 +66,10 @@ ] -# Remapping methods for backwards compatibility of legacy checkpoints - - def _remap_target(target: str) -> str: + """ + Remapping methods for backwards compatibility of legacy checkpoints + """ explicit = { "physicsnemo.models.dlwp_healpix_layers.healpix_encoder.UNetEncoder": "physicsnemo.models.dlwp_healpix.layers.UNetEncoder", "physicsnemo.models.dlwp_healpix_layers.healpix_decoder.UNetDecoder": "physicsnemo.models.dlwp_healpix.layers.UNetDecoder", @@ -139,6 +139,9 @@ def _remap_target(target: str) -> str: def _remap_obj(obj): + """ + Remapping of Dictionary and Hydra DictConfig objects to new targets. + """ from omegaconf import DictConfig, OmegaConf if isinstance(obj, DictConfig): diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py index b8b7342057..451193df6b 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py @@ -14,6 +14,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Implementation of the Deep Learning Weather Prediction (DLWP) convolutional and recurrent blocks on the HEALPix mesh. + +This module contains the implementation of the Deep Learning Weather Prediction (DLWP) convolutional and recurrent blocks on the HEALPix mesh. +The main classes are: +- ConvGRUBlock: A class for the DLWP convolutional GRU block on the HEALPix mesh. +- BasicConvBlock: A class for the DLWP basic convolutional block on the HEALPix mesh. +- ConvNeXtBlock: A class for the DLWP ConvNeXt block on the HEALPix mesh. +- DoubleConvNeXtBlock: A class for the DLWP double ConvNeXt block on the HEALPix mesh. +- Multi_SymmetricConvNeXtBlock: A class for the DLWP multi symmetric ConvNeXt block on the HEALPix mesh. +- SymmetricConvNeXtBlock: A class for the DLWP symmetric ConvNeXt block on the HEALPix mesh. +- TransposedConvUpsample: A class for the DLWP transposed convolutional upsample block on the HEALPix mesh. +- Interpolate: A class for the DLWP interpolate block on the HEALPix mesh. +""" + from typing import Callable, Sequence, Tuple, Union import torch @@ -31,10 +46,30 @@ class _LayerNormOverChannels(th.nn.Module): """Applies nn.LayerNorm over the channel dimension for (B, C, H, W) tensors.""" def __init__(self, channel_depth: int, eps: float = 1e-5): + """ + Parameters + ---------- + channel_depth: int + The number of channels in the input tensor + eps: float, optional + The epsilon value for the layer norm + """ super().__init__() self.norm = th.nn.LayerNorm(channel_depth, eps=eps) def forward(self, x): + """Forward pass of the _LayerNormOverChannels + + Parameters + ---------- + x: torch.Tensor + The input tensor + + Returns + ------- + torch.Tensor + The normed output tensor + """ x = x.permute(0, 2, 3, 1) x = self.norm(x) return x.permute(0, 3, 1, 2) @@ -693,6 +728,20 @@ def __init__( ) def forward(self, x, conditions_cln=None): + """Forward pass of the Multi_SymmetricConvNeXtBlock + + Parameters + ---------- + x: torch.Tensor + The input tensor + conditions_cln: torch.Tensor, optional + The conditions for the conditional layer normalization + + Returns + ------- + torch.Tensor + The output tensor + """ out = x for block in self.blocks: out = block(out, conditions_cln=conditions_cln) diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py index 0eaf2460c4..4c328b8e43 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py @@ -14,10 +14,25 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Implementation of the Deep Learning Weather Prediction (DLWP) constraints on the HEALPix mesh. + +This module contains the implementation of the Deep Learning Weather Prediction (DLWP) constraints on the HEALPix mesh. +The main classes are: +- NonnegativeConstraint: A class for the DLWP nonnegative constraint on the HEALPix mesh. +""" + import torch class NonnegativeConstraint(torch.nn.Module): + """Clamp selected output channels to be nonnegative in physical units. + + Applies a lower bound of zero (in unnormalized/physical space) to the + specified ``variables`` by clamping the corresponding normalized channels + to their per-variable, per-scaling threshold. + """ + def __init__( self, variables: list[str], diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py index a719c01540..454af91589 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py @@ -14,6 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Implementation of the Deep Learning Weather Prediction (DLWP) decoder on the HEALPix mesh. + +This class contains the implementation of the Deep Learning Weather Prediction (DLWP) decoder on the HEALPix mesh. +""" + from typing import Sequence import torch @@ -171,6 +177,24 @@ def _forward_layer_pass( skip_connection: torch.Tensor = None, conditions_cln: torch.Tensor = None, ) -> torch.Tensor: + """Helper function that performs the forward pass of a single layer of the decoder. + + Parameters + ---------- + layer: torch.nn.Module + The layer to forward pass + x: torch.Tensor + The input tensor + skip_connection: torch.Tensor, optional + The skip connection tensor + conditions_cln: torch.Tensor, optional + The conditional inputs for the normalization layers. + + Returns + ------- + torch.Tensor + The output tensor + """ if layer["upsamp"] is not None: up = layer["upsamp"](x) x = torch.cat([up, skip_connection], dim=self.channel_dim) diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py index 8fe0749b44..bbe874be83 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py @@ -14,6 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Implementation of the Deep Learning Weather Prediction (DLWP) encoder on the HEALPix mesh. + +This class contains the implementation of the Deep Learning Weather Prediction (DLWP) encoder on the HEALPix mesh. +""" + from typing import Sequence import torch @@ -141,6 +147,22 @@ def _forward_layer_pass( inp: torch.Tensor, conditions_cln: torch.Tensor = None, ) -> torch.Tensor: + """Helper function that performs the forward pass of a single layer of the encoder. + + Parameters + ---------- + layer_group: torch.nn.Module + The layer group to forward pass + inp: torch.Tensor + The input tensor + conditions_cln: torch.Tensor, optional + The conditional inputs for the normalization layers. + + Returns + ------- + torch.Tensor + The output tensor + """ interim_output = inp for layer in layer_group: if getattr(layer, "cln_enabled", False): diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index 30c541b755..697994269b 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -14,6 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Implementation of the Deep Learning Weather Prediction (DLWP) normalization on the HEALPix mesh. + +This class contains the implementation of the Deep Learning Weather Prediction (DLWP) normalization on the HEALPix mesh. +""" + from typing import List import torch as th @@ -41,6 +47,14 @@ def _cln_affine(x_norm, gamma_raw, beta, scale_center, n_faces): class ConditionalLayerNorm(th.nn.Module): + """LayerNorm whose affine (gamma/beta) parameters are predicted from a conditional input field. + + Normalizes the input over the channel dimension with no learnable affine + parameters, then predicts per-channel scale and shift from ``conditions`` + via a fused MLP and applies them, optionally recentered by + ``scale_center``. + """ + def __init__( self, condition_shape: int, @@ -114,6 +128,24 @@ def _make_mlp( out_dim: int, activation: th.nn.Module, ) -> th.nn.Sequential: + """Helper function that creates the MLP for the conditional layer normalization. + + Parameters + ---------- + in_dim: int + The input dimension + hidden_dims: List[int] + The hidden dimensions + out_dim: int + The output dimension + activation: th.nn.Module + The activation function + + Returns + ------- + th.nn.Sequential + The MLP + """ layers = [] for hdim in hidden_dims: layers.append(th.nn.Linear(in_dim, hdim)) diff --git a/test/datapipes/test_healpix_couple_zarr.py b/test/datapipes/test_healpix_couple_zarr.py index a5d5e32ec2..2a7f691350 100644 --- a/test/datapipes/test_healpix_couple_zarr.py +++ b/test/datapipes/test_healpix_couple_zarr.py @@ -145,7 +145,6 @@ def scaling_double_dict(): @import_or_fail("xarray") @nfsdata_or_fail def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.couplers import ( ConstantCoupler, ) @@ -290,7 +289,6 @@ def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): @import_or_fail("xarray") @nfsdata_or_fail def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.couplers import ( TrailingAverageCoupler, ) @@ -449,7 +447,6 @@ def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): def test_CoupledTimeSeriesDatasetZarr_initialization( dataset_path, scaling_dict, pytestconfig ): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( CoupledTimeSeriesDatasetZarr, ) @@ -571,7 +568,6 @@ def test_CoupledTimeSeriesDatasetZarr_initialization( def test_CoupledTimeSeriesDatasetZarr_get_constants( dataset_path, scaling_dict, constant_coupler_config, pytestconfig ): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( CoupledTimeSeriesDatasetZarr, ) @@ -876,7 +872,6 @@ def test_CoupledTimeSeriesDatasetZarr_get( def test_CoupledTimeSeriesDataModuleZarr_initialization( dataset_path, splits, scaling_double_dict, constant_coupler_config, pytestconfig ): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( CoupledTimeSeriesDataModuleZarr, ) @@ -940,7 +935,6 @@ def test_CoupledTimeSeriesDataModuleZarr_initialization( def test_CoupledTimeSeriesDataModuleZarr_get_constants( dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig ): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( CoupledTimeSeriesDataModuleZarr, ) @@ -1015,7 +1009,6 @@ def test_CoupledTimeSeriesDataModuleZarr_get_constants( def test_CoupledTimeSeriesDataModuleZarr_get_dataloaders( dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig ): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( CoupledTimeSeriesDataModuleZarr, ) diff --git a/test/metrics/test_hydrostatic_loss.py b/test/metrics/test_hydrostatic_loss.py index 4b7d660f44..a9152ab6de 100644 --- a/test/metrics/test_hydrostatic_loss.py +++ b/test/metrics/test_hydrostatic_loss.py @@ -14,10 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass -from pathlib import Path -from typing import Sequence - import numpy as np import pytest import torch @@ -28,6 +24,7 @@ xr = pytest.importorskip("xarray") + def test_constant_temperature(rtol: float = 1e-3, atol: float = 1e-3): R = 287 # J K^{-1} kg^{-1} g0 = 9.81 # m s^{-2} @@ -62,7 +59,12 @@ def test_constant_temperature(rtol: float = 1e-3, atol: float = 1e-3): Tv = constraint(x) print(Tv[0, 0, :, 0, 0]) - assert torch.allclose(T * torch.ones_like(Tv), Tv, rtol=rtol, atol=atol,) + assert torch.allclose( + T * torch.ones_like(Tv), + Tv, + rtol=rtol, + atol=atol, + ) @pytest.mark.parametrize("N", [10, 20]) diff --git a/test/models/dlwp_healpix/_cln_reference.py b/test/models/dlwp_healpix/_cln_reference.py index 41d17716b9..0dddf37e6a 100644 --- a/test/models/dlwp_healpix/_cln_reference.py +++ b/test/models/dlwp_healpix/_cln_reference.py @@ -17,14 +17,14 @@ # Reference (old) implementation of ConditionalLayerNorm for testing. # This is a copy of the original code before optimization. -from typing import List - import copy +from typing import List import torch as th try: from apex.normalization import FusedLayerNorm + _APEX_AVAILABLE = True except ImportError: _APEX_AVAILABLE = False @@ -49,8 +49,12 @@ def __init__( self.channel_depth = channel_depth self.hidden_dims = mlp_hidden_dims self.activation = activation if activation is not None else th.nn.Identity() - self.gamma_mlp = self._make_mlp(self.condition_shape, self.hidden_dims, self.channel_depth, self.activation) - self.beta_mlp = self._make_mlp(self.condition_shape, self.hidden_dims, self.channel_depth, self.activation) + self.gamma_mlp = self._make_mlp( + self.condition_shape, self.hidden_dims, self.channel_depth, self.activation + ) + self.beta_mlp = self._make_mlp( + self.condition_shape, self.hidden_dims, self.channel_depth, self.activation + ) self.n_faces = n_faces self.scale_center = scale_center @@ -64,10 +68,18 @@ def __init__( self.norm = th.nn.LayerNorm(channel_depth, elementwise_affine=False) elif norm_op == "apex": if not _APEX_AVAILABLE: - raise ImportError("Apex FusedLayerNorm requested but apex is not available") + raise ImportError( + "Apex FusedLayerNorm requested but apex is not available" + ) self.norm = FusedLayerNorm(channel_depth, elementwise_affine=False) - def _make_mlp(self, in_dim: int, hidden_dims: List[int], out_dim: int, activation: th.nn.Module) -> th.nn.Sequential: + def _make_mlp( + self, + in_dim: int, + hidden_dims: List[int], + out_dim: int, + activation: th.nn.Module, + ) -> th.nn.Sequential: layers = [] for hdim in hidden_dims: layers.append(th.nn.Linear(in_dim, hdim)) diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py index d1f9ee5596..8530f82c83 100644 --- a/test/models/dlwp_healpix/test_conditional_layer_norm.py +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -22,6 +22,7 @@ import pytest import torch from _cln_reference import ConditionalLayerNormReference + from physicsnemo.models.dlwp_healpix.layers.normalization import ConditionalLayerNorm from physicsnemo.nn import CappedGELU @@ -53,10 +54,13 @@ def _copy_old_to_new(old_cln, new_cln): old_sd = old_cln.state_dict() # Collect Linear layer indices from the old gamma MLP - gamma_linear_indices = sorted({ - int(k.split(".")[1]) - for k in old_sd if k.startswith("gamma_mlp.") and k.endswith(".weight") - }) + gamma_linear_indices = sorted( + { + int(k.split(".")[1]) + for k in old_sd + if k.startswith("gamma_mlp.") and k.endswith(".weight") + } + ) first_layer_idx = gamma_linear_indices[0] new_sd = {} @@ -79,10 +83,13 @@ def _copy_old_to_new(old_cln, new_cln): # Block-diagonal: [[gamma, 0], [0, beta]] out_old, in_old = gamma_val.shape zeros = torch.zeros_like(gamma_val) - new_sd[fused_key] = torch.cat([ - torch.cat([gamma_val, zeros], dim=1), - torch.cat([zeros, beta_val], dim=1), - ], dim=0) + new_sd[fused_key] = torch.cat( + [ + torch.cat([gamma_val, zeros], dim=1), + torch.cat([zeros, beta_val], dim=1), + ], + dim=0, + ) for key, value in old_sd.items(): if not key.startswith("gamma_mlp."): @@ -123,12 +130,14 @@ def test_old_vs_new_forward(n_cond, channels_last, scale_center): out_new = new_cln(x, cond) assert out_old.shape == out_new.shape - assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), \ + assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( f"Max diff: {(out_old - out_new).abs().max().item()}" + ) if channels_last: - assert out_new.is_contiguous(memory_format=torch.channels_last), \ + assert out_new.is_contiguous(memory_format=torch.channels_last), ( "Output should preserve channels_last format" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -162,11 +171,13 @@ def test_old_vs_new_backward(channels_last): out_new = new_cln(x_new, cond_new) out_new.sum().backward() - assert torch.allclose(x_old.grad, x_new.grad, atol=1e-4, rtol=1e-3), \ + assert torch.allclose(x_old.grad, x_new.grad, atol=1e-4, rtol=1e-3), ( f"Input grad max diff: {(x_old.grad - x_new.grad).abs().max().item()}" + ) - assert torch.allclose(cond_old.grad, cond_new.grad, atol=1e-4, rtol=1e-3), \ + assert torch.allclose(cond_old.grad, cond_new.grad, atol=1e-4, rtol=1e-3), ( f"Cond grad max diff: {(cond_old.grad - cond_new.grad).abs().max().item()}" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -192,8 +203,9 @@ def test_init_cln_to_zero_matches_layer_norm(channels_last): x_nhwc = x.permute(0, 2, 3, 1) out_ln = plain_ln(x_nhwc).permute(0, 3, 1, 2) - assert torch.allclose(out_cln, out_ln, atol=1e-5, rtol=1e-4), \ + assert torch.allclose(out_cln, out_ln, atol=1e-5, rtol=1e-4), ( f"Max diff: {(out_cln - out_ln).abs().max().item()}" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -253,17 +265,25 @@ def test_backward_channels_last_matches_contiguous(): cln.zero_grad() # Channels-last path - x_cl = x_base.clone().detach().to(memory_format=torch.channels_last).requires_grad_(True) + x_cl = ( + x_base.clone() + .detach() + .to(memory_format=torch.channels_last) + .requires_grad_(True) + ) cond_cl = cond_base.clone().detach().requires_grad_(True) out_cl = cln(x_cl, cond_cl) out_cl.sum().backward() - assert torch.allclose(out_cont, out_cl, atol=1e-5, rtol=1e-4), \ + assert torch.allclose(out_cont, out_cl, atol=1e-5, rtol=1e-4), ( f"Output max diff: {(out_cont - out_cl).abs().max().item()}" - assert torch.allclose(x_cont.grad, x_cl.grad, atol=1e-5, rtol=1e-4), \ + ) + assert torch.allclose(x_cont.grad, x_cl.grad, atol=1e-5, rtol=1e-4), ( f"Input grad max diff: {(x_cont.grad - x_cl.grad).abs().max().item()}" - assert torch.allclose(cond_cont.grad, cond_cl.grad, atol=1e-5, rtol=1e-4), \ + ) + assert torch.allclose(cond_cont.grad, cond_cl.grad, atol=1e-5, rtol=1e-4), ( f"Cond grad max diff: {(cond_cont.grad - cond_cl.grad).abs().max().item()}" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -287,8 +307,9 @@ def test_load_old_checkpoint(): assert out_new.shape == (12, C, 8, 8) assert torch.isfinite(out_new).all() - assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), \ + assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( f"Max diff after loading old checkpoint: {(out_old - out_new).abs().max().item()}" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -325,5 +346,6 @@ def test_load_old_checkpoint_with_capped_gelu(): out_old = old_cln(x, cond) out_new = new_cln(x, cond) - assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), \ + assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( f"Max diff after loading old CappedGELU checkpoint: {(out_old - out_new).abs().max().item()}" + ) From 522c3ac487be350847c5167debeae67c3253a324 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 8 Jul 2026 16:14:43 -0500 Subject: [PATCH 09/28] update tests for zarr dataloaders --- test/datapipes/test_healpix_couple_zarr.py | 83 +++++++++------------- test/datapipes/test_healpix_zarr.py | 57 ++++++--------- 2 files changed, 57 insertions(+), 83 deletions(-) diff --git a/test/datapipes/test_healpix_couple_zarr.py b/test/datapipes/test_healpix_couple_zarr.py index 2a7f691350..9767867644 100644 --- a/test/datapipes/test_healpix_couple_zarr.py +++ b/test/datapipes/test_healpix_couple_zarr.py @@ -17,15 +17,14 @@ import random import warnings from dataclasses import dataclass -from pathlib import Path import pytest import torch as th -from pytest_utils import import_or_fail, nfsdata_or_fail from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler from physicsnemo.distributed import DistributedManager +from test.conftest import requires_module omegaconf = pytest.importorskip("omegaconf") np = pytest.importorskip("numpy") @@ -35,11 +34,8 @@ @pytest.fixture -def dataset_path(): - data_dir = "/data/nfs/modulus-data/datasets/healpix/" - dataset_name = "healpix.zarr" - path = Path(data_dir, dataset_name) - return path +def dataset_path(nfs_data_dir): + return nfs_data_dir.joinpath("datasets/healpix/healpix.zarr") @pytest.fixture @@ -139,11 +135,10 @@ def scaling_double_dict(): return omegaconf.DictConfig(scaling) -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("pandas") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("pandas") +@requires_module("xarray") def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): from physicsnemo.datapipes.healpix.couplers import ( ConstantCoupler, @@ -283,11 +278,10 @@ def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("pandas") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("pandas") +@requires_module("xarray") def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): from physicsnemo.datapipes.healpix.couplers import ( TrailingAverageCoupler, @@ -440,10 +434,9 @@ def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("xarray") def test_CoupledTimeSeriesDatasetZarr_initialization( dataset_path, scaling_dict, pytestconfig ): @@ -561,10 +554,9 @@ def test_CoupledTimeSeriesDatasetZarr_initialization( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("xarray") def test_CoupledTimeSeriesDatasetZarr_get_constants( dataset_path, scaling_dict, constant_coupler_config, pytestconfig ): @@ -615,10 +607,9 @@ def test_CoupledTimeSeriesDatasetZarr_get_constants( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("xarray") def test_CoupledTimeSeriesDatasetZarr_len( dataset_path, scaling_dict, constant_coupler_config, pytestconfig ): @@ -684,10 +675,9 @@ def test_CoupledTimeSeriesDatasetZarr_len( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("xarray") def test_CoupledTimeSeriesDatasetZarr_get( dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig ): @@ -865,10 +855,9 @@ def test_CoupledTimeSeriesDatasetZarr_get( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("xarray") def test_CoupledTimeSeriesDataModuleZarr_initialization( dataset_path, splits, scaling_double_dict, constant_coupler_config, pytestconfig ): @@ -928,10 +917,9 @@ def test_CoupledTimeSeriesDataModuleZarr_initialization( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("xarray") def test_CoupledTimeSeriesDataModuleZarr_get_constants( dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig ): @@ -1004,8 +992,7 @@ def test_CoupledTimeSeriesDataModuleZarr_get_constants( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@nfsdata_or_fail +@requires_module("omegaconf") def test_CoupledTimeSeriesDataModuleZarr_get_dataloaders( dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig ): @@ -1055,8 +1042,7 @@ def test_CoupledTimeSeriesDataModuleZarr_get_dataloaders( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@nfsdata_or_fail +@requires_module("omegaconf") def test_CoupledTimeSeriesDataModuleZarr_get_coupled_vars( dataset_path, scaling_double_dict, @@ -1107,10 +1093,9 @@ def test_CoupledTimeSeriesDataModuleZarr_get_coupled_vars( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("xarray") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("xarray") def test_CoupledTimeSeriesDatasetZarr_next_integration( dataset_path, scaling_dict, pytestconfig ): diff --git a/test/datapipes/test_healpix_zarr.py b/test/datapipes/test_healpix_zarr.py index 05fc8d3c34..6b59ef8805 100644 --- a/test/datapipes/test_healpix_zarr.py +++ b/test/datapipes/test_healpix_zarr.py @@ -17,14 +17,13 @@ import importlib.util import random import warnings -from pathlib import Path import pytest -from pytest_utils import import_or_fail, nfsdata_or_fail from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler from physicsnemo.distributed import DistributedManager +from test.conftest import requires_module omegaconf = pytest.importorskip("omegaconf") np = pytest.importorskip("numpy") @@ -33,11 +32,8 @@ @pytest.fixture -def dataset_path(): - data_dir = "/data/nfs/modulus-data/datasets/healpix/" - dataset_name = "healpix.zarr" - path = Path(data_dir, dataset_name) - return path +def dataset_path(nfs_data_dir): + return nfs_data_dir.joinpath("datasets/healpix/healpix.zarr") @pytest.fixture @@ -89,9 +85,8 @@ def scaling_double_dict(): return omegaconf.DictConfig(scaling) -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") def test_TimeSeriesDataset_initialization( dataset_path, scaling_dict, @@ -276,11 +271,10 @@ def test_TimeSeriesDataset_initialization( assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("numpy") -@import_or_fail("zarr") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("numpy") +@requires_module("zarr") def test_TimeSeriesDataset_get_constants(dataset_path, scaling_dict, pytestconfig): from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( TimeSeriesDatasetZarr, @@ -338,9 +332,8 @@ def test_TimeSeriesDataset_get_constants(dataset_path, scaling_dict, pytestconfi zarr_ds.close() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") def test_TimeSeriesDataset_len(dataset_path, scaling_dict, pytestconfig): from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( TimeSeriesDatasetZarr, @@ -398,10 +391,9 @@ def test_TimeSeriesDataset_len(dataset_path, scaling_dict, pytestconfig): DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("numpy") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("numpy") def test_TimeSeriesDataset_get(dataset_path, scaling_double_dict, splits, pytestconfig): from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( TimeSeriesDatasetZarr, @@ -541,10 +533,9 @@ def test_TimeSeriesDataset_get(dataset_path, scaling_double_dict, splits, pytest assert (len(inputs) + 1) == len(timeseries_ds[0]) -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("zarr") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("zarr") def test_TimeSeriesDataModule_initialization( dataset_path, splits, scaling_double_dict, pytestconfig ): @@ -642,11 +633,10 @@ def test_TimeSeriesDataModule_initialization( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("netCDF4") -@import_or_fail("numpy") -@import_or_fail("zarr") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("numpy") +@requires_module("zarr") def test_TimeSeriesDataModule_get_constants( dataset_path, scaling_double_dict, splits, pytestconfig ): @@ -731,9 +721,8 @@ def test_TimeSeriesDataModule_get_constants( DistributedManager.cleanup() -@import_or_fail("omegaconf") -@import_or_fail("zarr") -@nfsdata_or_fail +@requires_module("omegaconf") +@requires_module("zarr") def test_TimeSeriesDataModule_get_dataloaders( dataset_path, scaling_double_dict, splits, pytestconfig ): From 2c508f2adcd9f8be2aad4e6b73bbcb370581940f Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Thu, 9 Jul 2026 14:00:17 -0500 Subject: [PATCH 10/28] Update dataloaeder tests --- .../healpix/base_timeseries_dataset_zarr.py | 61 +- .../healpix/coupledtimeseries_dataset.py | 13 +- physicsnemo/datapipes/healpix/couplers.py | 50 +- .../datapipes/healpix/timeseries_dataset.py | 9 +- test/datapipes/healpix/__init__.py | 17 + test/datapipes/healpix/conftest.py | 190 +++++ test/datapipes/{ => healpix}/test_healpix.py | 281 ++------ .../{ => healpix}/test_healpix_couple.py | 658 +++++++++--------- .../{ => healpix}/test_healpix_couple_zarr.py | 412 ++++++----- .../{ => healpix}/test_healpix_zarr.py | 266 ++++--- 10 files changed, 1117 insertions(+), 840 deletions(-) create mode 100644 test/datapipes/healpix/__init__.py create mode 100644 test/datapipes/healpix/conftest.py rename test/datapipes/{ => healpix}/test_healpix.py (67%) rename test/datapipes/{ => healpix}/test_healpix_couple.py (69%) rename test/datapipes/{ => healpix}/test_healpix_couple_zarr.py (77%) rename test/datapipes/{ => healpix}/test_healpix_zarr.py (74%) diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py index ada8c00d84..bb55abcd67 100644 --- a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py @@ -337,19 +337,35 @@ def _get_time_da( if "time" not in ds: raise KeyError(f"Dataset missing time. Dataset provided {dataset_path}") - if np.datetime64(start_date) < ds.time[0]: + # integer start/end dates are indices into the time array rather than + # dates themselves (see the int handling below), so they can't be + # compared against `ds.time` via `np.datetime64` + if ( + start_date is not None + and not isinstance(start_date, int) + and np.datetime64(start_date) < ds.time[0] + ): warnings.warn( f"Start date {start_date} is before first available date {ds.time[0].values}" ) - if ds.time[-1] < np.datetime64(end_date): + if ( + end_date is not None + and not isinstance(end_date, int) + and ds.time[-1] < np.datetime64(end_date) + ): warnings.warn( f"End date {end_date} is after last available date {ds.time[-1].values}" ) # used when we need all the dates to calculate things like offset indices self.time_da = ds.time.copy(deep=True) - # a list of dates that is available to fetch, used for things like inferencers - self.times = self.time_da.sel(time=slice(start_date, end_date)) + # a list of dates that is available to fetch, used for things like inferencers. + # integer start/end dates are positional indices into the time array rather + # than dates themselves, so they need `isel` instead of the label-based `sel` + if isinstance(start_date, int) or isinstance(end_date, int): + self.times = self.time_da.isel(time=slice(start_date, end_date)) + else: + self.times = self.time_da.sel(time=slice(start_date, end_date)) self.total_samples = self.times.shape[0] if start_date: @@ -481,27 +497,20 @@ def _get_scaling_da(self) -> None: f"Constant variables {missing_constants} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" ) - try: - self.constant_scaling = scaling_da.sel( - index=self.constant_variables - ).rename({"index": "channel_out"}) - self.constant_scaling = { - "mean": np.expand_dims( - self.constant_scaling["mean"].values.copy(), (1, 2, 3) - ), - "std": np.expand_dims( - self.constant_scaling["std"].values.copy(), (1, 2, 3) - ), - } - except (ValueError, KeyError): - missing = [ - m - for m in self.constant_variables - if m not in list(self.scaling.keys()) - ] - raise KeyError( - f"Constant channels {missing} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" - ) + # no try/except needed here: the missing_constants check above + # already guarantees every name in self.constant_variables is + # present in self.scaling, which is what backs scaling_da's index + self.constant_scaling = scaling_da.sel( + index=self.constant_variables + ).rename({"index": "channel_out"}) + self.constant_scaling = { + "mean": np.expand_dims( + self.constant_scaling["mean"].values.copy(), (1, 2, 3) + ), + "std": np.expand_dims( + self.constant_scaling["std"].values.copy(), (1, 2, 3) + ), + } def get_constants(self) -> np.ndarray: """Get constant fields used in dataset. @@ -627,4 +636,4 @@ def __getitem__( H = height W = width """ - pass + pass # pragma: no cover - abstract method body, always overridden diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index d63d81c089..ca6e51ff4a 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -221,13 +221,18 @@ def __getitem__(self, item): ) # for models with extra outputs + # an empty couplings list means zero coupled variables, rather than + # indexing into a non-existent first coupling + num_coupled_vars = len(self.couplings[0].variables) if self.couplings else 0 if len(self.ds["targets"].channel_out) != ( - len(self.ds["inputs"].channel_in) - len(self.couplings[0].variables) + len(self.ds["inputs"].channel_in) - num_coupled_vars ): + channel_slice = ( + slice(None, -num_coupled_vars) if num_coupled_vars else slice(None) + ) input_array = ( - input_array - - self.input_scaling["mean"][:, : -len(self.couplings[0].variables)] - ) / self.input_scaling["std"][:, : -len(self.couplings[0].variables)] + input_array - self.input_scaling["mean"][:, channel_slice] + ) / self.input_scaling["std"][:, channel_slice] else: input_array = ( input_array - self.input_scaling["mean"] diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 704133e7f2..7e4b5cefd3 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -45,7 +45,6 @@ def __init__( output_time_dim: int = 2, input_times: Sequence = [pd.Timedelta("24h"), pd.Timedelta("48h")], prepared_coupled_data: bool = True, - time_first: bool = True, ): """ Parameters @@ -74,9 +73,6 @@ def __init__( averages have already been calculated so that each time step denotes the right side of a averaging_window window. This is highly recommended for training, default True - time_first: boolean, optional - Whether the coupled data should be permuted to have the time dimension first - [T, B, C, F, H, W] rather than [B, F, T, C, H, W] """ # extract important meta data from ds self.ds = dataset @@ -96,7 +92,6 @@ def __init__( self.coupled_mode = False self.integrated_couplings = None self.ds_variable_indices = [] - self.time_first = time_first if not prepared_coupled_data: raise NotImplementedError("Data preparation not yet implemented") @@ -105,11 +100,15 @@ def __init__( self.use_zarr = False elif isinstance(self.ds, zr.Group): self.use_zarr = True + # iterate over self.variables in the outer loop so the selected + # indices follow the order of self.variables (matching the xarray + # path's `.sel(channel_in=self.variables)`), not the dataset's + # native channel_in order. Downstream code (e.g. coupled_scaling, + # built from self.variables order in set_scaling) assumes the + # coupled channel axis is ordered by self.variables. + channel_in = list(self.ds["channel_in"]) self.ds_variable_indices = [ - i - for i, ic in enumerate(self.ds["channel_in"]) - for v in self.variables - if ic == v + i for v in self.variables for i, ic in enumerate(channel_in) if ic == v ] else: raise TypeError( @@ -144,7 +143,7 @@ def compute_coupled_indices(self, interval, data_time_step): data_time_step: dataset timestep """ - pass + pass # pragma: no cover - abstract method body, always overridden def set_scaling(self, scaling_da): """ @@ -231,7 +230,7 @@ def set_coupled_fields(self, coupled_fields: th.tensor): The data to use when the dataloader requests coupled fields. Expected format is [B, F, T, C, H, W] """ - pass + pass # pragma: no cover - abstract method body, always overridden def _construct_integrated_couplings_from_dataset(self, batch, bsize): """ @@ -314,12 +313,11 @@ def construct_integrated_couplings( self.integrated_couplings[b, i, :, :, :] = coupling_temp.reshape( (self.timevar_dim,) + coupling_temp.shape[2:] ) - if self.time_first: - return self.integrated_couplings.transpose((1, 0, 2, 3, 4, 5)).astype( - "float32" - ) # cast to float for compatibility - else: - return self.integrated_couplings.astype("float32") + # permute to have the time dimension first [T, B, C, F, H, W] + # rather than [B, F, T, C, H, W], and cast to float for compatibility + return self.integrated_couplings.transpose((1, 0, 2, 3, 4, 5)).astype( + "float32" + ) class ConstantCoupler(BaseCoupler): @@ -428,10 +426,11 @@ def set_coupled_fields(self, coupled_fields: th.tensor): ) # broadcast the first time step to all the integration steps self.preset_coupled_fields[:, :, :, :, :, :] = coupled_fields[:, :, :1, :, :, :] - if self.time_first: - self.preset_coupled_fields = self.preset_coupled_fields.permute( - 2, 0, 3, 1, 4, 5 - ) + # permute to have the time dimension first [T, B, C, F, H, W] + # rather than [B, F, T, C, H, W] + self.preset_coupled_fields = self.preset_coupled_fields.permute( + 2, 0, 3, 1, 4, 5 + ) # flag for construct integrated coupling method to use this array self.coupled_mode = True @@ -585,9 +584,10 @@ def set_coupled_fields(self, coupled_fields: th.tensor): ] coupled_averaging_periods.append(th.concat(averaging_periods, dim=3)) self.preset_coupled_fields = th.concat(coupled_averaging_periods, dim=2) - if self.time_first: - self.preset_coupled_fields = self.preset_coupled_fields.permute( - 2, 0, 3, 1, 4, 5 - ) + # permute to have the time dimension first [T, B, C, F, H, W] + # rather than [B, F, T, C, H, W] + self.preset_coupled_fields = self.preset_coupled_fields.permute( + 2, 0, 3, 1, 4, 5 + ) # flag for construct integrated coupling method to use this array self.coupled_mode = True diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index b4c06be594..f06ce803c3 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -246,13 +246,16 @@ def _get_scaling_da(self): scaling_da = scaling_df.to_xarray().astype("float32") # REMARK: we remove the xarray overhead from these + # base (non-coupled) datasets have no 'couplings' attribute, and coupled + # datasets may be configured with an empty couplings list; treat both as + # zero coupled variables rather than indexing into an empty/missing list. + couplings = getattr(self, "couplings", []) + num_coupled_vars = len(couplings[0].variables) if couplings else 0 try: # 'if' statement used for cases where atmos model # includes diagnostic variables like tp6 and msl. # using 'channel_out' is still necessary for ocean models. - if len(self.ds.channel_out) != ( - len(self.ds.channel_in) - len(self.couplings[0].variables) - ): + if len(self.ds.channel_out) != (len(self.ds.channel_in) - num_coupled_vars): self.input_scaling = scaling_da.sel( index=self.ds.channel_in.values ).rename({"index": "channel_in"}) diff --git a/test/datapipes/healpix/__init__.py b/test/datapipes/healpix/__init__.py new file mode 100644 index 0000000000..c1b40ccc80 --- /dev/null +++ b/test/datapipes/healpix/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Test suite for the HEALPix datapipes.""" diff --git a/test/datapipes/healpix/conftest.py b/test/datapipes/healpix/conftest.py new file mode 100644 index 0000000000..510dd80510 --- /dev/null +++ b/test/datapipes/healpix/conftest.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared fixtures/helpers for the test_healpix*.py suite (classic/Zarr, +coupled/uncoupled). + +These describe the same on-disk NFS test dataset and the same +scaling/coupling configuration shapes across all four test_healpix*.py +modules in this package, so they're centralized here instead of being +duplicated (with minor, incidental drift) in each file. +""" + +from dataclasses import dataclass + +import pytest + + +@pytest.fixture +def data_dir(nfs_data_dir): + """Directory of the classic (non-Zarr) prebuilt test dataset.""" + return nfs_data_dir.joinpath("datasets/healpix") + + +@pytest.fixture +def dataset_name(): + return "healpix" + + +@pytest.fixture +def create_path(nfs_data_dir): + return nfs_data_dir.joinpath("datasets/healpix/merge") + + +@pytest.fixture +def dataset_path(nfs_data_dir): + """Path to the Zarr-backed version of the same test dataset.""" + return nfs_data_dir.joinpath("datasets/healpix/healpix.zarr") + + +@pytest.fixture +def splits(): + """Date ranges that fall within the small test dataset's actual time + range, suitable for exercising real dataloader construction (as opposed + to just checking that arbitrary split values are accepted). + """ + return { + "train_date_start": "1979-01-01", + "train_date_end": "1979-01-01T21:00", + "val_date_start": "1979-01-02", + "val_date_end": "1979-01-02T09:00", + "test_date_start": "1979-01-02T12:00", + "test_date_end": "1979-01-02T18:00", + } + + +@pytest.fixture +def scaling_dict(): + omegaconf = pytest.importorskip("omegaconf") + scaling = { + "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, + "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, + "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, + "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, + "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, + "z1000-12h": {"mean": 952.1435546875, "std": 895.7516479492188}, + "z250": {"mean": 101186.28125, "std": 5551.77978515625}, + "z500": {"mean": 55625.9609375, "std": 2681.712890625}, + "lsm": {"mean": 0, "std": 1}, + "z": {"mean": 0, "std": 1}, + "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, + "extra": {"mean": 0, "std": 0}, # doesn't appear in test dataset + } + return omegaconf.DictConfig(scaling) + + +@pytest.fixture +def scaling_double_dict(): + omegaconf = pytest.importorskip("omegaconf") + scaling = { + "t2m0": {"mean": 0, "std": 2}, + "t850": {"mean": 0, "std": 2}, + "tau300-700": {"mean": 0, "std": 2}, + "tcwv0": {"mean": 0, "std": 2}, + "z1000": {"mean": 0, "std": 2}, + "z1000-12h": {"mean": 0, "std": 2}, + "z250": {"mean": 0, "std": 2}, + "z500": {"mean": 0, "std": 2}, + "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, + "lsm": {"mean": 0, "std": 2}, + "z": {"mean": 0, "std": 2}, + "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset + } + return omegaconf.DictConfig(scaling) + + +@dataclass +class coupler_helper: + """Stand-in for a coupled module, exposing just what setup_coupling() + needs from it (output_variables, time_step).""" + + output_variables: list + time_step: str + + +@pytest.fixture +def constant_coupler_config(): + """A single-variable ConstantCoupler coupling config, in the list-of-dict + shape consumed by Coupled*Dataset(Zarr)/*DataModule(Zarr) `couplings=`.""" + return [ + { + "coupler": "ConstantCoupler", + "params": { + "batch_size": 1, + "variables": ["z250"], + "input_times": ["0h"], + "input_time_dim": 1, + "output_time_dim": 1, + "presteps": 0, + "prepared_coupled_data": True, + }, + } + ] + + +@pytest.fixture +def average_coupler_config(): + """A single-variable TrailingAverageCoupler coupling config, in the same + shape as `constant_coupler_config` but for a different coupler class.""" + return [ + { + "coupler": "TrailingAverageCoupler", + "params": { + "batch_size": 1, + "variables": ["z250"], + "input_times": ["6h"], + "averaging_window": "6h", + "input_time_dim": 1, + "output_time_dim": 1, + "presteps": 0, + "prepared_coupled_data": True, + }, + } + ] + + +def assert_shard_dataloaders(timeseries_dm): + """Shared assertions for `*DataModule(Zarr).{train,val,test}_dataloader`: + a single shard should get no distributed sampler, while multiple shards + should. Used identically across the classic/Zarr, coupled/uncoupled data + module tests. + """ + from torch.utils.data import DataLoader + from torch.utils.data.distributed import DistributedSampler + + train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) + assert train_sampler is None + assert isinstance(train_dataloader, DataLoader) + + val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) + assert val_sampler is None + assert isinstance(val_dataloader, DataLoader) + + test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) + assert test_sampler is None + assert isinstance(test_dataloader, DataLoader) + + train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) + assert isinstance(train_sampler, DistributedSampler) + assert isinstance(train_dataloader, DataLoader) + + val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) + assert isinstance(val_sampler, DistributedSampler) + assert isinstance(val_dataloader, DataLoader) + + test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) + assert isinstance(test_sampler, DistributedSampler) + assert isinstance(test_dataloader, DataLoader) diff --git a/test/datapipes/test_healpix.py b/test/datapipes/healpix/test_healpix.py similarity index 67% rename from test/datapipes/test_healpix.py rename to test/datapipes/healpix/test_healpix.py index 37aafc8fe8..10d8a7f48a 100644 --- a/test/datapipes/test_healpix.py +++ b/test/datapipes/healpix/test_healpix.py @@ -15,109 +15,20 @@ # limitations under the License. import random -import shutil import warnings from pathlib import Path import pytest -from torch.utils.data import DataLoader -from torch.utils.data.distributed import DistributedSampler from physicsnemo.distributed import DistributedManager from test.conftest import requires_module +from test.datapipes.healpix.conftest import assert_shard_dataloaders omegaconf = pytest.importorskip("omegaconf") np = pytest.importorskip("numpy") xr = pytest.importorskip("xarray") -@pytest.fixture -def data_dir(nfs_data_dir): - return nfs_data_dir.joinpath("datasets/healpix") - - -@pytest.fixture -def dataset_name(): - name = "healpix" - return name - - -@pytest.fixture -def create_path(nfs_data_dir): - return nfs_data_dir.joinpath("datasets/healpix/merge") - - -def delete_dataset(create_path, dataset_name): - """Helper that deletes a requested dataset at the specified location""" - dataset_path = f"{create_path}/{dataset_name}.zarr" - if Path(dataset_path).exists(): - shutil.rmtree(dataset_path) - - -@pytest.fixture -def scaling_dict(): - scaling = { - "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, - "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, - "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, - "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, - "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, - "z250": {"mean": 101186.28125, "std": 5551.77978515625}, - "z500": {"mean": 55625.9609375, "std": 2681.712890625}, - "lsm": {"mean": 0, "std": 1}, - "z": {"mean": 0, "std": 1}, - "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, - "extra": {"mean": 1, "std": 0}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - -@pytest.fixture -def scaling_double_dict(): - scaling = { - "t2m0": {"mean": 0, "std": 2}, - "t850": {"mean": 0, "std": 2}, - "tau300-700": {"mean": 0, "std": 2}, - "tcwv0": {"mean": 0, "std": 2}, - "z1000": {"mean": 0, "std": 2}, - "z250": {"mean": 0, "std": 2}, - "z500": {"mean": 0, "std": 2}, - "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, - "lsm": {"mean": 0, "std": 2}, - "z": {"mean": 0, "std": 2}, - "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - -@requires_module("omegaconf") -@requires_module("dask") -@requires_module("netCDF4") -def test_open_time_series_on_the_fly(create_path, pytestconfig): - from physicsnemo.datapipes.healpix.data_modules import ( - open_time_series_dataset_classic_on_the_fly, - ) - - variables = ["z500", "z1000"] - constants = {"lsm": "lsm"} - - ds = open_time_series_dataset_classic_on_the_fly( - directory=create_path, - input_variables=variables, - output_variables=variables, - constants=constants, - ) - assert isinstance(ds, xr.Dataset) - - test_var = variables[0] - base = xr.open_dataset(str(create_path.joinpath(f"{test_var}.nc"))) - ds_var = ds.inputs.sel(channel_in=test_var) - - assert ds_var.equals(base[test_var]) - ds.close() - base.close() - - @requires_module("omegaconf") def test_open_time_series(data_dir, dataset_name, pytestconfig): # check for failure of non-existant dataset @@ -125,7 +36,7 @@ def test_open_time_series(data_dir, dataset_name, pytestconfig): open_time_series_dataset_classic_prebuilt, ) - with pytest.raises(FileNotFoundError, match=("Dataset doesn't appear to exist at")): + with pytest.raises(FileNotFoundError, match=("Dataset doesn't exist at")): open_time_series_dataset_classic_prebuilt("/null_path", dataset_name) ds = open_time_series_dataset_classic_prebuilt(data_dir, dataset_name) @@ -134,72 +45,6 @@ def test_open_time_series(data_dir, dataset_name, pytestconfig): @requires_module("omegaconf") -@requires_module("dask") -@requires_module("netCDF4") -@requires_module("numpy") -def test_create_time_series(data_dir, dataset_name, create_path, pytestconfig): - from physicsnemo.datapipes.healpix.data_modules import ( - create_time_series_dataset_classic, - ) - - variables = ["z500", "z1000"] - constants = {"lsm": "lsm"} - scaling = {"z500": {"log_epsilon": 2}} - # check existing dataset - ds = create_time_series_dataset_classic( - src_directory="/null", - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=["null", "null"], - ) - assert isinstance(ds, xr.Dataset) - ds.close() - - # create new dataset - # open a base dataset to compare against - test_var = list(scaling.keys())[0] - base = xr.open_dataset(str(create_path.joinpath(f"{test_var}.nc"))) - # scale our test variable - base[test_var] = np.log(base[test_var] + scaling[test_var]["log_epsilon"]) - np.log( - scaling[test_var]["log_epsilon"] - ) - ds = create_time_series_dataset_classic( - src_directory=create_path, - dst_directory=create_path, - dataset_name=dataset_name, - input_variables=variables, - scaling=scaling, - ) - ds_var = ds.inputs.sel(channel_in=test_var) - - assert ds_var.equals(base[test_var]) - ds.close() - base.close() - - # delete the created file so we have a clean test for next time - delete_dataset(create_path, dataset_name) - - # and with constants - const = list(constants.keys())[0] - const_ds = xr.open_dataset(str(create_path.joinpath(f"{const}.nc"))) - ds = create_time_series_dataset_classic( - src_directory=create_path, - dst_directory=create_path, - dataset_name=dataset_name, - input_variables=variables, - scaling=scaling, - constants=constants, - ) - assert (const_ds[const] == ds.constants[0]).any() - ds.close() - const_ds.close() - - # delete the created file so we have a clean test for next time - delete_dataset(create_path, dataset_name) - - -@requires_module("omegaconf") -@requires_module("dask") @requires_module("netCDF4") def test_TimeSeriesDataset_initialization( data_dir, dataset_name, scaling_dict, pytestconfig @@ -247,6 +92,35 @@ def test_TimeSeriesDataset_initialization( scaling=invalid_scaling, ) + # check for failure when a channel_out variable has no scaling entry. + # dropping "z1000" from channel_in (but keeping it in channel_out) means + # len(channel_out) != len(channel_in), so input scaling is selected by + # channel_in (which succeeds), letting the target-scaling failure surface + # on its own + scaling_missing_target = omegaconf.DictConfig( + {k: v for k, v in scaling_dict.items() if k != "z1000"} + ) + zarr_ds_asymmetric = zarr_ds.sel( + channel_in=[v for v in zarr_ds.channel_in.values if v != "z1000"], + channel_out=zarr_ds.channel_out.values, + ) + with pytest.raises(KeyError, match=("Target channels ")): + timeseries_ds = TimeSeriesDataset( + dataset=zarr_ds_asymmetric, + scaling=scaling_missing_target, + ) + + # check for failure when a constant (channel_c) variable has no scaling + # entry, even though every channel_in/channel_out variable does + scaling_missing_constant = omegaconf.DictConfig( + {k: v for k, v in scaling_dict.items() if k != "z"} + ) + with pytest.raises(KeyError, match=("Constant channels ")): + timeseries_ds = TimeSeriesDataset( + dataset=zarr_ds, + scaling=scaling_missing_constant, + ) + # check for warning on batch size > 1 and forecast mode warnings.filterwarnings("error") with pytest.raises( @@ -475,37 +349,56 @@ def test_TimeSeriesDataset_get( @requires_module("omegaconf") -@requires_module("dask") @requires_module("netCDF4") def test_TimeSeriesDataModule_initialization( - data_dir, create_path, dataset_name, scaling_double_dict, pytestconfig + data_dir, create_path, dataset_name, scaling_double_dict, splits, pytestconfig ): from physicsnemo.datapipes.healpix.data_modules import ( TimeSeriesDataModule, ) variables = ["z500", "z1000"] - splits = { - "train_date_start": "1959-01-01", - "train_date_end": "1998-12-31T18:00", - "val_date_start": "1999-01-01", - "val_date_end": "2000-12-31T18:00", - "test_date_start": "2017-01-01", - "test_date_end": "2018-12-31T18:00", - } # open our test dataset ds_path = Path(data_dir, dataset_name + ".zarr") zarr_ds = xr.open_zarr(ds_path) - # test with an invalid mode - with pytest.raises(ValueError, match=("'data_format' must be one of")): - timeseries_dm = TimeSeriesDataModule( - src_directory=data_dir, - dst_directory=create_path, + # check for failure when a requested input variable isn't in the dataset + with pytest.raises(ValueError, match=("Input variables not found in dataset")): + TimeSeriesDataModule( + src_directory=create_path, + dst_directory=data_dir, + dataset_name=dataset_name, + input_variables=variables + ["DoesntExist"], + batch_size=1, + prebuilt_dataset=True, + scaling=scaling_double_dict, + ) + + # check for failure when a requested output variable isn't in the dataset + with pytest.raises(ValueError, match=("Output variables not found in dataset")): + TimeSeriesDataModule( + src_directory=create_path, + dst_directory=data_dir, dataset_name=dataset_name, + input_variables=variables, + output_variables=variables + ["DoesntExist"], batch_size=1, - data_format="null", + prebuilt_dataset=True, + scaling=scaling_double_dict, + ) + + # check for failure when a requested constant isn't in the dataset + with pytest.raises(ValueError, match=("Constants not found in dataset")): + TimeSeriesDataModule( + src_directory=create_path, + dst_directory=data_dir, + dataset_name=dataset_name, + input_variables=variables, + batch_size=1, + prebuilt_dataset=True, + scaling=scaling_double_dict, + constants={"lsm": "DoesntExist"}, ) # use the prebuilt dataset @@ -521,10 +414,12 @@ def test_TimeSeriesDataModule_initialization( ) assert isinstance(timeseries_dm, TimeSeriesDataModule) - # without the prebuilt dataset + # `prebuilt_dataset` is kept only for backwards compatibility; on-the-fly + # dataset generation has been removed, so `setup()` always opens an + # existing prebuilt dataset from `dst_directory` regardless of this flag timeseries_dm = TimeSeriesDataModule( src_directory=create_path, - dst_directory=create_path, + dst_directory=data_dir, dataset_name=dataset_name, input_variables=variables, batch_size=1, @@ -643,21 +538,13 @@ def test_TimeSeriesDataModule_get_constants( @requires_module("omegaconf") def test_TimeSeriesDataModule_get_dataloaders( - data_dir, create_path, dataset_name, scaling_double_dict, pytestconfig + data_dir, create_path, dataset_name, scaling_double_dict, splits, pytestconfig ): from physicsnemo.datapipes.healpix.data_modules import ( TimeSeriesDataModule, ) variables = ["z500", "z1000"] - splits = { - "train_date_start": "1979-01-01", - "train_date_end": "1979-01-01T21:00", - "val_date_start": "1979-01-02", - "val_date_end": "1979-01-02T09:00", - "test_date_start": "1979-01-02T12:00", - "test_date_end": "1979-01-02T18:00", - } # use the prebuilt dataset # Internally initializes DistributedManager @@ -673,29 +560,5 @@ def test_TimeSeriesDataModule_get_dataloaders( shuffle=False, ) - # with 1 shard should get no sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) - assert train_sampler is None - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) - assert val_sampler is None - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) - assert test_sampler is None - assert isinstance(test_dataloader, DataLoader) - - # with >1 shard should be distributed sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) - assert isinstance(train_sampler, DistributedSampler) - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) - assert isinstance(val_sampler, DistributedSampler) - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) - assert isinstance(test_sampler, DistributedSampler) - assert isinstance(test_dataloader, DataLoader) + assert_shard_dataloaders(timeseries_dm) DistributedManager.cleanup() diff --git a/test/datapipes/test_healpix_couple.py b/test/datapipes/healpix/test_healpix_couple.py similarity index 69% rename from test/datapipes/test_healpix_couple.py rename to test/datapipes/healpix/test_healpix_couple.py index 8a267aea33..c56990266c 100644 --- a/test/datapipes/test_healpix_couple.py +++ b/test/datapipes/healpix/test_healpix_couple.py @@ -15,18 +15,15 @@ # limitations under the License. import random -import shutil import warnings -from dataclasses import dataclass from pathlib import Path import pytest import torch as th -from torch.utils.data import DataLoader -from torch.utils.data.distributed import DistributedSampler from physicsnemo.distributed import DistributedManager from test.conftest import requires_module +from test.datapipes.healpix.conftest import assert_shard_dataloaders, coupler_helper omegaconf = pytest.importorskip("omegaconf") np = pytest.importorskip("numpy") @@ -34,73 +31,6 @@ xr = pytest.importorskip("xarray") -@pytest.fixture -def data_dir(nfs_data_dir): - return nfs_data_dir.joinpath("datasets/healpix") - - -@pytest.fixture -def dataset_name(): - name = "healpix" - return name - - -@pytest.fixture -def create_path(nfs_data_dir): - return nfs_data_dir.joinpath("datasets/healpix/merge") - - -@dataclass -class coupler_helper: - """helper class for setting up the couplers""" - - output_variables: list - time_step: str - - -def delete_dataset(create_path, dataset_name): - """Helper that deletes a requested dataset at the specified location""" - dataset_path = f"{create_path}/{dataset_name}.zarr" - if Path(dataset_path).exists(): - shutil.rmtree(dataset_path) - - -@pytest.fixture -def scaling_dict(): - scaling = { - "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, - "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, - "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, - "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, - "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, - "z250": {"mean": 101186.28125, "std": 5551.77978515625}, - "z500": {"mean": 55625.9609375, "std": 2681.712890625}, - "lsm": {"mean": 0, "std": 1}, - "z": {"mean": 0, "std": 1}, - "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, - "extra": {"mean": 0, "std": 0}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - -@pytest.fixture -def scaling_double_dict(): - scaling = { - "t2m0": {"mean": 0, "std": 2}, - "t850": {"mean": 0, "std": 2}, - "tau300-700": {"mean": 0, "std": 2}, - "tcwv0": {"mean": 0, "std": 2}, - "z1000": {"mean": 0, "std": 2}, - "z250": {"mean": 0, "std": 2}, - "z500": {"mean": 0, "std": 2}, - "lsm": {"mean": 0, "std": 2}, - "z": {"mean": 0, "std": 2}, - "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, - "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - @requires_module("omegaconf") @requires_module("netCDF4") @requires_module("pandas") @@ -153,8 +83,8 @@ def test_ConstantCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): output_variables=["not_coupled", "z500"], time_step="0h", ) - coupler.setup_coupling(mock_coupled_module) - assert coupler.coupled_channel_indices == [1] + with pytest.raises(ValueError, match=("Missing variables in coupled module")): + coupler.setup_coupling(mock_coupled_module) mock_coupled_module.output_variables = ["z500", "z1000"] coupler.setup_coupling(mock_coupled_module) @@ -183,7 +113,11 @@ def test_ConstantCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) assert np.array_equal(expected, coupler.coupled_scaling["std"]) - # test incorrect batch size + # `set_coupled_fields` derives its batch dimension from the shape of + # `coupled_fields` itself rather than validating it against the + # configured `batch_size` (an older physicsnemo implementation raised a + # ValueError here, but the ported DLESyM coupler intentionally adapts to + # whatever batch size is provided) coupler.coupled_channel_indices = [0, 1] coupled_fields_batch_size = batch_size * 2 coupled_fields_timedim = 2 @@ -195,8 +129,8 @@ def test_ConstantCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): coupler.spatial_dims[1], coupler.spatial_dims[2], ) - with pytest.raises(ValueError, match=("Batch size of coupled field 4 ")): - coupler.set_coupled_fields(coupled_fields) + coupler.set_coupled_fields(coupled_fields) + assert coupler.preset_coupled_fields.shape[1] == coupled_fields_batch_size coupled_fields_batch_size = batch_size coupled_fields_timedim = 4 @@ -217,21 +151,25 @@ def test_ConstantCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): assert coupler.coupled_mode assert list(coupler.construct_integrated_couplings().shape) == expected_shape - # verify that the data is being properly transformed - expected = coupled_fields[:, :, :, coupler.coupled_channel_indices, :, :].permute( - 2, 0, 3, 1, 4, 5 - ) - expected = expected[0, :, -1, :, :, :] - expected = expected.unsqueeze(0).unsqueeze(0) - expected = expected.repeat( - coupler.coupled_integration_dim, coupled_fields_batch_size, 1, 1, 1, 1 - ) + # verify that the data is being properly transformed: set_coupled_fields + # broadcasts the first time step of coupled_fields across the integration + # dim and permutes to [integration_dim, batch, timevar, face, height, width] + expected = coupled_fields[:, :, :, coupler.coupled_channel_indices, :, :] + expected = expected[:, :, :1, :, :, :] + expected = expected.permute(2, 0, 3, 1, 4, 5) assert th.equal(expected, coupler.construct_integrated_couplings()) # test coupler reset coupler.reset_coupler() assert coupler.coupled_mode is False + # once reset, batch/bsize are required to construct the couplings from + # the dataset rather than from preset (cached) fields + with pytest.raises( + ValueError, match=("batch and bsize must be provided when not in coupled_mode") + ): + coupler.construct_integrated_couplings() + # test loading from the dataset coupled_scaling = { "mean": np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)), @@ -244,6 +182,14 @@ def test_ConstantCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): ) assert np.array_equal(expected, coupled_field[0]) + # set_scaling raises if a coupled variable has no corresponding entry in + # the provided scaling values + scaling_missing_var = scaling_da.drop_sel(index="z1000") + with pytest.raises( + KeyError, match=("Coupled variable\\(s\\) not found in scaling values") + ): + coupler.set_scaling(scaling_missing_var) + zarr_ds.close() DistributedManager.cleanup() @@ -284,19 +230,6 @@ def test_TrailingAverageCoupler(data_dir, dataset_name, scaling_dict, pytestconf prepared_coupled_data=False, ) - # test fail when input times aren't evenly divisible by dataset dt - with pytest.raises(ValueError, match=("Coupled input times")): - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - averaging_window=averaging_window, - input_times=["30m"], - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - coupler = TrailingAverageCoupler( dataset=zarr_ds, batch_size=batch_size, @@ -309,11 +242,18 @@ def test_TrailingAverageCoupler(data_dir, dataset_name, scaling_dict, pytestconf ) assert isinstance(coupler, TrailingAverageCoupler) - # veryify averaging slices computed correctly mock_coupled_module = coupler_helper( output_variables=["not_coupled", "z500"], time_step="3h", ) + with pytest.raises(ValueError, match=("Missing variables in coupled module")): + coupler.setup_coupling(mock_coupled_module) + + # veryify averaging slices computed correctly + mock_coupled_module = coupler_helper( + output_variables=["z500", "z1000"], + time_step="3h", + ) coupler.setup_coupling(mock_coupled_module) averaging_window_max_indices = [ i // pd.Timedelta(mock_coupled_module.time_step) for i in input_times @@ -372,7 +312,11 @@ def test_TrailingAverageCoupler(data_dir, dataset_name, scaling_dict, pytestconf coupler.averaging_slices = averaging_slices coupler.coupled_channel_indices = [0, 1] - # test a mismatched batch size + # `set_coupled_fields` derives its batch dimension from the shape of + # `coupled_fields` itself rather than validating it against the + # configured `batch_size` (an older physicsnemo implementation raised a + # ValueError here, but the ported DLESyM coupler intentionally adapts to + # whatever batch size is provided) coupled_fields_batch_size = batch_size * 2 coupled_fields_timedim = 4 coupled_fields = th.rand( @@ -383,8 +327,8 @@ def test_TrailingAverageCoupler(data_dir, dataset_name, scaling_dict, pytestconf coupler.spatial_dims[1], coupler.spatial_dims[2], ) - with pytest.raises(ValueError, match=("Batch size of coupled field 4 ")): - coupler.set_coupled_fields(coupled_fields) + coupler.set_coupled_fields(coupled_fields) + assert coupler.preset_coupled_fields.shape[1] == coupled_fields_batch_size coupled_fields_batch_size = batch_size coupled_fields_timedim = 4 @@ -413,6 +357,143 @@ def test_TrailingAverageCoupler(data_dir, dataset_name, scaling_dict, pytestconf DistributedManager.cleanup() +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("pandas") +@requires_module("xarray") +def test_TrailingAverageCoupler_multiple_variables(data_dir, dataset_name): + """Regression test to ensure `set_coupled_fields` correctly averages and + threads through every coupled variable, not just the first. Gives each + coupled variable a distinct, constant value in the synthetic + `coupled_fields` tensor so that averaging over any time window still + yields that same constant, letting each variable's contribution to + `preset_coupled_fields` be checked in isolation and by position. + """ + from physicsnemo.datapipes.healpix.couplers import TrailingAverageCoupler + + variables = ["z500", "z1000", "z250"] + input_times = ["6h", "12h"] + input_time_dim = 2 + output_time_dim = 2 + batch_size = 2 + averaging_window = "6h" + ds_path = Path(data_dir, dataset_name + ".zarr") + zarr_ds = xr.open_zarr(ds_path) + + coupler = TrailingAverageCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + presteps=0, + averaging_window=averaging_window, + input_times=input_times, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ) + coupler.coupled_channel_indices = list(range(len(variables))) + + data_time_step = "3h" + averaging_window_max_indices = [ + i // pd.Timedelta(data_time_step) for i in input_times + ] + di = averaging_window_max_indices[0] + averaging_slices = [] + for j in range(coupler.coupled_integration_dim): + averaging_slices.append([]) + for i, r in enumerate(averaging_window_max_indices): + averaging_slices[j].append( + slice( + coupler.input_time_dim * j * di + i * di, + coupler.input_time_dim * j * di + r, + ) + ) + coupler.averaging_slices = averaging_slices + + # give each coupled variable a distinct, constant value across every + # other dimension so the averaging window doesn't change its value, and + # each variable's contribution can be checked independently of the + # others + channel_values = [100.0, 200.0, 300.0] + coupled_fields = th.empty( + batch_size, + coupler.spatial_dims[0], + 4, + len(variables), + coupler.spatial_dims[1], + coupler.spatial_dims[2], + ) + for i, value in enumerate(channel_values): + coupled_fields[:, :, :, i, :, :] = value + + coupler.set_coupled_fields(coupled_fields) + result = coupler.construct_integrated_couplings() + assert list(result.shape) == [ + coupler.coupled_integration_dim, + batch_size, + coupler.timevar_dim, + ] + list(coupler.spatial_dims) + + # timevar ordering groups by averaging period first, then by variable + # (matching the dataset-loading path in + # `_construct_integrated_couplings_from_dataset`/`construct_integrated_couplings`) + for period in range(len(input_times)): + for var_idx, value in enumerate(channel_values): + timevar_idx = period * len(variables) + var_idx + slice_result = result[:, :, timevar_idx, :, :, :] + assert th.allclose(slice_result, th.full_like(slice_result, value)), ( + f"coupled variable {variables[var_idx]!r} (value {value}) was not " + f"preserved at averaging period {period}, timevar index {timevar_idx}" + ) + + zarr_ds.close() + DistributedManager.cleanup() + + +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("pandas") +@requires_module("xarray") +def test_ConstantCoupler_multiple_variables_order(data_dir, dataset_name): + """Sanity check that the xarray-backed coupling path preserves the order + of `variables` as requested (rather than the dataset's native channel_in + order) when 3+ coupled variables are requested out of native order. This + mirrors the invariant that the Zarr-backed path also has to uphold (see + the analogous, previously-failing test in test_healpix_couple_zarr.py). + """ + from physicsnemo.datapipes.healpix.couplers import ConstantCoupler + + ds_path = Path(data_dir, dataset_name + ".zarr") + zarr_ds = xr.open_zarr(ds_path) + + native_order = list(zarr_ds.channel_in.values) + variables = ["z250", "z1000", "z500"] + assert [native_order.index(v) for v in variables] == sorted( + [native_order.index(v) for v in variables], reverse=True + ) + + batch_size = 2 + batch = {"time": slice(0, 2)} + coupler = ConstantCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + input_times=["0h"], + input_time_dim=1, + output_time_dim=1, + ) + coupler.compute_coupled_indices(interval=2, data_time_step="3h") + + coupled_field = coupler.construct_integrated_couplings( + batch=batch, bsize=batch_size + ) + for i, var in enumerate(variables): + expected_var = zarr_ds.sel(channel_in=var).inputs[:2].values + assert np.array_equal(expected_var, coupled_field[0][:, i]) + + zarr_ds.close() + DistributedManager.cleanup() + + @requires_module("omegaconf") @requires_module("netCDF4") @requires_module("xarray") @@ -527,7 +608,7 @@ def test_CoupledTimeSeriesDataset_initialization( @requires_module("netCDF4") @requires_module("xarray") def test_CoupledTimeSeriesDataset_get_constants( - data_dir, dataset_name, scaling_dict, pytestconfig + data_dir, dataset_name, scaling_dict, constant_coupler_config, pytestconfig ): from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( CoupledTimeSeriesDataset, @@ -539,25 +620,11 @@ def test_CoupledTimeSeriesDataset_get_constants( variables = ["z500", "z1000"] - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] timeseries_ds = CoupledTimeSeriesDataset( dataset=zarr_ds, input_variables=variables, scaling=scaling_dict, - couplings=constant_coupler, + couplings=constant_coupler_config, ) # constants are reshaped @@ -576,7 +643,7 @@ def test_CoupledTimeSeriesDataset_get_constants( @requires_module("netCDF4") @requires_module("xarray") def test_CoupledTimeSeriesDataset_len( - data_dir, dataset_name, scaling_dict, pytestconfig + data_dir, dataset_name, scaling_dict, constant_coupler_config, pytestconfig ): from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( CoupledTimeSeriesDataset, @@ -588,20 +655,6 @@ def test_CoupledTimeSeriesDataset_len( variables = ["z500", "z1000"] - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] # check forecast mode init_times = random.randint(1, len(zarr_ds.time.values)) timeseries_ds = CoupledTimeSeriesDataset( @@ -610,24 +663,12 @@ def test_CoupledTimeSeriesDataset_len( scaling=scaling_dict, batch_size=1, forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler, + couplings=constant_coupler_config, ) assert len(timeseries_ds) == init_times - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 2, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] + constant_coupler = constant_coupler_config.copy() + constant_coupler[0]["params"]["batch_size"] = 2 # check train mode timeseries_ds = CoupledTimeSeriesDataset( @@ -663,7 +704,7 @@ def test_CoupledTimeSeriesDataset_len( @requires_module("netCDF4") @requires_module("xarray") def test_CoupledTimeSeriesDataset_get( - data_dir, dataset_name, scaling_double_dict, pytestconfig + data_dir, dataset_name, scaling_double_dict, constant_coupler_config, pytestconfig ): from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( CoupledTimeSeriesDataset, @@ -674,25 +715,19 @@ def test_CoupledTimeSeriesDataset_get( zarr_ds = xr.open_zarr(ds_path) variables = list(zarr_ds.channel_out.to_numpy()) + # "z250" is coupled in below rather than being a regular input variable: + # channel_in scaling/slicing assumes channel_in == input_variables + + # coupled variables (with the coupled ones appended last), so a coupled + # variable can't also be one of input_variables + coupled_input_variables = [v for v in variables if v != "z250"] batch_size = 2 - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": batch_size, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] + constant_coupler = constant_coupler_config.copy() + constant_coupler[0]["params"]["batch_size"] = batch_size timeseries_ds = CoupledTimeSeriesDataset( dataset=zarr_ds, - input_variables=variables, + input_variables=coupled_input_variables, + output_variables=variables, scaling=scaling_double_dict, batch_size=batch_size, couplings=constant_coupler, @@ -731,7 +766,8 @@ def test_CoupledTimeSeriesDataset_get( # this time dropping incomplete so that we get a full sample sample timeseries_ds = CoupledTimeSeriesDataset( dataset=zarr_ds, - input_variables=variables, + input_variables=coupled_input_variables, + output_variables=variables, scaling=scaling_double_dict, batch_size=batch_size, drop_last=True, @@ -778,10 +814,43 @@ def test_CoupledTimeSeriesDataset_get( assert non_perturbed_inputs[0][0][0].shape == perturbed_inputs[0][0][0].shape assert not np.array_equal(non_perturbed_inputs[0][0][0], perturbed_inputs[0][0][0]) + # noise is also added to the integrated couplings themselves, not just + # the regular inputs, when couplings are present. `inputs_result` here is + # [inputs, constants, integrated_couplings] because the dataset carries a + # "constants" data variable which the base class picks up unconditionally + timeseries_ds = CoupledTimeSeriesDataset( + dataset=zarr_ds, + input_variables=coupled_input_variables, + output_variables=variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + couplings=constant_coupler, + ) + non_perturbed_sample = timeseries_ds[0][0] + assert len(non_perturbed_sample) == 3 + non_perturbed_couplings = non_perturbed_sample[2] + + timeseries_ds = CoupledTimeSeriesDataset( + dataset=zarr_ds, + input_variables=coupled_input_variables, + output_variables=variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + add_train_noise=True, + train_noise_params=noise_params, + couplings=constant_coupler, + ) + perturbed_couplings = timeseries_ds[0][0][2] + assert non_perturbed_couplings.shape == perturbed_couplings.shape + assert not np.array_equal(non_perturbed_couplings, perturbed_couplings) + # With insolation we get 1 extra channel timeseries_ds = CoupledTimeSeriesDataset( dataset=zarr_ds, - input_variables=variables, + input_variables=coupled_input_variables, + output_variables=variables, scaling=scaling_double_dict, batch_size=batch_size, drop_last=True, @@ -794,7 +863,8 @@ def test_CoupledTimeSeriesDataset_get( init_times = random.randint(1, len(zarr_ds.time.values)) timeseries_ds = CoupledTimeSeriesDataset( dataset=zarr_ds, - input_variables=variables, + input_variables=coupled_input_variables, + output_variables=variables, scaling=scaling_double_dict, batch_size=1, forecast_init_times=zarr_ds.time[:init_times], @@ -808,7 +878,8 @@ def test_CoupledTimeSeriesDataset_get( init_times = random.randint(1, len(zarr_ds.time.values)) timeseries_ds = CoupledTimeSeriesDataset( dataset=zarr_ds, - input_variables=variables, + input_variables=coupled_input_variables, + output_variables=variables, scaling=scaling_double_dict, batch_size=1, add_insolation=True, @@ -822,7 +893,8 @@ def test_CoupledTimeSeriesDataset_get( zarr_ds_no_const = zarr_ds.drop_vars("constants") timeseries_ds = CoupledTimeSeriesDataset( dataset=zarr_ds_no_const, - input_variables=variables, + input_variables=coupled_input_variables, + output_variables=variables, scaling=scaling_double_dict, batch_size=1, forecast_init_times=zarr_ds.time[:init_times], @@ -835,53 +907,66 @@ def test_CoupledTimeSeriesDataset_get( @requires_module("omegaconf") -@requires_module("dask") @requires_module("netCDF4") @requires_module("xarray") def test_CoupledTimeSeriesDataModule_initialization( - data_dir, create_path, dataset_name, scaling_double_dict, pytestconfig + data_dir, + create_path, + dataset_name, + scaling_double_dict, + splits, + constant_coupler_config, + pytestconfig, ): from physicsnemo.datapipes.healpix.data_modules import ( CoupledTimeSeriesDataModule, ) variables = ["z500", "z1000"] - splits = { - "train_date_start": "1959-01-01", - "train_date_end": "1998-12-31T18:00", - "val_date_start": "1999-01-01", - "val_date_end": "2000-12-31T18:00", - "test_date_start": "2017-01-01", - "test_date_end": "2018-12-31T18:00", - } - - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] + constant_coupler = constant_coupler_config # open our test dataset ds_path = Path(data_dir, dataset_name + ".zarr") zarr_ds = xr.open_zarr(ds_path) - # test with an invalid mode - with pytest.raises(ValueError, match=("'data_format' must be one of")): - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=data_dir, - dst_directory=create_path, + # check for failure when a requested input variable isn't in the dataset + with pytest.raises(ValueError, match=("Input variables not found in dataset")): + CoupledTimeSeriesDataModule( + src_directory=create_path, + dst_directory=data_dir, dataset_name=dataset_name, + input_variables=variables + ["DoesntExist"], + batch_size=1, + prebuilt_dataset=True, + scaling=scaling_double_dict, + couplings=constant_coupler, + ) + + # check for failure when a requested output variable isn't in the dataset + with pytest.raises(ValueError, match=("Output variables not found in dataset")): + CoupledTimeSeriesDataModule( + src_directory=create_path, + dst_directory=data_dir, + dataset_name=dataset_name, + input_variables=variables, + output_variables=variables + ["DoesntExist"], batch_size=1, - data_format="null", + prebuilt_dataset=True, + scaling=scaling_double_dict, + couplings=constant_coupler, + ) + + # check for failure when a requested constant isn't in the dataset + with pytest.raises(ValueError, match=("Constants not found in dataset")): + CoupledTimeSeriesDataModule( + src_directory=create_path, + dst_directory=data_dir, + dataset_name=dataset_name, + input_variables=variables, + batch_size=1, + prebuilt_dataset=True, + scaling=scaling_double_dict, + constants={"lsm": "DoesntExist"}, couplings=constant_coupler, ) @@ -899,10 +984,12 @@ def test_CoupledTimeSeriesDataModule_initialization( ) assert isinstance(timeseries_dm, CoupledTimeSeriesDataModule) - # without the prebuilt dataset + # `prebuilt_dataset` is kept only for backwards compatibility; on-the-fly + # dataset generation has been removed, so `setup()` always opens an + # existing prebuilt dataset from `dst_directory` regardless of this flag timeseries_dm = CoupledTimeSeriesDataModule( src_directory=create_path, - dst_directory=create_path, + dst_directory=data_dir, dataset_name=dataset_name, input_variables=variables, batch_size=1, @@ -947,7 +1034,12 @@ def test_CoupledTimeSeriesDataModule_initialization( @requires_module("netCDF4") @requires_module("xarray") def test_CoupledTimeSeriesDataModule_get_constants( - data_dir, create_path, dataset_name, scaling_double_dict, pytestconfig + data_dir, + create_path, + dataset_name, + scaling_double_dict, + constant_coupler_config, + pytestconfig, ): from physicsnemo.datapipes.healpix.data_modules import ( CoupledTimeSeriesDataModule, @@ -955,21 +1047,7 @@ def test_CoupledTimeSeriesDataModule_get_constants( variables = ["z500", "z1000"] constants = {"lsm": "lsm"} - - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] + constant_coupler = constant_coupler_config # No constants # Internally initializes DistributedManager @@ -1042,36 +1120,19 @@ def test_CoupledTimeSeriesDataModule_get_constants( @requires_module("omegaconf") def test_CoupledTimeSeriesDataModule_get_dataloaders( - data_dir, create_path, dataset_name, scaling_double_dict, pytestconfig + data_dir, + create_path, + dataset_name, + scaling_double_dict, + splits, + constant_coupler_config, + pytestconfig, ): from physicsnemo.datapipes.healpix.data_modules import ( CoupledTimeSeriesDataModule, ) variables = ["z500", "z1000"] - splits = { - "train_date_start": "1979-01-01", - "train_date_end": "1979-01-01T21:00", - "val_date_start": "1979-01-02", - "val_date_end": "1979-01-02T09:00", - "test_date_start": "1979-01-02T12:00", - "test_date_end": "1979-01-02T18:00", - } - - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] # use the prebuilt dataset # Internally initializes DistributedManager @@ -1085,60 +1146,28 @@ def test_CoupledTimeSeriesDataModule_get_dataloaders( scaling=scaling_double_dict, splits=splits, shuffle=False, - couplings=constant_coupler, + couplings=constant_coupler_config, ) - # with 1 shard should get no sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) - assert train_sampler is None - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) - assert val_sampler is None - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) - assert test_sampler is None - assert isinstance(test_dataloader, DataLoader) - - # with >1 shard should be distributed sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) - assert isinstance(train_sampler, DistributedSampler) - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) - assert isinstance(val_sampler, DistributedSampler) - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) - assert isinstance(test_sampler, DistributedSampler) - assert isinstance(test_dataloader, DataLoader) + assert_shard_dataloaders(timeseries_dm) DistributedManager.cleanup() @requires_module("omegaconf") def test_CoupledTimeSeriesDataModule_get_coupled_vars( - data_dir, create_path, dataset_name, scaling_double_dict, pytestconfig + data_dir, + create_path, + dataset_name, + scaling_double_dict, + constant_coupler_config, + average_coupler_config, + pytestconfig, ): from physicsnemo.datapipes.healpix.data_modules import ( CoupledTimeSeriesDataModule, ) variables = ["z500", "z1000"] - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] # Constant coupler # Internally initializes DistributedManager @@ -1150,7 +1179,7 @@ def test_CoupledTimeSeriesDataModule_get_coupled_vars( batch_size=1, prebuilt_dataset=True, scaling=scaling_double_dict, - couplings=constant_coupler, + couplings=constant_coupler_config, ) outvar = timeseries_dm._get_coupled_vars() @@ -1160,21 +1189,6 @@ def test_CoupledTimeSeriesDataModule_get_coupled_vars( assert expected == outvar - average_coupler = [ - { - "coupler": "TrailingAverageCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["6h"], - "averaging_window": "6h", - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] # Average coupler # Internally initializes DistributedManager timeseries_dm = CoupledTimeSeriesDataModule( @@ -1185,7 +1199,7 @@ def test_CoupledTimeSeriesDataModule_get_coupled_vars( batch_size=1, prebuilt_dataset=True, scaling=scaling_double_dict, - couplings=average_coupler, + couplings=average_coupler_config, ) outvar = timeseries_dm._get_coupled_vars() outvar.sort() @@ -1199,7 +1213,7 @@ def test_CoupledTimeSeriesDataModule_get_coupled_vars( @requires_module("netCDF4") @requires_module("xarray") def test_CoupledTimeSeriesDataset_next_integration( - data_dir, dataset_name, scaling_dict, pytestconfig + data_dir, dataset_name, scaling_dict, constant_coupler_config, pytestconfig ): from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( CoupledTimeSeriesDataset, @@ -1207,27 +1221,15 @@ def test_CoupledTimeSeriesDataset_next_integration( spatial_dims = [12, 32, 32] input_variables = ["z500", "z1000"] - coupled_channel_indices = [0, 1] + # length must match len(coupled_variables) (i.e. the coupler's timevar_dim); + # values are arbitrary since this synthetic setup bypasses setup_coupling() + coupled_channel_indices = [0] coupled_variables = ["z250"] num_variables = len(input_variables) input_time_dim = 1 output_time_dim = 1 batch_size = 1 - - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": coupled_variables, - "input_times": ["0h"], - "input_time_dim": input_time_dim, - "output_time_dim": output_time_dim, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] + constant_coupler = constant_coupler_config # open our test dataset ds_path = Path(data_dir, dataset_name + ".zarr") @@ -1270,12 +1272,11 @@ def test_CoupledTimeSeriesDataset_next_integration( spatial_dims[2], ) - expected_coupling = coupled_fields[:, :, :, coupled_channel_indices, :, :].permute( - 2, 0, 3, 1, 4, 5 - ) - expected_coupling = expected_coupling[0, :, -1, :, :, :] - expected_coupling = expected_coupling.unsqueeze(0).unsqueeze(0) - expected_coupling = expected_coupling.repeat(1, batch_size, 1, 1, 1, 1) + # mirrors set_coupled_fields: broadcast the first time step across the + # integration dim and permute to [integration_dim, batch, timevar, face, height, width] + expected_coupling = coupled_fields[:, :, :, coupled_channel_indices, :, :] + expected_coupling = expected_coupling[:, :, :1, :, :, :] + expected_coupling = expected_coupling.permute(2, 0, 3, 1, 4, 5) # need to grab at least 1 sample to properly intialize everything timeseries_ds[0] @@ -1292,8 +1293,15 @@ def test_CoupledTimeSeriesDataset_next_integration( assert np.array_equal(test_integration[3], expected_coupling) # I have absolutely no idea why a coupled dataset has the option for 0 couplings + # with no coupler, channel_in must exactly match input_variables (no extra + # coupled-variable channel appended), so re-select the dataset without + # the "z250" channel that `test_ds` carries for the coupled case above + test_ds_no_couplings = ds.sel( + channel_in=input_variables, + channel_out=input_variables, + ) timeseries_ds = CoupledTimeSeriesDataset( - dataset=test_ds, + dataset=test_ds_no_couplings, input_variables=input_variables, scaling=scaling_dict, batch_size=batch_size, @@ -1302,7 +1310,7 @@ def test_CoupledTimeSeriesDataset_next_integration( time_step="6h", drop_last=True, add_insolation=True, - forecast_init_times=test_ds.time[:init_times], + forecast_init_times=test_ds_no_couplings.time[:init_times], ) # need to grab at least 1 sample to properly intialize everything timeseries_ds[0] diff --git a/test/datapipes/test_healpix_couple_zarr.py b/test/datapipes/healpix/test_healpix_couple_zarr.py similarity index 77% rename from test/datapipes/test_healpix_couple_zarr.py rename to test/datapipes/healpix/test_healpix_couple_zarr.py index 9767867644..d596e7cd40 100644 --- a/test/datapipes/test_healpix_couple_zarr.py +++ b/test/datapipes/healpix/test_healpix_couple_zarr.py @@ -16,15 +16,13 @@ import random import warnings -from dataclasses import dataclass import pytest import torch as th -from torch.utils.data import DataLoader -from torch.utils.data.distributed import DistributedSampler from physicsnemo.distributed import DistributedManager from test.conftest import requires_module +from test.datapipes.healpix.conftest import assert_shard_dataloaders, coupler_helper omegaconf = pytest.importorskip("omegaconf") np = pytest.importorskip("numpy") @@ -33,106 +31,25 @@ zarr = pytest.importorskip("zarr") -@pytest.fixture -def dataset_path(nfs_data_dir): - return nfs_data_dir.joinpath("datasets/healpix/healpix.zarr") +def test_base_coupler_invalid_dataset_type(): + """Couplers dispatch on the runtime type of `dataset` to decide whether to use + the Zarr-native code path (`zarr.Group`) or the xarray-based one (`xr.Dataset`). + This dispatch logic is new (added to support the Zarr datapipes) and doesn't + require the NFS test dataset, so it can run unconditionally. + """ + from physicsnemo.datapipes.healpix.couplers import ConstantCoupler + fake_dataset = {"inputs": np.zeros((4, 12, 1, 4, 4))} -@pytest.fixture -def splits(): - split_dict = { - "train_date_start": "1979-01-01", - "train_date_end": "1979-01-01T21:00", - "val_date_start": "1979-01-02", - "val_date_end": "1979-01-02T09:00", - "test_date_start": "1979-01-02T12:00", - "test_date_end": "1979-01-02T18:00", - } - return split_dict - - -@pytest.fixture -def constant_coupler_config(): - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] - return constant_coupler - - -@pytest.fixture -def average_coupler_config(): - average_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] - return average_coupler - - -@dataclass -class coupler_helper: - """helper class for setting up the couplers""" - - output_variables: list - time_step: str - - -@pytest.fixture -def scaling_dict(): - scaling = { - "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, - "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, - "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, - "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, - "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, - "z1000-12h": {"mean": 952.1435546875, "std": 895.7516479492188}, - "z250": {"mean": 101186.28125, "std": 5551.77978515625}, - "z500": {"mean": 55625.9609375, "std": 2681.712890625}, - "lsm": {"mean": 0, "std": 1}, - "z": {"mean": 0, "std": 1}, - "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, - "extra": {"mean": 0, "std": 0}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - -@pytest.fixture -def scaling_double_dict(): - scaling = { - "t2m0": {"mean": 0, "std": 2}, - "t850": {"mean": 0, "std": 2}, - "tau300-700": {"mean": 0, "std": 2}, - "tcwv0": {"mean": 0, "std": 2}, - "z1000": {"mean": 0, "std": 2}, - "z1000-12h": {"mean": 0, "std": 2}, - "z250": {"mean": 0, "std": 2}, - "z500": {"mean": 0, "std": 2}, - "lsm": {"mean": 0, "std": 2}, - "z": {"mean": 0, "std": 2}, - "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, - "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) + with pytest.raises( + TypeError, + match=("Coupler only supports xarray Datasets or zarr Groups"), + ): + ConstantCoupler( + dataset=fake_dataset, + batch_size=1, + variables=["z500"], + ) @requires_module("omegaconf") @@ -155,7 +72,7 @@ def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): # open our test dataset zarr_ds = zarr.open(dataset_path) input_indices = [ - int(np.where(zarr_ds.channel_in[:] == ch)[0][0]) for ch in variables + int(np.where(zarr_ds["channel_in"][:] == ch)[0][0]) for ch in variables ] # test fail initialization @@ -238,26 +155,14 @@ def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): assert coupler.coupled_mode assert list(coupler.construct_integrated_couplings().shape) == expected_shape - # verify that the data is being properly transformed - expected = coupled_fields[:, :, :, coupler.coupled_channel_indices, :, :].permute( - 2, 0, 3, 1, 4, 5 - ) - expected = expected[0, :, -1:, :, :, :] - expected = expected.unsqueeze(0) - expected = expected.repeat( - coupler.coupled_integration_dim, 1, coupled_fields_batch_size, 1, 1, 1 - ) - result = coupler.construct_integrated_couplings() - assert th.equal(expected, coupler.construct_integrated_couplings()) - - # verify that dimensions aren't reordered when time_first is false - coupler.time_first = False - coupler.set_coupled_fields(coupled_fields) - # [T, B, C, F, H, W] - expected = expected.permute(1, 3, 0, 2, 4, 5) + # verify that the data is being properly transformed: set_coupled_fields + # broadcasts the first time step of coupled_fields across the integration + # dim and permutes to [integration_dim, batch, timevar, face, height, width] + expected = coupled_fields[:, :, :, coupler.coupled_channel_indices, :, :] + expected = expected[:, :, :1, :, :, :] + expected = expected.permute(2, 0, 3, 1, 4, 5) result = coupler.construct_integrated_couplings() assert th.equal(expected, result) - coupler.time_first = True # test coupler reset coupler.reset_coupler() @@ -268,7 +173,7 @@ def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): "mean": np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)), "std": np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)), } - expected = zarr_ds.inputs[:2][:, input_indices] + expected = zarr_ds["inputs"][:2][:, input_indices] expected = (expected - coupled_scaling["mean"]) / coupled_scaling["std"] coupled_field = coupler.construct_integrated_couplings( batch=batch, bsize=batch_size @@ -278,6 +183,88 @@ def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): DistributedManager.cleanup() +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("pandas") +@requires_module("xarray") +def test_ConstantCoupler_multiple_variables_order(dataset_path): + """Regression test for a bug where the Zarr-backed coupling path selected + coupled channels in the dataset's native channel_in order rather than the + order given by `variables`, while the xarray-backed path (`.sel`) always + honored the requested order. With 2+ coupled variables requested out of + their native dataset order, this silently misaligned `coupled_scaling` + (which is built from `variables` order) against the loaded data, applying + the wrong mean/std to the wrong channel. Only surfaces with 2+ variables + that aren't already in native dataset order, which none of the other + tests happen to exercise. + """ + from physicsnemo.datapipes.healpix.couplers import ConstantCoupler + + xr_ds = xr.open_zarr(dataset_path) + zarr_ds = zarr.open(dataset_path) + + # deliberately request the coupled variables in the opposite order from + # how they appear natively in the dataset's channel_in coordinate + # (native order is z500, ..., z1000, ..., z250) + native_order = list(xr_ds.channel_in.values) + variables = ["z250", "z1000", "z500"] + assert [native_order.index(v) for v in variables] == sorted( + [native_order.index(v) for v in variables], reverse=True + ) + + batch_size = 2 + batch = {"time": slice(0, 2)} + + coupler_xr = ConstantCoupler( + dataset=xr_ds, + batch_size=batch_size, + variables=variables, + input_times=["0h"], + input_time_dim=1, + output_time_dim=1, + ) + coupler_zarr = ConstantCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + input_times=["0h"], + input_time_dim=1, + output_time_dim=1, + ) + + interval = 2 + data_time_step = "3h" + coupler_xr.compute_coupled_indices(interval, data_time_step) + coupler_zarr.compute_coupled_indices(interval, data_time_step) + + coupled_xr = coupler_xr.construct_integrated_couplings( + batch=batch, bsize=batch_size + ) + coupled_zarr = coupler_zarr.construct_integrated_couplings( + batch=batch, bsize=batch_size + ) + + # the Zarr-backed and xarray-backed couplers must agree exactly, and both + # must match the raw data indexed in the order `variables` was given, not + # the dataset's native channel_in order + input_indices = [ + int(np.where(zarr_ds["channel_in"][:] == v)[0][0]) for v in variables + ] + expected = zarr_ds["inputs"][:2][:, input_indices] + # coupled_integration_dim is 1 here, so index into the leading dim to get + # the [batch, channel, face, height, width] slice that lines up with `expected` + assert np.array_equal(expected, coupled_xr[0]) + assert np.array_equal(expected, coupled_zarr[0]) + + # and, per-variable, each channel's data must match that specific + # variable's raw data, not another (coupled) variable's + for i, var in enumerate(variables): + var_index = int(np.where(zarr_ds["channel_in"][:] == var)[0][0]) + expected_var = zarr_ds["inputs"][:2][:, var_index] + assert np.array_equal(expected_var, coupled_zarr[0][:, i]) + assert np.array_equal(expected_var, coupled_xr[0][:, i]) + + @requires_module("omegaconf") @requires_module("netCDF4") @requires_module("pandas") @@ -313,19 +300,6 @@ def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): prepared_coupled_data=False, ) - # test fail when input times aren't evenly divisible by dataset dt - with pytest.raises(ValueError, match=("Coupled input times")): - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - averaging_window=averaging_window, - input_times=["30m"], - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - coupler = TrailingAverageCoupler( dataset=zarr_ds, batch_size=batch_size, @@ -338,6 +312,12 @@ def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): ) assert isinstance(coupler, TrailingAverageCoupler) + # Zarr-backed datasets store time as CF-encoded integers rather than + # numpy datetime64; verify the coupler decodes them via cftime into the + # same wall-clock times as the xarray view of the same store + expected_time_da = xr.open_zarr(dataset_path).time.values + assert np.array_equal(coupler.time_da, expected_time_da) + mock_coupled_module = coupler_helper( output_variables=["not_coupled", "z500"], time_step="3h", @@ -434,6 +414,96 @@ def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): DistributedManager.cleanup() +@requires_module("omegaconf") +@requires_module("netCDF4") +@requires_module("pandas") +@requires_module("xarray") +def test_TrailingAverageCoupler_multiple_variables(dataset_path): + """Regression test to ensure `set_coupled_fields` correctly averages and + threads through every coupled variable, not just the first. Gives each + coupled variable a distinct, constant value in the synthetic + `coupled_fields` tensor so that averaging over any time window still + yields that same constant, letting each variable's contribution to + `preset_coupled_fields` be checked in isolation and by position. + """ + from physicsnemo.datapipes.healpix.couplers import TrailingAverageCoupler + + variables = ["z500", "z1000-12h", "z250"] + input_times = ["6h", "12h"] + input_time_dim = 2 + output_time_dim = 2 + batch_size = 2 + averaging_window = "6h" + zarr_ds = zarr.open(dataset_path) + + coupler = TrailingAverageCoupler( + dataset=zarr_ds, + batch_size=batch_size, + variables=variables, + presteps=0, + averaging_window=averaging_window, + input_times=input_times, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ) + coupler.coupled_channel_indices = list(range(len(variables))) + + data_time_step = "3h" + averaging_window_max_indices = [ + i // pd.Timedelta(data_time_step) for i in input_times + ] + di = averaging_window_max_indices[0] + averaging_slices = [] + for j in range(coupler.coupled_integration_dim): + averaging_slices.append([]) + for i, r in enumerate(averaging_window_max_indices): + averaging_slices[j].append( + slice( + coupler.input_time_dim * j * di + i * di, + coupler.input_time_dim * j * di + r, + ) + ) + coupler.averaging_slices = averaging_slices + + # give each coupled variable a distinct, constant value across every + # other dimension so the averaging window doesn't change its value, and + # each variable's contribution can be checked independently of the + # others + channel_values = [100.0, 200.0, 300.0] + coupled_fields = th.empty( + batch_size, + coupler.spatial_dims[0], + 4, + len(variables), + coupler.spatial_dims[1], + coupler.spatial_dims[2], + ) + for i, value in enumerate(channel_values): + coupled_fields[:, :, :, i, :, :] = value + + coupler.set_coupled_fields(coupled_fields) + result = coupler.construct_integrated_couplings() + assert list(result.shape) == [ + coupler.coupled_integration_dim, + batch_size, + coupler.timevar_dim, + ] + list(coupler.spatial_dims) + + # timevar ordering groups by averaging period first, then by variable + # (matching the dataset-loading path in + # `_construct_integrated_couplings_from_dataset`/`construct_integrated_couplings`) + for period in range(len(input_times)): + for var_idx, value in enumerate(channel_values): + timevar_idx = period * len(variables) + var_idx + slice_result = result[:, :, timevar_idx, :, :, :] + assert th.allclose(slice_result, th.full_like(slice_result, value)), ( + f"coupled variable {variables[var_idx]!r} (value {value}) was not " + f"preserved at averaging period {period}, timevar index {timevar_idx}" + ) + + DistributedManager.cleanup() + + @requires_module("omegaconf") @requires_module("netCDF4") @requires_module("xarray") @@ -792,6 +862,39 @@ def test_CoupledTimeSeriesDatasetZarr_get( assert non_perturbed_inputs[0][0][0].shape == perturbed_inputs[0][0][0].shape assert not np.array_equal(non_perturbed_inputs[0][0][0], perturbed_inputs[0][0][0]) + # noise is also added to the integrated couplings themselves, not just + # the regular inputs, when couplings are present. With no constants and + # no insolation, inputs_result is [inputs, integrated_couplings] + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=batch2_constant_coupler, + ) + non_perturbed_sample = timeseries_ds[0][0] + assert len(non_perturbed_sample) == 2 + non_perturbed_couplings = non_perturbed_sample[1] + + timeseries_ds = CoupledTimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + add_train_noise=True, + train_noise_params=noise_params, + start_date=zarr_ds.time[0].values, + end_date=zarr_ds.time[-1].values, + couplings=batch2_constant_coupler, + ) + perturbed_couplings = timeseries_ds[0][0][1] + assert non_perturbed_couplings.shape == perturbed_couplings.shape + assert not np.array_equal(non_perturbed_couplings, perturbed_couplings) + # With insolation we get 1 extra channel timeseries_ds = CoupledTimeSeriesDatasetZarr( dataset_path=dataset_path, @@ -1014,31 +1117,7 @@ def test_CoupledTimeSeriesDataModuleZarr_get_dataloaders( couplings=constant_coupler_config, ) - # with 1 shard should get no sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) - assert train_sampler is None - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) - assert val_sampler is None - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) - assert test_sampler is None - assert isinstance(test_dataloader, DataLoader) - - # with >1 shard should be distributed sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) - assert isinstance(train_sampler, DistributedSampler) - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) - assert isinstance(val_sampler, DistributedSampler) - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) - assert isinstance(test_sampler, DistributedSampler) - assert isinstance(test_dataloader, DataLoader) + assert_shard_dataloaders(timeseries_dm) DistributedManager.cleanup() @@ -1105,7 +1184,9 @@ def test_CoupledTimeSeriesDatasetZarr_next_integration( spatial_dims = [12, 32, 32] input_variables = ["z500", "z1000"] - coupled_channel_indices = [0, 1] + # length must match len(coupled_variables) (i.e. the coupler's timevar_dim); + # values are arbitrary since this synthetic setup bypasses setup_coupling() + coupled_channel_indices = [0] coupled_variables = ["z250"] num_variables = len(input_variables) input_time_dim = 1 @@ -1167,12 +1248,11 @@ def test_CoupledTimeSeriesDatasetZarr_next_integration( spatial_dims[2], ) - expected_coupling = coupled_fields[:, :, :, coupled_channel_indices, :, :].permute( - 2, 0, 3, 1, 4, 5 - ) - expected_coupling = expected_coupling[0, :, -1, :, :, :] - expected_coupling = expected_coupling.unsqueeze(0).unsqueeze(0) - expected_coupling = expected_coupling.repeat(1, batch_size, 1, 1, 1, 1) + # mirrors set_coupled_fields: broadcast the first time step across the + # integration dim and permute to [integration_dim, batch, timevar, face, height, width] + expected_coupling = coupled_fields[:, :, :, coupled_channel_indices, :, :] + expected_coupling = expected_coupling[:, :, :1, :, :, :] + expected_coupling = expected_coupling.permute(2, 0, 3, 1, 4, 5) # need to grab at least 1 sample to properly intialize everything timeseries_ds[0] diff --git a/test/datapipes/test_healpix_zarr.py b/test/datapipes/healpix/test_healpix_zarr.py similarity index 74% rename from test/datapipes/test_healpix_zarr.py rename to test/datapipes/healpix/test_healpix_zarr.py index 6b59ef8805..ab111a7198 100644 --- a/test/datapipes/test_healpix_zarr.py +++ b/test/datapipes/healpix/test_healpix_zarr.py @@ -19,11 +19,10 @@ import warnings import pytest -from torch.utils.data import DataLoader -from torch.utils.data.distributed import DistributedSampler from physicsnemo.distributed import DistributedManager from test.conftest import requires_module +from test.datapipes.healpix.conftest import assert_shard_dataloaders omegaconf = pytest.importorskip("omegaconf") np = pytest.importorskip("numpy") @@ -31,58 +30,40 @@ zarr = pytest.importorskip("zarr") -@pytest.fixture -def dataset_path(nfs_data_dir): - return nfs_data_dir.joinpath("datasets/healpix/healpix.zarr") - - -@pytest.fixture -def splits(): - split_dict = { - "train_date_start": "1979-01-01", - "train_date_end": "1979-01-01T21:00", - "val_date_start": "1979-01-02", - "val_date_end": "1979-01-02T09:00", - "test_date_start": "1979-01-02T12:00", - "test_date_end": "1979-01-02T18:00", - } - return split_dict - - -@pytest.fixture -def scaling_dict(): - scaling = { - "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, - "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, - "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, - "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, - "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, - "z250": {"mean": 101186.28125, "std": 5551.77978515625}, - "z500": {"mean": 55625.9609375, "std": 2681.712890625}, - "lsm": {"mean": 0, "std": 1}, - "z": {"mean": 0, "std": 1}, - "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, - "extra": {"mean": 1, "std": 0}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - -@pytest.fixture -def scaling_double_dict(): - scaling = { - "t2m0": {"mean": 0, "std": 2}, - "t850": {"mean": 0, "std": 2}, - "tau300-700": {"mean": 0, "std": 2}, - "tcwv0": {"mean": 0, "std": 2}, - "z1000": {"mean": 0, "std": 2}, - "z250": {"mean": 0, "std": 2}, - "z500": {"mean": 0, "std": 2}, - "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, - "lsm": {"mean": 0, "std": 2}, - "z": {"mean": 0, "std": 2}, - "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) +def test_object_store_path_helpers(tmp_path): + """Zarr datasets may live on object stores (e.g. s3://, or fsspec-chained paths + like simplecache::s3://); such paths are recognized by path syntax alone and + skip the local filesystem existence check performed for plain local paths. + This behavior is unique to the Zarr-backed datapipes and doesn't require the + NFS test dataset, so it can run unconditionally. + """ + from physicsnemo.datapipes.healpix.base_timeseries_dataset_zarr import ( + _check_availability, + _is_object_store_path, + ) + + assert _is_object_store_path("s3://bucket/data.zarr") + assert _is_object_store_path("simplecache::s3://bucket/data.zarr") + assert not _is_object_store_path(str(tmp_path / "local.zarr")) + + # an existing local path passes without needing fsspec + existing_path = tmp_path / "exists.zarr" + existing_path.mkdir() + _check_availability(str(existing_path)) + + # a missing local path raises FileNotFoundError + with pytest.raises(FileNotFoundError, match=("Dataset not found at")): + _check_availability(str(tmp_path / "missing.zarr")) + + # object store paths bypass the local existence check as long as fsspec + # is installed, even when the remote resource doesn't actually exist + if importlib.util.find_spec("fsspec"): + _check_availability("s3://bucket-that-does-not-exist/missing.zarr") + else: + with pytest.raises( + ImportError, match=("fsspec is required to access object store paths") + ): + _check_availability("s3://bucket-that-does-not-exist/missing.zarr") @requires_module("omegaconf") @@ -103,7 +84,6 @@ def test_TimeSeriesDataset_initialization( bad_end_date = "2000-12-31" valid_end_date = "1979-01-02" - zarr_ds = zarr.open(dataset_path) time_da = xr.open_zarr(dataset_path).time # check for failure of invalid dataset path @@ -137,7 +117,7 @@ def test_TimeSeriesDataset_initialization( time_step="5h", scaling=scaling_dict, input_variables=input_variables, - forecast_init_times=zarr_ds.time[:2], + forecast_init_times=time_da[:2], batch_size=1, ) @@ -173,6 +153,38 @@ def test_TimeSeriesDataset_initialization( batch_size=1, ) + # check for failure when a requested variable isn't present in the dataset + # this validation is specific to the Zarr-backed dataset, which selects + # channels by name directly out of the store rather than relying on a + # pre-sliced xarray Dataset + with pytest.raises( + KeyError, + match=("Requested Input, coupled, or output variables not found in dataset"), + ): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=input_variables + ["DoesntExist"], + ) + + # check for warning when the configured data_time_step doesn't match the + # dataset's native cadence; this cross-check against the stored time + # coordinate only exists in the Zarr-backed dataset + warnings.filterwarnings("error") + with pytest.raises(UserWarning, match=("doesn't match configuration dt")): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + data_time_step="6h", + time_step="6h", + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=input_variables, + ) + warnings.resetwarnings() + # check for warning on batch size > 1 and forecast mode warnings.filterwarnings("error") with pytest.raises( @@ -185,7 +197,7 @@ def test_TimeSeriesDataset_initialization( dataset_path=dataset_path, scaling=scaling_dict, batch_size=2, - forecast_init_times=zarr_ds.time[:2], + forecast_init_times=time_da[:2], input_variables=input_variables, ) warnings.resetwarnings() @@ -270,6 +282,89 @@ def test_TimeSeriesDataset_initialization( ) assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) + # `start_date`/`end_date` can also be integer positional indices into the + # time array rather than dates (use a nonzero start_date since 0 is falsy + # and would take the same code path as "no start date provided") + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_dict, + start_date=2, + end_date=5, + input_variables=input_variables, + ) + assert timeseries_ds.start_index == 2 + assert timeseries_ds.total_samples == 3 + + # check for failure when a requested output variable has no scaling + # entry, even though it's present in the dataset (and in the input + # scaling, if it happens to also be an input variable) + scaling_missing_target = omegaconf.DictConfig( + {k: v for k, v in scaling_dict.items() if k != "z1000"} + ) + with pytest.raises(KeyError, match=("Target channels ")): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_missing_target, + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=["t2m0"], + output_variables=["t2m0", "z1000"], + ) + + # check for failure when a requested constant variable exists in the + # dataset but has no corresponding scaling entry + scaling_missing_constant = omegaconf.DictConfig( + {k: v for k, v in scaling_dict.items() if k != "z"} + ) + with pytest.raises(KeyError, match=("Constant variables ")): + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + scaling=scaling_missing_constant, + start_date=valid_start_date, + end_date=valid_end_date, + input_variables=input_variables, + constant_variables=["lsm", "z"], + ) + + +def test_TimeSeriesDataset_missing_time(tmp_path): + """The Zarr-backed dataset validates that the store has a `time` + variable before doing any time-based indexing; this check is unique to + the Zarr path (the classic path always derives its dataset from a + pre-opened xarray Dataset that already has a time dimension), so build a + minimal Zarr store without a time coordinate to exercise it directly. + """ + from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( + TimeSeriesDatasetZarr, + ) + + no_time_ds = xr.Dataset( + data_vars={ + "inputs": ( + ("t", "channel_in", "face", "height", "width"), + np.zeros((3, 1, 1, 2, 2), dtype="float32"), + ), + "targets": ( + ("t", "channel_out", "face", "height", "width"), + np.zeros((3, 1, 1, 2, 2), dtype="float32"), + ), + }, + coords={ + "channel_in": ["z500"], + "channel_out": ["z500"], + }, + ) + dataset_path = tmp_path / "no_time.zarr" + no_time_ds.to_zarr(dataset_path) + + with pytest.raises(KeyError, match=("Dataset missing time")): + TimeSeriesDatasetZarr( + dataset_path=str(dataset_path), + input_variables=["z500"], + start_date="1979-01-01", + end_date="1979-01-02", + ) + @requires_module("omegaconf") @requires_module("netCDF4") @@ -507,6 +602,37 @@ def test_TimeSeriesDataset_get(dataset_path, scaling_double_dict, splits, pytest ) assert (len(inputs)) + 2 == len(timeseries_ds[0][0]) + # train noise is applied directly to the inputs of the (non-coupled) Zarr + # dataset; this option doesn't exist on the classic (non-Zarr) TimeSeriesDataset + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + start_date=time_da[0], + end_date=time_da[-1], + ) + non_perturbed = timeseries_ds[0] + + noise_params = {"inputs": scaling_double_dict} + timeseries_ds = TimeSeriesDatasetZarr( + dataset_path=dataset_path, + input_variables=input_variables, + scaling=scaling_double_dict, + batch_size=batch_size, + drop_last=True, + add_train_noise=True, + train_noise_params=noise_params, + start_date=time_da[0], + end_date=time_da[-1], + ) + perturbed = timeseries_ds[0] + + # same shape, but the perturbed sample should differ from the un-perturbed one + assert non_perturbed[0][0].shape == perturbed[0][0].shape + assert not np.array_equal(non_perturbed[0][0], perturbed[0][0]) + # nothing should change with forecast mode other than getting just inputs init_times = random.randint(1, len(zarr_ds.time.values)) timeseries_ds = TimeSeriesDatasetZarr( @@ -743,29 +869,5 @@ def test_TimeSeriesDataModule_get_dataloaders( shuffle=False, ) - # with 1 shard should get no sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) - assert train_sampler is None - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) - assert val_sampler is None - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) - assert test_sampler is None - assert isinstance(test_dataloader, DataLoader) - - # with >1 shard should be distributed sampler - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) - assert isinstance(train_sampler, DistributedSampler) - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) - assert isinstance(val_sampler, DistributedSampler) - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) - assert isinstance(test_sampler, DistributedSampler) - assert isinstance(test_dataloader, DataLoader) + assert_shard_dataloaders(timeseries_dm) DistributedManager.cleanup() From 8135b1e1228b3fd11b51f45722040a2a0a68b8f8 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Fri, 10 Jul 2026 12:29:12 -0500 Subject: [PATCH 11/28] Update tests + small bugfixes --- .../models/dlwp_healpix/HEALPixRecUNet.py | 11 +- .../models/dlwp_healpix/HEALPixUNet.py | 11 +- .../dlwp_healpix/layers/normalization.py | 2 +- physicsnemo/nn/module/hpx/padding.py | 15 +- .../test_conditional_layer_norm.py | 341 ++++++-- .../dlwp_healpix/test_healpix_blocks.py | 437 +++++++++++ .../test_healpix_encoder_decoder.py | 675 ++++++++++++++++ .../dlwp_healpix/test_healpix_layers.py | 96 ++- .../test_healpix_recunet_model.py | 727 ++++++++++++++++++ .../dlwp_healpix/test_healpix_unet_model.py | 630 +++++++++++++++ .../test_layers_backward_compat.py | 222 ++++++ test/nn/module/healpix_helpers.py | 66 ++ test/nn/module/test_healpix.py | 320 +++++++- 13 files changed, 3481 insertions(+), 72 deletions(-) create mode 100644 test/models/dlwp_healpix/test_layers_backward_compat.py create mode 100644 test/nn/module/healpix_helpers.py diff --git a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py index 9127dec57f..6bacf565f3 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py @@ -400,9 +400,14 @@ def set_constraints(self, constraints: list[DictConfig] = None): Hydra instantiable constraint configurations. """ if constraints is not None: - self.constraints = [ - instantiate(constraints[constraint]) for constraint in constraints - ] + # Use a ModuleList (rather than a plain list) so the constraint + # modules are part of the module tree: their buffers (e.g. the + # non-persistent ``thresholds`` and ``var_indices`` in + # ``NonnegativeConstraint``) are then correctly moved by + # ``.to(device)`` along with the rest of the model. + self.constraints = torch.nn.ModuleList( + [instantiate(constraints[constraint]) for constraint in constraints] + ) def _reshape_outputs(self, outputs: torch.Tensor) -> torch.Tensor: r""" diff --git a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py index 0f46f5501f..479d0d5a82 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py @@ -417,9 +417,14 @@ def set_constraints(self, constraints: list[DictConfig] = None): Hydra instantiable constraint configurations. """ if constraints is not None: - self.constraints = [ - instantiate(constraints[constraint]) for constraint in constraints - ] + # Use a ModuleList (rather than a plain list) so the constraint + # modules are part of the module tree: their buffers (e.g. the + # non-persistent ``thresholds`` and ``var_indices`` in + # ``NonnegativeConstraint``) are then correctly moved by + # ``.to(device)`` along with the rest of the model. + self.constraints = torch.nn.ModuleList( + [instantiate(constraints[constraint]) for constraint in constraints] + ) def forward( self, diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index 697994269b..9d6f3c1c30 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -27,7 +27,7 @@ try: from apex.normalization import FusedLayerNorm - _APEX_AVAILABLE = True + _APEX_AVAILABLE = True # pragma: no cover except ImportError: _APEX_AVAILABLE = False diff --git a/physicsnemo/nn/module/hpx/padding.py b/physicsnemo/nn/module/hpx/padding.py index 0a97421945..935ebeeeb6 100644 --- a/physicsnemo/nn/module/hpx/padding.py +++ b/physicsnemo/nn/module/hpx/padding.py @@ -27,8 +27,13 @@ hpx_pad = importlib.import_module("earth2grid.healpix._padding").pad else: - def hpx_pad(*args, **kwargs): - """Dummy symbol for missing earth2grid backend.""" + def hpx_pad(*args, **kwargs): # pragma: no cover + """Dummy symbol for missing earth2grid backend. + + Only reachable when ``earth2grid`` is not installed, so it cannot be + exercised in an environment where it is (as is required to test the + rest of this module's accelerated path). + """ raise ImportError( ( "earth2grid is not installed, cannot use it as a backend for HEALPix padding.\n" @@ -390,7 +395,7 @@ def forward( Folded tensor of shape :math:`(B \cdot F, C, H, W)`. """ if not torch.compiler.is_compiling() and tensor.ndim != 5: - ValueError( + raise ValueError( f"HEALPixFoldFaces.forward requires 5D tensor, got {tensor.shape}" ) @@ -455,11 +460,11 @@ def forward( """ if not torch.compiler.is_compiling(): if tensor.ndim != 4: - ValueError( + raise ValueError( f"HEALPixUnfoldFaces.forward requires 4D tensor, got {tensor.shape}" ) if tensor.shape[0] % self.num_faces != 0: - ValueError( + raise ValueError( f"HEALPixUnfoldFaces.forward invalid batch size: {tensor.shape[0]}" ) diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py index 8530f82c83..da7bf6615b 100644 --- a/test/models/dlwp_healpix/test_conditional_layer_norm.py +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -23,22 +23,26 @@ import torch from _cln_reference import ConditionalLayerNormReference -from physicsnemo.models.dlwp_healpix.layers.normalization import ConditionalLayerNorm +from physicsnemo.models.dlwp_healpix.layers import normalization +from physicsnemo.models.dlwp_healpix.layers.normalization import ( + _APEX_AVAILABLE, + ConditionalLayerNorm, +) from physicsnemo.nn import CappedGELU -def _make_old_cln(condition_shape, channel_depth, **kwargs): +def _make_old_cln(condition_shape, channel_depth, device="cpu", **kwargs): """Instantiate the reference (old) implementation.""" return ConditionalLayerNormReference( condition_shape=condition_shape, channel_depth=channel_depth, **kwargs - ).cuda() + ).to(device) -def _make_new_cln(condition_shape, channel_depth, **kwargs): +def _make_new_cln(condition_shape, channel_depth, device="cpu", **kwargs): """Instantiate the optimized (new) implementation.""" return ConditionalLayerNorm( condition_shape=condition_shape, channel_depth=channel_depth, **kwargs - ).cuda() + ).to(device) def _copy_old_to_new(old_cln, new_cln): @@ -103,24 +107,36 @@ def _copy_old_to_new(old_cln, new_cln): new_cln.load_state_dict(new_sd) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def _assert_forward_is_finite( + cln, channel_depth, condition_shape, device, n_cond=1, n_faces=12, height=4, width=4 +): + """Run a minimal forward pass and assert the output is finite; used by + the legacy-checkpoint-loading tests to confirm a (possibly + partially-loaded) module still produces a usable output.""" + x = torch.randn(n_cond * n_faces, channel_depth, height, width, device=device) + cond = torch.randn(n_cond, condition_shape, device=device) + out = cln(x, cond) + assert torch.isfinite(out).all() + return out + + @pytest.mark.parametrize("n_cond", [1, 2, 4]) @pytest.mark.parametrize("channels_last", [False, True]) @pytest.mark.parametrize("scale_center", [0.0, 1.0]) -def test_old_vs_new_forward(n_cond, channels_last, scale_center): +def test_old_vs_new_forward(device, n_cond, channels_last, scale_center): """Verify optimized CLN matches reference implementation with block-diagonal weight mapping.""" C, H, W = 128, 16, 16 cond_shape = 32 B_nf = n_cond * 12 torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C, scale_center=scale_center) + old_cln = _make_old_cln(cond_shape, C, device=device, scale_center=scale_center) - new_cln = _make_new_cln(cond_shape, C, scale_center=scale_center) + new_cln = _make_new_cln(cond_shape, C, device=device, scale_center=scale_center) _copy_old_to_new(old_cln, new_cln) - x = torch.randn(B_nf, C, H, W, device="cuda") - cond = torch.randn(n_cond, cond_shape, device="cuda") + x = torch.randn(B_nf, C, H, W, device=device) + cond = torch.randn(n_cond, cond_shape, device=device) if channels_last: x = x.to(memory_format=torch.channels_last) @@ -140,9 +156,8 @@ def test_old_vs_new_forward(n_cond, channels_last, scale_center): ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("channels_last", [False, True]) -def test_old_vs_new_backward(channels_last): +def test_old_vs_new_backward(device, channels_last): """Verify gradients match between old and new implementations.""" C, H, W = 64, 8, 8 cond_shape = 16 @@ -150,12 +165,12 @@ def test_old_vs_new_backward(channels_last): B_nf = n_cond * 12 torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C) - new_cln = _make_new_cln(cond_shape, C) + old_cln = _make_old_cln(cond_shape, C, device=device) + new_cln = _make_new_cln(cond_shape, C, device=device) _copy_old_to_new(old_cln, new_cln) - x_base = torch.randn(B_nf, C, H, W, device="cuda") - cond_base = torch.randn(n_cond, cond_shape, device="cuda") + x_base = torch.randn(B_nf, C, H, W, device=device) + cond_base = torch.randn(n_cond, cond_shape, device=device) if channels_last: x_base = x_base.to(memory_format=torch.channels_last) @@ -180,20 +195,19 @@ def test_old_vs_new_backward(channels_last): ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("channels_last", [False, True]) -def test_init_cln_to_zero_matches_layer_norm(channels_last): +def test_init_cln_to_zero_matches_layer_norm(device, channels_last): """With scale_center=1.0 and init_cln_to_zero=True, CLN should behave like plain LayerNorm.""" C, H, W = 64, 8, 8 n_cond = 2 B_nf = n_cond * 12 torch.manual_seed(42) - cln = _make_new_cln(32, C, scale_center=1.0, init_cln_to_zero=True) - plain_ln = torch.nn.LayerNorm(C, elementwise_affine=False).cuda() + cln = _make_new_cln(32, C, device=device, scale_center=1.0, init_cln_to_zero=True) + plain_ln = torch.nn.LayerNorm(C, elementwise_affine=False).to(device) - x = torch.randn(B_nf, C, H, W, device="cuda") - cond = torch.randn(n_cond, 32, device="cuda") + x = torch.randn(B_nf, C, H, W, device=device) + cond = torch.randn(n_cond, 32, device=device) if channels_last: x = x.to(memory_format=torch.channels_last) @@ -208,9 +222,8 @@ def test_init_cln_to_zero_matches_layer_norm(channels_last): ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("channels_last", [False, True]) -def test_backward_gradients(channels_last): +def test_backward_gradients(device, channels_last): """Verify gradients flow through CLN and are finite.""" C, H, W = 64, 8, 8 cond_shape = 16 @@ -218,10 +231,10 @@ def test_backward_gradients(channels_last): B_nf = n_cond * 12 torch.manual_seed(42) - cln = _make_new_cln(cond_shape, C) + cln = _make_new_cln(cond_shape, C, device=device) - x = torch.randn(B_nf, C, H, W, device="cuda") - cond = torch.randn(n_cond, cond_shape, device="cuda") + x = torch.randn(B_nf, C, H, W, device=device) + cond = torch.randn(n_cond, cond_shape, device=device) if channels_last: x = x.to(memory_format=torch.channels_last) @@ -242,8 +255,7 @@ def test_backward_gradients(channels_last): assert torch.isfinite(p.grad).all(), f"Non-finite gradient for {name}" -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_backward_channels_last_matches_contiguous(): +def test_backward_channels_last_matches_contiguous(device): """Verify channels_last and contiguous inputs produce the same gradients.""" C, H, W = 64, 8, 8 cond_shape = 16 @@ -251,10 +263,10 @@ def test_backward_channels_last_matches_contiguous(): B_nf = n_cond * 12 torch.manual_seed(42) - cln = _make_new_cln(cond_shape, C) + cln = _make_new_cln(cond_shape, C, device=device) - x_base = torch.randn(B_nf, C, H, W, device="cuda") - cond_base = torch.randn(n_cond, cond_shape, device="cuda") + x_base = torch.randn(B_nf, C, H, W, device=device) + cond_base = torch.randn(n_cond, cond_shape, device=device) # Contiguous path x_cont = x_base.clone().detach().requires_grad_(True) @@ -286,20 +298,19 @@ def test_backward_channels_last_matches_contiguous(): ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_load_old_checkpoint(): +def test_load_old_checkpoint(device): """Verify new CLN can load old-format state dict via _load_from_state_dict.""" C, cond_shape = 64, 16 torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C) + old_cln = _make_old_cln(cond_shape, C, device=device) old_sd = old_cln.state_dict() - new_cln = _make_new_cln(cond_shape, C) + new_cln = _make_new_cln(cond_shape, C, device=device) new_cln.load_state_dict(old_sd, strict=False) # Verify outputs match after loading old checkpoint - x = torch.randn(12, C, 8, 8, device="cuda") - cond = torch.randn(1, cond_shape, device="cuda") + x = torch.randn(12, C, 8, 8, device=device) + cond = torch.randn(1, cond_shape, device=device) with torch.no_grad(): out_old = old_cln(x, cond) @@ -312,8 +323,7 @@ def test_load_old_checkpoint(): ) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_load_old_checkpoint_with_capped_gelu(): +def test_load_old_checkpoint_with_capped_gelu(device): """Verify strict loading of old CLN checkpoints that use CappedGELU activations.""" C, cond_shape = 64, 16 mlp_hidden_dims = [32, 32] @@ -323,6 +333,7 @@ def test_load_old_checkpoint_with_capped_gelu(): old_cln = _make_old_cln( cond_shape, C, + device=device, mlp_hidden_dims=mlp_hidden_dims, activation=activation, ) @@ -332,6 +343,7 @@ def test_load_old_checkpoint_with_capped_gelu(): new_cln = _make_new_cln( cond_shape, C, + device=device, mlp_hidden_dims=mlp_hidden_dims, activation=CappedGELU(), ) @@ -339,8 +351,8 @@ def test_load_old_checkpoint_with_capped_gelu(): assert not missing assert not unexpected - x = torch.randn(12, C, 8, 8, device="cuda") - cond = torch.randn(1, cond_shape, device="cuda") + x = torch.randn(12, C, 8, 8, device=device) + cond = torch.randn(1, cond_shape, device=device) with torch.no_grad(): out_old = old_cln(x, cond) @@ -349,3 +361,246 @@ def test_load_old_checkpoint_with_capped_gelu(): assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( f"Max diff after loading old CappedGELU checkpoint: {(out_old - out_new).abs().max().item()}" ) + + +@pytest.mark.parametrize("hidden_dims", [[], [64]]) +def test_old_vs_new_forward_hidden_dims_variation(device, hidden_dims): + """Verify the block-diagonal weight mapping generalizes to MLP depths other + than the default two hidden layers, including the degenerate zero-hidden-layer + case (a single Linear directly from ``condition_shape`` to the output). + """ + C, H, W = 32, 8, 8 + cond_shape = 16 + n_cond = 2 + B_nf = n_cond * 12 + + torch.manual_seed(42) + old_cln = _make_old_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) + new_cln = _make_new_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) + _copy_old_to_new(old_cln, new_cln) + + x = torch.randn(B_nf, C, H, W, device=device) + cond = torch.randn(n_cond, cond_shape, device=device) + + with torch.no_grad(): + out_old = old_cln(x, cond) + out_new = new_cln(x, cond) + + assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( + f"Max diff: {(out_old - out_new).abs().max().item()}" + ) + + +def test_norm_op_apex_raises_when_unavailable(): + """``norm_op='apex'`` must raise an informative ImportError when the + ``apex`` package isn't installed, rather than silently falling back or + failing with an unrelated error later on. + """ + if _APEX_AVAILABLE: + pytest.skip("apex is installed in this environment") + + with pytest.raises(ImportError, match="Apex FusedLayerNorm requested"): + ConditionalLayerNorm(condition_shape=16, channel_depth=8, norm_op="apex") + + +def test_norm_op_invalid_value_leaves_norm_unset(): + """``norm_op`` outside {"torch", "apex"} is not validated; document the + current behavior of silently skipping norm construction rather than + raising, so a regression (e.g. an accidental typo check) is caught. + """ + cln = ConditionalLayerNorm(condition_shape=16, channel_depth=8, norm_op="bogus") + assert not hasattr(cln, "norm") + + +def test_norm_op_apex_available_uses_fused_layer_norm(monkeypatch): + """When apex *is* available, ``norm_op='apex'`` must construct a + ``FusedLayerNorm`` instead of ``torch.nn.LayerNorm``. Exercised via a + stand-in class since apex isn't a hard dependency of this environment. + """ + + class _FakeFusedLayerNorm(torch.nn.Module): + def __init__(self, channel_depth, elementwise_affine=False): + super().__init__() + self.channel_depth = channel_depth + self.elementwise_affine = elementwise_affine + + def forward(self, x): + return x + + monkeypatch.setattr(normalization, "_APEX_AVAILABLE", True) + monkeypatch.setattr( + normalization, "FusedLayerNorm", _FakeFusedLayerNorm, raising=False + ) + + cln = ConditionalLayerNorm(condition_shape=16, channel_depth=8, norm_op="apex") + assert isinstance(cln.norm, _FakeFusedLayerNorm) + assert cln.norm.channel_depth == 8 + + +def test_make_mlp_skips_activation_when_falsy(): + """``_make_mlp`` only inserts an activation module between hidden Linear + layers when ``activation`` is truthy; verify the skip path directly since + the public constructor always coerces ``None`` to ``nn.Identity()`` + (itself truthy), making this otherwise unreachable through ``__init__``. + """ + cln = ConditionalLayerNorm(condition_shape=16, channel_depth=8, mlp_hidden_dims=[]) + mlp = cln._make_mlp(in_dim=4, hidden_dims=[8, 8], out_dim=2, activation=False) + assert all(isinstance(layer, torch.nn.Linear) for layer in mlp) + assert len(mlp) == 3 # one Linear per hidden dim, plus the output Linear + + +def test_cln_affine_eager_matches_compiled(monkeypatch, device): + """``_cln_affine`` is wrapped in ``@torch.compile``, which bypasses + coverage tracing of its body; run it once via the (functionally + identical) eager callable dynamo wraps to confirm equivalence and + exercise the underlying implementation directly. + """ + eager_cln_affine = normalization._cln_affine._torchdynamo_orig_callable + + C, H, W = 16, 4, 4 + cond_shape = 8 + n_cond = 2 + n_faces = 12 + + torch.manual_seed(0) + x_norm = torch.randn(n_cond * n_faces, H, W, C, device=device) + gamma_raw = torch.randn(n_cond, C, device=device) + beta = torch.randn(n_cond, C, device=device) + scale_center = 1.0 + + compiled_out = normalization._cln_affine( + x_norm, gamma_raw, beta, scale_center, n_faces + ) + eager_out = eager_cln_affine(x_norm, gamma_raw, beta, scale_center, n_faces) + + assert torch.allclose(compiled_out, eager_out) + + # also drive it end-to-end through ConditionalLayerNorm.forward with the + # module-level name patched to the eager callable + monkeypatch.setattr(normalization, "_cln_affine", eager_cln_affine) + cln = ConditionalLayerNorm( + condition_shape=cond_shape, channel_depth=C, mlp_hidden_dims=[] + ).to(device) + _assert_forward_is_finite( + cln, C, cond_shape, device, n_cond=n_cond, n_faces=n_faces, height=H, width=W + ) + + +def test_load_from_state_dict_ignores_malformed_activation_buffer_keys(device): + """The activation-buffer migration loop must gracefully skip legacy keys + that don't fit the expected ``"."`` shape (no dot, or a + non-numeric index), rather than raising. + """ + C, cond_shape = 16, 8 + old_cln = _make_old_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) + old_sd = old_cln.state_dict() + + # no dot after the gamma_mlp prefix at all + old_sd["gamma_mlp.malformed"] = torch.zeros(1) + # dot present, but the leading segment isn't a valid layer index + old_sd["gamma_mlp.bogus.cap"] = torch.zeros(1) + + new_cln = _make_new_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) + missing, unexpected = new_cln.load_state_dict(old_sd, strict=False) + + # both malformed keys are left untouched (never merged into the fused + # MLP), so they remain unexpected for the new module + assert "gamma_mlp.malformed" in unexpected + assert "gamma_mlp.bogus.cap" in unexpected + + _assert_forward_is_finite(new_cln, C, cond_shape, device) + + +def test_load_from_state_dict_activation_buffer_without_beta_counterpart(device): + """An activation submodule buffer (e.g. ``CappedGELU``'s ``cap``) present + on the gamma MLP but missing from the beta MLP (malformed/partial + checkpoint) must still migrate the gamma-side buffer, and must not raise + trying to remove the (absent) beta-side key. + """ + C, cond_shape = 16, 8 + mlp_hidden_dims = [8] + old_cln = _make_old_cln( + cond_shape, + C, + device=device, + mlp_hidden_dims=mlp_hidden_dims, + activation=CappedGELU(), + ) + old_sd = old_cln.state_dict() + assert "gamma_mlp.1.cap" in old_sd + assert "beta_mlp.1.cap" in old_sd + + # drop only the beta-side activation buffer + del old_sd["beta_mlp.1.cap"] + + new_cln = _make_new_cln( + cond_shape, + C, + device=device, + mlp_hidden_dims=mlp_hidden_dims, + activation=CappedGELU(), + ) + missing, unexpected = new_cln.load_state_dict(old_sd, strict=False) + + # the gamma-side buffer was still migrated into the fused MLP + assert "gamma_beta_mlp.1.cap" not in unexpected + assert "gamma_mlp.1.cap" not in unexpected + + _assert_forward_is_finite(new_cln, C, cond_shape, device) + + +def test_load_from_state_dict_skips_unmatched_gamma_key(device): + """A legacy state dict where a ``gamma_mlp`` key has no matching + ``beta_mlp`` counterpart (e.g. a malformed/partial checkpoint) must be + skipped rather than merged, and loading must not raise. + """ + C, cond_shape = 32, 16 + torch.manual_seed(42) + old_cln = _make_old_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) + old_sd = old_cln.state_dict() + + # drop the beta counterpart for the (only) gamma layer, simulating a + # malformed/partial legacy checkpoint + del old_sd["beta_mlp.0.weight"] + del old_sd["beta_mlp.0.bias"] + + new_cln = _make_new_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) + missing, unexpected = new_cln.load_state_dict(old_sd, strict=False) + + # the unmatched gamma key is left untouched (not merged into + # gamma_beta_mlp), so it shows up as unexpected for the new module, and + # the fused MLP's own parameters are all reported missing since nothing + # could be merged + assert "gamma_mlp.0.weight" in unexpected + assert "gamma_mlp.0.bias" in unexpected + assert any(k.startswith("gamma_beta_mlp.") for k in missing) + + # forward should still run without error using the (unloaded, default + # initialized) fused MLP weights + _assert_forward_is_finite(new_cln, C, cond_shape, device) + + +def test_load_new_format_checkpoint_roundtrip(device): + """A state dict already in the fused ``gamma_beta_mlp`` format (i.e. saved + from a ``ConditionalLayerNorm`` instance, with no legacy ``gamma_mlp``/ + ``beta_mlp`` keys) should load directly with no remapping and reproduce + identical outputs. + """ + C, cond_shape = 32, 16 + torch.manual_seed(42) + source = _make_new_cln(cond_shape, C, device=device) + state_dict = source.state_dict() + assert not any(k.startswith("gamma_mlp.") for k in state_dict) + + target = _make_new_cln(cond_shape, C, device=device) + missing, unexpected = target.load_state_dict(state_dict, strict=True) + assert not missing + assert not unexpected + + x = torch.randn(12, C, 4, 4, device=device) + cond = torch.randn(1, cond_shape, device=device) + with torch.no_grad(): + out_source = source(x, cond) + out_target = target(x, cond) + + assert torch.allclose(out_source, out_target, atol=1e-6, rtol=1e-5) diff --git a/test/models/dlwp_healpix/test_healpix_blocks.py b/test/models/dlwp_healpix/test_healpix_blocks.py index ae004e59b5..15405cba93 100644 --- a/test/models/dlwp_healpix/test_healpix_blocks.py +++ b/test/models/dlwp_healpix/test_healpix_blocks.py @@ -14,6 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # ruff: noqa: E402 +from functools import partial + import pytest import torch @@ -33,6 +35,38 @@ def generate_test_data(faces=12, channels=2, img_size=16, device="cpu"): return generate_test_data +def _cln_factory(cond_dim): + """Build a ``ConditionalLayerNorm`` factory bound to ``cond_dim``, ready + to pass as the ``conditional_layer_norm`` argument of a block.""" + from physicsnemo.models.dlwp_healpix.layers.normalization import ( + ConditionalLayerNorm, + ) + + return partial(ConditionalLayerNorm, condition_shape=cond_dim) + + +def _assert_dropout_present(module): + """Assert at least one ``Dropout2d`` was inserted somewhere in ``module``.""" + n_dropout = sum(1 for m in module.modules() if isinstance(m, torch.nn.Dropout2d)) + assert n_dropout > 0 + + +def _assert_dropout_eval_deterministic_train_stochastic(module, invar, out_shape): + """Dropout must be a no-op in ``eval()`` (repeated calls match) and + stochastic in ``train()`` (repeated calls with identical input diverge).""" + module.eval() + outvar_1 = module(invar) + outvar_2 = module(invar) + assert outvar_1.shape == out_shape + assert common.compare_output(outvar_1, outvar_2) + + module.train() + outvar_3 = module(invar) + outvar_4 = module(invar) + assert outvar_3.shape == out_shape + assert not common.compare_output(outvar_3, outvar_4) + + def test_ConvGRUBlock_initialization(device, test_data, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( ConvGRUBlock, @@ -64,6 +98,31 @@ def test_ConvGRUBlock_forward(device, test_data, pytestconfig): assert not common.compare_output(outvar_hist, outvar) +def test_ConvGRUBlock_reset(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + ConvGRUBlock, + ) + + in_channels = 2 + tensor_size = 16 + conv_gru_func = ConvGRUBlock(in_channels=in_channels).to(device) + + invar = test_data(img_size=tensor_size, device=device) + + # first call establishes the hidden state after being zero-initialized + first_call = conv_gru_func(invar) + + # subsequent call with the same input diverges because hidden state + # is now non-zero + second_call = conv_gru_func(invar) + assert not common.compare_output(first_call, second_call) + + # resetting the hidden state to zero should reproduce the first call + conv_gru_func.reset() + reset_call = conv_gru_func(invar) + assert common.compare_output(first_call, reset_call) + + def test_ConvNeXtBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( ConvNeXtBlock, @@ -167,6 +226,116 @@ def test_DoubleConvNeXtBlock_forward(device, test_data, pytestconfig): assert outvar.shape == out_shape +def test_DoubleConvNeXtBlock_dropout(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + DoubleConvNeXtBlock, + ) + + in_channels = 2 + out_channels = 1 + latent_channels = 2 + tensor_size = 16 + doubleconvnextblock = DoubleConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + latent_channels=latent_channels, + dropout=0.5, + ).to(device) + + # dropout is inserted after every conv/norm/activation step of both + # internal convnext blocks + _assert_dropout_present(doubleconvnextblock) + + invar = test_data(img_size=tensor_size, device=device) + out_shape = torch.Size([12, out_channels, tensor_size, tensor_size]) + _assert_dropout_eval_deterministic_train_stochastic( + doubleconvnextblock, invar, out_shape + ) + + +def test_DoubleConvNeXtBlock_conditional_layer_norm(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + DoubleConvNeXtBlock, + ) + + in_channels = 2 + out_channels = 1 + latent_channels = 2 + cond_dim = 4 + tensor_size = 16 + conditional_layer_norm = _cln_factory(cond_dim) + + doubleconvnextblock = DoubleConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + latent_channels=latent_channels, + conditional_layer_norm=conditional_layer_norm, + conditional_layer_norm_once=False, + ).to(device) + assert doubleconvnextblock.cln_enabled + + invar = test_data(img_size=tensor_size, device=device) + out_shape = torch.Size([12, out_channels, tensor_size, tensor_size]) + + conditions_a = torch.randn(1, cond_dim).to(device) + conditions_b = torch.randn(1, cond_dim).to(device) + + outvar_a = doubleconvnextblock(invar, conditions_cln=conditions_a) + outvar_b = doubleconvnextblock(invar, conditions_cln=conditions_b) + + assert outvar_a.shape == out_shape + # different conditions must produce different normalization affine + # parameters, and thus different outputs + assert not common.compare_output(outvar_a, outvar_b) + + +@pytest.mark.parametrize("conditional_layer_norm_enabled", [True, False]) +def test_DoubleConvNeXtBlock_conditional_layer_norm_once( + device, test_data, pytestconfig, conditional_layer_norm_enabled +): + from physicsnemo.models.dlwp_healpix.layers import ( + DoubleConvNeXtBlock, + ) + + in_channels = 2 + out_channels = 1 + latent_channels = 2 + cond_dim = 4 + tensor_size = 16 + + conditional_layer_norm = ( + _cln_factory(cond_dim) if conditional_layer_norm_enabled else None + ) + + doubleconvnextblock = DoubleConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + latent_channels=latent_channels, + conditional_layer_norm=conditional_layer_norm, + conditional_layer_norm_once=True, + ).to(device) + + invar = test_data(img_size=tensor_size, device=device) + out_shape = torch.Size([12, out_channels, tensor_size, tensor_size]) + + if conditional_layer_norm_enabled: + conditions_a = torch.randn(1, cond_dim).to(device) + conditions_b = torch.randn(1, cond_dim).to(device) + + outvar_a = doubleconvnextblock(invar, conditions_cln=conditions_a) + outvar_b = doubleconvnextblock(invar, conditions_cln=conditions_b) + + assert outvar_a.shape == out_shape + # entry norm is conditional, so different conditions must still + # change the output even though the internal norms have switched + # to plain (non-conditional) layer normalization + assert not common.compare_output(outvar_a, outvar_b) + else: + outvar = doubleconvnextblock(invar) + assert outvar.shape == out_shape + assert torch.isfinite(outvar).all() + + def test_SymmetricConvNeXtBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( SymmetricConvNeXtBlock, @@ -210,6 +379,175 @@ def test_SymmetricConvNeXtBlock_forward(device, test_data, pytestconfig): assert outvar.shape == out_shape +def test_SymmetricConvNeXtBlock_identity_skip(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + SymmetricConvNeXtBlock, + ) + + channels = 2 + latent_channels = 1 + tensor_size = 16 + # in_channels == out_channels triggers the identity skip_module path + symmetric_convnextblock = SymmetricConvNeXtBlock( + in_channels=channels, + out_channels=channels, + latent_channels=latent_channels, + ).to(device) + + invar = test_data(channels=channels, img_size=tensor_size, device=device) + out_shape = torch.Size([12, channels, tensor_size, tensor_size]) + + # identity skip_module must return the input unchanged + assert symmetric_convnextblock.skip_module(invar) is invar + + outvar = symmetric_convnextblock(invar) + assert outvar.shape == out_shape + + +def test_SymmetricConvNeXtBlock_no_skip_connection(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + SymmetricConvNeXtBlock, + ) + + in_channels = 2 + out_channels = 1 + latent_channels = 1 + tensor_size = 16 + + with_skip = SymmetricConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + latent_channels=latent_channels, + use_block_skip_connection=True, + ).to(device) + without_skip = SymmetricConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + latent_channels=latent_channels, + use_block_skip_connection=False, + ).to(device) + # share weights so the only difference between the two blocks is + # whether the skip connection is added. `use_block_skip_connection=False` + # never registers a `skip_module` submodule, so load non-strictly. + without_skip.load_state_dict(with_skip.state_dict(), strict=False) + + invar = test_data(channels=in_channels, img_size=tensor_size, device=device) + + out_with_skip = with_skip(invar) + out_without_skip = without_skip(invar) + + assert not common.compare_output(out_with_skip, out_without_skip) + assert common.compare_output( + out_with_skip, out_without_skip + with_skip.skip_module(invar) + ) + + +def test_SymmetricConvNeXtBlock_dropout(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + SymmetricConvNeXtBlock, + ) + + in_channels = 2 + latent_channels = 1 + tensor_size = 16 + symmetric_convnextblock = SymmetricConvNeXtBlock( + in_channels=in_channels, + latent_channels=latent_channels, + dropout=0.5, + ).to(device) + + _assert_dropout_present(symmetric_convnextblock) + + invar = test_data(img_size=tensor_size, device=device) + out_shape = torch.Size([12, 1, tensor_size, tensor_size]) + _assert_dropout_eval_deterministic_train_stochastic( + symmetric_convnextblock, invar, out_shape + ) + + +def test_SymmetricConvNeXtBlock_conditional_layer_norm(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + SymmetricConvNeXtBlock, + ) + + in_channels = 2 + latent_channels = 1 + cond_dim = 4 + tensor_size = 16 + conditional_layer_norm = _cln_factory(cond_dim) + + symmetric_convnextblock = SymmetricConvNeXtBlock( + in_channels=in_channels, + latent_channels=latent_channels, + conditional_layer_norm=conditional_layer_norm, + conditional_layer_norm_once=False, + ).to(device) + assert symmetric_convnextblock.cln_enabled + + invar = test_data(img_size=tensor_size, device=device) + out_shape = torch.Size([12, 1, tensor_size, tensor_size]) + + conditions_a = torch.randn(1, cond_dim).to(device) + conditions_b = torch.randn(1, cond_dim).to(device) + + outvar_a = symmetric_convnextblock(invar, conditions_cln=conditions_a) + outvar_b = symmetric_convnextblock(invar, conditions_cln=conditions_b) + + assert outvar_a.shape == out_shape + assert not common.compare_output(outvar_a, outvar_b) + + +@pytest.mark.parametrize("conditional_layer_norm_enabled", [True, False]) +def test_SymmetricConvNeXtBlock_conditional_layer_norm_once( + device, test_data, pytestconfig, conditional_layer_norm_enabled +): + from physicsnemo.models.dlwp_healpix.layers import ( + SymmetricConvNeXtBlock, + ) + from physicsnemo.models.dlwp_healpix.layers.normalization import ( + ConditionalLayerNorm, + ) + + in_channels = 2 + # latent_channels must be > 1: a single-channel LayerNorm always + # normalizes to exactly zero, which would erase the (still learnable) + # affine bias differences produced by the entry norm and mask the + # conditioning signal entirely. + latent_channels = 2 + cond_dim = 4 + tensor_size = 16 + + conditional_layer_norm = ( + _cln_factory(cond_dim) if conditional_layer_norm_enabled else None + ) + + symmetric_convnextblock = SymmetricConvNeXtBlock( + in_channels=in_channels, + latent_channels=latent_channels, + conditional_layer_norm=conditional_layer_norm, + conditional_layer_norm_once=True, + ).to(device) + + invar = test_data(img_size=tensor_size, device=device) + out_shape = torch.Size([12, 1, tensor_size, tensor_size]) + + if conditional_layer_norm_enabled: + assert isinstance(symmetric_convnextblock.entry_norm, ConditionalLayerNorm) + + conditions_a = torch.randn(1, cond_dim).to(device) + conditions_b = torch.randn(1, cond_dim).to(device) + + outvar_a = symmetric_convnextblock(invar, conditions_cln=conditions_a) + outvar_b = symmetric_convnextblock(invar, conditions_cln=conditions_b) + + assert outvar_a.shape == out_shape + assert not common.compare_output(outvar_a, outvar_b) + else: + outvar = symmetric_convnextblock(invar) + assert outvar.shape == out_shape + assert torch.isfinite(outvar).all() + + def test_Multi_SymmetricConvNeXtBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( Multi_SymmetricConvNeXtBlock, @@ -248,6 +586,105 @@ def test_Multi_SymmetricConvNeXtBlock_forward(device, test_data, pytestconfig): assert outvar.shape == out_shape +def test_Multi_SymmetricConvNeXtBlock_n_layers(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + Multi_SymmetricConvNeXtBlock, + SymmetricConvNeXtBlock, + ) + + in_channels = 2 + out_channels = 3 + latent_channels = 1 + n_layers = 3 + tensor_size = 16 + multi_symmetric_convnextblock = Multi_SymmetricConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + latent_channels=latent_channels, + n_layers=n_layers, + ).to(device) + + assert len(multi_symmetric_convnextblock.blocks) == n_layers + for block in multi_symmetric_convnextblock.blocks: + assert isinstance(block, SymmetricConvNeXtBlock) + # only the first block consumes in_channels, all subsequent blocks + # operate on out_channels + first_conv = multi_symmetric_convnextblock.blocks[0].convblock[0].layers[-1] + assert first_conv.in_channels == in_channels + later_conv = multi_symmetric_convnextblock.blocks[1].convblock[0].layers[-1] + assert later_conv.in_channels == out_channels + + invar = test_data(channels=in_channels, img_size=tensor_size, device=device) + out_shape = torch.Size([12, out_channels, tensor_size, tensor_size]) + + outvar = multi_symmetric_convnextblock(invar) + assert outvar.shape == out_shape + + +def test_Multi_SymmetricConvNeXtBlock_dropout(device, test_data, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + Multi_SymmetricConvNeXtBlock, + ) + + in_channels = 2 + latent_channels = 1 + n_layers = 2 + tensor_size = 16 + multi_symmetric_convnextblock = Multi_SymmetricConvNeXtBlock( + in_channels=in_channels, + latent_channels=latent_channels, + n_layers=n_layers, + dropout=0.5, + ).to(device) + + # dropout must be forwarded to every wrapped SymmetricConvNeXtBlock + _assert_dropout_present(multi_symmetric_convnextblock) + + invar = test_data(channels=in_channels, img_size=tensor_size, device=device) + out_shape = torch.Size([12, 1, tensor_size, tensor_size]) + _assert_dropout_eval_deterministic_train_stochastic( + multi_symmetric_convnextblock, invar, out_shape + ) + + +def test_Multi_SymmetricConvNeXtBlock_conditional_layer_norm( + device, test_data, pytestconfig +): + from physicsnemo.models.dlwp_healpix.layers import ( + Multi_SymmetricConvNeXtBlock, + ) + + in_channels = 2 + latent_channels = 1 + cond_dim = 4 + n_layers = 2 + tensor_size = 16 + conditional_layer_norm = _cln_factory(cond_dim) + + multi_symmetric_convnextblock = Multi_SymmetricConvNeXtBlock( + in_channels=in_channels, + latent_channels=latent_channels, + n_layers=n_layers, + conditional_layer_norm=conditional_layer_norm, + ).to(device) + assert multi_symmetric_convnextblock.cln_enabled + # cln must be propagated to every sub-block + for block in multi_symmetric_convnextblock.blocks: + assert block.cln_enabled + + invar = test_data(channels=in_channels, img_size=tensor_size, device=device) + out_shape = torch.Size([12, 1, tensor_size, tensor_size]) + + conditions_a = torch.randn(1, cond_dim).to(device) + conditions_b = torch.randn(1, cond_dim).to(device) + + outvar_a = multi_symmetric_convnextblock(invar, conditions_cln=conditions_a) + outvar_b = multi_symmetric_convnextblock(invar, conditions_cln=conditions_b) + + assert outvar_a.shape == out_shape + assert not common.compare_output(outvar_a, outvar_b) + + def test_BasicConvBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( BasicConvBlock, diff --git a/test/models/dlwp_healpix/test_healpix_encoder_decoder.py b/test/models/dlwp_healpix/test_healpix_encoder_decoder.py index 62622f8c16..93cd1caef3 100644 --- a/test/models/dlwp_healpix/test_healpix_encoder_decoder.py +++ b/test/models/dlwp_healpix/test_healpix_encoder_decoder.py @@ -14,11 +14,35 @@ # See the License for the specific language governing permissions and # limitations under the License. # ruff: noqa: E402 +import omegaconf +import pytest import torch from test import common +def _cln_conv_block(in_channels: int, cond_dim: int) -> omegaconf.DictConfig: + """Hydra-instantiable ``Multi_SymmetricConvNeXtBlock`` config with an + always-on conditional layer norm, used to exercise the ``per_level_cln`` + and ``conditions_cln`` wiring in ``UNetEncoder``/``UNetDecoder``. + + Must be an ``omegaconf.DictConfig`` (not a plain ``dict``) because + ``UNetEncoder``/``UNetDecoder`` access ``block_config.conditional_layer_norm`` + via attribute lookup when deciding whether to disable CLN for a given level. + """ + return omegaconf.DictConfig( + { + "_target_": "physicsnemo.models.dlwp_healpix.layers.Multi_SymmetricConvNeXtBlock", + "in_channels": in_channels, + "conditional_layer_norm": { + "_target_": "physicsnemo.models.dlwp_healpix.layers.ConditionalLayerNorm", + "_partial_": True, + "condition_shape": cond_dim, + }, + } + ) + + def test_UNetEncoder_initialize(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ConvNeXtBlock, UNetEncoder from physicsnemo.nn.module.hpx import HEALPixMaxPool @@ -134,6 +158,275 @@ def test_UNetEncoder_reset(device, pytestconfig): torch.cuda.empty_cache() +def test_UNetEncoder_per_level_cln_length_mismatch(device, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ConvNeXtBlock, UNetEncoder + from physicsnemo.nn.module.hpx import HEALPixMaxPool + + channels = 2 + n_channels = (16, 32, 64) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": channels, + } + down_sampling_block = { + "_target_": HEALPixMaxPool, + "pooling": 2, + } + + with pytest.raises(ValueError, match="per_level_cln must be a list of booleans"): + UNetEncoder( + conv_block=conv_block, + down_sampling_block=down_sampling_block, + n_channels=n_channels, + input_channels=channels, + # wrong length: 2 entries for 3 levels + per_level_cln=[True, False], + ) + + +def test_UNetEncoder_per_level_checkpointing_length_mismatch(device, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ConvNeXtBlock, UNetEncoder + from physicsnemo.nn.module.hpx import HEALPixMaxPool + + channels = 2 + n_channels = (16, 32, 64) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": channels, + } + down_sampling_block = { + "_target_": HEALPixMaxPool, + "pooling": 2, + } + + with pytest.raises( + ValueError, match="per_level_checkpointing must be a list of booleans" + ): + UNetEncoder( + conv_block=conv_block, + down_sampling_block=down_sampling_block, + n_channels=n_channels, + input_channels=channels, + # wrong length: 2 entries for 3 levels + per_level_checkpointing=[True, False], + ) + + +@pytest.mark.parametrize( + "mask", [[True, False, True], [False, True, False]], ids=["tft", "ftf"] +) +def test_UNetEncoder_per_level_cln_disables_expected_levels(device, mask, pytestconfig): + """Verify per_level_cln is honored independently at each level, not applied + uniformly to all levels of the encoder. + """ + from physicsnemo.models.dlwp_healpix.layers import UNetEncoder + from physicsnemo.nn.module.hpx import HEALPixMaxPool + + channels = 8 + cond_dim = 8 + n_cond = 2 + hw_size = 8 + n_channels = (8, 8, 8) + + conv_block = _cln_conv_block(channels, cond_dim) + down_sampling_block = { + "_target_": HEALPixMaxPool, + "pooling": 2, + } + + encoder = UNetEncoder( + conv_block=conv_block, + down_sampling_block=down_sampling_block, + n_channels=n_channels, + input_channels=channels, + per_level_cln=mask, + ).to(device) + + # instantiation-time check: every level's cln_enabled flag must exactly + # match the requested mask at that index + for n in range(len(n_channels)): + conv_module = encoder.encoder[n][-1] + assert conv_module.cln_enabled == mask[n], ( + f"level {n}: expected cln_enabled={mask[n]}, got {conv_module.cln_enabled}" + ) + + # forward-time check: call each level's conv module directly (bypassing + # the encoder's sequential level-to-level dependency, since an earlier + # CLN-sensitive level would otherwise contaminate the input to every + # later level regardless of that later level's own per_level_cln value) + for n in range(len(n_channels)): + conv_module = encoder.encoder[n][-1] + fixed_input = torch.rand([12 * n_cond, channels, hw_size, hw_size]).to(device) + cond_a = torch.randn(n_cond, cond_dim).to(device) + cond_b = torch.randn(n_cond, cond_dim).to(device) + + with torch.no_grad(): + out_a = conv_module(fixed_input, conditions_cln=cond_a) + out_b = conv_module(fixed_input, conditions_cln=cond_b) + + same = common.compare_output(out_a, out_b) + if mask[n]: + assert not same, f"level {n} should be sensitive to conditions_cln" + else: + assert same, f"level {n} should be unaffected by conditions_cln" + + del encoder + torch.cuda.empty_cache() + + +def test_UNetEncoder_forward_conditions_cln_required(device, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import UNetEncoder + from physicsnemo.nn.module.hpx import HEALPixMaxPool + + channels = 4 + cond_dim = 8 + hw_size = 16 + n_channels = (8, 8, 8) + + conv_block = _cln_conv_block(channels, cond_dim) + down_sampling_block = { + "_target_": HEALPixMaxPool, + "pooling": 2, + } + + # default per_level_cln (None) enables CLN at every level + encoder = UNetEncoder( + conv_block=conv_block, + down_sampling_block=down_sampling_block, + n_channels=n_channels, + input_channels=channels, + ).to(device) + + invar = torch.rand([12, channels, hw_size, hw_size]).to(device) + + with pytest.raises(ValueError, match="Conditional inputs are required"): + encoder(invar, conditions_cln=None) + + del encoder, invar + torch.cuda.empty_cache() + + +@pytest.mark.parametrize( + ("mask", "expected_calls"), + [([True, False, True], 2), ([False, True, False], 1)], + ids=["tft-2calls", "ftf-1call"], +) +def test_UNetEncoder_per_level_checkpointing_applied_only_at_configured_levels( + device, monkeypatch, mask, expected_calls, pytestconfig +): + """Verify per_level_checkpointing routes only the configured levels through + torch.utils.checkpoint.checkpoint(), not all (or none) of the encoder levels. + """ + from torch.utils.checkpoint import checkpoint as real_checkpoint + + from physicsnemo.models.dlwp_healpix.layers import ConvNeXtBlock, UNetEncoder + from physicsnemo.nn.module.hpx import HEALPixMaxPool + + channels = 2 + hw_size = 16 + n_channels = (8, 8, 8) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": channels, + } + down_sampling_block = { + "_target_": HEALPixMaxPool, + "pooling": 2, + } + + call_count = {"n": 0} + + def spy_checkpoint(*args, **kwargs): + call_count["n"] += 1 + return real_checkpoint(*args, **kwargs) + + monkeypatch.setattr( + "physicsnemo.models.dlwp_healpix.layers.healpix_encoder.checkpoint", + spy_checkpoint, + ) + + encoder = UNetEncoder( + conv_block=conv_block, + down_sampling_block=down_sampling_block, + n_channels=n_channels, + input_channels=channels, + per_level_checkpointing=mask, + ).to(device) + + invar = torch.rand([12, channels, hw_size, hw_size]).to(device) + encoder(invar) + + assert call_count["n"] == expected_calls, ( + f"mask={mask}: expected {expected_calls} checkpoint() call(s), " + f"got {call_count['n']}" + ) + + del encoder, invar + torch.cuda.empty_cache() + + +def test_UNetEncoder_per_level_checkpointing_matches_no_checkpointing( + device, pytestconfig +): + """A mixed (non-uniform) per_level_checkpointing pattern must produce the + same forward output and gradients as no checkpointing at all, using the + exact same model weights. + """ + from physicsnemo.models.dlwp_healpix.layers import ConvNeXtBlock, UNetEncoder + from physicsnemo.nn.module.hpx import HEALPixMaxPool + + channels = 2 + hw_size = 16 + b_size = 12 + n_channels = (8, 8, 8) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": channels, + } + down_sampling_block = { + "_target_": HEALPixMaxPool, + "pooling": 2, + } + + encoder = UNetEncoder( + conv_block=conv_block, + down_sampling_block=down_sampling_block, + n_channels=n_channels, + input_channels=channels, + per_level_checkpointing=[False, False, False], + ).to(device) + + invar = torch.rand([b_size, channels, hw_size, hw_size]).to(device) + invar.requires_grad_(True) + + baseline = encoder(invar) + sum(o.sum() for o in baseline).backward() + baseline_grad = invar.grad.clone() + baseline_out = [o.detach().clone() for o in baseline] + + invar.grad = None + encoder.zero_grad() + + # mixed pattern applied to the *same* model instance/weights + encoder.per_level_checkpointing = [True, False, True] + + checkpointed = encoder(invar) + sum(o.sum() for o in checkpointed).backward() + + for n in range(len(n_channels)): + assert common.compare_output(baseline_out[n], checkpointed[n]) + + assert torch.isfinite(invar.grad).all() + assert torch.allclose(baseline_grad, invar.grad, rtol=1e-4, atol=1e-5) + + del encoder, invar, baseline, checkpointed + torch.cuda.empty_cache() + + def test_UNetDecoder_initilization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( BasicConvBlock, # for the output layer @@ -377,3 +670,385 @@ def test_UNetDecoder_reset(device, pytestconfig): del decoder, outvar, invars torch.cuda.empty_cache() + + +def test_UNetDecoder_per_level_cln_length_mismatch(device, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + BasicConvBlock, + ConvNeXtBlock, + TransposedConvUpsample, + UNetDecoder, + ) + + in_channels = 2 + out_channels = 1 + n_channels = (64, 32, 16) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": in_channels, + } + up_sampling_block = { + "_target_": TransposedConvUpsample, + "in_channels": in_channels, + "out_channels": out_channels, + "upsampling": 2, + } + output_layer = { + "_target_": BasicConvBlock, + "in_channels": in_channels, + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + + with pytest.raises(ValueError, match="per_level_cln must be a list of booleans"): + UNetDecoder( + conv_block=conv_block, + up_sampling_block=up_sampling_block, + output_layer=output_layer, + recurrent_block=None, + n_channels=n_channels, + # wrong length: 2 entries for 3 levels + per_level_cln=[True, False], + ) + + +def test_UNetDecoder_per_level_checkpointing_length_mismatch(device, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + BasicConvBlock, + ConvNeXtBlock, + TransposedConvUpsample, + UNetDecoder, + ) + + in_channels = 2 + out_channels = 1 + n_channels = (64, 32, 16) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": in_channels, + } + up_sampling_block = { + "_target_": TransposedConvUpsample, + "in_channels": in_channels, + "out_channels": out_channels, + "upsampling": 2, + } + output_layer = { + "_target_": BasicConvBlock, + "in_channels": in_channels, + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + + with pytest.raises( + ValueError, match="per_level_checkpointing must be a list of booleans" + ): + UNetDecoder( + conv_block=conv_block, + up_sampling_block=up_sampling_block, + output_layer=output_layer, + recurrent_block=None, + n_channels=n_channels, + # wrong length: 2 entries for 3 levels + per_level_checkpointing=[True, False], + ) + + +@pytest.mark.parametrize( + "mask", [[True, False, True], [False, True, False]], ids=["tft", "ftf"] +) +def test_UNetDecoder_per_level_cln_disables_expected_levels(device, mask, pytestconfig): + """Verify per_level_cln is honored independently at each level, not applied + uniformly to all levels of the decoder. + """ + from physicsnemo.models.dlwp_healpix.layers import ( + BasicConvBlock, + TransposedConvUpsample, + UNetDecoder, + ) + + cond_dim = 8 + n_cond = 2 + hw_size = 8 + n_channels = (8, 8, 8) + out_channels = 4 + + conv_block = _cln_conv_block(n_channels[0], cond_dim) + up_sampling_block = { + "_target_": TransposedConvUpsample, + "in_channels": n_channels[0], + "out_channels": n_channels[0], + "upsampling": 2, + } + output_layer = { + "_target_": BasicConvBlock, + "in_channels": n_channels[0], + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + + decoder = UNetDecoder( + conv_block=conv_block, + up_sampling_block=up_sampling_block, + output_layer=output_layer, + recurrent_block=None, + n_channels=n_channels, + output_channels=out_channels, + per_level_cln=mask, + ).to(device) + + # instantiation-time check: every level's cln_enabled flag must exactly + # match the requested mask at that index + for n in range(len(n_channels)): + conv_module = decoder.decoder[n]["conv"] + assert conv_module.cln_enabled == mask[n], ( + f"level {n}: expected cln_enabled={mask[n]}, got {conv_module.cln_enabled}" + ) + + # forward-time check: call each level's conv module directly (bypassing + # the decoder's sequential level-to-level dependency, since an earlier + # CLN-sensitive level would otherwise contaminate the input to every + # later level regardless of that later level's own per_level_cln value) + for n in range(len(n_channels)): + conv_module = decoder.decoder[n]["conv"] + # level 0 has no skip-connection concat, so its conv sees n_channels[0] + # input channels; every subsequent level concatenates the upsampled + # skip connection, doubling the channel count + in_ch = n_channels[n] * 2 if n > 0 else n_channels[n] + fixed_input = torch.rand([12 * n_cond, in_ch, hw_size, hw_size]).to(device) + cond_a = torch.randn(n_cond, cond_dim).to(device) + cond_b = torch.randn(n_cond, cond_dim).to(device) + + with torch.no_grad(): + out_a = conv_module(fixed_input, conditions_cln=cond_a) + out_b = conv_module(fixed_input, conditions_cln=cond_b) + + same = common.compare_output(out_a, out_b) + if mask[n]: + assert not same, f"level {n} should be sensitive to conditions_cln" + else: + assert same, f"level {n} should be unaffected by conditions_cln" + + del decoder + torch.cuda.empty_cache() + + +def test_UNetDecoder_forward_conditions_cln_required(device, pytestconfig): + from physicsnemo.models.dlwp_healpix.layers import ( + BasicConvBlock, + TransposedConvUpsample, + UNetDecoder, + ) + + cond_dim = 8 + hw_size = 16 + n_channels = (8, 8, 8) + out_channels = 4 + + conv_block = _cln_conv_block(n_channels[0], cond_dim) + up_sampling_block = { + "_target_": TransposedConvUpsample, + "in_channels": n_channels[0], + "out_channels": n_channels[0], + "upsampling": 2, + } + output_layer = { + "_target_": BasicConvBlock, + "in_channels": n_channels[0], + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + + # default per_level_cln (None) enables CLN at every level + decoder = UNetDecoder( + conv_block=conv_block, + up_sampling_block=up_sampling_block, + output_layer=output_layer, + recurrent_block=None, + n_channels=n_channels, + output_channels=out_channels, + ).to(device) + + invars = [] + size = hw_size + for idx in range(len(n_channels) - 1, -1, -1): + invars.append(torch.rand([12, n_channels[idx], size, size]).to(device)) + size = size // 2 + + with pytest.raises(ValueError, match="Conditional inputs are required"): + decoder(invars, conditions_cln=None) + + del decoder, invars + torch.cuda.empty_cache() + + +@pytest.mark.parametrize( + ("mask", "expected_calls"), + [([True, False, True], 2), ([False, True, False], 1)], + ids=["tft-2calls", "ftf-1call"], +) +def test_UNetDecoder_per_level_checkpointing_applied_only_at_configured_levels( + device, monkeypatch, mask, expected_calls, pytestconfig +): + """Verify per_level_checkpointing routes only the configured levels through + torch.utils.checkpoint.checkpoint(), not all (or none) of the decoder levels. + """ + from torch.utils.checkpoint import checkpoint as real_checkpoint + + from physicsnemo.models.dlwp_healpix.layers import ( + BasicConvBlock, + ConvNeXtBlock, + TransposedConvUpsample, + UNetDecoder, + ) + + in_channels = 2 + out_channels = 1 + hw_size = 16 + n_channels = (8, 8, 8) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": in_channels, + } + up_sampling_block = { + "_target_": TransposedConvUpsample, + "in_channels": in_channels, + "out_channels": in_channels, + "upsampling": 2, + } + output_layer = { + "_target_": BasicConvBlock, + "in_channels": in_channels, + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + + call_count = {"n": 0} + + def spy_checkpoint(*args, **kwargs): + call_count["n"] += 1 + return real_checkpoint(*args, **kwargs) + + monkeypatch.setattr( + "physicsnemo.models.dlwp_healpix.layers.healpix_decoder.checkpoint", + spy_checkpoint, + ) + + decoder = UNetDecoder( + conv_block=conv_block, + up_sampling_block=up_sampling_block, + output_layer=output_layer, + recurrent_block=None, + n_channels=n_channels, + per_level_checkpointing=mask, + ).to(device) + + invars = [] + size = hw_size + for idx in range(len(n_channels) - 1, -1, -1): + invars.append(torch.rand([12, n_channels[idx], size, size]).to(device)) + size = size // 2 + + decoder(invars) + + assert call_count["n"] == expected_calls, ( + f"mask={mask}: expected {expected_calls} checkpoint() call(s), " + f"got {call_count['n']}" + ) + + del decoder, invars + torch.cuda.empty_cache() + + +def test_UNetDecoder_per_level_checkpointing_matches_no_checkpointing( + device, pytestconfig +): + """A mixed (non-uniform) per_level_checkpointing pattern must produce the + same forward output and gradients as no checkpointing at all, using the + exact same model weights. + """ + from physicsnemo.models.dlwp_healpix.layers import ( + BasicConvBlock, + ConvNeXtBlock, + TransposedConvUpsample, + UNetDecoder, + ) + + in_channels = 2 + out_channels = 1 + hw_size = 16 + b_size = 12 + n_channels = (8, 8, 8) + + conv_block = { + "_target_": ConvNeXtBlock, + "in_channels": in_channels, + } + up_sampling_block = { + "_target_": TransposedConvUpsample, + "in_channels": in_channels, + "out_channels": in_channels, + "upsampling": 2, + } + output_layer = { + "_target_": BasicConvBlock, + "in_channels": in_channels, + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + + decoder = UNetDecoder( + conv_block=conv_block, + up_sampling_block=up_sampling_block, + output_layer=output_layer, + recurrent_block=None, + n_channels=n_channels, + per_level_checkpointing=[False, False, False], + ).to(device) + + invars = [] + size = hw_size + for idx in range(len(n_channels) - 1, -1, -1): + invar = torch.rand([b_size, n_channels[idx], size, size]).to(device) + invar.requires_grad_(True) + invars.append(invar) + size = size // 2 + + baseline = decoder(invars) + baseline.sum().backward() + baseline_grads = [v.grad.clone() for v in invars] + baseline_out = baseline.detach().clone() + + for v in invars: + v.grad = None + decoder.zero_grad() + + # mixed pattern applied to the *same* model instance/weights + decoder.per_level_checkpointing = [True, False, True] + + checkpointed = decoder(invars) + checkpointed.sum().backward() + + assert common.compare_output(baseline_out, checkpointed) + + for v, baseline_grad in zip(invars, baseline_grads): + assert torch.isfinite(v.grad).all() + assert torch.allclose(baseline_grad, v.grad, rtol=1e-4, atol=1e-5) + + del decoder, invars, baseline, checkpointed + torch.cuda.empty_cache() diff --git a/test/models/dlwp_healpix/test_healpix_layers.py b/test/models/dlwp_healpix/test_healpix_layers.py index f6ad0135b6..e2e37ce866 100644 --- a/test/models/dlwp_healpix/test_healpix_layers.py +++ b/test/models/dlwp_healpix/test_healpix_layers.py @@ -19,17 +19,21 @@ import torch from test import common +from test.nn.module.healpix_helpers import MulX -class MulX(torch.nn.Module): - """Helper class that just multiplies the values of an input tensor""" +class KwargCapture(torch.nn.Module): + """Helper layer that records the kwargs it was instantiated with, so + tests can verify which kwargs HEALPixLayer strips before forwarding.""" - def __init__(self, multiplier: int = 1): - super(MulX, self).__init__() - self.multiplier = multiplier + captured_kwargs = None + + def __init__(self, **kwargs): + super().__init__() + KwargCapture.captured_kwargs = kwargs def forward(self, x): - return x * self.multiplier + return x HEALPixLayer_testdata = [ @@ -91,3 +95,83 @@ def test_HEALPixLayer_forward(device, multiplier, pytestconfig): del layer, outvar, invar torch.cuda.empty_cache() + + +def test_HEALPixLayer_strips_wrapper_kwargs(device, pytestconfig): + """`enable_nhwc` and `enable_healpixpad` are HEALPixLayer-only knobs and + must not be forwarded to the wrapped layer's constructor.""" + from physicsnemo.nn.module.hpx import HEALPixLayer + + HEALPixLayer( + layer=KwargCapture, + multiplier=5, + enable_nhwc=False, + enable_healpixpad=False, + ) + + assert "enable_nhwc" not in KwargCapture.captured_kwargs + assert "enable_healpixpad" not in KwargCapture.captured_kwargs + assert KwargCapture.captured_kwargs == {"multiplier": 5} + + +def test_HEALPixLayer_no_padding_when_kernel_size_small(device, pytestconfig): + """A conv layer with kernel_size == 1 has no spatial neighborhood to + stitch, so HEALPixLayer must not insert a HEALPixPadding submodule.""" + from physicsnemo.nn.module.hpx import HEALPixLayer + from physicsnemo.nn.module.hpx.padding import HEALPixPadding + + in_channels = 4 + out_channels = 3 + kernel_size = 1 + + layer = HEALPixLayer( + layer=torch.nn.Conv2d, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + device=device, + ) + + assert not any(isinstance(m, HEALPixPadding) for m in layer.layers) + + size = 4 + invar = torch.rand(24, in_channels, size, size, device=device) + outvar = layer(invar) + + assert outvar.shape == (24, out_channels, size, size) + + +def test_HEALPixLayer_conv_disables_native_padding(device, pytestconfig): + """When HEALPixLayer inserts its own HEALPixPadding submodule for a + kernel_size > 1 conv, the wrapped Conv2d's native padding must be + disabled (0) since padding is already applied upstream.""" + from physicsnemo.nn.module.hpx import HEALPixLayer + from physicsnemo.nn.module.hpx.padding import HEALPixPadding + + kernel_size = 3 + in_channels = 4 + out_channels = 3 + + layer = HEALPixLayer( + layer=torch.nn.Conv2d, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + device=device, + ) + + conv_submodules = [m for m in layer.layers if isinstance(m, torch.nn.Conv2d)] + padding_submodules = [m for m in layer.layers if isinstance(m, HEALPixPadding)] + + assert len(conv_submodules) == 1 + assert len(padding_submodules) == 1 + assert conv_submodules[0].padding == (0, 0) + + # kernel_size=3 with dilation=1 needs 1 pixel of context on each side. + assert padding_submodules[0].p == 1 + + size = 4 + invar = torch.rand(24, in_channels, size, size, device=device) + outvar = layer(invar) + # HEALPixPadding restores the spatial size that Conv2d's kernel consumes. + assert outvar.shape == (24, out_channels, size, size) diff --git a/test/models/dlwp_healpix/test_healpix_recunet_model.py b/test/models/dlwp_healpix/test_healpix_recunet_model.py index 3b60c98052..06a90139bd 100644 --- a/test/models/dlwp_healpix/test_healpix_recunet_model.py +++ b/test/models/dlwp_healpix/test_healpix_recunet_model.py @@ -127,6 +127,115 @@ def decoder_dict( return omegaconf.DictConfig(decoder) +@pytest.fixture +def cln_conv_block_dict(channels=4, cond_dim=8): + """``Multi_SymmetricConvNeXtBlock`` config with an always-on conditional + layer norm, used to exercise the ``conditions_cln`` wiring through + ``HEALPixRecUNet.forward``/``_initialize_hidden``. + """ + return omegaconf.DictConfig( + { + "_target_": "physicsnemo.models.dlwp_healpix.layers.Multi_SymmetricConvNeXtBlock", + "in_channels": channels, + "conditional_layer_norm": { + "_target_": "physicsnemo.models.dlwp_healpix.layers.ConditionalLayerNorm", + "_partial_": True, + "condition_shape": cond_dim, + }, + } + ) + + +@pytest.fixture +def cln_recurrent_block_dict(channels=4): + recurrent_block = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.ConvGRUBlock", + "in_channels": channels, + "kernel_size": 1, + "_recursive_": False, + } + return omegaconf.DictConfig(recurrent_block) + + +@pytest.fixture +def cln_up_sampling_block_dict(channels=4): + up_sampling_block = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.TransposedConvUpsample", + "in_channels": channels, + "out_channels": channels, + "activation": {"_target_": "physicsnemo.nn.CappedGELU", "cap_value": 10}, + "upsampling": 2, + } + return omegaconf.DictConfig(up_sampling_block) + + +@pytest.fixture +def cln_output_layer_dict(channels=4, out_channels=2): + output_layer = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.BasicConvBlock", + "in_channels": channels, + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + return omegaconf.DictConfig(output_layer) + + +@pytest.fixture +def cln_encoder_dict( + cln_conv_block_dict, down_sampling_block_dict, cln_recurrent_block_dict +): + """CLN-enabled encoder dict fixture (small, self-contained channel count).""" + encoder = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.UNetEncoder", + "conv_block": cln_conv_block_dict, + "down_sampling_block": down_sampling_block_dict, + "recurrent_block": cln_recurrent_block_dict, + "_recursive_": False, + "n_channels": [4, 4, 4], + "dilations": [1, 2, 4], + } + return encoder + + +@pytest.fixture +def cln_decoder_dict( + cln_conv_block_dict, + cln_up_sampling_block_dict, + cln_output_layer_dict, + cln_recurrent_block_dict, +): + """CLN-enabled decoder dict fixture (small, self-contained channel count).""" + decoder = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.UNetDecoder", + "conv_block": cln_conv_block_dict, + "up_sampling_block": cln_up_sampling_block_dict, + "recurrent_block": cln_recurrent_block_dict, + "output_layer": cln_output_layer_dict, + "_recursive_": False, + "n_channels": [4, 4, 4], + "dilations": [4, 2, 1], + } + return omegaconf.DictConfig(decoder) + + +@pytest.fixture +def coupling_data(): + # create dummy coupling data: a list (one entry per required coupling + # index) of (B, C_coupled, F, H, W) tensors, matching the shape expected + # by ``HEALPixRecUNet._reshape_inputs`` when ``couplings_time_first=True``. + def generate_coupling_data( + steps=1, channels=4, batch_size=8, img_size=16, device="cpu" + ): + return [ + torch.randn(batch_size, channels, 12, img_size, img_size).to(device) + for _ in range(steps) + ] + + return generate_coupling_data + + @pytest.fixture def test_data(): # create dummy data @@ -462,3 +571,621 @@ def test_HEALPixRecUNet_forward( del model, inputs torch.cuda.empty_cache() + + +def test_HEALPixRecUNet_forward_invalid_ndim( + device, encoder_dict, decoder_dict, pytestconfig +): + """``forward`` must reject prognostics that aren't shaped (B, F, T, C, H, W).""" + model = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=2, + output_channels=2, + n_constants=1, + decoder_input_channels=1, + input_time_dim=2, + output_time_dim=2, + delta_time="6h", + reset_cycle="6h", + ).to(device) + + bad_prognostics = torch.randn(2, 12, 2, 16, 16).to(device) + with pytest.raises(ValueError, match="expects prognostics shaped"): + model([bad_prognostics]) + + del model + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixRecUNet_forward_residual_prediction( + device, + encoder_dict, + decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """When ``residual_prediction=True`` prognostic channels must equal the + ``residual_prediction=False`` prognostics plus the raw prognostic input + for that step, while diagnostic channels are unaffected (verified with + identical weights so the only difference is the residual add). + """ + in_channels = 2 + out_channels = 3 # out_channels > in_channels gives a real diagnostic channel + n_constants = 2 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 2 # single integration step + presteps = 1 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + total_steps = presteps + 1 + x = test_data( + batch_size=batch_size, + time_dim=total_steps * input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, + time_dim=total_steps * output_time_dim, + img_size=size, + device=device, + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + model_no_residual = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + presteps=presteps, + delta_time="6h", + reset_cycle="6h", + residual_prediction=False, + ).to(device) + model_residual = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + presteps=presteps, + delta_time="6h", + reset_cycle="6h", + residual_prediction=True, + ).to(device) + model_residual.load_state_dict(model_no_residual.state_dict()) + + out_no_residual = model_no_residual(inputs) + out_residual = model_residual(inputs) + + residual_input = x[ + :, :, presteps * input_time_dim : (presteps + 1) * input_time_dim + ] + expected_prognostics = out_no_residual[:, :, :, :in_channels] + residual_input + assert common.compare_output( + out_residual[:, :, :, :in_channels], expected_prognostics + ) + assert common.compare_output( + out_residual[:, :, :, in_channels:], out_no_residual[:, :, :, in_channels:] + ) + + del model_no_residual, model_residual, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixRecUNet_forward_constraints( + device, + encoder_dict, + decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """``set_constraints`` should apply the configured constraint modules to + the output only for the targeted channel(s), leaving other channels + unaffected (verified with identical weights to an unconstrained model). + + Uses a single integration step: with more than one step, the + recurrent/autoregressive design feeds each step's (possibly clamped) + output back in as the next step's input, so clamping one channel would + legitimately also perturb the other channel's *later* outputs. + """ + in_channels = 2 + out_channels = 2 + n_constants = 2 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 2 + presteps = 1 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + total_steps = presteps + 1 + x = test_data( + batch_size=batch_size, + time_dim=total_steps * input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, + time_dim=total_steps * output_time_dim, + img_size=size, + device=device, + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + constraints = omegaconf.OmegaConf.create( + { + "non_negative": { + "_target_": "physicsnemo.models.dlwp_healpix.layers.NonnegativeConstraint", + "variables": ["var0"], + "channels": ["var0", "var1"], + "scaling": {"var0": {"mean": 0.0, "std": 1.0}}, + } + } + ) + + model_unconstrained = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + presteps=presteps, + delta_time="6h", + reset_cycle="6h", + ).to(device) + model_constrained = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + presteps=presteps, + delta_time="6h", + reset_cycle="6h", + constraints=constraints, + ).to(device) + model_constrained.load_state_dict(model_unconstrained.state_dict()) + + out_unconstrained = model_unconstrained(inputs) + out_constrained = model_constrained(inputs) + + assert out_unconstrained[:, :, :, 0].min() < 0 + assert torch.all(out_constrained[:, :, :, 0] >= 0) + assert common.compare_output( + out_unconstrained[:, :, :, 1], out_constrained[:, :, :, 1] + ) + + del model_unconstrained, model_constrained, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixRecUNet_forward_conditions_cln( + device, + cln_encoder_dict, + cln_decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """``conditions_cln`` must be required when the encoder/decoder are + CLN-enabled, and different per-step conditions must yield different + outputs (exercising both the warm-up path in ``_initialize_hidden`` and + the main integration loop in ``forward``). + """ + channels = 4 + out_channels = 2 + n_constants = 1 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 2 + presteps = 1 + cond_dim = 8 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + total_steps = presteps + 1 + x = test_data( + batch_size=batch_size, + time_dim=total_steps * input_time_dim, + channels=channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, + time_dim=total_steps * output_time_dim, + img_size=size, + device=device, + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + model = HEALPixRecUNet( + encoder=cln_encoder_dict, + decoder=cln_decoder_dict, + input_channels=channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + presteps=presteps, + delta_time="6h", + reset_cycle="6h", + # residual_prediction assumes input_channels <= output_channels + # (prognostics map onto a leading slice of the output channels); + # unrelated to what this test targets, so disable it here. + residual_prediction=False, + ).to(device) + + with pytest.raises(ValueError, match="Conditional inputs are required"): + model(inputs) + + conditions_a = [ + torch.randn(batch_size, cond_dim).to(device) + for _ in range(model.integration_steps) + ] + conditions_b = [ + torch.randn(batch_size, cond_dim).to(device) + for _ in range(model.integration_steps) + ] + + out_a = model(inputs, conditions_cln=conditions_a) + out_b = model(inputs, conditions_cln=conditions_b) + + assert not common.compare_output(out_a, out_b) + + del model, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixRecUNet_forward_couplings( + device, + encoder_dict, + decoder_dict, + test_data, + insolation_data, + constant_data, + coupling_data, + pytestconfig, +): + """A valid, non-empty ``couplings`` configuration should be accepted and + routed through ``_reshape_inputs``, ``_initialize_hidden`` and + ``forward`` without error, producing the expected output shape. Uses two + integration steps so both the warm-up (``step < presteps``) and + steady-state (``step >= presteps``) coupling branches in + ``_initialize_hidden`` are exercised. + """ + in_channels = 2 + out_channels = 2 + n_constants = 2 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 4 # 2 integration steps + presteps = 1 + batch_size = 2 + size = 16 + + couplings = [ + {"params": {"variables": ["v1", "v2"], "input_times": ["24h", "48h"]}}, + ] + + fix_random_seeds(seed=42) + model = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + presteps=presteps, + delta_time="6h", + reset_cycle="6h", + couplings=couplings, + ).to(device) + + assert model.coupled_channels == 4 + + total_steps = presteps + model.integration_steps + x = test_data( + batch_size=batch_size, + time_dim=total_steps * input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, + time_dim=total_steps * output_time_dim, + img_size=size, + device=device, + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + couplings_tensor = coupling_data( + steps=total_steps, + channels=model.coupled_channels, + batch_size=batch_size, + img_size=size, + device=device, + ) + inputs = [x, decoder_inputs, constants, couplings_tensor] + + output = model(inputs) + expected_shape = [batch_size, 12, output_time_dim, out_channels, size, size] + assert list(output.shape) == expected_shape + + del model, inputs + torch.cuda.empty_cache() + + +def test_HEALPixRecUNet_forward_is_diagnostic( + device, + encoder_dict, + decoder_dict, + test_data, + constant_data, + pytestconfig, +): + """When ``output_time_dim == 1`` and ``input_time_dim > 1`` the model + should switch to diagnostic mode: the usual + ``output_time_dim % input_time_dim == 0`` requirement is bypassed and the + output has a single time step regardless of ``input_time_dim``. + """ + in_channels = 2 + out_channels = 2 + n_constants = 1 + input_time_dim = 2 + output_time_dim = 1 + presteps = 1 + batch_size = 2 + size = 16 + + model = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=0, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + presteps=presteps, + delta_time="6h", + reset_cycle="6h", + # residual_prediction reshapes the input residual with the same + # (1 if is_diagnostic else input_time_dim) time factor as the + # decoder output, which is incompatible with is_diagnostic mode + # unless in_channels * input_time_dim == out_channels; disable it + # here since this test targets the diagnostic-mode output shape. + residual_prediction=False, + ).to(device) + + assert model.is_diagnostic + assert model.integration_steps == 1 + + total_steps = presteps + model.integration_steps + x = test_data( + batch_size=batch_size, + time_dim=total_steps * input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, constants] + + with torch.no_grad(): + output = model(inputs) + + expected_shape = [batch_size, 12, output_time_dim, out_channels, size, size] + assert list(output.shape) == expected_shape + + del model, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixRecUNet_reshape_inputs_enable_nhwc( + device, + encoder_dict, + decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """With ``enable_nhwc=True``, ``_reshape_inputs`` should return a tensor + in channels-last memory format. + """ + in_channels = 2 + out_channels = 2 + n_constants = 2 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 2 + size = 16 + batch_size = 2 + + model = HEALPixRecUNet( + encoder=encoder_dict, + decoder=decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + delta_time="6h", + reset_cycle="6h", + enable_nhwc=True, + ).to(device) + + x = test_data( + batch_size=batch_size, + time_dim=input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, time_dim=input_time_dim, img_size=size, device=device + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + + reshaped = model._reshape_inputs([x, decoder_inputs, constants], step=0) + assert reshaped.is_contiguous(memory_format=torch.channels_last) + + del model + torch.cuda.empty_cache() + + +def test_HEALPixRecUNet_backward_compat_arg_mapper(): + """Legacy (version ``0.1.0``) checkpoints used ``dlwp_healpix_layers`` + Hydra targets; ``_backward_compat_arg_mapper`` must remap them to the + current ``physicsnemo.models.dlwp_healpix.layers`` module, and must be a + no-op for any other (current) version. + """ + legacy_args = { + "encoder": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_encoder.UNetEncoder", + "conv_block": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.ConvNeXtBlock", + }, + }, + "decoder": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_decoder.UNetDecoder", + }, + "input_channels": 3, + } + + remapped = HEALPixRecUNet._backward_compat_arg_mapper("0.1.0", legacy_args) + assert ( + remapped["encoder"]["_target_"] + == "physicsnemo.models.dlwp_healpix.layers.UNetEncoder" + ) + assert ( + remapped["decoder"]["_target_"] + == "physicsnemo.models.dlwp_healpix.layers.UNetDecoder" + ) + assert ( + remapped["encoder"]["conv_block"]["_target_"] + == "physicsnemo.models.dlwp_healpix.layers.ConvNeXtBlock" + ) + assert remapped["input_channels"] == 3 + + # any other version is a no-op (the base class's default implementation) + unchanged = HEALPixRecUNet._backward_compat_arg_mapper("9.9.9", legacy_args) + assert unchanged == legacy_args + + +@torch.no_grad() +def test_HEALPixRecUNet_checkpoint( + device, + encoder_dict, + decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """MOD-008c: save one ``HEALPixRecUNet`` and restore its state into a + second, differently-initialized model via both ``Module.load`` and + ``Module.from_checkpoint``, verifying the forward outputs match.""" + in_channels = 2 + out_channels = 2 + n_constants = 2 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 4 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + x = test_data( + batch_size=batch_size, + time_dim=2 * input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, + time_dim=2 * output_time_dim, + img_size=size, + device=device, + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + # ``Module.save`` serializes the captured init args to JSON, so pass plain + # containers rather than ``DictConfig`` (which is not JSON-serializable). + encoder = omegaconf.OmegaConf.to_container( + omegaconf.OmegaConf.create(encoder_dict), resolve=True + ) + decoder = omegaconf.OmegaConf.to_container( + omegaconf.OmegaConf.create(decoder_dict), resolve=True + ) + + def build_model(): + return HEALPixRecUNet( + encoder=encoder, + decoder=decoder, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + delta_time="6h", + reset_cycle="6h", + ).to(device) + + model_1 = build_model() + model_2 = build_model() + + # Perturb model_2's weights so its output differs from model_1 before the + # checkpoint is loaded (validate_checkpoint asserts the initial mismatch). + with torch.no_grad(): + for param in model_2.parameters(): + param.add_(0.1) + + assert common.validate_checkpoint(model_1, model_2, (inputs,), rtol=1e-2) + + del model_1, model_2, inputs + torch.cuda.empty_cache() diff --git a/test/models/dlwp_healpix/test_healpix_unet_model.py b/test/models/dlwp_healpix/test_healpix_unet_model.py index 73b76ff686..b2b1790a3c 100644 --- a/test/models/dlwp_healpix/test_healpix_unet_model.py +++ b/test/models/dlwp_healpix/test_healpix_unet_model.py @@ -117,6 +117,97 @@ def generate_insolation_data(batch_size=8, time_dim=1, img_size=16, device="cpu" return generate_insolation_data +@pytest.fixture +def cln_conv_block_dict(channels=4, cond_dim=8): + """``Multi_SymmetricConvNeXtBlock`` config with an always-on conditional + layer norm, used to exercise the ``conditions_cln`` wiring through + ``HEALPixUNet.forward``. + """ + return omegaconf.DictConfig( + { + "_target_": "physicsnemo.models.dlwp_healpix.layers.Multi_SymmetricConvNeXtBlock", + "in_channels": channels, + "conditional_layer_norm": { + "_target_": "physicsnemo.models.dlwp_healpix.layers.ConditionalLayerNorm", + "_partial_": True, + "condition_shape": cond_dim, + }, + } + ) + + +@pytest.fixture +def cln_up_sampling_block_dict(channels=4): + up_sampling_block = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.TransposedConvUpsample", + "in_channels": channels, + "out_channels": channels, + "activation": {"_target_": "physicsnemo.nn.CappedGELU", "cap_value": 10}, + "upsampling": 2, + } + return omegaconf.DictConfig(up_sampling_block) + + +@pytest.fixture +def cln_output_layer_dict(channels=4, out_channels=2): + output_layer = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.BasicConvBlock", + "in_channels": channels, + "out_channels": out_channels, + "kernel_size": 1, + "dilation": 1, + "n_layers": 1, + } + return omegaconf.DictConfig(output_layer) + + +@pytest.fixture +def cln_unet_encoder_dict(cln_conv_block_dict, down_sampling_block_dict): + """CLN-enabled encoder dict fixture (small, self-contained channel count).""" + encoder = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.UNetEncoder", + "conv_block": cln_conv_block_dict, + "down_sampling_block": down_sampling_block_dict, + "_recursive_": False, + "n_channels": [4, 4, 4], + "dilations": [1, 2, 4], + } + return encoder + + +@pytest.fixture +def cln_unet_decoder_dict( + cln_conv_block_dict, cln_up_sampling_block_dict, cln_output_layer_dict +): + """CLN-enabled decoder dict fixture (small, self-contained channel count).""" + decoder = { + "_target_": "physicsnemo.models.dlwp_healpix.layers.UNetDecoder", + "conv_block": cln_conv_block_dict, + "up_sampling_block": cln_up_sampling_block_dict, + "output_layer": cln_output_layer_dict, + "_recursive_": False, + "n_channels": [4, 4, 4], + "dilations": [4, 2, 1], + } + return omegaconf.DictConfig(decoder) + + +@pytest.fixture +def coupling_data(): + # create dummy coupling data: a list (one entry per required coupling + # index) of (B, C_coupled, F, H, W) tensors, matching the shape expected + # by ``HEALPixUNet._reshape_inputs`` when ``couplings_time_first=True``. + def generate_coupling_data( + steps=1, channels=4, batch_size=8, img_size=16, device="cpu" + ): + return [ + torch.randn(batch_size, channels, 12, img_size, img_size).to(device) + for _ in range(steps) + ] + + return generate_coupling_data + + @pytest.fixture def unet_encoder_dict(conv_next_block_dict, down_sampling_block_dict): """Encoder dict fixture.""" @@ -375,3 +466,542 @@ def test_HEALPixUNet_forward( del inputs, model torch.cuda.empty_cache() + + +def test_HEALPixUNet_forward_invalid_ndim( + device, unet_encoder_dict, unet_decoder_dict, pytestconfig +): + """``forward`` must reject prognostics that aren't shaped (B, F, T, C, H, W).""" + model = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=2, + output_channels=2, + n_constants=1, + decoder_input_channels=1, + input_time_dim=2, + output_time_dim=2, + ).to(device) + + bad_prognostics = torch.randn(2, 12, 2, 16, 16).to(device) + with pytest.raises(ValueError, match="expects prognostics shaped"): + model([bad_prognostics]) + + del model + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixUNet_forward_output_only_last( + device, + unet_encoder_dict, + unet_decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """With ``output_only_last=True`` only the final integration step's + outputs should be returned, rather than all steps concatenated. + """ + in_channels = 2 + out_channels = 2 + n_constants = 1 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 4 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + x = test_data( + batch_size=batch_size, + time_dim=input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, time_dim=output_time_dim, img_size=size, device=device + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + model = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ).to(device) + + full_output = model(inputs) + last_output = model(inputs, output_only_last=True) + + expected_shape = [batch_size, 12, input_time_dim, out_channels, size, size] + assert list(last_output.shape) == expected_shape + assert common.compare_output(last_output, full_output[:, :, -input_time_dim:, ...]) + + del model, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixUNet_forward_residual_prediction( + device, + unet_encoder_dict, + unet_decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """When ``residual_prediction=True`` the model output must equal the + ``residual_prediction=False`` output plus the raw input tensor (verified + with identical weights, so the only difference is the residual add). + Requires ``in_channels == out_channels`` and a single integration step, + since ``HEALPixUNet`` adds the residual to the full (unsplit) decodings. + """ + channels = 2 + n_constants = 1 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 2 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + x = test_data( + batch_size=batch_size, + time_dim=input_time_dim, + channels=channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, time_dim=output_time_dim, img_size=size, device=device + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + model_no_residual = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=channels, + output_channels=channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + residual_prediction=False, + ).to(device) + model_residual = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=channels, + output_channels=channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + residual_prediction=True, + ).to(device) + model_residual.load_state_dict(model_no_residual.state_dict()) + + out_no_residual = model_no_residual(inputs) + out_residual = model_residual(inputs) + + assert common.compare_output(out_residual, out_no_residual + x) + + del model_no_residual, model_residual, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixUNet_forward_constraints( + device, + unet_encoder_dict, + unet_decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """``set_constraints`` should apply the configured constraint modules to + the output only for the targeted channel(s), leaving other channels + unaffected (verified with identical weights to an unconstrained model). + """ + in_channels = 2 + out_channels = 2 + n_constants = 1 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 2 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + x = test_data( + batch_size=batch_size, + time_dim=input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, time_dim=output_time_dim, img_size=size, device=device + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + constraints = omegaconf.OmegaConf.create( + { + "non_negative": { + "_target_": "physicsnemo.models.dlwp_healpix.layers.NonnegativeConstraint", + "variables": ["var0"], + "channels": ["var0", "var1"], + "scaling": {"var0": {"mean": 0.0, "std": 1.0}}, + } + } + ) + + model_unconstrained = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ).to(device) + model_constrained = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + constraints=constraints, + ).to(device) + model_constrained.load_state_dict(model_unconstrained.state_dict()) + + out_unconstrained = model_unconstrained(inputs) + out_constrained = model_constrained(inputs) + + # unconstrained "var0" should have some negative values to make this a + # meaningful test of the clamp + assert out_unconstrained[:, :, :, 0].min() < 0 + assert torch.all(out_constrained[:, :, :, 0] >= 0) + assert common.compare_output( + out_unconstrained[:, :, :, 1], out_constrained[:, :, :, 1] + ) + + del model_unconstrained, model_constrained, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixUNet_forward_conditions_cln( + device, + cln_unet_encoder_dict, + cln_unet_decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """``conditions_cln`` must be required when the encoder/decoder are + CLN-enabled, and different conditions must yield different outputs. + """ + channels = 4 + out_channels = 2 + n_constants = 1 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 2 + cond_dim = 8 + batch_size = 2 + size = 16 + + fix_random_seeds(seed=42) + x = test_data( + batch_size=batch_size, + time_dim=input_time_dim, + channels=channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, time_dim=output_time_dim, img_size=size, device=device + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + model = HEALPixUNet( + encoder=cln_unet_encoder_dict, + decoder=cln_unet_decoder_dict, + input_channels=channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ).to(device) + + with pytest.raises(ValueError, match="Conditional inputs are required"): + model(inputs) + + conditions_a = [ + torch.randn(batch_size, cond_dim).to(device) + for _ in range(model.integration_steps) + ] + conditions_b = [ + torch.randn(batch_size, cond_dim).to(device) + for _ in range(model.integration_steps) + ] + + out_a = model(inputs, conditions_cln=conditions_a) + out_b = model(inputs, conditions_cln=conditions_b) + + assert not common.compare_output(out_a, out_b) + + del model, inputs + torch.cuda.empty_cache() + + +@torch.no_grad() +def test_HEALPixUNet_forward_couplings( + device, + unet_encoder_dict, + unet_decoder_dict, + test_data, + insolation_data, + constant_data, + coupling_data, + pytestconfig, +): + """A valid, non-empty ``couplings`` configuration should be accepted and + routed through ``_reshape_inputs``/``forward`` without error, producing + the expected output shape. + """ + in_channels = 3 + out_channels = 3 + n_constants = 2 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 4 + batch_size = 2 + size = 16 + + couplings = [ + {"params": {"variables": ["v1", "v2"], "input_times": ["24h", "48h"]}}, + ] + + fix_random_seeds(seed=42) + model = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + couplings=couplings, + ).to(device) + + assert model.coupled_channels == 4 + + x = test_data( + batch_size=batch_size, + time_dim=input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + decoder_inputs = insolation_data( + batch_size=batch_size, time_dim=output_time_dim, img_size=size, device=device + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + couplings_tensor = coupling_data( + steps=model.integration_steps, + channels=model.coupled_channels, + batch_size=batch_size, + img_size=size, + device=device, + ) + inputs = [x, decoder_inputs, constants, couplings_tensor] + + output = model(inputs) + expected_shape = [batch_size, 12, output_time_dim, out_channels, size, size] + assert list(output.shape) == expected_shape + + del model, inputs + torch.cuda.empty_cache() + + +def test_HEALPixUNet_forward_is_diagnostic( + device, + unet_encoder_dict, + unet_decoder_dict, + test_data, + constant_data, + pytestconfig, +): + """When ``output_time_dim == 1`` and ``input_time_dim > 1`` the model + should switch to diagnostic mode: the usual + ``output_time_dim % input_time_dim == 0`` requirement is bypassed and the + output has a single time step regardless of ``input_time_dim``. + """ + in_channels = 2 + out_channels = 2 + n_constants = 1 + input_time_dim = 2 + output_time_dim = 1 + batch_size = 2 + size = 16 + + model = HEALPixUNet( + encoder=unet_encoder_dict, + decoder=unet_decoder_dict, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=0, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ).to(device) + + assert model.is_diagnostic + assert model.integration_steps == 1 + + x = test_data( + batch_size=batch_size, + time_dim=input_time_dim, + channels=in_channels, + img_size=size, + device=device, + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, constants] + + with torch.no_grad(): + output = model(inputs) + + expected_shape = [batch_size, 12, output_time_dim, out_channels, size, size] + assert list(output.shape) == expected_shape + + del model, inputs + torch.cuda.empty_cache() + + +def test_HEALPixUNet_backward_compat_arg_mapper(): + """Legacy (version ``0.1.0``) checkpoints used ``dlwp_healpix_layers`` + Hydra targets; ``_backward_compat_arg_mapper`` must remap them to the + current ``physicsnemo.models.dlwp_healpix.layers`` module, and must be a + no-op for any other (current) version. + """ + legacy_args = { + "encoder": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_encoder.UNetEncoder", + "conv_block": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.ConvNeXtBlock", + }, + }, + "decoder": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_decoder.UNetDecoder", + }, + "input_channels": 3, + } + + remapped = HEALPixUNet._backward_compat_arg_mapper("0.1.0", legacy_args) + assert ( + remapped["encoder"]["_target_"] + == "physicsnemo.models.dlwp_healpix.layers.UNetEncoder" + ) + assert ( + remapped["decoder"]["_target_"] + == "physicsnemo.models.dlwp_healpix.layers.UNetDecoder" + ) + assert ( + remapped["encoder"]["conv_block"]["_target_"] + == "physicsnemo.models.dlwp_healpix.layers.ConvNeXtBlock" + ) + assert remapped["input_channels"] == 3 + + # any other version is a no-op (the base class's default implementation) + unchanged = HEALPixUNet._backward_compat_arg_mapper("9.9.9", legacy_args) + assert unchanged == legacy_args + + +@torch.no_grad() +def test_HEALPixUNet_checkpoint( + device, + unet_encoder_dict, + unet_decoder_dict, + test_data, + insolation_data, + constant_data, + pytestconfig, +): + """MOD-008c: save one ``HEALPixUNet`` and restore its state into a second, + differently-initialized model via both ``Module.load`` and + ``Module.from_checkpoint``, verifying the forward outputs match.""" + in_channels = 3 + out_channels = 3 + n_constants = 2 + decoder_input_channels = 1 + input_time_dim = 2 + output_time_dim = 4 + size = 16 + + fix_random_seeds(seed=42) + x = test_data( + time_dim=input_time_dim, channels=in_channels, img_size=size, device=device + ) + decoder_inputs = insolation_data( + time_dim=output_time_dim, img_size=size, device=device + ) + constants = constant_data(channels=n_constants, img_size=size, device=device) + inputs = [x, decoder_inputs, constants] + + # ``Module.save`` serializes the captured init args to JSON, so pass plain + # containers rather than ``DictConfig`` (which is not JSON-serializable). + encoder = omegaconf.OmegaConf.to_container( + omegaconf.OmegaConf.create(unet_encoder_dict), resolve=True + ) + decoder = omegaconf.OmegaConf.to_container( + omegaconf.OmegaConf.create(unet_decoder_dict), resolve=True + ) + + def build_model(): + return HEALPixUNet( + encoder=encoder, + decoder=decoder, + input_channels=in_channels, + output_channels=out_channels, + n_constants=n_constants, + decoder_input_channels=decoder_input_channels, + input_time_dim=input_time_dim, + output_time_dim=output_time_dim, + ).to(device) + + model_1 = build_model() + model_2 = build_model() + + # Perturb model_2's weights so its output differs from model_1 before the + # checkpoint is loaded (validate_checkpoint asserts the initial mismatch). + with torch.no_grad(): + for param in model_2.parameters(): + param.add_(0.1) + + assert common.validate_checkpoint(model_1, model_2, (inputs,), rtol=1e-2) + + del model_1, model_2, inputs + torch.cuda.empty_cache() diff --git a/test/models/dlwp_healpix/test_layers_backward_compat.py b/test/models/dlwp_healpix/test_layers_backward_compat.py new file mode 100644 index 0000000000..83adc0fc40 --- /dev/null +++ b/test/models/dlwp_healpix/test_layers_backward_compat.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the legacy Hydra ``_target_`` remapping helpers in +``physicsnemo.models.dlwp_healpix.layers`` (``_remap_target``, ``_remap_obj``), +used by ``HEALPixUNet``/``HEALPixRecUNet`` to load ``0.1.0``-era checkpoints. +""" + +import omegaconf +import pytest + +from physicsnemo.models.dlwp_healpix.layers import ( + _legacy_hydra_targets_warning, + _remap_obj, + _remap_target, +) + + +def test_legacy_hydra_targets_warning_is_nonempty_string(): + assert isinstance(_legacy_hydra_targets_warning, str) + assert len(_legacy_hydra_targets_warning) > 0 + + +@pytest.mark.parametrize( + "target,expected", + [ + # explicit dict entries take precedence over the generic + # `healpix_encoder.`/`healpix_decoder.` prefix handling below + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_encoder.UNetEncoder", + "physicsnemo.models.dlwp_healpix.layers.UNetEncoder", + ), + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_decoder.UNetDecoder", + "physicsnemo.models.dlwp_healpix.layers.UNetDecoder", + ), + # `dlwp_healpix_layers.healpix_blocks.` -> generic class + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.ConvNeXtBlock", + "physicsnemo.models.dlwp_healpix.layers.ConvNeXtBlock", + ), + # `dlwp_healpix_layers.healpix_blocks.` -> AvgPool/MaxPool special-cased + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.AvgPool", + "physicsnemo.nn.HEALPixAvgPool", + ), + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.MaxPool", + "physicsnemo.nn.HEALPixMaxPool", + ), + # `dlwp_healpix_layers.healpix_encoder.`/`healpix_decoder.` with a + # non-explicit class name still maps generically + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_encoder.SomeHelper", + "physicsnemo.models.dlwp_healpix.layers.SomeHelper", + ), + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_decoder.SomeHelper", + "physicsnemo.models.dlwp_healpix.layers.SomeHelper", + ), + # `dlwp_healpix_layers.healpix_layers.` -> physicsnemo.nn + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_layers.HEALPixLayer", + "physicsnemo.nn.HEALPixLayer", + ), + # `dlwp_healpix_layers.normalization.` -> layers.normalization + ( + "physicsnemo.models.dlwp_healpix_layers.normalization.ConditionalLayerNorm", + "physicsnemo.models.dlwp_healpix.layers.normalization.ConditionalLayerNorm", + ), + # `dlwp_healpix_layers.healpix_constraints.` -> layers.healpix_constraints + ( + "physicsnemo.models.dlwp_healpix_layers.healpix_constraints.NonnegativeConstraint", + "physicsnemo.models.dlwp_healpix.layers.healpix_constraints.NonnegativeConstraint", + ), + # bare `dlwp_healpix_layers.` fallback, generic class + ( + "physicsnemo.models.dlwp_healpix_layers.SomeOtherClass", + "physicsnemo.models.dlwp_healpix.layers.SomeOtherClass", + ), + # bare `dlwp_healpix_layers.` fallback, AvgPool/MaxPool special-cased + ( + "physicsnemo.models.dlwp_healpix_layers.AvgPool", + "physicsnemo.nn.HEALPixAvgPool", + ), + ( + "physicsnemo.models.dlwp_healpix_layers.MaxPool", + "physicsnemo.nn.HEALPixMaxPool", + ), + # bare `dlwp_healpix_layers.` fallback, class name starting with + # "HEALPix" maps to physicsnemo.nn + ( + "physicsnemo.models.dlwp_healpix_layers.HEALPixFoldFaces", + "physicsnemo.nn.HEALPixFoldFaces", + ), + # new-style `dlwp_healpix.layers.healpix_blocks.` still gets + # normalized to the top-level `layers` module, generic class + ( + "physicsnemo.models.dlwp_healpix.layers.healpix_blocks.ConvNeXtBlock", + "physicsnemo.models.dlwp_healpix.layers.ConvNeXtBlock", + ), + # new-style `dlwp_healpix.layers.healpix_blocks.` AvgPool/MaxPool + ( + "physicsnemo.models.dlwp_healpix.layers.healpix_blocks.AvgPool", + "physicsnemo.nn.HEALPixAvgPool", + ), + ( + "physicsnemo.models.dlwp_healpix.layers.healpix_blocks.MaxPool", + "physicsnemo.nn.HEALPixMaxPool", + ), + # new-style `dlwp_healpix.layers.healpix_encoder.`/`healpix_decoder.` + ( + "physicsnemo.models.dlwp_healpix.layers.healpix_encoder.UNetEncoder", + "physicsnemo.models.dlwp_healpix.layers.UNetEncoder", + ), + ( + "physicsnemo.models.dlwp_healpix.layers.healpix_decoder.UNetDecoder", + "physicsnemo.models.dlwp_healpix.layers.UNetDecoder", + ), + # legacy activations module + ( + "physicsnemo.models.layers.activations.CappedGELU", + "physicsnemo.nn.CappedGELU", + ), + # anything unrecognized passes through unchanged + ( + "physicsnemo.models.dlwp_healpix.layers.UNetEncoder", + "physicsnemo.models.dlwp_healpix.layers.UNetEncoder", + ), + ( + "some.unrelated.module.Class", + "some.unrelated.module.Class", + ), + ], +) +def test_remap_target(target, expected): + assert _remap_target(target) == expected + + +def test_remap_obj_passthrough_for_scalars(): + assert _remap_obj(3) == 3 + assert _remap_obj("plain string") == "plain string" + assert _remap_obj(None) is None + + +def test_remap_obj_dict_remaps_target_key_only(): + obj = { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.ConvNeXtBlock", + "in_channels": 3, + "nested": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.normalization.ConditionalLayerNorm", + "channel_depth": 4, + }, + } + + remapped = _remap_obj(obj) + + assert ( + remapped["_target_"] == "physicsnemo.models.dlwp_healpix.layers.ConvNeXtBlock" + ) + # non-`_target_` keys are recursed into but otherwise left untouched + assert remapped["in_channels"] == 3 + assert ( + remapped["nested"]["_target_"] + == "physicsnemo.models.dlwp_healpix.layers.normalization.ConditionalLayerNorm" + ) + assert remapped["nested"]["channel_depth"] == 4 + + +def test_remap_obj_ignores_non_string_target_value(): + # a `_target_` key whose value isn't a string (e.g. already resolved to + # a class object) must be recursed into rather than passed to + # `_remap_target`, which only accepts strings + obj = {"_target_": 42} + assert _remap_obj(obj) == {"_target_": 42} + + +def test_remap_obj_list_remaps_each_element(): + obj = [ + {"_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.AvgPool"}, + {"_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.MaxPool"}, + "unchanged", + ] + + remapped = _remap_obj(obj) + + assert remapped[0]["_target_"] == "physicsnemo.nn.HEALPixAvgPool" + assert remapped[1]["_target_"] == "physicsnemo.nn.HEALPixMaxPool" + assert remapped[2] == "unchanged" + + +def test_remap_obj_dictconfig_roundtrips_through_container(): + cfg = omegaconf.OmegaConf.create( + { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_encoder.UNetEncoder", + "conv_block": { + "_target_": "physicsnemo.models.dlwp_healpix_layers.healpix_blocks.ConvNeXtBlock", + }, + } + ) + + remapped = _remap_obj(cfg) + + assert isinstance(remapped, omegaconf.DictConfig) + assert remapped._target_ == "physicsnemo.models.dlwp_healpix.layers.UNetEncoder" + assert ( + remapped.conv_block._target_ + == "physicsnemo.models.dlwp_healpix.layers.ConvNeXtBlock" + ) diff --git a/test/nn/module/healpix_helpers.py b/test/nn/module/healpix_helpers.py new file mode 100644 index 0000000000..08d2dcb5f3 --- /dev/null +++ b/test/nn/module/healpix_helpers.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared helpers for HEALPix layer tests (``test/nn/module/test_healpix.py`` +and ``test/models/dlwp_healpix/test_healpix_layers.py``).""" + +import torch + +Tensor = torch.Tensor + + +class MulX(torch.nn.Module): + """Helper layer that just multiplies the values of an input tensor. + + Used as a minimal, non-convolutional stand-in for the ``layer`` argument + of ``HEALPixLayer`` so wrapper behavior can be tested independently of + any real convolution. + """ + + def __init__(self, multiplier: int = 1): + super().__init__() + self.multiplier = multiplier + + def forward(self, x: Tensor) -> Tensor: + return x * self.multiplier + + +def distinct_face_tensor(size: int, channels: int = 1, device: str = "cpu") -> Tensor: + """Build a folded ``(12, C, H, W)`` tensor where every pixel encodes its + own ``(face, row, col)``, so HEALPix face-stitching logic can be checked + against independently-computed expected neighbor slices. + + Parameters + ---------- + size : int + Height/width of each (square) face. + channels : int, optional + Number of channels to replicate the pattern across, by default 1. + device : str, optional + Device to allocate the tensor on, by default "cpu". + + Returns + ------- + Tensor + Tensor of shape ``(12, channels, size, size)`` with values + ``face * 1000 + row * 10 + col``. + """ + faces = 12 + face_idx = torch.arange(faces, device=device).view(faces, 1, 1, 1) + row_idx = torch.arange(size, device=device).view(1, 1, size, 1) + col_idx = torch.arange(size, device=device).view(1, 1, 1, size) + pattern = face_idx * 1000 + row_idx * 10 + col_idx + return pattern.expand(faces, channels, size, size).float() diff --git a/test/nn/module/test_healpix.py b/test/nn/module/test_healpix.py index 595102d270..07c104aab9 100644 --- a/test/nn/module/test_healpix.py +++ b/test/nn/module/test_healpix.py @@ -28,6 +28,7 @@ ) from physicsnemo.nn.module.hpx.padding import ( HEALPixFoldFaces, + HEALPixPaddingv2, HEALPixUnfoldFaces, ) from physicsnemo.nn.module.hpx.tokenizer import ( @@ -35,17 +36,7 @@ ) from test import common from test.conftest import requires_module - - -class MulX(torch.nn.Module): - """Helper class that just multiplies the values of an input tensor.""" - - def __init__(self, multiplier: int = 1): - super().__init__() - self.multiplier = multiplier - - def forward(self, x): - return x * self.multiplier +from test.nn.module.healpix_helpers import MulX, distinct_face_tensor @pytest.fixture @@ -58,6 +49,36 @@ def generate_test_data(faces=12, channels=2, img_size=16, device="cpu"): return generate_test_data +def _padded_faces(pad_func, size, device): + """Run ``pad_func`` on a distinct-face pattern and return the padded + output alongside a per-face lookup of the (unpadded) input, so + correctness tests can slice out expected neighbor regions.""" + pattern = distinct_face_tensor(size=size, device=device) + out = pad_func(pattern) + faces = {i: pattern[i, 0] for i in range(12)} + return out, faces + + +def _arange_grid(size, offset, device): + """A square, value-distinguishable tensor for isolated corner-blend + tests, offset so the two operands of tl()/br() never overlap in value.""" + return offset + torch.arange( + size * size, device=device, dtype=torch.float32 + ).reshape(size, size) + + +def _make_pad_func(pad_cls, padding, device): + """Construct and move a ``HEALPixPadding``/``HEALPixPaddingv2`` instance + to ``device`` in one call, since every padding test needs this.""" + return pad_cls(padding=padding).to(device) + + +def _padded_shape(batch_faces, channels, size, padding): + """Expected output shape after padding a ``(batch_faces, channels, size, + size)`` tensor by ``padding`` on each side of the last two dims.""" + return torch.Size([batch_faces, channels, size + 2 * padding, size + 2 * padding]) + + @requires_module("earth2grid") def test_HEALPixFoldFaces_initialization(device, pytestconfig): fold_func = HEALPixFoldFaces() @@ -80,6 +101,88 @@ def test_HEALPixFoldFaces_forward(device, pytestconfig): assert fold_func(invar).stride() != outvar.stride() +@requires_module("earth2grid") +def test_HEALPixFoldFaces_forward_correctness(device, pytestconfig): + fold_func = HEALPixFoldFaces() + + batch, faces, channels, height, width = 2, 3, 2, 4, 5 + invar = torch.arange( + batch * faces * channels * height * width, device=device, dtype=torch.float32 + ).reshape(batch, faces, channels, height, width) + + outvar = fold_func(invar) + + for b in range(batch): + for f in range(faces): + assert torch.equal(outvar[b * faces + f], invar[b, f]) + + +@requires_module("earth2grid") +def test_HEALPixFoldFaces_forward_invalid_ndim(device, pytestconfig): + fold_func = HEALPixFoldFaces() + + invar = torch.randn(2, 3, 4, device=device) # only 3D, needs 5D + with pytest.raises(ValueError, match="requires 5D tensor"): + fold_func(invar) + + +@requires_module("earth2grid") +def test_HEALPixUnfoldFaces_forward_correctness(device, pytestconfig): + num_faces = 12 + unfold_func = HEALPixUnfoldFaces(num_faces=num_faces) + + batch, channels, height, width = 2, 2, 4, 5 + invar = torch.arange( + batch * num_faces * channels * height * width, + device=device, + dtype=torch.float32, + ).reshape(batch * num_faces, channels, height, width) + + outvar = unfold_func(invar) + + for b in range(batch): + for f in range(num_faces): + assert torch.equal(outvar[b, f], invar[b * num_faces + f]) + + +@requires_module("earth2grid") +def test_HEALPixUnfoldFaces_forward_invalid_ndim(device, pytestconfig): + unfold_func = HEALPixUnfoldFaces(num_faces=12) + + invar = torch.randn(2, 3, 4, device=device) # only 3D, needs 4D + with pytest.raises(ValueError, match="requires 4D tensor"): + unfold_func(invar) + + +@requires_module("earth2grid") +def test_HEALPixUnfoldFaces_forward_invalid_batch_size(device, pytestconfig): + unfold_func = HEALPixUnfoldFaces(num_faces=12) + + # batch_faces=13 is not a multiple of num_faces=12 + invar = torch.randn(13, 2, 4, 4, device=device) + with pytest.raises(ValueError, match="invalid batch size"): + unfold_func(invar) + + +@requires_module("earth2grid") +def test_HEALPixFoldFaces_UnfoldFaces_roundtrip(device, pytestconfig): + num_faces = 12 + fold_func = HEALPixFoldFaces() + unfold_func = HEALPixUnfoldFaces(num_faces=num_faces) + + batch, channels, height, width = 3, 2, 4, 4 + invar = torch.randn(batch, num_faces, channels, height, width, device=device) + + folded = fold_func(invar) + assert folded.shape == (batch * num_faces, channels, height, width) + + unfolded = unfold_func(folded) + assert torch.equal(unfolded, invar) + + refolded = fold_func(unfolded) + assert torch.equal(refolded, folded) + + @requires_module("earth2grid") def test_HEALPixUnfoldFaces_initialization(device, pytestconfig): unfold_func = HEALPixUnfoldFaces() @@ -134,6 +237,201 @@ def test_HEALPixPadding_forward(device, padding, pytestconfig): assert outvar.shape == out_size +@requires_module("earth2grid") +def test_HEALPixPadding_forward_invalid_ndim(device, pytestconfig): + pad_func = _make_pad_func(HEALPixPadding, 1, device) + + invar = torch.randn(2, 3, 4, device=device) # only 3D, needs 4D + with pytest.raises(ValueError, match="requires a 4D tensor"): + pad_func(invar) + + +@requires_module("earth2grid") +@pytest.mark.parametrize( + "pad_cls,batch_faces,channels", + [(HEALPixPadding, 12, 2), (HEALPixPaddingv2, 24, 3)], + ids=["HEALPixPadding", "HEALPixPaddingv2"], +) +def test_padding_forward_skips_validation_when_compiling( + device, monkeypatch, pad_cls, batch_faces, channels, pytestconfig +): + """When invoked from inside a compiled graph, ``torch.compiler.is_compiling()`` + reports ``True`` and shape validation (in ``HEALPixPadding``/``HEALPixFoldFaces``/ + ``HEALPixUnfoldFaces``) must be skipped without affecting correctness.""" + if pad_cls is HEALPixPaddingv2 and device == "cpu": + pytest.skip("HEALPixPaddingv2 requires a CUDA device") + + padding = 2 + size = 4 + pad_func = _make_pad_func(pad_cls, padding, device) + + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + invar = torch.randn(batch_faces, channels, size, size, device=device) + outvar = pad_func(invar) + + assert outvar.shape == _padded_shape(batch_faces, channels, size, padding) + + +@requires_module("earth2grid") +@pytest.mark.parametrize( + "pad_cls,batch_faces,channels", + [(HEALPixPadding, 12, 2), (HEALPixPaddingv2, 24, 3)], + ids=["HEALPixPadding", "HEALPixPaddingv2"], +) +def test_padding_forward_cuda_nvtx_skipped_without_cuda( + device, monkeypatch, pad_cls, batch_faces, channels, pytestconfig +): + """The nvtx range push/pop around the forward pass is gated on + ``torch.cuda.is_available()``, independent of the tensor's actual + device; verify the pass-through path when it reports unavailable.""" + if pad_cls is HEALPixPaddingv2 and device == "cpu": + pytest.skip("HEALPixPaddingv2 requires a CUDA device") + + padding = 2 + size = 4 + pad_func = _make_pad_func(pad_cls, padding, device) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + invar = torch.randn(batch_faces, channels, size, size, device=device) + outvar = pad_func(invar) + + assert outvar.shape == _padded_shape(batch_faces, channels, size, padding) + + +@requires_module("earth2grid") +def test_HEALPixPadding_forward_correctness_north(device, pytestconfig): + """Face 0 is a northern-hemisphere face padded via ``pn``, which rotates + its top/left neighbors. Verify against an independently computed expected + tensor built from the same documented neighbor/rotation contract.""" + padding = 2 + size = 4 + d = (-2, -1) + + pad_func = _make_pad_func(HEALPixPadding, padding, device) + out, f = _padded_faces(pad_func, size, device) + + p = padding + center = torch.cat((f[1].rot90(1, d)[-p:, :], f[0], f[4][:p, :]), dim=0) + left = torch.cat( + (f[2].rot90(2, d)[-p:, -p:], f[3].rot90(-1, d)[:, -p:], f[3][:p, -p:]), dim=0 + ) + right = torch.cat((f[1][-p:, :p], f[5][:, :p], f[8][:p, :p]), dim=0) + expected = torch.cat((left, center, right), dim=1) + + torch.testing.assert_close(out[0, 0], expected) + + +@requires_module("earth2grid") +def test_HEALPixPadding_forward_correctness_equator(device, pytestconfig): + """Face 4 is an equatorial face padded via ``pe``, including the tl()/br() + corner-blend helpers for its missing diagonal neighbors.""" + padding = 2 + size = 4 + + pad_func = _make_pad_func(HEALPixPadding, padding, device) + out, f = _padded_faces(pad_func, size, device) + + p = padding + tl_corner = pad_func.tl(f[0], f[3]) + br_corner = pad_func.br(f[11], f[8]) + + center = torch.cat((f[0][-p:, :], f[4], f[11][:p, :]), dim=0) + left = torch.cat((tl_corner[-p:, -p:], f[3][:, -p:], f[7][:p, -p:]), dim=0) + right = torch.cat((f[5][-p:, :p], f[8][:, :p], br_corner[:p, :p]), dim=0) + expected = torch.cat((left, center, right), dim=1) + + torch.testing.assert_close(out[4, 0], expected) + + +@requires_module("earth2grid") +def test_HEALPixPadding_forward_correctness_south(device, pytestconfig): + """Face 8 is a southern-hemisphere face padded via ``ps``, which rotates + its bottom/right neighbors.""" + padding = 2 + size = 4 + d = (-2, -1) + + pad_func = _make_pad_func(HEALPixPadding, padding, device) + out, f = _padded_faces(pad_func, size, device) + + p = padding + center = torch.cat((f[5][-p:, :], f[8], f[11].rot90(1, d)[:p, :]), dim=0) + left = torch.cat((f[0][-p:, -p:], f[4][:, -p:], f[11][:p, -p:]), dim=0) + right = torch.cat( + (f[9][-p:, :p], f[9].rot90(-1, d)[:, :p], f[10].rot90(2, d)[:p, :p]), dim=0 + ) + expected = torch.cat((left, center, right), dim=1) + + torch.testing.assert_close(out[8, 0], expected) + + +@requires_module("earth2grid") +def test_HEALPixPadding_tl_corner_blend(device, pytestconfig): + """Directly verify the tl() diagonal-corner blend formula in isolation.""" + padding = 3 + size = 5 + pad_func = _make_pad_func(HEALPixPadding, padding, device) + + top = _arange_grid(size, offset=0, device=device) + lft = _arange_grid(size, offset=1000, device=device) + + result = pad_func.tl(top, lft) + + p = padding + expected = torch.zeros_like(top)[..., :p, :p] + expected[..., -1, -1] = 0.5 * top[..., -1, 0] + 0.5 * lft[..., 0, -1] + for i in range(1, p): + expected[..., -i - 1, -i:] = top[..., -i - 1, :i] + expected[..., -i:, -i - 1] = lft[..., :i, -i - 1] + expected[..., -i - 1, -i - 1] = ( + 0.5 * top[..., -i - 1, 0] + 0.5 * lft[..., 0, -i - 1] + ) + + torch.testing.assert_close(result, expected) + + +@requires_module("earth2grid") +def test_HEALPixPadding_br_corner_blend(device, pytestconfig): + """Directly verify the br() diagonal-corner blend formula in isolation.""" + padding = 3 + size = 5 + pad_func = _make_pad_func(HEALPixPadding, padding, device) + + b = _arange_grid(size, offset=0, device=device) + r = _arange_grid(size, offset=1000, device=device) + + result = pad_func.br(b, r) + + p = padding + expected = torch.zeros_like(b)[..., :p, :p] + expected[..., 0, 0] = 0.5 * b[..., 0, -1] + 0.5 * r[..., -1, 0] + for i in range(1, p): + expected[..., :i, i] = r[..., -i:, i] + expected[..., i, :i] = b[..., i, -i:] + expected[..., i, i] = 0.5 * b[..., i, -1] + 0.5 * r[..., -1, i] + + torch.testing.assert_close(result, expected) + + +@requires_module("earth2grid") +def test_HEALPixPaddingv2_forward(device, pytestconfig): + """``HEALPixPaddingv2`` wraps the accelerated ``earth2grid`` padding kernel, + which requires an actual CUDA device regardless of earth2grid's presence.""" + if device == "cpu": + pytest.skip("HEALPixPaddingv2 requires a CUDA device") + + padding = 2 + size = 4 + pad_func = _make_pad_func(HEALPixPaddingv2, padding, device) + + invar = torch.randn(24, 3, size, size, device=device) + outvar = pad_func(invar) + + assert outvar.shape == _padded_shape(24, 3, size, padding) + + @requires_module("earth2grid") @pytest.mark.parametrize("multiplier", [2, 3, 4]) def test_HEALPixLayer_initialization(device, multiplier, pytestconfig): From 707c35b89e279ccfc661f51e00976564da460d08 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 20:02:37 +0000 Subject: [PATCH 12/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- test/pytest_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/pytest_utils.py b/test/pytest_utils.py index 1f00fded44..884a69b8a5 100644 --- a/test/pytest_utils.py +++ b/test/pytest_utils.py @@ -104,7 +104,10 @@ def _import_or_fail(module_names, config, min_versions=None): try: module = importlib.import_module(module_name) if hasattr(module, "__version__"): - if isinstance(module.__version__, str) or module.__version__ is None: + if ( + isinstance(module.__version__, str) + or module.__version__ is None + ): pytest.importorskip(module_name, min_version) elif isinstance(module.__version__, Version): version_check = Version(min_version) From 06a8c3aef59237f1fd29bea63eba9898e2277d62 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Fri, 10 Jul 2026 16:00:01 -0500 Subject: [PATCH 13/28] Make xarray zarr imports lazy --- .../healpix/base_timeseries_dataset_zarr.py | 9 +++++++-- .../datapipes/healpix/coupledtimeseries_dataset.py | 9 ++++++++- .../healpix/coupledtimeseries_dataset_zarr.py | 7 ++++++- physicsnemo/datapipes/healpix/couplers.py | 14 +++++++++++--- physicsnemo/datapipes/healpix/data_modules.py | 11 ++++++++--- .../datapipes/healpix/timeseries_dataset.py | 9 ++++++++- physicsnemo/metrics/climate/hydrostasy.py | 9 ++++++++- .../models/dlwp_healpix/layers/normalization.py | 13 ++++++++++--- test/metrics/test_hydrostatic_loss.py | 9 ++++++--- 9 files changed, 72 insertions(+), 18 deletions(-) diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py index bb55abcd67..e302d03948 100644 --- a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py @@ -31,15 +31,20 @@ import numpy as np import pandas as pd -import xarray as xr -import zarr from omegaconf import DictConfig, OmegaConf +from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.datapipe import Datapipe from physicsnemo.datapipes.meta import DatapipeMetaData logger = logging.getLogger(__name__) +# xarray/zarr ship in the ``datapipes-extras`` optional dependency group; load +# them lazily so the physicsnemo import graph carries no static dependency on +# them (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). +xr = OptionalImport("xarray") +zarr = OptionalImport("zarr") + def _is_object_store_path(path: str) -> bool: # pragma: no cover """Check if path is an object store path (contains :// or ::). diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index ca6e51ff4a..a7f1b9c21a 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -20,6 +20,8 @@ This class extends the TimeSeriesDataset to add the core functionality for coupling time series healpix data with external inputs from various earth system components. """ +from __future__ import annotations + import gc import logging import time @@ -29,9 +31,9 @@ import numpy as np import pandas as pd import torch -import xarray as xr from omegaconf import DictConfig, OmegaConf +from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.meta import DatapipeMetaData from physicsnemo.utils.insolation import insolation @@ -40,6 +42,11 @@ logger = logging.getLogger(__name__) +# xarray ships in the ``datapipes-extras`` optional dependency group; load it +# lazily so the physicsnemo import graph carries no static dependency on it +# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). +xr = OptionalImport("xarray") + @dataclass class MetaData(DatapipeMetaData): diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py index 73ade041ab..aba00d8343 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py @@ -27,9 +27,9 @@ import numpy as np import pandas as pd import torch -import zarr from omegaconf import DictConfig, OmegaConf +from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.meta import DatapipeMetaData from physicsnemo.utils.insolation import insolation @@ -39,6 +39,11 @@ logger = logging.getLogger(__name__) +# zarr ships in the ``datapipes-extras`` optional dependency group; load it +# lazily so the physicsnemo import graph carries no static dependency on it +# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). +zarr = OptionalImport("zarr") + @dataclass class MetaData(DatapipeMetaData): diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 7e4b5cefd3..32422fc051 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -14,6 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import logging from abc import ABC, abstractmethod from typing import Sequence @@ -22,11 +24,17 @@ import numpy as np import pandas as pd import torch as th -import xarray as xr -import zarr as zr + +from physicsnemo.core.version_check import OptionalImport logger = logging.getLogger(__name__) +# xarray/zarr ship in the ``datapipes-extras`` optional dependency group; load +# them lazily so the physicsnemo import graph carries no static dependency on +# them (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). +xr = OptionalImport("xarray") +zarr = OptionalImport("zarr") + class BaseCoupler(ABC): """ @@ -98,7 +106,7 @@ def __init__( if isinstance(self.ds, xr.Dataset): self.use_zarr = False - elif isinstance(self.ds, zr.Group): + elif isinstance(self.ds, zarr.Group): self.use_zarr = True # iterate over self.variables in the outer loop so the selected # indices follow the order of self.variables (matching the xarray diff --git a/physicsnemo/datapipes/healpix/data_modules.py b/physicsnemo/datapipes/healpix/data_modules.py index 931f5d20d6..c010a493a4 100644 --- a/physicsnemo/datapipes/healpix/data_modules.py +++ b/physicsnemo/datapipes/healpix/data_modules.py @@ -26,6 +26,8 @@ - CoupledTimeSeriesDataModule: A DataModule for loading and processing coupled time series healpix data. """ +from __future__ import annotations + # System modules import logging import warnings @@ -35,14 +37,12 @@ # numpy import numpy as np -# distributed stuff -import xarray as xr - # External modules from omegaconf import DictConfig from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler +from physicsnemo.core.version_check import OptionalImport from physicsnemo.distributed import DistributedManager from .coupledtimeseries_dataset import CoupledTimeSeriesDataset @@ -50,6 +50,11 @@ logger = logging.getLogger(__name__) +# xarray ships in the ``datapipes-extras`` optional dependency group; load it +# lazily so the physicsnemo import graph carries no static dependency on it +# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). +xr = OptionalImport("xarray") + def open_time_series_dataset_classic_prebuilt( directory: str, dataset_name: str, constants: bool = False, batch_size: int = 32 diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index f06ce803c3..b847da4cdb 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -21,6 +21,8 @@ It handles data loading, scaling, and time management. It is used by the TimeSeriesDataModule to set up the dataloaders for the time series healpix data. """ +from __future__ import annotations + import logging import time import warnings @@ -30,15 +32,20 @@ import numpy as np import pandas as pd import torch -import xarray as xr from omegaconf import DictConfig, OmegaConf +from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.datapipe import Datapipe from physicsnemo.datapipes.meta import DatapipeMetaData from physicsnemo.utils.insolation import insolation logger = logging.getLogger(__name__) +# xarray ships in the ``datapipes-extras`` optional dependency group; load it +# lazily so the physicsnemo import graph carries no static dependency on it +# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). +xr = OptionalImport("xarray") + @dataclass class MetaData(DatapipeMetaData): diff --git a/physicsnemo/metrics/climate/hydrostasy.py b/physicsnemo/metrics/climate/hydrostasy.py index db677099e5..4916ddb95b 100644 --- a/physicsnemo/metrics/climate/hydrostasy.py +++ b/physicsnemo/metrics/climate/hydrostasy.py @@ -19,10 +19,17 @@ import numpy as np import torch -import xarray as xr + +from physicsnemo.core.version_check import OptionalImport logger = logging.getLogger(__name__) +# xarray ships in the ``datapipes-extras`` optional dependency group and is only +# needed by the topography-loading helpers here; load it lazily so the +# physicsnemo import graph carries no static dependency on it (see +# CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). +xr = OptionalImport("xarray") + def _average_virtual_temperature_from_geopotential_height(z1, z2, p1, p2, R, g0): return g0 / (R * np.log(p1 / p2)) * (z2 - z1) diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index 9d6f3c1c30..2c1147445a 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -20,15 +20,22 @@ This class contains the implementation of the Deep Learning Weather Prediction (DLWP) normalization on the HEALPix mesh. """ +import importlib from typing import List import torch as th +# ``apex`` is an optional accelerator (not a core dependency). Import it +# dynamically so the physicsnemo import graph stays free of a static ``apex`` +# dependency (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003). The +# ``FusedLayerNorm`` symbol is only used when ``norm_op="apex"`` is explicitly +# requested, which is guarded by ``_APEX_AVAILABLE`` below. try: - from apex.normalization import FusedLayerNorm + FusedLayerNorm = importlib.import_module("apex.normalization").FusedLayerNorm - _APEX_AVAILABLE = True # pragma: no cover -except ImportError: + _APEX_AVAILABLE = True +except ImportError: # pragma: no cover - only exercised when apex is absent + FusedLayerNorm = None _APEX_AVAILABLE = False diff --git a/test/metrics/test_hydrostatic_loss.py b/test/metrics/test_hydrostatic_loss.py index a9152ab6de..fcc4019c2b 100644 --- a/test/metrics/test_hydrostatic_loss.py +++ b/test/metrics/test_hydrostatic_loss.py @@ -18,12 +18,15 @@ import pytest import torch -from physicsnemo.metrics.climate.hydrostasy import ( +# xarray is an optional dependency (datapipes-extras). Skip this module entirely +# when it is unavailable, before importing hydrostasy, rather than failing at +# collection time. +pytest.importorskip("xarray") + +from physicsnemo.metrics.climate.hydrostasy import ( # noqa: E402 HydrostaticBalance, ) -xr = pytest.importorskip("xarray") - def test_constant_temperature(rtol: float = 1e-3, atol: float = 1e-3): R = 287 # J K^{-1} kg^{-1} From f3149e348b9d7cc38fd13534d97e469a6f229737 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Mon, 13 Jul 2026 10:26:02 -0500 Subject: [PATCH 14/28] revert pytest_utils --- test/pytest_utils.py | 93 -------------------------------------------- 1 file changed, 93 deletions(-) diff --git a/test/pytest_utils.py b/test/pytest_utils.py index 884a69b8a5..51c21c2188 100644 --- a/test/pytest_utils.py +++ b/test/pytest_utils.py @@ -54,96 +54,3 @@ def modify_environment(*remove, **update): finally: env.update(restore_after) [env.pop(k, None) for k in purge_after] - - -import importlib -from functools import wraps - -import pytest -from packaging.version import Version - - -def import_or_fail( - module_names: str | list[str] | tuple, - min_versions: str | list[str] | tuple | None = None, -): - """Skip test if module is missing or below minimum version.""" - - def decorator(test_func): - @pytest.mark.usefixtures("pytestconfig") - @wraps(test_func) - def wrapper(*args, **kwargs): - pytestconfig = kwargs.get("pytestconfig") - if pytestconfig is None: - raise ValueError( - "pytestconfig must be passed when using import_or_fail." - ) - _import_or_fail(module_names, pytestconfig, min_versions) - return test_func(*args, **kwargs) - - return wrapper - - return decorator - - -def _import_or_fail(module_names, config, min_versions=None): - if not isinstance(module_names, (list, tuple)): - module_names = [module_names] - if min_versions is not None and not isinstance(min_versions, (list, tuple)): - min_versions = [min_versions] - - if min_versions is None: - min_versions = [None] * len(module_names) - elif len(min_versions) != len(module_names): - raise ValueError("module_names and min_versions must have the same length.") - - for module_name, min_version in zip(module_names, min_versions): - if config.getoption("--fail-on-missing-modules"): - __import__(module_name) - else: - try: - module = importlib.import_module(module_name) - if hasattr(module, "__version__"): - if ( - isinstance(module.__version__, str) - or module.__version__ is None - ): - pytest.importorskip(module_name, min_version) - elif isinstance(module.__version__, Version): - version_check = Version(min_version) - if module.__version__ < version_check: - pytest.skip( - f"{module_name} {module.__version__} is less than " - f"required version {version_check}" - ) - except ModuleNotFoundError: - pytest.importorskip(module_name, min_version) - - -def nfsdata_or_fail(test_func): - @pytest.mark.usefixtures("pytestconfig") - @wraps(test_func) - def wrapper(*args, **kwargs): - pytestconfig = kwargs.get("pytestconfig") - if pytestconfig is None: - raise ValueError("pytestconfig must be passed when using nfsdata_or_fail.") - _nfsdata_or_fail(pytestconfig) - return test_func(*args, **kwargs) - - return wrapper - - -def _nfsdata_or_fail(config): - nfs_data_dir_opt = config.getoption("--nfs-data-dir") - test_data_dir_env = os.environ.get("TEST_DATA_DIR") - if nfs_data_dir_opt and os.path.exists(nfs_data_dir_opt): - return - if test_data_dir_env and os.path.exists(test_data_dir_env): - return - for path in ("/data/nfs/physicsnemo-data", "/data/nfs/modulus-data"): - if os.path.exists(path): - return - pytest.skip( - "NFS volumes not set up with CI data repo. Run `make get-data` from the " - "root directory of the repo" - ) From b691d11b7449e09d636362d71b4ea5070d31dbc6 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Mon, 13 Jul 2026 11:15:39 -0500 Subject: [PATCH 15/28] Add greptile fixes for tests and docs --- .../healpix/coupledtimeseries_dataset.py | 1 + .../healpix/coupledtimeseries_dataset_zarr.py | 1 + physicsnemo/metrics/climate/healpix_loss.py | 16 ++++++++++++---- physicsnemo/metrics/climate/hydrostasy.py | 5 +++++ 4 files changed, 19 insertions(+), 4 deletions(-) diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index a7f1b9c21a..56d4abe268 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -320,6 +320,7 @@ def __getitem__(self, item): ) for c in self.couplings: for i, v in enumerate(c.variables): + # coupled fields are in the shape [T, B, C, F, H, W] integrated_couplings[i, :, :] += self.rng.normal( loc=0, scale=self.train_noise_params["couplings"][v]["std"], diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py index aba00d8343..dc4c8f7825 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py @@ -266,6 +266,7 @@ def __getitem__( ) for c in self.couplings: for i, v in enumerate(c.variables): + # coupled fields are in the shape [T, B, C, F, H, W] integrated_couplings[i, :, :] += self.rng.normal( loc=0, scale=self.train_noise_params["couplings"][v]["std"], diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index 169dfbc902..0f67772fcb 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -172,7 +172,7 @@ def setup(self, trainer): self.b = th.tensor(self.b, device=trainer.device) self.w = th.tensor(self.w, device=trainer.device) - def forward(self, prediction, target): + def forward(self, prediction, target, average_channels=True): """ Computes the MSE of model prediction and applies weights for zero and non-zero precipitation cases. @@ -182,13 +182,18 @@ def forward(self, prediction, target): The prediction tensor target: torch.Tensor The target tensor + average_channels: bool, optional + whether the mean of the channels should be taken """ weights_for_zero = th.ones_like(target) * self.weight_zero weights_for_nonzero = (th.ones_like(target) * self.weight_nonzero) * th.exp( self.b * target ) weights = th.where(target > 0, weights_for_nonzero, weights_for_zero) - loss = (th.mean(weights * (prediction - target) ** 2)) * self.w + if average_channels: + loss = (th.mean(weights * (prediction - target) ** 2)) * self.w + else: + loss = (weights * (prediction - target) ** 2) * self.w return loss @@ -701,6 +706,10 @@ def __init__( # Parameters for "almost fair CRPS" loss. See https://arxiv.org/html/2412.15832v1 self.coeff_eps = 1 - ((1 - alpha) / (n_members)) + if n_members < 2: + raise ValueError("n_members must be at least 2 for CRPS loss to be defined") + else: + self.n_members = n_members self.averaging_coeff = 1 / (2 * n_members * (n_members - 1)) # SHT utils: transform, grid reordering, output indexing @@ -765,8 +774,7 @@ def setup(self, trainer): if self.multiscale > 0: self.isht = self.isht.to(device=trainer.device) self.reorder_from_ring = self.reorder_from_ring.to(device=trainer.device) - if self.lsm_file is not None: - self.lsm_tensor = self.lsm_tensor.to(device=trainer.device) + self.lsm_tensor = self.lsm_tensor.to(device=trainer.device) def _apply_sht(self, x, face_dim, return_abs=True): """Apply SHT to a tensor diff --git a/physicsnemo/metrics/climate/hydrostasy.py b/physicsnemo/metrics/climate/hydrostasy.py index 4916ddb95b..c16d318559 100644 --- a/physicsnemo/metrics/climate/hydrostasy.py +++ b/physicsnemo/metrics/climate/hydrostasy.py @@ -297,10 +297,15 @@ def __init__( } # Get offset to map from q channels here to Tv channels # Relies on all pressure levels below a threshold to have q channels + self.q_index_offset = None for i, pl in enumerate(self.pressure_levels): if f"q{int(pl)}" in channels: self.q_index_offset = i break + if self.q_index_offset is None: + raise ValueError( + f"No q channels found in channels list that match the requested pressure levels {self.pressure_levels}" + ) # Create mapping for new tensor that holds only the constraint variables self.z_constraint_pressure_levels = { From 0a4ac20b40572958bf8aa70abaae31d2119e9f81 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Tue, 14 Jul 2026 17:15:15 -0500 Subject: [PATCH 16/28] Scope dlesym_updates to architecture only Remove the DLESyM datapipes/dataloader and training/loss changes from this branch, restoring those files to upstream. Those changes now live on their own branches (dlesym/datapipes, dlesym/training) for separate merge requests; this branch retains only the DLWP-HEALPix architecture changes (models, layers, HEALPix primitives), matching dlesym/architecture. --- physicsnemo/datapipes/healpix/__init__.py | 37 +- .../healpix/base_timeseries_dataset_zarr.py | 644 -------- .../healpix/coupledtimeseries_dataset.py | 45 +- .../healpix/coupledtimeseries_dataset_zarr.py | 392 ----- physicsnemo/datapipes/healpix/couplers.py | 565 ++++--- physicsnemo/datapipes/healpix/data_modules.py | 576 +++++-- .../datapipes/healpix/data_modules_zarr.py | 557 ------- .../datapipes/healpix/timeseries_dataset.py | 39 +- .../healpix/timeseries_dataset_zarr.py | 279 ---- physicsnemo/metrics/climate/healpix_loss.py | 827 ----------- physicsnemo/metrics/climate/hydrostasy.py | 900 ----------- test/datapipes/healpix/__init__.py | 17 - test/datapipes/healpix/conftest.py | 190 --- test/datapipes/healpix/test_healpix.py | 564 ------- test/datapipes/healpix/test_healpix_couple.py | 1319 ----------------- .../healpix/test_healpix_couple_zarr.py | 1288 ---------------- test/datapipes/healpix/test_healpix_zarr.py | 873 ----------- test/metrics/test_hydrostatic_loss.py | 172 --- 18 files changed, 749 insertions(+), 8535 deletions(-) delete mode 100644 physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py delete mode 100644 physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py delete mode 100644 physicsnemo/datapipes/healpix/data_modules_zarr.py delete mode 100644 physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py delete mode 100644 physicsnemo/metrics/climate/hydrostasy.py delete mode 100644 test/datapipes/healpix/__init__.py delete mode 100644 test/datapipes/healpix/conftest.py delete mode 100644 test/datapipes/healpix/test_healpix.py delete mode 100644 test/datapipes/healpix/test_healpix_couple.py delete mode 100644 test/datapipes/healpix/test_healpix_couple_zarr.py delete mode 100644 test/datapipes/healpix/test_healpix_zarr.py delete mode 100644 test/metrics/test_hydrostatic_loss.py diff --git a/physicsnemo/datapipes/healpix/__init__.py b/physicsnemo/datapipes/healpix/__init__.py index ee5b75e6cf..ae1a6ce394 100644 --- a/physicsnemo/datapipes/healpix/__init__.py +++ b/physicsnemo/datapipes/healpix/__init__.py @@ -14,39 +14,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Healpix DataPipes. - -This package contains the Healpix DataPipes for loading and processing healpix data. -It supports loading data from xarray Datasets or zarr Groups, and coupling it with a model. - -The main classes are: -- TimeSeriesDataModule: A DataModule for loading and processing time series data. -- CoupledTimeSeriesDataModule: A DataModule for loading and processing coupled time series data. -- TimeSeriesDataset: A dataset for loading and processing time series data. -- CoupledTimeSeriesDataset: A dataset for loading and processing coupled time series data. -- Zarr versions of the above classes. -- ConstantCoupler: A coupler that uses a constant value for the coupled data. -- TrailingAverageCoupler: A coupler that uses a trailing average for the coupled data. -""" - -from .coupledtimeseries_dataset import CoupledTimeSeriesDataset -from .coupledtimeseries_dataset_zarr import CoupledTimeSeriesDatasetZarr -from .couplers import ConstantCoupler, TrailingAverageCoupler -from .data_modules import CoupledTimeSeriesDataModule, TimeSeriesDataModule -from .data_modules_zarr import CoupledTimeSeriesDataModuleZarr, TimeSeriesDataModuleZarr -from .timeseries_dataset import TimeSeriesDataset -from .timeseries_dataset_zarr import TimeSeriesDatasetZarr - -__all__ = [ - "TimeSeriesDataModule", - "CoupledTimeSeriesDataModule", - "TimeSeriesDatasetZarr", - "CoupledTimeSeriesDatasetZarr", - "TimeSeriesDataset", - "CoupledTimeSeriesDataset", - "TimeSeriesDataModuleZarr", - "CoupledTimeSeriesDataModuleZarr", - "ConstantCoupler", - "TrailingAverageCoupler", -] +from .data_modules import CoupledTimeSeriesDataModule diff --git a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py deleted file mode 100644 index e302d03948..0000000000 --- a/physicsnemo/datapipes/healpix/base_timeseries_dataset_zarr.py +++ /dev/null @@ -1,644 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -BaseTimeSeriesDatasetZarr - Abstract base class for time series healpix datasets using Zarr storage. - -This class provides the core functionality for loading and processing time series healpix data -stored in Zarr format. It handles data loading, scaling, and time management. -Subclasses must implement the __getitem__ method to define specific data retrieval logic. -""" - -import importlib.util -import logging -import os -import warnings -from abc import ABC, abstractmethod -from typing import List, Optional, Sequence, Tuple, Union - -import numpy as np -import pandas as pd -from omegaconf import DictConfig, OmegaConf - -from physicsnemo.core.version_check import OptionalImport -from physicsnemo.datapipes.datapipe import Datapipe -from physicsnemo.datapipes.meta import DatapipeMetaData - -logger = logging.getLogger(__name__) - -# xarray/zarr ship in the ``datapipes-extras`` optional dependency group; load -# them lazily so the physicsnemo import graph carries no static dependency on -# them (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). -xr = OptionalImport("xarray") -zarr = OptionalImport("zarr") - - -def _is_object_store_path(path: str) -> bool: # pragma: no cover - """Check if path is an object store path (contains :// or ::). - - Parameters - ---------- - path : str - Path to check - - Returns - ------- - bool - True if path appears to be an object store path - """ - return "://" in str(path) or "::" in str(path) - - -def _check_availability(path: str) -> None: # pragma: no cover - """ - Check if path exists or fsspec is available for object store paths - - Parameters - ---------- - path : str - Path to check - - Raises - ------ - ImportError - If path is an object store path but fsspec is not available - FileNotFoundError - If the path is file and doesn't exist - """ - if _is_object_store_path(path): - if not importlib.util.find_spec("fsspec"): - raise ImportError( - f"fsspec is required to access object store paths like '{path}'. " - "Please install fsspec with: pip install fsspec" - ) - elif not os.path.exists(path): - raise FileNotFoundError(f"Dataset not found at specified location: {path}") - - -class BaseTimeSeriesDatasetZarr(Datapipe, ABC): - """Abstract base class for time series datasets using Zarr storage. - - This class provides the core functionality for loading and processing time series data - stored in Zarr format. It handles data loading, scaling, and time management. - Subclasses must implement the __getitem__ method to define specific data retrieval logic. - """ - - def __init__( - self, - dataset_path: str, - input_variables: Sequence, - output_variables: Sequence = None, - constant_variables: Sequence = None, - scaling: DictConfig = None, - input_time_dim: int = 1, - output_time_dim: int = 1, - data_time_step: Union[int, str] = "3h", - time_step: Union[int, str] = "6h", - gap: Union[int, str, None] = None, - batch_size: int = 32, - drop_last: bool = False, - add_insolation: bool = False, - forecast_init_times: Optional[Sequence] = None, - start_date: Optional[Union[int, str]] = None, - end_date: Optional[Union[int, str]] = None, - add_train_noise: bool = False, - train_noise_params: DictConfig = None, - train_noise_seed: int = 42, - meta: DatapipeMetaData = None, - ): - """Initialize base time series dataset. - - Parameters - ---------- - dataset_path : str - Path to the Zarr dataset - input_variables : Sequence - Variables to use as model inputs - output_variables : Sequence, optional - Variables to predict as outputs. If None, uses input_variables - constant_variables : Sequence, optional - Constant fields used as additional inputs - scaling : DictConfig, optional - Configuration for data scaling/normalization - input_time_dim : int, default=1 - Number of time steps in input sequence - output_time_dim : int, default=1 - Number of time steps to predict - data_time_step : Union[int, str], default="3h" - Either integer hours or a str interpretable by pandas - Time resolution of raw data - time_step : Union[int, str], default="6h" - Either integer hours or a str interpretable by pandas - Time step between predictions - gap : Union[int, str, None], optional - Either integer hours or a str interpretable by pandas - Time gap between input and output sequences - batch_size : int, default=32 - Number of samples per batch - drop_last : bool, default=False - Whether to drop last incomplete batch - add_insolation : bool, default=False - Whether to add solar insolation as input - forecast_init_times : Sequence, optional - A Sequence of pandas Timestamps - Specific times to initialize forecasts - start_date : Union[int, str], optional - Start date/index from which to load data - end_date : Union[int, str], optional - End date/index to which to load data - add_train_noise : bool, default=False - Whether to add train noise - train_noise_params : DictConfig, optional - Standard deviation of train noise - train_noise_seed : int, default=42 - Seed for train noise - meta : DatapipeMetaData, optional - Metadata for the datapipe - """ - Datapipe.__init__(self, meta=meta) - - self.dataset_path = dataset_path - self.scaling = OmegaConf.to_object(scaling) if scaling else None - self.input_time_dim = input_time_dim - self.output_time_dim = output_time_dim - self.data_time_step = self._convert_time_step(data_time_step) - self.time_step = self._convert_time_step(time_step) - self.gap = self._convert_time_step(gap if gap is not None else time_step) - self.batch_size = batch_size - self.drop_last = drop_last - self.add_insolation = add_insolation - self.forecast_init_times = forecast_init_times - self.forecast_mode = self.forecast_init_times is not None - self.input_variables = input_variables - self.output_variables = ( - input_variables if output_variables is None else output_variables - ) - self.constant_variables = constant_variables - self.all_variables = list( - set(self.input_variables).union(self.output_variables) - ) - self.all_scaling = None - - # Check if for fsspec if necessary and make sure path exist - _check_availability(dataset_path) - - self.ds = zarr.open(dataset_path) - - if ( - start_date is None or end_date is None - ) and self.forecast_init_times is None: - raise ValueError( - "Either start and end date or forecast_init_times must be provided" - ) - - # Validate channels exist - channels = set(self.input_variables).union(self.output_variables) - missing_channels = channels - set(self.ds["channel_in"][:]) - if len(missing_channels) > 0: - raise KeyError( - f"Requested Input, coupled, or output variables not found in dataset: {missing_channels}" - ) - - self._get_time_da(self.dataset_path, start_date, end_date) - - self.all_variable_indices = [ - int(np.where(self.ds["channel_in"][:] == ch)[0][0]) - for ch in self.all_variables - ] - - # Validate constants exist - if constant_variables: - missing_constants = set(constant_variables) - set(self.ds["channel_c"][:]) - if len(missing_constants) > 0: - raise KeyError( - f"Requested constants not found in dataset: {missing_constants}" - ) - - self.constant_variable_indices = ( - [ - int(np.where(self.ds["channel_c"][:] == ch)[0][0]) - for ch in self.constant_variables - ] - if self.constant_variables - else None - ) - self.input_variable_indices = [ - self.all_variables.index(inp_ch) for inp_ch in self.input_variables - ] - self.output_variable_indices = [ - self.all_variables.index(out_ch) for out_ch in self.output_variables - ] - - # Length of the data window needed for one sample - if self.forecast_mode: - self._window_length = self.interval * (self.input_time_dim - 1) + 1 - else: - self._window_length = ( - self.interval * (self.input_time_dim - 1) - + 1 - + (self.gap // self.data_time_step) - + self.interval * (self.output_time_dim - 1) - ) - - self._batch_window_length = self.batch_size + self._window_length - 1 - self._output_delay = self.interval * (self.input_time_dim - 1) + ( - self.gap // self.data_time_step - ) - - # Indices within a batch - self._input_indices = [ - list(range(n, n + self.interval * self.input_time_dim, self.interval)) - for n in range(self.batch_size) - ] - self._output_indices = [ - list( - range( - n + self._output_delay, - n + self.interval * self.output_time_dim + self._output_delay, - self.interval, - ) - ) - for n in range(self.batch_size) - ] - - self.spatial_dims = ( - self.ds["face"].shape[0], - self.ds["height"].shape[0], - self.ds["width"].shape[0], - ) - - # Cached values - self.lat = np.asarray(self.ds["lat"]) - self.lon = np.asarray(self.ds["lon"]) - self.input_scaling = None - self.target_scaling = None - self.constant_scaling = None - self.constants = None - - if self.scaling: - self._get_scaling_da() - if self.constant_variables: - self.constants = self.get_constants() - - self.add_train_noise = add_train_noise - self.train_noise_params = train_noise_params - if self.add_train_noise: - self.rng = np.random.default_rng(train_noise_seed) - - @staticmethod - def _convert_time_step(dt: Union[int, str]) -> pd.Timedelta: - """Convert time step specification to Timedelta. - - Parameters - ---------- - dt : Union[int, str] - Either integer hours or string time to convert to Timedelta - - Returns - ------- - pd.Timedelta - Converted time delta object - """ - return pd.Timedelta(hours=dt) if isinstance(dt, int) else pd.Timedelta(dt) - - def _get_time_da( - self, - dataset_path: str, - start_date: Optional[Union[int, str]], - end_date: Optional[Union[int, str]], - ) -> None: - """Load and decode time array from dataset. - - Sets up time-related attributes including total samples, start index, - and forecast initialization indices. - - Parameters - ---------- - dataset_path : str - Path to Zarr dataset - start_date : Optional[Union[int, str]] - Start date/index for data slice - end_date : Optional[Union[int, str]] - End date/index for data slice - """ - # Check if fsspec is available for object store paths - _check_availability(dataset_path) - - ds = xr.open_zarr(dataset_path) - - if "time" not in ds: - raise KeyError(f"Dataset missing time. Dataset provided {dataset_path}") - - # integer start/end dates are indices into the time array rather than - # dates themselves (see the int handling below), so they can't be - # compared against `ds.time` via `np.datetime64` - if ( - start_date is not None - and not isinstance(start_date, int) - and np.datetime64(start_date) < ds.time[0] - ): - warnings.warn( - f"Start date {start_date} is before first available date {ds.time[0].values}" - ) - if ( - end_date is not None - and not isinstance(end_date, int) - and ds.time[-1] < np.datetime64(end_date) - ): - warnings.warn( - f"End date {end_date} is after last available date {ds.time[-1].values}" - ) - - # used when we need all the dates to calculate things like offset indices - self.time_da = ds.time.copy(deep=True) - # a list of dates that is available to fetch, used for things like inferencers. - # integer start/end dates are positional indices into the time array rather - # than dates themselves, so they need `isel` instead of the label-based `sel` - if isinstance(start_date, int) or isinstance(end_date, int): - self.times = self.time_da.isel(time=slice(start_date, end_date)) - else: - self.times = self.time_da.sel(time=slice(start_date, end_date)) - self.total_samples = self.times.shape[0] - - if start_date: - if isinstance(start_date, int): - self.start_index = start_date - else: - self.start_index = int( - np.where(self.time_da == np.datetime64(start_date))[0][0] - ) - else: - self.start_index = 0 - - # Validate time stepping - if (self.time_step % self.data_time_step).total_seconds() != 0: - raise ValueError( - f"'time_step' must be a multiple of 'data_time_step' " - f"(got {self.time_step} and {self.data_time_step}" - ) - if (self.gap % self.data_time_step).total_seconds() != 0: - raise ValueError( - f"'gap' must be a multiple of 'data_time_step' " - f"(got {self.gap} and {self.data_time_step}" - ) - self.interval = self.time_step // self.data_time_step - - # Verify timestep matches data - ds_dt = pd.Timedelta((self.time_da[1] - self.time_da[0]).values) - if not (ds_dt == self.data_time_step): - warnings.warn( - f"Dataset dt {ds_dt} doesn't match configuration dt {self.data_time_step}. " - "This could be a configuration error or a dataset mismatch." - ) - - # Find indices of init times for forecast mode - if self.forecast_mode: - if self.batch_size != 1: - self.batch_size = 1 - warnings.warn( - "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`; " - "setting it now" - ) - self._forecast_init_indices = np.array( - [ - int(np.where(self.time_da == s)[0][0]) - for s in self.forecast_init_times - ], - dtype="int", - ) - ((self.input_time_dim - 1) * self.interval) - else: - self._forecast_init_indices = None - - def _get_scaling_da(self) -> None: - """Setup data scaling parameters. - - Processes scaling configuration and sets up scaling parameters for: - - Input variables - - Target/output variables - - All variables combined - - Constant fields - - Raises - ------ - KeyError - If scaling parameters are missing for any variables - """ - scaling_df = pd.DataFrame.from_dict(self.scaling).T - scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} - scaling_da = scaling_df.to_xarray().astype("float32") - - try: - self.input_scaling = scaling_da.sel(index=self.input_variables).rename( - {"index": "channel_in"} - ) - self.input_scaling = { - "mean": np.expand_dims( - self.input_scaling["mean"].values.copy(), (0, 2, 3, 4) - ), - "std": np.expand_dims( - self.input_scaling["std"].values.copy(), (0, 2, 3, 4) - ), - } - except (ValueError, KeyError): - missing = [ - m for m in self.input_variables if m not in list(self.scaling.keys()) - ] - raise KeyError( - f"Input channels {missing} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" - ) - - try: - self.target_scaling = scaling_da.sel(index=self.output_variables).rename( - {"index": "channel_out"} - ) - self.target_scaling = { - "mean": np.expand_dims( - self.target_scaling["mean"].values.copy(), (0, 2, 3, 4) - ), - "std": np.expand_dims( - self.target_scaling["std"].values.copy(), (0, 2, 3, 4) - ), - } - except (ValueError, KeyError): - missing = [ - m for m in self.output_variables if m not in list(self.scaling.keys()) - ] - raise KeyError( - f"Target channels {missing} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" - ) - - self.all_scaling = scaling_da.sel(index=self.all_variables).rename( - {"index": "channel_in"} - ) - self.all_scaling = { - "mean": np.expand_dims( - self.all_scaling["mean"].values.copy(), (0, 2, 3, 4) - ), - "std": np.expand_dims(self.all_scaling["std"].values.copy(), (0, 2, 3, 4)), - } - - if self.constant_variables: - # Check that all constant variables are present in scaling data - missing_constants = [ - var - for var in self.constant_variables - if var not in list(self.scaling.keys()) - ] - if missing_constants: - raise KeyError( - f"Constant variables {missing_constants} not found in the scaling config dict data.scaling ({list(self.scaling.keys())})" - ) - - # no try/except needed here: the missing_constants check above - # already guarantees every name in self.constant_variables is - # present in self.scaling, which is what backs scaling_da's index - self.constant_scaling = scaling_da.sel( - index=self.constant_variables - ).rename({"index": "channel_out"}) - self.constant_scaling = { - "mean": np.expand_dims( - self.constant_scaling["mean"].values.copy(), (1, 2, 3) - ), - "std": np.expand_dims( - self.constant_scaling["std"].values.copy(), (1, 2, 3) - ), - } - - def get_constants(self) -> np.ndarray: - """Get constant fields used in dataset. - - Returns - ------- - np.ndarray - Array of constant fields with shape [F, C, H, W] - where F=faces, C=channels, H=height, W=width - """ - if self.constants is not None: - return self.constants - - if self.constant_variables is None: - return None - - const = np.asarray(self.ds["constants"][self.constant_variable_indices]) - - if self.constant_scaling: - const = (const - self.constant_scaling["mean"]) / self.constant_scaling[ - "std" - ] - - self.constants = np.transpose(const, axes=(1, 0, 2, 3)) - return self.constants - - def _get_time_index(self, item: int) -> Tuple[Tuple[int, int], int]: - """Get time indices for specified sample. - - Parameters - ---------- - item : int - Sample index - - Returns - ------- - Tuple[Tuple[int, int], int] - ((start_index, end_index), batch_size) - Time window indices and actual batch size - """ - window_start_index = ( - self._forecast_init_indices[item] - if self.forecast_mode - else item * self.batch_size + self.start_index - ) - window_max_index = ( - window_start_index + self._window_length - if self.forecast_mode - else (item + 1) * self.batch_size + self._window_length + self.start_index - ) - if not self.drop_last and window_max_index > self.total_samples: - batch_size = self.batch_size - (window_max_index - self.total_samples) - else: - batch_size = self.batch_size - return (window_start_index, window_max_index), batch_size - - def _get_forecast_sol_times(self, item: int) -> np.ndarray: - """Get times for calculating solar insolation. - - Parameters - ---------- - item : int - Sample index - - Returns - ------- - np.ndarray - Array of timestamps for insolation calculation - """ - time_index, _ = self._get_time_index(item) - if self.forecast_mode: - timedeltas = ( - np.array(self._input_indices[0] + self._output_indices[0]) - ) * self.data_time_step - return self.time_da[time_index[0]].values + timedeltas - return self.time_da[slice(*time_index)].values - - def __len__(self) -> int: - """Get number of samples available in the dataset based on - timedeltas, gaps, start and end dates. - - Returns - ------- - int - Total number of available samples - """ - if self.forecast_mode: - return len(self._forecast_init_indices) - length = (self.total_samples - self._window_length + 1) / self.batch_size - if self.drop_last: - return int(np.floor(length)) - return int(np.ceil(length)) - - @abstractmethod - def __getitem__( - self, item: int - ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]]: - """Get requested sample - must be implemented by subclasses. - - Parameters - ---------- - item : int - Sample index - - Returns - ------- - Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]] - In forecast mode: List of input arrays - In training mode: Tuple of (input arrays, target array) - - Input arrays are in order: - - Model inputs [B, F, T, C, H, W] - - Insolation (if enabled) [B, F, T, 1, H, W] - - Constants (if provided) [F, C, H, W] - - Additional data (in subclasses) - - Target array has shape [B, F, T, C, H, W] - where: - B = batch size - F = faces - T = time steps - C = channels - H = height - W = width - """ - pass # pragma: no cover - abstract method body, always overridden diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index 56d4abe268..a0087bed00 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -14,12 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Dataset for coupling time series healpix data with external inputs from various earth system components. - -This class extends the TimeSeriesDataset to add the core functionality for coupling time series healpix data with external inputs from various earth system components. -""" - from __future__ import annotations import gc @@ -40,13 +34,10 @@ from . import couplers from .timeseries_dataset import TimeSeriesDataset -logger = logging.getLogger(__name__) - -# xarray ships in the ``datapipes-extras`` optional dependency group; load it -# lazily so the physicsnemo import graph carries no static dependency on it -# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). xr = OptionalImport("xarray") +logger = logging.getLogger(__name__) + @dataclass class MetaData(DatapipeMetaData): @@ -178,10 +169,6 @@ def __init__( self.curr_item = None # keeps track of current initialization def _get_scaling_da(self): - """ - Get the scaling for the coupled values from the scaling dictionary. - This is overridden to add the scaling for the coupled values. - """ scaling_df = pd.DataFrame.from_dict(self.scaling).T scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} scaling_da = scaling_df.to_xarray().astype("float32") @@ -192,9 +179,6 @@ def _get_scaling_da(self): super()._get_scaling_da() def __getitem__(self, item): - """ - Get the item from the dataset. This is overridden to add the coupled values. - """ # start range torch.cuda.nvtx.range_push("CoupledTimeSeriesDataset:__getitem__") @@ -227,24 +211,9 @@ def __getitem__(self, item): axis=2, ) - # for models with extra outputs - # an empty couplings list means zero coupled variables, rather than - # indexing into a non-existent first coupling - num_coupled_vars = len(self.couplings[0].variables) if self.couplings else 0 - if len(self.ds["targets"].channel_out) != ( - len(self.ds["inputs"].channel_in) - num_coupled_vars - ): - channel_slice = ( - slice(None, -num_coupled_vars) if num_coupled_vars else slice(None) - ) - input_array = ( - input_array - self.input_scaling["mean"][:, channel_slice] - ) / self.input_scaling["std"][:, channel_slice] - else: - input_array = ( - input_array - self.input_scaling["mean"] - ) / self.input_scaling["std"] - + input_array = (input_array - self.input_scaling["mean"]) / self.input_scaling[ + "std" + ] if not self.forecast_mode: # BAD NEWS: Indexing the array as commented out below causes unexpected behavior in target creation. # leaving this in here as a warning @@ -320,7 +289,6 @@ def __getitem__(self, item): ) for c in self.couplings: for i, v in enumerate(c.variables): - # coupled fields are in the shape [T, B, C, F, H, W] integrated_couplings[i, :, :] += self.rng.normal( loc=0, scale=self.train_noise_params["couplings"][v]["std"], @@ -366,9 +334,6 @@ def __getitem__(self, item): return inputs_result, targets def next_integration(self, model_outputs, constants): - """ - Fetch the next coupled integration step - """ inputs_result = [] # grab last few model outputs for re-initialization diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py deleted file mode 100644 index dc4c8f7825..0000000000 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset_zarr.py +++ /dev/null @@ -1,392 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Dataset for coupling time series healpix data stored in Zarr format with external inputs from various earth system components. - -This class extends the TimeSeriesDatasetZarr to add the core functionality for coupling time series healpix data stored in Zarr format with external inputs from various earth system components. -""" - -import logging -from dataclasses import dataclass -from typing import List, Optional, Sequence, Tuple, Union - -import numpy as np -import pandas as pd -import torch -from omegaconf import DictConfig, OmegaConf - -from physicsnemo.core.version_check import OptionalImport -from physicsnemo.datapipes.meta import DatapipeMetaData -from physicsnemo.utils.insolation import insolation - -from . import couplers -from .base_timeseries_dataset_zarr import _check_availability -from .timeseries_dataset_zarr import TimeSeriesDatasetZarr - -logger = logging.getLogger(__name__) - -# zarr ships in the ``datapipes-extras`` optional dependency group; load it -# lazily so the physicsnemo import graph carries no static dependency on it -# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). -zarr = OptionalImport("zarr") - - -@dataclass -class MetaData(DatapipeMetaData): - """Metadata for this datapipe""" - - name: str = "CoupledTimeSeries" - # Optimization - auto_device: bool = False - cuda_graphs: bool = False - # Parallel - ddp_sharding: bool = False - - -class CoupledTimeSeriesDatasetZarr(TimeSeriesDatasetZarr): - """Dataset for coupling time series data with external earth system components. - - This class extends the base time series functionality to include coupling with - external data sources like ocean models, land models, etc. It supports: - - Integration of coupled data during training/inference - - Addition of training noise to improve generalization - - Step-by-step integration for forecasting - """ - - def __init__( - self, - dataset_path: str, - scaling: DictConfig, - input_variables: Sequence, - output_variables: Sequence = None, - constant_variables: Sequence = None, - input_time_dim: int = 1, - output_time_dim: int = 1, - data_time_step: Union[int, str] = "3h", - time_step: Union[int, str] = "6h", - gap: Union[int, str, None] = None, - batch_size: int = 32, - drop_last: bool = False, - add_insolation: bool = False, - forecast_init_times: Optional[Sequence] = None, - start_date: Optional[Union[int, str]] = None, - end_date: Optional[Union[int, str]] = None, - couplings: Sequence = [], - meta: DatapipeMetaData = MetaData(), - add_train_noise: bool = False, - train_noise_params: DictConfig = None, - train_noise_seed: int = 42, - ): - """Initialize coupled time series dataset. - - Parameters - ---------- - couplings : Sequence - List of coupling configurations, each containing: - - coupler: Name of coupler class to use - - params: Parameters for the coupler - add_train_noise : bool, default=False - Whether to add noise during training - train_noise_params : DictConfig, optional - Configuration for training noise, containing: - - inputs: Dict mapping variable names to noise std - - couplings: Dict mapping variable names to noise std - train_noise_seed : int, default=42 - Random seed for noise generation - - Other parameters are same as BaseTimeSeriesDatasetZarr. - See base class for detailed parameter descriptions. - """ - self.coupled_variables = [] - for c in couplings: - self.coupled_variables.append(c["params"]["variables"]) - - # We setup couplers first so superclass can properly initialize - # and set the scaling - _check_availability(dataset_path) - self.ds = zarr.open(dataset_path) - self.couplings = [ - getattr(couplers, c["coupler"])( - self.ds, - **OmegaConf.to_object(DictConfig(c))["params"], - ) - for c in couplings - ] - - super().__init__( - dataset_path=dataset_path, - scaling=scaling, - input_variables=input_variables, - output_variables=output_variables, - constant_variables=constant_variables, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - data_time_step=data_time_step, - time_step=time_step, - gap=gap, - batch_size=batch_size, - drop_last=drop_last, - add_insolation=add_insolation, - forecast_init_times=forecast_init_times, - start_date=start_date, - end_date=end_date, - add_train_noise=add_train_noise, - train_noise_params=train_noise_params, - train_noise_seed=train_noise_seed, - meta=meta, - ) - - # calculate static indices for coupling - for c in self.couplings: - c.compute_coupled_indices(self.interval, self.data_time_step) - # keep track of integration steps - self.integration_step = ( - 1 # starts at 1 because first step is done by __getitem__ - ) - self.last_batch = None # keeps track of batch info for coupler - self.curr_item = None # keeps track of current initialization - - def _get_scaling_da(self) -> None: - """extend scaling to include coupling-specific additions. - - Extends base scaling setup to: - - Create scaling parameters for coupled variables - - Pass scaling info to coupler objects - """ - scaling_df = pd.DataFrame.from_dict(self.scaling).T - scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} - scaling_da = scaling_df.to_xarray().astype("float32") - - for c in self.couplings: - c.set_scaling(scaling_da) - super()._get_scaling_da() - - def __getitem__( - self, item: int - ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]]: - """Get a batch of coupled time series data. - - This implementation extends the base time series data loading to include: - 1. Integration of coupled data sources - 2. Addition of training noise if enabled - 3. Tracking of batch info for step-by-step integration - - Parameters - ---------- - item : int - Sample index - - Returns - ------- - Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]] - In forecast mode: List of input arrays - In training mode: Tuple of (input arrays, target array) - - Input arrays are in order: - - Model inputs [B, F, T, C, H, W] - - Insolation (if enabled) [B, F, T, 1, H, W] - - Constants (if provided) [F, C, H, W] - - Coupled inputs (if provided) [T, B, C, F, H, W] - - Target array has shape [B, F, T, C, H, W] - where: - B = batch size - F = faces - T = time steps - C = channels - H = height - W = width - - Raises - ------ - IndexError - If item is out of range - """ - - # check for valid item and adjust for negative indices - if item < 0: - item = len(self) + item - if item < 0 or item > len(self): - raise IndexError( - f"index {item} out of range for dataset with length {len(self)}" - ) - - # start range - torch.cuda.nvtx.range_push("CoupledTimeSeriesDataset:__getitem__") - - if self.forecast_mode: - inputs_result = super().__getitem__(item) - else: - inputs_result, targets = super().__getitem__(item) - - # used by the couplers to determine what time index to load - # see method "next_integration()" for details - time_index, this_batch = self._get_time_index(item) - batch = {"time": slice(*time_index)} - self.last_batch = batch - - torch.cuda.nvtx.range_push( - "CoupledTimeSeriesDataset:__getitem__:retrieve_coupled" - ) - # retrieve coupled inputs - if len(self.couplings) > 0: - integrated_couplings = np.concatenate( - [ - c.construct_integrated_couplings(batch, this_batch) - for c in self.couplings - ], - axis=2, - ) - torch.cuda.nvtx.range_pop() # CoupledTimeSeriesDataset:__getitem__:retrieve_coupled - - # Insolation - if self.add_insolation: - # update current item and reset integration_step counter for further integrations which need - # insolation but bypass this method see method "next_integration()" for details - self.curr_item = item - self.integration_step = 1 - - if not self.forecast_mode and self.add_train_noise: - torch.cuda.nvtx.range_push( - "CoupledTimeSeriesDataset:__getitem__:add_train_noise" - ) - for c in self.couplings: - for i, v in enumerate(c.variables): - # coupled fields are in the shape [T, B, C, F, H, W] - integrated_couplings[i, :, :] += self.rng.normal( - loc=0, - scale=self.train_noise_params["couplings"][v]["std"], - size=integrated_couplings[i, :, :].shape, - ) - torch.cuda.nvtx.range_pop() # CoupledTimeSeriesDataset:__getitem__:add_train_noise - - # append integrated couplings - if len(self.couplings) > 0: - inputs_result.append(integrated_couplings) - - torch.cuda.nvtx.range_pop() # CoupledTimeSeriesDataset:__getitem__ - if self.forecast_mode: - return inputs_result - - return inputs_result, targets - - def _get_next_insolation(self, time_offset: int) -> torch.Tensor: - """Calculate insolation for next integration step. - - Parameters - ---------- - time_offset : int - Time offset for insolation calculation - - Returns - ------- - torch.Tensor - Insolation tensor - """ - sol = torch.tensor( - insolation( - self._get_forecast_sol_times(self.curr_item) + time_offset, - self.lat, - self.lon, - )[:, None] - ) - decoder_inputs = np.empty( - (1, self.input_time_dim + self.output_time_dim, 1) + self.spatial_dims, - dtype="float32", - ) - decoder_inputs[0] = sol - return torch.tensor(decoder_inputs.transpose(0, 3, 1, 2, 4, 5)) - - def _get_next_couplings(self, offset: int) -> Optional[torch.Tensor]: - """Get coupled inputs for next integration step. - - Parameters - ---------- - offset : int - Integration offset - - Returns - ------- - Optional[torch.Tensor] - Coupled inputs tensor if couplings exist, None otherwise - """ - if not self.couplings: - return None - - # account for multiple integration steps without fetching data - batch = self.last_batch.copy() - batch = { - "time": slice( - self.last_batch["time"].start + offset, - self.last_batch["time"].stop + offset, - self.last_batch["time"].step, - ) - } - integrated_couplings = np.concatenate( - [ - c.construct_integrated_couplings( - batch=batch, - ) - for c in self.couplings - ], - axis=2, - ) - - return torch.tensor(integrated_couplings) - - def next_integration( - self, model_outputs: torch.Tensor, constants: torch.Tensor - ) -> List[torch.Tensor]: - """Get inputs for next integration step with coupling data. - - Parameters - ---------- - model_outputs : torch.Tensor - Model outputs from previous step [B, F, T, C, H, W] - constants : torch.Tensor - Constant fields [F, C, H, W] - - Returns - ------- - List[torch.Tensor] - List of input tensors for next step - """ - inputs_result = [] - - # Get prognostic inputs from model outputs - init_time_dim = len(self._input_indices[0]) - prognostic_inputs = model_outputs[:, :, -init_time_dim:] - inputs_result.append(prognostic_inputs) - - # Add insolation if needed - if self.add_insolation: - time_offset = self.time_step * self.output_time_dim * self.integration_step - inputs_result.append(self._get_next_insolation(time_offset)) - - # Add constants - inputs_result.append(constants) - - # Add coupled inputs if any - offset = self.interval * self.output_time_dim * self.integration_step - coupled_inputs = self._get_next_couplings(offset) - if coupled_inputs is not None: - inputs_result.append(coupled_inputs) - - # Increment integration step - self.integration_step += 1 - - return inputs_result diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 32422fc051..a0960b2e6d 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -17,30 +17,25 @@ from __future__ import annotations import logging -from abc import ABC, abstractmethod from typing import Sequence -import cftime import numpy as np import pandas as pd import torch as th from physicsnemo.core.version_check import OptionalImport -logger = logging.getLogger(__name__) - -# xarray/zarr ship in the ``datapipes-extras`` optional dependency group; load -# them lazily so the physicsnemo import graph carries no static dependency on -# them (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). xr = OptionalImport("xarray") -zarr = OptionalImport("zarr") + +logger = logging.getLogger(__name__) -class BaseCoupler(ABC): +class ConstantCoupler: """ - Base class for couplers used to interface two components of earth system. + coupler used to interface two component of earth system - This class contains common functionality shared by different coupler implementations. + constant coupler will take the the coupled field at integration time and + force the model with this field consistently """ def __init__( @@ -52,7 +47,7 @@ def __init__( input_time_dim: int = 2, output_time_dim: int = 2, input_times: Sequence = [pd.Timedelta("24h"), pd.Timedelta("48h")], - prepared_coupled_data: bool = True, + prepared_coupled_data=True, ): """ Parameters @@ -64,8 +59,7 @@ def __init__( forecasting batch size should be 1 variables: Sequence sequence of strings that indicate the coupled variable - names in the dataset. All names should be in the dataset with - an optional time component at the end, eg ttr-48h + names in the dataset presteps: int, optional the number of model steps used to initialize the hidden state. If not using a GRU, prestep is 0, default 0 @@ -77,15 +71,15 @@ def __init__( sequence of pandas Timedelta objects that indicate which times are to be coupled, default [pd.Timedelta("24h"), pd.Timedelta("48h")] prepared_coupled_data: boolean, optional - If True assumes data in dataset has been prepared appropriately for training: + If True assumes data in dataset has been prepared approiately for training: averages have already been calculated so that each time step denotes the right side of a averaging_window window. - This is highly recommended for training, default True + This is highly remcommended for training, default True """ # extract important meta data from ds self.ds = dataset self.batch_size = batch_size - self.spatial_dims = self.ds["inputs"].shape[2:] + self.spatial_dims = self.ds.inputs.shape[2:] self.variables = variables self.presteps = presteps self.input_time_dim = input_time_dim @@ -95,54 +89,38 @@ def __init__( self.output_channels = len(self.variables) * len(self.input_times) self.timevar_dim = self._compute_timevar_dim() self.coupled_inputs_shape = None - self.coupled_scaling = None + self.scaling_dict = None self._coupled_offsets = None self.coupled_mode = False self.integrated_couplings = None - self.ds_variable_indices = [] if not prepared_coupled_data: - raise NotImplementedError("Data preparation not yet implemented") - - if isinstance(self.ds, xr.Dataset): - self.use_zarr = False - elif isinstance(self.ds, zarr.Group): - self.use_zarr = True - # iterate over self.variables in the outer loop so the selected - # indices follow the order of self.variables (matching the xarray - # path's `.sel(channel_in=self.variables)`), not the dataset's - # native channel_in order. Downstream code (e.g. coupled_scaling, - # built from self.variables order in set_scaling) assumes the - # coupled channel axis is ordered by self.variables. - channel_in = list(self.ds["channel_in"]) - self.ds_variable_indices = [ - i for v in self.variables for i, ic in enumerate(channel_in) if ic == v - ] + logger.log( + logging.DEBUG, + "Assuming coupled data is not preprocessed preparing data.", + ) + self._prepare_coupled_data() else: - raise TypeError( - f"Coupler only supports xarray Datasets or zarr Groups, got {type(self.ds)}" + logger.log( + logging.DEBUG, + "**Assuming coupled data has been prepared properly, using coupled field[s] from " + 'dataset "as-is"**', ) + def _prepare_coupled_data(self): + # TODO: write function to lazily compute average as spcified in time scheme + raise NotImplementedError("Data preparation not yet implemented") + def _compute_coupled_integration_dim(self): - """ - Compute the coupled integration steps across the time dimension. - This is the number of presteps plus the maximum number of output times divided by the number of input times. - """ return self.presteps + max(self.output_time_dim // self.input_time_dim, 1) def _compute_timevar_dim(self): - """ - Compute the time variable dimension. - Multiple time integrations are stacked along the time dimension. - Thus it is the number of input times multiplied by the number of variables. - """ return len(self.input_times) * len(self.variables) - @abstractmethod def compute_coupled_indices(self, interval, data_time_step): """ - Called by CoupledDataset to compute static indices for training samples. - Must be implemented by subclasses as the logic varies between coupler types. + Called by CoupledDataset to compute static indices for training + samples Parameters ---------- @@ -151,25 +129,28 @@ def compute_coupled_indices(self, interval, data_time_step): data_time_step: dataset timestep """ - pass # pragma: no cover - abstract method body, always overridden + # create array of static coupled offstes that accompany each batch + self._coupled_offsets = np.empty( + [self.batch_size, self.coupled_integration_dim, len(self.input_times)] + ) + for b in range(self.batch_size): + for i in range(self.coupled_integration_dim): + self._coupled_offsets[b, i, :] = b + np.array( + [ts / data_time_step for ts in self.input_times] + ) + + self._coupled_offsets = self._coupled_offsets.astype(int) def set_scaling(self, scaling_da): """ - Called by CoupledDataset to compute static indices for training samples + Called by CoupledDataset to compute static indices for training + samples Parameters ---------- scaling_da: xarray.DataArray values used to scale input data, uses mean and std """ - # verify all the channels are there for scaling, this avoids an opaque - # "not all values found in index 'index'"" error that looks like its from hydra - missing_channels = set(self.variables) - set(scaling_da.index.values) - if len(missing_channels) > 0: - raise KeyError( - f"Coupled variable(s) not found in scaling values: {missing_channels}" - ) - coupled_scaling = scaling_da.sel(index=self.variables).rename( {"index": "channel_in"} ) @@ -179,58 +160,34 @@ def set_scaling(self, scaling_da): } def setup_coupling(self, coupled_module): - """ - Sets up the coupling between the coupled variables and the provided module - - Parameters - ---------- - coupled_module: physicsnemo.datapipes.healpix.TimeSeriesDataset - The module which this coupler will be coupled against. - """ - # To expedite the coupling process the coupled_forecast + # To expediate the coupling process the coupled_forecast # get proper channels from coupled component output output_channels = coupled_module.output_variables - # A bit convoluted. Some variable names are present in the dataset as is, - # Some prepared coupled variables are given a suffix for training associated - # with a time increment suach as a trailing average increment e.g. 'z1000-48H'. - # Some variables may have an additional suffix, e.g. 'z1000-3H-48H'. The final - # suffix (if it exists) is used to determine the coupling increment. + # A bit convoluted. Prepared coupled variables + # are given a suffix for training associated with their + # trailing average increment e.g. 'z1000-48H'. To extract + # thr proper field from the coupled model output, we see if + # we check if the coupled model output var is in self.variables. + # + # for example 'z1000' is in 'z1000-48H' channel_indices = [ i for i, oc in enumerate(output_channels) for v in self.variables - # extract everthing before the last "-" if there is one in the name - if (("-" not in v and oc == v) or (oc == "-".join(v.split("-")[:-1]))) + if oc == v.split("-")[0] ] - # check for missing variables - if len(self.variables) != len(channel_indices): - found_channels = [ - oc - for oc in output_channels - for v in self.variables - # extract everthing before the last - - if (("-" not in v and oc == v) or (oc == "-".join(v.split("-")[:-1]))) - ] - missing_channels = set(self.variables) - set(found_channels) - raise ValueError(f"Missing variables in coupled module: {missing_channels}") self.coupled_channel_indices = channel_indices def reset_coupler(self): - """Clear saved coupled fields, forces recomputation on next use. - - This method is called when the dataloader is reset, and it is used to clear the - cached coupled fields. This is necessary because the coupled fields are computed - on the fly, and they are not saved to the dataset. - """ self.coupled_mode = False self.integrated_couplings = None self.preset_coupled_fields = None - @abstractmethod def set_coupled_fields(self, coupled_fields: th.tensor): """ Set the data for the coupled field for the next iteration of the dataloader. - Must be implemented by subclasses as the processing logic varies. + Instead of loading data from the dataset the data from coupled_fields will + be returned instead. Parameters ---------- @@ -238,37 +195,26 @@ def set_coupled_fields(self, coupled_fields: th.tensor): The data to use when the dataloader requests coupled fields. Expected format is [B, F, T, C, H, W] """ - pass # pragma: no cover - abstract method body, always overridden - - def _construct_integrated_couplings_from_dataset(self, batch, bsize): - """ - Common logic for constructing integrated couplings from dataset. - Used by both ConstantCoupler and TrailingAverageCoupler. - """ - # reset integrated couplings - self.integrated_couplings = np.empty( - (bsize, self.coupled_integration_dim, self.timevar_dim) + self.spatial_dims - ) - - index_range = slice( - batch["time"].start, - batch["time"].start + self._coupled_offsets[-1, -1, -1] + 1, - ) - - # extract coupled variables - if self.use_zarr: - # Loading the contiguous time slice into memory and then pulling out the semi-random - # variable indices is quicker than trying to do this all at once. - ds_index_range = self.ds["inputs"][index_range] - ds_index_range = ds_index_range[:, self.ds_variable_indices] - else: - ds_index_range = ( - self.ds.inputs.sel(channel_in=self.variables) - .isel(time=index_range) - .compute() + if coupled_fields.shape[0] != self.batch_size: + raise ValueError( + f"Batch size of coupled field {coupled_fields.shape[0]} doesn't " + f" match configured batch size {self.batch_size}" ) - - return ds_index_range + # create buffer for coupling + coupled_fields = coupled_fields[ + :, :, :, self.coupled_channel_indices, :, : + ].permute(2, 0, 3, 1, 4, 5) + self.preset_coupled_fields = th.empty( + [self.coupled_integration_dim, self.batch_size, self.timevar_dim] + + list(self.spatial_dims) + ) + # we use a constant set of values so we just copy time 0 + for i in range(len(self.preset_coupled_fields)): + self.preset_coupled_fields[i, :, :, :, :, :] = coupled_fields[ + 0, :, -1, :, :, : + ] + # flag for construct integrated coupling method to use this array + self.coupled_mode = True def construct_integrated_couplings( self, @@ -293,162 +239,48 @@ def construct_integrated_couplings( if self.coupled_mode: return self.preset_coupled_fields else: - if (batch is None) or (bsize is None): - raise ValueError( - "batch and bsize must be provided when not in coupled_mode" - ) + # reset integrated couplings + self.integrated_couplings = np.empty( + (bsize, self.coupled_integration_dim, self.timevar_dim) + + self.spatial_dims + ) - ds_index_range = self._construct_integrated_couplings_from_dataset( - batch, bsize + index_range = slice( + batch["time"].start, + batch["time"].start + self._coupled_offsets[-1, -1, -1] + 1, ) - # Apply scaling if available - if self.coupled_scaling is not None: - ds_index_range -= self.coupled_scaling["mean"] - ds_index_range /= self.coupled_scaling["std"] + # extract coupled variables and scale lazily + input_array = self.ds.inputs.sel(channel_in=self.variables) + ds = (input_array - self.coupled_scaling["mean"]) / self.coupled_scaling[ + "std" + ] + # load before entering loop for efficiency + ds_index_range = ds.isel(time=index_range).load() # use static offsets to create integrated coupling array for b in range(bsize): for i in range(self.coupled_integration_dim): - if self.use_zarr: - coupling_temp = ds_index_range[ - self._coupled_offsets[b, i, :], : - ] - else: - coupling_temp = ds_index_range.isel( - time=self._coupled_offsets[b, i, :] - ).to_numpy() - self.integrated_couplings[b, i, :, :, :] = coupling_temp.reshape( - (self.timevar_dim,) + coupling_temp.shape[2:] + coupling_temp = ds_index_range.isel( + time=self._coupled_offsets[b, i, :] ) - # permute to have the time dimension first [T, B, C, F, H, W] - # rather than [B, F, T, C, H, W], and cast to float for compatibility + self.integrated_couplings[b, i, :, :, :] = ( + coupling_temp.to_numpy().reshape( + (self.timevar_dim,) + coupling_temp.shape[2:] + ) + ) + return self.integrated_couplings.transpose((1, 0, 2, 3, 4, 5)).astype( "float32" - ) - - -class ConstantCoupler(BaseCoupler): - """ - coupler used to interface two component of earth system - - constant coupler will take the the coupled field at integration time and - force the model with this field consistently - """ - - def __init__( - self, - dataset: xr.Dataset, - batch_size: int, - variables: Sequence, - presteps: int = 0, - input_time_dim: int = 2, - output_time_dim: int = 2, - input_times: Sequence = [pd.Timedelta("24h"), pd.Timedelta("48h")], - prepared_coupled_data=True, - ): - """ - Parameters - ---------- - dataset: xr.Dataset - xarray Dataset that holds coupled data - batch_size: int - number of batch size during training. - forecasting batch size should be 1 - variables: Sequence - sequence of strings that indicate the coupled variable - names in the dataset - presteps: int, optional - the number of model steps used to initialize the hidden state. - If not using a GRU, prestep is 0, default 0 - input_time_dim: int, optional - number of input times into the model, default 2 - output_time_dim: int, optional - number of output times for each model step, default 2 - input_times: Sequence, optional - sequence of pandas Timedelta objects that indicate which times are to be coupled, - default [pd.Timedelta("24h"), pd.Timedelta("48h")] - prepared_coupled_data: boolean, optional - If True assumes data in dataset has been prepared appropriately for training: - averages have already been calculated so that each time step denotes - the right side of a averaging_window window. - This is highly recommended for training, default True - """ - super().__init__( - dataset=dataset, - batch_size=batch_size, - variables=variables, - presteps=presteps, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - input_times=input_times, - prepared_coupled_data=prepared_coupled_data, - ) - - def compute_coupled_indices(self, interval, data_time_step): - """ - Called by CoupledDataset to compute static indices for training - samples - - Parameters - ---------- - interval: int - ratio of dataset timestep to model dt - data_time_step: - dataset timestep - """ - # create array of static coupled offsets that accompany each batch - self._coupled_offsets = np.empty( - [self.batch_size, self.coupled_integration_dim, len(self.input_times)] - ) - for b in range(self.batch_size): - for i in range(self.coupled_integration_dim): - self._coupled_offsets[b, i, :] = b + np.array( - [ts / data_time_step for ts in self.input_times] - ) - - self._coupled_offsets = self._coupled_offsets.astype(int) - - def set_coupled_fields(self, coupled_fields: th.tensor): - """ - Set the data for the coupled field for the next iteration of the dataloader. - Instead of loading data from the dataset the data from coupled_fields will - be returned instead. + ) # cast to float for compatability - Parameters - ---------- - coupled_fields: th.tensor - The data to use when the dataloader requests coupled fields. Expected - format is [B, F, T, C, H, W] - """ - # create buffer for coupling - coupled_fields = coupled_fields[:, :, :, self.coupled_channel_indices, :, :] - self.preset_coupled_fields = th.empty( - [ - coupled_fields.shape[0], - self.spatial_dims[0], - self.coupled_integration_dim, - self.timevar_dim, - ] - + list(self.spatial_dims[1:]) - ) - # broadcast the first time step to all the integration steps - self.preset_coupled_fields[:, :, :, :, :, :] = coupled_fields[:, :, :1, :, :, :] - # permute to have the time dimension first [T, B, C, F, H, W] - # rather than [B, F, T, C, H, W] - self.preset_coupled_fields = self.preset_coupled_fields.permute( - 2, 0, 3, 1, 4, 5 - ) - # flag for construct integrated coupling method to use this array - self.coupled_mode = True - -class TrailingAverageCoupler(BaseCoupler): +class TrailingAverageCoupler: """ - coupler used to interface two components of the earth system + coupler used to inferface two components of the earth system Trailing average coupler uses coupled input times as the right side of - an average that is taken over an "averaging_window" window size. + an averag that is taken over an "averaging_window" window size. """ def __init__( @@ -487,49 +319,55 @@ def __init__( sequence of pandas Timedelta objects that indicate which times are to be coupled, default [pd.Timedelta("24h"), pd.Timedelta("48h")] prepared_coupled_data: boolean, optional - If True assumes data in dataset has been prepared appropriately for training: + If True assumes data in dataset has been prepared approiately for training: averages have already been calculated so that each time step denotes the right side of a averaging_window window. - This is highly recommended for training, default True + This is highly remcommended for training, default True """ - super().__init__( - dataset=dataset, - batch_size=batch_size, - variables=variables, - presteps=presteps, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - input_times=input_times, - prepared_coupled_data=prepared_coupled_data, - ) - - # TrailingAverageCoupler-specific attributes + # extract important meta data from ds + self.ds = dataset + self.batch_size = batch_size + self.spatial_dims = self.ds.inputs.shape[2:] + self.variables = variables + self.presteps = presteps + self.input_time_dim = input_time_dim + self.output_time_dim = output_time_dim + self.input_times = [pd.Timedelta(t) for t in input_times] self.averaging_window = pd.Timedelta(averaging_window) + self.output_channels = len(self.variables) * len(self.input_times) + self._set_time_increments() + self.coupled_integration_dim = self._compute_coupled_integration_dim() + self.timevar_dim = self._compute_timevar_dim() + self.coupled_inputs_shape = None + self.scaling_dict = None + self._coupled_offsets = None + self.integrated_couplings = None + self.coupled_mode = False # if forecasting with another coupled model - if self.use_zarr: - cf_dates = cftime.num2pydate( - self.ds["time"][:], - units=self.ds["time"].attrs["units"], - calendar=self.ds["time"].attrs["calendar"], + if not prepared_coupled_data: + logger.log( + logging.DEBUG, + "Assuming coupled data is not preprocessed, averaging fields in as designed in" + "TrailingAverageCoupler. See docs for specifics.", ) - dates = [np.datetime64(date.isoformat()) for date in cf_dates] - self.time_da = np.asarray(dates) + self._prepare_coupled_data() else: - self.time_da = self.ds.time.values + logger.log( + logging.DEBUG, + "**Assuming coupled data has been prepared properly, using coupled field[s] from" + 'dataset "as-is"**', + ) def compute_coupled_indices(self, interval, data_time_step): """ Called by CoupledDataset to compute static indices for training samples - Parameters - ---------- - interval: int - ratio of dataset timestep to model dt - data_time_step: - dataset timestep + :param interval: int ratio of dataset timestep to model dt + :param data_time_step: dataset timestep """ - # create array of static coupled offsets that accompany each batch + + # create array of static coupled offstes that accompany each batch self._coupled_offsets = np.empty( [self.batch_size, self.coupled_integration_dim, len(self.input_times)] ) @@ -543,16 +381,58 @@ def compute_coupled_indices(self, interval, data_time_step): self._coupled_offsets = self._coupled_offsets.astype(int) + def _prepare_coupled_data(self): + # TODO: write function to lazily compute average as spcified in time scheme + raise NotImplementedError("Data preparation not yet implemented") + + def set_scaling(self, scaling_da): + coupled_scaling = scaling_da.sel(index=self.variables).rename( + {"index": "channel_in"} + ) + self.coupled_scaling = { + "mean": np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)), + "std": np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)), + } + + def _set_time_increments(self): + # get the dt of the dataset + dt = pd.Timedelta( + self.ds.time[1].values - self.ds.time[0].values + ).total_seconds() + # assert that the time increments are divisible by the dt of the dataset + if np.any([t.total_seconds() % dt != 0 for t in self.input_times]): + raise ValueError( + f"Coupled input times {self.input_times} " + f"({[t.total_seconds() for t in self.input_times]} in secs) are not divisible by dataset dt: {dt}" + ) + self.time_increments = [t.total_seconds() / dt for t in self.input_times] + + def _compute_timevar_dim(self): + return len(self.input_times) * len(self.variables) + + def _compute_coupled_integration_dim(self): + return self.presteps + max(self.output_time_dim // self.input_time_dim, 1) + def setup_coupling(self, coupled_module): - """ - Setup the trailing average specific coupling between the coupled variables and the provided module. - Precomputes the averaging slices for the coupled variables. - """ - # Call parent method first to set basic coupling - super().setup_coupling(coupled_module) + # To expediate the coupling process the coupled_forecast + # get proper channels from coupled component output + output_channels = coupled_module.output_variables + # A bit convoluted. Prepared coupled variables + # are given a suffix for training associated with their + # trailing average increment e.g. 'z1000-48H'. To extract + # thr proper field from the coupled model output, we see if + # we check if the coupled model output var is in self.variables. + # + # for example 'z1000' is in 'z1000-48H' + channel_indices = [ + i + for i, oc in enumerate(output_channels) + for v in self.variables + if oc == v.split("-")[0] + ] + self.coupled_channel_indices = channel_indices - # TrailingAverageCoupler-specific setup - # find averaging periods from component output + # find averaging periods from componenet output averaging_window_max_indices = [ i // pd.Timedelta(coupled_module.time_step) for i in self.input_times ] @@ -570,7 +450,12 @@ def setup_coupling(self, coupled_module): ) self.averaging_slices = averaging_slices - def set_coupled_fields(self, coupled_fields: th.tensor): + def reset_coupler(self): + self.coupled_mode = False + self.integrated_couplings = None + self.preset_coupled_fields = None + + def set_coupled_fields(self, coupled_fields): """ Set the data for the coupled field for the next iteration of the dataloader. Instead of loading data from the dataset the data from coupled_fields will @@ -582,6 +467,12 @@ def set_coupled_fields(self, coupled_fields: th.tensor): The data to use when the dataloader requests coupled fields. Expected format is [B, F, T, C, H, W] """ + if coupled_fields.shape[0] != self.batch_size: + raise ValueError( + f"Batch size of coupled field {coupled_fields.shape[0]} doesn't " + f" match configured batch size {self.batch_size}" + ) + coupled_fields = coupled_fields[:, :, :, self.coupled_channel_indices, :, :] # TODO: Now support output_time_dim =/= input_time_dim, but presteps need to be 0, will add support for presteps>0 coupled_averaging_periods = [] @@ -591,11 +482,67 @@ def set_coupled_fields(self, coupled_fields: th.tensor): for s in self.averaging_slices[j] ] coupled_averaging_periods.append(th.concat(averaging_periods, dim=3)) - self.preset_coupled_fields = th.concat(coupled_averaging_periods, dim=2) - # permute to have the time dimension first [T, B, C, F, H, W] - # rather than [B, F, T, C, H, W] - self.preset_coupled_fields = self.preset_coupled_fields.permute( - 2, 0, 3, 1, 4, 5 - ) + self.preset_coupled_fields = th.concat( + coupled_averaging_periods, dim=2 + ).permute(2, 0, 3, 1, 4, 5) # flag for construct integrated coupling method to use this array self.coupled_mode = True + + def construct_integrated_couplings( + self, + batch=None, + bsize=None, + ): + """ + Construct array of coupled inputs that includes values required for + model integration steps. + + :param batch: indices of dataset sample dimension associated with + current batch. + :param bsize: int batch size + """ + + if self.coupled_mode: + return self.preset_coupled_fields + else: + if (batch is None) or (bsize is None): + raise ValueError( + "Coupled fields must be set if no batch or batch size is provided" + ) + # reset integrated couplings + self.integrated_couplings = np.empty( + (bsize, self.coupled_integration_dim, self.timevar_dim) + + self.spatial_dims + ) + + index_range = slice( + batch["time"].start, + batch["time"].start + self._coupled_offsets[-1, -1, -1] + 1, + ) + + # extract coupled variables and scale lazily + ds = ( + self.ds.inputs.sel(channel_in=self.variables) + - self.coupled_scaling["mean"] + ) / self.coupled_scaling["std"] + # load before entering loop for efficiency + ds_index_range = ds.isel(time=index_range).load() + + # use static offsets to create integrated coupling array + for b in range(bsize): + for i in range(self.coupled_integration_dim): + # coupling_temp = \ + # ds.isel(time=batch["time"].start+self._coupled_offsets[b,i,:]) # changed i to 0 for debugging + # added to test speed, original is commented above + coupling_temp = ds_index_range.isel( + time=self._coupled_offsets[b, i, :] + ) + self.integrated_couplings[b, i, :, :, :] = ( + coupling_temp.to_numpy().reshape( + (self.timevar_dim,) + coupling_temp.shape[2:] + ) + ) + + return self.integrated_couplings.transpose((1, 0, 2, 3, 4, 5)).astype( + "float32" + ) # cast to float32 for pytroch compatibility. diff --git a/physicsnemo/datapipes/healpix/data_modules.py b/physicsnemo/datapipes/healpix/data_modules.py index c010a493a4..8d955bf50c 100644 --- a/physicsnemo/datapipes/healpix/data_modules.py +++ b/physicsnemo/datapipes/healpix/data_modules.py @@ -14,30 +14,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -TimeSeriesDataModule - DataModule that sets up the dataloaders for the time series healpix data. - -This class provides the core functionality for setting up the dataloaders for the time series healpix data. -Depending on the splits and forecast_init_times, it will set up the appropriate dataloaders for the training, validation, and test sets. -It also supports coupling the time series data with external inputs from various earth system components. - -The main classes are: -- TimeSeriesDataModule: A DataModule for loading and processing time series healpix data. -- CoupledTimeSeriesDataModule: A DataModule for loading and processing coupled time series healpix data. -""" - -from __future__ import annotations - # System modules import logging -import warnings +import os +import time from pathlib import Path -from typing import Optional, Sequence, Union +from typing import DefaultDict, Optional, Sequence, Union # numpy import numpy as np -# External modules +# distributed stuff +import torch from omegaconf import DictConfig from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler @@ -48,17 +36,155 @@ from .coupledtimeseries_dataset import CoupledTimeSeriesDataset from .timeseries_dataset import TimeSeriesDataset +xr = OptionalImport("xarray") + logger = logging.getLogger(__name__) -# xarray ships in the ``datapipes-extras`` optional dependency group; load it -# lazily so the physicsnemo import graph carries no static dependency on it -# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). -xr = OptionalImport("xarray") + +def _get_file_name(path, prefix, var, suffix): + """ + Helper that returns a fully formed path for a given variable + + Parameters + ---------- + path: str + The base path where the file is located + prefix: str + The prefix used for the filename + var: str + The variable stored in the file + suffix: str + The suffix used for the files + + Returns + ------- + str: The fully formed path + """ + return os.path.join(path, f"{prefix}{var}{suffix}.nc") + + +def open_time_series_dataset_classic_on_the_fly( + directory: str, + input_variables: Sequence, + output_variables: Optional[Sequence], + constants: Optional[DefaultDict] = None, + prefix: Optional[str] = None, + suffix: Optional[str] = None, + batch_size: int = 32, + scaling: Optional[DictConfig] = None, +) -> "xr.Dataset": + """ + Opens and merges multiple datasets that that contain individual variables + into a single dataset + + Parameters + ---------- + directory: str + The directory that contains the input datasets + input_variables: Sequence + The input variables to be merged into the new dataset + output_variables: Sequence, optional + The output variables to be merged into the new dataset + If no output variables are provided the input set is used + constants: DefaultDict, optional + A set of constants to add to the merged dataset + default None + prefix: str, optional + The prefix of the input datasets, default None + suffix: str, optional + The suffix of the input datasets, default None + batch_size: str, optional + The chunk size to use for the input datasets, default 32 + scaling: DictConfig, optional + Not used for open_time_series_dataset_classic_on_the_fly + + Returns + ------- + xr.Dataset: The merged dataset + """ + output_variables = output_variables or input_variables + all_variables = np.union1d(input_variables, output_variables) + prefix = prefix or "" + suffix = suffix or "" + + merge_time = time.time() + logger.info("merging input datasets") + + datasets = [] + remove_attrs = ["mean", "std"] if "LL" in prefix else ["varlev", "mean", "std"] + for variable in all_variables: + file_name = _get_file_name(directory, prefix, variable, suffix) + logger.debug("open nc dataset %s", file_name) + + ds = xr.open_dataset(file_name, autoclose=True) + + if "LL" in prefix: + ds = ds.rename({"lat": "height", "lon": "width"}) + ds = ds.isel({"height": slice(0, 180)}) + try: + ds = ds.isel(varlev=0) + except ValueError: + pass + + # remove unused + for attr in remove_attrs: + if attr in ds.indexes or attr in ds.variables: + ds = ds.drop(attr) + + # Rename variable + if "sample" in ds.variables or "sample" in ds.dims: + ds = ds.rename({"sample": "time"}) + + ds = ds.chunk({"time": batch_size}) + + # Change lat/lon to coordinates + try: + ds = ds.set_coords(["lat", "lon"]) + except (ValueError, KeyError): + pass + datasets.append(ds) + # Merge datasets + data = xr.merge(datasets, compat="override") + + # Convert to input/target array by merging along the variables + input_da = ( + data[list(input_variables)] + .to_array("channel_in", name="inputs") + .transpose("time", "channel_in", "face", "height", "width") + ) + target_da = ( + data[list(output_variables)] + .to_array("channel_out", name="targets") + .transpose("time", "channel_out", "face", "height", "width") + ) + + result = xr.Dataset() + result["inputs"] = input_da + result["targets"] = target_da + + # Get constants + if constants is not None: + constants_ds = [] + for name, var in constants.items(): + constants_ds.append( + xr.open_dataset( + _get_file_name(directory, prefix, name, suffix), autoclose=True + ).set_coords(["lat", "lon"])[var] + ) + constants_ds = xr.merge(constants_ds, compat="override") + constants_da = constants_ds.to_array("channel_c", name="constants").transpose( + "channel_c", "face", "height", "width" + ) + result["constants"] = constants_da + + logger.info("merged datasets in %0.1f s", time.time() - merge_time) + + return result def open_time_series_dataset_classic_prebuilt( directory: str, dataset_name: str, constants: bool = False, batch_size: int = 32 -) -> xr.Dataset: +) -> "xr.Dataset": """ Opens an existing dataset @@ -68,7 +194,7 @@ def open_time_series_dataset_classic_prebuilt( The directory that contains the dataset dataset_name: str The name of the dataset to open - constants: typing.DefaultDict, optional + constants: DefaultDict, optional Not used for open_time_series_dataset_classic_prebuilt, default False batch_size: str, optional The chunk size to use for the input datasets, default 32 @@ -76,21 +202,164 @@ def open_time_series_dataset_classic_prebuilt( Returns ------- xr.Dataset: The opened dataset - - Raises - ------ - FileNotFoundError: If the dataset doesn't exist at the specified path """ ds_path = Path(directory, dataset_name + ".zarr") if not ds_path.exists(): - raise FileNotFoundError(f"Dataset doesn't exist at {ds_path}") + raise FileNotFoundError(f"Dataset doesn't appear to exist at {ds_path}") result = xr.open_zarr(ds_path) return result +def create_time_series_dataset_classic( + src_directory: str, + dst_directory: str, + dataset_name: str, + input_variables: Sequence, + output_variables: Optional[Sequence] = None, + constants: Optional[DefaultDict] = None, + prefix: Optional[str] = None, + suffix: Optional[str] = None, + batch_size: int = 32, + scaling: Optional[DictConfig] = None, + overwrite: bool = False, +) -> "xr.Dataset": + """ + Opens and merges multiple datasets that that contain individual variables + into a single dataset + + Parameters + ---------- + src_directory: str + The directory that contains the input datasets + dst_directory: str + dataset_name: str + input_variables: Sequence + The input variables to be merged into the new dataset + output_variables: Sequence, optional + The output variables to be merged into the new dataset + If no output variables are provided the input set is used + constants: DefaultDict, optional + A set of constants to add to the merged dataset, default None + prefix: str, optional + The prefix of the input datasets, default None + suffix: str, optional + The suffix of the input datasets, default None + batch_size: str, optional + The chunk size to use for the input datasets, default 32 + scaling: DictConfig, optional + Scale factors applied to the listed variables, default None + overwrite: bool, optional + IF an existing dataset exists at the destination replace it, default False + + Returns + ------- + xr.Dataset: The merged dataset + """ + dst_zarr = os.path.join(dst_directory, dataset_name + ".zarr") + file_exists = os.path.exists(dst_zarr) + + if file_exists and not overwrite: + logger.info("opening input datasets") + return open_time_series_dataset_classic_prebuilt( + directory=dst_directory, + dataset_name=dataset_name, + constants=constants is not None, + ) + + output_variables = output_variables or input_variables + all_variables = np.union1d(input_variables, output_variables) + prefix = prefix or "" + suffix = suffix or "" + + merge_time = time.time() + logger.info("merging input datasets") + + datasets = [] + remove_attrs = ["varlev", "mean", "std"] + for variable in all_variables: + file_name = _get_file_name(src_directory, prefix, variable, suffix) + logger.debug("open nc dataset %s", file_name) + if "sample" in list(xr.open_dataset(file_name).sizes.keys()): + ds = xr.open_dataset(file_name).rename({"sample": "time"}) + else: + ds = xr.open_dataset(file_name) + if "varlev" in ds.dims: + ds = ds.isel(varlev=0) + + for attr in remove_attrs: + if attr in ds.indexes or attr in ds.variables: + ds = ds.drop(attr) + + # Rename variable + if "predictors" in list(ds.keys()): + ds = ds.rename({"predictors": variable}) + + # Change lat/lon to coordinates + try: + ds = ds.set_coords(["lat", "lon"]) + except (ValueError, KeyError): + pass + # Apply log scaling lazily + if ( + scaling + and variable in scaling + and scaling[variable].get("log_epsilon", None) is not None + ): + ds[variable] = np.log( + ds[variable] + scaling[variable]["log_epsilon"] + ) - np.log(scaling[variable]["log_epsilon"]) + datasets.append(ds) + # Merge datasets + data = xr.merge(datasets, compat="override") + + # Convert to input/target array by merging along the variables + input_da = ( + data[list(input_variables)] + .to_array("channel_in", name="inputs") + .transpose("time", "channel_in", "face", "height", "width") + ) + target_da = ( + data[list(output_variables)] + .to_array("channel_out", name="targets") + .transpose("time", "channel_out", "face", "height", "width") + ) + + result = xr.Dataset() + result["inputs"] = input_da + result["targets"] = target_da + + # Get constants + if constants is not None: + constants_ds = [] + for name, var in constants.items(): + constants_ds.append( + xr.open_dataset(_get_file_name(src_directory, prefix, name, suffix)) + .set_coords(["lat", "lon"])[var] + .astype(np.float32) + ) + constants_ds = xr.merge(constants_ds, compat="override") + constants_da = constants_ds.to_array("channel_c", name="constants").transpose( + "channel_c", "face", "height", "width" + ) + result["constants"] = constants_da + + logger.info("merged datasets in %0.1f s", time.time() - merge_time) + logger.info("writing unified dataset to file (takes long!)") + + # writing out + def _write_zarr(data, path): + write_job = data.to_zarr(path, compute=False, mode="w") + logger.info(f"writing dataset to {path}") + write_job.compute() + + _write_zarr(data=result, path=dst_zarr) + + return result + + class TimeSeriesDataModule: """pytorch-lightning module for complete model train, validation, and test data loading. Uses dlwp.data.data_loading.TimeSeriesDataset under-the-hood. Loaded data files follow the naming scheme @@ -99,12 +368,12 @@ class TimeSeriesDataModule: def __init__( self, - src_directory: str = ".", # kept for backwards compatibility + src_directory: str = ".", dst_directory: str = ".", dataset_name: str = "dataset", prefix: Optional[str] = None, suffix: Optional[str] = None, - data_format: str = "classic", # kept for backwards compatibility + data_format: str = "classic", batch_size: int = 32, drop_last: bool = False, input_variables: Optional[Sequence] = ["t2m"], @@ -120,15 +389,17 @@ def __init__( gap: Union[int, str, None] = None, shuffle: bool = True, add_insolation: bool = False, - cube_dim: int = 64, # kept for backwards compatibility + cube_dim: int = 64, num_workers: int = 4, pin_memory: bool = True, - prebuilt_dataset: bool = True, # kept for backwards compatibility + prebuilt_dataset: bool = True, forecast_init_times: Optional[Sequence] = None, ): """ Parameters ---------- + src_directory: str, optional + The directory containing data files per variable, default "." dst_directory: str, optional The directory containing joint data files, default "." dataset_name: str, optional @@ -137,6 +408,11 @@ def __init__( Prefix appended to all data files, default None suffix: str, optional Suffix appended to all data files, default None + data_format: str, optional + str indicating data schema. + 'classic': use classic DLWP file types. Loads .nc files, assuming + dimensions [sample, varlev, face, height, width] and + data variables 'predictors', 'lat', and 'lon'. batch_size: int, optional Size of batches to draw from data, defualt 32 drop_last: bool, optional @@ -173,10 +449,14 @@ def __init__( Option to shuffle the training data, default True add_insolation: bool, optional Option to add prescribed insolation as a decoder input feature, default True + cube_dim: int, optional + Number of points on the side of a cube face. Not currently used. num_workers: int, optional Number of parallel data loading workers, default 4 pin_memory: bool, optional Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True + prebuilt_dataset: bool, optional + Create a custom dataset for training. If False, the variables are gathered on the fly, default True forecast_init_times: Sequence, optional A Sequence of pandas Timestamps dictating the specific initialization times to produce inputs for. default None @@ -186,17 +466,12 @@ def __init__( NOT produce any target array. """ super().__init__() - - warnings.warn( - "TimeSeriesDataModule will be removed in a future release, please switch to using TimeSeriesDataModuleZarr", - DeprecationWarning, - stacklevel=2, - ) - + self.src_directory = src_directory self.dst_directory = dst_directory self.dataset_name = dataset_name self.prefix = prefix self.suffix = suffix + self.data_format = data_format self.batch_size = batch_size self.drop_last = drop_last self.input_variables = input_variables @@ -214,6 +489,7 @@ def __init__( self.cube_dim = cube_dim self.num_workers = num_workers self.pin_memory = pin_memory + self.prebuilt_dataset = prebuilt_dataset self.forecast_init_times = forecast_init_times self.train_dataset = None @@ -244,45 +520,93 @@ def get_constants(self) -> Optional[np.ndarray]: def setup(self) -> None: """Setup the datasets used for this DataModule""" + if self.data_format == "classic": + create_fn = create_time_series_dataset_classic + open_fn = ( + open_time_series_dataset_classic_prebuilt + if self.prebuilt_dataset + else open_time_series_dataset_classic_on_the_fly + ) + else: + raise ValueError("'data_format' must be one of ['classic']") + # make sure distributed manager is initalized if not DistributedManager.is_initialized(): DistributedManager.initialize() + dist = DistributedManager() - dataset = open_time_series_dataset_classic_prebuilt( - directory=self.dst_directory, - dataset_name=self.dataset_name, - constants=self.constants is not None, - batch_size=self.batch_size, - ) + if torch.distributed.is_initialized(): + if self.prebuilt_dataset: + if dist.rank == 0: + create_fn( + src_directory=self.src_directory, + dst_directory=self.dst_directory, + dataset_name=self.dataset_name, + input_variables=self.input_variables, + output_variables=self.output_variables, + constants=self.constants, + prefix=self.prefix, + suffix=self.suffix, + batch_size=self.dataset_batch_size, + scaling=self.scaling, + overwrite=False, + ) - # Verify that all input variables are available in the dataset - missing_input_vars = set(self.input_variables) - set(dataset.channel_in.values) - if missing_input_vars: - raise ValueError( - f"Input variables not found in dataset: {missing_input_vars}" - ) + # wait for rank 0 to complete, because then the files are guaranteed to exist + torch.distributed.barrier() - # Verify that all output variables are available in the dataset - missing_output_vars = set(self.output_variables) - set( - dataset.channel_out.values - ) - if missing_output_vars: - raise ValueError( - f"Output variables not found in dataset: {missing_output_vars}" - ) + dataset = open_fn( + directory=self.dst_directory, + dataset_name=self.dataset_name, + constants=self.constants is not None, + batch_size=self.batch_size, + ) + else: + dataset = open_fn( + input_variables=self.input_variables, + output_variables=self.output_variables, + directory=self.dst_directory, + constants=self.constants, + prefix=self.prefix, + batch_size=self.batch_size, + ) + else: + if self.prebuilt_dataset: + create_fn( + src_directory=self.src_directory, + dst_directory=self.dst_directory, + dataset_name=self.dataset_name, + input_variables=self.input_variables, + output_variables=self.output_variables, + constants=self.constants, + prefix=self.prefix, + suffix=self.suffix, + batch_size=self.dataset_batch_size, + scaling=self.scaling, + overwrite=False, + ) + + dataset = open_fn( + directory=self.dst_directory, + dataset_name=self.dataset_name, + constants=self.constants is not None, + batch_size=self.batch_size, + ) + else: + dataset = open_fn( + input_variables=self.input_variables, + output_variables=self.output_variables, + directory=self.dst_directory, + constants=self.constants, + prefix=self.prefix, + batch_size=self.batch_size, + ) dataset = dataset.sel( channel_in=self.input_variables, channel_out=self.output_variables, ) if self.constants is not None: - # Verify that all constants are available in the dataset - missing_constants = set(list(self.constants.values())) - set( - dataset.channel_c.values - ) - if missing_constants: - raise ValueError(f"Constants not found in dataset: {missing_constants}") - dataset = dataset.sel(channel_c=list(self.constants.values())) if self.splits is not None and self.forecast_init_times is None: @@ -477,12 +801,12 @@ class CoupledTimeSeriesDataModule(TimeSeriesDataModule): def __init__( self, - src_directory: str = ".", # kept for backwards compatibility + src_directory: str = ".", dst_directory: str = ".", dataset_name: str = "dataset", prefix: Optional[str] = None, suffix: Optional[str] = None, - data_format: str = "classic", # kept for backwards compatibility + data_format: str = "classic", batch_size: int = 32, drop_last: bool = False, input_variables: Optional[Sequence] = None, @@ -498,10 +822,10 @@ def __init__( gap: Union[int, str, None] = None, shuffle: bool = True, add_insolation: bool = False, - cube_dim: int = 64, # kept for backwards compatibility + cube_dim: int = 64, num_workers: int = 4, pin_memory: bool = True, - prebuilt_dataset: bool = True, # kept for backwards compatibility + prebuilt_dataset: bool = True, forecast_init_times: Optional[Sequence] = None, couplings: Sequence = None, add_train_noise: Optional[bool] = False, @@ -511,6 +835,8 @@ def __init__( """ Parameters ---------- + src_directory: str, optional + The directory containing data files per variable, default "." dst_directory: str, optional The directory containing joint data files, default "." dataset_name: str, optional @@ -519,6 +845,11 @@ def __init__( Prefix appended to all data files, default None suffix: str, optional Suffix appended to all data files, default None + data_format: str, optional + str indicating data schema. + 'classic': use classic DLWP file types. Loads .nc files, assuming + dimensions [sample, varlev, face, height, width] and + data variables 'predictors', 'lat', and 'lon'. batch_size: int, optional Size of batches to draw from data, defualt 32 drop_last: bool, optional @@ -561,6 +892,8 @@ def __init__( Number of parallel data loading workers, default 4 pin_memory: bool, optional Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True + prebuilt_dataset: bool, optional + Create a custom dataset for training. If False, the variables are gathered on the fly, default True forecast_init_times: Sequence, optional A Sequence of pandas Timestamps dictating the specific initialization times to produce inputs for. default None @@ -613,9 +946,6 @@ def __init__( ) def _get_coupled_vars(self): - """ " - Get the coupled variables from the couplings dictionary. - """ coupled_variables = [] for d in self.couplings: coupled_variables = coupled_variables + d["params"]["variables"] @@ -623,46 +953,94 @@ def _get_coupled_vars(self): def setup(self) -> None: """Setup the datasets used for this DataModule""" + if self.data_format == "classic": + create_fn = create_time_series_dataset_classic + open_fn = ( + open_time_series_dataset_classic_prebuilt + if self.prebuilt_dataset + else open_time_series_dataset_classic_on_the_fly + ) + else: + raise ValueError("'data_format' must be one of ['classic', 'zarr']") + coupled_variables = self._get_coupled_vars() # make sure distributed manager is initalized if not DistributedManager.is_initialized(): DistributedManager.initialize() dist = DistributedManager() - dataset = open_time_series_dataset_classic_prebuilt( - directory=self.dst_directory, - dataset_name=self.dataset_name, - constants=self.constants is not None, - batch_size=self.batch_size, - ) + if torch.distributed.is_initialized(): + if self.prebuilt_dataset: + if dist.rank == 0: + create_fn( + src_directory=self.src_directory, + dst_directory=self.dst_directory, + dataset_name=self.dataset_name, + input_variables=self.input_variables + coupled_variables, + output_variables=self.output_variables, + constants=self.constants, + prefix=self.prefix, + suffix=self.suffix, + batch_size=self.dataset_batch_size, + scaling=self.scaling, + overwrite=False, + ) - # Verify that all input variables are available in the dataset - missing_input_vars = set(self.input_variables) - set(dataset.channel_in.values) - if missing_input_vars: - raise ValueError( - f"Input variables not found in dataset: {missing_input_vars}" - ) + # wait for rank 0 to complete, because then the files are guaranteed to exist + torch.distributed.barrier() - # Verify that all output variables are available in the dataset - missing_output_vars = set(self.output_variables) - set( - dataset.channel_out.values - ) - if missing_output_vars: - raise ValueError( - f"Output variables not found in dataset: {missing_output_vars}" - ) + dataset = open_fn( + directory=self.dst_directory, + dataset_name=self.dataset_name, + constants=self.constants is not None, + batch_size=self.batch_size, + ) + else: + dataset = open_fn( + input_variables=self.input_variables + coupled_variables, + output_variables=self.output_variables, + directory=self.dst_directory, + constants=self.constants, + prefix=self.prefix, + batch_size=self.batch_size, + ) + else: + if self.prebuilt_dataset: + create_fn( + src_directory=self.src_directory, + dst_directory=self.dst_directory, + dataset_name=self.dataset_name, + input_variables=self.input_variables + coupled_variables, + output_variables=self.output_variables, + constants=self.constants, + prefix=self.prefix, + suffix=self.suffix, + batch_size=self.dataset_batch_size, + scaling=self.scaling, + overwrite=False, + ) + + dataset = open_fn( + directory=self.dst_directory, + dataset_name=self.dataset_name, + constants=self.constants is not None, + batch_size=self.batch_size, + ) + else: + dataset = open_fn( + input_variables=self.input_variables + coupled_variables, + output_variables=self.output_variables, + directory=self.dst_directory, + constants=self.constants, + prefix=self.prefix, + batch_size=self.batch_size, + ) dataset = dataset.sel( channel_in=self.input_variables + coupled_variables, channel_out=self.output_variables, ) if self.constants is not None: - # Verify that all constants are available in the dataset - missing_constants = set(list(self.constants.values())) - set( - dataset.channel_c.values - ) - if missing_constants: - raise ValueError(f"Constants not found in dataset: {missing_constants}") dataset = dataset.sel(channel_c=list(self.constants.values())) if self.splits is not None and self.forecast_init_times is None: diff --git a/physicsnemo/datapipes/healpix/data_modules_zarr.py b/physicsnemo/datapipes/healpix/data_modules_zarr.py deleted file mode 100644 index 988c03b738..0000000000 --- a/physicsnemo/datapipes/healpix/data_modules_zarr.py +++ /dev/null @@ -1,557 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -DataModule for loading and processing time series healpix data stored in Zarr format. - -This class provides the core functionality for setting up the dataloaders for the time series healpix data stored in Zarr format. -Data loading, scaling, and time management are handled by the TimeSeriesDatasetZarr and CoupledTimeSeriesDatasetZarr classes. -Zarr format is used to avoid Xarray overheads when loading and processing the data. -""" - -# System modules -import logging -import warnings -from pathlib import Path -from typing import Optional, Sequence, Union - -# numpy -import numpy as np - -# External modules -from omegaconf import DictConfig -from torch.utils.data import DataLoader -from torch.utils.data.distributed import DistributedSampler - -from physicsnemo.distributed import DistributedManager - -from .coupledtimeseries_dataset_zarr import CoupledTimeSeriesDatasetZarr -from .timeseries_dataset_zarr import TimeSeriesDatasetZarr - -logger = logging.getLogger(__name__) - - -class TimeSeriesDataModuleZarr: - """Module for complete model train, validation, and test data loading. Uses - dlwp.data.data_loading.TimeSeriesDataset under-the-hood. - """ - - def __init__( - self, - dataset_path: str, - batch_size: int = 32, - dataloader_batch_size: Optional[int] = None, - drop_last: bool = True, - input_variables: Optional[Sequence] = ["t2m"], - output_variables: Optional[Sequence] = None, - constant_variables: Optional[Sequence] = None, - scaling: Optional[DictConfig] = None, - splits: Optional[DictConfig] = None, - presteps: int = 0, - input_time_dim: int = 1, - output_time_dim: int = 1, - data_time_step: Union[int, str] = "3h", - time_step: Union[int, str] = "6h", - gap: Union[int, str, None] = None, - shuffle: bool = True, - add_insolation: bool = False, - num_workers: int = 4, - pin_memory: bool = True, - forecast_init_times: Optional[Sequence] = None, - add_train_noise: Optional[bool] = False, - train_noise_params: Optional[DictConfig] = None, - train_noise_seed: Optional[int] = 42, - ): - """ - Parameters - ---------- - dataset_path: str, optional - The path to the dataset, default "." - batch_size: int, optional - The number of sequential samples to load from the dataset to load, default 32 - dataloader_batch_size: int, optional - Passed to nn.DataLoader as batch_size. Used to assemble batches of samples from the dataloader - The total number of samples will be batch_size*dataloader_batch_size, default None - drop_last: bool, optional - Whether to drop the last batch if it is smaller than batch_size, it is - recommended to set this to true to avoid issues with mismatched sizes, default True - input_variables: Sequence, optional - List of input variable names, to be found in data file name, default "t2m" - output_variables: Sequence, optional - List of output variables names. If None, defaults to `input_variables`. default None - constant_variables: Sequence, optional - List of constant variables names. default None - scaling: DictConfig, optional - Dictionary containing scaling parameters for data variables, default None - splits: DictConfig - Dictionary with train/validation/test set start/end dates. - presteps: int, optional - Number of time steps to initialize recurrent hidden states. default 0 - input_time_dim: int, optional - Number of time steps in the input array, default 1 - output_time_dim: int, optional - Number of time steps in the target/ground truth array, default 1 - data_time_step: Union[int, str], optional - Either integer hours or a str interpretable by pandas: time between steps in the - original data time series, default "3h" - time_step: Union[int, str], optional - Either integer hours or a str interpretable by pandas: desired time between effective model - time steps, default "6h" - gap: Union[int, str], optional - either integer hours or a str interpretable by pandas: time step between the last input time and - the first output time. Defaults to `time_step`. - shuffle: bool, optional - Whether to shuffle the training data, default True - add_insolation: bool, optional - Whether to add prescribed insolation as a decoder input feature, default False - num_workers: int, optional - Number of parallel data loading workers, default 4 - pin_memory: bool, optional - Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True - forecast_init_times: Sequence, optional - A Sequence of pandas Timestamps dictating the specific initialization times - to produce inputs for. default None - Note: - - this is only applied to the test dataloader - - providing this parameter configures the data loader to only produce this number of samples, and - NOT produce any target array. - add_train_noise: bool, optional - Wether to add noise to the training data to inputs and integrated couplings to improve generalization, default False - train_noise_params: DictConfig, optional - Dictionary containing parameters for adding noise to the training data - train_noise_seed: int, optional - Seed for the random number generator for adding noise to the training data, default 42 - """ - super().__init__() - self.dataset_path = Path(dataset_path) - self.dataset_batch_size = batch_size - self.dataloader_batch_size = dataloader_batch_size - self.drop_last = drop_last - self.input_variables = input_variables - self.output_variables = output_variables or input_variables - self.constant_variables = constant_variables - self.constants = constant_variables - self.scaling = scaling - self.splits = splits - self.input_time_dim = input_time_dim + (presteps * input_time_dim) - self.output_time_dim = output_time_dim - self.data_time_step = data_time_step - self.time_step = time_step - self.gap = gap - self.shuffle = shuffle - self.add_insolation = add_insolation - self.num_workers = num_workers - self.pin_memory = pin_memory - self.forecast_init_times = forecast_init_times - self.add_train_noise = add_train_noise - self.train_noise_params = train_noise_params - self.train_noise_seed = train_noise_seed - - self.train_dataset = None - self.val_dataset = None - self.test_dataset = None - - self.collate_fn = None - - self.setup() - - def get_constants(self) -> Optional[np.ndarray]: - """Returns the constants used in this dataset - - Returns - ------- - np.ndarray: The list of constants, None if there are no constants - """ - if self.constant_variables is None: - return None - - return ( - self.train_dataset.get_constants() - if self.train_dataset is not None - else self.test_dataset.get_constants() - ) - - def _get_common_dataset_kwargs(self) -> dict: - """Get common keyword arguments for dataset creation""" - return { - "dataset_path": self.dataset_path, - "input_variables": self.input_variables, - "output_variables": self.output_variables, - "constant_variables": self.constant_variables, - "batch_size": self.dataset_batch_size, - "add_insolation": self.add_insolation, - "input_time_dim": self.input_time_dim, - "output_time_dim": self.output_time_dim, - "data_time_step": self.data_time_step, - "time_step": self.time_step, - "gap": self.gap, - "scaling": self.scaling, - } - - def _validate_setup_requirements(self) -> None: - """Validate that required setup conditions are met""" - if not self.dataset_path.exists(): - raise FileNotFoundError(f"Dataset path not found: {self.dataset_path}") - - if self.splits is None and self.forecast_init_times is None: - raise ValueError("Either splits or forecast_init_times must be provided") - - # sanity check if dates overlap - if self.splits: - train_test_overlap = ( - np.datetime64(self.splits["train_date_end"]) - >= np.datetime64(self.splits["test_date_start"]) - ) and ( - np.datetime64(self.splits["train_date_start"]) - <= np.datetime64(self.splits["test_date_end"]) - ) - if train_test_overlap: - warnings.warn("Training and test date ranges overlap") - - train_val_overlap = ( - np.datetime64(self.splits["train_date_end"]) - >= np.datetime64(self.splits["val_date_start"]) - ) and ( - np.datetime64(self.splits["train_date_start"]) - <= np.datetime64(self.splits["val_date_end"]) - ) - if train_val_overlap: - warnings.warn("Training and validation date ranges overlap") - - test_val_overlap = ( - np.datetime64(self.splits["test_date_end"]) - >= np.datetime64(self.splits["val_date_start"]) - ) and ( - np.datetime64(self.splits["test_date_start"]) - <= np.datetime64(self.splits["val_date_end"]) - ) - if test_val_overlap: - warnings.warn("Test and validation date ranges overlap") - - def _initialize_distributed_manager(self): - """Initialize distributed manager if not already initialized""" - if not DistributedManager.is_initialized(): - DistributedManager.initialize() - return DistributedManager() - - def _get_dataset_class(self): - """Get the dataset class to use for creating datasets""" - return TimeSeriesDatasetZarr - - def _create_datasets(self, common_kwargs: dict, dataset_class, dist) -> None: - """Create train, validation, and test datasets based on configuration""" - # forecast_init_times are provided just create the test dataset - if self.forecast_init_times is not None: - self.test_dataset = dataset_class( - drop_last=False, - **common_kwargs, - forecast_init_times=self.forecast_init_times, - ) - # splits are provided, use them to split the dataset - else: - self.train_dataset = dataset_class( - start_date=self.splits["train_date_start"], - end_date=self.splits["train_date_end"], - drop_last=self.drop_last, - add_train_noise=self.add_train_noise, - train_noise_params=self.train_noise_params, - train_noise_seed=self.train_noise_seed + int(dist.rank), - **common_kwargs, - ) - self.val_dataset = dataset_class( - start_date=self.splits["val_date_start"], - end_date=self.splits["val_date_end"], - drop_last=self.drop_last, - **common_kwargs, - ) - self.test_dataset = dataset_class( - start_date=self.splits["test_date_start"], - end_date=self.splits["test_date_end"], - drop_last=False, - **common_kwargs, - ) - - def setup(self) -> None: - """Setup the datasets used for this DataModule""" - self._validate_setup_requirements() - dist = self._initialize_distributed_manager() - - common_kwargs = self._get_common_dataset_kwargs() - dataset_class = self._get_dataset_class() - - self._create_datasets(common_kwargs, dataset_class, dist) - - def _base_dataloader( - self, dataset, num_shards=1, shard_id=0, shuffle=False, drop_last=False - ) -> DataLoader: - """Setup a dataloader with common functionality - - Parameters - ---------- - dataset: Dataset - The dataset to create the dataloader for - num_shards: int, optional - The total total number of distributed shards - default is 1 meaning distributed training is not being used - shard_id: int, optional - The shard number of this instance of the dataloader, default 0 - shuffle: bool, optional - Whether to shuffle the data, default False - - Returns - ------- - DataLoader: The configured dataloader - """ - sampler = None - drop_last = False - if num_shards > 1: - sampler = DistributedSampler( - dataset, - num_replicas=num_shards, - rank=shard_id, - shuffle=shuffle, - drop_last=drop_last, - ) - shuffle = False - drop_last = False - - loader = DataLoader( - dataset=dataset, - pin_memory=self.pin_memory, - num_workers=self.num_workers, - shuffle=shuffle, - drop_last=drop_last, - sampler=sampler, - collate_fn=self.collate_fn, - batch_size=self.dataloader_batch_size, - ) - - return loader, sampler - - def train_dataloader(self, num_shards=1, shard_id=0) -> DataLoader: - """Setup the training dataloader - - Parameters - ---------- - num_shards: int, optional - The total total number of distributed shards - default is 1 meaning distributed training is not being used - shard_id: int, optional - The shard number of this instance of the dataloader, default 0 - - Returns - ------- - DataLoader: The training dataloader - """ - return self._base_dataloader( - dataset=self.train_dataset, - num_shards=num_shards, - shard_id=shard_id, - shuffle=self.shuffle, - drop_last=True, - ) - - def val_dataloader(self, num_shards=1, shard_id=0) -> DataLoader: - """Setup the validation dataloader - - Parameters - ---------- - num_shards: int, optional - The total total number of distributed shards - default is 1 meaning distributed validation is not being used - shard_id: int, optional - The shard number of this instance of the dataloader, default 0 - - Returns - ------- - DataLoader: The validation dataloader - """ - return self._base_dataloader( - dataset=self.val_dataset, - num_shards=num_shards, - shard_id=shard_id, - shuffle=False, - drop_last=True, - ) - - def test_dataloader(self, num_shards=1, shard_id=0) -> DataLoader: - """Setup the test dataloader - - Parameters - ---------- - num_shards: int, optional - The total total number of distributed shards - default is 1 meaning distributed test is not being used - shard_id: int, optional - The shard number of this instance of the dataloader, default 0 - - Returns - ------- - DataLoader: The test dataloader - """ - return self._base_dataloader( - dataset=self.test_dataset, - num_shards=num_shards, - shard_id=shard_id, - shuffle=False, - drop_last=True, - ) - - -class CoupledTimeSeriesDataModuleZarr(TimeSeriesDataModuleZarr): - """ - Extension of TimeSeriesDataModule, designed for coupled models that take input from other - earth system components. - """ - - def __init__( - self, - dataset_path: str, - batch_size: int = 32, - dataloader_batch_size: Optional[int] = None, - drop_last: bool = True, - input_variables: Optional[Sequence] = None, - output_variables: Optional[Sequence] = None, - constant_variables: Optional[Sequence] = None, - scaling: Optional[DictConfig] = None, - splits: Optional[DictConfig] = None, - presteps: int = 0, - input_time_dim: int = 1, - output_time_dim: int = 1, - data_time_step: Union[int, str] = "3h", - time_step: Union[int, str] = "6h", - gap: Union[int, str, None] = None, - shuffle: bool = True, - add_insolation: bool = False, - num_workers: int = 4, - pin_memory: bool = True, - forecast_init_times: Optional[Sequence] = None, - couplings: Sequence = None, - add_train_noise: Optional[bool] = False, - train_noise_params: Optional[DictConfig] = None, - train_noise_seed: Optional[int] = 42, - ): - """ - Parameters - ---------- - dataset_path: str, optional - The path to the dataset, default "." - batch_size: int, optional - The number of sequential samples to load from the dataset to load, default 32 - dataloader_batch_size: int, optional - Passed to nn.DataLoader as batch_size. Used to assemble batches of samples from the dataloader - The total number of samples will be dataloader_batch_size*dataloader_batch_size, default None - drop_last: bool, optional - Whether to drop the last batch if it is smaller than batch_size, it is - recommended to set this to true to avoid issues with mismatched sizes, default True - input_variables: Sequence, optional - List of input variable names, to be found in data file name, default None - output_variables: Sequence, optional - List of output variables names. If None, defaults to `input_variables`. default None - constant_variables: Sequence, optional - List of constant variables names. default None - scaling: DictConfig, optional - Dictionary containing scaling parameters for data variables, default None - splits: DictConfig - Dictionary with train/validation/test set start/end dates. - presteps: int, optional - Number of time steps to initialize recurrent hidden states. default 0 - input_time_dim: int, optional - Number of time steps in the input array, default 1 - output_time_dim: int, optional - Number of time steps in the target/ground truth array, default 1 - data_time_step: Union[int, str], optional - Either integer hours or a str interpretable by pandas: time between steps in the - original data time series, default "3h" - time_step: Union[int, str], optional - Either integer hours or a str interpretable by pandas: desired time between effective model - time steps, default "6h" - gap: Union[int, str, None], optional - either integer hours or a str interpretable by pandas: time step between the last input time and - the first output time. default None. - shuffle: bool, optional - Whether to shuffle the training data, default True - add_insolation: bool, optional - Whether to add prescribed insolation as a decoder input feature, default False - num_workers: int, optional - Number of parallel data loading workers, default 4 - pin_memory: bool, optional - Whether pinned (page locked) memory should be used to store the tensors, improves GPU I/O, default True - forecast_init_times: Sequence, optional - A Sequence of pandas Timestamps dictating the specific initialization times - to produce inputs for. default None - Note: - - this is only applied to the test dataloader - - providing this parameter configures the data loader to only produce this number of samples, and - NOT produce any target array. - couplings: Sequence, optional - a Sequence of dictionaries that define the mechanics of couplings with other earth system - components. default None - add_train_noise: bool, optional - Wether to add noise to the training data to inputs and integrated couplings to improve generalization, default False - train_noise_params: DictConfig, optional - Dictionary containing parameters for adding noise to the training data - train_noise_seed: int, optional - Seed for the random number generator for adding noise to the training data, default 42 - """ - self.couplings = couplings - - super().__init__( - dataset_path, - batch_size, - dataloader_batch_size, - drop_last, - input_variables, - output_variables, - constant_variables, - scaling, - splits, - presteps, - input_time_dim, - output_time_dim, - data_time_step, - time_step, - gap, - shuffle, - add_insolation, - num_workers, - pin_memory, - forecast_init_times, - add_train_noise, - train_noise_params, - train_noise_seed, - ) - - def _get_coupled_vars(self): - """Get the coupled variables from the couplings""" - coupled_variables = [] - for d in self.couplings: - coupled_variables = coupled_variables + d["params"]["variables"] - return coupled_variables - - def _get_common_dataset_kwargs(self) -> dict: - """Get common keyword arguments for dataset creation (includes coupling-specific params)""" - base_kwargs = super()._get_common_dataset_kwargs() - base_kwargs.update( - { - "couplings": self.couplings, - } - ) - return base_kwargs - - def _get_dataset_class(self): - """Get the dataset class to use for creating datasets""" - return CoupledTimeSeriesDatasetZarr diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index b847da4cdb..9e2fbad3ef 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -14,13 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Dataset for sampling from continuous time-series data, compatible with pytorch data loading. - -This class provides the core functionality for sampling from continuous time-series data, compatible with pytorch data loading. -It handles data loading, scaling, and time management. It is used by the TimeSeriesDataModule to set up the dataloaders for the time series healpix data. -""" - from __future__ import annotations import logging @@ -33,19 +26,17 @@ import pandas as pd import torch from omegaconf import DictConfig, OmegaConf +from torch.utils.data import Dataset from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.datapipe import Datapipe from physicsnemo.datapipes.meta import DatapipeMetaData from physicsnemo.utils.insolation import insolation -logger = logging.getLogger(__name__) - -# xarray ships in the ``datapipes-extras`` optional dependency group; load it -# lazily so the physicsnemo import graph carries no static dependency on it -# (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). xr = OptionalImport("xarray") +logger = logging.getLogger(__name__) + @dataclass class MetaData(DatapipeMetaData): @@ -59,7 +50,7 @@ class MetaData(DatapipeMetaData): ddp_sharding: bool = False -class TimeSeriesDataset(Datapipe): +class TimeSeriesDataset(Dataset, Datapipe): """ Dataset for sampling from continuous time-series data, compatible with pytorch data loading. """ @@ -253,23 +244,13 @@ def _get_scaling_da(self): scaling_da = scaling_df.to_xarray().astype("float32") # REMARK: we remove the xarray overhead from these - # base (non-coupled) datasets have no 'couplings' attribute, and coupled - # datasets may be configured with an empty couplings list; treat both as - # zero coupled variables rather than indexing into an empty/missing list. - couplings = getattr(self, "couplings", []) - num_coupled_vars = len(couplings[0].variables) if couplings else 0 try: - # 'if' statement used for cases where atmos model - # includes diagnostic variables like tp6 and msl. - # using 'channel_out' is still necessary for ocean models. - if len(self.ds.channel_out) != (len(self.ds.channel_in) - num_coupled_vars): - self.input_scaling = scaling_da.sel( - index=self.ds.channel_in.values - ).rename({"index": "channel_in"}) - else: - self.input_scaling = scaling_da.sel( - index=self.ds.channel_out.values - ).rename({"index": "channel_in"}) + # we use channel_out instead of channel_in because + # the list of input channels may contain data fetched outside + # the datasets such as coupled fields + self.input_scaling = scaling_da.sel( + index=self.ds.channel_out.values + ).rename({"index": "channel_in"}) self.input_scaling = { "mean": np.expand_dims( self.input_scaling["mean"].to_numpy(), (0, 2, 3, 4) diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py b/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py deleted file mode 100644 index 79f08d3a93..0000000000 --- a/physicsnemo/datapipes/healpix/timeseries_dataset_zarr.py +++ /dev/null @@ -1,279 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Dataset for sampling from continuous time-series data stored in Zarr format, compatible with pytorch data loading. - -This class provides the core functionality for sampling from continuous time-series data stored in Zarr format, compatible with pytorch data loading. -It handles data loading, scaling, and time management. It is used by the TimeSeriesDataModule to set up the dataloaders for the time series healpix data stored in Zarr format. -Zarr format is used to avoid Xarray overheads when loading and processing the data. -""" - -import logging -from dataclasses import dataclass -from typing import List, Optional, Sequence, Tuple, Union - -import numpy as np -import torch -from omegaconf import DictConfig - -from physicsnemo.datapipes.meta import DatapipeMetaData -from physicsnemo.utils.insolation import insolation - -from .base_timeseries_dataset_zarr import BaseTimeSeriesDatasetZarr - -logger = logging.getLogger(__name__) - - -@dataclass -class MetaData(DatapipeMetaData): - """Metadata for this datapipe""" - - name: str = "TimeSeries" - # Optimization - auto_device: bool = False - cuda_graphs: bool = False - # Parallel - ddp_sharding: bool = False - - -class TimeSeriesDatasetZarr(BaseTimeSeriesDatasetZarr): - """Dataset for sampling from continuous time-series data stored in Zarr format. - - This class implements the basic time series dataset functionality without coupling - to external data sources. It provides data loading, scaling, and batching - capabilities for training and inference. - """ - - def __init__( - self, - dataset_path: str, - input_variables: Sequence, - output_variables: Sequence = None, - constant_variables: Sequence = None, - scaling: DictConfig = None, - input_time_dim: int = 1, - output_time_dim: int = 1, - data_time_step: Union[int, str] = "3h", - time_step: Union[int, str] = "6h", - gap: Union[int, str, None] = None, - batch_size: int = 32, - drop_last: bool = False, - add_insolation: bool = False, - forecast_init_times: Optional[Sequence] = None, - start_date: Optional[Union[int, str]] = None, - end_date: Optional[Union[int, str]] = None, - add_train_noise: bool = False, - train_noise_params: DictConfig = None, - train_noise_seed: int = 42, - meta: DatapipeMetaData = MetaData(), - ): - """Initialize time series dataset. - - Parameters are same as BaseTimeSeriesDatasetZarr. - See base class for detailed parameter descriptions. - """ - super().__init__( - dataset_path=dataset_path, - input_variables=input_variables, - output_variables=output_variables, - constant_variables=constant_variables, - scaling=scaling, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - data_time_step=data_time_step, - time_step=time_step, - gap=gap, - batch_size=batch_size, - drop_last=drop_last, - add_insolation=add_insolation, - forecast_init_times=forecast_init_times, - start_date=start_date, - end_date=end_date, - add_train_noise=add_train_noise, - train_noise_params=train_noise_params, - train_noise_seed=train_noise_seed, - meta=meta, - ) - - def __getitem__( - self, item: int - ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]]: - """Get a batch of time series data. - - This implementation provides the basic time series data loading functionality: - 1. Load data window for batch - 2. Apply scaling - 3. Extract input and target sequences - 4. Add optional insolation and constant inputs - 5. Format dimensions appropriately - - Parameters - ---------- - item : int - Sample index - - Returns - ------- - Union[List[np.ndarray], Tuple[List[np.ndarray], np.ndarray]] - In forecast mode: List of input arrays - In training mode: Tuple of (input arrays, target array) - - Input arrays are in order: - - Model inputs [B, F, T, C, H, W] - - Insolation (if enabled) [B, F, T, 1, H, W] - - Constants (if provided) [F, C, H, W] - - Target array has shape [B, F, T, C, H, W] - where: - B = batch size - F = faces - T = time steps - C = channels - H = height - W = width - - Raises - ------ - IndexError - If item is out of range - """ - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__") - - if item < 0: - item = len(self) + item - if item < 0 or item > len(self): - raise IndexError( - f"index {item} out of range for dataset with length {len(self)}" - ) - - # remark: load first then normalize - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_batch") - time_index, this_batch = self._get_time_index(item) - - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_staging_data") - # data from the input and target arrays overlap, to avoid 2 seperate loads - # we load both and then slice it later - staging_ds = self.ds["inputs"][slice(*time_index)] - staging_ds = staging_ds[:, self.all_variable_indices] - torch.cuda.nvtx.range_pop() - - # we scale this dataset to avoid doing twice the work when we scale as both - # input and output - # we do the scaling as in place operations to avoid creating temp arrays - # that result in a lot of data movement. This is around 4x faster than using - # standard operations - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:scale_batch") - if self.all_scaling is not None: - staging_ds -= self.all_scaling["mean"] - staging_ds /= self.all_scaling["std"] - - torch.cuda.nvtx.range_pop() - - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_input") - input_array = staging_ds[:, self.input_variable_indices] - torch.cuda.nvtx.range_pop() - - if not self.forecast_mode: - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:load_target") - target_array = staging_ds[:, self.output_variable_indices] - torch.cuda.nvtx.range_pop() - torch.cuda.nvtx.range_pop() - - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:process_batch") - - # Calculate insolation if needed - if self.add_insolation: - sol = insolation( - self._get_forecast_sol_times(item), - self.lat, - self.lon, - )[:, None] - decoder_inputs = np.empty( - (this_batch, self.input_time_dim + self.output_time_dim, 1) - + self.spatial_dims, - dtype="float32", - ) - - # Get buffers for the batches, which we'll fill in iteratively. - inputs = np.empty( - (this_batch, self.input_time_dim, len(self.input_variables)) - + self.spatial_dims, - dtype="float32", - ) - if not self.forecast_mode: - targets = np.empty( - (this_batch, self.output_time_dim, len(self.output_variables)) - + self.spatial_dims, - dtype="float32", - ) - - # Iterate over valid sample windows - torch.cuda.nvtx.range_push( - "TimeSeriesDataset:__getitem__:copy_inputs_targets_insolation" - ) - for sample in range(this_batch): - inputs[sample] = input_array[self._input_indices[sample]] - if not self.forecast_mode: - targets[sample] = target_array[self._output_indices[sample]] - if self.add_insolation: - decoder_inputs[sample] = ( - sol - if self.forecast_mode - else sol[self._input_indices[sample] + self._output_indices[sample]] - ) - torch.cuda.nvtx.range_pop() - - if not self.forecast_mode and self.add_train_noise: - torch.cuda.nvtx.range_push("TimeSeriesDataset:__getitem__:add_train_noise") - # Iterate over C: inputs.shape = [B, T, C, F, H, W] - for i in range(inputs.shape[2]): - inputs[:, :, i] += self.rng.normal( - loc=0, - scale=self.train_noise_params["inputs"][self.input_variables[i]][ - "std" - ], - size=inputs[:, :, i].shape, - ) - torch.cuda.nvtx.range_pop() - - inputs_result = [inputs] - torch.cuda.nvtx.range_push( - "CoupledTimeSeriesDataset:__getitem__:add_insolation" - ) - if self.add_insolation: - inputs_result.append(decoder_inputs) - torch.cuda.nvtx.range_pop() - - # Transpose dimensions to match expected format - inputs_result = [ - np.transpose(x, axes=(0, 3, 1, 2, 4, 5)) for x in inputs_result - ] - - if self.constant_variables: - # Add constants - inputs_result.append(self.constants) - torch.cuda.nvtx.range_pop() - - if self.forecast_mode: - torch.cuda.nvtx.range_pop() - return inputs_result - - # Transpose targets to match input format - targets = np.transpose(targets, axes=(0, 3, 1, 2, 4, 5)) - - torch.cuda.nvtx.range_pop() - return inputs_result, targets diff --git a/physicsnemo/metrics/climate/healpix_loss.py b/physicsnemo/metrics/climate/healpix_loss.py index 0f67772fcb..ce7f5b6d69 100644 --- a/physicsnemo/metrics/climate/healpix_loss.py +++ b/physicsnemo/metrics/climate/healpix_loss.py @@ -22,10 +22,6 @@ from physicsnemo.core.version_check import OptionalImport xr = OptionalImport("xarray") -earth2grid = OptionalImport( - "earth2grid", "Install earth2grid from https://github.com/earth2grid/earth2grid" -) -cuhpx = OptionalImport("cuhpx", "Install cuhpx from https://github.com/NVIDIA/cuhpx") """ Custom dlwp compatible loss classes that allow for more sophisticated training optimization. @@ -136,67 +132,6 @@ def forward(self, prediction, target, average_channels=True): return d -class ConditionalWeightLoss(th.nn.MSELoss): - """ - Conditional loss for precipitation diagnostic model. - (Total 6hr precipitation is the only output field.) - """ - - def __init__( - self, - weight=(0.01, 1.0), - b=None, - w=1, - ): - """ - Parameters - ----------- - weight: tuple of floats - weight[0] is used when the target precipitation value is zero - weight[1] is used for all non-zero precipitation - b: float - Exponential scaling factor used to define weighting curve for non-zero precip. weights - w: float - Final scaling factor applied to loss. - """ - - super().__init__() - self.weight_zero = weight[0] - self.weight_nonzero = weight[1] - self.b = b - self.device = None - self.w = w - - def setup(self, trainer): - """Setup function that moves tensors to the same device as the model""" - self.b = th.tensor(self.b, device=trainer.device) - self.w = th.tensor(self.w, device=trainer.device) - - def forward(self, prediction, target, average_channels=True): - """ - Computes the MSE of model prediction and applies weights for zero and non-zero precipitation cases. - - Parameters - ----------- - prediction: torch.tensor - The prediction tensor - target: torch.Tensor - The target tensor - average_channels: bool, optional - whether the mean of the channels should be taken - """ - weights_for_zero = th.ones_like(target) * self.weight_zero - weights_for_nonzero = (th.ones_like(target) * self.weight_nonzero) * th.exp( - self.b * target - ) - weights = th.where(target > 0, weights_for_nonzero, weights_for_zero) - if average_channels: - loss = (th.mean(weights * (prediction - target) ** 2)) * self.w - else: - loss = (weights * (prediction - target) ** 2) * self.w - return loss - - class OceanMSE(th.nn.MSELoss): """ Ocean MSE class offers impementaion for MSE loss weighted by a land-sea-mask field. @@ -321,765 +256,3 @@ def forward(self, prediction, target, average_channels=True): return th.sum(ocean_mean_err) / self.lsm_sum else: return ocean_mean_err / self.lsm_var_sum - - -class WeightedCRPSLoss(th.nn.MSELoss): - """ - Probabilistic loss function that allows for user defined weighting of variables when calculating CRPS. - """ - - def __init__( - self, - weights: Sequence = [], - n_members: int = 2, - alpha: float = 0.95, - mean_penalty: float = 0.0, - lsm_file: str = None, - open_dict: dict = {"engine": "zarr"}, - selection_dict: dict = {"channel_c": "land_sea_mask"}, - multiscale: float = 0.0, - masked_processing: bool = False, - ): - """ - Parameters - ---------- - weights: Sequence - list of floats that determine weighting of variable loss, assumed to be - in order consistent with order of model output channels - n_members: int - number of ensemble members in the model output - alpha: float - hyperparamter for approximating fair CRPS loss. between 0 and 1, 1 corresponds to a fair CRPS loss. - mean_penalty: float - weight for the penalty constraining the global mean of the ensemble to be close to the target mean - if 0, no penalty is applied - lsm_file: str - land-sea-mask file - open_dict: dict, optional - dictionary that store land-sea-mask file information - selection_dict: dict, optional - dictionary that store channel selection information - multiscale: float, optional - weight for the multiscale CRPS loss. Default is 0, no multiscale loss is applied. - masked_processing: bool, optional - whether masked pixels should be excluded from the processing. Default is False, no masked processing is applied. - """ - super().__init__() - self.loss_weights = th.tensor(weights) - if n_members < 2: - raise ValueError("n_members must be at least 2 for CRPS loss to be defined") - else: - self.n_members = n_members - self.device = None - self.mean_penalty = mean_penalty - self.multiscale = multiscale - self.scales = [4, 16, 32] - self.masked_processing = masked_processing - - if lsm_file is not None: - self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel( - selection_dict - ) - self.lsm_tensor = 1 - th.tensor( - np.expand_dims(self.lsm_ds.values, (0, 2, 3)) - ) - else: - self.lsm_tensor = th.ones( - 1, 1, 1, 1, 1, 1 - ) # Spoof the tensor dimensions for broadcasting - - # Parameters for "almost fair CRPS" loss. See https://arxiv.org/html/2412.15832v1 - self.coeff_eps = 1 - ((1 - alpha) / (n_members)) - self.averaging_coeff = 1 / (2 * n_members * (n_members - 1)) - - # For n>2, will use pairwise distance to copmute [NxN] distance matrix - # Diagonal elements of (prediciton - target) matrix are zeroed out to avoid double counting - self.pdist = th.nn.PairwiseDistance(p=1) - self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye( - self.n_members - ) # Mask to zero out diagonal elements - - def setup(self, trainer): - """ - pushes constants to cuda device - """ - - if len(trainer.output_variables) != len(self.loss_weights): - raise ValueError("Length of outputs and loss_weights is not the same!") - - self.loss_weights = self.loss_weights.to(device=trainer.device) - self.averaging_coeff = th.tensor(self.averaging_coeff, device=trainer.device) - self.coeff_eps = th.tensor(self.coeff_eps, device=trainer.device) - self.pdist = self.pdist.to(device=trainer.device) - self.diag_mask = self.diag_mask.to(device=trainer.device) - self.lsm_tensor = self.lsm_tensor.to(device=trainer.device) - - def _2member_crps(self, prediction, target, lsm_tensor): - diff_target = th.abs(prediction - target.unsqueeze(0)).sum( - dim=0 - ) # [B, F, T, C, H, W] - diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] - crps = self.averaging_coeff * ( - diff_target - self.coeff_eps * diff_ensemble - ) # [B, F, T, C, H, W] - crps *= lsm_tensor - return crps - - def _pool(self, tensor, scale): - shape = tensor.shape - h, w = shape[-2:] - pooled = th.nn.functional.avg_pool2d( - tensor.reshape(shape[0], -1, h, w), scale, scale - ) - return pooled.reshape(*shape[:-2], h // scale, w // scale) - - def _masked_pool(self, tensor, scale, mask): - """ - Pools a tensor while excluding masked pixels - - Parameters - ---------- - tensor: torch.Tensor - The tensor to pool - scale: int - The scale factor for the pooling - mask: torch.Tensor - The mask to apply to the tensor, masked should only contain 1s and 0s - - Returns - ------- - torch.Tensor - The pooled tensor - torch.Tensor - The pooled mask - """ - # Apply the mask to the tensor - masked_values = tensor * mask - - # Compute the non-masked values for the pooled tensor - pooled_values = self._pool(masked_values, scale) - - # pool the mask to use as a weight for the pooled tensor - pooled_mask = self._pool(mask, scale) - - pooled_tensor = pooled_values / (pooled_mask + 1e-8) # Avoid division by zero - return pooled_tensor, pooled_mask - - def forward(self, prediction, target, average_channels=True): - """ - Forward pass of the WeightedCRPSLoss - Computes the CRPS loss for the model prediction and target. - - Parameters - ---------- - prediction: torch.Tensor - The prediction tensor shape [Cond*B, F, T, C, H, W] where Cond is the number of ensemble members - target: torch.Tensor - The target tensor shape [B, F, T, C, H, W] - average_channels: bool, optional - whether the mean of the channels should be taken - """ - - # Unfold ensemble dimension from batch dimension to have shape [Cond, B, F, T, C, H, W] - b, f, t, c, h, w = target.shape - prediction = prediction.view(self.n_members, b, f, t, c, h, w) - - # checks for dimensions - if not prediction.shape[1:] == target.shape: - raise ValueError( - f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}" - ) - - if not prediction.shape[0] == self.n_members: - raise ValueError( - f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}" - ) - - n = self.n_members - - # Manual Cast - prediction = prediction.to(th.float32) - target = target.to(th.float32) - - # Apply channel weights across channel dims - prediction *= self.loss_weights[None, None, None, None, :, None, None] - target *= self.loss_weights[None, None, None, :, None, None] - - if n == 2: - # Use faster explicit implementation - crps = self._2member_crps(prediction, target, self.lsm_tensor) - - if average_channels: - loss = crps.mean() - else: - loss = crps.mean(dim=(0, 1, 2, 4, 5)) - - if self.mean_penalty > 0: - # the fraction of valid pixels in land sea mask - if self.masked_processing and self.lsm_tensor.numel() > 1: - valid_pixels = self.lsm_tensor.numel() - else: - valid_pixels = f * h * w - - prediction_count = n * b * t * valid_pixels - target_count = b * t * valid_pixels - ens_global_means = (prediction * self.lsm_tensor).sum( - dim=(0, 1, 2, 3, 5, 6) - ) / prediction_count - target_global_means = (target * self.lsm_tensor).sum( - dim=(0, 1, 2, 4, 5) - ) / target_count - - bias_penalty = self.mean_penalty * th.abs( - ens_global_means - target_global_means - ) - if average_channels: - loss += bias_penalty.mean() - else: - loss += bias_penalty - - # spatial multiscale loss - if self.multiscale > 0.0: - crps_scales = 0 - for scale in self.scales: - if self.masked_processing and self.lsm_tensor.numel() > 1: - masked_pred, masked_lsm = self._masked_pool( - prediction, scale, self.lsm_tensor - ) - masked_tar, _ = self._masked_pool( - target, scale, self.lsm_tensor - ) - crps_scale = self._2member_crps( - masked_pred, masked_tar, masked_lsm - ) - else: - pred, tar, lsm = ( - self._pool(prediction, scale), - self._pool(target, scale), - self._pool(self.lsm_tensor, scale), - ) - crps_scale = self._2member_crps(pred, tar, lsm) - - if average_channels: - crps_scale = crps_scale.mean() - else: - crps_scale = crps_scale.mean(dim=(0, 1, 2, 4, 5)) - crps_scales += crps_scale - - crps_scales = crps_scales / len(self.scales) - loss += self.multiscale * crps_scales - - return loss - else: - # Do mean penalty - bias_penalty = 0 - if self.mean_penalty > 0: - # the fraction of valid pixels in land sea mask - if self.masked_processing and self.lsm_tensor.numel() > 1: - valid_pixels = self.lsm_tensor.numel() - else: - valid_pixels = f * h * w - - prediction_count = n * b * t * valid_pixels - target_count = b * t * valid_pixels - ens_global_means = (prediction * self.lsm_tensor).sum( - dim=(0, 1, 2, 3, 5, 6) - ) / prediction_count - target_global_means = (target * self.lsm_tensor).sum( - dim=(0, 1, 2, 4, 5) - ) / target_count - - bias_penalty = self.mean_penalty * th.abs( - ens_global_means - target_global_means - ) - if average_channels: - bias_penalty = bias_penalty.mean() - - # zero out land and determine number of valid pixels - if self.masked_processing and self.lsm_tensor.numel() > 1: - prediction = prediction * self.lsm_tensor - target = target * self.lsm_tensor - valid_pixels = b * t * self.lsm_tensor.sum() - else: - valid_pixels = b * f * t * h * w - - # Use pairwise distance method - if not average_channels: - # Move channels to first dimension and exclude that dimension from the reductions - prediction = prediction.permute( - 4, 0, 1, 2, 3, 5, 6 - ) # [C, Cond, B, F, T, H, W] - target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] - - prediction = prediction.reshape(c, n, -1) # [C, Cond, ...] - target = target.unsqueeze(1).reshape( - c, 1, -1 - ) # [C, 1, ...] (second dim will broadcast across ensemble) - - diff = self.pdist(prediction, target) # [C, Cond] - dist_matrix = self.pdist( - prediction.unsqueeze(1), prediction.unsqueeze(2) - ) # [C, Cond, Cond] - - diff_terms = self.diag_mask[None, ...] * ( - diff.unsqueeze(1) + diff.unsqueeze(2) - ) # [C, Cond, Cond], diagonal elements zeroed out - crps = ( - self.averaging_coeff - * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1, 2)) - / valid_pixels - ) - else: - prediction = prediction.reshape(n, -1) - target = target.unsqueeze(0).reshape( - 1, -1 - ) # [1, ...] (first dim will broadcast across ensemble) - diff = self.pdist(prediction, target) # [Cond] - dist_matrix = self.pdist( - prediction.unsqueeze(1), prediction.unsqueeze(0) - ) # [Cond, Cond] - - diff_terms = self.diag_mask * ( - diff.unsqueeze(0) + diff.unsqueeze(1) - ) # [Cond, Cond], diagonal elements zeroed out - crps = ( - self.averaging_coeff - * (diff_terms - self.coeff_eps * dist_matrix).sum() - / valid_pixels - ) - - return crps + bias_penalty - - -class WeightedCRPSLossSpectral(th.nn.MSELoss): - """ - Probabilistic loss function that allows for user defined weighting of variables when calculating CRPS. - """ - - def __init__( - self, - weights: Sequence = [], - n_members: int = 2, - alpha: float = 0.95, - lambda_spec: float = 0.1, - nside: int = 64, - lmax: int = 3 * 64 - 1, - mmax: int = 3 * 64 - 1, - multiscale: float = 0.0, - lsm_file: str = None, - open_dict: dict = {"engine": "zarr"}, - selection_dict: dict = {"channel_c": "land_sea_mask"}, - ): - """ - Parameters - ---------- - weights: Sequence - list of floats that determine weighting of variable loss, assumed to be - in order consistent with order of model output channels - n_members: int - number of ensemble members in the model output - alpha: float - hyperparamter for approximating fair CRPS loss. between 0 and 1, 1 corresponds to a fair CRPS loss. - lambda_spec: float - weight for the spectral loss. Default is 0, no spectral loss is applied. - nside: int - nside for the HEALPix grid. Default is 64. - lmax: int - lmax for the SHT. Default is 3*nside - 1. - mmax: int - mmax for the SHT. Default is 3*nside - 1. - multiscale: float, optional - weight for the multiscale loss. Default is 0, no multiscale loss is applied. - lsm_file: str - path to the lsm file. Default is None, no lsm is applied. - open_dict: dict - dictionary of keyword arguments for xarray.open_dataset. Default is {"engine": "zarr"}. - selection_dict: dict - dictionary of keyword arguments for xarray.open_dataset. Default is {"channel_c": "land_sea_mask"}. - """ - super().__init__() - self.loss_weights = th.tensor(weights) - self.n_members = n_members - self.device = None - self.lambda_spec = lambda_spec - self.multiscale = multiscale - - # Parameters for "almost fair CRPS" loss. See https://arxiv.org/html/2412.15832v1 - self.coeff_eps = 1 - ((1 - alpha) / (n_members)) - if n_members < 2: - raise ValueError("n_members must be at least 2 for CRPS loss to be defined") - else: - self.n_members = n_members - self.averaging_coeff = 1 / (2 * n_members * (n_members - 1)) - - # SHT utils: transform, grid reordering, output indexing - self.lmax = lmax - self.mmax = mmax - self.nside = nside - self.sht = cuhpx.SHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights="ring") - src_grid = earth2grid.healpix.Grid( - level=int(np.log2(nside)), pixel_order=earth2grid.healpix.HEALPIX_PAD_XY - ) - tar_grid = earth2grid.healpix.Grid( - level=int(np.log2(nside)), pixel_order=earth2grid.healpix.PixelOrder.RING - ) - self.reorder_to_ring = earth2grid.get_regridder(src_grid, tar_grid).to( - th.float32 - ) - if self.multiscale > 0: - self.scales = [200, 400, 800, 1600] # in units of km - self.isht = cuhpx.iSHTCUDA( - nside=nside, lmax=lmax, mmax=mmax, quad_weights="ring" - ) - self.reorder_from_ring = earth2grid.get_regridder(tar_grid, src_grid).to( - th.float32 - ) - - self.lsm_file = lsm_file - if lsm_file is not None: - self.lsm_ds = xr.open_dataset(lsm_file, **open_dict).constants.sel( - selection_dict - ) - self.lsm_tensor = 1 - th.tensor( - np.expand_dims(self.lsm_ds.values, (0, 2, 3)) - ) - else: - self.lsm_tensor = th.ones( - 1, 1, 1, 1, 1, 1 - ) # Spoof the tensor dimensions for broadcasting - - # For n>2, will use pairwise distance to copmute [NxN] distance matrix - # Diagonal elements of (prediciton - target) matrix are zeroed out to avoid double counting - self.pdist = th.nn.PairwiseDistance(p=1) - self.diag_mask = th.ones(self.n_members, self.n_members) - th.eye( - self.n_members - ) # Mask to zero out diagonal elements - - def setup(self, trainer): - """ - pushes constants to cuda device - """ - - if len(trainer.output_variables) != len(self.loss_weights): - raise ValueError("Length of outputs and loss_weights is not the same!") - - self.loss_weights = self.loss_weights.to(device=trainer.device) - self.averaging_coeff = th.tensor(self.averaging_coeff, device=trainer.device) - self.coeff_eps = th.tensor(self.coeff_eps, device=trainer.device) - self.reorder_to_ring = self.reorder_to_ring.to(device=trainer.device) - self.sht = self.sht.to(device=trainer.device) - self.pdist = self.pdist.to(device=trainer.device) - self.diag_mask = self.diag_mask.to(device=trainer.device) - - if self.multiscale > 0: - self.isht = self.isht.to(device=trainer.device) - self.reorder_from_ring = self.reorder_from_ring.to(device=trainer.device) - self.lsm_tensor = self.lsm_tensor.to(device=trainer.device) - - def _apply_sht(self, x, face_dim, return_abs=True): - """Apply SHT to a tensor - Reshape to [..., F*H*W], reorder to ring, apply SHT - If return_abs is True, return the absolute value of the SHT (real**2 + imag**2) - - Parameters - ---------- - x: torch.Tensor - The tensor to apply SHT to - face_dim: int - The dimension of the tensor corrsponding to HEALPix faces - return_abs: bool, optional - Whether to return the absolute value of the SHT (real**2 + imag**2) - """ - x = th.movedim(x, face_dim, -3) - if x.shape[-3:] != (12, self.nside, self.nside): - face, height, width = x.shape[-3:] - raise ValueError( - f"Shape of input tensor should be [..., 12, ..., {self.nside}, {self.nside}] with F (12) in position {face_dim}," - f"got {x.shape}" - ) - - x = x.reshape(*x.shape[:-3], -1) - x = self.reorder_to_ring( - x.contiguous() - ) # contiguous needed for channels first format in validation loop - x = self.sht(x) - if return_abs: - x = x.real**2 + x.imag**2 - return x - - def _apply_isht(self, x, face_dim): - """Apply inverse SHT to a tensor shape [..., l, m] - Inverse transform, reorder from ring, Reshape to [..., F, H, W], move face dim appropriately - """ - - x = self.isht(x) # [..., l, m] -> [..., F*H*W] - x = self.reorder_from_ring(x) - x = x.reshape( - *x.shape[:-1], 12, self.nside, self.nside - ) # [..., F*H*W] -> [..., F, H, W] - x = th.movedim(x, -3, face_dim) # [..., F, H, W] -> [..., F, ..., H, W] - return x - - def _l_filter(self, scale, device="cuda"): - """Return a spherical gaussian filter of scale `scale` (in units of km)""" - scale_radians = scale / 6371.0 - ell = th.arange(self.lmax, device=device, dtype=th.float32) - return th.exp(-0.5 * ell * (ell + 1) * (scale_radians**2)) - - def forward(self, prediction, target, average_channels=True): - """ - Forward pass of the WeightedCRPSLoss - Computes the CRPS loss for the model prediction and target. - - Parameters - ---------- - prediction: torch.Tensor - The prediction tensor shape [Cond*B, F, T, C, H, W] where Cond is the number of ensemble members - target: torch.Tensor - The target tensor shape [B, F, T, C, H, W] - average_channels: bool, optional - whether the mean of the channels should be taken - """ - - # Unfold ensemble dimension from batch dimension to have shape [Cond, B, F, T, C, H, W] - b, f, t, c, h, w = target.shape - prediction = prediction.view(self.n_members, b, f, t, c, h, w) - - # checks for dimensions - if not prediction.shape[1:] == target.shape: - raise ValueError( - f"Shape of prediction should match shape of target along non-ensemble dimensions, got {prediction.shape} and {target.shape}" - ) - - if not prediction.shape[0] == self.n_members: - raise ValueError( - f"Shape of prediction should have ensemble dimension of size {self.n_members}, got {prediction.shape[0]}" - ) - - n = self.n_members - - # Manual cast - prediction = prediction.to(th.float32) - target = target.to(th.float32) - - # Apply channel weights across channel dims - prediction *= self.loss_weights[None, None, None, None, :, None, None] - target *= self.loss_weights[None, None, None, :, None, None] - - if n == 2: - # Use faster explicit implementation - diff_target = th.abs(prediction - target.unsqueeze(0)).sum( - dim=0 - ) # [B, F, T, C, H, W] - diff_ensemble = th.abs(prediction[0] - prediction[1]) # [B, F, T, C, H, W] - crps = self.averaging_coeff * ( - diff_target - self.coeff_eps * diff_ensemble - ) # [B, F, T, C, H, W] - - if average_channels: - loss = crps.mean() - else: - loss = crps.mean(dim=(0, 1, 2, 4, 5)) - - if self.lambda_spec > 0: - with th.amp.autocast("cuda", enabled=False): - # # Reorder predictions: [N, B, F, T, C, H, W] -> [N, B, T, C, F*H*W] - # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) - - # # Reorder targets: [B, F, T, C, H, W] -> [B, T, C, F*H*W] - # tar_ring = self.reorder_to_ring(target.permute(0, 2, 3, 1, 4, 5).reshape(b, t, c, f*h*w)) - - # # Compute SHT of predictions and targets - # sht_pred = self.sht(pred_ring) # [N, B, T, C, l, m] - # sht_tar = self.sht(tar_ring) # [B, T, C, l, m] - # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 - # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 - - sht_pred = self._apply_sht(prediction, face_dim=2, return_abs=True) - sht_tar = self._apply_sht(target, face_dim=1, return_abs=True) - - diff_sht_target = th.abs(sht_pred - sht_tar.unsqueeze(0)).sum( - dim=(0, 4, 5) - ) # [B, T, C] - diff_sht_ensemble = th.abs(sht_pred[0] - sht_pred[1]).sum( - dim=(-1, -2) - ) # [B, T, C] - crps_sht = self.averaging_coeff * ( - diff_sht_target - self.coeff_eps * diff_sht_ensemble - ) # [B, T, C] - - # Compute spectral afCRPS - if average_channels: - spec_loss = crps_sht.mean() - else: - spec_loss = crps_sht.mean(dim=(0, 1)) - - loss += self.lambda_spec * spec_loss - - if self.multiscale > 0: - for scale in self.scales: - l_filter = self._l_filter(scale, device=prediction.device) - with th.amp.autocast("cuda", enabled=False): - sht_pred = self._apply_sht( - prediction, face_dim=2, return_abs=False - ) - sht_tar = self._apply_sht(target, face_dim=1, return_abs=False) - - l_filter_pred = l_filter[ - None, None, None, None, :, None - ] # [1, 1, 1, 1, lmax, 1] - l_filter_tar = l_filter[ - None, None, None, :, None - ] # [1, 1, 1, lmax, 1] - - pred_smooth = self._apply_isht( - l_filter_pred * sht_pred, face_dim=2 - ) - tar_smooth = self._apply_isht( - l_filter_tar * sht_tar, face_dim=1 - ) - - diff_target = th.abs(pred_smooth - tar_smooth.unsqueeze(0)).sum( - dim=0 - ) # [B, F, T, C, H, W] - diff_ensemble = th.abs( - pred_smooth[0] - pred_smooth[1] - ) # [B, F, T, C, H, W] - crps = self.averaging_coeff * ( - diff_target - self.coeff_eps * diff_ensemble - ) # [B, F, T, C, H, W] - - crps *= self.lsm_tensor - - if average_channels: - loss += self.multiscale * crps.mean() - else: - loss += self.multiscale * crps.mean(dim=(0, 1, 2, 4, 5)) - - return loss - - else: - # Use pairwise distance method - if not average_channels: - # Move channels to first dimension and exclude that dimension from the reductions - prediction = prediction.permute( - 4, 0, 1, 2, 3, 5, 6 - ) # [C, Cond, B, F, T, H, W] - target = target.permute(3, 0, 1, 2, 4, 5) # [C, B, F, T, H, W] - - pred = prediction.reshape(c, n, -1) # [C, Cond, ...] - tar = target.unsqueeze(1).reshape( - c, 1, -1 - ) # [C, 1, ...] (second dim will broadcast across ensemble) - - diff = self.pdist(pred, tar) # [C, Cond] - dist_matrix = self.pdist( - pred.unsqueeze(1), pred.unsqueeze(2) - ) # [C, Cond, Cond] - - diff_terms = self.diag_mask[None, ...] * ( - diff.unsqueeze(1) + diff.unsqueeze(2) - ) # [C, Cond, Cond], diagonal elements zeroed out - loss = ( - self.averaging_coeff - * (diff_terms - self.coeff_eps * dist_matrix).sum(dim=(1, 2)) - / (b * f * t * h * w) - ) - - if self.lambda_spec > 0: - with th.amp.autocast("cuda", enabled=False): - # # Reorder predictions: [C, Cond, B, F, T, H, W] -> [C, Cond, B, T, F*H*W] - # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 2, 4, 3, 5, 6).reshape(c, n, b, t, f*h*w)) - - # # Reorder targets: [C, B, F, T, H, W] -> [C, B, T, F*H*W] - # tar_ring = self.reorder_to_ring(target.permute(0, 1, 3, 2, 4, 5).reshape(c, b, t, f*h*w)) - - # # Compute SHT of predictions and targets - # sht_pred = self.sht(pred_ring).reshape(c, n, -1) # [C, Cond, B, T, l, m] -> [C, Cond, ...] - # sht_tar = self.sht(tar_ring).unsqueeze(1).reshape(c, 1, -1) # [C, B, T, l, m] -> [C, 1, ...] (second dim will broadcast across ensemble) - # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 - # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 - - sht_pred = self._apply_sht( - prediction, face_dim=3, return_abs=True - ).reshape(c, n, -1) - sht_tar = ( - self._apply_sht(target, face_dim=2, return_abs=True) - .unsqueeze(1) - .reshape(c, 1, -1) - ) - - diff = self.pdist(sht_pred, sht_tar) # [C, Cond] - dist_matrix = self.pdist( - sht_pred.unsqueeze(1), sht_pred.unsqueeze(2) - ) # [C, Cond, Cond] - - diff_terms = self.diag_mask[None, ...] * ( - diff.unsqueeze(1) + diff.unsqueeze(2) - ) # [C, Cond, Cond], diagonal elements zeroed out - loss += ( - self.lambda_spec - * self.averaging_coeff - * (diff_terms - self.coeff_eps * dist_matrix).sum( - dim=(1, 2) - ) - / (b * t * self.lmax * self.mmax) - ) - - else: - pred = prediction.reshape(n, -1) - tar = target.unsqueeze(0).reshape( - 1, -1 - ) # [1, ...] (first dim will broadcast across ensemble) - diff = self.pdist(pred, tar) # [Cond] - dist_matrix = self.pdist( - pred.unsqueeze(1), pred.unsqueeze(0) - ) # [Cond, Cond] - - diff_terms = self.diag_mask * ( - diff.unsqueeze(0) + diff.unsqueeze(1) - ) # [Cond, Cond], diagonal elements zeroed out - loss = ( - self.averaging_coeff - * (diff_terms - self.coeff_eps * dist_matrix).sum() - / (b * f * c * t * h * w) - ) - - if self.lambda_spec > 0: - with th.amp.autocast("cuda", enabled=False): - # # Reorder predictions: [Cond, B, F, T, C, H, W] -> [Cond, B, T, C, F*H*W] - # pred_ring = self.reorder_to_ring(prediction.permute(0, 1, 3, 4, 2, 5, 6).reshape(n, b, t, c, f*h*w)) - - # # Reorder targets: [B, F, T, C, H, W] -> [B, T, C, F*H*W] - # tar_ring = self.reorder_to_ring(target.permute(0, 2, 3, 1, 4, 5).reshape(b, t, c, f*h*w)) - - # # Compute SHT of predictions and targets - # sht_pred = self.sht(pred_ring).reshape(n, -1) # [Cond, B, T, C, l, m] -> [Cond, ...] - # sht_tar = self.sht(tar_ring).unsqueeze(0).reshape(1, -1) # [B, T, C, l, m] -> [1, ...] (first dim will broadcast across ensemble) - # sht_pred = sht_pred.real ** 2 + sht_pred.imag ** 2 - # sht_tar = sht_tar.real ** 2 + sht_tar.imag ** 2 - sht_pred = self._apply_sht( - prediction, face_dim=2, return_abs=True - ).reshape(n, -1) - sht_tar = ( - self._apply_sht(target, face_dim=1, return_abs=True) - .unsqueeze(0) - .reshape(1, -1) - ) - - diff = self.pdist(sht_pred, sht_tar) # [Cond] - dist_matrix = self.pdist( - sht_pred.unsqueeze(1), sht_pred.unsqueeze(0) - ) # [Cond, Cond] - - diff_terms = self.diag_mask * ( - diff.unsqueeze(0) + diff.unsqueeze(1) - ) # [Cond, Cond], diagonal elements zeroed out - loss += ( - self.lambda_spec - * self.averaging_coeff - * (diff_terms - self.coeff_eps * dist_matrix).sum() - / (b * t * self.lmax * self.mmax) - ) - - return loss diff --git a/physicsnemo/metrics/climate/hydrostasy.py b/physicsnemo/metrics/climate/hydrostasy.py deleted file mode 100644 index c16d318559..0000000000 --- a/physicsnemo/metrics/climate/hydrostasy.py +++ /dev/null @@ -1,900 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import logging -from typing import Dict, Sequence - -import numpy as np -import torch - -from physicsnemo.core.version_check import OptionalImport - -logger = logging.getLogger(__name__) - -# xarray ships in the ``datapipes-extras`` optional dependency group and is only -# needed by the topography-loading helpers here; load it lazily so the -# physicsnemo import graph carries no static dependency on it (see -# CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003/EXT-004). -xr = OptionalImport("xarray") - - -def _average_virtual_temperature_from_geopotential_height(z1, z2, p1, p2, R, g0): - return g0 / (R * np.log(p1 / p2)) * (z2 - z1) - - -def _virtual_temperature_from_geopotential_height(z1, z2, p1, p2, T1, R, g0): - return (2.0 * g0) / (R * np.log(p1 / p2)) * (z2 - z1) - T1 - - -class HydrostaticBalance(torch.nn.Module): - """Derive virtual temperature at pressure levels from geopotential height. - - Given geopotential heights (Z) at a set of pressure levels and a known - virtual temperature anchor, integrates the hypsometric equation outward - from the anchor level to recover virtual temperature at every other - level. - """ - - def __init__( - self, - z_pressure_levels: Dict[int, float], - anchor_z_channel: int, - anchor_T_channel: int, - R: float, - g0: float, - ) -> None: - super().__init__() - self.R = R - self.g0 = g0 - - if anchor_z_channel not in z_pressure_levels: - raise ValueError( - f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" - ) - - # Sort channels by pressure levels for ease of use later - self.z_pressure_levels = dict( - sorted(z_pressure_levels.items(), key=lambda item: item[1]) - ) - self.anchor_z_channel = anchor_z_channel - self.anchor_T_channel = anchor_T_channel - self.anchor_pressure = z_pressure_levels[anchor_z_channel] - self.anchor_z_index = list(self.z_pressure_levels.keys()).index( - anchor_z_channel - ) - self.z_channels = list(self.z_pressure_levels.keys()) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - - Parameters - ---------- - x : torch.Tensor - Input tensor that contains the geopotential heights (Z) at the channels and pressure levels specified in the constructor - [B, F, C, H, W] is the format - - - Returns - ------- - torch.Tensor - - """ - Tv_size = [ - len(self.z_pressure_levels) if i == 2 else s for i, s in enumerate(x.size()) - ] - Tv = torch.empty(Tv_size, dtype=x.dtype, layout=x.layout, device=x.device) - Tv[:, :, self.anchor_z_index, ...] = x[:, :, self.anchor_T_channel, ...] - - # Go down in index (up in vertical level) from anchor - for i in range(self.anchor_z_index, 0, -1): - z_channel = self.z_channels[i] - z_channel_m1 = self.z_channels[i - 1] - zi = x[:, :, z_channel, ...] - zim1 = x[:, :, z_channel_m1, ...] - Tv[:, :, i - 1, ...] = _virtual_temperature_from_geopotential_height( - zi, - zim1, - self.z_pressure_levels[z_channel], - self.z_pressure_levels[z_channel_m1], - Tv[:, :, i, ...], - self.R, - self.g0, - ) - - # Go up in index (down in vertical level) from anchor - for i in range(self.anchor_z_index, len(self.z_pressure_levels) - 1): - z_channel = self.z_channels[i] - z_channel_p1 = self.z_channels[i + 1] - zi = x[:, :, z_channel, ...] - zip1 = x[:, :, z_channel_p1, ...] - Tv[:, :, i + 1, ...] = _virtual_temperature_from_geopotential_height( - zi, - zip1, - self.z_pressure_levels[z_channel], - self.z_pressure_levels[z_channel_p1], - Tv[:, :, i, ...], - self.R, - self.g0, - ) - - return Tv - - -class DifferentialHydrostaticBalanceConstraint(torch.nn.Module): - """Compute layer-average virtual temperatures implied by geopotential height and model output. - - For each pair of adjacent pressure levels, derives the layer-mean virtual - temperature implied by the hypsometric equation from geopotential height - (``Tv_avg``) alongside the layer-mean virtual temperature predicted - directly by the model (``Tv_model_avg``), so the two can be compared as a - hydrostatic-balance constraint. - """ - - def __init__( - self, - z_pressure_levels: Dict[int, float], - Tv_pressure_levels: Dict[int, float], - anchor_z_channel: int, - anchor_T_channel: int, - R: float, - g0: float, - ) -> None: - super().__init__() - self.R = R - self.g0 = g0 - - if anchor_z_channel not in z_pressure_levels.keys(): - raise ValueError( - f"anchor_z_channel ({anchor_z_channel}) not in z_pressure_levels ({z_pressure_levels.keys()})" - ) - if anchor_T_channel not in Tv_pressure_levels.keys(): - raise ValueError( - f"anchor_T_channel ({anchor_T_channel}) not in Tv_pressure_levels ({Tv_pressure_levels.keys()})" - ) - - # Sort channels by pressure levels for ease of use later - self.z_pressure_levels = dict( - sorted(z_pressure_levels.items(), key=lambda item: item[1]) - ) - self.Tv_pressure_levels = dict( - sorted(Tv_pressure_levels.items(), key=lambda item: item[1]) - ) - self.anchor_z_channel = anchor_z_channel - self.anchor_T_channel = anchor_T_channel - self.anchor_pressure = z_pressure_levels[anchor_z_channel] - self.anchor_z_index = list(self.z_pressure_levels.keys()).index( - anchor_z_channel - ) - self.anchor_T_index = list(self.Tv_pressure_levels.keys()).index( - anchor_T_channel - ) - self.z_channels = list(self.z_pressure_levels.keys()) - self.Tv_channels = list(self.Tv_pressure_levels.keys()) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - - Parameters - ---------- - x : torch.Tensor - Input tensor that contains the geopotential heights (Z) at the channels and pressure levels specified in the constructor - [B, F, C, H, W] is the format - - - Returns - ------- - torch.Tensor - - """ - Tv_avg_size = [ - len(self.z_pressure_levels) - 1 if i == 2 else s - for i, s in enumerate(x.size()) - ] - Tv_avg = torch.empty( - Tv_avg_size, dtype=x.dtype, layout=x.layout, device=x.device - ) - - # Go up in index (down in vertical level) from anchor - # TODO: remove constraint that first z_channel is 0 - for i in range(len(self.z_pressure_levels) - 1): - z_channel = self.z_channels[i] - z_channel_p1 = self.z_channels[i + 1] - zi = x[:, :, z_channel, ...] - zip1 = x[:, :, z_channel_p1, ...] - Tv_avg[:, :, i, ...] = ( - _average_virtual_temperature_from_geopotential_height( - zi, - zip1, - self.z_pressure_levels[z_channel], - self.z_pressure_levels[z_channel_p1], - self.R, - self.g0, - ) - ) - - Tv_model_avg_size = [ - len(self.Tv_pressure_levels) - 1 if i == 2 else s - for i, s in enumerate(x.size()) - ] - Tv_model_avg = torch.empty( - Tv_model_avg_size, dtype=x.dtype, layout=x.layout, device=x.device - ) - - # Go up in index (down in vertical level) from anchor - for i in range(len(self.Tv_pressure_levels) - 1): - Tv_channel = self.Tv_channels[i] - Tv_channel_p1 = self.Tv_channels[i + 1] - Tvi = x[:, :, Tv_channel, ...] - Tvip1 = x[:, :, Tv_channel_p1, ...] - Tv_model_avg[:, :, i, ...] = 0.5 * (Tvi + Tvip1) - - return Tv_avg, Tv_model_avg - - -class WeightedMSEWithHydrostasy(torch.nn.MSELoss): - """ - Loss object that adds a differential Hydrostatic balance constraint in addition to - user defined weighting of variables when calculating MSE - """ - - def __init__( - self, - hPa_levels: Sequence[int], - channels: Sequence[str], - weights: Sequence, - alpha: Sequence[float], # K - scaling: Dict[str, Dict[str, float]], - src_directory: str, - dst_directory: str, - dataset_name: str, - data_format: str, - surface_geopotential_mean: float = -597.7115478515625, - surface_geopotential_std: float = 55658.21484375, - R: float = 287, # J K^{-1} kg^{-1} - g0: float = 9.81, # m s^{-2} - topography_masking: bool = True, - ): - """ - Parameters - ---------- - weights: Sequence - list of floats that determine weighting of variable loss, assumed to be - in order consistent with order of model output channels - """ - super().__init__() - self.loss_weights = torch.tensor(weights) - self.device = None - self.g0 = g0 - self.topography_masking = topography_masking - - # Get channel index to pressure level mapping - self.pressure_levels = sorted(hPa_levels) - self.z_pressure_levels = { - channels.index(f"z{int(pl)}"): pl - for i, pl in enumerate(self.pressure_levels) - } - self.T_pressure_levels = { - channels.index(f"t{int(pl)}"): pl - for i, pl in enumerate(self.pressure_levels) - } - self.q_pressure_levels = { - channels.index(f"q{int(pl)}"): pl - for i, pl in enumerate(self.pressure_levels) - if f"q{int(pl)}" in channels - } - # Get offset to map from q channels here to Tv channels - # Relies on all pressure levels below a threshold to have q channels - self.q_index_offset = None - for i, pl in enumerate(self.pressure_levels): - if f"q{int(pl)}" in channels: - self.q_index_offset = i - break - if self.q_index_offset is None: - raise ValueError( - f"No q channels found in channels list that match the requested pressure levels {self.pressure_levels}" - ) - - # Create mapping for new tensor that holds only the constraint variables - self.z_constraint_pressure_levels = { - i: pl for i, pl in enumerate(self.pressure_levels) - } - self.Tv_constraint_pressure_levels = { - len(self.z_constraint_pressure_levels) + i: pl - for i, pl in enumerate(self.pressure_levels) - } - - # Get scaling weights - self.z_mean = torch.Tensor( - [ - scaling[f"z{int(pl)}"]["mean"] - for i, pl in enumerate(self.pressure_levels) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - self.z_std = torch.Tensor( - [scaling[f"z{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] - ).reshape((1, 1, 1, -1, 1, 1)) - self.T_mean = torch.Tensor( - [ - scaling[f"t{int(pl)}"]["mean"] - for i, pl in enumerate(self.pressure_levels) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - self.T_std = torch.Tensor( - [scaling[f"t{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] - ).reshape((1, 1, 1, -1, 1, 1)) - self.q_mean = torch.Tensor( - [ - scaling[f"q{int(pl)}"]["mean"] - for i, pl in enumerate(self.q_pressure_levels.values()) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - self.q_std = torch.Tensor( - [ - scaling[f"q{int(pl)}"]["std"] - for i, pl in enumerate(self.q_pressure_levels.values()) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - - # Set per level alphas - if len(alpha) != len(hPa_levels) - 1: - raise AssertionError( - f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels) - 1}], got {len(alpha)}" - ) - self.alpha = torch.Tensor(alpha).reshape((1, 1, -1, 1, 1)) - - # Molecular weight ratio factor of water vapor to air - self.Mw_ratio = 28.97 / 18.016 - 1.0 # 0.6078 - - # Create the constraint - # TODO: remove anchor levels since it's not needed for the - # differential constraint - self.constraint = DifferentialHydrostaticBalanceConstraint( - self.z_constraint_pressure_levels, - self.Tv_constraint_pressure_levels, - 0, - len(self.z_pressure_levels), - R, - self.g0, - ) - - self.num_z_levels = len(self.z_constraint_pressure_levels) - self.num_Tv_levels = len(self.Tv_constraint_pressure_levels) - self.z_level_mapping = torch.tensor(list(self.z_pressure_levels.keys())) - self.T_level_mapping = torch.tensor(list(self.T_pressure_levels.keys())) - self.q_level_mapping = torch.tensor(list(self.q_pressure_levels.keys())) - - # Get topography information - ds = xr.open_zarr(f"{src_directory}{dataset_name}.zarr") - self.topography = ( - surface_geopotential_std * ds.constants[1, :, :, :].values - + surface_geopotential_mean - ) / self.g0 - - self.topography = torch.tensor( - self.topography[np.newaxis, :, np.newaxis, :, :], dtype=torch.float - ) - logger.info( - f"Min/Max topography (m): {self.topography.min()}/{self.topography.max()}" - ) - - def setup(self, trainer): - """ - pushes weights to cuda device - """ - - if len(trainer.output_variables) + len(self.z_pressure_levels) - 1 != len( - self.loss_weights - ): - raise ValueError("Length of outputs and loss_weights is not the same!") - - self.loss_weights = self.loss_weights.to(device=trainer.device) - - # Move means and stds - self.z_mean = self.z_mean.to(device=trainer.device) - self.z_std = self.z_std.to(device=trainer.device) - self.T_mean = self.T_mean.to(device=trainer.device) - self.T_std = self.T_std.to(device=trainer.device) - self.q_mean = self.q_mean.to(device=trainer.device) - self.q_std = self.q_std.to(device=trainer.device) - - # Move alphas - self.alpha = self.alpha.to(device=trainer.device) - - # Move indexing arrays for CUDA graphs - self.z_level_mapping = self.z_level_mapping.to(device=trainer.device) - self.T_level_mapping = self.T_level_mapping.to(device=trainer.device) - self.q_level_mapping = self.q_level_mapping.to(device=trainer.device) - - # Move topography - self.topography = self.topography.to(device=trainer.device) - - def scale(self, x): - """ - Scale inputs to physical values and compute virtual temperature - Tensors are expected to be in the shape [N, B, F, C, H, W] - """ - # N, B, F, C, H, W = x.shape - N, F, B, C, H, W = x.shape - C_scaled = self.num_z_levels + self.num_Tv_levels - x_scaled = torch.zeros( - # (N, B, F, C_scaled, H, W), device=x.device, dtype=torch.float - (N, F, B, C_scaled, H, W), - device=x.device, - dtype=torch.float, - ) - # Get scaled geopotential heights - x_scaled[:, :, :, : self.num_z_levels, :, :] = ( - x[:, :, :, self.z_level_mapping, :, :] * self.z_std + self.z_mean - ) / self.g0 # divide by g0 for heights - # Get scaled temperatures - x_scaled[:, :, :, self.num_z_levels :, :, :] = ( - x[:, :, :, self.T_level_mapping, :, :] * self.T_std + self.T_mean - ) - # Add q correction to get virtual temperature for levels with non-zero q - x_scaled[ - :, - :, - :, - (self.num_z_levels + self.q_index_offset) :, - :, - :, - ] *= 1.0 + self.Mw_ratio * ( - x[:, :, :, self.q_level_mapping, :, :] * self.q_std + self.q_mean - ) - - # transpose B dim to before F - x_scaled = x_scaled.transpose(1, 2) - - # Combine N and B dimensions and return - return x_scaled.reshape((-1, F, C_scaled, H, W)) - - def error_histogram(self, prediction, bins, accumulator=None): - """Accumulate a per-level histogram of absolute hydrostatic-balance error. - - Parameters - ---------- - prediction: torch.Tensor - Model prediction tensor of shape [N, F, B, C, H, W]. - bins: int or torch.Tensor - Either the number of equal-width bins to use, or an existing - tensor of per-level bin edges to reuse. - accumulator: torch.Tensor, optional - Existing per-level histogram counts to add to. If None, a new - zero-initialized accumulator is created. - - Returns - ------- - tuple[torch.Tensor, torch.Tensor] - The updated ``(accumulator, bin_edges)``. - """ - N, F, B, C, H, W = tuple(prediction.shape) - - if not (prediction.ndim == 6): - raise AssertionError("Expected predictions to have 6 dimensions") - - # Scale to physical units and compute virtual temperature - x = self.scale(prediction) - Tv_avg, Tv_model_avg = self.constraint(x) - Tv_error = Tv_avg - Tv_model_avg - # Mask out error in regions below the surface - if self.topography_masking: - Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 - - vlevels = Tv_error.shape[2] - if accumulator is None: - if isinstance(bins, int): - accumulator = torch.zeros( - (vlevels, bins), dtype=torch.float32, device=prediction.device - ) - else: - accumulator = torch.zeros( - (vlevels, bins.shape[1] - 1), - dtype=torch.float32, - device=prediction.device, - ) - if isinstance(bins, int): - bin_edges = torch.zeros( - (vlevels, bins + 1), - dtype=prediction.dtype, - device=prediction.device, - ) - else: - bin_edges = bins - - for level_idx in range(vlevels): - hist, be = torch.histogram( - torch.absolute(Tv_error[:, :, level_idx, :, :]), - bins=bins if isinstance(bins, int) else bin_edges[level_idx, :], - ) - accumulator[level_idx, :] += hist - bin_edges[level_idx, :] = be - - return accumulator, bin_edges - - def forward(self, prediction, target, average_channels=True): - """ - Forward pass of the WeightedMSE pass - Tensors are expected to be in the shape [N, F, B, C, H, W] - - Parameters - ---------- - prediction: torch.Tensor - The prediction tensor - target: torch.Tensor - The target tensor - average_channels: bool, optional - whether the mean of the channels should be taken - """ - - # Need to scale back to physical units here so disable autocast - # and explicitly cast to float32 - with torch.cuda.amp.autocast(enabled=False): - prediction = prediction.float() - target = target.float() - - N, F, B, C, H, W = tuple(prediction.shape) - - if not (prediction.ndim == 6 and target.ndim == 6): - raise AssertionError("Expected predictions to have 6 dimensions") - - # Scale to physical units and compute virtual temperature - x = self.scale(prediction) - Tv_avg, Tv_model_avg = self.constraint(x) - Tv_error = ((Tv_avg - Tv_model_avg) / self.alpha) ** 2 - - # Mask out error in regions below the surface - if self.topography_masking: - Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 - - # Compute the error tolerant loss - Tv_loss = (Tv_error / (1 + torch.exp(1 - Tv_error))).mean(dim=(0, 1, 3, 4)) - - data_loss = ((target - prediction) ** 2).mean(dim=(0, 1, 2, 4, 5)) - d = torch.concatenate((data_loss, Tv_loss)) * self.loss_weights - - if average_channels: - return torch.mean(d) - else: - return d - - -class LossWithHydrostasy(torch.nn.MSELoss): - """ - Loss object that adds a differential Hydrostatic balance constraint in addition to - user defined data loss function. - """ - - def __init__( - self, - data_loss: torch.nn.MSELoss, - hPa_levels: Sequence[int], - channels: Sequence[str], - weights: Sequence, - alpha: Sequence[float], # K - scaling: Dict[str, Dict[str, float]], - src_directory: str, - dst_directory: str, - dataset_name: str, - data_format: str, - surface_geopotential_mean: float = -597.7115478515625, - surface_geopotential_std: float = 55658.21484375, - R: float = 287, # J K^{-1} kg^{-1} - g0: float = 9.81, # m s^{-2} - topography_masking: bool = True, - ): - """ - Parameters - ---------- - weights: Sequence - list of floats that determine weighting of hydrostatic loss terms, assumed to be - in order of increasing pressure levels - """ - super().__init__() - self.data_loss = data_loss - self.loss_weights = torch.tensor(weights) - self.device = None - self.g0 = g0 - self.topography_masking = topography_masking - - # Get channel index to pressure level mapping - self.pressure_levels = sorted(hPa_levels) - self.z_pressure_levels = { - channels.index(f"z{int(pl)}"): pl - for i, pl in enumerate(self.pressure_levels) - } - self.T_pressure_levels = { - channels.index(f"t{int(pl)}"): pl - for i, pl in enumerate(self.pressure_levels) - } - self.q_pressure_levels = { - channels.index(f"q{int(pl)}"): pl - for i, pl in enumerate(self.pressure_levels) - if f"q{int(pl)}" in channels - } - # Get offset to map from q channels here to Tv channels - # Relies on all pressure levels below a threshold to have q channels - for i, pl in enumerate(self.pressure_levels): - if f"q{int(pl)}" in channels: - self.q_index_offset = i - break - - # Create mapping for new tensor that holds only the constraint variables - self.z_constraint_pressure_levels = { - i: pl for i, pl in enumerate(self.pressure_levels) - } - self.Tv_constraint_pressure_levels = { - len(self.z_constraint_pressure_levels) + i: pl - for i, pl in enumerate(self.pressure_levels) - } - - # Get scaling weights - self.z_mean = torch.Tensor( - [ - scaling[f"z{int(pl)}"]["mean"] - for i, pl in enumerate(self.pressure_levels) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - self.z_std = torch.Tensor( - [scaling[f"z{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] - ).reshape((1, 1, 1, -1, 1, 1)) - self.T_mean = torch.Tensor( - [ - scaling[f"t{int(pl)}"]["mean"] - for i, pl in enumerate(self.pressure_levels) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - self.T_std = torch.Tensor( - [scaling[f"t{int(pl)}"]["std"] for i, pl in enumerate(self.pressure_levels)] - ).reshape((1, 1, 1, -1, 1, 1)) - self.q_mean = torch.Tensor( - [ - scaling[f"q{int(pl)}"]["mean"] - for i, pl in enumerate(self.q_pressure_levels.values()) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - self.q_std = torch.Tensor( - [ - scaling[f"q{int(pl)}"]["std"] - for i, pl in enumerate(self.q_pressure_levels.values()) - ] - ).reshape((1, 1, 1, -1, 1, 1)) - - # Set per level alphas - if len(alpha) != len(hPa_levels) - 1: - raise AssertionError( - f"Incorrect number of alpha values. Expected len(hPa_levels)-1 [{len(hPa_levels) - 1}], got {len(alpha)}" - ) - self.alpha = torch.Tensor(alpha).reshape((1, 1, -1, 1, 1)) - - # Molecular weight ratio factor of water vapor to air - self.Mw_ratio = 28.97 / 18.016 - 1.0 # 0.6078 - - # Create the constraint - # TODO: remove anchor levels since it's not needed for the - # differential constraint - self.constraint = DifferentialHydrostaticBalanceConstraint( - self.z_constraint_pressure_levels, - self.Tv_constraint_pressure_levels, - 0, - len(self.z_pressure_levels), - R, - self.g0, - ) - - self.num_z_levels = len(self.z_constraint_pressure_levels) - self.num_Tv_levels = len(self.Tv_constraint_pressure_levels) - self.z_level_mapping = torch.tensor(list(self.z_pressure_levels.keys())) - self.T_level_mapping = torch.tensor(list(self.T_pressure_levels.keys())) - self.q_level_mapping = torch.tensor(list(self.q_pressure_levels.keys())) - - # Get topography information - ds = xr.open_zarr(f"{src_directory}{dataset_name}.zarr") - self.topography = ( - surface_geopotential_std * ds.constants.sel(channel_c="z").values - + surface_geopotential_mean - ) / self.g0 - - self.topography = torch.tensor( - self.topography[np.newaxis, :, np.newaxis, :, :], dtype=torch.float - ) - logger.info( - f"Min/Max topography (m): {self.topography.min()}/{self.topography.max()}" - ) - - def setup(self, trainer): - """ - pushes weights to cuda device - """ - - # Call setup for data loss first - self.data_loss.setup(trainer) - - if len(self.z_pressure_levels) - 1 != len(self.loss_weights): - raise ValueError( - "Length of loss_weights is not one less than number of pressure levels!" - ) - - self.loss_weights = self.loss_weights.to(device=trainer.device) - - # Move means and stds - self.z_mean = self.z_mean.to(device=trainer.device) - self.z_std = self.z_std.to(device=trainer.device) - self.T_mean = self.T_mean.to(device=trainer.device) - self.T_std = self.T_std.to(device=trainer.device) - self.q_mean = self.q_mean.to(device=trainer.device) - self.q_std = self.q_std.to(device=trainer.device) - - # Move alphas - self.alpha = self.alpha.to(device=trainer.device) - - # Move indexing arrays for CUDA graphs - self.z_level_mapping = self.z_level_mapping.to(device=trainer.device) - self.T_level_mapping = self.T_level_mapping.to(device=trainer.device) - self.q_level_mapping = self.q_level_mapping.to(device=trainer.device) - - # Move topography - self.topography = self.topography.to(device=trainer.device) - - def scale(self, x): - """ - Scale inputs to physical values and compute virtual temperature - Tensors are expected to be in the shape [N, B, F, C, H, W] - """ - # N, B, F, C, H, W = x.shape - N, F, B, C, H, W = x.shape - C_scaled = self.num_z_levels + self.num_Tv_levels - x_scaled = torch.zeros( - # (N, B, F, C_scaled, H, W), device=x.device, dtype=torch.float - (N, F, B, C_scaled, H, W), - device=x.device, - dtype=torch.float, - ) - # Get scaled geopotential heights - x_scaled[:, :, :, : self.num_z_levels, :, :] = ( - x[:, :, :, self.z_level_mapping, :, :] * self.z_std + self.z_mean - ) / self.g0 # divide by g0 for heights - # Get scaled temperatures - x_scaled[:, :, :, self.num_z_levels :, :, :] = ( - x[:, :, :, self.T_level_mapping, :, :] * self.T_std + self.T_mean - ) - # Add q correction to get virtual temperature for levels with non-zero q - x_scaled[ - :, - :, - :, - (self.num_z_levels + self.q_index_offset) :, - :, - :, - ] *= 1.0 + self.Mw_ratio * ( - x[:, :, :, self.q_level_mapping, :, :] * self.q_std + self.q_mean - ) - - # transpose B dim to before F - x_scaled = x_scaled.transpose(1, 2) - - # Combine N and B dimensions and return - return x_scaled.reshape((-1, F, C_scaled, H, W)) - - def error_histogram(self, prediction, bins, accumulator=None): - """Accumulate a per-level histogram of absolute hydrostatic-balance error. - - Parameters - ---------- - prediction: torch.Tensor - Model prediction tensor of shape [N, F, B, C, H, W]. - bins: int or torch.Tensor - Either the number of equal-width bins to use, or an existing - tensor of per-level bin edges to reuse. - accumulator: torch.Tensor, optional - Existing per-level histogram counts to add to. If None, a new - zero-initialized accumulator is created. - - Returns - ------- - tuple[torch.Tensor, torch.Tensor] - The updated ``(accumulator, bin_edges)``. - """ - N, F, B, C, H, W = tuple(prediction.shape) - - if not (prediction.ndim == 6): - raise AssertionError("Expected predictions to have 6 dimensions") - - # Scale to physical units and compute virtual temperature - x = self.scale(prediction) - Tv_avg, Tv_model_avg = self.constraint(x) - Tv_error = Tv_avg - Tv_model_avg - # Mask out error in regions below the surface - if self.topography_masking: - Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 - - vlevels = Tv_error.shape[2] - if accumulator is None: - if isinstance(bins, int): - accumulator = torch.zeros( - (vlevels, bins), dtype=torch.float32, device=prediction.device - ) - else: - accumulator = torch.zeros( - (vlevels, bins.shape[1] - 1), - dtype=torch.float32, - device=prediction.device, - ) - if isinstance(bins, int): - bin_edges = torch.zeros( - (vlevels, bins + 1), - dtype=prediction.dtype, - device=prediction.device, - ) - else: - bin_edges = bins - - for level_idx in range(vlevels): - hist, be = torch.histogram( - torch.absolute(Tv_error[:, :, level_idx, :, :]), - bins=bins if isinstance(bins, int) else bin_edges[level_idx, :], - ) - accumulator[level_idx, :] += hist - bin_edges[level_idx, :] = be - - return accumulator, bin_edges - - def forward(self, prediction, target, average_channels=True): - """ - Forward pass of LossWithHydrostasy - Tensors are expected to be in the shape [N, F, B, C, H, W] - - Parameters - ---------- - prediction: torch.Tensor - The prediction tensor - target: torch.Tensor - The target tensor - average_channels: bool, optional - whether the mean of the channels should be taken - """ - - # Need to scale back to physical units here so disable autocast - # and explicitly cast to float32 - with torch.cuda.amp.autocast(enabled=False): - prediction = prediction.float() - target = target.float() - - if not (prediction.ndim == 6 and target.ndim == 6): - raise AssertionError("Expected predictions to have 6 dimensions") - - # Scale to physical units and compute virtual temperature - x = self.scale(prediction) - Tv_avg, Tv_model_avg = self.constraint(x) - Tv_error = ((Tv_avg - Tv_model_avg) / self.alpha) ** 2 - - # Mask out error in regions below the surface - if self.topography_masking: - Tv_error[x[:, :, 1 : self.num_z_levels, :, :] < self.topography] = 0.0 - - # Compute the error tolerant loss - Tv_loss = self.loss_weights * ( - Tv_error / (1 + torch.exp(1 - Tv_error)) - ).mean(dim=(0, 1, 3, 4)) - - # Compute data loss - data_loss = self.data_loss( - prediction, target, average_channels=average_channels - ) - - if average_channels: - return data_loss + torch.mean(Tv_loss) - else: - return torch.concatenate((data_loss, Tv_loss)) diff --git a/test/datapipes/healpix/__init__.py b/test/datapipes/healpix/__init__.py deleted file mode 100644 index c1b40ccc80..0000000000 --- a/test/datapipes/healpix/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Test suite for the HEALPix datapipes.""" diff --git a/test/datapipes/healpix/conftest.py b/test/datapipes/healpix/conftest.py deleted file mode 100644 index 510dd80510..0000000000 --- a/test/datapipes/healpix/conftest.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Shared fixtures/helpers for the test_healpix*.py suite (classic/Zarr, -coupled/uncoupled). - -These describe the same on-disk NFS test dataset and the same -scaling/coupling configuration shapes across all four test_healpix*.py -modules in this package, so they're centralized here instead of being -duplicated (with minor, incidental drift) in each file. -""" - -from dataclasses import dataclass - -import pytest - - -@pytest.fixture -def data_dir(nfs_data_dir): - """Directory of the classic (non-Zarr) prebuilt test dataset.""" - return nfs_data_dir.joinpath("datasets/healpix") - - -@pytest.fixture -def dataset_name(): - return "healpix" - - -@pytest.fixture -def create_path(nfs_data_dir): - return nfs_data_dir.joinpath("datasets/healpix/merge") - - -@pytest.fixture -def dataset_path(nfs_data_dir): - """Path to the Zarr-backed version of the same test dataset.""" - return nfs_data_dir.joinpath("datasets/healpix/healpix.zarr") - - -@pytest.fixture -def splits(): - """Date ranges that fall within the small test dataset's actual time - range, suitable for exercising real dataloader construction (as opposed - to just checking that arbitrary split values are accepted). - """ - return { - "train_date_start": "1979-01-01", - "train_date_end": "1979-01-01T21:00", - "val_date_start": "1979-01-02", - "val_date_end": "1979-01-02T09:00", - "test_date_start": "1979-01-02T12:00", - "test_date_end": "1979-01-02T18:00", - } - - -@pytest.fixture -def scaling_dict(): - omegaconf = pytest.importorskip("omegaconf") - scaling = { - "t2m0": {"mean": 287.8665771484375, "std": 14.86227798461914}, - "t850": {"mean": 281.2710266113281, "std": 12.04991626739502}, - "tau300-700": {"mean": 61902.72265625, "std": 2559.8408203125}, - "tcwv0": {"mean": 24.034976959228516, "std": 16.411935806274414}, - "z1000": {"mean": 952.1435546875, "std": 895.7516479492188}, - "z1000-12h": {"mean": 952.1435546875, "std": 895.7516479492188}, - "z250": {"mean": 101186.28125, "std": 5551.77978515625}, - "z500": {"mean": 55625.9609375, "std": 2681.712890625}, - "lsm": {"mean": 0, "std": 1}, - "z": {"mean": 0, "std": 1}, - "tp6": {"mean": 1, "std": 0, "log_epsilon": 1e-6}, - "extra": {"mean": 0, "std": 0}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - -@pytest.fixture -def scaling_double_dict(): - omegaconf = pytest.importorskip("omegaconf") - scaling = { - "t2m0": {"mean": 0, "std": 2}, - "t850": {"mean": 0, "std": 2}, - "tau300-700": {"mean": 0, "std": 2}, - "tcwv0": {"mean": 0, "std": 2}, - "z1000": {"mean": 0, "std": 2}, - "z1000-12h": {"mean": 0, "std": 2}, - "z250": {"mean": 0, "std": 2}, - "z500": {"mean": 0, "std": 2}, - "tp6": {"mean": 0, "std": 2, "log_epsilon": 1e-6}, - "lsm": {"mean": 0, "std": 2}, - "z": {"mean": 0, "std": 2}, - "extra": {"mean": 0, "std": 2}, # doesn't appear in test dataset - } - return omegaconf.DictConfig(scaling) - - -@dataclass -class coupler_helper: - """Stand-in for a coupled module, exposing just what setup_coupling() - needs from it (output_variables, time_step).""" - - output_variables: list - time_step: str - - -@pytest.fixture -def constant_coupler_config(): - """A single-variable ConstantCoupler coupling config, in the list-of-dict - shape consumed by Coupled*Dataset(Zarr)/*DataModule(Zarr) `couplings=`.""" - return [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["0h"], - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] - - -@pytest.fixture -def average_coupler_config(): - """A single-variable TrailingAverageCoupler coupling config, in the same - shape as `constant_coupler_config` but for a different coupler class.""" - return [ - { - "coupler": "TrailingAverageCoupler", - "params": { - "batch_size": 1, - "variables": ["z250"], - "input_times": ["6h"], - "averaging_window": "6h", - "input_time_dim": 1, - "output_time_dim": 1, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] - - -def assert_shard_dataloaders(timeseries_dm): - """Shared assertions for `*DataModule(Zarr).{train,val,test}_dataloader`: - a single shard should get no distributed sampler, while multiple shards - should. Used identically across the classic/Zarr, coupled/uncoupled data - module tests. - """ - from torch.utils.data import DataLoader - from torch.utils.data.distributed import DistributedSampler - - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=1) - assert train_sampler is None - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=1) - assert val_sampler is None - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=1) - assert test_sampler is None - assert isinstance(test_dataloader, DataLoader) - - train_dataloader, train_sampler = timeseries_dm.train_dataloader(num_shards=2) - assert isinstance(train_sampler, DistributedSampler) - assert isinstance(train_dataloader, DataLoader) - - val_dataloader, val_sampler = timeseries_dm.val_dataloader(num_shards=2) - assert isinstance(val_sampler, DistributedSampler) - assert isinstance(val_dataloader, DataLoader) - - test_dataloader, test_sampler = timeseries_dm.test_dataloader(num_shards=2) - assert isinstance(test_sampler, DistributedSampler) - assert isinstance(test_dataloader, DataLoader) diff --git a/test/datapipes/healpix/test_healpix.py b/test/datapipes/healpix/test_healpix.py deleted file mode 100644 index 10d8a7f48a..0000000000 --- a/test/datapipes/healpix/test_healpix.py +++ /dev/null @@ -1,564 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import random -import warnings -from pathlib import Path - -import pytest - -from physicsnemo.distributed import DistributedManager -from test.conftest import requires_module -from test.datapipes.healpix.conftest import assert_shard_dataloaders - -omegaconf = pytest.importorskip("omegaconf") -np = pytest.importorskip("numpy") -xr = pytest.importorskip("xarray") - - -@requires_module("omegaconf") -def test_open_time_series(data_dir, dataset_name, pytestconfig): - # check for failure of non-existant dataset - from physicsnemo.datapipes.healpix.data_modules import ( - open_time_series_dataset_classic_prebuilt, - ) - - with pytest.raises(FileNotFoundError, match=("Dataset doesn't exist at")): - open_time_series_dataset_classic_prebuilt("/null_path", dataset_name) - - ds = open_time_series_dataset_classic_prebuilt(data_dir, dataset_name) - assert isinstance(ds, xr.Dataset) - ds.close() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -def test_TimeSeriesDataset_initialization( - data_dir, dataset_name, scaling_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.timeseries_dataset import TimeSeriesDataset - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # check for failure of timestep not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'time_step' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - data_time_step="2h", - time_step="5h", - scaling=scaling_dict, - ) - - # check for failure of gap not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'gap' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - data_time_step="2h", - time_step="6h", - gap="3h", - scaling=scaling_dict, - ) - - # check for failure of invalid scaling variable on input - invalid_scaling = omegaconf.DictConfig( - { - "bogosity": {"mean": 0, "std": 42}, - } - ) - with pytest.raises(KeyError, match=("Input channels ")): - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - data_time_step="3h", - time_step="6h", - scaling=invalid_scaling, - ) - - # check for failure when a channel_out variable has no scaling entry. - # dropping "z1000" from channel_in (but keeping it in channel_out) means - # len(channel_out) != len(channel_in), so input scaling is selected by - # channel_in (which succeeds), letting the target-scaling failure surface - # on its own - scaling_missing_target = omegaconf.DictConfig( - {k: v for k, v in scaling_dict.items() if k != "z1000"} - ) - zarr_ds_asymmetric = zarr_ds.sel( - channel_in=[v for v in zarr_ds.channel_in.values if v != "z1000"], - channel_out=zarr_ds.channel_out.values, - ) - with pytest.raises(KeyError, match=("Target channels ")): - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds_asymmetric, - scaling=scaling_missing_target, - ) - - # check for failure when a constant (channel_c) variable has no scaling - # entry, even though every channel_in/channel_out variable does - scaling_missing_constant = omegaconf.DictConfig( - {k: v for k, v in scaling_dict.items() if k != "z"} - ) - with pytest.raises(KeyError, match=("Constant channels ")): - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_missing_constant, - ) - - # check for warning on batch size > 1 and forecast mode - warnings.filterwarnings("error") - with pytest.raises( - UserWarning, - match=( - "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`" - ), - ): - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_dict, - batch_size=2, - forecast_init_times=zarr_ds.time[:2], - ) - - # test no scaling - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - ) - assert isinstance(timeseries_ds, TimeSeriesDataset) - - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_dict, - ) - assert isinstance(timeseries_ds, TimeSeriesDataset) - - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:2], - ) - assert isinstance(timeseries_ds, TimeSeriesDataset) - - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:2], - data_time_step="3h", - time_step="6h", - ) - assert isinstance(timeseries_ds, TimeSeriesDataset) - zarr_ds.close() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("numpy") -def test_TimeSeriesDataset_get_constants( - data_dir, dataset_name, scaling_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.timeseries_dataset import TimeSeriesDataset - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_dict, - ) - - # constants are reshaped - expected = np.transpose(zarr_ds.constants.values, axes=(1, 0, 2, 3)) - outvar = timeseries_ds.get_constants() - assert np.array_equal( - expected, - outvar, - ) - zarr_ds.close() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -def test_TimeSeriesDataset_len(data_dir, dataset_name, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.timeseries_dataset import TimeSeriesDataset - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # check forecast mode - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - ) - assert len(timeseries_ds) == init_times - - # check train mode - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - data_time_step="3h", - time_step="9h", - scaling=scaling_dict, - batch_size=2, - ) - # Window length of 3 for one sample size - assert len(timeseries_ds) == (len(zarr_ds.time.values) - 2) // 2 - - # check train mode - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - data_time_step="3h", - time_step="9h", - scaling=scaling_dict, - batch_size=2, - drop_last=True, - ) - assert len(timeseries_ds) == (len(zarr_ds.time.values) - 2) // 2 - zarr_ds.close() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("numpy") -def test_TimeSeriesDataset_get( - data_dir, dataset_name, scaling_double_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.timeseries_dataset import TimeSeriesDataset - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - batch_size = 2 - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_double_dict, - batch_size=batch_size, - ) - - # check for invalid index - invalid_idx = len(zarr_ds.targets) + 1 - with pytest.raises( - IndexError, match=(f"index {invalid_idx} out of range for dataset with length") - ): - inputs, targets = timeseries_ds[invalid_idx] - - inputs, targets = timeseries_ds[0] - - # make sure number of targets is correct - assert len(targets) == batch_size - - # check target data - # need to transpose - targets_expected = zarr_ds.targets[batch_size].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # check for negative index - inputs, targets = timeseries_ds[-1] - targets_expected = zarr_ds.targets[12].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - - # we're not dropping incomplete elements by default - assert len(targets) == 0 - - # this time dropping incomplete so that we get a full sample sample - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - ) - - inputs, targets = timeseries_ds[-1] - targets_expected = zarr_ds.targets[-1 - batch_size].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # With insolation we get 1 extra channel - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_insolation=True, - ) - assert (len(inputs)) + 1 == len(timeseries_ds[0][0]) - - # nothing should change with forecast mode other than getting just inputs - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_double_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - ) - inputs = timeseries_ds[0] - - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # insolation adds 1 extra channel - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds, - scaling=scaling_double_dict, - batch_size=1, - add_insolation=True, - forecast_init_times=zarr_ds.time[:init_times], - ) - assert (len(inputs)) + 1 == len(timeseries_ds[0]) - - # No constants in input data - init_times = random.randint(1, len(zarr_ds.time.values)) - zarr_ds_no_const = zarr_ds.drop_vars("constants") - timeseries_ds = TimeSeriesDataset( - dataset=zarr_ds_no_const, - scaling=scaling_double_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - ) - assert len(inputs) == (len(timeseries_ds[0]) + 1) - zarr_ds.close() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -def test_TimeSeriesDataModule_initialization( - data_dir, create_path, dataset_name, scaling_double_dict, splits, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules import ( - TimeSeriesDataModule, - ) - - variables = ["z500", "z1000"] - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # check for failure when a requested input variable isn't in the dataset - with pytest.raises(ValueError, match=("Input variables not found in dataset")): - TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables + ["DoesntExist"], - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - ) - - # check for failure when a requested output variable isn't in the dataset - with pytest.raises(ValueError, match=("Output variables not found in dataset")): - TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - output_variables=variables + ["DoesntExist"], - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - ) - - # check for failure when a requested constant isn't in the dataset - with pytest.raises(ValueError, match=("Constants not found in dataset")): - TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants={"lsm": "DoesntExist"}, - ) - - # use the prebuilt dataset - # Internally initializes DistributedManager - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - ) - assert isinstance(timeseries_dm, TimeSeriesDataModule) - - # `prebuilt_dataset` is kept only for backwards compatibility; on-the-fly - # dataset generation has been removed, so `setup()` always opens an - # existing prebuilt dataset from `dst_directory` regardless of this flag - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=False, - scaling=scaling_double_dict, - ) - assert isinstance(timeseries_dm, TimeSeriesDataModule) - - # with init times - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - forecast_init_times=zarr_ds.time[:2], - ) - assert isinstance(timeseries_dm, TimeSeriesDataModule) - - # with splits - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - splits=omegaconf.DictConfig(splits), - ) - assert isinstance(timeseries_dm, TimeSeriesDataModule) - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("numpy") -def test_TimeSeriesDataModule_get_constants( - data_dir, create_path, dataset_name, scaling_double_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules import ( - TimeSeriesDataModule, - ) - - variables = ["z500", "z1000"] - constants = {"lsm": "lsm"} - - # No constants - # Internally initializes DistributedManager - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants=None, - ) - - assert timeseries_dm.get_constants() is None - - # just lsm as constant - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants=constants, - ) - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # dividing by 2 due to scaling - expected = ( - np.transpose( - zarr_ds.constants.sel(channel_c=list(constants.keys())).values, - axes=(1, 0, 2, 3), - ) - / 2.0 - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - - # with splits we're doing forecasting and get - # constants from train instead of test dataset - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants=constants, - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -def test_TimeSeriesDataModule_get_dataloaders( - data_dir, create_path, dataset_name, scaling_double_dict, splits, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules import ( - TimeSeriesDataModule, - ) - - variables = ["z500", "z1000"] - - # use the prebuilt dataset - # Internally initializes DistributedManager - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - splits=splits, - shuffle=False, - ) - - assert_shard_dataloaders(timeseries_dm) - DistributedManager.cleanup() diff --git a/test/datapipes/healpix/test_healpix_couple.py b/test/datapipes/healpix/test_healpix_couple.py deleted file mode 100644 index c56990266c..0000000000 --- a/test/datapipes/healpix/test_healpix_couple.py +++ /dev/null @@ -1,1319 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import random -import warnings -from pathlib import Path - -import pytest -import torch as th - -from physicsnemo.distributed import DistributedManager -from test.conftest import requires_module -from test.datapipes.healpix.conftest import assert_shard_dataloaders, coupler_helper - -omegaconf = pytest.importorskip("omegaconf") -np = pytest.importorskip("numpy") -pd = pytest.importorskip("pandas") -xr = pytest.importorskip("xarray") - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_ConstantCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.couplers import ( - ConstantCoupler, - ) - - variables = ["z500", "z1000"] - input_times = ["0h"] - input_time_dim = 1 - output_time_dim = 1 - presteps = 0 - batch_size = 2 - batch = {"time": slice(0, 2)} - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # test fail initialization - with pytest.raises( - NotImplementedError, match=("Data preparation not yet implemented") - ): - coupler = ConstantCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - prepared_coupled_data=False, - ) - - coupler = ConstantCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - assert isinstance(coupler, ConstantCoupler) - - # check setting coupled variable indices - mock_coupled_module = coupler_helper( - output_variables=["not_coupled", "z500"], - time_step="0h", - ) - with pytest.raises(ValueError, match=("Missing variables in coupled module")): - coupler.setup_coupling(mock_coupled_module) - - mock_coupled_module.output_variables = ["z500", "z1000"] - coupler.setup_coupling(mock_coupled_module) - assert coupler.coupled_channel_indices == [0, 1] - - interval = 2 - data_time_step = "3h" - coupler.compute_coupled_indices(interval, data_time_step) - coupled_integration_dim = presteps + max(output_time_dim // input_time_dim, 1) - expected = np.empty([batch_size, coupled_integration_dim, len(input_times)]) - for b in range(batch_size): - for i in range(coupled_integration_dim): - expected[b, i, :] = b + np.array( - [pd.Timedelta(ts) / pd.Timedelta(data_time_step) for ts in input_times] - ) - expected = expected.astype(int) - assert np.array_equal(expected, coupler._coupled_offsets) - - scaling_df = pd.DataFrame.from_dict(omegaconf.OmegaConf.to_object(scaling_dict)).T - scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} - scaling_da = scaling_df.to_xarray().astype("float32") - coupler.set_scaling(scaling_da) - coupled_scaling = scaling_da.sel(index=variables).rename({"index": "channel_in"}) - expected = np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["mean"]) - expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["std"]) - - # `set_coupled_fields` derives its batch dimension from the shape of - # `coupled_fields` itself rather than validating it against the - # configured `batch_size` (an older physicsnemo implementation raised a - # ValueError here, but the ported DLESyM coupler intentionally adapts to - # whatever batch size is provided) - coupler.coupled_channel_indices = [0, 1] - coupled_fields_batch_size = batch_size * 2 - coupled_fields_timedim = 2 - coupled_fields = th.rand( - coupled_fields_batch_size, - coupler.spatial_dims[0], - coupled_fields_timedim, - len(coupler.coupled_channel_indices), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - coupler.set_coupled_fields(coupled_fields) - assert coupler.preset_coupled_fields.shape[1] == coupled_fields_batch_size - - coupled_fields_batch_size = batch_size - coupled_fields_timedim = 4 - expected_shape = [ - coupler.coupled_integration_dim, - coupled_fields_batch_size, - coupler.timevar_dim, - ] + list(coupler.spatial_dims) - coupled_fields = th.rand( - coupled_fields_batch_size, - coupler.spatial_dims[0], - coupled_fields_timedim, - len(coupler.coupled_channel_indices), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - coupler.set_coupled_fields(coupled_fields) - assert coupler.coupled_mode - assert list(coupler.construct_integrated_couplings().shape) == expected_shape - - # verify that the data is being properly transformed: set_coupled_fields - # broadcasts the first time step of coupled_fields across the integration - # dim and permutes to [integration_dim, batch, timevar, face, height, width] - expected = coupled_fields[:, :, :, coupler.coupled_channel_indices, :, :] - expected = expected[:, :, :1, :, :, :] - expected = expected.permute(2, 0, 3, 1, 4, 5) - assert th.equal(expected, coupler.construct_integrated_couplings()) - - # test coupler reset - coupler.reset_coupler() - assert coupler.coupled_mode is False - - # once reset, batch/bsize are required to construct the couplings from - # the dataset rather than from preset (cached) fields - with pytest.raises( - ValueError, match=("batch and bsize must be provided when not in coupled_mode") - ): - coupler.construct_integrated_couplings() - - # test loading from the dataset - coupled_scaling = { - "mean": np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)), - "std": np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)), - } - expected = zarr_ds.sel(channel_in=variables).inputs[:2].values - expected = (expected - coupled_scaling["mean"]) / coupled_scaling["std"] - coupled_field = coupler.construct_integrated_couplings( - batch=batch, bsize=batch_size - ) - assert np.array_equal(expected, coupled_field[0]) - - # set_scaling raises if a coupled variable has no corresponding entry in - # the provided scaling values - scaling_missing_var = scaling_da.drop_sel(index="z1000") - with pytest.raises( - KeyError, match=("Coupled variable\\(s\\) not found in scaling values") - ): - coupler.set_scaling(scaling_missing_var) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_TrailingAverageCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.couplers import ( - TrailingAverageCoupler, - ) - - variables = ["z500", "z1000"] - input_times = ["6h", "12h"] - input_time_dim = 2 - output_time_dim = 2 - presteps = 0 - batch_size = 2 - averaging_window = "6h" - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # test fail initialization when trying to prepare data - with pytest.raises( - NotImplementedError, match=("Data preparation not yet implemented") - ): - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - averaging_window=averaging_window, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - prepared_coupled_data=False, - ) - - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - averaging_window=averaging_window, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - assert isinstance(coupler, TrailingAverageCoupler) - - mock_coupled_module = coupler_helper( - output_variables=["not_coupled", "z500"], - time_step="3h", - ) - with pytest.raises(ValueError, match=("Missing variables in coupled module")): - coupler.setup_coupling(mock_coupled_module) - - # veryify averaging slices computed correctly - mock_coupled_module = coupler_helper( - output_variables=["z500", "z1000"], - time_step="3h", - ) - coupler.setup_coupling(mock_coupled_module) - averaging_window_max_indices = [ - i // pd.Timedelta(mock_coupled_module.time_step) for i in input_times - ] - dt = averaging_window_max_indices[0] - # assumes only 1 integration step, otherwise would be wrong - expected_slices = [[]] - for i, window_end in enumerate(averaging_window_max_indices): - expected_slices[0].append(slice(i * dt, window_end)) - assert expected_slices == coupler.averaging_slices - - interval = 2 - data_time_step = "3h" - coupler.compute_coupled_indices(interval, data_time_step) - coupled_integration_dim = presteps + max(output_time_dim // input_time_dim, 1) - expected = np.empty([batch_size, coupled_integration_dim, len(input_times)]) - for b in range(batch_size): - for i in range(coupled_integration_dim): - expected[b, i, :] = ( - b - + (input_time_dim * i + 1) * interval - + np.array( - [ - pd.Timedelta(ts) / pd.Timedelta(data_time_step) - for ts in input_times - ] - ) - ) - expected = expected.astype(int) - assert np.array_equal(expected, coupler._coupled_offsets) - - scaling_df = pd.DataFrame.from_dict(omegaconf.OmegaConf.to_object(scaling_dict)).T - scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} - scaling_da = scaling_df.to_xarray().astype("float32") - coupler.set_scaling(scaling_da) - coupled_scaling = scaling_da.sel(index=variables).rename({"index": "channel_in"}) - expected = np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["mean"]) - expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["std"]) - - averaging_window_max_indices = [ - i // pd.Timedelta(data_time_step) for i in coupler.input_times - ] - di = averaging_window_max_indices[0] - averaging_slices = [] - for j in range(coupler.coupled_integration_dim): - averaging_slices.append([]) - for i, r in enumerate(averaging_window_max_indices): - averaging_slices[j].append( - slice( - coupler.input_time_dim * j * di + i * di, - coupler.input_time_dim * j * di + r, - ) - ) - coupler.averaging_slices = averaging_slices - coupler.coupled_channel_indices = [0, 1] - - # `set_coupled_fields` derives its batch dimension from the shape of - # `coupled_fields` itself rather than validating it against the - # configured `batch_size` (an older physicsnemo implementation raised a - # ValueError here, but the ported DLESyM coupler intentionally adapts to - # whatever batch size is provided) - coupled_fields_batch_size = batch_size * 2 - coupled_fields_timedim = 4 - coupled_fields = th.rand( - coupled_fields_batch_size, - coupler.spatial_dims[0], - coupled_fields_timedim, - len(coupler.coupled_channel_indices), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - coupler.set_coupled_fields(coupled_fields) - assert coupler.preset_coupled_fields.shape[1] == coupled_fields_batch_size - - coupled_fields_batch_size = batch_size - coupled_fields_timedim = 4 - expected_shape = [ - coupler.coupled_integration_dim, - coupled_fields_batch_size, - coupler.timevar_dim, - ] + list(coupler.spatial_dims) - coupled_fields = th.rand( - coupled_fields_batch_size, - coupler.spatial_dims[0], - coupled_fields_timedim, - len(coupler.coupled_channel_indices), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - coupler.set_coupled_fields(coupled_fields) - assert list(coupler.preset_coupled_fields.shape) == expected_shape - - # check reset - assert coupler.coupled_mode - coupler.reset_coupler() - assert coupler.coupled_mode is False - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_TrailingAverageCoupler_multiple_variables(data_dir, dataset_name): - """Regression test to ensure `set_coupled_fields` correctly averages and - threads through every coupled variable, not just the first. Gives each - coupled variable a distinct, constant value in the synthetic - `coupled_fields` tensor so that averaging over any time window still - yields that same constant, letting each variable's contribution to - `preset_coupled_fields` be checked in isolation and by position. - """ - from physicsnemo.datapipes.healpix.couplers import TrailingAverageCoupler - - variables = ["z500", "z1000", "z250"] - input_times = ["6h", "12h"] - input_time_dim = 2 - output_time_dim = 2 - batch_size = 2 - averaging_window = "6h" - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=0, - averaging_window=averaging_window, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - coupler.coupled_channel_indices = list(range(len(variables))) - - data_time_step = "3h" - averaging_window_max_indices = [ - i // pd.Timedelta(data_time_step) for i in input_times - ] - di = averaging_window_max_indices[0] - averaging_slices = [] - for j in range(coupler.coupled_integration_dim): - averaging_slices.append([]) - for i, r in enumerate(averaging_window_max_indices): - averaging_slices[j].append( - slice( - coupler.input_time_dim * j * di + i * di, - coupler.input_time_dim * j * di + r, - ) - ) - coupler.averaging_slices = averaging_slices - - # give each coupled variable a distinct, constant value across every - # other dimension so the averaging window doesn't change its value, and - # each variable's contribution can be checked independently of the - # others - channel_values = [100.0, 200.0, 300.0] - coupled_fields = th.empty( - batch_size, - coupler.spatial_dims[0], - 4, - len(variables), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - for i, value in enumerate(channel_values): - coupled_fields[:, :, :, i, :, :] = value - - coupler.set_coupled_fields(coupled_fields) - result = coupler.construct_integrated_couplings() - assert list(result.shape) == [ - coupler.coupled_integration_dim, - batch_size, - coupler.timevar_dim, - ] + list(coupler.spatial_dims) - - # timevar ordering groups by averaging period first, then by variable - # (matching the dataset-loading path in - # `_construct_integrated_couplings_from_dataset`/`construct_integrated_couplings`) - for period in range(len(input_times)): - for var_idx, value in enumerate(channel_values): - timevar_idx = period * len(variables) + var_idx - slice_result = result[:, :, timevar_idx, :, :, :] - assert th.allclose(slice_result, th.full_like(slice_result, value)), ( - f"coupled variable {variables[var_idx]!r} (value {value}) was not " - f"preserved at averaging period {period}, timevar index {timevar_idx}" - ) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_ConstantCoupler_multiple_variables_order(data_dir, dataset_name): - """Sanity check that the xarray-backed coupling path preserves the order - of `variables` as requested (rather than the dataset's native channel_in - order) when 3+ coupled variables are requested out of native order. This - mirrors the invariant that the Zarr-backed path also has to uphold (see - the analogous, previously-failing test in test_healpix_couple_zarr.py). - """ - from physicsnemo.datapipes.healpix.couplers import ConstantCoupler - - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - native_order = list(zarr_ds.channel_in.values) - variables = ["z250", "z1000", "z500"] - assert [native_order.index(v) for v in variables] == sorted( - [native_order.index(v) for v in variables], reverse=True - ) - - batch_size = 2 - batch = {"time": slice(0, 2)} - coupler = ConstantCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - input_times=["0h"], - input_time_dim=1, - output_time_dim=1, - ) - coupler.compute_coupled_indices(interval=2, data_time_step="3h") - - coupled_field = coupler.construct_integrated_couplings( - batch=batch, bsize=batch_size - ) - for i, var in enumerate(variables): - expected_var = zarr_ds.sel(channel_in=var).inputs[:2].values - assert np.array_equal(expected_var, coupled_field[0][:, i]) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataset_initialization( - data_dir, dataset_name, scaling_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( - CoupledTimeSeriesDataset, - ) - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - variables = ["z500", "z1000"] - - # check for failure of timestep not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'time_step' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - data_time_step="2h", - time_step="5h", - scaling=scaling_dict, - ) - - # check for failure of gap not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'gap' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - data_time_step="2h", - time_step="6h", - gap="3h", - scaling=scaling_dict, - ) - - # check for failure of invalid scaling variable on input - invalid_scaling = omegaconf.DictConfig( - { - "bogosity": {"mean": 0, "std": 42}, - } - ) - with pytest.raises(KeyError, match=("Input channels ")): - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - data_time_step="3h", - time_step="6h", - scaling=invalid_scaling, - ) - - # check for warning on batch size > 1 and forecast mode - warnings.filterwarnings("error") - with pytest.raises( - UserWarning, - match=( - "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`" - ), - ): - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - batch_size=2, - forecast_init_times=zarr_ds.time[:2], - ) - - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - add_train_noise=True, - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDataset) - - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDataset) - - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:2], - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDataset) - - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:2], - data_time_step="3h", - time_step="6h", - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDataset) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataset_get_constants( - data_dir, dataset_name, scaling_dict, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( - CoupledTimeSeriesDataset, - ) - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - variables = ["z500", "z1000"] - - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - couplings=constant_coupler_config, - ) - - # constants are reshaped - expected = np.transpose(zarr_ds.constants.values, axes=(1, 0, 2, 3)) - outvar = timeseries_ds.get_constants() - assert np.array_equal( - expected, - outvar, - ) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataset_len( - data_dir, dataset_name, scaling_dict, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( - CoupledTimeSeriesDataset, - ) - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - variables = ["z500", "z1000"] - - # check forecast mode - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler_config, - ) - assert len(timeseries_ds) == init_times - - constant_coupler = constant_coupler_config.copy() - constant_coupler[0]["params"]["batch_size"] = 2 - - # check train mode - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - data_time_step="3h", - time_step="9h", - scaling=scaling_dict, - batch_size=2, - couplings=constant_coupler, - ) - # Window length of 3 for one sample size - assert len(timeseries_ds) == (len(zarr_ds.time.values) - 2) // 2 - - # check train mode - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - data_time_step="3h", - time_step="9h", - scaling=scaling_dict, - batch_size=2, - drop_last=True, - couplings=constant_coupler, - ) - assert len(timeseries_ds) == (len(zarr_ds.time.values) - 2) // 2 - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataset_get( - data_dir, dataset_name, scaling_double_dict, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( - CoupledTimeSeriesDataset, - ) - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - variables = list(zarr_ds.channel_out.to_numpy()) - # "z250" is coupled in below rather than being a regular input variable: - # channel_in scaling/slicing assumes channel_in == input_variables + - # coupled variables (with the coupled ones appended last), so a coupled - # variable can't also be one of input_variables - coupled_input_variables = [v for v in variables if v != "z250"] - - batch_size = 2 - constant_coupler = constant_coupler_config.copy() - constant_coupler[0]["params"]["batch_size"] = batch_size - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=batch_size, - couplings=constant_coupler, - ) - - # check for invalid index - invalid_idx = len(zarr_ds.targets) + 1 - with pytest.raises( - IndexError, match=(f"index {invalid_idx} out of range for dataset with length") - ): - inputs, targets = timeseries_ds[invalid_idx] - - inputs, targets = timeseries_ds[0] - - # make sure number of targets is correct - assert len(targets) == batch_size - - # check target data - # need to transpose - targets_expected = zarr_ds.targets[batch_size].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # check for negative index - inputs, targets = timeseries_ds[-1] - targets_expected = zarr_ds.targets[12].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - - # we're not dropping incomplete elements by default - assert len(targets) == 0 - - # this time dropping incomplete so that we get a full sample sample - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - couplings=constant_coupler, - ) - - inputs, targets = timeseries_ds[-1] - targets_expected = zarr_ds.targets[-1 - batch_size].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # without couplings - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - couplings=[], - ) - non_perturbed_inputs = timeseries_ds - assert len(non_perturbed_inputs[0][0]) == 2 # just inputs and targets - - # wihtout couplings but with noise - noise_params = { - "inputs": scaling_double_dict, - "couplings": scaling_double_dict, - } - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_train_noise=True, - train_noise_params=noise_params, - couplings=[], - ) - perturbed_inputs = timeseries_ds - # The first input will be the same sample, with perturbation it should have - # different values - assert non_perturbed_inputs[0][0][0].shape == perturbed_inputs[0][0][0].shape - assert not np.array_equal(non_perturbed_inputs[0][0][0], perturbed_inputs[0][0][0]) - - # noise is also added to the integrated couplings themselves, not just - # the regular inputs, when couplings are present. `inputs_result` here is - # [inputs, constants, integrated_couplings] because the dataset carries a - # "constants" data variable which the base class picks up unconditionally - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - couplings=constant_coupler, - ) - non_perturbed_sample = timeseries_ds[0][0] - assert len(non_perturbed_sample) == 3 - non_perturbed_couplings = non_perturbed_sample[2] - - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_train_noise=True, - train_noise_params=noise_params, - couplings=constant_coupler, - ) - perturbed_couplings = timeseries_ds[0][0][2] - assert non_perturbed_couplings.shape == perturbed_couplings.shape - assert not np.array_equal(non_perturbed_couplings, perturbed_couplings) - - # With insolation we get 1 extra channel - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_insolation=True, - couplings=constant_coupler, - ) - assert (len(inputs)) + 1 == len(timeseries_ds[0][0]) - - # nothing should change with forecast mode other than getting just inputs - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler, - ) - inputs = timeseries_ds[0] - - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # insolation adds 1 extra channel - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=1, - add_insolation=True, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler, - ) - assert (len(inputs)) + 1 == len(timeseries_ds[0]) - - # No constants in input data - init_times = random.randint(1, len(zarr_ds.time.values)) - zarr_ds_no_const = zarr_ds.drop_vars("constants") - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds_no_const, - input_variables=coupled_input_variables, - output_variables=variables, - scaling=scaling_double_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler, - ) - assert len(inputs) == (len(timeseries_ds[0]) + 1) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataModule_initialization( - data_dir, - create_path, - dataset_name, - scaling_double_dict, - splits, - constant_coupler_config, - pytestconfig, -): - from physicsnemo.datapipes.healpix.data_modules import ( - CoupledTimeSeriesDataModule, - ) - - variables = ["z500", "z1000"] - constant_coupler = constant_coupler_config - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # check for failure when a requested input variable isn't in the dataset - with pytest.raises(ValueError, match=("Input variables not found in dataset")): - CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables + ["DoesntExist"], - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - couplings=constant_coupler, - ) - - # check for failure when a requested output variable isn't in the dataset - with pytest.raises(ValueError, match=("Output variables not found in dataset")): - CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - output_variables=variables + ["DoesntExist"], - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - couplings=constant_coupler, - ) - - # check for failure when a requested constant isn't in the dataset - with pytest.raises(ValueError, match=("Constants not found in dataset")): - CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants={"lsm": "DoesntExist"}, - couplings=constant_coupler, - ) - - # use the prebuilt dataset - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - couplings=constant_coupler, - ) - assert isinstance(timeseries_dm, CoupledTimeSeriesDataModule) - - # `prebuilt_dataset` is kept only for backwards compatibility; on-the-fly - # dataset generation has been removed, so `setup()` always opens an - # existing prebuilt dataset from `dst_directory` regardless of this flag - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=False, - scaling=scaling_double_dict, - couplings=constant_coupler, - ) - assert isinstance(timeseries_dm, CoupledTimeSeriesDataModule) - - # with init times - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - forecast_init_times=zarr_ds.time[:2], - couplings=constant_coupler, - ) - assert isinstance(timeseries_dm, CoupledTimeSeriesDataModule) - - # with splits - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - splits=omegaconf.DictConfig(splits), - couplings=constant_coupler, - ) - assert isinstance(timeseries_dm, CoupledTimeSeriesDataModule) - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataModule_get_constants( - data_dir, - create_path, - dataset_name, - scaling_double_dict, - constant_coupler_config, - pytestconfig, -): - from physicsnemo.datapipes.healpix.data_modules import ( - CoupledTimeSeriesDataModule, - ) - - variables = ["z500", "z1000"] - constants = {"lsm": "lsm"} - constant_coupler = constant_coupler_config - - # No constants - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants=None, - couplings=constant_coupler, - ) - - assert timeseries_dm.get_constants() is None - - # just lsm as constant - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants=constants, - couplings=constant_coupler, - ) - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - zarr_ds = xr.open_zarr(ds_path) - - # divide by 2 due to scaling - expected = ( - np.transpose( - zarr_ds.constants.sel(channel_c=list(constants.keys())).values, - axes=(1, 0, 2, 3), - ) - / 2.0 - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - - # with splits we're doing forecasting and get - # constants from train instead of test dataset - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - constants=constants, - couplings=constant_coupler, - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -def test_CoupledTimeSeriesDataModule_get_dataloaders( - data_dir, - create_path, - dataset_name, - scaling_double_dict, - splits, - constant_coupler_config, - pytestconfig, -): - from physicsnemo.datapipes.healpix.data_modules import ( - CoupledTimeSeriesDataModule, - ) - - variables = ["z500", "z1000"] - - # use the prebuilt dataset - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - splits=splits, - shuffle=False, - couplings=constant_coupler_config, - ) - - assert_shard_dataloaders(timeseries_dm) - DistributedManager.cleanup() - - -@requires_module("omegaconf") -def test_CoupledTimeSeriesDataModule_get_coupled_vars( - data_dir, - create_path, - dataset_name, - scaling_double_dict, - constant_coupler_config, - average_coupler_config, - pytestconfig, -): - from physicsnemo.datapipes.healpix.data_modules import ( - CoupledTimeSeriesDataModule, - ) - - variables = ["z500", "z1000"] - - # Constant coupler - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - couplings=constant_coupler_config, - ) - - outvar = timeseries_dm._get_coupled_vars() - outvar.sort() - expected = ["z250"] - expected.sort() - - assert expected == outvar - - # Average coupler - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=data_dir, - dataset_name=dataset_name, - input_variables=variables, - batch_size=1, - prebuilt_dataset=True, - scaling=scaling_double_dict, - couplings=average_coupler_config, - ) - outvar = timeseries_dm._get_coupled_vars() - outvar.sort() - - assert expected == outvar - - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataset_next_integration( - data_dir, dataset_name, scaling_dict, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( - CoupledTimeSeriesDataset, - ) - - spatial_dims = [12, 32, 32] - input_variables = ["z500", "z1000"] - # length must match len(coupled_variables) (i.e. the coupler's timevar_dim); - # values are arbitrary since this synthetic setup bypasses setup_coupling() - coupled_channel_indices = [0] - coupled_variables = ["z250"] - num_variables = len(input_variables) - input_time_dim = 1 - output_time_dim = 1 - batch_size = 1 - constant_coupler = constant_coupler_config - - # open our test dataset - ds_path = Path(data_dir, dataset_name + ".zarr") - ds = xr.open_zarr(ds_path) - init_times = random.randint(1, len(ds.time.values)) - # channels need to be subselected before being handed over - test_ds = ds.sel( - channel_in=input_variables + coupled_variables, - channel_out=input_variables, - ) - - timeseries_ds = CoupledTimeSeriesDataset( - dataset=test_ds, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=batch_size, - couplings=constant_coupler, - data_time_step="6h", - time_step="6h", - drop_last=True, - add_insolation=True, - forecast_init_times=test_ds.time[:init_times], - ) - - test_model_outputs = th.rand( - 1, - spatial_dims[0], - output_time_dim, - num_variables, - spatial_dims[1], - spatial_dims[2], - ) - constants = np.transpose(ds.constants.values, axes=(1, 0, 2, 3)) - coupled_fields = th.rand( - batch_size, - spatial_dims[0], - input_time_dim + output_time_dim, - len(input_variables), - spatial_dims[1], - spatial_dims[2], - ) - - # mirrors set_coupled_fields: broadcast the first time step across the - # integration dim and permute to [integration_dim, batch, timevar, face, height, width] - expected_coupling = coupled_fields[:, :, :, coupled_channel_indices, :, :] - expected_coupling = expected_coupling[:, :, :1, :, :, :] - expected_coupling = expected_coupling.permute(2, 0, 3, 1, 4, 5) - - # need to grab at least 1 sample to properly intialize everything - timeseries_ds[0] - # hacky way to setup the indices since we don't actually have any coupled fields - timeseries_ds.couplings[0].coupled_channel_indices = coupled_channel_indices - - # set the coupled fields - timeseries_ds.couplings[0].set_coupled_fields(coupled_fields) - test_integration = timeseries_ds.next_integration(test_model_outputs, constants) - # test to make sure prognostics are used, constants stay the same, and couplings - # are what we set - assert np.array_equal(test_integration[0], test_model_outputs[:, :, -1:]) - assert np.array_equal(test_integration[2], constants) - assert np.array_equal(test_integration[3], expected_coupling) - - # I have absolutely no idea why a coupled dataset has the option for 0 couplings - # with no coupler, channel_in must exactly match input_variables (no extra - # coupled-variable channel appended), so re-select the dataset without - # the "z250" channel that `test_ds` carries for the coupled case above - test_ds_no_couplings = ds.sel( - channel_in=input_variables, - channel_out=input_variables, - ) - timeseries_ds = CoupledTimeSeriesDataset( - dataset=test_ds_no_couplings, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=batch_size, - couplings=[], - data_time_step="6h", - time_step="6h", - drop_last=True, - add_insolation=True, - forecast_init_times=test_ds_no_couplings.time[:init_times], - ) - # need to grab at least 1 sample to properly intialize everything - timeseries_ds[0] - test_integration = timeseries_ds.next_integration(test_model_outputs, constants) - assert np.array_equal(test_integration[0], test_model_outputs[:, :, -1:]) - assert np.array_equal(test_integration[2], constants) diff --git a/test/datapipes/healpix/test_healpix_couple_zarr.py b/test/datapipes/healpix/test_healpix_couple_zarr.py deleted file mode 100644 index d596e7cd40..0000000000 --- a/test/datapipes/healpix/test_healpix_couple_zarr.py +++ /dev/null @@ -1,1288 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import random -import warnings - -import pytest -import torch as th - -from physicsnemo.distributed import DistributedManager -from test.conftest import requires_module -from test.datapipes.healpix.conftest import assert_shard_dataloaders, coupler_helper - -omegaconf = pytest.importorskip("omegaconf") -np = pytest.importorskip("numpy") -pd = pytest.importorskip("pandas") -xr = pytest.importorskip("xarray") -zarr = pytest.importorskip("zarr") - - -def test_base_coupler_invalid_dataset_type(): - """Couplers dispatch on the runtime type of `dataset` to decide whether to use - the Zarr-native code path (`zarr.Group`) or the xarray-based one (`xr.Dataset`). - This dispatch logic is new (added to support the Zarr datapipes) and doesn't - require the NFS test dataset, so it can run unconditionally. - """ - from physicsnemo.datapipes.healpix.couplers import ConstantCoupler - - fake_dataset = {"inputs": np.zeros((4, 12, 1, 4, 4))} - - with pytest.raises( - TypeError, - match=("Coupler only supports xarray Datasets or zarr Groups"), - ): - ConstantCoupler( - dataset=fake_dataset, - batch_size=1, - variables=["z500"], - ) - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_ConstantCoupler(dataset_path, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.couplers import ( - ConstantCoupler, - ) - - variables = ["z500", "z1000"] - input_times = ["0h"] - input_time_dim = 1 - output_time_dim = 1 - presteps = 0 - batch_size = 2 - batch = {"time": slice(0, 2)} - - # open our test dataset - zarr_ds = zarr.open(dataset_path) - input_indices = [ - int(np.where(zarr_ds["channel_in"][:] == ch)[0][0]) for ch in variables - ] - - # test fail initialization - with pytest.raises( - NotImplementedError, match=("Data preparation not yet implemented") - ): - coupler = ConstantCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - prepared_coupled_data=False, - ) - - coupler = ConstantCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - assert isinstance(coupler, ConstantCoupler) - - # check setting coupled variable indices - mock_coupled_module = coupler_helper( - output_variables=["not_coupled", "z500"], - time_step="0h", - ) - with pytest.raises(ValueError, match=("Missing variables in coupled module")): - coupler.setup_coupling(mock_coupled_module) - - mock_coupled_module.output_variables = ["z500", "z1000"] - coupler.setup_coupling(mock_coupled_module) - assert coupler.coupled_channel_indices == [0, 1] - - interval = 2 - data_time_step = "3h" - coupler.compute_coupled_indices(interval, data_time_step) - coupled_integration_dim = presteps + max(output_time_dim // input_time_dim, 1) - expected = np.empty([batch_size, coupled_integration_dim, len(input_times)]) - for b in range(batch_size): - for i in range(coupled_integration_dim): - expected[b, i, :] = b + np.array( - [pd.Timedelta(ts) / pd.Timedelta(data_time_step) for ts in input_times] - ) - expected = expected.astype(int) - assert np.array_equal(expected, coupler._coupled_offsets) - - scaling_df = pd.DataFrame.from_dict(omegaconf.OmegaConf.to_object(scaling_dict)).T - scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} - scaling_da = scaling_df.to_xarray().astype("float32") - coupler.set_scaling(scaling_da) - coupled_scaling = scaling_da.sel(index=variables).rename({"index": "channel_in"}) - expected = np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["mean"]) - expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["std"]) - - coupled_fields_batch_size = batch_size - coupled_fields_timedim = 4 - expected_shape = [ - coupler.coupled_integration_dim, - coupled_fields_batch_size, - coupler.timevar_dim, - ] + list(coupler.spatial_dims) - coupled_fields = th.rand( - coupled_fields_batch_size, - coupler.spatial_dims[0], - coupled_fields_timedim, - len(coupler.coupled_channel_indices), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - coupler.set_coupled_fields(coupled_fields) - assert coupler.coupled_mode - assert list(coupler.construct_integrated_couplings().shape) == expected_shape - - # verify that the data is being properly transformed: set_coupled_fields - # broadcasts the first time step of coupled_fields across the integration - # dim and permutes to [integration_dim, batch, timevar, face, height, width] - expected = coupled_fields[:, :, :, coupler.coupled_channel_indices, :, :] - expected = expected[:, :, :1, :, :, :] - expected = expected.permute(2, 0, 3, 1, 4, 5) - result = coupler.construct_integrated_couplings() - assert th.equal(expected, result) - - # test coupler reset - coupler.reset_coupler() - assert coupler.coupled_mode is False - - # test loading from the dataset - coupled_scaling = { - "mean": np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)), - "std": np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)), - } - expected = zarr_ds["inputs"][:2][:, input_indices] - expected = (expected - coupled_scaling["mean"]) / coupled_scaling["std"] - coupled_field = coupler.construct_integrated_couplings( - batch=batch, bsize=batch_size - ) - assert np.array_equal(expected, coupled_field[0]) - - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_ConstantCoupler_multiple_variables_order(dataset_path): - """Regression test for a bug where the Zarr-backed coupling path selected - coupled channels in the dataset's native channel_in order rather than the - order given by `variables`, while the xarray-backed path (`.sel`) always - honored the requested order. With 2+ coupled variables requested out of - their native dataset order, this silently misaligned `coupled_scaling` - (which is built from `variables` order) against the loaded data, applying - the wrong mean/std to the wrong channel. Only surfaces with 2+ variables - that aren't already in native dataset order, which none of the other - tests happen to exercise. - """ - from physicsnemo.datapipes.healpix.couplers import ConstantCoupler - - xr_ds = xr.open_zarr(dataset_path) - zarr_ds = zarr.open(dataset_path) - - # deliberately request the coupled variables in the opposite order from - # how they appear natively in the dataset's channel_in coordinate - # (native order is z500, ..., z1000, ..., z250) - native_order = list(xr_ds.channel_in.values) - variables = ["z250", "z1000", "z500"] - assert [native_order.index(v) for v in variables] == sorted( - [native_order.index(v) for v in variables], reverse=True - ) - - batch_size = 2 - batch = {"time": slice(0, 2)} - - coupler_xr = ConstantCoupler( - dataset=xr_ds, - batch_size=batch_size, - variables=variables, - input_times=["0h"], - input_time_dim=1, - output_time_dim=1, - ) - coupler_zarr = ConstantCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - input_times=["0h"], - input_time_dim=1, - output_time_dim=1, - ) - - interval = 2 - data_time_step = "3h" - coupler_xr.compute_coupled_indices(interval, data_time_step) - coupler_zarr.compute_coupled_indices(interval, data_time_step) - - coupled_xr = coupler_xr.construct_integrated_couplings( - batch=batch, bsize=batch_size - ) - coupled_zarr = coupler_zarr.construct_integrated_couplings( - batch=batch, bsize=batch_size - ) - - # the Zarr-backed and xarray-backed couplers must agree exactly, and both - # must match the raw data indexed in the order `variables` was given, not - # the dataset's native channel_in order - input_indices = [ - int(np.where(zarr_ds["channel_in"][:] == v)[0][0]) for v in variables - ] - expected = zarr_ds["inputs"][:2][:, input_indices] - # coupled_integration_dim is 1 here, so index into the leading dim to get - # the [batch, channel, face, height, width] slice that lines up with `expected` - assert np.array_equal(expected, coupled_xr[0]) - assert np.array_equal(expected, coupled_zarr[0]) - - # and, per-variable, each channel's data must match that specific - # variable's raw data, not another (coupled) variable's - for i, var in enumerate(variables): - var_index = int(np.where(zarr_ds["channel_in"][:] == var)[0][0]) - expected_var = zarr_ds["inputs"][:2][:, var_index] - assert np.array_equal(expected_var, coupled_zarr[0][:, i]) - assert np.array_equal(expected_var, coupled_xr[0][:, i]) - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_TrailingAverageCoupler(dataset_path, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.couplers import ( - TrailingAverageCoupler, - ) - - variables = ["z500", "z1000-12h"] - input_times = ["6h", "12h"] - input_time_dim = 2 - output_time_dim = 2 - presteps = 0 - batch_size = 2 - averaging_window = "6h" - # open our test dataset - zarr_ds = zarr.open(dataset_path) - - # test fail initialization when trying to prepare data - with pytest.raises( - NotImplementedError, match=("Data preparation not yet implemented") - ): - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - averaging_window=averaging_window, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - prepared_coupled_data=False, - ) - - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=presteps, - averaging_window=averaging_window, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - assert isinstance(coupler, TrailingAverageCoupler) - - # Zarr-backed datasets store time as CF-encoded integers rather than - # numpy datetime64; verify the coupler decodes them via cftime into the - # same wall-clock times as the xarray view of the same store - expected_time_da = xr.open_zarr(dataset_path).time.values - assert np.array_equal(coupler.time_da, expected_time_da) - - mock_coupled_module = coupler_helper( - output_variables=["not_coupled", "z500"], - time_step="3h", - ) - with pytest.raises(ValueError, match=("Missing variables in coupled module")): - coupler.setup_coupling(mock_coupled_module) - - # veryify averaging slices computed correctly - mock_coupled_module = coupler_helper( - output_variables=["z500", "z1000"], - time_step="3h", - ) - coupler.setup_coupling(mock_coupled_module) - averaging_window_max_indices = [ - i // pd.Timedelta(mock_coupled_module.time_step) for i in input_times - ] - dt = averaging_window_max_indices[0] - # assumes only 1 integration step, otherwise would be wrong - expected_slices = [[]] - for i, window_end in enumerate(averaging_window_max_indices): - expected_slices[0].append(slice(i * dt, window_end)) - assert expected_slices == coupler.averaging_slices - - interval = 2 - data_time_step = "3h" - coupler.compute_coupled_indices(interval, data_time_step) - coupled_integration_dim = presteps + max(output_time_dim // input_time_dim, 1) - expected = np.empty([batch_size, coupled_integration_dim, len(input_times)]) - for b in range(batch_size): - for i in range(coupled_integration_dim): - expected[b, i, :] = ( - b - + (input_time_dim * i + 1) * interval - + np.array( - [ - pd.Timedelta(ts) / pd.Timedelta(data_time_step) - for ts in input_times - ] - ) - ) - expected = expected.astype(int) - assert np.array_equal(expected, coupler._coupled_offsets) - - scaling_df = pd.DataFrame.from_dict(omegaconf.OmegaConf.to_object(scaling_dict)).T - scaling_df.loc["zeros"] = {"mean": 0.0, "std": 1.0} - scaling_da = scaling_df.to_xarray().astype("float32") - coupler.set_scaling(scaling_da) - coupled_scaling = scaling_da.sel(index=variables).rename({"index": "channel_in"}) - expected = np.expand_dims(coupled_scaling["mean"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["mean"]) - expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) - assert np.array_equal(expected, coupler.coupled_scaling["std"]) - - averaging_window_max_indices = [ - i // pd.Timedelta(data_time_step) for i in coupler.input_times - ] - di = averaging_window_max_indices[0] - averaging_slices = [] - for j in range(coupler.coupled_integration_dim): - averaging_slices.append([]) - for i, r in enumerate(averaging_window_max_indices): - averaging_slices[j].append( - slice( - coupler.input_time_dim * j * di + i * di, - coupler.input_time_dim * j * di + r, - ) - ) - coupler.averaging_slices = averaging_slices - coupler.coupled_channel_indices = [0, 1] - - coupled_fields_batch_size = batch_size - coupled_fields_timedim = 4 - expected_shape = [ - coupler.coupled_integration_dim, - coupled_fields_batch_size, - coupler.timevar_dim, - ] + list(coupler.spatial_dims) - coupled_fields = th.rand( - coupled_fields_batch_size, - coupler.spatial_dims[0], - coupled_fields_timedim, - len(coupler.coupled_channel_indices), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - coupler.set_coupled_fields(coupled_fields) - assert list(coupler.preset_coupled_fields.shape) == expected_shape - - # check reset - assert coupler.coupled_mode - coupler.reset_coupler() - assert coupler.coupled_mode is False - - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("pandas") -@requires_module("xarray") -def test_TrailingAverageCoupler_multiple_variables(dataset_path): - """Regression test to ensure `set_coupled_fields` correctly averages and - threads through every coupled variable, not just the first. Gives each - coupled variable a distinct, constant value in the synthetic - `coupled_fields` tensor so that averaging over any time window still - yields that same constant, letting each variable's contribution to - `preset_coupled_fields` be checked in isolation and by position. - """ - from physicsnemo.datapipes.healpix.couplers import TrailingAverageCoupler - - variables = ["z500", "z1000-12h", "z250"] - input_times = ["6h", "12h"] - input_time_dim = 2 - output_time_dim = 2 - batch_size = 2 - averaging_window = "6h" - zarr_ds = zarr.open(dataset_path) - - coupler = TrailingAverageCoupler( - dataset=zarr_ds, - batch_size=batch_size, - variables=variables, - presteps=0, - averaging_window=averaging_window, - input_times=input_times, - input_time_dim=input_time_dim, - output_time_dim=output_time_dim, - ) - coupler.coupled_channel_indices = list(range(len(variables))) - - data_time_step = "3h" - averaging_window_max_indices = [ - i // pd.Timedelta(data_time_step) for i in input_times - ] - di = averaging_window_max_indices[0] - averaging_slices = [] - for j in range(coupler.coupled_integration_dim): - averaging_slices.append([]) - for i, r in enumerate(averaging_window_max_indices): - averaging_slices[j].append( - slice( - coupler.input_time_dim * j * di + i * di, - coupler.input_time_dim * j * di + r, - ) - ) - coupler.averaging_slices = averaging_slices - - # give each coupled variable a distinct, constant value across every - # other dimension so the averaging window doesn't change its value, and - # each variable's contribution can be checked independently of the - # others - channel_values = [100.0, 200.0, 300.0] - coupled_fields = th.empty( - batch_size, - coupler.spatial_dims[0], - 4, - len(variables), - coupler.spatial_dims[1], - coupler.spatial_dims[2], - ) - for i, value in enumerate(channel_values): - coupled_fields[:, :, :, i, :, :] = value - - coupler.set_coupled_fields(coupled_fields) - result = coupler.construct_integrated_couplings() - assert list(result.shape) == [ - coupler.coupled_integration_dim, - batch_size, - coupler.timevar_dim, - ] + list(coupler.spatial_dims) - - # timevar ordering groups by averaging period first, then by variable - # (matching the dataset-loading path in - # `_construct_integrated_couplings_from_dataset`/`construct_integrated_couplings`) - for period in range(len(input_times)): - for var_idx, value in enumerate(channel_values): - timevar_idx = period * len(variables) + var_idx - slice_result = result[:, :, timevar_idx, :, :, :] - assert th.allclose(slice_result, th.full_like(slice_result, value)), ( - f"coupled variable {variables[var_idx]!r} (value {value}) was not " - f"preserved at averaging period {period}, timevar index {timevar_idx}" - ) - - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDatasetZarr_initialization( - dataset_path, scaling_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( - CoupledTimeSeriesDatasetZarr, - ) - - # open our test dataset - time_da = xr.open_zarr(dataset_path).time.values - input_variables = ["z500", "z1000"] - valid_start_date = "1979-01-01" - valid_end_date = "1979-01-02" - - # check for failure of timestep not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'time_step' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - data_time_step="2h", - time_step="5h", - scaling=scaling_dict, - forecast_init_times=time_da[:2], - ) - - # check for failure of gap not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'gap' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - data_time_step="2h", - time_step="6h", - gap="3h", - scaling=scaling_dict, - batch_size=1, - forecast_init_times=time_da[:2], - ) - - # check for failure of invalid scaling variable on input - invalid_scaling = omegaconf.DictConfig( - { - "bogosity": {"mean": 0, "std": 42}, - } - ) - with pytest.raises(KeyError, match=("Input channels ")): - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - data_time_step="3h", - time_step="6h", - scaling=invalid_scaling, - batch_size=1, - forecast_init_times=time_da[:2], - ) - - # check for warning on batch size > 1 and forecast mode - warnings.filterwarnings("error") - with pytest.raises( - UserWarning, - match=( - "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`" - ), - ): - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=2, - forecast_init_times=time_da[:2], - ) - warnings.resetwarnings() - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - add_train_noise=True, - batch_size=1, - forecast_init_times=time_da[:2], - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - start_date=valid_start_date, - end_date=valid_end_date, - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=time_da[:2], - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=time_da[:2], - data_time_step="3h", - time_step="6h", - ) - assert isinstance(timeseries_ds, CoupledTimeSeriesDatasetZarr) - - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDatasetZarr_get_constants( - dataset_path, scaling_dict, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( - CoupledTimeSeriesDatasetZarr, - ) - - input_variables = ["z500", "z1000"] - constant_variables = ["lsm"] - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - constant_indices = [ - int(np.where(zarr_ds.channel_c[:] == ch)[0][0]) for ch in constant_variables - ] - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_dict, - couplings=constant_coupler_config, - forecast_init_times=zarr_ds.time[:2], - ) - assert timeseries_ds.get_constants() is None - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=constant_variables, - batch_size=1, - scaling=scaling_dict, - couplings=constant_coupler_config, - forecast_init_times=zarr_ds.time[:2], - ) - - # constants are reshaped - expected = np.transpose( - zarr_ds.constants.values[constant_indices], axes=(1, 0, 2, 3) - ) - outvar = timeseries_ds.get_constants() - assert np.array_equal( - expected, - outvar, - ) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDatasetZarr_len( - dataset_path, scaling_dict, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( - CoupledTimeSeriesDatasetZarr, - ) - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - - variables = ["z500", "z1000"] - batch_size = 2 - - # check forecast mode - init_times = random.randint(1, zarr_ds.time.shape[0]) - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler_config, - ) - assert len(timeseries_ds) == init_times - - batch2_coupler = constant_coupler_config.copy() - batch2_coupler[0]["params"]["batch_size"] = 2 - - # get the last index that's evenly divisible by 3 (9h / 3h) - last_index = (zarr_ds.time.shape[0] // 3) * 3 - 1 - - # check train mode - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=variables, - data_time_step="3h", - time_step="9h", - scaling=scaling_dict, - batch_size=batch_size, - couplings=batch2_coupler, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[last_index - 1].values, - ) - # Window length of 3 for one sample size - assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 3) // batch_size - - # drop incomplete last window - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=variables, - data_time_step="3h", - time_step="9h", - scaling=scaling_dict, - batch_size=batch_size, - drop_last=True, - couplings=batch2_coupler, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[last_index - 1].values, - ) - assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 4) // batch_size - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDatasetZarr_get( - dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( - CoupledTimeSeriesDatasetZarr, - ) - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - - input_variables = list(zarr_ds.channel_out.values) - constant_variables = ["lsm"] - batch_size = 2 - batch2_constant_coupler = constant_coupler_config.copy() - batch2_constant_coupler[0]["params"]["batch_size"] = 2 - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=batch2_constant_coupler, - ) - - # check for invalid index - invalid_idx = len(zarr_ds.targets) + 1 - with pytest.raises( - IndexError, match=(f"index {invalid_idx} out of range for dataset with length") - ): - inputs, targets = timeseries_ds[invalid_idx] - - inputs, targets = timeseries_ds[0] - - # make sure number of targets is correct - assert len(targets) == batch_size - - # check target data - # need to transpose - targets_expected = zarr_ds.targets[batch_size].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # check for negative index - inputs, targets = timeseries_ds[-1] - targets_expected = zarr_ds.targets[12].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - - # we're not dropping incomplete elements by default - assert len(targets) == 0 - - # this time dropping incomplete so that we get a full sample - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=batch2_constant_coupler, - ) - - inputs, targets = timeseries_ds[-1] - targets_expected = zarr_ds.targets[-1 - batch_size].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # without couplings - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=constant_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=[], - ) - non_perturbed_inputs = timeseries_ds - assert len(non_perturbed_inputs[0][0]) == 2 # just inputs and targets - - # without couplings but with noise - noise_params = { - "inputs": scaling_double_dict, - "couplings": scaling_double_dict, - } - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=constant_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_train_noise=True, - train_noise_params=noise_params, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=[], - ) - perturbed_inputs = timeseries_ds - # The first input will be the same sample, with perturbation it should have - # different values - assert non_perturbed_inputs[0][0][0].shape == perturbed_inputs[0][0][0].shape - assert not np.array_equal(non_perturbed_inputs[0][0][0], perturbed_inputs[0][0][0]) - - # noise is also added to the integrated couplings themselves, not just - # the regular inputs, when couplings are present. With no constants and - # no insolation, inputs_result is [inputs, integrated_couplings] - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=batch2_constant_coupler, - ) - non_perturbed_sample = timeseries_ds[0][0] - assert len(non_perturbed_sample) == 2 - non_perturbed_couplings = non_perturbed_sample[1] - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_train_noise=True, - train_noise_params=noise_params, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=batch2_constant_coupler, - ) - perturbed_couplings = timeseries_ds[0][0][1] - assert non_perturbed_couplings.shape == perturbed_couplings.shape - assert not np.array_equal(non_perturbed_couplings, perturbed_couplings) - - # With insolation we get 1 extra channel - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=None, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_insolation=True, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=batch2_constant_coupler, - ) - assert (len(inputs)) + 1 == len(timeseries_ds[0][0]) - - # nothing should change with forecast mode other than getting just inputs - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=None, - scaling=scaling_double_dict, - batch_size=1, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[-1].values, - couplings=constant_coupler_config, - ) - inputs = timeseries_ds[0] - - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # insolation adds 1 extra channel - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=None, - scaling=scaling_double_dict, - batch_size=1, - add_insolation=True, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler_config, - ) - assert (len(inputs)) + 1 == len(timeseries_ds[0]) - - # Constants + insolation is 2 extra channels - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=constant_variables, - scaling=scaling_double_dict, - batch_size=1, - add_insolation=True, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler_config, - ) - assert len(inputs) + 2 == len(timeseries_ds[0]) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataModuleZarr_initialization( - dataset_path, splits, scaling_double_dict, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( - CoupledTimeSeriesDataModuleZarr, - ) - - input_variables = ["z500", "z1000"] - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - - # test for invalid path - with pytest.raises(FileNotFoundError, match=("Dataset path not found")): - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path="DoesntExist", - input_variables=input_variables, - batch_size=1, - couplings=constant_coupler_config, - ) - - # use the prebuilt dataset - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - splits=omegaconf.DictConfig(splits), - couplings=constant_coupler_config, - ) - assert isinstance(timeseries_dm, CoupledTimeSeriesDataModuleZarr) - - # with init times - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - forecast_init_times=zarr_ds.time[:2], - couplings=constant_coupler_config, - ) - assert isinstance(timeseries_dm, CoupledTimeSeriesDataModuleZarr) - - # with splits - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - splits=omegaconf.DictConfig(splits), - couplings=constant_coupler_config, - ) - assert isinstance(timeseries_dm, CoupledTimeSeriesDataModuleZarr) - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataModuleZarr_get_constants( - dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( - CoupledTimeSeriesDataModuleZarr, - ) - - variables = ["z500", "z1000"] - constants = ["lsm"] - - # No constants - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=variables, - batch_size=1, - scaling=scaling_double_dict, - splits=splits, - constant_variables=None, - couplings=constant_coupler_config, - ) - - assert timeseries_dm.get_constants() is None - - # just lsm as constant - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=variables, - batch_size=1, - scaling=scaling_double_dict, - splits=splits, - constant_variables=constants, - couplings=constant_coupler_config, - ) - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - - # divide by 2 due to scaling - expected = ( - np.transpose( - zarr_ds.constants.sel(channel_c=constants).values, - axes=(1, 0, 2, 3), - ) - / 2.0 - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - - # with splits we're doing forecasting and get - # constants from train instead of test dataset - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=variables, - batch_size=1, - scaling=scaling_double_dict, - splits=splits, - constant_variables=constants, - couplings=constant_coupler_config, - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -def test_CoupledTimeSeriesDataModuleZarr_get_dataloaders( - dataset_path, scaling_double_dict, splits, constant_coupler_config, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( - CoupledTimeSeriesDataModuleZarr, - ) - - input_variables = ["z500", "z1000"] - - # use the prebuilt dataset - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - splits=splits, - shuffle=False, - couplings=constant_coupler_config, - ) - - assert_shard_dataloaders(timeseries_dm) - DistributedManager.cleanup() - - -@requires_module("omegaconf") -def test_CoupledTimeSeriesDataModuleZarr_get_coupled_vars( - dataset_path, - scaling_double_dict, - splits, - constant_coupler_config, - average_coupler_config, - pytestconfig, -): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( - CoupledTimeSeriesDataModuleZarr, - ) - - input_variables = ["z500", "z1000"] - - # Constant coupler - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - splits=splits, - couplings=constant_coupler_config, - ) - - outvar = timeseries_dm._get_coupled_vars() - outvar.sort() - expected = ["z250"] - expected.sort() - - assert expected == outvar - - # Average coupler - # Internally initializes DistributedManager - timeseries_dm = CoupledTimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - splits=splits, - couplings=average_coupler_config, - ) - outvar = timeseries_dm._get_coupled_vars() - outvar.sort() - - assert expected == outvar - - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDatasetZarr_next_integration( - dataset_path, scaling_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset_zarr import ( - CoupledTimeSeriesDatasetZarr, - ) - - spatial_dims = [12, 32, 32] - input_variables = ["z500", "z1000"] - # length must match len(coupled_variables) (i.e. the coupler's timevar_dim); - # values are arbitrary since this synthetic setup bypasses setup_coupling() - coupled_channel_indices = [0] - coupled_variables = ["z250"] - num_variables = len(input_variables) - input_time_dim = 1 - output_time_dim = 1 - batch_size = 1 - - constant_coupler = [ - { - "coupler": "ConstantCoupler", - "params": { - "batch_size": 1, - "variables": coupled_variables, - "input_times": ["0h"], - "input_time_dim": input_time_dim, - "output_time_dim": output_time_dim, - "presteps": 0, - "prepared_coupled_data": True, - }, - } - ] - - # open our test dataset - ds = xr.open_zarr(dataset_path) - init_times = random.randint(1, len(ds.time.values)) - # channels need to be subselected before being handed over - test_ds = ds.sel( - channel_in=input_variables + coupled_variables, - channel_out=input_variables, - ) - - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=batch_size, - couplings=constant_coupler, - data_time_step="6h", - time_step="6h", - drop_last=True, - add_insolation=True, - forecast_init_times=test_ds.time[:init_times], - ) - - test_model_outputs = th.rand( - 1, - spatial_dims[0], - output_time_dim, - num_variables, - spatial_dims[1], - spatial_dims[2], - ) - constants = np.transpose(ds.constants.values, axes=(1, 0, 2, 3)) - coupled_fields = th.rand( - batch_size, - spatial_dims[0], - input_time_dim + output_time_dim, - len(input_variables), - spatial_dims[1], - spatial_dims[2], - ) - - # mirrors set_coupled_fields: broadcast the first time step across the - # integration dim and permute to [integration_dim, batch, timevar, face, height, width] - expected_coupling = coupled_fields[:, :, :, coupled_channel_indices, :, :] - expected_coupling = expected_coupling[:, :, :1, :, :, :] - expected_coupling = expected_coupling.permute(2, 0, 3, 1, 4, 5) - - # need to grab at least 1 sample to properly intialize everything - timeseries_ds[0] - # hacky way to setup the indices since we don't actually have any coupled fields - timeseries_ds.couplings[0].coupled_channel_indices = coupled_channel_indices - - # set the coupled fields - timeseries_ds.couplings[0].set_coupled_fields(coupled_fields) - test_integration = timeseries_ds.next_integration(test_model_outputs, constants) - # test to make sure prognostics are used, constants stay the same, and couplings - # are what we set - assert np.array_equal(test_integration[0], test_model_outputs[:, :, -1:]) - assert np.array_equal(test_integration[2], constants) - assert np.array_equal(test_integration[3], expected_coupling) - - # I have absolutely no idea why a coupled dataset has the option for 0 couplings - timeseries_ds = CoupledTimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=batch_size, - couplings=[], - data_time_step="6h", - time_step="6h", - drop_last=True, - add_insolation=True, - forecast_init_times=test_ds.time[:init_times], - ) - # need to grab at least 1 sample to properly intialize everything - timeseries_ds[0] - test_integration = timeseries_ds.next_integration(test_model_outputs, constants) - assert np.array_equal(test_integration[0], test_model_outputs[:, :, -1:]) - assert np.array_equal(test_integration[2], constants) diff --git a/test/datapipes/healpix/test_healpix_zarr.py b/test/datapipes/healpix/test_healpix_zarr.py deleted file mode 100644 index ab111a7198..0000000000 --- a/test/datapipes/healpix/test_healpix_zarr.py +++ /dev/null @@ -1,873 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import importlib.util -import random -import warnings - -import pytest - -from physicsnemo.distributed import DistributedManager -from test.conftest import requires_module -from test.datapipes.healpix.conftest import assert_shard_dataloaders - -omegaconf = pytest.importorskip("omegaconf") -np = pytest.importorskip("numpy") -xr = pytest.importorskip("xarray") -zarr = pytest.importorskip("zarr") - - -def test_object_store_path_helpers(tmp_path): - """Zarr datasets may live on object stores (e.g. s3://, or fsspec-chained paths - like simplecache::s3://); such paths are recognized by path syntax alone and - skip the local filesystem existence check performed for plain local paths. - This behavior is unique to the Zarr-backed datapipes and doesn't require the - NFS test dataset, so it can run unconditionally. - """ - from physicsnemo.datapipes.healpix.base_timeseries_dataset_zarr import ( - _check_availability, - _is_object_store_path, - ) - - assert _is_object_store_path("s3://bucket/data.zarr") - assert _is_object_store_path("simplecache::s3://bucket/data.zarr") - assert not _is_object_store_path(str(tmp_path / "local.zarr")) - - # an existing local path passes without needing fsspec - existing_path = tmp_path / "exists.zarr" - existing_path.mkdir() - _check_availability(str(existing_path)) - - # a missing local path raises FileNotFoundError - with pytest.raises(FileNotFoundError, match=("Dataset not found at")): - _check_availability(str(tmp_path / "missing.zarr")) - - # object store paths bypass the local existence check as long as fsspec - # is installed, even when the remote resource doesn't actually exist - if importlib.util.find_spec("fsspec"): - _check_availability("s3://bucket-that-does-not-exist/missing.zarr") - else: - with pytest.raises( - ImportError, match=("fsspec is required to access object store paths") - ): - _check_availability("s3://bucket-that-does-not-exist/missing.zarr") - - -@requires_module("omegaconf") -@requires_module("netCDF4") -def test_TimeSeriesDataset_initialization( - dataset_path, - scaling_dict, - pytestconfig, -): - from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( - TimeSeriesDatasetZarr, - ) - - input_variables = ["t2m0", "t850", "z500"] - - bad_start_date = "1900-01-01" - valid_start_date = "1979-01-01" - bad_end_date = "2000-12-31" - valid_end_date = "1979-01-02" - - time_da = xr.open_zarr(dataset_path).time - - # check for failure of invalid dataset path - # Check if fsspec is available for object store paths, optional dependency - if not importlib.util.find_spec("fsspec"): - # If fsspec is not available, expect an ImportError - with pytest.raises( - ImportError, match=("fsspec is required to access object store paths") - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path="s3://physicsnemo-data/datasets/healpix/healpix.zarr", - scaling=scaling_dict, - input_variables=input_variables, - ) - - # If path doesn't exist, expect a FileNotFoundError - with pytest.raises(FileNotFoundError, match=("Dataset not found at")): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path="/boguspath.zarr", - scaling=scaling_dict, - input_variables=input_variables, - ) - - # check for failure of timestep not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'time_step' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - data_time_step="2h", - time_step="5h", - scaling=scaling_dict, - input_variables=input_variables, - forecast_init_times=time_da[:2], - batch_size=1, - ) - - # check for failure of gap not being a multiple of datatime step - with pytest.raises( - ValueError, match=("'gap' must be a multiple of 'data_time_step' ") - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - data_time_step="2h", - time_step="6h", - gap="3h", - scaling=scaling_dict, - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=input_variables, - ) - - # check for failure of invalid scaling variable on input - invalid_scaling = omegaconf.DictConfig( - { - "bogosity": {"mean": 0, "std": 42}, - } - ) - with pytest.raises(KeyError, match=("Input channels ")): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - data_time_step="3h", - time_step="6h", - scaling=invalid_scaling, - forecast_init_times=time_da[:10], - input_variables=input_variables, - batch_size=1, - ) - - # check for failure when a requested variable isn't present in the dataset - # this validation is specific to the Zarr-backed dataset, which selects - # channels by name directly out of the store rather than relying on a - # pre-sliced xarray Dataset - with pytest.raises( - KeyError, - match=("Requested Input, coupled, or output variables not found in dataset"), - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=input_variables + ["DoesntExist"], - ) - - # check for warning when the configured data_time_step doesn't match the - # dataset's native cadence; this cross-check against the stored time - # coordinate only exists in the Zarr-backed dataset - warnings.filterwarnings("error") - with pytest.raises(UserWarning, match=("doesn't match configuration dt")): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - data_time_step="6h", - time_step="6h", - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=input_variables, - ) - warnings.resetwarnings() - - # check for warning on batch size > 1 and forecast mode - warnings.filterwarnings("error") - with pytest.raises( - UserWarning, - match=( - "providing 'forecast_init_times' to TimeSeriesDataset requires `batch_size=1`" - ), - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - batch_size=2, - forecast_init_times=time_da[:2], - input_variables=input_variables, - ) - warnings.resetwarnings() - - # check for no dates provided - with pytest.raises( - ValueError, - match=("Either start and end date or forecast_init_times must be provided"), - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - batch_size=2, - input_variables=input_variables, - ) - - # check for out of range dates - warnings.filterwarnings("error") - with pytest.raises( - UserWarning, - match=(f"Start date {bad_start_date} is before first available date"), - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - batch_size=2, - start_date=bad_start_date, - end_date=valid_end_date, - input_variables=input_variables, - ) - - with pytest.raises( - UserWarning, match=(f"End date {bad_end_date} is after last available date") - ): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - batch_size=2, - start_date=valid_start_date, - end_date=bad_end_date, - input_variables=input_variables, - ) - warnings.resetwarnings() - - # test no scaling - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=input_variables, - ) - assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) - - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=input_variables, - ) - assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) - - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - batch_size=1, - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=input_variables, - ) - assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) - - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - batch_size=1, - start_date=valid_start_date, - end_date=valid_end_date, - data_time_step="3h", - time_step="6h", - input_variables=input_variables, - ) - assert isinstance(timeseries_ds, TimeSeriesDatasetZarr) - - # `start_date`/`end_date` can also be integer positional indices into the - # time array rather than dates (use a nonzero start_date since 0 is falsy - # and would take the same code path as "no start date provided") - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_dict, - start_date=2, - end_date=5, - input_variables=input_variables, - ) - assert timeseries_ds.start_index == 2 - assert timeseries_ds.total_samples == 3 - - # check for failure when a requested output variable has no scaling - # entry, even though it's present in the dataset (and in the input - # scaling, if it happens to also be an input variable) - scaling_missing_target = omegaconf.DictConfig( - {k: v for k, v in scaling_dict.items() if k != "z1000"} - ) - with pytest.raises(KeyError, match=("Target channels ")): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_missing_target, - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=["t2m0"], - output_variables=["t2m0", "z1000"], - ) - - # check for failure when a requested constant variable exists in the - # dataset but has no corresponding scaling entry - scaling_missing_constant = omegaconf.DictConfig( - {k: v for k, v in scaling_dict.items() if k != "z"} - ) - with pytest.raises(KeyError, match=("Constant variables ")): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - scaling=scaling_missing_constant, - start_date=valid_start_date, - end_date=valid_end_date, - input_variables=input_variables, - constant_variables=["lsm", "z"], - ) - - -def test_TimeSeriesDataset_missing_time(tmp_path): - """The Zarr-backed dataset validates that the store has a `time` - variable before doing any time-based indexing; this check is unique to - the Zarr path (the classic path always derives its dataset from a - pre-opened xarray Dataset that already has a time dimension), so build a - minimal Zarr store without a time coordinate to exercise it directly. - """ - from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( - TimeSeriesDatasetZarr, - ) - - no_time_ds = xr.Dataset( - data_vars={ - "inputs": ( - ("t", "channel_in", "face", "height", "width"), - np.zeros((3, 1, 1, 2, 2), dtype="float32"), - ), - "targets": ( - ("t", "channel_out", "face", "height", "width"), - np.zeros((3, 1, 1, 2, 2), dtype="float32"), - ), - }, - coords={ - "channel_in": ["z500"], - "channel_out": ["z500"], - }, - ) - dataset_path = tmp_path / "no_time.zarr" - no_time_ds.to_zarr(dataset_path) - - with pytest.raises(KeyError, match=("Dataset missing time")): - TimeSeriesDatasetZarr( - dataset_path=str(dataset_path), - input_variables=["z500"], - start_date="1979-01-01", - end_date="1979-01-02", - ) - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("numpy") -@requires_module("zarr") -def test_TimeSeriesDataset_get_constants(dataset_path, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( - TimeSeriesDatasetZarr, - ) - - input_variables = ["z500", "z1000"] - constant_variables = ["lsm"] - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - constant_indices = [ - int(np.where(zarr_ds.channel_c[:] == ch)[0][0]) for ch in constant_variables - ] - - # Constant that isn't in dataset - with pytest.raises(KeyError, match=("Requested constants not found in dataset")): - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_dict, - constant_variables=["DoesntExist"], - forecast_init_times=zarr_ds.time[:2], - ) - - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - constant_variables=None, - forecast_init_times=zarr_ds.time[:2], - batch_size=1, - ) - - assert timeseries_ds.get_constants() is None - - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - constant_variables=constant_variables, - forecast_init_times=zarr_ds.time[:2], - batch_size=1, - ) - - # constants are reshaped - expected = np.transpose( - zarr_ds.constants.values[constant_indices], axes=(1, 0, 2, 3) - ) - outvar = timeseries_ds.get_constants() - assert np.array_equal( - expected, - outvar, - ) - zarr_ds.close() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -def test_TimeSeriesDataset_len(dataset_path, scaling_dict, pytestconfig): - from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( - TimeSeriesDatasetZarr, - ) - - input_variables = ["z500", "z1000"] - batch_size = 2 - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - - # check forecast mode - init_times = random.randint(1, zarr_ds.time.shape[0]) - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - ) - assert len(timeseries_ds) == init_times - - # get the last index that's evenly divisible by 3 (9h / 3h) - last_index = (zarr_ds.time.shape[0] // 3) * 3 - 1 - - # check train mode - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - data_time_step="3h", - time_step="9h", - input_variables=input_variables, - scaling=scaling_dict, - batch_size=batch_size, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[last_index - 1].values, - ) - # Window length of 3 for one sample size - assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 3) // batch_size - - # drop incomplete last window - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - data_time_step="3h", - time_step="9h", - input_variables=input_variables, - scaling=scaling_dict, - batch_size=batch_size, - drop_last=True, - start_date=zarr_ds.time[0].values, - end_date=zarr_ds.time[last_index - 1].values, - ) - assert len(timeseries_ds) == (zarr_ds.time.shape[0] - 4) // batch_size - - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("numpy") -def test_TimeSeriesDataset_get(dataset_path, scaling_double_dict, splits, pytestconfig): - from physicsnemo.datapipes.healpix.timeseries_dataset_zarr import ( - TimeSeriesDatasetZarr, - ) - - input_variables = ["z500", "z1000"] - constant_variables = ["lsm"] - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - time_da = zarr_ds.time.values - - batch_size = 2 - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - start_date=splits["train_date_start"], - end_date=splits["test_date_start"], - ) - - # check for invalid index - invalid_idx = len(zarr_ds.targets) + 1 - with pytest.raises( - IndexError, match=(f"index {invalid_idx} out of range for dataset with length") - ): - inputs, targets = timeseries_ds[invalid_idx] - - inputs, targets = timeseries_ds[0] - - # make sure number of targets is correct - assert len(targets) == batch_size - - # check target data - # need to transpose - targets_expected = ( - zarr_ds.targets[batch_size] - .transpose("face", "channel_out", "height", "width") - .sel(channel_out=input_variables) - ) - - # scale target data - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # check for negative index - inputs, targets = timeseries_ds[-1] - targets_expected = zarr_ds.targets[12].transpose( - "face", "channel_out", "height", "width" - ) - targets_expected = targets_expected.to_numpy() / 2 - - # we're not dropping incomplete elements by default - assert len(targets) == 0 - - # this time dropping incomplete so that we get a full sample sample - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - start_date=time_da[0], - end_date=time_da[-1], - ) - - inputs, targets = timeseries_ds[-1] - targets_expected = ( - zarr_ds.targets[-1 - batch_size] - .transpose("face", "channel_out", "height", "width") - .sel(channel_out=input_variables) - ) - targets_expected = targets_expected.to_numpy() / 2 - assert np.array_equal(targets[0][:, 0, :, :], targets_expected) - - # With insolation we get 1 extra tensor - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_insolation=True, - start_date=time_da[0], - end_date=time_da[-1], - ) - - # verify underlying data doesn't change - targets_expected = ( - zarr_ds.targets[2] - .transpose("face", "channel_out", "height", "width") - .sel(channel_out=input_variables) - ) - targets_expected = targets_expected.to_numpy() / 2 - result = timeseries_ds[0] - assert (len(inputs)) + 1 == len(result[0]) - assert np.array_equal(result[1][0][:, 0, :, :], targets_expected) - - # With insolation and constants we get 2 extra tensors - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - constant_variables=constant_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_insolation=True, - start_date=time_da[0], - end_date=time_da[-1], - ) - assert (len(inputs)) + 2 == len(timeseries_ds[0][0]) - - # train noise is applied directly to the inputs of the (non-coupled) Zarr - # dataset; this option doesn't exist on the classic (non-Zarr) TimeSeriesDataset - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - start_date=time_da[0], - end_date=time_da[-1], - ) - non_perturbed = timeseries_ds[0] - - noise_params = {"inputs": scaling_double_dict} - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=batch_size, - drop_last=True, - add_train_noise=True, - train_noise_params=noise_params, - start_date=time_da[0], - end_date=time_da[-1], - ) - perturbed = timeseries_ds[0] - - # same shape, but the perturbed sample should differ from the un-perturbed one - assert non_perturbed[0][0].shape == perturbed[0][0].shape - assert not np.array_equal(non_perturbed[0][0], perturbed[0][0]) - - # nothing should change with forecast mode other than getting just inputs - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times].values, - ) - inputs = timeseries_ds[0] - - assert len(inputs) == 1 - - # insolation adds 1 extra tensor, same as above but using forecast mode - init_times = random.randint(1, len(zarr_ds.time.values)) - timeseries_ds = TimeSeriesDatasetZarr( - dataset_path=dataset_path, - input_variables=input_variables, - scaling=scaling_double_dict, - batch_size=1, - add_insolation=True, - forecast_init_times=zarr_ds.time[:init_times].values, - ) - assert (len(inputs) + 1) == len(timeseries_ds[0]) - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("zarr") -def test_TimeSeriesDataModule_initialization( - dataset_path, splits, scaling_double_dict, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( - TimeSeriesDataModuleZarr, - ) - - input_variables = ["z500", "z1000"] - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - - # test for invalid path - with pytest.raises(FileNotFoundError, match=("Dataset path not found")): - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path="DoesntExist", - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - ) - - # test for missing times - with pytest.raises( - ValueError, match=("Either splits or forecast_init_times must be provided") - ): - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - ) - - # test for overlapping dates - warnings.filterwarnings("error") - with pytest.raises( - UserWarning, match=("Training and validation date ranges overlap") - ): - bad_splits = splits.copy() - bad_splits["val_date_start"] = splits["train_date_start"] - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - splits=bad_splits, - scaling=scaling_double_dict, - ) - warnings.resetwarnings() - - warnings.filterwarnings("error") - with pytest.raises(UserWarning, match=("Training and test date ranges overlap")): - bad_splits = splits.copy() - bad_splits["test_date_start"] = splits["train_date_start"] - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - splits=bad_splits, - scaling=scaling_double_dict, - ) - warnings.resetwarnings() - - warnings.filterwarnings("error") - with pytest.raises(UserWarning, match=("Test and validation date ranges overlap")): - bad_splits = splits.copy() - bad_splits["val_date_end"] = splits["test_date_start"] - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - splits=bad_splits, - scaling=scaling_double_dict, - ) - warnings.resetwarnings() - - # with init times - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - forecast_init_times=zarr_ds.time[:2], - ) - assert isinstance(timeseries_dm, TimeSeriesDataModuleZarr) - - # with splits - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - splits=omegaconf.DictConfig(splits), - ) - assert isinstance(timeseries_dm, TimeSeriesDataModuleZarr) - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("netCDF4") -@requires_module("numpy") -@requires_module("zarr") -def test_TimeSeriesDataModule_get_constants( - dataset_path, scaling_double_dict, splits, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( - TimeSeriesDataModuleZarr, - ) - - input_variables = ["z500", "z1000"] - constants = ["lsm"] - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - forecast_times = zarr_ds.time[:2] - - # No constants - # Internally initializes DistributedManager - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - constant_variables=None, - forecast_init_times=forecast_times, - ) - - assert timeseries_dm.get_constants() is None - - # Constant that isn't in dataset - with pytest.raises(KeyError, match=("Requested constants not found in dataset")): - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - constant_variables={"missing": "missing"}, - forecast_init_times=forecast_times, - ) - - # just lsm as constant - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - constant_variables=constants, - forecast_init_times=forecast_times, - ) - - # open our test dataset - zarr_ds = xr.open_zarr(dataset_path) - - # dividing by 2 due to scaling - expected = ( - np.transpose( - zarr_ds.constants.sel(channel_c=constants).values, - axes=(1, 0, 2, 3), - ) - / 2.0 - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - - # with splits we're doing forecasting and get - # constants from train instead of test dataset - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - constant_variables=constants, - splits=splits, - ) - - assert np.array_equal( - timeseries_dm.get_constants(), - expected, - ) - zarr_ds.close() - DistributedManager.cleanup() - - -@requires_module("omegaconf") -@requires_module("zarr") -def test_TimeSeriesDataModule_get_dataloaders( - dataset_path, scaling_double_dict, splits, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules_zarr import ( - TimeSeriesDataModuleZarr, - ) - - input_variables = ["z500", "z1000"] - - # use the prebuilt dataset - # Internally initializes DistributedManager - timeseries_dm = TimeSeriesDataModuleZarr( - dataset_path=dataset_path, - input_variables=input_variables, - batch_size=1, - scaling=scaling_double_dict, - splits=splits, - shuffle=False, - ) - - assert_shard_dataloaders(timeseries_dm) - DistributedManager.cleanup() diff --git a/test/metrics/test_hydrostatic_loss.py b/test/metrics/test_hydrostatic_loss.py deleted file mode 100644 index fcc4019c2b..0000000000 --- a/test/metrics/test_hydrostatic_loss.py +++ /dev/null @@ -1,172 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-FileCopyrightText: All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import numpy as np -import pytest -import torch - -# xarray is an optional dependency (datapipes-extras). Skip this module entirely -# when it is unavailable, before importing hydrostasy, rather than failing at -# collection time. -pytest.importorskip("xarray") - -from physicsnemo.metrics.climate.hydrostasy import ( # noqa: E402 - HydrostaticBalance, -) - - -def test_constant_temperature(rtol: float = 1e-3, atol: float = 1e-3): - R = 287 # J K^{-1} kg^{-1} - g0 = 9.81 # m s^{-2} - T = 273.15 # K - z_pressure_levels = {0: 850.0, 1: 500.0, 2: 250.0, 3: 50.0} - z_pressure_levels = { - c: p * 100 for c, p in z_pressure_levels.items() - } # convert hPa to Pa - anchor_z_channel = 0 - anchor_T_channel = 4 - - # Create HydrostaticBalance constraint object - constraint = HydrostaticBalance( - z_pressure_levels, anchor_z_channel, anchor_T_channel, R, g0 - ) - - x = torch.zeros((1, 12, 5, 1, 1)) - # Set Z - z = ( - R - * T - / g0 - * torch.log( - z_pressure_levels[anchor_z_channel] - / torch.Tensor(list(z_pressure_levels.values())) - ) - ) - x[:, :, 0:4, :, :] = z.view(1, 1, -1, 1, 1) - # Set Tv at 850 hPa - x[:, :, 4, :, :] = T + torch.zeros_like(x[:, :, 4, :, :]) - - Tv = constraint(x) - - print(Tv[0, 0, :, 0, 0]) - assert torch.allclose( - T * torch.ones_like(Tv), - Tv, - rtol=rtol, - atol=atol, - ) - - -@pytest.mark.parametrize("N", [10, 20]) -def test_constant_lapse_rate(N: int, rtol: float = 1e-3, atol: float = 1e-3): - R = 287 # J K^{-1} kg^{-1} - g0 = 9.81 # m s^{-2} - T0 = 273.15 # K - z0 = 0 # m - Gamma = 9.8 / 1000 # K m^{-1} - p = np.logspace(np.log10(850.0), np.log10(50.0), num=N) - # z_pressure_levels = {0: 850., 1: 500., 2: 250., 3: 50.} - # z_pressure_levels = {i: 850. - 50.*i for i in range(17)} - z_pressure_levels = {i: p[i] for i in range(p.shape[0])} - z_pressure_levels = { - c: p * 100 for c, p in z_pressure_levels.items() - } # convert hPa to Pa - anchor_z_channel = 0 - anchor_T_channel = len(z_pressure_levels) - - # Create HydrostaticBalance constraint object - constraint = HydrostaticBalance( - z_pressure_levels, anchor_z_channel, anchor_T_channel, R, g0 - ) - - x = torch.zeros((1, 12, anchor_T_channel + 1, 1, 1)) - # Set T - T = T0 * torch.pow( - torch.Tensor(list(z_pressure_levels.values())) - / z_pressure_levels[anchor_z_channel], - Gamma * R / g0, - ) - print("T: ", T) - z = z0 + (T0 - T) / Gamma - print("z: ", z) - x[:, :, 0:anchor_T_channel, :, :] = z.view(1, 1, -1, 1, 1) - # Set Tv at 850 hPa - x[:, :, anchor_T_channel, :, :] = T[0] + torch.zeros_like( - x[:, :, anchor_T_channel, :, :] - ) - - Tv = constraint(x) - - print(Tv[0, 0, :, 0, 0]) - # assert torch.allclose(T * torch.ones_like(Tv), Tv) - return p, z, T, Tv[0, 0, constraint.z_channels, 0, 0] - - -@pytest.mark.parametrize("N", [10, 20]) -def test_dual_lapse_rate(N: int, rtol: float = 1e-3, atol: float = 1e-3): - R = 287 # J K^{-1} kg^{-1} - g0 = 9.81 # m s^{-2} - T0 = 273.15 # K - z0 = 0 # m - Gamma1 = 9.8 / 1000 # K m^{-1} - Gamma2 = Gamma1 / 2.0 - zc = 10000 # m - p = np.logspace(np.log10(850.0), np.log10(50.0), num=N) - # z_pressure_levels = {0: 850., 1: 500., 2: 250., 3: 50.} - # z_pressure_levels = {i: 850. - 50.*i for i in range(17)} - z_pressure_levels = {i: p[i] for i in range(p.shape[0])} - z_pressure_levels = { - c: p * 100 for c, p in z_pressure_levels.items() - } # convert hPa to Pa - anchor_z_channel = 0 - anchor_T_channel = len(z_pressure_levels) - - # Create HydrostaticBalance constraint object - constraint = HydrostaticBalance( - z_pressure_levels, anchor_z_channel, anchor_T_channel, R, g0 - ) - - x = torch.zeros((1, 12, anchor_T_channel + 1, 1, 1)) - # Set T - T1 = T0 * torch.pow( - torch.Tensor(list(z_pressure_levels.values())) - / z_pressure_levels[anchor_z_channel], - Gamma1 * R / g0, - ) - z1 = z0 + (T0 - T1) / Gamma1 - - Tc = T0 - Gamma1 * (zc - z0) - Pc = z_pressure_levels[anchor_z_channel] * (Tc / T0) ** (g0 / (Gamma1 * R)) - T2 = Tc * torch.pow( - torch.Tensor(list(z_pressure_levels.values())) / Pc, Gamma2 * R / g0 - ) - z2 = zc + (Tc - T2) / Gamma2 - - z = torch.zeros_like(z1) - z = torch.where(z1 < zc, z1, z2) - T = torch.where(z1 < zc, T1, T2) - - x[:, :, 0:anchor_T_channel, :, :] = z.view(1, 1, -1, 1, 1) - # Set Tv at 850 hPa - x[:, :, anchor_T_channel, :, :] = T[0] + torch.zeros_like( - x[:, :, anchor_T_channel, :, :] - ) - - Tv = constraint(x) - - print(Tv[0, 0, :, 0, 0]) - # assert torch.allclose(T * torch.ones_like(Tv), Tv) - return p, z, T, Tv[0, 0, constraint.z_channels, 0, 0] From 46d4942c3ea2408d65e175bfbab483da2ef9af81 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 15 Jul 2026 10:55:36 -0500 Subject: [PATCH 17/28] fix gitignore, update import for healpix_layers --- .gitignore | 1 - .../dlwp_healpix/layers/healpix_blocks.py | 115 +++++++++--------- 2 files changed, 57 insertions(+), 59 deletions(-) diff --git a/.gitignore b/.gitignore index a742871a4a..d999b7f1d4 100644 --- a/.gitignore +++ b/.gitignore @@ -125,7 +125,6 @@ venv/ ENV/ env.bak/ venv.bak/ -.venv-dlesym/ # Spyder project settings .spyderproject diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py index 451193df6b..6448ef3791 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py @@ -32,7 +32,6 @@ from typing import Callable, Sequence, Tuple, Union import torch -import torch as th from physicsnemo.nn.module.hpx import HEALPixLayer @@ -42,7 +41,7 @@ # # Helper: standard LayerNorm over channel dimension for (B, C, H, W) # -class _LayerNormOverChannels(th.nn.Module): +class _LayerNormOverChannels(torch.nn.Module): """Applies nn.LayerNorm over the channel dimension for (B, C, H, W) tensors.""" def __init__(self, channel_depth: int, eps: float = 1e-5): @@ -55,7 +54,7 @@ def __init__(self, channel_depth: int, eps: float = 1e-5): The epsilon value for the layer norm """ super().__init__() - self.norm = th.nn.LayerNorm(channel_depth, eps=eps) + self.norm = torch.nn.LayerNorm(channel_depth, eps=eps) def forward(self, x): """Forward pass of the _LayerNormOverChannels @@ -80,7 +79,7 @@ def forward(self, x): # -class ConvGRUBlock(th.nn.Module): +class ConvGRUBlock(torch.nn.Module): """Class that implements a Convolutional GRU Code modified from https://github.com/happyjin/ConvGRU-pytorch/blob/master/convGRU.py @@ -88,7 +87,7 @@ class ConvGRUBlock(th.nn.Module): def __init__( self, - geometry_layer: th.nn.Module = HEALPixLayer, + geometry_layer: torch.nn.Module = HEALPixLayer, in_channels: int = 3, kernel_size: int = 1, enable_nhwc: bool = False, @@ -129,7 +128,7 @@ def __init__( enable_nhwc=enable_nhwc, enable_healpixpad=enable_healpixpad, ) - self.h = th.zeros(1, 1, 1, 1) + self.h = torch.zeros(1, 1, 1, 1) def forward(self, inputs: Sequence) -> Sequence: """Forward pass of the ConvGRUBlock @@ -145,17 +144,17 @@ def forward(self, inputs: Sequence) -> Sequence: Result of the forward pass """ if inputs.shape != self.h.shape: - self.h = th.zeros_like(inputs) - combined = th.cat([inputs, self.h], dim=1) + self.h = torch.zeros_like(inputs) + combined = torch.cat([inputs, self.h], dim=1) combined_conv = self.conv_gates(combined) - gamma, beta = th.split(combined_conv, self.channels, dim=1) - reset_gate = th.sigmoid(gamma) - update_gate = th.sigmoid(beta) + gamma, beta = torch.split(combined_conv, self.channels, dim=1) + reset_gate = torch.sigmoid(gamma) + update_gate = torch.sigmoid(beta) - combined = th.cat([inputs, reset_gate * self.h], dim=1) + combined = torch.cat([inputs, reset_gate * self.h], dim=1) cc_cnm = self.conv_can(combined) - cnm = th.tanh(cc_cnm) + cnm = torch.tanh(cc_cnm) h_next = (1 - update_gate) * self.h + update_gate * cnm self.h = h_next @@ -164,7 +163,7 @@ def forward(self, inputs: Sequence) -> Sequence: def reset(self): """Reset the update gates""" - self.h = th.zeros_like(self.h) + self.h = torch.zeros_like(self.h) # @@ -172,19 +171,19 @@ def reset(self): # -class BasicConvBlock(th.nn.Module): +class BasicConvBlock(torch.nn.Module): """Convolution block consisting of n subsequent convolutions and activations""" def __init__( self, - geometry_layer: th.nn.Module = HEALPixLayer, + geometry_layer: torch.nn.Module = HEALPixLayer, in_channels: int = 3, out_channels: int = 1, kernel_size: int = 3, dilation: int = 1, n_layers: int = 1, latent_channels: int = None, - activation: th.nn.Module = None, + activation: torch.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, ): @@ -230,7 +229,7 @@ def __init__( ) if activation is not None: convblock.append(activation) - self.convblock = th.nn.Sequential(*convblock) + self.convblock = torch.nn.Sequential(*convblock) def forward(self, x): """Forward pass of the BasicConvBlock @@ -248,14 +247,14 @@ def forward(self, x): return self.convblock(x) -class ConvNeXtBlock(th.nn.Module): +class ConvNeXtBlock(torch.nn.Module): """Class implementing a modified ConvNeXt network as described in https://arxiv.org/pdf/2201.03545.pdf and shown in figure 4 """ def __init__( self, - geometry_layer: th.nn.Module = HEALPixLayer, + geometry_layer: torch.nn.Module = HEALPixLayer, in_channels: int = 3, latent_channels: int = 1, out_channels: int = 1, @@ -263,7 +262,7 @@ def __init__( dilation: int = 1, n_layers: int = 1, # not used, but required for hydra instantiation upscale_factor: int = 4, - activation: th.nn.Module = None, + activation: torch.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, ): @@ -346,7 +345,7 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) - self.convblock = th.nn.Sequential(*convblock) + self.convblock = torch.nn.Sequential(*convblock) def forward(self, x): """Forward pass of the ConvNextBlock @@ -364,7 +363,7 @@ def forward(self, x): return self.skip_module(x) + self.convblock(x) -class DoubleConvNeXtBlock(th.nn.Module): +class DoubleConvNeXtBlock(torch.nn.Module): """Modification of ConvNeXtBlock block this time putting two sequentially in a single block with the number of channels in the middle being the number of latent channels @@ -372,7 +371,7 @@ class DoubleConvNeXtBlock(th.nn.Module): def __init__( self, - geometry_layer: th.nn.Module = HEALPixLayer, + geometry_layer: torch.nn.Module = HEALPixLayer, in_channels: int = 3, out_channels: int = 1, kernel_size: int = 3, @@ -380,7 +379,7 @@ def __init__( n_layers: int = 1, # not used, but required for hydra instantiation upscale_factor: int = 4, latent_channels: int = 1, - activation: th.nn.Module = None, + activation: torch.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, conditional_layer_norm: Callable = None, @@ -484,7 +483,7 @@ def __init__( if activation is not None: convblock1.append(activation) if dropout > 0.0: - convblock1.append(th.nn.Dropout2d(p=dropout)) + convblock1.append(torch.nn.Dropout2d(p=dropout)) # 1x1 convolution establishing increased channels convblock1.append( @@ -515,7 +514,7 @@ def __init__( if activation is not None: convblock1.append(activation) if dropout > 0.0: - convblock1.append(th.nn.Dropout2d(p=dropout)) + convblock1.append(torch.nn.Dropout2d(p=dropout)) # 1x1 convolution returning to latent channels convblock1.append( @@ -532,8 +531,8 @@ def __init__( if activation is not None: convblock1.append(activation) if dropout > 0.0: - convblock1.append(th.nn.Dropout2d(p=dropout)) - self.convblock1 = th.nn.ModuleList(convblock1) + convblock1.append(torch.nn.Dropout2d(p=dropout)) + self.convblock1 = torch.nn.ModuleList(convblock1) # 2nd ConNeXt block, takes the output of the first convnext block convblock2 = [] @@ -557,7 +556,7 @@ def __init__( if activation is not None: convblock2.append(activation) if dropout > 0.0: - convblock2.append(th.nn.Dropout2d(p=dropout)) + convblock2.append(torch.nn.Dropout2d(p=dropout)) # 1x1 convolution establishing increased channels convblock2.append( @@ -587,7 +586,7 @@ def __init__( if activation is not None: convblock2.append(activation) if dropout > 0.0: - convblock2.append(th.nn.Dropout2d(p=dropout)) + convblock2.append(torch.nn.Dropout2d(p=dropout)) # 1x1 convolution reducing to output channels convblock2.append( @@ -604,8 +603,8 @@ def __init__( if activation is not None: convblock2.append(activation) if dropout > 0.0: - convblock2.append(th.nn.Dropout2d(p=dropout)) - self.convblock2 = th.nn.ModuleList(convblock2) + convblock2.append(torch.nn.Dropout2d(p=dropout)) + self.convblock2 = torch.nn.ModuleList(convblock2) def forward(self, x, conditions_cln=None): """Forward pass of the DoubleConvNextBlock @@ -662,14 +661,14 @@ def forward(self, x, conditions_cln=None): return x2_residual + x1 -class Multi_SymmetricConvNeXtBlock(th.nn.Module): +class Multi_SymmetricConvNeXtBlock(torch.nn.Module): """ Wrapper for SymmetricConvNeXtBlock that allows serial linking of blocks. """ def __init__( self, - geometry_layer: th.nn.Module = HEALPixLayer, + geometry_layer: torch.nn.Module = HEALPixLayer, in_channels: int = 3, latent_channels: int = 1, out_channels: int = 1, @@ -677,7 +676,7 @@ def __init__( dilation: int = 1, upscale_factor: int = 4, n_layers: int = 1, - activation: th.nn.Module = None, + activation: torch.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, dropout: float = 0.0, @@ -701,7 +700,7 @@ def __init__( super().__init__() # Create a ModuleList to store complete blocks - self.blocks = th.nn.ModuleList() + self.blocks = torch.nn.ModuleList() # flag for conditional layer normalization self.cln_enabled = conditional_layer_norm is not None @@ -748,7 +747,7 @@ def forward(self, x, conditions_cln=None): return out -class SymmetricConvNeXtBlock(th.nn.Module): +class SymmetricConvNeXtBlock(torch.nn.Module): """Another modification of ConvNeXtBlock block this time using 4 layers and adding a layer that instead of going from in_channels to latent*upscale channesl goes to latent channels first @@ -756,7 +755,7 @@ class SymmetricConvNeXtBlock(th.nn.Module): def __init__( self, - geometry_layer: th.nn.Module = HEALPixLayer, + geometry_layer: torch.nn.Module = HEALPixLayer, in_channels: int = 3, latent_channels: int = 1, out_channels: int = 1, @@ -764,12 +763,12 @@ def __init__( dilation: int = 1, n_layers: int = 1, # not used, but required for hydra instantiation upscale_factor: int = 4, - activation: th.nn.Module = None, + activation: torch.nn.Module = None, enable_nhwc: bool = False, use_block_skip_connection: bool = True, enable_healpixpad: bool = False, dropout: float = 0.0, - conditional_layer_norm: th.nn.Module = None, + conditional_layer_norm: torch.nn.Module = None, conditional_layer_norm_once: bool = False, ): """ @@ -799,7 +798,7 @@ def __init__( Whether or not to use block-level skip connection dropout: float, optional Dropout probability to apply after the first convolution - conditional_layer_norm: th.nn.Module, optional + conditional_layer_norm: torch.nn.Module, optional conditional layer normalization. If None, no conditional layer normalization is applied. conditional_layer_norm_once: bool, optional @@ -820,7 +819,7 @@ def __init__( self.skip_module = lambda x: x else: self.skip_module = geometry_layer( - layer=th.nn.Conv2d, + layer=torch.nn.Conv2d, in_channels=in_channels, out_channels=out_channels, kernel_size=1, @@ -849,7 +848,7 @@ def __init__( # 3x3: in → latent convblock.append( geometry_layer( - layer=th.nn.Conv2d, + layer=torch.nn.Conv2d, in_channels=in_channels, out_channels=int(latent_channels), kernel_size=kernel_size, @@ -866,12 +865,12 @@ def __init__( if activation is not None: convblock.append(activation) if dropout > 0.0: - convblock.append(th.nn.Dropout2d(p=dropout)) + convblock.append(torch.nn.Dropout2d(p=dropout)) # 1x1: latent → latent * upscale convblock.append( geometry_layer( - layer=th.nn.Conv2d, + layer=torch.nn.Conv2d, in_channels=int(latent_channels), out_channels=int(latent_channels * upscale_factor), kernel_size=1, @@ -897,12 +896,12 @@ def __init__( if activation is not None: convblock.append(activation) if dropout > 0.0: - convblock.append(th.nn.Dropout2d(p=dropout)) + convblock.append(torch.nn.Dropout2d(p=dropout)) # 1x1: upscale → latent convblock.append( geometry_layer( - layer=th.nn.Conv2d, + layer=torch.nn.Conv2d, in_channels=int(latent_channels * upscale_factor), out_channels=int(latent_channels), kernel_size=1, @@ -922,12 +921,12 @@ def __init__( if activation is not None: convblock.append(activation) if dropout > 0.0: - convblock.append(th.nn.Dropout2d(p=dropout)) + convblock.append(torch.nn.Dropout2d(p=dropout)) # 3x3: latent → out (no norm on this one, following convnext) convblock.append( geometry_layer( - layer=th.nn.Conv2d, + layer=torch.nn.Conv2d, in_channels=int(latent_channels), out_channels=out_channels, kernel_size=kernel_size, @@ -939,9 +938,9 @@ def __init__( if activation is not None: convblock.append(activation) if dropout > 0.0: - convblock.append(th.nn.Dropout2d(p=dropout)) + convblock.append(torch.nn.Dropout2d(p=dropout)) - self.convblock = th.nn.ModuleList(convblock) + self.convblock = torch.nn.ModuleList(convblock) def forward(self, x, conditions_cln=None): """ @@ -983,18 +982,18 @@ def forward(self, x, conditions_cln=None): # -class TransposedConvUpsample(th.nn.Module): +class TransposedConvUpsample(torch.nn.Module): """This class provides a wrapper for a HEALPix (or other) tensor data around the torch.nn.ConvTranspose2d class. """ def __init__( self, - geometry_layer: th.nn.Module = HEALPixLayer, + geometry_layer: torch.nn.Module = HEALPixLayer, in_channels: int = 3, out_channels: int = 1, upsampling: int = 2, - activation: th.nn.Module = None, + activation: torch.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, ): @@ -1033,7 +1032,7 @@ def __init__( ) if activation is not None: upsampler.append(activation) - self.upsampler = th.nn.Sequential(*upsampler) + self.upsampler = torch.nn.Sequential(*upsampler) def forward(self, x): """Forward pass of the TransposedConvUpsample layer @@ -1056,7 +1055,7 @@ def forward(self, x): # -class Interpolate(th.nn.Module): +class Interpolate(torch.nn.Module): """Helper class that handles interpolation This is done as a class so that scale and mode can be stored """ @@ -1071,7 +1070,7 @@ def __init__(self, scale_factor: Union[int, Tuple], mode: str = "nearest"): Interpolation mode used for upsampling, passed to torch.nn.functional.interpolate """ super().__init__() - self.interp = th.nn.functional.interpolate + self.interp = torch.nn.functional.interpolate self.scale_factor = scale_factor self.mode = mode From 4931876076bae4c621a2aae29921cf6ab22fe6cb Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 15 Jul 2026 10:55:59 -0500 Subject: [PATCH 18/28] Add clarification to per level healpix unet options --- .../models/dlwp_healpix/layers/healpix_decoder.py | 9 +++++++++ .../models/dlwp_healpix/layers/healpix_encoder.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py index 454af91589..e91ce5aa59 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py @@ -76,6 +76,15 @@ def __init__( per_level_checkpointing: list[bool], optional If the checkpointing should be applied to each level of the decoder If None, the checkpointing will not be applied + + per level options are lists of booleans of the same length as n_channels, + if only one value is provided, it will be applied to all levels. The level layout is a + mirror of the encoder. The first value in the list will be applied to the lowest level + in the decoder and the last value in the list will be applied to the highest level in the decoder. + Example: + n_channels = [16, 32, 64] + per_level_cln = [True, False, False] will apply CLN to the lowest level (16 channels) and not the second and third levels. + Not this is the opposite of the encoder. """ super().__init__() self.channel_dim = 1 diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py index bbe874be83..a026e4f443 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py @@ -73,6 +73,15 @@ def __init__( per_level_checkpointing: list[bool] | None, optional If the checkpointing should be applied to each level of the encoder If None, the checkpointing will not be applied + + per level options are lists of booleans of the same length as n_channels, + if only one value is provided, it will be applied to all levels. The first + level is the highest (top) level in the unet encoder and the last level is the lowest (bottom) level. + The value in the list will be applied to the top level, the second value to the next level down, etc. + Example: + n_channels = [16, 32, 64] + per_level_cln = [True, False, False] will apply CLN to the highest level (16 channels) and not the second and third levels. + """ super().__init__() self.n_channels = n_channels From 2c41b4dd14ed41ebd38c010a10c00f76e59959c2 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 15 Jul 2026 12:12:31 -0500 Subject: [PATCH 19/28] remove checkpoint conversion for experimental models --- .../models/dlwp_healpix/HEALPixRecUNet.py | 6 - .../models/dlwp_healpix/HEALPixUNet.py | 6 - .../dlwp_healpix/layers/normalization.py | 158 +-------------- .../test_conditional_layer_norm.py | 188 ++---------------- 4 files changed, 21 insertions(+), 337 deletions(-) diff --git a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py index 6bacf565f3..6204d567db 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py @@ -113,12 +113,6 @@ class HEALPixRecUNet(Module): "0.1.0": _legacy_hydra_targets_warning, } - # Allows ``Module.from_checkpoint(..., override_args=...)`` to substitute - # a corrected ``encoder``/``decoder`` config. Used by checkpoint conversion - # utilities to fix up legacy checkpoints (e.g. the ``per_level_cln`` - # decoder bug) without having to hand-edit the model source. - _overridable_args: set = {"encoder", "decoder"} - @classmethod def _backward_compat_arg_mapper( cls, version: str, args: Dict[str, Any] diff --git a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py index 479d0d5a82..6ec2c7903a 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py @@ -108,12 +108,6 @@ class HEALPixUNet(Module): "0.1.0": _legacy_hydra_targets_warning, } - # Allows ``Module.from_checkpoint(..., override_args=...)`` to substitute - # a corrected ``encoder``/``decoder`` config. Used by checkpoint conversion - # utilities to fix up legacy checkpoints (e.g. the ``per_level_cln`` - # decoder bug) without having to hand-edit the model source. - _overridable_args: set = {"encoder", "decoder"} - @classmethod def _backward_compat_arg_mapper( cls, version: str, args: Dict[str, Any] diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index 2c1147445a..af5241be22 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -20,23 +20,16 @@ This class contains the implementation of the Deep Learning Weather Prediction (DLWP) normalization on the HEALPix mesh. """ -import importlib from typing import List import torch as th -# ``apex`` is an optional accelerator (not a core dependency). Import it -# dynamically so the physicsnemo import graph stays free of a static ``apex`` -# dependency (see CODING_STANDARDS/EXTERNAL_IMPORTS.md, EXT-003). The -# ``FusedLayerNorm`` symbol is only used when ``norm_op="apex"`` is explicitly -# requested, which is guarded by ``_APEX_AVAILABLE`` below. -try: - FusedLayerNorm = importlib.import_module("apex.normalization").FusedLayerNorm +from physicsnemo.core.version_check import OptionalImport - _APEX_AVAILABLE = True -except ImportError: # pragma: no cover - only exercised when apex is absent - FusedLayerNorm = None - _APEX_AVAILABLE = False +# ``apex`` is an optional accelerator (not a core dependency). Import it lazily +# via ``OptionalImport``. If ``apex`` is missing, ``OptionalImport`` raises +# an ``ImportError`` with an install hint at that point. +apex_normalization = OptionalImport("apex.normalization") @th.compile @@ -122,11 +115,9 @@ def __init__( if norm_op == "torch": self.norm = th.nn.LayerNorm(channel_depth, elementwise_affine=False) elif norm_op == "apex": - if not _APEX_AVAILABLE: - raise ImportError( - "Apex FusedLayerNorm requested but apex is not available, please install it from https://github.com/NVIDIA/apex" - ) - self.norm = FusedLayerNorm(channel_depth, elementwise_affine=False) + self.norm = apex_normalization.FusedLayerNorm( + channel_depth, elementwise_affine=False + ) def _make_mlp( self, @@ -162,139 +153,6 @@ def _make_mlp( layers.append(th.nn.Linear(in_dim, out_dim)) return th.nn.Sequential(*layers) - def _load_from_state_dict( - self, - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, - ): - """Backward compatibility: merge old separate gamma_mlp/beta_mlp into fused gamma_beta_mlp. - - Old MLPs had hidden_dims [h1, h2, ...] and output dim C. - New fused MLP has hidden_dims [2*h1, 2*h2, ...] and output dim 2*C. - - For the first Linear (condition_shape → 2*h1), we vertically concatenate: - new_weight = cat([gamma_weight, beta_weight], dim=0) - - For subsequent Linear layers (2*h_i → 2*h_{i+1} or 2*h_last → 2*C), - we build a block-diagonal weight matrix: - new_weight = [[gamma_weight, 0 ], - [0, beta_weight]] - - Biases are always concatenated: cat([gamma_bias, beta_bias]). - - Activation submodule buffers (e.g. CappedGELU ``cap``) are copied from - the gamma MLP to the fused MLP at the same sequential index. - """ - gamma_prefix = prefix + "gamma_mlp." - beta_prefix = prefix + "beta_mlp." - fused_prefix = prefix + "gamma_beta_mlp." - - has_old_keys = any(k.startswith(gamma_prefix) for k in state_dict) - - if has_old_keys: - # Collect all Linear layer indices from the old gamma MLP - gamma_layer_indices = set() - for k in state_dict: - if k.startswith(gamma_prefix): - layer_key = k[len(gamma_prefix) :] - parts = layer_key.split(".") - if len(parts) == 2 and parts[1] in ("weight", "bias"): - gamma_layer_indices.add(int(parts[0])) - first_layer_idx = min(gamma_layer_indices) - - keys_to_remove = [] - keys_to_add = {} - - # handle weights and biases for the gamma MLP - for k in list(state_dict.keys()): - if k.startswith(gamma_prefix): - layer_key = k[len(gamma_prefix) :] # e.g. "0.weight" - parts = layer_key.split(".") - if len(parts) != 2 or parts[1] not in ("weight", "bias"): - continue - idx = int(parts[0]) - param_type = parts[1] - fused_key = fused_prefix + layer_key - beta_key = beta_prefix + layer_key - - if beta_key not in state_dict: - continue - - gamma_val = state_dict[k] - beta_val = state_dict[beta_key] - - if param_type == "bias": - # Biases are always concatenated - keys_to_add[fused_key] = th.cat([gamma_val, beta_val], dim=0) - elif idx == first_layer_idx: - # First layer: input dim is shared (condition_shape), - # just concatenate along output dim - keys_to_add[fused_key] = th.cat([gamma_val, beta_val], dim=0) - else: - # Hidden→hidden or hidden→output: block-diagonal - # gamma_val: (out_old, in_old), beta_val: (out_old, in_old) - # result: (2*out_old, 2*in_old) - out_old, in_old = gamma_val.shape - zeros = th.zeros( - out_old, - in_old, - dtype=gamma_val.dtype, - device=gamma_val.device, - ) - keys_to_add[fused_key] = th.cat( - [ - th.cat([gamma_val, zeros], dim=1), - th.cat([zeros, beta_val], dim=1), - ], - dim=0, - ) - - keys_to_remove.append(k) - if beta_key not in keys_to_remove: - keys_to_remove.append(beta_key) - - # handle activation submodule buffers - for k in list(state_dict.keys()): - if not k.startswith(gamma_prefix): - continue - layer_key = k[len(gamma_prefix) :] # e.g. "0.weight" - parts = layer_key.split(".", 1) - if len(parts) != 2: - continue - try: - int(parts[0]) - except ValueError: - continue - if parts[1] in ("weight", "bias"): - continue - - fused_key = fused_prefix + layer_key - keys_to_add[fused_key] = state_dict[k] - keys_to_remove.append(k) - beta_key = beta_prefix + layer_key - if beta_key in state_dict and beta_key not in keys_to_remove: - keys_to_remove.append(beta_key) - - for k in keys_to_remove: - if k in state_dict: - del state_dict[k] - state_dict.update(keys_to_add) - - super()._load_from_state_dict( - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, - ) - def forward(self, x: th.Tensor, conditions: th.Tensor) -> th.Tensor: """ Parameters diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py index da7bf6615b..e5a3913b06 100644 --- a/test/models/dlwp_healpix/test_conditional_layer_norm.py +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib.util import os import sys @@ -24,11 +25,7 @@ from _cln_reference import ConditionalLayerNormReference from physicsnemo.models.dlwp_healpix.layers import normalization -from physicsnemo.models.dlwp_healpix.layers.normalization import ( - _APEX_AVAILABLE, - ConditionalLayerNorm, -) -from physicsnemo.nn import CappedGELU +from physicsnemo.models.dlwp_healpix.layers.normalization import ConditionalLayerNorm def _make_old_cln(condition_shape, channel_depth, device="cpu", **kwargs): @@ -110,9 +107,7 @@ def _copy_old_to_new(old_cln, new_cln): def _assert_forward_is_finite( cln, channel_depth, condition_shape, device, n_cond=1, n_faces=12, height=4, width=4 ): - """Run a minimal forward pass and assert the output is finite; used by - the legacy-checkpoint-loading tests to confirm a (possibly - partially-loaded) module still produces a usable output.""" + """Run a minimal forward pass and assert the output is finite.""" x = torch.randn(n_cond * n_faces, channel_depth, height, width, device=device) cond = torch.randn(n_cond, condition_shape, device=device) out = cln(x, cond) @@ -298,71 +293,6 @@ def test_backward_channels_last_matches_contiguous(device): ) -def test_load_old_checkpoint(device): - """Verify new CLN can load old-format state dict via _load_from_state_dict.""" - C, cond_shape = 64, 16 - torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C, device=device) - old_sd = old_cln.state_dict() - - new_cln = _make_new_cln(cond_shape, C, device=device) - new_cln.load_state_dict(old_sd, strict=False) - - # Verify outputs match after loading old checkpoint - x = torch.randn(12, C, 8, 8, device=device) - cond = torch.randn(1, cond_shape, device=device) - - with torch.no_grad(): - out_old = old_cln(x, cond) - out_new = new_cln(x, cond) - - assert out_new.shape == (12, C, 8, 8) - assert torch.isfinite(out_new).all() - assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( - f"Max diff after loading old checkpoint: {(out_old - out_new).abs().max().item()}" - ) - - -def test_load_old_checkpoint_with_capped_gelu(device): - """Verify strict loading of old CLN checkpoints that use CappedGELU activations.""" - C, cond_shape = 64, 16 - mlp_hidden_dims = [32, 32] - activation = CappedGELU() - - torch.manual_seed(42) - old_cln = _make_old_cln( - cond_shape, - C, - device=device, - mlp_hidden_dims=mlp_hidden_dims, - activation=activation, - ) - old_sd = old_cln.state_dict() - assert any(k.endswith(".cap") for k in old_sd if k.startswith("gamma_mlp.")) - - new_cln = _make_new_cln( - cond_shape, - C, - device=device, - mlp_hidden_dims=mlp_hidden_dims, - activation=CappedGELU(), - ) - missing, unexpected = new_cln.load_state_dict(old_sd, strict=True) - assert not missing - assert not unexpected - - x = torch.randn(12, C, 8, 8, device=device) - cond = torch.randn(1, cond_shape, device=device) - - with torch.no_grad(): - out_old = old_cln(x, cond) - out_new = new_cln(x, cond) - - assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( - f"Max diff after loading old CappedGELU checkpoint: {(out_old - out_new).abs().max().item()}" - ) - - @pytest.mark.parametrize("hidden_dims", [[], [64]]) def test_old_vs_new_forward_hidden_dims_variation(device, hidden_dims): """Verify the block-diagonal weight mapping generalizes to MLP depths other @@ -394,12 +324,13 @@ def test_old_vs_new_forward_hidden_dims_variation(device, hidden_dims): def test_norm_op_apex_raises_when_unavailable(): """``norm_op='apex'`` must raise an informative ImportError when the ``apex`` package isn't installed, rather than silently falling back or - failing with an unrelated error later on. + failing with an unrelated error later on. The error is raised by + ``OptionalImport`` when the fused-norm symbol is accessed. """ - if _APEX_AVAILABLE: + if importlib.util.find_spec("apex") is not None: pytest.skip("apex is installed in this environment") - with pytest.raises(ImportError, match="Apex FusedLayerNorm requested"): + with pytest.raises(ImportError, match="Missing optional dependency: apex"): ConditionalLayerNorm(condition_shape=16, channel_depth=8, norm_op="apex") @@ -415,7 +346,8 @@ def test_norm_op_invalid_value_leaves_norm_unset(): def test_norm_op_apex_available_uses_fused_layer_norm(monkeypatch): """When apex *is* available, ``norm_op='apex'`` must construct a ``FusedLayerNorm`` instead of ``torch.nn.LayerNorm``. Exercised via a - stand-in class since apex isn't a hard dependency of this environment. + stand-in ``apex.normalization`` module since apex isn't a hard dependency + of this environment. """ class _FakeFusedLayerNorm(torch.nn.Module): @@ -427,10 +359,10 @@ def __init__(self, channel_depth, elementwise_affine=False): def forward(self, x): return x - monkeypatch.setattr(normalization, "_APEX_AVAILABLE", True) - monkeypatch.setattr( - normalization, "FusedLayerNorm", _FakeFusedLayerNorm, raising=False - ) + class _FakeApexNormalization: + FusedLayerNorm = _FakeFusedLayerNorm + + monkeypatch.setattr(normalization, "apex_normalization", _FakeApexNormalization) cln = ConditionalLayerNorm(condition_shape=16, channel_depth=8, norm_op="apex") assert isinstance(cln.norm, _FakeFusedLayerNorm) @@ -486,100 +418,6 @@ def test_cln_affine_eager_matches_compiled(monkeypatch, device): ) -def test_load_from_state_dict_ignores_malformed_activation_buffer_keys(device): - """The activation-buffer migration loop must gracefully skip legacy keys - that don't fit the expected ``"."`` shape (no dot, or a - non-numeric index), rather than raising. - """ - C, cond_shape = 16, 8 - old_cln = _make_old_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) - old_sd = old_cln.state_dict() - - # no dot after the gamma_mlp prefix at all - old_sd["gamma_mlp.malformed"] = torch.zeros(1) - # dot present, but the leading segment isn't a valid layer index - old_sd["gamma_mlp.bogus.cap"] = torch.zeros(1) - - new_cln = _make_new_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) - missing, unexpected = new_cln.load_state_dict(old_sd, strict=False) - - # both malformed keys are left untouched (never merged into the fused - # MLP), so they remain unexpected for the new module - assert "gamma_mlp.malformed" in unexpected - assert "gamma_mlp.bogus.cap" in unexpected - - _assert_forward_is_finite(new_cln, C, cond_shape, device) - - -def test_load_from_state_dict_activation_buffer_without_beta_counterpart(device): - """An activation submodule buffer (e.g. ``CappedGELU``'s ``cap``) present - on the gamma MLP but missing from the beta MLP (malformed/partial - checkpoint) must still migrate the gamma-side buffer, and must not raise - trying to remove the (absent) beta-side key. - """ - C, cond_shape = 16, 8 - mlp_hidden_dims = [8] - old_cln = _make_old_cln( - cond_shape, - C, - device=device, - mlp_hidden_dims=mlp_hidden_dims, - activation=CappedGELU(), - ) - old_sd = old_cln.state_dict() - assert "gamma_mlp.1.cap" in old_sd - assert "beta_mlp.1.cap" in old_sd - - # drop only the beta-side activation buffer - del old_sd["beta_mlp.1.cap"] - - new_cln = _make_new_cln( - cond_shape, - C, - device=device, - mlp_hidden_dims=mlp_hidden_dims, - activation=CappedGELU(), - ) - missing, unexpected = new_cln.load_state_dict(old_sd, strict=False) - - # the gamma-side buffer was still migrated into the fused MLP - assert "gamma_beta_mlp.1.cap" not in unexpected - assert "gamma_mlp.1.cap" not in unexpected - - _assert_forward_is_finite(new_cln, C, cond_shape, device) - - -def test_load_from_state_dict_skips_unmatched_gamma_key(device): - """A legacy state dict where a ``gamma_mlp`` key has no matching - ``beta_mlp`` counterpart (e.g. a malformed/partial checkpoint) must be - skipped rather than merged, and loading must not raise. - """ - C, cond_shape = 32, 16 - torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) - old_sd = old_cln.state_dict() - - # drop the beta counterpart for the (only) gamma layer, simulating a - # malformed/partial legacy checkpoint - del old_sd["beta_mlp.0.weight"] - del old_sd["beta_mlp.0.bias"] - - new_cln = _make_new_cln(cond_shape, C, device=device, mlp_hidden_dims=[]) - missing, unexpected = new_cln.load_state_dict(old_sd, strict=False) - - # the unmatched gamma key is left untouched (not merged into - # gamma_beta_mlp), so it shows up as unexpected for the new module, and - # the fused MLP's own parameters are all reported missing since nothing - # could be merged - assert "gamma_mlp.0.weight" in unexpected - assert "gamma_mlp.0.bias" in unexpected - assert any(k.startswith("gamma_beta_mlp.") for k in missing) - - # forward should still run without error using the (unloaded, default - # initialized) fused MLP weights - _assert_forward_is_finite(new_cln, C, cond_shape, device) - - def test_load_new_format_checkpoint_roundtrip(device): """A state dict already in the fused ``gamma_beta_mlp`` format (i.e. saved from a ``ConditionalLayerNorm`` instance, with no legacy ``gamma_mlp``/ From 130533d768de19a9c1a401dcf10be1f1dc9d820d Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 15 Jul 2026 12:20:04 -0500 Subject: [PATCH 20/28] update docstrings with new options --- physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py | 8 ++++++++ physicsnemo/models/dlwp_healpix/HEALPixUNet.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py index 6204d567db..1c6267a921 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py @@ -92,6 +92,14 @@ class HEALPixRecUNet(Module): Enable CUDA HEALPix padding when available. couplings : list, optional Optional coupling specifications appended to the input feature channels. + residual_prediction : bool, optional + If ``True``, the model will predict the residual of the input and the output + and add it to the output. If ``False``, the model will predict the output directly. + couplings_time_first : bool, optional + If ``True``, the couplings will be passed to the model in the time dimension first. + If ``False``, the couplings will be passed to the model in the channel dimension first. + constraints : list[DictConfig], optional + Optional constraints to be applied to the model outputs. Forward ------- diff --git a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py index 6ec2c7903a..4004b84fba 100644 --- a/physicsnemo/models/dlwp_healpix/HEALPixUNet.py +++ b/physicsnemo/models/dlwp_healpix/HEALPixUNet.py @@ -85,6 +85,14 @@ class HEALPixUNet(Module): Enable CUDA HEALPix padding when the optional dependency is installed. couplings : list, optional Optional coupling specifications appended to the input feature channels. + residual_prediction : bool, optional + If ``True``, the model will predict the residual of the input and the output + and add it to the output. If ``False``, the model will predict the output directly. + couplings_time_first : bool, optional + If ``True``, the couplings will be passed to the model in the time dimension first. + If ``False``, the couplings will be passed to the model in the channel dimension first. + constraints : list[DictConfig], optional + Optional constraints to be applied to the model outputs. Forward ------- From b5ea37117f0154d5523d532951806f745eefa36b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:22:12 +0000 Subject: [PATCH 21/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py | 4 ++-- physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py index e91ce5aa59..68f70b896c 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py @@ -76,9 +76,9 @@ def __init__( per_level_checkpointing: list[bool], optional If the checkpointing should be applied to each level of the decoder If None, the checkpointing will not be applied - + per level options are lists of booleans of the same length as n_channels, - if only one value is provided, it will be applied to all levels. The level layout is a + if only one value is provided, it will be applied to all levels. The level layout is a mirror of the encoder. The first value in the list will be applied to the lowest level in the decoder and the last value in the list will be applied to the highest level in the decoder. Example: diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py index a026e4f443..dee98855c0 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py @@ -75,13 +75,13 @@ def __init__( If None, the checkpointing will not be applied per level options are lists of booleans of the same length as n_channels, - if only one value is provided, it will be applied to all levels. The first + if only one value is provided, it will be applied to all levels. The first level is the highest (top) level in the unet encoder and the last level is the lowest (bottom) level. The value in the list will be applied to the top level, the second value to the next level down, etc. Example: n_channels = [16, 32, 64] per_level_cln = [True, False, False] will apply CLN to the highest level (16 channels) and not the second and third levels. - + """ super().__init__() self.n_channels = n_channels From 45712e7c9d83bf730734ed325e6f4ea9c1097025 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Wed, 15 Jul 2026 12:33:54 -0500 Subject: [PATCH 22/28] remove 'import torch as th' idiom --- .../dlwp_healpix/layers/normalization.py | 40 ++++++++++--------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/physicsnemo/models/dlwp_healpix/layers/normalization.py b/physicsnemo/models/dlwp_healpix/layers/normalization.py index af5241be22..e1819a2b60 100644 --- a/physicsnemo/models/dlwp_healpix/layers/normalization.py +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -22,7 +22,7 @@ from typing import List -import torch as th +import torch from physicsnemo.core.version_check import OptionalImport @@ -32,7 +32,7 @@ apex_normalization = OptionalImport("apex.normalization") -@th.compile +@torch.compile def _cln_affine(x_norm, gamma_raw, beta, scale_center, n_faces): """Fused affine transform: expand gamma/beta across faces and apply to normalized input.""" C = gamma_raw.shape[-1] @@ -46,7 +46,7 @@ def _cln_affine(x_norm, gamma_raw, beta, scale_center, n_faces): return gamma * x_norm + beta -class ConditionalLayerNorm(th.nn.Module): +class ConditionalLayerNorm(torch.nn.Module): """LayerNorm whose affine (gamma/beta) parameters are predicted from a conditional input field. Normalizes the input over the channel dimension with no learnable affine @@ -60,7 +60,7 @@ def __init__( condition_shape: int, channel_depth: int, mlp_hidden_dims: List[int] = [128, 128], - activation: th.nn.Module = None, + activation: torch.nn.Module = None, eps: float = 1e-5, n_faces: int = 12, norm_op: str = "torch", @@ -98,7 +98,7 @@ def __init__( self.condition_shape = condition_shape self.channel_depth = channel_depth self.hidden_dims = mlp_hidden_dims - self.activation = activation if activation is not None else th.nn.Identity() + self.activation = activation if activation is not None else torch.nn.Identity() self.gamma_beta_mlp = self._make_mlp( self.condition_shape, [2 * h for h in self.hidden_dims], @@ -113,7 +113,7 @@ def __init__( self.gamma_beta_mlp[-1].bias.data.zero_() if norm_op == "torch": - self.norm = th.nn.LayerNorm(channel_depth, elementwise_affine=False) + self.norm = torch.nn.LayerNorm(channel_depth, elementwise_affine=False) elif norm_op == "apex": self.norm = apex_normalization.FusedLayerNorm( channel_depth, elementwise_affine=False @@ -124,8 +124,8 @@ def _make_mlp( in_dim: int, hidden_dims: List[int], out_dim: int, - activation: th.nn.Module, - ) -> th.nn.Sequential: + activation: torch.nn.Module, + ) -> torch.nn.Sequential: """Helper function that creates the MLP for the conditional layer normalization. Parameters @@ -136,39 +136,39 @@ def _make_mlp( The hidden dimensions out_dim: int The output dimension - activation: th.nn.Module + activation: torch.nn.Module The activation function Returns ------- - th.nn.Sequential + torch.nn.Sequential The MLP """ layers = [] for hdim in hidden_dims: - layers.append(th.nn.Linear(in_dim, hdim)) + layers.append(torch.nn.Linear(in_dim, hdim)) if activation: layers.append(activation) in_dim = hdim - layers.append(th.nn.Linear(in_dim, out_dim)) - return th.nn.Sequential(*layers) + layers.append(torch.nn.Linear(in_dim, out_dim)) + return torch.nn.Sequential(*layers) - def forward(self, x: th.Tensor, conditions: th.Tensor) -> th.Tensor: + def forward(self, x: torch.Tensor, conditions: torch.Tensor) -> torch.Tensor: """ Parameters ---------- - x : th.Tensor + x : torch.Tensor Input tensor of shape: (B, C, H, W) - conditions : th.Tensor + conditions : torch.Tensor Conditioning tensor of shape (B*n_cond, cond_dim) Returns ------- - th.Tensor + torch.Tensor Normalized and conditioned tensor of shape: (B, C, H, W) """ - is_channels_last = x.is_contiguous(memory_format=th.channels_last) + is_channels_last = x.is_contiguous(memory_format=torch.channels_last) # LayerNorm on last dim: permute to (B, H, W, C) x_nhwc = x.permute(0, 2, 3, 1) @@ -185,6 +185,8 @@ def forward(self, x: th.Tensor, conditions: th.Tensor) -> th.Tensor: # Return to NCHW logical layout, preserving channels_last memory format if input was if is_channels_last: - return result.permute(0, 3, 1, 2).contiguous(memory_format=th.channels_last) + return result.permute(0, 3, 1, 2).contiguous( + memory_format=torch.channels_last + ) else: return result.permute(0, 3, 1, 2) From c891eb7e53d281b1ef97ce494d3fdf4c4bd0cc01 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Fri, 17 Jul 2026 10:40:46 -0500 Subject: [PATCH 23/28] Remove experimental features and notes --- physicsnemo/datapipes/healpix/couplers.py | 11 +- .../dlwp_healpix/layers/healpix_blocks.py | 141 +----------------- .../dlwp_healpix/test_healpix_blocks.py | 100 ------------- 3 files changed, 13 insertions(+), 239 deletions(-) diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index a0960b2e6d..0b693e27cc 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -201,6 +201,8 @@ def set_coupled_fields(self, coupled_fields: th.tensor): f" match configured batch size {self.batch_size}" ) # create buffer for coupling + # coupled_channel_indices covers selecting multiple instances of the same variable, + # e.g. z1000-24H and z1000-48H. coupled_fields = coupled_fields[ :, :, :, self.coupled_channel_indices, :, : ].permute(2, 0, 3, 1, 4, 5) @@ -208,11 +210,10 @@ def set_coupled_fields(self, coupled_fields: th.tensor): [self.coupled_integration_dim, self.batch_size, self.timevar_dim] + list(self.spatial_dims) ) - # we use a constant set of values so we just copy time 0 - for i in range(len(self.preset_coupled_fields)): - self.preset_coupled_fields[i, :, :, :, :, :] = coupled_fields[ - 0, :, -1, :, :, : - ] + + # we use a constant set of values so we just broadcast time 0 + self.preset_coupled_fields[:, :, :, :, :, :] = coupled_fields[:1, :, :, :, :, :] + # flag for construct integrated coupling method to use this array self.coupled_mode = True diff --git a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py index 6448ef3791..db45530d21 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py @@ -37,43 +37,6 @@ from .normalization import ConditionalLayerNorm - -# -# Helper: standard LayerNorm over channel dimension for (B, C, H, W) -# -class _LayerNormOverChannels(torch.nn.Module): - """Applies nn.LayerNorm over the channel dimension for (B, C, H, W) tensors.""" - - def __init__(self, channel_depth: int, eps: float = 1e-5): - """ - Parameters - ---------- - channel_depth: int - The number of channels in the input tensor - eps: float, optional - The epsilon value for the layer norm - """ - super().__init__() - self.norm = torch.nn.LayerNorm(channel_depth, eps=eps) - - def forward(self, x): - """Forward pass of the _LayerNormOverChannels - - Parameters - ---------- - x: torch.Tensor - The input tensor - - Returns - ------- - torch.Tensor - The normed output tensor - """ - x = x.permute(0, 2, 3, 1) - x = self.norm(x) - return x.permute(0, 3, 1, 2) - - # # RECURRENT BLOCKS # @@ -383,7 +346,6 @@ def __init__( enable_nhwc: bool = False, enable_healpixpad: bool = False, conditional_layer_norm: Callable = None, - conditional_layer_norm_once: bool = False, dropout: float = 0.0, ): """ @@ -441,25 +403,6 @@ def __init__( enable_healpixpad=enable_healpixpad, ) - # check if we're applying a layer norm at the beginning - # we've got two ConvNeXt equivalent blocks in the layer, so have two entry points - # TODO: conditional layer norm once is doing two things, it's applying a norm on block entry - # and switch from conditional to non-conditional layer norm. This is not ideal and should be fixed once we determine - # what works best. - if conditional_layer_norm_once: - if conditional_layer_norm is not None: - # Conditional norm at the beginning of the block - self.entry_norm1 = conditional_layer_norm(channel_depth=in_channels) - self.entry_norm2 = conditional_layer_norm(channel_depth=latent_channels) - else: - # Regular layer normalization at the beginning of the block - self.entry_norm1 = _LayerNormOverChannels(channel_depth=in_channels) - self.entry_norm2 = _LayerNormOverChannels(channel_depth=latent_channels) - else: - # No normalization at the beginning - self.entry_norm1 = None - self.entry_norm2 = None - # 1st ConvNeXt block, the output of this one remains internal convblock1 = [] # 3x3 convolution establishing latent channels channels @@ -476,7 +419,7 @@ def __init__( ) # Apply batch norm and conditional layer norm if needed - if conditional_layer_norm is not None and not conditional_layer_norm_once: + if conditional_layer_norm is not None: cln = conditional_layer_norm(channel_depth=int(latent_channels)) convblock1.append(cln) @@ -499,13 +442,7 @@ def __init__( ) # Apply layer norm if needed - if conditional_layer_norm_once: - convblock1.append( - _LayerNormOverChannels( - channel_depth=int(latent_channels * upscale_factor) - ) - ) - elif conditional_layer_norm is not None: + if conditional_layer_norm is not None: cln = conditional_layer_norm( channel_depth=int(latent_channels * upscale_factor) ) @@ -549,7 +486,7 @@ def __init__( ) ) # Apply batch norm and conditional layer norm if needed - if conditional_layer_norm is not None and not conditional_layer_norm_once: + if conditional_layer_norm is not None: cln = conditional_layer_norm(channel_depth=int(latent_channels)) convblock2.append(cln) @@ -571,13 +508,7 @@ def __init__( ) ) # Apply layer norm if needed - if conditional_layer_norm_once: - convblock2.append( - _LayerNormOverChannels( - channel_depth=int(latent_channels * upscale_factor) - ) - ) - elif conditional_layer_norm is not None: + if conditional_layer_norm is not None: cln = conditional_layer_norm( channel_depth=int(latent_channels * upscale_factor) ) @@ -622,18 +553,9 @@ def forward(self, x, conditions_cln=None): conditions for the conditional layer normalization """ - # TODO: performance of skip connectioni hasn't been compared - # check cln(x) vs. cln(x1_residual + x) in the future # save residual for the first block x1_residual = self.skip_module1(x) - # entry norm for the first block - if self.entry_norm1 is not None: - if conditions_cln is not None: - x = self.entry_norm1(x, conditions=conditions_cln) - else: - x = self.entry_norm1(x) - # internal convnext result for layer in self.convblock1: if isinstance(layer, ConditionalLayerNorm): @@ -645,13 +567,6 @@ def forward(self, x, conditions_cln=None): # save residual for the second block x2_residual = self.skip_module2(x1) - # entry norm for the second block - if self.entry_norm2 is not None: - if conditions_cln is not None: - x1 = self.entry_norm2(x1, conditions=conditions_cln) - else: - x1 = self.entry_norm2(x1) - # return second convnext result for layer in self.convblock2: if isinstance(layer, ConditionalLayerNorm): @@ -681,7 +596,6 @@ def __init__( enable_healpixpad: bool = False, dropout: float = 0.0, conditional_layer_norm: Callable = None, - conditional_layer_norm_once: bool = False, ): """ Parameters @@ -692,10 +606,6 @@ def __init__( Callable for physicsnemo.models.dlwp_healpix_layers.normalization.ConditionalLayerNorm. Callable can be passed in by setting _partial_ to True in hydra config. If None, conditional layer normalization is not applied. - conditional_layer_norm_once: bool, optional - Whether or not to apply conditional layer normalization only once. If True, - the conditional layer normalization is applied only once, otherwise it is applied - for each block. """ super().__init__() @@ -722,7 +632,6 @@ def __init__( conditional_layer_norm=conditional_layer_norm if conditional_layer_norm is not None else None, - conditional_layer_norm_once=conditional_layer_norm_once, ), ) @@ -769,7 +678,6 @@ def __init__( enable_healpixpad: bool = False, dropout: float = 0.0, conditional_layer_norm: torch.nn.Module = None, - conditional_layer_norm_once: bool = False, ): """ Parameters @@ -801,10 +709,6 @@ def __init__( conditional_layer_norm: torch.nn.Module, optional conditional layer normalization. If None, no conditional layer normalization is applied. - conditional_layer_norm_once: bool, optional - Whether or not to apply conditional layer normalization only once. If True, - the conditional layer normalization is applied only once, otherwise it is applied - for each block. """ super().__init__() @@ -827,21 +731,6 @@ def __init__( enable_healpixpad=enable_healpixpad, ) - # check if we're applying a layer norm at the beginning - # TODO: conditional layer norm once is doing two things, it's applying a norm on block entry - # and switch from conditional to non-conditional layer norm. This is not ideal and should be fixed once we determine - # what works best. - if conditional_layer_norm_once: - if conditional_layer_norm is not None: - # Conditional norm at the beginning of the block - self.entry_norm = conditional_layer_norm(channel_depth=in_channels) - else: - # Regular layer normalization at the beginning of the block - self.entry_norm = _LayerNormOverChannels(channel_depth=in_channels) - else: - # No normalization at the beginning - self.entry_norm = None - # Collect conv->norm->activation->dropout operations in list for sequential execution convblock = [] @@ -858,7 +747,7 @@ def __init__( ) ) # Apply layer norm if needed - if conditional_layer_norm is not None and not conditional_layer_norm_once: + if conditional_layer_norm is not None: cln = conditional_layer_norm(channel_depth=int(latent_channels)) convblock.append(cln) @@ -881,13 +770,7 @@ def __init__( ) # Apply layer norm if needed - if conditional_layer_norm_once: - convblock.append( - _LayerNormOverChannels( - channel_depth=int(latent_channels * upscale_factor) - ) - ) - elif conditional_layer_norm is not None: + if conditional_layer_norm is not None: cln = conditional_layer_norm( channel_depth=int(latent_channels * upscale_factor) ) @@ -912,9 +795,7 @@ def __init__( ) # Apply layer norm if needed - if conditional_layer_norm_once: - convblock.append(_LayerNormOverChannels(channel_depth=int(latent_channels))) - elif conditional_layer_norm is not None: + if conditional_layer_norm is not None: cln = conditional_layer_norm(channel_depth=int(latent_channels)) convblock.append(cln) @@ -957,17 +838,9 @@ def forward(self, x, conditions_cln=None): result of the forward pass """ - # TODO: performance of skip connectioni hasn't been compared - # check cln(x) vs. cln(x1_residual + x) in the future # Save residual residual = self.skip_module(x) if self.use_block_skip_connection else 0 - if self.entry_norm is not None: - if conditions_cln is not None: - x = self.entry_norm(x, conditions=conditions_cln) - else: - x = self.entry_norm(x) - for layer in self.convblock: if isinstance(layer, ConditionalLayerNorm): x = layer(x, conditions=conditions_cln) diff --git a/test/models/dlwp_healpix/test_healpix_blocks.py b/test/models/dlwp_healpix/test_healpix_blocks.py index 15405cba93..88385abb3a 100644 --- a/test/models/dlwp_healpix/test_healpix_blocks.py +++ b/test/models/dlwp_healpix/test_healpix_blocks.py @@ -270,7 +270,6 @@ def test_DoubleConvNeXtBlock_conditional_layer_norm(device, test_data, pytestcon out_channels=out_channels, latent_channels=latent_channels, conditional_layer_norm=conditional_layer_norm, - conditional_layer_norm_once=False, ).to(device) assert doubleconvnextblock.cln_enabled @@ -289,53 +288,6 @@ def test_DoubleConvNeXtBlock_conditional_layer_norm(device, test_data, pytestcon assert not common.compare_output(outvar_a, outvar_b) -@pytest.mark.parametrize("conditional_layer_norm_enabled", [True, False]) -def test_DoubleConvNeXtBlock_conditional_layer_norm_once( - device, test_data, pytestconfig, conditional_layer_norm_enabled -): - from physicsnemo.models.dlwp_healpix.layers import ( - DoubleConvNeXtBlock, - ) - - in_channels = 2 - out_channels = 1 - latent_channels = 2 - cond_dim = 4 - tensor_size = 16 - - conditional_layer_norm = ( - _cln_factory(cond_dim) if conditional_layer_norm_enabled else None - ) - - doubleconvnextblock = DoubleConvNeXtBlock( - in_channels=in_channels, - out_channels=out_channels, - latent_channels=latent_channels, - conditional_layer_norm=conditional_layer_norm, - conditional_layer_norm_once=True, - ).to(device) - - invar = test_data(img_size=tensor_size, device=device) - out_shape = torch.Size([12, out_channels, tensor_size, tensor_size]) - - if conditional_layer_norm_enabled: - conditions_a = torch.randn(1, cond_dim).to(device) - conditions_b = torch.randn(1, cond_dim).to(device) - - outvar_a = doubleconvnextblock(invar, conditions_cln=conditions_a) - outvar_b = doubleconvnextblock(invar, conditions_cln=conditions_b) - - assert outvar_a.shape == out_shape - # entry norm is conditional, so different conditions must still - # change the output even though the internal norms have switched - # to plain (non-conditional) layer normalization - assert not common.compare_output(outvar_a, outvar_b) - else: - outvar = doubleconvnextblock(invar) - assert outvar.shape == out_shape - assert torch.isfinite(outvar).all() - - def test_SymmetricConvNeXtBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( SymmetricConvNeXtBlock, @@ -480,7 +432,6 @@ def test_SymmetricConvNeXtBlock_conditional_layer_norm(device, test_data, pytest in_channels=in_channels, latent_channels=latent_channels, conditional_layer_norm=conditional_layer_norm, - conditional_layer_norm_once=False, ).to(device) assert symmetric_convnextblock.cln_enabled @@ -497,57 +448,6 @@ def test_SymmetricConvNeXtBlock_conditional_layer_norm(device, test_data, pytest assert not common.compare_output(outvar_a, outvar_b) -@pytest.mark.parametrize("conditional_layer_norm_enabled", [True, False]) -def test_SymmetricConvNeXtBlock_conditional_layer_norm_once( - device, test_data, pytestconfig, conditional_layer_norm_enabled -): - from physicsnemo.models.dlwp_healpix.layers import ( - SymmetricConvNeXtBlock, - ) - from physicsnemo.models.dlwp_healpix.layers.normalization import ( - ConditionalLayerNorm, - ) - - in_channels = 2 - # latent_channels must be > 1: a single-channel LayerNorm always - # normalizes to exactly zero, which would erase the (still learnable) - # affine bias differences produced by the entry norm and mask the - # conditioning signal entirely. - latent_channels = 2 - cond_dim = 4 - tensor_size = 16 - - conditional_layer_norm = ( - _cln_factory(cond_dim) if conditional_layer_norm_enabled else None - ) - - symmetric_convnextblock = SymmetricConvNeXtBlock( - in_channels=in_channels, - latent_channels=latent_channels, - conditional_layer_norm=conditional_layer_norm, - conditional_layer_norm_once=True, - ).to(device) - - invar = test_data(img_size=tensor_size, device=device) - out_shape = torch.Size([12, 1, tensor_size, tensor_size]) - - if conditional_layer_norm_enabled: - assert isinstance(symmetric_convnextblock.entry_norm, ConditionalLayerNorm) - - conditions_a = torch.randn(1, cond_dim).to(device) - conditions_b = torch.randn(1, cond_dim).to(device) - - outvar_a = symmetric_convnextblock(invar, conditions_cln=conditions_a) - outvar_b = symmetric_convnextblock(invar, conditions_cln=conditions_b) - - assert outvar_a.shape == out_shape - assert not common.compare_output(outvar_a, outvar_b) - else: - outvar = symmetric_convnextblock(invar) - assert outvar.shape == out_shape - assert torch.isfinite(outvar).all() - - def test_Multi_SymmetricConvNeXtBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( Multi_SymmetricConvNeXtBlock, From 14e380d9aa124b519a73c4eb3bf8ee8ca2dd693b Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Fri, 17 Jul 2026 14:06:15 -0500 Subject: [PATCH 24/28] Add scaling for diagnostic vars back in --- .../healpix/coupledtimeseries_dataset.py | 21 +- .../datapipes/healpix/timeseries_dataset.py | 30 ++- .../test_healpix_diagnostic_outputs.py | 245 ++++++++++++++++++ 3 files changed, 284 insertions(+), 12 deletions(-) create mode 100644 test/datapipes/test_healpix_diagnostic_outputs.py diff --git a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py index a0087bed00..be375fe315 100644 --- a/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/coupledtimeseries_dataset.py @@ -211,9 +211,24 @@ def __getitem__(self, item): axis=2, ) - input_array = (input_array - self.input_scaling["mean"]) / self.input_scaling[ - "std" - ] + # for models with extra outputs + # an empty couplings list means zero coupled variables, rather than + # indexing into a non-existent first coupling + num_coupled_vars = len(self.couplings[0].variables) if self.couplings else 0 + if len(self.ds["targets"].channel_out) != ( + len(self.ds["inputs"].channel_in) - num_coupled_vars + ): + channel_slice = ( + slice(None, -num_coupled_vars) if num_coupled_vars else slice(None) + ) + input_array = ( + input_array - self.input_scaling["mean"][:, channel_slice] + ) / self.input_scaling["std"][:, channel_slice] + else: + input_array = ( + input_array - self.input_scaling["mean"] + ) / self.input_scaling["std"] + if not self.forecast_mode: # BAD NEWS: Indexing the array as commented out below causes unexpected behavior in target creation. # leaving this in here as a warning diff --git a/physicsnemo/datapipes/healpix/timeseries_dataset.py b/physicsnemo/datapipes/healpix/timeseries_dataset.py index 9e2fbad3ef..3fb8884034 100644 --- a/physicsnemo/datapipes/healpix/timeseries_dataset.py +++ b/physicsnemo/datapipes/healpix/timeseries_dataset.py @@ -26,7 +26,6 @@ import pandas as pd import torch from omegaconf import DictConfig, OmegaConf -from torch.utils.data import Dataset from physicsnemo.core.version_check import OptionalImport from physicsnemo.datapipes.datapipe import Datapipe @@ -50,9 +49,9 @@ class MetaData(DatapipeMetaData): ddp_sharding: bool = False -class TimeSeriesDataset(Dataset, Datapipe): +class TimeSeriesDataset(Datapipe): """ - Dataset for sampling from continuous time-series data, compatible with pytorch data loading. + Datapipe for sampling from continuous time-series data, compatible with pytorch data loading. """ def __init__( @@ -244,13 +243,26 @@ def _get_scaling_da(self): scaling_da = scaling_df.to_xarray().astype("float32") # REMARK: we remove the xarray overhead from these + # base (non-coupled) datasets have no 'couplings' attribute, and coupled + # datasets may be configured with an empty couplings list; treat both as + # zero coupled variables rather than indexing into an empty/missing list. + couplings = getattr(self, "couplings", []) + num_coupled_vars = len(couplings[0].variables) if couplings else 0 try: - # we use channel_out instead of channel_in because - # the list of input channels may contain data fetched outside - # the datasets such as coupled fields - self.input_scaling = scaling_da.sel( - index=self.ds.channel_out.values - ).rename({"index": "channel_in"}) + # for models with extra (diagnostic) outputs the number of output + # channels no longer matches the number of prognostic input channels + # (input channels minus any coupled fields). In that case scale the + # inputs by 'channel_in'; otherwise 'channel_out' is still needed + # because the input channel list may reference data fetched outside + # the dataset such as coupled fields. + if len(self.ds.channel_out) != (len(self.ds.channel_in) - num_coupled_vars): + self.input_scaling = scaling_da.sel( + index=self.ds.channel_in.values + ).rename({"index": "channel_in"}) + else: + self.input_scaling = scaling_da.sel( + index=self.ds.channel_out.values + ).rename({"index": "channel_in"}) self.input_scaling = { "mean": np.expand_dims( self.input_scaling["mean"].to_numpy(), (0, 2, 3, 4) diff --git a/test/datapipes/test_healpix_diagnostic_outputs.py b/test/datapipes/test_healpix_diagnostic_outputs.py new file mode 100644 index 0000000000..11cd3e9ad0 --- /dev/null +++ b/test/datapipes/test_healpix_diagnostic_outputs.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Self-contained tests for HEALPix dataloader handling of models with extra +(diagnostic) outputs. + +A model has *extra outputs* when the number of target channels does not match +the number of prognostic input channels (i.e. ``channel_in`` minus any coupled +fields). In that case the input scaling must be selected from ``channel_in`` +(rather than ``channel_out``) so its size matches the loaded input array, and +the coupled portion of that scaling must be sliced off before normalizing the +(prognostic-only) inputs. + +These tests build tiny in-memory xarray datasets so they exercise the real +``TimeSeriesDataset``/``CoupledTimeSeriesDataset`` code paths without requiring +the NFS-backed test dataset. +""" + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("xarray") +omegaconf = pytest.importorskip("omegaconf") +import xarray as xr # noqa: E402 + +from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( # noqa: E402 + CoupledTimeSeriesDataset, +) +from physicsnemo.datapipes.healpix.timeseries_dataset import ( # noqa: E402 + TimeSeriesDataset, +) + +# per-variable (mean, std); values are distinct so we can tell which scaling +# entry (and thus which channel axis) was selected. +_SCALING = { + "a": (1.0, 2.0), + "b": (3.0, 4.0), + "cpl": (10.0, 5.0), + "diag": (7.0, 8.0), +} + +_FACE, _HEIGHT, _WIDTH = 2, 3, 4 + + +def _scaling_config(variables): + return omegaconf.OmegaConf.create( + {v: {"mean": _SCALING[v][0], "std": _SCALING[v][1]} for v in variables} + ) + + +def _make_dataset(channel_in, channel_out, n_time=8): + """Build a minimal classic-style HEALPix dataset. + + Input channel ``i`` is filled with the constant ``i + 1`` so that the + normalized value is deterministic and easy to assert against. + """ + time = pd.date_range("1979-01-01", periods=n_time, freq="3h") + + inputs = np.empty( + (n_time, len(channel_in), _FACE, _HEIGHT, _WIDTH), dtype="float32" + ) + for i in range(len(channel_in)): + inputs[:, i] = float(i + 1) + + targets = np.zeros( + (n_time, len(channel_out), _FACE, _HEIGHT, _WIDTH), dtype="float32" + ) + + return xr.Dataset( + data_vars={ + "inputs": ( + ("time", "channel_in", "face", "height", "width"), + inputs, + ), + "targets": ( + ("time", "channel_out", "face", "height", "width"), + targets, + ), + }, + coords={ + "time": time, + "channel_in": list(channel_in), + "channel_out": list(channel_out), + "face": np.arange(_FACE), + "height": np.arange(_HEIGHT), + "width": np.arange(_WIDTH), + }, + ) + + +def _constant_coupler_config(variables, batch_size): + return [ + { + "coupler": "ConstantCoupler", + "params": { + "batch_size": batch_size, + "variables": list(variables), + "input_times": ["0h"], + "input_time_dim": 1, + "output_time_dim": 1, + "presteps": 0, + "prepared_coupled_data": True, + }, + } + ] + + +def _expected_normalized(channel): + """Normalized value of the constant-filled input for ``channel``. + + Channel ``a`` is the first input channel (constant 1.0), ``b`` the second + (constant 2.0). + """ + value = {"a": 1.0, "b": 2.0}[channel] + mean, std = _SCALING[channel] + return (value - mean) / std + + +def test_timeseries_scaling_uses_channel_out_without_diagnostics(): + """No diagnostic outputs: input scaling is selected from channel_out.""" + ds = _make_dataset(["a", "b"], ["a", "b"]) + dset = TimeSeriesDataset( + ds, + scaling=_scaling_config(["a", "b"]), + input_time_dim=1, + output_time_dim=1, + data_time_step="3h", + time_step="3h", + batch_size=2, + ) + assert dset.input_scaling["mean"].shape[1] == 2 + np.testing.assert_allclose(dset.input_scaling["mean"].ravel(), [1.0, 3.0]) + np.testing.assert_allclose(dset.input_scaling["std"].ravel(), [2.0, 4.0]) + + +def test_timeseries_scaling_uses_channel_in_with_diagnostics(): + """Diagnostic output ('diag' in channel_out only) forces channel_in + scaling, so its size matches the loaded input array.""" + ds = _make_dataset(["a", "b"], ["a", "b", "diag"]) + dset = TimeSeriesDataset( + ds, + scaling=_scaling_config(["a", "b", "diag"]), + input_time_dim=1, + output_time_dim=1, + data_time_step="3h", + time_step="3h", + batch_size=2, + ) + # input scaling covers channel_in (a, b) -- not channel_out (a, b, diag) + assert dset.input_scaling["mean"].shape[1] == 2 + np.testing.assert_allclose(dset.input_scaling["mean"].ravel(), [1.0, 3.0]) + np.testing.assert_allclose(dset.input_scaling["std"].ravel(), [2.0, 4.0]) + # target scaling still covers all output channels + assert dset.target_scaling["mean"].shape[1] == 3 + + +def test_timeseries_getitem_normalizes_with_channel_in_scaling(): + """With diagnostic outputs, __getitem__ normalizes the full input array + against the channel_in scaling and returns channel_in-many channels.""" + ds = _make_dataset(["a", "b"], ["a", "b", "diag"]) + dset = TimeSeriesDataset( + ds, + scaling=_scaling_config(["a", "b", "diag"]), + input_time_dim=1, + output_time_dim=1, + data_time_step="3h", + time_step="3h", + batch_size=2, + add_insolation=False, + ) + inputs_result, _ = dset[0] + inputs = inputs_result[0] # [B, F, T, C, H, W] + assert inputs.shape[3] == 2 + np.testing.assert_allclose(inputs[:, :, :, 0], _expected_normalized("a")) + np.testing.assert_allclose(inputs[:, :, :, 1], _expected_normalized("b")) + + +def test_coupled_getitem_slices_coupled_scaling_with_diagnostics(): + """Coupled model with an extra output: input scaling is taken from + channel_in (prognostic + coupled) and the coupled tail is sliced off so the + prognostic inputs are normalized correctly.""" + ds = _make_dataset(["a", "b", "cpl"], ["a", "b", "diag"]) + dset = CoupledTimeSeriesDataset( + ds, + scaling=_scaling_config(["a", "b", "cpl", "diag"]), + input_variables=["a", "b"], + output_variables=["a", "b", "diag"], + input_time_dim=1, + output_time_dim=1, + data_time_step="3h", + time_step="3h", + batch_size=2, + add_insolation=False, + couplings=_constant_coupler_config(["cpl"], batch_size=2), + ) + # diagnostic path -> input scaling covers channel_in (a, b, cpl) + assert dset.input_scaling["mean"].shape[1] == 3 + + inputs_result, _ = dset[0] + inputs = inputs_result[0] # [B, F, T, C, H, W] + # only the prognostic inputs are returned, normalized by their own scaling + assert inputs.shape[3] == 2 + np.testing.assert_allclose(inputs[:, :, :, 0], _expected_normalized("a")) + np.testing.assert_allclose(inputs[:, :, :, 1], _expected_normalized("b")) + + +def test_coupled_getitem_without_diagnostics_uses_channel_out_scaling(): + """Coupled model with no extra outputs: channel_out scaling is used and no + slicing is applied.""" + ds = _make_dataset(["a", "b", "cpl"], ["a", "b"]) + dset = CoupledTimeSeriesDataset( + ds, + scaling=_scaling_config(["a", "b", "cpl"]), + input_variables=["a", "b"], + output_variables=["a", "b"], + input_time_dim=1, + output_time_dim=1, + data_time_step="3h", + time_step="3h", + batch_size=2, + add_insolation=False, + couplings=_constant_coupler_config(["cpl"], batch_size=2), + ) + # non-diagnostic path -> input scaling covers channel_out (a, b) + assert dset.input_scaling["mean"].shape[1] == 2 + + inputs_result, _ = dset[0] + inputs = inputs_result[0] + assert inputs.shape[3] == 2 + np.testing.assert_allclose(inputs[:, :, :, 0], _expected_normalized("a")) + np.testing.assert_allclose(inputs[:, :, :, 1], _expected_normalized("b")) From c20b3a0626d7d9ebc3b48e9ccffaa2c6af7a66fe Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Fri, 17 Jul 2026 14:18:10 -0500 Subject: [PATCH 25/28] remove check for batch size in coupler to account for CRPS inferencing --- physicsnemo/datapipes/healpix/couplers.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 0b693e27cc..95fb09df33 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -468,12 +468,8 @@ def set_coupled_fields(self, coupled_fields): The data to use when the dataloader requests coupled fields. Expected format is [B, F, T, C, H, W] """ - if coupled_fields.shape[0] != self.batch_size: - raise ValueError( - f"Batch size of coupled field {coupled_fields.shape[0]} doesn't " - f" match configured batch size {self.batch_size}" - ) - + # coupled_channel_indices covers selecting multiple instances of the same variable, + # e.g. z1000-24H and z1000-48H. coupled_fields = coupled_fields[:, :, :, self.coupled_channel_indices, :, :] # TODO: Now support output_time_dim =/= input_time_dim, but presteps need to be 0, will add support for presteps>0 coupled_averaging_periods = [] From 4dffe3fb2c0759f138146d3122256b3ba3cb496b Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Fri, 17 Jul 2026 14:36:31 -0500 Subject: [PATCH 26/28] Have coupler adapt to batch size instead of error, add missing coupler tests. --- physicsnemo/datapipes/healpix/couplers.py | 8 +- test/datapipes/test_healpix_couplers.py | 111 ++++++++++++++++++++++ 2 files changed, 113 insertions(+), 6 deletions(-) create mode 100644 test/datapipes/test_healpix_couplers.py diff --git a/physicsnemo/datapipes/healpix/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index 95fb09df33..38a4bd60d2 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -195,11 +195,7 @@ def set_coupled_fields(self, coupled_fields: th.tensor): The data to use when the dataloader requests coupled fields. Expected format is [B, F, T, C, H, W] """ - if coupled_fields.shape[0] != self.batch_size: - raise ValueError( - f"Batch size of coupled field {coupled_fields.shape[0]} doesn't " - f" match configured batch size {self.batch_size}" - ) + # create buffer for coupling # coupled_channel_indices covers selecting multiple instances of the same variable, # e.g. z1000-24H and z1000-48H. @@ -207,7 +203,7 @@ def set_coupled_fields(self, coupled_fields: th.tensor): :, :, :, self.coupled_channel_indices, :, : ].permute(2, 0, 3, 1, 4, 5) self.preset_coupled_fields = th.empty( - [self.coupled_integration_dim, self.batch_size, self.timevar_dim] + [self.coupled_integration_dim, coupled_fields.shape[0], self.timevar_dim] + list(self.spatial_dims) ) diff --git a/test/datapipes/test_healpix_couplers.py b/test/datapipes/test_healpix_couplers.py new file mode 100644 index 0000000000..18b564c040 --- /dev/null +++ b/test/datapipes/test_healpix_couplers.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Self-contained tests for the HEALPix couplers. + +``ConstantCoupler.set_coupled_fields`` no longer validates the batch dimension +of the provided fields against the configured ``batch_size`` -- the caller is +now responsible for providing consistently shaped fields, and the coupler sizes +its buffer from the field it is given. These tests exercise that behavior with +tiny in-memory datasets so they run without the NFS-backed test dataset. +""" + +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("xarray") +import torch # noqa: E402 +import xarray as xr # noqa: E402 + +from physicsnemo.datapipes.healpix.couplers import ConstantCoupler # noqa: E402 + +_FACE, _HEIGHT, _WIDTH = 2, 3, 4 + + +def _make_coupler_dataset(channel_in, n_time=8): + """Minimal dataset; the coupler only reads ``inputs.shape[2:]`` from it to + determine its spatial dimensions.""" + return xr.Dataset( + data_vars={ + "inputs": ( + ("time", "channel_in", "face", "height", "width"), + np.zeros( + (n_time, len(channel_in), _FACE, _HEIGHT, _WIDTH), + dtype="float32", + ), + ) + }, + coords={ + "time": pd.date_range("1979-01-01", periods=n_time, freq="3h"), + "channel_in": list(channel_in), + "face": np.arange(_FACE), + "height": np.arange(_HEIGHT), + "width": np.arange(_WIDTH), + }, + ) + + +def _make_coupler(batch_size): + ds = _make_coupler_dataset(["c0", "c1", "x"]) + coupler = ConstantCoupler( + dataset=ds, + batch_size=batch_size, + variables=["c0", "c1"], + input_times=["0h"], + input_time_dim=1, + output_time_dim=1, + presteps=0, + ) + # normally assigned by setup_coupling(); the two coupled channels map to the + # two coupled variables. + coupler.coupled_channel_indices = [0, 1] + return coupler + + +def _coupled_fields(batch): + """Coupled field in the expected ``[B, F, T, C, H, W]`` layout. + + ``ConstantCoupler`` broadcasts the first time step and sizes its buffer from + the provided field, so the caller is responsible for providing a + consistently shaped field. + """ + return torch.rand(batch, _FACE, batch, len([0, 1]), _HEIGHT, _WIDTH) + + +def test_set_coupled_fields_adapts_to_provided_batch_size(): + """The old ``batch_size`` mismatch ``ValueError`` was removed: providing a + field whose batch differs from the configured ``batch_size`` no longer + raises, and the buffer is sized from the provided field.""" + configured_batch_size = 2 + coupler = _make_coupler(batch_size=configured_batch_size) + + provided_batch = configured_batch_size * 2 + coupler.set_coupled_fields(_coupled_fields(provided_batch)) + + assert coupler.coupled_mode + # buffer batch dim reflects the provided field, not the configured batch_size + assert coupler.preset_coupled_fields.shape[1] == provided_batch + # in coupled mode, construct_integrated_couplings returns the preset buffer + assert coupler.construct_integrated_couplings() is coupler.preset_coupled_fields + + +def test_set_coupled_fields_matches_configured_batch_size(): + """Sanity: a field whose batch equals the configured batch_size still + works.""" + coupler = _make_coupler(batch_size=3) + coupler.set_coupled_fields(_coupled_fields(3)) + assert coupler.preset_coupled_fields.shape[1] == 3 From 0658f78fa85d4b6db0c955af5bccdb3eba167a23 Mon Sep 17 00:00:00 2001 From: David Pruitt Date: Fri, 17 Jul 2026 16:11:57 -0500 Subject: [PATCH 27/28] Update tests --- .../test_conditional_layer_norm.py | 128 +++++++++--------- 1 file changed, 65 insertions(+), 63 deletions(-) diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py index e5a3913b06..949d476fac 100644 --- a/test/models/dlwp_healpix/test_conditional_layer_norm.py +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -28,63 +28,64 @@ from physicsnemo.models.dlwp_healpix.layers.normalization import ConditionalLayerNorm -def _make_old_cln(condition_shape, channel_depth, device="cpu", **kwargs): - """Instantiate the reference (old) implementation.""" +def _make_reference_cln(condition_shape, channel_depth, device="cpu", **kwargs): + """Instantiate the reference implementation.""" return ConditionalLayerNormReference( condition_shape=condition_shape, channel_depth=channel_depth, **kwargs ).to(device) -def _make_new_cln(condition_shape, channel_depth, device="cpu", **kwargs): - """Instantiate the optimized (new) implementation.""" +def _make_optimized_cln(condition_shape, channel_depth, device="cpu", **kwargs): + """Instantiate the optimized implementation.""" return ConditionalLayerNorm( condition_shape=condition_shape, channel_depth=channel_depth, **kwargs ).to(device) -def _copy_old_to_new(old_cln, new_cln): - """Copy old separate gamma/beta MLP weights into new fused MLP using block-diagonal structure. +def _copy_reference_to_optimized(reference_cln, optimized_cln): + """Copy the reference implementation's separate gamma/beta MLP weights into the + optimized fused MLP using block-diagonal structure. - Old: gamma_mlp and beta_mlp each have hidden_dims [h1, h2] and output C. - New: gamma_beta_mlp has hidden_dims [2*h1, 2*h2] and output 2*C. + Reference: gamma_mlp and beta_mlp each have hidden_dims [h1, h2] and output C. + Optimized: gamma_beta_mlp has hidden_dims [2*h1, 2*h2] and output 2*C. Layer 0 (condition_shape → 2*h1): vertical cat of gamma/beta weights. Layer i>0 (2*h_{i-1} → 2*h_i or 2*C): block-diagonal [[gamma, 0], [0, beta]]. Biases: always concatenated. """ - old_sd = old_cln.state_dict() + reference_sd = reference_cln.state_dict() - # Collect Linear layer indices from the old gamma MLP + # Collect Linear layer indices from the reference gamma MLP gamma_linear_indices = sorted( { int(k.split(".")[1]) - for k in old_sd + for k in reference_sd if k.startswith("gamma_mlp.") and k.endswith(".weight") } ) first_layer_idx = gamma_linear_indices[0] - new_sd = {} - for key in old_sd: + optimized_sd = {} + for key in reference_sd: if key.startswith("norm."): - new_sd[key] = old_sd[key] + optimized_sd[key] = reference_sd[key] for idx in gamma_linear_indices: for param in ("weight", "bias"): - gamma_val = old_sd[f"gamma_mlp.{idx}.{param}"] - beta_val = old_sd[f"beta_mlp.{idx}.{param}"] + gamma_val = reference_sd[f"gamma_mlp.{idx}.{param}"] + beta_val = reference_sd[f"beta_mlp.{idx}.{param}"] fused_key = f"gamma_beta_mlp.{idx}.{param}" if param == "bias": - new_sd[fused_key] = torch.cat([gamma_val, beta_val], dim=0) + optimized_sd[fused_key] = torch.cat([gamma_val, beta_val], dim=0) elif idx == first_layer_idx: # First layer: shared input dim, just cat along output dim - new_sd[fused_key] = torch.cat([gamma_val, beta_val], dim=0) + optimized_sd[fused_key] = torch.cat([gamma_val, beta_val], dim=0) else: # Block-diagonal: [[gamma, 0], [0, beta]] - out_old, in_old = gamma_val.shape + out_dim, in_dim = gamma_val.shape zeros = torch.zeros_like(gamma_val) - new_sd[fused_key] = torch.cat( + optimized_sd[fused_key] = torch.cat( [ torch.cat([gamma_val, zeros], dim=1), torch.cat([zeros, beta_val], dim=1), @@ -92,16 +93,16 @@ def _copy_old_to_new(old_cln, new_cln): dim=0, ) - for key, value in old_sd.items(): + for key, value in reference_sd.items(): if not key.startswith("gamma_mlp."): continue layer_key = key[len("gamma_mlp.") :] parts = layer_key.split(".", 1) if len(parts) != 2 or parts[1] in ("weight", "bias"): continue - new_sd[f"gamma_beta_mlp.{layer_key}"] = value + optimized_sd[f"gamma_beta_mlp.{layer_key}"] = value - new_cln.load_state_dict(new_sd) + optimized_cln.load_state_dict(optimized_sd) def _assert_forward_is_finite( @@ -118,17 +119,17 @@ def _assert_forward_is_finite( @pytest.mark.parametrize("n_cond", [1, 2, 4]) @pytest.mark.parametrize("channels_last", [False, True]) @pytest.mark.parametrize("scale_center", [0.0, 1.0]) -def test_old_vs_new_forward(device, n_cond, channels_last, scale_center): +def test_optimized_vs_reference_forward(device, n_cond, channels_last, scale_center): """Verify optimized CLN matches reference implementation with block-diagonal weight mapping.""" C, H, W = 128, 16, 16 cond_shape = 32 B_nf = n_cond * 12 torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C, device=device, scale_center=scale_center) + reference_cln = _make_reference_cln(cond_shape, C, device=device, scale_center=scale_center) - new_cln = _make_new_cln(cond_shape, C, device=device, scale_center=scale_center) - _copy_old_to_new(old_cln, new_cln) + optimized_cln = _make_optimized_cln(cond_shape, C, device=device, scale_center=scale_center) + _copy_reference_to_optimized(reference_cln, optimized_cln) x = torch.randn(B_nf, C, H, W, device=device) cond = torch.randn(n_cond, cond_shape, device=device) @@ -137,32 +138,33 @@ def test_old_vs_new_forward(device, n_cond, channels_last, scale_center): x = x.to(memory_format=torch.channels_last) with torch.no_grad(): - out_old = old_cln(x, cond) - out_new = new_cln(x, cond) + out_ref = reference_cln(x, cond) + out_opt = optimized_cln(x, cond) - assert out_old.shape == out_new.shape - assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( - f"Max diff: {(out_old - out_new).abs().max().item()}" + assert out_ref.shape == out_opt.shape + assert torch.allclose(out_ref, out_opt, atol=1e-5, rtol=1e-4), ( + f"Max diff: {(out_ref - out_opt).abs().max().item()}" ) if channels_last: - assert out_new.is_contiguous(memory_format=torch.channels_last), ( + assert out_opt.is_contiguous(memory_format=torch.channels_last), ( "Output should preserve channels_last format" ) @pytest.mark.parametrize("channels_last", [False, True]) -def test_old_vs_new_backward(device, channels_last): - """Verify gradients match between old and new implementations.""" +def test_optimized_vs_reference_backward(device, channels_last): + """Verify gradients match between the optimized implementation and the + reference implementation.""" C, H, W = 64, 8, 8 cond_shape = 16 n_cond = 2 B_nf = n_cond * 12 torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C, device=device) - new_cln = _make_new_cln(cond_shape, C, device=device) - _copy_old_to_new(old_cln, new_cln) + reference_cln = _make_reference_cln(cond_shape, C, device=device) + optimized_cln = _make_optimized_cln(cond_shape, C, device=device) + _copy_reference_to_optimized(reference_cln, optimized_cln) x_base = torch.randn(B_nf, C, H, W, device=device) cond_base = torch.randn(n_cond, cond_shape, device=device) @@ -170,23 +172,23 @@ def test_old_vs_new_backward(device, channels_last): if channels_last: x_base = x_base.to(memory_format=torch.channels_last) - x_old = x_base.clone().detach().requires_grad_(True) - cond_old = cond_base.clone().detach().requires_grad_(True) - x_new = x_base.clone().detach().requires_grad_(True) - cond_new = cond_base.clone().detach().requires_grad_(True) + x_ref = x_base.clone().detach().requires_grad_(True) + cond_ref = cond_base.clone().detach().requires_grad_(True) + x_opt = x_base.clone().detach().requires_grad_(True) + cond_opt = cond_base.clone().detach().requires_grad_(True) - out_old = old_cln(x_old, cond_old) - out_old.sum().backward() + out_ref = reference_cln(x_ref, cond_ref) + out_ref.sum().backward() - out_new = new_cln(x_new, cond_new) - out_new.sum().backward() + out_opt = optimized_cln(x_opt, cond_opt) + out_opt.sum().backward() - assert torch.allclose(x_old.grad, x_new.grad, atol=1e-4, rtol=1e-3), ( - f"Input grad max diff: {(x_old.grad - x_new.grad).abs().max().item()}" + assert torch.allclose(x_ref.grad, x_opt.grad, atol=1e-4, rtol=1e-3), ( + f"Input grad max diff: {(x_ref.grad - x_opt.grad).abs().max().item()}" ) - assert torch.allclose(cond_old.grad, cond_new.grad, atol=1e-4, rtol=1e-3), ( - f"Cond grad max diff: {(cond_old.grad - cond_new.grad).abs().max().item()}" + assert torch.allclose(cond_ref.grad, cond_opt.grad, atol=1e-4, rtol=1e-3), ( + f"Cond grad max diff: {(cond_ref.grad - cond_opt.grad).abs().max().item()}" ) @@ -198,7 +200,7 @@ def test_init_cln_to_zero_matches_layer_norm(device, channels_last): B_nf = n_cond * 12 torch.manual_seed(42) - cln = _make_new_cln(32, C, device=device, scale_center=1.0, init_cln_to_zero=True) + cln = _make_optimized_cln(32, C, device=device, scale_center=1.0, init_cln_to_zero=True) plain_ln = torch.nn.LayerNorm(C, elementwise_affine=False).to(device) x = torch.randn(B_nf, C, H, W, device=device) @@ -226,7 +228,7 @@ def test_backward_gradients(device, channels_last): B_nf = n_cond * 12 torch.manual_seed(42) - cln = _make_new_cln(cond_shape, C, device=device) + cln = _make_optimized_cln(cond_shape, C, device=device) x = torch.randn(B_nf, C, H, W, device=device) cond = torch.randn(n_cond, cond_shape, device=device) @@ -258,7 +260,7 @@ def test_backward_channels_last_matches_contiguous(device): B_nf = n_cond * 12 torch.manual_seed(42) - cln = _make_new_cln(cond_shape, C, device=device) + cln = _make_optimized_cln(cond_shape, C, device=device) x_base = torch.randn(B_nf, C, H, W, device=device) cond_base = torch.randn(n_cond, cond_shape, device=device) @@ -294,7 +296,7 @@ def test_backward_channels_last_matches_contiguous(device): @pytest.mark.parametrize("hidden_dims", [[], [64]]) -def test_old_vs_new_forward_hidden_dims_variation(device, hidden_dims): +def test_optimized_vs_reference_forward_hidden_dims_variation(device, hidden_dims): """Verify the block-diagonal weight mapping generalizes to MLP depths other than the default two hidden layers, including the degenerate zero-hidden-layer case (a single Linear directly from ``condition_shape`` to the output). @@ -305,19 +307,19 @@ def test_old_vs_new_forward_hidden_dims_variation(device, hidden_dims): B_nf = n_cond * 12 torch.manual_seed(42) - old_cln = _make_old_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) - new_cln = _make_new_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) - _copy_old_to_new(old_cln, new_cln) + reference_cln = _make_reference_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) + optimized_cln = _make_optimized_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) + _copy_reference_to_optimized(reference_cln, optimized_cln) x = torch.randn(B_nf, C, H, W, device=device) cond = torch.randn(n_cond, cond_shape, device=device) with torch.no_grad(): - out_old = old_cln(x, cond) - out_new = new_cln(x, cond) + out_ref = reference_cln(x, cond) + out_opt = optimized_cln(x, cond) - assert torch.allclose(out_old, out_new, atol=1e-5, rtol=1e-4), ( - f"Max diff: {(out_old - out_new).abs().max().item()}" + assert torch.allclose(out_ref, out_opt, atol=1e-5, rtol=1e-4), ( + f"Max diff: {(out_ref - out_opt).abs().max().item()}" ) @@ -426,11 +428,11 @@ def test_load_new_format_checkpoint_roundtrip(device): """ C, cond_shape = 32, 16 torch.manual_seed(42) - source = _make_new_cln(cond_shape, C, device=device) + source = _make_optimized_cln(cond_shape, C, device=device) state_dict = source.state_dict() assert not any(k.startswith("gamma_mlp.") for k in state_dict) - target = _make_new_cln(cond_shape, C, device=device) + target = _make_optimized_cln(cond_shape, C, device=device) missing, unexpected = target.load_state_dict(state_dict, strict=True) assert not missing assert not unexpected From 25dde1252f892aad915a06083467eb0587a490ae Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:13:46 +0000 Subject: [PATCH 28/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../test_conditional_layer_norm.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/test/models/dlwp_healpix/test_conditional_layer_norm.py b/test/models/dlwp_healpix/test_conditional_layer_norm.py index 949d476fac..4559d560e3 100644 --- a/test/models/dlwp_healpix/test_conditional_layer_norm.py +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -126,9 +126,13 @@ def test_optimized_vs_reference_forward(device, n_cond, channels_last, scale_cen B_nf = n_cond * 12 torch.manual_seed(42) - reference_cln = _make_reference_cln(cond_shape, C, device=device, scale_center=scale_center) + reference_cln = _make_reference_cln( + cond_shape, C, device=device, scale_center=scale_center + ) - optimized_cln = _make_optimized_cln(cond_shape, C, device=device, scale_center=scale_center) + optimized_cln = _make_optimized_cln( + cond_shape, C, device=device, scale_center=scale_center + ) _copy_reference_to_optimized(reference_cln, optimized_cln) x = torch.randn(B_nf, C, H, W, device=device) @@ -200,7 +204,9 @@ def test_init_cln_to_zero_matches_layer_norm(device, channels_last): B_nf = n_cond * 12 torch.manual_seed(42) - cln = _make_optimized_cln(32, C, device=device, scale_center=1.0, init_cln_to_zero=True) + cln = _make_optimized_cln( + 32, C, device=device, scale_center=1.0, init_cln_to_zero=True + ) plain_ln = torch.nn.LayerNorm(C, elementwise_affine=False).to(device) x = torch.randn(B_nf, C, H, W, device=device) @@ -307,8 +313,12 @@ def test_optimized_vs_reference_forward_hidden_dims_variation(device, hidden_dim B_nf = n_cond * 12 torch.manual_seed(42) - reference_cln = _make_reference_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) - optimized_cln = _make_optimized_cln(cond_shape, C, device=device, mlp_hidden_dims=hidden_dims) + reference_cln = _make_reference_cln( + cond_shape, C, device=device, mlp_hidden_dims=hidden_dims + ) + optimized_cln = _make_optimized_cln( + cond_shape, C, device=device, mlp_hidden_dims=hidden_dims + ) _copy_reference_to_optimized(reference_cln, optimized_cln) x = torch.randn(B_nf, C, H, W, device=device)