Converting WGS84 to Local Grid with Python

#!/usr/bin/env python3
"""
WGS84 to Local Grid Converter for UAV/Survey Pipelines
Handles chunked streaming I/O, incremental Parquet output, and deterministic validation.
"""

import argparse
import logging
import sys
from pathlib import Path
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import pyproj
from pyproj.exceptions import ProjError

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s"
)

def setup_transformer(src_epsg: int, dst_epsg: int) -> pyproj.Transformer:
    """Initialize deterministic transformer with explicit axis ordering."""
    src_crs = pyproj.CRS.from_epsg(src_epsg)
    dst_crs = pyproj.CRS.from_epsg(dst_epsg)
    return pyproj.Transformer.from_crs(src_crs, dst_crs, always_xy=True)

def transform_and_validate(chunk: pd.DataFrame, transformer: pyproj.Transformer, tolerance: float) -> pd.DataFrame:
    """Apply transformation, check bounds/NaN, and flag outliers."""
    required = {"lon", "lat", "alt"}
    missing = required - set(chunk.columns)
    if missing:
        raise ValueError(f"Missing required columns: {missing}")

    # Extract to contiguous float64 arrays for vectorized transformation
    lons = chunk["lon"].to_numpy(dtype=np.float64)
    lats = chunk["lat"].to_numpy(dtype=np.float64)
    alts = chunk["alt"].to_numpy(dtype=np.float64)

    # Execute transformation
    x, y, z = transformer.transform(lons, lats, alts)

    # Use .to_numpy() for the id column to avoid index-alignment surprises
    original_id = chunk["id"].to_numpy() if "id" in chunk.columns else np.arange(len(chunk))
    out = pd.DataFrame({
        "easting": x,
        "northing": y,
        "elevation": z,
        "original_id": original_id,
    })

    # Validation: NaN propagation
    nan_mask = out.isna().any(axis=1)
    out.loc[nan_mask, "status"] = "FAIL_NAN"
    out.loc[~nan_mask, "status"] = "PASS"

    # Validation: Bounds & Tolerance (assumes positive UTM/northing/easting)
    bounds_mask = (out["easting"] < 0) | (out["northing"] < 0)
    out.loc[bounds_mask, "status"] = "FAIL_BOUNDS"

    # Optional: reference drift check if 'ref_easting'/'ref_northing' exist.
    # Persist the reference columns so downstream RMSE checks can read them back.
    if {"ref_easting", "ref_northing"}.issubset(chunk.columns):
        out["ref_easting"] = chunk["ref_easting"].to_numpy()
        out["ref_northing"] = chunk["ref_northing"].to_numpy()
        dx = np.abs(out["easting"] - out["ref_easting"])
        dy = np.abs(out["northing"] - out["ref_northing"])
        drift = np.sqrt(dx**2 + dy**2)
        out.loc[(drift > tolerance) & (out["status"] == "PASS"), "status"] = "FAIL_DRIFT"

    return out

def main():
    parser = argparse.ArgumentParser(description="WGS84 to Local Grid Converter")
    parser.add_argument("--input-csv", type=Path, required=True, help="Source CSV/Parquet path")
    parser.add_argument("--target-epsg", type=int, required=True, help="Destination EPSG code")
    parser.add_argument("--chunk-size", type=int, default=50000, help="Rows per memory chunk")
    parser.add_argument("--tolerance-m", type=float, default=0.02, help="Max allowable drift (m)")
    parser.add_argument("--output-dir", type=Path, default=Path("./output"), help="Export directory")
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)
    out_path = args.output_dir / "transformed_grid.parquet"

    try:
        transformer = setup_transformer(4326, args.target_epsg)
    except ProjError as e:
        logging.error(f"CRS initialization failed: {e}. Verify PROJ_DATA contains required grid files.")
        sys.exit(1)

    logging.info(f"Initializing pipeline: {args.input_csv} -> EPSG:{args.target_epsg}")

    # Chunked reader: CSV via pandas' chunksize; Parquet via pyarrow row-group
    # batches (pd.read_parquet has no chunksize argument).
    suffix = args.input_csv.suffix.lower()
    if suffix == ".parquet":
        parquet_file = pq.ParquetFile(args.input_csv)
        chunk_iter = (batch.to_pandas() for batch in parquet_file.iter_batches(batch_size=args.chunk_size))
    else:
        chunk_iter = pd.read_csv(args.input_csv, chunksize=args.chunk_size)

    # Stream each validated chunk straight to a single Parquet file. A
    # ParquetWriter appends row groups; re-calling to_parquet() would instead
    # overwrite the file on every chunk and lose all but the last.
    writer = None
    try:
        for i, chunk in enumerate(chunk_iter):
            try:
                transformed = transform_and_validate(chunk, transformer, args.tolerance_m)
                table = pa.Table.from_pandas(transformed, preserve_index=False)
                if writer is None:
                    writer = pq.ParquetWriter(out_path, table.schema)
                writer.write_table(table)
                logging.info(f"Chunk {i+1} processed. Rows: {len(transformed)} | Status: {transformed['status'].value_counts().to_dict()}")
            except Exception as e:
                logging.error(f"Chunk {i+1} failed: {e}")
                continue
    finally:
        if writer is not None:
            writer.close()

    logging.info(f"Pipeline complete. Output written to {out_path}")

if __name__ == "__main__":
    main()

# Post-processing validation snippet
df = pd.read_parquet("output/transformed_grid.parquet")
valid = df[df["status"] == "PASS"]
rmse = np.sqrt(((valid["easting"] - valid["ref_easting"])**2 + 
                (valid["northing"] - valid["ref_northing"])**2).mean())
logging.info(f"Post-transform RMSE: {rmse:.4f} m")