From b3611aa53bc2db8ee40ba39901c95e2a99cbec04 Mon Sep 17 00:00:00 2001 From: Maria Frediani Date: Mon, 20 Jul 2026 08:55:05 -0600 Subject: [PATCH] Support wrfinput and nested WPS grid spacing --- README.md | 9 ++- fire_preprocess/cli.py | 4 +- fire_preprocess/namelist.py | 86 +++++++++++++++++++++++-- fire_preprocess/wrf_domain.py | 56 +++++++++++------ tests/__init__.py | 1 + tests/test_namelist.py | 114 ++++++++++++++++++++++++++++++++++ tests/test_wrf_domain.py | 98 +++++++++++++++++++++++++++++ 7 files changed, 339 insertions(+), 29 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/test_namelist.py create mode 100644 tests/test_wrf_domain.py diff --git a/README.md b/README.md index 4eddeca..b94efd0 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # fire_preprocess A standalone Python tool that adds fire-specific input fields (`NFUEL_CAT` and `ZSF`) directly to -WPS output files (`geo_em` or `met_em`) for use with the Community Fire Behavior Model +WRF/WPS netCDF files (`geo_em`, `met_em`, or `wrfinput`) for use with the Community Fire Behavior Model (https://github.com/NCAR/fire_behavior), bypassing the traditional workflow that requires converting data to WPS geogrid binary format and editing `GEOGRID.TBL`. @@ -78,7 +78,7 @@ All arguments can also be supplied via a YAML config file (see below). | Argument | Default | Description | |---|---|---| -| `--wps-files` | — | WPS output file (`geo_em`/`met_em`) or glob pattern (e.g. `'met_em.d02.*.nc'`) | +| `--wps-files` | — | WRF/WPS file (`geo_em`/`met_em`/`wrfinput`) or glob pattern (e.g. `'met_em.d02.*.nc'`) | | `--fuel` | — | LANDFIRE fuel-category GeoTIFF (NFUEL_CAT source) | | `--zsf` | — | High-resolution terrain DEM GeoTIFF (ZSF source) | | `--namelist` | `namelist.wps` | Path to `namelist.wps` | @@ -147,3 +147,8 @@ reprojects to the WRF fire subgrid using the projection parameters and `subgrid_ratio_x/y` values from `namelist.wps`, and writes `NFUEL_CAT` and `ZSF` into the existing WPS netCDF files. The rest of the workflow (`real.exe` → WRF) is unchanged. + +For nested WPS configurations, `dx` and `dy` are normally scalar values for the +parent domain. When processing a child domain selected with `--domain`, the tool +derives the child-domain grid spacing from `parent_grid_ratio`. If `dx` and `dy` +are supplied as per-domain arrays, those explicit values are used instead. diff --git a/fire_preprocess/cli.py b/fire_preprocess/cli.py index 9490199..588ac3b 100644 --- a/fire_preprocess/cli.py +++ b/fire_preprocess/cli.py @@ -76,7 +76,7 @@ def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="fire_preprocess", description=( - "Add CFBM fire fields (NFUEL_CAT, ZSF) directly to WPS netCDF files,\n" + "Add CFBM fire fields (NFUEL_CAT, ZSF) directly to WRF/WPS netCDF files,\n" "bypassing GEOGRID.TBL editing and geogrid binary format conversion." ), formatter_class=argparse.RawDescriptionHelpFormatter, @@ -98,7 +98,7 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument( "--wps-files", default=None, metavar="PATH", dest="wps_files", - help="Path to a WPS output file (geo_em/met_em) or a glob pattern (e.g. 'met_em.d01.*.nc')", + help="Path to a WRF/WPS file (geo_em/met_em/wrfinput) or a glob pattern (e.g. 'met_em.d01.*.nc')", ) parser.add_argument( "--zsf", default=None, metavar="GEOTIFF", diff --git a/fire_preprocess/namelist.py b/fire_preprocess/namelist.py index 07ebe73..d803498 100644 --- a/fire_preprocess/namelist.py +++ b/fire_preprocess/namelist.py @@ -2,6 +2,11 @@ import f90nml +def _is_sequence(value): + """Return True for per-domain namelist arrays.""" + return isinstance(value, (list, tuple)) + + def read_namelist_wps(path): """Return a flat dict of all namelist.wps parameters (share + geogrid merged).""" nml = f90nml.read(path) @@ -13,11 +18,75 @@ def read_namelist_wps(path): def _scalar(value, index=0): """Return a scalar from a value that may be a list (per-domain array).""" - if isinstance(value, (list, tuple)): + if _is_sequence(value): return value[index] return value +def _effective_grid_spacing(params, domain_index, axis): + """Return domain grid spacing, deriving nests from parent_grid_ratio. + + WPS normally stores dx/dy as scalar parent-domain spacings. For nested + domains, the child spacing is the parent spacing divided by the refinement + ratio along the parent chain. + """ + value = params[axis] + if _is_sequence(value): + return float(_scalar(value, domain_index)) + + if domain_index == 0: + return float(value) + + parent_ids = params.get("parent_id") + parent_ratios = params.get("parent_grid_ratio") + if parent_ids is None or parent_ratios is None: + raise ValueError( + f"Cannot derive nested-domain {axis}: parent_id and " + "parent_grid_ratio are required when dx/dy are scalar." + ) + + parent_index = int(_scalar(parent_ids, domain_index)) - 1 + if parent_index == domain_index: + raise ValueError(f"Domain {domain_index + 1} cannot be its own parent.") + + parent_spacing = _effective_grid_spacing(params, parent_index, axis) + ratio = float(_scalar(parent_ratios, domain_index)) + return parent_spacing / ratio + + +def _parent_grid_info(params, domain_index): + """Return nested-domain parent metadata, or None for the root domain.""" + if domain_index == 0: + return None + + parent_index = int(_scalar(params["parent_id"], domain_index)) - 1 + ratio = int(_scalar(params["parent_grid_ratio"], domain_index)) + i_start = int(_scalar(params["i_parent_start"], domain_index)) + j_start = int(_scalar(params["j_parent_start"], domain_index)) + return { + "parent_index": parent_index, + "parent_grid_ratio": ratio, + "i_parent_start": i_start, + "j_parent_start": j_start, + } + + +def _mass_offset_from_root(params, domain_index, axis): + """Return nested-domain SW mass-point offset from the root mass grid.""" + if domain_index == 0: + return 0.0 + + parent = _parent_grid_info(params, domain_index) + parent_index = parent["parent_index"] + parent_spacing = _effective_grid_spacing(params, parent_index, axis) + start_key = "i_parent_start" if axis == "dx" else "j_parent_start" + + # WPS parent starts are 1-based parent-domain mass-grid indices. Convert + # them to a zero-based offset, then accumulate through any deeper nests. + parent_offset = _mass_offset_from_root(params, parent_index, axis) + return parent_offset + (parent[start_key] - 1) * parent_spacing + + def get_fire_subgrid_ratios(path, domain_index=0): """Return (sr_x, sr_y) fire subgrid refinement ratios from namelist.wps. @@ -45,6 +114,8 @@ def get_domain_params(path, domain_index=0): e_we = int(_scalar(params["e_we"], domain_index)) e_sn = int(_scalar(params["e_sn"], domain_index)) + root_e_we = int(_scalar(params["e_we"], 0)) + root_e_sn = int(_scalar(params["e_sn"], 0)) proj_map = {"lambert": 1, "polar": 2, "mercator": 3, "lat-lon": 6, "latlong": 6} raw_proj = str(params.get("map_proj", "lambert")).strip().lower().strip("'\"") @@ -57,10 +128,15 @@ def get_domain_params(path, domain_index=0): "stand_lon": float(params.get("stand_lon", 0.0)), "ref_lat": float(params.get("ref_lat", 0.0)), "ref_lon": float(params.get("ref_lon", 0.0)), - "ref_x": float(params.get("ref_x", (e_we + 1) / 2.0)), - "ref_y": float(params.get("ref_y", (e_sn + 1) / 2.0)), - "dx": float(_scalar(params["dx"], domain_index)), - "dy": float(_scalar(params["dy"], domain_index)), + "ref_x": float(params.get("ref_x", (root_e_we + 1) / 2.0)), + "ref_y": float(params.get("ref_y", (root_e_sn + 1) / 2.0)), + "root_dx": _effective_grid_spacing(params, 0, "dx"), + "root_dy": _effective_grid_spacing(params, 0, "dy"), + "x_mass_offset_from_root": _mass_offset_from_root(params, domain_index, "dx"), + "y_mass_offset_from_root": _mass_offset_from_root(params, domain_index, "dy"), + "dx": _effective_grid_spacing(params, domain_index, "dx"), + "dy": _effective_grid_spacing(params, domain_index, "dy"), "e_we": e_we, "e_sn": e_sn, + "parent": _parent_grid_info(params, domain_index), } diff --git a/fire_preprocess/wrf_domain.py b/fire_preprocess/wrf_domain.py index ded3a4f..829a279 100644 --- a/fire_preprocess/wrf_domain.py +++ b/fire_preprocess/wrf_domain.py @@ -58,19 +58,33 @@ def build_wrf_crs(map_proj, truelat1, truelat2, stand_lon): return CRS.from_proj4(proj_str) -def _sw_corner_from_wps(ds, crs): - """Return (x_sw_mass, y_sw_mass) in WRF proj coords from WPS arrays.""" - if "XLAT_M" not in ds.variables or "XLONG_M" not in ds.variables: - raise KeyError("netCDF file is missing XLAT_M / XLONG_M variables") +def _mass_latlon_names(ds): + """Return mass-grid latitude/longitude variable names for WPS or WRF files.""" + candidates = ( + ("XLAT_M", "XLONG_M"), # geo_em/met_em + ("XLAT", "XLONG"), # wrfinput + ) + for lat_name, lon_name in candidates: + if lat_name in ds.variables and lon_name in ds.variables: + return lat_name, lon_name + raise KeyError("netCDF file is missing mass-grid latitude/longitude variables") + + +def _sw_corner_from_file(ds, crs): + """Return mass-grid geometry from WPS or WRF latitude/longitude arrays.""" + lat_name, lon_name = _mass_latlon_names(ds) + lat = ds.variables[lat_name] + lon = ds.variables[lon_name] - lat_sw = float(ds.variables["XLAT_M"][0, 0, 0]) - lon_sw = float(ds.variables["XLONG_M"][0, 0, 0]) + lat_sw = float(lat[0, 0, 0]) + lon_sw = float(lon[0, 0, 0]) # Use WRF's sphere for the geographic source CRS so the transformer is # internally consistent with the projection definition. geo_crs = CRS.from_proj4(f"+proj=longlat +a={WRF_SPHERE_RADIUS} +b={WRF_SPHERE_RADIUS} +no_defs") transformer = Transformer.from_crs(geo_crs, crs, always_xy=True) - return transformer.transform(lon_sw, lat_sw) + x_sw_mass, y_sw_mass = transformer.transform(lon_sw, lat_sw) + return x_sw_mass, y_sw_mass, lat.shape[1], lat.shape[2] def _sw_corner_from_namelist(params, crs): @@ -80,26 +94,29 @@ def _sw_corner_from_namelist(params, crs): x_ref, y_ref = transformer.transform(params["ref_lon"], params["ref_lat"]) - # ref_x/ref_y are 1-indexed staggered grid positions that map to ref_lat/ref_lon. - # Convert to 0-indexed mass point offset. + # ref_x/ref_y describe the root WPS grid. Convert that reference location + # to the root-domain SW mass point, then add the selected nest's accumulated + # parent-grid offset. i_ref = params["ref_x"] - 1.0 # staggered i → subtract 0.5 for mass point j_ref = params["ref_y"] - 1.0 # staggered j i_ref_mass = i_ref - 0.5 j_ref_mass = j_ref - 0.5 - x_sw_mass = x_ref - i_ref_mass * params["dx"] - y_sw_mass = y_ref - j_ref_mass * params["dy"] + root_x_sw_mass = x_ref - i_ref_mass * params["root_dx"] + root_y_sw_mass = y_ref - j_ref_mass * params["root_dy"] + x_sw_mass = root_x_sw_mass + params.get("x_mass_offset_from_root", 0.0) + y_sw_mass = root_y_sw_mass + params.get("y_mass_offset_from_root", 0.0) return x_sw_mass, y_sw_mass def read_domain_from_file(file_path, sr_x, sr_y, namelist_params=None): - """Build a WRFDomain by reading a WPS netCDF file. + """Build a WRFDomain by reading a WPS or WRF input netCDF file. Args: - file_path: path to any met_em*.nc or geo_em*.nc file for the domain + file_path: path to any met_em*.nc, geo_em*.nc, or wrfinput* file sr_x, sr_y: fire subgrid ratios from namelist.wps namelist_params: optional dict from namelist.get_domain_params(), used - as fallback if XLAT_M/XLONG_M are absent + as fallback if mass-grid lat/lon arrays are absent Returns: WRFDomain dataclass @@ -117,14 +134,13 @@ def read_domain_from_file(file_path, sr_x, sr_y, namelist_params=None): crs = build_wrf_crs(map_proj, truelat1, truelat2, stand_lon) try: - x_sw_mass, y_sw_mass = _sw_corner_from_wps(ds, crs) - ny = ds.variables["XLAT_M"].shape[1] - nx = ds.variables["XLAT_M"].shape[2] - except KeyError: + x_sw_mass, y_sw_mass, ny, nx = _sw_corner_from_file(ds, crs) + except KeyError as exc: if namelist_params is None: raise RuntimeError( - "XLAT_M/XLONG_M not found in file and no namelist_params supplied" - ) + "Mass-grid latitude/longitude variables not found in file " + "and no namelist_params supplied" + ) from exc x_sw_mass, y_sw_mass = _sw_corner_from_namelist(namelist_params, crs) nx = namelist_params["e_we"] - 1 ny = namelist_params["e_sn"] - 1 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..29f6495 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for fire_preprocess.""" diff --git a/tests/test_namelist.py b/tests/test_namelist.py new file mode 100644 index 0000000..5792ae8 --- /dev/null +++ b/tests/test_namelist.py @@ -0,0 +1,114 @@ +import os +import tempfile +import unittest + +from fire_preprocess.namelist import get_domain_params + + +class NamelistDomainParamsTest(unittest.TestCase): + def _write_namelist(self, text): + handle = tempfile.NamedTemporaryFile( + "w", suffix=".wps", delete=False, dir=os.environ.get("TMPDIR") + ) + self.addCleanup(lambda: os.path.exists(handle.name) and os.unlink(handle.name)) + with handle: + handle.write(text) + return handle.name + + def test_scalar_dx_dy_are_refined_by_parent_grid_ratio(self): + path = self._write_namelist( + """ +&share + max_dom = 2, +/ +&geogrid + parent_id = 1, 1, + parent_grid_ratio = 1, 3, + i_parent_start = 1, 180, + j_parent_start = 1, 150, + e_we = 802, 1201, + e_sn = 802, 1201, + dx = 300, + dy = 300, + map_proj = 'lambert', + ref_lat = 37.30717, + ref_lon = -107.4509, + truelat1 = 30.0, + truelat2 = 50.0, + stand_lon = -105.5, +/ +""" + ) + + params = get_domain_params(path, domain_index=1) + + self.assertEqual(params["dx"], 100.0) + self.assertEqual(params["dy"], 100.0) + self.assertEqual(params["x_mass_offset_from_root"], 179 * 300.0) + self.assertEqual(params["y_mass_offset_from_root"], 149 * 300.0) + + def test_explicit_dx_dy_arrays_are_preserved(self): + path = self._write_namelist( + """ +&share + max_dom = 2, +/ +&geogrid + parent_id = 1, 1, + parent_grid_ratio = 1, 3, + i_parent_start = 1, 180, + j_parent_start = 1, 150, + e_we = 802, 1201, + e_sn = 802, 1201, + dx = 300, 90, + dy = 300, 90, + map_proj = 'lambert', + ref_lat = 37.30717, + ref_lon = -107.4509, + truelat1 = 30.0, + truelat2 = 50.0, + stand_lon = -105.5, +/ +""" + ) + + params = get_domain_params(path, domain_index=1) + + self.assertEqual(params["dx"], 90.0) + self.assertEqual(params["dy"], 90.0) + + def test_recursive_nested_spacing(self): + path = self._write_namelist( + """ +&share + max_dom = 3, +/ +&geogrid + parent_id = 1, 1, 2, + parent_grid_ratio = 1, 3, 5, + i_parent_start = 1, 180, 20, + j_parent_start = 1, 150, 30, + e_we = 802, 1201, 501, + e_sn = 802, 1201, 501, + dx = 300, + dy = 300, + map_proj = 'lambert', + ref_lat = 37.30717, + ref_lon = -107.4509, + truelat1 = 30.0, + truelat2 = 50.0, + stand_lon = -105.5, +/ +""" + ) + + params = get_domain_params(path, domain_index=2) + + self.assertEqual(params["dx"], 20.0) + self.assertEqual(params["dy"], 20.0) + self.assertEqual(params["x_mass_offset_from_root"], 179 * 300.0 + 19 * 100.0) + self.assertEqual(params["y_mass_offset_from_root"], 149 * 300.0 + 29 * 100.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_wrf_domain.py b/tests/test_wrf_domain.py new file mode 100644 index 0000000..4ad74eb --- /dev/null +++ b/tests/test_wrf_domain.py @@ -0,0 +1,98 @@ +import os +import tempfile +import unittest + +import netCDF4 as nc + +from fire_preprocess.wrf_domain import read_domain_from_file + + +class WRFDomainFileTest(unittest.TestCase): + def _make_file(self, lat_name, lon_name): + handle = tempfile.NamedTemporaryFile( + suffix=".nc", delete=False, dir=os.environ.get("TMPDIR") + ) + handle.close() + self.addCleanup(lambda: os.path.exists(handle.name) and os.unlink(handle.name)) + + with nc.Dataset(handle.name, "w") as ds: + ds.createDimension("Time", 1) + ds.createDimension("south_north", 3) + ds.createDimension("west_east", 4) + lat = ds.createVariable(lat_name, "f4", ("Time", "south_north", "west_east")) + lon = ds.createVariable(lon_name, "f4", ("Time", "south_north", "west_east")) + lat[:] = 37.0 + lon[:] = -107.0 + ds.MAP_PROJ = 6 + ds.TRUELAT1 = 0.0 + ds.TRUELAT2 = 0.0 + ds.STAND_LON = 0.0 + ds.DX = 100.0 + ds.DY = 100.0 + return handle.name + + def _make_file_without_coordinates(self): + handle = tempfile.NamedTemporaryFile( + suffix=".nc", delete=False, dir=os.environ.get("TMPDIR") + ) + handle.close() + self.addCleanup(lambda: os.path.exists(handle.name) and os.unlink(handle.name)) + + with nc.Dataset(handle.name, "w") as ds: + ds.MAP_PROJ = 6 + ds.TRUELAT1 = 0.0 + ds.TRUELAT2 = 0.0 + ds.STAND_LON = 0.0 + ds.DX = 100.0 + ds.DY = 100.0 + return handle.name + + def test_reads_wrfinput_mass_coordinates(self): + path = self._make_file("XLAT", "XLONG") + + domain = read_domain_from_file(path, sr_x=4, sr_y=4) + + self.assertEqual(domain.nx, 4) + self.assertEqual(domain.ny, 3) + self.assertEqual(domain.dx, 100.0) + self.assertEqual(domain.dy, 100.0) + self.assertEqual(domain.sr_x, 4) + self.assertEqual(domain.sr_y, 4) + + def test_preserves_wps_mass_coordinates(self): + path = self._make_file("XLAT_M", "XLONG_M") + + domain = read_domain_from_file(path, sr_x=2, sr_y=2) + + self.assertEqual(domain.nx, 4) + self.assertEqual(domain.ny, 3) + self.assertEqual(domain.sr_x, 2) + self.assertEqual(domain.sr_y, 2) + + def test_uses_namelist_fallback_without_mass_coordinates(self): + path = self._make_file_without_coordinates() + namelist_params = { + "ref_lon": -107.0, + "ref_lat": 37.0, + "ref_x": 3.0, + "ref_y": 4.0, + "root_dx": 100.0, + "root_dy": 100.0, + "x_mass_offset_from_root": 200.0, + "y_mass_offset_from_root": 300.0, + "e_we": 5, + "e_sn": 6, + } + + domain = read_domain_from_file( + path, sr_x=4, sr_y=4, namelist_params=namelist_params + ) + + self.assertEqual(domain.nx, 4) + self.assertEqual(domain.ny, 5) + self.assertEqual(domain.sr_x, 4) + self.assertEqual(domain.sr_y, 4) + + +if __name__ == "__main__": + unittest.main()