diff --git a/.gitignore b/.gitignore index 330b89ca..451bb15b 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,7 @@ grib2 *.nc GRIBFILE.* *.o -FILE:* \ No newline at end of file +FILE:* +WPS_GEOG +install +QNWFA_QNIFA_QNBCA_SIGMA_MONTHLY.dat diff --git a/link_grib.py b/link_grib.py new file mode 100755 index 00000000..daa4253f --- /dev/null +++ b/link_grib.py @@ -0,0 +1,57 @@ +#! /usr/bin/python + +# This script creates symbolic links for a list of grib files. + +import sys +import glob +import os + +def link_grib_file(filelist): + + # This function creates the links for a list of grib files. + + n = 0 + for ifile in filelist: + if ifile!=".": + file_suffix = file_alphabet(n) + os.symlink(ifile, "GRIBFILE."+file_suffix) + n = n + 1 + +def file_alphabet(number): + + # This function converts a number of a file to its three-letter number. + + alpha = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z'] + + if number>26*26*26: + print("Error: ran out of grib file suffixes.") + sys.exit(1) + else: + i1 = number%26 + i2 = (number - i1)//26%26 + i3 = (number - i2*26 - i1)//26//26%26 + + return alpha[i3] + alpha[i2] + alpha[i1] + +if __name__=="__main__": + + varlist = sys.argv[1:] + + if len(varlist)==1 or len(varlist)==2 and varlist[1]=='.': + os.system("rm -f GRIBFILE.???") + filelist = glob.glob(varlist[0]) + link_grib_file(filelist) + elif len(varlist)>1: + os.system("rm -f GRIBFILE.???") + filelist = varlist + link_grib_file(filelist) + elif len(varlist)==0: + print("") + print("Please provide some GRIB data to link") + print("usage: ./link_grib.py path_to_grib_data/*") + print("") + + + + + diff --git a/ungrib/py/.gitignore b/ungrib/py/.gitignore new file mode 100644 index 00000000..0d20b648 --- /dev/null +++ b/ungrib/py/.gitignore @@ -0,0 +1 @@ +*.pyc diff --git a/ungrib/py/MIGRATION.md b/ungrib/py/MIGRATION.md new file mode 100644 index 00000000..44dee1cc --- /dev/null +++ b/ungrib/py/MIGRATION.md @@ -0,0 +1,27 @@ +# Fortran ungrib → `ungrib_py` + +## Done here + +- Vtable (GRIB2 columns), ecCodes GRIB2 scan, common level types vs `rd_grib2.F`. +- Grids: `regular_ll`, Lambert, polar stereographic, Mercator, **NCEP GDT 3.32769** (`ncep_32769`, WPS `igrid=6`). +- WPS intermediate **v5** (`out_format` WPS) aligned with `metgrid/src/read_met_module.F`. +- Parallel **GRIBFILE.*** workers, merge by time, CLI **`--workers`**. +- Minimal **`namelist.wps`** (`start`/`end`, `interval_seconds`, `prefix`, `pmin`, …); many `&ungrib` options still ignored. + +## Not done / gaps + +- **`rrpr` / `datint` / `PFILE:`** second pass — Python writes final `FILE:` from merged messages. +- **GRIB1** — partial; may differ from `rd_grib1.F` on edge cases. +- **Other gaps**: JMA per-field grid refresh, Gaussian `igrid=4`, UM soil specials, `add_lvls`, EC record quirks, regression harness vs `ungrib.exe`. + +## Test + +```bash +cd ungrib/py && pip install -r requirements.txt && pytest tests/ -q +``` + +Integration: same `namelist.wps` / `Vtable` / `GRIBFILE.*` as Fortran; compare with `rd_intermediate.exe` or slab diffs. + +## Practical success + +Same field names/levels and **metgrid-acceptable** v5 files as Fortran for the same GRIB2 + Vtable on supported grids; speedup when I/O allows multiple workers and split inputs. diff --git a/ungrib/py/README.md b/ungrib/py/README.md new file mode 100644 index 00000000..09274b70 --- /dev/null +++ b/ungrib/py/README.md @@ -0,0 +1,42 @@ +# Python ungrib (`ungrib/py`) + +Parallel Python port of WPS **ungrib**: GRIB2 → WPS intermediate v5 for **metgrid**, with clearer structure than `ungrib/src/` Fortran. + +## Fortran ungrib (reference) + +| Piece | Role | +|-------|------| +| `ungrib.F` | Main: `GRIBFILE.*` loop, storage, `output`, `rrpr` / `datint` | +| `read_namelist.F` | `namelist.wps` (`&share`, `&ungrib`) | +| `parse_table.F`, `table.F` | **Vtable** → WPS names / levels | +| `rd_grib1.F`, `rd_grib2.F` | Decode GRIB, match Vtable, grid info | +| `output.F` | WPS intermediate (v5 when `out_format='WPS'`) | + +## This package (`ungrib_py`) + +- One worker per **`GRIBFILE.*`**; parent merges by validity time; one **`prefix:YYYY-MM-DD_HH`** file per step (hourly pattern matches Fortran). +- **GRIB2 only** via **ecCodes** (`pip install eccodes`; OS `libeccodes` often required). +- Scope: Vtable-driven extraction and v5 I/O for common grids—not full Fortran parity (`rrpr`/`datint`/GRIB1 edge cases, etc.). Details: **`MIGRATION.md`**. + +## Quick start + +```bash +pip install --user -r /path/to/WPS/ungrib/py/requirements.txt +ln -sf /path/to/WPS/ungrib/py/ungrib.py ./ungrib.py && chmod +x ./ungrib.py +./ungrib.py --workers 16 +``` + +Run from the directory containing `namelist.wps`, `Vtable`, and `GRIBFILE.AAA`, … Optional: `pip install /path/to/WPS/ungrib/py` for the `ungrib-py` console script. + +## Layout + +| Path | Role | +|------|------| +| `ungrib.py` | Launcher (`sys.path` + CLI); symlink like `ungrib.exe` | +| `ungrib_py/vtable.py` | Vtable parse / GRIB2 match | +| `ungrib_py/grib2.py` | ecCodes scan, levels, `extract_file` | +| `ungrib_py/intermediate.py` | WPS v5 writer | +| `ungrib_py/namelist_wps.py` | `&share`, `&ungrib` | +| `ungrib_py/parallel_ungrib.py` | Pool + merge + write | +| `ungrib_py/cli.py` | CLI | +| `tests/` | `pytest` | diff --git a/ungrib/py/pyproject.toml b/ungrib/py/pyproject.toml new file mode 100644 index 00000000..05a3918e --- /dev/null +++ b/ungrib/py/pyproject.toml @@ -0,0 +1,21 @@ +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[project] +name = "ungrib-py" +version = "0.1.0" +description = "Parallel GRIB2 to WPS intermediate (Python ungrib core)" +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "numpy>=1.22", + "eccodes>=1.7.0", +] + +[project.scripts] +ungrib-py = "ungrib_py.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["ungrib_py*"] diff --git a/ungrib/py/requirements.txt b/ungrib/py/requirements.txt new file mode 100644 index 00000000..18057f91 --- /dev/null +++ b/ungrib/py/requirements.txt @@ -0,0 +1,3 @@ +numpy>=1.22 +eccodes>=1.7.0 +pytest>=7.0 diff --git a/ungrib/py/tests/conftest.py b/ungrib/py/tests/conftest.py new file mode 100644 index 00000000..249be81d --- /dev/null +++ b/ungrib/py/tests/conftest.py @@ -0,0 +1,16 @@ +import sys +from pathlib import Path + +_root = Path(__file__).resolve().parents[1] +if str(_root) not in sys.path: + sys.path.insert(0, str(_root)) + + + + + + + + + + diff --git a/ungrib/py/tests/test_ecmwf_soil151.py b/ungrib/py/tests/test_ecmwf_soil151.py new file mode 100644 index 00000000..c116b989 --- /dev/null +++ b/ungrib/py/tests/test_ecmwf_soil151.py @@ -0,0 +1,48 @@ +from pathlib import Path + +import eccodes +import pytest + +from ungrib_py.grib2 import compute_wps_level, extract_file +from ungrib_py.vtable import parse_vtable + +_REPO = Path(__file__).resolve().parents[3] +_IFS_GRIB = _REPO / "ungrib_out" / "data" / "ifs.2026041106.f000.grib2" +_IFS_VT = _REPO / "ungrib_out" / "Vtable" + + +def test_compute_wps_level_soil_type151_200100(): + + # This function checks GRIB2 fixed-surface 151 (ECMWF IFS soil layer) maps to WPS surface code 200100. + + if not _IFS_GRIB.is_file(): + pytest.skip("ungrib_out sample IFS GRIB not present") + with _IFS_GRIB.open("rb") as fh: + while True: + gid = eccodes.codes_grib_new_from_file(fh) + if gid is None: + break + try: + if eccodes.codes_get(gid, "shortName") != "sot": + continue + assert compute_wps_level(gid, 50000.0) == 200100.0 + return + finally: + eccodes.codes_release(gid) + pytest.fail("expected at least one sot field in sample GRIB") + + +@pytest.mark.skipif(not _IFS_GRIB.is_file() or not _IFS_VT.is_file(), reason="IFS sample or Vtable missing") +def test_extract_ifs_soil_st_sm_four_layers(): + + # This function checks ECMWF IFS layer-151 soil temperature and moisture decode to four ST/SM fields each. + + _, vt_out = parse_vtable(_IFS_VT) + out, _ = extract_file(str(_IFS_GRIB), vt_out, pmin=50000.0) + assert out + bucket = next(iter(out.values())) + st = sorted({k[1] for k in bucket if k[1].startswith("ST")}) + sm = sorted({k[1] for k in bucket if k[1].startswith("SM")}) + assert st == ["ST000007", "ST007028", "ST028100", "ST100289"] + assert sm == ["SM000007", "SM007028", "SM028100", "SM100289"] + diff --git a/ungrib/py/tests/test_grib2_alias.py b/ungrib/py/tests/test_grib2_alias.py new file mode 100644 index 00000000..985baebb --- /dev/null +++ b/ungrib/py/tests/test_grib2_alias.py @@ -0,0 +1,96 @@ +import numpy as np + +from ungrib_py.grib2 import ensure_soilt000_for_metgrid +from ungrib_py.intermediate import FieldSlab + + +def test_ensure_soilt000_from_soilt001(): + + # This function checks SOILT000 is created from SOILT001 when missing. + + lev = 200100.0 + slab = np.asfortranarray(np.ones((2, 3), dtype=np.float32)) + src = FieldSlab( + name="SOILT001", + units="K".ljust(25), + desc="soil t", + level=lev, + data=slab, + ) + bucket: dict[tuple[float, str], FieldSlab] = {(lev, "SOILT001"): src} + ensure_soilt000_for_metgrid(bucket) + assert (lev, "SOILT000") in bucket + dup = bucket[(lev, "SOILT000")] + assert dup.name.strip() == "SOILT000" + assert np.array_equal(dup.data, src.data) + + +def test_ensure_soilt000_prefers_shallowest(): + + # This function checks SOILT000 is taken from SOILT002 when SOILT001 is absent. + + lev = 200100.0 + a = np.asfortranarray(np.full((2, 3), 2.0, dtype=np.float32)) + b = np.asfortranarray(np.full((2, 3), 3.0, dtype=np.float32)) + bucket = { + (lev, "SOILT002"): FieldSlab( + name="SOILT002", + units="K".ljust(25), + desc="x", + level=lev, + data=a, + ), + (lev, "SOILT006"): FieldSlab( + name="SOILT006", + units="K".ljust(25), + desc="y", + level=lev, + data=b, + ), + } + ensure_soilt000_for_metgrid(bucket) + assert (lev, "SOILT000") in bucket + assert np.array_equal(bucket[(lev, "SOILT000")].data, a) + + +def test_ensure_soilt000_skips_if_present(): + + # This function checks an existing SOILT000 is not overwritten. + + lev = 200100.0 + a = np.asfortranarray(np.zeros((2, 3), dtype=np.float32)) + b = np.asfortranarray(np.ones((2, 3), dtype=np.float32)) + existing = FieldSlab( + name="SOILT000", + units="K".ljust(25), + desc="x", + level=lev, + data=a, + ) + fld = FieldSlab( + name="SOILT001", + units="K".ljust(25), + desc="y", + level=lev, + data=b, + ) + bucket = {(lev, "SOILT000"): existing, (lev, "SOILT001"): fld} + ensure_soilt000_for_metgrid(bucket) + assert bucket[(lev, "SOILT000")] is existing + + +def test_ensure_soilt000_noop_without_soilt(): + + # This function checks that buckets without SOILT* fields are unchanged. + + bucket: dict[tuple[float, str], FieldSlab] = { + (200100.0, "TT"): FieldSlab( + name="TT", + units="K".ljust(25), + desc="t", + level=100000.0, + data=np.asfortranarray(np.ones((2, 3), dtype=np.float32)), + ) + } + ensure_soilt000_for_metgrid(bucket) + assert len(bucket) == 1 diff --git a/ungrib/py/tests/test_intermediate.py b/ungrib/py/tests/test_intermediate.py new file mode 100644 index 00000000..2ca189bb --- /dev/null +++ b/ungrib/py/tests/test_intermediate.py @@ -0,0 +1,70 @@ +import struct + +import numpy as np + +from ungrib_py.intermediate import FieldSlab, MapInfo, sort_fields_fortran_order, write_wps_intermediate_v5 +from ungrib_py.vtable import VtableEntry + +def _read_fortran_record_be_int32(fh) -> int: + + # This function reads one sequential record (big-endian length tags) payload as big-endian int32. + + n = int.from_bytes(fh.read(4), byteorder="big") + payload = fh.read(n) + fh.read(4) + return struct.unpack(">i", payload)[0] + +def test_write_then_read_version_and_dims(tmp_path): + + # This function checks WPS v5 intermediate writes a leading version record of 5 (big-endian payload). + + path = tmp_path / "out.bin" + m = MapInfo( + source="test".ljust(32), + igrid=0, + nx=2, + ny=3, + startloc="SWCORNER", + lat1=20.0, + lon1=-120.0, + dx=1.0, + dy=-1.0, + lov=0.0, + truelat1=0.0, + truelat2=0.0, + r_earth_km=6371.0, + grid_wind=True, + centerlat=0.0, + centerlon=0.0, + ) + slab = np.arange(6, dtype=np.float32).reshape(2, 3, order="F") + fld = FieldSlab( + name="TT", + units="K", + desc="Temperature", + level=50000.0, + data=slab, + ) + write_wps_intermediate_v5(str(path), "2020-01-01_00:00:00", m, [fld]) + with open(path, "rb") as f: + assert _read_fortran_record_be_int32(f) == 5 + + +def test_sort_matches_fortran_level_then_vtable_order(): + + # This function checks output order: descending level, then first Vtable row index per name. + + rows = [ + VtableEntry(0, 0, 0.0, 0.0, "AA ", "u".ljust(25), "d".ljust(46), 0, 0, 0, 100, 0), + VtableEntry(0, 0, 0.0, 0.0, "BB ", "u".ljust(25), "d".ljust(46), 0, 0, 1, 100, 0), + ] + sa = np.float32(0) + sb = np.float32(0) + items = [ + (100.0, "BB", FieldSlab("BB", "u", "d", 100.0, np.array([[sb]]))), + (500.0, "AA", FieldSlab("AA", "u", "d", 500.0, np.array([[sa]]))), + (500.0, "BB", FieldSlab("BB", "u", "d", 500.0, np.array([[sb]]))), + (100.0, "AA", FieldSlab("AA", "u", "d", 100.0, np.array([[sa]]))), + ] + got = [f.name.strip() for f in sort_fields_fortran_order(items, rows)] + assert got == ["AA", "BB", "AA", "BB"] diff --git a/ungrib/py/tests/test_namelist.py b/ungrib/py/tests/test_namelist.py new file mode 100644 index 00000000..607dfb24 --- /dev/null +++ b/ungrib/py/tests/test_namelist.py @@ -0,0 +1,41 @@ +from pathlib import Path + +from ungrib_py.namelist_wps import parse_namelist_wps + +def test_parse_minimal_namelist(tmp_path: Path): + + # This function verifies parse_namelist_wps on a minimal share and ungrib namelist. + + p = tmp_path / "namelist.wps" + p.write_text( + """ +&share + start_date = '2020-01-01_00:00:00', + end_date = '2020-01-01_06:00:00', + interval_seconds = 21600, + debug_level = 0, +/ +&ungrib + out_format = 'WPS', + prefix = 'FILE', + ordered_by_date = .true., + pmin = 100., +/ +""", + encoding="ascii", + ) + n = parse_namelist_wps(p) + assert n.hstart.startswith("2020-01-01") + assert n.interval_seconds == 21600 + assert n.out_format == "WPS" + assert n.pmin == 100.0 + + + + + + + + + + diff --git a/ungrib/py/tests/test_pressure_levels.py b/ungrib/py/tests/test_pressure_levels.py new file mode 100644 index 00000000..aa0c18ca --- /dev/null +++ b/ungrib/py/tests/test_pressure_levels.py @@ -0,0 +1,32 @@ +import numpy as np + +from ungrib_py.intermediate import FieldSlab +from ungrib_py.parallel_ungrib import count_distinct_pressure_levels_pa + + +def test_count_distinct_pressure_levels_excludes_surface(): + + # This function checks that 200100-style codes are not counted as isobaric levels. + + slab = FieldSlab( + name="TT", + units="K".ljust(25), + desc="t", + level=85000.0, + data=np.asfortranarray(np.ones((2, 3), dtype=np.float32)), + ) + sfc = FieldSlab( + name="TT", + units="K".ljust(25), + desc="2m", + level=200100.0, + data=np.asfortranarray(np.ones((2, 3), dtype=np.float32)), + ) + merged = { + "2025-07-13_00:00:00": { + (85000.0, "TT"): slab, + (70000.0, "UU"): slab, + (200100.0, "TT"): sfc, + } + } + assert count_distinct_pressure_levels_pa(merged) == 2 diff --git a/ungrib/py/tests/test_soilhgt_rrpr.py b/ungrib/py/tests/test_soilhgt_rrpr.py new file mode 100644 index 00000000..5975817a --- /dev/null +++ b/ungrib/py/tests/test_soilhgt_rrpr.py @@ -0,0 +1,269 @@ +import numpy as np + +from ungrib_py.grib2 import ( + drop_fields_blank_vtable_desc, + ensure_hgt_from_geopt, + ensure_soilhgt_from_soilgeo, +) +from ungrib_py.intermediate import FieldSlab +from ungrib_py.vtable import BLANK, VtableEntry + + +def test_soilhgt_derived_from_soilgeo(): + + # This function checks rrpr-style SOILHGT = SOILGEO / 9.81. + + slab = np.asfortranarray(np.full((2, 3), 98.1, dtype=np.float32)) + geo = FieldSlab( + name="SOILGEO", + units="m2 s-2".ljust(25), + desc="x", + level=200100.0, + data=slab, + ) + bucket: dict[tuple[float, str], FieldSlab] = {(200100.0, "SOILGEO"): geo} + ensure_soilhgt_from_soilgeo(bucket) + assert (200100.0, "SOILHGT") in bucket + assert (200100.0, "SOILGEO") not in bucket + h = bucket[(200100.0, "SOILHGT")] + assert abs(float(h.data.flat[0]) - 10.0) < 1e-5 + + +def test_soilhgt_skips_when_present(): + + # This function checks existing SOILHGT is not overwritten. + + a = np.asfortranarray(np.ones((2, 3), dtype=np.float32)) + existing = FieldSlab( + name="SOILHGT", + units="m".ljust(25), + desc="y", + level=200100.0, + data=a, + ) + geo = FieldSlab( + name="SOILGEO", + units="m2 s-2".ljust(25), + desc="z", + level=200100.0, + data=np.asfortranarray(np.full((2, 3), 98.1, dtype=np.float32)), + ) + bucket = {(200100.0, "SOILHGT"): existing, (200100.0, "SOILGEO"): geo} + ensure_soilhgt_from_soilgeo(bucket) + assert bucket[(200100.0, "SOILHGT")] is existing + assert (200100.0, "SOILGEO") not in bucket + + +def test_soilhgt_from_soilgeo_level_one(): + + # This function matches rrpr hybrid-level storage at level 1. + + slab = np.asfortranarray(np.full((2, 3), 98.1, dtype=np.float32)) + geo = FieldSlab( + name="SOILGEO", + units="m2 s-2".ljust(25), + desc="x", + level=1.0, + data=slab, + ) + bucket = {(1.0, "SOILGEO"): geo} + ensure_soilhgt_from_soilgeo(bucket) + assert (200100.0, "SOILHGT") in bucket + assert (1.0, "SOILGEO") not in bucket + + +def test_soilhgt_from_geopt_surface(): + + # This function checks GEOPT at surface is accepted like SOILGEO (rrpr divides by g). + + slab = np.asfortranarray(np.full((2, 3), 98.1, dtype=np.float32)) + geo = FieldSlab( + name="GEOPT", + units="m2 s-2".ljust(25), + desc="gp", + level=200100.0, + data=slab, + ) + bucket = {(200100.0, "GEOPT"): geo} + ensure_soilhgt_from_soilgeo(bucket) + assert (200100.0, "GEOPT") not in bucket + assert abs(float(bucket[(200100.0, "SOILHGT")].data.flat[0]) - 10.0) < 1e-5 + + +def test_soilhgt_from_hgt_meters(): + + # This function checks surface HGT in meters can be copied when geopotential is absent. + + slab = np.asfortranarray(np.full((2, 3), 500.0, dtype=np.float32)) + h = FieldSlab( + name="HGT", + units="m".ljust(25), + desc="h", + level=200100.0, + data=slab, + ) + bucket = {(200100.0, "HGT"): h} + ensure_soilhgt_from_soilgeo(bucket) + assert abs(float(bucket[(200100.0, "SOILHGT")].data.flat[0]) - 500.0) < 1e-5 + + +def test_hgt_from_geopt_pressure_levels(): + + # rrpr: HGT = GEOPT / g on isobaric levels (metgrid reads HGT as GHT). + + slab = np.asfortranarray(np.full((2, 2), 9810.0, dtype=np.float32)) + geopt = FieldSlab( + name="GEOPT", + units="m2 s-2".ljust(25), + desc="gp", + level=50000.0, + data=slab, + ) + bucket: dict[tuple[float, str], FieldSlab] = {(50000.0, "GEOPT"): geopt} + ensure_hgt_from_geopt(bucket) + assert (50000.0, "GEOPT") not in bucket + assert (50000.0, "HGT") in bucket + assert abs(float(bucket[(50000.0, "HGT")].data.flat[0]) - 1000.0) < 1e-3 + + +def test_hgt_from_geopt_uses_vtable_hgt_row(): + + # Same Vtable as ungrib (e.g. ungrib_out/Vtable): HGT row supplies units/desc for derived slab. + + hgt_row = VtableEntry( + g1_param=156, + g1_level_type=100, + level1=float(BLANK), + level2=float(BLANK), + name="HGT".ljust(9)[:9], + units="m".ljust(25)[:25], + desc="Custom height desc".ljust(46)[:46], + g2_discipline=0, + g2_category=3, + g2_parameter=5, + g2_level_type=100, + g2_pdt=0, + ) + slab = np.asfortranarray(np.full((2, 2), 9810.0, dtype=np.float32)) + geopt = FieldSlab( + name="GEOPT", + units="m2 s-2".ljust(25), + desc="gp", + level=50000.0, + data=slab, + ) + bucket: dict[tuple[float, str], FieldSlab] = {(50000.0, "GEOPT"): geopt} + ensure_hgt_from_geopt(bucket, vtable_rows=[hgt_row]) + h = bucket[(50000.0, "HGT")] + assert "Custom height desc" in h.desc + assert h.units.strip() == "m" + + +def test_geopt_dropped_when_hgt_present(): + + g = FieldSlab( + name="GEOPT", + units="m2 s-2".ljust(25), + desc="gp", + level=85000.0, + data=np.asfortranarray(np.ones((2, 2), dtype=np.float32)), + ) + h = FieldSlab( + name="HGT", + units="m".ljust(25), + desc="Height".ljust(46), + level=85000.0, + data=np.asfortranarray(np.full((2, 2), 99.0, dtype=np.float32)), + ) + bucket = {(85000.0, "GEOPT"): g, (85000.0, "HGT"): h} + ensure_hgt_from_geopt(bucket) + assert (85000.0, "GEOPT") not in bucket + assert bucket[(85000.0, "HGT")] is h + + +def test_drop_fields_only_when_all_vtable_rows_blank_desc(): + + dewpt = VtableEntry( + g1_param=168, + g1_level_type=1, + level1=0.0, + level2=0.0, + name="DEWPT".ljust(9)[:9], + units="K".ljust(25)[:25], + desc=" " * 46, + g2_discipline=0, + g2_category=0, + g2_parameter=6, + g2_level_type=103, + g2_pdt=0, + ) + tt = VtableEntry( + g1_param=168, + g1_level_type=1, + level1=0.0, + level2=0.0, + name="TT".ljust(9)[:9], + units="K".ljust(25)[:25], + desc="Temperature".ljust(46)[:46], + g2_discipline=0, + g2_category=0, + g2_parameter=0, + g2_level_type=103, + g2_pdt=0, + ) + uu_sfc = VtableEntry( + g1_param=165, + g1_level_type=1, + level1=0.0, + level2=0.0, + name="UU".ljust(9)[:9], + units="m s-1".ljust(25)[:25], + desc=" " * 46, + g2_discipline=0, + g2_category=0, + g2_parameter=2, + g2_level_type=103, + g2_pdt=0, + ) + uu_pl = VtableEntry( + g1_param=131, + g1_level_type=100, + level1=float(BLANK), + level2=float(BLANK), + name="UU".ljust(9)[:9], + units="m s-1".ljust(25)[:25], + desc="U".ljust(46)[:46], + g2_discipline=0, + g2_category=0, + g2_parameter=2, + g2_level_type=100, + g2_pdt=0, + ) + slab = np.asfortranarray(np.ones((2, 2), dtype=np.float32)) + bucket: dict[tuple[float, str], FieldSlab] = { + (200100.0, "DEWPT"): FieldSlab( + name="DEWPT", + units="K".ljust(25), + desc="x", + level=200100.0, + data=slab, + ), + (200100.0, "TT"): FieldSlab( + name="TT", + units="K".ljust(25), + desc="y", + level=200100.0, + data=slab, + ), + (200100.0, "UU"): FieldSlab( + name="UU", + units="m s-1".ljust(25), + desc="z", + level=200100.0, + data=slab, + ), + } + drop_fields_blank_vtable_desc(bucket, [dewpt, tt, uu_sfc, uu_pl]) + assert (200100.0, "DEWPT") not in bucket + assert (200100.0, "TT") in bucket + assert (200100.0, "UU") in bucket diff --git a/ungrib/py/tests/test_vtable.py b/ungrib/py/tests/test_vtable.py new file mode 100644 index 00000000..d32a84e4 --- /dev/null +++ b/ungrib/py/tests/test_vtable.py @@ -0,0 +1,51 @@ +from pathlib import Path + +import pytest + +from ungrib_py.vtable import VtableEntry, parse_vtable, match_entry + +VTABLE_GFS = Path(__file__).resolve().parents[2] / "Variable_Tables" / "Vtable.GFS" + +@pytest.mark.skipif(not VTABLE_GFS.is_file(), reason="Vtable.GFS not in tree") +def test_parse_vtable_gfs(): + + # This function verifies parse_vtable reads the repository Vtable.GFS when present. + + rows, out_rows = parse_vtable(VTABLE_GFS) + assert len(rows) >= 10 + assert len(out_rows) >= 10 + first = rows[0] + assert isinstance(first, VtableEntry) + assert first.name.strip() == "TT" + +def test_match_gfs_temperature_100hPa(): + + # This function verifies Vtable match_entry for GFS-like temperature on isobaric surfaces. + + row = VtableEntry( + g1_param=11, + g1_level_type=100, + level1=-88.0, + level2=-99.0, + name="TT ", + units="K ", + desc="Temperature", + g2_discipline=0, + g2_category=0, + g2_parameter=0, + g2_level_type=100, + g2_pdt=0, + ) + m = match_entry([row], 0, 0, 0, 100, 0) + assert m is row + assert match_entry([row], 0, 0, 1, 100, 0) is None + + + + + + + + + + diff --git a/ungrib/py/ungrib.py b/ungrib/py/ungrib.py new file mode 100755 index 00000000..e34a157d --- /dev/null +++ b/ungrib/py/ungrib.py @@ -0,0 +1,23 @@ +#!/usr/bin/env python3 + +# Add sibling package ungrib_py to sys.path when this file is run or symlinked (Fortran-style workflow). + +import sys +from pathlib import Path + +_root = Path(__file__).resolve().parent +if str(_root) not in sys.path: + sys.path.insert(0, str(_root)) + +from ungrib_py.cli import main + +if __name__ == "__main__": + raise SystemExit(main()) + + + + + + + + diff --git a/ungrib/py/ungrib_py/__init__.py b/ungrib/py/ungrib_py/__init__.py new file mode 100644 index 00000000..41f32e10 --- /dev/null +++ b/ungrib/py/ungrib_py/__init__.py @@ -0,0 +1,3 @@ +# Parallel GRIB2 -> WPS intermediate (Python port of ungrib core). + +__version__ = "0.1.0" diff --git a/ungrib/py/ungrib_py/__main__.py b/ungrib/py/ungrib_py/__main__.py new file mode 100644 index 00000000..562a5c0f --- /dev/null +++ b/ungrib/py/ungrib_py/__main__.py @@ -0,0 +1,6 @@ +import sys + +from ungrib_py.cli import main + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/ungrib/py/ungrib_py/cli.py b/ungrib/py/ungrib_py/cli.py new file mode 100644 index 00000000..1f7ca533 --- /dev/null +++ b/ungrib/py/ungrib_py/cli.py @@ -0,0 +1,65 @@ +# Command-line interface. + +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +from ungrib_py.namelist_wps import parse_namelist_wps +from ungrib_py.parallel_ungrib import run_ungrib_parallel + + +def main(argv: list[str] | None = None) -> int: + + # This function parses CLI options, reads namelist.wps, and runs parallel ungrib. + + p = argparse.ArgumentParser( + description="Parallel Python ungrib: GRIB2 -> WPS intermediate (v5)." + ) + p.add_argument( + "--namelist", + default="namelist.wps", + help="Path to namelist.wps (default: ./namelist.wps)", + ) + p.add_argument( + "--vtable", + default="Vtable", + help="Path to Vtable (default: ./Vtable)", + ) + p.add_argument( + "--workers", + type=int, + default=None, + help="Process pool size (default: min(32, CPU count))", + ) + p.add_argument( + "--cwd", + default=".", + help="Working directory for GRIBFILE.* and outputs (default: .)", + ) + args = p.parse_args(argv) + cwd = Path(args.cwd).resolve() + namelist_path = Path(args.namelist) + if not namelist_path.is_absolute(): + namelist_path = cwd / namelist_path + vtable_path = Path(args.vtable) + if not vtable_path.is_absolute(): + vtable_path = cwd / vtable_path + os.chdir(cwd) + nml = parse_namelist_wps(namelist_path) + written = run_ungrib_parallel( + cwd, + nml, + vtable_path, + max_workers=args.workers, + ) + print(f"Wrote {len(written)} intermediate file(s).") + for w in written: + print(" ", w) + return 0 + + + + + diff --git a/ungrib/py/ungrib_py/grib2.py b/ungrib/py/ungrib_py/grib2.py new file mode 100644 index 00000000..fe08c14d --- /dev/null +++ b/ungrib/py/ungrib_py/grib2.py @@ -0,0 +1,1009 @@ +# GRIB2 extraction using ecCodes (Vtable-driven). + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path +import struct + +import numpy as np + +import eccodes + +from ungrib_py.intermediate import FieldSlab, MapInfo +from ungrib_py.vtable import BLANK, SPLAT, VtableEntry, match_entry + +_GDT32769_MAPGRID: tuple[int, ...] = ( + 1, + 1, + 4, + 1, + 4, + 1, + 4, + 4, + 4, + 4, + 4, + -4, + 4, + 1, + -4, + 4, + 4, + 4, + 1, + 4, + 4, +) + +_SOILT_SHALLOW_FIRST: tuple[str, ...] = ( + "SOILT001", + "SOILT002", + "SOILT006", + "SOILT018", + "SOILT054", + "SOILT162", + "SOILT486", + "SOILT999", +) + +_GRAVITY_M_S2: float = 9.81 + + +def _is_pressure_level_pa(lev: float) -> bool: + + # This function returns True if lev is a plausible isobaric pressure in Pa for WPS intermediate fields. + + lf = float(lev) + return 50.0 <= lf <= 120000.0 + + +def _delete_surface_geopotential_sources(bucket: dict[tuple[float, str], FieldSlab]) -> None: + + # This function removes surface SOILGEO/GEOPT slabs from the bucket when they should not reach metgrid. + + for k in list(bucket.keys()): + lev, nm = k + if nm.strip() not in ("SOILGEO", "GEOPT"): + continue + lf = float(lev) + if abs(lf - 200100.0) < 0.5 or abs(lf - 1.0) < 0.5: + del bucket[k] + + +def _hgt_field_meta_from_vtable(rows: list[VtableEntry] | None) -> tuple[str, str, str]: + + # This function returns (name, units, desc) for derived HGT from GEOPT, preferring the Vtable HGT row. + + if rows: + for r in rows: + if r.name.strip() == "HGT" and r.desc.strip(): + return (r.name[:9].ljust(9)[:9], r.units[:25], r.desc[:46]) + return ("HGT".ljust(9)[:9], "m".ljust(25)[:25], "Height".ljust(46)[:46]) + + +def ensure_hgt_from_geopt( + bucket: dict[tuple[float, str], FieldSlab], + *, + vtable_rows: list[VtableEntry] | None = None, +) -> None: + + # This function adds HGT slabs from isobaric GEOPT by dividing geopotential by gravity (Fortran rrpr). + + name9, units25, desc46 = _hgt_field_meta_from_vtable(vtable_rows) + for key in list(bucket.keys()): + lev, name = key + if name.strip() != "GEOPT": + continue + if not _is_pressure_level_pa(lev): + continue + hgt_key = (float(lev), "HGT") + if hgt_key in bucket: + del bucket[key] + continue + geopt = bucket[key] + arr = np.asarray(geopt.data, dtype=np.float64) + hgt = np.asfortranarray((arr / _GRAVITY_M_S2).astype(np.float32)) + bucket[hgt_key] = FieldSlab( + name=name9, + units=units25, + desc=desc46, + level=float(lev), + data=hgt, + ) + del bucket[key] + + +def drop_fields_blank_vtable_desc( + bucket: dict[tuple[float, str], FieldSlab], + all_rows: list[VtableEntry], +) -> None: + + # This function drops fields whose Vtable description is all blanks, matching Fortran output.F iflag=2. + + by_name: dict[str, list[VtableEntry]] = {} + for r in all_rows: + by_name.setdefault(r.name.strip(), []).append(r) + drop_names = { + nm + for nm, rs in by_name.items() + if rs and all(not x.desc.strip() for x in rs) + } + for key in list(bucket.keys()): + if key[1].strip() in drop_names: + del bucket[key] + + +def ensure_soilt000_for_metgrid(bucket: dict[tuple[float, str], FieldSlab]) -> None: + + # This function ensures SOILT000 exists by copying the shallowest SOILT* slab when METGRID.TBL expects it. + + levels = {k[0] for k in bucket} + for lev in levels: + if (lev, "SOILT000") in bucket: + continue + for name in _SOILT_SHALLOW_FIRST: + key = (lev, name) + if key not in bucket: + continue + src = bucket[key] + bucket[(lev, "SOILT000")] = FieldSlab( + name="SOILT000".ljust(9)[:9], + units=src.units, + desc=src.desc, + level=lev, + data=np.asfortranarray(src.data.copy()), + ) + break + + +def ensure_soilhgt_from_soilgeo(bucket: dict[tuple[float, str], FieldSlab]) -> None: + + # This function fills SOILHGT at 200100 from SOILGEO/GEOPT or surface HGT when SOILHGT is missing. + + for (lev, name) in bucket: + if name.strip() != "SOILHGT": + continue + if abs(float(lev) - 200100.0) < 0.5: + _delete_surface_geopotential_sources(bucket) + return + + def find_at_surface(names: tuple[str, ...]) -> FieldSlab | None: + + # This function returns the first bucket field matching names at surface level codes. + + for (lev, name), fld in bucket.items(): + if name.strip() not in names: + continue + lf = float(lev) + if abs(lf - 200100.0) < 0.5 or abs(lf - 1.0) < 0.5: + return fld + return None + + geo_like = find_at_surface(("SOILGEO", "GEOPT")) + if geo_like is not None: + arr = np.asarray(geo_like.data, dtype=np.float64) + hgt = np.asfortranarray((arr / _GRAVITY_M_S2).astype(np.float32)) + bucket[(200100.0, "SOILHGT")] = FieldSlab( + name="SOILHGT".ljust(9)[:9], + units="m".ljust(25)[:25], + desc="Terrain field of source analysis".ljust(46)[:46], + level=200100.0, + data=hgt, + ) + _delete_surface_geopotential_sources(bucket) + return + + hgt_sfc = find_at_surface(("HGT",)) + if hgt_sfc is None: + return + arr = np.asarray(hgt_sfc.data, dtype=np.float64) + if float(np.nanmax(np.abs(arr))) > 12000.0: + return + bucket[(200100.0, "SOILHGT")] = FieldSlab( + name="SOILHGT".ljust(9)[:9], + units="m".ljust(25)[:25], + desc="Terrain field of source analysis".ljust(46)[:46], + level=200100.0, + data=np.asfortranarray(arr.astype(np.float32)), + ) + + +def ecmwf_surface_geopotential_vtable_row(gid: int, rows: list[VtableEntry]) -> VtableEntry | None: + + # This function picks SOILGEO or GEOPT Vtable row for ECMWF surface geopotential when GRIB2 columns are blank. + + pid = _try_get(gid, "paramId") + pid_ok = False + if pid is not None: + try: + pid_ok = int(pid) == 129 + except (TypeError, ValueError): + pid_ok = False + sn = str(_try_get(gid, "shortName", "") or "").strip().lower() + if not pid_ok and sn != "z": + return None + if _iget(gid, "discipline") != 0: + return None + if type_of_first_fixed_surface_int(gid) != 1: + return None + soilgeo = [r for r in rows if r.name.strip() == "SOILGEO"] + if soilgeo: + return soilgeo[0] + geopt = [r for r in rows if r.name.strip() == "GEOPT"] + if geopt: + return geopt[0] + return None + + +def _iget(gid: int, key: str) -> int: + + # This function returns an integer ecCodes key value for a GRIB handle. + + return int(eccodes.codes_get(gid, key)) + +def _fget(gid: int, key: str) -> float: + + # This function returns a floating-point ecCodes key value for a GRIB handle. + + return float(eccodes.codes_get(gid, key)) + +def _sget(gid: int, key: str) -> str: + + # This function returns a string ecCodes key value for a GRIB handle. + + return str(eccodes.codes_get(gid, key)) + +def _try_get(gid: int, key: str, default: object | None = None) -> object | None: + + # This function returns a GRIB key value or default if the key is missing. + + try: + return eccodes.codes_get(gid, key) + except eccodes.KeyValueNotFoundError: + return default + +_TYPE_OF_FIRST_FIXED_SURFACE_STR: dict[str, int] = { + + # This table maps ecCodes string codetable labels to WMO GRIB2 Table 4.5 codes. + + "sfc": 1, + "surface": 1, + "msl": 101, + "pl": 100, + "pt": 100, + "hl": 103, + "hhl": 118, + "sol": 106, + "soil": 106, + "sot": 106, + "pv": 109, + "lcl": 212, + "lcd": 213, +} + +def type_of_first_fixed_surface_int(gid: int) -> int: + + # This function returns typeOfFirstFixedSurface as an integer (some GRIBs expose it as a string). + + try: + return int(eccodes.codes_get(gid, "typeOfFirstFixedSurface", int)) + except (TypeError, ValueError, eccodes.CodesInternalError, eccodes.GribInternalError): + pass + raw = _try_get(gid, "typeOfFirstFixedSurface") + if raw is None: + raise ValueError("typeOfFirstFixedSurface is missing from GRIB message") + if isinstance(raw, bool): + return int(raw) + if isinstance(raw, int): + return raw + if isinstance(raw, float): + return int(raw) + s = str(raw).strip().lower() + if s.isdigit(): + return int(s) + if s in _TYPE_OF_FIRST_FIXED_SURFACE_STR: + return _TYPE_OF_FIRST_FIXED_SURFACE_STR[s] + raise ValueError(f"unsupported typeOfFirstFixedSurface string {raw!r}") + +def earth_radius_km(gid: int) -> float: + + # This function returns spherical Earth radius in km from GRIB shape metadata. + + shape = _try_get(gid, "shapeOfTheEarth") + if shape is None: + return 6371.229 + shape = int(shape) + if shape == 0: + return 6367.47 + if shape == 1: + try: + sc = int(eccodes.codes_get(gid, "scaleFactorOfRadiusOfSphericalEarth")) + val = int(eccodes.codes_get(gid, "scaledValueOfRadiusOfSphericalEarth")) + return (val / (10**sc)) * 0.001 + except eccodes.KeyValueNotFoundError: + return 6371.229 + if shape == 6: + return 6371.229 + if shape == 8: + return 6371.2 + return 6371.229 + +def _wps_source32(s: str) -> str: + + # This function pads a source label to 32 characters like Fortran map%source. + + return s.ljust(32)[:32] + + +def _year_from_gid(gid: int) -> int: + + # This function returns the validity year from GRIB date keys (for NCEP RAP vs RUC). + + vdate = _try_get(gid, "validityDate", _try_get(gid, "dataDate")) + if vdate is None: + return 2000 + return int(vdate) // 10000 + + +def guess_source(gid: int) -> str: + + # This function matches rd_grib2.F map%source assignment from centre and generating process. + + gp = _try_get(gid, "generatingProcessIdentifier") + if gp is None: + return _wps_source32("Unknown GRIB2") + try: + gp_i = int(gp) + except (TypeError, ValueError): + return _wps_source32("Unknown GRIB2") + + c = _try_get(gid, "centre") + if isinstance(c, str): + cs = c.strip().lower() + if cs in ("ecmf", "ecmwf"): + return _wps_source32("ECMWF") + if cs == "edzw": + return _wps_source32("DWD") + try: + c_i = int(cs) + except ValueError: + return _wps_source32("unknown model and orig center") + elif isinstance(c, (int, float)): + c_i = int(c) + else: + return _wps_source32("unknown model and orig center") + + if c_i == 7: + yr = _year_from_gid(gid) + if gp_i == 81: + return _wps_source32("NCEP GFS Analysis") + if gp_i == 82: + return _wps_source32("NCEP GFS GDAS/FNL") + if gp_i == 83: + return _wps_source32("NCEP HRRR Model") + if gp_i == 84: + return _wps_source32("NCEP MESO NAM Model") + if gp_i == 89: + return _wps_source32("NCEP NMM ") + if gp_i == 96: + return _wps_source32("NCEP GFS Model") + if gp_i in (86, 100): + return _wps_source32("NCEP RUC Model") + if gp_i == 101: + return _wps_source32("NCEP RUC Model") + if gp_i == 105: + if yr > 2011: + return _wps_source32("NCEP RAP Model") + return _wps_source32("NCEP RUC Model") + if gp_i == 107: + return _wps_source32("NCEP GEFS") + if gp_i == 109: + return _wps_source32("NCEP RTMA") + if gp_i == 140: + return _wps_source32("NCEP NARR") + if gp_i == 44: + return _wps_source32("NCEP SST Analysis") + if gp_i == 70: + return _wps_source32("GFDL Hurricane Model") + if gp_i == 80: + return _wps_source32("NCEP GFS Ensemble") + if gp_i == 111: + return _wps_source32("NCEP NMMB Model") + if gp_i == 112: + return _wps_source32("NCEP WRF-NMM Model") + if gp_i == 116: + return _wps_source32("NCEP WRF-ARW Model") + if gp_i == 129: + return _wps_source32("NCEP GODAS") + if gp_i == 197: + return _wps_source32("NCEP CDAS CFSV2") + if gp_i == 25: + return _wps_source32("NCEP SNOW COVER ANALYSIS") + return _wps_source32("unknown model from NCEP") + if c_i == 57: + if gp_i == 87: + return _wps_source32("AFWA AGRMET") + return _wps_source32("AFWA") + if c_i == 58: + return _wps_source32("US Navy FNOC") + if c_i == 59: + if gp_i == 125: + return _wps_source32("NOAA GSD Rapid Refresh Model") + if gp_i == 83: + return _wps_source32("NOAA GSD HRRR Model") + if gp_i == 105: + return _wps_source32("NOAA GSD") + return _wps_source32("NOAA GSD") + if c_i == 60: + return _wps_source32("NCAR") + if c_i == 98: + return _wps_source32("ECMWF") + if c_i == 34: + return _wps_source32("JMA") + if c_i in (74, 75): + return _wps_source32("UKMO") + if c_i in (78, 79): + return _wps_source32("DWD") + return _wps_source32("unknown model and orig center") + +def validity_hdate(gid: int) -> str: + + # This function formats GRIB validity date and time as a 19-character hdate string. + + vdate = _try_get(gid, "validityDate", _try_get(gid, "date")) + vtime = _try_get(gid, "validityTime", _try_get(gid, "time")) + if vdate is None: + vdate = _try_get(gid, "dataDate") + if vtime is None: + vtime = _try_get(gid, "dataTime", 0) + vdate = int(vdate) + vtime = int(vtime) + y, m, d = vdate // 10000, (vdate % 10000) // 100, vdate % 100 + hh = vtime // 100 + mm = vtime % 100 + ss = 0 + if vtime > 2400: + hh = vtime // 10000 + mm = (vtime % 10000) // 100 + ss = vtime % 100 + return f"{y:04d}-{m:02d}-{d:02d}_{hh:02d}:{mm:02d}:{ss:02d}" + +def grid_relative_wind(gid: int) -> bool: + + # This function reports whether winds are grid-relative from GRIB metadata. + + v = _try_get(gid, "uvRelativeToGrid") + if v is not None: + return bool(int(v)) + flags = _try_get(gid, "resolutionAndComponentFlags") + if flags is None: + return True + flags = int(flags) + bit4 = (flags >> 3) & 1 + return bit4 != 0 + +def _grib2_section_bytes(msg: bytes, section_number: int) -> bytes | None: + + # This function returns the raw bytes of a GRIB2 section with the given section number. + + pos = 16 + while pos + 5 <= len(msg): + sec_len = struct.unpack_from(">I", msg, pos)[0] + if sec_len < 5 or pos + sec_len > len(msg): + return None + if msg[pos + 4] == section_number: + return msg[pos : pos + sec_len] + pos += sec_len + return None + +def _ncep32769_template_values(template_bytes: bytes) -> tuple[int, ...]: + + # This function unpacks GDT 3.32769 template integers using the NCEP g2lib mapgrid octet map. + + idx = 0 + out: list[int] = [] + for w in _GDT32769_MAPGRID: + n = abs(w) + chunk = template_bytes[idx : idx + n] + if len(chunk) != n: + raise ValueError("truncated GDT 3.32769 template bytes") + if w == 1: + out.append(chunk[0]) + else: + out.append(struct.unpack(">i", chunk)[0]) + idx += n + if idx != len(template_bytes): + raise ValueError("GDT 3.32769 template byte length mismatch") + return tuple(out) + +def _map_info_ncep_rotated_latlon(gid: int, src: str, r_earth: float, gw: bool, startloc: str) -> MapInfo: + + # This function builds MapInfo for NCEP rotated lat-lon (GDT 3.32769) like rd_grib2.F map%igrid=6. + + msg = eccodes.codes_get_message(gid) + sec3 = _grib2_section_bytes(msg, 3) + if sec3 is None or len(sec3) < 14: + raise ValueError("missing GRIB2 section 3 for ncep_32769") + gdt = struct.unpack_from(">H", sec3, 12)[0] + if gdt != 32769: + raise ValueError(f"expected GDT 32769 in section 3, got {gdt}") + tpl = sec3[14:] + vals = _ncep32769_template_values(tpl) + basic = vals[9] + if basic not in (0, 255): + raise ValueError( + "GDT 3.32769 with non-zero basic angle is unsupported in ungrib_py (extend _map_info_ncep_rotated_latlon)" + ) + di_raw = vals[16] + dj_raw = vals[17] + la2_raw = vals[14] + lo2_raw = vals[15] + ni = _iget(gid, "Ni") + nj = _iget(gid, "Nj") + lat1 = _fget(gid, "latitudeOfFirstGridPointInDegrees") + lon1 = _fget(gid, "longitudeOfFirstGridPointInDegrees") + di_deg = float(di_raw) / 1e9 + dj_deg = float(dj_raw) / 1e9 + center_lat = float(la2_raw) / 1e6 + center_lon = float(lo2_raw) / 1e6 + return MapInfo( + source=src, + igrid=6, + nx=ni, + ny=nj, + startloc=startloc, + lat1=lat1, + lon1=lon1, + dx=di_deg, + dy=dj_deg, + lov=0.0, + truelat1=0.0, + truelat2=0.0, + r_earth_km=r_earth, + grid_wind=gw, + centerlat=center_lat, + centerlon=center_lon, + ) + +def map_info_from_grib(gid: int) -> MapInfo: + + # This function builds WPS MapInfo from ecCodes grid description for supported projections. + + gt = _sget(gid, "gridType") + src = guess_source(gid) + r_earth = earth_radius_km(gid) + gw = grid_relative_wind(gid) + startloc = "SWCORNER" + + if gt == "regular_ll": + ni = _iget(gid, "Ni") + nj = _iget(gid, "Nj") + lat1 = _fget(gid, "latitudeOfFirstGridPointInDegrees") + lon1 = _fget(gid, "longitudeOfFirstGridPointInDegrees") + dlon = _fget(gid, "iDirectionIncrementInDegrees") + dlat = _fget(gid, "jDirectionIncrementInDegrees") + j_scan_pos = _try_get(gid, "jScansPositively") + if j_scan_pos is not None and int(j_scan_pos) == 0 and dlat > 0: + dlat = -dlat + return MapInfo( + source=src, + igrid=0, + nx=ni, + ny=nj, + startloc=startloc, + lat1=lat1, + lon1=lon1, + dx=dlon, + dy=dlat, + lov=0.0, + truelat1=0.0, + truelat2=0.0, + r_earth_km=r_earth, + grid_wind=gw, + centerlat=0.0, + centerlon=0.0, + ) + + if gt == "lambert": + ni = _iget(gid, "Ni") + nj = _iget(gid, "Nj") + lat1 = _fget(gid, "latitudeOfFirstGridPointInDegrees") + lon1 = _fget(gid, "longitudeOfFirstGridPointInDegrees") + lov = _fget(gid, "LoVInDegrees") + t1 = _fget(gid, "Latin1InDegrees") + t2 = _fget(gid, "Latin2InDegrees") + dx_m = _fget(gid, "DxInMetres") + dy_m = _fget(gid, "DyInMetres") + return MapInfo( + source=src, + igrid=3, + nx=ni, + ny=nj, + startloc=startloc, + lat1=lat1, + lon1=lon1, + dx=dx_m * 0.001, + dy=dy_m * 0.001, + lov=lov, + truelat1=t1, + truelat2=t2, + r_earth_km=r_earth, + grid_wind=gw, + centerlat=0.0, + centerlon=0.0, + ) + + if gt == "polar_stereographic": + ni = _iget(gid, "Ni") + nj = _iget(gid, "Nj") + lat1 = _fget(gid, "latitudeOfFirstGridPointInDegrees") + lon1 = _fget(gid, "longitudeOfFirstGridPointInDegrees") + lov = _fget(gid, "orientationOfTheGridInDegrees") + dx_m = _fget(gid, "DxInMetres") + dy_m = _fget(gid, "DyInMetres") + return MapInfo( + source=src, + igrid=5, + nx=ni, + ny=nj, + startloc=startloc, + lat1=lat1, + lon1=lon1, + dx=dx_m * 0.001, + dy=dy_m * 0.001, + lov=lov, + truelat1=60.0, + truelat2=91.0, + r_earth_km=r_earth, + grid_wind=gw, + centerlat=0.0, + centerlon=0.0, + ) + + if gt == "mercator": + ni = _iget(gid, "Ni") + nj = _iget(gid, "Nj") + lat1 = _fget(gid, "latitudeOfFirstGridPointInDegrees") + lon1 = _fget(gid, "longitudeOfFirstGridPointInDegrees") + dlon = _fget(gid, "iDirectionIncrementInDegrees") + dlat = _fget(gid, "jDirectionIncrementInDegrees") + latin = _try_get(gid, "LaDInDegrees", _try_get(gid, "LatinInDegrees", 0.0)) + latin = float(latin or 0.0) + return MapInfo( + source=src, + igrid=1, + nx=ni, + ny=nj, + startloc=startloc, + lat1=lat1, + lon1=lon1, + dx=dlon, + dy=dlat, + lov=0.0, + truelat1=latin, + truelat2=0.0, + r_earth_km=r_earth, + grid_wind=gw, + centerlat=0.0, + centerlon=0.0, + ) + + if gt == "ncep_32769": + return _map_info_ncep_rotated_latlon(gid, src, r_earth, gw, startloc) + + raise ValueError(f"unsupported gridType={gt!r} for Python ungrib (extend map_info_from_grib)") + +def scaled_surface_value(gid: int, which: str) -> float: + + # This function decodes scaled first or second fixed surface values from GRIB2. + + if which == "first": + sf = int(_try_get(gid, "scaleFactorOfFirstFixedSurface", 0) or 0) + sv = int(_try_get(gid, "scaledValueOfFirstFixedSurface", 0) or 0) + else: + sf = int(_try_get(gid, "scaleFactorOfSecondFixedSurface", 0) or 0) + sv = int(_try_get(gid, "scaledValueOfSecondFixedSurface", 0) or 0) + return float(sv) * (10.0 ** (-sf)) + +def pressure_level_pa(gid: int) -> float: + + # This function returns isobaric pressure in Pa for WPS intermediate (Fortran rd_grib2 uses Pa; pmin in namelist.wps is Pa). + # ECMWF ERA5 GRIB2 often has typeOfLevel=isobaricInhPa with ecCodes key "level" in hPa. + + tol = str(_try_get(gid, "typeOfLevel", "") or "") + if "InhPa" in tol: + lev = _try_get(gid, "level") + if lev is not None: + return float(lev) * 100.0 + if "InPa" in tol and "InhPa" not in tol: + lev = _try_get(gid, "level") + if lev is not None: + return float(lev) + return scaled_surface_value(gid, "first") + +def compute_wps_level(gid: int, pmin: float) -> float | None: + + # This function maps GRIB2 vertical type to WPS intermediate level codes, honoring pmin. + + t = type_of_first_fixed_surface_int(gid) + if t == 100: + pa = pressure_level_pa(gid) + if pa < pmin: + return None + return pa + if t in (105, 118, 150): + return scaled_surface_value(gid, "first") + if t == 104: + return float(_iget(gid, "scaledValueOfFirstFixedSurface")) + if t == 101: + return 201300.0 + if t == 103: + depth = scaled_surface_value(gid, "first") + if depth in (2.0, 10.0, 1000.0): + return 200100.0 + return None + if 206 <= t <= 234 or 242 <= t <= 254 or t in (200, 10): + return 200100.0 + if t in (106, 1, 151): + return 200100.0 + if t == 6: + return 200100.0 + if t == 7: + return 200100.0 + return None + +def soil_depths_cm(gid: int) -> tuple[float, float]: + + # This function returns soil layer depths in centimeters for GRIB2 type-106 levels. + + g1 = scaled_surface_value(gid, "first") + g2 = scaled_surface_value(gid, "second") + if g1 > 1e6 or g2 > 1e6: + sv1 = float(_iget(gid, "scaledValueOfFirstFixedSurface")) + sv2 = float(_iget(gid, "scaledValueOfSecondFixedSurface")) + return sv1 * 100.0, sv2 * 100.0 + return 100.0 * g1, 100.0 * g2 + +def g1_level_type_from_ecmwf_type_of_level(type_of_level: str) -> int | None: + + # This function maps ECMWF-style GRIB1 typeOfLevel strings to WPS Vtable g1_level_type codes. + + match type_of_level: + case "isobaricInhPa" | "isobaricInPa": + return 100 + case "surface": + return 1 + case "heightAboveGround": + return 103 + case "depthBelowLandLayer": + return 112 + case _: + return None + + +def _grib1_row_matches_level_value(row: VtableEntry, level_val: float) -> bool: + + # This function checks Vtable level1/level2 against a GRIB1 scalar level (pressure hPa or meters). + + if abs(row.level1 - SPLAT) < 1e-6: + return True + return abs(row.level1 - level_val) < 0.01 + + +def pick_vtable_row_grib1(gid: int, rows: list[VtableEntry]) -> VtableEntry | None: + + # This function selects the Vtable row for a GRIB edition-1 message using g1_param and g1_level_type. + + pid = _try_get(gid, "paramId") + param = int(pid) if pid is not None else _iget(gid, "indicatorOfParameter") + tol = str(_sget(gid, "typeOfLevel")) + if tol == "depthBelowLandLayer": + g1 = float(_try_get(gid, "topLevel", 0) or 0) + g2 = float(_try_get(gid, "bottomLevel", 0) or 0) + for r in rows: + if r.g1_param == BLANK or r.g1_param != param: + continue + if r.g1_level_type != 112: + continue + l1_ok = abs(r.level1 - SPLAT) < 1e-6 or abs(r.level1 - g1) < 0.01 + if not l1_ok: + continue + if g2 > 1e6: + l2_ok = r.level2 > SPLAT and abs(r.level1 - g1) < 0.01 + else: + l2_ok = r.level2 <= SPLAT or abs(r.level2 - g2) < 0.01 + if l2_ok: + return r + return None + g1_lt = g1_level_type_from_ecmwf_type_of_level(tol) + if g1_lt is None: + return None + level_val = float(_try_get(gid, "level", 0) or 0) + for r in rows: + if r.g1_param == BLANK or r.g1_param != param: + continue + if r.g1_level_type == BLANK: + continue + if r.g1_level_type != g1_lt: + continue + if not _grib1_row_matches_level_value(r, level_val): + continue + return r + return None + + +def compute_wps_level_grib1(gid: int, pmin: float) -> float | None: + + # This function maps GRIB1 vertical coordinates to WPS intermediate level codes, honoring pmin. + + tol = str(_sget(gid, "typeOfLevel")) + lev = float(_try_get(gid, "level", 0) or 0) + if tol == "isobaricInhPa": + lev_pa = lev * 100.0 + if lev_pa < pmin: + return None + return lev_pa + if tol == "isobaricInPa": + if lev < pmin: + return None + return lev + if tol == "surface": + return 200100.0 + if tol == "heightAboveGround": + if lev in (2.0, 10.0, 1000.0): + return 200100.0 + return None + if tol == "depthBelowLandLayer": + return 200100.0 + return None + + +def pick_vtable_row_ecmwf_soil_layer151( + gid: int, + rows: list[VtableEntry], + discipline: int, + cat: int, + param: int, + pdt: int, +) -> VtableEntry | None: + + # This function picks the ECMWF IFS soil Vtable row for GRIB2 type-of-first-fixed-surface 151 using layer index. + + layer_start = scaled_surface_value(gid, "first") + for strict_pdt in (True, False): + candidates: list[VtableEntry] = [] + for r in rows: + if r.g2_discipline != discipline or r.g2_category != cat or r.g2_parameter != param: + continue + if r.g2_level_type != 151: + continue + if strict_pdt and r.g2_pdt != pdt: + continue + candidates.append(r) + if not candidates: + continue + candidates.sort(key=lambda x: (x.level1, x.level2)) + idx = int(round(layer_start)) + if 0 <= idx < len(candidates): + return candidates[idx] + return None + + +def pick_vtable_row(gid: int, rows: list[VtableEntry]) -> VtableEntry | None: + + # This function selects the Vtable row for a GRIB message, including soil depth matching. + + discipline = _iget(gid, "discipline") + cat = _iget(gid, "parameterCategory") + param = _iget(gid, "parameterNumber") + first_surface = type_of_first_fixed_surface_int(gid) + pdt = _iget(gid, "productDefinitionTemplateNumber") + if first_surface == 106: + g1, g2 = soil_depths_cm(gid) + for strict_pdt in (True, False): + for r in rows: + if r.g2_discipline != discipline or r.g2_category != cat or r.g2_parameter != param: + continue + if r.g2_level_type != 106: + continue + if strict_pdt and r.g2_pdt != pdt: + continue + l1_ok = abs(r.level1 - SPLAT) < 1e-6 or abs(r.level1 - g1) < 0.01 + l2_ok = r.level2 <= SPLAT or abs(r.level2 - g2) < 0.01 + if l1_ok and l2_ok: + return r + return None + if first_surface == 151: + m151 = pick_vtable_row_ecmwf_soil_layer151(gid, rows, discipline, cat, param, pdt) + if m151 is not None: + return m151 + m = match_entry(rows, discipline, cat, param, first_surface, pdt) + if m is not None: + return m + return ecmwf_surface_geopotential_vtable_row(gid, rows) + + +def values_to_slab(values: np.ndarray, ni: int, nj: int) -> np.ndarray: + + # This function reshapes ecCodes 1-D values to a Fortran-order (ni, nj) slab. + + arr = values.reshape((nj, ni), order="C") + return np.asfortranarray(arr.T.astype(np.float32)) + + +def extract_file( + path: str | Path, + vtable_rows: list[VtableEntry], + *, + pmin: float, +) -> tuple[dict[str, dict[tuple[float, str], FieldSlab]], MapInfo | None]: + + # This function scans one GRIB1/GRIB2 file and collects Vtable-matched fields grouped by hdate. + + out: dict[str, dict[tuple[float, str], FieldSlab]] = {} + map_info: MapInfo | None = None + p = Path(path) + if not p.is_file(): + return out, None + with p.open("rb") as fh: + while True: + gid = eccodes.codes_grib_new_from_file(fh) + if gid is None: + break + try: + edition = _iget(gid, "editionNumber") + if edition == 2: + row = pick_vtable_row(gid, vtable_rows) + if row is None: + continue + level = compute_wps_level(gid, pmin) + elif edition == 1: + row = pick_vtable_row_grib1(gid, vtable_rows) + if row is None: + continue + level = compute_wps_level_grib1(gid, pmin) + else: + continue + if level is None: + continue + hdate = validity_hdate(gid) + mi = map_info_from_grib(gid) + if map_info is None: + map_info = mi + ni = _iget(gid, "Ni") + nj = _iget(gid, "Nj") + vals = np.asarray(eccodes.codes_get_values(gid), dtype=np.float64) + slab = values_to_slab(vals, ni, nj) + if (ni, nj) != (mi.nx, mi.ny): + continue + fld = FieldSlab( + name=row.name[:9].strip().ljust(9)[:9], + units=row.units[:25], + desc=row.desc[:46], + level=float(level), + data=slab, + ) + bucket = out.setdefault(hdate, {}) + bucket[(float(level), fld.name.strip())] = fld + finally: + eccodes.codes_release(gid) + for bucket in out.values(): + ensure_soilt000_for_metgrid(bucket) + ensure_soilhgt_from_soilgeo(bucket) + ensure_hgt_from_geopt(bucket, vtable_rows=vtable_rows) + return out, map_info + +def iter_grib_filenames() -> Iterator[str]: + + # This function yields GRIBFILE.AAA through GRIBFILE.ZZZ names in lexicographic order. + + letters = [chr(c) for c in range(ord("A"), ord("Z") + 1)] + for a in letters: + for b in letters: + for c in letters: + yield f"GRIBFILE.{a}{b}{c}" + + + diff --git a/ungrib/py/ungrib_py/intermediate.py b/ungrib/py/ungrib_py/intermediate.py new file mode 100644 index 00000000..9f987d05 --- /dev/null +++ b/ungrib/py/ungrib_py/intermediate.py @@ -0,0 +1,283 @@ +# WPS intermediate format version 5 writer (Fortran unformatted sequential). +# +# On-disk layout matches WPS `ungrib/src/output.F` (WPS format) and `metgrid/src/read_met_module.F` +# for `fg_data % version == 5`: metgrid reads each field as five sequential unformatted records. +# +# Endianness: payloads use big-endian float32 / int32 (Fortran `-fconvert=big-endian`). Record +# delimiters are big-endian uint32 byte counts (`-frecord-marker=4`), consistent with typical +# WPS `configure` / `arch/configure.defaults` gfortran builds. +# +# Per field, records are: +# 1) int32 version code 5 (4-byte payload). +# 2) Mixed header, 156 bytes: character*24 hdate, real xfcst, character*32 map_source, +# character*9 field, character*25 units, character*46 desc, real xlvl, int32 nx, ny, iproj. +# Here `iproj` is the ungrib `map%igrid` code (0 lat-lon, 1 Mercator, 3 Lambert, 4 Gaussian, +# 5 polar stereographic, 6 Cassini). +# 3) Map extension: character*8 startloc then a projection-dependent list of real32 values +# (see `_map_extension_bytes_be` and `_MAP_EXT_PAYLOAD_LEN`). +# 4) Logical `is_wind_grid_rel`: written as a 4-byte big-endian integer 0 or 1 (gfortran +# commonly uses 4-byte LOGICAL in unformatted sequential I/O; metgrid accepts this pattern). +# 5) real32 array slab(nx, ny) in Fortran array element order, big-endian IEEE float32. + +from __future__ import annotations + +import struct +from typing import BinaryIO, NamedTuple, Sequence + +import numpy as np + +from ungrib_py.vtable import VtableEntry + +WPS_V5_VERSION: int = 5 + +WPS_V5_MIXED_HEADER_BYTES: int = 156 + +_MAP_EXT_PAYLOAD_LEN: dict[int, int] = { + 0: 28, + 1: 32, + 3: 40, + 4: 28, + 5: 36, + 6: 36, +} + +def _pad_str(s: str, n: int) -> bytes: + + # This function encodes a string to ASCII bytes padded or truncated to length n. + + b = s.encode("ascii", errors="replace")[:n] + return b + b" " * (n - len(b)) + +def _fortran_unformatted_record(fh: BinaryIO, payload: bytes) -> None: + + # This function writes one gfortran sequential unformatted record using big-endian 32-bit length tags. + + n = len(payload) + len_be = struct.pack(">I", n) + fh.write(len_be) + fh.write(payload) + fh.write(len_be) + +def _be_i32(v: int) -> bytes: + + # This function encodes a 32-bit signed integer in big-endian form for WPS I/O (-fconvert=big-endian). + + return struct.pack(">i", v) + +def _be_f32(v: float) -> bytes: + + # This function encodes a 32-bit IEEE float in big-endian form for WPS I/O (-fconvert=big-endian). + + return struct.pack(">f", v) + +class MapInfo(NamedTuple): + + # Projection metadata for WPS intermediate version 5. + + source: str + igrid: int + nx: int + ny: int + startloc: str + lat1: float + lon1: float + dx: float + dy: float + lov: float + truelat1: float + truelat2: float + r_earth_km: float + grid_wind: bool + centerlat: float + centerlon: float + +class FieldSlab(NamedTuple): + name: str + units: str + desc: str + level: float + data: np.ndarray + +def write_wps_intermediate_v5( + path: str, + hdate19: str, + map_info: MapInfo, + fields: Sequence[FieldSlab], + *, + xfcst: float = 0.0, +) -> None: + + # This function writes WPS intermediate format version 5 for one timestep to path. + + h24 = _pad_str(hdate19[:19], 24).decode("ascii") + src32 = _pad_str(map_info.source[:32], 32).decode("ascii") + + with open(path, "wb") as fh: + for fld in fields: + name9 = _pad_str(fld.name.strip()[:9], 9).decode("ascii") + units25 = _pad_str(fld.units.strip()[:25], 25).decode("ascii") + desc46 = _pad_str(fld.desc.strip()[:46], 46).decode("ascii") + slab = np.asfortranarray( + np.asarray(fld.data, dtype=np.float32).reshape(map_info.nx, map_info.ny) + ) + + ver_b = _be_i32(WPS_V5_VERSION) + if len(ver_b) != 4: + raise AssertionError("WPS v5 version record must be 4 bytes") + + hdr_b = _mixed_header_bytes_be( + h24, + float(xfcst), + src32, + name9, + units25, + desc46, + float(fld.level), + int(map_info.nx), + int(map_info.ny), + int(map_info.igrid), + ) + if len(hdr_b) != WPS_V5_MIXED_HEADER_BYTES: + raise AssertionError( + f"WPS v5 mixed header must be {WPS_V5_MIXED_HEADER_BYTES} bytes, got {len(hdr_b)}" + ) + + ext_b = _map_extension_bytes_be(map_info) + exp_len = _MAP_EXT_PAYLOAD_LEN.get(map_info.igrid) + if exp_len is None: + raise ValueError(f"unsupported igrid={map_info.igrid} for WPS v5 writer") + if len(ext_b) != exp_len: + raise AssertionError( + f"map extension for igrid={map_info.igrid} must be {exp_len} bytes, got {len(ext_b)}" + ) + + wind_b = _be_i32(1 if map_info.grid_wind else 0) + if len(wind_b) != 4: + raise AssertionError("grid_wind record must be 4 bytes") + + slab_b = _slab_bytes_be(slab) + if len(slab_b) != map_info.nx * map_info.ny * 4: + raise AssertionError("slab payload size must be nx*ny*4 bytes") + + _fortran_unformatted_record(fh, ver_b) + _fortran_unformatted_record(fh, hdr_b) + _fortran_unformatted_record(fh, ext_b) + _fortran_unformatted_record(fh, wind_b) + _fortran_unformatted_record(fh, slab_b) + +def _mixed_header_bytes_be( + h24: str, + xfcst: float, + src32: str, + name9: str, + units25: str, + desc46: str, + level: float, + nx: int, + ny: int, + iproj: int, +) -> bytes: + + # This function builds the mixed-type header record with big-endian reals and integers. + + buf = bytearray() + buf.extend(_pad_str(h24, 24)) + buf.extend(_be_f32(xfcst)) + buf.extend(_pad_str(src32, 32)) + buf.extend(_pad_str(name9, 9)) + buf.extend(_pad_str(units25, 25)) + buf.extend(_pad_str(desc46, 46)) + buf.extend(_be_f32(level)) + buf.extend(_be_i32(nx)) + buf.extend(_be_i32(ny)) + buf.extend(_be_i32(iproj)) + return bytes(buf) + +def _map_extension_bytes_be(m: MapInfo) -> bytes: + + # This function builds the map projection extension record with big-endian floats. + + start_b = _pad_str(m.startloc[:8], 8) + if m.igrid == 0: + buf = bytearray() + buf.extend(start_b) + for v in (m.lat1, m.lon1, m.dy, m.dx, m.r_earth_km): + buf.extend(_be_f32(v)) + return bytes(buf) + if m.igrid == 1: + buf = bytearray() + buf.extend(start_b) + for v in (m.lat1, m.lon1, m.dx, m.dy, m.truelat1, m.r_earth_km): + buf.extend(_be_f32(v)) + return bytes(buf) + if m.igrid == 3: + buf = bytearray() + buf.extend(start_b) + for v in (m.lat1, m.lon1, m.dx, m.dy, m.lov, m.truelat1, m.truelat2, m.r_earth_km): + buf.extend(_be_f32(v)) + return bytes(buf) + if m.igrid == 5: + buf = bytearray() + buf.extend(start_b) + for v in (m.lat1, m.lon1, m.dx, m.dy, m.lov, m.truelat1, m.r_earth_km): + buf.extend(_be_f32(v)) + return bytes(buf) + if m.igrid == 4: + buf = bytearray() + buf.extend(start_b) + for v in (m.lat1, m.lon1, m.dx, m.dy, m.r_earth_km): + buf.extend(_be_f32(v)) + return bytes(buf) + if m.igrid == 6: + buf = bytearray() + buf.extend(start_b) + # Fortran output.F writes map%dy before map%dx; metgrid read_met reads dx then dy (same file order). + for v in (m.lat1, m.lon1, m.dy, m.dx, m.centerlat, m.centerlon, m.r_earth_km): + buf.extend(_be_f32(v)) + return bytes(buf) + raise ValueError(f"unsupported igrid={m.igrid} for WPS v5 writer") + +def _slab_bytes_be(slab: np.ndarray) -> bytes: + + # This function serializes a Fortran-order (nx, ny) float32 slab as big-endian floats. + + be = np.asfortranarray(slab).astype(np.dtype(">f4"), order="K", copy=True) + return be.tobytes("F") + +def intermediate_filename_datelen(interval_seconds: int) -> int: + + # This function returns output filename date substring length from interval_seconds. + + if interval_seconds % 3600 == 0: + return 13 + if interval_seconds % 60 == 0: + return 16 + return 19 + +def hdate_slice(hdate: str, datelen: int) -> str: + + # This function returns the first datelen characters of an hdate string. + + return hdate[:datelen] + +def sort_fields_fortran_order( + items: list[tuple[float, str, FieldSlab]], + vtable_rows: list[VtableEntry], +) -> list[FieldSlab]: + + # This function orders fields like WPS ungrib output.F: descending level (get_plvls), then Vtable row order. + + name_first_index: dict[str, int] = {} + for i, row in enumerate(vtable_rows): + nm = row.name.strip() + if nm not in name_first_index: + name_first_index[nm] = i + levels = sorted({t[0] for t in items}, reverse=True) + out: list[FieldSlab] = [] + for lev in levels: + at_level = [(t[1], t[2]) for t in items if t[0] == lev] + at_level.sort( + key=lambda pair: (name_first_index.get(pair[0].strip(), 10**9), pair[0].strip()), + ) + out.extend(slab for _, slab in at_level) + return out diff --git a/ungrib/py/ungrib_py/namelist_wps.py b/ungrib/py/ungrib_py/namelist_wps.py new file mode 100644 index 00000000..efc7fabb --- /dev/null +++ b/ungrib/py/ungrib_py/namelist_wps.py @@ -0,0 +1,130 @@ +# Minimal namelist.wps reader for ungrib. + +from __future__ import annotations + +import re +from pathlib import Path +from typing import NamedTuple + +class UngribNamelist(NamedTuple): + # Parsed &share / &ungrib keys; debug_level and ordered_by_date are not used by ungrib_py yet. + hstart: str + hend: str + interval_seconds: int + prefix: str + out_format: str + ordered_by_date: bool + debug_level: int + pmin: float + +def _first_int(text: str, name: str, default: int) -> int: + + # This function reads the first integer assignment for name from a namelist section string. + + m = re.search(rf"{name}\s*=\s*(-?\d+)", text, re.I) + return int(m.group(1)) if m else default + +def _first_float(text: str, name: str, default: float) -> float: + + # This function reads the first float assignment for name from a namelist section string. + + m = re.search(rf"{name}\s*=\s*(-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?)", text, re.I) + return float(m.group(1)) if m else default + +def _first_bool(text: str, name: str, default: bool) -> bool: + + # This function reads the first Fortran logical assignment for name from a section string. + + m = re.search(rf"{name}\s*=\s*\.(true|false)\.", text, re.I) + if not m: + return default + return m.group(1).lower() == "true" + +def _quoted_str(text: str, name: str, default: str) -> str: + + # This function reads a single- or double-quoted string assignment for name. + + m = re.search(rf"{name}\s*=\s*'([^']*)'", text, re.I) + if m: + return m.group(1) + m2 = re.search(rf'{name}\s*=\s*"([^"]*)"', text, re.I) + return m2.group(1) if m2 else default + +def _extract_section(full: str, section: str) -> str: + + # This function returns the body text of a Fortran namelist section without delimiters. + + m = re.search(rf"&{section}\s(.*?)^\s*/\s*$", full, re.I | re.S | re.M) + return m.group(1) if m else "" + +def _build_hdate_from_components( + year: int, month: int, day: int, hour: int, minute: int, second: int +) -> str: + + # This function formats WPS-style hdate components into a 19-character timestamp string. + + return f"{year:04d}-{month:02d}-{day:02d}_{hour:02d}:{minute:02d}:{second:02d}" + +def parse_namelist_wps(path: str | Path) -> UngribNamelist: + + # This function parses namelist.wps share and ungrib sections needed by ungrib. + + full = Path(path).read_text(encoding="utf-8", errors="replace") + share = _extract_section(full, "share") + ungrib = _extract_section(full, "ungrib") + + start_date = _quoted_str(share, "start_date", "") + end_date = _quoted_str(share, "end_date", "") + + if not start_date.strip() or start_date.strip().startswith("0000"): + sy = _first_int(share, "start_year", 0) + sm = _first_int(share, "start_month", 0) + sd = _first_int(share, "start_day", 0) + sh = _first_int(share, "start_hour", 0) + smin = _first_int(share, "start_minute", 0) + ss = _first_int(share, "start_second", 0) + hstart = _build_hdate_from_components(sy, sm, sd, sh, smin, ss) + else: + hstart = start_date.strip()[:19].ljust(19, "0")[:19] + + if not end_date.strip() or end_date.strip().startswith("0000"): + ey = _first_int(share, "end_year", 0) + em = _first_int(share, "end_month", 0) + ed = _first_int(share, "end_day", 0) + eh = _first_int(share, "end_hour", 0) + emin = _first_int(share, "end_minute", 0) + es = _first_int(share, "end_second", 0) + hend = _build_hdate_from_components(ey, em, ed, eh, emin, es) + else: + hend = end_date.strip()[:19].ljust(19, "0")[:19] + + interval_seconds = _first_int(share, "interval_seconds", 0) + debug_level = _first_int(share, "debug_level", 0) + prefix = _quoted_str(ungrib, "prefix", "FILE") + out_format = _quoted_str(ungrib, "out_format", "WPS").strip().upper()[:3] + ordered_by_date = _first_bool(ungrib, "ordered_by_date", True) + pmin = _first_float(ungrib, "pmin", 100.0) + + if interval_seconds <= 0: + raise ValueError("namelist.wps: interval_seconds must be > 0") + + return UngribNamelist( + hstart=hstart, + hend=hend, + interval_seconds=interval_seconds, + prefix=prefix, + out_format=out_format, + ordered_by_date=ordered_by_date, + debug_level=debug_level, + pmin=pmin, + ) + + + + + + + + + + diff --git a/ungrib/py/ungrib_py/parallel_ungrib.py b/ungrib/py/ungrib_py/parallel_ungrib.py new file mode 100644 index 00000000..f908823b --- /dev/null +++ b/ungrib/py/ungrib_py/parallel_ungrib.py @@ -0,0 +1,151 @@ +# Parallel GRIB-file processing and WPS intermediate output. + +from __future__ import annotations + +import os +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path + +from ungrib_py.grib2 import ( + drop_fields_blank_vtable_desc, + ensure_hgt_from_geopt, + ensure_soilhgt_from_soilgeo, + ensure_soilt000_for_metgrid, + extract_file, + iter_grib_filenames, +) +from ungrib_py.intermediate import ( + FieldSlab, + MapInfo, + hdate_slice, + intermediate_filename_datelen, + sort_fields_fortran_order, + write_wps_intermediate_v5, +) +from ungrib_py.namelist_wps import UngribNamelist +from ungrib_py.vtable import VtableEntry, parse_vtable + +def discover_grib_files(cwd: Path) -> list[str]: + + # This function lists GRIBFILE.AAA-style names that exist under cwd. + + return [name for name in iter_grib_filenames() if (cwd / name).is_file()] + +def _worker_extract( + args: tuple[str, str, str, float], +) -> tuple[str, dict[str, dict[tuple[float, str], FieldSlab]], MapInfo | None]: + + # This function runs extract_file for one GRIB path inside a worker process. + + cwd_s, fname, vtable_path, pmin = args + cwd = Path(cwd_s) + rows, _ = parse_vtable(vtable_path) + data, mi = extract_file(cwd / fname, rows, pmin=pmin) + return fname, data, mi + +def count_distinct_pressure_levels_pa( + merged: dict[str, dict[tuple[float, str], FieldSlab]], +) -> int: + + # This function counts distinct isobaric levels in Pa (WPS convention), excluding surface/special codes. + + levels: set[float] = set() + for bucket in merged.values(): + for lev, _name in bucket: + lf = float(lev) + if 50.0 <= lf <= 120000.0: + levels.add(round(lf, 1)) + return len(levels) + +def merge_by_time( + parts: list[ + tuple[ + dict[str, dict[tuple[float, str], FieldSlab]], + MapInfo | None, + ] + ], + hstart: str, + hend: str, + all_vtable_rows: list[VtableEntry], +) -> tuple[dict[str, dict[tuple[float, str], FieldSlab]], MapInfo]: + + # This function merges worker dicts by hdate, filters the time window, and checks the grid. + + merged: dict[str, dict[tuple[float, str], FieldSlab]] = {} + map_ref: MapInfo | None = None + for blob, mi in parts: + if mi is not None: + if map_ref is None: + map_ref = mi + elif (mi.nx, mi.ny, mi.igrid) != (map_ref.nx, map_ref.ny, map_ref.igrid): + raise ValueError( + "Inconsistent grid dimensions or projection between GRIB inputs." + ) + for hdate, fields in blob.items(): + if hdate < hstart or hdate > hend: + continue + bucket = merged.setdefault(hdate, {}) + bucket.update(fields) + if map_ref is None: + raise ValueError("No GRIB2 fields matched the Vtable in the given date range.") + for bucket in merged.values(): + ensure_soilt000_for_metgrid(bucket) + ensure_soilhgt_from_soilgeo(bucket) + ensure_hgt_from_geopt(bucket, vtable_rows=all_vtable_rows) + drop_fields_blank_vtable_desc(bucket, all_vtable_rows) + return merged, map_ref + +def run_ungrib_parallel( + cwd: Path, + nml: UngribNamelist, + vtable_path: Path, + *, + grib_files: list[str] | None = None, + max_workers: int | None = None, +) -> list[str]: + + # This function processes GRIB files in parallel and writes WPS intermediate v5 outputs. + + if nml.out_format[:2] != "WP": + raise NotImplementedError( + f"out_format={nml.out_format!r}: only WPS intermediate (version 5) is implemented." + ) + files = grib_files if grib_files is not None else discover_grib_files(cwd) + if not files: + raise FileNotFoundError( + "No GRIBFILE.AAA-style inputs found in working directory." + ) + vtable_abs = str(vtable_path.resolve()) + vtable_rows, _ = parse_vtable(vtable_path) + workers = max_workers or min(32, max(1, (os.cpu_count() or 4))) + tasks = [(str(cwd.resolve()), fn, vtable_abs, float(nml.pmin)) for fn in files] + results: list[ + tuple[dict[str, dict[tuple[float, str], FieldSlab]], MapInfo | None] + ] = [] + with ProcessPoolExecutor(max_workers=workers) as ex: + futs = [ex.submit(_worker_extract, t) for t in tasks] + for fut in as_completed(futs): + _, data, mi = fut.result() + results.append((data, mi)) + merged, map_info = merge_by_time(results, nml.hstart, nml.hend, vtable_rows) + n_pressure = count_distinct_pressure_levels_pa(merged) + print( + "ungrib_py: distinct isobaric levels in intermediate output: " + f"{n_pressure} (set num_metgrid_levels = {n_pressure} in namelist.input &domains " + "to match met_em / BOTTOM-TOP_GRID_DIMENSION)" + ) + datelen = intermediate_filename_datelen(nml.interval_seconds) + written: list[str] = [] + for hdate in sorted(merged.keys()): + items = [(k[0], k[1], v) for k, v in merged[hdate].items()] + fields = sort_fields_fortran_order(items, vtable_rows) + out_name = f"{nml.prefix}:{hdate_slice(hdate, datelen)}" + out_path = cwd / out_name + write_wps_intermediate_v5(str(out_path), hdate, map_info, fields) + written.append(str(out_path)) + return written + + + + + diff --git a/ungrib/py/ungrib_py/vtable.py b/ungrib/py/ungrib_py/vtable.py new file mode 100644 index 00000000..a9d0de16 --- /dev/null +++ b/ungrib/py/ungrib_py/vtable.py @@ -0,0 +1,152 @@ +# Parse WPS Vtable files (GRIB2-aware lines). + +from __future__ import annotations + +from pathlib import Path +from typing import NamedTuple + +BLANK = -99 +SPLAT = -88 + +class VtableEntry(NamedTuple): + + # One non-comment data row after the header dashes. + + g1_param: int + g1_level_type: int + level1: float + level2: float + name: str + units: str + desc: str + g2_discipline: int + g2_category: int + g2_parameter: int + g2_level_type: int + g2_pdt: int + +def _split_pipe_line(line: str) -> list[str]: + + # This function splits a Vtable text line on pipes and strips each cell. + + parts = line.split("|") + return [p.strip() for p in parts] + +def _parse_int(tok: str, *, allow_blank: bool = False) -> int: + + # This function parses an integer Vtable token, optionally treating blanks as BLANK. + + t = tok.strip() + if not t or t == " ": + if allow_blank: + return BLANK + raise ValueError(f"expected integer, got empty in {tok!r}") + if "*" in t: + raise ValueError(f"unexpected wildcard in integer field: {tok!r}") + return int(t) + +def _parse_float_level(tok: str) -> float: + + # This function parses a Vtable level field as a float, blank, or SPLAT wildcard. + + t = tok.strip() + if not t: + return float(BLANK) + if "*" in t: + return float(SPLAT) + return float(t) + +def parse_vtable(path: str | Path) -> tuple[list[VtableEntry], list[VtableEntry]]: + + # This function parses a WPS Vtable file into all rows and output-eligible rows. + + text = Path(path).read_text(encoding="utf-8", errors="replace").splitlines() + i = 0 + while i < len(text) and not text[i].lstrip().startswith("-"): + i += 1 + if i >= len(text): + raise ValueError("Vtable: no header line starting with '-'") + i += 1 + rows: list[VtableEntry] = [] + while i < len(text): + raw = text[i].rstrip("\n") + i += 1 + if not raw.strip() or raw.lstrip().startswith("#"): + continue + if raw.lstrip().startswith("-"): + while i < len(text): + nxt = text[i].rstrip("\n") + i += 1 + if not nxt.strip() or nxt.lstrip().startswith("#"): + continue + if nxt.lstrip().startswith("-"): + break + continue + parts = _split_pipe_line(raw) + ncols = len(parts) + if ncols < 11: + raise ValueError(f"Vtable line needs >= 11 |-columns, got {ncols}: {raw!r}") + g1_param = _parse_int(parts[0], allow_blank=True) + g1_level_type = _parse_int(parts[1], allow_blank=True) + level1 = _parse_float_level(parts[2]) + level2 = _parse_float_level(parts[3]) + name = parts[4].strip() + if not name: + raise ValueError(f"missing field name: {raw!r}") + units = parts[5].strip() if len(parts) > 5 else "" + desc = parts[6].strip() if len(parts) > 6 else "" + g2_discipline = _parse_int(parts[7], allow_blank=True) + g2_category = _parse_int(parts[8], allow_blank=True) + g2_parameter = _parse_int(parts[9], allow_blank=True) + g2_level_type = _parse_int(parts[10], allow_blank=True) + g2_pdt = 0 + if ncols >= 12 and parts[11].strip(): + g2_pdt = _parse_int(parts[11], allow_blank=True) + if g2_pdt == BLANK: + g2_pdt = 0 + row = VtableEntry( + g1_param=g1_param, + g1_level_type=g1_level_type, + level1=level1, + level2=level2, + name=name[:9].ljust(9)[:9], + units=units[:25].ljust(25)[:25], + desc=desc[:46].ljust(46)[:46], + g2_discipline=g2_discipline, + g2_category=g2_category, + g2_parameter=g2_parameter, + g2_level_type=g2_level_type, + g2_pdt=g2_pdt, + ) + rows.append(row) + output_rows = [r for r in rows if r.desc.strip()] + return rows, output_rows + +def match_entry( + rows: list[VtableEntry], + discipline: int, + cat: int, + param: int, + first_surface: int, + pdt: int, +) -> VtableEntry | None: + + # This function returns the first Vtable row matching GRIB2 keys or None if none match. + + for r in rows: + if r.g2_discipline == BLANK or r.g2_category == BLANK or r.g2_parameter == BLANK: + continue + if ( + discipline == r.g2_discipline + and cat == r.g2_category + and param == r.g2_parameter + and first_surface == r.g2_level_type + and pdt == r.g2_pdt + ): + return r + return None + + + + +