From 8936168aad9d11228b69b215edaa772afa59149a Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Mon, 17 Mar 2025 10:02:45 +1100
Subject: [PATCH 01/16] Add generate_bottom_N.py

---
 .../generate_bottom_N.py                      | 234 ++++++++++++++++++
 1 file changed, 234 insertions(+)
 create mode 100644 external_tidal_generation/generate_bottom_N.py

diff --git a/external_tidal_generation/generate_bottom_N.py b/external_tidal_generation/generate_bottom_N.py
new file mode 100644
index 0000000..a932772
--- /dev/null
+++ b/external_tidal_generation/generate_bottom_N.py
@@ -0,0 +1,234 @@
+# Copyright 2025 ACCESS-NRI and contributors. See the top-level COPYRIGHT file for details.
+# SPDX-License-Identifier: Apache-2.0
+
+# =========================================================================================
+# Generate a bottom-Xm-averaged stratification from WOA climatology.
+#
+# This script creates a 3D stratification field from WOA temperature and salinity
+# data, computes the buoyancy frequency (N) on a mid-depth grid, and then averages N
+# over the bottom DEPTH_THRESHOLD meters of the water column for each (lat, lon) column.
+#
+# Usage:
+#    python3 generate_bottom_N.py --depth-threshold 500 \
+#         --temp-file /path/to/WOA_decav_t00_04.nc \
+#         --sal-file /path/to/WOA_decav_s00_04.nc \
+#         --output Nbot_freq.nc
+#
+# Contact:
+#    Minghang Li <Minghang.Li1@anu.edu.au>
+#    This script was originally developed by Luwei Yang <Luwei.Yang@anu.edu.au>
+#
+# Dependencies:
+#   - gsw
+#   - xarray
+#   - xesmf
+#   - numpy
+# =========================================================================================
+
+from pathlib import Path
+import sys
+import argparse
+import os
+import warnings
+
+warnings.filterwarnings("ignore")
+
+import xarray as xr
+import numpy as np
+import gsw
+import xesmf as xe
+
+path_root = Path(__file__).parents[1]
+sys.path.append(str(path_root))
+
+from scripts_common import get_provenance_metadata, md5sum
+
+
+def load_woa_data(temp_file: str, sal_file: str) -> (xr.DataArray, xr.DataArray):
+    """Load the WOA climatological temperature and salinity datasets."""
+    temp_ds = xr.open_dataset(temp_file, decode_times=False)
+    sal_ds = xr.open_dataset(sal_file, decode_times=False)
+
+    sea_water_temp = temp_ds.t_an.squeeze().drop_vars("time", errors="ignore")
+    sea_water_sal = sal_ds.s_an.squeeze().drop_vars("time", errors="ignore")
+    return sea_water_temp, sea_water_sal
+
+
+def compute_pressure(depth: xr.DataArray, lat: xr.DataArray, lon: xr.DataArray):
+    """
+    Compute a 3D pressure field (in dbar) from 1D depth and lat arrays.
+    gsw.p_from_z expects negative depth.
+    """
+    depth_3d = depth.expand_dims({"lat": lat, "lon": lon})
+    lat_3d = lat.expand_dims({"depth": depth, "lon": lon})
+    pressure_3d = xr.apply_ufunc(
+        lambda dep, la: gsw.p_from_z(-dep, la),
+        depth_3d,
+        lat_3d,
+        input_core_dims=[["depth", "lat", "lon"]] * 2,
+        output_core_dims=[["depth", "lat", "lon"]],
+        vectorize=True,
+    )
+    return pressure_3d
+
+
+def compute_stratification(
+    sea_water_temp: xr.DataArray,
+    sea_water_sal: xr.DataArray,
+    pressure_3d: xr.DataArray,
+    depth: xr.DataArray,
+    lat: xr.DataArray,
+) -> (xr.DataArray, xr.DataArray):
+    """
+    Compute squared buoyancy frequency and mid pressures, then derive N and mid-depth.
+    """
+    N2_3D, p_mid = xr.apply_ufunc(
+        gsw.Nsquared,
+        sea_water_sal,
+        sea_water_temp,
+        pressure_3d,
+        input_core_dims=[["depth", "lat", "lon"]] * 3,
+        output_core_dims=[["depth_mid", "lat", "lon"]] * 2,
+        vectorize=True,
+    )
+    # Compute N from N square
+    N_3D = np.sqrt(N2_3D)
+
+    # Compute mid-depth
+    lat_broad, _ = xr.broadcast(lat, p_mid)
+    depth_mid = -gsw.z_from_p(p_mid, lat_broad)  # Now depth_mid > 0
+    return N_3D, depth_mid
+
+
+def compute_bottom_average(
+    N_3D: xr.DataArray,
+    depth_mid: xr.DataArray,
+    sea_water_temp: xr.DataArray,
+    depth: xr.DataArray,
+    depth_threshold: float,
+) -> xr.DataArray:
+    """
+    Create a mask for the bottom depth_threshold meters, apply it to N_3D and average over depth_mid.
+    """
+    # Broadcast depth
+    depth_array = sea_water_temp * 0 + depth
+    max_depth = depth_array.max(dim="depth", skipna=True)
+    bottom_threshold = max_depth - depth_threshold
+
+    # Create a mask
+    mask_bottom = xr.where(
+        (depth_mid >= bottom_threshold) & (depth_mid < max_depth),
+        1,
+        np.nan,
+    )
+    N_3D_bottom = N_3D * mask_bottom
+
+    # average depth_mid
+    N_3D_ave = N_3D_bottom.mean(dim="depth_mid", skipna=True) / (2 * np.pi)
+    return N_3D_ave
+
+
+def regrid_data(
+    N_3D_ave: xr.DataArray, lon: xr.DataArray, lat: xr.DataArray
+) -> xr.Dataset:
+    """
+    Regrid the averaged stratification data to a target grid (using the original lon/lat).
+    """
+    mask_array = ~np.isnan(N_3D_ave.values)
+    ds = xr.Dataset(
+        data_vars={
+            "N_3D_ave": (("lat", "lon"), N_3D_ave.values),
+            "mask": (("lat", "lon"), mask_array),
+        },
+        coords={"lon": lon, "lat": lat},
+    )
+
+    ds_out = xr.Dataset(
+        {
+            "lat": (["lat"], lat.values),
+            "lon": (["lon"], lon.values),
+        }
+    )
+
+    # Regrid
+    regridder = xe.Regridder(ds, ds_out, "bilinear", extrap_method="inverse_dist")
+    ds_out = regridder(ds)
+
+    Nbot_data = xr.Dataset(
+        {"Nbot": (("lat", "lon"), ds_out["N_3D_ave"].values)},
+        coords={"lon": lon, "lat": lat},
+    )
+    return Nbot_data
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate bottom-Xm averaged stratification from WOA climatology."
+    )
+    parser.add_argument(
+        "--depth-threshold",
+        type=float,
+        default=500,
+        help="Bottom threshold in meters over which to average (default: 500).",
+    )
+    parser.add_argument(
+        "--temp-file",
+        type=str,
+        required=True,
+        help="Path to the WOA temperature NetCDF file.",
+    )
+    parser.add_argument(
+        "--sal-file",
+        type=str,
+        required=True,
+        help="Path to the WOA salinity NetCDF file.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="Nbot_freq.nc",
+        help="Output NetCDF filename (default: Nbot_freq.nc).",
+    )
+    args = parser.parse_args()
+
+    # Load Data
+    sea_water_temp, sea_water_sal = load_woa_data(args.temp_file, args.sal_file)
+    lon = sea_water_temp.lon
+    lat = sea_water_temp.lat
+    depth = sea_water_temp.depth
+
+    # Compute 3D Pressure
+    pressure_3d = compute_pressure(depth, lat, lon)
+
+    # Compute Stratification
+    N_3D, depth_mid = compute_stratification(
+        sea_water_temp, sea_water_sal, pressure_3d, depth, lat
+    )
+
+    # Compute bottom average of N over the bottom DEPTH_THRESHOLD meters
+    N_3D_ave = compute_bottom_average(
+        N_3D, depth_mid, sea_water_temp, depth, args.depth_threshold
+    )
+
+    # Regrid the averaged field onto the original grid
+    Nbot_data = regrid_data(N_3D_ave, lon, lat)
+
+    # Add metadata
+    this_file = os.path.normpath(__file__)
+    runcmd = (
+        f"python3 {os.path.basename(this_file)} "
+        f"--temp-file={args.temp_file} "
+        f"--sal-file={args.sal_file} "
+        f"--depth-threshold={args.depth_threshold} "
+        f"--output={args.output}"
+    )
+
+    global_attrs = {"history": get_provenance_metadata(this_file, runcmd)}
+    Nbot_data.attrs.update(global_attrs)
+
+    # Write output to a NetCDF file.
+    Nbot_data.to_netcdf(args.output)
+
+
+if __name__ == "__main__":
+    main()

From e55729858a614d5efcbefd3f1d54374315463d5b Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Mon, 17 Mar 2025 16:11:18 +1100
Subject: [PATCH 02/16] enable both bottom-Xm-averaged N and depth-averaged N

---
 .../generate_bottom_N.py                      | 60 +++++++++++--------
 1 file changed, 34 insertions(+), 26 deletions(-)

diff --git a/external_tidal_generation/generate_bottom_N.py b/external_tidal_generation/generate_bottom_N.py
index a932772..d1751d5 100644
--- a/external_tidal_generation/generate_bottom_N.py
+++ b/external_tidal_generation/generate_bottom_N.py
@@ -2,16 +2,20 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # =========================================================================================
-# Generate a bottom-Xm-averaged stratification from WOA climatology.
+# Generate an N field (buoyancy frequency) from WOA climatology.
 #
 # This script creates a 3D stratification field from WOA temperature and salinity
 # data, computes the buoyancy frequency (N) on a mid-depth grid, and then averages N
 # over the bottom DEPTH_THRESHOLD meters of the water column for each (lat, lon) column.
 #
+# - If --depth-threshold > 0, compute a bottom-X meters average.
+# - If --depth-threshold == 0 (default), compute a full water-column (depth-averaged) mean.
+#
 # Usage:
-#    python3 generate_bottom_N.py --depth-threshold 500 \
-#         --temp-file /path/to/WOA_decav_t00_04.nc \
-#         --sal-file /path/to/WOA_decav_s00_04.nc \
+#    python3 generate_bottom_N.py \
+#         --temp-file /path/to/woa_decav_t00_04.nc \
+#         --sal-file /path/to/woa_decav_s00_04.nc \
+#         [--depth-threshold 500.0] \
 #         --output Nbot_freq.nc
 #
 # Contact:
@@ -30,7 +34,6 @@
 import argparse
 import os
 import warnings
-
 warnings.filterwarnings("ignore")
 
 import xarray as xr
@@ -54,7 +57,7 @@ def load_woa_data(temp_file: str, sal_file: str) -> (xr.DataArray, xr.DataArray)
     return sea_water_temp, sea_water_sal
 
 
-def compute_pressure(depth: xr.DataArray, lat: xr.DataArray, lon: xr.DataArray):
+def compute_pressure(depth: xr.DataArray, lat: xr.DataArray, lon: xr.DataArray) -> xr.DataArray:
     """
     Compute a 3D pressure field (in dbar) from 1D depth and lat arrays.
     gsw.p_from_z expects negative depth.
@@ -76,7 +79,6 @@ def compute_stratification(
     sea_water_temp: xr.DataArray,
     sea_water_sal: xr.DataArray,
     pressure_3d: xr.DataArray,
-    depth: xr.DataArray,
     lat: xr.DataArray,
 ) -> (xr.DataArray, xr.DataArray):
     """
@@ -105,26 +107,32 @@ def compute_bottom_average(
     depth_mid: xr.DataArray,
     sea_water_temp: xr.DataArray,
     depth: xr.DataArray,
-    depth_threshold: float,
+    depth_threshold: float = 0.0,
 ) -> xr.DataArray:
     """
-    Create a mask for the bottom depth_threshold meters, apply it to N_3D and average over depth_mid.
+    Compute either a bottom-X meters average if depth_threshold>0,
+    or a full-column (depth-averaged) mean if depth_threshold==0.
     """
-    # Broadcast depth
-    depth_array = sea_water_temp * 0 + depth
-    max_depth = depth_array.max(dim="depth", skipna=True)
-    bottom_threshold = max_depth - depth_threshold
-
-    # Create a mask
-    mask_bottom = xr.where(
-        (depth_mid >= bottom_threshold) & (depth_mid < max_depth),
-        1,
-        np.nan,
-    )
-    N_3D_bottom = N_3D * mask_bottom
+    if depth_threshold > 0:
+        # Broadcast depth
+        depth_array = sea_water_temp * 0 + depth
+        max_depth = depth_array.max(dim="depth", skipna=True)
+        bottom_threshold = max_depth - depth_threshold
+
+        # Create a mask
+        mask_bottom = xr.where(
+            (depth_mid >= bottom_threshold) & (depth_mid < max_depth),
+            1,
+            np.nan,
+        )
+        N_3D_bottom = N_3D * mask_bottom
+
+        # average depth_mid
+        N_3D_ave = N_3D_bottom.mean(dim="depth_mid", skipna=True) / (2 * np.pi)
+    else:
+        # full depth average
+        N_3D_ave = N_3D.mean(dim="depth_mid", skipna=True) / (2 * np.pi)
 
-    # average depth_mid
-    N_3D_ave = N_3D_bottom.mean(dim="depth_mid", skipna=True) / (2 * np.pi)
     return N_3D_ave
 
 
@@ -168,7 +176,7 @@ def main():
     parser.add_argument(
         "--depth-threshold",
         type=float,
-        default=500,
+        default=0.0,
         help="Bottom threshold in meters over which to average (default: 500).",
     )
     parser.add_argument(
@@ -202,10 +210,10 @@ def main():
 
     # Compute Stratification
     N_3D, depth_mid = compute_stratification(
-        sea_water_temp, sea_water_sal, pressure_3d, depth, lat
+        sea_water_temp, sea_water_sal, pressure_3d, lat
     )
 
-    # Compute bottom average of N over the bottom DEPTH_THRESHOLD meters
+    # Compute average N (bottom-X or full column)
     N_3D_ave = compute_bottom_average(
         N_3D, depth_mid, sea_water_temp, depth, args.depth_threshold
     )

From 9890a66fe4b6c235b51ab8eb58773579cee8fdd2 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Mon, 17 Mar 2025 17:05:12 +1100
Subject: [PATCH 03/16] Add md5sum for input files

---
 external_tidal_generation/generate_bottom_N.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/external_tidal_generation/generate_bottom_N.py b/external_tidal_generation/generate_bottom_N.py
index d1751d5..238acee 100644
--- a/external_tidal_generation/generate_bottom_N.py
+++ b/external_tidal_generation/generate_bottom_N.py
@@ -34,6 +34,7 @@
 import argparse
 import os
 import warnings
+
 warnings.filterwarnings("ignore")
 
 import xarray as xr
@@ -57,7 +58,9 @@ def load_woa_data(temp_file: str, sal_file: str) -> (xr.DataArray, xr.DataArray)
     return sea_water_temp, sea_water_sal
 
 
-def compute_pressure(depth: xr.DataArray, lat: xr.DataArray, lon: xr.DataArray) -> xr.DataArray:
+def compute_pressure(
+    depth: xr.DataArray, lat: xr.DataArray, lon: xr.DataArray
+) -> xr.DataArray:
     """
     Compute a 3D pressure field (in dbar) from 1D depth and lat arrays.
     gsw.p_from_z expects negative depth.
@@ -232,6 +235,14 @@ def main():
     )
 
     global_attrs = {"history": get_provenance_metadata(this_file, runcmd)}
+
+    # add md5 hashes for input files
+    file_hashes = [
+        f"{args.temp_file} (md5 hash: {md5sum(args.temp_file)})",
+        f"{args.sal_file} (md5 hash: {md5sum(args.sal_file)})",
+    ]
+    global_attrs["inputFile"] = ", ".join(file_hashes)
+
     Nbot_data.attrs.update(global_attrs)
 
     # Write output to a NetCDF file.

From b9218e2f6fbcaaf1a55ef8bd46d7f40b1130e00d Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Wed, 2 Apr 2025 12:49:19 +1100
Subject: [PATCH 04/16] generate bottom roughness using polyfit method, Jayne
 et al (2001)

---
 .../generate_bottom_roughness_polyfit.py      | 354 ++++++++++++++++++
 1 file changed, 354 insertions(+)
 create mode 100644 external_tidal_generation/generate_bottom_roughness_polyfit.py

diff --git a/external_tidal_generation/generate_bottom_roughness_polyfit.py b/external_tidal_generation/generate_bottom_roughness_polyfit.py
new file mode 100644
index 0000000..d175244
--- /dev/null
+++ b/external_tidal_generation/generate_bottom_roughness_polyfit.py
@@ -0,0 +1,354 @@
+#!/usr/bin/env python3
+# Copyright 2025 ACCESS-NRI and contributors.
+# SPDX-License-Identifier: Apache-2.0
+
+# =========================================================================================
+# Compute bottom roughness (h2) on an ocean model grid by fitting a
+# polynomial to high-resolution bathymetry (1/240 degree) using MPI.
+# Jayne, Steven R., and Louis C. St. Laurent. "Parameterizing tidal dissipation over rough topography." Geophysical Research Letters 28.5 (2001): 811-814.
+#
+# Usage:
+#    mpirun -n <ranks> python3 generate_bottom_roughness_polyfit.py \
+#         --topo-file /path/to/topog.nc \
+#         --grid-file /path/to/ocean_static.nc \
+#         --mask-file /path/to/ocean_mask.nc \
+#         --output h2.nc
+#
+# Contact:
+#    - Minghang Li <Minghang.Li1@anu.edu.au>
+#
+# Dependencies:
+#   - xarray
+#   - numpy
+#   - mpi4py
+# =========================================================================================
+from pathlib import Path
+import os
+import sys
+import argparse
+import numpy as np
+import xarray as xr
+from mpi4py import MPI
+import subprocess
+
+path_root = Path(__file__).parents[1]
+sys.path.append(str(path_root))
+
+from scripts_common import get_provenance_metadata, md5sum
+
+
+def polyfit_roughness(H: np.ndarray, xv: np.ndarray, yv: np.ndarray) -> float:
+    """
+    Fit a polynomial of the form:
+        H(x,y) = a + b*x + c*y + d*(x*y)
+    to the 2D topography array H, and compute the RMS
+    of the residuals (h2).
+    """
+    H1 = H.ravel()
+    valid = ~np.isnan(H1)
+    # At least 6 valid points
+    if np.count_nonzero(valid) < 6:
+        return np.nan
+
+    X1 = xv.ravel()
+    Y1 = yv.ravel()
+    X = np.column_stack(
+        [np.ones(np.sum(valid)), X1[valid], Y1[valid], (X1 * Y1)[valid]]
+    )
+    Y_valid = H1[valid]
+    try:
+        coeffs, *_ = np.linalg.lstsq(X, Y_valid, rcond=None)
+
+        H_fit = (
+            coeffs[0]
+            + coeffs[1] * X1[valid]
+            + coeffs[2] * Y1[valid]
+            + coeffs[3] * (X1[valid] * Y1[valid])
+        )
+        res = Y_valid - H_fit
+        h2 = np.sqrt(np.mean(res**2))
+    except Exception as e:
+        h2 = np.nan
+    return h2
+
+
+def compute_h2_poly_cell(
+    lon_min: float,
+    lon_max: float,
+    lat_min: float,
+    lat_max: float,
+    topo_da: xr.DataArray,
+) -> float:
+    """
+    Given bounding coordinates (lon_min, lon_max, lat_min, lat_max),
+    select the corresponding sub-region from the high-res bathymetry
+    and fit a polynomial (H = a + b*x + c*y + d*x*y) to that
+    sub-region.
+    """
+    sub_da = topo_da.sel(
+        lon=slice(lon_min, lon_max),
+        lat=slice(lat_min, lat_max),
+    )
+
+    H = sub_da.values
+    x = sub_da.lon.values
+    y = sub_da.lat.values
+    xv, yv = np.meshgrid(x, y)
+    h2 = polyfit_roughness(H, xv, yv)
+    return h2
+
+
+def load_topo(
+    path: str, chunk_lat: int = 800, chunk_lon: int = 1600
+) -> (xr.DataArray, float):
+    """
+    Load a high-resolution bathymetry file
+    """
+    ds = xr.open_dataset(path, chunks={"lat": chunk_lat, "lon": chunk_lon})
+    depth = ds["z"]
+    depth = ds["z"].where(ds["z"] < 0, np.nan)
+    new_lon = xr.where(depth.lon >= 80, depth.lon - 360, depth.lon)
+    depth = depth.assign_coords(lon=new_lon).sortby("lon")
+    return depth
+
+
+def load_model_grids(path: str):
+    """
+    Load the ocean model grid information from a static file.
+    """
+    ds = xr.open_dataset(path)
+    xh = ds.xh
+    yh = ds.yh
+    ds.close()
+    return xh, yh
+
+
+def compute_edges_from_centers(centers: np.ndarray) -> np.ndarray:
+    """
+    Given an array of cell center coordinates, compute the
+    corresponding cell edges.
+    """
+    N = len(centers)
+    edges = np.empty(N + 1, dtype=centers.dtype)
+    edges[0] = centers[0] - 0.5 * (centers[1] - centers[0])
+    for i in range(1, N):
+        edges[i] = 0.5 * (centers[i - 1] + centers[i])
+    edges[N] = centers[N - 1] + 0.5 * (centers[N - 1] - centers[N - 2])
+    return edges
+
+
+def evaluate_roughness(
+    topo_da: xr.DataArray,
+    ocean_mask: xr.DataArray,
+    xedges: np.ndarray,
+    yedges: np.ndarray,
+    NxCells: int,
+    NyCells: int,
+    comm: MPI.Comm,
+) -> np.ndarray:
+    """
+    Distribute roughness computations across all MPI ranks, and
+    gather the final 2D array of h2 values on rank 0.
+    """
+
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+
+    total_rows = NyCells
+    block_size = total_rows // size
+    rem = total_rows % size
+
+    y_start = rank * block_size + min(rank, rem)
+    y_count = block_size + (1 if rank < rem else 0)
+    y_end = y_start + y_count
+
+    if rank == 0:
+        print(
+            f"[Rank {rank}] Domain partition: total rows {total_rows}, "
+            f"Rank {rank} covers rows {y_start} to {y_end - 1}."
+        )
+
+    # Initilise local h2 array
+    # local_h2 = np.zeros((y_count, NxCells), dtype=np.float32)
+    local_h2 = np.full((y_count, NxCells), np.nan, dtype=np.float32)
+    # Compute h2 for each rank
+    for j in range(y_start, y_end):
+        for i in range(NxCells):
+            # Skip land cells
+            if ocean_mask[j, i] == 0:
+                continue
+            h2_val = compute_h2_poly_cell(
+                xedges[i], xedges[i + 1], yedges[j], yedges[j + 1], topo_da
+            )
+            local_h2[j - y_start, i] = h2_val
+
+        if (j - y_start) % 3 == 0:
+            print(
+                f"[Rank {rank}] Processed global row {j} (local index {j - y_start})",
+                flush=True,
+            )
+
+    # Collect all local data to rank 0
+    local_1d = local_h2.ravel()
+    local_size = local_1d.size
+
+    # Gather sizes from all ranks
+    all_sizes = comm.gather(local_size, root=0)
+
+    if rank == 0:
+        total_size = sum(all_sizes)
+        global_1d = np.empty(total_size, dtype=np.float32)
+        displs = np.insert(np.cumsum(all_sizes), 0, 0)[:-1]
+    else:
+        global_1d = None
+        displs = None
+
+    comm.Gatherv(sendbuf=local_1d, recvbuf=(global_1d, (all_sizes, displs)), root=0)
+
+    # On rank 0, reshape back into a 2D array
+    if rank == 0:
+        final_h2 = np.full((NyCells, NxCells), np.nan, dtype=np.float32)
+        offset = 0
+        for r in range(size):
+            block = all_sizes[r]
+            sub_data = global_1d[offset : offset + block]
+            offset += block
+
+            block_rows = block // NxCells
+            r_start = r * block_size + min(r, rem)
+            final_h2[r_start : r_start + block_rows, :] = sub_data.reshape(
+                block_rows, NxCells
+            )
+
+        return final_h2
+    else:
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compute bottom roughness (h2) via polynomial fit with a high-resolution topography."
+    )
+    parser.add_argument(
+        "--topo-file",
+        type=str,
+        required=True,
+        help="Path to a high-resolution topography file.",
+    )
+    parser.add_argument(
+        "--grid-file", type=str, required=True, help="Path to the ocean static file."
+    )
+    parser.add_argument(
+        "--grid-mask", type=str, required=True, help="Path to the ocean mask file."
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="h2.nc",
+        help="Output roughness filename (default: h2.nc).",
+    )
+    parser.add_argument(
+        "--chunk-lat",
+        type=int,
+        default=800,
+        help="Dask chunk size along lat dimension (default:800).",
+    )
+    parser.add_argument(
+        "--chunk-lon",
+        type=int,
+        default=1600,
+        help="Dask chunk size along lon dimension (default:1600).",
+    )
+    args = parser.parse_args()
+
+    # create communicator
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+
+    # Rank 0 loads data and initializes shared info
+    if rank == 0:
+        topo_da = load_topo(
+            args.topo_file, chunk_lat=args.chunk_lat, chunk_lon=args.chunk_lon
+        )
+        ocean_mask = xr.open_dataset(args.grid_mask).mask
+        xh, yh = load_model_grids(args.grid_file)
+        xedges = compute_edges_from_centers(xh)
+        yedges = compute_edges_from_centers(yh)
+        NxCells = len(xedges) - 1
+        NyCells = len(yedges) - 1
+    else:
+        # Placeholders on other ranks
+        topo_da = None
+        ocean_mask = None
+        xedges = None
+        yedges = None
+        NxCells = None
+        NyCells = None
+        xh = None
+        yh = None
+
+    topo_da = comm.bcast(topo_da, root=0)
+    ocean_mask = comm.bcast(ocean_mask, root=0)
+    NxCells = comm.bcast(NxCells, root=0)
+    NyCells = comm.bcast(NyCells, root=0)
+    xedges = comm.bcast(xedges, root=0)
+    yedges = comm.bcast(yedges, root=0)
+    xh = comm.bcast(xh, root=0)
+    yh = comm.bcast(yh, root=0)
+
+    final_h2 = evaluate_roughness(
+        topo_da=topo_da,
+        ocean_mask=ocean_mask,
+        xedges=xedges,
+        yedges=yedges,
+        NxCells=NxCells,
+        NyCells=NyCells,
+        comm=comm,
+    )
+
+    # Rank 0 writes results to file
+    if rank == 0:
+        h2_out = xr.Dataset(
+            {"h2": (("yh", "xh"), final_h2)},
+            coords={
+                "yh": yh,
+                "xh": xh,
+            },
+            attrs={
+                "long_name": "Polynomial-fit roughness (h2) per model cell",
+                "units": "m",
+            },
+        )
+        # Add metadata
+        this_file = os.path.normpath(__file__)
+        runcmd = (
+            f"mpirun -n $PBS_NCPUS"
+            f"python3 {os.path.basename(this_file)} "
+            f"--topo-file={args.topo_file}"
+            f"--chunk_lat={args.chunk_lat}"
+            f"--chunk_lon={args.chunk_lon}"
+            f"--grid-file={args.grid_file} "
+            f"--grid-mask={args.grid_mask} "
+            f"--output={args.output}"
+        )
+
+        try:
+            history = get_provenance_metadata(this_file, runcmd)
+        except subprocess.CalledProcessError:
+            history = "Provenance metadata unavailable (not a git repo?)"
+        global_attrs = {"history": history}
+        # add md5 hashes for input files
+        file_hashes = [
+            f"{args.topo_file} (md5 hash: {md5sum(args.topo_file)})",
+            f"{args.grid_file} (md5 hash: {md5sum(args.grid_file)})",
+            f"{args.grid_mask} (md5 hash: {md5sum(args.grid_mask)})",
+        ]
+        global_attrs["inputFile"] = ", ".join(file_hashes)
+
+        h2_out.attrs.update(global_attrs)
+
+        h2_out.to_netcdf(args.output)
+
+
+if __name__ == "__main__":
+    main()

From 3922c53d4b8b52ce15ccc599ccc26b92891dc159 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Fri, 4 Apr 2025 08:59:59 +1100
Subject: [PATCH 05/16] tidy a bit

---
 .../generate_bottom_roughness_polyfit.py         | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/external_tidal_generation/generate_bottom_roughness_polyfit.py b/external_tidal_generation/generate_bottom_roughness_polyfit.py
index d175244..7351a89 100644
--- a/external_tidal_generation/generate_bottom_roughness_polyfit.py
+++ b/external_tidal_generation/generate_bottom_roughness_polyfit.py
@@ -46,8 +46,8 @@ def polyfit_roughness(H: np.ndarray, xv: np.ndarray, yv: np.ndarray) -> float:
     """
     H1 = H.ravel()
     valid = ~np.isnan(H1)
-    # At least 6 valid points
-    if np.count_nonzero(valid) < 6:
+    # At least 4 valid points
+    if np.count_nonzero(valid) < 4:
         return np.nan
 
     X1 = xv.ravel()
@@ -105,7 +105,6 @@ def load_topo(
     Load a high-resolution bathymetry file
     """
     ds = xr.open_dataset(path, chunks={"lat": chunk_lat, "lon": chunk_lon})
-    depth = ds["z"]
     depth = ds["z"].where(ds["z"] < 0, np.nan)
     new_lon = xr.where(depth.lon >= 80, depth.lon - 360, depth.lon)
     depth = depth.assign_coords(lon=new_lon).sortby("lon")
@@ -185,7 +184,6 @@ def evaluate_roughness(
         if (j - y_start) % 3 == 0:
             print(
                 f"[Rank {rank}] Processed global row {j} (local index {j - y_start})",
-                flush=True,
             )
 
     # Collect all local data to rank 0
@@ -239,7 +237,7 @@ def main():
         "--grid-file", type=str, required=True, help="Path to the ocean static file."
     )
     parser.add_argument(
-        "--grid-mask", type=str, required=True, help="Path to the ocean mask file."
+        "--mask-file", type=str, required=True, help="Path to the ocean mask file."
     )
     parser.add_argument(
         "--output",
@@ -270,7 +268,7 @@ def main():
         topo_da = load_topo(
             args.topo_file, chunk_lat=args.chunk_lat, chunk_lon=args.chunk_lon
         )
-        ocean_mask = xr.open_dataset(args.grid_mask).mask
+        ocean_mask = xr.open_dataset(args.mask_file).mask
         xh, yh = load_model_grids(args.grid_file)
         xedges = compute_edges_from_centers(xh)
         yedges = compute_edges_from_centers(yh)
@@ -327,8 +325,8 @@ def main():
             f"--topo-file={args.topo_file}"
             f"--chunk_lat={args.chunk_lat}"
             f"--chunk_lon={args.chunk_lon}"
-            f"--grid-file={args.grid_file} "
-            f"--grid-mask={args.grid_mask} "
+            f"--grid-file={args.grid_file}"
+            f"--mask-file={args.mask_file}"
             f"--output={args.output}"
         )
 
@@ -341,7 +339,7 @@ def main():
         file_hashes = [
             f"{args.topo_file} (md5 hash: {md5sum(args.topo_file)})",
             f"{args.grid_file} (md5 hash: {md5sum(args.grid_file)})",
-            f"{args.grid_mask} (md5 hash: {md5sum(args.grid_mask)})",
+            f"{args.mask_file} (md5 hash: {md5sum(args.mask_file)})",
         ]
         global_attrs["inputFile"] = ", ".join(file_hashes)
 

From fd629bdeae938714a81258c324902d6a74a20bca Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 8 Apr 2025 13:47:15 +1000
Subject: [PATCH 06/16] refactor generate_bottom_roughness_polyfit.py

---
 .../generate_bottom_roughness_polyfit.py      | 238 +++++++++---------
 1 file changed, 124 insertions(+), 114 deletions(-)

diff --git a/external_tidal_generation/generate_bottom_roughness_polyfit.py b/external_tidal_generation/generate_bottom_roughness_polyfit.py
index 7351a89..75a0950 100644
--- a/external_tidal_generation/generate_bottom_roughness_polyfit.py
+++ b/external_tidal_generation/generate_bottom_roughness_polyfit.py
@@ -3,16 +3,20 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # =========================================================================================
-# Compute bottom roughness (h2) on an ocean model grid by fitting a
-# polynomial to high-resolution bathymetry (1/240 degree) using MPI.
-# Jayne, Steven R., and Louis C. St. Laurent. "Parameterizing tidal dissipation over rough topography." Geophysical Research Letters 28.5 (2001): 811-814.
+# Compute bottom roughness on an ocean model grid by fitting a
+# polynomial to high-resolution bathymetry (1/240 degree) using mpi.
+#
+# Reference:
+# Jayne, Steven R., and Louis C. St. Laurent.
+# "Parameterizing tidal dissipation over rough topography."
+# Geophysical Research Letters 28.5 (2001): 811-814.
 #
 # Usage:
 #    mpirun -n <ranks> python3 generate_bottom_roughness_polyfit.py \
 #         --topo-file /path/to/topog.nc \
-#         --grid-file /path/to/ocean_static.nc \
-#         --mask-file /path/to/ocean_mask.nc \
-#         --output h2.nc
+#         --hgrid-file /path/to/ocean_static.nc \
+#         --regrid-mask-file /path/to/ocean_mask.nc \
+#         --output output.nc
 #
 # Contact:
 #    - Minghang Li <Minghang.Li1@anu.edu.au>
@@ -26,9 +30,10 @@
 import os
 import sys
 import argparse
+from mpi4py import MPI
 import numpy as np
 import xarray as xr
-from mpi4py import MPI
+
 import subprocess
 
 path_root = Path(__file__).parents[1]
@@ -42,7 +47,7 @@ def polyfit_roughness(H: np.ndarray, xv: np.ndarray, yv: np.ndarray) -> float:
     Fit a polynomial of the form:
         H(x,y) = a + b*x + c*y + d*(x*y)
     to the 2D topography array H, and compute the RMS
-    of the residuals (h2).
+    of the residuals (hrms).
     """
     H1 = H.ravel()
     valid = ~np.isnan(H1)
@@ -50,6 +55,7 @@ def polyfit_roughness(H: np.ndarray, xv: np.ndarray, yv: np.ndarray) -> float:
     if np.count_nonzero(valid) < 4:
         return np.nan
 
+    # prepare X matrix for least-squares fitting
     X1 = xv.ravel()
     Y1 = yv.ravel()
     X = np.column_stack(
@@ -66,13 +72,14 @@ def polyfit_roughness(H: np.ndarray, xv: np.ndarray, yv: np.ndarray) -> float:
             + coeffs[3] * (X1[valid] * Y1[valid])
         )
         res = Y_valid - H_fit
-        h2 = np.sqrt(np.mean(res**2))
+        hrms = np.sqrt(np.mean(res**2))
     except Exception as e:
-        h2 = np.nan
-    return h2
+        hrms = np.nan
 
+    return hrms
 
-def compute_h2_poly_cell(
+
+def compute_hrms_poly_cell(
     lon_min: float,
     lon_max: float,
     lat_min: float,
@@ -94,66 +101,42 @@ def compute_h2_poly_cell(
     x = sub_da.lon.values
     y = sub_da.lat.values
     xv, yv = np.meshgrid(x, y)
-    h2 = polyfit_roughness(H, xv, yv)
-    return h2
+    hrms = polyfit_roughness(H, xv, yv)
+    return hrms
 
 
-def load_topo(
-    path: str, chunk_lat: int = 800, chunk_lon: int = 1600
-) -> (xr.DataArray, float):
+def load_topo(path: str, chunk_lat: int = 800, chunk_lon: int = 1600) -> xr.DataArray:
     """
     Load a high-resolution bathymetry file
     """
     ds = xr.open_dataset(path, chunks={"lat": chunk_lat, "lon": chunk_lon})
     depth = ds["z"].where(ds["z"] < 0, np.nan)
-    new_lon = xr.where(depth.lon >= 80, depth.lon - 360, depth.lon)
-    depth = depth.assign_coords(lon=new_lon).sortby("lon")
     return depth
 
 
-def load_model_grids(path: str):
-    """
-    Load the ocean model grid information from a static file.
-    """
-    ds = xr.open_dataset(path)
-    xh = ds.xh
-    yh = ds.yh
-    ds.close()
-    return xh, yh
-
-
-def compute_edges_from_centers(centers: np.ndarray) -> np.ndarray:
-    """
-    Given an array of cell center coordinates, compute the
-    corresponding cell edges.
-    """
-    N = len(centers)
-    edges = np.empty(N + 1, dtype=centers.dtype)
-    edges[0] = centers[0] - 0.5 * (centers[1] - centers[0])
-    for i in range(1, N):
-        edges[i] = 0.5 * (centers[i - 1] + centers[i])
-    edges[N] = centers[N - 1] + 0.5 * (centers[N - 1] - centers[N - 2])
-    return edges
+def load_regrid_ocean_mask(path: str) -> xr.Dataset:
+    interp_ocean_mask = xr.open_dataset(path)
+    return interp_ocean_mask
 
 
 def evaluate_roughness(
     topo_da: xr.DataArray,
     ocean_mask: xr.DataArray,
-    xedges: np.ndarray,
-    yedges: np.ndarray,
-    NxCells: int,
-    NyCells: int,
+    hgrid_xc: np.ndarray,
+    hgrid_yc: np.ndarray,
+    nx: int,
+    ny: int,
     comm: MPI.Comm,
 ) -> np.ndarray:
     """
     Distribute roughness computations across all MPI ranks, and
-    gather the final 2D array of h2 values on rank 0.
+    gather the final 2D array of hrms values on rank 0.
     """
 
     rank = comm.Get_rank()
     size = comm.Get_size()
 
-    total_rows = NyCells
+    total_rows = ny
     block_size = total_rows // size
     rem = total_rows % size
 
@@ -167,27 +150,43 @@ def evaluate_roughness(
             f"Rank {rank} covers rows {y_start} to {y_end - 1}."
         )
 
-    # Initilise local h2 array
-    # local_h2 = np.zeros((y_count, NxCells), dtype=np.float32)
-    local_h2 = np.full((y_count, NxCells), np.nan, dtype=np.float32)
-    # Compute h2 for each rank
+    local_hrms = np.full((y_count, nx), np.nan, dtype=np.float32)
+
+    # Compute hrms for each rank
     for j in range(y_start, y_end):
-        for i in range(NxCells):
+        for i in range(nx):
             # Skip land cells
             if ocean_mask[j, i] == 0:
                 continue
-            h2_val = compute_h2_poly_cell(
-                xedges[i], xedges[i + 1], yedges[j], yedges[j + 1], topo_da
-            )
-            local_h2[j - y_start, i] = h2_val
+
+            this_lon_corners = [
+                hgrid_xc[j, i],
+                hgrid_xc[j, i + 1],
+                hgrid_xc[j + 1, i],
+                hgrid_xc[j + 1, i + 1],
+            ]
+            this_lat_corners = [
+                hgrid_yc[j, i],
+                hgrid_yc[j, i + 1],
+                hgrid_yc[j + 1, i],
+                hgrid_yc[j + 1, i + 1],
+            ]
+
+            lon_min = np.min(this_lon_corners)
+            lon_max = np.max(this_lon_corners)
+            lat_min = np.min(this_lat_corners)
+            lat_max = np.max(this_lat_corners)
+
+            hrms_val = compute_hrms_poly_cell(lon_min, lon_max, lat_min, lat_max, topo_da)
+            local_hrms[j - y_start, i] = hrms_val
 
         if (j - y_start) % 3 == 0:
             print(
-                f"[Rank {rank}] Processed global row {j} (local index {j - y_start})",
+                f"[Rank {rank}] Processed row {j} (local index {j - y_start})",
             )
 
     # Collect all local data to rank 0
-    local_1d = local_h2.ravel()
+    local_1d = local_hrms.ravel()
     local_size = local_1d.size
 
     # Gather sizes from all ranks
@@ -205,27 +204,26 @@ def evaluate_roughness(
 
     # On rank 0, reshape back into a 2D array
     if rank == 0:
-        final_h2 = np.full((NyCells, NxCells), np.nan, dtype=np.float32)
+        final_hrms = np.full((ny, nx), np.nan, dtype=np.float32)
         offset = 0
         for r in range(size):
             block = all_sizes[r]
             sub_data = global_1d[offset : offset + block]
             offset += block
 
-            block_rows = block // NxCells
+            block_rows = block // nx
             r_start = r * block_size + min(r, rem)
-            final_h2[r_start : r_start + block_rows, :] = sub_data.reshape(
-                block_rows, NxCells
+            final_hrms[r_start : r_start + block_rows, :] = sub_data.reshape(
+                block_rows, nx
             )
-
-        return final_h2
+        return final_hrms
     else:
         return None
 
 
 def main():
     parser = argparse.ArgumentParser(
-        description="Compute bottom roughness (h2) via polynomial fit with a high-resolution topography."
+        description="Compute bottom roughness via polynomial fit with a high-resolution topography."
     )
     parser.add_argument(
         "--topo-file",
@@ -234,16 +232,28 @@ def main():
         help="Path to a high-resolution topography file.",
     )
     parser.add_argument(
-        "--grid-file", type=str, required=True, help="Path to the ocean static file."
+        "--regrid-mask-file",
+        type=str,
+        required=True,
+        help="Path to the regrid ocean mask file.",
     )
     parser.add_argument(
-        "--mask-file", type=str, required=True, help="Path to the ocean mask file."
+        "--method",
+        type=str,
+        default="conservative_normed",
+        help="Regridding method (e.g., bilinear, conservative, conservative_normed)",
+    )
+    parser.add_argument(
+        "--interpolated-field-name",
+        type=str,
+        default="hrms_interpolated",
+        help="The field name for the interpolated roughness (default: hrms_interpolated).",
     )
     parser.add_argument(
         "--output",
         type=str,
-        default="h2.nc",
-        help="Output roughness filename (default: h2.nc).",
+        default="interpolated_hrms.nc",
+        help="Output roughness filename (default: interpolated_hrms.nc).",
     )
     parser.add_argument(
         "--chunk-lat",
@@ -263,60 +273,60 @@ def main():
     comm = MPI.COMM_WORLD
     rank = comm.Get_rank()
 
-    # Rank 0 loads data and initializes shared info
     if rank == 0:
         topo_da = load_topo(
             args.topo_file, chunk_lat=args.chunk_lat, chunk_lon=args.chunk_lon
         )
-        ocean_mask = xr.open_dataset(args.mask_file).mask
-        xh, yh = load_model_grids(args.grid_file)
-        xedges = compute_edges_from_centers(xh)
-        yedges = compute_edges_from_centers(yh)
-        NxCells = len(xedges) - 1
-        NyCells = len(yedges) - 1
+        regrid_ocean_mask = load_regrid_ocean_mask(args.regrid_mask_file)
+        lon_b = regrid_ocean_mask.lon_b.values
+        lat_b = regrid_ocean_mask.lat_b.values
+        ny = lon_b.shape[0] - 1
+        nx = lon_b.shape[1] - 1
+
     else:
-        # Placeholders on other ranks
         topo_da = None
-        ocean_mask = None
-        xedges = None
-        yedges = None
-        NxCells = None
-        NyCells = None
-        xh = None
-        yh = None
+        regrid_ocean_mask = None
+        lon_b = None
+        lat_b = None
+        ny = None
+        nx = None
 
+    # Broadcast shared datasets to all ranks
     topo_da = comm.bcast(topo_da, root=0)
-    ocean_mask = comm.bcast(ocean_mask, root=0)
-    NxCells = comm.bcast(NxCells, root=0)
-    NyCells = comm.bcast(NyCells, root=0)
-    xedges = comm.bcast(xedges, root=0)
-    yedges = comm.bcast(yedges, root=0)
-    xh = comm.bcast(xh, root=0)
-    yh = comm.bcast(yh, root=0)
-
-    final_h2 = evaluate_roughness(
+    regrid_ocean_mask = comm.bcast(regrid_ocean_mask, root=0)
+    lon_b = comm.bcast(lon_b, root=0)
+    lat_b = comm.bcast(lat_b, root=0)
+    nx = comm.bcast(nx, root=0)
+    ny = comm.bcast(ny, root=0)
+
+    final_hrms = evaluate_roughness(
         topo_da=topo_da,
-        ocean_mask=ocean_mask,
-        xedges=xedges,
-        yedges=yedges,
-        NxCells=NxCells,
-        NyCells=NyCells,
+        ocean_mask=regrid_ocean_mask[args.method],
+        hgrid_xc=lon_b,
+        hgrid_yc=lat_b,
+        nx=nx,
+        ny=ny,
         comm=comm,
     )
 
     # Rank 0 writes results to file
     if rank == 0:
-        h2_out = xr.Dataset(
-            {"h2": (("yh", "xh"), final_h2)},
+        hrms_out = xr.Dataset(
+            {
+                f"{args.interpolated_field_name}": (("y", "x"), final_hrms),
+            },
             coords={
-                "yh": yh,
-                "xh": xh,
+                "lon": (("y", "x"), regrid_ocean_mask.lon.values),
+                "lat": (("y", "x"), regrid_ocean_mask.lat.values),
+                "lon_b": (("y_b", "x_b"), regrid_ocean_mask.lon_b.values),
+                "lat_b": (("y_b", "x_b"), regrid_ocean_mask.lat_b.values),
             },
             attrs={
-                "long_name": "Polynomial-fit roughness (h2) per model cell",
+                "long_name": "Polynomial-fit roughness per interpolated grid cell",
                 "units": "m",
             },
         )
+
         # Add metadata
         this_file = os.path.normpath(__file__)
         runcmd = (
@@ -325,28 +335,28 @@ def main():
             f"--topo-file={args.topo_file}"
             f"--chunk_lat={args.chunk_lat}"
             f"--chunk_lon={args.chunk_lon}"
-            f"--grid-file={args.grid_file}"
-            f"--mask-file={args.mask_file}"
+            f"--regrid-mask-file={args.regrid_mask_file}"
+            f"--method={args.method}"
+            f"--interpolated-field-name={args.interpolated_field_name}"
             f"--output={args.output}"
         )
 
-        try:
-            history = get_provenance_metadata(this_file, runcmd)
-        except subprocess.CalledProcessError:
-            history = "Provenance metadata unavailable (not a git repo?)"
+        history = get_provenance_metadata(this_file, runcmd)
+
         global_attrs = {"history": history}
+        
         # add md5 hashes for input files
         file_hashes = [
             f"{args.topo_file} (md5 hash: {md5sum(args.topo_file)})",
-            f"{args.grid_file} (md5 hash: {md5sum(args.grid_file)})",
-            f"{args.mask_file} (md5 hash: {md5sum(args.mask_file)})",
+            f"{args.regrid_mask_file} (md5 hash: {md5sum(args.regrid_mask_file)})",
         ]
         global_attrs["inputFile"] = ", ".join(file_hashes)
 
-        h2_out.attrs.update(global_attrs)
+        hrms_out.attrs.update(global_attrs)
 
-        h2_out.to_netcdf(args.output)
+        hrms_out.to_netcdf(args.output)
 
 
 if __name__ == "__main__":
     main()
+

From 9e590f59690b92710bcee73224549dde24e094b5 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 8 Apr 2025 13:48:27 +1000
Subject: [PATCH 07/16] Add regrid_ocean_mask_to_global_grid.py to do the esmf
 regridding from model grid to a normal global grid

---
 .../regrid_ocean_mask_to_global_grid.py       | 169 ++++++++++++++++++
 1 file changed, 169 insertions(+)
 create mode 100644 external_tidal_generation/regrid_ocean_mask_to_global_grid.py

diff --git a/external_tidal_generation/regrid_ocean_mask_to_global_grid.py b/external_tidal_generation/regrid_ocean_mask_to_global_grid.py
new file mode 100644
index 0000000..4a58e50
--- /dev/null
+++ b/external_tidal_generation/regrid_ocean_mask_to_global_grid.py
@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+# Copyright 2025 ACCESS-NRI and contributors.
+# SPDX-License-Identifier: Apache-2.0
+
+# =========================================================================================
+# Regrid ocean mask onto a regular grid.
+#
+# This script regrids an ocean mask from the model grid to a new, regular grid
+# using the xESMF regridding library. It loads an ocean mosaic grid (ocean_hgrid.nc),
+# and its associated ocean mask (ocean_mask.nc), then creates a target grid
+# at a user-specified resolution (--dx (deg) and --dy (deg)). The mask is subsequently
+# interpolated onto this new grid using a regridding method (e.g., conservative_normed).
+# The final regridded mask is saved as a netcdf file.
+#
+# Usage:
+# python3 regrid_mask.py \
+#   --hgrid /path/to/ocean_hgrid.nc \
+#   --mask /path/to/ocean_mask.nc \
+#   --dx 0.25 \
+#   --dy 0.25 \
+#   --method conservative_normed \
+#   --output regridded_output.nc
+#
+# Contact:
+#    - Minghang Li <Minghang.Li1@anu.edu.au>
+#
+# Dependencies:
+#   - xarray
+#   - numpy
+#   - xesmf
+# =========================================================================================
+import os
+import sys
+from pathlib import Path
+import argparse
+import xarray as xr
+import xesmf as xe
+
+path_root = Path(__file__).parents[1]
+sys.path.append(str(path_root))
+
+from scripts_common import get_provenance_metadata, md5sum
+
+
+def load_dataset(path: str) -> xr.Dataset:
+    """
+    Load an input dataset from a netcdf file.
+    """
+    ds = xr.open_dataset(path)
+    return ds
+
+
+def create_mask_ds(
+    ocean_mask: xr.Dataset,
+    hgrid_x: xr.DataArray,
+    hgrid_y: xr.DataArray,
+    hgrid_xc: xr.DataArray,
+    hgrid_yc: xr.DataArray,
+) -> xr.Dataset:
+    """
+    Loads the ocean mask and converts it into a Dataset formatted
+    for use in xesmf regridding.
+    `data` is a dummy variable that contains the mask values.
+    """
+    ocean_mask_ds = xr.Dataset(
+        data_vars={
+            "data": (("y", "x"), ocean_mask.mask.values),
+            "mask": (("y", "x"), ocean_mask.mask.values),
+        },
+        coords={
+            "lon": (("y", "x"), hgrid_x.values),
+            "lat": (("y", "x"), hgrid_y.values),
+            "lon_b": (("y_b", "x_b"), hgrid_xc.values),
+            "lat_b": (("y_b", "x_b"), hgrid_yc.values),
+        },
+    )
+
+    return ocean_mask_ds
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Regrid ocean mask onto a regular target grid."
+    )
+    parser.add_argument(
+        "--hgrid", type=str, required=True, help="Path to ocean_hgrid.nc"
+    )
+    parser.add_argument("--mask", type=str, required=True, help="Path to ocean_mask.nc")
+    parser.add_argument(
+        "--dx",
+        type=float,
+        default=0.25,
+        help="Target grid resolution in longitude (degrees)",
+    )
+    parser.add_argument(
+        "--dy",
+        type=float,
+        default=0.25,
+        help="Target grid resolution in latitude (degrees)",
+    )
+    parser.add_argument(
+        "--method",
+        type=str,
+        default="conservative_normed",
+        help="Regridding method (e.g., bilinear, conservative, conservative_normed)",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="regridded_mask.nc",
+        help="Output netcdf file name",
+    )
+
+    args = parser.parse_args()
+
+    # Load model grid
+    hgrid = load_dataset(args.hgrid)
+    hgrid_x = hgrid.x[1::2, 1::2]
+    hgrid_y = hgrid.y[1::2, 1::2]
+    hgrid_xc = hgrid.x[::2, ::2]
+    hgrid_yc = hgrid.y[::2, ::2]
+
+    # Load ocean mask
+    ocean_mask = load_dataset(args.mask)
+
+    ocean_mask_ds = create_mask_ds(ocean_mask, hgrid_x, hgrid_y, hgrid_xc, hgrid_yc)
+
+    # Create a regular global grid for regridding
+    ds_mask_out = xe.util.grid_global(args.dx, args.dy)
+
+    # Create regridder and regrid
+    regridder = xe.Regridder(
+        ocean_mask_ds, ds_mask_out, method=args.method, periodic=True
+    )
+    ds_mask_out[args.method] = regridder(ocean_mask_ds["data"])
+
+    ds_mask_out.attrs["description"] = (
+        f"Ocean mask regridded onto a target grid ({args.dy}, {args.dx}) using xesmf"
+    )
+
+    this_file = os.path.normpath(__file__)
+    runcmd = (
+        f"mpirun -n $PBS_NCPUS"
+        f"python3 {os.path.basename(this_file)} "
+        f"--hgrid={args.hgrid}"
+        f"--mask={args.mask}"
+        f"--dx={args.dx}"
+        f"--dy={args.dy}"
+        f"--method={args.method}"
+        f"--output={args.output}"
+    )
+
+    history = get_provenance_metadata(this_file, runcmd)
+
+    global_attrs = {"history": history}
+
+    # add md5 hashes for input files
+    file_hashes = [
+        f"{args.hgrid} (md5 hash: {md5sum(args.hgrid)})",
+        f"{args.mask} (md5 hash: {md5sum(args.mask)})",
+    ]
+    global_attrs["inputFile"] = ", ".join(file_hashes)
+
+    ds_mask_out.attrs.update(global_attrs)
+    ds_mask_out.to_netcdf(args.output)
+
+
+if __name__ == "__main__":
+    main()

From f54cfa2104d1d38fcfb3a591e36481a18000c3ba Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 8 Apr 2025 13:49:02 +1000
Subject: [PATCH 08/16] Add regrid_interped_bottom_roughness_to_model_grid.py

---
 ...interped_bottom_roughness_to_model_grid.py | 151 ++++++++++++++++++
 1 file changed, 151 insertions(+)
 create mode 100644 external_tidal_generation/regrid_interped_bottom_roughness_to_model_grid.py

diff --git a/external_tidal_generation/regrid_interped_bottom_roughness_to_model_grid.py b/external_tidal_generation/regrid_interped_bottom_roughness_to_model_grid.py
new file mode 100644
index 0000000..7a4d9f5
--- /dev/null
+++ b/external_tidal_generation/regrid_interped_bottom_roughness_to_model_grid.py
@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+# Copyright 2025 ACCESS-NRI and contributors.
+# SPDX-License-Identifier: Apache-2.0
+
+# =========================================================================================
+# Regrid interpolated bottom roughness onto the ocean model grid
+#
+# This script regrids an interpolated bottom roughness field onto the ocean
+# model grid using the xesmf library. It loads the ocean mosaic grid (ocean_hgrid.nc),
+# ocean mask (ocean_mask.nc), and the interpolated bottom roughness field (interpolated_h2.nc).
+#
+# Usage:
+#    python3 regrid_bottom_roughness_to_model_grid.py \
+#         --hgrid /path/to/ocean_hgrid.nc \
+#         --mask /path/to/ocean_mask.nc \
+#         --interpolated-file /path/to/interpolated_h2.nc \
+#         --method conservative_normed \
+#         --output h2.nc
+#
+# Contact:
+#    - Minghang Li <Minghang.Li1@anu.edu.au>
+#
+# Dependencies:
+#   - xarray
+#   - numpy
+#   - xesmf
+# =========================================================================================
+import os
+import sys
+from pathlib import Path
+import argparse
+import xarray as xr
+import xesmf as xe
+
+path_root = Path(__file__).parents[1]
+sys.path.append(str(path_root))
+
+from scripts_common import get_provenance_metadata, md5sum
+
+
+def load_dataset(path: str) -> xr.Dataset:
+    """
+    Load an input dataset from a netcdf file.
+    """
+    ds = xr.open_dataset(path)
+    return ds
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Regrid bottom roughness onto a target ocean model grid."
+    )
+    parser.add_argument(
+        "--hgrid", type=str, required=True, help="Path to ocean mosaic ocean_hgrid.nc."
+    )
+    parser.add_argument(
+        "--mask",
+        type=str,
+        required=True,
+        help="Path to ocean_mask.nc.",
+    )
+    parser.add_argument(
+        "--interpolated-file",
+        type=str,
+        required=True,
+        help="Path to an interpolated field, such as interpolated_h2.nc.",
+    )
+    parser.add_argument(
+        "--method",
+        type=str,
+        default="conservative_normed",
+        help="Regridding method (e.g., bilinear, conservative, conservative_normed)",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="regridded_output.nc",
+        help="Output netcdf file name",
+    )
+    args = parser.parse_args()
+
+    # Load model grid
+    hgrid = load_dataset(args.hgrid)
+    hgrid_x = hgrid.x[1::2, 1::2]
+    hgrid_y = hgrid.y[1::2, 1::2]
+    hgrid_xc = hgrid.x[::2, ::2]
+    hgrid_yc = hgrid.y[::2, ::2]
+
+    # Load ocean mask
+    ocean_mask = load_dataset(args.mask)
+
+    # Load interpolated bottom roughness
+    interpolated_h2 = load_dataset(args.interpolated_file)
+
+    # Build the target field on the model grid for regridding with mask information
+    target_ds = xr.Dataset(
+        data_vars={"mask": (("y", "x"), ocean_mask.mask.values)},
+        coords={
+            "lon": (("y", "x"), hgrid_x.values),
+            "lat": (("y", "x"), hgrid_y.values),
+            "lon_b": (("y_b", "x_b"), hgrid_xc.values),
+            "lat_b": (("y_b", "x_b"), hgrid_yc.values),
+        },
+    )
+
+    # Create regridder and regrid
+    regridder = xe.Regridder(
+        interpolated_h2, target_ds, method=args.method, periodic=True
+    )
+
+    # Apply the regridder to the interped bottom roughness
+    target_ds["h"] = regridder(interpolated_h2[args.regrid_field_name])
+
+    target_ds = target_ds["h"].drop_vars(["lon_b", "lat_b", "mask"])
+
+    h2 = xr.Dataset({"h2": target_ds["h"] ** 2})
+
+    h2.attrs["description"] = (
+        "Squared bottom roughness amplitude computed using a polynomial fitting method "
+        "following Jayne & Laurent (2001)."
+    )
+
+    this_file = os.path.normpath(__file__)
+    runcmd = (
+        f"mpirun -n $PBS_NCPUS"
+        f"python3 {os.path.basename(this_file)} "
+        f"--hgrid={args.hgrid}"
+        f"--mask={args.mask}"
+        f"--interpolated-file={args.interpolated_file}"
+        f"--method={args.method}"
+        f"--output={args.output}"
+    )
+
+    history = get_provenance_metadata(this_file, runcmd)
+
+    global_attrs = {"history": history}
+
+    # add md5 hashes for input files
+    file_hashes = [
+        f"{args.hgrid} (md5 hash: {md5sum(args.hgrid)})",
+        f"{args.mask} (md5 hash: {md5sum(args.mask)})",
+        f"{args.interpolated_file} (md5 hash: {md5sum(args.interpolated_file)})",
+    ]
+    global_attrs["inputFile"] = ", ".join(file_hashes)
+
+    h2.attrs.update(global_attrs)
+    h2.to_netcdf(args.output)
+
+
+if __name__ == "__main__":
+    main()

From 1465f09d98df25acd3fcb849ed4c0b2c584349e8 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 8 Apr 2025 13:51:40 +1000
Subject: [PATCH 09/16] tidy

---
 .../generate_bottom_roughness_polyfit.py                 | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/external_tidal_generation/generate_bottom_roughness_polyfit.py b/external_tidal_generation/generate_bottom_roughness_polyfit.py
index 75a0950..4906d62 100644
--- a/external_tidal_generation/generate_bottom_roughness_polyfit.py
+++ b/external_tidal_generation/generate_bottom_roughness_polyfit.py
@@ -14,7 +14,7 @@
 # Usage:
 #    mpirun -n <ranks> python3 generate_bottom_roughness_polyfit.py \
 #         --topo-file /path/to/topog.nc \
-#         --hgrid-file /path/to/ocean_static.nc \
+#         --hgrid-file /path/to/ocean_hgrid.nc \
 #         --regrid-mask-file /path/to/ocean_mask.nc \
 #         --output output.nc
 #
@@ -177,7 +177,9 @@ def evaluate_roughness(
             lat_min = np.min(this_lat_corners)
             lat_max = np.max(this_lat_corners)
 
-            hrms_val = compute_hrms_poly_cell(lon_min, lon_max, lat_min, lat_max, topo_da)
+            hrms_val = compute_hrms_poly_cell(
+                lon_min, lon_max, lat_min, lat_max, topo_da
+            )
             local_hrms[j - y_start, i] = hrms_val
 
         if (j - y_start) % 3 == 0:
@@ -344,7 +346,7 @@ def main():
         history = get_provenance_metadata(this_file, runcmd)
 
         global_attrs = {"history": history}
-        
+
         # add md5 hashes for input files
         file_hashes = [
             f"{args.topo_file} (md5 hash: {md5sum(args.topo_file)})",
@@ -359,4 +361,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-

From 5319eb6830679e65ae0bd020c1d65b1cac82c34c Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Thu, 10 Apr 2025 10:18:19 +1000
Subject: [PATCH 10/16] More tidy up

---
 .../generate_bottom_roughness_polyfit.py      | 26 +++++++++----------
 ...polated_bottom_roughness_to_model_grid.py} | 10 ++++---
 2 files changed, 19 insertions(+), 17 deletions(-)
 rename external_tidal_generation/{regrid_interped_bottom_roughness_to_model_grid.py => regrid_interpolated_bottom_roughness_to_model_grid.py} (93%)

diff --git a/external_tidal_generation/generate_bottom_roughness_polyfit.py b/external_tidal_generation/generate_bottom_roughness_polyfit.py
index 4906d62..b2f8c7f 100644
--- a/external_tidal_generation/generate_bottom_roughness_polyfit.py
+++ b/external_tidal_generation/generate_bottom_roughness_polyfit.py
@@ -34,8 +34,6 @@
 import numpy as np
 import xarray as xr
 
-import subprocess
-
 path_root = Path(__file__).parents[1]
 sys.path.append(str(path_root))
 
@@ -233,6 +231,18 @@ def main():
         required=True,
         help="Path to a high-resolution topography file.",
     )
+    parser.add_argument(
+        "--chunk-lat",
+        type=int,
+        default=800,
+        help="Dask chunk size along lat dimension (default:800).",
+    )
+    parser.add_argument(
+        "--chunk-lon",
+        type=int,
+        default=1600,
+        help="Dask chunk size along lon dimension (default:1600).",
+    )
     parser.add_argument(
         "--regrid-mask-file",
         type=str,
@@ -257,18 +267,6 @@ def main():
         default="interpolated_hrms.nc",
         help="Output roughness filename (default: interpolated_hrms.nc).",
     )
-    parser.add_argument(
-        "--chunk-lat",
-        type=int,
-        default=800,
-        help="Dask chunk size along lat dimension (default:800).",
-    )
-    parser.add_argument(
-        "--chunk-lon",
-        type=int,
-        default=1600,
-        help="Dask chunk size along lon dimension (default:1600).",
-    )
     args = parser.parse_args()
 
     # create communicator
diff --git a/external_tidal_generation/regrid_interped_bottom_roughness_to_model_grid.py b/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
similarity index 93%
rename from external_tidal_generation/regrid_interped_bottom_roughness_to_model_grid.py
rename to external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
index 7a4d9f5..33d379d 100644
--- a/external_tidal_generation/regrid_interped_bottom_roughness_to_model_grid.py
+++ b/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
@@ -65,6 +65,12 @@ def main():
         required=True,
         help="Path to an interpolated field, such as interpolated_h2.nc.",
     )
+    parser.add_argument(
+        "--interpolated-field-name",
+        type=str,
+        default="hrms_interpolated",
+        help="The field name for the interpolated roughness (default: hrms_interpolated).",
+    )
     parser.add_argument(
         "--method",
         type=str,
@@ -109,9 +115,7 @@ def main():
     )
 
     # Apply the regridder to the interped bottom roughness
-    target_ds["h"] = regridder(interpolated_h2[args.regrid_field_name])
-
-    target_ds = target_ds["h"].drop_vars(["lon_b", "lat_b", "mask"])
+    target_ds["h"] = regridder(interpolated_h2[args.interpolated_field_name])
 
     h2 = xr.Dataset({"h2": target_ds["h"] ** 2})
 

From cbfbcd3e3a413fb45bd4a9f0826ee2ef9b8fb80b Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Wed, 16 Apr 2025 13:46:19 +1000
Subject: [PATCH 11/16] Refactor a bit and use distances in meters for the
 calculation

---
 .../generate_bottom_roughness_polyfit.py      | 264 +++++++++++++-----
 1 file changed, 197 insertions(+), 67 deletions(-)

diff --git a/external_tidal_generation/generate_bottom_roughness_polyfit.py b/external_tidal_generation/generate_bottom_roughness_polyfit.py
index b2f8c7f..ffb573c 100644
--- a/external_tidal_generation/generate_bottom_roughness_polyfit.py
+++ b/external_tidal_generation/generate_bottom_roughness_polyfit.py
@@ -40,6 +40,44 @@
 from scripts_common import get_provenance_metadata, md5sum
 
 
+def load_topo(path: str, chunk_lat: int = 800, chunk_lon: int = 1600) -> xr.DataArray:
+    """
+    Load a high-resolution bathymetry file
+    """
+    ds = xr.open_dataset(path, chunks={"lat": chunk_lat, "lon": chunk_lon})
+    depth = ds["z"].where(ds["z"] < 0, np.nan)
+    return depth
+
+
+def project_lon(
+    lon: float,
+) -> float:
+    """
+    Project longitude to the range [-280, 80].
+    """
+
+    return ((lon + 280) % 360) - 280
+
+
+def align_lon_coords(da: xr.DataArray) -> xr.DataArray:
+    """
+    Align high resolution topography lon coord to the range [-280, 80].
+    """
+
+    da = da.assign_coords(lon=project_lon(da.lon))
+    da = da.sortby("lon")
+
+    return da
+
+
+def load_dataset(path: str) -> xr.Dataset:
+    """
+    Load an input dataset.
+    """
+    ds = xr.open_dataset(path)
+    return ds
+
+
 def polyfit_roughness(H: np.ndarray, xv: np.ndarray, yv: np.ndarray) -> float:
     """
     Fit a polynomial of the form:
@@ -99,22 +137,16 @@ def compute_hrms_poly_cell(
     x = sub_da.lon.values
     y = sub_da.lat.values
     xv, yv = np.meshgrid(x, y)
-    hrms = polyfit_roughness(H, xv, yv)
-    return hrms
-
-
-def load_topo(path: str, chunk_lat: int = 800, chunk_lon: int = 1600) -> xr.DataArray:
-    """
-    Load a high-resolution bathymetry file
-    """
-    ds = xr.open_dataset(path, chunks={"lat": chunk_lat, "lon": chunk_lon})
-    depth = ds["z"].where(ds["z"] < 0, np.nan)
-    return depth
 
+    R = 6.371229e6  # earth radius in meters
+    deg2rad = np.pi / 180
+    lat0 = np.mean(y)  # a reference lat
+    lon0 = np.mean(x)  # a reference lon
 
-def load_regrid_ocean_mask(path: str) -> xr.Dataset:
-    interp_ocean_mask = xr.open_dataset(path)
-    return interp_ocean_mask
+    xlon = (xv - lon0) * np.cos(yv * deg2rad) * deg2rad * R
+    ylat = (yv - lat0) * deg2rad * R
+    hrms = polyfit_roughness(H, xlon, ylat)
+    return hrms
 
 
 def evaluate_roughness(
@@ -244,10 +276,13 @@ def main():
         help="Dask chunk size along lon dimension (default:1600).",
     )
     parser.add_argument(
-        "--regrid-mask-file",
+        "--mask-file",
         type=str,
         required=True,
-        help="Path to the regrid ocean mask file.",
+        help="Path to the ocean mask file.",
+    )
+    parser.add_argument(
+        "--hgrid", type=str, required=True, help="Path to ocean_hgrid.nc"
     )
     parser.add_argument(
         "--method",
@@ -256,10 +291,10 @@ def main():
         help="Regridding method (e.g., bilinear, conservative, conservative_normed)",
     )
     parser.add_argument(
-        "--interpolated-field-name",
-        type=str,
-        default="hrms_interpolated",
-        help="The field name for the interpolated roughness (default: hrms_interpolated).",
+        "--agg-factor",
+        type=int,
+        default=1,
+        help="Coare factor. Eg, 1 for original grid (eg 0.25deg); 2 for 0.5deg or 4 for 1deg resolution (default: 1 means original model grid).",
     )
     parser.add_argument(
         "--output",
@@ -274,36 +309,127 @@ def main():
     rank = comm.Get_rank()
 
     if rank == 0:
+        # Load high-resolution bathymetry
         topo_da = load_topo(
             args.topo_file, chunk_lat=args.chunk_lat, chunk_lon=args.chunk_lon
         )
-        regrid_ocean_mask = load_regrid_ocean_mask(args.regrid_mask_file)
-        lon_b = regrid_ocean_mask.lon_b.values
-        lat_b = regrid_ocean_mask.lat_b.values
-        ny = lon_b.shape[0] - 1
-        nx = lon_b.shape[1] - 1
+        ocean_mask = load_dataset(args.mask_file)
+        if "mask" not in ocean_mask:
+            raise KeyError("Missing variable 'mask' in ocean_mask file!")
+
+        # Convert lon coords to a [-280, 80] to match the range of the model grid
+        topo_da = align_lon_coords(topo_da)
+
+        # Load hgrid
+        hgrid = load_dataset(args.hgrid)
+        hgrid_x = hgrid.x[1::2, 1::2]
+        hgrid_y = hgrid.y[1::2, 1::2]
+        hgrid_xc = hgrid.x[::2, ::2]
+        hgrid_yc = hgrid.y[::2, ::2]
+
+        # Load model ocean mask and tweak coords
+        ocean_mask = ocean_mask.drop_vars(["geolon_t", "geolat_t"])
+        ocean_mask = ocean_mask.rename({"ny": "lat", "nx": "lon"})
+
+        ocean_mask = ocean_mask.assign_coords(
+            {
+                "lon": (("lat", "lon"), hgrid_x.values),
+                "lat": (("lat", "lon"), hgrid_y.values),
+                "lon_b": (("lat_b", "lon_b"), hgrid_xc.values),
+                "lat_b": (("lat_b", "lon_b"), hgrid_yc.values),
+            }
+        )
+
+        # check lon range for high resolution topo
+        print(topo_da.lon.values)
+
+        # Get the full boundary values of the ocean mask
+        lon_b_full = ocean_mask.lon_b.values
+        lat_b_full = ocean_mask.lat_b.values
+        ocean_mask_full = ocean_mask["mask"].values
+
+        # Get the dims of the full ocean mask
+        ny_full = lon_b_full.shape[0] - 1
+        nx_full = lon_b_full.shape[1] - 1
+
+        fac = args.agg_factor
+        if fac > 1:
+            # Compute number of coarse cells
+            nC_y = ny_full // fac
+            nC_x = nx_full // fac
+
+            # Crop boundaries for cell edges
+            ny_edge = nC_y * fac + 1
+            nx_edge = nC_x * fac + 1
+            lon_b_crop = lon_b_full[:ny_edge, :nx_edge]
+            lat_b_crop = lat_b_full[:ny_edge, :nx_edge]
+
+            # Coarse using striding
+            lon_b_coarse = lon_b_crop[::fac, ::fac]
+            lat_b_coarse = lat_b_crop[::fac, ::fac]
+
+            # Compute coarse cell centers by averaging the four surrounding boundaries.
+            lon_coarse = 0.25 * (
+                lon_b_coarse[:-1, :-1]
+                + lon_b_coarse[:-1, 1:]
+                + lon_b_coarse[1:, :-1]
+                + lon_b_coarse[1:, 1:]
+            )
+            lat_coarse = 0.25 * (
+                lat_b_coarse[:-1, :-1]
+                + lat_b_coarse[:-1, 1:]
+                + lat_b_coarse[1:, :-1]
+                + lat_b_coarse[1:, 1:]
+            )
+
+            # Crop ocean_mask at cell center
+            ny_crop = (ny_full // fac) * fac
+            nx_crop = (nx_full // fac) * fac
+            ocean_mask_cropped = ocean_mask_full[:ny_crop, :nx_crop]
+            # ocean_mask_coarse = ocean_mask_cropped.reshape(ny_crop // fac, fac,nx_crop // fac, fac).max(axis=(1, 3))
+            ocean_mask_coarse = ocean_mask_cropped[::fac, ::fac]
+
+            # Update dimensions
+            ny = ocean_mask_coarse.shape[0]
+            nx = ocean_mask_coarse.shape[1]
+
+        else:
+            # No coarsen needed
+            lon_b_coarse = lon_b_full
+            lat_b_coarse = lat_b_full
+            ocean_mask_coarse = ocean_mask_full
+            ny = ny_full
+            nx = nx_full
+
+            lon_coarse = ocean_mask.lon.values
+            lat_coarse = ocean_mask.lat.values
+
+        print(f"[Rank 0] coarsen grid dimensions: {ny} cells in y, {nx} cells in x.")
 
     else:
         topo_da = None
-        regrid_ocean_mask = None
-        lon_b = None
-        lat_b = None
+        lon_b_coarse = None
+        lat_b_coarse = None
+        ocean_mask_coarse = None
+        lon_coarse = None
+        lat_coarse = None
         ny = None
         nx = None
 
-    # Broadcast shared datasets to all ranks
+    # Broadcast fields and grid sizes to all ranks.
     topo_da = comm.bcast(topo_da, root=0)
-    regrid_ocean_mask = comm.bcast(regrid_ocean_mask, root=0)
-    lon_b = comm.bcast(lon_b, root=0)
-    lat_b = comm.bcast(lat_b, root=0)
-    nx = comm.bcast(nx, root=0)
+    lon_b_coarse = comm.bcast(lon_b_coarse, root=0)
+    lat_b_coarse = comm.bcast(lat_b_coarse, root=0)
+    ocean_mask_coarse = comm.bcast(ocean_mask_coarse, root=0)
     ny = comm.bcast(ny, root=0)
+    nx = comm.bcast(nx, root=0)
 
+    # Evaluate HRMS on the grid.
     final_hrms = evaluate_roughness(
         topo_da=topo_da,
-        ocean_mask=regrid_ocean_mask[args.method],
-        hgrid_xc=lon_b,
-        hgrid_yc=lat_b,
+        ocean_mask=ocean_mask_coarse,
+        hgrid_xc=lon_b_coarse,
+        hgrid_yc=lat_b_coarse,
         nx=nx,
         ny=ny,
         comm=comm,
@@ -311,50 +437,54 @@ def main():
 
     # Rank 0 writes results to file
     if rank == 0:
-        hrms_out = xr.Dataset(
-            {
-                f"{args.interpolated_field_name}": (("y", "x"), final_hrms),
-            },
-            coords={
-                "lon": (("y", "x"), regrid_ocean_mask.lon.values),
-                "lat": (("y", "x"), regrid_ocean_mask.lat.values),
-                "lon_b": (("y_b", "x_b"), regrid_ocean_mask.lon_b.values),
-                "lat_b": (("y_b", "x_b"), regrid_ocean_mask.lat_b.values),
-            },
-            attrs={
-                "long_name": "Polynomial-fit roughness per interpolated grid cell",
-                "units": "m",
-            },
-        )
+        if args.agg_factor == 1:
+            # No regridding and no coarsen,
+            # hence output results on the original model grid
+            h2_out = xr.Dataset(
+                {
+                    "h2": (("y", "x"), final_hrms**2),
+                },
+                coords={
+                    "lon": (("y", "x"), lon_coarse),
+                    "lat": (("y", "x"), lat_coarse),
+                    "lon_b": (("y_b", "x_b"), lon_b_coarse),
+                    "lat_b": (("y_b", "x_b"), lat_b_coarse),
+                },
+                attrs={
+                    "long_name": (
+                        "Polynomial-fit roughness square h^2 per model grid cell "
+                        "following Jayne & Laurent (2001)"
+                    ),
+                    "units": "m",
+                },
+            )
+            h2_out["h2_0"] = h2_out["h2"].fillna(0.0)
 
-        # Add metadata
+        # Add provenance metadata and MD5 hashes for input files.
         this_file = os.path.normpath(__file__)
         runcmd = (
-            f"mpirun -n $PBS_NCPUS"
-            f"python3 {os.path.basename(this_file)} "
-            f"--topo-file={args.topo_file}"
-            f"--chunk_lat={args.chunk_lat}"
-            f"--chunk_lon={args.chunk_lon}"
-            f"--regrid-mask-file={args.regrid_mask_file}"
-            f"--method={args.method}"
-            f"--interpolated-field-name={args.interpolated_field_name}"
+            f"mpirun -n $PBS_NCPUS python3 {os.path.basename(this_file)} "
+            f"--topo-file={args.topo_file} "
+            f"--hgrid={args.hgrid} "
+            f"--chunk-lat={args.chunk_lat} "
+            f"--chunk-lon={args.chunk_lon} "
+            f"--mask-file={args.mask_file} "
+            f"--agg-factor={args.agg_factor} "
             f"--output={args.output}"
         )
 
         history = get_provenance_metadata(this_file, runcmd)
-
         global_attrs = {"history": history}
-
-        # add md5 hashes for input files
         file_hashes = [
             f"{args.topo_file} (md5 hash: {md5sum(args.topo_file)})",
-            f"{args.regrid_mask_file} (md5 hash: {md5sum(args.regrid_mask_file)})",
+            f"{args.hgrid} (md5 hash: {md5sum(args.hgrid)})",
+            f"{args.mask_file} (md5 hash: {md5sum(args.mask_file)})",
         ]
         global_attrs["inputFile"] = ", ".join(file_hashes)
+        h2_out.attrs.update(global_attrs)
 
-        hrms_out.attrs.update(global_attrs)
-
-        hrms_out.to_netcdf(args.output)
+        h2_out.to_netcdf(args.output)
+        print(f"[Rank 0] HRMS output written to: {args.output}")
 
 
 if __name__ == "__main__":

From 3b603a861b24e4ed807b230ba7e9f5e0cf04d8b2 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Wed, 16 Apr 2025 13:48:22 +1000
Subject: [PATCH 12/16] typo

---
 external_tidal_generation/generate_bottom_roughness_polyfit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/external_tidal_generation/generate_bottom_roughness_polyfit.py b/external_tidal_generation/generate_bottom_roughness_polyfit.py
index ffb573c..7151b71 100644
--- a/external_tidal_generation/generate_bottom_roughness_polyfit.py
+++ b/external_tidal_generation/generate_bottom_roughness_polyfit.py
@@ -294,7 +294,7 @@ def main():
         "--agg-factor",
         type=int,
         default=1,
-        help="Coare factor. Eg, 1 for original grid (eg 0.25deg); 2 for 0.5deg or 4 for 1deg resolution (default: 1 means original model grid).",
+        help="Coarse factor. Eg, 1 for original grid (eg 0.25deg); 2 for 0.5deg or 4 for 1deg resolution (default: 1 means original model grid).",
     )
     parser.add_argument(
         "--output",

From c12c2e8d49e4c210d4e547a5357501689969fde7 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 29 Apr 2025 11:34:46 +1000
Subject: [PATCH 13/16] Add h2_0 by filling 0.0 with nan

---
 .../regrid_interpolated_bottom_roughness_to_model_grid.py        | 1 +
 1 file changed, 1 insertion(+)

diff --git a/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py b/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
index 33d379d..e515e3b 100644
--- a/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
+++ b/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
@@ -118,6 +118,7 @@ def main():
     target_ds["h"] = regridder(interpolated_h2[args.interpolated_field_name])
 
     h2 = xr.Dataset({"h2": target_ds["h"] ** 2})
+    h2['h2_0'] = h2['h2'].fillna(0.)
 
     h2.attrs["description"] = (
         "Squared bottom roughness amplitude computed using a polynomial fitting method "

From 7e9f70da691cec25cee13027d9bef18f8e502a5d Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 29 Apr 2025 11:35:26 +1000
Subject: [PATCH 14/16] Add generate_tide_amplitude.py script

---
 .../generate_tide_amplitude.py                | 264 ++++++++++++++++++
 1 file changed, 264 insertions(+)
 create mode 100644 external_tidal_generation/generate_tide_amplitude.py

diff --git a/external_tidal_generation/generate_tide_amplitude.py b/external_tidal_generation/generate_tide_amplitude.py
new file mode 100644
index 0000000..b40a90f
--- /dev/null
+++ b/external_tidal_generation/generate_tide_amplitude.py
@@ -0,0 +1,264 @@
+#!/usr/bin/env python3
+# Copyright 2025 ACCESS-NRI and contributors.
+# SPDX-License-Identifier: Apache-2.0
+
+# =========================================================================================
+# Compute the barotropic tidal amplitude from TPXO10-atlas v2
+# using eight primary constituents (M2, S2, N2, K2, K1, O1, P1 and Q1)
+#
+# Usage:
+#    python3 generate_tide_amplitude.py \
+#       --hgrid /path/to/ocean_hgrid.nc \
+#       --mask /path/to/ocean_mask.nc \
+#       --data-path /path/to/TPXO10/ \
+#       --method  conservative_normed \
+#       --output  tideamp.nc
+#
+# Contact:
+#    - Minghang Li <Minghang.Li1@anu.edu.au>
+#
+# Dependencies:
+#   - xarray
+#   - numpy
+#   - xesmf
+# =========================================================================================
+from pathlib import Path
+import sys
+import os
+import argparse
+import numpy as np
+import xarray as xr
+import xesmf as xe
+
+path_root = Path(__file__).parents[1]
+sys.path.append(str(path_root))
+
+from scripts_common import get_provenance_metadata, md5sum
+
+
+PRIMARY_CONSTITUENTS = ["m2", "s2", "n2", "k2", "k1", "o1", "p1", "q1"]
+
+
+def load_dataset(path: Path) -> xr.Dataset:
+    """
+    Load an input dataset from a netcdf file.
+    """
+    return xr.open_dataset(path)
+
+
+def interp_complex(arr, lon: xr.DataArray, lat: xr.DataArray) -> xr.DataArray:
+    """
+    Linear interpolation for complex DataArrays.
+    """
+    re = arr.real.interp(x=lon, y=lat, method="linear")
+    im = arr.imag.interp(x=lon, y=lat, method="linear")
+    return (re + 1j * im).astype("complex64")
+
+
+def compute_constituent_speed(
+    trans_file: Path,
+    hu: xr.DataArray,
+    hv: xr.DataArray,
+    lon_z: xr.DataArray,
+    lat_z: xr.DataArray,
+    lon_u: xr.DataArray,
+    lat_v: xr.DataArray,
+) -> xr.DataArray:
+    """
+    Barotropic speed amplitude (m/s) on centre points
+    for a single tidal constituent.
+    """
+    ds = load_dataset(trans_file)
+
+    # complex transports from cm^2/s to m^2/s
+    Uc = (ds.uRe + 1j * ds.uIm).astype("complex64") * 1e-4
+    Vc = (ds.vRe + 1j * ds.vIm).astype("complex64") * 1e-4
+
+    # unify dimensions to (y,x)
+    dim_map = {"ny": "y", "nx": "x"}
+    Uc = Uc.rename(dim_map)
+    Vc = Vc.rename(dim_map)
+    hu = hu.rename(dim_map)
+    hv = hv.rename(dim_map)
+    lon_z = lon_z.rename({"nx": "x"})
+    lat_z = lat_z.rename({"ny": "y"})
+    lon_u = lon_u.rename({"nx": "x"})
+    lat_v = lat_v.rename({"ny": "y"})
+
+    # compute velocities by transport / depth
+    u_vel = (Uc / hu).assign_coords({"x": lon_u, "y": lat_z})
+    v_vel = (Vc / hv).assign_coords({"x": lon_z, "y": lat_v})
+
+    # build the 2-D center grids in (y,x)
+    lat_c, lon_c = xr.broadcast(lat_z[:-1], lon_z[:-1])
+
+    # interpolate real/imag and calculate speed
+    u_c = interp_complex(u_vel, lon_c, lat_c)
+    v_c = interp_complex(v_vel, lon_c, lat_c)
+    speed_vals = np.sqrt(np.abs(u_c) ** 2 + np.abs(v_c) ** 2)
+
+    # construct da
+    speed = xr.DataArray(
+        speed_vals,
+        dims=("y", "x"),
+        coords={
+            "lat": (("y", "x"), lat_c.values),
+            "lon": (("y", "x"), lon_c.values),
+        },
+        name="speed",
+    )
+
+    return speed
+
+
+def compute_tideamp(data_path: Path) -> xr.DataArray:
+    """
+    Compute tidal amplitude (speed) from eight primary constituents
+    """
+    # TPXO10 grids
+    grid = load_dataset(data_path / "grid_tpxo10atlas_v2.nc")
+    # depth at u and v points
+    hu = grid["hu"].where(grid["hu"] > 0)
+    hv = grid["hv"].where(grid["hv"] > 0)
+
+    speed_sq_sum = None
+
+    for cons in PRIMARY_CONSTITUENTS:
+        fname = data_path / f"u_{cons}_tpxo10_atlas_30_v2.nc"
+        print(f"Processing {fname}...")
+        speed = compute_constituent_speed(
+            fname, hu, hv, grid.lon_z, grid.lat_z, grid.lon_u, grid.lat_v
+        )
+        speed_sq_sum = speed**2 if speed_sq_sum is None else speed_sq_sum + speed**2
+
+    # rms over all 8 constituents
+    da = np.sqrt(speed_sq_sum).rename("tideamp")
+    ds = da.to_dataset()
+
+    # boundary grids for xesmf interpolation
+    lon_ub = grid.lon_u.rename({"nx": "x"})
+    lat_vb = grid.lat_v.rename({"ny": "y"})
+    lat_b2d, lon_b2d = xr.broadcast(lat_vb, lon_ub)
+
+    ds = ds.assign_coords(
+        {
+            "lon_b": (("y_b", "x_b"), lon_b2d.values),
+            "lat_b": (("y_b", "x_b"), lat_b2d.values),
+        }
+    )
+
+    ds.tideamp.attrs.update(
+        units="m/s",
+        long_name="Barotropic tidal current speed (8 constituents)",
+        source="TPXO10-atlas v2",
+        note="RMS of M2, S2, N2, K2, K1, O1, P1, Q1",
+    )
+
+    return ds
+
+
+def regrid(
+    hgrid_path: Path,
+    mask_path: Path,
+    tideamp: xr.Dataset,
+    method: str = "conservative_normed",
+) -> xr.Dataset:
+    """
+    Regrid tideamp onto the model grid using xesmf
+    """
+    # Load model grid
+    hgrid = load_dataset(hgrid_path)
+    hgrid_x = hgrid.x[1::2, 1::2]
+    hgrid_y = hgrid.y[1::2, 1::2]
+    hgrid_xc = hgrid.x[::2, ::2]
+    hgrid_yc = hgrid.y[::2, ::2]
+
+    # load model ocean mask
+    ocean_mask = load_dataset(mask_path)
+
+    # construct target dataset for tideamp
+    target_ds = xr.Dataset(
+        data_vars={"mask": (("y", "x"), ocean_mask.mask.values)},
+        coords={
+            "lon": (("y", "x"), hgrid_x.values),
+            "lat": (("y", "x"), hgrid_y.values),
+            "lon_b": (("y_b", "x_b"), hgrid_xc.values),
+            "lat_b": (("y_b", "x_b"), hgrid_yc.values),
+        },
+    )
+
+    regridder = xe.Regridder(tideamp, target_ds, method=method, periodic=True)
+
+    target_ds["tideamp"] = regridder(tideamp["tideamp"])
+
+    target_ds["tideamp"] = target_ds["tideamp"].fillna(0.0)
+
+    return target_ds
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compute tidal amplitude from TPXO10-atlas v2 and regrid onto a model grid."
+    )
+    parser.add_argument(
+        "--hgrid", type=str, required=True, help="Path to ocean mosaic ocean_hgrid.nc."
+    )
+    parser.add_argument(
+        "--mask",
+        type=str,
+        required=True,
+        help="Path to ocean_mask.nc.",
+    )
+    parser.add_argument(
+        "--method",
+        type=str,
+        default="conservative_normed",
+        help="Regridding method (e.g., bilinear, conservative, conservative_normed)",
+    )
+    parser.add_argument(
+        "--data-path",
+        type=str,
+        required=True,
+        help="Directory containing TPXO10 files.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        required=True,
+        help="Output netcdf file name.",
+    )
+
+    args = parser.parse_args()
+
+    tideamp_tmp = compute_tideamp(Path(args.data_path))
+
+    tideamp = regrid(Path(args.hgrid), Path(args.mask), tideamp_tmp, args.method)
+
+    # Add provenance metadata and MD5 hashes for input files.
+    this_file = os.path.normpath(__file__)
+    runcmd = (
+        f"python3 {os.path.basename(this_file)} "
+        f"--hgrid={args.hgrid} "
+        f"--mask={args.mask} "
+        f"--method={args.method} "
+        f"--data-path={args.data_path} "
+        f"--output={args.output} "
+    )
+
+    history = get_provenance_metadata(this_file, runcmd)
+    global_attrs = {"history": history}
+
+    # add md5 hashes for input files
+    file_hashes = [
+        f"{args.target_grid} (md5 hash: {md5sum(args.hgrid)})",
+        f"{args.mask} (md5 hash: {md5sum(args.mask)})",
+    ]
+    global_attrs["inputFile"] = ", ".join(file_hashes)
+    tideamp.attrs.update(global_attrs)
+
+    tideamp.to_netcdf(args.out)
+    print(f"Complete {args.out}!")
+
+
+if __name__ == "__main__":
+    main()

From c38ce44643194f3a284dbe5ba309fdf9e8e23265 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 29 Apr 2025 11:37:54 +1000
Subject: [PATCH 15/16] black
 regrid_interpolated_bottom_roughness_to_model_grid.py

---
 .../regrid_interpolated_bottom_roughness_to_model_grid.py       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py b/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
index e515e3b..b6e6cd4 100644
--- a/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
+++ b/external_tidal_generation/regrid_interpolated_bottom_roughness_to_model_grid.py
@@ -118,7 +118,7 @@ def main():
     target_ds["h"] = regridder(interpolated_h2[args.interpolated_field_name])
 
     h2 = xr.Dataset({"h2": target_ds["h"] ** 2})
-    h2['h2_0'] = h2['h2'].fillna(0.)
+    h2["h2_0"] = h2["h2"].fillna(0.0)
 
     h2.attrs["description"] = (
         "Squared bottom roughness amplitude computed using a polynomial fitting method "

From dde053d2dc18f1919e1b1b977826429ec90d8003 Mon Sep 17 00:00:00 2001
From: minghangli-uni <24727729+minghangli-uni@users.noreply.github.com>
Date: Tue, 29 Apr 2025 13:42:51 +1000
Subject: [PATCH 16/16] fix

---
 external_tidal_generation/generate_tide_amplitude.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/external_tidal_generation/generate_tide_amplitude.py b/external_tidal_generation/generate_tide_amplitude.py
index b40a90f..cc4fefc 100644
--- a/external_tidal_generation/generate_tide_amplitude.py
+++ b/external_tidal_generation/generate_tide_amplitude.py
@@ -250,14 +250,14 @@ def main():
 
     # add md5 hashes for input files
     file_hashes = [
-        f"{args.target_grid} (md5 hash: {md5sum(args.hgrid)})",
+        f"{args.hgrid} (md5 hash: {md5sum(args.hgrid)})",
         f"{args.mask} (md5 hash: {md5sum(args.mask)})",
     ]
     global_attrs["inputFile"] = ", ".join(file_hashes)
     tideamp.attrs.update(global_attrs)
 
-    tideamp.to_netcdf(args.out)
-    print(f"Complete {args.out}!")
+    tideamp.to_netcdf(args.output)
+    print(f"Complete {args.output}!")
 
 
 if __name__ == "__main__":