WILLIAMSON TESTS

#%%
from pathlib import Path
import sys
import numpy as np

SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent / "Scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

import williamson_cube as wc

# ============================================================
# Williamson TC2: Global steady-state nonlinear zonal geostrophic flow
# Clean MITgcm initializer for one-layer spherical shallow water
# ============================================================

# Grid
NX = 1440
NY = 720
DLON_RAD = 2.0 * np.pi / NX
DLAT_RAD = np.pi / NY

# Physical constants
R_EARTH = 6.371e6
OMEGA = 7.292e-5
G = 9.81

# Standard Williamson TC2 constants
DAY = 86400.0
U0 = 2.0 * np.pi * R_EARTH / (12.0 * DAY)   # one revolution in 12 days
GH0 = 2.94e4                                # m^2/s^2
H0 = GH0 / G                                # mean layer depth in meters

# Rotation angle. The job scripts patch ALPHA_RAD in the run-local copy
# before this generator writes the initial fields.
ALPHA_RAD = 0.0

# Output files
BATHY_FILE = "bathymetry_flat4000.bin"
ETA_FILE = "eta_init.bin"
U_FILE = "u_init.bin"
V_FILE = "v_init.bin"
FCORI_C_FILE = "fCoriC.bin"
FCORI_G_FILE = "fCoriG.bin"
FCORI_COS_FILE = "fCorCs.bin"

def write_field(base_dir: Path, name: str, field: np.ndarray) -> None:
    field.astype(">f4").tofile(base_dir / name)
    print(f"Wrote {name}: shape={field.shape}, min={field.min():.6e}, max={field.max():.6e}")

def grid_cell_centers():
    """
    C-cell centers:
      lon_c = (i+1/2) dlon
      lat_c = -pi/2 + (j+1/2) dlat
    """
    lon_c = (np.arange(NX, dtype=np.float64) + 0.5) * DLON_RAD
    lat_c = -0.5 * np.pi + (np.arange(NY, dtype=np.float64) + 0.5) * DLAT_RAD
    return np.meshgrid(lon_c, lat_c, indexing="xy")

def grid_u_faces():
    """
    U faces:
      lon_u = i dlon
      lat_u = cell-center latitude
    """
    lon_u = np.arange(NX, dtype=np.float64) * DLON_RAD
    lat_u = -0.5 * np.pi + (np.arange(NY, dtype=np.float64) + 0.5) * DLAT_RAD
    return np.meshgrid(lon_u, lat_u, indexing="xy")

def grid_v_faces():
    """
    V faces:
      lon_v = cell-center longitude
      lat_v = -pi/2 + j dlat
    """
    lon_v = (np.arange(NX, dtype=np.float64) + 0.5) * DLON_RAD
    lat_v = -0.5 * np.pi + np.arange(NY, dtype=np.float64) * DLAT_RAD
    return np.meshgrid(lon_v, lat_v, indexing="xy")

def grid_corners():
    """
    Grid corners used by MITgcm's fCoriG read path.
    """
    lon_g = np.arange(NX, dtype=np.float64) * DLON_RAD
    lat_g = -0.5 * np.pi + np.arange(NY, dtype=np.float64) * DLAT_RAD
    return np.meshgrid(lon_g, lat_g, indexing="xy")

def rotated_mu(lon_rad: np.ndarray, lat_rad: np.ndarray, alpha_rad: float) -> np.ndarray:
    return (
        -np.cos(lon_rad) * np.cos(lat_rad) * np.sin(alpha_rad)
        + np.sin(lat_rad) * np.cos(alpha_rad)
    )

def write_flat_bathymetry(base_dir: Path) -> None:
    bathy = np.full((NY, NX), -H0, dtype=">f4")
    bathy.tofile(base_dir / BATHY_FILE)
    print(f"Wrote {BATHY_FILE} with constant depth {-H0:.6f} m")

def write_rotated_coriolis(base_dir: Path, alpha_rad: float = ALPHA_RAD) -> None:
    lon_c, lat_c = grid_cell_centers()
    lon_g, lat_g = grid_corners()

    mu_c = rotated_mu(lon_c, lat_c, alpha_rad)
    mu_g = rotated_mu(lon_g, lat_g, alpha_rad)

    write_field(base_dir, FCORI_C_FILE, 2.0 * OMEGA * mu_c)
    write_field(base_dir, FCORI_G_FILE, 2.0 * OMEGA * mu_g)
    write_field(
        base_dir,
        FCORI_COS_FILE,
        2.0 * OMEGA * np.sqrt(np.maximum(0.0, 1.0 - mu_c * mu_c)),
    )

def make_tc2_eta(alpha_rad: float = ALPHA_RAD):
    """
    Cell-center free-surface anomaly eta = h - H0,
    with h from the Williamson TC2 analytic height field.
    """
    lon_c, lat_c = grid_cell_centers()

    mu = (
        -np.cos(lon_c) * np.cos(lat_c) * np.sin(alpha_rad)
        + np.sin(lat_c) * np.cos(alpha_rad)
    )

    gh = GH0 - (R_EARTH * OMEGA * U0 + 0.5 * U0**2) * mu**2
    h = gh / G
    eta = h - H0

    return eta

def make_tc2_u(alpha_rad: float = ALPHA_RAD):
    """
    Zonal-face velocity from the paper formula.
    """
    lon_u, lat_u = grid_u_faces()

    u = U0 * (
        np.cos(lat_u) * np.cos(alpha_rad)
        + np.cos(lon_u) * np.sin(lat_u) * np.sin(alpha_rad)
    )

    return u

def make_tc2_v(alpha_rad: float = ALPHA_RAD):
    """
    Meridional-face velocity from the paper formula.
    """
    lon_v, lat_v = grid_v_faces()

    v = -U0 * np.sin(lon_v) * np.sin(alpha_rad)

    # Pole rows are coordinate-singular V-faces on the lat-lon grid.
    # Set them to zero for a clean MITgcm initialization.
    v[0, :] = 0.0
    v[-1, :] = 0.0

    return v

def main():
    base_dir = Path(__file__).resolve().parent
    if not (base_dir / ".use_latlon_grid").exists():
        wc.generate_tc2(base_dir, alpha_rad=ALPHA_RAD)
        print("")
        print("Done.")
        print(f"ALPHA_RAD = {ALPHA_RAD:.16g}")
        print(f"ALPHA_DEG = {np.rad2deg(ALPHA_RAD):.16g}")
        print(f"U0        = {U0:.12f} m/s")
        print(f"H0        = {H0:.12f} m")
        print(f"GH0       = {GH0:.12f} m^2/s^2")
        return

    write_flat_bathymetry(base_dir)
    write_rotated_coriolis(base_dir, alpha_rad=ALPHA_RAD)
    eta = make_tc2_eta(alpha_rad=ALPHA_RAD)
    u = make_tc2_u(alpha_rad=ALPHA_RAD)
    v = make_tc2_v(alpha_rad=ALPHA_RAD)

    write_field(base_dir, ETA_FILE, eta)
    write_field(base_dir, U_FILE, u)
    write_field(base_dir, V_FILE, v)

    print("")
    print("Done.")
    print(f"ALPHA_RAD = {ALPHA_RAD:.16g}")
    print(f"ALPHA_DEG = {np.rad2deg(ALPHA_RAD):.16g}")
    print(f"U0        = {U0:.12f} m/s")
    print(f"H0        = {H0:.12f} m")
    print(f"GH0       = {GH0:.12f} m^2/s^2")

if __name__ == "__main__":
    main()

# %%

#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path
import os
import sys

SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent / "Scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

import williamson_fields as wf

NX = wf.NX
NY = wf.NY
R_EARTH = wf.R_EARTH
G = wf.G
H0 = wf.TC4_H0
U0 = float(os.environ.get("TC4_U0_VALUE", wf.TC4_U0))

BATHY_FILE = "bathymetry_flat_tc4.bin"
ETA_FILE = "eta_init.bin"
U_FILE = "u_init.bin"
V_FILE = "v_init.bin"


def make_bathymetry() -> object:
    return wf.flat_bathymetry(H0)


def make_eta() -> object:
    lon, lat = wf.grid_cell_centers()
    return wf.tc4_eta(lon, lat, time_s=0.0, u0=U0)


def make_velocity_fields() -> tuple[object, object]:
    lon_u, lat_u = wf.grid_u_faces()
    lon_v, lat_v = wf.grid_v_faces()
    u = wf.tc4_u(lon_u, lat_u, time_s=0.0, u0=U0)
    v = wf.tc4_v(lon_v, lat_v, time_s=0.0, u0=U0)
    v[0, :] = 0.0
    v[-1, :] = 0.0
    return u, v


def main() -> None:
    base_dir = Path(__file__).resolve().parent
    bathy = make_bathymetry()
    eta = make_eta()
    u, v = make_velocity_fields()

    wf.write_field(base_dir, BATHY_FILE, bathy)
    wf.write_field(base_dir, ETA_FILE, eta)
    wf.write_field(base_dir, U_FILE, u)
    wf.write_field(base_dir, V_FILE, v)
    print("")
    print("Done.")
    print("Williamson TC4 translating-low initial state generated.")
    print(f"U0        = {U0:.12f} m/s")
    print(f"H0        = {H0:.12f} m")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path
import sys

SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent / "Scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

import williamson_fields as wf

NX = wf.NX
NY = wf.NY
R_EARTH = wf.R_EARTH
G = wf.G
H0 = wf.TC6_H0
ALPHA_RAD = 0.0

BATHY_FILE = "bathymetry_flat_tc6.bin"
ETA_FILE = "eta_init.bin"
U_FILE = "u_init.bin"
V_FILE = "v_init.bin"


def make_bathymetry() -> object:
    return wf.flat_bathymetry(H0)


def make_eta(alpha_rad: float = ALPHA_RAD) -> object:
    del alpha_rad
    lon, lat = wf.grid_cell_centers()
    return wf.tc6_eta(lon, lat)


def make_velocity_fields(alpha_rad: float = ALPHA_RAD) -> tuple[object, object]:
    del alpha_rad
    lon_u, lat_u = wf.grid_u_faces()
    lon_v, lat_v = wf.grid_v_faces()
    u = wf.tc6_u(lon_u, lat_u)
    v = wf.tc6_v(lon_v, lat_v)
    v[0, :] = 0.0
    v[-1, :] = 0.0
    return u, v


def main() -> None:
    base_dir = Path(__file__).resolve().parent
    bathy = make_bathymetry()
    eta = make_eta()
    u, v = make_velocity_fields()

    wf.write_field(base_dir, BATHY_FILE, bathy)
    wf.write_field(base_dir, ETA_FILE, eta)
    wf.write_field(base_dir, U_FILE, u)
    wf.write_field(base_dir, V_FILE, v)
    print("")
    print("Done.")
    print(f"H0          = {H0:.12f} m")
    print(f"wave number = {wf.TC6_WAVE_NUMBER}")
    print(f"K           = {wf.TC6_K:.12e} s^-1")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path
import sys

import numpy as np

SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent / "Scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

import williamson_fields as wf
import williamson_cube as wc

NX = wf.NX
NY = wf.NY
R_EARTH = wf.R_EARTH
G = wf.G
H0 = wf.TC7_REFERENCE_DEPTH
ALPHA_RAD = 0.0

BATHY_FILE = "bathymetry_tc7.bin"
ETA_FILE = "eta_init.bin"
U_FILE = "u_init.bin"
V_FILE = "v_init.bin"
POLAR_TAPER_START_DEG = 85.0


def load_fields() -> tuple[object, object, object, object]:
    return wf.load_tc7_fields(Path(__file__).resolve().parent)


def polar_velocity_taper() -> object:
    lat = -0.5 * np.pi + (np.arange(NY, dtype=np.float64) + 0.5) * (np.pi / NY)
    abs_lat_deg = np.abs(np.rad2deg(lat))
    taper = np.ones(NY, dtype=np.float64)
    cap = abs_lat_deg > POLAR_TAPER_START_DEG
    x = (abs_lat_deg[cap] - POLAR_TAPER_START_DEG) / (90.0 - POLAR_TAPER_START_DEG)
    taper[cap] = np.cos(0.5 * np.pi * np.clip(x, 0.0, 1.0)) ** 2
    return taper[:, None]


def apply_polar_velocity_taper(u: object, v: object) -> tuple[object, object]:
    taper = polar_velocity_taper()
    u_tapered = np.asarray(u, dtype=np.float64) * taper
    v_tapered = np.asarray(v, dtype=np.float64) * taper
    v_tapered[0, :] = 0.0
    v_tapered[-1, :] = 0.0
    return u_tapered, v_tapered


def make_bathymetry() -> object:
    _eta, _u, _v, bathy = load_fields()
    return bathy


def make_eta(alpha_rad: float = ALPHA_RAD) -> object:
    del alpha_rad
    eta, u, v, _bathy = load_fields()
    eta, _u, _v = wf.preprocess_tc7_analysis_fields(eta, u, v)
    return eta


def make_velocity_fields(alpha_rad: float = ALPHA_RAD) -> tuple[object, object]:
    del alpha_rad
    eta, u, v, _bathy = load_fields()
    _eta, u, v = wf.preprocess_tc7_analysis_fields(eta, u, v)
    return apply_polar_velocity_taper(u, v)


def main() -> None:
    base_dir = Path(__file__).resolve().parent
    raw_file = wf.resolve_tc7_raw_file(base_dir)
    if not (base_dir / ".use_latlon_grid").exists():
        wc.generate_tc7(base_dir)
        print("")
        print("Done.")
        print(f"TC7 input source: {raw_file}")
        print(
            "Applied TC7 large-scale filter: "
            f"zonal wavenumber <= {wf.TC7_LONGITUDE_WAVENUMBER_CUTOFF}, "
            f"{wf.TC7_SHAPIRO_SMOOTHING_PASSES} Shapiro passes"
        )
        return

    eta, u, v, bathy = load_fields()
    eta, u, v = wf.preprocess_tc7_analysis_fields(eta, u, v)
    u, v = apply_polar_velocity_taper(u, v)
    wf.write_field(base_dir, BATHY_FILE, bathy)
    wf.write_field(base_dir, ETA_FILE, eta)
    wf.write_field(base_dir, U_FILE, u)
    wf.write_field(base_dir, V_FILE, v)
    print("")
    print("Done.")
    print(f"TC7 input source: {raw_file}")
    print(
        "Applied TC7 large-scale filter: "
        f"zonal wavenumber <= {wf.TC7_LONGITUDE_WAVENUMBER_CUTOFF}, "
        f"{wf.TC7_SHAPIRO_SMOOTHING_PASSES} Shapiro passes"
    )
    print(f"Applied polar velocity taper poleward of {POLAR_TAPER_START_DEG:g} deg")


if __name__ == "__main__":
    main()