Add Option to Disable Vertical Shift in Nuth and Kääb Coregistration

xdem/coreg/affine.py

@@ -454,6 +454,7 @@ def _nuth_kaab_iteration_step(
     aspect: NDArrayf,
     res: tuple[int, int],
     params_fit_bin: InFitOrBinDict,
+    apply_vertical_shift: bool = True,
 ) -> tuple[tuple[float, float, float], float]:
     """
     Iteration step of Nuth and Kääb (2011), passed to the iterate_method function.
@@ -466,6 +467,7 @@ def _nuth_kaab_iteration_step(
     :param slope_tan: Array of slope tangent.
     :param aspect: Array of aspect.
     :param res: Resolution of DEM.
+    :param apply_vertical_shift: Whether to apply the vertical shift or not (default is True).
     """
 
     # Calculate the elevation difference with offsets
@@ -497,7 +499,7 @@ def _nuth_kaab_iteration_step(
     new_coords_offsets = (
         coords_offsets[0] + easting_offset * res[0],
         coords_offsets[1] + northing_offset * res[1],
-        float(vshift),
+        float(vshift) if apply_vertical_shift else 0,
     )
 
     # Compute statistic on offset to know if it reached tolerance
@@ -520,6 +522,7 @@ def nuth_kaab(
     params_random: InRandomDict,
     z_name: str,
     weights: NDArrayf | None = None,
+    apply_vertical_shift: bool = True,
     **kwargs: Any,
 ) -> tuple[tuple[float, float, float], int]:
     """
@@ -565,7 +568,14 @@ def nuth_kaab(
     res = _res(transform)
     # Iterate through method of Nuth and Kääb (2011) until tolerance or max number of iterations is reached
     assert sub_aux_vars is not None  # Mypy: dictionary cannot be None here
-    constant_inputs = (sub_dh_interpolator, sub_aux_vars["slope_tan"], sub_aux_vars["aspect"], res, params_fit_or_bin)
+    constant_inputs = (
+        sub_dh_interpolator,
+        sub_aux_vars["slope_tan"],
+        sub_aux_vars["aspect"],
+        res,
+        params_fit_or_bin,
+        apply_vertical_shift,
+    )
     final_offsets = _iterate_method(
         method=_nuth_kaab_iteration_step,
         iterating_input=initial_offset,
@@ -1287,6 +1297,7 @@ def __init__(
         bin_sizes: int | dict[str, int | Iterable[float]] = 72,
         bin_statistic: Callable[[NDArrayf], np.floating[Any]] = np.nanmedian,
         subsample: int | float = 5e5,
+        apply_vertical_shift: bool = True,
     ) -> None:
         """
         Instantiate a new Nuth and Kääb (2011) coregistration object.
@@ -1299,8 +1310,11 @@ def __init__(
         :param bin_sizes: Size (if integer) or edges (if iterable) for binning variables later passed in .fit().
         :param bin_statistic: Statistic of central tendency (e.g., mean) to apply during the binning.
         :param subsample: Subsample the input for speed-up. <1 is parsed as a fraction. >1 is a pixel count.
+        :param apply_vertical_shift: Whether to apply the vertical shift or not (default is True).
         """
 
+        self.apply_vertical_shift = apply_vertical_shift
+
         # Input checks
         _check_inputs_bin_before_fit(
             bin_before_fit=bin_before_fit, fit_optimizer=fit_optimizer, bin_sizes=bin_sizes, bin_statistic=bin_statistic
@@ -1390,6 +1404,7 @@ def _fit_rst_pts(
             params_fit_or_bin=params_fit_or_bin,
             max_iterations=self._meta["inputs"]["iterative"]["max_iterations"],
             tolerance=self._meta["inputs"]["iterative"]["tolerance"],
+            apply_vertical_shift=self.apply_vertical_shift,
         )
 
         # Write output to class