Simplify the canonical data loading boundary

AGENTS.md

@@ -183,7 +183,6 @@ For vizualization or other long-running tasks:
    * `samudra_multi.py`: samudra-multi encoder → processor → decoder architecture supporting multi-scale training
    * `samudra_mini.py`: SamudraMini single PerceiverIO model for lightweight training-shape experiments
    * `base.py`: Abstract base model class with common functionality (residual predictions, masking, gradient detaching)
-   * `corrector.py`: Optional corrector network for error correction (will likely be deprecated soon)
    * `modules/`: Reusable building blocks including `unet_backbone.py`, `encoder.py` (PerceiverEncoder), `blocks.py` (ConvNext blocks), `activations.py`, and `augment_input.py`
 
 2. **Time Stepping** (`src/samudra/stepper.py`)
@@ -196,7 +195,7 @@ For vizualization or other long-running tasks:
    * `TrainData` and `InferenceDataset` classes for training/eval
    * Supports time-based train/validation splits
    * Variables include temperature (`thetao`), salinity (`so`), u/v velocities, sea surface height (`zos`), and surface heat flux (`hfds`)
-   * Data normalization via the `Normalize` multiton (only used in the Corrector and Aggregator, should be deprecated)
+   * Data normalization via the `Normalize` multiton for aggregation and output utilities
 
 4. **Training Loop** (`src/samudra/train.py`)
    * Distributed training support via PyTorch DDP
@@ -268,7 +267,6 @@ src/samudra/
 │   ├── samudra.py        # Samudra (ConvNeXt U-Net)
 │   ├── samudra_multi.py  # samudra-multi (encoder-processor-decoder)
 │   ├── samudra_mini.py   # SamudraMini (single PerceiverIO)
-│   ├── corrector.py      # Corrector network
 │   └── modules/          # Reusable blocks (unet_backbone, encoder, blocks, activations)
 ├── aggregator/
 │   ├── main.py           # Base Aggregator class

configs/data/llc.yaml

@@ -20,7 +20,6 @@ sources:
     data_location: LLC.zarr
     data_means_location: LLC_means.zarr
     data_stds_location: LLC_stds.zarr
-static_data_vars: []
 concurrent_compute: false
 loading:
   type: cpu

configs/data/om4.yaml

@@ -20,7 +20,6 @@ sources:
     data_location: OM4.zarr
     data_means_location: OM4_means.zarr
     data_stds_location: OM4_stds.zarr
-static_data_vars: []
 concurrent_compute: true
 loading:
   type: cpu

configs/samudra_multi_om4/eval.yaml

@@ -31,7 +31,6 @@ experiment:
     project: default
 
 data:
-  static_data_vars: []
   concurrent_compute: true
   sources:
     # First source is used for inference. Reorder to eval a different scale.

configs/samudra_multi_om4/train.yaml

@@ -34,7 +34,6 @@ experiment:
     mode: disabled
     project: default
 data:
-  static_data_vars: []
   concurrent_compute: true
   loading:
     type: cpu

configs/samudra_om4/model.yaml

@@ -23,7 +23,3 @@ unet:
 
   down_sampling_block: "avg_pool"
   up_sampling_block: "zonally_periodic_upsample"
-
-corrector:
-  non_negative_corrector_names: null
-  ocean_heat_corrector: false

configs/samudra_om4_v1/model.yaml

@@ -23,7 +23,3 @@ unet:
 
   down_sampling_block: "avg_pool"
   up_sampling_block: "bilinear_upsample"
-
-corrector:
-  non_negative_corrector_names: null
-  ocean_heat_corrector: false

configs/samudra_om4_v2/model.yaml

@@ -23,7 +23,3 @@ unet:
 
   down_sampling_block: "avg_pool"
   up_sampling_block: "zonally_periodic_upsample"
-
-corrector:
-  non_negative_corrector_names: null
-  ocean_heat_corrector: false

configs/samudra_om4_v2_highres/model.yaml

@@ -24,7 +24,3 @@ unet:
 
   down_sampling_block: "avg_pool"
   up_sampling_block: "zonally_periodic_upsample"
-
-corrector:
-  non_negative_corrector_names: null
-  ocean_heat_corrector: false

configs/test/train_samudra_om4_v2_resume.yaml

@@ -81,7 +81,3 @@ model:
 
     down_sampling_block: "avg_pool"
     up_sampling_block: "zonally_periodic_upsample"
-
-  corrector:
-    non_negative_corrector_names: null
-    ocean_heat_corrector: false

data/README.md

@@ -196,7 +196,7 @@ These files live on the OSN pod:
 - Data: https://nyu1.osn.mghpcc.org/m2lines-pubs/FOMO/raw/om4_5daily.zarr
 - Gaussian Grid: https://nyu1.osn.mghpcc.org/m2lines-pubs/FOMO/raw/grids/gaussian_grid_180_by_360.zarr
 - Mosaic File: https://nyu1.osn.mghpcc.org/m2lines-pubs/FOMO/raw/grids/ocean_hgrid.zarr
-- Static Data File: https://nyu1.osn.mghpcc.org/m2lines-pubs/FOMO/raw/ocean_static_no_mask_table.zarr
+- Native Grid File: https://nyu1.osn.mghpcc.org/m2lines-pubs/FOMO/raw/ocean_static_no_mask_table.zarr
 
 These can be directly input to our pre-processing scripts (see usage above).
 
@@ -237,4 +237,4 @@ mamba install pytest dask pre-commit
 Before you edit the code make sure all tests pass by running pytest from the root level of this repository
 ```
 pytest
-```
+```

data/ocean_preprocessing/__main__.py

@@ -269,7 +269,6 @@ def om4(
         native_grid_path: str,
         nc_mosaic_path: str,
         target_grid_path: str,
-        ocean_static_path: str | None = None,
         spatial_filter_scale: None | int = None,
     ) -> None:
         """Process the OM4 oceans dataset (the ocean component of CMIP).
@@ -279,16 +278,15 @@ def om4(
         1. Model-specific preprocessing (interpolate velocities to tracer grid)
         2. Account for partial depths (if --account-for-partial-depths)
         3. Validation of preprocessed dataset (unless --skip-validation)
-        4. Optionally merge static ocean variables (if ocean_static_path provided)
-        5. Vector rotation to zonal/meridional coordinates
-        6. QC plots for rotated vectors (unless --skip-plots)
-        7. Spatial filtering with Gaussian kernel (unless --skip-spatial-filtering)
-        8. Horizontal regridding to target resolution (unless --skip-regridding)
-        9. Type conversion to float32
-        10. Restore variable attributes and add provenance metadata
-        11. Optionally flatten 3D variables by depth level (unless --skip-flattening)
-        12. Drop extraneous dimensions (we only need 'x', 'y', 'time' and maybe 'lev').
-        13. Write output to Zarr store in Zarr v2 format.
+        4. Vector rotation to zonal/meridional coordinates
+        5. QC plots for rotated vectors (unless --skip-plots)
+        6. Spatial filtering with Gaussian kernel (unless --skip-spatial-filtering)
+        7. Horizontal regridding to target resolution (unless --skip-regridding)
+        8. Type conversion to float32
+        9. Restore variable attributes and add provenance metadata
+        10. Optionally flatten 3D variables by depth level (unless --skip-flattening)
+        11. Drop extraneous dimensions (we only need 'x', 'y', 'time' and maybe 'lev').
+        12. Write output to Zarr store in Zarr v2 format.
 
         > NOTE: We pre-determine the spatial filtering scale based on the target grid!
         > We make informed guesses for what the scale should be. If you'd like to
@@ -313,10 +311,6 @@ def om4(
                 'gaussian_grid_360_by_720.zarr'). This file defines the output grid resolution
                 and is used for horizontal regridding (unless --skip-regridding). The basename
                 is also used to determine the appropriate spatial filter scale.
-            ocean_static_path: Optional path to a Zarr file containing static ocean variables
-                on the native grid. If provided, variables 'wet' (renamed to 'sea_surface_fraction')
-                and 'hfgeou' (geothermal heat flux) will be added to the processed dataset.
-                Default is None (no static variables added).
             spatial_filter_scale: Optional integer to override the automatic spatial filter
                 scale determination. When spatial filtering is performed (not --skip-spatial-filtering),
                 this value will be used instead of the scale inferred from the target grid name.
@@ -343,19 +337,6 @@ def om4(
             logger.info("validating preprocessing.")
             ds_processed_validate(ds_processed, deep=True)
 
-        if ocean_static_path is not None:
-            logger.info("adding static variables.")
-            ocean_static_names = ["wet", "hfgeou"]
-            ocean_static_renaming = {
-                "xh": "x",
-                "yh": "y",
-                "wet": "sea_surface_fraction",
-            }
-            ds_static = xr.open_zarr(ocean_static_path, consolidated=True, chunks={})[
-                ocean_static_names
-            ].rename(ocean_static_renaming)
-            ds_processed = xr.merge([ds_processed, ds_static])
-
         saved_attrs = {}
         for var in ds_processed.data_vars:
             saved_attrs[var] = ds_processed[var].attrs

data/ocean_preprocessing/make_norm_datasets.py

@@ -37,14 +37,10 @@
     # This ensures mean/std calculations only consider ocean points
     ds_masked = ds.copy()
 
-    # Variables that should not be masked (masks themselves, indices, static variables)
+    # Variables that should not be masked (masks themselves and indices)
     skip_vars = set()
     for var in ds.data_vars:
-        if (
-            var.startswith("mask_")
-            or var.startswith("idepth_")
-            or var in ["sea_surface_fraction", "hfgeou"]
-        ):
+        if var.startswith("mask_") or var.startswith("idepth_"):
             skip_vars.add(var)
 
     for var in ds.data_vars:

src/samudra/aggregator/inference/main.py

@@ -7,8 +7,8 @@
 import wandb
 import xarray as xr
 
-from samudra.constants import TensorMap
-from samudra.utils.data import Normalize, get_aggregator_dicts
+from samudra.constants import DataLayout
+from samudra.utils.data import BatchPreprocessor, get_aggregator_dicts
 from samudra.utils.output import ModelInferenceOutput
 from samudra.utils.wandb import Metrics, MetricsDict
 
@@ -30,8 +30,8 @@ def __init__(
         area_weights: torch.Tensor,
         wet: torch.Tensor,
         num_prognostic_channels: int,
-        normalize: Normalize,
-        tensor_map: TensorMap,
+        preprocessor: BatchPreprocessor,
+        data_layout: DataLayout,
         record_step_20: bool = True,
         log_global_mean_time_series: bool = True,
         log_global_mean_norm_time_series: bool = True,
@@ -47,8 +47,8 @@ def __init__(
             area_weights: Area weights for the data.
             wet: Wet mask for the data.
             num_prognostic_channels: Number of prognostic channels in the data.
-            normalize: Normalization helper for prognostic channels.
-            tensor_map: Mapping from prognostic variables to tensor channels.
+            preprocessor: Normalization helper for prognostic channels.
+            data_layout: Mapping from prognostic variables to tensor channels.
             record_step_20: Whether to record the mean of the 20th steps.
             log_global_mean_time_series: Whether to log global mean time series metrics.
             log_global_mean_norm_time_series: Whether to log the normalized global mean
@@ -103,8 +103,8 @@ def __init__(
             if name not in ["mean", "mean_norm"]
         }
         self._n_timesteps_seen = 0
-        self._normalize = normalize
-        self._tensor_map = tensor_map
+        self._preprocessor = preprocessor
+        self._data_layout = data_layout
         self.num_prognostic_channels = num_prognostic_channels
         self.hist = hist
         self.wet = wet
@@ -123,8 +123,8 @@ def record_batch(self, data: ModelInferenceOutput):
         assert data.prediction.shape[0] == total_len // (self.hist + 1)
         target_norm_dict, target_unnorm_dict = get_aggregator_dicts(
             data.target,
-            normalize=self._normalize,
-            tensor_map=self._tensor_map,
+            preprocessor=self._preprocessor,
+            data_layout=self._data_layout,
             wet=self.wet,
             long_rollout=True,
             input_type="prognostic",
@@ -133,8 +133,8 @@ def record_batch(self, data: ModelInferenceOutput):
         )
         gen_norm_dict, gen_unnorm_dict = get_aggregator_dicts(
             data.prediction,
-            normalize=self._normalize,
-            tensor_map=self._tensor_map,
+            preprocessor=self._preprocessor,
+            data_layout=self._data_layout,
             wet=self.wet,
             long_rollout=True,
             input_type="prognostic",
@@ -177,8 +177,8 @@ def record_initial_prognostic(
 
         data_norm_dict, data_unnorm_dict = get_aggregator_dicts(
             initial_prognostic,
-            normalize=self._normalize,
-            tensor_map=self._tensor_map,
+            preprocessor=self._preprocessor,
+            data_layout=self._data_layout,
             wet=self.wet,
             long_rollout=True,
             input_type="input",

src/samudra/aggregator/loss.py

@@ -5,19 +5,19 @@
 
 import torch
 
-from samudra.constants import TensorMap
+from samudra.constants import DataLayout
 
 
 def get_depth_loss_dict(
     label: str,
     loss_per_channel: torch.Tensor,
     *,
-    tensor_map: TensorMap,
+    data_layout: DataLayout,
 ) -> dict[str, torch.Tensor]:
     metrics = {}
-    for depth in tensor_map.DEPTH_SET:
+    for depth in data_layout.depths:
         metrics[f"{label}/loss/depth/depth_{depth}_loss"] = loss_per_channel[
-            tensor_map.DP_3D_IDX[depth]
+            data_layout.depth_indices[depth]
         ].mean()
     return metrics
 
@@ -26,12 +26,12 @@ def get_variable_loss_dict(
     label: str,
     loss_per_channel: torch.Tensor,
     *,
-    tensor_map: TensorMap,
+    data_layout: DataLayout,
 ) -> dict[str, torch.Tensor]:
     metrics = {}
-    for variable in tensor_map.VAR_SET:
+    for variable in data_layout.variables:
         metrics[f"{label}/loss/variable/{variable}_loss"] = loss_per_channel[
-            tensor_map.VAR_3D_IDX[variable]
+            data_layout.variable_indices[variable]
         ].mean()
     return metrics
 
@@ -40,23 +40,23 @@ def get_channel_loss_dict(
     label: str,
     loss_per_channel: torch.Tensor,
     *,
-    tensor_map: TensorMap,
+    data_layout: DataLayout,
     loss_name: str = "loss",
 ) -> dict[str, torch.Tensor]:
-    return get_channel_dict(label, loss_name, loss_per_channel, tensor_map=tensor_map)
+    return get_channel_dict(label, loss_name, loss_per_channel, data_layout=data_layout)
 
 
 def get_channel_loss_scale_dict(
     label: str,
     loss_scale_per_channel: torch.Tensor,
     *,
-    tensor_map: TensorMap,
+    data_layout: DataLayout,
 ) -> dict[str, torch.Tensor]:
     return get_channel_dict(
         label,
         "loss_scale",
         loss_scale_per_channel,
-        tensor_map=tensor_map,
+        data_layout=data_layout,
     )
 
 
@@ -65,9 +65,9 @@ def get_channel_dict(
     measure: str,
     per_channel: torch.Tensor,
     *,
-    tensor_map: TensorMap,
+    data_layout: DataLayout,
 ) -> dict[str, torch.Tensor]:
     metrics = {}
-    for i, channel in enumerate(tensor_map.prognostic_var_names):
+    for i, channel in enumerate(data_layout.prognostic_var_names):
         metrics[f"{prefix}/{measure}/channel/{channel}_{measure}"] = per_channel[i]
     return metrics