spatialcv 0.1.0__py3-none-any.whl

spatialcv/cluster.py

@@ -0,0 +1,143 @@
+"""Cluster-based splitters: k-means in geographic space or in feature space."""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Sequence
+
+import numpy as np
+from sklearn.cluster import KMeans
+from sklearn.utils import check_random_state
+
+from ._base import SpatialSplitter, assign_groups_to_folds, check_n_splits
+from ._coords import to_unit_sphere
+
+
+def _kmeans_folds(
+    data: np.ndarray,
+    n_splits: int,
+    n_clusters: Optional[int],
+    random_state: Any,
+    balance: bool,
+) -> np.ndarray:
+    k = n_splits if n_clusters is None else int(n_clusters)
+    if k < n_splits:
+        raise ValueError(f"n_clusters ({k}) must be >= n_splits ({n_splits})")
+    if len(data) < k:
+        raise ValueError(f"{len(data)} samples cannot form {k} clusters")
+    rng = check_random_state(random_state)
+    km = KMeans(n_clusters=k, n_init=10, random_state=rng.randint(np.iinfo(np.int32).max))
+    clusters = km.fit_predict(data)
+    _, inverse, counts = np.unique(clusters, return_inverse=True, return_counts=True)
+    inverse = np.asarray(inverse).reshape(-1)
+    if k == n_splits and len(counts) == n_splits:
+        return inverse.astype(int)
+    fold_of_cluster = assign_groups_to_folds(counts, n_splits, rng, balance)
+    return fold_of_cluster[inverse]
+
+
+class SpatialClusterKFold(SpatialSplitter):
+    """K-fold cross-validation on k-means clusters of the sample coordinates.
+
+    Samples are grouped by k-means on their coordinates (on 3D unit vectors for lon/lat,
+    so clusters respect the sphere) and whole clusters form the test folds
+    (Brenning 2012; spatial cluster CV in ``sperrorest`` and ``spatialsample``).
+
+    Parameters
+    ----------
+    n_splits : int, default=5
+    n_clusters : int, optional
+        Number of k-means clusters. Defaults to ``n_splits`` (one cluster per fold). With
+        more clusters than folds, clusters are combined into folds (see ``balance``).
+    balance : bool, default=True
+        When ``n_clusters > n_splits``, combine clusters so folds have similar sizes.
+    coords, coord_columns, metric
+        See :class:`SpatialBlockKFold`.
+    random_state : int, RandomState instance or None
+    """
+
+    def __init__(
+        self,
+        n_splits: int = 5,
+        *,
+        n_clusters: Optional[int] = None,
+        balance: bool = True,
+        coords: Any = None,
+        coord_columns: Optional[Sequence[Any]] = None,
+        metric: str = "euclidean",
+        random_state: Any = None,
+    ) -> None:
+        self.n_splits = check_n_splits(n_splits)
+        self.n_clusters = n_clusters
+        self.balance = balance
+        self.coords = coords
+        self.coord_columns = coord_columns
+        self.metric = metric
+        self.random_state = random_state
+
+    def _fold_labels(self, X: Any, y: Any = None, groups: Any = None) -> np.ndarray:
+        xy = self._get_coords(X)
+        data = to_unit_sphere(xy) if self.metric == "haversine" else xy
+        return _kmeans_folds(data, self.n_splits, self.n_clusters, self.random_state, self.balance)
+
+
+class EnvironmentalBlockKFold(SpatialSplitter):
+    """K-fold cross-validation on k-means clusters in *feature* space.
+
+    Samples with similar predictor values end up in the same fold, so each test fold
+    represents environmental conditions that are under-represented in its training data.
+    This is the "environmental blocking" of Roberts et al. (2017) and blockCV's
+    ``cv_cluster(..., r = covariates)`` (Valavi et al. 2019).
+
+    Parameters
+    ----------
+    n_splits : int, default=5
+    n_clusters : int, optional
+        Defaults to ``n_splits``.
+    features : sequence of column names or positions, optional
+        Columns of ``X`` to cluster on; all columns by default.
+    standardize : bool, default=True
+        Scale each feature to zero mean and unit variance before clustering.
+    balance : bool, default=True
+    random_state : int, RandomState instance or None
+    """
+
+    def __init__(
+        self,
+        n_splits: int = 5,
+        *,
+        n_clusters: Optional[int] = None,
+        features: Optional[Sequence[Any]] = None,
+        standardize: bool = True,
+        balance: bool = True,
+        random_state: Any = None,
+    ) -> None:
+        self.n_splits = check_n_splits(n_splits)
+        self.n_clusters = n_clusters
+        self.features = features
+        self.standardize = standardize
+        self.balance = balance
+        self.random_state = random_state
+
+    def _fold_labels(self, X: Any, y: Any = None, groups: Any = None) -> np.ndarray:
+        if X is None:
+            raise ValueError("EnvironmentalBlockKFold needs the feature matrix X")
+        if self.features is not None:
+            cols = list(self.features)
+            if hasattr(X, "iloc"):
+                sub = X[cols] if all(c in X.columns for c in cols) else X.iloc[:, cols]
+                data = np.asarray(sub, dtype=float)
+            else:
+                data = np.asarray(X)[:, cols].astype(float)
+        else:
+            if hasattr(X, "select_dtypes"):
+                X = X.select_dtypes(include=[np.number])
+            data = np.asarray(X, dtype=float)
+        if data.ndim == 1:
+            data = data[:, None]
+        if not np.all(np.isfinite(data)):
+            raise ValueError("feature matrix contains NaN or infinite values")
+        if self.standardize:
+            std = data.std(axis=0)
+            std[std == 0] = 1.0
+            data = (data - data.mean(axis=0)) / std
+        return _kmeans_folds(data, self.n_splits, self.n_clusters, self.random_state, self.balance)

spatialcv/datasets.py

@@ -0,0 +1,101 @@
+"""Synthetic spatially autocorrelated data for examples and tests."""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Tuple
+
+import numpy as np
+from sklearn.utils import check_random_state
+
+
+def gaussian_random_field(
+    coords: Any,
+    correlation_range: float,
+    *,
+    n_components: int = 500,
+    random_state: Any = None,
+) -> np.ndarray:
+    """Unit-variance stationary Gaussian random field sampled at ``coords``.
+
+    Uses random Fourier features (Rahimi & Recht 2007) to approximate a Gaussian
+    covariance whose practical range (correlation 0.05) is ``correlation_range``.
+    """
+    rng = check_random_state(random_state)
+    xy = np.asarray(coords, dtype=float)
+    length = correlation_range / math.sqrt(2.0 * math.log(20.0))
+    omega = rng.normal(scale=1.0 / length, size=(2, n_components))
+    phase = rng.uniform(0.0, 2.0 * math.pi, size=n_components)
+    return np.asarray(math.sqrt(2.0 / n_components) * np.cos(xy @ omega + phase).sum(axis=1))
+
+
+def make_spatial_regression(
+    n_samples: int = 500,
+    *,
+    n_features: int = 4,
+    extent: float = 10_000.0,
+    correlation_range: float = 1_500.0,
+    noise: float = 0.2,
+    sampling: str = "clustered",
+    n_clusters: int = 8,
+    prediction_points: Any = None,
+    random_state: Any = None,
+) -> Tuple[np.ndarray, ...]:
+    """Regression data whose features *and* residuals are spatially autocorrelated.
+
+    Samples lie in the square ``[0, extent]^2``. Each feature is an independent Gaussian
+    random field; the target is a non-linear function of the features plus another
+    random field the features cannot explain, plus white noise.
+
+    Parameters
+    ----------
+    sampling : {"clustered", "uniform"}, default="clustered"
+        ``"clustered"`` draws samples around ``n_clusters`` random centres (a Thomas
+        process), the typical situation of field data, where random CV is most optimistic.
+    prediction_points : array-like (m, 2), optional
+        Extra locations drawn from the *same* fields. Their features and target are
+        returned too, so the true map accuracy can be compared with CV estimates.
+
+    Returns
+    -------
+    coords : ndarray (n_samples, 2)
+    X : ndarray (n_samples, n_features)
+    y : ndarray (n_samples,)
+    X_pred, y_pred : ndarray
+        Only when ``prediction_points`` is given.
+    """
+    rng = check_random_state(random_state)
+    if sampling == "uniform":
+        coords = rng.uniform(0.0, extent, size=(n_samples, 2))
+    elif sampling == "clustered":
+        centres = rng.uniform(0.1 * extent, 0.9 * extent, size=(n_clusters, 2))
+        which = rng.randint(n_clusters, size=n_samples)
+        coords = centres[which] + rng.normal(scale=0.05 * extent, size=(n_samples, 2))
+        coords = np.clip(coords, 0.0, extent)
+    else:
+        raise ValueError(f"sampling must be 'clustered' or 'uniform', got {sampling!r}")
+    pred = np.empty((0, 2)) if prediction_points is None else np.asarray(prediction_points, float)
+    allxy = np.vstack([coords, pred])
+    X = np.column_stack(
+        [
+            gaussian_random_field(allxy, correlation_range, random_state=rng)
+            for _ in range(n_features)
+        ]
+    )
+    signal = np.sin(X[:, 0]) + 0.5 * X[:, 1 % n_features] ** 2
+    if n_features > 2:
+        signal = signal + 0.5 * X[:, 2]
+    residual = gaussian_random_field(allxy, correlation_range, random_state=rng)
+    y = signal + residual + noise * rng.normal(size=len(allxy))
+    if prediction_points is None:
+        return coords, X, y
+    n = n_samples
+    return coords, X[:n], y[:n], X[n:], y[n:]
+
+
+def make_prediction_grid(extent: float = 10_000.0, n_per_side: int = 40) -> np.ndarray:
+    """Regular grid of cell centres covering ``[0, extent]^2`` (prediction locations)."""
+    step = extent / n_per_side
+    g = (np.arange(n_per_side) + 0.5) * step
+    gx, gy = np.meshgrid(g, g)
+    return np.column_stack([gx.ravel(), gy.ravel()])

spatialcv/diagnostics.py

@@ -0,0 +1,151 @@
+"""Diagnostics: how far are CV test points from their training data?
+
+Compares the nearest-neighbour distance distributions of Milà et al. (2022):
+
+* ``Gij``: prediction location -> nearest training sample (what the model faces in use);
+* ``Gjstar``: CV test sample -> nearest sample in its training set (what CV measures);
+* ``Gj``: training sample -> nearest other training sample (ordinary random LOO CV).
+
+If ``Gjstar`` sits well to the left of ``Gij`` the CV estimate is optimistic. The
+Wasserstein distance ``W`` (Linnenbrink et al. 2024) summarises the mismatch.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Sequence
+
+import numpy as np
+
+from ._coords import (
+    as_coords,
+    check_metric,
+    nearest_distance,
+    nearest_neighbour_loo,
+    resolve_coords,
+)
+from .nndm import wasserstein_distance
+
+
+@dataclass
+class FoldDiagnostics:
+    """Result of :func:`fold_diagnostics`.
+
+    Attributes
+    ----------
+    Gjstar : ndarray
+        Test -> training nearest-neighbour distance of every test sample in every split.
+    Gj : ndarray
+        LOO nearest-neighbour distance among all samples.
+    Gij : ndarray or None
+        Prediction -> training distance (None without ``prediction_points``).
+    W : float or None
+        Wasserstein distance between ``Gjstar`` and ``Gij``.
+    W_random : float or None
+        Wasserstein distance between ``Gj`` and ``Gij`` (what random LOO would give).
+    n_train, n_test : list of int
+        Sizes of every split.
+    """
+
+    Gjstar: np.ndarray
+    Gj: np.ndarray
+    Gij: Optional[np.ndarray]
+    W: Optional[float]
+    W_random: Optional[float]
+    n_train: List[int]
+    n_test: List[int]
+
+    def summary(self) -> Dict[str, Any]:
+        """Medians and W statistics as a plain dict."""
+        out: Dict[str, Any] = {
+            "n_splits": len(self.n_test),
+            "median_train_size": float(np.median(self.n_train)),
+            "median_test_size": float(np.median(self.n_test)),
+            "median_Gjstar": float(np.median(self.Gjstar)),
+            "median_Gj": float(np.median(self.Gj)),
+        }
+        if self.Gij is not None:
+            out["median_Gij"] = float(np.median(self.Gij))
+            out["W"] = self.W
+            out["W_random"] = self.W_random
+        return out
+
+    def __str__(self) -> str:
+        s = self.summary()
+        lines = [f"FoldDiagnostics ({s['n_splits']} splits)"]
+        for k, v in s.items():
+            if k == "n_splits":
+                continue
+            lines.append(f"  {k:<18} {v:.4g}" if isinstance(v, float) else f"  {k:<18} {v}")
+        return "\n".join(lines)
+
+    def plot(self, ax: Any = None) -> Any:
+        """ECDFs of the three distance distributions (requires matplotlib)."""
+        from .plotting import _pyplot
+
+        plt = _pyplot()
+        if ax is None:
+            _, ax = plt.subplots(figsize=(6, 4))
+        series = [
+            (self.Gj, "random LOO (Gj)", "0.55"),
+            (self.Gjstar, "this CV (Gj*)", "tab:blue"),
+        ]
+        if self.Gij is not None:
+            series.append((self.Gij, "prediction (Gij)", "tab:red"))
+        for values, label, color in series:
+            x = np.sort(values)
+            ax.step(x, np.arange(1, x.size + 1) / x.size, where="post", label=label, color=color)
+        ax.set_xlabel("nearest-neighbour distance")
+        ax.set_ylabel("ECDF")
+        ax.legend(frameon=False)
+        return ax
+
+
+def fold_diagnostics(
+    cv: Any,
+    X: Any = None,
+    y: Any = None,
+    groups: Any = None,
+    *,
+    prediction_points: Any = None,
+    coords: Any = None,
+    coord_columns: Optional[Sequence[Any]] = None,
+    metric: Optional[str] = None,
+) -> FoldDiagnostics:
+    """Nearest-neighbour distance diagnostics for any splitter.
+
+    Parameters
+    ----------
+    cv : cross-validator
+        Any object with ``split`` (spatialcv or scikit-learn splitters).
+    X, y, groups
+        Passed to ``cv.split``.
+    prediction_points : array-like (m, 2), optional
+        Where the model will predict. Enables ``Gij`` and the W statistics.
+    coords, coord_columns, metric
+        Sample coordinates; default to the splitter's own settings, then to ``X``.
+    """
+    metric = str(metric or getattr(cv, "metric", "euclidean"))
+    check_metric(metric)
+    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
+    gjstar: List[np.ndarray] = []
+    n_train: List[int] = []
+    n_test: List[int] = []
+    for train, test in cv.split(data, y, groups):
+        train = np.asarray(train)
+        test = np.asarray(test)
+        gjstar.append(nearest_distance(xy[test], xy[train], metric))
+        n_train.append(len(train))
+        n_test.append(len(test))
+    Gjstar = np.concatenate(gjstar) if gjstar else np.empty(0)
+    Gj = nearest_neighbour_loo(xy, metric)
+    Gij = W = W_random = None
+    if prediction_points is not None:
+        Gij = nearest_distance(as_coords(prediction_points, metric), xy, metric)
+        W = wasserstein_distance(Gjstar, Gij)
+        W_random = wasserstein_distance(Gj, Gij)
+    return FoldDiagnostics(Gjstar, Gj, Gij, W, W_random, n_train, n_test)