m2lines · fomo-bot · Jul 24, 2026
diff --git a/AGENTS.md b/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

diff --git a/configs/data/llc.yaml b/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

diff --git a/configs/data/om4.yaml b/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
diff --git a/configs/samudra_multi_om4/eval.yaml b/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.

diff --git a/configs/samudra_multi_om4/train.yaml b/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

diff --git a/configs/samudra_om4/model.yaml b/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
diff --git a/configs/samudra_om4_v1/model.yaml b/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
diff --git a/configs/samudra_om4_v2/model.yaml b/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
diff --git a/configs/samudra_om4_v2_highres/model.yaml b/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
diff --git a/configs/test/train_samudra_om4_v2_resume.yaml b/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
diff --git a/data/README.md b/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
-```
+```
diff --git a/data/ocean_preprocessing/__main__.py b/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

diff --git a/data/ocean_preprocessing/make_norm_datasets.py b/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: