Add regridding from low resolution ML (.npy) prediction to high resolution restart

README.md

@@ -3,7 +3,6 @@
 Update NEMO ocean model restart files with machine learning predictions.
 
 ## Installation
-
 Start by installing a virtual environment and then:
 
 ```bash
@@ -13,11 +12,22 @@ pip install .
 ## Usage
 
 ```bash
-nemo-spinup-restart --restart_path /path/to/restarts \
-                    --radical RESTART_NAME \
-                    --mask_file /path/to/mask.nc \
-                    --prediction_path /path/to/predictions \
-                    --ocean_terms ocean_terms.yaml
+nemo-restart 
+  --restart_path /path/to/restarts \
+  --radical RESTART_NAME \
+  --mask_file /path/to/mask.nc \
+  --prediction_path /path/to/predictions \
+  --ocean_terms ocean_terms.yaml
+
+nemo-upscale upscale \\
+  --predictions-dir ./predictions \
+  --coarse-template ./1deg/restart_template.nc \
+  --coarse-mask ./1deg/mesh_mask.nc \
+  --coarse-namelist ./1deg/namelist_cfg \
+  --fine-template ./025deg/restart_template.nc \
+  --fine-mask ./025deg/mesh_mask.nc \
+  --output-dir ./generated \
+  --name C2
 ```
 
 ### Required Arguments

algorithm.md

@@ -0,0 +1,68 @@
+
+Key insight: Pre-existing template lat and lon values do not matter. Mesh mask is the authoritative grid definition and we overwrite the coordinates with mask.glamt and mask.gphit values. 
+
+---
+
+## Regrid.py Algorithm
+
+### `upscale_predictions()`
+1. Load numpy predictions (toce, soce, ssh) at specified time index
+2. Create coarse restart: `create_restart_from_predictions()`
+3. Regrid to fine resolution: `regrid_restart()`
+4. Return path to fine restart file
+
+### `create_restart_from_predictions()`
+1. Load template restart and mesh mask
+2. Extract depth from mask (`gdept_0`)
+3. Compute potential density using NEMO equation of state
+4. Populate restart fields: tb/tn, sb/sn, sshb/sshn, rhop
+5. Update metadata and save to NetCDF
+
+### `regrid_restart()`
+
+**Phase 1: Load & Prepare**
+1. Load coarse/fine restarts and masks
+2. Assign lon/lat coords to restarts from masks (glamt/gphit)
+3. Save intermediate: `restart_lr.nc`, `restart_hr_template.nc`
+
+**Phase 2: Extrapolate Coarse**
+4. Call `extrapolate_to_land(restart_lr, mask_lr)`
+5. Save intermediate: `restart_lr_extrap.nc`
+
+**Phase 3: Regrid**
+6. Create xESMF regridder (bilinear, nearest_s2d extrapolation)
+7. Apply regridding: `restart_hr = regridder(restart_lr_extrap)`
+8. Save intermediate: unmasked fine restart
+
+**Phase 4: Cleanup Coordinates**
+9. Rename lat→nav_lat, lon→nav_lon
+10. Drop x, y coordinate variables
+
+**Phase 5: Apply Fine Mask**
+11. Prepare fine mask (align nav_lev, drop x/y/time)
+12. Mask all variables (set land to 0.0)
+
+**Phase 6: Finalize**
+13. Zero all velocities (ub/un/vb/vn)
+14. Copy time metadata from coarse (kt, ndastp, adatrj, ntime)
+15. Copy timestep from fine template (rdt)
+16. Reorder variables to match template
+17. Update file metadata and save
+
+### `extrapolate_to_land()`
+1. Prepare mask (squeeze, drop x/y/time_counter, align nav_lev)
+2. Create 2D surface mask
+3. Apply mask to set land to NaN (3D and 4D variables)
+4. Extrapolate along x-dimension (nearest neighbor)
+5. Extrapolate along y-dimension (nearest neighbor)
+6. Return extrapolated restart
+
+**Note:** No vertical extrapolation performed (~12k NaNs remain in deep ocean)
+
+### Key Design Choices
+- **Coordinates from mask:** Always overwrite restart coords with glamt/gphit  
+- **xESMF requirements:** Needs lon/lat named coordinates for geographic interpolation
+- **Dimension-based masking:** `.where()` broadcasts on dimension names (y, x), not coord values.
+- **Conservative velocities:** Zero out for NEMO to recompute from density
+- **Two-stage extrapolation:** Fill land before regridding to avoid NaN propagation
+

examples/diffusion.md

@@ -0,0 +1,42 @@
+# Diffusion state restart file generation and evaluation
+
+## Data setup 
+
+TODO: Host data on spiritx
+
+## Regrid and upscale
+
+We need 1-degree, 0.25-degree restart file and mesh mask references. 
+
+```bash
+
+REFERENCE=./data/reference
+
+nemo-upscale upscale \
+    --predictions-dir $REFERENCE/diffusion_states/chamon_C2_clean/ \
+    --coarse-template $REFERENCE/1deg_restart/DINO_00000002_restart.nc \
+    --coarse-mask     $REFERENCE/1deg_restart/mesh_mask.nc \
+    --coarse-namelist $REFERENCE/1deg_restart/namelist_cfg \
+    --fine-template   $REFERENCE/025deg_restart/restart25_arch/DINO_10800000_restart.nc \
+    --fine-mask       $REFERENCE/025deg_restart/restart25_arch/mesh_mask.nc \
+    --name            C2 \
+    --time-index      4 \
+    --output-dir      ./generated
+```
+
+## Evaluate using spinup-evaluation
+
+Install spinup-evaluation as normal and create a gen-setup.yaml that looks as so:
+
+```gen-setup.yaml
+
+mesh_mask: $REFERENCE/025deg_restart/restart25_arch/mesh_mask.nc
+restart_files: 'generated_restart_C2'
+```
+
+```bash
+spinup-eval \
+  --sim-path ./generated/ \
+  --config ./gen-setup.yaml \
+  --mode restart
+```

examples/gen-setup.yaml

@@ -0,0 +1,5 @@
+# DINO-setup.yaml
+
+mesh_mask: /Users/matt/work/nemo/spinup-data/Gorce-data/Dinonline/restart0/mesh_mask_new.nc
+
+restart_files: 'generated_restart_C2_fine'

pyproject.toml

@@ -31,6 +31,8 @@ dependencies = [
     "matplotlib",
     "netcdf4",
     "pyyaml",
+    "xesmf",
+    "f90nml",
 ]
 
 [project.optional-dependencies]
@@ -41,7 +43,8 @@ dev = [
 ]
 
 [project.scripts]
-nemo-spinup-restart = "nemo_spinup_restart.cli:main"
+nemo-restart = "nemo_spinup_restart.cli:main"
+nemo-upscale = "nemo_spinup_restart.regrid_cli:main"   
 
 [tool.setuptools.packages.find]
 where = ["src"]