004 feature simple regridder

OceanOSSE/gridding/regridder_swap.py

@@ -0,0 +1,184 @@
+# ===================================================================
+# Copyright 2025 National Oceanography Centre
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied. See the License for the specific language governing
+# permissions and limitations under the License.
+# ===================================================================
+
+"""
+sampler_nearest_neighbour.py
+
+Description: Sampling module for OceanOSSE package.
+
+Created By: OceanOSSE Development Team (NOC, UK)
+"""
+
+# -- Import Dependencies -- #
+from __future__ import annotations
+
+import logging
+
+import xarray as xr
+import numpy as np
+
+from OceanOSSE.utils import import_class
+from OceanOSSE.gridding.regridder import Regridder
+
+logger = logging.getLogger(__name__)
+
+
+class SwapRegridder(Regridder):
+    """
+    Regridding class for synthetic ocean observations onto
+    the original model grid using climatology and exchanging profiles.
+
+    Parameters
+    ----------
+    target_grid : xarray.Dataset or None, optional
+        Dataset describing the target grid (coordinates, masks, etc.).
+    """
+
+    def __init__(
+        self,
+        target_grid: xr.Dataset | None = None,
+    ) -> None:
+        if target_grid is not None and not isinstance(target_grid, xr.Dataset):
+            raise TypeError("``target_grid`` must be an xarray.Dataset or None.")
+        self.target_grid = target_grid
+        # In future load these names from config
+        self.varlist = ['votemper', 'vosaline']
+        self.upper = 2000 # upper ocean threshold
+
+    def __repr__(self) -> str:
+        has_grid = self._target_grid is not None
+        return f"{type(self).__name__}(target_grid={'<Dataset>' if has_grid else None})"
+
+
+    def from_config(cls, config: dict) -> Self:
+        """
+        Construct a Regridder from the from the `[regridding]` table of
+        the .toml configuration file.
+
+        Parameters
+        ----------
+        config : dict
+            Configuration dictionary containing input parameters from .toml
+            configuration file.
+
+        Returns
+        -------
+        Self
+            Initialised Regridder instance.
+        """
+        return self
+
+
+    def regrid(self, ds_profile: xr.Dataset) -> xr.Dataset:
+        """
+        Regrid the synthetic observation dataset into the target grid.
+
+        Parameters
+        ----------
+        ds_profile : xarray.Dataset
+            Model profile dataset.
+
+        Returns
+        -------
+        xarray.Dataset
+            Dataset of synthetic observations placed into target grid.
+        """
+        ds_zeros = self.initialise_anomaly(ds_profile)
+
+        ds_anom = self.insert_profiles(ds_profile, ds_zeros)
+        # select data above 2000 m
+        ds_anom = ds_anom.where(ds_anom.depth > self.upper)
+
+        ds_out = ds_anom + self.target_grid
+
+        return ds_out
+
+
+    def initialise_anomaly(self, ds_profile):
+        """
+        Make a dataset of zeros on the same grid as the climatology.
+
+        Parameters
+        ----------
+        ds_profile : xarray.Dataset
+            Model profile dataset.
+
+        Returns
+        -------
+        xarray.Dataset
+            zeros on target grid.
+        """
+        # initialise a zero array to sum profile anomalies on
+        zero_anomaly = np.zeros((list(self.target_grid.sizes.values())))
+        ds_zeros = self.target_grid.copy(deep=True)
+
+        for var in self.varlist:
+            if var not in list(ds_profile.keys()):
+                raise ValueError(var + " is not in the profile dataset.")
+            if var not in list(self.target_grid.keys()):
+                raise ValueError(var + " is not in the climatology dataset.")
+
+
+            ds_zeros[var].data = zero_anomaly.copy()
+
+        return ds_zeros
+
+
+    def insert_profiles(self, ds_profile, ds_anom):
+        """
+        Regrid the synthetic observation dataset into the target grid.
+
+        Parameters
+        ----------
+        ds_profile : xarray.Dataset
+            Model profile dataset.
+        ds_anom : xarray.Dataset
+            zeros on target grid.
+
+        Returns
+        -------
+        xarray.Dataset
+            Dataset of summed anomaly model profiles.
+        """
+
+        # calculate month for profiles to put them on climatology
+        ds_profile = ds_profile.assign_coords(
+            month=("profile_id", 
+            ds_profile.t.dt.strftime("%m").astype(int).data))
+
+        n_profile = len(ds_profile.coords['profile_id'])
+
+        # Use indices in synthetic profile set to replace data in the 
+        # climatology with model data
+
+        for var in self.varlist:
+            # loop over profiles
+            for p in range(n_profile):
+                i_ind = ds_profile.coords['i'][p].to_numpy()
+                j_ind = ds_profile.coords['j'][p].to_numpy()
+                t_ind = ds_profile.coords['t'][p].to_numpy()
+                ps = ds_profile.coords['profile_id'][p].to_numpy()
+
+                profile = ds_profile[var].isel(profile_id=ps)
+
+                ds_anom[var].loc[
+                    dict(
+                        t=ds_profile.t.sel(profile_id=ps),
+                        j=ds_profile.j.sel(profile_id=ps),
+                        i=ds_profile.i.sel(profile_id=ps))
+                    ] += profile.data
+
+        return ds_anom
+
+

tests/unit/test_regridder.py

@@ -0,0 +1,183 @@
+# ===================================================================
+# Copyright 2025 National Oceanography Centre
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied. See the License for the specific language governing
+# permissions and limitations under the License.
+# ===================================================================
+"""
+test_sampler.py
+
+Description:
+This module includes unit tests for extracting profiles.
+
+Author:
+Benjamin Barton ([email protected])
+"""
+
+import pytest
+import datetime as dt
+import numpy as np
+import xarray as xr
+
+from OceanOSSE.gridding.regridder_swap import SwapRegridder
+
+
+def test_regrid(construct_ds, construct_profile_ds):
+    """
+    Test replacing profiles in climatology with model data.
+    """
+    ds_profile = construct_profile_ds
+    ds = construct_ds
+    ds_clim = climatology(ds)
+
+    regrid_data = SwapRegridder(ds_clim)
+    ds_model = regrid_data.regrid(ds_profile)
+
+    assert ((ds_model != ds_clim) 
+            & (ds_model.isel(i=3, j=5, t=31) 
+            == ds_profile.isel(profile_id=0)))
+
+def test_climatology(construct_ds):
+    """
+    Test producing a daily climatology.
+    """
+    ds = construct_ds
+    clim = climatology(ds)
+
+    ts = ds.votemper.mean(dim=["d", "j", "i"])
+    clim_mean = clim.votemper.mean(dim=["d", "j", "i"])
+
+    st_date1 = dt.datetime(2020, 5, 1)
+    st_date2 = dt.datetime(2021, 5, 1)
+    test_date1 = np.array([st_date1 + dt.timedelta(days=x) for x in range(31)])
+    test_date2 = np.array([st_date2 + dt.timedelta(days=x) for x in range(31)])
+    test_sec1 = np.array([(x - st_date1).days for x in test_date1])
+    test_sec2 = np.array([(x - st_date1).days for x in test_date2])
+    test_temp1 = 15 - (0 * 0.4) + (0 * 0.2) - (0 * 0.2) + (test_sec1 * 0.000005)
+    test_temp2 = 15 - (0 * 0.4) + (0 * 0.2) - (0 * 0.2) + (test_sec2 * 0.000005)
+    test_temp = np.sum(test_temp1 + test_temp2) / (2 * 31)
+    clim_day = clim.votemper.sel(t='2020-05-01').isel(d=0, j=0, i=0)
+
+    assert (np.isclose(ts.mean().to_numpy(), clim_mean.mean().to_numpy(), atol=1e-8)
+             & (clim_day.to_numpy() == test_temp))
+
+
+@pytest.fixture
+def construct_ds() -> xr.Dataset:
+    """
+    Build a dataset for testing.
+    """
+    lat = np.arange(0, 8)
+    lon = np.arange(0, 10)
+    depth = np.arange(0, 150, 10)
+    st_date = dt.datetime(2020, 5, 1)
+    num_days = 730
+    model_dates = np.array([st_date + dt.timedelta(days=x) for x in range(num_days)])
+    model_day = np.array([x for x in range(num_days)])
+
+    # Broadcast to 4D (time, depth, lat, lon)
+    t, d, y, x = np.meshgrid(model_day, depth, lat, lon, indexing='ij')
+
+    votemper = 15 - (y * 0.4) + (x * 0.2) - (d * 0.2) + (t * 0.000005)
+    vosaline = 33 + (y * 0.4) - (x * 0.2) + (d * 0.2) + (t * 0.000005)
+
+    # Build dataset
+    ds = xr.Dataset(
+        {
+            "votemper": (("t", "d", "j", "i"), votemper),
+            "vosaline": (("t", "d", "j", "i"), vosaline),
+
+            "lat": (("j", "i"), y[0, 0, :, :]),
+            "lon": (("j", "i"), x[0, 0, :, :]),
+            "depth": (("d", "j", "i"), d[0, :, :, :]),
+            "time": (("t"), t[:, 0, 0, 0])
+        },
+        coords={
+            "d": depth,
+            "j": lat,
+            "i": lon,
+            "t": model_dates
+        },
+    )
+
+    return ds
+
+
+@pytest.fixture
+def construct_profile_ds()  -> xr.Dataset:
+    d = np.arange(0, 150, 10)
+    profile_id = np.arange(2)
+
+    j = np.array([5, 6])
+    i = np.array([3, 8])
+
+    depth = np.tile(d[:, None], (1, profile_id.size))
+
+    # Time coordinate
+    st_date = dt.datetime(2020, 5, 1)
+    time = np.array([
+        dt.datetime(2020, 6, 1),
+        dt.datetime(2020, 7, 2),
+    ])
+    time_day = np.array([(x - st_date).days for x in time])
+
+    votemper = 15 - depth * 0.02 - j[None, :] * 0.2 + i[None, :] * 0.1 + (time_day * 0.000005)
+    vosaline = 33 + depth * 0.02 + j[None, :] * 0.2 - i[None, :] * 0.1 + (time_day * 0.000005)
+
+
+    ds = xr.Dataset(
+        data_vars={
+            "votemper": (("d", "profile_id"), votemper),
+            "vosaline": (("d", "profile_id"), vosaline),
+            "lat": (("profile_id",), j),
+            "lon": (("profile_id",), i),
+            "depth": (("d", "profile_id"), depth),
+        },
+        coords={
+            "d": d,
+            "profile_id": profile_id,
+            "t": (("profile_id",), time),
+            "j": (("profile_id",), j),
+            "i": (("profile_id",), i),
+        },
+    )
+
+    return ds
+
+
+def climatology(ds):
+    """
+    Calculate the climatology of the target grid.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Input time varying dataset.
+
+    Returns
+    -------
+    xarray.Dataset
+        Dataset of monthly means.
+    """
+    ds = ds.assign_coords(
+        month=("t", ds.t.dt.strftime("%m").astype(int).data)
+    )
+    # calculate climatology
+    ds_clim = ds.groupby('month').mean()
+
+    # tile the climatology data back over full time series
+    ds_clim_full = ds_clim.sel(month=ds.month)
+
+    # Remove not needed time dim from variables
+    for v in ["lat", "lon", "depth"]:
+        ds_clim_full[v] = ds_clim_full[v].isel(t=0, drop=True)
+    #ds_clim_full = ds_clim_full.drop_vars('month')
+
+    return ds_clim_full