diff --git a/.gitignore b/.gitignore index bcdc980b..f1168bdc 100644 --- a/.gitignore +++ b/.gitignore @@ -128,3 +128,8 @@ docs/source/sg_execution_times.rst *.vtu *.vtr *.vtk + +uv.lock +test.ipynb +benchmarking_local_krige.py +benchmarking_local_krige_results.csv \ No newline at end of file diff --git a/dev/local_kriging_numpy.py b/dev/local_kriging_numpy.py new file mode 100644 index 00000000..55079a68 --- /dev/null +++ b/dev/local_kriging_numpy.py @@ -0,0 +1,373 @@ +""" +Pure Python/NumPy prototype for *local* (moving-neighborhood) Kriging. +""" + +from __future__ import annotations + +import json + +import numpy as np +from scipy.spatial import cKDTree +from scipy.spatial.distance import cdist +from gstools import config + + +if config._GSTOOLS_CORE_AVAIL: + from gstools_core import calc_field_krige_local as calc_field_krige_local_gsc + + +# model classes GSTools-Core's CovModelSpec (covmodel_spec.rs) currently +# implements; the JSON "type" tag is just the lowercased class name for +# all three +_CORE_MODEL_TYPES = ("Gaussian", "Exponential", "Matern") + + +def _cov_model_to_json(model): + """ + Serialize a :any:`gstools.CovModel` to the JSON wire format the + GSTools-Core Rust backend expects (``CovModelSpec`` in + ``covmodel_spec.rs``: a ``{"type": ..., ...}``-tagged enum). + + Only ``var``/``len_scale``/``nugget`` (and, for Matern, ``nu``) travel + over the wire + """ + if model.name not in _CORE_MODEL_TYPES: + raise NotImplementedError( + f"calc_field_krige_local: GSTools-Core does not (yet) support " + f"the '{model.name}' covariance model; supported: " + f"{_CORE_MODEL_TYPES}." + ) + spec = { + "type": model.name.lower(), + "var": model.var, + "len_scale": model.len_scale, + "nugget": model.nugget, + } + if model.name == "Matern": + spec["nu"] = model.nu + return json.dumps(spec) + + +def calc_field_krige_local( + cond_pos, + cond_val, + target_pos, + model, + cond_err, + drift_cond, + drift_target, + unbiased, + exact, + local_radius, + num_threads=None, +): + if config.USE_GSTOOLS_CORE and config._GSTOOLS_CORE_AVAIL: + calc_field_krige_local = calc_field_krige_local_gsc + model_arg = _cov_model_to_json(model) + else: + calc_field_krige_local = calc_field_krige_local_python + model_arg = model + return calc_field_krige_local( + cond_pos, + cond_val, + target_pos, + model_arg, + cond_err, + drift_cond, + drift_target, + unbiased, + exact, + local_radius, + num_threads=num_threads, + ) + +def local_krige( + krige, + pos, + local_radius, + mesh_type="unstructured", + ext_drift=None, + return_var=True, +): + """ + Evaluate local (moving-neighborhood) kriging at the given positions. + + Parameters + ---------- + krige : :any:`gstools.krige.Krige` + An already-conditioned Krige instance (Simple/Ordinary/Universal/ + ExtDrift/Detrended). Its ``model``, ``cond_pos``, ``cond_val``, + drift functions, external drift, ``unbiased``, ``exact`` and + ``cond_err`` settings are reused as-is. ``pseudo_inv`` is *not* + used here (see notes in :any:`calc_field_krige_local`). + pos : :class:`list` + Position tuple of the target points (x, [y, z]). + local_radius : :class:`float` + Search radius (in the model's isometrized/isotropic distance): + every conditioning point within this distance of a target + point enters its local kriging system. This is the *only* + neighborhood-selection criterion -- there is no separate + neighbor-count cap. + mesh_type : :class:`str`, optional + 'structured' / 'unstructured'. Default: "unstructured" + ext_drift : :class:`numpy.ndarray` or :any:`None`, optional + External drift values at the target positions (only relevant + for :any:`gstools.krige.ExtDrift`). Same semantics as + ``Krige.__call__``'s ``ext_drift`` argument. + return_var : :class:`bool`, optional + Whether to also return the kriging error variance. Default: True + + Returns + ------- + field : :class:`numpy.ndarray` + The local-kriging field (raw, i.e. *not* post-processed with + mean/normalizer/trend -- see :any:`Krige.post_field` for that). + krige_var : :class:`numpy.ndarray`, optional + The kriging error variance (only if ``return_var`` is True). + """ + model = krige.model # Kovarianzmodell + iso_cond = krige._krige_pos # (dim, cond_no), isometrized cond. pos. + iso_targ, shape = krige.pre_pos(pos, mesh_type) # (dim, pnt_cnt), isometrized target pos. + pnt_cnt = iso_targ.shape[1] + cond_no = krige.cond_no + + # cond values already normalized/detrended/zero-mean (without padding) + cond_val = krige._krige_cond[:cond_no] + + cond_err = krige.cond_err + cond_err_arr = ( + np.full(cond_no, cond_err, dtype=np.double) + if np.isscalar(cond_err) + else np.asarray(cond_err, dtype=np.double) + ) + + # drift terms as plain arrays, evaluated once -- see docstrings there + drift_cond = _calc_cond_drift(krige) + ext_drift = krige._pre_ext_drift(pnt_cnt, ext_drift) + drift_target = _calc_target_drift(krige, iso_targ, ext_drift, pnt_cnt) + + # ================================================================ + # RUST BOUNDARY: in production this becomes a single call to + # gstools_core.krige.calc_field_krige_local(...) + # ================================================================ + field, error = calc_field_krige_local( + cond_pos=iso_cond, + cond_val=cond_val, + target_pos=iso_targ, + model=model, # stand-in for &CovModelCore + cond_err=cond_err_arr, + drift_cond=drift_cond, + drift_target=drift_target, + unbiased=krige.unbiased, + exact=krige.exact, + local_radius=local_radius, + num_threads=None, + ) + # ================================================================ + # END RUST BOUNDARY + # ================================================================ + + field = np.reshape(field, shape) + if return_var: + krige_var = np.reshape(np.maximum(model.sill - error, 0), shape) + return field, krige_var + return field + + +def _calc_cond_drift(krige): + """ + Evaluate all drift terms at the conditioning points *once*, as a + plain array -- mirrors ``Krige.cond_drift`` (there cached on the + instance; here just computed fresh per :any:`local_krige` call, + since this prototype has no persistent object to cache on). + + Internal (functional) drift and external drift are unified into a + single ``(p, cond_no)`` array, ``p = int_drift_no + ext_drift_no``: + external drift is already array-shaped, so it is simply stacked + below the evaluated functional drift rows. This way Universal- and + External-Drift-Kriging need no callback at all on the Rust side -- + the local solver only ever indexes into ``drift_cond`` with the + neighbor indices. + + Parameters + ---------- + krige : :any:`gstools.krige.Krige` + + Returns + ------- + :class:`numpy.ndarray` + Shape ``(drift_no, cond_no)``. + """ + cond_pos = np.asarray(krige.cond_pos) # (dim, cond_no), original frame + int_functions = krige.drift_functions + int_no = len(int_functions) + ext_no = krige.ext_drift_no + drift_no = int_no + ext_no + + drift_cond = np.empty((drift_no, krige.cond_no), dtype=np.double) + for i, f in enumerate(int_functions): + drift_cond[i] = f(*cond_pos) + if ext_no > 0: + drift_cond[int_no:] = krige.cond_ext_drift + return drift_cond + + +def _calc_target_drift(krige, iso_targ, ext_drift, pnt_cnt): + """ + Evaluate all drift terms at the target positions *once*, as a + plain array -- mirrors ``Krige._calc_target_drift``: computed a + single time for the whole call, instead of re-evaluating drift + functions for every neighbor inside the (future Rust) solve loop. + + Parameters + ---------- + krige : :any:`gstools.krige.Krige` + iso_targ : (dim, pnt_cnt) ndarray + Isometrized target positions. + ext_drift : :class:`numpy.ndarray` + Preprocessed external drift values at the target positions + (output of ``krige._pre_ext_drift``). + pnt_cnt : :class:`int` + + Returns + ------- + :class:`numpy.ndarray` + Shape ``(drift_no, pnt_cnt)``. + """ + int_functions = krige.drift_functions + int_no = len(int_functions) + ext_no = krige.ext_drift_no + drift_no = int_no + ext_no + + drift_target = np.empty((drift_no, pnt_cnt), dtype=np.double) + if int_no > 0: + # drift functions live in the original (anisotropic) frame + targ_pos = krige.model.anisometrize(iso_targ) + for i, f in enumerate(int_functions): + drift_target[i] = f(*targ_pos) + if ext_no > 0: + drift_target[int_no:] = ext_drift + return drift_target + + +def calc_field_krige_local_python( + cond_pos, + cond_val, + target_pos, + model, + cond_err, + drift_cond, + drift_target, + unbiased, + exact, + local_radius, + num_threads=None, +): + """ + NumPy stand-in for the planned Rust extension function + ``gstools_core.krige.calc_field_krige_local`` (see module docstring + for the exact Rust signature this mirrors). + + Everything in this function -- building the KD-tree once, then per + target point finding neighbors within ``local_radius``, assembling + and solving the local kriging system, and accumulating field/error + -- is what moves into ``gstools-core/src/krige.rs``. There the + ``for j in range(pnt_cnt)`` loop below becomes a ``rayon`` + ``par_iter``. + + Note this function only ever receives plain arrays and scalars, no + Python callables and no ``Krige``/``CovModel`` Python objects with + behaviour beyond ``covariance``/``cov_nugget`` -- ``model`` stands + in for the future ``&CovModelCore`` (a lightweight, serializable + model spec on the Rust side; ``as_core_model(model)`` would build + it in the Python prep step). + + Unlike global kriging, this does *not* use ``Krige``'s + ``pseudo_inv``/``pseudo_inv_type``: local systems are typically + small and (with nugget > 0) well conditioned, so this solves + directly (``numpy.linalg.solve``, i.e. the NumPy equivalent of the + LDL^T decomposition sketched for the Rust core) instead of forming + a pseudo-inverse. This matches the planned Rust signature, which + has no ``pseudo_inv`` parameter at all. + + Parameters + ---------- + cond_pos : (d, n) ndarray -- isometrized conditioning positions + cond_val : (n,) ndarray -- unpadded, centered conditioning values + target_pos : (d, m) ndarray -- isometrized target positions + model : :any:`gstools.CovModel` -- stand-in for ``&CovModelCore`` + cond_err : (n,) ndarray + drift_cond : (p, n) ndarray -- may have p == 0 + drift_target : (p, m) ndarray + unbiased : :class:`bool` + exact : :class:`bool` + local_radius : :class:`float` + Search radius (isometrized distance). Every conditioning point + within this distance of a target point enters its local + system -- the only neighborhood-selection criterion, no count + cap. + num_threads : :class:`int` or :any:`None` + Unused by this sequential NumPy stand-in; kept for signature + parity with the future Rust call. + + Returns + ------- + field : (m,) ndarray -- raw, not yet post-processed + error : (m,) ndarray -- raw RHS . w, not yet ``sill - error`` + """ + pnt_cnt = target_pos.shape[1] + drift_no = drift_cond.shape[0] + pad = drift_no + int(unbiased) + + tree = cKDTree(cond_pos.T) + # Diese Querry ist halt nur abhängig vom Radius. Es könnte also sein dass kein Nachbar gefunden wird + # Damit müsste man irgendwie noch umgehen + neighbor_lists = tree.query_ball_point(target_pos.T, r=local_radius) + + cov_fct = model.cov_nugget if exact else model.covariance + + field = np.empty(pnt_cnt, dtype=np.double) + error = np.empty(pnt_cnt, dtype=np.double) + + for j in range(pnt_cnt): # <- becomes `par_iter` over target points in Rust + nbrs = np.asarray(neighbor_lists[j], dtype=np.intp) + k = nbrs.size + if k == 0: + raise ValueError( + f"local_krige: no conditioning points within " + f"local_radius={local_radius} for target point {j}. " + "Increase local_radius." + ) + sys_size = k + pad + nbr_pos = cond_pos[:, nbrs] # (d, k) + + # --- local LHS (kriging matrix) --- + lhs = np.zeros((sys_size, sys_size), dtype=np.double) + lhs[:k, :k] = model.covariance(cdist(nbr_pos.T, nbr_pos.T)) + lhs[np.diag_indices(k)] += cond_err[nbrs] + if unbiased: + lhs[k, :k] = 1 + lhs[:k, k] = 1 + if drift_no: + lhs[-drift_no:, :k] = drift_cond[:, nbrs] + lhs[:k, -drift_no:] = drift_cond[:, nbrs].T + + # --- local RHS (kriging vector) --- + rhs = np.empty(sys_size, dtype=np.double) + dists = np.linalg.norm(nbr_pos - target_pos[:, j : j + 1], axis=0) + rhs[:k] = cov_fct(dists) + if unbiased: + rhs[k] = 1 + if drift_no: + rhs[-drift_no:] = drift_target[:, j] + + # --- solve & accumulate --- + try: + weights = np.linalg.solve(lhs, rhs) + except np.linalg.LinAlgError: + weights = np.linalg.lstsq(lhs, rhs, rcond=None)[0] + local_cond = np.concatenate([cond_val[nbrs], np.zeros(pad)]) + field[j] = local_cond @ weights + error[j] = rhs @ weights + + return field, error diff --git a/dev/test_local_kriging_numpy.py b/dev/test_local_kriging_numpy.py new file mode 100644 index 00000000..5ae4c91d --- /dev/null +++ b/dev/test_local_kriging_numpy.py @@ -0,0 +1,243 @@ +""" +Validation script for the pure NumPy local-kriging draft. + +Checks the core assumption of local_kriging_numpy.local_krige: with a +``local_radius`` big enough to cover the whole domain (so every +conditioning point is a "neighbor" of every target point), it must +reproduce global kriging exactly (up to floating point noise), for +Simple, Ordinary and Universal kriging, with and without +anisotropy/rotation. + +Run with: + uv run pytest dev/test_local_kriging_numpy.py -v +""" + +import numpy as np +import pytest + +import gstools as gs +from local_kriging_numpy import local_krige + +# radius comfortably bigger than any pairwise distance the test domains +# below can produce -> every conditioning point is always a neighbor +FULL_RADIUS = 1e6 + + +def make_cond(dim, seed=1): + rng = np.random.RandomState(seed) + n = 20 + cond_pos = [rng.uniform(0, 10, n) for _ in range(dim)] + cond_val = rng.uniform(-1, 1, n) + return cond_pos, cond_val + + +def make_targets(dim, seed=2): + rng = np.random.RandomState(seed) + n = 30 + return [rng.uniform(-2, 12, n) for _ in range(dim)] + + +MODELS = [gs.Gaussian, gs.Exponential, gs.Matern] +DIMS = [1, 2, 3] + + +@pytest.mark.parametrize("Model", MODELS) +@pytest.mark.parametrize("dim", DIMS) +def test_simple_matches_global(Model, dim): + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = Model(dim=dim, var=1.5, len_scale=3, nugget=0.1) + krige = gs.krige.Simple(model, cond_pos, cond_val, mean=0.3) + + field_g, var_g = krige(pos, return_var=True, post_process=False, store=False) + field_l, var_l = local_krige(krige, pos, local_radius=FULL_RADIUS) + + np.testing.assert_allclose(field_l, field_g, atol=1e-8, rtol=1e-6) + np.testing.assert_allclose(var_l, var_g, atol=1e-8, rtol=1e-6) + + +@pytest.mark.parametrize("Model", MODELS) +@pytest.mark.parametrize("dim", DIMS) +def test_ordinary_matches_global(Model, dim): + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = Model(dim=dim, var=1.5, len_scale=3, nugget=0.1) + krige = gs.krige.Ordinary(model, cond_pos, cond_val) + + field_g, var_g = krige(pos, return_var=True, post_process=False, store=False) + field_l, var_l = local_krige(krige, pos, local_radius=FULL_RADIUS) + + np.testing.assert_allclose(field_l, field_g, atol=1e-8, rtol=1e-6) + np.testing.assert_allclose(var_l, var_g, atol=1e-8, rtol=1e-6) + + +@pytest.mark.parametrize("Model", MODELS) +@pytest.mark.parametrize("dim", DIMS) +def test_universal_matches_global(Model, dim): + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = Model(dim=dim, var=1.5, len_scale=3, nugget=0.1) + krige = gs.krige.Universal(model, cond_pos, cond_val, drift_functions="linear") + + field_g, var_g = krige(pos, return_var=True, post_process=False, store=False) + field_l, var_l = local_krige(krige, pos, local_radius=FULL_RADIUS) + + np.testing.assert_allclose(field_l, field_g, atol=1e-6, rtol=1e-6) + np.testing.assert_allclose(var_l, var_g, atol=1e-6, rtol=1e-6) + + +def test_anisotropic_rotated_matches_global(): + dim = 2 + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = gs.Gaussian( + dim=dim, var=2.0, len_scale=3, anis=0.5, angles=0.7, nugget=0.05 + ) + krige = gs.krige.Ordinary(model, cond_pos, cond_val) + + field_g, var_g = krige(pos, return_var=True, post_process=False, store=False) + field_l, var_l = local_krige(krige, pos, local_radius=FULL_RADIUS) + + np.testing.assert_allclose(field_l, field_g, atol=1e-8, rtol=1e-6) + np.testing.assert_allclose(var_l, var_g, atol=1e-8, rtol=1e-6) + + +def test_exact_matches_global(): + dim = 2 + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = gs.Gaussian(dim=dim, var=2.0, len_scale=3, nugget=0.2) + krige = gs.krige.Ordinary(model, cond_pos, cond_val, exact=True) + + field_g, var_g = krige(pos, return_var=True, post_process=False, store=False) + field_l, var_l = local_krige(krige, pos, local_radius=FULL_RADIUS) + + np.testing.assert_allclose(field_l, field_g, atol=1e-8, rtol=1e-6) + np.testing.assert_allclose(var_l, var_g, atol=1e-8, rtol=1e-6) + + +def test_reproduces_cond_values_at_cond_pos(): + """Sanity check independent of global kriging: with exact=True the + local field must reproduce the conditioning values at the + conditioning points themselves (like normal kriging does).""" + dim = 2 + cond_pos, cond_val = make_cond(dim) + model = gs.Gaussian(dim=dim, var=2.0, len_scale=3, nugget=0.0) + krige = gs.krige.Ordinary(model, cond_pos, cond_val, exact=True) + + field_l, _ = local_krige(krige, cond_pos, local_radius=4.0) + np.testing.assert_allclose(field_l, cond_val, atol=1e-6) + + +def test_error_decreases_as_radius_increases(): + """Not a strict correctness check, but a sanity check: as + local_radius grows towards the full domain, the local field should + converge monotonically (in RMSE) towards the global field.""" + dim = 2 + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = gs.Gaussian(dim=dim, var=1.5, len_scale=2, nugget=0.1) + krige = gs.krige.Ordinary(model, cond_pos, cond_val) + field_g, _ = krige(pos, return_var=True, post_process=False, store=False) + + rmses = [] + for r in (3.5, 5.0, 7.0, 9.0, FULL_RADIUS): + field_l, _ = local_krige(krige, pos, local_radius=r) + rmses.append(np.sqrt(np.mean((field_l - field_g) ** 2))) + + assert rmses[-1] < 1e-6 + assert all(a >= b - 1e-12 for a, b in zip(rmses, rmses[1:])) + + +def test_local_radius_restricts_neighborhood(): + """A small radius must actually cut down which neighbors are used: + field should differ from the full-radius case and the error + variance should be larger (less information nearby).""" + dim = 2 + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = gs.Gaussian(dim=dim, var=1.5, len_scale=3, nugget=0.1) + krige = gs.krige.Ordinary(model, cond_pos, cond_val) + + field_unrestricted, var_unrestricted = local_krige( + krige, pos, local_radius=FULL_RADIUS + ) + field_restricted, var_restricted = local_krige( + krige, pos, local_radius=3.5 + ) + + assert not np.allclose(field_unrestricted, field_restricted) + assert np.all(var_restricted >= var_unrestricted - 1e-12) + + +def test_local_radius_raises_when_no_neighbor_in_radius(): + dim = 2 + cond_pos, cond_val = make_cond(dim) + model = gs.Gaussian(dim=dim, var=1.5, len_scale=3, nugget=0.1) + krige = gs.krige.Ordinary(model, cond_pos, cond_val) + + far_away = [np.array([1e6]), np.array([1e6])] + with pytest.raises(ValueError, match="no conditioning points within"): + local_krige(krige, far_away, local_radius=1.0) + + +def test_ext_drift_matches_global(): + dim = 2 + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = gs.Gaussian(dim=dim, var=1.5, len_scale=3, nugget=0.1) + + rng = np.random.RandomState(3) + cond_ext = rng.uniform(-1, 1, len(cond_val)) + targ_ext = rng.uniform(-1, 1, len(pos[0])) + + krige = gs.krige.ExtDrift(model, cond_pos, cond_val, cond_ext) + + field_g, var_g = krige( + pos, ext_drift=targ_ext, return_var=True, post_process=False, store=False + ) + field_l, var_l = local_krige( + krige, pos, local_radius=FULL_RADIUS, ext_drift=targ_ext + ) + + np.testing.assert_allclose(field_l, field_g, atol=1e-6, rtol=1e-6) + np.testing.assert_allclose(var_l, var_g, atol=1e-6, rtol=1e-6) + + +def test_ext_drift_and_functional_drift_combined(): + """Sanity check for the unified drift array: functional (Universal) + drift together with external drift in the same system, compared + directly against a global Krige base instance configured the same + way (drift_functions + ext_drift).""" + dim = 2 + cond_pos, cond_val = make_cond(dim) + pos = make_targets(dim) + model = gs.Gaussian(dim=dim, var=1.5, len_scale=3, nugget=0.1) + + rng = np.random.RandomState(4) + cond_ext = rng.uniform(-1, 1, len(cond_val)) + targ_ext = rng.uniform(-1, 1, len(pos[0])) + + krige = gs.krige.Krige( + model, + cond_pos, + cond_val, + drift_functions="linear", + ext_drift=cond_ext, + ) + + field_g, var_g = krige( + pos, ext_drift=targ_ext, return_var=True, post_process=False, store=False + ) + field_l, var_l = local_krige( + krige, pos, local_radius=FULL_RADIUS, ext_drift=targ_ext + ) + + np.testing.assert_allclose(field_l, field_g, atol=1e-6, rtol=1e-6) + np.testing.assert_allclose(var_l, var_g, atol=1e-6, rtol=1e-6) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/src/gstools/krige/__init__.py b/src/gstools/krige/__init__.py index 66d03246..e2aa3cdf 100644 --- a/src/gstools/krige/__init__.py +++ b/src/gstools/krige/__init__.py @@ -18,6 +18,14 @@ """ from gstools.krige.base import Krige +from gstools.krige.local import ( + LocalDetrended, + LocalExtDrift, + LocalKrige, + LocalOrdinary, + LocalSimple, + LocalUniversal, +) from gstools.krige.methods import ( Detrended, ExtDrift, @@ -26,4 +34,4 @@ Universal, ) -__all__ = ["Krige", "Simple", "Ordinary", "Universal", "ExtDrift", "Detrended"] +__all__ = ["Krige", "LocalKrige", "Simple","LocalSimple", "Ordinary", "LocalOrdinary", "Universal", "LocalUniversal", "ExtDrift", "LocalExtDrift", "Detrended", "LocalDetrended"] diff --git a/src/gstools/krige/local.py b/src/gstools/krige/local.py new file mode 100644 index 00000000..91a2b409 --- /dev/null +++ b/src/gstools/krige/local.py @@ -0,0 +1,889 @@ +""" +GStools subpackage providing a prototype class for local kriging. + +.. currentmodule:: gstools.krige.local + +The following classes are provided + +.. autosummary:: + LocalKrige +""" + +import json + +import numpy as np +from gstools_core import calc_field_krige_local as calc_field_krige_local_gsc + +from gstools.krige.base import Krige +from gstools.krige.tools import set_condition +from gstools.tools.geometric import rotated_main_axes +from gstools.variogram import vario_estimate + +__all__ = [ + "LocalKrige", + "LocalSimple", + "LocalOrdinary", + "LocalUniversal", + "LocalExtDrift", + "LocalDetrended", +] + + +# model classes GSTools-Core's CovModelSpec (covmodel_spec.rs) currently +# implements; the JSON "type" tag is just the lowercased class name for +# all four +_CORE_MODEL_TYPES = ("Gaussian", "Exponential", "Matern", "Spherical") + + +def _cov_model_to_json(model): + """ + Serialize a :any:`gstools.CovModel` to the JSON wire format the + GSTools-Core Rust backend expects (``CovModelSpec`` in + ``covmodel_spec.rs``: a ``{"type": ..., ...}``-tagged enum). + + Only ``var``/``len_scale``/``nugget`` (and, for Matern, ``nu``) travel + over the wire. + """ + if model.name not in _CORE_MODEL_TYPES: + raise NotImplementedError( + f"LocalKrige: GSTools-Core does not (yet) support " + f"the '{model.name}' covariance model; supported: " + f"{_CORE_MODEL_TYPES}." + ) + spec = { + "type": model.name.lower(), + "var": model.var, + "len_scale": model.len_scale, + "nugget": model.nugget, + } + if model.name == "Matern": + spec["nu"] = model.nu + return json.dumps(spec) + + +def _calc_field_krige_local( + cond_pos, + cond_val, + target_pos, + model, + cond_err, + drift_cond, + drift_target, + unbiased, + exact, + local_radius, + num_threads=None, +): + """A wrapper function for calling the local krige algorithm (Rust).""" + return calc_field_krige_local_gsc( + cond_pos, + cond_val, + target_pos, + _cov_model_to_json(model), + cond_err, + drift_cond, + drift_target, + unbiased, + exact, + local_radius, + num_threads=num_threads, + ) + + +class LocalKrige(Krige): + """ + Prototype class for local (moving-neighborhood) kriging. + + + Parameters + ---------- + model : :any:`CovModel` + Covariance Model used for kriging. + cond_pos : :class:`list` + tuple, containing the given condition positions (x, [y, z]) + cond_val : :class:`numpy.ndarray` + the values of the conditions (nan values will be ignored) + local_radius : :class:`float` + Search radius (in the model's isometrized/isotropic distance): + every conditioning point within this distance of a target + point enters its local kriging system. + drift_functions : :class:`list` of :any:`callable`, :class:`str` or :class:`int` + Either a list of callable functions, an integer representing + the polynomial order of the drift or one of the following strings: + + * "linear" : regional linear drift (equals order=1) + * "quadratic" : regional quadratic drift (equals order=2) + + ext_drift : :class:`numpy.ndarray` or :any:`None`, optional + the external drift values at the given cond. positions. + mean : :class:`float`, optional + mean value used to shift normalized conditioning data. + Could also be a callable. The default is None. + normalizer : :any:`None` or :any:`Normalizer`, optional + Normalizer to be applied to the input data to gain normality. + The default is None. + trend : :any:`None` or :class:`float` or :any:`callable`, optional + A callable trend function. Should have the signature: f(x, [y, z, ...]) + This is used for detrended kriging, where the trended is subtracted + from the conditions before kriging is applied. + This can be used for regression kriging, where the trend function + is determined by an external regression algorithm. + If no normalizer is applied, this behaves equal to 'mean'. + The default is None. + unbiased : :class:`bool`, optional + Whether the kriging weights should sum up to 1, so the estimator + is unbiased. If unbiased is `False` and no drifts are given, + this results in simple kriging. + Default: True + exact : :class:`bool`, optional + Whether the interpolator should reproduce the exact input values. + If `False`, `cond_err` is interpreted as measurement error + at the conditioning points and the result will be more smooth. + Default: False + cond_err : :class:`str`, :class :class:`float` or :class:`list`, optional + The measurement error at the conditioning points. + Either "nugget" to apply the model-nugget, a single value applied to + all points or an array with individual values for each point. + The "exact=True" variant only works with "cond_err='nugget'". + Default: "nugget" + fit_normalizer : :class:`bool`, optional + Whether to fit the data-normalizer to the given conditioning data. + Default: False + fit_variogram : :class:`bool`, optional + Whether to fit the given variogram model to the data. + Directional variogram fitting is triggered by setting + any anisotropy factor of the model to anything unequal 1 + but the main axes of correlation are taken from the model + rotation angles. If the model is a spatio-temporal latlon + model, this will raise an error. + This assumes the sill to be the data variance and with + standard bins provided by the :any:`standard_bins` routine. + Default: False + """ + + def __init__( + self, + model, + cond_pos, + cond_val, + local_radius, + drift_functions=None, + ext_drift=None, + mean=None, + normalizer=None, + trend=None, + unbiased=True, + exact=False, + cond_err="nugget", + fit_normalizer=False, + fit_variogram=False, + ): + self.local_radius = local_radius + self._cond_drift = None + self._target_drift = None + # super().__init__ calls set_condition and set_drift_functions + super().__init__( + model, + cond_pos, + cond_val, + drift_functions=drift_functions, + ext_drift=ext_drift, + mean=mean, + normalizer=normalizer, + trend=trend, + unbiased=unbiased, + exact=exact, + cond_err=cond_err, + fit_normalizer=fit_normalizer, + fit_variogram=fit_variogram, + ) + + def __call__( + self, + pos=None, + mesh_type="unstructured", + ext_drift=None, + only_mean=False, + return_var=True, + post_process=True, + store=True, + ): + """ + Generate the local kriging field. + + The field is saved as `self.field` and is also returned. + The error variance is saved as `self.krige_var` and is also returned. + + Parameters + ---------- + pos : :class:`list`, optional + the position tuple, containing main direction and transversal + directions (x, [y, z]) + mesh_type : :class:`str`, optional + 'structured' / 'unstructured' + ext_drift : :class:`numpy.ndarray` or :any:`None`, optional + the external drift values at the given positions (only for EDK) + only_mean : :class:`bool`, optional + Whether to only calculate the mean of the kriging field. + Not implemented yet for local kriging -- accepted here (and + has to be `False`) only so :any:`LocalKrige` is a drop-in + replacement for :any:`Krige` in places like :any:`CondSRF`, + which always call with ``only_mean=False``. + Default: `False` + return_var : :class:`bool`, optional + Whether to return the variance along with the field. + Default: `True` + post_process : :class:`bool`, optional + Whether to apply mean, normalizer and trend to the field. + Default: `True` + store : :class:`str` or :class:`bool` or :class:`list`, optional + Whether to store kriging fields (True/False) with default name + or with specified names. + The default is :any:`True` for default names + ["field", "krige_var"]. + + Returns + ------- + field : :class:`numpy.ndarray` + the kriged field + krige_var : :class:`numpy.ndarray`, optional + the kriging error variance (if return_var is True) + """ + if only_mean: + raise NotImplementedError( + "LocalKrige: only_mean=True is not implemented yet." + ) + fld_cnt = 2 if return_var else 1 + name, save = self.get_store_config(store, None, fld_cnt) + + iso_targ, shape = self.pre_pos(pos, mesh_type) + pnt_cnt = iso_targ.shape[1] + + # cond values already normalized/detrended/zero-mean (unpadded) + cond_val = self._krige_cond[: self.cond_no] + cond_err = self.cond_err + cond_err_arr = ( + np.full(self.cond_no, cond_err, dtype=np.double) + if np.isscalar(cond_err) + else np.asarray(cond_err, dtype=np.double) + ) + + ext_drift = self._pre_ext_drift(pnt_cnt, ext_drift) + drift_target = self._calc_target_drift(iso_targ, ext_drift, pnt_cnt) + + # call the local kriging routine (Rust or NumPy fallback) + field, error = _calc_field_krige_local( + cond_pos=self._krige_pos, + cond_val=cond_val, + target_pos=iso_targ, + model=self.model, + cond_err=cond_err_arr, + drift_cond=self._cond_drift, + drift_target=drift_target, + unbiased=self.unbiased, + exact=self.exact, + local_radius=self.local_radius, + ) + + field = np.reshape(field, shape) + field = self.post_field(field, name[0], post_process, save[0]) + if return_var: # care about the estimated error variance + krige_var = np.reshape( + np.maximum(self.model.sill - error, 0), shape + ) + krige_var = self.post_field(krige_var, name[1], False, save[1]) + return field, krige_var + return field + + def set_condition( + self, + cond_pos=None, + cond_val=None, + ext_drift=None, + cond_err=None, + fit_normalizer=False, + fit_variogram=False, + ): + """Set the conditions for kriging. + + This method could also be used to update the kriging setup, when + properties were changed. Then you can call it without arguments. + + This function is the same as :any:`Krige.set_condition`, but it also precomputes + the drift terms at the conditioning points for the local kriging solve. + (last line of this function) + TODO: This function should be refactored to avoid code duplication with the base class. + + Parameters + ---------- + cond_pos : :class:`list`, optional + the position tuple of the conditions (x, [y, z]). Default: current. + cond_val : :class:`numpy.ndarray`, optional + the values of the conditions (nan values will be ignored). + Default: current. + ext_drift : :class:`numpy.ndarray` or :any:`None`, optional + the external drift values at the given conditions (only for EDK) + For multiple external drifts, the first dimension + should be the index of the drift term. When passing `None`, the + extisting external drift will be used. + cond_err : :class:`str`, :class :class:`float`, :class:`list`, optional + The measurement error at the conditioning points. + Either "nugget" to apply the model-nugget, a single value applied + to all points or an array with individual values for each point. + The measurement error has to be <= nugget. + The "exact=True" variant only works with "cond_err='nugget'". + Default: "nugget" + fit_normalizer : :class:`bool`, optional + Whether to fit the data-normalizer to the given conditioning data. + Default: False + fit_variogram : :class:`bool`, optional + Whether to fit the given variogram model to the data. + Directional variogram fitting is triggered by setting + any anisotropy factor of the model to anything unequal 1 + but the main axes of correlation are taken from the model + rotation angles. If the model is a spatio-temporal latlon + model, this will raise an error. + This assumes the sill to be the data variance and with + standard bins provided by the :any:`standard_bins` routine. + Default: False + """ + # only use existing external drift, if no new positions are given + ext_drift = ( + self._cond_ext_drift + if (ext_drift is None and cond_pos is None) + else ext_drift + ) + # use existing values or set default + cond_pos = self._cond_pos if cond_pos is None else cond_pos + cond_val = self._cond_val if cond_val is None else cond_val + cond_err = self._cond_err if cond_err is None else cond_err + cond_err = "nugget" if cond_err is None else cond_err # default + if cond_pos is None or cond_val is None: + raise ValueError("Krige.set_condition: missing cond_pos/cond_val.") + # correctly format cond_pos and cond_val + self._cond_pos, self._cond_val = set_condition( + cond_pos, cond_val, self.dim + ) + if fit_normalizer: # fit normalizer to detrended data + self.normalizer.fit(self.cond_val - self.cond_trend) + if fit_variogram: # fitting model to empirical variogram of data + # normalize field + if self.model.latlon and self.model.temporal: + msg = "Krige: can't fit variogram for spatio-temporal latlon data." + raise ValueError(msg) + field = self.normalizer.normalize(self.cond_val - self.cond_trend) + field -= self.cond_mean + sill = np.var(field) + if self.model.is_isotropic: + emp_vario = vario_estimate( + self.cond_pos, + field, + latlon=self.model.latlon, + geo_scale=self.model.geo_scale, + ) + else: + axes = rotated_main_axes(self.model.dim, self.model.angles) + emp_vario = vario_estimate( + self.cond_pos, field, direction=axes + ) + # set the sill to the field variance + self.model.fit_variogram(*emp_vario, sill=sill) + # set the measurement errors + self.cond_err = cond_err + # set the external drift values and the conditioning points + self._cond_ext_drift = self._pre_ext_drift( + self.cond_no, ext_drift, set_cond=True + ) + # upate the internal kriging settings + self._krige_pos = self.model.isometrize(self.cond_pos) + # krige pos are the unrotated and isotropic condition positions + # local kriging: no global kriging matrix -- pre-compute the drift + # terms at the conditioning points instead, for the local solve + self._cond_drift = self._calc_cond_drift() + + def _calc_cond_drift(self): + """ + Evaluate all drift terms at the conditioning points once. + + Internal (functional) drift and external drift are unified into + a single ``(drift_no, cond_no)`` array: external drift is + already array-shaped, so it is simply stacked below the + evaluated functional drift rows. + + Returns + ------- + :class:`numpy.ndarray` + Shape ``(drift_no, cond_no)``. + """ + cond_pos = np.asarray(self.cond_pos) + int_functions = self.drift_functions + int_no = len(int_functions) + ext_no = self.ext_drift_no + drift_no = int_no + ext_no + + drift_cond = np.empty((drift_no, self.cond_no), dtype=np.double) + for i, f in enumerate(int_functions): + drift_cond[i] = f(*cond_pos) + if ext_no > 0: + drift_cond[int_no:] = self.cond_ext_drift + return drift_cond + + def _calc_target_drift(self, iso_targ, ext_drift, pnt_cnt): + """ + Evaluate all drift terms at the target positions once. + + Mirrors :any:`_calc_cond_drift`, but for the (isometrized) target + positions of the current evaluation call instead of the + conditioning points. Internal (functional) drift and external + drift are unified into a single ``(drift_no, pnt_cnt)`` array, + the same layout :any:`_calc_cond_drift` uses for the + conditioning points. + + Parameters + ---------- + iso_targ : :class:`numpy.ndarray` + Isometrized target positions, shape ``(dim, pnt_cnt)``. + ext_drift : :class:`numpy.ndarray` + Preprocessed external drift values at the target positions + (output of :any:`Krige._pre_ext_drift`), shape + ``(ext_drift_no, pnt_cnt)``. + pnt_cnt : :class:`int` + Number of target points. + + Returns + ------- + :class:`numpy.ndarray` + Shape ``(drift_no, pnt_cnt)``. + """ + int_functions = self.drift_functions + int_no = len(int_functions) + ext_no = self.ext_drift_no + drift_no = int_no + ext_no + + drift_target = np.empty((drift_no, pnt_cnt), dtype=np.double) + if int_no > 0: + # drift functions live in the original (anisotropic) frame + targ_pos = self.model.anisometrize(iso_targ) + for i, f in enumerate(int_functions): + drift_target[i] = f(*targ_pos) + if ext_no > 0: + drift_target[int_no:] = ext_drift + self._target_drift = drift_target + return drift_target + + @property + def local_radius(self): + """:class:`float`: The search radius for local kriging.""" + return self._local_radius + + @local_radius.setter + def local_radius(self, value): + value = float(value) + if value <= 0: + raise ValueError("LocalKrige: local_radius must be > 0.") + self._local_radius = value + + +class LocalSimple(LocalKrige): + """ + Local simple kriging. + + Local simple kriging is used to interpolate data with a given mean, + using only the conditioning points within ``local_radius`` of each + target point. See :any:`LocalKrige` and :any:`gstools.krige.Simple`. + + Parameters + ---------- + model : :any:`CovModel` + Covariance Model used for kriging. + cond_pos : :class:`list` + tuple, containing the given condition positions (x, [y, z]) + cond_val : :class:`numpy.ndarray` + the values of the conditions (nan values will be ignored) + local_radius : :class:`float` + Search radius (in the model's isometrized/isotropic distance): + every conditioning point within this distance of a target + point enters its local kriging system. + mean : :class:`float`, optional + mean value used to shift normalized conditioning data. + Could also be a callable. The default is None. + normalizer : :any:`None` or :any:`Normalizer`, optional + Normalizer to be applied to the input data to gain normality. + The default is None. + trend : :any:`None` or :class:`float` or :any:`callable`, optional + A callable trend function. Should have the signature: f(x, [y, z, ...]) + This is used for detrended kriging, where the trended is subtracted + from the conditions before kriging is applied. + This can be used for regression kriging, where the trend function + is determined by an external regression algorithm. + If no normalizer is applied, this behaves equal to 'mean'. + The default is None. + exact : :class:`bool`, optional + Whether the interpolator should reproduce the exact input values. + If `False`, `cond_err` is interpreted as measurement error + at the conditioning points and the result will be more smooth. + Default: False + cond_err : :class:`str`, :class :class:`float` or :class:`list`, optional + The measurement error at the conditioning points. + Either "nugget" to apply the model-nugget, a single value applied to + all points or an array with individual values for each point. + The measurement error has to be <= nugget. + The "exact=True" variant only works with "cond_err='nugget'". + Default: "nugget" + fit_normalizer : :class:`bool`, optional + Whether to fit the data-normalizer to the given conditioning data. + Default: False + fit_variogram : :class:`bool`, optional + Whether to fit the given variogram model to the data. + Default: False + """ + + def __init__( + self, + model, + cond_pos, + cond_val, + local_radius, + mean=0.0, + normalizer=None, + trend=None, + exact=False, + cond_err="nugget", + fit_normalizer=False, + fit_variogram=False, + ): + super().__init__( + model, + cond_pos, + cond_val, + local_radius, + mean=mean, + normalizer=normalizer, + trend=trend, + unbiased=False, + exact=exact, + cond_err=cond_err, + fit_normalizer=fit_normalizer, + fit_variogram=fit_variogram, + ) + + +class LocalOrdinary(LocalKrige): + """ + Local ordinary kriging. + + Local ordinary kriging is used to interpolate data and estimate a + proper mean, using only the conditioning points within + ``local_radius`` of each target point. See :any:`LocalKrige` and + :any:`gstools.krige.Ordinary`. + + Parameters + ---------- + model : :any:`CovModel` + Covariance Model used for kriging. + cond_pos : :class:`list` + tuple, containing the given condition positions (x, [y, z]) + cond_val : :class:`numpy.ndarray` + the values of the conditions (nan values will be ignored) + local_radius : :class:`float` + Search radius (in the model's isometrized/isotropic distance): + every conditioning point within this distance of a target + point enters its local kriging system. + normalizer : :any:`None` or :any:`Normalizer`, optional + Normalizer to be applied to the input data to gain normality. + The default is None. + trend : :any:`None` or :class:`float` or :any:`callable`, optional + A callable trend function. Should have the signature: f(x, [y, z, ...]) + This is used for detrended kriging, where the trended is subtracted + from the conditions before kriging is applied. + This can be used for regression kriging, where the trend function + is determined by an external regression algorithm. + If no normalizer is applied, this behaves equal to 'mean'. + The default is None. + exact : :class:`bool`, optional + Whether the interpolator should reproduce the exact input values. + If `False`, `cond_err` is interpreted as measurement error + at the conditioning points and the result will be more smooth. + Default: False + cond_err : :class:`str`, :class :class:`float` or :class:`list`, optional + The measurement error at the conditioning points. + Either "nugget" to apply the model-nugget, a single value applied to + all points or an array with individual values for each point. + The measurement error has to be <= nugget. + The "exact=True" variant only works with "cond_err='nugget'". + Default: "nugget" + fit_normalizer : :class:`bool`, optional + Whether to fit the data-normalizer to the given conditioning data. + Default: False + fit_variogram : :class:`bool`, optional + Whether to fit the given variogram model to the data. + Default: False + """ + + def __init__( + self, + model, + cond_pos, + cond_val, + local_radius, + normalizer=None, + trend=None, + exact=False, + cond_err="nugget", + fit_normalizer=False, + fit_variogram=False, + ): + super().__init__( + model, + cond_pos, + cond_val, + local_radius, + normalizer=normalizer, + trend=trend, + exact=exact, + cond_err=cond_err, + fit_normalizer=fit_normalizer, + fit_variogram=fit_variogram, + ) + + +class LocalUniversal(LocalKrige): + """ + Local universal kriging. + + Local universal kriging is used to interpolate given data with a + variable mean determined by a functional drift, using only the + conditioning points within ``local_radius`` of each target point. + See :any:`LocalKrige` and :any:`gstools.krige.Universal`. + + This estimator is set to be unbiased by default. + + Parameters + ---------- + model : :any:`CovModel` + Covariance Model used for kriging. + cond_pos : :class:`list` + tuple, containing the given condition positions (x, [y, z]) + cond_val : :class:`numpy.ndarray` + the values of the conditions (nan values will be ignored) + local_radius : :class:`float` + Search radius (in the model's isometrized/isotropic distance): + every conditioning point within this distance of a target + point enters its local kriging system. + drift_functions : :class:`list` of :any:`callable`, :class:`str` or :class:`int` + Either a list of callable functions, an integer representing + the polynomial order of the drift or one of the following strings: + + * "linear" : regional linear drift (equals order=1) + * "quadratic" : regional quadratic drift (equals order=2) + + normalizer : :any:`None` or :any:`Normalizer`, optional + Normalizer to be applied to the input data to gain normality. + The default is None. + trend : :any:`None` or :class:`float` or :any:`callable`, optional + A callable trend function. Should have the signature: f(x, [y, z, ...]) + This is used for detrended kriging, where the trended is subtracted + from the conditions before kriging is applied. + This can be used for regression kriging, where the trend function + is determined by an external regression algorithm. + If no normalizer is applied, this behaves equal to 'mean'. + The default is None. + exact : :class:`bool`, optional + Whether the interpolator should reproduce the exact input values. + If `False`, `cond_err` is interpreted as measurement error + at the conditioning points and the result will be more smooth. + Default: False + cond_err : :class:`str`, :class :class:`float` or :class:`list`, optional + The measurement error at the conditioning points. + Either "nugget" to apply the model-nugget, a single value applied to + all points or an array with individual values for each point. + The measurement error has to be <= nugget. + The "exact=True" variant only works with "cond_err='nugget'". + Default: "nugget" + fit_normalizer : :class:`bool`, optional + Whether to fit the data-normalizer to the given conditioning data. + Default: False + fit_variogram : :class:`bool`, optional + Whether to fit the given variogram model to the data. + Default: False + """ + + def __init__( + self, + model, + cond_pos, + cond_val, + local_radius, + drift_functions, + normalizer=None, + trend=None, + exact=False, + cond_err="nugget", + fit_normalizer=False, + fit_variogram=False, + ): + super().__init__( + model, + cond_pos, + cond_val, + local_radius, + drift_functions=drift_functions, + normalizer=normalizer, + trend=trend, + exact=exact, + cond_err=cond_err, + fit_normalizer=fit_normalizer, + fit_variogram=fit_variogram, + ) + + +class LocalExtDrift(LocalKrige): + """ + Local external drift kriging (local EDK). + + Local external drift kriging is used to interpolate given data with + a variable mean determined by an external drift, using only the + conditioning points within ``local_radius`` of each target point. + See :any:`LocalKrige` and :any:`gstools.krige.ExtDrift`. + + This estimator is set to be unbiased by default. + + Parameters + ---------- + model : :any:`CovModel` + Covariance Model used for kriging. + cond_pos : :class:`list` + tuple, containing the given condition positions (x, [y, z]) + cond_val : :class:`numpy.ndarray` + the values of the conditions (nan values will be ignored) + local_radius : :class:`float` + Search radius (in the model's isometrized/isotropic distance): + every conditioning point within this distance of a target + point enters its local kriging system. + ext_drift : :class:`numpy.ndarray` + the external drift values at the given condition positions. + normalizer : :any:`None` or :any:`Normalizer`, optional + Normalizer to be applied to the input data to gain normality. + The default is None. + trend : :any:`None` or :class:`float` or :any:`callable`, optional + A callable trend function. Should have the signature: f(x, [y, z, ...]) + This is used for detrended kriging, where the trended is subtracted + from the conditions before kriging is applied. + This can be used for regression kriging, where the trend function + is determined by an external regression algorithm. + If no normalizer is applied, this behaves equal to 'mean'. + The default is None. + exact : :class:`bool`, optional + Whether the interpolator should reproduce the exact input values. + If `False`, `cond_err` is interpreted as measurement error + at the conditioning points and the result will be more smooth. + Default: False + cond_err : :class:`str`, :class :class:`float` or :class:`list`, optional + The measurement error at the conditioning points. + Either "nugget" to apply the model-nugget, a single value applied to + all points or an array with individual values for each point. + The measurement error has to be <= nugget. + The "exact=True" variant only works with "cond_err='nugget'". + Default: "nugget" + fit_normalizer : :class:`bool`, optional + Whether to fit the data-normalizer to the given conditioning data. + Default: False + fit_variogram : :class:`bool`, optional + Whether to fit the given variogram model to the data. + Default: False + """ + + def __init__( + self, + model, + cond_pos, + cond_val, + local_radius, + ext_drift, + normalizer=None, + trend=None, + exact=False, + cond_err="nugget", + fit_normalizer=False, + fit_variogram=False, + ): + super().__init__( + model, + cond_pos, + cond_val, + local_radius, + ext_drift=ext_drift, + normalizer=normalizer, + trend=trend, + exact=exact, + cond_err=cond_err, + fit_normalizer=fit_normalizer, + fit_variogram=fit_variogram, + ) + + +class LocalDetrended(LocalKrige): + """ + Local detrended simple kriging. + + In local detrended kriging, the data is detrended before + interpolation by local simple kriging with zero mean, using only + the conditioning points within ``local_radius`` of each target + point. See :any:`LocalKrige` and :any:`gstools.krige.Detrended`. + + This is just a shortcut for local simple kriging with a given trend + function, zero mean and no normalizer. + + Parameters + ---------- + model : :any:`CovModel` + Covariance Model used for kriging. + cond_pos : :class:`list` + tuple, containing the given condition positions (x, [y, z]) + cond_val : :class:`numpy.ndarray` + the values of the conditions (nan values will be ignored) + local_radius : :class:`float` + Search radius (in the model's isometrized/isotropic distance): + every conditioning point within this distance of a target + point enters its local kriging system. + trend_function : :any:`callable` + The callable trend function. Should have the signature: f(x, [y, z]) + exact : :class:`bool`, optional + Whether the interpolator should reproduce the exact input values. + If `False`, `cond_err` is interpreted as measurement error + at the conditioning points and the result will be more smooth. + Default: False + cond_err : :class:`str`, :class :class:`float` or :class:`list`, optional + The measurement error at the conditioning points. + Either "nugget" to apply the model-nugget, a single value applied to + all points or an array with individual values for each point. + The measurement error has to be <= nugget. + The "exact=True" variant only works with "cond_err='nugget'". + Default: "nugget" + fit_variogram : :class:`bool`, optional + Whether to fit the given variogram model to the data. + Default: False + """ + + def __init__( + self, + model, + cond_pos, + cond_val, + local_radius, + trend_function, + exact=False, + cond_err="nugget", + fit_variogram=False, + ): + super().__init__( + model, + cond_pos, + cond_val, + local_radius, + trend=trend_function, + unbiased=False, + exact=exact, + cond_err=cond_err, + fit_variogram=fit_variogram, + ) diff --git a/tests/test_local_krige.py b/tests/test_local_krige.py new file mode 100644 index 00000000..d61d056a --- /dev/null +++ b/tests/test_local_krige.py @@ -0,0 +1,306 @@ +""" +This is the unittest of the LocalKrige prototype (moving-neighborhood +kriging, see gstools.krige.local). It mirrors test_krige.py, using the +Local* counterparts of the Krige convenience classes. +""" + +import unittest + +import numpy as np + +import gstools as gs + + +def trend(*xyz): + return xyz[0] + + +class TestLocalKrige(unittest.TestCase): + def setUp(self): + # GSTools-Core's local kriging currently only supports these + # covariance models (see LocalKrige._cov_model_to_json / the + # Rust CovModelSpec) -- unlike test_krige.py, no Spherical here + self.cov_models = [gs.Gaussian, gs.Exponential] + self.dims = range(1, 4) + self.data = np.array( + [ + [0.3, 1.2, 0.5, 0.47], + [1.9, 0.6, 1.0, 0.56], + [1.1, 3.2, 1.5, 0.74], + [3.3, 4.4, 2.0, 1.47], + [4.7, 3.8, 2.5, 1.74], + ] + ) + # x, y, z components for the condition position + self.cond_pos = (self.data[:, 0], self.data[:, 1], self.data[:, 2]) + # condition values + self.cond_val = self.data[:, 3] + self.cond_err = np.array([0.01, 0.0, 0.1, 0.05, 0]) + # the arithmetic mean of the conditions + self.mean = np.mean(self.cond_val) + # a moderate grid: unlike global kriging, local kriging solves one + # small system per target point, and (for now) does so + # sequentially (see GSTools-Core's local_krige.rs) -- the + # 51x61x71 grid test_krige.py uses would take tens of seconds + # per call here, so this is deliberately much smaller + self.x = np.linspace(0, 5, 11) + self.y = np.linspace(0, 6, 13) + self.z = np.linspace(0, 7, 15) + self.pos = (self.x, self.y, self.z) + self.grids = [self.x] + self.grids.append(np.meshgrid(self.x, self.y, indexing="ij")) + self.grids.append(np.meshgrid(self.x, self.y, self.z, indexing="ij")) + self.grid_shape = [11, 13, 15] + # a radius spanning the whole domain: every conditioning point is + # then a "neighbor" of every target point, so local kriging + # should reproduce the corresponding global kriging result (up + # to solver round-off, since the local solver does a direct LU + # solve instead of a pseudo-inverse) + self.big_radius = 1e3 + + def _check(self, local_krige, global_krige, dim): + """Shared checks for a Local* class against its global counterpart.""" + # structured and unstructured evaluation have to agree + field_1, __ = local_krige.unstructured(self.grids[dim - 1]) + field_1 = field_1.reshape(self.grid_shape[:dim]) + field_2, __ = local_krige.structured(self.pos[:dim]) + self.assertAlmostEqual( + np.max(np.abs(field_1 - field_2)), 0.0, places=2 + ) + # with a radius spanning the whole domain, local kriging has to + # match global kriging + field_g, __ = global_krige.unstructured(self.grids[dim - 1]) + field_g = field_g.reshape(self.grid_shape[:dim]) + np.testing.assert_allclose(field_1, field_g, atol=1e-6) + # both reproduce the conditioning values themselves + field_cond, __ = local_krige.unstructured(self.cond_pos[:dim]) + for i, val in enumerate(self.cond_val): + self.assertAlmostEqual(field_cond[i], val, places=2) + + def test_simple(self): + for Model in self.cov_models: + for dim in self.dims: + model = Model( + dim=dim, + var=2, + len_scale=2, + anis=[0.9, 0.8], + angles=[2, 1, 0.5], + ) + simple = gs.krige.Simple( + model, self.cond_pos[:dim], self.cond_val, self.mean + ) + local_simple = gs.krige.LocalSimple( + model, + self.cond_pos[:dim], + self.cond_val, + self.big_radius, + self.mean, + ) + self._check(local_simple, simple, dim) + + def test_ordinary(self): + for trend_func in [None, trend]: + for Model in self.cov_models: + for dim in self.dims: + model = Model( + dim=dim, + var=5, + len_scale=10, + anis=[0.9, 0.8], + angles=[2, 1, 0.5], + ) + ordinary = gs.krige.Ordinary( + model, + self.cond_pos[:dim], + self.cond_val, + trend=trend_func, + ) + local_ordinary = gs.krige.LocalOrdinary( + model, + self.cond_pos[:dim], + self.cond_val, + self.big_radius, + trend=trend_func, + ) + self._check(local_ordinary, ordinary, dim) + + def test_universal(self): + # "quadratic" -> too few conditioning points + for drift in ["linear", 0, 1, trend]: + for Model in self.cov_models: + for dim in self.dims: + model = Model( + dim=dim, + var=2, + len_scale=10, + anis=[0.9, 0.8], + angles=[2, 1, 0.5], + ) + universal = gs.krige.Universal( + model, self.cond_pos[:dim], self.cond_val, drift + ) + local_universal = gs.krige.LocalUniversal( + model, + self.cond_pos[:dim], + self.cond_val, + self.big_radius, + drift, + ) + self._check(local_universal, universal, dim) + + def test_detrended(self): + for Model in self.cov_models: + for dim in self.dims: + model = Model( + dim=dim, + var=2, + len_scale=10, + anis=[0.5, 0.2], + angles=[0.4, 0.2, 0.1], + ) + detrended = gs.krige.Detrended( + model, self.cond_pos[:dim], self.cond_val, trend + ) + local_detrended = gs.krige.LocalDetrended( + model, + self.cond_pos[:dim], + self.cond_val, + self.big_radius, + trend, + ) + self._check(local_detrended, detrended, dim) + + def test_extdrift(self): + rng = np.random.RandomState(42) + for dim in self.dims: + cond_drift = rng.normal(size=len(self.cond_val)) + target_drift = rng.normal(size=int(np.prod(self.grid_shape[:dim]))) + for Model in self.cov_models: + model = Model( + dim=dim, + var=2, + len_scale=10, + anis=[0.5, 0.2], + angles=[0.4, 0.2, 0.1], + ) + extdrift = gs.krige.ExtDrift( + model, self.cond_pos[:dim], self.cond_val, cond_drift + ) + local_extdrift = gs.krige.LocalExtDrift( + model, + self.cond_pos[:dim], + self.cond_val, + self.big_radius, + cond_drift, + ) + field_1, __ = local_extdrift.unstructured( + self.grids[dim - 1], ext_drift=target_drift + ) + field_1 = field_1.reshape(self.grid_shape[:dim]) + field_2, __ = local_extdrift.structured( + self.pos[:dim], ext_drift=target_drift + ) + self.assertAlmostEqual( + np.max(np.abs(field_1 - field_2)), 0.0, places=2 + ) + field_g, __ = extdrift.unstructured( + self.grids[dim - 1], ext_drift=target_drift + ) + field_g = field_g.reshape(self.grid_shape[:dim]) + np.testing.assert_allclose(field_1, field_g, atol=1e-6) + field_cond, __ = local_extdrift.unstructured( + self.cond_pos[:dim], ext_drift=cond_drift + ) + for i, val in enumerate(self.cond_val): + self.assertAlmostEqual(field_cond[i], val, places=2) + + def test_error(self): + for Model in self.cov_models: + for dim in self.dims: + model = Model( + dim=dim, + var=5, + len_scale=10, + nugget=0.1, + anis=[0.9, 0.8], + angles=[2, 1, 0.5], + ) + ordinary = gs.krige.LocalOrdinary( + model, + self.cond_pos[:dim], + self.cond_val, + self.big_radius, + exact=False, + cond_err=self.cond_err, + ) + __, err = ordinary(self.cond_pos[:dim]) + # when the given measurement error is 0, the kriging-var + # should equal the nugget of the model + self.assertAlmostEqual(err[1], model.nugget, places=2) + self.assertAlmostEqual(err[4], model.nugget, places=2) + + def test_raise(self): + # no cond_pos/cond_val given + self.assertRaises( + ValueError, gs.krige.LocalKrige, gs.Stable(), None, None, 1.0 + ) + # local_radius has to be strictly positive + self.assertRaises( + ValueError, + gs.krige.LocalKrige, + gs.Gaussian(), + self.cond_pos, + self.cond_val, + 0, + ) + self.assertRaises( + ValueError, + gs.krige.LocalKrige, + gs.Gaussian(), + self.cond_pos, + self.cond_val, + -1.0, + ) + + def test_no_neighbors_raises(self): + # a target point far outside of local_radius has no neighbors + # at all -- this has to be reported, not silently mis-evaluated + model = gs.Gaussian(dim=2, var=1, len_scale=1) + local = gs.krige.LocalOrdinary( + model, self.cond_pos[:2], self.cond_val, 1e-3 + ) + self.assertRaises(ValueError, local, ([100.0], [100.0])) + + def test_smaller_radius_gives_similar_result(self): + # a denser, random conditioning setup: with a small radius only a + # few neighbors enter each local system, so the result differs + # from the global one -- but that difference should shrink + # monotonically as the radius grows, and vanish once the radius + # spans the whole domain + rng = np.random.RandomState(2) + cond_pos = (rng.uniform(0, 20, 60), rng.uniform(0, 20, 60)) + cond_val = rng.normal(size=60) + 0.05 * cond_pos[0] + target_pos = (rng.uniform(2, 18, 25), rng.uniform(2, 18, 25)) + + model = gs.Gaussian(dim=2, var=1.5, len_scale=4, nugget=0.05) + ordinary = gs.krige.Ordinary(model, cond_pos, cond_val) + field_g, __ = ordinary.unstructured(target_pos) + + prev_diff = np.inf + for radius in [5.0, 8.0, 12.0, self.big_radius]: + local = gs.krige.LocalOrdinary(model, cond_pos, cond_val, radius) + field_l, __ = local.unstructured(target_pos) + diff = np.max(np.abs(field_l - field_g)) + corr = np.corrcoef(field_l, field_g)[0, 1] + # results should get closer to the global one (or at least + # not get worse) as the radius grows + self.assertLessEqual(diff, prev_diff + 1e-8) + self.assertGreater(corr, 0.9) + prev_diff = diff + # with a radius spanning the whole domain, it should match closely + self.assertAlmostEqual(prev_diff, 0.0, places=4) + + +if __name__ == "__main__": + unittest.main()