Simplify the canonical data loading boundary

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

src/samudra/aggregator/main.py

@@ -18,23 +18,23 @@
 from samudra.aggregator.validate.reduced import MeanAggregator
 from samudra.aggregator.validate.snapshot import SnapshotAggregator
 from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator
-from samudra.constants import TensorMap
-from samudra.utils.data import Normalize
+from samudra.constants import DataLayout
+from samudra.utils.data import BatchPreprocessor
 
 
 class Aggregator:
     @staticmethod
-    def get_train_aggregator(tensor_map: TensorMap) -> TrainAggregator:
-        return TrainAggregator(tensor_map)
+    def get_train_aggregator(data_layout: DataLayout) -> TrainAggregator:
+        return TrainAggregator(data_layout)
 
     @staticmethod
     def get_validation_aggregator(
         metadata: dict[str, dict[str, str]],
         hist: int,
         area_weights: torch.Tensor,
         num_prognostic_channels: int,
-        tensor_map: TensorMap,
-        normalize: Normalize,
+        data_layout: DataLayout,
+        preprocessor: BatchPreprocessor,
         *,
         include_image_aggregators: bool = True,
     ) -> ValidateAggregator:
@@ -53,8 +53,8 @@ def get_validation_aggregator(
             val_aggregators,
             hist=hist,
             num_prognostic_channels=num_prognostic_channels,
-            tensor_map=tensor_map,
-            normalize=normalize,
+            data_layout=data_layout,
+            preprocessor=preprocessor,
         )
 
     @staticmethod
@@ -65,8 +65,8 @@ def get_inline_inference_aggregator(
         area_weights: torch.Tensor,
         wet: torch.Tensor,
         num_prognostic_channels: int,
-        tensor_map: TensorMap,
-        normalize: Normalize,
+        data_layout: DataLayout,
+        preprocessor: BatchPreprocessor,
         channel_mean_names: list[str] | None = None,
     ) -> InferenceEvaluatorAggregator:
         return InferenceEvaluatorAggregator(
@@ -76,8 +76,8 @@ def get_inline_inference_aggregator(
             area_weights=area_weights,
             wet=wet,
             num_prognostic_channels=num_prognostic_channels,
-            normalize=normalize,
-            tensor_map=tensor_map,
+            preprocessor=preprocessor,
+            data_layout=data_layout,
             record_step_20=(n_timesteps > 20),
             log_global_mean_time_series=False,
             log_global_mean_norm_time_series=False,
@@ -92,8 +92,8 @@ def get_standalone_inference_aggregator(
         area_weights: torch.Tensor,
         wet: torch.Tensor,
         num_prognostic_channels: int,
-        tensor_map: TensorMap,
-        normalize: Normalize,
+        data_layout: DataLayout,
+        preprocessor: BatchPreprocessor,
         channel_mean_names: list[str] | None = None,
     ) -> InferenceEvaluatorAggregator:
         return InferenceEvaluatorAggregator(
@@ -103,8 +103,8 @@ def get_standalone_inference_aggregator(
             area_weights=area_weights,
             wet=wet,
             num_prognostic_channels=num_prognostic_channels,
-            normalize=normalize,
-            tensor_map=tensor_map,
+            preprocessor=preprocessor,
+            data_layout=data_layout,
             record_step_20=(n_timesteps > 20),
             log_global_mean_time_series=True,
             log_global_mean_norm_time_series=True,

src/samudra/aggregator/train.py

@@ -10,7 +10,7 @@
     get_depth_loss_dict,
     get_variable_loss_dict,
 )
-from samudra.constants import TensorMap
+from samudra.constants import DataLayout
 from samudra.utils.distributed import all_reduce_mean
 from samudra.utils.output import TrainBatchOutput
 from samudra.utils.wandb import Metrics
@@ -19,8 +19,8 @@
 class TrainAggregator:
     """Aggregates train statistics for an epoch."""
 
-    def __init__(self, tensor_map: TensorMap):
-        self.tensor_map = tensor_map
+    def __init__(self, data_layout: DataLayout):
+        self.data_layout = data_layout
         self._n_batches = 0
         self._loss = torch.tensor(torch.nan)
         self._loss_per_channel = torch.tensor(torch.nan)
@@ -41,13 +41,13 @@ def get_logs(self, label: str = "train") -> Metrics:
 
         loss_per_channel = self._loss_per_channel / self._n_batches
         depth_loss_dict = get_depth_loss_dict(
-            label, loss_per_channel, tensor_map=self.tensor_map
+            label, loss_per_channel, data_layout=self.data_layout
         )
         var_loss_dict = get_variable_loss_dict(
-            label, loss_per_channel, tensor_map=self.tensor_map
+            label, loss_per_channel, data_layout=self.data_layout
         )
         channel_loss_dict = get_channel_loss_dict(
-            label, loss_per_channel, tensor_map=self.tensor_map
+            label, loss_per_channel, data_layout=self.data_layout
         )
         logs = {
             f"{label}/mean/loss": loss,

src/samudra/aggregator/validate/main.py

@@ -7,8 +7,8 @@
 
 from samudra.aggregator.train import TrainAggregator
 from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator
-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 ValBatchOutput
 from samudra.utils.wandb import Metrics, MetricsDict
 
@@ -22,14 +22,14 @@ def __init__(
         hist: int,
         num_prognostic_channels: int,
         *,
-        tensor_map: TensorMap,
-        normalize: Normalize,
+        data_layout: DataLayout,
+        preprocessor: BatchPreprocessor,
     ):
-        super().__init__(tensor_map)
+        super().__init__(data_layout)
         self._aggregators = aggregators
         self.hist = hist
         self.num_prognostic_channels = num_prognostic_channels
-        self.normalize = normalize
+        self.preprocessor = preprocessor
 
     # TODO(jder): we could remove this by moving from inheritance
     # to composition with the TrainAggregator functionality.
@@ -46,7 +46,7 @@ def record_validation_batch(self, batch: ValBatchOutput):
         if not self._aggregators:
             return
 
-        # Translate the GridContext mask by removing history.
+        # Translate the BatchGrid mask by removing history.
         target_data = batch.target_data  # [B, C*(hist+1), H, W]
         wet = batch.ctx.label_mask  # [C*(hist+1), H, W]
         assert wet.shape == target_data.shape[1:], (
@@ -66,8 +66,8 @@ def record_validation_batch(self, batch: ValBatchOutput):
         assert target_data.shape[1] == self.num_prognostic_channels
         target_data_dict, target_data_unnorm_dict = get_aggregator_dicts(
             target_data,
-            normalize=self.normalize,
-            tensor_map=self.tensor_map,
+            preprocessor=self.preprocessor,
+            data_layout=self.data_layout,
             wet=wet,
             long_rollout=False,
             input_type="prognostic",
@@ -77,8 +77,8 @@ def record_validation_batch(self, batch: ValBatchOutput):
 
         gen_data_dict, gen_data_unnorm_dict = get_aggregator_dicts(
             batch.gen_data,
-            normalize=self.normalize,
-            tensor_map=self.tensor_map,
+            preprocessor=self.preprocessor,
+            data_layout=self.data_layout,
             wet=wet,
             long_rollout=False,
             input_type="prognostic",
@@ -87,8 +87,8 @@ def record_validation_batch(self, batch: ValBatchOutput):
         )
         input_data_dict, input_data_unnorm_dict = get_aggregator_dicts(
             batch.input_data,
-            normalize=self.normalize,
-            tensor_map=self.tensor_map,
+            preprocessor=self.preprocessor,
+            data_layout=self.data_layout,
             wet=wet,
             long_rollout=False,
             input_type="input",