Stratified Random Sampling for Forest Plots: Production-Grade Python Implementation

import geopandas as gpd
import numpy as np
from shapely.validation import make_valid

def prepare_strata(strata_gdf: gpd.GeoDataFrame, target_epsg: int = 6933) -> gpd.GeoDataFrame:
    """
    Enforce equal-area projection, clean topology, and validate strata boundaries.
    """
    if strata_gdf.crs is None:
        raise ValueError("Input GeoDataFrame must have an assigned CRS.")
    
    # Explicit projection to equal-area system
    strata_gdf = strata_gdf.to_crs(epsg=target_epsg)
    
    # Topology validation and zero-buffer cleanup for self-intersections
    strata_gdf["geometry"] = strata_gdf.geometry.apply(lambda geom: make_valid(geom))
    strata_gdf["geometry"] = strata_gdf.geometry.buffer(0)
    
    # Remove degenerate geometries (zero area after clipping/cleaning)
    strata_gdf = strata_gdf[strata_gdf.geometry.area > 0.0].copy()
    strata_gdf["strata_area_m2"] = strata_gdf.geometry.area
    
    if strata_gdf.empty:
        raise RuntimeError("All strata geometries collapsed to zero area after projection/cleaning.")
        
    return strata_gdf

def allocate_plots_neymann(
    total_n: int, 
    areas: np.ndarray, 
    std_devs: np.ndarray
) -> np.ndarray:
    """
    Compute Neyman allocation with floor constraints and remainder distribution.
    """
    if total_n <= 0:
        raise ValueError("Total plot count must be positive.")
    if np.any(areas <= 0):
        raise ValueError("All strata areas must be strictly positive.")
        
    # Weighted allocation numerator
    weights = areas * std_devs
    total_weight = np.sum(weights)
    
    if total_weight == 0:
        raise ValueError("Sum of weighted strata is zero. Check standard deviation inputs.")
        
    # Raw allocation
    raw_alloc = (total_n * weights) / total_weight
    
    # Floor constraint to prevent zero plots in critical strata
    floor_alloc = np.floor(raw_alloc).astype(int)
    floor_alloc = np.maximum(floor_alloc, 1)
    
    # Remainder distribution
    allocated = np.sum(floor_alloc)
    remainder = total_n - allocated
    
    if remainder > 0:
        fractional_parts = raw_alloc - np.floor(raw_alloc)
        # Assign leftover plots to strata with highest fractional remainders
        top_indices = np.argsort(fractional_parts)[::-1][:remainder]
        floor_alloc[top_indices] += 1
        
    return floor_alloc

from shapely.geometry import Point

def generate_stratified_points(
    strata_gdf: gpd.GeoDataFrame,
    allocations: np.ndarray,
    plot_radius: float,
    seed: int = 42,
) -> gpd.GeoDataFrame:
    """
    Generate random plot centers within strata, enforcing edge buffers.
    `allocations[i]` is interpreted positionally and must align with the
    row order of `strata_gdf`.
    """
    if len(allocations) != len(strata_gdf):
        raise ValueError("allocations length must match the number of strata rows.")

    rng = np.random.default_rng(seed)
    points_data = []

    # Iterate positionally so allocations[i] always matches the i-th stratum,
    # regardless of the GeoDataFrame's index labels.
    for i, (label, row) in enumerate(strata_gdf.iterrows()):
        n_plots = int(allocations[i])
        if n_plots <= 0:
            continue
        geom = row.geometry

        # Create interior buffer to prevent edge overlap
        valid_geom = geom.buffer(-plot_radius)
        if valid_geom.is_empty or valid_geom.area == 0:
            raise ValueError(f"Stratum {label} too small for plot radius {plot_radius}m.")

        # Rejection sampling for uniform distribution within polygon
        minx, miny, maxx, maxy = valid_geom.bounds
        generated = 0
        attempts  = 0
        max_attempts = n_plots * 500

        while generated < n_plots and attempts < max_attempts:
            x = rng.uniform(minx, maxx)
            y = rng.uniform(miny, maxy)
            pt = Point(x, y)
            attempts += 1
            if valid_geom.contains(pt):
                points_data.append({
                    "strata_id": label,
                    "plot_radius_m": plot_radius,
                    "geometry": pt,
                })
                generated += 1

        if generated < n_plots:
            raise RuntimeError(
                f"Failed to generate {n_plots} points in stratum {label} "
                f"after {max_attempts} attempts."
            )

    return gpd.GeoDataFrame(points_data, crs=strata_gdf.crs)

def validate_sampling_pipeline(
    strata_gdf: gpd.GeoDataFrame,
    plots_gdf: gpd.GeoDataFrame,
    expected_total: int
) -> dict:
    """
    Run automated spatial and statistical validation checks.
    """
    report = {"status": "PASS", "errors": []}
    
    # Check CRS alignment
    if strata_gdf.crs != plots_gdf.crs:
        report["status"] = "FAIL"
        report["errors"].append("CRS mismatch between strata and plots.")
        
    # Check plot containment
    contained = plots_gdf.geometry.apply(lambda p: strata_gdf.geometry.contains(p).any())
    if not contained.all():
        report["status"] = "FAIL"
        report["errors"].append(f"{(~contained).sum()} plots fall outside valid strata boundaries.")
        
    # Check allocation sum
    if len(plots_gdf) != expected_total:
        report["status"] = "FAIL"
        report["errors"].append(f"Plot count {len(plots_gdf)} != expected {expected_total}.")
        
    # Validate geometry topology
    invalid_mask = ~plots_gdf.geometry.is_valid
    if invalid_mask.any():
        report["status"] = "FAIL"
        report["errors"].append(f"{invalid_mask.sum()} invalid plot geometries detected.")
        
    return report