DLESYM pdates

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

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 = []

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)

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]