spatialcv 0.1.0__py3-none-any.whl

spatialcv/nndm.py

@@ -0,0 +1,470 @@
+"""Nearest-neighbour distance matching: NNDM LOO CV and kNNDM CV.
+
+Both methods choose cross-validation folds so that the distribution of nearest-neighbour
+distances *from CV test points to their training points* matches the distribution of
+distances *from the prediction locations to the training points*. Random CV is then
+"as hard" as the real prediction task, no harder and no easier.
+
+* NNDM: Milà, C., Mateu, J., Pebesma, E., Meyer, H. (2022). Nearest neighbour distance
+  matching Leave-One-Out Cross-Validation for map validation. *Methods in Ecology and
+  Evolution* 13, 1304-1316. https://doi.org/10.1111/2041-210X.13851
+* kNNDM: Linnenbrink, J., Milà, C., Ludwig, M., Meyer, H. (2024). kNNDM CV: k-fold
+  nearest-neighbour distance matching cross-validation for map accuracy estimation.
+  *Geoscientific Model Development* 17, 5897-5912.
+  https://doi.org/10.5194/gmd-17-5897-2024
+
+The implementations follow the reference R code in the CAST package (``nndm()`` and
+``knndm()``, geographical space).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple
+
+import numpy as np
+from sklearn.cluster import KMeans
+from sklearn.utils import check_random_state
+
+from ._base import SpatialSplitter, check_n_splits
+from ._coords import (
+    as_coords,
+    check_metric,
+    nearest_distance,
+    nearest_neighbour_loo,
+    pairwise_distances,
+    to_planar,
+    to_unit_sphere,
+)
+
+
+def wasserstein_distance(u: Any, v: Any) -> float:
+    """First Wasserstein distance between two 1-D samples: the area between their ECDFs.
+
+    Equals ``scipy.stats.wasserstein_distance`` and the W statistic
+    (``twosamples::wass_stat``) used by Linnenbrink et al. (2024).
+    """
+    a = np.sort(np.asarray(u, dtype=float).ravel())
+    b = np.sort(np.asarray(v, dtype=float).ravel())
+    if a.size == 0 or b.size == 0:
+        raise ValueError("both samples must be non-empty")
+    values = np.sort(np.concatenate([a, b]))
+    deltas = np.diff(values)
+    cdf_a = np.searchsorted(a, values[:-1], side="right") / a.size
+    cdf_b = np.searchsorted(b, values[:-1], side="right") / b.size
+    return float(np.sum(np.abs(cdf_a - cdf_b) * deltas))
+
+
+# ---------------------------------------------------------------------------------------
+# NNDM LOO
+# ---------------------------------------------------------------------------------------
+@dataclass
+class NNDMResult:
+    """Output of :func:`nndm`.
+
+    Attributes
+    ----------
+    train_indices : list of ndarray
+        Training indices of the LOO split that tests sample ``i``.
+    test_indices : list of ndarray
+        ``[i]`` for every sample.
+    Gij : ndarray
+        Nearest-neighbour distance from every prediction point to the training points.
+    Gj : ndarray
+        Nearest-neighbour distance between training points (ordinary LOO).
+    Gjstar : ndarray
+        Nearest-neighbour distance from each test point to its NNDM training set.
+    phi : float
+        Distance up to which the distributions were matched.
+    """
+
+    train_indices: List[np.ndarray]
+    test_indices: List[np.ndarray]
+    Gij: np.ndarray
+    Gj: np.ndarray
+    Gjstar: np.ndarray
+    phi: float
+    min_train: float
+
+    @property
+    def W(self) -> float:
+        """Wasserstein distance between ``Gjstar`` and ``Gij``."""
+        return wasserstein_distance(self.Gjstar, self.Gij)
+
+
+def nndm(
+    train_coords: Any,
+    prediction_coords: Any,
+    *,
+    phi: Optional[float] = None,
+    min_train: float = 0.5,
+    metric: str = "euclidean",
+) -> NNDMResult:
+    """Nearest Neighbour Distance Matching LOO folds (Milà et al. 2022).
+
+    Starting from ordinary LOO, the algorithm walks through the nearest-neighbour
+    distances ``r`` in increasing order. Whenever the LOO empirical CDF ``Gj*(r)`` would
+    still be at or above the prediction CDF ``Gij(r)`` after removing the nearest
+    neighbour of test point ``j``, that neighbour is removed from ``j``'s training set
+    (as long as more than ``min_train`` of the data remain). It stops at ``r > phi``.
+
+    Parameters
+    ----------
+    train_coords : array-like (n, 2)
+        Training sample locations.
+    prediction_coords : array-like (m, 2)
+        Locations where the model will predict (e.g. a regular sample of the area of
+        interest). Use a few thousand points.
+    phi : float, optional
+        Range of spatial autocorrelation; distances beyond it are not matched. Defaults to
+        the largest observed distance (CAST's ``phi = "max"``).
+    min_train : float, default=0.5
+        Minimum proportion of samples each LOO training set must keep, in [0, 1).
+    metric : {"euclidean", "haversine"}
+    """
+    check_metric(metric)
+    tpts = as_coords(train_coords, metric)
+    ppts = as_coords(prediction_coords, metric)
+    if not 0.0 <= min_train < 1.0:
+        raise ValueError(f"min_train must be in [0, 1), got {min_train}")
+    n = len(tpts)
+    if n < 3:
+        raise ValueError("NNDM needs at least three training points")
+    if len(ppts) == 0:
+        raise ValueError("prediction_coords is empty")
+
+    tdist = pairwise_distances(tpts, tpts, metric)
+    np.fill_diagonal(tdist, np.inf)
+    Gij = nearest_distance(ppts, tpts, metric)
+    Gij_sorted = np.sort(Gij)
+    Gj = tdist.min(axis=1)
+    if phi is None:
+        finite = tdist[np.isfinite(tdist)]
+        phi = float(max(Gij.max(), finite.max()))
+    elif not phi > 0:
+        raise ValueError(f"phi must be positive, got {phi}")
+    active = np.isfinite(tdist)
+    Gjstar = Gj.copy()
+
+    def ecdf_pred(r: float) -> float:
+        return float(np.searchsorted(Gij_sorted, r, side="right")) / Gij_sorted.size
+
+    rmin = float(Gjstar.min())
+    jmin = int(np.flatnonzero(Gjstar == rmin)[0])
+    while rmin <= phi:
+        cond_cdf = (np.count_nonzero(Gjstar <= rmin) - 1) / n >= ecdf_pred(rmin)
+        cond_size = np.count_nonzero(active[jmin]) / n > min_train
+        if cond_cdf and cond_size:
+            kmin = np.flatnonzero(tdist[jmin] == rmin)
+            tdist[jmin, kmin] = np.inf
+            active[jmin, kmin] = False
+            Gjstar[jmin] = tdist[jmin].min()
+            rmin = float(Gjstar[Gjstar >= rmin].min())
+        elif not np.any(Gjstar > rmin):
+            break
+        else:
+            rmin = float(Gjstar[Gjstar > rmin].min())
+        if not np.isfinite(rmin):
+            break
+        jmin = int(np.flatnonzero(Gjstar == rmin)[0])
+
+    indices = np.arange(n)
+    train = [indices[active[i]] for i in range(n)]
+    test = [np.array([i]) for i in range(n)]
+    return NNDMResult(train, test, Gij, Gj, Gjstar, float(phi), float(min_train))
+
+
+class NNDMLeaveOneOut(SpatialSplitter):
+    """Nearest Neighbour Distance Matching leave-one-out CV (Milà et al. 2022).
+
+    A scikit-learn splitter around :func:`nndm`. Test sets are single samples (so
+    ``cross_val_predict`` works); training sets drop neighbours until the CV distance
+    distribution matches the prediction-to-training distance distribution.
+
+    Parameters
+    ----------
+    prediction_points : array-like (m, 2)
+        Locations where the model will be applied, in the same coordinates as the samples.
+    phi : float, optional
+        Autocorrelation range; defaults to the maximum distance.
+    min_train : float, default=0.5
+    coords, coord_columns, metric
+        See :class:`SpatialBlockKFold`.
+    """
+
+    def __init__(
+        self,
+        prediction_points: Any,
+        *,
+        phi: Optional[float] = None,
+        min_train: float = 0.5,
+        coords: Any = None,
+        coord_columns: Optional[Sequence[Any]] = None,
+        metric: str = "euclidean",
+    ) -> None:
+        self.prediction_points = prediction_points
+        self.phi = phi
+        self.min_train = min_train
+        self.coords = coords
+        self.coord_columns = coord_columns
+        self.metric = metric
+
+    def fit_nndm(self, X: Any = None) -> NNDMResult:
+        """Run the NNDM algorithm and return the full result (distances included)."""
+        return nndm(
+            self._get_coords(X),
+            self.prediction_points,
+            phi=self.phi,
+            min_train=self.min_train,
+            metric=self.metric,
+        )
+
+    def get_n_splits(self, X: Any = None, y: Any = None, groups: Any = None) -> int:
+        return len(self._get_coords(X))
+
+    def split(
+        self, X: Any, y: Any = None, groups: Any = None
+    ) -> Iterator[Tuple[np.ndarray, np.ndarray]]:
+        res = self.fit_nndm(X)
+        yield from zip(res.train_indices, res.test_indices)
+
+
+# ---------------------------------------------------------------------------------------
+# kNNDM
+# ---------------------------------------------------------------------------------------
+@dataclass
+class KNNDMResult:
+    """Output of :func:`knndm`.
+
+    Attributes
+    ----------
+    labels : ndarray (n,)
+        Test fold (0 .. n_splits-1) of each training sample.
+    Gij, Gj, Gjstar : ndarray
+        Prediction-to-training, LOO and kNNDM CV nearest-neighbour distances.
+    W : float
+        Wasserstein distance between ``Gjstar`` and ``Gij`` of the chosen configuration.
+    n_clusters : int or None
+        Number of clusters of the chosen configuration (None when random CV was chosen).
+    method : str
+        ``"hierarchical"``, ``"kmeans"`` or ``"random"``.
+    W_by_n_clusters : dict
+        W statistic of every candidate that satisfied ``maxp``.
+    """
+
+    labels: np.ndarray
+    Gij: np.ndarray
+    Gj: np.ndarray
+    Gjstar: np.ndarray
+    W: float
+    n_clusters: Optional[int]
+    method: str
+    W_by_n_clusters: Dict[int, float] = field(default_factory=dict)
+
+    @property
+    def train_indices(self) -> List[np.ndarray]:
+        return [np.flatnonzero(self.labels != f) for f in range(int(self.labels.max()) + 1)]
+
+    @property
+    def test_indices(self) -> List[np.ndarray]:
+        return [np.flatnonzero(self.labels == f) for f in range(int(self.labels.max()) + 1)]
+
+
+def cv_nearest_distances(coords: np.ndarray, labels: np.ndarray, metric: str) -> np.ndarray:
+    """For each sample, the distance to the nearest sample in a *different* fold."""
+    out = np.empty(len(coords))
+    for f in np.unique(labels):
+        test = labels == f
+        out[test] = nearest_distance(coords[test], coords[~test], metric)
+    return out
+
+
+def _merge_clusters(
+    clusters: np.ndarray, order_score: np.ndarray, n_splits: int, n: int
+) -> np.ndarray:
+    """CAST's rule: big clusters (>= n/k) get their own fold, the rest are dealt
+    round-robin to the remaining folds in the order of their 1st principal component."""
+    ids, inverse, counts = np.unique(clusters, return_inverse=True, return_counts=True)
+    inverse = np.asarray(inverse).reshape(-1)
+    centroid_score = np.bincount(inverse, weights=order_score) / counts
+    order = np.argsort(centroid_score, kind="stable")
+    fold_of = np.full(len(ids), -1)
+    next_fold = 0
+    for c in order:
+        if counts[c] >= n / n_splits and next_fold < n_splits:
+            fold_of[c] = next_fold
+            next_fold += 1
+    free = np.setdiff1d(np.arange(n_splits), fold_of[fold_of >= 0])
+    rest = [c for c in order if fold_of[c] < 0]
+    if rest and free.size:
+        for i, c in enumerate(rest):
+            fold_of[c] = free[i % free.size]
+    return fold_of[inverse]
+
+
+def knndm(
+    train_coords: Any,
+    prediction_coords: Any,
+    n_splits: int = 10,
+    *,
+    maxp: float = 0.5,
+    clustering: str = "hierarchical",
+    metric: str = "euclidean",
+    n_candidates: int = 100,
+    random_state: Any = None,
+) -> KNNDMResult:
+    """k-fold Nearest Neighbour Distance Matching (Linnenbrink et al. 2024).
+
+    The training points are clustered into ``q`` clusters for a log-spaced grid of
+    ``q`` between ``n_splits`` and ``n - 2``. For every ``q``, clusters holding at least
+    ``n / n_splits`` samples form a fold on their own and the remaining clusters, ordered
+    along the first principal component of the coordinates, are merged round-robin into
+    the other folds. Configurations where a fold exceeds ``maxp`` of the data are
+    discarded; the one minimising the Wasserstein distance ``W`` between the CV distances
+    ``Gj*`` and the prediction distances ``Gij`` is returned. If ordinary random CV already
+    matches (W of LOO distances below 1e-4), random folds are returned.
+
+    Parameters
+    ----------
+    train_coords, prediction_coords : array-like (n, 2), (m, 2)
+    n_splits : int, default=10
+        Number of folds ``k``.
+    maxp : float, default=0.5
+        Maximum fraction of samples in one fold, in [1/k, 1].
+    clustering : {"hierarchical", "kmeans"}, default="hierarchical"
+        Ward hierarchical clustering (CAST's default ``ward.D2``) or k-means.
+    metric : {"euclidean", "haversine"}
+    n_candidates : int, default=100
+        Length of the grid of cluster numbers (before removing duplicates).
+    random_state : int, RandomState instance or None
+        Used by k-means and by the random-CV fallback.
+    """
+    check_metric(metric)
+    k = check_n_splits(n_splits)
+    tpts = as_coords(train_coords, metric)
+    ppts = as_coords(prediction_coords, metric)
+    n = len(tpts)
+    if not (1.0 / k <= maxp <= 1.0):
+        raise ValueError(f"maxp must be between 1/n_splits and 1, got {maxp}")
+    if clustering not in ("hierarchical", "kmeans"):
+        raise ValueError(f"clustering must be 'hierarchical' or 'kmeans', got {clustering!r}")
+    if n < k + 2:
+        raise ValueError(f"kNNDM needs at least n_splits + 2 = {k + 2} training points")
+    if len(ppts) == 0:
+        raise ValueError("prediction_coords is empty")
+    rng = check_random_state(random_state)
+
+    Gj = nearest_neighbour_loo(tpts, metric)
+    Gij = nearest_distance(ppts, tpts, metric)
+
+    if wasserstein_distance(Gj, Gij) < 1e-4:
+        pool = np.tile(np.arange(k), int(np.ceil(n / k)))
+        labels = rng.choice(pool, size=n, replace=False)
+        Gjstar = cv_nearest_distances(tpts, labels, metric)
+        return KNNDMResult(
+            labels, Gij, Gj, Gjstar, wasserstein_distance(Gjstar, Gij), None, "random"
+        )
+
+    grid = np.unique(np.round(np.exp(np.linspace(np.log(k), np.log(n - 2), n_candidates))))
+    grid = grid.astype(int)
+
+    planar = to_planar(tpts, metric)
+    centred = planar - planar.mean(axis=0)
+    _, _, vt = np.linalg.svd(centred, full_matrices=False)
+    pc1 = centred @ vt[0]
+
+    if clustering == "hierarchical":
+        from scipy.cluster.hierarchy import cut_tree, linkage
+        from scipy.spatial.distance import squareform
+
+        if metric == "haversine":
+            dmat = pairwise_distances(tpts, tpts, metric)
+            np.fill_diagonal(dmat, 0.0)
+            Z = linkage(squareform(dmat, checks=False), method="ward")
+        else:
+            Z = linkage(tpts, method="ward")
+        cuts = np.asarray(cut_tree(Z, n_clusters=grid))
+        cluster_sets = [cuts[:, j] for j in range(len(grid))]
+    else:
+        data = to_unit_sphere(tpts) if metric == "haversine" else tpts
+        cluster_sets = []
+        for q in grid:
+            km = KMeans(
+                n_clusters=int(q), n_init=1, random_state=rng.randint(np.iinfo(np.int32).max)
+            )
+            cluster_sets.append(km.fit_predict(data))
+
+    W_by_q: Dict[int, float] = {}
+    best: Optional[Tuple[float, int, np.ndarray, np.ndarray]] = None
+    for q, clusters in zip(grid, cluster_sets):
+        labels = _merge_clusters(np.asarray(clusters), pc1, k, n)
+        if np.any(labels < 0):
+            continue
+        if np.max(np.bincount(labels, minlength=k)) / n > maxp:
+            continue
+        Gjstar_q = cv_nearest_distances(tpts, labels, metric)
+        W_q = wasserstein_distance(Gjstar_q, Gij)
+        W_by_q[int(q)] = W_q
+        if best is None or W_q < best[0]:
+            best = (W_q, int(q), labels, Gjstar_q)
+    if best is None:
+        raise ValueError(
+            f"No cluster configuration satisfies maxp={maxp}; increase maxp or n_splits."
+        )
+    W_best, q_best, labels_best, Gjstar_best = best
+    return KNNDMResult(labels_best, Gij, Gj, Gjstar_best, W_best, q_best, clustering, W_by_q)
+
+
+class KNNDMKFold(SpatialSplitter):
+    """k-fold Nearest Neighbour Distance Matching CV (Linnenbrink et al. 2024).
+
+    A scikit-learn splitter around :func:`knndm`: whole folds are held out and the
+    folds are chosen so that CV test-to-training distances resemble prediction-to-training
+    distances.
+
+    Parameters
+    ----------
+    prediction_points : array-like (m, 2)
+    n_splits : int, default=10
+    maxp : float, default=0.5
+    clustering : {"hierarchical", "kmeans"}, default="hierarchical"
+    coords, coord_columns, metric
+        See :class:`SpatialBlockKFold`.
+    random_state : int, RandomState instance or None
+    """
+
+    def __init__(
+        self,
+        prediction_points: Any,
+        n_splits: int = 10,
+        *,
+        maxp: float = 0.5,
+        clustering: str = "hierarchical",
+        coords: Any = None,
+        coord_columns: Optional[Sequence[Any]] = None,
+        metric: str = "euclidean",
+        random_state: Any = None,
+    ) -> None:
+        self.prediction_points = prediction_points
+        self.n_splits = check_n_splits(n_splits)
+        self.maxp = maxp
+        self.clustering = clustering
+        self.coords = coords
+        self.coord_columns = coord_columns
+        self.metric = metric
+        self.random_state = random_state
+
+    def fit_knndm(self, X: Any = None) -> KNNDMResult:
+        """Run kNNDM and return the full result (distances and W included)."""
+        return knndm(
+            self._get_coords(X),
+            self.prediction_points,
+            self.n_splits,
+            maxp=self.maxp,
+            clustering=self.clustering,
+            metric=self.metric,
+            random_state=self.random_state,
+        )
+
+    def _fold_labels(self, X: Any, y: Any = None, groups: Any = None) -> np.ndarray:
+        return self.fit_knndm(X).labels

spatialcv/plotting.py

@@ -0,0 +1,83 @@
+"""Fold maps (optional matplotlib dependency: ``pip install spatialcv[plot]``)."""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Sequence
+
+import numpy as np
+
+from ._coords import resolve_coords
+
+
+def _pyplot() -> Any:
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError as exc:  # pragma: no cover - exercised only without matplotlib
+        raise ImportError("Plotting requires matplotlib: pip install 'spatialcv[plot]'") from exc
+    return plt
+
+
+def plot_folds(
+    cv: Any,
+    X: Any = None,
+    y: Any = None,
+    groups: Any = None,
+    *,
+    fold: Optional[int] = None,
+    coords: Any = None,
+    coord_columns: Optional[Sequence[Any]] = None,
+    ax: Any = None,
+    cmap: str = "tab10",
+    s: float = 12.0,
+) -> Any:
+    """Map the folds of a splitter.
+
+    With ``fold=None`` every sample is coloured by the split in which it is tested (best
+    for partitioning splitters). With ``fold=i`` split ``i`` is shown as training (grey),
+    test (colour) and excluded points (hollow, e.g. inside a buffer).
+
+    Returns the matplotlib ``Axes``.
+    """
+    plt = _pyplot()
+    metric = getattr(cv, "metric", "euclidean")
+    if coords is None and coord_columns is None:
+        coords = getattr(cv, "coords", None)
+        coord_columns = getattr(cv, "coord_columns", None)
+    xy = resolve_coords(X, coords, coord_columns, metric)
+    data = X if X is not None else xy
+    if ax is None:
+        _, ax = plt.subplots(figsize=(6, 6))
+    splits = list(cv.split(data, y, groups))
+    if fold is None:
+        label = np.full(len(xy), -1)
+        for i, (_, test) in enumerate(splits):
+            label[np.asarray(test)] = i
+        colors = plt.get_cmap(cmap)
+        for i in range(len(splits)):
+            m = label == i
+            if m.any():
+                ax.scatter(xy[m, 0], xy[m, 1], s=s, color=colors(i % 10), label=f"fold {i}")
+        if (label < 0).any():
+            m = label < 0
+            ax.scatter(xy[m, 0], xy[m, 1], s=s, color="0.8", label="never tested")
+        if len(splits) <= 12:
+            ax.legend(frameon=False, fontsize="small", markerscale=1.5)
+        ax.set_title(f"{type(cv).__name__}: test folds")
+    else:
+        train, test = (np.asarray(a) for a in splits[fold])
+        excluded = np.setdiff1d(np.arange(len(xy)), np.union1d(train, test))
+        ax.scatter(xy[train, 0], xy[train, 1], s=s, color="0.6", label="train")
+        if excluded.size:
+            ax.scatter(
+                xy[excluded, 0],
+                xy[excluded, 1],
+                s=s,
+                facecolors="none",
+                edgecolors="tab:orange",
+                label="excluded",
+            )
+        ax.scatter(xy[test, 0], xy[test, 1], s=s * 1.5, color="tab:blue", label="test")
+        ax.legend(frameon=False, fontsize="small")
+        ax.set_title(f"{type(cv).__name__}: split {fold}")
+    ax.set_aspect("equal", adjustable="datalim")
+    return ax