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/couplers.py b/physicsnemo/datapipes/healpix/couplers.py index a0960b2e6d..38a4bd60d2 100644 --- a/physicsnemo/datapipes/healpix/couplers.py +++ b/physicsnemo/datapipes/healpix/couplers.py @@ -195,24 +195,21 @@ 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. 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] + [self.coupled_integration_dim, coupled_fields.shape[0], 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 @@ -467,12 +464,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 = [] 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/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py b/physicsnemo/models/dlwp_healpix/HEALPixRecUNet.py index 20553b5e24..1c6267a921 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 @@ -83,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 ------- @@ -145,6 +162,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 +200,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 +228,9 @@ def __init__( enable_healpixpad=self.enable_healpixpad, ) + self.constraints = None + self.set_constraints(constraints) + @property def integration_steps(self): r""" @@ -297,7 +322,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 +388,29 @@ 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: + # 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""" Reshape decoder output back to explicit time and channel dimensions. @@ -396,7 +444,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 +499,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 +553,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 +594,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..4004b84fba 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 @@ -77,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 ------- @@ -139,6 +155,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 +185,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 +213,9 @@ def __init__( enable_healpixpad=self.enable_healpixpad, ) + self.constraints = None + self.set_constraints(constraints) + @property def integration_steps(self): r""" @@ -283,7 +307,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 +409,31 @@ 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: + # 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, + inputs: Sequence, + output_only_last: bool = False, + conditions_cln=None, + ) -> torch.Tensor: r""" Forward pass of the HEALPix UNet. @@ -426,10 +476,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/__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 84f6463c3b..dec4636954 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", @@ -62,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", @@ -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": @@ -127,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 32e6ffaac2..db45530d21 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_blocks.py @@ -14,12 +14,29 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Sequence, Tuple, Union +""" +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 from physicsnemo.nn.module.hpx import HEALPixLayer +from .normalization import ConditionalLayerNorm + # # RECURRENT BLOCKS # @@ -328,6 +345,8 @@ def __init__( activation: torch.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, + conditional_layer_norm: Callable = None, + dropout: float = 0.0, ): """ Parameters: @@ -355,9 +374,11 @@ 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( @@ -369,8 +390,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( @@ -396,8 +417,17 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply batch norm and conditional layer norm if needed + if conditional_layer_norm is not None: + 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(torch.nn.Dropout2d(p=dropout)) + # 1x1 convolution establishing increased channels convblock1.append( geometry_layer( @@ -410,8 +440,19 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply layer norm if needed + if 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(torch.nn.Dropout2d(p=dropout)) + # 1x1 convolution returning to latent channels convblock1.append( geometry_layer( @@ -426,7 +467,9 @@ def __init__( ) if activation is not None: convblock1.append(activation) - self.convblock1 = torch.nn.Sequential(*convblock1) + if dropout > 0.0: + 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 = [] @@ -442,8 +485,16 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + # Apply batch norm and conditional layer norm if needed + if conditional_layer_norm is not None: + 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(torch.nn.Dropout2d(p=dropout)) + # 1x1 convolution establishing increased channels convblock2.append( geometry_layer( @@ -456,8 +507,18 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + # Apply layer norm if needed + if 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(torch.nn.Dropout2d(p=dropout)) + # 1x1 convolution reducing to output channels convblock2.append( geometry_layer( @@ -472,9 +533,11 @@ def __init__( ) if activation is not None: convblock2.append(activation) - self.convblock2 = torch.nn.Sequential(*convblock2) + if dropout > 0.0: + convblock2.append(torch.nn.Dropout2d(p=dropout)) + self.convblock2 = torch.nn.ModuleList(convblock2) - def forward(self, x): + def forward(self, x, conditions_cln=None): """Forward pass of the DoubleConvNextBlock Parameters @@ -486,16 +549,36 @@ def forward(self, x): ------- torch.Tensor result of the forward pass + conditions_cln: torch.Tensor, optional + conditions for the conditional layer normalization """ + + # save residual for the first block + x1_residual = self.skip_module1(x) + # internal convnext result - x1 = self.skip_module1(x) + self.convblock1(x) + 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) + # return second convnext result - return self.skip_module2(x1) + self.convblock2(x1) + 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(torch.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__( @@ -511,22 +594,28 @@ def __init__( activation: torch.nn.Module = None, enable_nhwc: bool = False, enable_healpixpad: bool = False, + dropout: float = 0.0, + conditional_layer_norm: Callable = None, ): """ 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. """ super().__init__() # Create a ModuleList to store complete blocks self.blocks = torch.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,13 +628,31 @@ 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, + ), ) - def forward(self, x): + 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) + out = block(out, conditions_cln=conditions_cln) return out @@ -567,7 +674,10 @@ def __init__( upscale_factor: int = 4, 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: torch.nn.Module = None, ): """ Parameters @@ -592,24 +702,39 @@ 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: torch.nn.Module, optional + conditional layer normalization. If None, + no conditional layer normalization is applied. """ + super().__init__() - if in_channels == int(latent_channels): - self.skip_module = lambda x: x # Identity-function required in forward pass - 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, - ) + 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 - # 1st ConvNeXt block, the output of this one remains internal + if use_block_skip_connection: + if in_channels == int(out_channels): + self.skip_module = lambda x: x + 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, + ) + + # 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, @@ -621,9 +746,17 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + # Apply layer norm if needed + if 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) - # 1x1 convolution establishing increased channels + if dropout > 0.0: + convblock.append(torch.nn.Dropout2d(p=dropout)) + + # 1x1: latent → latent * upscale convblock.append( geometry_layer( layer=torch.nn.Conv2d, @@ -635,9 +768,20 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply layer norm if needed + if 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(torch.nn.Dropout2d(p=dropout)) + + # 1x1: upscale → latent convblock.append( geometry_layer( layer=torch.nn.Conv2d, @@ -649,14 +793,23 @@ def __init__( enable_healpixpad=enable_healpixpad, ) ) + + # Apply layer norm if needed + if 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(torch.nn.Dropout2d(p=dropout)) + + # 3x3: latent → out (no norm on this one, following convnext) convblock.append( geometry_layer( layer=torch.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 +818,36 @@ def __init__( ) if activation is not None: convblock.append(activation) - self.convblock = torch.nn.Sequential(*convblock) + if dropout > 0.0: + convblock.append(torch.nn.Dropout2d(p=dropout)) - def forward(self, x): - """Forward pass of the SymmetricConvNextBlock + self.convblock = torch.nn.ModuleList(convblock) + 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) + + # Save residual + residual = self.skip_module(x) if self.use_block_skip_connection else 0 + + for layer in self.convblock: + if isinstance(layer, ConditionalLayerNorm): + x = layer(x, conditions=conditions_cln) + else: + x = layer(x) + + return x + residual # @@ -708,7 +874,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 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..4c328b8e43 --- /dev/null +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_constraints.py @@ -0,0 +1,80 @@ +# 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. + +""" +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], + 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.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..68f70b896c 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_decoder.py @@ -14,11 +14,18 @@ # 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 from hydra.utils import instantiate from omegaconf import DictConfig +from torch.utils.checkpoint import checkpoint class UNetDecoder(torch.nn.Module): @@ -36,6 +43,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 +70,52 @@ 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 + + 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 # 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 +131,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 +150,6 @@ def __init__( enable_healpixpad=enable_healpixpad, ) - # Recurrent module if recurrent_block is not None: rec_module = instantiate( config=recurrent_block, @@ -121,8 +170,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 +179,49 @@ 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: + """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) + 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 +229,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 +238,22 @@ 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..dee98855c0 100644 --- a/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py +++ b/physicsnemo/models/dlwp_healpix/layers/healpix_encoder.py @@ -14,11 +14,18 @@ # 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 from hydra.utils import instantiate from omegaconf import DictConfig +from torch.utils.checkpoint import checkpoint class UNetEncoder(torch.nn.Module): @@ -35,6 +42,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 +67,50 @@ 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 + + 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 + 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 +124,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 +150,74 @@ 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: + """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): + 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..e1819a2b60 --- /dev/null +++ b/physicsnemo/models/dlwp_healpix/layers/normalization.py @@ -0,0 +1,192 @@ +# 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. + +""" +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 + +from physicsnemo.core.version_check import OptionalImport + +# ``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") + + +@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] + 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(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 + 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, + channel_depth: int, + mlp_hidden_dims: List[int] = [128, 128], + activation: torch.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 torch.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 = torch.nn.LayerNorm(channel_depth, elementwise_affine=False) + elif norm_op == "apex": + self.norm = apex_normalization.FusedLayerNorm( + channel_depth, elementwise_affine=False + ) + + def _make_mlp( + self, + in_dim: int, + hidden_dims: List[int], + out_dim: int, + activation: torch.nn.Module, + ) -> torch.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: torch.nn.Module + The activation function + + Returns + ------- + torch.nn.Sequential + The MLP + """ + layers = [] + for hdim in hidden_dims: + layers.append(torch.nn.Linear(in_dim, hdim)) + if activation: + layers.append(activation) + in_dim = hdim + layers.append(torch.nn.Linear(in_dim, out_dim)) + return torch.nn.Sequential(*layers) + + def forward(self, x: torch.Tensor, conditions: torch.Tensor) -> torch.Tensor: + """ + Parameters + ---------- + x : torch.Tensor + Input tensor of shape: (B, C, H, W) + conditions : torch.Tensor + Conditioning tensor of shape (B*n_cond, cond_dim) + + Returns + ------- + torch.Tensor + Normalized and conditioned tensor of shape: (B, C, H, W) + """ + + 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) + 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=torch.channels_last + ) + else: + return result.permute(0, 3, 1, 2) 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/datapipes/test_healpix.py b/test/datapipes/test_healpix.py deleted file mode 100644 index 37aafc8fe8..0000000000 --- a/test/datapipes/test_healpix.py +++ /dev/null @@ -1,701 +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 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 - -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 - from physicsnemo.datapipes.healpix.data_modules import ( - open_time_series_dataset_classic_prebuilt, - ) - - with pytest.raises(FileNotFoundError, match=("Dataset doesn't appear to 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("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 -): - 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 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("dask") -@requires_module("netCDF4") -def test_TimeSeriesDataModule_initialization( - data_dir, create_path, dataset_name, scaling_double_dict, 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, - dataset_name=dataset_name, - batch_size=1, - data_format="null", - ) - - # 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) - - # without the prebuilt dataset - timeseries_dm = TimeSeriesDataModule( - src_directory=create_path, - dst_directory=create_path, - 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, 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 - 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, - ) - - # 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/datapipes/test_healpix_couple.py b/test/datapipes/test_healpix_couple.py deleted file mode 100644 index 8a267aea33..0000000000 --- a/test/datapipes/test_healpix_couple.py +++ /dev/null @@ -1,1311 +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 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 - -omegaconf = pytest.importorskip("omegaconf") -np = pytest.importorskip("numpy") -pd = pytest.importorskip("pandas") -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") -@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", - ) - coupler.setup_coupling(mock_coupled_module) - assert coupler.coupled_channel_indices == [1] - - 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"]) - - # test incorrect batch size - 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], - ) - with pytest.raises(ValueError, match=("Batch size of coupled field 4 ")): - coupler.set_coupled_fields(coupled_fields) - - 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).unsqueeze(0) - expected = expected.repeat( - coupler.coupled_integration_dim, coupled_fields_batch_size, 1, 1, 1, 1 - ) - assert th.equal(expected, coupler.construct_integrated_couplings()) - - # 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.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]) - - 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, - ) - - # 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) - - # veryify averaging slices computed correctly - mock_coupled_module = coupler_helper( - output_variables=["not_coupled", "z500"], - 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] - - # test a mismatched batch size - 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], - ) - with pytest.raises(ValueError, match=("Batch size of coupled field 4 ")): - coupler.set_coupled_fields(coupled_fields) - - 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("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, 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"] - - 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, - ) - - # 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, 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"] - - 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( - dataset=zarr_ds, - input_variables=variables, - scaling=scaling_dict, - batch_size=1, - forecast_init_times=zarr_ds.time[:init_times], - couplings=constant_coupler, - ) - 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, - }, - } - ] - - # 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, 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()) - - 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, - }, - } - ] - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_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=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]) - - # With insolation we get 1 extra channel - timeseries_ds = CoupledTimeSeriesDataset( - dataset=zarr_ds, - input_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=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=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=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("dask") -@requires_module("netCDF4") -@requires_module("xarray") -def test_CoupledTimeSeriesDataModule_initialization( - data_dir, create_path, dataset_name, scaling_double_dict, 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, - }, - } - ] - - # 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, - dataset_name=dataset_name, - batch_size=1, - data_format="null", - 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) - - # without the prebuilt dataset - timeseries_dm = CoupledTimeSeriesDataModule( - src_directory=create_path, - dst_directory=create_path, - 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, pytestconfig -): - from physicsnemo.datapipes.healpix.data_modules import ( - CoupledTimeSeriesDataModule, - ) - - 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, - }, - } - ] - - # 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, 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 - 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, - ) - - # 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() - - -@requires_module("omegaconf") -def test_CoupledTimeSeriesDataModule_get_coupled_vars( - data_dir, create_path, dataset_name, scaling_double_dict, 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 - 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, - ) - - outvar = timeseries_dm._get_coupled_vars() - outvar.sort() - expected = ["z250"] - expected.sort() - - 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( - 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, - ) - 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, pytestconfig -): - from physicsnemo.datapipes.healpix.coupledtimeseries_dataset import ( - CoupledTimeSeriesDataset, - ) - - 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_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], - ) - - 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 = CoupledTimeSeriesDataset( - dataset=test_ds, - 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_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 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")) diff --git a/test/models/dlwp_healpix/_cln_reference.py b/test/models/dlwp_healpix/_cln_reference.py new file mode 100644 index 0000000000..0dddf37e6a --- /dev/null +++ b/test/models/dlwp_healpix/_cln_reference.py @@ -0,0 +1,105 @@ +# 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. + +import copy +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: + # 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) + + 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..4559d560e3 --- /dev/null +++ b/test/models/dlwp_healpix/test_conditional_layer_norm.py @@ -0,0 +1,456 @@ +# 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 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 import normalization +from physicsnemo.models.dlwp_healpix.layers.normalization import ConditionalLayerNorm + + +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_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_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. + + 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. + """ + reference_sd = reference_cln.state_dict() + + # Collect Linear layer indices from the reference gamma MLP + gamma_linear_indices = sorted( + { + int(k.split(".")[1]) + for k in reference_sd + if k.startswith("gamma_mlp.") and k.endswith(".weight") + } + ) + first_layer_idx = gamma_linear_indices[0] + + optimized_sd = {} + for key in reference_sd: + if key.startswith("norm."): + optimized_sd[key] = reference_sd[key] + + for idx in gamma_linear_indices: + for param in ("weight", "bias"): + 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": + 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 + optimized_sd[fused_key] = torch.cat([gamma_val, beta_val], dim=0) + else: + # Block-diagonal: [[gamma, 0], [0, beta]] + out_dim, in_dim = gamma_val.shape + zeros = torch.zeros_like(gamma_val) + optimized_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 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 + optimized_sd[f"gamma_beta_mlp.{layer_key}"] = value + + optimized_cln.load_state_dict(optimized_sd) + + +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.""" + 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_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) + 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 + ) + _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) + + if channels_last: + x = x.to(memory_format=torch.channels_last) + + with torch.no_grad(): + out_ref = reference_cln(x, cond) + out_opt = optimized_cln(x, cond) + + 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_opt.is_contiguous(memory_format=torch.channels_last), ( + "Output should preserve channels_last format" + ) + + +@pytest.mark.parametrize("channels_last", [False, True]) +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) + 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) + + if channels_last: + x_base = x_base.to(memory_format=torch.channels_last) + + 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_ref = reference_cln(x_ref, cond_ref) + out_ref.sum().backward() + + out_opt = optimized_cln(x_opt, cond_opt) + out_opt.sum().backward() + + 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_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()}" + ) + + +@pytest.mark.parametrize("channels_last", [False, True]) +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_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) + cond = torch.randn(n_cond, 32, device=device) + + 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.parametrize("channels_last", [False, True]) +def test_backward_gradients(device, 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_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) + + 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}" + + +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 + n_cond = 2 + B_nf = n_cond * 12 + + torch.manual_seed(42) + 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) + + # 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.parametrize("hidden_dims", [[], [64]]) +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). + """ + C, H, W = 32, 8, 8 + cond_shape = 16 + n_cond = 2 + 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 + ) + _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_ref = reference_cln(x, cond) + out_opt = optimized_cln(x, cond) + + assert torch.allclose(out_ref, out_opt, atol=1e-5, rtol=1e-4), ( + f"Max diff: {(out_ref - out_opt).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. The error is raised by + ``OptionalImport`` when the fused-norm symbol is accessed. + """ + if importlib.util.find_spec("apex") is not None: + pytest.skip("apex is installed in this environment") + + with pytest.raises(ImportError, match="Missing optional dependency: apex"): + 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 ``apex.normalization`` module 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 + + 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) + 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_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_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_optimized_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..88385abb3a 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,68 @@ 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, + ).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) + + def test_SymmetricConvNeXtBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( SymmetricConvNeXtBlock, @@ -210,6 +331,123 @@ 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, + ).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) + + def test_Multi_SymmetricConvNeXtBlock_initialization(device, pytestconfig): from physicsnemo.models.dlwp_healpix.layers import ( Multi_SymmetricConvNeXtBlock, @@ -248,6 +486,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):