Download all files from this example.

Great Barrier Reef testcase

This page gathers the shared configuration for the GBR testcase. The full example is split across several scripts because the workflow is relatively large:

  1. params.py : transversal parameters

  2. download.py : downloads static datasets such as bathymetry, reefs, land polygons.

  3. mesh.py : builds the horizontal mesh.

  4. dl_data.py: downloads time-dependent forcings such as ocean currents, temperature, salinity, wind, heat fluxes, and precipitation.

  5. prepro.py : interpolates and converts those datasets to SLIM4 inputs.

  6. run.py : runs the 3D hydrodynamic simulation.

A final gbrtools.py contains the remaining function that don’t obviously belong to any part of the workflow

The parameters below define the common domain, time window, paths, vertical discretization, boundary tags, and time-stepping options used by the workflow.

# general
bounds = [137.0, 159.0, -32.0, -6.0]  # minlon, maxlon, minlat, maxlat
initial_time = "2024-01-01 00:00:00"
final_time = "2024-01-31 23:59:59"

# paths
nc_data_dir = "data/netcdf"
prepro_data_dir = "prepro"
output_directory = "out"
mesh_file = "data/gbr.msh"

# extrusion settings
min_depth = 1.0
n_sigma = 10
depth_sigma = 120.0
z_levels = 20
n_moving_max = 10

# boundary tags
open_tags = ["coral_sea", "torres_strait"]
closed_tags = ["coast"]

# temporal parameters
dt_mult = 3
check_dt = 24 * 3600
export_dt = 3600
ramp = 86400
restart = False

Download this section.

Download static data

This is the first step of the Great Barrier Reef example workflow. It gathers the static data used by the following scripts: coral-reef polygons, coastline and land polygons, bathymetric rasters, and a small TPXO-like tidal dataset with the expected file layout.

The tidal files created here contain zero amplitudes and phases. They are only meant to document and test the expected TPXO format; real tidal forcing can be used by replacing these files with data downloaded from TPXO.

import os
import urllib.request
import zipfile

import numpy as np
import xarray as xr
from gbrtools import poly_to_line, shapefile_to_flat_gpkg
from osgeo import gdal
from slim4.netcdf_to_abin import tpxo2abin

We define the working directories and the geographic bounding box used to clip vector datasets to the model region.

os.makedirs("tmp", exist_ok=True)
os.makedirs("data", exist_ok=True)
os.makedirs("data/netcdf/tpxo", exist_ok=True)

bounds = [137.0, 159.0, -32.0, -6.0]
def extract_progress(file, outpath):
    """Extract a zip archive while printing a short progress message.

    Parameters
    ----------
    file : str
        Path to the zip archive.
    outpath : str
        Directory where the archive contents are extracted.
    """
    print(f"Extracting '{file}'... ", end="")
    with zipfile.ZipFile(file, "r") as zip_ref:
        zip_ref.extractall(outpath)
    print("done")

We download coral-reef polygons and convert them to a flattened GeoPackage clipped to the model bounding box. These polygons are later used both for mesh refinement and for the reef-density roughness parameterization.

if not os.path.exists("tmp/WCMC008_CoralReefs2018_v4_1.zip"):
    urllib.request.urlretrieve(
        "https://wcmc.io/WCMC_008", "tmp/WCMC008_CoralReefs2018_v4_1.zip"
    )
extract_progress("tmp/WCMC008_CoralReefs2018_v4_1.zip", "tmp/reefs")

input_shp = "tmp/reefs/14_001_WCMC008_CoralReefs2018_v4_1/01_Data/WCMC008_CoralReef2018_Py_v4_1.shp"
output_gpkg = "data/reefs.gpkg"
output_layer_name = "reefs"
shapefile_to_flat_gpkg(input_shp, output_gpkg, output_layer_name, bounds)

We download the bathymetry datasets covering the Great Barrier Reef, Torres Strait, and the Gulf of Papua. The original files are stored in data/bathymetry and reused if already present.

GBR 30 m bathymetry Source: https://ecat.ga.gov.au/geonetwork/srv/eng/catalog.search#/metadata/115066

if not os.path.exists("tmp/Great Barrier Reef Bathymetry 2020 30m.zip"):
    urllib.request.urlretrieve(
        "https://files.ausseabed.gov.au/survey/Great%20Barrier%20Reef%20Bathymetry%202020%2030m.zip",
        "tmp/Great Barrier Reef Bathymetry 2020 30m.zip",
    )
extract_progress("tmp/Great Barrier Reef Bathymetry 2020 30m.zip", "data/bathymetry")

Torres Strait 30 m bathymetry Source: https://ecat.ga.gov.au/geonetwork/srv/eng/catalog.search#/metadata/144348

if not os.path.exists(
    "tmp/Australian Bathymetry Topography (Torres Strait) 2023 30m.zip"
):
    urllib.request.urlretrieve(
        "https://files.ausseabed.gov.au/survey/Australian%20Bathymetry%20Topography%20(Torres%20Strait)%202023%2030m.zip",
        "tmp/Australian Bathymetry Topography (Torres Strait) 2023 30m.zip",
    )
extract_progress(
    "tmp/Australian Bathymetry Topography (Torres Strait) 2023 30m.zip",
    "data/bathymetry",
)

GBR 100 m bathymetry Source: https://ecat.ga.gov.au/geonetwork/srv/eng/catalog.search#/metadata/133163

if not os.path.exists("tmp/Great Barrier Reef Bathymetry 2020 100m.zip"):
    urllib.request.urlretrieve(
        "https://files.ausseabed.gov.au/survey/Great%20Barrier%20Reef%20Bathymetry%202020%20100m.zip",
        "tmp/Great Barrier Reef Bathymetry 2020 100m.zip",
    )
extract_progress("tmp/Great Barrier Reef Bathymetry 2020 100m.zip", "data/bathymetry")

Gulf of Papua 100 m bathymetry Source: https://figshare.com/articles/dataset/Gulf_of_Papua_Bathymetry_raster_dataset/11986797/1

if not os.path.exists("data/bathymetry/GulfOfPapua_100m_bathymetry.tif"):
    urllib.request.urlretrieve(
        "https://ndownloader.figshare.com/files/22014357",
        "data/bathymetry/GulfOfPapua_100m_bathymetry.tif",
    )

We merge the bathymetry rasters into two products: a 0.001-degree raster used for the mesh-size field and a coarser 0.01-degree raster useful for fast prototyping.

files = [
    "GulfOfPapua_100m_bathymetry.tif",
    "Great_Barrier_Reef_2020_100m_MSL_cog.tif",
    "Australian_Bathymetry_Topography__Torres_Strait__2023_30m_MSL_cog.tif",
    "Great_Barrier_Reef_D_2020_30m_MSL_cog.tif",
    "Great_Barrier_Reef_C_2020_30m_MSL_cog.tif",
    "Great_Barrier_Reef_B_2020_30m_MSL_cog.tif",
    "Great_Barrier_Reef_A_2020_30m_MSL_cog.tif",
]
files = ["data/bathymetry/" + f for f in files]

warp_options = gdal.WarpOptions(
    format="GTiff",
    xRes=0.001,
    yRes=0.001,
    outputType=gdal.GDT_Float32,
    resampleAlg=gdal.GRA_Bilinear,
    warpMemoryLimit=4000,
    creationOptions=[
        "COMPRESS=DEFLATE",
        "PREDICTOR=2",
        "ZLEVEL=6",
        "BIGTIFF=YES",
        "NUM_THREADS=ALL_CPUS",
    ],
    multithread=True,
)
os.makedirs("data/bathymetry/merged/", exist_ok=True)
ds = gdal.Warp("data/bathymetry/merged/merge_d100.tif", files, options=warp_options)
ds = gdal.Warp(
    "data/bathymetry/merged/merge_d1000.tif",
    "data/bathymetry/merged/merge_d100.tif",
    xRes=0.01,
    yRes=0.01,
    resampleAlg="average",
)
ds = None

We download global land polygons and convert them to coastlines. The coastline layer is used by mesh.py to close the domain and to control mesh refinement close to land.

if not os.path.exists("tmp/land-polygons-split-4326.zip"):
    urllib.request.urlretrieve(
        "https://osmdata.openstreetmap.de/download/land-polygons-split-4326.zip",
        "tmp/land-polygons-split-4326.zip",
    )
extract_progress("tmp/land-polygons-split-4326.zip", "tmp/land_polys")
input_shp = "tmp/land_polys/land-polygons-split-4326/land_polygons.shp"
land_gpkg = "data/contour/land_polygons.gpkg"
os.makedirs("data/contour", exist_ok=True)
shapefile_to_flat_gpkg(input_shp, land_gpkg, "land", bounds)
coastlines_gpkg = "data/contour/coastlines.gpkg"
poly_to_line(land_gpkg, coastlines_gpkg, "coast")

We create a dummy TPXO-like dataset with the expected variables, dimensions, and file names. Elevation and velocity amplitudes/phases are prescribed as simple functions of lon/lat so that their magnitudes remain roughly realistic.

This allows the full preprocessing and runtime workflow to be tested without downloading a real TPXO atlas, while keeping the exact same data path and file structure as in a production setup.

nx, ny, nc = 751, 1141, 2
con_labels = np.array(["m2  ", "s2  "], dtype="S4")
lon = np.linspace(135.0, 160.0, nx, dtype=np.float32)
lat = np.linspace(-33.0, -5.0, ny, dtype=np.float32)
lon_u, lat_u = np.meshgrid(lon, lat, indexing="ij")
lon_v, lat_v = np.meshgrid(lon, lat, indexing="ij")
lon_z, lat_z = np.meshgrid(lon, lat, indexing="ij")
zeros_data = np.zeros((nc, nx, ny), dtype="float32")

ha = np.array(
    [-0.0252 * lon_z - 0.0216 * lat_z + 3.91, -0.0476 * lon_z - 0.0173 * lat_z + 7.18]
)
ua = np.array([-3.32 * lon_u - 1.29 * lat_u + 488, -1.62 * lon_u - 0.592 * lat_u + 239])
va = np.array([-1.6 * lon_v - 0.704 * lat_v + 234, -1.77 * lon_v - 0.748 * lat_v + 258])
hp = np.array(
    [-11.5 * lon_z - 1.26 * lat_z + 1.7e03, -14.2 * lon_z - 5.94 * lat_z + 2.01e03]
)
up = np.array(
    [-25.2 * lon_u - 8.75 * lat_u + 3.77e03, -25.2 * lon_u - 11.3 * lat_u + 3.71e03]
)
vp = np.array(
    [-23.9 * lon_v - 20.9 * lat_v + 3.26e03, -13.5 * lon_v + 2.84 * lat_v + 2.3e03]
)
hp = (hp + 180) % 360 - 180
up = (up + 180) % 360 - 180
vp = (vp + 180) % 360 - 180

ds_h = xr.Dataset(
    data_vars={
        "lon_z": (
            ("nx", "ny"),
            lon_z,
            {"units": "degree_east", "long_name": "longitude of Z nodes"},
        ),
        "lat_z": (
            ("nx", "ny"),
            lat_z,
            {"units": "degree_north", "long_name": "latitude of Z nodes"},
        ),
        "con": (("nc"), con_labels, {"long_name": "tidal constituents"}),
        "ha": (
            ("nc", "nx", "ny"),
            ha,
            {
                "units": "meter",
                "long_name": "Tidal elevation amplitude at Z-nodes",
                "field": "amplitude, scalar",
            },
        ),
        "hp": (
            ("nc", "nx", "ny"),
            hp,
            {
                "units": "degree GMT",
                "long_name": "Tidal elevation phase at Z-nodes",
                "field": "phase, scalar",
            },
        ),
    }
)
ds_h.to_netcdf(
    "data/netcdf/tpxo/h_tpxo_zone.nc",
    format="NETCDF4",
    encoding={var: {"_FillValue": np.nan} for var in ds_h.data_vars if var != "con"},
)

ds_uv = xr.Dataset(
    data_vars={
        "lon_u": (
            ("nx", "ny"),
            lon_u,
            {"units": "degree_east", "long_name": "longitude of U nodes"},
        ),
        "lat_u": (
            ("nx", "ny"),
            lat_u,
            {"units": "degree_north", "long_name": "latitude of U nodes"},
        ),
        "lon_v": (
            ("nx", "ny"),
            lon_v,
            {"units": "degree_east", "long_name": "longitude of V nodes"},
        ),
        "lat_v": (
            ("nx", "ny"),
            lat_v,
            {"units": "degree_north", "long_name": "latitude of V nodes"},
        ),
        "con": (("nc"), con_labels, {"long_name": "tidal constituents"}),
        "ua": (
            ("nc", "nx", "ny"),
            ua,
            {
                "units": "cm/s",
                "long_name": "Tidal WE velocity amplitude at U-nodes",
                "field": "vector, W->E",
            },
        ),
        "Ua": (
            ("nc", "nx", "ny"),
            zeros_data,
            {
                "units": "m^2/s",
                "long_name": "Tidal WE transport amplitude at U-nodes",
                "field": "vector, W->E",
            },
        ),
        "up": (
            ("nc", "nx", "ny"),
            up,
            {
                "units": "degree GMT",
                "long_name": "Tidal WE transport/velocity phase at U-nodes",
                "field": "phase, scalar",
            },
        ),
        "va": (
            ("nc", "nx", "ny"),
            va,
            {
                "units": "cm/s",
                "long_name": "Tidal SN velocity amplitude at V-nodes",
                "field": "vector, S->N",
            },
        ),
        "Va": (
            ("nc", "nx", "ny"),
            zeros_data,
            {
                "units": "m^2/s",
                "long_name": "Tidal SN transport amplitude at V-nodes",
                "field": "vector, S->N",
            },
        ),
        "vp": (
            ("nc", "nx", "ny"),
            vp,
            {
                "units": "degree GMT",
                "long_name": "Tidal SN transport/velocity phase at V-nodes",
                "field": "phase, scalar",
            },
        ),
    }
)

ds_uv.to_netcdf(
    "data/netcdf/tpxo/u_tpxo_zone.nc",
    format="NETCDF4",
    encoding={var: {"_FillValue": np.nan} for var in ds_uv.data_vars if var != "con"},
)

h_path = tpxo2abin("data/netcdf/tpxo/h_tpxo_zone.nc")
u_path = tpxo2abin("data/netcdf/tpxo/u_tpxo_zone.nc")

# For compatibility with case-insensitive filesystems, SLIM expects transport
# variables to be named hua/hva instead of Ua/Va.
for varname in ["Ua.abin", "Ua.ajson", "Va.abin", "Va.ajson"]:
    newname = "h" + varname.lower()
    os.replace(f"{u_path}/{varname}", f"{u_path}/{newname}")

Download this section.

Generate the horizontal mesh

This is the second step of the Great Barrier Reef workflow. It uses the static datasets downloaded by download.py to build an unstructured horizontal mesh.

The mesh-size field is refined near reefs, coastlines, and narrow openings, while also respecting a bathymetry-based stability constraint. The resulting mesh is written to data/gbr.msh and is used by the preprocessing step.

import numpy as np
import pyogrio
import pyproj
import seamsh
import shapely
from osgeo import osr

CurveType = seamsh.geometry.CurveType
gmsh = seamsh.gmsh.gmsh

The computational mesh is generated in a projected coordinate system, while some of the the input boundary curves are specified in longitude/latitude coordinates.

proj = osr.SpatialReference()
proj.ImportFromProj4("+ellps=WGS84 +proj=utm +zone=55 +south")

projlonlat = osr.SpatialReference()
projlonlat.ImportFromEPSG(4326)
projlonlat.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)

The bathymetry raster and reef polygons are used by the mesh-size field. The reef distance query is cached because it is called many times during meshing.

bath_interp = seamsh.field.Raster("data/bathymetry/merged/merge_d100.tif")
reefs = "data/reefs.gpkg"


def reefs_dist(xy, projection):
    if not hasattr(reefs_dist, "tree"):
        gdf = pyogrio.read_dataframe(reefs)
        gdf = gdf.to_crs(pyproj.CRS.from_wkt(proj.ExportToWkt()))
        polys = gdf.explode(index_parts=False).geometry.to_numpy()
        reefs_dist.tree = shapely.STRtree(polys)
    x, y = xy.T
    out = np.full(x.shape[0], np.nan, dtype="float64")
    pts = shapely.points(x, y)
    [pt_idx, poly_idx], dist = reefs_dist.tree.query_nearest(pts, return_distance=True)
    out[pt_idx] = dist
    return out

We define two open-boundary curves that close the offshore side of the domain: one through the Coral Sea and one through Torres Strait. Coastline curves are read from the GeoPackage created in download.py.

coral_sea = np.array(
    [
        (152.909557, -26.515829),
        (154.342362, -25.847186),
        (155.297566, -25.297944),
        (155.934369, -24.255180),
        (156.181130, -23.077095),
        (156.006009, -20.784606),
        (155.361246, -19.312000),
        (154.676684, -18.420477),
        (153.355319, -16.482209),
        (152.081714, -14.094200),
        (150.999149, -11.761911),
        (150.541448, -10.551986),
    ]
)

torres_strait = np.array(
    [
        (142.306795, -9.097290),
        (142.408286, -9.383851),
        (142.489876, -9.674392),
        (142.547586, -9.978864),
        (142.573456, -10.14403),
        (142.557536, -10.37686),
        (142.529676, -10.60372),
        (142.444106, -10.86441),
        (142.396345, -10.98381),
    ]
)

domain = seamsh.geometry.Domain(proj)
domain.add_boundary_curve(coral_sea, "coral_sea", projlonlat, CurveType.BSPLINE)
domain.add_boundary_curve(torres_strait, "torres_strait", projlonlat, CurveType.BSPLINE)
domain.add_boundary_curves_shp("data/contour/coastlines.gpkg", "physical")

coast_dist = seamsh.field.Distance(domain, 200, ["coast"])
coral_dist = seamsh.field.Distance(domain, 200, ["coral_sea"])
torres_strait_dist = seamsh.field.Distance(domain, 200, ["torres_strait"])

The following functions define the spatially varying target mesh size. The final size is the result of geometric refinement constraints and a minimum size imposed by the local water depth and target time step.

def expogrowth(d, dnear: float, dgrowth: float, growthfactor: float):
    return np.maximum(np.exp((np.log(growthfactor) / dgrowth) * (d - dnear)), 1)


def mesh_size(x, projection):
    dreef = reefs_dist(x, projection)
    dreef = np.clip(dreef, 0, 1e5)

    lc = 10000  # largest mesh size
    f = np.ones(x.shape[0])

    eg = expogrowth(dreef, 600, 2000, 1.5)
    freef = 200 / lc * eg
    f = np.minimum(f, freef)

    c_dist = coral_dist(x, projection)
    fopen = 1 / expogrowth(c_dist, lc, lc, 2)
    f = np.maximum(f, fopen)

    ts_dist = torres_strait_dist(x, projection)
    fopen = 1 / expogrowth(ts_dist, 4000, 4000, 2) * (4000 / lc)
    f = np.maximum(f, fopen)

    size = f * lc
    size = np.minimum(size, lc)

    bath = -bath_interp(x, projection)
    bath = np.clip(bath, 2, np.inf)
    g = 9.81
    dtmax = 7

    sizemin = np.sqrt(12 * g * bath) * dtmax
    sizemin = np.clip(np.clip(sizemin, 200, np.inf), 0, lc)
    sizemax = np.sqrt(12 * g * bath) * dtmax * 5

    size = np.minimum(size, sizemax)
    size = np.maximum(size, sizemin)
    return size * 10

Finally, we coarsen the input boundaries according to the same size field and call gmsh through seamsh. The resulting mesh is the horizontal mesh used by prepro.py.

coarse_initial = seamsh.geometry.coarsen_boundaries(
    domain, (145.440, -14.746), projlonlat, mesh_size
)
seamsh.gmsh.mesh(coarse_initial, "data/gbr.msh", mesh_size, version=4.1, smoothness=0.3)
# seamsh.gmsh.convert_to_gis("data/gbr.msh", proj, "data/gbr.gpkg")

Download this section.

Download time-varying forcing data

This is the third step of the Great Barrier Reef workflow. It downloads the time-dependent ocean and atmospheric fields used by prepro.py: 3D currents, temperature, salinity, sea surface elevation, wind, heat fluxes, precipitation, and sea-surface temperature.

Most products are downloaded through NCSS/THREDDS requests in monthly or shorter chunks. The SST product is downloaded through the Copernicus Marine client. Existing files are checked before downloading so interrupted runs can be resumed safely.

import os
import sys
import time
import threading
from datetime import datetime, timedelta
from typing import List, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed
import copernicusmarine

import requests
import xarray as xr
import numpy as np


def _get_next_date(date: datetime) -> datetime:
    return datetime(date.year + date.month // 12, date.month % 12 + 1, 1)


def generate_months_from_timestamps(
    start_ts: float, end_ts: float, chunk: str | int = "month"
):
    """Split a time interval into monthly or shorter download chunks.

    Parameters
    ----------
    start_ts, end_ts : float
        Start and end times as UTC timestamps.
    chunk : {"month", "week"} or int, optional
        Chunking strategy. Integer values are interpreted as a number of days.

    Returns
    -------
    months : list of str
        Month tags in ``YYYYMM`` format, one per chunk.
    intervals : list of tuple of str
        Start and end date tags in ``YYYYMMDD`` format, one per chunk.
    """
    start = datetime(*time.gmtime(start_ts)[:6])
    end = datetime(*time.gmtime(end_ts)[:6])

    months = []
    intervals = []

    # normalize start to midnight
    start = datetime(start.year, start.month, start.day)
    end = datetime(end.year, end.month, end.day)

    if chunk == "month":
        cur = start
        while cur <= end:
            months.append(cur.strftime("%Y%m"))
            last = min(end, _get_next_date(cur) - timedelta(days=1))
            intervals.append((cur.strftime("%Y%m%d"), last.strftime("%Y%m%d")))
            cur = _get_next_date(cur)

    else:
        if chunk == "week":
            step = 7
        elif isinstance(chunk, int) and chunk > 0:
            step = chunk
        else:
            raise ValueError(f"Invalid chunk: {chunk}")
        cur = start
        while cur <= end:
            month_end = _get_next_date(cur) - timedelta(days=1)
            last = min(end, month_end, cur + timedelta(days=step - 1))
            months.append(cur.strftime("%Y%m"))
            intervals.append((cur.strftime("%Y%m%d"), last.strftime("%Y%m%d")))
            cur = last + timedelta(days=1)

    return months, intervals


def file_exists_and_probably_ok(path: str) -> bool:
    """Check whether a downloaded netCDF file can be reused.

    Parameters
    ----------
    path : str
        Path to the candidate netCDF file.

    Returns
    -------
    bool
        ``True`` if the file exists, can be opened by xarray, and does not end
        with an all-NaN time slice; otherwise ``False``. Invalid files are
        removed so they can be downloaded again.
    """
    if not os.path.exists(path):
        return False
    try:
        with xr.open_dataset(path) as dsdisk:
            for variable in list(dsdisk.keys()):
                diskshape = dsdisk[variable].shape
                if (
                    len(dsdisk[variable].dims) > 0
                    and dsdisk[variable].dims[0] == "time"
                ):
                    if np.all(np.isnan(dsdisk[variable][-1])):
                        print(
                            f"{path} : All NANS in variable {variable}, re-downloading"
                        )
                        os.remove(path)
                        return False
            print(f"{path} already present")
            return True
    except Exception:
        print(f"{path} corrupted, re-downloading")
        os.remove(path)
        return False


def _month_dataset_url(variable: str, month: str, dataset="barra2r2") -> str:
    """Build the THREDDS NCSS URL for a monthly source file.

    Parameters
    ----------
    variable : str
        Name of the variable to download.
    month : str
        Month tag in ``YYYYMM`` format.
    dataset : {"barra2r2", "barra2c2", "bran"}, optional
        Source dataset identifier.

    Returns
    -------
    str
        THREDDS NCSS endpoint for the requested source file.
    """
    THREDDS_BASE = "https://thredds.nci.org.au/thredds/ncss/grid"
    if dataset == "barra2r2":
        dspath = (
            "ob53/output/reanalysis/AUS-11/BOM/ERA5/historical/hres/BARRA-R2/v1/1hr"
        )
        return f"{THREDDS_BASE}/{dspath}/{variable}/latest/{variable}_AUS-11_ERA5_historical_hres_BOM_BARRA-R2_v1_1hr_{month}-{month}.nc"
    if dataset == "barra2c2":
        dspath = (
            "ob53/output/reanalysis/AUST-04/BOM/ERA5/historical/hres/BARRA-C2/v1/1hr"
        )
        return f"{THREDDS_BASE}/{dspath}/{variable}/latest/{variable}_AUST-04_ERA5_historical_hres_BOM_BARRA-C2_v1_1hr_{month}-{month}.nc"
    if dataset == "bran":
        dspath = "gb6/BRAN/BRAN2023/daily"
        return f"{THREDDS_BASE}/{dspath}/ocean_{variable}_{month[0:4]}_{month[4:6]}.nc"
    else:
        raise ValueError(f"Unknown dataset: {dataset}")


def _interval_to_iso_utc(interval: Tuple[str, str]) -> Tuple[str, str]:
    start_tag, end_tag = interval
    start_iso = datetime.strptime(start_tag, "%Y%m%d").strftime("%Y-%m-%dT00:00:00Z")
    end_iso = datetime.strptime(end_tag, "%Y%m%d").strftime("%Y-%m-%dT23:59:59Z")
    return start_iso, end_iso


class MultiProgress:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._lines: List[str] = []
        self._initialized = False
        self._interactive = sys.stdout.isatty()

    def add(self, text: str) -> int:
        with self._lock:
            idx = len(self._lines)
            self._lines.append(text)
            if self._interactive:
                sys.stdout.write(text + "\n")
                sys.stdout.flush()
                self._initialized = True
            else:
                print(text, flush=True)

            return idx

    def update(self, idx: int, text: str) -> None:
        with self._lock:
            self._lines[idx] = text
            if not self._interactive:
                return
            n = len(self._lines)
            # Move cursor to the first progress line
            sys.stdout.write(f"\x1b[{n}F")
            for line in self._lines:
                sys.stdout.write("\x1b[2K")  # clear entire line
                sys.stdout.write(line + "\n")
            sys.stdout.flush()

    def finish(self) -> None:
        if self._interactive:
            sys.stdout.flush()


def _format_progress_line(
    filename: str,
    downloaded: int,
    total_size: int,
    elapsed: float,
    status: str,
) -> str:
    """Format one human-readable download-progress line.

    Parameters
    ----------
    filename : str
        Destination file name.
    downloaded : int
        Number of bytes already written.
    total_size : int
        Expected total size in bytes, or zero if unknown.
    elapsed : float
        Elapsed time in seconds.
    status : str
        Short status label.

    Returns
    -------
    str
        Progress line suitable for terminal output.
    """
    name = os.path.basename(filename)
    elapsed = max(elapsed, 1e-12)
    speed_mbs = downloaded / elapsed / 2**20
    if total_size > 0:
        percent = 100.0 * downloaded / total_size
        remaining = max(total_size - downloaded, 0)
        eta = remaining / max(downloaded / elapsed, 1e-12) if downloaded > 0 else 0.0
        return f"{name}: {status:11s} {percent:5.1f}% {downloaded / 2**20:8.1f}/{total_size / 2**20:8.1f} MB {speed_mbs:6.1f} MB/s ETA {eta:6.1f}s"
    return f"{name}: {status:11s} {downloaded / 2**20:8.1f} MB {speed_mbs:6.1f} MB/s"


def _download_with_progress(
    response: requests.Response,
    dst_file: str,
    progress: MultiProgress,
    progress_idx: int,
    chunk_size: int = 1024 * 1024,
) -> None:
    """Stream a response to disk while updating progress output.

    Parameters
    ----------
    response : requests.Response
        Open streaming HTTP response.
    dst_file : str
        Final output file path.
    progress : MultiProgress
        Progress manager.
    progress_idx : int
        Index of the progress line to update.
    chunk_size : int, optional
        Size of chunks read from the HTTP stream, in bytes.
    """
    tmp_file = dst_file + ".part"
    total_size = int(response.headers.get("Content-Length", 0))
    downloaded = 0
    t0 = time.time()
    last_update = 0.0

    with open(tmp_file, "wb") as f:
        for chunk in response.iter_content(chunk_size=chunk_size):
            if not chunk:
                continue
            f.write(chunk)
            downloaded += len(chunk)
            now = time.time()
            if now - last_update >= 0.2:
                progress.update(
                    progress_idx,
                    _format_progress_line(
                        dst_file, downloaded, total_size, now - t0, "downloading"
                    ),
                )
                last_update = now

    os.replace(tmp_file, dst_file)
    progress.update(
        progress_idx,
        _format_progress_line(
            dst_file, downloaded, total_size, time.time() - t0, "done"
        ),
    )


def _download_ncss_subset(
    variable: str,
    dataset: str,
    month: str,
    interval: Tuple[str, str],
    bounds,
    dst_file: str,
    progress: MultiProgress,
    progress_idx: int,
    timeout: int = 300,
) -> str:
    minlon, maxlon, minlat, maxlat = bounds
    time_start, time_end = _interval_to_iso_utc(interval)

    url = _month_dataset_url(variable, month, dataset)
    params = {
        "var": variable,
        "west": float(minlon),
        "east": float(maxlon),
        "south": float(minlat),
        "north": float(maxlat),
        "time_start": time_start,
        "time_end": time_end,
        "accept": "netcdf4",
    }

    progress.update(progress_idx, f"{os.path.basename(dst_file)}: connecting")
    with requests.get(url, params=params, stream=True, timeout=timeout) as response:
        response.raise_for_status()
        content_type = response.headers.get("Content-Type", "")
        if "html" in content_type.lower():
            text_preview = response.text[:1000]
            raise RuntimeError(
                f"Server returned HTML instead of netCDF for {response.url}\nContent-Type: {content_type}\nResponse preview:\n{text_preview}"
            )
        _download_with_progress(response, dst_file, progress, progress_idx)
    return dst_file


def _download_one_task(task, progress: MultiProgress) -> str:
    dst_file = task["dst_file"]
    progress_idx = progress.add(f"{os.path.basename(dst_file)}: queued")
    try:
        return _download_ncss_subset(
            variable=task["variable"],
            dataset=task["dataset"],
            month=task["month"],
            interval=task["interval"],
            bounds=task["bounds"],
            dst_file=dst_file,
            progress=progress,
            progress_idx=progress_idx,
        )
    except Exception as exc:
        tmp_file = dst_file + ".part"
        if os.path.exists(tmp_file):
            os.remove(tmp_file)
        progress.update(progress_idx, f"{os.path.basename(dst_file)}: FAILED - {exc}")
        return ""


def download_variables(
    dst_dir: str,
    dataset: str,
    variables: List[str],
    bounds,
    time_from: float,
    time_to: float,
    chunk: str | int = "month",
    max_workers: int = 4,
) -> List[str]:
    """Download one or more variables from a gridded THREDDS dataset.

    Parameters
    ----------
    dst_dir : str
        Directory where downloaded netCDF files are stored.
    dataset : {"barra2r2", "barra2c2", "bran"}
        Source dataset identifier.
    variables : list of str
        Variables to download.
    bounds : sequence of float
        Bounding box ``[minlon, maxlon, minlat, maxlat]``.
    time_from, time_to : float
        Download interval as UTC timestamps.
    chunk : {"month", "week"} or int, optional
        Chunking strategy. Integer values are interpreted as a number of days.
    max_workers : int, optional
        Number of concurrent downloads.

    Returns
    -------
    list of str
        Sorted list of downloaded or reused files.
    """
    os.makedirs(dst_dir, exist_ok=True)
    months, intervals = generate_months_from_timestamps(time_from, time_to, chunk)

    all_files: List[str] = []
    tasks = []

    for variable in variables:
        for month, interval in zip(months, intervals):
            start_tag, end_tag = interval
            dst_file = os.path.join(dst_dir, f"{variable}_{start_tag}_{end_tag}.nc")

            if file_exists_and_probably_ok(dst_file):
                all_files.append(dst_file)
                continue

            tasks.append(
                {
                    "variable": variable,
                    "dataset": dataset,
                    "month": month,
                    "interval": interval,
                    "bounds": bounds,
                    "dst_file": dst_file,
                }
            )

    if not tasks:
        return sorted(all_files)

    progress = MultiProgress()
    results: List[str] = []

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            executor.submit(_download_one_task, task, progress) for task in tasks
        ]
        for future in as_completed(futures):
            results.append(future.result())

    progress.finish()
    all_files.extend(results)
    return sorted(all_files)


def download_sst(
    dst_dir: str,
    variables: List[str],
    bounds,
    time_from: float,
    time_to: float,
    chunk: str | int = "month",
    username: str | None = None,
    password: str | None = None,
) -> List[str]:
    """Download sea-surface-temperature products from Copernicus Marine.

    Parameters
    ----------
    dst_dir : str
        Directory where downloaded netCDF files are stored.
    variables : list of str
        SST variables to request.
    bounds : sequence of float
        Bounding box ``[minlon, maxlon, minlat, maxlat]``.
    time_from, time_to : float
        Download interval as UTC timestamps.
    chunk : {"month", "week"} or int, optional
        Chunking strategy. Kept consistent with ``download_variables``.
    username, password : str, optional
        Copernicus Marine credentials.

    Returns
    -------
    list of str
        Sorted list of downloaded or reused files.
    """
    os.makedirs(dst_dir, exist_ok=True)
    months, intervals = generate_months_from_timestamps(time_from, time_to, chunk)
    all_files: List[str] = []
    for variable in variables:
        for month, interval in zip(months, intervals):
            start_tag, end_tag = interval
            fname = f"{variable}_{start_tag}_{end_tag}.nc"
            dst_file = os.path.join(dst_dir, fname)

            if file_exists_and_probably_ok(dst_file):
                all_files.append(dst_file)
                continue

            copernicusmarine.subset(
                dataset_id="METOFFICE-GLO-SST-L4-REP-OBS-SST",
                variables=[f"{variable}"],
                start_datetime=f"{start_tag[0:4]}-{start_tag[4:6]}-{start_tag[6:8]}",
                end_datetime=f"{end_tag[0:4]}-{end_tag[4:6]}-{end_tag[6:8]}",
                minimum_longitude=137.0,
                maximum_longitude=159.0,
                minimum_latitude=-32.0,
                maximum_latitude=-6.0,
                output_directory=dst_dir,
                output_filename=fname,
                username=username,
                password=password,
                overwrite=True,
            )
            all_files.append(dst_file)

    return sorted(all_files)

When executed as a script, we download all forcing products required by the preprocessing step.

if __name__ == "__main__":
    from slim4.tools import parse_time
    import params

    bounds = params.bounds
    time_from = parse_time(params.initial_time)
    time_to = parse_time(params.final_time)

    files = download_variables(
        "data/netcdf/hydro3d",
        "bran",
        ["u", "v", "temp", "salt"],
        bounds,
        time_from,
        time_to,
        chunk="month",
        max_workers=5,
    )

    files = download_variables(
        "data/netcdf/hydro2d",
        "bran",
        ["eta_t"],
        bounds,
        time_from,
        time_to,
        chunk="month",
        max_workers=5,
    )

    files = download_variables(
        "data/netcdf/wind",
        "barra2c2",
        [
            "uas",  # east wind 10m
            "vas",  # notrh wind 10m
            "psl",  # sea level pressure
        ],
        bounds,
        time_from,
        time_to,
        chunk=16,
        max_workers=5,
    )

    files = download_variables(
        "data/netcdf/atm",
        "barra2c2",
        [
            "hurs",  # surface relative humidity [%]
            "hfls",  # surface upward latent heat flux [W/m²]
            "hfss",  # surface upward sensible heat flux [W/m²]
            "rsds",  # surface downwelling shortwave radiation [W/m²]
            "rsus",  # surface upwelling shortwave radiation [W/m²]
            "rlds",  # surface downwelling longwave radiation [W/m²]
            "rlus",  # surface upwelling longwave radiation [W/m²]
            "pr",  # precipitation [kg/m²/s]
            "tas",  # near-surface air temp [K]
            "ts",  # surface skin temp [K]
        ],
        bounds,
        time_from,
        time_to,
        chunk=16,
        max_workers=5,
    )

    files = download_sst(
        "data/netcdf/sst", ["analysed_sst"], bounds, time_from, time_to
    )

Download this section.

Preprocess data for SLIM4

This is the fourth step of the Great Barrier Reef workflow. It converts the static and time-varying datasets prepared by the previous scripts into SLIM4 data files on the model mesh.

The script builds the 3D mesh and bathymetry, computes static fields such as Coriolis and reef density, and interpolates atmospheric, tidal, ocean-current, temperature, and salinity products onto the grids and boundaries needed by run.py.

import os
import numpy as np
from slim4 import (
    slim,
    slim2d,
    slim3d,
    pre3d,
    tpxo,
    mpi,
    particles,
)
from slim4.tools import parse_time, VectorRotation, read_polygons_shp
from slim4.netcdf_to_abin import nc2abin
from slim4.pre import Struct, clip_inclusive
from gbrtools import bathy_on_mesh

import params as param


def log(*args, **kwargs) -> None:
    if mpi.rank > 0:
        return
    print(*args, **kwargs)


def smith_banke(uv_wind: np.ndarray):
    """Compute wind stress with the Smith and Banke drag law.

    Parameters
    ----------
    uv_wind : ndarray of shape (2, n)
        Wind velocity components in m/s.

    Returns
    -------
    ndarray of shape (2, n)
        Wind stress components.
    """
    atm_density = 1.25
    f0, f1, f2 = 0, 6.3e-4, 6.6e-5
    norm = np.hypot(*uv_wind)
    coeff = f0 + f1 * norm + f2 * norm**2
    return atm_density * coeff[None, :] * uv_wind


tstart = param.initial_time
tend = param.final_time

timelaps = np.array([parse_time(tstart), parse_time(tend)])
time_tides = np.arange(timelaps[0], timelaps[-1] + 3599, 3600)
time_check = np.arange(timelaps[0], timelaps[-1], param.check_dt)

base_dir = param.nc_data_dir + "/"
prepro_dir = param.prepro_data_dir + "/"
if mpi.rank == 0:
    os.makedirs(prepro_dir + "check", exist_ok=True)
mpi.barrier()


def get_files(path, vars):
    path = os.path.join(base_dir, path)
    r = []
    for var in vars:
        ncfiles = os.listdir(path)
        ncfiles = [f for f in ncfiles if f.startswith(var)]
        ncfiles = [os.path.join(path, f) for f in ncfiles if f.endswith(".nc")]
        for f in ncfiles:
            if not os.path.exists(f[:-3]):
                print(f"Converting {f[:-3]} to abin format...")
                nc2abin(f)
        ncfiles = [f[:-3] for f in ncfiles]
        ncfiles.sort()
        if len(ncfiles) == 0:
            raise FileNotFoundError(f"No file found for variable {var} in path {path}")
        r.append(ncfiles)

    return r


def get_struct(path, vars):
    """Load 2D/3D forcing variables as a merged SLIM4 Struct.

    Parameters
    ----------
    path : str
        Subdirectory of ``base_dir`` containing the files.
    vars : list of str
        Variables to load.

    Returns
    -------
    Struct
        Merged structure containing the requested variables.
    """
    flists = get_files(path, vars)
    structs = [Struct.from_files(flist, [var]) for (flist, var) in zip(flists, vars)]
    return Struct.merge_structs(structs)

For the bathymetry source used for preprocessing, the preferred approach is to provide a list of rasters, which are loaded in order of priority (earlier files take precedence where they overlap). This preserves the highest available resolution locally.

Single merged rasters are convenient for prototyping but less accurate.

tif = [
    "data/bathymetry/" + s
    for s in [
        "Great_Barrier_Reef_A_2020_30m_MSL_cog.tif",
        "Great_Barrier_Reef_B_2020_30m_MSL_cog.tif",
        "Great_Barrier_Reef_C_2020_30m_MSL_cog.tif",
        "Great_Barrier_Reef_D_2020_30m_MSL_cog.tif",
        "Australian_Bathymetry_Topography__Torres_Strait__2023_30m_MSL_cog.tif",
        "Great_Barrier_Reef_2020_100m_MSL_cog.tif",
        "GulfOfPapua_100m_bathymetry.tif",
    ]
]

# Alternative (faster, lower fidelity):
# tif = "data/bathymetry/merged/merge_d100.tif"   # good compromise
# tif = "data/bathymetry/merged/merge_d1000.tif"  # very fast prototyping

Next, we read and (temporarily) partition the horizontal mesh generated by mesh.py. This 2D partition is only used during preprocessing. After extrusion, the mesh is repartitioned in 3D so that both the horizontal triangles and vertical prisms are balanced across MPI ranks.

Bathymetry is interpolated onto the mesh, clipped to this minimum depth, and smoothed in shallow regions. The mesh is then extruded using a hybrid vertical coordinate. Because the computational cost depends on the number of vertical layers, the final partitioning is performed on the 3D mesh to properly balance the workload.

def mesh_and_bathy():
    mesh_file = param.mesh_file
    if mpi.size > 1:
        pfile = prepro_dir + "/mesh/${rank}/mesh.msh"
        mesh_file = slim2d.partition_mesh(param.mesh_file, pfile=pfile)
    mesh2d0 = slim2d.Mesh(mesh_file)

    lon0, lat0 = slim.proj_transform(mesh2d0.projection, "lonlat", mesh2d0.x)[
        [[0], [1]]
    ]
    tides0 = tpxo.GridTPXO(
        mesh2d0, base_dir + "tpxo/h_tpxo_zone", base_dir + "tpxo/u_tpxo_zone", npass=30
    )
    tidal_amplitude0 = (
        tides0.tides["h"]
        .approx_amplitude_at_points(
            np.array([lon0.ravel(), lat0.ravel()]).T, timelaps[-1]
        )
        .reshape(lon0.shape)
    )
    min_depth0 = tidal_amplitude0 + param.min_depth

    bath_dg0 = bathy_on_mesh(mesh2d0, tif, nodal=False, cg=True, rule_depth=2)
    bath_dg0 = np.maximum(bath_dg0, min_depth0)
    deep_areas = bath_dg0 > param.depth_sigma
    bath_deep = bath_dg0[deep_areas]
    bath_dg0[deep_areas] = param.depth_sigma

    bath_dg0 = mesh2d0.smooth_bath(bath_dg0, coefficient=0.4, minimum_delta=0.3)
    # bath_dg0 = mesh2d0.smooth_bath_haney(bath_dg0, coefficient=1.5, minimum_delta=0.3)
    bath_dg0[deep_areas] = bath_deep
    bath0 = slim.DataStatic(mesh2d0, "DG", 1, value=bath_dg0)

    shallow_vmap = pre3d.VMapSongHaidvogel(param.n_sigma, 4, 0.5, 5)
    vmap = pre3d.VMapHybrid.create_exponential_deep(
        param.depth_sigma, shallow_vmap, param.z_levels, True, bath0
    )
    mesh0, z_, _ = pre3d.extrude(mesh2d0, bath0, vmap)
    mesh0.save_mesh(z_, prepro_dir + "mesh3d")
    slim3d.partition_mesh(
        param.mesh_file, prepro_dir + "mesh3d", prepro_dir + "mesh3d_split"
    )


mesh_and_bathy()

We load the repartitioned 3D mesh and recompute the real bathymetry on the final 2D mesh. This field is kept separately from the smoothed model bathymetry and is later used for diagnostics and velocity clipping.

mesh, z, bath = slim3d.Mesh.load(
    prepro_dir + "mesh3d_split", n_moving_max=param.n_moving_max
)
mesh2d = mesh.mesh2d
realbath = bathy_on_mesh(mesh2d, tif, nodal=False, cg=True, rule_depth=2)
realbath = np.maximum(realbath, -1.0)
deep_areas = realbath > param.depth_sigma
bath_deep = realbath[deep_areas]
realbath[deep_areas] = param.depth_sigma
realbath = mesh2d.smooth_bath(realbath, coefficient=2.0, minimum_delta=0.3)
realbath[deep_areas] = bath_deep
slim.DataFile.write(
    prepro_dir + "bath", [("bath", bath.eval_mesh()), ("realbath", realbath)], "DG"
)
mesh2d.create_xdmf(param.output_directory + "/bath", cg=False)
mesh2d.append_to_xdmf(
    param.output_directory + "/bath",
    0,
    0.0,
    [("bath", bath.eval_mesh()), ("realbath", realbath)],
)
lon, lat = slim.proj_transform(mesh2d.projection, "lonlat", mesh2d.x)[[[0], [1]]]

bounds = lon.min(), lon.max(), lat.min(), lat.max()
mpi.barrier()

We prepare static 2D fields used by the simulation: the Coriolis parameter and reef-density-based bottom roughness.

def prepro_coriolis():
    """Compute and write the Coriolis parameter on the 2D mesh."""
    cor = 2 * (2 * np.pi / 86400) * np.sin(np.deg2rad(lat))
    slim.DataFile.write(prepro_dir + "cor", [("cor", cor)], "DG")
    mpi.barrier()
    log("coriolis done")


def prepro_reef_denisty():
    """Compute reef coverage density on mesh elements and write it to disk."""
    reef_shp = "data/reefs.gpkg"
    polymesh = particles.PolygonMeshIntersection(
        mesh2d, *read_polygons_shp(mesh2d, reef_shp)
    )
    reef_density = polymesh.get_density().reshape((1, 3, -1))
    reef_density = np.clip(reef_density, 0, 1)
    reef_density_cg = mesh2d.gather_dg_to_cg("mean_area", reef_density)
    sreef_density = reef_density_cg[0][mesh2d.element_nodes][:, :, None].T

    slim.DataFile.write(
        prepro_dir + "reef_density", [("reef_density", sreef_density)], "DG"
    )
    mesh2d.create_xdmf(param.output_directory + "/reef_density")
    mesh2d.append_to_xdmf(
        param.output_directory + "/reef_density",
        0,
        0,
        [("reef_density", sreef_density)],
    )
    mpi.barrier()
    log("Reef density done")

Next, we prepare the 2D time-dependent surface and open-boundary forcings, including wind surface stress, heat fluxes, precipitation and sea level

def prepro_wind():
    """Interpolate wind, rotate it to mesh coordinates, and compute wind stress."""
    log("Processing WIND velocity...")
    wind_vars = ["uas", "vas"]
    wind_vars.append("psl")  # if use atmospheric pressure

    wind = get_struct("wind", wind_vars)
    csr = wind.compute_csr_cg_2d(mesh2d)
    rot_wind = VectorRotation(mesh2d, wind.get_lonlat_grid(csr))
    with slim.DataFileWriter(
        prepro_dir + "slim4_wind", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(wind.get_time(*timelaps)):
            [(_, windx), (_, windy)] = wind.get_data(csr, t, ["uas", "vas"])
            windx, windy = rot_wind.rotate(windx, windy)
            ws = np.vstack([windx, windy])
            tau = smith_banke(ws)
            wind_data = [("wind", ws), ("tau", tau)]
            wind_data += wind.get_data(csr, t, ["psl"])
            writer.append(wind_data, t)

    mpi.barrier()


def prepro_elevation():
    """Combine ocean-model sea level with tidal elevation on open boundaries."""
    log("Processing Sea Surface Elevation...")
    tides = tpxo.GridTPXO(
        mesh2d, base_dir + "tpxo/h_tpxo_zone", base_dir + "tpxo/u_tpxo_zone"
    )
    eta = get_struct("hydro2d", ["eta_t"])

    # tides typically have a finer grid, so interpolate eta on the tide grid
    eta.set_reinterp_coordinates(tides.lon, tides.lat)
    csr = eta.compute_csr_cg_2d(mesh2d, param.open_tags)
    tides.set_points(eta.get_lonlat_grid(csr))
    with slim.DataFileWriter(
        prepro_dir + "slim4_eta", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(time_tides):
            [(_, h)] = eta.get_data(csr, t, ["eta_t"])
            h += tides.get_eta(t)
            writer.append([("h", h)], t)

    mpi.barrier()


def rel_hu_from_dew(dew, tair):
    """Calculate relative humidity from dew point and air temperature.

    Parameters
    ----------
        dew : NDarray
            Dewpoint temperature in °C.
        tair : NDarray
            Air temperature in °C.
    Returns
    -------
        rel_hum : NDarray
            Relative humidity as a ratio between 0 and 1.
    """
    b = 17.625
    c = 243.04  # degree Celsius
    return np.exp(b * c * (dew - tair) / ((c + tair) * (c + dew)))


def dewpoint(rel_hum, temp):
    """
    Parameters:
        rel_hum -- relative humidity as a ratio between 0 and 1.
        temp -- air temperature in degree Celsius
    """
    b = 17.625
    c = 243.04  # degree Celsius
    gamma = (b * temp) / (c + temp) + np.log(rel_hum)
    dew_point = (c * gamma) / (b - gamma)
    return dew_point


def vapor_pressure(temp):
    """Calculate saturation vapor pressure from temperature.

    Parameters
    ----------
        temp : NDarray
            Temperature in °C.
    Returns
    -------
        vapor_press : NDarray
            Saturation vapor pressure in Pa.
    """
    return 6.133523 * np.exp(17.502 * temp / (240.97 + temp))


def prepro_atmo():
    """Prepare atmospheric heat-flux and humidity fields for the model."""
    log("Processing atmospheric data...")
    atmo_vars1 = [
        "hfls",  # surface upward latent heat flux [W/m²]
        "hfss",  # surface upward sensible heat flux [W/m²]
        "rsds",  # surface downwelling shortwave radiation [W/m²]
        "rlus",  # surface upwelling longwave radiation [W/m²]
    ]
    atmo_vars2 = [
        "hurs",  # surface relative humidity [%]
        "tas",  # near-surface air temp [K]
    ]
    atmo1 = get_struct("atm", atmo_vars1)
    atmo2 = get_struct("atm", atmo_vars2)
    sst = get_struct("sst", ["analysed_sst"])
    [atmo1, atmo2, sst], time_combined = Struct.reinterp_at_finest_grid(
        [atmo1, atmo2, sst], bounds=bounds
    )

    time_combined = clip_inclusive(time_combined, *timelaps)
    csr = atmo1.compute_csr_cg_2d(mesh2d)
    atmo2.compute_csr_cg_2d(mesh2d)
    sst.compute_csr_cg_2d(mesh2d)

    with slim.DataFileWriter(
        prepro_dir + "slim4_atmo", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer_atmo:
        for i, t in enumerate(time_combined):
            [
                (_, data_lat_heat_flux_up),
                (_, data_sens_heat_flux_up),
                (_, shortwave_down),
                (_, shortwave_up),
                (_, longwave_down),
                (_, longwave_up),
            ] = atmo1.get_data(csr, t, ["hfls", "hfss", "rsds", "rsus", "rlds", "rlus"])

            rho = 1000
            shortwave_pen_ratio = 0.76 * 0.302
            lat_heat_flux_wm2 = -data_lat_heat_flux_up  # [W/m2]
            sens_heat_flux_wm2 = -data_sens_heat_flux_up  # [W/m2]
            longwave_heat_fluw_wm2 = longwave_down - longwave_up  # [W/m2]
            shorwave_heat_fluw_wm2 = shortwave_down - shortwave_up

            pen_flux = (
                shorwave_heat_fluw_wm2 * shortwave_pen_ratio / (4180 * rho)
            )  # [K m/s]

            surf_flux = (
                lat_heat_flux_wm2
                + sens_heat_flux_wm2
                + longwave_heat_fluw_wm2
                + shorwave_heat_fluw_wm2 * (1 - shortwave_pen_ratio)
            ) / (4180 * rho)  # [K m/s]
            surf_flux = np.vstack((surf_flux, np.zeros_like(surf_flux)))
            pdist = 19.1  # Fitted to match average absorbance
            pen_dist = np.full_like(pen_flux, pdist)

            # Compute mixing ratios used for latent heat flux corrections
            [(_, data_airtemp_K)] = atmo2.get_data(csr, t, ["tas"])
            [(_, data_relhum_percent)] = atmo2.get_data(csr, t, ["hurs"])
            [(_, data_sst_K)] = sst.get_data(csr, t, ["analysed_sst"])

            airtemp_C = data_airtemp_K - 273.15  # [°C]
            sst_C = data_sst_K - 273.15  # [°C]

            # [(_, data_dewpoint_K)] = atmo2.get_data(csr, t, ["d2m"])
            # dewpoint_C = data_dewpoint_K - 273.15  # [°C]
            # rel_hum = rel_hu_from_dew(dewpoint_C, airtemp_C) # [0-1]
            rel_hum = data_relhum_percent / 100  # [0-1]

            vapor_press_tair = vapor_pressure(airtemp_C)
            vapor_press_sst = vapor_pressure(sst_C)
            air_mix_ratio = (
                rel_hum * 0.622 * vapor_press_tair / (1013 - 0.378 * vapor_press_tair)
            )
            sat_mix_ratio = (
                0.98 * 0.622 * vapor_press_sst / (1013 - 0.378 * vapor_press_sst)
            )

            writer_atmo.append(
                [
                    ("sst", sst_C),
                    ("airtemp", airtemp_C),
                    ("surf_flux", surf_flux),
                    ("pen_flux", pen_flux),
                    ("q_lat_wm2", lat_heat_flux_wm2),
                    ("q_sens_wm2", sens_heat_flux_wm2),
                    ("pen_dist", pen_dist),
                    ("air_mix_ratio", air_mix_ratio),
                    ("sat_mix_ratio", sat_mix_ratio),
                ],
                t,
            )

    mpi.barrier()


def prepro_freshwater_flux():
    """Convert precipitation and evaporation products to freshwater flux."""
    log("Processing Freshwater Flux...")
    atmo_vars1 = [
        "hfls",  # surface upward latent heat flux [W/m²]
        "pr",  # precipitation [kg/m²/s]
    ]
    atmo_vars2 = [
        "ts",  # surface skin temp [K]
    ]
    atmo1 = get_struct("atm", atmo_vars1)
    atmo2 = get_struct("atm", atmo_vars2)
    [atmo1, atmo2], time_combined = Struct.reinterp_at_finest_grid(
        [atmo1, atmo2], bounds=bounds
    )

    time_combined = clip_inclusive(time_combined, *timelaps)

    csr = atmo1.compute_csr_cg_2d(mesh2d)
    atmo2.compute_csr_cg_2d(mesh2d)
    with slim.DataFileWriter(
        prepro_dir + "slim4_fw", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(time_combined):
            [(_, data_tp), (_, data_lat_heat_flux)] = atmo1.get_data(
                csr, t, ["pr", "hfls"]
            )
            [(_, data_skintemp_K)] = atmo2.get_data(csr, t, ["ts"])
            lat_heat_flux_wm2 = -data_lat_heat_flux  # [W/m2]

            data_skintemp_K = data_skintemp_K
            rho = 1000
            q_lat = lat_heat_flux_wm2 / rho  # [Wm/kg]
            prec = data_tp / rho  # convert from kg/m2/s to m/s
            lat_heat_vap = (
                1.91846e6 * data_skintemp_K**2 / (data_skintemp_K - 33.91) ** 2
            )
            evap = q_lat / lat_heat_vap
            qwater = prec + evap  # [m/s]
            writer.append([("fw", qwater)], t)

    mpi.barrier()

We prepare the 3D initial conditions and open-boundary conditions. Currents are combined with tidal velocities and rotated to mesh coordinates, while temperature and salinity are interpolated from the external ocean product.

def prepro_currents():
    """Prepare 3D velocity initial and open-boundary conditions."""
    log("Processing Currents...")
    uvgrid = get_struct("hydro3d", ["u", "v"])
    tides = tpxo.GridTPXO(
        mesh2d, base_dir + "tpxo/h_tpxo_zone", base_dir + "tpxo/u_tpxo_zone"
    )
    ut = tides.tides["u"]
    vt = tides.tides["v"]

    # Initial condition
    csr_init = uvgrid.compute_csr_cg_3d(mesh, z)
    pnts = uvgrid.get_lonlat_grid(csr_init)
    rot_uv = VectorRotation(mesh2d, pnts)
    ampu, phau = ut.compute_amp_pha(pnts.T)
    ampv, phav = vt.compute_amp_pha(pnts.T)

    path_init = f"{prepro_dir}/slim4_uv_init"
    with slim.DataFileWriter(
        path_init, "CG_INTERP", csr_init.columns, csr_init.row_start, csr_init.weights
    ) as writer_init:
        for i, t in enumerate(timelaps[0:1]):
            [(_, u_base), (_, v_base)] = uvgrid.get_data(csr_init, t, ["u", "v"])
            u_tide = ut.at_points_precomputed(ampu, phau, t)
            v_tide = vt.at_points_precomputed(ampv, phav, t)
            u = u_base + u_tide * 0.01
            v = v_base + v_tide * 0.01
            u, v = rot_uv.rotate(u, v)
            uv = np.vstack([u, v])
            writer_init.append([("u", uv)], t)

    # Boundary conditions
    csr = uvgrid.compute_csr_cg_3d(mesh, z, boundaries=param.open_tags)
    pnts = uvgrid.get_lonlat_grid(csr)
    rot_uv = VectorRotation(mesh2d, pnts)
    ampu, phau = ut.compute_amp_pha(pnts.T)
    ampv, phav = vt.compute_amp_pha(pnts.T)

    with slim.DataFileWriter(
        prepro_dir + "slim4_uv", "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(time_tides):
            [(_, u_base), (_, v_base)] = uvgrid.get_data(csr, t, ["u", "v"])
            u_tide = ut.at_points_precomputed(ampu, phau, t)
            v_tide = vt.at_points_precomputed(ampv, phav, t)
            u = u_base + u_tide * 0.01
            v = v_base + v_tide * 0.01
            u, v = rot_uv.rotate(u, v)
            uv = np.vstack([u, v])
            writer.append([("u", uv)], t)

    mpi.barrier()


def prepro_temp_sal():
    """Prepare temperature and salinity initial and boundary conditions."""
    log("Processing Temperature and Salinity...")
    tsgrid = get_struct("hydro3d", ["temp", "salt"])

    # Initial condition
    csr_init = tsgrid.compute_csr_cg_3d(mesh, z)

    path_init = f"{prepro_dir}/slim4_temp_sal_init"
    with slim.DataFileWriter(
        path_init, "CG_INTERP", csr_init.columns, csr_init.row_start, csr_init.weights
    ) as writer_init:
        for i, t in enumerate(timelaps[0:1]):
            [(_, temp), (_, salt)] = tsgrid.get_data(csr_init, t, ["temp", "salt"])
            writer_init.append([("temp_sal", np.vstack([temp, salt]))], t)

    # Boundary condition
    csr = tsgrid.compute_csr_cg_3d(mesh, z, param.open_tags)

    path_open = f"{prepro_dir}/slim4_temp_sal_open"
    with slim.DataFileWriter(
        path_open, "CG_INTERP", csr.columns, csr.row_start, csr.weights
    ) as writer:
        for i, t in enumerate(clip_inclusive(tsgrid.time(), *timelaps)):
            [(_, temp), (_, salt)] = tsgrid.get_data(csr, t, ["temp", "salt"])
            writer.append([("temp_sal", np.vstack([temp, salt]))], t)

    mpi.barrier()

Finally, we execute all preprocessing steps.

prepro_coriolis()
prepro_reef_denisty()
prepro_wind()
prepro_elevation()
prepro_atmo()
prepro_freshwater_flux()
prepro_currents()
prepro_temp_sal()
log("Preprocessing Done =)")

Download this section.

Run the Great Barrier Reef simulation

This is the final step of the workflow. It loads the preprocessed mesh and forcing files produced by prepro.py, configures the 3D hydrodynamic and tracer problem, and advances the simulation in time.

The model uses open-boundary elevation, velocity, temperature, and salinity; atmospheric forcing; reef-dependent bottom roughness; turbulence closure; Smagorinsky horizontal viscosity; Okubo tracer diffusivity; and a realistic temperature-salinity equation of state.


from time import perf_counter as clock
import numpy as np
import slim4
from slim4 import slim, slim3d, mpi
from slim4.tools import xdmf_abin_shrink, parse_time, format_time, dhms
import params as param


def log(*args, **kwargs) -> None:
    if mpi.rank > 0:
        return
    print(*args, **kwargs)

We gather MPI information, define the preprocessing directory, and parse the initial time used for ramps and initial conditions.

comm_rank = mpi.rank
comm_size = mpi.size
prepro_dir = param.prepro_data_dir + "/"

# Time interval
tstart, tend = param.initial_time, param.final_time
t_init = parse_time(tstart)

We load the partitioned 3D mesh and bathymetry created by prepro.py.

log("Loading mesh file...")
ishybrid = param.n_sigma > 0 and (
    (isinstance(param.z_levels, int) and param.z_levels > 0)
    or (isinstance(param.z_levels, list) and len(param.z_levels) != 0)
)
mesh, z0, bath2d = slim3d.Mesh.load(
    prepro_dir + "mesh3d_split", n_moving_max=param.n_moving_max
)
mesh2d = mesh.mesh2d
nt = mesh2d.x.shape[-1]

We create the 3D problem and add the combined temperature-salinity tracer. The two tracer components are later exposed as named subvectors for output and for the equation of state.

# Initialize problems
g = 9.80616
rho_0 = 1000
n_elem = mesh.z_layers[-1]
run = slim3d.Problem(mesh, g, bath2d, rho_0)
problem_ts = run.add_tracer(2, "temp_sal", limiter_type="past_bound")
sw2d = run.problemUV2d
sw3d = run.problemUV

We load time-dependent open-boundary data. The 3D velocity is converted to a transport by multiplying by the moving-layer thickness associated with the boundary elevation and bathymetry. A ramp is applied at the start of the run.

# Load data for boundary conditions
DT = param.ramp
u_open = slim.DataFile(prepro_dir + "slim4_uv", "u", mesh)
u_init = slim.DataFile(prepro_dir + "slim4_uv_init", "u", mesh)
eta_open = slim.DataFile(prepro_dir + "slim4_eta", "h", mesh2d)
hu_open = slim3d.DataHMultiply(u_open, eta_open, bath2d, z0)
hu_open.set_ramp(t_init, DT)
eta_open.set_ramp(t_init, DT)
ts_open = slim.DataFile(prepro_dir + "slim4_temp_sal_open", "temp_sal", mesh)
ts_init = slim.DataFile(prepro_dir + "slim4_temp_sal_init", "temp_sal", mesh)

We apply wall boundaries and open boundaries. At open boundaries, elevation, transport, temperature, and salinity are all prescribed from the preprocessed files.

# Set boundary conditions
for tag in param.closed_tags:
    run.set_wall_boundary(tag)
for tag in param.open_tags:
    run.set_open_boundary(
        tag, eta=eta_open, u=hu_open, t={"temp_sal": ts_open}, uv_is_transport=True
    )

We initialize the prognostic fields: moving vertical coordinates, 3D velocity, and temperature-salinity tracers. The initial temperature and salinity fields are read from the preprocessed ocean product.

# Create and initialize variables
z = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv = sw3d.new_vector()
ts_tracers = problem_ts.new_vector()
temperature = ts_tracers.get_sub(0)
temperature.set_name("temperature")
salinity = ts_tracers.get_sub(1)
salinity.set_name("salinity")

# Set initial conditions for temperature and salinity
ts_tracers.set(ts_init.eval_mesh(t_init))

We apply wind stress, and optionally atmospheric pressure, from the atmospheric forcing files.

# Wind and atmospheric pressure
tau = slim.DataFile(prepro_dir + "slim4_wind", "tau", mesh2d)
wind = slim.DataFile(prepro_dir + "slim4_wind", "wind", mesh2d)
atm = slim.DataFile(prepro_dir + "slim4_wind", "psl", mesh2d)
tau.set_ramp(t_init, DT)
atm.set_ramp(t_init, DT)
run.set_surface_stress(tau=tau)
sw2d.set_atmospheric_pressure(atm)

We apply surface and penetrating heat fluxes when requested. Additional temperature-reaction terms correct the fluxes using sea-surface temperature, air temperature, humidity, and latent/sensible heat components.

# Load data
surf_heat_open = slim.DataFile(prepro_dir + "slim4_atmo", "surf_flux", mesh2d)
pen_heat_open = slim.DataFile(prepro_dir + "slim4_atmo", "pen_flux", mesh2d)
pen_dist = slim.DataFile(prepro_dir + "slim4_atmo", "pen_dist", mesh2d)

# Surface water and heat fluxes
surf_heat_open.set_ramp(t_init, DT)
pen_heat_open.set_ramp(t_init, DT)

# Set surface heat fluxes
run.set_surface_flux("temp_sal", surf_heat_open)
run.set_penetrating_flux(pen_heat_open, pen_dist, temperature, z)

# Add corrections to the surface heat fluxes :
sst = slim.DataFile(prepro_dir + "slim4_atmo", "sst", mesh2d)
airtemp = slim.DataFile(prepro_dir + "slim4_atmo", "airtemp", mesh2d)
air_mix_ratio = slim.DataFile(prepro_dir + "slim4_atmo", "air_mix_ratio", mesh2d)
sat_mix_ratio = slim.DataFile(prepro_dir + "slim4_atmo", "sat_mix_ratio", mesh2d)
q_lat = slim.DataFile(prepro_dir + "slim4_atmo", "q_lat_wm2", mesh2d)
q_sens = slim.DataFile(prepro_dir + "slim4_atmo", "q_sens_wm2", mesh2d)
reaction_fluxes = slim3d.ReactionTemperatureCorrection(
    mesh,
    sst,
    air_mix_ratio,
    sat_mix_ratio,
    q_lat,
    q_sens,
    airtemp,
    wind,
    z,
    temperature,
)
run.add_reaction(reaction_fluxes)

Static and parameterized physical terms: Coriolis, reef-dependent bottom roughness, turbulence, horizontal viscosity, tracer diffusivity, the equation of state, and velocity clipping.

# Coriolis
cor = slim4.slim.DataFile(prepro_dir + "cor", "cor", mesh2d)
run.problemUV.set_coriolis(cor)

# Disipation
reef_density = slim.DataFile(
    prepro_dir + "reef_density", "reef_density", mesh2d
).eval_mesh()
gamma = slim.DataStatic(mesh2d, "DG", 1)
base_z0b = 0.003  # roughness length for sandy bottom
hard_z0b = 0.2  # roughness length for hard bottom
gamma.set(base_z0b + reef_density * (hard_z0b - base_z0b))
sw3d.set_dissipation(gamma, log_profile=True)

# Turbulence closure for vertical viscosity and diffusivity
run.add_turbulence(z0b=gamma)

# Horizontal viscosity (SMAGO) and diffusivity (OKUBO)
smago_coeff = slim.DataStatic.create_constant(mesh, value=0.02)
okubo_coeff = slim.DataStatic.create_constant(mesh, value=2e-4)
sw3d.set_horizontal_viscosity(smago_coeff, "smago")
for name, problem_tracer in run.problem_tracer.items():
    problem_tracer.set_horizontal_diffusivity(okubo_coeff, "okubo")

# State equation for the density
run.set_state_equation_jackett(temperature, salinity)

# Clip velocity
max_uv = 3.0  # Max horizontal speed
realbath = slim.DataFile(prepro_dir + "bath", "realbath", mesh2d)
run.set_clipping_velocity(
    max_uv,
    slim.ScalarArrayDevice.create_zeros([1, 1, n_elem]),
    realbath=realbath,
)

run.set_clipping_velocity(
    3.0, slim.ScalarArrayDevice.create_zeros([1, 1, n_elem]), realbath=realbath
)

We estimate stable 3D and 2D time steps from the CFL condition, then choose a model time step that divides the requested export interval exactly.

# Compute timesteps
rho = run.get_gpu_tmpvars().rho
run.problemUV.update_density(z, rho)
dt3d, dt2d = run.compute_cfl(u_max=max_uv, drho=rho.get())
dt3d *= param.dt_mult
nsub = int(np.ceil(param.export_dt / dt3d))
dt = param.export_dt / nsub
M = dt / dt2d
log(f"dt3d = {dt:.3f}, M = {M:.2f}")

We initialize the temporal loop. In restart mode, previous save files are read and XDMF/abin outputs are shrunk so new results append cleanly.

# Initialize values for the temporal loop and load previous data if restarting
checkdt = 24
if param.restart:
    t, iout, _ = run.load(
        param.output_directory + "/saves", z, uv, [temperature, salinity]
    )
    xdmf_abin_shrink(
        param.output_directory + "/slim3D", divmod(iout, checkdt)[0] * checkdt + 1
    )
    xdmf_abin_shrink(
        param.output_directory + "/slim2D", divmod(iout, checkdt)[0] * checkdt + 1
    )
else:
    t = t_init
    iout = 0
iiter = 0

We prepare the XDMF output files and write the initial state unless this is a restart. The export options include density and velocity-clipping diagnostics.

# Create xdmf files and define exports
path3d = param.output_directory + "/slim3D"
path2d = param.output_directory + "/slim2D"
export_kwargs = dict(
    Huv=False,
    turbulence=False,
    rho=True,
    clip_loc=run.uv_clip_loc,
    nitperex=nsub,
)
if not param.restart:
    mesh.create_xdmf(z0, path3d, cg=True)
    mesh2d.create_xdmf(path2d, cg=True)
    run.do_export(
        path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
    )

mpi.barrier()
log("starting iterations")

We run the simulation. At each export interval, fields are written to disk, basic diagnostics are printed, and periodic save files are created for restart.

# Temporal loop
tic = toc = clock()
log()
while t + dt <= parse_time(tend):
    t = slim3d.iterate_impl(run, t, dt, M, 1, z, uv, [ts_tracers])
    # t = slim3d.iterate_midpoint(run, t, dt, M, 1, z, uv, [ts_tracers])

    periter = (clock() - tic) / (iiter + 1)
    if mpi.size == 1:
        log(
            f"\033[FIter {iiter + 1:d}, Time {format_time(t)}, Clock Time {dhms(clock() - tic)}, {periter:.3f} s per iter ({dt / periter:.1f}x)"
        )

    if ((iiter + 1) % nsub) == 0 or mpi.max_idx(uv.has_nan()):
        iout += 1
        run.do_export(
            path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
        )
        if max_uv is not None:
            run.uv_clip_loc.memset(0)

        uvmax = mpi.max_scalar(np.linalg.norm(uv.get(), axis=0).max())
        w_max = mpi.max_scalar(np.abs(run.get_gpu_tmpvars().w_gpu.get()).max())
        temp_cpu = temperature.get()
        sali_cpu = salinity.get()
        temp_min = mpi.min_scalar(temp_cpu.min())
        temp_max = mpi.max_scalar(temp_cpu.max())
        sali_min = mpi.min_scalar(sali_cpu.min())
        sali_max = mpi.max_scalar(sali_cpu.max())
        log(f"Output {iout}, Time {format_time(t)}, Clock time {dhms(clock() - tic)}")
        log(f"max uv = {uvmax:.2f}, max w = {w_max:.2f}")
        log(f"{temp_min:.3f} < temperature < {temp_max:.3f}")
        log(f"{sali_min:.3f} < salinity < {sali_max:.3f}")
        toc = clock()
        if iout % checkdt == 0:
            run.save(
                t,
                iout,
                iiter,
                param.output_directory + "/saves",
                z,
                uv,
                [temperature, salinity],
            )
    iiter += 1

log("Hydro done =)")

Download this section.

GBR tools

This contans additional functions used in the GBR testcase

from typing import List
import numpy as np
from slim4 import slim2d, tools
from osgeo import ogr


def recursive_rule(depth=0):
    if depth == 0:
        points = np.array(
            [
                [2 / 3, 1 / 6],
                [1 / 6, 2 / 3],
                [1 / 6, 1 / 6],
            ]
        )
        weights = np.array([1 / 6, 1 / 6, 1 / 6])
        # points = np.array([[1 / 3, 1 / 3]])
        # weights = np.array([1 / 2])
        return points, weights
    else:
        points, weights = recursive_rule(depth - 1)
        points = np.vstack(
            [
                points,
                points + np.array([1, 0]),
                points + np.array([0, 1]),
                np.array([1, 1]) - points,
            ]
        )
        weights = np.hstack([weights, weights, weights, weights])
        return points / 2, weights / 4


# bathymetry
def bathy_on_mesh(
    mesh: slim2d.Mesh,
    filename: str | List[str],
    nodal=True,
    cg=True,
    rule_depth=1,
):
    if nodal:
        xn, yn = mesh.xnodes.astype(np.float64)
        bath_cg = -tools.interpolate_tif(filename, xn, yn, mesh.projection)
        bath_dg = bath_cg[mesh.element_nodes][:, :, None].T
        return bath_dg
    else:
        x, y = mesh.x.astype(np.float64)
        points, weights = recursive_rule(rule_depth)
        xsi, eta = points[:, 0], points[:, 1]
        phi = np.array([1 - xsi - eta, xsi, eta])

        xp = np.einsum("jn,jp->np", x, phi, optimize="optimal")
        yp = np.einsum("jn,jp->np", y, phi, optimize="optimal")
        bath_data = -tools.interpolate_tif(filename, xp, yp, mesh.projection)
        bath_mass = np.einsum(
            "ep,np,p->en", bath_data, phi, weights, optimize="optimal"
        )
        imass = 24 * np.eye(3) - 6
        bath_dg = np.einsum("ej,ij->ie", bath_mass, imass, optimize="optimal")[
            None, :, :
        ]

        if cg:
            bath_cg = mesh.gather_dg_to_cg("mean_area", bath_dg)
            bath_dg = bath_cg[0][mesh.element_nodes][:, :, None].T
        return bath_dg


def shapefile_to_flat_gpkg(input_shp, output_gpkg, output_layer_name, bounds):
    """
    This is mostly ChatGPT generated code
    """
    minlon, maxlon, minlat, maxlat = bounds

    src_ds = ogr.Open(input_shp)
    if src_ds is None:
        raise RuntimeError(f"Could not open {input_shp}")

    src_lyr = src_ds.GetLayer()
    src_srs = src_lyr.GetSpatialRef()

    # Build bbox polygon
    ring = ogr.Geometry(ogr.wkbLinearRing)
    ring.AddPoint(minlon, minlat)
    ring.AddPoint(maxlon, minlat)
    ring.AddPoint(maxlon, maxlat)
    ring.AddPoint(minlon, maxlat)
    ring.AddPoint(minlon, minlat)

    bbox_geom = ogr.Geometry(ogr.wkbPolygon)
    bbox_geom.AddGeometry(ring)

    # Coarse filter
    src_lyr.SetSpatialFilterRect(minlon, minlat, maxlon, maxlat)
    drv = ogr.GetDriverByName("GPKG")
    try:
        drv.DeleteDataSource(output_gpkg)
    except Exception:
        pass

    def polygon_parts(geom):
        if geom is None:
            return
        name = geom.GetGeometryName().upper()
        if name == "POLYGON":
            yield geom.Clone()
        elif name == "MULTIPOLYGON":
            for i in range(geom.GetGeometryCount()):
                part = geom.GetGeometryRef(i)
                if part is not None:
                    yield part.Clone()

    all_parts = ogr.Geometry(ogr.wkbMultiPolygon)
    count = 0
    for src_feat in src_lyr:
        geom = src_feat.GetGeometryRef()
        if geom is None:
            continue
        for part in polygon_parts(geom):
            if part.Within(bbox_geom):
                all_parts.AddGeometry(part)
                count += 1
    all_parts = all_parts.UnionCascaded()

    dst_ds = drv.CreateDataSource(output_gpkg)
    dst_lyr = dst_ds.CreateLayer(
        output_layer_name, srs=src_srs, geom_type=ogr.wkbMultiPolygon
    )
    if count > 0:
        dst_defn = dst_lyr.GetLayerDefn()
        out_feat = ogr.Feature(dst_defn)
        out_feat.SetGeometry(all_parts)

        if dst_lyr.CreateFeature(out_feat) != 0:
            raise RuntimeError("Failed to create output feature")
        out_feat = None
    src_ds = None
    dst_ds = None
    print(f"Wrote {output_gpkg} with {count} polygons in one feature")


def poly_to_line(gpkg_in, gpkg_out, physical_name):
    src = ogr.Open(gpkg_in)
    lyr = src.GetLayer()
    drv = ogr.GetDriverByName("GPKG")
    dst = drv.CreateDataSource(gpkg_out)
    out = dst.CreateLayer("lines", lyr.GetSpatialRef(), ogr.wkbMultiLineString)
    out.CreateField(ogr.FieldDefn("physical", ogr.OFTString))
    defn = out.GetLayerDefn()
    for f in lyr:
        g = f.GetGeometryRef()
        b = g.Boundary()
        if b.GetGeometryType() == ogr.wkbLineString:
            m = ogr.Geometry(ogr.wkbMultiLineString)
            m.AddGeometry(b)
            b = m
        nf = ogr.Feature(defn)
        nf.SetGeometry(b)
        nf.SetField("physical", physical_name)
        out.CreateFeature(nf)

Download this section.