venus/dbus-tides/tools/build_coastal_grid.py

#!/usr/bin/env python3
"""
Build the pre-computed coastal tidal constituent grid.

This script runs on a DEVELOPMENT MACHINE where pyTMD is installed.
It extracts tidal constituents at every 0.25-degree ocean point across
the configured regions, masks land, and writes a gzipped JSON file.

Usage:
    pip install -r requirements-dev.txt
    python build_coastal_grid.py [--resolution 0.25] [--model GOT4.10_nc]

Output: ../constituents/coastal_grid.json.gz
"""

import argparse
import gzip
import json
import os
import sys
import time

try:
    import numpy as np
    import pyTMD
except ImportError:
    print("ERROR: pyTMD is not installed. Run: pip install -r requirements-dev.txt")
    sys.exit(1)

REGIONS = [
    ('us_east',    24.0, 45.0, -82.0, -66.0),
    ('bahamas',    20.0, 28.0, -80.0, -72.0),
    ('caribbean',  10.0, 24.0, -88.0, -59.0),
    ('gulf',       24.0, 31.0, -98.0, -82.0),
    ('us_west',    32.0, 49.0, -125.0, -117.0),
]


def main():
    parser = argparse.ArgumentParser(
        description='Build coastal tidal constituent grid')
    parser.add_argument('--resolution', type=float, default=0.25,
                        help='Grid resolution in degrees (default: 0.25)')
    parser.add_argument('--model', type=str, default='GOT4.10_nc',
                        help='pyTMD database model name (default: GOT4.10_nc)')
    parser.add_argument('--output', type=str,
                        default='../constituents/coastal_grid.json.gz',
                        help='Output file path')
    args = parser.parse_args()

    res = args.resolution
    print(f"Building coastal grid at {res}-degree resolution")
    print(f"Model: {args.model}")
    print(f"Regions: {[r[0] for r in REGIONS]}")
    print()

    all_coords = set()
    for name, lat_min, lat_max, lon_min, lon_max in REGIONS:
        lat = lat_min
        while lat <= lat_max + 1e-9:
            lon = lon_min
            while lon <= lon_max + 1e-9:
                all_coords.add((round(lat, 4), round(lon, 4)))
                lon += res
            lat += res
        print(f"  {name}: {lat_min}N-{lat_max}N, "
              f"{abs(lon_max)}W-{abs(lon_min)}W")

    coords = sorted(all_coords)
    print(f"\nTotal grid points to extract: {len(coords)}")

    lats = np.array([c[0] for c in coords])
    lons_neg = np.array([c[1] for c in coords])
    lons_360 = lons_neg % 360

    print("Loading tidal model...")
    t0 = time.time()

    try:
        model = pyTMD.io.model().from_database(args.model)
    except Exception as e:
        print(f"ERROR loading model from database: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        ds = pyTMD.io.GOT.open_mfdataset(model.z.model_file)
    except Exception as e:
        print(f"ERROR opening model files: {e}", file=sys.stderr)
        sys.exit(1)

    const_names = list(ds.data_vars)
    print(f"Constituents in model: {const_names}")

    unique_lats = np.unique(lats)
    unique_lons = np.unique(lons_360)

    print(f"Interpolating on {len(unique_lats)} unique lats x "
          f"{len(unique_lons)} unique lons...")

    ds_interp = ds.interp(y=unique_lats, x=unique_lons)

    elapsed = time.time() - t0
    print(f"Interpolation complete in {elapsed:.1f}s")

    lat_to_idx = {float(v): i for i, v in enumerate(unique_lats)}
    lon_to_idx = {float(v): i for i, v in enumerate(unique_lons)}

    points = []
    n_land = 0
    for i in range(len(coords)):
        lat_i = lat_to_idx[float(lats[i])]
        lon_i = lon_to_idx[float(lons_360[i])]

        amplitudes = []
        phases = []
        all_nan = True

        for cname in const_names:
            val = complex(ds_interp[cname].values[lat_i, lon_i])
            if np.isnan(val):
                amplitudes.append(0.0)
                phases.append(0.0)
            else:
                all_nan = False
                amp_cm = np.abs(val)
                amp_m = amp_cm / 100.0
                phase_deg = (-np.angle(val, deg=True)) % 360
                amplitudes.append(round(float(amp_m), 6))
                phases.append(round(float(phase_deg), 4))

        if all_nan or max(amplitudes) < 1e-7:
            n_land += 1
            continue

        points.append({
            'lat': float(lats[i]),
            'lon': float(lons_neg[i]),
            'amp': amplitudes,
            'phase': phases,
        })

    print(f"\nOcean points: {len(points)}")
    print(f"Land points (masked): {n_land}")

    grid_data = {
        'regions': [r[0] for r in REGIONS],
        'resolution_deg': res,
        'constituents': [str(n).upper() for n in const_names],
        'points': points,
    }

    print(f"\nWriting {args.output}...")
    json_str = json.dumps(grid_data, separators=(',', ':'))

    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)

    if args.output.endswith('.gz'):
        with gzip.open(args.output, 'wt', encoding='utf-8') as f:
            f.write(json_str)
    else:
        with open(args.output, 'w') as f:
            f.write(json_str)

    size = os.path.getsize(args.output)
    print(f"Done: {size / 1024:.0f} KB compressed")
    print(f"Uncompressed JSON: {len(json_str) / 1024:.0f} KB")


if __name__ == '__main__':
    main()