coastal_example.py

"""
examples/coastal_example.py
===========================
Demonstrates Inverse Path Distance Weighting (IPDW) for coastal salinity
interpolation.

Scenario
--------
A rectangular bay is bounded on the north and south by land.  A contiguous
land barrier cuts across the centre of the bay from west to east, dividing
it into two sub-basins.  Salinity measurements are scattered across both
sub-basins.

With standard Euclidean-distance IDW, salinity values "bleed through" the
land barrier.  IPDW uses path distances (as the fish swims) and therefore
respects the barrier, giving physically correct isolines.

Run
---
    python examples/coastal_example.py
"""

import sys
import pathlib

# Allow running the script directly from the repo root
sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))

import numpy as np
import geopandas as gpd
import matplotlib.pyplot as plt
from shapely.geometry import Point, box
from rasterio.transform import from_bounds

from ipdw import costraster_gen, pathdist_gen, ipdw_interp, ipdw


# ---------------------------------------------------------------------------
# 1.  Synthetic study area
# ---------------------------------------------------------------------------
#
# Grid: 60 × 80 cells, 1 km resolution → 60 km × 80 km bay
# Land barrier: rows 28–32 (y ≈ 28–32 km from south), columns 0–55
#
RESOLUTION = 1_000          # metres per cell
WIDTH_KM   = 80             # east–west extent
HEIGHT_KM  = 60             # north–south extent

# CRS: UTM zone 32N (metres) – any projected CRS works
CRS = "EPSG:32632"

# ---------------------------------------------------------------------------
# 2.  Cost raster (build manually here; normally use costraster_gen)
# ---------------------------------------------------------------------------
rows = int(HEIGHT_KM * 1_000 / RESOLUTION)   # 60
cols = int(WIDTH_KM  * 1_000 / RESOLUTION)   # 80

cost_array = np.ones((rows, cols), dtype=np.float64)   # all water

# North and south land borders (2 cells thick)
cost_array[:2,  :] = 10_000.0
cost_array[-2:, :] = 10_000.0

# Contiguous east–west barrier in the middle (rows 27–33, cols 0–55)
cost_array[27:34, :56] = 10_000.0

transform = from_bounds(0, 0, WIDTH_KM * 1_000, HEIGHT_KM * 1_000, cols, rows)

# ---------------------------------------------------------------------------
# 3.  Synthetic salinity observations
# ---------------------------------------------------------------------------
# South sub-basin: high salinity (marine water comes in from the east gap)
# North sub-basin: low salinity (freshwater input from the north-west)
rng = np.random.default_rng(42)

def _add_noise(vals, sigma=0.5):
    return vals + rng.normal(0, sigma, len(vals))

# South basin observation points  (y < 27 km → row > 33)
south_x = rng.uniform(2_000, 78_000, 20)
south_y = rng.uniform(2_000, 25_000, 20)
south_sal = _add_noise(30 + (south_x / 80_000) * 4)   # 30–34 ppt gradient

# North basin observation points  (y > 33 km → row < 27)
north_x = rng.uniform(2_000, 78_000, 20)
north_y = rng.uniform(35_000, 58_000, 20)
north_sal = _add_noise(20 + (north_x / 80_000) * 3)   # 20–23 ppt gradient

all_x   = np.concatenate([south_x, north_x])
all_y   = np.concatenate([south_y, north_y])
all_sal = np.concatenate([south_sal, north_sal])

training_gdf = gpd.GeoDataFrame(
    {"salinity": all_sal},
    geometry=[Point(x, y) for x, y in zip(all_x, all_y)],
    crs=CRS,
)

# ---------------------------------------------------------------------------
# 4.  IPDW interpolation
# ---------------------------------------------------------------------------
RANGE = 20_000   # 20 km search radius

print("Running IPDW …")
result_ipdw = ipdw(
    training_gdf,
    cost_array,
    transform,
    range_val=RANGE,
    param_col="salinity",
    dist_power=1,
    progress=True,
)

# ---------------------------------------------------------------------------
# 5.  Standard IDW for comparison  (Euclidean distances via scipy)
# ---------------------------------------------------------------------------
from scipy.interpolate import NearestNDInterpolator
from scipy.spatial import cKDTree

def euclidean_idw(pts_gdf, param_col, grid_x, grid_y, power=2, max_dist=None):
    """Simple Euclidean IDW on a flat grid."""
    obs_xy = np.array([(g.x, g.y) for g in pts_gdf.geometry])
    obs_v  = pts_gdf[param_col].values.astype(float)
    tree   = cKDTree(obs_xy)

    query_xy = np.column_stack([grid_x.ravel(), grid_y.ravel()])
    dists, idx = tree.query(query_xy, k=len(obs_xy))

    if max_dist is not None:
        dists = np.where(dists <= max_dist, dists, np.inf)

    weights = np.where(dists == 0, 1e20, 1.0 / dists ** power)
    wsum    = weights.sum(axis=1, keepdims=True)
    wsum    = np.where(wsum == 0, np.nan, wsum)
    result  = (weights * obs_v[idx]).sum(axis=1) / wsum[:, 0]

    return result.reshape(grid_x.shape)

# Build grid cell centres
cell_x = np.array([transform.c + (c + 0.5) * transform.a for c in range(cols)])
cell_y = np.array([transform.f + (r + 0.5) * transform.e for r in range(rows)])
grid_x, grid_y = np.meshgrid(cell_x, cell_y)

print("Running Euclidean IDW …")
result_idw = euclidean_idw(training_gdf, "salinity", grid_x, grid_y,
                           power=2, max_dist=RANGE)

# Mask land cells in both results
land_mask = cost_array >= 10_000
result_ipdw_plot = np.where(land_mask, np.nan, result_ipdw)
result_idw_plot  = np.where(land_mask, np.nan, result_idw)

# ---------------------------------------------------------------------------
# 6.  Plot
# ---------------------------------------------------------------------------
vmin = float(np.nanmin([result_ipdw_plot, result_idw_plot]))
vmax = float(np.nanmax([result_ipdw_plot, result_idw_plot]))

extent_km = [0, WIDTH_KM, 0, HEIGHT_KM]

fig, axes = plt.subplots(1, 3, figsize=(15, 5), constrained_layout=True)

for ax, data, title in zip(
    axes,
    [result_ipdw_plot, result_idw_plot, result_idw_plot - result_ipdw_plot],
    ["IPDW", "Euclidean IDW", "IDW − IPDW"],
):
    if title == "IDW − IPDW":
        vmin_d = float(np.nanmin(result_idw_plot - result_ipdw_plot))
        vmax_d = float(np.nanmax(result_idw_plot - result_ipdw_plot))
        lim = max(abs(vmin_d), abs(vmax_d))
        im = ax.imshow(data, origin="upper", extent=extent_km,
                       cmap="RdBu_r", vmin=-lim, vmax=lim, aspect="auto")
    else:
        im = ax.imshow(data, origin="upper", extent=extent_km,
                       cmap="viridis", vmin=vmin, vmax=vmax, aspect="auto")

    # Overlay training points
    ax.scatter(
        [g.x / 1_000 for g in training_gdf.geometry],
        [g.y / 1_000 for g in training_gdf.geometry],
        c="white", s=20, edgecolors="k", linewidths=0.5, zorder=5,
    )
    ax.set_title(title)
    ax.set_xlabel("Easting (km)")
    ax.set_ylabel("Northing (km)")
    plt.colorbar(im, ax=ax, shrink=0.8, label="Salinity (ppt)")

fig.suptitle(
    "Inverse Path Distance Weighting vs Euclidean IDW\n"
    "(land barrier runs west–east across the centre)",
    fontsize=13,
)

out_path = pathlib.Path(__file__).parent / "coastal_example_output.png"
fig.savefig(out_path, dpi=150)
print(f"Figure saved to {out_path}")
plt.show()