Python Implementation of Spatial Threat Modeling

from __future__ import annotations

from typing import Dict, Iterable

import numpy as np


class SpatialThreatModel:
    """Release-boundary guard for privacy-preserving GIS pipelines.

    Combines density-aware epsilon calibration, coordinate-precision
    clamping, L2-clipped secure gradient aggregation, and a spatial
    autocorrelation stability gate. Every output that crosses a
    compliance boundary is expected to pass `is_release_safe`.
    """

    def __init__(
        self,
        base_epsilon: float = 1.0,
        min_epsilon: float = 0.1,
        max_epsilon: float = 5.0,
        max_precision: int = 6,
        clip_norm: float = 1.0,
        autocorr_tolerance: float = 0.15,
    ) -> None:
        self.base_epsilon = base_epsilon
        self.min_epsilon = min_epsilon
        self.max_epsilon = max_epsilon
        self.max_precision = max_precision  # ~0.11 m at 6 dp
        self.clip_norm = clip_norm
        self.autocorr_tolerance = autocorr_tolerance

    def dynamic_epsilon(self, population_density: np.ndarray) -> np.ndarray:
        """Scale epsilon by density.

        Dense cells get a larger epsilon (less noise — crowds already
        give k-anonymity); sparse cells get a smaller epsilon (more
        noise — re-identification risk is highest there).
        """
        density = np.asarray(population_density, dtype=float)
        # log1p(0) == 0, so guard the denominator before normalising.
        scaling = np.log1p(density) / np.log1p(float(np.max(density)) + 1.0)
        eps = self.base_epsilon * scaling
        return np.clip(eps, self.min_epsilon, self.max_epsilon)

    def clamp_precision(self, coords: np.ndarray) -> np.ndarray:
        """Truncate lon/lat to the configured precision floor.

        Sub-meter precision in non-critical telemetry is a leak that
        survives aggregation, so it is rounded away before any threat
        logic runs.
        """
        coords = np.asarray(coords, dtype=float)
        if coords.ndim != 2 or coords.shape[1] != 2:
            raise ValueError("coords must be an (N, 2) lon/lat array")
        return np.round(coords, self.max_precision)

    def secure_gradient_aggregate(
        self,
        local_gradients: Dict[str, np.ndarray],
        epsilon: float,
        delta: float = 1e-5,
    ) -> np.ndarray:
        """L2-clip, average, then add Gaussian noise for (epsilon, delta)-DP.

        Clipping bounds any single node's pull on the mean — the core
        defence against gradient inversion and aggregation poisoning in
        a federated round.
        """
        clipped = []
        for grad in local_gradients.values():
            norm = float(np.linalg.norm(grad))
            if norm > self.clip_norm:
                grad = grad * (self.clip_norm / norm)
            clipped.append(np.asarray(grad, dtype=float))

        n = max(len(clipped), 1)
        aggregated = np.sum(clipped, axis=0) / n

        # Sensitivity of the mean over n disjoint contributors.
        sensitivity = self.clip_norm / n
        sigma = (sensitivity * np.sqrt(2.0 * np.log(1.25 / delta))) / epsilon
        return aggregated + np.random.normal(0.0, sigma, size=aggregated.shape)

    @staticmethod
    def _morans_i_proxy(coords: np.ndarray, values: np.ndarray) -> float:
        """Cheap spatial autocorrelation proxy in roughly [-1, 1].

        For production gating substitute `esda.Moran(values, weights)`;
        this closed form compares each value to its 3 nearest
        neighbours and avoids a heavy dependency in the hot path.
        """
        from scipy.spatial import cKDTree

        values = np.asarray(values, dtype=float)
        if len(values) < 4:
            return 0.0
        tree = cKDTree(coords)
        _, idx = tree.query(coords, k=4)  # self + 3 neighbours
        neighbour_means = values[idx[:, 1:]].mean(axis=1)
        centered = values - values.mean()
        neighbour_centered = neighbour_means - values.mean()
        denom = float((centered ** 2).sum())
        if denom == 0.0:
            return 0.0
        return float((centered * neighbour_centered).sum() / denom)

    def is_release_safe(
        self,
        original_coords: np.ndarray,
        perturbed_coords: np.ndarray,
        values: np.ndarray,
    ) -> bool:
        """Gate a release on cluster-structure preservation.

        Returns True only if perturbation kept Moran's I within the
        tolerance band — i.e. the noise did not shred clinically or
        financially meaningful clusters.
        """
        before = self._morans_i_proxy(self.clamp_precision(original_coords), values)
        after = self._morans_i_proxy(self.clamp_precision(perturbed_coords), values)
        return abs(before - after) <= self.autocorr_tolerance


def _run_validation() -> None:
    rng = np.random.default_rng(7)
    model = SpatialThreatModel()

    # 1. Sparse zones must never receive more budget than dense zones.
    density = np.array([1.0, 10.0, 1000.0])
    eps = model.dynamic_epsilon(density)
    assert eps[0] <= eps[1] <= eps[2]
    assert np.all(eps >= model.min_epsilon) and np.all(eps <= model.max_epsilon)

    # 2. Precision is clamped to the configured floor.
    fine = np.array([[12.3456789, -1.2345678]])
    clamped = model.clamp_precision(fine)
    assert np.allclose(clamped, np.round(fine, 6))

    # 3. Tighter epsilon must inject strictly more noise.
    grads = {f"n{i}": rng.normal(size=8) for i in range(5)}
    spread = lambda e: float(np.std(model.secure_gradient_aggregate(grads, epsilon=e)))
    assert spread(0.5) > spread(5.0)

    # 4. A faithful (near-identity) perturbation passes the gate;
    #    destroying structure fails it.
    coords = rng.uniform(0, 1, size=(60, 2))
    values = coords[:, 0] + coords[:, 1]  # smooth, highly autocorrelated
    gentle = coords + rng.normal(0, 1e-4, size=coords.shape)
    assert model.is_release_safe(coords, gentle, values) is True
    assert model.is_release_safe(coords, rng.uniform(0, 1, size=coords.shape), values) is False

    print("SpatialThreatModel: all assertions passed")


if __name__ == "__main__":
    _run_validation()