diff --git a/components/elm/tools/landgen/README.md b/components/elm/tools/landgen/README.md
index 680ca729ec70..fa0c95559956 100644
--- a/components/elm/tools/landgen/README.md
+++ b/components/elm/tools/landgen/README.md
@@ -1 +1,239 @@
-# This is the readme file for landgen
\ No newline at end of file
+# landgen
+
+A Python package that preprocesses source land data to an unstructured grid for use by global models. Landgen has been developed to provide complete, self-consistent land data for `mksurfdata`, as part of the [E3SM Land Model (ELM)](https://e3sm.org) toolchain.
+
+## Table of Contents
+
+- [Version](#version)
+- [Overview](#overview)
+- [Requirements](#requirements)
+- [Installation](#installation)
+- [Configuration](#configuration)
+- [Usage](#usage)
+- [Outputs](#outputs)
+- [Visualizing output data](#visualizing)
+- [Contributing](#contributing)
+- [Authors](#authors)
+- [License](#license)
+
+---
+
+## Version
+
+Current development version 0.1.0, testing with an ~10km HEALPix grid
+
+---
+
+## Overview
+
+`landgen` processes raw source datasets (land type, soil, topography, human, atmosphere, etc.) and regrids them onto a target land grid (e.g., HEALPix ~1 km). Output NetCDF files are consumed by `mksurfdata` to generate ELM surface data files. Processing is parallelized via Python `multiprocessing` and is designed to run on HPC clusters with SLURM.
+
+---
+
+## Requirements
+
+- Linux/Unix (HPC cluster recommended)
+- [Conda](https://docs.conda.io/) (for environment management)
+- Python 3.x (managed via the conda environment)
+
+All Python dependencies are specified in the conda environment defined by `landgen_env.yml`.
+
+---
+
+## Installation
+
+### 1. Set up the conda environment
+
+Note that this needs to be done only upon first cloning the repository or after updating a local copy of the repository to ensure that all dependencies are included in the environment.
+
+Run the `load_landgen_env.sh` script from the `landgen` directory. This creates (or updates) the `landgen_env` conda environment and installs `landgen` in editable mode:
+
+```bash
+./load_landgen_env.sh
+```
+
+Or step by step:
+
+```bash
+module load conda
+conda env create -f landgen_env.yml # first time
+# conda env update -n landgen_env -f landgen_env.yml # subsequent updates
+conda activate landgen_env
+pip install -e .
+```
+
+The step by step version activates the environment in the current shell, but the environment does not need to be activated in any interactive shell because `landgen` cannot be run directly from the command line due to its high resource usage. Running `landgen` must be done via the `sbatch` SLURM command (see [Usage](#usage) below). Furthermore, code development does not require the environment to be active, only that the installation of `landgen` in editable mode has been included in the environment (via `pip install -e .`). This is why the `load_landgen_env.sh` script is not sourced in the first example.
+
+See `load_landgen_env.sh` for commands related to adding python packages to the environment and updating the environment and its associated `landgen_env.yml` definition file.
+
+### 2. Verify installation
+
+```bash
+conda activate landgen_env
+python -m landgen --help
+```
+
+---
+
+## Configuration
+
+Copy the provided input file template and edit it for your run:
+
+```bash
+cp config_template.json config.json
+```
+
+Key fields in `config.json`:
+
+| Field | Description |
+|---|---|
+| `start_year` | First year to process |
+| `end_year` | Last year to process (can be less than `start_year` to process backwards) |
+| `source_data_path` | Full root path to directory holding the raw input datasets |
+| `landgen_grid_path` | Path to the target land grid NetCDF file (including file name), relative to `source_data_path`. See [Grid definition](#grid) for more details. |
+| `out_path` | Full path to directory for output files |
+| `decomp_box_size_degrees` | Spatial decomposition box size in degrees |
+| `modules` | List of processing modules to run (see [Modules](#modules) for details) |
+
+### Grid definition
+
+
+
+### Modules
+
+| Module Name | Output file | Description |
+|---|---|---|
+| `topography` | `landgen_topography.nc` | Terrain elevation and related fields |
+| `land_type` | `landgen_land_type.nc` | Land cover, crop, urban, lake, ice, wetland, management, vegetation characteristics |
+| `soil` | `landgen_soil.nc` | Soil properties |
+| `human` | `landgen_human.nc` | Human datasets (e.g., gdp, population) |
+| `atmosphere` | `landgen_atmosphere.nc` | Atmospheric forcing-related land properties |
+
+Each module in the `config.json` `modules` list has a unique set of input parameter fields, with a few common ones required for each module. The common required fields are:
+
+| Required module input fields | Sub-fields | Description |
+|---|---|---|
+| `name` | none | Module name that matches the corresponding python file (e.g., entry "name": "topography" matches topography.py) |
+| `params` | `active`
`out_fname` | Set `"active": true` to enable a module.
`out_fname` is the generic name of the output netcdf file associated with this module. If a module outputs multiple years, there will be a file for each year with the year appended before the `.nc` extension
Module-specific fields are defined as sub-fields within `params` and can be single values or sets of values (or nested sets).|
+
+Additional module-specific input parameters (in `params`) are tabulated below for each module.
+
+| `topography` `params` | Sub-fields | Description |
+|---|---|---|
+| `TBD` | `TBD` | Terrain elevation and related fields |
+
+| `land_type` `params` | Sub-fields | Description |
+|---|---|---|
+| `sumbod_run` | `landcover`
`crop`
`urban`
`lake`
`ice`
`wetland`
`management`
`veg_char` | Set submodule name to `true` to enable it |
+| `sumbod_dyn` | `landcover`
`crop`
`urban`
`lake`
`ice`
`wetland`
`management`
`veg_char` | Set submodule name to `true` to enable multi-year processing |
+| `lc_rs_path` | none | Full path to directory holding source `landcover` data files |
+| `lc_rs_name` | none | Name of source `landcover` data. Used to determine how to process the land cover data. Currently, `modis` is the only supported value and the files are downloaded as needed and not stored in `lc_rs_path` because they are so large. |
+| `crop_path` | none | Full path to directory holding source `crop` data files |
+| `urban_path` | none | Full path to directory holding source `urban` data files |
+| `lake_path` | none | Full path to directory holding source `lake` data files |
+| `ice_path` | none | Full path to directory holding source `ice` data files |
+| `wetland_path` | none | Full path to directory holding source `wetland` data files |
+| `harvest_path` | none | Full path to directory holding source `management` harvest data files |
+| `harvest_name` | none | Name of harvest data file |
+| `grazing_path` | none | Full path to directory holding source `management` grazing data files |
+| `grazing_names` | `pasture`
`rangeland` | Names of grazing data files. |
+| `veg_char_path` | none | Full path to directory holding source `veg_char` data files |
+
+
+
+| `soil` `params` | Sub-fields | Description |
+|---|---|---|
+| `TBD` | `TBD` | Soil properties |
+
+| `human` `params` | Sub-fields | Description |
+|---|---|---|
+| `TBD` | `TBD` | Human datasets (e.g., gdp, population) |
+
+
+| `atmosphere` `params` | Sub-fields | Description |
+|---|---|---|
+| `TBD` | `TBD` | Atmospheric forcing-related land properties |
+
+
+---
+
+## Usage
+
+`landgen` cannot be run locally in interactive shell mode due to its high resource usage.
+
+
+### SLURM batch job (recommended for HPC)
+
+First ensure that the environment and code are up to date:
+```bash
+./load_landgen.sh
+```
+
+Edit `submit_landgen.sh` as needed, then from within the same directory where `submit_landgen.sh` is located submit:
+
+```bash
+sbatch submit_landgen.sh
+```
+
+`submit_landgen.sh` assumes that it has been called from the directory in which it is located in order to find the default `config.json` input file. If either of these is not the case (calling location or default input file), then the user must modify the `SCRIPT_DIR` definition and/or the `INPUT_FILE` definition lines in `submit_landgen.sh`.
+
+`submit_landgen.sh` uses a single node in exclusive mode and sets the number of worker processes to `SRUN_CPUS_PER_TASK` as defined in the script. Adjust `SRUN_CPUS_PER_TASK` in the script to change the number of cores used.
+
+All [outputs](#outputs) are written to `out_path` as defined in the input file (see [Configuration](#configuration))
+
+---
+
+## Outputs
+
+Each module has a specified netcdf output data file (see [Modules](#modules) for setting output file names). If a module output multiple years, each year will be written to its own file. The output files contain processed data on the grid specified by `landgen_grid_path`. These output data represent a self-consistent set of land variables for use by global models.
+
+Will need to create the list of variables for each module (see the confluence page for the mksurfdata api). Maybe do this via agent after all the ouput data structurs are developed.
+
+## Visualizing output data
+
+`plot_landgen.py` is both a library for use within `landgen` and a standalone script that can be called independently to plot data from the module output netcdf files. First load the environement and landgen code into the working shell:
+
+```bash
+source load_landgen_env.sh
+```
+
+Then run plot_landgen:
+
+```bash
+python -m landgen.plot_landgen [options]
+```
+
+| Argument | Description |
+|---|---|
+| `-h`, `--help` | Show help message and exit |
+| `file_path` | Path to the input landgen NetCDF file (including file name) |
+| `out_path` | Output directory for plot images |
+| `year` | Calendar year to plot |
+| `--varnames` | One or more variable names to plot; default plots all non-geometry variables |
+| `--layers` | Layer selector: single integer, JSON list (e.g. `[0,1,2]`), or JSON dict by variable (e.g. `{"pct_pft":[0,1]}`); default plots all layers |
+| `--plot-type` | `scatter` (colored points, default) or `rendered` (filled cell polygons) |
+| `--file-type` | Output image format; only `png` is currently supported |
+| `--colormap` | Matplotlib colormap name; default `viridis` |
+| `--ll-limits` | Spatial subset as `MIN_LAT MAX_LAT MIN_LON MAX_LON`; default uses full extent |
+| `--scale-limits` | Colorscale bounds as `VMIN VMAX`; default uses data min/max |
+
+## Contributing
+
+This package is developed as part of the E3SM project.
+
+Need to add specific contribution guide for landgen and mksurfdata. Also becauase this dev is separate from E3SM, the contribution guidelines do not necessarily apply.
+
+---
+
+## Authors
+
+- Alan Di Vittorio (avdivittorio@lbl.gov)
+- Eva Sinha (eva.sinha@pnnl.gov)
+
+---
+
+## License
+
+Actually the e3sm license will not explicitly apply when we put this in a separate repo. So need to set up a new one.
+
+See the E3SM [LICENSE](../../../../LICENSE) file. Need to point to or copy the license file.
diff --git a/components/elm/tools/landgen/config.json b/components/elm/tools/landgen/config_template.json
similarity index 61%
rename from components/elm/tools/landgen/config.json
rename to components/elm/tools/landgen/config_template.json
index e799b7cce5f5..8c683db2afaa 100644
--- a/components/elm/tools/landgen/config.json
+++ b/components/elm/tools/landgen/config_template.json
@@ -3,7 +3,9 @@
"end_year": 2015,
"source_data_path": "/global/cfs/cdirs/e3sm/landgen/rawdata",
"landgen_grid_path": "healpix/10km/land_domain_file_healpix_10km.nc",
- "out_path": "/global/cfs/cdirs/e3sm/landgen/output",
+ "ocean_shapefile_path": "ne_10m_ocean/ne_10m_ocean.shp",
+ "out_path": "/global/cfs/cdirs/e3sm/landgen/output_test",
+ "decomp_box_size_degrees": 10,
"modules": [
{
"name": "topography",
@@ -17,6 +19,26 @@
"params": {
"active": true,
"out_fname": "landgen_land_type.nc",
+ "submod_run": {
+ "landcover": true,
+ "crop": false,
+ "urban": false,
+ "lake": false,
+ "ice": false,
+ "wetland": false,
+ "management": false,
+ "veg_char": false
+ },
+ "submod_dyn": {
+ "landcover": true,
+ "crop": true,
+ "urban": false,
+ "lake": false,
+ "ice": false,
+ "wetland": false,
+ "management": true,
+ "veg_char": false
+ },
"lc_rs_path": "/global/cfs/cdirs/e3sm/landgen/rawdata/modis",
"lc_rs_name": "modis",
"crop_path": "",
@@ -24,14 +46,14 @@
"lake_path": "",
"ice_path": "",
"wetland_path": "",
- "harvest_path": "/global/cfs/cdirs/e3sm/landgen/rawdata/LUH2/LUH2_v2h",
- "harvest_name": "transitions.nc",
+ "management_path": "/global/cfs/cdirs/e3sm/landgen/rawdata/LUH2/LUH2_v2h",
+ "management_name": "transitions.nc",
"grazing_path": "/global/cfs/cdirs/e3sm/landgen/rawdata/HYDE3.5/original/gbc2025_7apr_base/NetCDF",
"grazing_names": {
"pasture": "pasture.nc",
"rangeland": "rangeland.nc"
},
- "assoc_path": ""
+ "veg_char_path": ""
}
},
{
diff --git a/components/elm/tools/landgen/generate_code_flowchart.py b/components/elm/tools/landgen/generate_code_flowchart.py
new file mode 100755
index 000000000000..246df41ec681
--- /dev/null
+++ b/components/elm/tools/landgen/generate_code_flowchart.py
@@ -0,0 +1,803 @@
+"""
+generate_flowchart.py
+Generates a PowerPoint (.pptx) flowchart of the landgen code structure.
+
+Run this script from the upper-level landgen directory (.../tools/landgen/)
+ and it will write landgen_flowchart.pptx in the same directory.
+ You can specify a custom output path as an optional argument.
+
+Usage:
+ python generate_flowchart.py # writes landgen_flowchart.pptx
+ python generate_flowchart.py # writes to a custom path
+
+Requires:
+ conda install python-pptx
+"""
+
+import sys
+import argparse
+import ast
+import json
+from pathlib import Path
+from pptx import Presentation
+from pptx.util import Inches, Pt, Emu
+from pptx.enum.text import PP_ALIGN
+from pptx.dml.color import RGBColor
+from pptx.enum.shapes import MSO_AUTO_SHAPE_TYPE
+from pptx.oxml import parse_xml
+from pptx.oxml.ns import qn, nsdecls
+from lxml import etree
+
+# ---------------------------------------------------------------------------
+# Colour palette
+# ---------------------------------------------------------------------------
+C_ENTRY = RGBColor(0x1F, 0x49, 0x7D) # dark blue – entry point
+C_CORE = RGBColor(0x2E, 0x75, 0xB6) # mid blue – core orchestration
+C_MODULE = RGBColor(0x70, 0xAD, 0x47) # green – top-level modules
+C_SUB = RGBColor(0xFF, 0xC0, 0x00) # amber – sub-steps
+C_IO = RGBColor(0xED, 0x7D, 0x31) # orange – I/O utilities
+C_DATA = RGBColor(0x7B, 0x0E, 0xA8) # purple – shared data structures
+C_PARALLEL = RGBColor(0xC0, 0x50, 0x4D) # red – parallel workers
+C_NOTIMPL = RGBColor(0xA5, 0xA5, 0xA5) # grey – not yet implemented
+C_WHITE = RGBColor(0xFF, 0xFF, 0xFF)
+C_BLACK = RGBColor(0x00, 0x00, 0x00)
+C_ARROW = RGBColor(0x40, 0x40, 0x40)
+
+SLIDE_W = Inches(20)
+SLIDE_H = Inches(14)
+
+# Font sizes – adjust these to make all text larger or smaller
+FS_TITLE = Pt(22) # slide title
+FS_LABEL = Pt(11) # section labels / IO header
+FS_BODY = Pt(10) # main box text
+FS_SMALL = Pt(9.5) # smaller boxes (substeps, workers, IO)
+FS_ARROW = Pt(9) # arrow labels
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+SRC_DIR = SCRIPT_DIR / "src" / "landgen"
+CONFIG_PATH = SCRIPT_DIR / "config.json"
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _rgb_to_hex(rgb: RGBColor) -> str:
+ return f"{rgb[0]:02X}{rgb[1]:02X}{rgb[2]:02X}"
+
+
+def add_box(slide, left, top, width, height,
+ text, fill_color, text_color=C_WHITE,
+ font_size=FS_BODY, bold=False, shape_type="rect"):
+ """Add a rounded-rectangle shape and return it."""
+
+ shape = slide.shapes.add_shape(
+ MSO_AUTO_SHAPE_TYPE.ROUNDED_RECTANGLE,
+ left, top, width, height
+ )
+
+ # solid fill via python-pptx API (avoids duplicate prstGeom XML corruption)
+ shape.fill.solid()
+ shape.fill.fore_color.rgb = fill_color
+
+ # thin white border
+ shape.line.color.rgb = C_WHITE
+ shape.line.width = Pt(0.75)
+
+ # text
+ tf = shape.text_frame
+ tf.word_wrap = True
+ p = tf.paragraphs[0]
+ p.alignment = PP_ALIGN.CENTER
+ run = p.add_run()
+ run.text = text
+ run.font.size = font_size
+ run.font.bold = bold
+ run.font.color.rgb = text_color
+
+ return shape
+
+
+def add_arrow(slide, x1, y1, x2, y2, color=C_ARROW, label=None):
+ """Draw a straight connector with an arrowhead at (x2,y2).
+
+ Builds the from a complete XML template so the element has
+ no block (which can trigger PowerPoint repair on blank
+ presentations) and no python-pptx connector-API side-effects.
+ """
+ # Bounding box of the line segment
+ left = int(min(x1, x2))
+ top = int(min(y1, y2))
+ cx = int(abs(x2 - x1)) or 1 # OOXML requires cx >= 1
+ cy = int(abs(y2 - y1)) or 1
+
+ # flip attributes encode which diagonal the line runs along
+ flip_attrs = ''
+ if x2 < x1:
+ flip_attrs += ' flipH="1"'
+ if y2 < y1:
+ flip_attrs += ' flipV="1"'
+
+ hex_color = _rgb_to_hex(color)
+ sp_tree = slide.shapes._spTree
+ shape_id = sp_tree._next_shape_id
+
+ xml = (
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ f''
+ )
+
+ cxnSp = parse_xml(xml)
+ sp_tree.append(cxnSp)
+
+ if label:
+ mx = (x1 + x2) / 2
+ my = (y1 + y2) / 2
+ tb = slide.shapes.add_textbox(mx, my - Inches(0.12), Inches(0.9), Inches(0.2))
+ tf = tb.text_frame
+ p = tf.paragraphs[0]
+ run = p.add_run()
+ run.text = label
+ run.font.size = FS_ARROW
+ run.font.color.rgb = C_ARROW
+
+ return cxnSp
+
+
+def _read_ast(py_path: Path):
+ """Parse a Python file and return its AST, or None if unavailable."""
+ try:
+ return ast.parse(py_path.read_text(), filename=str(py_path))
+ except Exception:
+ return None
+
+
+def _module_path(module_name: str) -> Path:
+ return SRC_DIR / f"{module_name}.py"
+
+
+def _module_exists(module_name: str) -> bool:
+ return _module_path(module_name).exists()
+
+
+def _top_level_functions(py_path: Path):
+ """Return top-level function names from a Python file."""
+ tree = _read_ast(py_path)
+ if tree is None:
+ return []
+ return [n.name for n in tree.body if isinstance(n, ast.FunctionDef)]
+
+
+def _discover_top_modules():
+ """Read top-level module names from config.json modules list."""
+ if CONFIG_PATH.exists():
+ try:
+ cfg = json.loads(CONFIG_PATH.read_text())
+ names = [m.get("name") for m in cfg.get("modules", []) if m.get("name")]
+ if names:
+ return names
+ except Exception:
+ pass
+
+ # fallback if config is unavailable
+ fallback = []
+ for p in sorted(SRC_DIR.glob("*.py")):
+ stem = p.stem
+ if stem.startswith("_"):
+ continue
+ if stem in {"__init__", "__main__", "tools", "shared_data", "landgen_io", "plot_landgen"}:
+ continue
+ fallback.append(stem)
+ return fallback
+
+
+def _discover_land_type_submodules():
+ """Parse land_type.py for importlib.import_module('landgen.') calls."""
+ land_type_path = _module_path("land_type")
+ tree = _read_ast(land_type_path)
+ if tree is None:
+ return []
+
+ names = []
+ for node in ast.walk(tree):
+ if not isinstance(node, ast.Call):
+ continue
+ fn = node.func
+ if not (isinstance(fn, ast.Attribute) and fn.attr == "import_module"):
+ continue
+ if not node.args:
+ continue
+ arg0 = node.args[0]
+ if isinstance(arg0, ast.Constant) and isinstance(arg0.value, str):
+ mod = arg0.value
+ if mod.startswith("landgen."):
+ short = mod.split(".")[-1]
+ if short not in names:
+ names.append(short)
+ return names
+
+
+def _module_has_parallel_work(module_name: str) -> bool:
+ """Heuristic: module uses mp.Pool or defines *_process worker functions."""
+ py_path = _module_path(module_name)
+ tree = _read_ast(py_path)
+ if tree is None:
+ return False
+
+ for node in ast.walk(tree):
+ if isinstance(node, ast.FunctionDef) and node.name.endswith("_process"):
+ return True
+ if isinstance(node, ast.Call):
+ fn = node.func
+ if isinstance(fn, ast.Attribute) and fn.attr == "Pool":
+ return True
+ return False
+
+
+def _discover_landcover_deps():
+ """Return relative-imported module names from landcover.py."""
+ landcover_path = _module_path("landcover")
+ tree = _read_ast(landcover_path)
+ if tree is None:
+ return []
+
+ deps = []
+ for node in ast.walk(tree):
+ if isinstance(node, ast.ImportFrom) and node.level >= 1:
+ for alias in node.names:
+ name = alias.name
+ if name not in deps and _module_exists(name):
+ deps.append(name)
+ return deps
+
+
+def _discover_io_items():
+ """Build IO row items from parsed source functions."""
+ items = []
+
+ # Show remote-sensing helper when present.
+ rs_path = _module_path("landcover_remote_sensing")
+ if rs_path.exists():
+ rs_funcs = [f for f in _top_level_functions(rs_path) if not f.startswith("_")]
+ if rs_funcs:
+ items.append(
+ "landcover_remote_\nsensing.py\n" + " / ".join(f"{f}()" for f in rs_funcs[:2])
+ )
+
+ io_path = _module_path("landgen_io")
+ io_funcs = [
+ f for f in _top_level_functions(io_path)
+ if f.startswith(("read_", "write_", "regrid_", "load_", "set_decomp_"))
+ ]
+ for f in io_funcs[:6]:
+ items.append(f"{f}()")
+
+ return items
+
+
+def _brief_doc(node):
+ """Return a short single-line summary from a function docstring."""
+ doc = ast.get_docstring(node)
+ if not doc:
+ return "No docstring"
+ first = doc.strip().splitlines()[0].strip()
+ return first[:110]
+
+
+def _function_summaries(module_name: str):
+ """Return [(func_name, brief_doc), ...] for top-level functions in a module."""
+ py_path = _module_path(module_name)
+ tree = _read_ast(py_path)
+ if tree is None:
+ return []
+
+ out = []
+ for node in tree.body:
+ if isinstance(node, ast.FunctionDef) and not node.name.startswith("_"):
+ out.append((node.name, _brief_doc(node)))
+ return out
+
+
+def _discover_landgen_import_aliases(module_name: str, tree):
+ """Map local import aliases to landgen module names for a module AST."""
+ alias_to_module = {}
+
+ for node in tree.body:
+ if isinstance(node, ast.ImportFrom):
+ # from . import landgen_io as io
+ if node.level >= 1 and node.module is None:
+ for alias in node.names:
+ target = alias.name
+ if _module_exists(target):
+ alias_to_module[alias.asname or alias.name] = target
+ # from .landcover_remote_sensing import foo (function/class import)
+ # skip these for module-call matching.
+
+ elif isinstance(node, ast.Import):
+ # import landgen.landcover_remote_sensing as lc_rs
+ for alias in node.names:
+ name = alias.name
+ if name.startswith("landgen."):
+ target = name.split(".")[-1]
+ if _module_exists(target):
+ alias_to_module[alias.asname or target] = target
+
+ # Never treat self-module calls as cross-module.
+ alias_to_module = {k: v for k, v in alias_to_module.items() if v != module_name}
+ return alias_to_module
+
+
+def _discover_process_calls(module_name: str):
+ """Return all calls from *_process to functions in other landgen modules."""
+ py_path = _module_path(module_name)
+ tree = _read_ast(py_path)
+ if tree is None:
+ return []
+
+ alias_to_module = _discover_landgen_import_aliases(module_name, tree)
+ module_funcs_cache = {
+ mod: set(_top_level_functions(_module_path(mod)))
+ for mod in set(alias_to_module.values())
+ }
+
+ process_funcs = [
+ n for n in tree.body
+ if isinstance(n, ast.FunctionDef) and n.name.endswith("_process")
+ ]
+ if not process_funcs:
+ return []
+
+ calls = []
+ for fn_node in process_funcs:
+ for node in ast.walk(fn_node):
+ if not isinstance(node, ast.Call):
+ continue
+ # Only include calls like .() where alias is an imported
+ # landgen module and func is defined in that target module.
+ if not isinstance(node.func, ast.Attribute):
+ continue
+ if not isinstance(node.func.value, ast.Name):
+ continue
+ base = node.func.value.id
+ func = node.func.attr
+ if base not in alias_to_module:
+ continue
+ if func.startswith("_"):
+ continue
+
+ target_module = alias_to_module[base]
+ target_funcs = module_funcs_cache.get(target_module, set())
+ if func not in target_funcs:
+ continue
+
+ name = f"{base}.{func}"
+ if name not in calls:
+ calls.append(name)
+
+ return calls
+
+
+def build_functions_slide(prs, module_name: str, title_color=C_IO):
+ """Add a slide listing top-level functions and brief descriptions."""
+ slide_layout = prs.slide_layouts[6] # blank
+ slide = prs.slides.add_slide(slide_layout)
+
+ func_fs = Pt(18)
+
+ # title
+ tb = slide.shapes.add_textbox(Inches(0.2), Inches(0.1), Inches(19.6), Inches(0.5))
+ p = tb.text_frame.paragraphs[0]
+ p.alignment = PP_ALIGN.CENTER
+ run = p.add_run()
+ run.text = f"{module_name}.py – Function Summary"
+ run.font.size = Pt(26)
+ run.font.bold = True
+ run.font.color.rgb = title_color
+
+ # subtitle
+ st = slide.shapes.add_textbox(Inches(0.3), Inches(0.65), Inches(19.2), Inches(0.35))
+ sp = st.text_frame.paragraphs[0]
+ sr = sp.add_run()
+ sr.text = f"Auto-parsed from src/landgen/{module_name}.py"
+ sr.font.size = Pt(18)
+ sr.font.color.rgb = C_BLACK
+
+ funcs = _function_summaries(module_name)
+ if not funcs:
+ b = add_box(slide, Inches(0.5), Inches(1.3), Inches(19.0), Inches(0.8),
+ f"No top-level functions found in {module_name}.py", C_NOTIMPL,
+ text_color=C_WHITE, font_size=FS_BODY)
+ return slide
+
+ lines = [f"{name}() – {desc}" for name, desc in funcs]
+
+ # Two-column list layout
+ mid = (len(lines) + 1) // 2
+ cols = [lines[:mid], lines[mid:]]
+ lefts = [Inches(0.4), Inches(10.2)]
+ top = Inches(1.1)
+ width = Inches(9.4)
+ height = Inches(12.4)
+
+ for i, col_lines in enumerate(cols):
+ box = slide.shapes.add_textbox(lefts[i], top, width, height)
+ tf = box.text_frame
+ tf.word_wrap = True
+ tf.clear()
+ for j, line in enumerate(col_lines):
+ p = tf.paragraphs[0] if j == 0 else tf.add_paragraph()
+ p.text = line
+ p.level = 0
+ for r in p.runs:
+ r.font.size = func_fs
+
+ return slide
+
+
+def build_overview_index_slide(prs):
+ """Add an index slide summarizing discovered structure and where details live."""
+ slide_layout = prs.slide_layouts[6] # blank
+ slide = prs.slides.add_slide(slide_layout)
+
+ # title
+ tb = slide.shapes.add_textbox(Inches(0.2), Inches(0.1), Inches(19.6), Inches(0.5))
+ p = tb.text_frame.paragraphs[0]
+ p.alignment = PP_ALIGN.CENTER
+ run = p.add_run()
+ run.text = "landgen – Overview Index"
+ run.font.size = FS_TITLE
+ run.font.bold = True
+ run.font.color.rgb = C_ENTRY
+
+ top_modules = _discover_top_modules()
+ submods = _discover_land_type_submodules()
+ parallel = [m for m in submods if _module_exists(m) and _module_has_parallel_work(m)]
+
+ io_funcs = [name for name, _ in _function_summaries("landgen_io")][:12]
+ tools_funcs = [name for name, _ in _function_summaries("tools")][:10]
+ plot_funcs = [name for name, _ in _function_summaries("plot_landgen")][:10]
+
+ sections = [
+ (
+ "Execution Flow",
+ "Entry: __main__.py -> landgen.main()\n"
+ "Top-level modules: " + (", ".join(top_modules) if top_modules else "none"),
+ C_CORE,
+ ),
+ (
+ "land_type Submodules",
+ "_process_single_year imports: " + (", ".join(submods) if submods else "none") + "\n"
+ "Parallel workers: " + (", ".join(parallel) if parallel else "none"),
+ C_SUB,
+ ),
+ (
+ "landgen_io.py (Slide 3)",
+ "Functions: " + (", ".join(io_funcs) if io_funcs else "none"),
+ C_IO,
+ ),
+ (
+ "tools.py (Slide 4) and plot_landgen.py (Slide 5)",
+ "tools.py: " + (", ".join(tools_funcs) if tools_funcs else "none") + "\n"
+ "plot_landgen.py: " + (", ".join(plot_funcs) if plot_funcs else "none"),
+ C_ENTRY,
+ ),
+ ]
+
+ y = Inches(0.9)
+ h = Inches(2.85)
+ for title, body, color in sections:
+ add_box(
+ slide,
+ Inches(0.5),
+ y,
+ Inches(19.0),
+ h,
+ f"{title}\n{body}",
+ color,
+ text_color=C_WHITE,
+ font_size=FS_SMALL,
+ )
+ y += h + Inches(0.2)
+
+ return slide
+
+
+# ---------------------------------------------------------------------------
+# Layout constants (all in Inches, converted via Inches())
+# ---------------------------------------------------------------------------
+
+def build_slide(prs):
+ slide_layout = prs.slide_layouts[6] # blank
+ slide = prs.slides.add_slide(slide_layout)
+
+ slide_fs = Pt(18)
+
+ # ---- title ----
+ tb = slide.shapes.add_textbox(Inches(0.2), Inches(0.1),
+ Inches(19.6), Inches(0.5))
+ tf = tb.text_frame
+ p = tf.paragraphs[0]
+ p.alignment = PP_ALIGN.CENTER
+ run = p.add_run()
+ run.text = "landgen – Code Structure Flowchart"
+ run.font.size = Pt(26)
+ run.font.bold = True
+ run.font.color.rgb = C_ENTRY
+
+ # ---- legend ----
+ legend_items = [
+ (C_ENTRY, "Entry point"),
+ (C_CORE, "Core orchestration"),
+ (C_MODULE, "Top-level module"),
+ (C_SUB, "Sub-step (sequential)"),
+ (C_PARALLEL, "Parallel worker"),
+ (C_IO, "I/O utility"),
+ (C_DATA, "Shared data structure"),
+ (C_NOTIMPL, "Not yet implemented"),
+ ]
+ lx = Inches(0.2)
+ ly = Inches(0.7)
+ for i, (col, lbl) in enumerate(legend_items):
+ add_box(slide, lx + Inches(i * 2.4), ly, Inches(2.2), Inches(0.55),
+ lbl, col, font_size=slide_fs)
+
+ # =========================================================
+ # ROW 1 – Entry point
+ # =========================================================
+ r1y = Inches(1.5)
+ bh = Inches(0.95) # box height
+ bw_sm = Inches(3.1)
+ bw_med = Inches(4.3)
+ bw_lg = Inches(3.6)
+
+ # __main__.py
+ main_x = Inches(7.9)
+ b_main = add_box(slide, main_x, r1y, bw_med, bh,
+ "__main__.py\npython -m landgen ",
+ C_ENTRY, font_size=slide_fs, bold=True)
+
+ # config.json (to the right)
+ cfg_x = Inches(13.3)
+ b_cfg = add_box(slide, cfg_x, r1y, bw_sm, bh,
+ "config.json\n(start/end year, paths, modules)",
+ C_IO, font_size=slide_fs)
+
+ add_arrow(slide,
+ cfg_x, r1y + bh / 2,
+ main_x + bw_med, r1y + bh / 2)
+
+ # =========================================================
+ # ROW 2 – landgen.main()
+ # =========================================================
+ r2y = r1y + bh + Inches(0.65)
+ core_x = Inches(7.3)
+ bw_core = Inches(5.2)
+ bh_core = Inches(1.6)
+ b_core = add_box(slide, core_x, r2y, bw_core, bh_core,
+ "landgen.py · main(config_path)\n"
+ "Load config · start mp.Manager · create GridData\n"
+ "Loop modules → importlib.import_module(name).run()",
+ C_CORE, font_size=slide_fs, bold=True)
+
+ # arrow __main__ → main
+ add_arrow(slide,
+ main_x + bw_med / 2, r1y + bh,
+ core_x + bw_core / 2, r2y)
+
+ # =========================================================
+ # ROW 3 – Shared data (right side)
+ # =========================================================
+ r3y = r2y
+ sd_x = Inches(13.0)
+ bw_sd = Inches(6.6)
+ b_sd = add_box(slide, sd_x, r3y, bw_sd, Inches(2.15),
+ "shared_data.py – Multiprocessing shared structures\n"
+ "GridData / GridManager (cell ids, lon/lat, landfrac)\n"
+ "TopoData / TopoManager (elevation, slope, fmax …)\n"
+ "LtData / LtManager (pct_pft, harvest_frac, …)",
+ C_DATA, font_size=slide_fs)
+
+ add_arrow(slide,
+ core_x + bw_core, r2y + bh_core / 2,
+ sd_x, r3y + Inches(1.075))
+
+ # =========================================================
+ # ROW 4 – Top-level modules (5 boxes across)
+ # =========================================================
+ r4y = r3y + Inches(2.35) + Inches(0.75)
+ top_modules = _discover_top_modules()
+ modules = []
+ for name in top_modules:
+ col = C_MODULE if _module_exists(name) else C_NOTIMPL
+ funcs = _top_level_functions(_module_path(name)) if _module_exists(name) else []
+ text = f"{name}.py\nrun()"
+ if name == "topography":
+ text += "\n[must run first;\nsets landfrac]"
+ elif name == "land_type":
+ text += "\n[main land-type\norchestrator]"
+ elif funcs:
+ extras = [f for f in funcs if f not in {"run"} and not f.startswith("_")]
+ if extras:
+ text += "\n" + ", ".join(extras[:2])
+ modules.append((text, col))
+
+ if not modules:
+ modules = [("No modules\nfound", C_NOTIMPL)]
+ mod_bw = Inches(3.6)
+ mod_gap = Inches(0.35)
+ total_mw = len(modules) * mod_bw + (len(modules) - 1) * mod_gap
+ mod_x0 = (SLIDE_W - total_mw) / 2
+ mod_bh = Inches(1.25)
+
+ mod_boxes = []
+ for i, (txt, col) in enumerate(modules):
+ bx = mod_x0 + i * (mod_bw + mod_gap)
+ b = add_box(slide, bx, r4y, mod_bw, mod_bh, txt, col,
+ font_size=slide_fs)
+ mod_boxes.append((bx, b))
+ # Branch arrows should originate from the bottom edge of landgen.main box.
+ add_arrow(slide,
+ core_x + bw_core / 2, r2y + bh_core,
+ bx + mod_bw / 2, r4y)
+
+ # =========================================================
+ # ROW 5 – land_type internals: _process_single_year loop
+ # =========================================================
+ r5y = r4y + mod_bh + Inches(0.8)
+ land_type_idx = 0
+ for i, name in enumerate(top_modules):
+ if name == "land_type":
+ land_type_idx = i
+ break
+ lt_x = mod_x0 + land_type_idx * (mod_bw + mod_gap)
+ bw_lt = Inches(3.4)
+ bh_lt = Inches(0.9)
+ b_loop = add_box(slide, lt_x, r5y, bw_lt, bh_lt,
+ "Loop years (start_year → end_year)\n"
+ "_process_single_year(year, …)",
+ C_CORE, font_size=slide_fs)
+ add_arrow(slide,
+ lt_x + bw_lt / 2, r4y + mod_bh,
+ lt_x + bw_lt / 2, r5y)
+
+ # =========================================================
+ # ROW 6 – Sub-steps inside _process_single_year
+ # =========================================================
+ r6y = r5y + bh_lt + Inches(0.8)
+ submod_names = _discover_land_type_submodules()
+ substeps = []
+ for name in submod_names:
+ exists = _module_exists(name)
+ parallel = _module_has_parallel_work(name) if exists else False
+ col = C_SUB if exists else C_NOTIMPL
+ txt = f"{name}.py\nrun()"
+ if parallel:
+ txt += " [parallel]"
+ if not exists:
+ txt += "\n[not yet impl.]"
+ substeps.append((txt, col, name, parallel, exists))
+
+ if not substeps:
+ substeps = [("No sub-steps\nfound", C_NOTIMPL, "", False, False)]
+
+ # Single-row substep layout: compact boxes so all modules fit on one row.
+ ss_cols = max(1, len(substeps))
+ ss_bw = Inches(2.05)
+ ss_bh = Inches(1.0)
+ ss_gap_x = Inches(0.12)
+ ss_gap_y = Inches(0.0)
+ row_w = ss_cols * ss_bw + (ss_cols - 1) * ss_gap_x
+ ss_x0 = (SLIDE_W - row_w) / 2
+
+ substep_centers = {}
+
+ for i, (txt, col, name, parallel, exists) in enumerate(substeps):
+ row = i // ss_cols
+ col_i = i % ss_cols
+ bx = ss_x0 + col_i * (ss_bw + ss_gap_x)
+ by = r6y + row * (ss_bh + ss_gap_y)
+ add_box(slide, bx, by, ss_bw, ss_bh, txt, col, font_size=slide_fs)
+ substep_centers[name] = (bx + ss_bw / 2, by + ss_bh)
+ add_arrow(slide,
+ lt_x + bw_lt / 2, r5y + bh_lt,
+ bx + ss_bw / 2, by)
+
+ # =========================================================
+ # ROW 7 – Parallel worker call detail
+ # =========================================================
+ n_sub_rows = (len(substeps) + ss_cols - 1) // ss_cols
+ r7y = r6y + n_sub_rows * ss_bh + max(0, n_sub_rows - 1) * ss_gap_y + Inches(0.55)
+
+ worker_modules = [name for _, _, name, parallel, exists in substeps if parallel and exists]
+ if worker_modules:
+ w_cols = min(3, len(worker_modules))
+ w_bw = Inches(5.9)
+ # Auto-size worker boxes so all discovered call lines fit.
+ shown_worker_modules = worker_modules[:3]
+ max_lines = 1
+ for mod_name in shown_worker_modules:
+ calls = _discover_process_calls(mod_name)
+ # Account for visual wrapping of long call names.
+ wrapped_call_lines = sum(max(1, (len(c) + 33) // 34) for c in calls) if calls else 1
+ n_lines = 1 + wrapped_call_lines
+ if n_lines > max_lines:
+ max_lines = n_lines
+ worker_line_h = 0.30
+ worker_pad_h = 0.62
+ w_bh = Inches(max(2.0, worker_pad_h + max_lines * worker_line_h))
+ w_gap = Inches(0.35)
+ total_ww = w_cols * w_bw + (w_cols - 1) * w_gap
+ w_x0 = (SLIDE_W - total_ww) / 2
+
+ for i, mod_name in enumerate(shown_worker_modules):
+ wx = w_x0 + i * (w_bw + w_gap)
+ proc_name = f"{mod_name}_process"
+ calls = _discover_process_calls(mod_name)
+ call_lines = "\n".join([f"· {c}()" for c in calls]) if calls else "· parallel chunk work"
+ txt = f"{proc_name}()\n{call_lines}"
+
+ add_box(slide, wx, r7y, w_bw, w_bh, txt, C_PARALLEL, font_size=slide_fs)
+
+ if mod_name in substep_centers:
+ sx, sy = substep_centers[mod_name]
+ add_arrow(slide, sx, sy, wx + w_bw / 2, r7y)
+
+ return slide
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+ parser = argparse.ArgumentParser(
+ prog="generate_code_flowchart.py",
+ description=(
+ "Generate a PowerPoint (.pptx) flowchart of the landgen code structure."
+ ),
+ epilog=(
+ "Examples:\n"
+ " python generate_code_flowchart.py\n"
+ " python generate_code_flowchart.py custom_flowchart.pptx"
+ ),
+ formatter_class=argparse.RawTextHelpFormatter,
+ )
+ parser.add_argument(
+ "out_path",
+ nargs="?",
+ default="landgen_flowchart.pptx",
+ help="Output .pptx file path (default: landgen_flowchart.pptx)",
+ )
+ args = parser.parse_args()
+
+ out_path = Path(args.out_path)
+
+ prs = Presentation()
+ prs.slide_width = SLIDE_W
+ prs.slide_height = SLIDE_H
+
+ build_slide(prs)
+ build_functions_slide(prs, "landgen_io", title_color=C_IO)
+ build_functions_slide(prs, "tools", title_color=C_CORE)
+ build_functions_slide(prs, "plot_landgen", title_color=C_ENTRY)
+
+ prs.save(out_path)
+ print(f"Saved: {out_path}")
+
+
+if __name__ == "__main__":
+ main()
+
diff --git a/components/elm/tools/landgen/landgen_env.yml b/components/elm/tools/landgen/landgen_env.yml
index 7f124061f320..9764ea340321 100644
--- a/components/elm/tools/landgen/landgen_env.yml
+++ b/components/elm/tools/landgen/landgen_env.yml
@@ -5,7 +5,13 @@ dependencies:
- _openmp_mutex=4.5=20_gnu
- _python_abi3_support=1.0=hd8ed1ab_2
- affine=2.4.0=pyhd8ed1ab_1
+ - aiobotocore=3.7.0=pyhcf101f3_0
+ - aiohttp=3.13.5=pyhf64b827_0
+ - aioitertools=0.13.0=pyhd8ed1ab_0
+ - aiosignal=1.4.0=pyhd8ed1ab_0
- alsa-lib=1.2.15.3=hb03c661_0
+ - async-timeout=5.0.1=pyhcf101f3_2
+ - attr=2.5.2=hb03c661_1
- attrs=26.1.0=pyhcf101f3_0
- aws-c-auth=0.10.1=h2d2dd48_2
- aws-c-cal=0.9.13=h2c9d079_1
@@ -25,15 +31,20 @@ dependencies:
- azure-storage-blobs-cpp=12.16.0=hdd73cc9_1
- azure-storage-common-cpp=12.12.0=ha7a2c86_1
- azure-storage-files-datalake-cpp=12.14.0=h52c5a47_1
+ - backports.zstd=1.5.0=py314h680f03e_0
- blosc=1.21.6=he440d0b_1
+ - botocore=1.43.0=pyhd8ed1ab_0
+ - bounded-pool-executor=0.0.3=pyhd8ed1ab_0
- brotli=1.2.0=hed03a55_1
- brotli-bin=1.2.0=hb03c661_1
+ - brotli-python=1.2.0=py314h3de4e8d_1
- bzip2=1.0.8=hda65f42_9
- c-ares=1.34.6=hb03c661_0
- - ca-certificates=2026.4.22=hbd8a1cb_0
+ - ca-certificates=2026.5.20=hbd8a1cb_0
- cairo=1.18.4=he90730b_1
- - certifi=2026.4.22=pyhd8ed1ab_0
+ - certifi=2026.5.20=pyhd8ed1ab_0
- cftime=1.6.5=py314hc02f841_1
+ - charset-normalizer=3.4.7=pyhd8ed1ab_0
- click=8.3.2=pyhc90fa1f_0
- click-plugins=1.1.1.2=pyhd8ed1ab_0
- cligj=0.7.2=pyhd8ed1ab_2
@@ -44,6 +55,7 @@ dependencies:
- cyrus-sasl=2.1.28=hac629b4_1
- dbus=1.16.2=h24cb091_1
- double-conversion=3.4.0=hecca717_0
+ - earthaccess=0.18.0=pyhc364b38_0
- fastparquet=2025.12.0=py314hc02f841_0
- font-ttf-dejavu-sans-mono=2.37=hab24e00_0
- font-ttf-inconsolata=3.000=h77eed37_0
@@ -62,11 +74,18 @@ dependencies:
- giflib=5.2.2=hd590300_0
- glog=0.7.1=hbabe93e_0
- graphite2=1.3.14=hecca717_2
+ - h2=4.3.0=pyhcf101f3_0
- harfbuzz=14.2.0=h6083320_0
- hdf4=4.2.15=h2a13503_7
- hdf5=2.1.0=nompi_hd4fcb43_104
+ - hpack=4.1.0=pyhd8ed1ab_0
+ - hyperframe=6.1.0=pyhd8ed1ab_0
- icu=78.3=h33c6efd_0
+ - idna=3.13=pyhcf101f3_0
- importlib-metadata=8.8.0=pyhcf101f3_0
+ - importlib-resources=7.1.0=pyhd8ed1ab_0
+ - importlib_resources=7.1.0=pyhd8ed1ab_0
+ - jmespath=1.1.0=pyhcf101f3_1
- json-c=0.18=h6688a6e_0
- keyutils=1.6.3=hb9d3cd8_0
- kiwisolver=1.5.0=py314h97ea11e_0
@@ -82,6 +101,7 @@ dependencies:
- libarrow-compute=23.0.1=h53684a4_9_cpu
- libarrow-dataset=23.0.1=h635bf11_9_cpu
- libarrow-substrait=23.0.1=hb4dd7c2_9_cpu
+ - libattr=2.5.2=hb03c661_1
- libblas=3.11.0=6_h4a7cf45_openblas
- libbrotlicommon=1.2.0=hb03c661_1
- libbrotlidec=1.2.0=hb03c661_1
@@ -105,6 +125,7 @@ dependencies:
- libgcc=15.2.0=he0feb66_18
- libgcc-ng=15.2.0=h69a702a_18
- libgdal-core=3.12.3=he63569f_3
+ - libgdal-hdf4=3.12.3=hf70aa56_3
- libgfortran=15.2.0=h69a702a_18
- libgfortran5=15.2.0=h68bc16d_18
- libgl=1.7.0=ha4b6fd6_2
@@ -161,11 +182,14 @@ dependencies:
- libxslt=1.1.43=h711ed8c_1
- libzip=1.11.2=h6991a6a_0
- libzlib=1.3.2=h25fd6f3_2
+ - lxml=6.1.0=py314hae3bed6_0
- lz4-c=1.10.0=h5888daf_1
- lzo=2.10=h280c20c_1002
- matplotlib=3.10.8=py314hdafbbf9_0
- matplotlib-base=3.10.8=py314h1194b4b_0
- minizip=4.0.10=h05a5f5f_0
+ - multidict=6.7.1=pyh62beb40_0
+ - multimethod=2.0.2=pyhd8ed1ab_0
- munkres=1.1.4=pyhd8ed1ab_1
- muparser=2.3.5=h5888daf_0
- ncurses=6.5=h2d0b736_3
@@ -182,6 +206,7 @@ dependencies:
- pillow=12.2.0=py314h8ec4b1a_0
- pip=26.0.1=pyh145f28c_0
- pixman=0.46.4=h54a6638_1
+ - pqdm=0.2.0=pyhd8ed1ab_1
- proj=9.7.1=he0df7b0_3
- prometheus-cpp=1.3.0=ha5d0236_0
- pthread-stubs=0.4=hb9d3cd8_1002
@@ -190,9 +215,12 @@ dependencies:
- pyparsing=3.3.2=pyhcf101f3_0
- pyproj=3.7.2=py314h68799e9_3
- pyside6=6.11.0=py314h3987850_2
+ - pysocks=1.7.1=pyha55dd90_7
- python=3.14.4=habeac84_100_cp314
+ - python-cmr=0.13.0=pyhff2d567_1
- python-dateutil=2.9.0.post0=pyhe01879c_2
- python-gil=3.14.4=h4df99d1_100
+ - python-pptx=1.0.2=pyhcf101f3_1
- python_abi=3.14=8_cp314
- qhull=2020.2=h434a139_5
- qt6-main=6.11.0=pl5321h16c4a6b_4
@@ -200,6 +228,7 @@ dependencies:
- re2=2025.11.05=h5301d42_1
- readline=8.3=h853b02a_0
- s2n=1.7.1=h1cbb8d7_1
+ - s3fs=2026.3.0=pyhd8ed1ab_0
- scipy=1.17.1=py314hf07bd8e_0
- setuptools=82.0.1=pyh332efcf_0
- shapely=2.1.2=py314hbe3edd8_2
@@ -207,12 +236,18 @@ dependencies:
- snappy=1.2.2=h03e3b7b_1
- snuggs=1.4.7=pyhd8ed1ab_2
- sqlite=3.53.0=h04a0ce9_0
+ - tenacity=9.1.4=pyhcf101f3_0
+ - tinynetrc=1.3.1=pyhd8ed1ab_0
- tk=8.6.13=noxft_h366c992_103
- tornado=6.5.5=py314h5bd0f2a_0
+ - tqdm=4.67.3=pyh8f84b5b_0
+ - typing-extensions=4.15.0=h396c80c_0
+ - typing_extensions=4.15.0=pyhcf101f3_0
- tzdata=2025c=hc9c84f9_1
- unicodedata2=17.0.1=py314h5bd0f2a_0
- uriparser=0.9.8=hac33072_0
- wayland=1.25.0=hd6090a7_0
+ - wrapt=2.2.1=py314h5bd0f2a_0
- xarray=2026.4.0=pyhc364b38_0
- xcb-util=0.4.1=h4f16b4b_2
- xcb-util-cursor=0.1.6=hb03c661_0
@@ -222,6 +257,7 @@ dependencies:
- xcb-util-wm=0.4.2=hb711507_0
- xerces-c=3.3.0=hd9031aa_1
- xkeyboard-config=2.47=hb03c661_0
+ - xlsxwriter=3.2.9=pyhd8ed1ab_0
- xorg-libice=1.1.2=hb9d3cd8_0
- xorg-libsm=1.2.6=he73a12e_0
- xorg-libx11=1.8.13=he1eb515_0
@@ -244,11 +280,8 @@ dependencies:
- zstd=1.5.7=hb78ec9c_6
- pip:
- aiohappyeyeballs==2.6.1
- - aiohttp==3.13.5
- - aiosignal==1.4.0
- asttokens==3.0.1
- cartopy==0.25.0
- - charset-normalizer==3.4.7
- click-default-group==1.2.4
- cmocean==4.0.3
- comm==0.2.3
@@ -260,21 +293,18 @@ dependencies:
- frozenlist==1.8.0
- geographiclib==2.1
- geovista==0.5.3
- - idna==3.13
- imageio==2.37.3
- ipython==9.12.0
- ipython-pygments-lexers==1.1.1
- ipywidgets==8.1.8
- jedi==0.19.2
- jupyterlab-widgets==3.0.16
- - landgen==0.1.0
- lazy-loader==0.5
- markdown-it-py==4.0.0
- matplotlib-inline==0.2.1
- mdurl==0.1.2
- more-itertools==11.0.2
- msgpack==1.1.2
- - multidict==6.7.1
- parso==0.8.6
- patsy==1.0.2
- pexpect==4.9.0
@@ -299,14 +329,12 @@ dependencies:
- scooby==0.11.1
- stack-data==0.6.3
- statsmodels==0.14.6
- - tqdm==4.67.3
- traitlets==5.14.3
- trame==3.12.0
- trame-client==3.11.4
- trame-common==1.1.3
- trame-server==3.10.0
- trame-vuetify==3.2.1
- - typing-extensions==4.15.0
- uraster==0.1.7
- urllib3==2.6.3
- vtk==9.6.1
@@ -314,4 +342,4 @@ dependencies:
- widgetsnbextension==4.0.15
- wslink==2.5.6
- yarl==1.23.0
-prefix: /global/homes/e/esinha/.conda/envs/landgen_env
+prefix: /global/homes/a/avdivitt/.conda/envs/landgen_env
diff --git a/components/elm/tools/landgen/landgen_flowchart.pptx b/components/elm/tools/landgen/landgen_flowchart.pptx
new file mode 100644
index 000000000000..1e47b69675e1
Binary files /dev/null and b/components/elm/tools/landgen/landgen_flowchart.pptx differ
diff --git a/components/elm/tools/landgen/load_landgen_env.sh b/components/elm/tools/landgen/load_landgen_env.sh
index 3c9187e5c315..e7b763f873ba 100755
--- a/components/elm/tools/landgen/load_landgen_env.sh
+++ b/components/elm/tools/landgen/load_landgen_env.sh
@@ -4,15 +4,32 @@
# first load the conda module
module load conda
-# first time only!
-# create a the landgen_env conda environment from the .yml file
-#conda env create -f landgen_env.yml
-
-# to update the enviroment
-conda env update -n landgen_env -f landgen_env.yml
+# create the landgen_env conda environment from the .yml file if it does not exist,
+# otherwise update the existing environment
+if conda env list | grep -q "^landgen_env "; then
+ conda env update -n landgen_env -f landgen_env.yml
+else
+ conda env create -f landgen_env.yml
+fi
# activate the landgen_env conda environment
conda activate landgen_env
+# install the landgen package in editable mode (run from the landgen directory)
+# this is separate from conda env update so that the env can be updated without pip trying to load landgen
+pip install -e .
+
+########## the following commands are to be run as needed, not in this script ##########
+
# install a new python package
-#conda install
\ No newline at end of file
+#conda install
+
+# remove a python package
+#conda remove
+
+# to update the .yml file with new packages added to the environment
+# (landgen is removed automatically since it is not on PyPI and is installed locally via pip install -e .)
+#conda env export -n landgen_env | grep -v "landgen==" > landgen_env.yml
+
+# to update the enviroment with a new .yml file (e.g. after pulling changes from git)
+#conda env update -n landgen_env -f landgen_env.yml
\ No newline at end of file
diff --git a/components/elm/tools/landgen/src/landgen/__main__.py b/components/elm/tools/landgen/src/landgen/__main__.py
index daaeb6be2794..2842d3a96bae 100644
--- a/components/elm/tools/landgen/src/landgen/__main__.py
+++ b/components/elm/tools/landgen/src/landgen/__main__.py
@@ -2,10 +2,41 @@
from landgen import landgen
import sys
-from pathlib import Path
+import multiprocessing as mp
-if len(sys.argv) != 2:
- print(f"Usage: python {Path(__file__).name} ")
- sys.exit(1)
-print('Executing as standalone script')
-landgen.main(sys.argv[1])
\ No newline at end of file
+HELP = """\
+Usage: python -m landgen
+
+Preprocess source land data to a target grid for use by mksurfdata.
+
+Arguments:
+ path/config_name.json Path to and name of the JSON configuration file.
+ Do not need to specify the full path if the config file is in the current working directory.
+
+Options:
+ -h, --help Show this message and exit.
+
+Base example call:
+ python -m landgen config.json
+
+This package is designed to be run on a cluster with SLURM, using a single node and multiple cores.
+It can also be run on a local machine with multiple cores, but the user must adapt the load and submit scripts accordingly.
+It is not designed for distributed execution across multiple nodes.
+
+To execute this package using SLURM:
+ source load_landgen_env.sh # create/update the landgen environment and install the current landgen package in editable mode
+ sbatch submit_landgen.sh # submit the SLURM job with the appropriate configuration file specified in submit_landgen.sh
+
+"""
+
+if __name__ == '__main__':
+ # use forkserver start method for multiprocessing to avoid issues with fork on some platforms
+ mp.set_start_method('forkserver')
+ if len(sys.argv) == 2 and sys.argv[1] in ('-h', '--help'):
+ print(HELP)
+ sys.exit(0)
+ if len(sys.argv) != 2:
+ print(f"Usage: python -m landgen (use --help for more info)")
+ sys.exit(1)
+ print('Executing as standalone script')
+ landgen.main(sys.argv[1])
diff --git a/components/elm/tools/landgen/src/landgen/atmosphere.py b/components/elm/tools/landgen/src/landgen/atmosphere.py
index 33e0193dcb06..1bbae26823ca 100644
--- a/components/elm/tools/landgen/src/landgen/atmosphere.py
+++ b/components/elm/tools/landgen/src/landgen/atmosphere.py
@@ -8,8 +8,11 @@
import multiprocessing as mp
import importlib
+import logging
from pathlib import Path
+logger = logging.getLogger('landgen')
+
########## define helper functions for land_type run() here
@@ -33,9 +36,9 @@
## output
-def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager):
+def run(active, out_fname, com_config_dict, out_grid_data, manager, decomp_indices, decomp_ll_limits):
if active is False:
- print(f"Skipping atmosphere module")
+ logger.info("Skipping atmosphere module")
return
# extract common parameters from shared config dict
@@ -46,7 +49,7 @@ def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager
out_path = com_config_dict['out_path']
landfrac = out_grid_data.get_landfrac()
- print(f"Processing atmosphere module with parameters:")
+ logger.info("Processing atmosphere module")
# todo: print the parameters here
# processing code for atmosphere
diff --git a/components/elm/tools/landgen/src/landgen/harvest.py b/components/elm/tools/landgen/src/landgen/harvest.py
deleted file mode 100644
index fe0bbd16c24b..000000000000
--- a/components/elm/tools/landgen/src/landgen/harvest.py
+++ /dev/null
@@ -1,159 +0,0 @@
-# harvest.py
-# this module processes harvest and grazing data for the landgen workflow
-# the output is a complete harvest distribution
-# and includes some data associated with particular harvests
-
-# run() function is the main entry point for this module, and will be called by process_single_year in land_type.py
-
-import multiprocessing as mp
-from pathlib import Path
-from . import shared_data
-from . import landgen_io
-# import normalize_cell # not created yet
-import pandas as pd
-import os
-
-########## define helper functions for harvest run() here
-
-##### harvest_process()
-
-## arguments
-# lc_data: land cover data structure that is passed between modules
-# year: the year for which to process the land cover data
-# source_data_path: base path to the source data
-# landgen_grid_path: path from source_data_path and the filename of the landgen grid
-# out_path: base path for the output data; this is needed to read in the previous year's harvest data
-
-## output
-
-def harvest_process(lt_year_data, year, harvest_path, harvest_name, grazing_path, grazing_names,
- com_config_dict, out_grid_data, ll_limits, cell_ids,
- global_mesh_df, man_lock, grid_lock, lt_lock):
-
- print(f"Processing harvest and grazing module year {year} with parameters:")
-
- # each worker writes its temp files to a unique subdirectory to avoid collisions
- min_lat, max_lat, min_lon, max_lon = ll_limits
- tmp_dir = (
- Path(com_config_dict['out_path'])
- / 'tmp'
- / f"harvest_{year}_{min_lat:.0f}_{min_lon:.0f}"
- )
-
- # --- read source data (full global arrays; slicing happens inside regrid) ---
- harvest_data = landgen_io.read_luh2_harvest(year, harvest_path, harvest_name)
- grazing_data = landgen_io.read_hyde_grazing(year, grazing_path, grazing_names)
-
- # --- regrid and store harvest variables into lt_year_data.harvest_frac ---
- # LUH2_HARVEST_VARS order matches the n_harvest=5 dimension in LtData:
- # index 0: primf_harv, 1: primn_harv, 2: secmf_harv, 3: secyf_harv, 4: secnf_harv
- for i, varname in enumerate(landgen_io.LUH2_HARVEST_VARS):
- regridded = landgen_io.regrid_to_landgen_grid(
- harvest_data[varname],
- harvest_data['lat'],
- harvest_data['lon'],
- cell_ids, ll_limits,
- global_mesh_df,
- tmp_dir / varname,
- varname,
- )
- with lt_lock:
- lt_year_data.harvest_frac[cell_ids, i] = regridded
-
- # --- regrid and store grazing variables into lt_year_data.grazing_frac ---
- # grazing_names dict order matches n_grazing=2 dimension in LtData:
- # index 0: pasture (grazing_land), index 1: rangeland
- for i, category in enumerate(grazing_names.keys()):
- regridded = landgen_io.regrid_to_landgen_grid(
- grazing_data[category],
- grazing_data['lat'],
- grazing_data['lon'],
- cell_ids, ll_limits,
- global_mesh_df,
- tmp_dir / category,
- category,
- )
- with lt_lock:
- lt_year_data.grazing_frac[cell_ids, i] = regridded
-
- return
-
-
-########## run()
-
-## called by land_type.process_single_year() for each year, and this is where the multiprocessing happens for the landcover module
-## this sets up the pool and calls the harvest_process() function for each chunk of data
-
-def run(lt_year_data, year, prev_year, harvest_path, harvest_name, grazing_path, grazing_names,
- com_config_dict, out_grid_data, manager, grid_manager, lt_manager):
-
- print(f"Processing harvest module with parameters:")
- # todo: print the parameters here
-
- # load the global HEALPix mesh parquet once here so worker processes
- # don't each re-read the 37 MB file; pass the DataFrame into each chunk tuple
- global_parquet_path = (
- Path(com_config_dict['source_data_path'])
- / Path(com_config_dict['landgen_grid_path']).parent
- / 'merged_land_cells.parquet'
- )
- global_mesh_df = pd.read_parquet(global_parquet_path)
- print(f" Loaded HEALPix mesh: {len(global_mesh_df)} cells from {global_parquet_path}")
-
- # number of available cpu cores (set by SBATCH during job submission)
- omp_threads_str = os.environ.get('OMP_NUM_THREADS')
-
- if omp_threads_str is not None:
- try:
- # Convert the string value to an integer
- omp_threads_int = int(omp_threads_str)
- print(f"OMP_NUM_THREADS is set to: {omp_threads_int}")
- except ValueError:
- print(f"OMP_NUM_THREADS is set to an invalid integer value: {omp_threads_str}")
- else:
- print("OMP_NUM_THREADS environment variable is not set.")
- # If not set, set to total cores on the node
- omp_threads_int = mp.cpu_count()
- print(f"Using total cores: {omp_threads_int}, but this may fail if "
- "SBATCH --cpus-per-task is set to a lower number or SBATCH --exclusive is not set")
-
- # set up the pool and call the harvest_process() function for each chunk of data
- # chunks are defined by the lat-lon limits and corresponding landgen grid cell ids for the chunk;
- # these are created in land_type.process_single_year() and passed to this run() function as lists?
- # there are more chunks than cpus; the pool will manage this for efficiency because chunks vary in size
- # the results will be stored directly in the lt_year_data shared structure
-
- # get the manager locks for the shared data structures
- # all locks come from the main mp.Manager() (SyncManager); custom managers don't support Lock()
- man_lock = manager.Lock()
- grid_lock = manager.Lock()
- lt_lock = manager.Lock()
-
- # Build data_chunks: one tuple per spatial chunk covering the globe in 10x10 degree boxes.
- # For each chunk, filter the global mesh to cells whose centroid lat/lon falls within the box.
- from . import land_type as _lt
- decomp_box_size_degrees = 10
- chunk_ll_limits = _lt.calc_ll_limits(decomp_box_size_degrees)
-
- data_chunks = []
- for ll in chunk_ll_limits:
- min_lat, max_lat, min_lon, max_lon = ll
- mask = (
- (global_mesh_df['lat'] >= min_lat) & (global_mesh_df['lat'] < max_lat) &
- (global_mesh_df['lon'] >= min_lon) & (global_mesh_df['lon'] < max_lon)
- )
- chunk_cell_ids = global_mesh_df.loc[mask, 'cellid'].values
- if len(chunk_cell_ids) == 0:
- continue # skip ocean-only or empty chunks
- data_chunks.append((
- lt_year_data, year, harvest_path, harvest_name, grazing_path, grazing_names,
- com_config_dict, out_grid_data, ll, chunk_cell_ids,
- global_mesh_df, man_lock, grid_lock, lt_lock,
- ))
-
- print(f" Submitting {len(data_chunks)} harvest chunks to pool of {omp_threads_int} workers")
-
- with mp.Pool(processes=omp_threads_int) as pool:
- pool.starmap(harvest_process, data_chunks)
-
- return
\ No newline at end of file
diff --git a/components/elm/tools/landgen/src/landgen/human.py b/components/elm/tools/landgen/src/landgen/human.py
index ec427bf10fed..227889562e0b 100644
--- a/components/elm/tools/landgen/src/landgen/human.py
+++ b/components/elm/tools/landgen/src/landgen/human.py
@@ -8,8 +8,11 @@
import multiprocessing as mp
import importlib
+import logging
from pathlib import Path
+logger = logging.getLogger('landgen')
+
########## define helper functions for land_type run() here
@@ -33,9 +36,9 @@
## output
-def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager):
+def run(active, out_fname, com_config_dict, out_grid_data, manager, decomp_indices, decomp_ll_limits):
if active is False:
- print(f"Skipping human module")
+ logger.info("Skipping human module")
return
# extract common parameters from shared config dict
@@ -46,7 +49,7 @@ def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager
out_path = com_config_dict['out_path']
landfrac = out_grid_data.get_landfrac()
- print(f"Processing human module with parameters:")
+ logger.info("Processing human module")
# todo: print the parameters here
# processing code for human
diff --git a/components/elm/tools/landgen/src/landgen/land_type.py b/components/elm/tools/landgen/src/landgen/land_type.py
index 2b5959ea54da..98d6de200e35 100644
--- a/components/elm/tools/landgen/src/landgen/land_type.py
+++ b/components/elm/tools/landgen/src/landgen/land_type.py
@@ -5,46 +5,24 @@
# run() function is the main entry point for this module, and will be called by landgen.py
-import multiprocessing as mp
import importlib
+import logging
from pathlib import Path
-from . import shared_data
-from .shared_data import LtData, LtManager
+from .shared_data import LtData
import numpy as np
+from . import landgen_io
+from . import tools
+
+logger = logging.getLogger('landgen')
########## define helper functions for land_type here
-def calc_ll_limits(size_degrees):
- """Calculate and return a list of tuples (min_lat, max_lat, min_lon, max_lon)
- for each size_degrees x size_degrees chunk covering the globe.
-
- Latitude spans -90 to 90 degrees.
- Longitude spans -180 to 180 degrees.
-
- Returns:
- list of tuples: [(min_lat, max_lat, min_lon, max_lon), ...]
- """
- ll_limits = []
- if 90 % size_degrees != 0 or 180 % size_degrees != 0:
- raise ValueError(
- f"size_degrees ({size_degrees}) must evenly divide "
- f"both 90 (latitude half-range) and 180 (longitude half-range)."
- )
- lat = -90.0
- while lat < 90.0:
- max_lat = min(lat + size_degrees, 90.0)
- lon = -180.0
- while lon < 180.0:
- max_lon = min(lon + size_degrees, 180.0)
- ll_limits.append((lat, max_lat, lon, max_lon))
- lon = max_lon
- lat = max_lat
- return ll_limits
-
-##### process_single_year()
-def _process_single_year(lt_year_data, year, prev_year, out_fname, lc_rs_path, lc_rs_name, crop_path, urban_path,
- lake_path, ice_path, wetland_path, harvest_path, harvest_name, grazing_path, grazing_names, assoc_path, com_config_dict, out_grid_data,
- manager, grid_manager, lt_manager):
+
+##### _process_single_year()
+def _process_single_year(lt_year_data, year, prev_year, submod_run, submod_dyn, out_fname, lc_rs_path, lc_rs_name, crop_path, urban_path,
+ lake_path, ice_path, wetland_path, harvest_path, harvest_name, grazing_path, grazing_names,
+ veg_char_path, com_config_dict, out_grid_data,
+ manager, decomp_indices, decomp_ll_limits):
"""Process land type data for a single year."""
# arguments
@@ -54,95 +32,81 @@ def _process_single_year(lt_year_data, year, prev_year, out_fname, lc_rs_path, l
# other arguments are described below for the run() function
- # data chunks are based on 10x10 degree lat-lon boxes (648 chunks)
- # 15x15 degree box gives 288 chunks, 30x30 box gives 72 chunks
- # todo: add this decomp box size to the config file and to com_config_dict
- decomp_box_size_degrees = 10
- # get the lat-lon limits for the landgen grid
- # ll_limits = list(float) of [(min_lat, max_lat, min_lon, max_lon),(min_lat, max_lat, min_lon, max_lon),... for each chunk]
- chunk_ll_limits = calc_ll_limits(decomp_box_size_degrees)
-
-
- # todo: set the multiprocessing chunk info here for all run functions
- # need the lat/lon limits for the chunk, and the corresponding landgen grid cell ids for the chunk
- # do this by lat-lon because all source raw data are on lat-lon grids, but at different resolutions
- # and since the source data are several and varied,
- # it will be faster to read in just the corresponding landgren grid cell info,
- # even though it is non-sequential
- # the chunks won't be equal size on the landgen grid, but will be consistent,
- # and will be varied in size based on input source raw data res
- # so put the lat/lon limits and cell ids in corresponding chunks and use the process queue or pool for efficiency
- # arguments for each run function below will include data chunk list with the lat/lon limits and cell ids for the chunk
- # or the entire list is created here
- # the chunked data are a list of tuples with each argument;
- # an example is that the static arguments here will be repeated in each tuple,
- # and the lat/lon and cell ids will be different for each chunk
- # this info will have to be passed to the run functions below
-
- # ll_limits: the full list of (min_lat, max_lat, min_lon, max_lon) tuples for all chunks
- ll_limits = chunk_ll_limits
- # cell_ids: all land cell ids from the shared grid data structure
- cell_ids = out_grid_data.get_cell_id()
-
+## todo: need to figure out how to use static data that has already been processed for the first year
+# maybe: if prev_year is not None and a submodule is static (submod_dyn==false) then use data from previous year file
+## todo change the order below, such that static data are processed first
# Process landcover
- # derive prev_fname from out_fname and prev_year by inserting the year before the file extension
- # e.g. landgen_land_type.nc -> landgen_land_type_2009.nc
- if prev_year is not None:
- stem, suffix = out_fname.rsplit('.', 1)
- prev_fname = f"{stem}_{prev_year}.{suffix}"
- else:
- prev_fname = None
- # each module's run function calls the multiple processes because these modules need to be done sequentially
- landcover = importlib.import_module('landgen.landcover')
- lc_data = landcover.run(lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
- com_config_dict, out_grid_data, ll_limits, cell_ids, manager, grid_manager, lt_manager)
-
- # Process crop data - adjust lc crop area
- crop = importlib.import_module('landgen.crop')
- lc_data = crop.run(lt_year_data, year, prev_year, crop_path, com_config_dict, out_grid_data, ll_limits, cell_ids,
- manager, grid_manager, lt_manager)
-
- # Process urban data - adjust lc urban area
- urban = importlib.import_module('landgen.urban')
- lc_data = urban.run(lt_year_data, year, prev_year, urban_path, com_config_dict, out_grid_data, ll_limits, cell_ids,
- manager, grid_manager, lt_manager)
-
- # Process lake data - adjust lc lake area
- lake = importlib.import_module('landgen.lake')
- lc_data = lake.run(lt_year_data, year, prev_year, lake_path, com_config_dict, out_grid_data, ll_limits, cell_ids,
- manager, grid_manager, lt_manager)
-
- # Process ice data - adjust lc ice area
- ice = importlib.import_module('landgen.ice')
- lc_data = ice.run(lt_year_data, year, prev_year, ice_path, com_config_dict, out_grid_data, ll_limits, cell_ids,
- manager, grid_manager, lt_manager)
-
- # Process wetland data - adjust lc wetland area
- # (may not be needed as the main source is currently the modis cover data;
- # can allow for this in the future)
- #wetland = importlib.import_module('wetland')
- #lc_data = wetland.run(lt_year_data, year, prev_year, wetland_path, com_config_dict, out_grid_data, ll_limits, cell_ids, manager, grid_manager, lt_manager)
-
- # Process harvest/grazing data - adjust harvest/grazing area
- harvest = importlib.import_module('landgen.harvest')
- lc_data = harvest.run(lt_year_data, year, prev_year, harvest_path, harvest_name, grazing_path, grazing_names,
- com_config_dict, out_grid_data, manager, grid_manager, lt_manager)
+ if submod_run['landcover']:
+ # derive prev_fname from out_fname and prev_year by inserting the year before the file extension
+ # e.g. landgen_land_type.nc -> landgen_land_type_2009.nc
+ if prev_year is not None:
+ stem, suffix = out_fname.rsplit('.', 1)
+ prev_fname = f"{stem}_{prev_year}.{suffix}"
+ else:
+ prev_fname = None
+ # each module's run function calls the multiple processes because these modules need to be done sequentially
+ landcover = importlib.import_module('landgen.landcover')
+ landcover.run(lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
+ com_config_dict, out_grid_data, decomp_indices, decomp_ll_limits, manager)
+
+ if submod_run['crop']:
+ # Process crop data - adjust lc crop area
+ crop = importlib.import_module('landgen.crop')
+ lc_data = crop.run(lt_year_data, year, prev_year, crop_path, com_config_dict, out_grid_data,
+ decomp_indices, decomp_ll_limits, manager)
+
+ if submod_run['urban']:
+ # Process urban data - adjust lc urban area
+ urban = importlib.import_module('landgen.urban')
+ lc_data = urban.run(lt_year_data, year, prev_year, urban_path, com_config_dict, out_grid_data,
+ decomp_indices, decomp_ll_limits, manager)
+
+ if submod_run['lake']:
+ # Process lake data - adjust lc lake area
+ lake = importlib.import_module('landgen.lake')
+ lc_data = lake.run(lt_year_data, year, prev_year, lake_path, com_config_dict, out_grid_data,
+ decomp_indices, decomp_ll_limits, manager)
+
+ if submod_run['ice']:
+ # Process ice data - adjust lc ice area
+ ice = importlib.import_module('landgen.ice')
+ lc_data = ice.run(lt_year_data, year, prev_year, ice_path, com_config_dict, out_grid_data,
+ decomp_indices, decomp_ll_limits, manager)
+
+ #if submod_run['wetland']:
+ # Process wetland data - adjust lc wetland area
+ # (may not be needed as the main source is currently the modis cover data;
+ # can allow for this in the future)
+ #wetland = importlib.import_module('landgen.wetland')
+ #lc_data = wetland.run(lt_year_data, year, prev_year, wetland_path, com_config_dict, out_grid_data,
+ # decomp_indices, decomp_ll_limits, manager)
+ pass
+
+ if submod_run['management']:
+ #todo: update this with more efficient decomp and generalized reading and chunking
+ # Process harvest/grazing data - adjust harvest/grazing area
+ management = importlib.import_module('landgen.management')
+ lc_data = management.run(lt_year_data, year, prev_year, harvest_path, harvest_name, grazing_path,
+ grazing_names, com_config_dict, out_grid_data, decomp_indices, decomp_ll_limits)
# Normalize cell
- normalize_cell = importlib.import_module('landgen.normalize_cell')
- lc_data = normalize_cell.fill_land(lt_year_data, out_grid_data, ll_limits, cell_ids, manager, grid_manager, lt_manager) # fill_land
- lc_data = normalize_cell.reconcile_ocean(lt_year_data, out_grid_data, ll_limits, cell_ids, manager, grid_manager, lt_manager) # reconcile_ocean
-
- # Process veg-associated data
- veg_assoc = importlib.import_module('landgen.veg_assoc')
- lc_data = veg_assoc.run(lt_year_data, year, prev_year, assoc_path, com_config_dict, out_grid_data, ll_limits, cell_ids,
- manager, grid_manager, lt_manager)
+ #normalize_cell = importlib.import_module('landgen.normalize_cell')
+ #lc_data = normalize_cell.fill_land(lt_year_data, out_grid_data, decomp_indices, decomp_ll_limits,
+ # manager) # fill_land
+ #lc_data = normalize_cell.reconcile_ocean(lt_year_data, out_grid_data, decomp_indices, decomp_ll_limits,
+ # manager) # reconcile_ocean
+
+ if submod_run['veg_char']:
+ # Process veg-associated data
+ veg_char = importlib.import_module('landgen.veg_char')
+ lc_data = veg_char.run(lt_year_data, year, prev_year, veg_char_path, com_config_dict, out_grid_data,
+ decomp_indices, decomp_ll_limits, manager)
# Ensure consistency
- consistency = importlib.import_module('landgen.consistency')
- lc_data = consistency.run(lt_year_data, year, out_grid_data, ll_limits, cell_ids, manager, grid_manager, lt_manager)
+ #consistency = importlib.import_module('landgen.consistency')
+ #lc_data = consistency.run(lt_year_data, year, out_grid_data, decomp_indices, decomp_ll_limits, manager)
return
@@ -155,23 +119,36 @@ def _process_single_year(lt_year_data, year, prev_year, out_fname, lc_rs_path, l
# the rest of the params set in the config file
# com_config_dict: the shared dictionary for the common parameters for all modules
# out_grid_data: the shared data structure for the landgen grid data
+# manager: the multiprocessing manager for the shared data structures
+# grid_manager: the multiprocessing manager for the shared data structure for the landgen grid data
+# decomp_indices: the list of cell index chunks for parallel processing
+# decomp_ll_limits: the list of lat/lon limits for each chunk for parallel processing
+
+# Note that chunks are not equal in size
## output
-def run(active, out_fname, lc_rs_path, lc_rs_name, crop_path, urban_path, lake_path, ice_path,
- wetland_path, harvest_path, harvest_name, grazing_path, grazing_names, assoc_path, com_config_dict, out_grid_data, manager, grid_manager):
+def run(active, submod_run, submod_dyn, out_fname, lc_rs_path, lc_rs_name, crop_path, urban_path, lake_path, ice_path,
+ wetland_path, harvest_path, harvest_name, grazing_path, grazing_names, veg_char_path,
+ com_config_dict, out_grid_data, manager, decomp_indices, decomp_ll_limits):
if active is False:
- print(f"Skipping land_type module")
+ logger.info("Skipping land_type module")
return
# set up the land_type module shared data structure
# this holds only one year of data, so write it each year
- lt_manager = LtManager()
- lt_manager.start()
- lt_year_data = lt_manager.LtData()
- lt_year_data.allocate()
-
- print(f"Processing land_type module with parameters:")
+ lt_year_data = LtData()
+ # get the actual number of land cells from out_grid_data
+ n_cells = out_grid_data.num_cells
+ print(f" Allocating LtData for {n_cells} land cells")
+ lt_year_data.allocate(n_cells=n_cells)
+
+ #lt_manager = LtManager()
+ #lt_manager.start()
+ #lt_year_data = lt_manager.LtData()
+ #lt_year_data.allocate()
+
+ logger.info("Processing land_type module")
# todo: print the parameters here
# extract common parameters from shared config dict
@@ -181,30 +158,60 @@ def run(active, out_fname, lc_rs_path, lc_rs_name, crop_path, urban_path, lake_p
# processing code for land_type
years = np.arange(start_year, end_year + 1)
- output_file = Path(out_path) / out_fname
+ #output_file = Path(out_path) / out_fname
prev_year = None
# 1. Loop over desired years
for year in years:
# 2. Process single year
- print(f" Processing year: {year}")
- _process_single_year(lt_year_data, year, prev_year, out_fname, lc_rs_path, lc_rs_name, crop_path, urban_path,
- lake_path, ice_path, wetland_path, harvest_path, harvest_name, grazing_path, grazing_names, assoc_path, com_config_dict, out_grid_data,
- manager, grid_manager, lt_manager)
-
-
-
- # append this year's data to the output file
- # todo: can we prepend data is going backwards in time, so that the output file is in chronological order?
- # otherwise need to write it after all years are processed
+ logger.info(f"Processing year: {year}")
+ _process_single_year(lt_year_data, year, prev_year, submod_run, submod_dyn, out_fname,
+ lc_rs_path, lc_rs_name, crop_path, urban_path, lake_path, ice_path,
+ wetland_path, harvest_path, harvest_name, grazing_path, grazing_names,
+ veg_char_path, com_config_dict, out_grid_data, manager,
+ decomp_indices, decomp_ll_limits)
+
+ # no - would have to read in while file to reverse the order - append this year's data to the output file
+ # these data may need to be appended chunk by chunk if memory is an issue, but try writing the whole year at once first
+ # todo: can write each year, then combine at end in proper order
+
+ # set timevars in shared_data for each data class
+ # Variables to write to output NetCDF:
+ # pct_pft: landcover percentages [n_cells, n_pfts]
+ # pct_ocean: ocean percentage [n_cells]
+ # harvest_frac: harvest fractions from LUH2 [n_cells, n_harvest=10]
+ # grazing_frac: grazing fractions from HYDE3.5 [n_cells, n_grazing=2]
+ # Variables with time dimension (for annual concatenation with ncrcat):
+ # All of the above vary by year
+ varnames = ['pct_pft', 'pct_ocean', 'harvest_frac', 'grazing_frac']
+ timevars = ['pct_pft', 'harvest_frac', 'grazing_frac']
+
+ # insert _ before the extension (or at the end if no extension)
+ out_fname_p = Path(out_fname)
+ out_fname_year = f"{out_fname_p.stem}_{year}{out_fname_p.suffix}"
+
+ landgen_io.write_module_netcdf(out_grid_data, lt_year_data, out_path, out_fname_year,
+ year=year, timevars=timevars, varnames=varnames, ll_limits=None)
prev_year = year
+ # todo: combine the annual files into one file in the correct time order
+ # can use xarray.open_mfdataset(sorted_files) or ncrcat
+ # actually, write individual year files
+
+ ## todo: this is temporary for testing? or maybe not?
+ # just plot the start year for now
+ #plot_fname_year = f"{out_fname_year.stem}_{start_year}{out_fname_year.suffix}"
+ ncdf_path = Path(out_path) / out_fname_year
+ print_layers = [0, 1]
+ tools.plot_module_netcdf(ncdf_path, out_path, start_year, varnames=varnames, layers=print_layers,
+ plot_type='scatter', file_type='png',
+ colormap='viridis', ll_limits=None)
# free the module-specific shared data structure
lt_year_data = None
- lt_manager.shutdown()
+ #lt_manager.shutdown()
return
diff --git a/components/elm/tools/landgen/src/landgen/landcover.py b/components/elm/tools/landgen/src/landgen/landcover.py
index 22e857dd4b06..ea16a8fcb5fe 100644
--- a/components/elm/tools/landgen/src/landgen/landcover.py
+++ b/components/elm/tools/landgen/src/landgen/landcover.py
@@ -6,14 +6,26 @@
# run() function is the main entry point for this module, and will be called by process_single_year in land_type.py
import multiprocessing as mp
-#import importlib
+import logging
+import shutil
+import tempfile
+from datetime import datetime
from pathlib import Path
from . import shared_data
-import landcover_remote_sensing # not created yet
-import transitions # not created yet
-import normalize_cell # not created yet
-import pandas as pd
+from . import landgen_io
+from . import tools
+from . import landcover_remote_sensing as lc_rs
+#from . import transitions # not created yet
+#from . import normalize_cell # not created yet
import os
+import numpy as np
+
+logger = logging.getLogger('landgen')
+resource_logger = logging.getLogger('ClusterMonitor')
+
+def _landcover_process_star(args):
+ """Module-level wrapper so imap_unordered can unpack the args tuple."""
+ return landcover_process(*args)
########## define helper functions for landcover run() here
@@ -31,79 +43,170 @@
## output
-def landcover_process(lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
- com_config_dict, out_grid_data, ll_limits, cell_ids,
- man_lock, grid_lock, lt_lock):
-
- print(f"Processing landcover module year {year} with parameters:")
- # todo: print the parameters here
-
- # todo: use the with man_lock:, with grid_lock:, with lt_lock: syntax for accessing each managed data structure
-
- # todo: probably need to add lai path to land_type params and pass it through to here
-
- lc_rs_data = None
- prev_lt_data = None
- climate_data = None
- lai_data = None
- elm_data = None
-
- # todo: these data below will be read in based on ll_limits
-
- if lc_rs.use_lc_rs(year):
- # read modis cover data
- # reading both the igbp cover data and the veg continuous fields data
- lc_rs_data = lc_rs.read(year, com_config_dict['source_data_path'], lc_rs_path, lc_rs_name)
- else:
- # read and process previous year's landgen land type data and transitions to calculate this year's land cover distribution
- if prev_year is not None:
- # read previous year's landgen land type data
- prev_out_file = Path(com_config_dict['out_path']) / prev_fname
- if prev_out_file.exists():
- print(f"Reading previous year landgen land type data from {prev_out_file}")
- # todo: define this in a helper function
- prev_lt_data = read_prev_lt(prev_year, prev_out_file)
- else:
- print(f"Error: Previous year output file {prev_out_file} does not exist; cannot read previous year landgen land type data.")
- sys.exit(1)
-
- # Calculate this year's land cover distribution using the previous year's data and the transitions
- # these calculations are based on landgen land type outputs
- temp_lt_data = transitions.run(prev_lt_data, year, prev_year)
- # convert this year's land cover distribution to the lc rs classes
- lc_rs_data = lc_rs.convert_landgen_to_lc_rs(temp_lt_data, lc_rs_name)
+def landcover_process(year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
+ com_config_dict, out_grid_data, cell_indices, ll_limits
+ ):
+
+ # todo: need to sort out printing from multiple processes
+ #print(f"Processing landcover module year {year} with parameters:")
+ # todo: print the parameters here?
+
+ # Determine a scratch base directory for worker-local temporary files.
+ # Priority: $SCRATCH (NERSC/HPC) -> $TMPDIR -> system default (usually /tmp).
+ # Each worker gets its own subdirectory to avoid cross-worker file collisions.
+ scratch_base = os.environ.get('SCRATCH') or os.environ.get('TMPDIR') or tempfile.gettempdir()
+ tmp_dir = Path(tempfile.mkdtemp(dir=scratch_base, prefix=f'landcover_{year}_'))
+
+ # Re-attach logging in this worker process. In forkserver mode each worker
+ # starts as a clean process with no logging handlers configured.
+ log_path = com_config_dict.get('log_path')
+ if log_path:
+ tools.init_worker_logging(log_path)
+
+ # Redirect uraster's auto-created log files (utility.log, uraster.log, etc.)
+ # from the process CWD into the run output directory.
+ out_path = com_config_dict.get('out_path')
+ if out_path:
+ tools.redirect_uraster_logs(out_path)
+
+ worker_id = f"{mp.current_process().name} (pid {os.getpid()})"
+ logger.info(f"[landcover_process] START {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} worker={worker_id} year={year} chunk_size={len(cell_indices)}")
+
+ try:
+
+ # first write the mesh file for this chunk
+ # first write the mesh file for this chunk
+ mesh_file = Path(tmp_dir) / 'mesh.geojson'
+ landgen_io.write_mesh_to_geojson(out_grid_data, mesh_file, cell_indices)
+
+ # create a local data structure just for this chunk; only allocate the
+ # fields this module actually writes to minimise per-chunk memory usage.
+ # copy_from() silently skips None fields, so a partial allocation is safe.
+ lt_chunk_data = shared_data.LtData()
+ n = len(cell_indices)
+ lt_chunk_data.cell_idx = np.array(cell_indices, dtype=np.int64)
+ lt_chunk_data.pct_pft = np.zeros((n, shared_data.n_pfts_default), dtype=np.float64)
+ lt_chunk_data.pct_ocean = np.zeros(n, dtype=np.float64)
+
+ # todo: probably need to add lai path to land_type params and pass it through to here
+
+ lc_rs_data = None
+ prev_lt_data = None
+ climate_data = None
+ lai_data = None
+ elm_data = None
+
+ if lc_rs.use_lc_rs(year, lc_rs_name):
+ # read modis cover data - use tmp_dir because the data are downloaded (not stored)
+ # reading both the igbp cover data and the veg continuous fields data
+ lc_rs_geotiffs = lc_rs.read_to_geotiff(year, lc_rs_name, tmp_dir, cell_indices, ll_limits)
+ if not lc_rs_geotiffs:
+ logger.warning(
+ "landcover_process: no remote-sensing GeoTIFFs found; skipping chunk "
+ f"year={year} ll_limits={ll_limits} n_cells={len(cell_indices)}"
+ )
+ return lt_chunk_data
+
+ ### todo: this does not need to happen each year, see water_class in lc_rs module
+ # also, this should be generalized to other lc_rs_name values
+ # Replace water_class raster with an ocean-only mask derived from
+ # shapefile overlap before any regridding is performed.
+ if 'water_class' in lc_rs_geotiffs:
+ lc_rs.set_ocean(
+ lc_rs_name=lc_rs_name,
+ water_class_tif=lc_rs_geotiffs['water_class'],
+ source_data_path=com_config_dict['source_data_path'],
+ ocean_shapefile_path=com_config_dict.get('ocean_shapefile_path', ''),
+ )
else:
- print(f"Error: No previous year data available for year {year}, and _use_lc_rs is False; cannot process land cover data.")
- sys.exit(1)
-
- # get climate data (1900-2020, four historical periods; and cmip 6 future scenarios, 1km) need to pick the correct period
- # if year < 1900, then use the 1900 climate data
- # todo: probably define this here becase these are specific data
- climate_data = read_climate_data(year, com_config_dict['source_data_path'], lai_path)
-
- # get lai data for splitting tree/grass/shrub; this is based on li et al 1km lai data
- # these data do have short timeseries? then need to select appropriate year
- #todo: this can be in a utils module because other modules need to read these source data
- lai_data = read_lai_data(year, com_config_dict['source_data_path'], lai_path)
-
- # todo: use uraster to convert lc_rs_data and climate data and lai data to the landgen grid
-
- # convert lc_rs_data to the elm land types; this is igbp to generic elm land type mapping
- # also use the veg continuous fields data; can set modis to elm mapping file name here and read it based on lc_rs_name
- elm_data = lc_rs.convert_lc_rs_to_elm(lc_rs_data, lc_rs_name)
-
- # split tree/grass/shrub pfts based on cliamte data and li et al 1km lai data
- elm_data = split_tree_grass_shrub(elm_data, climate_data, lai_data)
-
- # normalize cell by adjusting the land cover distribution to fill the cell land area and reconciling with ocean data (landfrac)
- elm_data = normalize_cell.fill_land(elm_data, landfrac) # fill_land
- elm_data = normalize_cell.reconcile_ocean(elm_data, landfrac) # reconcile_ocean
-
- # now put elm data into lt_year_data
-
- return
-
-
+ # read and process previous year's landgen land type data and transitions to calculate this year's land cover distribution
+ if prev_year is not None:
+ # read previous year's landgen land type data
+ prev_out_file = Path(com_config_dict['out_path']) / prev_fname
+ if prev_out_file.exists():
+ #print(f"Reading previous year landgen land type data from {prev_out_file}")
+ # todo: define this in a helper function
+ prev_lt_data = read_prev_lt(prev_year, prev_out_file)
+ else:
+ raise FileNotFoundError(
+ f"Previous year output file does not exist: {prev_out_file}")
+
+ #Calculate this year's land cover distribution using the previous year's data and the transitions
+ # these calculations are based on landgen land type outputs
+ #temp_lt_data = transitions.run(prev_lt_data, year, prev_year)
+
+ ## todo: is this necessary? maybe we should leave this as output classes; these are also on landgen grid
+ ## can we check for appropriate types using the climate data below back to 1900, without converting to orignal lc classes?
+ # convert this year's land cover distribution to the lc rs classes
+ #lc_rs_data = lc_rs.convert_landgen_to_lc_rs(temp_lt_data, lc_rs_name)
+ else:
+ raise ValueError(f"Previous year is None for year {year}, and use_lc_rs is False; cannot process land cover data.")
+
+## todo: deal with generic elm pfts from modis data first, at modis resolution,
+# and then convert to landgen grid and split tree/grass/shrub based on climate and lai data
+# this is because the modis data are finer than 1km res, so we can do more explicit processing
+# the finer vcf data should be applied to the lc data to do this.
+
+ # get climate data (1900-2020, four historical periods; and cmip 6 future scenarios, 1km) need to pick the correct period
+ # if year < 1900, then use the 1900 climate data
+ # todo: probably define this here becase these are specific data
+ #climate_data = read_climate_data(year, com_config_dict['source_data_path'], lai_path)
+
+ # get lai data for splitting tree/grass/shrub; this is based on li et al 1km lai data
+ # these data do have short timeseries? then need to select appropriate year
+ #todo: this can be in a utils module because other modules need to read these source data
+ #lai_data = read_lai_data(year, com_config_dict['source_data_path'], lai_path)
+
+ #######
+ # todo: use uraster to convert lc_rs_data and climate data and lai data to the landgen grid
+
+ # regrid each lc_rs variable to the landgen mesh grid
+ # lc_rs_data: dict {varname: 1D np.ndarray (n_cells,)}; stack into 2D array (n_vars, n_cells)
+ lc_rs_data = {}
+ for varname, tif_path in lc_rs_geotiffs.items():
+ lc_rs_data[varname] = landgen_io.regrid_to_mesh(
+ mesh_file, {varname: tif_path}, cell_indices, out_grid_data,
+ out_type='data', remap_method=3
+ )
+ # 2D array: rows = variables (same order as lc_rs_geotiffs), cols = cells
+ #import numpy as np
+ #lc_rs_varnames = list(lc_rs_data.keys())
+ #lc_rs_array = np.stack([lc_rs_data[v] for v in lc_rs_varnames], axis=0) # (n_vars, n_cells)
+
+ # convert lc_rs_data to the elm land types; this is igbp to generic elm land type mapping
+ # also use the veg continuous fields data; can set modis to elm mapping file name here and read it based on lc_rs_name
+ #elm_data = lc_rs.convert_lc_rs_to_elm(lc_rs_data, lc_rs_name)
+
+ # split tree/grass/shrub pfts based on cliamte data and li et al 1km lai data
+ #elm_data = split_tree_grass_shrub(elm_data, climate_data, lai_data)
+
+ # normalize cell by adjusting the land cover distribution to fill the cell land area and reconciling with ocean data (landfrac)
+ #elm_data = normalize_cell.fill_land(elm_data, landfrac) # fill_land
+ #elm_data = normalize_cell.reconcile_ocean(elm_data, landfrac) # reconcile_ocean
+
+ # now put elm data into lt_year_data
+
+ # put some lc_rs data into lt_chunk_data for testing
+ # map lc_rs varnames -> lt_chunk_data fields
+ # LC_Type1: IGBP land cover class per cell -> stored in pct_pft row 0 as a placeholder
+ # (full pft distribution is computed later in convert_lc_rs_to_elm)
+ # VCF variables stored for later use in split_tree_grass_shrub
+ if 'LC_Type1' in lc_rs_data:
+ lt_chunk_data.pct_pft[:, 0] = lc_rs_data['LC_Type1']
+ if 'water_class' in lc_rs_data:
+ # set_ocean() rewrites the water_class tiffs to a binary mask where
+ # 100=ocean and 0=non-ocean.
+ ocn_mask = (lc_rs_data['water_class'] == 100)
+ lt_chunk_data.pct_ocean[ocn_mask] = 100
+ if 'Percent_Tree_Cover' in lc_rs_data:
+ lt_chunk_data.pct_pft[:, 1] = lc_rs_data['Percent_Tree_Cover']
+
+ return lt_chunk_data
+
+ finally:
+ logger.info(f"[landcover_process] FINISH {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} worker={worker_id} year={year} chunk_size={len(cell_indices)}")
+ # clean up worker temp dir regardless of success or failure
+ shutil.rmtree(tmp_dir, ignore_errors=True)
@@ -113,42 +216,46 @@ def landcover_process(lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_
## this sets up the pool and calls the landcover_process() function for each chunk of data
def run(lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
- com_config_dict, out_grid_data, ll_limits, cell_ids,
- manager, grid_manager, lt_manager):
+ com_config_dict, out_grid_data, decomp_indices, decomp_ll_limits,
+ manager):
- print(f"Processing landcover module with parameters:")
+ logger.info(f"Processing landcover module")
# todo: print the parameters here
- # number of available cpu cores (set by SBATCH during job submission)
- omp_threads_str = os.environ.get('OMP_NUM_THREADS')
+ # Determine the number of worker processes to use.
+ # Priority: SRUN_CPUS_PER_TASK (set explicitly in submit script via srun)
+ # -> SLURM_CPUS_PER_TASK (set by SLURM when --cpus-per-task is in the sbatch directives)
+ # -> SLURM_CPUS_ON_NODE (total CPUs allocated to this job on this node; always set by SLURM)
+ # -> mp.cpu_count() (all logical cores; safe for local runs)
+ in_slurm = os.environ.get('SLURM_JOB_ID') is not None
+
+ cpus_avail_int = (
+ tools.parse_cpu_env('SRUN_CPUS_PER_TASK') or
+ tools.parse_cpu_env('SLURM_CPUS_PER_TASK') or
+ tools.parse_cpu_env('SLURM_CPUS_ON_NODE') or
+ mp.cpu_count()
+ )
- if omp_threads_str is not None:
- try:
- # Convert the string value to an integer
- omp_threads_int = int(omp_threads_str)
- print(f"OMP_NUM_THREADS is set to: {omp_threads_int}")
- except ValueError:
- print(f"OMP_NUM_THREADS is set to an invalid integer value: {omp_threads_str}")
+ if in_slurm:
+ logger.info(f"Running under SLURM (job {os.environ['SLURM_JOB_ID']}): using {cpus_avail_int} workers "
+ f"(SRUN_CPUS_PER_TASK={os.environ.get('SRUN_CPUS_PER_TASK')}, "
+ f"SLURM_CPUS_PER_TASK={os.environ.get('SLURM_CPUS_PER_TASK')}, "
+ f"SLURM_CPUS_ON_NODE={os.environ.get('SLURM_CPUS_ON_NODE')})")
else:
- print("OMP_NUM_THREADS environment variable is not set.")
- # If not set, set to total cores on the node
- omp_threads_int = mp.cpu_count()
- print(f"Using total cores: {omp_threads_int}, but this may fail if "
- "SBATCH --cpus-per-task is set to a lower number or SBATCH --exclusive is not set")
+ logger.info(f"Running locally: using {cpus_avail_int} workers (mp.cpu_count())")
# set up the pool and call the landcover_process() function for each chunk of data
# chunks are defined by the lat-lon limits and corresponding landgen grid cell ids for the chunk;
- # these are created in land_type.process_single_year() and passed to this run() function as lists?
+ # these are created in land_type.process_single_year() and passed to this run() function
# there are more chunks than cpus; the pool will manage this for efficiency because chunks vary in size
# the results will be stored directly in the lt_year_data shared structure
# get the manager locks for the shared data structures
- # all locks come from the main mp.Manager() (SyncManager); custom managers don't support Lock()
- man_lock = manager.Lock()
- grid_lock = manager.Lock()
- lt_lock = manager.Lock()
+ # using data-specific locks, watch out for deadlocks.
+ #man_lock = manager.Lock()
+ #grid_lock = manager.Lock()
## todo: figure out the data to pass here
# each chunk is a tuple of the arguments for landcover_process, residing in a list
@@ -158,34 +265,41 @@ def run(lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
# (lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
# com_config_dict, out_grid_data, ll_limits2, cell_ids2, man_lock, grid_lock, lt_lock), etc]
- # load HEALPix mesh to map ll_limits chunks to cell ids
- global_parquet_path = (
- Path(com_config_dict['source_data_path'])
- / Path(com_config_dict['landgen_grid_path']).parent
- / 'merged_land_cells.parquet'
- )
- global_mesh_df = pd.read_parquet(global_parquet_path)
-
+## first try the return/copy approach to get this working
+# have the process function return an LtData object for the chunk, and then copy the data into the shared lt_year_data structure in this run() function
+# this is simpler to code and avoids potential issues with multiple processes writing to the same shared data structure, but may be less efficient because of the copying step
+## an alternative is to redefine lt_year_data as a numpy array of individual cell data structures of numpy arrays for each variable,
+#. define this as shared memory,
+#. and then each worker can write directly to the appropriate cells in the shared lt_year_data structure based on the cell ids for the chunk
+
+ ## todo: check that this is correct
+ # sort decomp_indices and decomp_ll_limits together, largest chunks first,
+ # so the pool receives the most expensive work items early (improves load balancing)
+ sorted_pairs = sorted(zip(decomp_indices, decomp_ll_limits),
+ key=lambda pair: len(pair[0]), reverse=True)
+ decomp_indices, decomp_ll_limits = zip(*sorted_pairs) if sorted_pairs else ([], [])
+
+ # create the list of chunk data
data_chunks = []
- for ll in ll_limits:
- min_lat, max_lat, min_lon, max_lon = ll
- mask = (
- (global_mesh_df['lat'] >= min_lat) & (global_mesh_df['lat'] < max_lat) &
- (global_mesh_df['lon'] >= min_lon) & (global_mesh_df['lon'] < max_lon)
- )
- chunk_cell_ids = global_mesh_df.loc[mask, 'cellid'].values
- if len(chunk_cell_ids) == 0:
- continue # skip ocean-only or empty chunks
+ for cidx in range(len(decomp_indices)):
+ # this list includes only cells in the landgen grid file
data_chunks.append((
- lt_year_data, year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
- com_config_dict, out_grid_data, ll, chunk_cell_ids,
- man_lock, grid_lock, lt_lock,
+ year, prev_year, prev_fname, lc_rs_path, lc_rs_name,
+ com_config_dict, out_grid_data, decomp_indices[cidx], decomp_ll_limits[cidx],
))
- print(f" Submitting {len(data_chunks)} landcover chunks to pool of {omp_threads_int} workers")
+ logger.info(f"Submitting {len(data_chunks)} landcover chunks to pool of {cpus_avail_int} workers")
+ resource_logger.info(f"In landcover submodule run():\n"
+ f"Submitting {len(data_chunks)} landcover chunks to pool of {cpus_avail_int} workers")
- with mp.Pool(processes=omp_threads_int) as pool:
- pool.starmap(landcover_process, data_chunks)
+ # Use imap_unordered so each chunk result is copied into lt_year_data and
+ # discarded as soon as it arrives, keeping only ~1 chunk result in memory
+ # at a time instead of accumulating all 360 simultaneously.
+ updated_vars = ['pct_pft', 'pct_ocean']
+ with mp.Pool(processes=cpus_avail_int) as pool:
+ for chunk_result in pool.imap_unordered(_landcover_process_star, data_chunks):
+ lt_year_data.copy_from(chunk_result, updated_vars)
+ del chunk_result
return
diff --git a/components/elm/tools/landgen/src/landgen/landcover_remote_sensing.py b/components/elm/tools/landgen/src/landgen/landcover_remote_sensing.py
new file mode 100644
index 000000000000..3a4f22a0b5c4
--- /dev/null
+++ b/components/elm/tools/landgen/src/landgen/landcover_remote_sensing.py
@@ -0,0 +1,225 @@
+# landcover_remote_sensing.py
+# this module defines the landcover source data
+# including schema mapping files
+
+# the goal is to be able to add different source data definitions here
+# and use a generic api to landcover.py
+
+import json
+import logging
+from pathlib import Path
+
+import numpy as np
+import rasterio
+from osgeo import ogr
+from rasterio.features import rasterize
+from rasterio.warp import transform_geom
+
+from . import landgen_io
+
+logger = logging.getLogger('landgen')
+
+
+# exclude the caspian sea from the ocean mask because it is inlcluded in the hydrolakes data
+# for now keep the kerch straight and sea of azov; these will be dealt with later, maybe in mksurfdata
+# these are vertices to exclude sea of azov and kerch straight in one polygon:
+# [ [34.4, 45.2], [34.4, 46.5], 37.6, 47.5], [40.1, 47.5], [38.3, 45.2], [36.8, 45.2], [36.7, 45.16], [36.4, 45.1], [35.5, 45.2], [34.4, 45.2] ]
+# coordinates are in [lon,lat] order
+_OCEAN_EXCLUSION_POLYGONS_LONLAT = [
+ {
+ 'name': 'Caspian Sea',
+ 'geometry': {
+ 'type': 'Polygon',
+ 'coordinates': [[
+ [46.0, 36.0],
+ [55.8, 36.0],
+ [55.8, 47.6],
+ [46.0, 47.6],
+ [46.0, 36.0],
+ ]],
+ },
+ }
+]
+
+#### set some parameters for the landcover remote sensing data here
+
+#------- modis -----------------------------------
+modis_name = 'modis'
+modis_years = list(range(2001, 2024+1)) # 2001-2024
+# MCD12Q1 is the land cover type product, MOD44B is the vegetation continuous fields product
+modis_products = ['MCD12Q1.061', 'MOD44B.061', 'MOD44W.061'] # specify the modis products to use; adjust as needed
+# MCD12Q1: LC_Type1 is the IGBP land cover type, LW is the land/water mask; not using QC
+modis_lc_variable_names = ['LC_Type1']
+# MOD44B: water cell has value==200; not using Quality
+modis_vcf_variable_names = ['Percent_Tree_Cover', 'Percent_NonTree_Vegetation', 'Percent_NonVegetated']
+# MOD44W: use for ocean mask: 0=shallow ocean, 6=moderate ocean, 7=deep ocean
+# the ocean data are static, so only need to read for one year; use the same for all years
+# use year 2000 because it is most consistent with the source static data
+# this can have only one variable name, even though it is a list
+modis_wat_variable_names = ['seven_class']
+water_year = 2000
+
+# note that there are no modis tiles for these products for:
+# easter island (ll_limtis=(-27.363585, -26.944359, -109.599609, -109.160156))
+# sao pedro e sao paulo islands (ll_limtis=(0.74606, 1.044512, -29.53125, -29.179688))
+# random southern ocean ice island (ll_limtis=(-54.628738, -54.244919, 3.28125, 3.534031))
+
+# the water cells in MOD44B are from MOD44W, and are the highest res land/water mask;
+# LW is from MOD44W, with water cell where 2 or more MOD44W cells are water
+
+
+# todo: dictionary for modis to elm land type mapping
+
+#------- use_lc_rs -------------------------------
+def use_lc_rs(year, lc_rs_name):
+ """Return True if remote sensing data identified by lc_rs_name is available for year."""
+ if lc_rs_name == modis_name:
+ return year in modis_years
+ else:
+ raise ValueError(f"use_lc_rs: unknown lc_rs_name '{lc_rs_name}'")
+
+
+### maybe pass com_confg_dict to this read to get more control
+# e.g., the water data only need to be read once
+
+#------- read ------------------------------------
+def read_to_geotiff(year, lc_rs_name, lc_rs_path, cell_indices, ll_limits):
+ """Read land cover data from remote sensing source identified by lc_rs_name for the given year."""
+
+ if lc_rs_name == modis_name:
+ # read modis cover data - lc_rs_path is the temp dir because the data are downloaded as needed
+ # reading the igbp cover data and writing the geotiffs for this chunk
+ lc_files = landgen_io.read_modis_ll_to_geotiff(year, lc_rs_path, modis_products[0],
+ modis_lc_variable_names, ll_limits=ll_limits)
+ # reading the vcf data and writing the geotiffs for this chunk
+ vcf_files = landgen_io.read_modis_ll_to_geotiff(year, lc_rs_path, modis_products[1],
+ modis_vcf_variable_names, ll_limits=ll_limits)
+ # reading the water mask data and writing the geotiffs for this chunk
+ # todo: only do this once, but how to determine when? and read only one year
+ wat_files = landgen_io.read_modis_ll_to_geotiff(water_year, lc_rs_path, modis_products[2],
+ modis_wat_variable_names, ll_limits=ll_limits)
+ # Present a generic key to downstream code while preserving MODIS variable naming internally.
+ water_files = {}
+ if modis_wat_variable_names[0] in wat_files:
+ water_files['water_class'] = wat_files[modis_wat_variable_names[0]]
+ out_files = {**lc_files, **vcf_files, **water_files}
+ else:
+ raise ValueError(f"read: unknown lc_rs_name '{lc_rs_name}'")
+
+ return out_files
+
+
+def set_ocean(lc_rs_name, water_class_tif, source_data_path, ocean_shapefile_path):
+ """Write GeoTIFF binary ocean mask.
+
+ Pixels are set to 100 only when both conditions are true:
+ 1) original water_class values identify ocean
+ 2) any part of the pixel intersects the ocean polygon shapefile
+
+ All other pixels are set to 0. The input TIFF is replaced in place.
+
+ Args:
+ lc_rs_name (str): Remote-sensing source identifier.
+ water_class_tif (str|Path): Path to water_class GeoTIFF.
+ source_data_path (str|Path): Root source-data path from config.
+ ocean_shapefile_path (str|Path): Shapefile path relative to source_data_path.
+
+ Returns:
+ Path: Path to the rewritten TIFF (same as input path).
+ """
+
+ tif_path = Path(water_class_tif)
+ if not tif_path.exists():
+ raise FileNotFoundError(f"set_ocean: water_class GeoTIFF not found: {tif_path}")
+
+ with rasterio.open(tif_path) as src:
+ arr = src.read(1)
+ profile = src.profile.copy()
+ out_shape = arr.shape
+ transform = src.transform
+ target_crs = src.crs
+
+ # ocean shapefile should apply to all rs data
+ if not ocean_shapefile_path:
+ raise ValueError(
+ "set_ocean: ocean_shapefile_path is empty; set this in config.json "
+ "as a path relative to source_data_path"
+ )
+
+ shp_path = Path(source_data_path) / Path(ocean_shapefile_path)
+ if not shp_path.exists():
+ raise FileNotFoundError(f"set_ocean: ocean shapefile not found: {shp_path}")
+
+ ds = ogr.Open(str(shp_path))
+ if ds is None:
+ raise RuntimeError(f"set_ocean: could not open ocean shapefile: {shp_path}")
+ layer = ds.GetLayer(0)
+ layer_srs = layer.GetSpatialRef()
+ src_crs_wkt = layer_srs.ExportToWkt() if layer_srs is not None else None
+
+ geometries = []
+ layer.ResetReading()
+ for feat in layer:
+ geom = feat.GetGeometryRef()
+ if geom is None:
+ continue
+ gjson = json.loads(geom.ExportToJson())
+ if src_crs_wkt and target_crs is not None:
+ gjson = transform_geom(src_crs_wkt, target_crs.to_string(), gjson)
+ geometries.append(gjson)
+ ds = None
+
+ if not geometries:
+ raise RuntimeError(f"set_ocean: no geometries found in {shp_path}")
+
+ ocean_touch_mask = rasterize(
+ geometries,
+ out_shape=out_shape,
+ transform=transform,
+ fill=0,
+ default_value=1,
+ all_touched=True,
+ dtype=np.uint8,
+ ).astype(bool)
+
+ exclusion_geometries = []
+ for feature in _OCEAN_EXCLUSION_POLYGONS_LONLAT:
+ exclusion_geom = feature['geometry']
+ if target_crs is not None:
+ exclusion_geom = transform_geom('EPSG:4326', target_crs.to_string(), exclusion_geom)
+ exclusion_geometries.append(exclusion_geom)
+
+ exclusion_mask = rasterize(
+ exclusion_geometries,
+ out_shape=out_shape,
+ transform=transform,
+ fill=0,
+ default_value=1,
+ all_touched=True,
+ dtype=np.uint8,
+ ).astype(bool)
+
+ ocean_touch_mask &= ~exclusion_mask
+
+ # --------- modis -----------------------------------
+ # 1) original seven_class value is in {0, 5, 6, 7}, see above for class definitions
+ if lc_rs_name == modis_name:
+ class_mask = np.isin(arr, [0, 5, 6, 7])
+ else:
+ raise ValueError(f"set_ocean: unknown lc_rs_name '{lc_rs_name}'")
+
+ out_arr = np.where(class_mask & ocean_touch_mask, 100, 0).astype(np.uint8)
+
+ profile.update(dtype=rasterio.uint8, count=1, nodata=0)
+ with rasterio.open(tif_path, 'w', **profile) as dst:
+ dst.write(out_arr, 1)
+
+ logger.info(
+ f"set_ocean: rewrote {tif_path} using ocean polygons from {shp_path}; "
+ f"excluded={','.join(f['name'] for f in _OCEAN_EXCLUSION_POLYGONS_LONLAT)}; "
+ f"ocean_pixels={int(np.count_nonzero(out_arr == 100))}, total_pixels={out_arr.size}"
+ )
+ return tif_path
+
+
+
\ No newline at end of file
diff --git a/components/elm/tools/landgen/src/landgen/landgen.py b/components/elm/tools/landgen/src/landgen/landgen.py
index e128f93da434..4afdcc9dc16d 100644
--- a/components/elm/tools/landgen/src/landgen/landgen.py
+++ b/components/elm/tools/landgen/src/landgen/landgen.py
@@ -6,58 +6,148 @@
import multiprocessing as mp
import importlib
import json
-import sys
+import os
from pathlib import Path
+import sys
+import logging
+from datetime import datetime
from . import shared_data
-from .shared_data import GridData, GridManager
+from . import landgen_io
+from . import tools
+import threading
+
+
def load_config(config_path):
with open(config_path, 'r') as f:
return json.load(f)
def main(config_path):
- # todo: need to deal with landfrac data structure
- landfrac = None
- config = load_config(config_path)
- modules = config.get('modules', [])
-
- # get the common parameters for all modules and store in a shared dictionary
- temp_dict = {
- 'start_year': config.get('start_year', 2015),
- 'end_year': config.get('end_year', 2015),
- 'source_data_path': config.get('source_data_path', ''),
- 'landgen_grid_path': config.get('landgen_grid_path', ''),
- 'out_path': config.get('out_path', ''),
- }
- manager = mp.Manager()
- com_config_dict = manager.dict(temp_dict)
-
- # create the shared landgen out grid shared data structure
- grid_manager = GridManager()
- grid_manager.start()
- out_grid_data = grid_manager.GridData()
- out_grid_data.allocate()
-
- ## todo: read in the grid file and set the values in out_grid_data
-
- for mod in modules:
- name = mod['name']
- params = mod.get('params', {})
- try:
- module = importlib.import_module(f'landgen.{name}')
- if hasattr(module, 'run'):
- print(f"Running module: {name}")
- run_list = [*params.values(), com_config_dict, out_grid_data, manager, grid_manager]
- module.run(*run_list)
- else:
- print(f"Module {name} does not have a 'run' function.")
- except ImportError as e:
- print(f"Could not import module {name}: {e}")
-
-
- # free the shared memory
- com_config_dict = None
- out_grid_data = None
- manager.shutdown()
- grid_manager.shutdown()
- return
+ # todo: need to deal with landfrac data structure
+ landfrac = None
+ config = load_config(config_path)
+
+ # set up the shared logger before anything else so all modules can use it
+ out_path = Path(config.get('out_path', '.'))
+ out_path.mkdir(parents=True, exist_ok=True)
+ timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
+ job_id = os.environ.get('SLURM_JOB_ID', 'local')
+ log_name = f'landgen_{timestamp}_{job_id}.log'
+ logger = tools.setup_logger('landgen', out_path / log_name)
+ logger.info(f"landgen started at {timestamp} — config: {config_path}")
+ logger.info(f"log file: {out_path / log_name}")
+ modules = config.get('modules', [])
+
+ # set up the cluster resource logger
+ # the default interval is 1 minute (60 seconds)
+ # adjust as needed by changing'interval_sec': ### below in kwargs
+ resource_log_name = f'resource_monitor_{timestamp}_{job_id}.log'
+ resource_logger = tools.setup_logger('ClusterMonitor', out_path / resource_log_name)
+ stop_event = threading.Event()
+ resource_monitor_thread = threading.Thread(
+ target=tools.monitor_cluster_resources,
+ kwargs={'interval_sec': 60.0, 'stop_event': stop_event},
+ daemon=True)
+ resource_monitor_thread.start()
+
+ # get the common parameters for all modules and store in a shared dictionary
+ temp_dict = {
+ 'start_year': config.get('start_year', 2015),
+ 'end_year': config.get('end_year', 2015),
+ 'source_data_path': config.get('source_data_path', ''),
+ 'landgen_grid_path': config.get('landgen_grid_path', ''),
+ 'ocean_shapefile_path': config.get('ocean_shapefile_path', ''),
+ 'out_path': config.get('out_path', ''),
+ 'decomp_box_size_degrees': config.get('decomp_box_size_degrees', 10),
+ 'log_path': str(out_path / log_name),
+ }
+ manager = mp.Manager()
+ com_config_dict = manager.dict(temp_dict)
+
+ ## todo: delete if not using grid manager
+ # for now just set up a structure for use by the master proc
+ # the only thing to be updated in this structure is the landfrac
+ # if using the manager, need to create set/get functions for the variables
+ # create the shared landgen out grid shared data structure
+ #grid_manager = shared_data.GridManager()
+ #grid_manager.start()
+ #out_grid_data = grid_manager.GridData()
+ #out_grid_data.allocate()
+
+ # do the decomposition of the landgen mesh here for all modules
+ # these are passed by reference to the run() functions for each module
+ # default data chunks are based on 10x10 degree lat-lon boxes (648 chunks)
+ # 15x15 degree box gives 288 chunks, 30x30 box gives 72 chunks
+ # Note that chunks are not equal in size
+
+ # these are lists of tuples with each tuple defining a chunk, and are paired in order
+ # decomp_indices: indices within each chunk for the landgen grid file variables
+ # decomp_ll_limits = list(float) of [(min_lat, max_lat, min_lon, max_lon),... for each chunk]
+ # these are based on the vertices of the cells in decomp_indices to ensure full coverage
+ # the chunk_file is written, but not used; it is for diagnostics
+ # note that indices are 0-based in these arrays
+ decomp_indices = []
+ decomp_ll_limits = []
+ mesh_nc_path = Path(com_config_dict['source_data_path']) / com_config_dict['landgen_grid_path']
+
+ # load all mesh cells from the NetCDF domain file and fill out_grid_data
+ mesh = landgen_io.load_mesh_nc(mesh_nc_path) # loads all cells (no indices/ll_limits filter)
+ out_grid_data = shared_data.GridData()
+ out_grid_data.allocate(n_cells=mesh['cellid'].shape[0], n_vertices=mesh['lon_v'].shape[1])
+ out_grid_data.cell_id[:] = mesh['cellid']
+ out_grid_data.lon_cen[:] = mesh['lon']
+ out_grid_data.lat_cen[:] = mesh['lat']
+ out_grid_data.cell_area[:] = mesh['area']
+ out_grid_data.lon_vtx[:, :] = mesh['lon_v'] # shape (n_cells, n_vertices)
+ out_grid_data.lat_vtx[:, :] = mesh['lat_v'] # shape (n_cells, n_vertices)
+ # landfrac is initialised to 1 by allocate(); updated later by landcover? module
+
+ chunk_file = landgen_io.set_decomp_cell_idx_ll_limits(out_grid_data, decomp_indices, decomp_ll_limits,
+ com_config_dict['decomp_box_size_degrees'], com_config_dict['out_path'])
+
+ try:
+ for mod in modules:
+ name = mod['name']
+ params = mod.get('params', {})
+ try:
+ module = importlib.import_module(f'landgen.{name}')
+ if hasattr(module, 'run'):
+ logger.info(f"Running module: {name}")
+ module.run(**params,
+ com_config_dict=com_config_dict,
+ out_grid_data=out_grid_data,
+ manager=manager,
+ decomp_indices=decomp_indices,
+ decomp_ll_limits=decomp_ll_limits)
+ else:
+ logger.warning(f"Module {name} does not have a 'run' function.")
+ except ImportError as e:
+ logger.error(f"Could not import module {name}: {e}; skipping.")
+
+ except Exception as e:
+ logger.exception(f"ERROR in landgen: {e}")
+ raise
+
+ finally:
+ # remove uraster side logs only from run/submit cwd locations.
+ # keep copies in out_path for diagnostics.
+ uraster_log_names = ('extract.log', 'intersect.log', 'uraster.log', 'utility.log')
+ cleanup_dirs = [Path.cwd()]
+ submit_dir = os.environ.get('SLURM_SUBMIT_DIR')
+ if submit_dir:
+ cleanup_dirs.append(Path(submit_dir))
+ for d in cleanup_dirs:
+ for name in uraster_log_names:
+ try:
+ (d / name).unlink(missing_ok=True)
+ except Exception:
+ pass
+
+ manager.shutdown()
+ stop_event.set()
+ resource_monitor_thread.join()
+ resource_logger.info("Cluster resource monitor thread stopped.")
+ timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
+ logger.info(f"landgen finished at {timestamp}")
+
+ return
diff --git a/components/elm/tools/landgen/src/landgen/landgen_io.py b/components/elm/tools/landgen/src/landgen/landgen_io.py
index 83464d411231..c1be2f4f54d5 100644
--- a/components/elm/tools/landgen/src/landgen/landgen_io.py
+++ b/components/elm/tools/landgen/src/landgen/landgen_io.py
@@ -1,24 +1,204 @@
# landgen_io.py
# Utility functions for reading harvest data from LUH2 and grazing data from HYDE3.5
+import logging
+import time
import xarray as xr
import numpy as np
from pathlib import Path
-import pandas as pd
+
+logger = logging.getLogger('landgen')
import rasterio
from rasterio.transform import from_bounds
import json
-import shapely.wkb
from uraster.classes.uraster import uraster as URaster
+import earthaccess
+import glob
+import os
+import math
+import re
+from pyproj import Proj
+from osgeo import gdal
+
+
+#--------------------------------------------------------------------------
+def _earthaccess_search_with_retry(max_retries=3, delay_sec=10, **kwargs):
+ """Call earthaccess.search_data with retries on transient server errors (HTTP 5xx)."""
+ for attempt in range(1, max_retries + 1):
+ try:
+ return earthaccess.search_data(**kwargs)
+ except (RuntimeError, Exception) as exc:
+ msg = str(exc)
+ if attempt < max_retries and ('500' in msg or 'Internal Error' in msg or 'Server Error' in msg):
+ logger.warning(
+ f"_earthaccess_search_with_retry: attempt {attempt}/{max_retries} failed "
+ f"({msg!r}); retrying in {delay_sec}s"
+ )
+ time.sleep(delay_sec)
+ else:
+ raise
+
+
+#--------------------------------------------------------------------------
+def set_decomp_cell_idx_ll_limits(grid_data, decomp_indices, decomp_ll_limits,
+ chunk_size_degrees=10, out_dir=None):
+ """
+ Populate chunk_indices and chunk_ll_limits in-place from a GridData structure.
+ Also write a companion .npz file mapping chunk lat-lon limits to cell indices.
+
+ Approx. cell-centre lat/lon assigns each cell to a chunk box. Vertex coordinates
+ lon_vtx/lat_vtx determine the tight actual bounding box of the cells in that chunk,
+ which is what callers (e.g. write_latlon_to_geotiff) need for source-data
+ slicing — the fixed chunk box may extend into ocean where no source data exist.
+
+ Chunks with no land cells are skipped, so len(chunk_indices) may be less
+ than the total number of chunk boxes.
+
+ Based on HealPix mesh, but any mesh with the same variables in the same format should work.
+
+ Args:
+ grid_data (GridData): GridData object with lat_cen, lon_cen (~ cell centres)
+ and lat_vtx, lon_vtx (vertex coordinates,
+ shape n_cells x n_vertices) already populated.
+ chunk_indices (list): Output — populated in-place. Each element is a
+ 1D tuple of HEALPix cellid values (int64)
+ identifying the cells in that chunk.
+ chunk_ll_limits (list): Output — populated in-place. Each element is a
+ (min_lat, max_lat, min_lon, max_lon) tuple of the
+ tight vertex bounding box for that chunk's cells.
+ Matched by position to chunk_indices.
+ chunk_size_degrees (int): Size of the decomposition boxes in degrees.
+ out_dir (str|Path): Directory to write the companion .spatial_index.npz file.
+ """
+ if out_dir is None:
+ raise ValueError("set_decomp_cell_idx_ll_limits: out_dir is required")
+
+ lat = grid_data.lat_cen.astype(np.float64) # (n_cells,)
+ lon = grid_data.lon_cen.astype(np.float64) # (n_cells,)
+ lon_v = grid_data.lon_vtx.astype(np.float64) # (n_cells, n_vertices)
+ lat_v = grid_data.lat_vtx.astype(np.float64) # (n_cells, n_vertices)
+
+ decomp_indices.clear()
+ decomp_ll_limits.clear()
+
+ index = {}
+ for min_lat, max_lat, min_lon, max_lon in calc_ll_limits(chunk_size_degrees):
+ mask = (lat >= min_lat) & (lat < max_lat) & (lon >= min_lon) & (lon < max_lon)
+ indices = np.where(mask)[0]
+ if indices.size == 0:
+ continue # skip ocean-only boxes
+
+ # tight bounding box from actual vertex extents of the cells in this chunk
+ lon_v_chunk = lon_v[indices] # (n_chunk_cells, n_vertices)
+ lat_v_chunk = lat_v[indices]
+ decomp_indices.append(tuple(indices.tolist()))
+ decomp_ll_limits.append((
+ float(lat_v_chunk.min()), float(lat_v_chunk.max()),
+ float(lon_v_chunk.min()), float(lon_v_chunk.max()),
+ ))
+
+ # companion file stores NC row indices (used by load_mesh_nc via ds.isel())
+ key = f"{min_lat:.0f}_{max_lat:.0f}_{min_lon:.0f}_{max_lon:.0f}"
+ index[key] = indices
+
+ # write companion file
+ out_path = Path(out_dir) / 'spatial_index.npz'
+ np.savez_compressed(out_path, **index)
+
+ logger.info(f"set_decomp_cell_idx_ll_limits: built {len(decomp_indices)} chunks")
+ logger.info(f"set_decomp_cell_idx_ll_limits: chunks include {sum(len(t) for t in decomp_indices)} cells")
+ logger.info(f"set_decomp_cell_idx_ll_limits: check against grid_data.num_cells or netcdf file to ensure all cells are included")
+ logger.info(f"set_decomp_cell_idx_ll_limits: wrote {len(index)} chunks to {out_path}")
+
+ return out_path
+
+#--------------------------------------------------------------------------
+def load_mesh_nc(nc_path_name, indices=None, ll_limits=None):
+ """
+ Load HEALPix mesh data from a NetCDF domain file.
+
+ Args:
+ nc_path_name (str|Path): Path (and name) to the domain NetCDF file.
+ indices (tuple|None): 1D tuple of NC indices to load,
+ as returned in chunk_indices by
+ set_decomp_cell_idx_ll_limits(). If None,
+ check for ll_limits.
+ ll_limits (tuple|None): (min_lat, max_lat, min_lon, max_lon) tuple
+ to select cells by mapping file. If also None,
+ no lat/lon filtering is applied. This option
+ uses the companion .npz spatial mapping file
+ created by set_decomp_cell_idx_ll_limits().
+
+ Returns:
+ dict with keys 'cellid', 'lon_v', 'lat_v', 'lon', 'lat', 'area' (all np.ndarray).
+ """
+ nc_path = Path(nc_path_name)
+ if not nc_path.exists():
+ raise FileNotFoundError(f"load_mesh_nc: Mesh NetCDF file not found: {nc_path}")
+
+ ds = xr.open_dataset(nc_path, decode_times=False)
+
+ if indices is not None:
+ cell_dim = ds['lat'].dims[0]
+ subset = ds.isel({cell_dim: indices})
+ elif ll_limits is not None:
+ min_lat, max_lat, min_lon, max_lon = ll_limits
+ key = f"{min_lat:.0f}_{max_lat:.0f}_{min_lon:.0f}_{max_lon:.0f}"
+ idx_path = nc_path.with_suffix('.spatial_index.npz')
+ if not idx_path.exists():
+ raise FileNotFoundError(
+ f"load_mesh_nc: Spatial index not found: {idx_path}. "
+ f"Run set_decomp_cell_idx_ll_limits('{nc_path}', decomp_indices, decomp_ll_limits, chunk_size_degrees) first."
+ )
+ indices_ll = np.load(idx_path)[key]
+ if indices_ll.size == 0:
+ ds.close()
+ raise ValueError(f"load_mesh_nc: No mesh cells found within ll_limits {ll_limits}.")
+ cell_dim = ds['lat'].dims[0]
+ subset = ds.isel({cell_dim: indices_ll})
+ else:
+ subset = ds
+
+ cellid = subset['cellid'].values.astype(np.int64)
+ lon_v = subset['xv'].values.astype(np.float64)
+ lat_v = subset['yv'].values.astype(np.float64)
+ lon = subset['lon'].values.astype(np.float64)
+ lat = subset['lat'].values.astype(np.float64)
+ area = subset['area'].values.astype(np.float64)
+ ds.close()
+
+ logger.info(f"load_mesh_nc: loaded {len(cellid)} cells from {nc_path}")
+ return {'cellid': cellid, 'lon_v': lon_v, 'lat_v': lat_v, 'lon': lon, 'lat': lat, 'area': area}
+
-# Default harvest variable names from LUH2 transitions.nc
-LUH2_HARVEST_VARS = [
- 'primf_harv', # wood harvest area from primary forest land
- 'primn_harv', # wood harvest area from primary non forest land
- 'secmf_harv', # wood harvest area from secondary mature forest land
- 'secyf_harv', # wood harvest area from secondary young forest land
- 'secnf_harv', # wood harvest area from secondary non forest land
-]
+
+#--------------------------------------------------------------------------
+def calc_ll_limits(size_degrees):
+ """Calculate and return a list of tuples (min_lat, max_lat, min_lon, max_lon)
+ for each size_degrees x size_degrees chunk covering the globe.
+
+ Latitude spans -90 to 90 degrees.
+ Longitude spans -180 to 180 degrees.
+
+ Returns:
+ list of tuples: [(min_lat, max_lat, min_lon, max_lon), ...]
+ """
+ ll_limits = []
+ if 90 % size_degrees != 0 or 180 % size_degrees != 0:
+ raise ValueError(
+ f"calc_ll_limits: size_degrees ({size_degrees}) must evenly divide "
+ f"both 90 (latitude half-range) and 180 (longitude half-range)."
+ )
+ lat = -90.0
+ while lat < 90.0:
+ max_lat = min(lat + size_degrees, 90.0)
+ lon = -180.0
+ while lon < 180.0:
+ max_lon = min(lon + size_degrees, 180.0)
+ ll_limits.append((lat, max_lat, lon, max_lon))
+ lon = max_lon
+ lat = max_lat
+ return ll_limits
#--------------------------------------------------------------------------
def _get_year_idx(time_values, year, ncfile, time_units=None):
@@ -69,7 +249,7 @@ def _to_cal(offset): return offset + epoch_year
actual_year = time_as_years[year_idx]
if abs(actual_year - year) > 1:
raise ValueError(
- f"Requested year {year} not found in {ncfile} "
+ f"_get_year_idx: Requested year {year} not found in {ncfile} "
f"(closest available: {actual_year:.0f})"
)
return year_idx
@@ -84,98 +264,412 @@ def _to_cal(offset): return offset
if abs(actual_year - year) > 1:
raise ValueError(
- f"Requested year {year} not found in {ncfile} "
+ f"_get_year_idx: Requested year {year} not found in {ncfile} "
f"(closest available: {actual_year:.0f})"
)
return year_idx
#--------------------------------------------------------------------------
-def read_luh2_harvest(year, harvest_path, harvest_name, variable_names=None):
+def _get_modis_tile_idx(lon, lat):
+ """
+ Calculate the MODIS tile indices (h, v) for a given longitude and latitude.
+ The MODIS Sinusoidal grid divides the globe into 36 horizontal (h) and 18 vertical (v) tiles,
+ each approximately 10 degrees in size.
+ The tile indices start at (h=0, v=0) in the upper left corner (90N, 180W) and increase to the right and downward.
+
+ Returns:
+ tuple: (h, v) tile indices for the given longitude and latitude.
+ """
+ # Standard MODIS Sinusoidal parameters
+ R = 6371007.181
+ W = 2 * math.pi * R
+ T = W / 36
+
+ # Project Lat/Lon to Sinusoidal
+ modis_grid = Proj(f'+proj=sinu +R={R} +nadgrids=@null +wktext')
+ x, y = modis_grid(lon, lat)
+
+ # Calculate tile indices
+ h = int((x + W / 2) / T)
+ v = int((W / 4 - y) / T)
+ return h, v
+
+
+#--------------------------------------------------------------------------
+def _get_modis_tile_idxs_ll(ll_limits):
"""
- Read LUH2 harvest variables for a given year.
+ Get the MODIS tile indices for a lat/lon bounding box.
Args:
- year (int): Year to extract. LUH2 covers 850-2015.
- harvest_path (str or Path): Directory containing the LUH2 NetCDF file
- harvest_name (str): Filename of the LUH2 NetCDF file
- variable_names (list or None): Variables to extract. Defaults to all 5 LUH2_HARVEST_VARS if None.
+ ll_limits (tuple/list): 4-element (min_lat, max_lat, min_lon, max_lon).
Returns:
- dict: {varname: 2D np.ndarray shape (lat, lon)} for the requested year.
- lat/lon coordinate arrays are included as 'lat' and 'lon' keys.
+ a cvs string of indices in modis name format: 'h##v##, ...' for the tiles that intersect the bounding box.
"""
- if variable_names is None:
- variable_names = LUH2_HARVEST_VARS
+ if ll_limits is None or len(ll_limits) != 4:
+ raise ValueError(
+ f"_get_modis_tile_idxs_ll: ll_limits must be a 4-element (min_lat, max_lat, min_lon, max_lon), got {ll_limits}"
+ )
- ncfile = Path(harvest_path) / harvest_name
- if not ncfile.exists():
- raise FileNotFoundError(f"LUH2 harvest file not found: {ncfile}")
+ min_lat, max_lat, min_lon, max_lon = [float(x) for x in ll_limits]
+ if not (-90.0 <= min_lat <= 90.0 and -90.0 <= max_lat <= 90.0):
+ raise ValueError(
+ f"_get_modis_tile_idxs_ll: invalid latitude bounds {ll_limits}. "
+ "Expected (min_lat, max_lat, min_lon, max_lon) with lat in [-90, 90]."
+ )
+ if not (-180.0 <= min_lon <= 180.0 and -180.0 <= max_lon <= 180.0):
+ raise ValueError(
+ f"_get_modis_tile_idxs_ll: invalid longitude bounds {ll_limits}. "
+ "Expected (min_lat, max_lat, min_lon, max_lon) with lon in [-180, 180]."
+ )
+ if min_lat > max_lat or min_lon > max_lon:
+ raise ValueError(
+ f"_get_modis_tile_idxs_ll: non-monotonic bounds {ll_limits}. "
+ "Expected min <= max for both lat and lon."
+ )
- ds = xr.open_dataset(ncfile, decode_times=False)
+ corners = [(min_lon, min_lat), (min_lon, max_lat), (max_lon, min_lat), (max_lon, max_lat)]
+ tile_idxs = [_get_modis_tile_idx(lon, lat) for lon, lat in corners]
+ min_h = min(h for h, v in tile_idxs)
+ max_h = max(h for h, v in tile_idxs)
+ min_v = min(v for h, v in tile_idxs)
+ max_v = max(v for h, v in tile_idxs)
+ tile_idxs = [(h, v) for h in range(min_h, max_h + 1) for v in range(min_v, max_v + 1)]
+ tile_strs = sorted(f"h{h:02d}v{v:02d}" for h, v in set(tile_idxs))
+ return ", ".join(tile_strs)
- # LUH2 time axis is 'years since 850-01-01'; values are offsets from 850, not calendar years
- time_units = ds['time'].attrs.get('units', None)
- year_idx = _get_year_idx(ds['time'].values, year, ncfile, time_units=time_units)
- print(f" read_luh2_harvest: reading year {year} (time index {year_idx}) from {ncfile}")
- out = {'lat': ds['lat'].values, 'lon': ds['lon'].values}
- for v in variable_names:
- if v not in ds:
- raise KeyError(f"Variable '{v}' not found in {ncfile}. "
- f"Available variables: {list(ds.data_vars)}")
- out[v] = ds[v].isel(time=year_idx).values # shape: (lat, lon)
+#--------------------------------------------------------------------------
+def _granule_hv_tag(granule):
+ """Extract MODIS h##v## tile tag from an earthaccess granule-like object."""
+ texts = [str(granule)]
+ try:
+ links = granule.data_links() if hasattr(granule, 'data_links') else []
+ texts.extend([str(x) for x in links])
+ except Exception:
+ pass
+ for txt in texts:
+ m = re.search(r'h\d{2}v\d{2}', txt, flags=re.IGNORECASE)
+ if m:
+ return m.group(0).lower()
+ return None
- ds.close()
- return out
#--------------------------------------------------------------------------
-def read_hyde_grazing(year, grazing_path, grazing_names):
+def read_modis_ll_to_geotiff(year, dir_path, product, variable_names=None, ll_limits=None):
"""
- Read HYDE3.5 grazing data for a given year.
-
- grazing_names is a dict mapping a grazing category label to a NetCDF filename,
- as specified in config.json. Each file is expected to contain a single variable
- whose name matches the file stem (e.g. 'pasture.nc' -> variable 'pasture').
+ Download MODIS HDF tiles for a given year, mosaic them, and convert each
+ requested variable to a GeoTIFF file in dir_path.
Args:
- year (int): Year to extract. HYDE3.5 baseline covers 10000 BCE - 2023 CE.
- grazing_path (str or Path): Directory containing the HYDE3.5 NetCDF files
- grazing_names (dict): Mapping of {category_label: filename}
-
+ year (int): Year to extract.
+ dir_path (str or Path): Directory for downloaded HDF files and output GeoTIFFs.
+ If called from a worker process, use a unique per-worker
+ temp directory to avoid file collisions.
+ The caller is responsible for removing this directory
+ after processing is complete.
+ product (str): MODIS product identifier (e.g. 'MCD12Q1.061').
+ variable_names (list): Variable (subdataset) names to convert. Must be provided.
+ ll_limits (tuple/list or None): 4-element (min_lat, max_lat, min_lon, max_lon).
+ When given, only the lat/lon rows/columns that overlap this region are read.
+ The slice is inclusive — any grid cell whose coordinate falls within
+ [min_lat, max_lat] x [min_lon, max_lon] is included.
Returns:
- dict: {category_label: 2D np.ndarray shape (lat, lon)} for the requested year,
- plus 'lat' and 'lon' coordinate arrays (taken from the first file read).
+ dict: {varname: Path} mapping each variable name to the Path of the GeoTIFF
+ written to dir_path.
"""
- if not isinstance(grazing_names, dict):
- raise TypeError(
- f"grazing_names must be a dict (e.g. {{'pasture': 'pasture.nc', "
- f"'rangeland': 'rangeland.nc'}}), got {type(grazing_names).__name__}. "
- f"Please update config.json accordingly."
+
+ ###### need an Earthdata username and password to download MODIS data;
+ # user must set up a .netrc file in their home directory (~/.netrc) with their credentials on one line, e.g.:
+ # machine urs.earthdata.nasa.gov login password
+ # and set the file permissions to user read/write only (chmod 600 ~/.netrc) to protect their credentials
+
+ dir_path = Path(dir_path)
+ output_dir = dir_path # GeoTIFFs are written to the same directory as the downloaded HDF files
+
+ # Validate and unpack ll_limits as (min_lat, max_lat, min_lon, max_lon).
+ # Earthaccess uses a different ordering for bounding_box (lon, lat, lon, lat).
+ if ll_limits is not None:
+ min_lat, max_lat, min_lon, max_lon = ll_limits
+ if not (-90.0 <= float(min_lat) <= 90.0 and -90.0 <= float(max_lat) <= 90.0):
+ raise ValueError(
+ f"read_modis_ll_to_geotiff: invalid latitude bounds in ll_limits {ll_limits}. "
+ "Expected (min_lat, max_lat, min_lon, max_lon)."
+ )
+ if not (-180.0 <= float(min_lon) <= 180.0 and -180.0 <= float(max_lon) <= 180.0):
+ raise ValueError(
+ f"read_modis_ll_to_geotiff: invalid longitude bounds in ll_limits {ll_limits}. "
+ "Expected (min_lat, max_lat, min_lon, max_lon)."
+ )
+ if float(min_lat) > float(max_lat) or float(min_lon) > float(max_lon):
+ raise ValueError(
+ f"read_modis_ll_to_geotiff: non-monotonic ll_limits {ll_limits}. "
+ "Expected min <= max for both lat and lon."
+ )
+ else:
+ min_lat = max_lat = min_lon = max_lon = None
+
+ # latest_date and earliest_date are in 'YYYY-MM-DD' format
+ latest_date = f"{year}-12-31"
+ earliest_date = f"{year}-01-01"
+
+ # Parse product short name and version from 'MCD12Q1.061' -> ('MCD12Q1', '061')
+ product_parts = product.split('.')
+ short_name = product_parts[0]
+ version = product_parts[1] if len(product_parts) > 1 else '061'
+
+ # Download HDF granules via NASA Earthdata Cloud (earthaccess).
+ # The old LP DAAC HTTPS server (e4ftl01.cr.usgs.gov) decommissioned MODIS
+ # directories in 2024; earthaccess is the current NASA-recommended approach.
+ # Credentials are read from ~/.netrc (machine urs.earthdata.nasa.gov).
+ logger.info(
+ f"read_modis_ll_to_geotiff: searching Earthdata for {product} year={year} "
+ f"ll_limits(min_lat,max_lat,min_lon,max_lon)={ll_limits}"
+ )
+
+ try:
+ earthaccess.login(strategy='netrc')
+ except Exception as e:
+ raise RuntimeError(
+ f"read_modis_ll_to_geotiff: Earthdata login failed. "
+ f"Ensure ~/.netrc contains credentials for urs.earthdata.nasa.gov: {e}"
+ ) from e
+
+ search_kwargs = {
+ 'short_name': short_name,
+ 'version': version,
+ 'temporal': (earliest_date, latest_date),
+ }
+ if ll_limits is not None:
+ # earthaccess bounding_box: (min_lon, min_lat, max_lon, max_lat)
+ # Round to 6 decimal places to avoid scientific notation (e.g. 3.68e-14)
+ # which the CMR API rejects as an invalid bounding_box value.
+ search_kwargs['bounding_box'] = (
+ round(float(min_lon), 6),
+ round(float(min_lat), 6),
+ round(float(max_lon), 6),
+ round(float(max_lat), 6),
+ )
+ logger.info(
+ "read_modis_ll_to_geotiff: Earthaccess bounding_box(min_lon,min_lat,max_lon,max_lat)="
+ f"{search_kwargs['bounding_box']}"
+ )
+
+ granules = _earthaccess_search_with_retry(**search_kwargs)
+
+ # CMR bbox filtering can return zero granules for very small boxes even when
+ # intersecting MODIS tiles exist. Retry by tile tags as a robust fallback.
+ if not granules and ll_limits is not None:
+ target_tiles = {
+ t.strip().lower()
+ for t in _get_modis_tile_idxs_ll(ll_limits).split(',')
+ if t.strip()
+ }
+ logger.warning(
+ "read_modis_ll_to_geotiff: no granules from bbox search; retrying tile-filtered search "
+ f"for tiles={sorted(target_tiles)}, {product}, "
+ f"bounding_box(min_lon,min_lat,max_lon,max_lat)={search_kwargs['bounding_box']}"
)
+ broad_granules = _earthaccess_search_with_retry(
+ short_name=short_name,
+ version=version,
+ temporal=(earliest_date, latest_date),
+ )
+ granules = [
+ g for g in broad_granules
+ if (_granule_hv_tag(g) in target_tiles)
+ ]
+ logger.info(
+ f"read_modis_ll_to_geotiff: tile-filtered search found, {product}, "
+ f"{len(granules)} granule(s) from {len(broad_granules)} broad candidates"
+ f"bounding_box(min_lon,min_lat,max_lon,max_lat)={search_kwargs['bounding_box']}"
+ )
+
+ if not granules:
+ logger.warning(
+ f"read_modis_ll_to_geotiff: no granules found for {product} "
+ f"temporal={earliest_date}..{latest_date} bounding_box={search_kwargs.get('bounding_box')} "
+ f"ll_limits={ll_limits}; skipping this chunk for {product}"
+ )
+ return {}
+ logger.info(f"read_modis_ll_to_geotiff: found {len(granules)} granules for {product}; downloading to {dir_path}")
+
+ try:
+ downloaded = earthaccess.download(granules, local_path=str(dir_path))
+ except Exception as e:
+ raise RuntimeError(
+ f"read_modis_ll_to_geotiff: MODIS download failed for product '{product}', year {year}: {e}"
+ ) from e
+
+ # Collect downloaded HDF files for this specific product/version only.
+ # The same temp directory is reused across product calls, so we must not
+ # mix, for example, MCD12Q1 tiles into a MOD44B read.
+ hdf_files = []
+ for f in sorted(dir_path.glob('*.[hH][dD][fF]')):
+ f_name_upper = f.name.upper()
+ if not f_name_upper.startswith(f"{short_name.upper()}."):
+ continue
+ if len(product_parts) > 1 and f".{version.upper()}." not in f_name_upper:
+ continue
+ hdf_files.append(str(f))
+ if not hdf_files:
+ raise RuntimeError(
+ f"read_modis_ll_to_geotiff: no HDF files for {product} found in {dir_path} after download"
+ )
+ logger.info(
+ f"read_modis_ll_to_geotiff: {len(hdf_files)} HDF tile(s) for {product} to process"
+ )
+
+ # Build per-variable GeoTIFFs using GDAL directly
+ # gdal.BuildVRT mosaics the per-tile subdatasets; gdal.Warp reprojects and
+ # clips to ll_limits. This avoids pymodis's requirement for .hdf.xml sidecars.
- grazing_path = Path(grazing_path)
out = {}
+ for var in variable_names:
+ var_norm = ''.join(ch for ch in var.lower() if ch.isalnum())
+ # Find the matching HDF4_EOS subdataset path in each tile
+ sds_paths = []
+ for hdf_file in hdf_files:
+ ds_hdf = gdal.Open(hdf_file)
+ if ds_hdf is None:
+ raise RuntimeError(
+ f"read_modis_ll_to_geotiff: GDAL could not open {hdf_file}"
+ )
+ matched = None
+ for sds_name, sds_desc in ds_hdf.GetSubDatasets():
+ sds_name_norm = ''.join(ch for ch in sds_name.lower() if ch.isalnum())
+ sds_desc_norm = ''.join(ch for ch in sds_desc.lower() if ch.isalnum())
+ if var_norm in sds_name_norm or var_norm in sds_desc_norm:
+ matched = sds_name
+ break
+ ds_hdf = None # close
+ if matched is None:
+ raise RuntimeError(
+ f"read_modis_ll_to_geotiff: variable '{var}' not found in {hdf_file}"
+ )
+ sds_paths.append(matched)
+
+ output_tif = output_dir / f"{var}_{year}.tif"
+ warp_kwargs = dict(
+ format='GTiff',
+ dstSRS='EPSG:4326',
+ resampleAlg=gdal.GRA_NearestNeighbour,
+ )
+ if ll_limits is not None:
+ min_lat, max_lat, min_lon, max_lon = ll_limits
+ warp_kwargs['outputBounds'] = (min_lon, min_lat, max_lon, max_lat)
+ warp_kwargs['outputBoundsSRS'] = 'EPSG:4326'
+
+ if len(sds_paths) == 1:
+ # Single tile — Warp directly from the subdataset path
+ gdal.Warp(str(output_tif), sds_paths[0], **warp_kwargs)
+ else:
+ # Multiple tiles — build an in-memory VRT mosaic, then Warp
+ vrt_path = str(output_dir / f"{var}_{year}_mosaic.vrt")
+ vrt = gdal.BuildVRT(vrt_path, sds_paths)
+ vrt.FlushCache()
+ vrt = None # close before Warp reads it
+ try:
+ gdal.Warp(str(output_tif), vrt_path, **warp_kwargs)
+ finally:
+ Path(vrt_path).unlink(missing_ok=True)
+
+ if not output_tif.exists():
+ raise RuntimeError(
+ f"read_modis_ll_to_geotiff: GeoTIFF not created for {product} '{var}' at {output_tif}"
+ )
+ logger.info(f"read_modis_ll_to_geotiff: wrote {product} {output_tif}")
+ out[var] = output_tif
+
+ return out
+
+
+
+
+
+###todo: finish and test read_netcdf general function
+
+
+#--------------------------------------------------------------------------
+def read_netcdf_ll(year, file_path_name, variable_names=None, ll_limits=None):
+ """
+ Read variables from a NetCDF file for a given year.
+
+ Args:
+ year (int): Year to extract.
+ file_path_name (str or Path): Full path to the NetCDF file.
+ variable_names (list or None): Variables to extract. Must be provided.
+ ll_limits (tuple/list or None): 4-element (min_lat, max_lat, min_lon, max_lon).
+ When given, only the lat/lon rows/columns that
+ overlap this region are read. The slice is
+ inclusive — any grid cell whose coordinate falls
+ within [min_lat, max_lat] x [min_lon, max_lon]
+ is included, so cells that straddle the boundary
+ are never dropped.
+
+ Returns:
+ dict: {varname: 2D np.ndarray shape (lat, lon)} for the requested year,
+ plus 'lat' and 'lon' coordinate arrays (possibly subsetted).
+ """
+
+ ncfile = Path(file_path_name)
+ if not ncfile.exists():
+ raise FileNotFoundError(f"NetCDF file not found: {ncfile}")
+
+ ds = xr.open_dataset(ncfile, decode_times=False)
+
+ if variable_names is None:
+ # Raise an error
+ # consider reading all variables in file instead of raising an error?
+ raise KeyError(f"read_netcdf_ll: Variable names must be provided in the json input file for {ncfile}. "
+ f"Available variables: {list(ds.data_vars)}")
+
+ time_units = ds['time'].attrs.get('units', None)
+ # _get_year_idx() handles both 'years since' and 'days since' patterns,
+ # as well as the case of no time units (assumed calendar years)
+ # add cases to _get_year_idx as needed if other time unit patterns are encountered in source data
+ year_idx = _get_year_idx(ds['time'].values, year, ncfile, time_units=time_units)
+ logger.info(f"read_netcdf_ll: reading year {year} (time index {year_idx}) from {ncfile}")
+
+ ####todo: this currently assumes that the variables lat/lon are ~cell centers
+ # need to allow for other variable names to represent these coordinates
+
+ # if ll_limits=None, do not subset the data; otherwise subset by ll_limits
+ if ll_limits is not None:
+ min_lat, max_lat, min_lon, max_lon = ll_limits
+ lat_vals = ds['lat'].values
+ lon_vals = ds['lon'].values
+
+ # add a one-cell buffer so cells that straddle the boundary are included
+ lat_step = float(abs(lat_vals[1] - lat_vals[0])) if lat_vals.size > 1 else 0.0
+ lon_step = float(abs(lon_vals[1] - lon_vals[0])) if lon_vals.size > 1 else 0.0
+
+ lat_mask = (lat_vals >= min_lat - lat_step) & (lat_vals <= max_lat + lat_step)
+ lon_mask = (lon_vals >= min_lon - lon_step) & (lon_vals <= max_lon + lon_step)
- first = True
- for category, filename in grazing_names.items():
- ncfile = grazing_path / filename
- if not ncfile.exists():
- raise FileNotFoundError(f"HYDE3.5 grazing file not found: {ncfile}")
- # variable name is the file stem, e.g. 'pasture.nc' -> 'pasture'
- varname = Path(filename).stem
- ds = xr.open_dataset(ncfile, decode_times=False)
- time_units = ds['time'].attrs.get('units', None)
- year_idx = _get_year_idx(ds['time'].values, year, ncfile, time_units=time_units)
- print(f" read_hyde_grazing: reading '{varname}' year {year} "
- f"(time index {year_idx}) from {ncfile}")
- if first:
- out['lat'] = ds['lat'].values
- out['lon'] = ds['lon'].values
- first = False
- out[category] = ds[varname].isel(time=year_idx).values # shape: (lat, lon)
- ds.close()
+ lat_idx = np.where(lat_mask)[0]
+ lon_idx = np.where(lon_mask)[0]
+ if lat_idx.size == 0 or lon_idx.size == 0:
+ raise ValueError(
+ f"read_netcdf_ll: No grid cells found within ll_limits {ll_limits} in {ncfile}. "
+ f"lat range: [{lat_vals.min():.2f}, {lat_vals.max():.2f}], "
+ f"lon range: [{lon_vals.min():.2f}, {lon_vals.max():.2f}]"
+ )
+
+ lat_dim = ds['lat'].dims[0]
+ lon_dim = ds['lon'].dims[0]
+ ds = ds.isel({lat_dim: lat_idx, lon_dim: lon_idx})
+
+ out = {'lat': ds['lat'].values, 'lon': ds['lon'].values}
+ for v in variable_names:
+ if v not in ds:
+ raise KeyError(f"read_netcdf_ll: Variable '{v}' not found in {ncfile}. "
+ f"Available variables: {list(ds.data_vars)}")
+ out[v] = ds[v].isel(time=year_idx).values # shape: (lat, lon)
+
+ ds.close()
return out
#--------------------------------------------------------------------------
@@ -200,15 +694,21 @@ def write_latlon_to_geotiff(data_2d, lat, lon, ll_limits, tmp_path):
# add a 1-cell buffer on each side to avoid edge interpolation artefacts
lat_step = float(lat[1] - lat[0])
lon_step = float(lon[1] - lon[0])
- lat_mask = (lat >= min_lat - abs(lat_step)) & (lat <= max_lat + abs(lat_step))
- lon_mask = (lon >= min_lon - abs(lon_step)) & (lon <= max_lon + abs(lon_step))
+ # 1-cell buffer for floating-point safety — ensures the outermost raster
+ # pixels are not accidentally clipped by strict inequality comparisons.
+ # Chunk boundary artifacts are avoided upstream by using tight vertex
+ # bounding boxes (from set_decomp_cell_idx_ll_limits) as ll_limits, so a
+ # large buffer is not needed here.
+ buffer_cells = 1
+ lat_mask = (lat >= min_lat - buffer_cells * abs(lat_step)) & (lat <= max_lat + buffer_cells * abs(lat_step))
+ lon_mask = (lon >= min_lon - buffer_cells * abs(lon_step)) & (lon <= max_lon + buffer_cells * abs(lon_step))
lat_idx = np.where(lat_mask)[0]
lon_idx = np.where(lon_mask)[0]
if lat_idx.size == 0 or lon_idx.size == 0:
raise ValueError(
- f"No source grid cells found within ll_limits {ll_limits}. "
+ f"write_latlon_to_geotiff: No source grid cells found within ll_limits {ll_limits}. "
f"lat range: [{lat.min():.2f}, {lat.max():.2f}], "
f"lon range: [{lon.min():.2f}, {lon.max():.2f}]"
)
@@ -219,10 +719,12 @@ def write_latlon_to_geotiff(data_2d, lat, lon, ll_limits, tmp_path):
chunk_data = data_2d[np.ix_(lat_idx, lon_idx)].astype(np.float32)
# rasterio uses (west, south, east, north) bounds
- west = float(chunk_lon[0]) - abs(lon_step) / 2.0
- east = float(chunk_lon[-1]) + abs(lon_step) / 2.0
- south = float(chunk_lat[0]) - abs(lat_step) / 2.0
- north = float(chunk_lat[-1]) + abs(lat_step) / 2.0
+ # Use min/max rather than first/last element so this is correct for both
+ # south-to-north (e.g. some grids) and north-to-south (HYDE3.5, LUH2) lat arrays.
+ west = float(chunk_lon.min()) - abs(lon_step) / 2.0
+ east = float(chunk_lon.max()) + abs(lon_step) / 2.0
+ south = float(chunk_lat.min()) - abs(lat_step) / 2.0
+ north = float(chunk_lat.max()) + abs(lat_step) / 2.0
n_rows, n_cols = chunk_data.shape
transform = from_bounds(west, south, east, north, n_cols, n_rows)
@@ -239,56 +741,64 @@ def write_latlon_to_geotiff(data_2d, lat, lon, ll_limits, tmp_path):
crs='EPSG:4326',
transform=transform,
nodata=np.nan,
- ) as dst:
+ ) as dst:
# rasterio band 1 is row-ordered north-to-south; flip if lat is south-to-north
if chunk_lat[0] < chunk_lat[-1]:
dst.write(np.flipud(chunk_data), 1)
else:
dst.write(chunk_data, 1)
- print(f" write_latlon_to_geotiff: wrote {n_rows}x{n_cols} chunk to {tmp_path}")
+ logger.info(f"write_latlon_to_geotiff: wrote {n_rows}x{n_cols} chunk to {tmp_path}")
return tmp_path
#--------------------------------------------------------------------------
-def write_chunk_mesh_to_geojson(global_mesh_df, cell_ids, tmp_path):
+def write_mesh_to_geojson(out_grid_data, tmp_path, cell_indices=None):
"""
- Filter the global HEALPix mesh DataFrame to only the cells in cell_ids
- and write a chunk-sized GeoJSON file for use as uraster source mesh.
+ Write mesh cells from a GridData object to a GeoJSON file.
- GeoJSON is used (rather than Parquet) because it is always supported by
- GDAL/OGR without additional plugins (unlike the Parquet/DuckDB driver).
-
- The caller should load the global parquet once (e.g. in run()) and pass
- the resulting DataFrame here to avoid re-reading the 37 MB file for every
- variable and chunk.
+ Builds polygon geometries directly from the vertex coordinate arrays
+ (lon_vtx, lat_vtx) stored in out_grid_data
Args:
- global_mesh_df (pd.DataFrame): Full merged_land_cells DataFrame, already
- loaded. Must have 'cellid' (int) and
- 'geometry' (WKB bytes) columns.
- cell_ids (array-like): 1D array of integer cellid values for chunk.
- tmp_path (str|Path): Full path of the output GeoJSON file to write.
+ out_grid_data (GridData): Populated grid geometry object whose arrays
+ have been filled (e.g. via load_mesh_nc).
+ tmp_path (str|Path): Full path of the output GeoJSON file to write.
+ Parent directories are created if needed.
+ cell_indices (tuple|array-like|None): 0-based indices into the
+ out_grid_data arrays selecting which cells to
+ write. Pass None to write all cells.
Returns:
- Path: Path to the written GeoJSON.
+ Path: Path to the written GeoJSON file.
"""
- chunk_df = global_mesh_df[global_mesh_df['cellid'].isin(cell_ids)]
+ if cell_indices is None:
+ idx = np.arange(out_grid_data.num_cells, dtype=np.intp)
+ else:
+ idx = np.asarray(cell_indices, dtype=np.intp)
- if chunk_df.empty:
- raise ValueError(
- f"No mesh cells found for the provided cell_ids. "
- f"First few cell_ids: {list(cell_ids[:5])}"
- )
+ if idx.size == 0:
+ raise ValueError("write_mesh_to_geojson: cell_indices is empty — no cells to write")
+
+ lon_vtx = out_grid_data.lon_vtx[idx] # (n, n_vertices)
+ lat_vtx = out_grid_data.lat_vtx[idx] # (n, n_vertices)
- # build GeoJSON manually from WKB geometry column
features = []
- for _, row in chunk_df.iterrows():
- geom = shapely.wkb.loads(row['geometry'])
+ for i in range(len(idx)):
+ # GeoJSON polygon ring: list of [lon, lat] pairs, closed (first == last)
+ # it is more efficient to store the cell indices than the cell ids
+ # but i think the property still need to be called 'cellid' for uraster to recognize it
+ ring = [[float(lon_vtx[i, v]), float(lat_vtx[i, v])]
+ for v in range(lon_vtx.shape[1])]
+ ring.append(ring[0]) # close the ring
+
features.append({
'type': 'Feature',
- 'geometry': geom.__geo_interface__,
- 'properties': {'cellid': int(row['cellid'])},
+ 'geometry': {
+ 'type': 'Polygon',
+ 'coordinates': [ring],
+ },
+ 'properties': {'cellid': int(idx[i])},
})
geojson = {'type': 'FeatureCollection', 'features': features}
@@ -297,108 +807,239 @@ def write_chunk_mesh_to_geojson(global_mesh_df, cell_ids, tmp_path):
with open(tmp_path, 'w') as f:
json.dump(geojson, f)
- print(f" write_chunk_mesh_to_geojson: wrote {len(features)} cells to {tmp_path}")
+ logger.info(f"write_mesh_to_geojson: wrote {len(features)} cells to {tmp_path}")
return tmp_path
-
#--------------------------------------------------------------------------
-def regrid_to_landgen_grid(data_2d, src_lat, src_lon, cell_ids, ll_limits,
- global_mesh_df, tmp_dir, varname,
- remap_method=3):
+def regrid_to_mesh(mesh_file_path, var_file_path, cell_indices, out_grid_data,
+ out_type='path', remap_method=3):
"""
- Regrid a single 2D source variable onto the landgen HEALPix grid cells
- for one spatial chunk, using uraster.
-
- Workflow:
- 1. Slice source data to ll_limits and write to a temp GeoTIFF.
- 2. Filter global_mesh_df to cell_ids and write a chunk parquet.
- 3. Run uraster with iFlag_remap_method=3 (weighted average).
- 4. Read the uraster output GeoJSON, extract 'mean' per cellid.
- 5. Return a 1D array aligned to cell_ids order; cells with no overlap get 0.
- 6. Clean up all temp files.
+ Regrid a single raster variable onto the landgen mesh cells for one chunk,
+ using uraster. Inputs are already-written files (mesh GeoJSON + source
+ raster GeoTIFF), so no temp-file management is needed here.
Args:
- data_2d (np.ndarray): 2D source array shape (n_lat, n_lon).
- src_lat (np.ndarray): 1D source latitude array.
- src_lon (np.ndarray): 1D source longitude array.
- cell_ids (array-like): 1D array of integer cellid values for this chunk.
- ll_limits (tuple): (min_lat, max_lat, min_lon, max_lon).
- global_mesh_df (pd.DataFrame): Pre-loaded merged_land_cells DataFrame.
- Load once in run() and pass here to avoid
- re-reading the 37 MB parquet for every variable.
- tmp_dir (str|Path): Directory for temporary files (unique per worker).
- varname (str): Variable name, used for temp file naming only.
- remap_method (int): uraster iFlag_remap_method
- (1=nearest, 2=nearest, 3=weighted average).
- Default 3 is correct for area fraction data.
+ mesh_file_path (str|Path): Full path to the chunk mesh GeoJSON written
+ by write_mesh_to_geojson.
+ var_file_path (dict): One-element dict {'': Path} pointing to
+ the source raster GeoTIFF for this variable.
+ cell_indices (tuple): 0-based indices into out_grid_data arrays for
+ the cells in this chunk.
+ out_grid_data (GridData): Grid geometry object; its cell_id array is used
+ to map cellid values to output order when
+ out_type='data'.
+ out_type (str): 'path' — return Path of the uraster output GeoJSON.
+ 'data' — return 1D np.ndarray of mean values
+ aligned to cell_indices order.
+ remap_method (int): uraster iFlag_remap_method
+ (1=nearest, 2=nearest, 3=weighted average).
+ Default 3 is correct for area-fraction data.
Returns:
- np.ndarray: 1D float64 array of length len(cell_ids), regridded values
- in the same order as cell_ids.
+ Path if out_type='path': Path to the uraster output GeoJSON.
+ np.ndarray if out_type='data': 1D float64 array of length
+ len(cell_indices), regridded values in the same order
+ as cell_indices.
"""
- tmp_dir = Path(tmp_dir)
- tmp_dir.mkdir(parents=True, exist_ok=True)
-
- # unique suffix to avoid collisions between parallel workers
- suffix = f"{varname}_{ll_limits[0]:.0f}_{ll_limits[2]:.0f}"
- tmp_raster = tmp_dir / f"src_{suffix}.tif"
- tmp_mesh_in = tmp_dir / f"mesh_in_{suffix}.geojson" # GeoJSON: universal GDAL support
- tmp_mesh_out= tmp_dir / f"mesh_out_{suffix}.geojson"
-
+ if out_type not in ('path', 'data'):
+ raise ValueError(f"regrid_to_mesh: out_type must be 'path' or 'data', got '{out_type}'")
+
+ mesh_path = Path(mesh_file_path)
+ varname = next(iter(var_file_path.keys()))
+ raster_path = Path(next(iter(var_file_path.values())))
+
+ if not mesh_path.exists():
+ raise FileNotFoundError(f"regrid_to_mesh: mesh file not found: {mesh_path}")
+ if not raster_path.exists():
+ raise FileNotFoundError(f"regrid_to_mesh: raster file not found: {raster_path}")
+
+ # output file sits next to the mesh file
+ out_path = mesh_path.parent / f"{varname}.geojson"
+
+ config = {
+ 'sFilename_source_mesh': str(mesh_path),
+ 'aFilename_source_raster': [str(raster_path)],
+ 'sFilename_target_mesh': str(out_path),
+ 'iFlag_remap_method': remap_method,
+ 'sField_unique_id': 'cellid',
+ 'iFlag_global': 0,
+ 'iFlag_polar': 0,
+ }
try:
- # 1. write source chunk as GeoTIFF
- write_latlon_to_geotiff(data_2d, src_lat, src_lon, ll_limits, tmp_raster)
-
- # 2. write filtered chunk mesh as GeoJSON (parquet requires libgdal-arrow-parquet)
- write_chunk_mesh_to_geojson(global_mesh_df, cell_ids, tmp_mesh_in)
-
- # 3. run uraster
- config = {
- 'sFilename_source_mesh': str(tmp_mesh_in),
- 'aFilename_source_raster': [str(tmp_raster)],
- 'sFilename_target_mesh': str(tmp_mesh_out),
- 'iFlag_remap_method': remap_method,
- 'sField_unique_id': 'cellid',
- 'iFlag_global': 0, # chunk is regional, not global
- 'iFlag_polar': 0,
- }
processor = URaster(config)
processor.setup()
processor.run_remap()
+ except Exception as e:
+ raise RuntimeError(
+ f"regrid_to_mesh: uraster failed for '{varname}' "
+ f"(mesh={mesh_path}, raster={raster_path}): {e}"
+ ) from e
+
+ if not out_path.exists():
+ raise RuntimeError(
+ f"regrid_to_mesh: uraster completed but output not found: {out_path}")
+
+ if out_type == 'path':
+ return out_path
+
+ # out_type == 'data': read output GeoJSON and align values to cell_indices order
+ with open(out_path, 'r') as f:
+ geojson = json.load(f)
+
+ result_map = {}
+ for feature in geojson['features']:
+ props = feature['properties']
+ cid = int(props['cellid'])
+ val = props.get('mean', None)
+ result_map[cid] = float(val) if (val is not None and not np.isnan(val)) else 0.0
+
+ # cellid in the GeoJSON equals the row index (written by write_mesh_to_geojson)
+ # this is more efficient than looking up cell id values
+ out = np.array([result_map.get(int(idx), 0.0) for idx in np.asarray(cell_indices, dtype=np.intp)],
+ dtype=np.float64)
+
+ logger.info(f"regrid_to_mesh: '{varname}' -> {len(out)} cells, "
+ f"non-zero: {np.count_nonzero(out)}")
+ return out
- # 4. read output GeoJSON and extract 'mean' per cellid
- with open(tmp_mesh_out, 'r') as f:
- geojson = json.load(f)
-
- # build a cellid -> mean value lookup from the uraster output
- # cells with no raster overlap are missing the 'mean' key entirely;
- # default those to 0.0
- result_map = {}
- for feature in geojson['features']:
- props = feature['properties']
- cid = int(props['cellid'])
- val = props.get('mean', None)
- result_map[cid] = float(val) if (val is not None and not np.isnan(val)) else 0.0
-
- # 5. align to cell_ids order; missing cells default to 0
- out = np.array([result_map.get(int(cid), 0.0) for cid in cell_ids],
- dtype=np.float64)
-
- print(f" regrid_to_landgen_grid: '{varname}' -> {len(out)} cells, "
- f"non-zero: {np.count_nonzero(out)}")
- return out
-
- finally:
- # 6. clean up temp files regardless of success or failure
- for f in [tmp_raster, tmp_mesh_in, tmp_mesh_out]:
- try:
- Path(f).unlink(missing_ok=True)
- except Exception:
- pass
- # uraster also creates a '_fixed' copy of the input mesh; clean that up too
- try:
- Path(str(tmp_mesh_in).replace('.geojson', '_fixed.geojson')).unlink(missing_ok=True)
- except Exception:
- pass
-#--------------------------------------------------------------------------
\ No newline at end of file
+##### todo: need to set units an other attributres on the output variables here; currently just copying the raw values without metadata
+
+#--------------------------------------------------------------------------
+def write_module_netcdf(out_grid_data, module_data, out_path, file_name,
+ year=None, timevars=None, varnames=None, ll_limits=None):
+ """
+ Write a NetCDF file containing grid geometry from out_grid_data plus
+ selected variables from one module data object (TopoData, LtData, etc.).
+
+ Grid variables written (from out_grid_data):
+ cell_id (coordinate), cell_area, landfrac, lon_cen, lat_cen,
+ lon_vtx, lat_vtx
+
+ Module variables written:
+ all non-None array attributes of module_data, or only those listed in
+ varnames if provided. cell_idx is always skipped (internal bookkeeping).
+ Variables in timevars get a leading unlimited 'time' dimension so that
+ per-year files can be concatenated later with ncrcat or
+ xarray.open_mfdataset(..., concat_dim='time').
+
+ Spatial subsetting:
+ If ll_limits=(min_lat, max_lat, min_lon, max_lon) is given, only the
+ cells whose centre coordinates (out_grid_data.lat_cen / lon_cen) fall
+ within the bounding box are written. Otherwise all cells are written.
+
+ Args:
+ out_grid_data (GridData): Populated grid geometry object.
+ module_data (object): Any per-cell data object (TopoData, LtData,
+ …) — NOT a BaseManager-derived class.
+ out_path (str|Path): Directory in which to write the file.
+ file_name (str): NetCDF filename (including extension).
+ year (int|None): Calendar year for this record. When given,
+ a 'time' coordinate with value=year is added
+ and all timevars receive a leading time dim.
+ timevars (list[str]|None): Variable names that receive a leading
+ unlimited 'time' dimension. Ignored when
+ year=None. None → no variables get a time dim.
+ varnames (list[str]|None): Variables to write from module_data.
+ None → write all non-None array attributes.
+ ll_limits (tuple|None): (min_lat, max_lat, min_lon, max_lon) spatial
+ subset. None → write all cells.
+
+ Returns:
+ Path: Path to the written NetCDF file.
+ """
+ out_path = Path(out_path)
+ out_path.mkdir(parents=True, exist_ok=True)
+ nc_path = out_path / file_name
+
+ n_cells = out_grid_data.num_cells
+ time_set = set(timevars) if (timevars and year is not None) else set()
+
+ # --- determine which cells to write ---
+ if ll_limits is not None:
+ min_lat, max_lat, min_lon, max_lon = ll_limits
+ mask = (
+ (out_grid_data.lat_cen >= min_lat) & (out_grid_data.lat_cen <= max_lat) &
+ (out_grid_data.lon_cen >= min_lon) & (out_grid_data.lon_cen <= max_lon)
+ )
+ idx = np.where(mask)[0]
+ else:
+ idx = np.arange(n_cells, dtype=np.intp)
+
+ n_out = len(idx)
+ if n_out == 0:
+ raise ValueError(f"write_module_netcdf: no cells found within ll_limits {ll_limits}")
+
+ # --- collect variables to write from module_data ---
+ _skip = {'lock', 'cell_idx'}
+ if varnames is None:
+ varnames_to_write = [
+ k for k, v in vars(module_data).items()
+ if k not in _skip and isinstance(v, np.ndarray)
+ ]
+ else:
+ varnames_to_write = list(varnames)
+
+ # --- build xarray Dataset ---
+ cell_dim = 'cell'
+ vtx_dim = 'vertex'
+ time_dim = 'time'
+
+ # time coordinate (size-1 so the unlimited dim exists for concatenation)
+ coords = {
+ 'cell_id': xr.DataArray(out_grid_data.cell_id[idx], dims=[cell_dim],
+ attrs={'long_name': 'cell identifier'}),
+ }
+ if year is not None:
+ coords['time'] = xr.DataArray(
+ np.array([year], dtype=np.int32), dims=[time_dim],
+ attrs={'long_name': 'year', 'units': 'calendar year'}
+ )
+
+ data_vars = {
+ 'cell_area': xr.DataArray(out_grid_data.cell_area[idx], dims=[cell_dim],
+ attrs={'long_name': 'cell area'}),
+ 'landfrac': xr.DataArray(out_grid_data.landfrac[idx], dims=[cell_dim],
+ attrs={'long_name': 'land fraction'}),
+ 'lon_cen': xr.DataArray(out_grid_data.lon_cen[idx], dims=[cell_dim],
+ attrs={'long_name': 'cell-centre longitude', 'units': 'degrees_east'}),
+ 'lat_cen': xr.DataArray(out_grid_data.lat_cen[idx], dims=[cell_dim],
+ attrs={'long_name': 'cell-centre latitude', 'units': 'degrees_north'}),
+ 'lon_vtx': xr.DataArray(out_grid_data.lon_vtx[idx], dims=[cell_dim, vtx_dim],
+ attrs={'long_name': 'vertex longitudes', 'units': 'degrees_east'}),
+ 'lat_vtx': xr.DataArray(out_grid_data.lat_vtx[idx], dims=[cell_dim, vtx_dim],
+ attrs={'long_name': 'vertex latitudes', 'units': 'degrees_north'}),
+ }
+
+ # add module variables
+ for name in varnames_to_write:
+ val = getattr(module_data, name, None)
+ if val is None or not isinstance(val, np.ndarray):
+ continue
+ arr = val[idx] if val.shape[0] == n_cells else val
+
+ # build spatial dimension names (no time prefix yet)
+ if arr.ndim == 1:
+ dims = [cell_dim]
+ elif arr.ndim == 2:
+ dims = [cell_dim, f'{name}_dim1']
+ elif arr.ndim == 3:
+ dims = [cell_dim, f'{name}_dim1', f'{name}_dim2']
+ else:
+ dims = [cell_dim] + [f'{name}_dim{i}' for i in range(1, arr.ndim)]
+
+ if name in time_set:
+ # prepend unlimited time dim with a size-1 axis
+ arr = arr[np.newaxis, ...] # (1, cell, ...)
+ dims = [time_dim] + dims
+ data_vars[name] = xr.DataArray(arr, dims=dims)
+
+ ds = xr.Dataset(data_vars=data_vars, coords=coords)
+ # declare 'time' as the unlimited dimension so ncrcat / concat works
+ encoding = {time_dim: {'unlimited': True}} if year is not None else {}
+ ds.to_netcdf(nc_path, unlimited_dims=[time_dim] if year is not None else [])
+
+ logger.info(f"write_module_netcdf: wrote {n_out} cells, "
+ f"{len(varnames_to_write)} module vars to {nc_path}"
+ + (f" (year {year}, {len(time_set)} timevars)" if year is not None else ""))
+ return nc_path
diff --git a/components/elm/tools/landgen/src/landgen/management.py b/components/elm/tools/landgen/src/landgen/management.py
new file mode 100644
index 000000000000..4e6725799001
--- /dev/null
+++ b/components/elm/tools/landgen/src/landgen/management.py
@@ -0,0 +1,227 @@
+# management.py
+# this module processes harvest and grazing data for the landgen workflow
+
+# run() function is the main entry point for this module, and will be called by process_single_year in land_type.py
+
+import multiprocessing as mp
+import logging
+from pathlib import Path
+from .shared_data import LtData
+from . import landgen_io
+from . import tools
+# import normalize_cell # not created yet
+import numpy as np
+import os
+import traceback
+import time
+import shutil
+import tempfile
+
+logger = logging.getLogger('landgen')
+
+def _management_process_star(args):
+ """Module-level wrapper so imap_unordered can unpack the args tuple."""
+ return management_process(*args)
+
+########## define some module-specific constants here
+
+# Default harvest variable names from LUH2 transitions.nc
+LUH2_HARVEST_VARS = [
+ 'primf_harv', # wood harvest area from primary forest land
+ 'primn_harv', # wood harvest area from primary non forest land
+ 'secmf_harv', # wood harvest area from secondary mature forest land
+ 'secyf_harv', # wood harvest area from secondary young forest land
+ 'secnf_harv', # wood harvest area from secondary non forest land
+ 'primf_bioh', # wood harvest biomass carbon from primary forest land
+ 'primn_bioh', # wood harvest biomass carbon from primary non forest land
+ 'secmf_bioh', # wood harvest biomass carbon from secondary mature forest land
+ 'secyf_bioh', # wood harvest biomass carbon from secondary young forest land
+ 'secnf_bioh', # wood harvest biomass carbon from secondary non forest land
+]
+
+########## define helper functions for management run() here
+
+##### management_process()
+
+## arguments
+# lc_data: land cover data structure that is passed between modules
+# year: the year for which to process the land cover data
+# source_data_path: base path to the source data
+# landgen_grid_path: path from source_data_path and the filename of the landgen grid
+# out_path: base path for the output data; this is needed to read in the previous year's management data
+
+## output
+
+def management_process(year, harvest_path, harvest_name, grazing_path, grazing_names,
+ com_config_dict, out_grid_data, ll_limits, row_indices):
+ """Compute regridded harvest/grazing for one spatial chunk.
+ Each worker reads its own source data (simple starmap approach like landcover.py).
+ Returns chunk LtData object with cell_idx, harvest_frac, and grazing_frac populated.
+ """
+ t0 = time.time()
+ try:
+ return _management_process_impl(
+ year, harvest_path, harvest_name, grazing_path, grazing_names,
+ com_config_dict, ll_limits, row_indices, out_grid_data
+ )
+ except Exception:
+ print(f"ERROR in management_process chunk {ll_limits} year {year}:\n{traceback.format_exc()}", flush=True)
+ raise
+ finally:
+ elapsed = time.time() - t0
+ print(f" chunk {ll_limits} year {year}: {elapsed:.1f}s", flush=True)
+
+def _management_process_impl(year, harvest_path, harvest_name, grazing_path, grazing_names,
+ com_config_dict, ll_limits, row_indices, out_grid_data):
+ """Worker implementation: reads source data, regrids using modular workflow, returns chunk LtData.
+ Each worker does its own I/O (simple starmap approach like landcover.py).
+ Uses same workflow as landcover.py: write mesh once, then regrid each variable.
+ """
+
+ # each worker writes its temp files to a unique subdirectory to avoid collisions
+ scratch_base = os.environ.get('SCRATCH') or os.environ.get('TMPDIR') or tempfile.gettempdir()
+ tmp_dir = Path(tempfile.mkdtemp(dir=scratch_base, prefix=f'management_{year}_'))
+
+ # Read source data (each worker reads its own copy)
+ harvest_data = landgen_io.read_netcdf_ll(year, Path(harvest_path) / harvest_name, LUH2_HARVEST_VARS, ll_limits)
+ grazing_data = {}
+ for stem, grazing_name in grazing_names.items():
+ grazing_data[stem] = landgen_io.read_netcdf_ll(year, Path(grazing_path) / grazing_name, [stem], ll_limits)
+
+ # Create chunk-sized LtData object
+ n_chunk_cells = len(row_indices)
+ n_harvest = len(LUH2_HARVEST_VARS)
+ n_grazing = len(grazing_names)
+
+ chunk_lt_data = LtData()
+ chunk_lt_data.cell_idx = np.array(row_indices, dtype=np.int64)
+ chunk_lt_data.harvest_frac = np.zeros((n_chunk_cells, n_harvest), dtype=np.float64)
+ chunk_lt_data.grazing_frac = np.zeros((n_chunk_cells, n_grazing), dtype=np.float64)
+
+ try:
+ # Write mesh once per chunk (same approach as landcover.py)
+ mesh_file = tmp_dir / 'mesh.geojson'
+ landgen_io.write_mesh_to_geojson(out_grid_data, mesh_file, row_indices)
+
+ # --- regrid harvest variables ---
+ # LUH2_HARVEST_VARS order matches the n_harvest=10 dimension in LtData:
+ # index 0: primf_harv, 1: primn_harv, 2: secmf_harv, 3: secyf_harv, 4: secnf_harv
+ # 5: primf_bioh, 6: primn_bioh, 7: secmf_bioh, 8: secyf_bioh, 9: secnf_bioh
+ # Note that the biomass carbon variables (primf_bioh, etc) are not currently used in mksurfdat, but we regrid them here for completeness and potential future use.
+ for i, varname in enumerate(LUH2_HARVEST_VARS):
+ # Write source data to GeoTIFF
+ src_tif = tmp_dir / f"{varname}.tif"
+ landgen_io.write_latlon_to_geotiff(
+ harvest_data[varname],
+ harvest_data['lat'],
+ harvest_data['lon'],
+ ll_limits,
+ src_tif
+ )
+ # Regrid using modular function
+ regridded = landgen_io.regrid_to_mesh(
+ mesh_file, {varname: src_tif},
+ row_indices, out_grid_data,
+ out_type='data'
+ )
+ chunk_lt_data.harvest_frac[:, i] = regridded
+
+ # --- regrid grazing variables ---
+ # HYDE3.5 data is in km² per source grid cell. After area-weighted regridding
+ # to HEALPix the result is still in km². Divide by the (constant) HEALPix cell
+ # area to convert to a dimensionless fraction (0-1).
+ # Extract from out_grid_data (more direct than via DataFrame)
+ cell_area_km2 = out_grid_data.cell_area[row_indices[0]] / 1_000_000 # m² → km²
+
+ # Use stems as keys (e.g., 'pasture' not 'pasture.nc') to match grazing_data dict
+ grazing_stems = [Path(name).stem for name in grazing_names.keys()]
+ for i, stem in enumerate(grazing_stems):
+ # Write source data to GeoTIFF
+ src_tif = tmp_dir / f"{stem}.tif"
+ landgen_io.write_latlon_to_geotiff(
+ grazing_data[stem][stem],
+ grazing_data[stem]['lat'],
+ grazing_data[stem]['lon'],
+ ll_limits,
+ src_tif
+ )
+ # Regrid using modular function
+ regridded = landgen_io.regrid_to_mesh(
+ mesh_file, {stem: src_tif},
+ row_indices, out_grid_data,
+ out_type='data'
+ )
+ regridded = regridded / cell_area_km2 # km² → fraction
+ np.clip(regridded, 0.0, 1.0, out=regridded) # clamp rounding artefacts
+ chunk_lt_data.grazing_frac[:, i] = regridded
+
+ return chunk_lt_data
+
+ finally:
+ # Clean up temp files
+ shutil.rmtree(tmp_dir, ignore_errors=True)
+
+
+########## run()
+
+## called by land_type.process_single_year() for each year, and this is where the multiprocessing happens for the landcover module
+## this sets up the pool and calls the management_process() function for each chunk of data
+
+def run(lt_year_data, year, prev_year, harvest_path, harvest_name, grazing_path, grazing_names,
+ com_config_dict, out_grid_data, decomp_indices, decomp_ll_limits):
+
+ print(f"Processing management module with parameters:")
+ # todo: print the parameters here
+
+ # Determine the number of worker processes to use.
+ # Priority: SRUN_CPUS_PER_TASK -> SLURM_CPUS_PER_TASK -> SLURM_CPUS_ON_NODE -> mp.cpu_count()
+ in_slurm = os.environ.get('SLURM_JOB_ID') is not None
+
+ omp_threads_int = (
+ tools.parse_cpu_env('SRUN_CPUS_PER_TASK') or
+ tools.parse_cpu_env('SLURM_CPUS_PER_TASK') or
+ tools.parse_cpu_env('SLURM_CPUS_ON_NODE') or
+ mp.cpu_count()
+ )
+
+ if in_slurm:
+ logger.info(f"Running under SLURM (job {os.environ['SLURM_JOB_ID']}): using {omp_threads_int} workers "
+ f"(SRUN_CPUS_PER_TASK={os.environ.get('SRUN_CPUS_PER_TASK')}, "
+ f"SLURM_CPUS_PER_TASK={os.environ.get('SLURM_CPUS_PER_TASK')}, "
+ f"SLURM_CPUS_ON_NODE={os.environ.get('SLURM_CPUS_ON_NODE')})")
+ else:
+ logger.info(f"Running locally: using {omp_threads_int} workers (mp.cpu_count())")
+
+ # Build data_chunks from the pre-computed decomp_indices / decomp_ll_limits
+ # passed in from landgen.py via land_type.py. These were produced by
+ # set_decomp_cell_idx_ll_limits(), which computes tight vertex bounding boxes
+ # per chunk — exactly the ll_limits that write_latlon_to_geotiff needs so the
+ # raster slice fully covers every polygon in the chunk.
+ #
+ # decomp_indices contains 0-based row indices into out_grid_data arrays.
+ # write_mesh_to_geojson writes each cell's row index as the GeoJSON cellid property,
+ # so regrid_to_mesh can look up uraster results directly by index.
+ data_chunks = []
+ for row_indices, ll in zip(decomp_indices, decomp_ll_limits):
+ if len(row_indices) == 0:
+ continue # skip empty (ocean-only) chunks
+ # Each tuple contains all args for management_process (starmap approach)
+ data_chunks.append((
+ year, harvest_path, harvest_name, grazing_path, grazing_names,
+ com_config_dict, out_grid_data, ll, row_indices
+ ))
+
+ # Sort largest chunks first (most cells = slowest) so they are dispatched
+ # immediately and don't create a long tail at the end of the job.
+ data_chunks.sort(key=lambda t: len(t[8]), reverse=True) # row_indices is at index 8
+
+ n_chunks = len(data_chunks)
+ print(f" Submitting {n_chunks} management chunks to pool of {omp_threads_int} workers")
+
+ updated_vars = ['harvest_frac', 'grazing_frac']
+ with mp.Pool(processes=omp_threads_int) as pool:
+ for chunk_lt_data in pool.imap_unordered(_management_process_star, data_chunks):
+ lt_year_data.copy_from(chunk_lt_data, updated_vars)
+ del chunk_lt_data
+
+ return
diff --git a/components/elm/tools/landgen/src/landgen/plot_landgen.py b/components/elm/tools/landgen/src/landgen/plot_landgen.py
new file mode 100644
index 000000000000..64abfd628456
--- /dev/null
+++ b/components/elm/tools/landgen/src/landgen/plot_landgen.py
@@ -0,0 +1,320 @@
+"""Plot utilities for landgen NetCDF outputs.
+
+This module can be imported from other landgen code and can also be run as a
+standalone CLI module via:
+
+ python -m landgen.plot_landgen [options]
+"""
+
+import argparse
+import json
+import logging
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+import matplotlib
+
+matplotlib.use('Agg') # non-interactive backend — safe for HPC/batch
+import matplotlib.pyplot as plt
+from matplotlib.collections import PolyCollection
+
+try:
+ import cartopy.crs as ccrs
+ import cartopy.feature as cfeature
+ _HAS_CARTOPY = True
+except ImportError:
+ _HAS_CARTOPY = False
+
+
+# Grid geometry fields written by write_module_netcdf() — excluded from
+# auto-detect, except for landfrac or when explicitly requested by varnames.
+_GRID_FIELDS = frozenset({'lon_cen', 'lat_cen', 'lon_vtx', 'lat_vtx', 'cell_area'})
+
+
+def _plot_one_var(ax, ds, varname, layer, year, mask, lon_cen, lat_cen,
+ plot_type, colormap, ll_limits, _log, scale_limits=None):
+ """Draw a single variable onto an existing Axes object."""
+ da = ds[varname]
+ dims = list(da.dims)
+
+ # select time slice if present
+ if 'time' in dims:
+ time_vals = ds['time'].values
+ t_idx = int(np.argmin(np.abs(time_vals - year)))
+ da = da.isel(time=t_idx)
+ dims = list(da.dims)
+
+ # select layer / collapse extra dims
+ layer_applied = da.ndim >= 2
+ if da.ndim >= 2:
+ if da.ndim > 2:
+ extra_dims = dims[2:]
+ _log.warning(
+ f"plot_module_netcdf: '{varname}' has extra dimensions {extra_dims} "
+ f"beyond (cell, {dims[1]}); collapsing each to index 0."
+ )
+ da = da.isel({dims[1]: layer})
+ while da.ndim > 1:
+ da = da.isel({list(da.dims)[1]: 0})
+
+ values = da.values.astype(np.float64)[mask]
+
+ if scale_limits is not None:
+ vmin, vmax = float(scale_limits[0]), float(scale_limits[1])
+ else:
+ vmin, vmax = float(np.nanmin(values)), float(np.nanmax(values))
+ norm = plt.Normalize(vmin=vmin, vmax=vmax)
+ cmap = plt.colormaps[colormap]
+
+ # map extent
+ if ll_limits is not None:
+ min_lat, max_lat, min_lon, max_lon = ll_limits
+ else:
+ pad = 1.0
+ min_lon = float(lon_cen.min()) - pad
+ max_lon = float(lon_cen.max()) + pad
+ min_lat = float(lat_cen.min()) - pad
+ max_lat = float(lat_cen.max()) + pad
+
+ if _HAS_CARTOPY:
+ ax.set_extent([min_lon, max_lon, min_lat, max_lat], crs=ccrs.PlateCarree())
+ ax.add_feature(cfeature.COASTLINE, linewidth=0.5, zorder=5)
+ else:
+ ax.set_xlim(min_lon, max_lon)
+ ax.set_ylim(min_lat, max_lat)
+ ax.set_aspect('equal')
+ ax.set_xlabel('Longitude')
+ ax.set_ylabel('Latitude')
+ layer_suffix = f", layer {layer}" if layer_applied else ""
+ ax.set_title(f"{varname} (year {year}{layer_suffix})")
+
+ geo_kw = {'transform': ccrs.PlateCarree()} if _HAS_CARTOPY else {}
+
+ if plot_type == 'scatter':
+ mappable = ax.scatter(lon_cen, lat_cen, c=values, cmap=cmap, norm=norm,
+ s=1, linewidths=0, **geo_kw)
+ else: # 'rendered'
+ lon_vtx = ds['lon_vtx'].values[mask]
+ lat_vtx = ds['lat_vtx'].values[mask]
+ verts = np.stack([lon_vtx, lat_vtx], axis=-1)
+ mappable = PolyCollection(verts, array=values, cmap=cmap, norm=norm,
+ linewidths=0, **geo_kw)
+ ax.add_collection(mappable)
+
+ return mappable
+
+
+def _get_layer_indices(ds, vname, layers):
+ """Return (layer_indices, is_multilayer) for a variable in ds."""
+ da_tmp = ds[vname]
+ dims_tmp = list(da_tmp.dims)
+ if 'time' in dims_tmp:
+ da_tmp = da_tmp.isel(time=0)
+ dims_tmp = list(da_tmp.dims)
+ if da_tmp.ndim >= 2:
+ n_layers = da_tmp.sizes[dims_tmp[1]]
+ if layers is None:
+ idxs = list(range(n_layers))
+ elif isinstance(layers, int):
+ idxs = [layers]
+ else:
+ idxs = list(layers)
+ return idxs, True
+ return [0], False
+
+
+def _layers_for_var(vname, layers):
+ """Resolve per-variable layer selection.
+
+ layers can be:
+ - None
+ - int
+ - iterable of ints
+ - dict mapping varname -> (None | int | iterable of ints)
+ """
+ if isinstance(layers, dict):
+ return layers.get(vname, None)
+ return layers
+
+
+def plot_module_netcdf(file_path, out_path, year, varnames=None, layers=None,
+ plot_type='scatter', file_type='png',
+ colormap='viridis', ll_limits=None, scale_limits=None):
+ """Plot one or more variables from a NetCDF file written by landgen."""
+ _log = logging.getLogger('landgen')
+
+ if plot_type not in ('scatter', 'rendered'):
+ _log.error(f"plot_module_netcdf: plot_type must be 'scatter' or 'rendered', got '{plot_type}'")
+ raise ValueError(f"plot_module_netcdf: plot_type must be 'scatter' or 'rendered', got '{plot_type}'")
+ if file_type != 'png':
+ _log.error(f"plot_module_netcdf: file_type must be 'png', got '{file_type}'")
+ raise ValueError(f"plot_module_netcdf: file_type must be 'png', got '{file_type}'")
+
+ file_path = Path(file_path)
+ if not file_path.exists():
+ _log.error(f"plot_module_netcdf: file not found: {file_path}")
+ raise FileNotFoundError(f"plot_module_netcdf: file not found: {file_path}")
+
+ out_path = Path(out_path)
+ out_path.mkdir(parents=True, exist_ok=True)
+
+ ds = xr.open_dataset(file_path, decode_times=False)
+
+ cell_dim = 'cell'
+ if varnames is None:
+ varnames = [
+ v for v in ds.data_vars
+ if v not in _GRID_FIELDS and cell_dim in ds[v].dims
+ ]
+ if not varnames:
+ ds.close()
+ msg = f"plot_module_netcdf: no plottable data variables found in {file_path}"
+ _log.error(msg)
+ raise ValueError(msg)
+ _log.info(f"plot_module_netcdf: auto-detected variables: {varnames}")
+ else:
+ missing = [v for v in varnames if v not in ds]
+ if missing:
+ ds.close()
+ msg = (f"plot_module_netcdf: variables {missing} not found in {file_path}. "
+ f"Available: {list(ds.data_vars)}")
+ _log.error(msg)
+ raise KeyError(msg)
+
+ for v in varnames:
+ if 'time' in ds[v].dims:
+ time_vals = ds['time'].values
+ t_idx = int(np.argmin(np.abs(time_vals - year)))
+ if int(time_vals[t_idx]) != year:
+ ds.close()
+ msg = (f"plot_module_netcdf: year {year} not found in {file_path} "
+ f"(closest: {int(time_vals[t_idx])})")
+ _log.error(msg)
+ raise ValueError(msg)
+ break
+
+ lon_cen = ds['lon_cen'].values
+ lat_cen = ds['lat_cen'].values
+
+ if ll_limits is not None:
+ min_lat, max_lat, min_lon, max_lon = ll_limits
+ mask = (
+ (lat_cen >= min_lat) & (lat_cen <= max_lat) &
+ (lon_cen >= min_lon) & (lon_cen <= max_lon)
+ )
+ else:
+ mask = np.ones(len(lon_cen), dtype=bool)
+
+ if not mask.any():
+ ds.close()
+ msg = f"plot_module_netcdf: no cells remain after applying ll_limits {ll_limits}"
+ _log.error(msg)
+ raise ValueError(msg)
+
+ lon_cen = lon_cen[mask]
+ lat_cen = lat_cen[mask]
+
+ out_files = []
+ for vname in varnames:
+ var_layers = _layers_for_var(vname, layers)
+ layer_indices, is_multilayer = _get_layer_indices(ds, vname, var_layers)
+ for layer_idx in layer_indices:
+ if _HAS_CARTOPY:
+ fig, ax = plt.subplots(figsize=(12, 6),
+ subplot_kw={'projection': ccrs.PlateCarree()})
+ else:
+ fig, ax = plt.subplots(figsize=(12, 6))
+ try:
+ mappable = _plot_one_var(ax, ds, vname, layer_idx, year, mask,
+ lon_cen, lat_cen, plot_type, colormap,
+ ll_limits, _log, scale_limits=scale_limits)
+ fig.colorbar(mappable, ax=ax, label=vname)
+ if is_multilayer:
+ out_file = out_path / f"{file_path.stem}_{vname}_layer{layer_idx}_{year}.png"
+ else:
+ out_file = out_path / f"{file_path.stem}_{vname}_{year}.png"
+ fig.savefig(out_file, dpi=300, bbox_inches='tight')
+ out_files.append(out_file)
+ _log.info(f"plot_module_netcdf: wrote {out_file}")
+ except Exception as e:
+ _log.warning(f"plot_module_netcdf: skipping '{vname}' layer {layer_idx}: {e}")
+ finally:
+ plt.close(fig)
+ ds.close()
+ return out_files
+
+
+def _parse_layers_arg(layers_raw):
+ """Parse CLI layers argument.
+
+ Accepted examples:
+ "0"
+ "[0,1]"
+ '{"pct_pft": [0,1], "pct_ocean": null}'
+ """
+ if layers_raw is None:
+ return None
+ txt = layers_raw.strip()
+ if txt == '':
+ return None
+ try:
+ parsed = json.loads(txt)
+ except json.JSONDecodeError:
+ if txt.lstrip('-').isdigit():
+ return int(txt)
+ raise ValueError(
+ "Invalid --layers value. Use JSON (e.g. '[0,1]' or '{\"pct_pft\":[0,1]}') "
+ "or a single integer like '0'."
+ )
+ return parsed
+
+
+def main(argv=None):
+ """CLI entry point for standalone plotting."""
+ parser = argparse.ArgumentParser(
+ description="Plot variables from a landgen NetCDF file."
+ )
+ parser.add_argument('file_path', help='Input NetCDF file path')
+ parser.add_argument('out_path', help='Output directory for plots')
+ parser.add_argument('year', type=int, help='Calendar year to plot')
+ parser.add_argument('--varnames', nargs='*', default=None,
+ help='Optional variable names to plot; default=None prints all variables except grid geometry fields')
+ parser.add_argument('--layers', default=None,
+ help="Layer selector: int, JSON list, or JSON dict by variable; default=None prints all layers")
+ parser.add_argument('--plot-type', choices=['scatter', 'rendered'], default='scatter',
+ help="Plot type: 'scatter' for colored points, 'rendered' for cell polygons; default='scatter'")
+ parser.add_argument('--file-type', choices=['png'], default='png', help="Output file type; only 'png' is supported")
+ parser.add_argument('--colormap', default='viridis', help="Colormap for plots; default='viridis'")
+ parser.add_argument('--ll-limits', nargs=4, type=float, metavar=('MIN_LAT', 'MAX_LAT', 'MIN_LON', 'MAX_LON'),
+ default=None, help="Latitude and longitude limits for plots; default=None uses full extent")
+ parser.add_argument('--scale-limits', nargs=2, type=float, metavar=('VMIN', 'VMAX'),
+ default=None, help="Colorscale min and max; default=None uses data min/max")
+ args = parser.parse_args(argv)
+
+ logging.basicConfig(
+ level=logging.INFO,
+ format='%(asctime)s %(levelname)-8s %(message)s',
+ datefmt='%Y-%m-%d %H:%M:%S'
+ )
+
+ layers = _parse_layers_arg(args.layers)
+
+ result = plot_module_netcdf(
+ file_path=args.file_path,
+ out_path=args.out_path,
+ year=args.year,
+ varnames=args.varnames,
+ layers=layers,
+ plot_type=args.plot_type,
+ file_type=args.file_type,
+ colormap=args.colormap,
+ ll_limits=tuple(args.ll_limits) if args.ll_limits is not None else None,
+ scale_limits=tuple(args.scale_limits) if args.scale_limits is not None else None,
+ )
+
+ print(result)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/components/elm/tools/landgen/src/landgen/shared_data.py b/components/elm/tools/landgen/src/landgen/shared_data.py
index 9ff16d7dc46a..ed1f3c92c3af 100644
--- a/components/elm/tools/landgen/src/landgen/shared_data.py
+++ b/components/elm/tools/landgen/src/landgen/shared_data.py
@@ -8,6 +8,7 @@
# the shared custom data structures are registered in main()
import multiprocessing as mp
+from multiprocessing import Lock
from multiprocessing.managers import BaseManager
import numpy as np
@@ -17,7 +18,7 @@
# LtData dimensions
n_pfts_default = 51
-n_harvest_default = 5
+n_harvest_default = 10
n_grazing_default = 2
n_elev_default = 61
n_elev_edges_default = 62
@@ -46,7 +47,7 @@
# multiprocessing workers. All arrays are float64 except cell_id (int64).
# This can be used for any grid, but we currently use a healpix grid
# n_cells: number of 'land' cells in the landgen out grid
-# lon_xy and lat_xy are the cell 'center' coordinates
+# lon_cen and lat_cen are the cell 'center' coordinates
# lon_vtx and lat_vtx are the vertex coordinates for each cell
# n_vertices: number of vertices per cell (4 for quadrilateral cells)
# landfrac is determined by the topography processing, so initialize to 1 here
@@ -55,12 +56,11 @@ class GridData:
"""Per-cell grid geometry container (cell ids, coordinates, landfrac)."""
def __init__(self):
-
self.num_cells = None # int - number of 'land' cells in out grid
self.num_vertices = None # int - number of vertices per cell
self.cell_id = None # int64 [n_cells]
- self.lon_xy = None # float64 [n_cells]
- self.lat_xy = None # float64 [n_cells]
+ self.lon_cen = None # float64 [n_cells]
+ self.lat_cen = None # float64 [n_cells]
self.cell_area = None # float64 [n_cells]
self.landfrac = None # float64 [n_cells]
self.lon_vtx = None # float64 [n_cells, n_vertices]
@@ -72,8 +72,8 @@ def allocate(self, n_cells=n_cells_default, n_vertices=n_vertices_default):
self.num_cells = n_cells
self.num_vertices = n_vertices
self.cell_id = np.zeros(n_cells, dtype=np.int64)
- self.lon_xy = np.zeros(n_cells, dtype=np.float64)
- self.lat_xy = np.zeros(n_cells, dtype=np.float64)
+ self.lon_cen = np.zeros(n_cells, dtype=np.float64)
+ self.lat_cen = np.zeros(n_cells, dtype=np.float64)
self.cell_area = np.zeros(n_cells, dtype=np.float64)
self.landfrac = np.ones(n_cells, dtype=np.float64)
self.lon_vtx = np.zeros((n_cells, n_vertices), dtype=np.float64)
@@ -81,8 +81,8 @@ def allocate(self, n_cells=n_cells_default, n_vertices=n_vertices_default):
# --- getter methods (needed for proxy access via GridManager) ---
def get_cell_id(self): return self.cell_id
- def get_lon_xy(self): return self.lon_xy
- def get_lat_xy(self): return self.lat_xy
+ def get_lon_cen(self): return self.lon_cen
+ def get_lat_cen(self): return self.lat_cen
def get_cell_area(self): return self.cell_area
def get_landfrac(self): return self.landfrac
def get_lon_vtx(self): return self.lon_vtx
@@ -91,8 +91,8 @@ def get_num_cells(self): return self.num_cells
# --- setter methods (needed for proxy access via GridManager) ---
def set_cell_id(self, v): self.cell_id = v
- def set_lon_xy(self, v): self.lon_xy = v
- def set_lat_xy(self, v): self.lat_xy = v
+ def set_lon_cen(self, v): self.lon_cen = v
+ def set_lat_cen(self, v): self.lat_cen = v
def set_cell_area(self, v): self.cell_area = v
def set_landfrac(self, v): self.landfrac = v
def set_lon_vtx(self, v): self.lon_vtx = v
@@ -119,6 +119,7 @@ class TopoData:
"""Per-cell topography data container for the landgen workflow."""
def __init__(self):
+ self.cell_idx = None # int64 [n_cells]
self.topo = None # float64 [n_cells] - topographic height
self.std_elev = None # float32 [n_cells] - standard deviation of elevation
self.slope = None # float64 [n_cells] - mean slope
@@ -132,6 +133,7 @@ def __init__(self):
def allocate(self, n_cells=n_cells_default, n_levslp=n_levslp_default):
"""Allocate all arrays given dimension sizes."""
g = n_cells
+ self.cell_idx = np.zeros(g, dtype=np.int64)
self.topo = np.zeros(g, dtype=np.float64)
self.std_elev = np.zeros(g, dtype=np.float32)
self.slope = np.zeros(g, dtype=np.float64)
@@ -162,7 +164,7 @@ class TopoManager(BaseManager):
# These data are on the landgen grid defined by GridData
# n_cells: number of 'land' cells in the landgen grid
# n_pfts: number of plant functional types (51, includes bare and crop functional types)
-# n_harvest: number of harvest types (5: luh categories)
+# n_harvest: number of harvest types (10: luh categories)
# n_grazing: number of grazing types (2: pasture (grass, intensive) and rangeland)
# n_elev: number of elevation bins for glacier cover (currently 61, may change)
# n_elev_edges: number of elevation bin edges for glacier cover (n_elev + 1)
@@ -177,8 +179,8 @@ class LtData:
"""Per-cell landcover data container for the landgen workflow."""
def __init__(self):
-
# 1-D grid arrays [n_cells]
+ self.cell_idx = None # int64 [n_cells]
self.pct_ocean = None # float64
self.lake_depth = None # float64
self.lake_depth_mask = None # float64
@@ -246,6 +248,7 @@ def allocate(self, n_cells=n_cells_default, n_pfts=n_pfts_default, n_harvest=n_h
n_density=n_density_default, n_month=n_month_default, n_levurb=n_levurb_default,
n_rad=n_rad_default, n_solar=n_solar_default, n_vocveg=n_vocveg_default):
"""Allocate all arrays given dimension sizes."""
+ self.cell_idx = np.zeros(n_cells, dtype=np.int64)
self.pct_ocean = np.zeros(n_cells, dtype=np.float64)
self.lake_depth = np.zeros(n_cells, dtype=np.float64)
self.lake_depth_mask = np.zeros(n_cells, dtype=np.float64)
@@ -302,6 +305,70 @@ def allocate(self, n_cells=n_cells_default, n_pfts=n_pfts_default, n_harvest=n_h
self.cv_wall = np.zeros((n_cells, n_levurb), dtype=np.float64)
self.cv_improad = np.zeros((n_cells, n_levurb), dtype=np.float64)
+## todo: delete these if we are not using the manager proxy
+ # getter methods for manager proxy
+ def get_harvest_frac(self): return self.harvest_frac
+ def get_harvest_mass(self): return self.harvest_mass
+ def get_grazing_frac(self): return self.grazing_frac
+ def get_pct_pft(self): return self.pct_pft
+ def get_pct_ocean(self): return self.pct_ocean
+ def get_pct_lake(self): return self.pct_lake
+ def get_pct_wetland(self): return self.pct_wetland
+ def get_pct_glacier(self): return self.pct_glacier
+ def get_pct_urban(self): return self.pct_urban
+
+ # setter methods for manager proxy
+ def set_harvest_frac(self, cell_ids, i, values):
+ self.harvest_frac[cell_ids, i] = values
+ def set_harvest_mass(self, cell_ids, i, values):
+ self.harvest_mass[cell_ids, i] = values
+ def set_grazing_frac(self, cell_ids, i, values):
+ self.grazing_frac[cell_ids, i] = values
+ def set_pct_pft(self, cell_ids, i, values):
+ self.pct_pft[cell_ids, i] = values
+##
+
+ def copy_from(self, source, varnames):
+ """Copy listed variable values from a chunk LtData object into self.
+
+ Uses source.cell_idx to identify which global cell positions in self
+ receive the data. Works for 1-D, 2-D, and 3-D arrays alike because
+ numpy fancy indexing on the first axis handles all shapes uniformly.
+
+ bin_centers and bin_edges have no n_cells dimension; they are copied
+ directly (full array assignment) rather than via cell_idx.
+
+ Args:
+ source (LtData): Chunk object whose data are to be merged in.
+ source.cell_idx must already be set to the
+ global cell indices corresponding to source's
+ local 0..n_cells-1 positions.
+ varnames (list[str]): Names of LtData attributes to copy.
+
+ Raises:
+ AttributeError: If a name in varnames is not an attribute of LtData.
+ ValueError: If source.cell_idx is None (not set).
+ """
+ if source.cell_idx is None:
+ raise ValueError("copy_from: source.cell_idx is None; call set_cell_idx() before copy_from().")
+
+ # variables that have no n_cells first dimension
+ _no_cell_dim = {'bin_centers', 'bin_edges'}
+
+ idx = source.cell_idx # 1-D int64 array of global cell positions
+
+ for name in varnames:
+ if not hasattr(self, name):
+ raise AttributeError(f"copy_from: LtData has no attribute '{name}'.")
+ if not hasattr(source, name):
+ raise AttributeError(f"copy_from: source LtData has no attribute '{name}'.")
+ src_val = getattr(source, name)
+ if src_val is None:
+ continue # source variable was never set; skip silently
+ if name in _no_cell_dim:
+ getattr(self, name)[:] = src_val
+ else:
+ getattr(self, name)[idx] = src_val
# ---------------------------------------------------------------------------
# LtManager: custom BaseManager that can vend LtData proxy objects to
@@ -311,8 +378,3 @@ class LtManager(BaseManager):
pass
LtManager.register('LtData', LtData)
-
-
-
-
-
diff --git a/components/elm/tools/landgen/src/landgen/soil.py b/components/elm/tools/landgen/src/landgen/soil.py
index bdaa474c0bf5..7099d31bce08 100644
--- a/components/elm/tools/landgen/src/landgen/soil.py
+++ b/components/elm/tools/landgen/src/landgen/soil.py
@@ -7,8 +7,11 @@
import multiprocessing as mp
import importlib
+import logging
from pathlib import Path
+logger = logging.getLogger('landgen')
+
########## define helper functions for land_type run() here
@@ -32,9 +35,9 @@
## output
-def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager):
+def run(active, out_fname, com_config_dict, out_grid_data, manager, decomp_indices, decomp_ll_limits):
if active is False:
- print(f"Skipping soil module")
+ logger.info("Skipping soil module")
return
# extract common parameters from shared config dict
@@ -46,7 +49,7 @@ def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager
# landfrac is stored in out_grid_data (set by topography module)
landfrac = out_grid_data.get_landfrac()
- print(f"Processing soil module with parameters:")
+ logger.info("Processing soil module")
# todo: print the parameters here
# processing code for soil
diff --git a/components/elm/tools/landgen/src/landgen/tools.py b/components/elm/tools/landgen/src/landgen/tools.py
new file mode 100644
index 000000000000..50c25c743e97
--- /dev/null
+++ b/components/elm/tools/landgen/src/landgen/tools.py
@@ -0,0 +1,225 @@
+# Utility functions for landgen
+
+import logging
+import multiprocessing
+import os
+import sys
+import threading
+import time
+import psutil
+import numpy as np
+from pathlib import Path
+
+from .plot_landgen import plot_module_netcdf
+
+
+#------- shared logger setup ----------------------------------------------------
+def setup_logger(name, log_path, level=logging.INFO):
+ """
+ Configure a named logger to write to log_path and to stdout.
+ Call once (e.g. from main()) for each logger you need. Any module can
+ then obtain the same logger with:
+
+ import logging
+ logger = logging.getLogger('')
+
+ Each call is idempotent: if handlers are already attached the logger is
+ returned unchanged, so calling setup_logger() twice is safe.
+
+ Args:
+ name (str): Logger name, e.g. 'landgen' or 'ClusterMonitor'.
+ log_path (Path): Full path to the log file (unique per run).
+ level (int): Logging level (default logging.INFO).
+
+ Returns:
+ logging.Logger: The configured logger.
+ """
+ logger = logging.getLogger(name)
+ if logger.handlers:
+ return logger # already configured
+
+ logger.setLevel(level)
+ logger.propagate = False # don't double-log via the root logger
+
+ fmt = logging.Formatter('%(asctime)s %(levelname)-8s %(message)s',
+ datefmt='%Y-%m-%d %H:%M:%S')
+
+ # file handler — each run gets its own file (mode='w')
+ fh = logging.FileHandler(log_path, mode='w')
+ fh.setFormatter(fmt)
+ logger.addHandler(fh)
+
+ # console handler
+ ch = logging.StreamHandler(sys.stdout)
+ ch.setFormatter(fmt)
+ logger.addHandler(ch)
+
+ return logger
+
+
+def init_worker_logging(log_path, logger_name='landgen'):
+ """
+ Attach a FileHandler (append mode) to the named logger inside a worker
+ process. Must be called at the start of any function that runs inside a
+ multiprocessing Pool, because forkserver/spawn workers start with a clean
+ logging state — the parent's handlers are not inherited.
+
+ Idempotent: does nothing if the logger already has handlers.
+
+ Args:
+ log_path (str | Path): Path to the shared log file.
+ logger_name (str): Logger name (default 'landgen').
+ """
+ worker_logger = logging.getLogger(logger_name)
+ if worker_logger.handlers:
+ return # already configured in this process
+ worker_logger.setLevel(logging.INFO)
+ worker_logger.propagate = False
+ fmt = logging.Formatter('%(asctime)s %(levelname)-8s %(message)s',
+ datefmt='%Y-%m-%d %H:%M:%S')
+ fh = logging.FileHandler(log_path, mode='a')
+ fh.setFormatter(fmt)
+ worker_logger.addHandler(fh)
+
+
+def redirect_uraster_logs(out_path):
+ """
+ Redirect uraster's auto-created log files from the process CWD to out_path.
+
+ uraster calls logging.FileHandler(f"{module_name}.log") at module-level
+ import time (no directory prefix), so the files land in whatever the CWD is
+ when the worker process first imports uraster. This function walks all
+ active loggers, finds FileHandlers whose log file is not already inside
+ out_path, and replaces them with append-mode handlers in out_path.
+
+ Call this in each worker process after uraster has been imported.
+
+ Args:
+ out_path (str | Path): Directory to write uraster log files into.
+ """
+ out_path = Path(out_path)
+ out_path.mkdir(parents=True, exist_ok=True)
+ out_str = str(out_path.resolve())
+
+ for obj in logging.Logger.manager.loggerDict.values():
+ if not isinstance(obj, logging.Logger):
+ continue # skip PlaceHolder entries
+ for handler in list(obj.handlers):
+ if not isinstance(handler, logging.FileHandler):
+ continue
+ current = Path(handler.baseFilename).resolve()
+ if str(current).startswith(out_str):
+ continue # already in the right place
+ new_path = out_path / current.name
+ new_handler = logging.FileHandler(new_path, mode='a')
+ if handler.formatter:
+ new_handler.setFormatter(handler.formatter)
+ new_handler.setLevel(handler.level)
+ obj.removeHandler(handler)
+ handler.close()
+ obj.addHandler(new_handler)
+
+
+#------- HPC / multiprocessing helpers -----------------------------------------
+def parse_cpu_env(varname):
+ """
+ Return the integer value of an environment variable, or None if unset or
+ not a valid integer. Logs a warning via the 'landgen' logger on bad values.
+
+ Args:
+ varname (str): Environment variable name to read.
+
+ Returns:
+ int | None
+ """
+ _log = logging.getLogger('landgen')
+ val = os.environ.get(varname)
+ if val is None:
+ return None
+ try:
+ return int(val)
+ except ValueError:
+ _log.warning(f"{varname} has invalid integer value '{val}'; ignoring.")
+ return None
+
+
+#------- monitoring computational resources for debugging and performance tuning ------
+def monitor_cluster_resources(interval_sec=60.0, stop_event=None):
+ """Periodically logs aggregated CPU and memory usage of the entire process tree."""
+ parent_pid = os.getpid()
+
+ try:
+ parent_proc = psutil.Process(parent_pid)
+ except psutil.NoSuchProcess:
+ return
+
+ _resource_logger = logging.getLogger('ClusterMonitor')
+ _resource_logger.info(f"Starting resource monitor thread (Interval: {interval_sec}s)...")
+
+ # Use logical CPUs so the reported capacity matches psutil cpu_percent(),
+ # which is measured in units of one logical CPU.
+ # On Perlmutter: 128 physical cores, 256 logical (hyperthreaded).
+ node_cores = psutil.cpu_count(logical=True) or psutil.cpu_count(logical=False) or 1
+
+ # Prime cpu_percent() for all current processes. psutil.Process.cpu_percent()
+ # always returns 0.0 on the first call — it only establishes the baseline
+ # timestamp. Without priming here, every process would show 0% on the first
+ # monitoring interval.
+ known_processes = {}
+ try:
+ for proc in [parent_proc] + parent_proc.children(recursive=True):
+ try:
+ proc.cpu_percent(interval=None) # first call; always returns 0 — primes the counter
+ known_processes[proc.pid] = proc
+ except (psutil.NoSuchProcess, psutil.AccessDenied):
+ pass
+ except psutil.NoSuchProcess:
+ pass
+
+ while not stop_event.is_set():
+ # Wait for the interval so cpu_percent() has a meaningful delta to measure
+ stop_event.wait(interval_sec)
+ if stop_event.is_set():
+ break
+
+ try:
+ # Gather current active process tree
+ current_procs = [parent_proc] + parent_proc.children(recursive=True)
+ except psutil.NoSuchProcess:
+ continue
+
+ total_mem_bytes = 0
+ total_cpu_pct = 0.0
+ active_count = 0
+ new_process_cache = {}
+
+ for proc in current_procs:
+ pid = proc.pid
+ try:
+ if proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE:
+ # Reuse cached psutil.Process object so cpu_percent() delta
+ # is measured from the previous interval (not from now).
+ # New processes (not yet primed) will return 0 this interval
+ # but will report accurately on the next one.
+ tracked_proc = known_processes.get(pid, proc)
+ new_process_cache[pid] = tracked_proc
+
+ total_mem_bytes += tracked_proc.memory_info().rss
+ total_cpu_pct += tracked_proc.cpu_percent(interval=None)
+ active_count += 1
+ except (psutil.NoSuchProcess, psutil.AccessDenied):
+ continue
+
+ # Drop dead processes from cache
+ known_processes = new_process_cache
+
+ total_mem_gb = total_mem_bytes / (1024 ** 3)
+ # cpu_percent() returns percent of one core, so divide by 100 to get core-equivalents
+ cores_utilized = total_cpu_pct / 100.0
+
+ _resource_logger.info(
+ f"[SYSTEM MONITOR] Processes: {active_count} | "
+ f"Equivalent logical cores utilized: {cores_utilized:.1f}/{node_cores} | "
+ f"Memory in use: {total_mem_gb:.2f} GB"
+ )
+
diff --git a/components/elm/tools/landgen/src/landgen/topography.py b/components/elm/tools/landgen/src/landgen/topography.py
index bba680cfd68b..5cfe612f00e5 100644
--- a/components/elm/tools/landgen/src/landgen/topography.py
+++ b/components/elm/tools/landgen/src/landgen/topography.py
@@ -8,11 +8,14 @@
import multiprocessing as mp
import importlib
+import logging
import sys
from pathlib import Path
from . import shared_data
from .shared_data import TopoData, TopoManager
+logger = logging.getLogger('landgen')
+
########## define helper functions for land_type run() here
@@ -31,18 +34,16 @@
## output
-def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager):
+def run(active, out_fname, com_config_dict, out_grid_data, manager, decomp_indices, decomp_ll_limits):
if active is False:
- print(f"Topography processing is not active, but reading in landfrac data for other active modules.")
+ logger.info("Topography processing is not active, but reading in landfrac data for other active modules.")
output_file = Path(com_config_dict['out_path']) / out_fname
if output_file.exists():
- print(f"Output file {output_file} exists; reading in existing data.")
+ logger.info(f"Output file {output_file} exists; reading in existing data.")
## todo: need to define read_landfrac to return the data in the correct format for the shared data structure
- ## may need to add elevation to GridData class and here
- out_grid_data.landfrac = read_landfrac(output_file)
- else:
- print(f"Error: Output file {output_file} does not exist; set active to True to process topography data.")
- sys.exit(1)
+ # out_grid_data.set_landfrac(read_landfrac(output_file))
+ raise FileNotFoundError(
+ f"Output file {output_file} does not exist; set active to True to process topography data.")
# set up the topography module shared data structure
topo_manager = TopoManager()
@@ -50,7 +51,7 @@ def run(active, out_fname, com_config_dict, out_grid_data, manager, grid_manager
topo_out_data = topo_manager.TopoData()
topo_out_data.allocate()
- print(f"Processing topography module with parameters:")
+ logger.info("Processing topography module")
# todo: print the parameters here
diff --git a/components/elm/tools/landgen/src/submit_landgen.sh b/components/elm/tools/landgen/src/submit_landgen.sh
deleted file mode 100644
index e69de29bb2d1..000000000000
diff --git a/components/elm/tools/landgen/submit_landgen.sh b/components/elm/tools/landgen/submit_landgen.sh
old mode 100644
new mode 100755
index 8e39733326dd..64fc31fbeafe
--- a/components/elm/tools/landgen/submit_landgen.sh
+++ b/components/elm/tools/landgen/submit_landgen.sh
@@ -1,18 +1,77 @@
#!/bin/bash
-#SBATCH --job-name=parallel_python
-#SBATCH --nodes=1 # Ensure single node
-#SBATCH --ntasks=1 # Run one task (master process)
-#SBATCH --exclusive # Ensure exclusive access to the node (uses all cpus)
-##SBATCH --mem=16G # Request memory (adjust as needed)
-#SBATCH --time=00:30:00
+#SBATCH --job-name=landgen
+#SBATCH --nodes=1 # Ensure single node
+#SBATCH --ntasks=1 # Run one task (master process)
+#SBATCH --cpus-per-task=128
+#SBATCH --exclusive # Ensure exclusive access to the node (uses all cpus)
+#SBATCH --mem=0 # 0=Request all memory in this node (adjust as needed)
+#SBATCH --time=02:00:00
#SBATCH --account e3sm
-#SBATCH --partition=regular
+#SBATCH --qos=regular
#SBATCH --constraint=cpu
-# Tell Python/OpenMP to use the requested number of threads
-export OMP_NUM_THREADS=${SLURM_CPUS_ON_NODE}
-export MKL_NUM_THREADS=${SLURM_CPUS_ON_NODE}
+# type: sbatch submit_landgen.sh to submit the job to SLURM (in the directory where this script is located)
-# Run the landgen package
-# This command will run the __main__.py file inside the landgen package/directory
-python -m landgen
\ No newline at end of file
+# srun belpow will run the __main__.py file inside the landgen package/directory
+# with config.json as the default input configuration file
+# The configuration file should be located in the same directory as this script
+# Update the configuration file path if it's located elsewhere or has a different name
+# Note that config_template.json is a template and should be copied to config.json, or another name,
+# and edited with the desired settings before running this script
+
+# with --exclusive and --nodes=1 and --ntasks=1, the job will have access to all cores on the node
+# and SLURM_CPUS_PER_TASK will be set to the total number of cores on the node
+# this allows multiprocessing to use all cores without oversubscribing,
+# when combined with OMP_NUM_THREADS=1 and MKL_NUM_THREADS=1 below
+
+# but to automatically adjust cores based on the node, we can use srun instead, with SRUN_CPUS_PER_TASK
+# so launch with srun below and then the code looks for SRUN_CPUS_PER_TASK instead of SLURM_CPUS_PER_TASK
+# to reduce the number of cores, set SRUN_CPUS_PER_TASK below to the desired number
+# and keep --exclusive to ensure the entire node is reserved for this job
+
+# perlmutter has a regular and debug queue for cpus
+# perlmutter has 128 physical cores, with 256 logical (2 hyperthreads per core), and 512G per node
+
+# Calculate the number of workers to use:
+# SLURM_CPUS_ON_NODE counts logical CPUs (hyperthreaded), so on Perlmutter:
+# SLURM_CPUS_ON_NODE = 256 logical = 128 physical x 2 hyperthreads
+# physical = 256 / 2 = 128
+PHYS_CPUS=$(( SLURM_CPUS_ON_NODE / 2 ))
+HALF_CPUS=$(( PHYS_CPUS / 2 ))
+P90_CPUS=$(( (PHYS_CPUS * 90 + 99) / 100 )) # 90% of physical cores, rounded up
+HALF_LOG_CPUS=$(( (SLURM_CPUS_ON_NODE * 50 + 99) / 100 )) # 50% of logical cores, rounded up
+P90_LOG_CPUS=$(( (SLURM_CPUS_ON_NODE * 90 + 99) / 100 )) # 90% of logical cores, rounded up
+
+# set the srun cpus to use per task
+export SRUN_CPUS_PER_TASK=$P90_LOG_CPUS
+
+# Tell Python math and OpenMP to use the requested number of threads
+# set these to 1 so that each process uses one thread on each core
+# If these are >1 they allow multiple threads per process, which asks for #process * #thread cores
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+
+#### Activate the landgen conda environment
+
+module load conda
+conda activate landgen_env
+
+#### Run the landgen package
+
+# Use SLURM_SUBMIT_DIR (directory where sbatch was invoked),
+# which resolves to the SLURM spool directory at runtime.
+
+SCRIPT_DIR="${SLURM_SUBMIT_DIR}"
+INPUT_FILE="config.json"
+
+# Redirect stdout and stderr to the run output directory.
+# #SBATCH --output cannot reference shell variables, so we use exec to replace
+# this script's file descriptors after reading out_path from config.json.
+# The default slurm-JOBID.out in the submit dir will be created but left empty.
+OUT_PATH=$(python -c "import json,sys; print(json.load(open('${SCRIPT_DIR}/${INPUT_FILE}')).get('out_path','.'))")
+mkdir -p "${OUT_PATH}"
+exec > "${OUT_PATH}/slurm-${SLURM_JOB_ID}.out" 2>&1
+# Remove the now-empty default SLURM output file from the submit directory
+rm -f "${SLURM_SUBMIT_DIR}/slurm-${SLURM_JOB_ID}.out"
+
+srun python -m landgen "${SCRIPT_DIR}/${INPUT_FILE}"