#!/usr/bin/env python3 # Extract the deepest available SOCOML pH and aragonite value in each SOCOML # grid cell and output on the native EPSG:4326 grid. # # IBCSO is used only to: # - record median seafloor depth within each output cell # - exclude cells deeper than the deepest SOCOML level + configured buffer # # Inputs: # input/SOCOML_Float_GRID_v2025.nc # input/*.tif <-- IBCSO bathymetry # # Outputs: # output/SOCOML_pH_seafloor.tif # output/SOCOML_Aragonite_seafloor.tif # # Install requirements: # pip install numpy xarray netCDF4 rasterio from pathlib import Path import os import numpy as np import rasterio import xarray as xr from rasterio.enums import Resampling from rasterio.transform import from_origin from rasterio.warp import reproject # ========================= # CONFIGURATION # ========================= INPUT_DIR = Path("input") OUTPUT_DIR = Path("output") INPUT_NETCDF = INPUT_DIR / "SOCOML_Float_GRID_v2025.nc" TARGET_CRS = "EPSG:4326" NODATA = -999.0 SEAFLOOR_BUFFER_M = 200 NUM_THREADS = max(1, os.cpu_count() or 1) OUTPUTS = { "pH": { "uncer": "pH_UNCER", "file": "SOCOML_pH_seafloor.tif", "bands": [ "seafloor_pH", "seafloor_pH_uncertainty", "SOCOML_depth_m", "IBCSO_depth_m", ], }, "OMEGA_Ar": { "uncer": "OMEGA_Ar_UNCER", "file": "SOCOML_Aragonite_seafloor.tif", "bands": [ "seafloor_Ar", "seafloor_Ar_uncertainty", "SOCOML_depth_m", "IBCSO_depth_m", ], }, } # ========================= # FUNCTIONS # ========================= def find_ibcso(): tifs = list(INPUT_DIR.glob("*.tif")) + list(INPUT_DIR.glob("*.tiff")) if len(tifs) != 1: raise RuntimeError(f"Expected exactly one TIFF in {INPUT_DIR}; found {len(tifs)}.") return tifs[0] def get_nodata(var): nodata = var.attrs.get("_FillValue", var.attrs.get("missing_value", var.encoding.get("_FillValue"))) if nodata is None: raise RuntimeError(f"NoData not defined for {var.name}.") return float(nodata) def get_coord(ds, names): for name in names: if name in ds.variables: return ds[name].values.astype(np.float64) raise RuntimeError(f"Could not find coordinate: {', '.join(names)}") def socoml_grid(ds): lon = get_coord(ds, ["LONGITUDE", "longitude", "lon"]) lat = get_coord(ds, ["LATITUDE", "latitude", "lat"]) xres = abs(np.median(np.diff(lon))) yres = abs(np.median(np.diff(lat))) if not np.allclose(abs(np.diff(lon)), xres) or not np.allclose(abs(np.diff(lat)), yres): raise RuntimeError("SOCOML longitude/latitude grid is not regularly spaced.") transform = from_origin(lon.min() - xres / 2, lat.max() + yres / 2, xres, yres) return transform, len(lon), len(lat) def median_ibcso(path, transform, width, height): with rasterio.open(path) as src: if src.crs is None: raise RuntimeError("IBCSO raster has no CRS.") if src.nodata is None: raise RuntimeError("IBCSO raster has no NoData value.") print(f"IBCSO CRS: {src.crs}") print(f"IBCSO NoData: {src.nodata}") bathy = np.full((height, width), np.nan, dtype=np.float32) reproject( rasterio.band(src, 1), bathy, src_transform=src.transform, src_crs=src.crs, src_nodata=src.nodata, dst_transform=transform, dst_crs=TARGET_CRS, dst_nodata=np.nan, resampling=Resampling.med, num_threads=NUM_THREADS, ) ocean = np.isfinite(bathy) & (bathy < 0) depth = np.where(ocean, -bathy, np.nan).astype(np.float32) return depth, ocean def source_slice(var, index): nodata = get_nodata(var) arr = var.isel(depth_level=index).transpose("lat", "lon").values.astype(np.float32) return np.where(np.isfinite(arr) & (arr != nodata), arr, NODATA).astype(np.float32) def write_tif(path, arrays, names, transform, width, height): profile = { "driver": "GTiff", "width": width, "height": height, "count": 4, "dtype": "float32", "crs": TARGET_CRS, "transform": transform, "nodata": NODATA, "compress": "DEFLATE", "tiled": True, "BIGTIFF": "IF_SAFER", } with rasterio.open(path, "w", **profile) as dst: for i, (arr, name) in enumerate(zip(arrays, names), 1): dst.write(arr, i) dst.set_band_description(i, name) # ========================= # MAIN # ========================= def main(): if not INPUT_NETCDF.exists(): raise FileNotFoundError(f"NetCDF not found: {INPUT_NETCDF}") OUTPUT_DIR.mkdir(exist_ok=True) ibcso_path = find_ibcso() print(f"SOCOML: {INPUT_NETCDF}") print(f"IBCSO : {ibcso_path}") with xr.open_dataset(INPUT_NETCDF, mask_and_scale=False) as ds: transform, width, height = socoml_grid(ds) depths = ds["DEPTH"].values.astype(np.float32) max_ibcso_depth = depths[-1] + SEAFLOOR_BUFFER_M print(f"Output grid: {width} x {height} ({TARGET_CRS})") print(f"Transform: {transform}") print(f"SOCOML depth range: {depths[0]:g}-{depths[-1]:g} m") print(f"Maximum accepted IBCSO depth: {max_ibcso_depth:g} m") ibcso_depth, ocean = median_ibcso(ibcso_path, transform, width, height) ibcso_ok = ocean & (ibcso_depth <= max_ibcso_depth) for value_name, cfg in OUTPUTS.items(): uncer_name = cfg["uncer"] filename = cfg["file"] print(f"\nProcessing {filename}") value = np.full((height, width), NODATA, dtype=np.float32) uncer = np.full((height, width), NODATA, dtype=np.float32) socoml_depth = np.full((height, width), NODATA, dtype=np.float32) assigned = np.zeros((height, width), dtype=bool) for i in range(len(depths) - 1, -1, -1): depth = depths[i] layer_value = source_slice(ds[value_name], i) fill = ibcso_ok & ~assigned & (layer_value != NODATA) if not np.any(fill): continue layer_uncer = source_slice(ds[uncer_name], i) value[fill] = layer_value[fill] socoml_depth[fill] = depth assigned[fill] = True uncer_ok = fill & (layer_uncer != NODATA) uncer[uncer_ok] = layer_uncer[uncer_ok] print(f" {depth:g} m -> {np.count_nonzero(fill):,} cells") value[~assigned] = NODATA uncer[~assigned] = NODATA socoml_depth[~assigned] = NODATA output_ibcso_depth = np.where(assigned, ibcso_depth, NODATA).astype(np.float32) print(f"Output cells: {np.count_nonzero(assigned):,}") print(f"IBCSO > {max_ibcso_depth:g} m excluded: " f"{np.count_nonzero(ocean & (ibcso_depth > max_ibcso_depth)):,}") output_path = OUTPUT_DIR / filename write_tif( output_path, [value, uncer, socoml_depth, output_ibcso_depth], cfg["bands"], transform, width, height, ) print(f"Created: {output_path}") print("\nDone.") if __name__ == "__main__": main()