Update BlockKriging functions to use the n-nearest observations for i…

…nterpolation (#24)

* added merge to mergeplg so that submodules work properly

* added backup variogram if variogram estimation fail

* added neighbourhood kriging functionality for additive block kriging

* added neighborhood kriging functionality for KED

* adapted tests to account for neighbourhood block kriging

* Make IDW additive use radolan IDW code

* added comparison of adjusted fields to CML observations

src/mergeplg/__init__.py

@@ -7,6 +7,6 @@
 
 __version__ = "0.0.0"
 
-__all__ = ["__version__", "io", "radolan"]
+__all__ = ["__version__", "io", "radolan", "merge"]
 
-from . import io, radolan
+from . import io, merge, radolan

src/mergeplg/merge.py

@@ -457,7 +457,16 @@ def update(self, da_rad, da_cml=None, da_gauge=None):
         """
         self.update_(da_rad, da_cml=da_cml, da_gauge=da_gauge)
 
-    def adjust(self, da_rad, da_cml=None, da_gauge=None):
+    def adjust(
+        self,
+        da_rad,
+        da_cml=None,
+        da_gauge=None,
+        p=2,
+        idw_method="radolan",
+        nnear=8,
+        max_distance=60000,
+    ):
         """Adjust radar field for one time step.
 
         Adjust radar field for one time step. The function assumes that the
@@ -475,6 +484,14 @@ def adjust(self, da_rad, da_cml=None, da_gauge=None):
         da_gauge: xarray.DataArray
             Gauge observations. Must contain the coordinates for the rain gauge
             positions (lat, lon) as well as the projected coordinates (x, y).
+        p: float
+            IDW interpolation parameter
+        idw_method: str
+            by default "radolan"
+        nnear: int
+            number of neighbours to use for interpolation
+        max_distance: float
+            max distance allowed interpolation distance
 
         Returns
         -------
@@ -507,6 +524,10 @@ def adjust(self, da_rad, da_cml=None, da_gauge=None):
                 xr.where(da_rad_t > 0, da_rad_t, np.nan),  # function skips nan
                 diff[keep],
                 x0[keep, :],
+                p=p,
+                idw_method=idw_method,
+                nnear=nnear,
+                max_distance=max_distance,
             )
 
             # Replace nan with original radar data (so that da_rad nan is kept)
@@ -551,7 +572,9 @@ def update(self, da_rad, da_cml=None, da_gauge=None):
             da_rad, self.discretization, da_cml=da_cml, da_gauge=da_gauge
         )
 
-    def adjust(self, da_rad, da_cml=None, da_gauge=None, variogram="exponential"):
+    def adjust(
+        self, da_rad, da_cml=None, da_gauge=None, variogram="exponential", n_closest=8
+    ):
         """Adjust radar field for one time step.
 
         Adjust radar field for one time step. The function assumes that the
@@ -572,6 +595,8 @@ def adjust(self, da_rad, da_cml=None, da_gauge=None, variogram="exponential"):
         variogram: function or str
             If function: Must return expected variance given distance between
             observations. If string: Must be a valid variogram type in pykrige.
+        n_closest: int
+            Number of closest links to use for interpolation
 
         Returns
         -------
@@ -615,6 +640,7 @@ def adjust(self, da_rad, da_cml=None, da_gauge=None, variogram="exponential"):
                 diff[keep],
                 x0[keep, :],
                 variogram,
+                diff[keep].size - 1 if diff[keep].size <= n_closest else n_closest,
             )
 
             # Replace nan with original radar data (so that da_rad nan is kept)
@@ -663,13 +689,7 @@ def update(self, da_rad, da_cml=None, da_gauge=None):
         )
 
     def adjust(
-        self,
-        da_rad,
-        da_cml=None,
-        da_gauge=None,
-        variogram="exponential",
-        transform=None,
-        backtransform=None,
+        self, da_rad, da_cml=None, da_gauge=None, variogram="exponential", n_closest=8
     ):
         """Adjust radar field for one time step.
 
@@ -695,10 +715,8 @@ def adjust(
         variogram: function
             If function: Must return expected variance given distance between
             observations. If string: Must be a valid variogram type in pykrige.
-        transform: function
-            Transform rainfall distribution to Gaussian distributed data
-        backtransform: function
-            Gaussian distributed data to rainfall distribution
+        n_closest: int
+            Number of closest links to use for interpolation
 
         Returns
         -------
@@ -717,17 +735,11 @@ def adjust(
 
         # Check that that there is enough observations
         if keep.size > self.min_obs_:
-            # If transformation functions not provided, estimate it from obs.
-            if (transform is None) | (backtransform is None):
-                # Estimate Gamma distribution
-                param = merge_functions.estimate_transformation(obs[keep])
-                transform, backtransform, self.gamma_param = param
-
             # If variogram provided as string, estimate from ground obs.
             if isinstance(variogram, str):
                 # Estimate variogram
                 param = merge_functions.estimate_variogram(
-                    obs=transform(obs[keep]),
+                    obs=obs[keep],
                     x0=x0[keep],
                 )
 
@@ -743,13 +755,14 @@ def adjust(
             adjusted = merge_functions.merge_ked_blockkriging(
                 xr.where(da_rad_t > 0, da_rad_t, np.nan),  # function skips nan
                 rad[keep],
-                transform(obs)[keep],
+                obs[keep],
                 x0[keep],
                 variogram,
+                obs[keep].size - 1 if obs[keep].size <= n_closest else n_closest,
             )
 
             # Replace nan with original radar data (so that da_rad nan is kept)
-            adjusted = xr.where(np.isnan(adjusted), da_rad_t, backtransform(adjusted))
+            adjusted = xr.where(np.isnan(adjusted), da_rad_t, adjusted)
 
             # Re-assign timestamp and return
             return adjusted.assign_coords(time=time)

src/mergeplg/merge_functions.py

@@ -8,10 +8,13 @@
 import pykrige
 import xarray as xr
 from scipy import stats
-from sklearn.neighbors import KNeighborsRegressor
 
+from .radolan import idw
 
-def merge_additive_idw(da_rad, cml_diff, x0):
+
+def merge_additive_idw(
+    da_rad, cml_diff, x0, p=2, idw_method="radolan", nnear=8, max_distance=60000
+):
     """Merge CML and radar using an additive approach and the CML midpoint.
 
     Merges the CML and radar field by interpolating the difference between
@@ -26,6 +29,14 @@ def merge_additive_idw(da_rad, cml_diff, x0):
     x0: numpy.array
         Coordinates of CML midpoints given as [[cml_1_y, cml_1_x], ..
         [cml_n_y, cml_n_x] using the same order as cml_diff.
+    p: float
+        IDW interpolation parameter
+    idw_method: str
+        by default "radolan"
+    nnear: int
+        number of neighbours to use for interpolation
+    max_distance: float
+        max distance allowed interpolation distance
 
     Returns
     -------
@@ -49,14 +60,17 @@ def merge_additive_idw(da_rad, cml_diff, x0):
             [ygrid[~mask].reshape(-1, 1), xgrid[~mask].reshape(-1, 1)]
         )
 
-        # IDW interpolator kdtree, only supports IDW p=1
-        idw_interpolator = KNeighborsRegressor(
-            n_neighbors=cml_diff.size if cml_diff.size <= 8 else 8,
-            weights="distance",  # Use distance for setting weights
+        # IDW interpolator invdisttree
+        idw_interpolator = idw.Invdisttree(x0)
+        estimate = idw_interpolator(
+            q=coord_pred,
+            z=cml_diff,
+            nnear=cml_diff.size if cml_diff.size <= nnear else nnear,
+            p=p,
+            idw_method=idw_method,
+            max_distance=max_distance,
         )
-        idw_interpolator.fit(x0, cml_diff)
 
-        estimate = idw_interpolator.predict(coord_pred)
         shift[~mask] = estimate
 
     # create xarray object similar to ds_rad
@@ -78,7 +92,7 @@ def merge_additive_idw(da_rad, cml_diff, x0):
     return ds_rad_out.adjusted
 
 
-def merge_additive_blockkriging(da_rad, cml_diff, x0, variogram):
+def merge_additive_blockkriging(da_rad, cml_diff, x0, variogram, n_closest):
     """Merge CML and radar using an additive block kriging.
 
     Marges the provided radar field in ds_rad to CML observations by
@@ -98,6 +112,8 @@ def merge_additive_blockkriging(da_rad, cml_diff, x0, variogram):
     variogram: function
         A user defined python function defining the variogram. Takes a distance
         h and returns the expected variance.
+    n_closest: int
+        Number of closest links to use for interpolation
 
     Returns
     -------
@@ -124,9 +140,6 @@ def merge_additive_blockkriging(da_rad, cml_diff, x0, variogram):
     mat[-1, :-1] = np.ones(cov_block.shape[1])  # non-bias condition
     mat[:-1, -1] = np.ones(cov_block.shape[0])  # lagrange multipliers
 
-    # Calc the inverse, only dependent on geometry
-    a_inv = np.linalg.pinv(mat)
-
     # Skip radar pixels with np.nan
     mask = np.isnan(da_rad.data)
 
@@ -143,8 +156,15 @@ def merge_additive_blockkriging(da_rad, cml_diff, x0, variogram):
         delta_y = x0[:, 0] - ygrid_t[i]
         lengths = np.sqrt(delta_x**2 + delta_y**2)
 
+        # Get the n closest links
+        indices = np.argpartition(lengths.min(axis=1), n_closest)[:n_closest]
+        ind_mat = np.append(indices, mat.shape[0] - 1)
+
+        # Calc the inverse, only dependent on geometry
+        a_inv = np.linalg.pinv(mat[np.ix_(ind_mat, ind_mat)])
+
         # Estimate expected variance for all links
-        target = variogram(lengths).mean(axis=1)
+        target = variogram(lengths[indices]).mean(axis=1)
 
         # Add non bias condition
         target = np.append(target, 1)
@@ -153,7 +173,7 @@ def merge_additive_blockkriging(da_rad, cml_diff, x0, variogram):
         w = (a_inv @ target)[:-1]
 
         # Estimate rainfall amounts at location i
-        estimate[i] = cml_diff @ w
+        estimate[i] = cml_diff[indices] @ w
 
     # Store shift values
     shift[~mask] = estimate
@@ -177,7 +197,7 @@ def merge_additive_blockkriging(da_rad, cml_diff, x0, variogram):
     return ds_rad_out.adjusted_rainfall
 
 
-def merge_ked_blockkriging(da_rad, cml_rad, cml_obs, x0, variogram):
+def merge_ked_blockkriging(da_rad, cml_rad, cml_obs, x0, variogram, n_closest):
     """Merge CML and radar using an additive block kriging.
 
     Marges the provided radar field in ds_rad to CML observations by
@@ -199,6 +219,8 @@ def merge_ked_blockkriging(da_rad, cml_rad, cml_obs, x0, variogram):
     variogram: function
         A user defined python function defining the variogram. Takes a distance
         h and returns the expected variance.
+    n_closest: int
+        Number of closest links to use for interpolation
 
     Returns
     -------
@@ -225,9 +247,6 @@ def merge_ked_blockkriging(da_rad, cml_rad, cml_obs, x0, variogram):
     mat[:-2, -2] = np.ones(cov_block.shape[0])  # lagrange multipliers
     mat[:-2, -1] = cml_rad  # Radar drift
 
-    # Calc the inverse, only dependent on geometry (and radar for KED)
-    a_inv = np.linalg.pinv(mat)
-
     # Skip radar pixels with np.nan
     mask = np.isnan(da_rad.data)
 
@@ -245,7 +264,14 @@ def merge_ked_blockkriging(da_rad, cml_rad, cml_obs, x0, variogram):
         delta_y = x0[:, 0] - ygrid_t[i]
         lengths = np.sqrt(delta_x**2 + delta_y**2)
 
-        target = variogram(lengths).mean(axis=1)
+        # Get the n closest links
+        indices = np.argpartition(lengths.min(axis=1), n_closest)[:n_closest]
+        ind_mat = np.append(indices, [mat.shape[0] - 2, mat.shape[0] - 1])
+
+        # Calc the inverse, only dependent on geometry
+        a_inv = np.linalg.pinv(mat[np.ix_(ind_mat, ind_mat)])
+
+        target = variogram(lengths[indices]).mean(axis=1)
 
         target = np.append(target, 1)  # non bias condition
         target = np.append(target, rad_field_t[i])  # radar value
@@ -254,7 +280,7 @@ def merge_ked_blockkriging(da_rad, cml_rad, cml_obs, x0, variogram):
         w = (a_inv @ target)[:-2]
 
         # its then the sum of the CML values (eq 8, see paragraph after eq 15)
-        estimate[i] = cml_obs @ w
+        estimate[i] = cml_obs[indices] @ w
 
     rain[~mask] = estimate
 
@@ -476,19 +502,30 @@ def estimate_variogram(obs, x0, variogram_model="exponential"):
     if len(x0.shape) > 2:
         x0 = x0[:, :, int(x0.shape[1] / 2)]
 
-    # Fit variogram using pykrige
-    ok = pykrige.OrdinaryKriging(
-        x0[:, 1],  # x coordinate
-        x0[:, 0],  # y coordinate
-        obs,
-        variogram_model=variogram_model,
-    )
+    try:
+        # Fit variogram using pykrige
+        ok = pykrige.OrdinaryKriging(
+            x0[:, 1],  # x coordinate
+            x0[:, 0],  # y coordinate
+            obs,
+            variogram_model=variogram_model,
+        )
+
+        # construct variogram using pykrige
+        def variogram(h):
+            return ok.variogram_function(ok.variogram_model_parameters, h)
+
+        # Return variogram and parameters
+        return variogram, [ok.variogram_model_parameters, ok.variogram_function]
+
+    # If an error occurs just use a linear variogram
+    except ValueError:
 
-    # construct variogram using pykrige
-    def variogram(h):
-        return ok.variogram_function(ok.variogram_model_parameters, h)
+        def variogram(h):
+            return h
 
-    return variogram, [ok.variogram_model_parameters, ok.variogram_function]
+        # Return the linear variogram
+        return variogram, [1, variogram]
 
 
 def estimate_transformation(obs):