Fix and harden elevation-uncertainty propagation

- Enable the sqrt(2) "same precision" correction (Hugonnet 2022, Eq. 7-8) in the
  heteroscedasticity step; it was commented out, so precision_of_other="same" was a no-op.
  Default "finer" is unchanged, and an unsupported value now raises.
- Add an opt-in precoreg flag to co-register before inferring the error structure, so it is
  estimated on the aligned residual rather than the raw inputs. Default False (unchanged).
- Skip a non-converging Monte Carlo simulation with a warning instead of aborting the whole
  run, with a minimum-successful-simulations guard.
- Add tests on the Longyearbyen sample data (sqrt(2) scaling, precoreg equivalence and
  determinism, diverged-simulation skip, and the variogram-plotting invariant).

Author: dshean

tests/test_uncertainty/test_uncertainty_notebook.py

@@ -0,0 +1,206 @@
+"""
+Tests for the elevation-uncertainty module, derived from the experimental notebook
+``~/src/vantor/notebooks/xdem_uncertainty_test.ipynb`` (Longyearbyen sample data).
+
+They target three fixes to the uncertainty propagation:
+  1. the sqrt(2) "same precision" correction (Hugonnet et al., 2022, Eq. 7-8),
+  2. pre-coregistration before inferring the error structure,
+  3. correct variogram plotting via the (already public) variogram model function.
+
+Cropped sample data (``get_path_test``) and a small ``nsim`` are used for CI speed.
+"""
+
+from __future__ import annotations
+
+import logging
+import warnings
+from typing import Any, Callable
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import xdem
+from xdem import coreg
+from xdem.dem import DEM
+from xdem.uncertainty.uncertainty import _infer_uncertainty, _propag_uncertainty_coreg
+
+SEED = 42
+
+
+def _finite_array(elev: object) -> np.ndarray:
+    """Return a 1D float array of values from a Raster or PointCloud error field (masked -> NaN)."""
+    data = elev.data  # type: ignore[attr-defined]
+    if np.ma.isMaskedArray(data):
+        arr = data.filled(np.nan)
+    else:
+        arr = np.asarray(data, dtype=float)
+    return np.asarray(arr, dtype=float).ravel()
+
+
+def _load_dems() -> tuple[DEM, DEM]:
+    fn_ref = xdem.examples.get_path_test("longyearbyen_ref_dem")
+    fn_tba = xdem.examples.get_path_test("longyearbyen_tba_dem")
+    return DEM(fn_ref), DEM(fn_tba)
+
+
+class TestUncertaintyNotebook:
+
+    def test_precision_of_other_sqrt2(self) -> None:
+        """precision_of_other='same' divides the inferred error by sqrt(2) vs 'finer' (Eq. 7-8)."""
+        pytest.importorskip("skgstat")
+        warnings.filterwarnings("ignore", category=UserWarning)
+        dem_ref, dem_tba = _load_dems()
+
+        het_finer, _ = _infer_uncertainty(dem_ref, dem_tba, precision_of_other="finer", random_state=SEED)
+        het_same, _ = _infer_uncertainty(dem_ref, dem_tba, precision_of_other="same", random_state=SEED)
+
+        sig_finer = _finite_array(het_finer[0])
+        sig_same = _finite_array(het_same[0])
+
+        m = np.isfinite(sig_finer) & np.isfinite(sig_same)
+        assert m.sum() > 0
+        # Same random_state -> identical binning/fit, so the only difference is the sqrt(2) scaling.
+        np.testing.assert_allclose(sig_same[m], sig_finer[m] / np.sqrt(2), rtol=1e-5)
+
+    def test_precision_of_other_invalid_raises(self) -> None:
+        """An unsupported precision (e.g. a coarser 'worse' other) raises instead of silently acting like 'finer'."""
+        pytest.importorskip("skgstat")
+        dem_ref, dem_tba = _load_dems()
+        with pytest.raises(ValueError, match="precision_of_other"):
+            _infer_uncertainty(dem_ref, dem_tba, precision_of_other="worse")  # type: ignore[arg-type]
+
+    def test_precoreg_equivalence(self) -> None:
+        """precoreg=True equals a manual fit->apply->propagate(precoreg=False) with a matched RNG stream."""
+        pytest.importorskip("skgstat")
+        warnings.filterwarnings("ignore", category=UserWarning)
+        dem_ref, dem_tba = _load_dems()
+        method = coreg.LZD()
+        nsim = 5
+
+        # Auto: precoreg performs the initial fit+apply internally, consuming the shared rng first.
+        auto = _propag_uncertainty_coreg(
+            reference_elev=dem_ref,
+            to_be_aligned_elev=dem_tba,
+            coreg_method=method,
+            nsim=nsim,
+            error_applied_to="ref",
+            precoreg=True,
+            random_state=SEED,
+        )[0]
+
+        # Manual mirror: advance an rng with the SAME initial fit + apply, then propagate with
+        # precoreg=False passing the *advanced* generator so the random stream continues identically.
+        rng = np.random.default_rng(SEED)
+        c0 = method.copy()
+        c0.fit(reference_elev=dem_ref, to_be_aligned_elev=dem_tba, inlier_mask=None, random_state=rng)
+        dem_tba_align = c0.apply(dem_tba)
+        manual = _propag_uncertainty_coreg(
+            reference_elev=dem_ref,
+            to_be_aligned_elev=dem_tba_align,
+            coreg_method=method,
+            nsim=nsim,
+            error_applied_to="ref",
+            precoreg=False,
+            random_state=rng,
+        )[0]
+
+        pd.testing.assert_frame_equal(auto, manual, check_exact=False, rtol=1e-6, atol=1e-6)
+
+    def test_precoreg_deterministic(self) -> None:
+        """Two precoreg=True runs with the same seed give identical reports."""
+        pytest.importorskip("skgstat")
+        warnings.filterwarnings("ignore", category=UserWarning)
+        dem_ref, dem_tba = _load_dems()
+        method = coreg.NuthKaab()
+        nsim = 4
+
+        kw = dict(
+            reference_elev=dem_ref,
+            to_be_aligned_elev=dem_tba,
+            coreg_method=method,
+            nsim=nsim,
+            error_applied_to="tba",
+            precoreg=True,
+            random_state=123,
+        )
+        r1 = _propag_uncertainty_coreg(**kw)[0]
+        r2 = _propag_uncertainty_coreg(**kw)[0]
+        pd.testing.assert_frame_equal(r1, r2, check_exact=False, rtol=0, atol=0)
+
+    def test_diverged_simulation_is_skipped(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        caplog: Any,
+        assert_and_allow_log: Callable[..., None],
+    ) -> None:
+        """A simulation whose coregistration fails to converge is skipped (with a warning) rather than
+        aborting the whole Monte Carlo run; the remaining simulations still yield a finite report.
+
+        A real divergence (e.g. NuthKaab exhausting its subsample on a small/pre-aligned extent) is
+        numerics- and version-dependent, so the failure is forced deterministically here instead.
+        """
+        pytest.importorskip("skgstat")
+        warnings.filterwarnings("ignore", category=UserWarning)
+        dem_ref, dem_tba = _load_dems()
+
+        # Make exactly one simulation's coregistration raise, independent of platform/library versions.
+        original_fit = coreg.NuthKaab.fit
+        state = {"n": 0}
+
+        def flaky_fit(self: coreg.NuthKaab, *args: object, **kwargs: object) -> object:
+            state["n"] += 1
+            if state["n"] == 2:
+                raise ValueError("forced non-convergence for test")
+            return original_fit(self, *args, **kwargs)
+
+        monkeypatch.setattr(coreg.NuthKaab, "fit", flaky_fit)
+
+        with caplog.at_level(logging.WARNING):
+            summary = _propag_uncertainty_coreg(
+                reference_elev=dem_ref,
+                to_be_aligned_elev=dem_tba,
+                coreg_method=coreg.NuthKaab(),
+                nsim=4,
+                error_applied_to="tba",
+                precoreg=False,
+                random_state=SEED,
+            )[0]
+
+        # The skip is expected: confirm it was logged (and allow it past the global log-warning collector).
+        assert_and_allow_log(caplog, level=logging.WARNING, match="skipped")
+        for k in ("tx", "ty", "tz"):
+            assert np.isfinite(summary.loc[k, "std"])
+
+    def test_variogram_plotting_invariant(self) -> None:
+        """Fitted variogram (rises 0->sill) and correlation (falls 1->0) satisfy gamma = sill*(1 - rho).
+
+        This is the hack-free basis for plotting the empirical + fitted variogram together:
+        ``plot_variogram(corr_out[0], [xdem.spatialstats.get_variogram_model_func(corr_out[1])])``.
+        """
+        pytest.importorskip("skgstat")
+        warnings.filterwarnings("ignore", category=UserWarning)
+        dem_ref, dem_tba = _load_dems()
+
+        _, corr_out = _infer_uncertainty(dem_ref, dem_tba, random_state=SEED)
+        df_emp, params, corr_func = corr_out
+
+        # Reconstruct the variogram model function from the already-public helper (the correct input for
+        # plot_variogram, which expects a variogram rather than the returned correlation function).
+        vario_func = xdem.spatialstats.get_variogram_model_func(params)
+        total_sill = float(params["psill"].sum())
+        assert total_sill > 0
+        max_range = float(params["range"].max())
+
+        h = np.array([0.0, max_range, 5.0 * max_range])
+        rho = np.asarray(corr_func(h), dtype=float)
+        gamma = np.asarray(vario_func(h), dtype=float)
+
+        # Correlation starts at 1 and decays to ~0; variogram starts at 0 and rises to ~sill.
+        np.testing.assert_allclose(rho[0], 1.0, atol=1e-6)
+        np.testing.assert_allclose(gamma[0], 0.0, atol=1e-9 + 1e-6 * total_sill)
+        assert rho[-1] < 1e-2
+        np.testing.assert_allclose(gamma[-1], total_sill, rtol=1e-2)
+
+        # Exact invariant linking the two representations (what makes the plot render correctly).
+        np.testing.assert_allclose(gamma, total_sill * (1.0 - rho), rtol=1e-9, atol=1e-12)

xdem/uncertainty/uncertainty.py

@@ -81,6 +81,7 @@ def _propag_uncertainty_coreg(
     coreg_method: coreg.Coreg,
     nsim: int = 30,
     error_applied_to: Literal["ref", "tba"] = "tba",
+    precoreg: bool = False,
     inlier_mask: Raster | NDArrayb = None,
     random_state: int | np.random.Generator | None = None,
     kwargs_coreg_fit: dict[str, Any] | None = None,
@@ -96,6 +97,16 @@ def _propag_uncertainty_coreg(
     :param to_be_aligned_elev: To-be-aligned elevation.
     :param coreg_method: Coregistration method.
     :param nsim: Number of simulations to perform.
+    :param error_applied_to: Which input the simulated error field is applied to ("ref" or "tba").
+    :param precoreg: If True, co-register the inputs once before inferring the error structure and
+        running the simulations, so the error structure is estimated from the aligned residual rather
+        than from the raw (mis-aligned) inputs (which would otherwise conflate true error with
+        misalignment, terrain and real change). Requires an affine coregistration method. With
+        precoreg=True the reported per-parameter mean is the *residual* transform (~0 after alignment),
+        and the pre-alignment's own uncertainty is not itself propagated; this is a single pre-alignment
+        pass. Defaults to False (infer from the raw inputs, the previous behaviour). Note that
+        precoreg=True and False are not comparable for a given seed, as the pre-coregistration consumes
+        the random state.
     :param inlier_mask: Inlier mask (valid = True).
     :param random_state: Random state.
     :param kwargs_coreg_fit: Keyword arguments passed to `Coreg.fit`.
@@ -111,6 +122,23 @@ def _propag_uncertainty_coreg(
     # Define random state
     rng = np.random.default_rng(random_state)
 
+    # Optionally co-register once before inferring the error structure, then operate on the aligned
+    # inputs. Inferring from un-aligned inputs conflates true error with misalignment/terrain/real change;
+    # pre-aligning de-contaminates the inferred structure. This is a single pre-alignment pass (full
+    # iterative re-estimation, converging in 1-2 iterations, is a future enhancement). Only an affine
+    # method can be applied here, and the reported per-parameter mean becomes the residual transform.
+    if precoreg:
+        logging.info("Pre-coregistering inputs before inferring uncertainty...")
+        c_init = coreg_method.copy()
+        c_init.fit(
+            reference_elev=reference_elev,
+            to_be_aligned_elev=to_be_aligned_elev,
+            inlier_mask=inlier_mask,
+            random_state=rng,
+            **kwargs_coreg_fit,
+        )
+        to_be_aligned_elev = c_init.apply(to_be_aligned_elev)
+
     # First, infer uncertainty
     if error_applied_to == "ref":
         source_elev, other_elev = reference_elev, to_be_aligned_elev
@@ -146,16 +174,34 @@ def _propag_uncertainty_coreg(
             ref_elev = reference_elev
             tba_elev = to_be_aligned_elev + error_field
 
-        # Run coreg fit
+        # Run coreg fit. A single simulation can occasionally fail to converge (e.g. a subsampling
+        # method such as NuthKaab can exhaust its valid subsample on a small extent); skip it with a
+        # warning rather than aborting the whole Monte Carlo run.
         logging.info(f"  Running coregistration fit...")
         c = coreg_method.copy()  # Avoid carrying over the state over multiple simulations
-        c.fit(reference_elev=ref_elev, to_be_aligned_elev=tba_elev, inlier_mask=inlier_mask, random_state=rng,
-              **kwargs_coreg_fit)
+        try:
+            c.fit(reference_elev=ref_elev, to_be_aligned_elev=tba_elev, inlier_mask=inlier_mask, random_state=rng,
+                  **kwargs_coreg_fit)
+        except Exception as err:
+            logging.warning(f"  Simulation {i+1} of {nsim} failed to converge and was skipped: {err}")
+            continue
         df_it = _postproc_coreg_metadata(c)
         df_it["nsim"] = i + 1
         list_df.append(df_it)
         list_coreg.append(c)
 
+    # Require at least two successful simulations to estimate a standard deviation
+    if len(list_df) < 2:
+        raise RuntimeError(
+            f"Only {len(list_df)} of {nsim} simulations succeeded; cannot estimate uncertainty. The "
+            "coregistration repeatedly failed to converge (e.g. a subsampling method such as NuthKaab "
+            "on a small extent). Try subsample=1 (via kwargs_coreg_fit), a larger extent, a more robust "
+            "method (LZD/ICP), or precoreg=False."
+        )
+    if len(list_df) < nsim:
+        logging.warning(f"{nsim - len(list_df)} of {nsim} simulations were skipped after failing to "
+                        f"converge; uncertainty estimated from {len(list_df)} simulations.")
+
     # Finally, estimate errors for all the translations/rotations in the simulations
     df = pd.concat(list_df, ignore_index=True)
     t_r_names = ["tx", "ty", "tz", "rx", "ry", "rz"]
@@ -313,21 +359,22 @@ def _infer_uncertainty(
              Tuple of (Empirical variogram dataframe, Model parameters dataframe, Spatial error correlation function).
     """
 
+    # Validate the precision assumption: only 'finer' or 'same' are invertible from the difference alone.
+    if precision_of_other not in ("finer", "same"):
+        raise ValueError(
+            f"`precision_of_other` must be 'finer' or 'same', got {precision_of_other!r}. A coarser "
+            "('worse') other dataset is not supported, as the source error cannot be isolated from the "
+            "elevation difference alone. To characterize the less precise dataset, pass it as "
+            "`source_elev` and the more precise one as `other_elev` (precision_of_other='finer')."
+        )
+
     # Summarize approach steps
     approach_dict = {
         "H2022": {"heterosc": True, "multi_range": True},
         "R2009": {"heterosc": False, "multi_range": True},
         "Basic": {"heterosc": False, "multi_range": False},
     }
 
-    # # Difference the two datasets
-    # dh = _difference(source_elev, other_elev)
-
-    # # If the precision of the other Raster is the same, divide the dh values by sqrt(2)
-    # # See Equation 7 and 8 of Hugonnet et al. (2022)
-    # if precision_of_other == "same":
-    #     dh = dh / np.sqrt(2)
-
     logging.info(f"Starting heteroscedasticity inference.")
     # Heteroscedasticity
     sig_dh, df_bin, fun_bin = _infer_heteroscedasticity(
@@ -339,6 +386,7 @@ def _infer_uncertainty(
         z_name=z_name,
         subsample_hetesc=subsample_hetesc,
         spread_statistic=spread_estimator,
+        precision_of_other=precision_of_other,
     )
 
     logging.info(f"Starting spatial correlation inference.")
@@ -349,6 +397,7 @@ def _infer_uncertainty(
         inlier_mask=stable_terrain,
         errors=sig_dh,
         estimator=variogram_estimator,
+        precision_of_other=precision_of_other,
         random_state=random_state,
         list_models=vario_model,
         subsample=subsample_pairs_vario,
@@ -367,6 +416,8 @@ def _infer_heteroscedasticity(
     vector_mask_mode: Literal["inside", "outside"] = "inside",
     # Whether to infer a variable error (default) or constant
     heterosc: bool = True,
+    # Precision of the other dataset relative to the source (Hugonnet 2022, Eq. 7-8)
+    precision_of_other: Literal["finer", "same"] = "finer",
     # Heteroscedastic predictors
     hetesc_vars: (
         tuple[Raster | np.ndarray | str, ...]
@@ -452,6 +503,13 @@ def _infer_heteroscedasticity(
     # Elevation difference of the subsample
     dvalues_fit = rp1_fit - rp2_fit
 
+    # If the other dataset is of similar (not finer) precision, the difference variance is the sum of
+    # both errors (var(dh) = 2*sigma^2); divide by sqrt(2) to recover the single-dataset error and avoid
+    # double-counting. See Hugonnet et al. (2022), Eq. 7-8. Applied before binning so it propagates to
+    # both the heteroscedastic and the constant-error paths below.
+    if precision_of_other == "same":
+        dvalues_fit = dvalues_fit / np.sqrt(2)
+
     # 3) Perform binning and function fit on array inputs
 
     # 3A) If heteroscedastic, perform binning and fit
@@ -583,6 +641,7 @@ def _infer_spatial_correlation(
     vector_mask_mode: Literal["inside", "outside"] = "inside",
     errors: NDArrayf | Raster | None = None,
     estimator: Literal["matheron", "cressie", "genton", "dowd"] = "dowd",
+    precision_of_other: Literal["finer", "same"] = "finer",
     sampling: Literal["loglag", "random_xy"] = "loglag",
     subsample: int | float = 1,
     random_state: int | np.random.Generator | None = None,
@@ -654,6 +713,11 @@ def _infer_spatial_correlation(
     # Difference and standardize
     logging.info(f"  Step 2: Standardizing elevation differences...")
     dh_vals = rp1 - rp2
+    # Same-precision correction (Hugonnet 2022, Eq. 7-8), applied to the raw difference for consistency
+    # with the heteroscedasticity step and the legacy dem.py path. The returned correlation function is
+    # scale-invariant (rho = cov/total_sill), so this only affects the reported variogram magnitude.
+    if precision_of_other == "same":
+        dh_vals = dh_vals / np.sqrt(2)
     if errors is not None:
         dh_vals = dh_vals / aux_e["err"]