From b8e35546bc3d60a64373dece02767e6696e55070 Mon Sep 17 00:00:00 2001 From: MHBalsmeier Date: Tue, 7 Apr 2026 23:09:40 +0200 Subject: [PATCH 1/6] link_grib.py included --- link_grib.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100755 link_grib.py 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("") + + + + + From 0fea6a255b2948cbb71d162a0c13e51357d0bf60 Mon Sep 17 00:00:00 2001 From: MHBalsmeier Date: Wed, 8 Apr 2026 00:27:06 +0200 Subject: [PATCH 2/6] install and WPS_GEOG in .gitignore --- .gitignore | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 330b89ca..ace8f070 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,6 @@ grib2 *.nc GRIBFILE.* *.o -FILE:* \ No newline at end of file +FILE:* +WPS_GEOG +install From 1dbdf4ad2d862688f1a77011ac28b95ff1cc58d4 Mon Sep 17 00:00:00 2001 From: MHBalsmeier Date: Sat, 11 Apr 2026 22:22:30 +0200 Subject: [PATCH 3/6] first running version of python version of ungrib WPS_GEOG, install and QNWFA_QNIFA_QNBCA_SIGMA_MONTHLY.dat added to main .gitignore --- .gitignore | 1 + ungrib/py/.gitignore | 1 + ungrib/py/MIGRATION.md | 75 ++ ungrib/py/README.md | 103 +++ ungrib/py/pyproject.toml | 21 + ungrib/py/requirements.txt | 3 + ungrib/py/tests/conftest.py | 16 + ungrib/py/tests/test_ecmwf_soil151.py | 48 ++ ungrib/py/tests/test_grib2_alias.py | 96 +++ ungrib/py/tests/test_intermediate.py | 70 ++ ungrib/py/tests/test_namelist.py | 41 ++ ungrib/py/tests/test_pressure_levels.py | 32 + ungrib/py/tests/test_soilhgt_rrpr.py | 269 +++++++ ungrib/py/tests/test_vtable.py | 51 ++ ungrib/py/ungrib.py | 23 + ungrib/py/ungrib_py/__init__.py | 13 + ungrib/py/ungrib_py/__main__.py | 16 + ungrib/py/ungrib_py/cli.py | 69 ++ ungrib/py/ungrib_py/grib2.py | 900 ++++++++++++++++++++++++ ungrib/py/ungrib_py/intermediate.py | 302 ++++++++ ungrib/py/ungrib_py/namelist_wps.py | 129 ++++ ungrib/py/ungrib_py/parallel_ungrib.py | 150 ++++ ungrib/py/ungrib_py/vtable.py | 157 +++++ 23 files changed, 2586 insertions(+) create mode 100644 ungrib/py/.gitignore create mode 100644 ungrib/py/MIGRATION.md create mode 100644 ungrib/py/README.md create mode 100644 ungrib/py/pyproject.toml create mode 100644 ungrib/py/requirements.txt create mode 100644 ungrib/py/tests/conftest.py create mode 100644 ungrib/py/tests/test_ecmwf_soil151.py create mode 100644 ungrib/py/tests/test_grib2_alias.py create mode 100644 ungrib/py/tests/test_intermediate.py create mode 100644 ungrib/py/tests/test_namelist.py create mode 100644 ungrib/py/tests/test_pressure_levels.py create mode 100644 ungrib/py/tests/test_soilhgt_rrpr.py create mode 100644 ungrib/py/tests/test_vtable.py create mode 100755 ungrib/py/ungrib.py create mode 100644 ungrib/py/ungrib_py/__init__.py create mode 100644 ungrib/py/ungrib_py/__main__.py create mode 100644 ungrib/py/ungrib_py/cli.py create mode 100644 ungrib/py/ungrib_py/grib2.py create mode 100644 ungrib/py/ungrib_py/intermediate.py create mode 100644 ungrib/py/ungrib_py/namelist_wps.py create mode 100644 ungrib/py/ungrib_py/parallel_ungrib.py create mode 100644 ungrib/py/ungrib_py/vtable.py diff --git a/.gitignore b/.gitignore index ace8f070..451bb15b 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ GRIBFILE.* FILE:* WPS_GEOG install +QNWFA_QNIFA_QNBCA_SIGMA_MONTHLY.dat 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..c094fc65 --- /dev/null +++ b/ungrib/py/MIGRATION.md @@ -0,0 +1,75 @@ +# Migration plan: Fortran ungrib → Python (`ungrib_py`) + +This document lists tasks to reach feature parity with `ungrib.exe`, notes what is already done in `ungrib_py`, and explains how to test the Python code. + +## Task outline + +### Phase A — Core decode and output (done in this port) + +1. **Vtable parsing** — Port `parse_table.F` logic for GRIB2 columns (discipline, category, parameter, first fixed surface / level type, product definition template). *Status: implemented for standard 11–12 column Vtables.* +2. **GRIB2 read path** — Replace `rd_grib2.F` + NCEP g2lib with **ecCodes** message iteration. *Status: implemented for common level types (pressure, mean sea level, height above ground, surface, soil 106, etc.) aligned with `rd_grib2.F`.* +3. **Grid metadata** — Port `gridinfo` population from GRIB2 Section 3 (lat/lon, Lambert, polar stereographic, Mercator). *Status: lat/lon (`regular_ll`) and Lambert/polar/Mercator via ecCodes `gridType`.* +4. **WPS intermediate writer** — Port `output.F` records for `out_format='WPS'`, version **5** (matches `metgrid/src/read_met_module.F`). *Status: implemented.* + +### WPS v5 file structure (metgrid) + +Metgrid reads each time-varying field as **five** Fortran sequential unformatted records: version `5`; **156-byte** mixed header (`hdate`, `xfcst`, `map_source`, `field`, `units`, `desc`, `xlvl`, `nx`, `ny`, `iproj`); projection extension (length depends on `iproj` / `map%igrid`); **4-byte** grid-relative wind flag; **`nx*ny*4`** bytes of big-endian `real32` slab data in Fortran order. The Python writer documents this in `ungrib_py/intermediate.py` and asserts payload sizes so the **layout stays aligned** with Fortran WPS builds that use gfortran **big-endian** unformatted I/O. **Floating-point values** may differ slightly from Fortran (ecCodes vs g2lib); **structure** (record order, lengths, endianness) is what metgrid relies on. +5. **namelist.wps** — Read `start_date` / `end_date`, `interval_seconds`, `prefix`, `out_format`, `ordered_by_date`, `pmin`. *Status: minimal reader; many `&ungrib` options still ignored.* + +### Phase B — Parallelism and workflow (done in this port) + +6. **Parallel file processing** — One worker per `GRIBFILE.*`; merge fields by validity time; write one output file per timestep in range. *Status: implemented (`parallel_ungrib.py`).* +7. **Configurable worker count** — CLI `--workers` (recommended 8–32 depending on filesystem). *Status: implemented.* + +### Phase C — Parity gaps (future work) + +8. **GRIB1** — Vtable GRIB1 columns and ECMWF-style `typeOfLevel` matching in `extract_file`. *Implemented for common cases; edge cases may differ from `rd_grib1.F`.* +9. **`rrpr` / `datint` / `PFILE:` pipeline** — Second pass, temporal interpolation, temp file handling from `ungrib.F`. *Not started; Python tool writes final `FILE:` directly from merged messages.* +10. **Edge cases** — JMA per-field grid refresh, Gaussian grid as `igrid=4`, Cassini `igrid=6`, UM soil specials, `ec_rec_len` EC paths, `add_lvls` / vertical interpolation. *Partial or not started.* +11. **Regression harness** — Automated byte-compare or field-compare against `ungrib.exe` on a reference dataset. *Not started; see testing below.* + +## How to test the Python code + +### 1. Unit tests (no GRIB files) + +From `ungrib/py` with a virtualenv and dependencies installed: + +```bash +cd ungrib/py +pip install -r requirements.txt +pytest tests/ -q +``` + +The `tests/conftest.py` file prepends the `ungrib/py` directory to `sys.path`, so `pytest` finds the `ungrib_py` package without setting `PYTHONPATH`. + +These cover Vtable parsing, namelist parsing, intermediate-file record layout, and pure-Python helpers. + +### 2. Integration test (with GRIB2 + ecCodes) + +1. Build or obtain WPS `namelist.wps` and link `Vtable` (e.g. `Vtable.GFS`) in a run directory. +2. Link or copy `GRIBFILE.AAA`, … as for the Fortran workflow. +3. Run (no `PYTHONPATH` if you use the launcher): + + ```bash + cd your_run_directory + ln -sf /path/to/WPS/ungrib/py/ungrib.py ./ungrib.py + chmod +x ./ungrib.py + ./ungrib.py --workers 8 + ``` + +4. Compare outputs: + - **Inventory**: Run WPS `rd_intermediate.exe` (from a full WPS build) on both Fortran and Python `FILE:` outputs and compare field lists and headers. + - **Numerical**: For a few fields, use Python or NCL to read slabs (same shape) and compare max abs error; expect small differences if bitmap/missing-value handling differs. + +### 3. Optional: ecCodes sanity check + +```bash +python -c "import eccodes; print('eccodes OK')" +``` + +If import fails, install the system ecCodes library (e.g. `libeccodes-dev` on Debian/Ubuntu) before `pip install eccodes`. + +## Success criteria (practical) + +- For a **GRIB2 lat-lon** (e.g. GFS on a regular grid) run with a standard Vtable, `ungrib_py` produces intermediate files that **metgrid accepts** and that list the same WPS field names and levels as Fortran ungrib for the same inputs. +- Wall time improves when using **multiple workers** and **multiple `GRIBFILE.*` chunks**, subject to storage bandwidth. diff --git a/ungrib/py/README.md b/ungrib/py/README.md new file mode 100644 index 00000000..75b8fa19 --- /dev/null +++ b/ungrib/py/README.md @@ -0,0 +1,103 @@ +# Python ungrib (`ungrib/py`) + +This directory contains a modern Python reimplementation of the WPS **ungrib** program, designed for **parallel I/O** over input GRIB files and clearer structure than the legacy Fortran stack in `ungrib/src/`. + +## How the original Fortran ungrib is organized + +The classic ungrib executable is driven by `ungrib.F` and supporting modules under `ungrib/src/`. At a high level: + +| Area | Role | +|------|------| +| **`ungrib.F`** | Main program: loops over `GRIBFILE.AAA`, `GRIBFILE.AAB`, … reads GRIB records, fills per-timestep storage, calls `output`, then runs post-pass steps (`rrpr`, `datint`, `file_delete`). | +| **`read_namelist.F`** | Reads `namelist.wps` (`&share`, `&ungrib`): simulation window, `interval_seconds`, `prefix`, `out_format`, `ordered_by_date`, optional level interpolation (`add_lvls`, `new_plvl`, …). | +| **`parse_table.F`**, **`table.F`** | Reads the **`Vtable`** file: maps GRIB (GRIB1 or GRIB2) metadata to WPS field names, units, descriptions, and output priorities. | +| **`rd_grib1.F`**, **`rd_grib2.F`** | Decode one GRIB message at a time; match against the Vtable; populate `gridinfo` and `storage_module`. | +| **`gridinfo.F`** | Map projection metadata (`map%igrid`, `nx`, `ny`, corners, `dx`/`dy`, earth radius, grid-relative winds). | +| **`output.F`** | Writes **WPS intermediate format** (version **5** when `out_format = 'WPS'`): unformatted sequential Fortran records consumed by **metgrid**. | +| **`rrpr.F`**, **`datint.F`**, **`file_delete.F`** | Second pass and temporal filling; cleanup of temporary `PFILE:` data. | +| **`Variable_Tables/`** | Example Vtables (`Vtable.GFS`, …) shipped with WPS. | +| **`src/ngl/`** | NCEP GRIB1/GRIB2 library (w3, g2, …) used by the Fortran reader. | + +Data flow (conceptual): + +```mermaid +flowchart LR + subgraph inputs + NL[namelist.wps] + VT[Vtable] + GF[GRIBFILE.*] + end + subgraph fortran_ungrib + RD[rd_grib1/2] + ST[storage] + OUT[output.F] + RR[rrpr / datint] + end + subgraph outputs + INT[FILE:YYYY-MM-DD_HH] + end + NL --> RD + VT --> RD + GF --> RD + RD --> ST + ST --> OUT + OUT --> RR + RR --> INT +``` + +## This Python package (`ungrib_py`) + +- **Parallelism**: Each worker processes one **GRIB input file** (`GRIBFILE.*`). Results are merged by **validity time** in the parent process, then **one intermediate file per timestep** is written (same naming pattern as Fortran: `prefix:YYYY-MM-DD_HH` for hourly data). +- **GRIB decoding**: **GRIB Edition 2** via **ecCodes** (`pip install eccodes`). GRIB1 is not implemented here yet. +- **Parity**: The Fortran tool’s full behavior (all projections, soil-level corner cases, `rrpr`, `datint`, EC non-standard record lengths, etc.) is large. This port focuses on a **maintainable core**: Vtable-driven GRIB2 extraction and WPS intermediate version **5** output for common grids. + +See **`MIGRATION.md`** for the staged migration plan, testing instructions, and known gaps. + +## Quick start (same idea as `ungrib.exe`) + +You do **not** need a virtual environment or **`PYTHONPATH`** if you run the top-level launcher script **`ungrib.py`**. That script lives next to the `ungrib_py` package and prepends its directory to `sys.path`, so it works when **symlinked** into your WPS run directory—like linking the old compiled binary. + +**1. Install Python dependencies once** (system Python 3.10+ is fine; use `--user` if you prefer not to touch the system site-packages): + +```bash +pip install --user -r /path/to/WPS/ungrib/py/requirements.txt +``` + +You may need OS packages for ecCodes (e.g. `libeccodes-dev` on Debian/Ubuntu) before `pip install eccodes` succeeds. + +**2. Link the launcher into your working directory** (same pattern as `ln -s .../ungrib.exe .`): + +```bash +ln -sf /path/to/WPS/ungrib/py/ungrib.py ./ungrib.py +chmod +x ./ungrib.py +``` + +**3. Run it** from that directory (where `namelist.wps`, `Vtable`, and `GRIBFILE.*` live): + +```bash +./ungrib.py --help +./ungrib.py --workers 16 +``` + +Alternatively, call the interpreter explicitly (no execute bit needed): + +```bash +python3 /path/to/WPS/ungrib/py/ungrib.py --workers 16 +``` + +**Optional:** `pip install /path/to/WPS/ungrib/py` installs the `ungrib-py` console entry point on your `PATH`; then you do not need a symlink (but the symlink matches the old `ungrib.exe` workflow). + +Expect `Vtable` and `namelist.wps` in the current working directory (same convention as Fortran ungrib). Input files must follow `GRIBFILE.AAA`, `GRIBFILE.AAB`, … naming. + +## Layout of `ungrib/py` + +| Path | Purpose | +|------|---------| +| `ungrib.py` | Launcher: sets `sys.path` and runs the CLI (symlink this like the old `ungrib.exe`). | +| `ungrib_py/vtable.py` | Parse WPS Vtable (GRIB2 columns). | +| `ungrib_py/grib2.py` | Scan messages with ecCodes; match Vtable; compute WPS level codes. | +| `ungrib_py/intermediate.py` | Write WPS intermediate format version 5 (Fortran unformatted records). | +| `ungrib_py/namelist_wps.py` | Minimal `namelist.wps` reader (`&share`, `&ungrib`). | +| `ungrib_py/parallel_ungrib.py` | Process pool over GRIB files; merge; write outputs. | +| `ungrib_py/cli.py` | Command-line entry point. | +| `tests/` | Unit tests (no GRIB fixture required for most tests). | 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..a63d5ffd --- /dev/null +++ b/ungrib/py/ungrib_py/__init__.py @@ -0,0 +1,13 @@ +# 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..088315ff --- /dev/null +++ b/ungrib/py/ungrib_py/__main__.py @@ -0,0 +1,16 @@ +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..542a0235 --- /dev/null +++ b/ungrib/py/ungrib_py/cli.py @@ -0,0 +1,69 @@ +# 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..22b0bb94 --- /dev/null +++ b/ungrib/py/ungrib_py/grib2.py @@ -0,0 +1,900 @@ +# GRIB2 extraction using ecCodes (Vtable-driven). + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path + +import numpy as np + +import eccodes + +from ungrib_py.intermediate import FieldSlab, MapInfo +from ungrib_py.vtable import BLANK, SPLAT, VtableEntry, match_entry + +_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: + + # Isobaric levels in WPS intermediate use Pa; match rrpr "plvl < 200000" and count_distinct_pressure_levels_pa. + + lf = float(lev) + return 50.0 <= lf <= 120000.0 + + +def _delete_surface_geopotential_sources(bucket: dict[tuple[float, str], FieldSlab]) -> None: + + # SOILGEO / surface GEOPT are not written by Fortran WPS output when desc is blank; drop for 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]: + + # Prefer the HGT row from the same Vtable as Fortran/Python ungrib (e.g. ungrib_out/Vtable). + + 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: + + # Fortran rrpr.F: HGT (m) = GEOPT (m^2 s^-2) / g on isobaric levels; metgrid reads HGT as GHT (METGRID.TBL). + + 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: + + # Fortran output.F (iflag=2): fields whose Vtable description is all blanks are not written to intermediate. + + 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: + + # METGRID.TBL ICON fills derived SOILT level 0 from SOILT000(200100). Vtable.ICONp uses SOILT001 for the top layer. + # If SOILT000 is absent, that output level is skipped and NUM_METGRID_SOIL_LEVELS can be 7 vs namelist 8. + # Duplicate the shallowest present SOILT* slab as SOILT000 when needed (also covers SOILT001 match failures). + + 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: + + # Fortran ungrib runs rrpr.F: SOILHGT (m) from SOILGEO (m^2 s^-2) when SOILHGT is missing. + # Also accept surface GEOPT (same units) or surface HGT already in meters (some Vtables). + + 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: + 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: + + # ERA5 / ECMWF surface geopotential is often paramId 129 (or shortName=z); Vtable.ERA-interim.pl has no GRIB2 columns so match_entry fails. + # Prefer SOILGEO, then GEOPT, when discipline 0 / surface type 1 matches typical orography GRIB2. + + 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 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, + ) + + 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..24bdac5e --- /dev/null +++ b/ungrib/py/ungrib_py/intermediate.py @@ -0,0 +1,302 @@ +# 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. With -fconvert=big-endian, + # gfortran also uses big-endian 32-bit record length tags (see WPS arch/configure.defaults). + + 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 = _pad_str(m.startloc[:8], 8).decode("ascii") + if m.igrid == 0: + buf = bytearray() + buf.extend(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + for v in (m.lat1, m.lon1, m.dx, m.dy, 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 + +def sort_fields_for_output( + items: list[tuple[float, str, FieldSlab]], + vtable_rows: list[VtableEntry], +) -> list[FieldSlab]: + + # This function sorts field items for WPS intermediate output (alias for sort_fields_fortran_order). + + return sort_fields_fortran_order(items, vtable_rows) + + + + + + + + + + diff --git a/ungrib/py/ungrib_py/namelist_wps.py b/ungrib/py/ungrib_py/namelist_wps.py new file mode 100644 index 00000000..b773d341 --- /dev/null +++ b/ungrib/py/ungrib_py/namelist_wps.py @@ -0,0 +1,129 @@ +# Minimal namelist.wps reader for ungrib. + +from __future__ import annotations + +import re +from pathlib import Path +from typing import NamedTuple + +class UngribNamelist(NamedTuple): + 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..f7c1ec1d --- /dev/null +++ b/ungrib/py/ungrib_py/parallel_ungrib.py @@ -0,0 +1,150 @@ +# 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.vtable import VtableEntry +from ungrib_py.intermediate import ( + FieldSlab, + MapInfo, + hdate_slice, + intermediate_filename_datelen, + sort_fields_for_output, + write_wps_intermediate_v5, +) +from ungrib_py.namelist_wps import UngribNamelist +from ungrib_py.vtable import 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, 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, 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, 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): + _fname, 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_for_output(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..e15530e8 --- /dev/null +++ b/ungrib/py/ungrib_py/vtable.py @@ -0,0 +1,157 @@ +# 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 + + + + + + + + + + From 453ee397422d74c7c65560cca3cd7bac291ace8f Mon Sep 17 00:00:00 2001 From: MHBalsmeier Date: Sun, 12 Apr 2026 10:20:07 +0200 Subject: [PATCH 4/6] simplifications in ungrib_py --- ungrib/py/ungrib_py/intermediate.py | 27 +++++++------------------- ungrib/py/ungrib_py/namelist_wps.py | 1 + ungrib/py/ungrib_py/parallel_ungrib.py | 15 +++++++------- 3 files changed, 16 insertions(+), 27 deletions(-) diff --git a/ungrib/py/ungrib_py/intermediate.py b/ungrib/py/ungrib_py/intermediate.py index 24bdac5e..6f34ab19 100644 --- a/ungrib/py/ungrib_py/intermediate.py +++ b/ungrib/py/ungrib_py/intermediate.py @@ -198,40 +198,40 @@ def _map_extension_bytes_be(m: MapInfo) -> bytes: # This function builds the map projection extension record with big-endian floats. - start = _pad_str(m.startloc[:8], 8).decode("ascii") + start_b = _pad_str(m.startloc[:8], 8) if m.igrid == 0: buf = bytearray() - buf.extend(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + 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(_pad_str(start, 8)) + buf.extend(start_b) for v in (m.lat1, m.lon1, m.dx, m.dy, m.centerlat, m.centerlon, m.r_earth_km): buf.extend(_be_f32(v)) return bytes(buf) @@ -282,19 +282,6 @@ def sort_fields_fortran_order( out.extend(slab for _, slab in at_level) return out -def sort_fields_for_output( - items: list[tuple[float, str, FieldSlab]], - vtable_rows: list[VtableEntry], -) -> list[FieldSlab]: - - # This function sorts field items for WPS intermediate output (alias for sort_fields_fortran_order). - - return sort_fields_fortran_order(items, vtable_rows) - - - - - diff --git a/ungrib/py/ungrib_py/namelist_wps.py b/ungrib/py/ungrib_py/namelist_wps.py index b773d341..efc7fabb 100644 --- a/ungrib/py/ungrib_py/namelist_wps.py +++ b/ungrib/py/ungrib_py/namelist_wps.py @@ -7,6 +7,7 @@ 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 diff --git a/ungrib/py/ungrib_py/parallel_ungrib.py b/ungrib/py/ungrib_py/parallel_ungrib.py index f7c1ec1d..f3a62e0d 100644 --- a/ungrib/py/ungrib_py/parallel_ungrib.py +++ b/ungrib/py/ungrib_py/parallel_ungrib.py @@ -14,17 +14,16 @@ extract_file, iter_grib_filenames, ) -from ungrib_py.vtable import VtableEntry from ungrib_py.intermediate import ( FieldSlab, MapInfo, hdate_slice, intermediate_filename_datelen, - sort_fields_for_output, + sort_fields_fortran_order, write_wps_intermediate_v5, ) from ungrib_py.namelist_wps import UngribNamelist -from ungrib_py.vtable import parse_vtable +from ungrib_py.vtable import VtableEntry, parse_vtable def discover_grib_files(cwd: Path) -> list[str]: @@ -34,7 +33,7 @@ def discover_grib_files(cwd: Path) -> list[str]: def _worker_extract( args: tuple[str, str, str, float], -) -> tuple[str, dict, MapInfo | None]: +) -> tuple[str, dict[str, dict[tuple[float, str], FieldSlab]], MapInfo | None]: # This function runs extract_file for one GRIB path inside a worker process. @@ -115,11 +114,13 @@ def run_ungrib_parallel( 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, MapInfo | None]] = [] + 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): - _fname, data, mi = fut.result() + _, 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) @@ -132,7 +133,7 @@ def run_ungrib_parallel( 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_for_output(items, vtable_rows) + 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) From 5b217306ae889b52df3f366b220b063ed4d3d311 Mon Sep 17 00:00:00 2001 From: MHBalsmeier Date: Sun, 12 Apr 2026 12:38:56 +0200 Subject: [PATCH 5/6] RAP support added to ungrib_py --- ungrib/py/ungrib_py/grib2.py | 111 ++++++++++++++++++++++++++++ ungrib/py/ungrib_py/intermediate.py | 3 +- 2 files changed, 113 insertions(+), 1 deletion(-) diff --git a/ungrib/py/ungrib_py/grib2.py b/ungrib/py/ungrib_py/grib2.py index 22b0bb94..aedd9863 100644 --- a/ungrib/py/ungrib_py/grib2.py +++ b/ungrib/py/ungrib_py/grib2.py @@ -4,6 +4,7 @@ from collections.abc import Iterator from pathlib import Path +import struct import numpy as np @@ -12,6 +13,30 @@ 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", @@ -457,6 +482,89 @@ def grid_relative_wind(gid: int) -> bool: 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. @@ -580,6 +688,9 @@ def map_info_from_grib(gid: int) -> MapInfo: 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: diff --git a/ungrib/py/ungrib_py/intermediate.py b/ungrib/py/ungrib_py/intermediate.py index 6f34ab19..2f6ea4d0 100644 --- a/ungrib/py/ungrib_py/intermediate.py +++ b/ungrib/py/ungrib_py/intermediate.py @@ -232,7 +232,8 @@ def _map_extension_bytes_be(m: MapInfo) -> bytes: if m.igrid == 6: buf = bytearray() buf.extend(start_b) - for v in (m.lat1, m.lon1, m.dx, m.dy, m.centerlat, m.centerlon, m.r_earth_km): + # 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") From 7d897c047291a96fd05b2c99ec2065882f65ff39 Mon Sep 17 00:00:00 2001 From: MHBalsmeier Date: Sun, 12 Apr 2026 12:48:36 +0200 Subject: [PATCH 6/6] several simplifications to ungrib_py --- ungrib/py/MIGRATION.md | 80 ++++------------- ungrib/py/README.md | 115 ++++++------------------- ungrib/py/ungrib_py/__init__.py | 10 --- ungrib/py/ungrib_py/__main__.py | 10 --- ungrib/py/ungrib_py/cli.py | 6 +- ungrib/py/ungrib_py/grib2.py | 36 ++++---- ungrib/py/ungrib_py/intermediate.py | 9 +- ungrib/py/ungrib_py/parallel_ungrib.py | 12 +-- ungrib/py/ungrib_py/vtable.py | 5 -- 9 files changed, 68 insertions(+), 215 deletions(-) diff --git a/ungrib/py/MIGRATION.md b/ungrib/py/MIGRATION.md index c094fc65..44dee1cc 100644 --- a/ungrib/py/MIGRATION.md +++ b/ungrib/py/MIGRATION.md @@ -1,75 +1,27 @@ -# Migration plan: Fortran ungrib → Python (`ungrib_py`) +# Fortran ungrib → `ungrib_py` -This document lists tasks to reach feature parity with `ungrib.exe`, notes what is already done in `ungrib_py`, and explains how to test the Python code. +## Done here -## Task outline +- 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. -### Phase A — Core decode and output (done in this port) +## Not done / gaps -1. **Vtable parsing** — Port `parse_table.F` logic for GRIB2 columns (discipline, category, parameter, first fixed surface / level type, product definition template). *Status: implemented for standard 11–12 column Vtables.* -2. **GRIB2 read path** — Replace `rd_grib2.F` + NCEP g2lib with **ecCodes** message iteration. *Status: implemented for common level types (pressure, mean sea level, height above ground, surface, soil 106, etc.) aligned with `rd_grib2.F`.* -3. **Grid metadata** — Port `gridinfo` population from GRIB2 Section 3 (lat/lon, Lambert, polar stereographic, Mercator). *Status: lat/lon (`regular_ll`) and Lambert/polar/Mercator via ecCodes `gridType`.* -4. **WPS intermediate writer** — Port `output.F` records for `out_format='WPS'`, version **5** (matches `metgrid/src/read_met_module.F`). *Status: implemented.* +- **`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`. -### WPS v5 file structure (metgrid) - -Metgrid reads each time-varying field as **five** Fortran sequential unformatted records: version `5`; **156-byte** mixed header (`hdate`, `xfcst`, `map_source`, `field`, `units`, `desc`, `xlvl`, `nx`, `ny`, `iproj`); projection extension (length depends on `iproj` / `map%igrid`); **4-byte** grid-relative wind flag; **`nx*ny*4`** bytes of big-endian `real32` slab data in Fortran order. The Python writer documents this in `ungrib_py/intermediate.py` and asserts payload sizes so the **layout stays aligned** with Fortran WPS builds that use gfortran **big-endian** unformatted I/O. **Floating-point values** may differ slightly from Fortran (ecCodes vs g2lib); **structure** (record order, lengths, endianness) is what metgrid relies on. -5. **namelist.wps** — Read `start_date` / `end_date`, `interval_seconds`, `prefix`, `out_format`, `ordered_by_date`, `pmin`. *Status: minimal reader; many `&ungrib` options still ignored.* - -### Phase B — Parallelism and workflow (done in this port) - -6. **Parallel file processing** — One worker per `GRIBFILE.*`; merge fields by validity time; write one output file per timestep in range. *Status: implemented (`parallel_ungrib.py`).* -7. **Configurable worker count** — CLI `--workers` (recommended 8–32 depending on filesystem). *Status: implemented.* - -### Phase C — Parity gaps (future work) - -8. **GRIB1** — Vtable GRIB1 columns and ECMWF-style `typeOfLevel` matching in `extract_file`. *Implemented for common cases; edge cases may differ from `rd_grib1.F`.* -9. **`rrpr` / `datint` / `PFILE:` pipeline** — Second pass, temporal interpolation, temp file handling from `ungrib.F`. *Not started; Python tool writes final `FILE:` directly from merged messages.* -10. **Edge cases** — JMA per-field grid refresh, Gaussian grid as `igrid=4`, Cassini `igrid=6`, UM soil specials, `ec_rec_len` EC paths, `add_lvls` / vertical interpolation. *Partial or not started.* -11. **Regression harness** — Automated byte-compare or field-compare against `ungrib.exe` on a reference dataset. *Not started; see testing below.* - -## How to test the Python code - -### 1. Unit tests (no GRIB files) - -From `ungrib/py` with a virtualenv and dependencies installed: - -```bash -cd ungrib/py -pip install -r requirements.txt -pytest tests/ -q -``` - -The `tests/conftest.py` file prepends the `ungrib/py` directory to `sys.path`, so `pytest` finds the `ungrib_py` package without setting `PYTHONPATH`. - -These cover Vtable parsing, namelist parsing, intermediate-file record layout, and pure-Python helpers. - -### 2. Integration test (with GRIB2 + ecCodes) - -1. Build or obtain WPS `namelist.wps` and link `Vtable` (e.g. `Vtable.GFS`) in a run directory. -2. Link or copy `GRIBFILE.AAA`, … as for the Fortran workflow. -3. Run (no `PYTHONPATH` if you use the launcher): - - ```bash - cd your_run_directory - ln -sf /path/to/WPS/ungrib/py/ungrib.py ./ungrib.py - chmod +x ./ungrib.py - ./ungrib.py --workers 8 - ``` - -4. Compare outputs: - - **Inventory**: Run WPS `rd_intermediate.exe` (from a full WPS build) on both Fortran and Python `FILE:` outputs and compare field lists and headers. - - **Numerical**: For a few fields, use Python or NCL to read slabs (same shape) and compare max abs error; expect small differences if bitmap/missing-value handling differs. - -### 3. Optional: ecCodes sanity check +## Test ```bash -python -c "import eccodes; print('eccodes OK')" +cd ungrib/py && pip install -r requirements.txt && pytest tests/ -q ``` -If import fails, install the system ecCodes library (e.g. `libeccodes-dev` on Debian/Ubuntu) before `pip install eccodes`. +Integration: same `namelist.wps` / `Vtable` / `GRIBFILE.*` as Fortran; compare with `rd_intermediate.exe` or slab diffs. -## Success criteria (practical) +## Practical success -- For a **GRIB2 lat-lon** (e.g. GFS on a regular grid) run with a standard Vtable, `ungrib_py` produces intermediate files that **metgrid accepts** and that list the same WPS field names and levels as Fortran ungrib for the same inputs. -- Wall time improves when using **multiple workers** and **multiple `GRIBFILE.*` chunks**, subject to storage bandwidth. +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 index 75b8fa19..09274b70 100644 --- a/ungrib/py/README.md +++ b/ungrib/py/README.md @@ -1,103 +1,42 @@ # Python ungrib (`ungrib/py`) -This directory contains a modern Python reimplementation of the WPS **ungrib** program, designed for **parallel I/O** over input GRIB files and clearer structure than the legacy Fortran stack in `ungrib/src/`. +Parallel Python port of WPS **ungrib**: GRIB2 → WPS intermediate v5 for **metgrid**, with clearer structure than `ungrib/src/` Fortran. -## How the original Fortran ungrib is organized +## Fortran ungrib (reference) -The classic ungrib executable is driven by `ungrib.F` and supporting modules under `ungrib/src/`. At a high level: +| 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'`) | -| Area | Role | -|------|------| -| **`ungrib.F`** | Main program: loops over `GRIBFILE.AAA`, `GRIBFILE.AAB`, … reads GRIB records, fills per-timestep storage, calls `output`, then runs post-pass steps (`rrpr`, `datint`, `file_delete`). | -| **`read_namelist.F`** | Reads `namelist.wps` (`&share`, `&ungrib`): simulation window, `interval_seconds`, `prefix`, `out_format`, `ordered_by_date`, optional level interpolation (`add_lvls`, `new_plvl`, …). | -| **`parse_table.F`**, **`table.F`** | Reads the **`Vtable`** file: maps GRIB (GRIB1 or GRIB2) metadata to WPS field names, units, descriptions, and output priorities. | -| **`rd_grib1.F`**, **`rd_grib2.F`** | Decode one GRIB message at a time; match against the Vtable; populate `gridinfo` and `storage_module`. | -| **`gridinfo.F`** | Map projection metadata (`map%igrid`, `nx`, `ny`, corners, `dx`/`dy`, earth radius, grid-relative winds). | -| **`output.F`** | Writes **WPS intermediate format** (version **5** when `out_format = 'WPS'`): unformatted sequential Fortran records consumed by **metgrid**. | -| **`rrpr.F`**, **`datint.F`**, **`file_delete.F`** | Second pass and temporal filling; cleanup of temporary `PFILE:` data. | -| **`Variable_Tables/`** | Example Vtables (`Vtable.GFS`, …) shipped with WPS. | -| **`src/ngl/`** | NCEP GRIB1/GRIB2 library (w3, g2, …) used by the Fortran reader. | - -Data flow (conceptual): - -```mermaid -flowchart LR - subgraph inputs - NL[namelist.wps] - VT[Vtable] - GF[GRIBFILE.*] - end - subgraph fortran_ungrib - RD[rd_grib1/2] - ST[storage] - OUT[output.F] - RR[rrpr / datint] - end - subgraph outputs - INT[FILE:YYYY-MM-DD_HH] - end - NL --> RD - VT --> RD - GF --> RD - RD --> ST - ST --> OUT - OUT --> RR - RR --> INT -``` - -## This Python package (`ungrib_py`) - -- **Parallelism**: Each worker processes one **GRIB input file** (`GRIBFILE.*`). Results are merged by **validity time** in the parent process, then **one intermediate file per timestep** is written (same naming pattern as Fortran: `prefix:YYYY-MM-DD_HH` for hourly data). -- **GRIB decoding**: **GRIB Edition 2** via **ecCodes** (`pip install eccodes`). GRIB1 is not implemented here yet. -- **Parity**: The Fortran tool’s full behavior (all projections, soil-level corner cases, `rrpr`, `datint`, EC non-standard record lengths, etc.) is large. This port focuses on a **maintainable core**: Vtable-driven GRIB2 extraction and WPS intermediate version **5** output for common grids. +## This package (`ungrib_py`) -See **`MIGRATION.md`** for the staged migration plan, testing instructions, and known gaps. +- 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 (same idea as `ungrib.exe`) - -You do **not** need a virtual environment or **`PYTHONPATH`** if you run the top-level launcher script **`ungrib.py`**. That script lives next to the `ungrib_py` package and prepends its directory to `sys.path`, so it works when **symlinked** into your WPS run directory—like linking the old compiled binary. - -**1. Install Python dependencies once** (system Python 3.10+ is fine; use `--user` if you prefer not to touch the system site-packages): +## Quick start ```bash pip install --user -r /path/to/WPS/ungrib/py/requirements.txt -``` - -You may need OS packages for ecCodes (e.g. `libeccodes-dev` on Debian/Ubuntu) before `pip install eccodes` succeeds. - -**2. Link the launcher into your working directory** (same pattern as `ln -s .../ungrib.exe .`): - -```bash -ln -sf /path/to/WPS/ungrib/py/ungrib.py ./ungrib.py -chmod +x ./ungrib.py -``` - -**3. Run it** from that directory (where `namelist.wps`, `Vtable`, and `GRIBFILE.*` live): - -```bash -./ungrib.py --help +ln -sf /path/to/WPS/ungrib/py/ungrib.py ./ungrib.py && chmod +x ./ungrib.py ./ungrib.py --workers 16 ``` -Alternatively, call the interpreter explicitly (no execute bit needed): +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. -```bash -python3 /path/to/WPS/ungrib/py/ungrib.py --workers 16 -``` - -**Optional:** `pip install /path/to/WPS/ungrib/py` installs the `ungrib-py` console entry point on your `PATH`; then you do not need a symlink (but the symlink matches the old `ungrib.exe` workflow). - -Expect `Vtable` and `namelist.wps` in the current working directory (same convention as Fortran ungrib). Input files must follow `GRIBFILE.AAA`, `GRIBFILE.AAB`, … naming. +## Layout -## Layout of `ungrib/py` - -| Path | Purpose | -|------|---------| -| `ungrib.py` | Launcher: sets `sys.path` and runs the CLI (symlink this like the old `ungrib.exe`). | -| `ungrib_py/vtable.py` | Parse WPS Vtable (GRIB2 columns). | -| `ungrib_py/grib2.py` | Scan messages with ecCodes; match Vtable; compute WPS level codes. | -| `ungrib_py/intermediate.py` | Write WPS intermediate format version 5 (Fortran unformatted records). | -| `ungrib_py/namelist_wps.py` | Minimal `namelist.wps` reader (`&share`, `&ungrib`). | -| `ungrib_py/parallel_ungrib.py` | Process pool over GRIB files; merge; write outputs. | -| `ungrib_py/cli.py` | Command-line entry point. | -| `tests/` | Unit tests (no GRIB fixture required for most tests). | +| 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/ungrib_py/__init__.py b/ungrib/py/ungrib_py/__init__.py index a63d5ffd..41f32e10 100644 --- a/ungrib/py/ungrib_py/__init__.py +++ b/ungrib/py/ungrib_py/__init__.py @@ -1,13 +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 index 088315ff..562a5c0f 100644 --- a/ungrib/py/ungrib_py/__main__.py +++ b/ungrib/py/ungrib_py/__main__.py @@ -4,13 +4,3 @@ 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 index 542a0235..1f7ca533 100644 --- a/ungrib/py/ungrib_py/cli.py +++ b/ungrib/py/ungrib_py/cli.py @@ -9,6 +9,7 @@ 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. @@ -62,8 +63,3 @@ def main(argv: list[str] | None = None) -> int: - - - - - diff --git a/ungrib/py/ungrib_py/grib2.py b/ungrib/py/ungrib_py/grib2.py index aedd9863..fe08c14d 100644 --- a/ungrib/py/ungrib_py/grib2.py +++ b/ungrib/py/ungrib_py/grib2.py @@ -53,7 +53,7 @@ def _is_pressure_level_pa(lev: float) -> bool: - # Isobaric levels in WPS intermediate use Pa; match rrpr "plvl < 200000" and count_distinct_pressure_levels_pa. + # 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 @@ -61,7 +61,7 @@ def _is_pressure_level_pa(lev: float) -> bool: def _delete_surface_geopotential_sources(bucket: dict[tuple[float, str], FieldSlab]) -> None: - # SOILGEO / surface GEOPT are not written by Fortran WPS output when desc is blank; drop for metgrid. + # 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 @@ -74,7 +74,7 @@ def _delete_surface_geopotential_sources(bucket: dict[tuple[float, str], FieldSl def _hgt_field_meta_from_vtable(rows: list[VtableEntry] | None) -> tuple[str, str, str]: - # Prefer the HGT row from the same Vtable as Fortran/Python ungrib (e.g. ungrib_out/Vtable). + # This function returns (name, units, desc) for derived HGT from GEOPT, preferring the Vtable HGT row. if rows: for r in rows: @@ -89,7 +89,7 @@ def ensure_hgt_from_geopt( vtable_rows: list[VtableEntry] | None = None, ) -> None: - # Fortran rrpr.F: HGT (m) = GEOPT (m^2 s^-2) / g on isobaric levels; metgrid reads HGT as GHT (METGRID.TBL). + # 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()): @@ -120,7 +120,7 @@ def drop_fields_blank_vtable_desc( all_rows: list[VtableEntry], ) -> None: - # Fortran output.F (iflag=2): fields whose Vtable description is all blanks are not written to intermediate. + # 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: @@ -137,9 +137,7 @@ def drop_fields_blank_vtable_desc( def ensure_soilt000_for_metgrid(bucket: dict[tuple[float, str], FieldSlab]) -> None: - # METGRID.TBL ICON fills derived SOILT level 0 from SOILT000(200100). Vtable.ICONp uses SOILT001 for the top layer. - # If SOILT000 is absent, that output level is skipped and NUM_METGRID_SOIL_LEVELS can be 7 vs namelist 8. - # Duplicate the shallowest present SOILT* slab as SOILT000 when needed (also covers SOILT001 match failures). + # 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: @@ -159,10 +157,10 @@ def ensure_soilt000_for_metgrid(bucket: dict[tuple[float, str], FieldSlab]) -> N ) break + def ensure_soilhgt_from_soilgeo(bucket: dict[tuple[float, str], FieldSlab]) -> None: - # Fortran ungrib runs rrpr.F: SOILHGT (m) from SOILGEO (m^2 s^-2) when SOILHGT is missing. - # Also accept surface GEOPT (same units) or surface HGT already in meters (some Vtables). + # 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": @@ -172,6 +170,9 @@ def ensure_soilhgt_from_soilgeo(bucket: dict[tuple[float, str], FieldSlab]) -> N 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 @@ -208,10 +209,10 @@ def find_at_surface(names: tuple[str, ...]) -> FieldSlab | None: data=np.asfortranarray(arr.astype(np.float32)), ) + def ecmwf_surface_geopotential_vtable_row(gid: int, rows: list[VtableEntry]) -> VtableEntry | None: - # ERA5 / ECMWF surface geopotential is often paramId 129 (or shortName=z); Vtable.ERA-interim.pl has no GRIB2 columns so match_entry fails. - # Prefer SOILGEO, then GEOPT, when discipline 0 / surface type 1 matches typical orography GRIB2. + # 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 @@ -235,6 +236,7 @@ def ecmwf_surface_geopotential_vtable_row(gid: int, rows: list[VtableEntry]) -> return geopt[0] return None + def _iget(gid: int, key: str) -> int: # This function returns an integer ecCodes key value for a GRIB handle. @@ -888,6 +890,7 @@ def pick_vtable_row_ecmwf_soil_layer151( 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. @@ -921,6 +924,7 @@ def pick_vtable_row(gid: int, rows: list[VtableEntry]) -> VtableEntry | 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. @@ -928,6 +932,7 @@ def values_to_slab(values: np.ndarray, ni: int, nj: int) -> np.ndarray: 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], @@ -1002,10 +1007,3 @@ def iter_grib_filenames() -> Iterator[str]: - - - - - - - diff --git a/ungrib/py/ungrib_py/intermediate.py b/ungrib/py/ungrib_py/intermediate.py index 2f6ea4d0..9f987d05 100644 --- a/ungrib/py/ungrib_py/intermediate.py +++ b/ungrib/py/ungrib_py/intermediate.py @@ -50,8 +50,7 @@ def _pad_str(s: str, n: int) -> bytes: def _fortran_unformatted_record(fh: BinaryIO, payload: bytes) -> None: - # This function writes one gfortran sequential unformatted record. With -fconvert=big-endian, - # gfortran also uses big-endian 32-bit record length tags (see WPS arch/configure.defaults). + # This function writes one gfortran sequential unformatted record using big-endian 32-bit length tags. n = len(payload) len_be = struct.pack(">I", n) @@ -282,9 +281,3 @@ def sort_fields_fortran_order( ) out.extend(slab for _, slab in at_level) return out - - - - - - diff --git a/ungrib/py/ungrib_py/parallel_ungrib.py b/ungrib/py/ungrib_py/parallel_ungrib.py index f3a62e0d..f908823b 100644 --- a/ungrib/py/ungrib_py/parallel_ungrib.py +++ b/ungrib/py/ungrib_py/parallel_ungrib.py @@ -58,7 +58,12 @@ def count_distinct_pressure_levels_pa( return len(levels) def merge_by_time( - parts: list[tuple[dict, MapInfo | None]], + parts: list[ + tuple[ + dict[str, dict[tuple[float, str], FieldSlab]], + MapInfo | None, + ] + ], hstart: str, hend: str, all_vtable_rows: list[VtableEntry], @@ -144,8 +149,3 @@ def run_ungrib_parallel( - - - - - diff --git a/ungrib/py/ungrib_py/vtable.py b/ungrib/py/ungrib_py/vtable.py index e15530e8..a9d0de16 100644 --- a/ungrib/py/ungrib_py/vtable.py +++ b/ungrib/py/ungrib_py/vtable.py @@ -150,8 +150,3 @@ def match_entry( - - - - -