spatialsample 0.1.0__py3-none-any.whl

spatialsample/designs.py

@@ -0,0 +1,562 @@
+"""Classical and field designs: simple random, systematic, stratified, Poisson-disk,
+cluster, two-stage and transect sampling."""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Callable, Dict, Hashable, List, Mapping, Optional, Sequence, Union
+
+import numpy as np
+
+from ._frame import as_domain
+from ._types import Domain, FloatArray, SeedLike, as_rng
+from .bas import bas
+from .grts import grts
+from .lpm import lpm
+from .probabilities import allocate, inclusion_probabilities, systematic_pps
+from .raster import RasterMask
+from .sample import Sample
+
+__all__ = [
+    "srs",
+    "systematic",
+    "stratified",
+    "stratified_pixels",
+    "poisson_disk",
+    "cluster",
+    "two_stage",
+    "points_along_lines",
+    "random_transects",
+]
+
+
+def srs(domain: Any, n: int, seed: SeedLike = None) -> Sample:
+    """Simple random sample: ``n`` independent uniform locations in the domain."""
+    dom = as_domain(domain)
+    if n < 1:
+        raise ValueError("n must be >= 1")
+    xy = dom.random_points(n, as_rng(seed))
+    density = n / dom.area
+    return Sample(
+        xy=xy,
+        pi=np.full(n, density),
+        weight=np.full(n, 1.0 / density),
+        design="srs",
+        meta={"n": n, "area": dom.area},
+    )
+
+
+def systematic(
+    domain: Any,
+    n: Optional[int] = None,
+    spacing: Optional[float] = None,
+    shape: str = "square",
+    rotation: Union[float, str, None] = 0.0,
+    seed: SeedLike = None,
+) -> Sample:
+    """Systematic grid with a uniformly random start.
+
+    Parameters
+    ----------
+    n:
+        Target sample size; the spacing is chosen so that the expected number
+        of points in the domain is ``n``. The realised size varies slightly.
+    spacing:
+        Distance between neighbouring grid points (overrides ``n``).
+    shape:
+        ``"square"`` or ``"hex"`` (triangular lattice, hexagonal cells).
+    rotation:
+        Grid rotation in degrees, or ``"random"`` for a uniformly random angle.
+    """
+    dom = as_domain(domain)
+    rng = as_rng(seed)
+    shape = shape.lower()
+    if shape not in ("square", "hex"):
+        raise ValueError("shape must be 'square' or 'hex'")
+    if spacing is None:
+        if n is None or n < 1:
+            raise ValueError("give n >= 1 or spacing")
+        cell = dom.area / n
+        spacing = math.sqrt(cell) if shape == "square" else math.sqrt(2.0 * cell / math.sqrt(3.0))
+    if spacing <= 0:
+        raise ValueError("spacing must be positive")
+    s = float(spacing)
+    row_h = s if shape == "square" else s * math.sqrt(3.0) / 2.0
+    cell_area = s * row_h
+    if rotation == "random":
+        theta = rng.uniform(0.0, 180.0 if shape == "square" else 60.0)
+    else:
+        theta = float(rotation or 0.0)
+    x0, y0, x1, y1 = dom.bounds
+    cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
+    r = math.hypot(x1 - x0, y1 - y0) / 2 + 2 * s
+    off = np.array([rng.uniform(0, s), rng.uniform(0, row_h)])
+    ni = int(math.ceil(r / s)) + 2
+    nj = int(math.ceil(r / row_h)) + 2
+    ii, jj = np.meshgrid(np.arange(-ni, ni + 1), np.arange(-nj, nj + 1))
+    gx = ii * s + off[0]
+    if shape == "hex":
+        gx = gx + (jj % 2) * (s / 2.0)
+    gy = jj * row_h + off[1]
+    t = math.radians(theta)
+    ct, st = math.cos(t), math.sin(t)
+    px = cx + gx.ravel() * ct - gy.ravel() * st
+    py = cy + gx.ravel() * st + gy.ravel() * ct
+    pts = np.column_stack([px, py])
+    pts = pts[dom.contains(pts)]
+    # serpentine (boustrophedon) visiting order in the grid frame
+    lx = (pts[:, 0] - cx) * ct + (pts[:, 1] - cy) * st
+    ly = -(pts[:, 0] - cx) * st + (pts[:, 1] - cy) * ct
+    row = np.round((ly - off[1]) / row_h).astype(np.int64)
+    key = np.where(row % 2 == 0, lx, -lx)
+    order_idx = np.lexsort((key, row))
+    pts = pts[order_idx]
+    m = len(pts)
+    density = 1.0 / cell_area
+    return Sample(
+        xy=pts,
+        pi=np.full(m, density),
+        weight=np.full(m, cell_area),
+        design=f"systematic-{shape}",
+        meta={
+            "n": n,
+            "spacing": s,
+            "rotation": theta,
+            "area": dom.area,
+            "expected_n": dom.area / cell_area,
+        },
+    )
+
+
+_CONTINUOUS: Dict[str, Callable[..., Sample]] = {"srs": srs, "bas": bas, "grts": grts, "lpm": lpm}
+
+
+def stratified(
+    strata: Any,
+    n: Union[int, Mapping[Hashable, int]],
+    allocation: str = "proportional",
+    sd: Optional[Mapping[Hashable, float]] = None,
+    min_per_stratum: int = 1,
+    method: str = "srs",
+    seed: SeedLike = None,
+) -> Sample:
+    """Stratified sample over polygons or raster classes.
+
+    Parameters
+    ----------
+    strata:
+        Mapping ``label -> domain`` (Regions, RasterMasks or GeoJSON), or a
+        :class:`~spatialsample.RasterMask` whose distinct class values become
+        strata.
+    n:
+        Total sample size, or an explicit mapping ``label -> n_h``.
+    allocation:
+        ``"proportional"`` (to stratum area), ``"equal"`` or ``"neyman"``
+        (proportional to area times ``sd[label]``).
+    min_per_stratum:
+        Minimum sites per stratum when allocating a total.
+    method:
+        Within-stratum design: ``"srs"``, ``"bas"``, ``"grts"`` or ``"lpm"``.
+    """
+    rng = as_rng(seed)
+    if isinstance(strata, RasterMask):
+        doms: Dict[Hashable, Domain] = {c: strata.select([c]) for c in strata.class_counts()}
+    elif isinstance(strata, Mapping) and strata.get("type") is None:
+        doms = {k: as_domain(v) for k, v in strata.items()}
+    else:
+        raise TypeError("strata must be a mapping label -> domain or a RasterMask")
+    if method not in _CONTINUOUS:
+        raise ValueError(f"method must be one of {sorted(_CONTINUOUS)}")
+    areas = {k: d.area for k, d in doms.items()}
+    if isinstance(n, Mapping):
+        alloc = {k: int(n.get(k, 0)) for k in doms}
+    else:
+        alloc = allocate(int(n), areas, allocation, sd=sd, min_per_stratum=min_per_stratum)
+    parts: List[Sample] = []
+    labels: List[Any] = []
+    for k, dom in doms.items():
+        nh = alloc[k]
+        if nh <= 0:
+            continue
+        sub = _CONTINUOUS[method](dom, nh, seed=rng)
+        parts.append(sub)
+        labels.extend([k] * len(sub))
+    if not parts:
+        raise ValueError("allocation produced an empty sample")
+    xy = np.vstack([p.xy for p in parts])
+    pi = np.concatenate([p.pi for p in parts])
+    return Sample(
+        xy=xy,
+        pi=pi,
+        weight=1.0 / pi,
+        design=f"stratified-{method}",
+        stratum=np.array(labels, dtype=object),
+        meta={
+            "n": int(sum(alloc.values())),
+            "allocation": alloc,
+            "stratum_area": areas,
+            "area": float(sum(areas.values())),
+        },
+    )
+
+
+def stratified_pixels(
+    mask: RasterMask,
+    n: Union[int, Mapping[Hashable, int]],
+    allocation: str = "proportional",
+    sd: Optional[Mapping[Hashable, float]] = None,
+    min_per_stratum: int = 1,
+    seed: SeedLike = None,
+) -> Sample:
+    """Stratified random sample of *pixels* (without replacement) by map class.
+
+    This is the design recommended by Olofsson et al. (2014) for accuracy
+    assessment and area estimation of land-cover and change maps: strata are
+    the map classes, sample units are pixels, and each pixel in stratum ``h``
+    has inclusion probability ``n_h / N_h``.
+    """
+    rng = as_rng(seed)
+    counts = mask.class_counts()
+    if isinstance(n, Mapping):
+        alloc = {k: int(n.get(k, 0)) for k in counts}
+    else:
+        alloc = allocate(
+            int(n), counts, allocation, sd=sd, min_per_stratum=min_per_stratum, caps=counts
+        )
+    centres, rows, cols = mask.cell_centers()
+    values = mask.data[rows, cols]
+    take: List[np.ndarray] = []
+    pis: List[np.ndarray] = []
+    labs: List[Any] = []
+    for k, nh in alloc.items():
+        if nh <= 0:
+            continue
+        idx = np.flatnonzero(values == k)
+        if nh > len(idx):
+            raise ValueError(f"stratum {k!r} has only {len(idx)} pixels")
+        sel = rng.choice(idx, size=nh, replace=False)
+        take.append(sel)
+        pis.append(np.full(nh, nh / len(idx)))
+        labs.extend([k] * nh)
+    sel_all = np.concatenate(take)
+    pi = np.concatenate(pis)
+    return Sample(
+        xy=centres[sel_all],
+        pi=pi,
+        weight=1.0 / pi,
+        design="stratified-pixels",
+        kind="finite",
+        stratum=np.array(labs, dtype=object),
+        index=sel_all,
+        attrs={"row": rows[sel_all], "col": cols[sel_all]},
+        meta={
+            "n": int(pi.size),
+            "allocation": alloc,
+            "stratum_size": counts,
+            "population_size": int(sum(counts.values())),
+            "cell_area": mask.cell_area,
+        },
+    )
+
+
+def poisson_disk(
+    domain: Any,
+    n: int,
+    min_dist: Optional[float] = None,
+    max_tries: int = 30000,
+    seed: SeedLike = None,
+) -> Sample:
+    """Minimum-distance (simple sequential inhibition) sample.
+
+    Uniform candidate points are accepted in turn if they lie at least
+    ``min_dist`` from every point already accepted. The default distance is
+    ``0.6 * sqrt(area / n)``, comfortably below the random-packing limit
+    (about ``0.83 * sqrt(area / n)``). Inclusion densities are reported as
+    ``n / area``, which is exact only away from the boundary; use GRTS, BAS or
+    LPM when design-based inference with exact probabilities matters.
+    """
+    dom = as_domain(domain)
+    rng = as_rng(seed)
+    if n < 1:
+        raise ValueError("n must be >= 1")
+    d = float(min_dist) if min_dist is not None else 0.6 * math.sqrt(dom.area / n)
+    d2 = d * d
+    acc = np.empty((n, 2))
+    k = 0
+    tried = 0
+    while k < n and tried < max_tries:
+        cand = dom.random_points(min(512, max_tries - tried), rng)
+        for c in cand:
+            tried += 1
+            if k == 0 or np.min(((acc[:k] - c) ** 2).sum(axis=1)) >= d2:
+                acc[k] = c
+                k += 1
+                if k == n:
+                    break
+    if k < n:
+        raise ValueError(f"only {k} of {n} points could be placed {d:g} apart; lower min_dist or n")
+    density = n / dom.area
+    return Sample(
+        xy=acc,
+        pi=np.full(n, density),
+        weight=np.full(n, 1.0 / density),
+        design="poisson-disk",
+        meta={"n": n, "min_dist": d, "area": dom.area, "approximate_pi": True},
+    )
+
+
+_LAYOUTS: Dict[str, FloatArray] = {
+    # US Forest Inventory and Analysis plot: centre plus three subplots at 120 degrees.
+    "fia": np.array(
+        [[0.0, 0.0]]
+        + [[math.sin(math.radians(a)), math.cos(math.radians(a))] for a in (0.0, 120.0, 240.0)]
+    ),
+    "square": np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) - 0.5,
+    "line": np.array([[0.0, k] for k in range(5)]) - np.array([0.0, 2.0]),
+    "cross": np.array([[0.0, 0.0], [1.0, 0.0], [-1.0, 0.0], [0.0, 1.0], [0.0, -1.0]]),
+}
+
+
+def cluster(
+    domain: Any,
+    n_clusters: int,
+    spacing: float,
+    layout: Union[str, Sequence[Sequence[float]]] = "fia",
+    rotation: Union[float, str] = 0.0,
+    centers: str = "srs",
+    seed: SeedLike = None,
+) -> Sample:
+    """Cluster sample: a fixed pattern of sub-plots around each sampled centre.
+
+    Centres are drawn with ``centers`` (``"srs"``, ``"bas"`` or ``"grts"``) and
+    every centre carries the sub-plot ``layout`` scaled by ``spacing``:
+    ``"fia"`` (centre + 3 at 120 degrees), ``"square"`` (2 x 2), ``"line"``
+    (5 in a row), ``"cross"`` (5) or an explicit list of unit offsets.
+    Sub-plots that fall outside the domain are dropped. Each sub-plot location
+    has inclusion density ``n_clusters / area`` (exact away from the edge);
+    analyse cluster means, not individual sub-plots, for variance estimation.
+    """
+    dom = as_domain(domain)
+    rng = as_rng(seed)
+    offs = _LAYOUTS[layout] if isinstance(layout, str) else np.asarray(layout, dtype=float)
+    centre_design = _CONTINUOUS.get(centers)
+    if centre_design is None or centers == "lpm":
+        raise ValueError("centers must be 'srs', 'bas' or 'grts'")
+    base = centre_design(dom, n_clusters, seed=rng)
+    cs = base.xy
+    m = len(offs)
+    if rotation == "random":
+        ang = rng.uniform(0, 2 * math.pi, len(cs))
+    else:
+        ang = np.full(len(cs), math.radians(float(rotation)))
+    ca, sa = np.cos(ang)[:, None], np.sin(ang)[:, None]
+    ox = offs[None, :, 0] * spacing
+    oy = offs[None, :, 1] * spacing
+    px = cs[:, 0:1] + ox * ca - oy * sa
+    py = cs[:, 1:2] + ox * sa + oy * ca
+    xy = np.column_stack([px.ravel(), py.ravel()])
+    cid = np.repeat(np.arange(len(cs)), m)
+    sub = np.tile(np.arange(m), len(cs))
+    keep = dom.contains(xy)
+    density = n_clusters / dom.area
+    k = int(keep.sum())
+    return Sample(
+        xy=xy[keep],
+        pi=np.full(k, density),
+        weight=np.full(k, 1.0 / density),
+        design=f"cluster-{centers}",
+        attrs={"cluster": cid[keep], "subplot": sub[keep]},
+        meta={"n_clusters": n_clusters, "subplots": m, "spacing": spacing, "area": dom.area},
+    )
+
+
+def two_stage(
+    psus: Mapping[Hashable, Any],
+    n_psu: int,
+    m_per_psu: int,
+    pps: bool = True,
+    seed: SeedLike = None,
+) -> Sample:
+    """Two-stage sample: primary units (polygons), then random sites inside them.
+
+    First stage: ``n_psu`` primary sampling units (PSUs, e.g. villages, census
+    tracts, forest compartments) are drawn without replacement, with
+    probability proportional to area (``pps=True``, randomised systematic
+    PPS) or with equal probability. Second stage: ``m_per_psu`` uniform sites
+    in each selected PSU. With PPS and a fixed ``m_per_psu`` the design is
+    self-weighting (every site represents the same area).
+    """
+    rng = as_rng(seed)
+    doms = {k: as_domain(v) for k, v in psus.items()}
+    keys = list(doms)
+    if not 1 <= n_psu <= len(keys):
+        raise ValueError("n_psu must be between 1 and the number of PSUs")
+    areas = np.array([doms[k].area for k in keys])
+    sizes = areas if pps else np.ones(len(keys))
+    pi1 = inclusion_probabilities(sizes, n_psu)
+    perm = rng.permutation(len(keys))
+    chosen = perm[systematic_pps(pi1[perm], rng)]
+    xy_parts: List[FloatArray] = []
+    dens: List[np.ndarray] = []
+    lab: List[Any] = []
+    for c in chosen:
+        k = keys[c]
+        pts = doms[k].random_points(m_per_psu, rng)
+        xy_parts.append(pts)
+        dens.append(np.full(m_per_psu, pi1[c] * m_per_psu / areas[c]))
+        lab.extend([k] * m_per_psu)
+    pi = np.concatenate(dens)
+    return Sample(
+        xy=np.vstack(xy_parts),
+        pi=pi,
+        weight=1.0 / pi,
+        design="two-stage",
+        attrs={
+            "psu": np.array(lab, dtype=object),
+            "psu_pi": np.repeat(pi1[chosen], m_per_psu),
+        },
+        meta={"n_psu": n_psu, "m_per_psu": m_per_psu, "pps": pps, "area": float(areas.sum())},
+    )
+
+
+def _as_lines(lines: Any) -> List[FloatArray]:
+    if isinstance(lines, Mapping) or hasattr(lines, "__geo_interface__"):
+        obj = lines if isinstance(lines, Mapping) else lines.__geo_interface__
+        out: List[FloatArray] = []
+        _collect_lines(obj, out)
+        if not out:
+            raise ValueError("no LineString geometry found")
+        return out
+    arr = [np.asarray(ln, dtype=float)[:, :2] for ln in lines]
+    for a in arr:
+        if a.ndim != 2 or len(a) < 2:
+            raise ValueError("each line needs at least two vertices")
+    return arr
+
+
+def _collect_lines(obj: Mapping[str, Any], out: List[FloatArray]) -> None:
+    t = obj.get("type")
+    if t == "FeatureCollection":
+        for f in obj.get("features", []):
+            _collect_lines(f, out)
+    elif t == "Feature" and obj.get("geometry"):
+        _collect_lines(obj["geometry"], out)
+    elif t == "GeometryCollection":
+        for g in obj.get("geometries", []):
+            _collect_lines(g, out)
+    elif t == "LineString":
+        out.append(np.asarray(obj["coordinates"], dtype=float)[:, :2])
+    elif t == "MultiLineString":
+        out.extend(np.asarray(c, dtype=float)[:, :2] for c in obj["coordinates"])
+
+
+def points_along_lines(
+    lines: Any,
+    n: Optional[int] = None,
+    spacing: Optional[float] = None,
+    seed: SeedLike = None,
+) -> Sample:
+    """Systematic points along lines (roads, rivers, hedgerows, transects).
+
+    The lines are chained end to end; points are placed every ``spacing``
+    length units from a uniformly random start in ``[0, spacing)``, so each
+    location along the network has the same inclusion density ``1/spacing``
+    per unit length. Give ``n`` to derive ``spacing = total_length / n``.
+    """
+    rng = as_rng(seed)
+    segs = _as_lines(lines)
+    lens = [np.hypot(*np.diff(s, axis=0).T) for s in segs]
+    totals = np.array([float(L.sum()) for L in lens])
+    total = float(totals.sum())
+    if total <= 0:
+        raise ValueError("lines have zero length")
+    if spacing is None:
+        if n is None or n < 1:
+            raise ValueError("give n >= 1 or spacing")
+        spacing = total / n
+    s = float(spacing)
+    pos = rng.uniform(0, s) + s * np.arange(int(math.floor(total / s)) + 1)
+    pos = pos[pos < total]
+    starts = np.concatenate([[0.0], np.cumsum(totals)[:-1]])
+    line_id = np.searchsorted(np.cumsum(totals), pos, side="right")
+    line_id = np.minimum(line_id, len(segs) - 1)
+    xy = np.empty((len(pos), 2))
+    chain = pos - starts[line_id]
+    for k in range(len(segs)):
+        sel = line_id == k
+        if not sel.any():
+            continue
+        cum = np.concatenate([[0.0], np.cumsum(lens[k])])
+        xy[sel, 0] = np.interp(chain[sel], cum, segs[k][:, 0])
+        xy[sel, 1] = np.interp(chain[sel], cum, segs[k][:, 1])
+    m = len(pos)
+    return Sample(
+        xy=xy,
+        pi=np.full(m, 1.0 / s),
+        weight=np.full(m, s),
+        design="line-systematic",
+        attrs={"line": line_id, "chainage": chain},
+        meta={"spacing": s, "total_length": total},
+        lines=segs,
+    )
+
+
+def random_transects(
+    domain: Any,
+    n: int,
+    length: float,
+    bearing: Optional[float] = None,
+    spacing: Optional[float] = None,
+    seed: SeedLike = None,
+) -> Sample:
+    """Random straight transects centred on uniform random points.
+
+    Each transect has the given ``length`` and a uniformly random bearing (or
+    a fixed ``bearing`` in degrees clockwise from north). With ``spacing`` the
+    sample contains observation points every ``spacing`` along each transect
+    (points outside the domain are dropped); otherwise it contains the centres.
+    The transect lines are kept in ``Sample.lines``.
+    """
+    dom = as_domain(domain)
+    rng = as_rng(seed)
+    if n < 1 or length <= 0:
+        raise ValueError("n must be >= 1 and length > 0")
+    centres = dom.random_points(n, rng)
+    if bearing is None:
+        az = rng.uniform(0.0, 180.0, n)
+    else:
+        az = np.full(n, float(bearing))
+    t = np.radians(az)
+    dx, dy = np.sin(t) * length / 2, np.cos(t) * length / 2
+    lines = [
+        np.array([[c[0] - a, c[1] - b], [c[0] + a, c[1] + b]]) for c, a, b in zip(centres, dx, dy)
+    ]
+    if spacing is None:
+        xy = centres
+        tid = np.arange(n)
+        chain = np.full(n, length / 2)
+    else:
+        offs = np.arange(0.0, length + 1e-9, float(spacing))
+        frac = offs / length
+        xy_l, tid_l, ch_l = [], [], []
+        for k, ln in enumerate(lines):
+            p = ln[0] + frac[:, None] * (ln[1] - ln[0])
+            xy_l.append(p)
+            tid_l.append(np.full(len(p), k))
+            ch_l.append(offs)
+        xy = np.vstack(xy_l)
+        tid = np.concatenate(tid_l)
+        chain = np.concatenate(ch_l)
+        keep = dom.contains(xy)
+        xy, tid, chain = xy[keep], tid[keep], chain[keep]
+    density = n / dom.area
+    m = len(xy)
+    return Sample(
+        xy=xy,
+        pi=np.full(m, density),
+        weight=np.full(m, 1.0 / density),
+        design="random-transects",
+        attrs={"transect": tid, "chainage": chain, "bearing": az[tid]},
+        meta={"n": n, "length": length, "area": dom.area},
+        lines=lines,
+    )