Point Density Metrics in Python LiDAR Workflows

"""
point_density.py — Grid-based point density metrics for LAS/LAZ files.
Requires: laspy>=2.4, numpy>=1.24, rasterio>=1.3, pyproj>=3.6
"""

from __future__ import annotations

import logging
from pathlib import Path

import laspy
import numpy as np
import pyproj
import rasterio
from rasterio.transform import from_origin

log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

# ── Classification masks ─────────────────────────────────────────────────────
# ASPRS LAS 1.4 classes to exclude before density computation.
# 7 = Low point (noise), 12 = Overlap, 18 = High noise, 22 = Withheld
EXCLUDE_CLASSES: frozenset[int] = frozenset({7, 12, 18, 22})


def validate_las_crs(file_path: Path) -> pyproj.CRS:
    """Parse and validate the CRS embedded in LAS VLRs.

    Raises:
        ValueError: If no VLRs are found, CRS cannot be parsed, or the
                    projection is not metric (geographic CRS rejected).
    """
    with laspy.open(file_path) as f:
        if not f.header.vlrs:
            raise ValueError(
                f"{file_path.name}: No VLRs found — CRS metadata is missing. "
                "Re-export the file with a WKT or GeoTIFF CRS VLR."
            )

        crs = f.header.parse_crs()
        if crs is None:
            raise ValueError(
                f"{file_path.name}: Could not parse a CRS from VLRs. "
                "Inspect the VLR records with laspy.open().header.vlrs."
            )

        if not crs.is_projected:
            raise ValueError(
                f"{file_path.name}: Geographic CRS '{crs.name}' detected. "
                "Density calculations require a planar metric projection. "
                "Reproject to e.g. EPSG:32632 (UTM 32N) first."
            )

        log.info("CRS validated: %s (EPSG:%s)", crs.name, crs.to_epsg())
        return crs


def compute_grid_density(
    file_path: Path,
    cell_size: float = 1.0,
    chunk_size: int = 5_000_000,
    exclude_classes: frozenset[int] = EXCLUDE_CLASSES,
) -> tuple[np.ndarray, tuple[float, float, float, float]]:
    """Compute per-cell point density via two-pass chunked I/O.

    Pass 1: determine bounding box without loading Z or attributes.
    Pass 2: accumulate filtered XY counts into a uint32 grid.

    Args:
        file_path:       Path to a LAS or LAZ file.
        cell_size:       Grid cell edge length in the file's projection units
                         (metres for metric CRS). Default 1.0 m.
        chunk_size:      Points read per I/O chunk. 5 M ≈ 120 MB RAM peak.
        exclude_classes: Set of ASPRS classification codes to ignore.

    Returns:
        density_grid: uint32 NumPy array, shape (rows, cols), raw counts.
        bounds:       (x_min, y_min, x_max, y_max) in projection units.
    """
    log.info("Pass 1/2: scanning bounding box — %s", file_path.name)
    x_min, y_min = np.inf, np.inf
    x_max, y_max = -np.inf, -np.inf

    with laspy.open(file_path) as f:
        for chunk in f.chunk_iterator(chunk_size):
            cls = np.array(chunk.classification, dtype=np.uint8)
            keep = ~np.isin(cls, list(exclude_classes))
            if not keep.any():
                continue
            x, y = np.asarray(chunk.x)[keep], np.asarray(chunk.y)[keep]
            x_min, y_min = min(x_min, x.min()), min(y_min, y.min())
            x_max, y_max = max(x_max, x.max()), max(y_max, y.max())

    if not np.isfinite(x_min):
        raise RuntimeError(
            "No valid points remain after classification filtering. "
            "Check EXCLUDE_CLASSES against the file's actual class codes."
        )

    cols = max(1, int(np.ceil((x_max - x_min) / cell_size)))
    rows = max(1, int(np.ceil((y_max - y_min) / cell_size)))
    log.info("Grid: %d rows × %d cols at %.2f m cell size", rows, cols, cell_size)

    density_grid = np.zeros((rows, cols), dtype=np.uint32)

    log.info("Pass 2/2: accumulating density counts")
    with laspy.open(file_path) as f:
        for chunk in f.chunk_iterator(chunk_size):
            cls = np.array(chunk.classification, dtype=np.uint8)
            keep = ~np.isin(cls, list(exclude_classes))
            if not keep.any():
                continue
            x, y = np.asarray(chunk.x)[keep], np.asarray(chunk.y)[keep]

            col_idx = np.clip(
                ((x - x_min) / cell_size).astype(np.intp), 0, cols - 1
            )
            row_idx = np.clip(
                ((y - y_min) / cell_size).astype(np.intp), 0, rows - 1
            )
            np.add.at(density_grid, (row_idx, col_idx), 1)

    return density_grid, (x_min, y_min, x_max, y_max)


def normalize_density(
    density_grid: np.ndarray, cell_size: float
) -> np.ndarray:
    """Convert raw count grid to points per square metre (pts/m²)."""
    return density_grid.astype(np.float32) / (cell_size ** 2)


def generate_density_stats(
    density_pts_m2: np.ndarray, threshold: float = 8.0
) -> dict:
    """Compute QA statistics over non-empty cells.

    Args:
        density_pts_m2: Normalized density grid (pts/m²).
        threshold:       Minimum acceptable density for the project spec.

    Returns:
        Dictionary of descriptive statistics and a low-density flag.
    """
    valid = density_pts_m2[density_pts_m2 > 0]
    total_cells = density_pts_m2.size

    if len(valid) == 0:
        log.warning("All cells are empty — check CRS and classification filter.")
        return {
            "mean_pts_m2": 0.0, "median_pts_m2": 0.0,
            "p5_pts_m2": 0.0, "p95_pts_m2": 0.0,
            "std_dev": 0.0, "min_pts_m2": 0.0, "max_pts_m2": 0.0,
            "coverage_pct": 0.0, "low_density_cells_pct": 100.0,
            "passes_threshold": False,
        }

    low_count = int((density_pts_m2 > 0) & (density_pts_m2 < threshold)).sum()

    stats = {
        "mean_pts_m2":          float(np.mean(valid)),
        "median_pts_m2":        float(np.median(valid)),
        "p5_pts_m2":            float(np.percentile(valid, 5)),
        "p95_pts_m2":           float(np.percentile(valid, 95)),
        "std_dev":              float(np.std(valid)),
        "min_pts_m2":           float(valid.min()),
        "max_pts_m2":           float(valid.max()),
        "coverage_pct":         float(len(valid) / total_cells * 100),
        "low_density_cells_pct": float(low_count / total_cells * 100),
        "passes_threshold":      float(np.mean(valid)) >= threshold,
    }
    log.info(
        "Density — mean: %.2f, p5: %.2f, p95: %.2f pts/m² | "
        "coverage: %.1f%% | passes %.1f threshold: %s",
        stats["mean_pts_m2"], stats["p5_pts_m2"], stats["p95_pts_m2"],
        stats["coverage_pct"], threshold, stats["passes_threshold"],
    )
    return stats


def export_density_geotiff(
    density_pts_m2: np.ndarray,
    bounds: tuple[float, float, float, float],
    cell_size: float,
    crs: pyproj.CRS,
    output_path: Path,
) -> None:
    """Write the density raster to a cloud-optimised GeoTIFF.

    Uses a north-up transform derived from the grid's top-left corner.
    NoData value is -9999.0 (empty cells).
    """
    x_min, y_min, _x_max, y_max = bounds
    transform = from_origin(x_min, y_max, cell_size, cell_size)

    # Flip vertically: row 0 in the array is y_min; GeoTIFF expects y_max at top
    flipped = np.flipud(density_pts_m2.copy())
    flipped[flipped == 0] = -9999.0

    with rasterio.open(
        output_path,
        mode="w",
        driver="GTiff",
        height=density_pts_m2.shape[0],
        width=density_pts_m2.shape[1],
        count=1,
        dtype="float32",
        crs=crs.to_wkt(),
        transform=transform,
        nodata=-9999.0,
        compress="deflate",
        tiled=True,
        blockxsize=256,
        blockysize=256,
    ) as dst:
        dst.write(flipped, 1)

    log.info("Density raster written to %s", output_path)


# ── Orchestrator ─────────────────────────────────────────────────────────────

def run_density_pipeline(
    input_path: Path,
    output_dir: Path,
    cell_size: float = 1.0,
    threshold: float = 8.0,
    chunk_size: int = 5_000_000,
) -> dict:
    """End-to-end density pipeline: validate → filter → bin → export.

    Returns:
        QA statistics dictionary; raises on CRS or empty-cloud errors.
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    crs = validate_las_crs(input_path)

    raw_grid, bounds = compute_grid_density(
        input_path, cell_size=cell_size, chunk_size=chunk_size
    )
    density = normalize_density(raw_grid, cell_size)
    stats = generate_density_stats(density, threshold=threshold)

    geotiff_path = output_dir / (input_path.stem + "_density.tif")
    export_density_geotiff(density, bounds, cell_size, crs, geotiff_path)

    if not stats["passes_threshold"]:
        log.warning(
            "FAIL: mean density %.2f pts/m² is below project threshold %.2f pts/m²",
            stats["mean_pts_m2"], threshold,
        )

    return stats


if __name__ == "__main__":
    import json, sys
    result = run_density_pipeline(
        input_path=Path(sys.argv[1]),
        output_dir=Path(sys.argv[2]) if len(sys.argv) > 2 else Path("."),
        cell_size=float(sys.argv[3]) if len(sys.argv) > 3 else 1.0,
        threshold=float(sys.argv[4]) if len(sys.argv) > 4 else 8.0,
    )
    print(json.dumps(result, indent=2))

import rasterio
import numpy as np

def validate_density_output(
    geotiff_path,
    stats: dict,
    min_coverage_pct: float = 90.0,
    max_low_density_pct: float = 5.0,
) -> None:
    """Assert that density outputs meet project acceptance criteria."""

    # 1. Raster is readable and has a valid CRS
    with rasterio.open(geotiff_path) as src:
        assert src.crs is not None, "GeoTIFF has no CRS"
        assert src.nodata == -9999.0, f"Unexpected nodata: {src.nodata}"
        data = src.read(1)

    # 2. Coverage: at least 90% of the survey area has returns
    valid_cells = data[data != -9999.0]
    coverage = len(valid_cells) / data.size * 100
    assert coverage >= min_coverage_pct, (
        f"Coverage {coverage:.1f}% < required {min_coverage_pct}%"
    )

    # 3. No negative density values (would indicate a count overflow)
    assert valid_cells.min() >= 0.0, "Negative density values detected"

    # 4. Mean density passes the project threshold
    assert stats["passes_threshold"], (
        f"Mean density {stats['mean_pts_m2']:.2f} pts/m² below threshold"
    )

    # 5. Low-density cells are within acceptable fraction
    assert stats["low_density_cells_pct"] <= max_low_density_pct, (
        f"Low-density cells {stats['low_density_cells_pct']:.1f}% "
        f"> allowed {max_low_density_pct}%"
    )

    print("All density validation checks passed.")