diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..fe419ba --- /dev/null +++ b/.dockerignore @@ -0,0 +1,7 @@ +target +slurm +docs +.github +.claude +test_data +Dockerfile diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 9fd45e0..cbef49f 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -16,6 +16,8 @@ jobs: steps: - uses: actions/checkout@v4 + - name: Install dependencies + run: sudo apt-get update && sudo apt-get install -y python3-dev libnetcdf-dev libhdf5-dev libsqlite3-dev - name: Build run: cargo build --verbose - name: Run tests diff --git a/Cargo.lock b/Cargo.lock index 5e662e0..a76ec4b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -169,6 +169,7 @@ dependencies = [ "rusqlite", "serde", "serde_json", + "tempfile", "thiserror 1.0.69", "toml", "zarrs", @@ -427,6 +428,16 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "event-listener" version = "5.4.1" @@ -460,6 +471,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fastrand" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" + [[package]] name = "find-msvc-tools" version = "0.1.5" @@ -830,6 +847,12 @@ dependencies = [ "cc", ] +[[package]] +name = "linux-raw-sys" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" + [[package]] name = "lock_api" version = "0.4.14" @@ -1386,6 +1409,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustversion" version = "1.0.22" @@ -1546,6 +1582,19 @@ version = "0.12.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +[[package]] +name = "tempfile" +version = "3.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0136791f7c95b1f6dd99f9cc786b91bb81c3800b639b3478e561ddb7be95e5f1" +dependencies = [ + "fastrand", + "getrandom 0.4.2", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "1.0.69" diff --git a/Cargo.toml b/Cargo.toml index 62e5522..0a4077f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,9 @@ toml = "0.8" dirs = "5" zarrs = { version = "0.23", features = ["filesystem"], optional = true } +[dev-dependencies] +tempfile = "3" + [[bin]] name = "bmi-driver" path = "src/main.rs" diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..1a67a76 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,14 @@ +FROM awiciroh/ngiab AS ngiab_build +WORKDIR /build +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | bash -s -- -y +RUN dnf install -y python3.11-devel netcdf-devel hdf5-devel sqlite-devel g++ +COPY . . +RUN cargo build --release + +FROM awiciroh/ngiab AS final +COPY --from=ngiab_build /build/target/release/bmi-driver /usr/local/bin/bmi-driver +VOLUME /data +RUN ln -s /ngen/ngen/data /data +ENV OMP_NUM_THREADS=1 +ENTRYPOINT [ "/usr/local/bin/bmi-driver" ] +CMD ["/data"] diff --git a/README.md b/README.md index 5e09286..bd35591 100644 --- a/README.md +++ b/README.md @@ -2,56 +2,23 @@ A parallel BMI (Basic Model Interface) orchestrator that dynamically loads and couples C, Fortran, and Python hydrological models. It reads NetCDF forcing data, resolves model dependencies, handles automatic unit conversion, and runs across many geographic locations in parallel. -The binary is named `bmi-driver`. - ## Building -### Default (all adapters) - -```bash -cargo build --release -``` - -By default this builds with C, Fortran, and Python adapter support. The build will check for required system libraries and print clear install instructions if anything is missing. - -### C adapter only - -```bash -cargo build --release --no-default-features -``` - -Builds with only the C shared library adapter and the built-in SLOTH dummy model. This has the fewest system dependencies (only `libnetcdf`). - -### Selecting features - -```bash -cargo build --release --no-default-features --features fortran -cargo build --release --no-default-features --features python -cargo build --release --no-default-features --features fortran,python -``` - -- **`fortran`** -- Enables the `bmi_fortran` adapter for loading Fortran shared libraries. No additional build dependencies; a Fortran runtime must be available at execution time. -- **`python`** -- Enables the `bmi_python` adapter for loading Python BMI models via [PyO3](https://pyo3.rs). Requires Python development headers at build time (e.g., `python3-dev`). The Python environment must have the target BMI packages installed at runtime. - -### Install - ```bash -cargo install --path . +cargo build --release # All adapters (C, Fortran, Python, Zarr) +cargo build --release --no-default-features # C adapter only (fewest deps) +cargo install --path . # Install to ~/.cargo/bin/ ``` -This installs `bmi-driver` to `~/.cargo/bin/`. +Feature flags: `fortran`, `python`, `zarr` (all enabled by default). ### System dependencies | Library | Purpose | Install | |---------|---------|---------| -| `libnetcdf` | Reading NetCDF forcing files | `apt install libnetcdf-dev` / `dnf install netcdf-devel` / `brew install netcdf` | -| `pkg-config` | Library discovery at build time | `apt install pkg-config` / `dnf install pkgconfig` / `brew install pkg-config` | -| Python dev headers | Only with `python` feature (default) | `apt install python3-dev` / `dnf install python3-devel` / `brew install python3` | - -SQLite is bundled and compiled from source automatically -- no system `libsqlite3` package is needed. - -Model shared libraries (`.so`) must be accessible at the paths specified in the config. +| `libnetcdf` | Reading NetCDF forcing files | `apt install libnetcdf-dev` | +| `pkg-config` | Library discovery at build time | `apt install pkg-config` | +| Python dev headers | Only with `python` feature | `apt install python3-dev` | ## Usage @@ -59,13 +26,13 @@ Model shared libraries (`.so`) must be accessible at the paths specified in the bmi-driver [OPTIONS] ``` -### Options - | Flag | Description | |------|-------------| | `-j, --jobs N` | Number of parallel worker processes (default: CPU count) | -| `--units` | Print all unit conversion info and exit without running | +| `--units` | Print all unit conversion info and exit | | `--minify` | Strip unused fields from realization.json and exit | +| `--node-start N` | First location index for this node (multi-node/SLURM) | +| `--node-count N` | Number of locations for this node (0 = all remaining) | ### Data directory layout @@ -74,78 +41,42 @@ bmi-driver [OPTIONS] config/ realization.json # Model chain configuration *.gpkg # GeoPackage with location IDs - config/cat_config/ - /{{id}}.input # Per-location model configs + cat_config//{{id}} # Per-location model configs forcings/ forcings.nc # NetCDF forcing data outputs/ - bmi-driver/ # Output CSVs written here - .csv + bmi-driver/ # Output files written here ``` -The GeoPackage must contain a `divides` table with a `divide_id` text column listing all location IDs. - -### How it runs - -1. **Parent process** reads config, queries location IDs from the GeoPackage, checks for unmapped model inputs (offering to auto-add suggested mappings), and prints active unit conversions. -2. Parent divides locations into chunks and spawns N worker subprocesses. -3. **Each worker** iterates its assigned locations: initializes models, runs all timesteps, writes a CSV per location, and reports progress counts to the parent via stdout. -4. Parent displays nested progress bars (overall + per-worker) using the reported counts. - -### Library usage - -```rust -use bmi_driver::{ModelRunner, BmiError}; - -fn main() -> Result<(), BmiError> { - let mut runner = ModelRunner::from_config("config/realization.json")?; - runner.initialize("cat-123")?; - runner.run()?; - - let outputs = runner.main_outputs()?; - println!("Completed {} timesteps", outputs.len()); - - runner.finalize()?; - Ok(()) -} -``` +The GeoPackage must contain a `divides` table with a `divide_id` text column. ## Configuration -The config file is `/config/realization.json`. It defines the model chain, forcing data source, and time range. The format is compatible with [ngen](https://github.com/NOAA-OWP/ngen) realization configs -- extra ngen-specific fields are silently ignored. - -### Minimal example +Expects `/config/realization.json` defining the model chain, forcing source, and time range. Compatible with [ngen](https://github.com/NOAA-OWP/ngen) realization configs. ```json { "global": { - "formulations": [ - { - "name": "bmi_multi", - "params": { - "main_output_variable": "Q_OUT", - "output_variables": ["Q_OUT"], - "modules": [ - { - "name": "bmi_c", - "params": { - "model_type_name": "CFE", - "library_file": "/path/to/libcfebmi.so", - "init_config": "./config/cat_config/CFE/{{id}}.ini", - "registration_function": "register_bmi_cfe", - "main_output_variable": "Q_OUT", - "variables_names_map": { - "atmosphere_water__liquid_equivalent_precipitation_rate": "precip_rate" - } - } + "formulations": [{ + "name": "bmi_multi", + "params": { + "main_output_variable": "Q_OUT", + "modules": [{ + "name": "bmi_c", + "params": { + "model_type_name": "CFE", + "library_file": "/path/to/libcfebmi.so", + "init_config": "./config/cat_config/CFE/{{id}}.ini", + "registration_function": "register_bmi_cfe", + "main_output_variable": "Q_OUT", + "variables_names_map": { + "atmosphere_water__liquid_equivalent_precipitation_rate": "precip_rate" } - ] - } + } + }] } - ], - "forcing": { - "path": "./forcings/forcings.nc" - } + }], + "forcing": { "path": "./forcings/forcings.nc" } }, "time": { "start_time": "2010-01-01 00:00:00", @@ -155,313 +86,21 @@ The config file is `/config/realization.json`. It defines the model ch } ``` -### Top-level fields - -| Field | Required | Description | -|-------|----------|-------------| -| `global.formulations` | yes | Array with one `bmi_multi` formulation entry | -| `global.forcing.path` | yes | Path to NetCDF forcing file (relative to data_dir) | -| `time.start_time` | yes | Simulation start, `"YYYY-MM-DD HH:MM:SS"` | -| `time.end_time` | yes | Simulation end, `"YYYY-MM-DD HH:MM:SS"` | -| `time.output_interval` | no | Output interval in seconds (default: 3600) | - -### Formulation fields - -| Field | Required | Description | -|-------|----------|-------------| -| `name` | yes | Must be `"bmi_multi"` for multi-model orchestration | -| `params.modules` | yes | Array of module configs (see below) | -| `params.main_output_variable` | no | Primary output variable name | -| `params.output_variables` | no | Variables to include in output CSV. If empty, all model outputs are written. | - -### Module fields (shared across all adapter types) - -| Field | Required | Description | -|-------|----------|-------------| -| `name` | yes | Adapter type: `"bmi_c"`, `"bmi_fortran"`, `"bmi_python"`, or `"bmi_c++"` | -| `params.model_type_name` | yes | Model identifier (e.g., `"CFE"`, `"NoahOWP"`, `"SLOTH"`) | -| `params.init_config` | yes | Path to model config file. `{{id}}` is replaced with the location ID at runtime. | -| `params.main_output_variable` | yes | The primary output variable of this model | -| `params.library_file` | depends | Path to shared library. Required for C and Fortran adapters. | -| `params.registration_function` | no | Name of the C registration function (default: `"register_bmi"`) | -| `params.variables_names_map` | no | Maps model input names to source variable names (from forcings or upstream models) | -| `params.model_params` | no | Key-value parameters set on the model after initialization via `set_value()` | -| `params.python_type` | no | Python class path for the Python adapter (see below) | - -## Module types - -### bmi_c -- C shared libraries - -Loads a C shared library via `dlopen` with `RTLD_GLOBAL` (so symbols are shared across libraries). Looks up the registration function by name, calls it to populate a struct of BMI function pointers, then calls through those pointers for all BMI operations. - -```json -{ - "name": "bmi_c", - "params": { - "model_type_name": "CFE", - "library_file": "/path/to/libcfebmi.so", - "init_config": "./config/cat_config/CFE/{{id}}.ini", - "registration_function": "register_bmi_cfe", - "main_output_variable": "Q_OUT", - "variables_names_map": { - "atmosphere_water__liquid_equivalent_precipitation_rate": "QINSUR", - "water_potential_evaporation_flux": "EVAPOTRANS" - } - } -} -``` - -The shared library must export a registration function matching this signature: - -```c -Bmi* register_bmi_cfe(Bmi* model); -``` - -This function fills in the `Bmi` struct's function pointers and returns the pointer. The default function name is `register_bmi` if `registration_function` is omitted. - -### bmi_fortran -- Fortran shared libraries - -**Requires:** `--features fortran` - -Loads Fortran BMI models. The adapter supports two modes: - -1. **Single library** -- the `.so` implements both the registration function and the BMI interface directly. -2. **With middleware** -- a separate C-to-Fortran middleware library translates C-calling-convention calls to Fortran. The runner auto-detects a middleware library if one exists alongside the model library. - -The adapter passes an opaque handle pointer as the first argument to all BMI calls, matching the convention used by Fortran BMI implementations. - -```json -{ - "name": "bmi_fortran", - "params": { - "model_type_name": "NoahOWP", - "library_file": "/path/to/libsurfacebmi.so", - "init_config": "./config/cat_config/NOAH-OWP-M/{{id}}.input", - "main_output_variable": "QINSUR", - "variables_names_map": { - "PRCPNONC": "precip_rate", - "SFCTMP": "TMP_2maboveground", - "SFCPRS": "PRES_surface", - "Q2": "SPFH_2maboveground", - "UU": "UGRD_10maboveground", - "VV": "VGRD_10maboveground", - "LWDN": "DLWRF_surface", - "SOLDN": "DSWRF_surface" - } - } -} -``` - -### bmi_python -- Python BMI classes - -**Requires:** `--features python` - -Loads a Python class that implements the BMI interface. Values are exchanged via numpy arrays. There are two ways to specify the Python class: - -**Option 1: `python_type` (preferred)** -- a dotted module path where the last component is the class name: - -```json -{ - "name": "bmi_python", - "params": { - "model_type_name": "LSTM", - "python_type": "bmi_lstm.bmi_LSTM", - "init_config": "./config/cat_config/LSTM/{{id}}.yml", - "main_output_variable": "land_surface_water__runoff_volume_flux", - "variables_names_map": { - "atmosphere_water__liquid_equivalent_precipitation_rate": "precip_rate", - "land_surface_air__temperature": "TMP_2maboveground" - } - } -} -``` - -This imports `bmi_lstm` and instantiates `bmi_LSTM()`. The package must be installed or on `PYTHONPATH`. - -**Option 2: `library_file` + `registration_function`** -- loads a `.py` file directly: - -```json -{ - "name": "bmi_python", - "params": { - "model_type_name": "MyModel", - "library_file": "/path/to/my_model.py", - "registration_function": "MyBmiClass", - "init_config": "./config/cat_config/MyModel/{{id}}.yml", - "main_output_variable": "output_var" - } -} -``` - -This reads the file, adds its parent directory to `sys.path`, imports the module by filename, and instantiates `MyBmiClass()`. - -### bmi_c++ (SLOTH) -- Built-in dummy model - -The `bmi_c++` adapter name with `model_type_name: "SLOTH"` activates the built-in SLOTH (Simple Lightweight Output Testing Handler) model. SLOTH returns constant values and requires no shared library. It is configured entirely through `model_params`. - -Variables are defined using a special string format: - -``` -"variable_name(count,type,units,location)": value -``` - -| Component | Description | -|-----------|-------------| -| `variable_name` | The BMI output variable name | -| `count` | Array size (typically `1` for scalar) | -| `type` | Data type: `double`, `float`, or `int` | -| `units` | Unit string (e.g., `m`, `1`, `mm/s`) | -| `location` | BMI grid location (e.g., `node`) | -| `value` | Constant value returned for all timesteps | - -```json -{ - "name": "bmi_c++", - "params": { - "model_type_name": "SLOTH", - "init_config": "/dev/null", - "main_output_variable": "z", - "model_params": { - "sloth_ice_fraction_schaake(1,double,m,node)": 0.0, - "sloth_ice_fraction_xinanjiang(1,double,1,node)": 0.0, - "sloth_soil_moisture_profile(1,double,1,node)": 0.0 - } - } -} -``` - -## Model dependency resolution - -Models are loaded in dependency order, not config order. The runner uses `variables_names_map` to determine dependencies: - -- The **keys** are the model's input variable names. -- The **values** are source variable names (from forcings or upstream model outputs). - -The runner performs an iterative topological sort: - -1. Start with all forcing variables as available. -2. Find a module whose `variables_names_map` values are all satisfied by available variables. -3. Load that module, add its output variables to the available set. -4. Repeat until all modules are loaded. - -This means modules can be listed in any order in the config. If a circular dependency or missing variable is detected, the runner reports an error. - -## Unit conversion - -The runner automatically converts units between source variables and model inputs. When a model input is mapped to a source, the runner reads both sides' unit strings and builds a linear conversion (`output = input * scale + offset`). - -### Supported conversions - -| Category | Examples | -|----------|----------| -| Length | m, mm, cm, km, ft, in | -| Pressure | Pa, kPa, hPa, mb, atm, bar | -| Temperature | K, C, F (with offset handling for K/C/F conversions) | -| Rates | mm/s, mm/h, m/s, m s^-1, mm s-1 | -| Mass flux | kg m^-2 s^-1 to/from mm/s (assumes water density 1000 kg/m^3) | -| Dimensionless | `""`, `"1"`, `"-"`, `"m/m"`, `"none"` | - -### Notation variants - -All of these are recognized as equivalent: - -- `mm s^-1`, `mm/s`, `mm s-1` -- `kg m^-2 s^-1`, `kg/m^2/s`, `kg m-2 s-1` -- `W m^-2`, `W/m^2`, `W m-2` - -Use `--units` to print all active conversions for debugging: - -```bash -bmi-driver --units -``` - -If units are unknown or incompatible, the runner falls back to an identity conversion (no change) and prints a warning. - -## Variable aliases - -The runner knows common aliases between AORC forcing field names and CSDMS standard names. When it detects unmapped model inputs that match an available variable under a different name, it suggests adding the mapping. - -| Short name | CSDMS standard name | -|------------|-------------------| -| `precip_rate` | `atmosphere_water__liquid_equivalent_precipitation_rate` | -| `TMP_2maboveground` | `land_surface_air__temperature` | -| `UGRD_10maboveground` | `land_surface_wind__x_component_of_velocity` | -| `VGRD_10maboveground` | `land_surface_wind__y_component_of_velocity` | -| `DLWRF_surface` | `land_surface_radiation__incoming_longwave_flux` | -| `DSWRF_surface` | `land_surface_radiation__incoming_shortwave_flux` | -| `PRES_surface` | `land_surface_air__pressure` | -| `SPFH_2maboveground` | `land_surface_air__specific_humidity` | -| `APCP_surface` | `land_surface_water__precipitation_volume_flux` | - -On the first run, if unmapped inputs are found, the runner prompts: - -``` -Found unmapped model inputs that match available variables: - [1] CFE: "atmosphere_water__liquid_equivalent_precipitation_rate" <- "QINSUR" - -Add these mappings to realization.json? [y/N] -``` - -Answering `y` updates the config file in place. - -## Config minification - -Use `--minify` to strip a realization.json down to only the fields bmi-driver reads: - -```bash -bmi-driver --minify -``` - -This removes: -- ngen-specific fields (`routing`, `forcing.provider`, `allow_exceed_end_time`, `fixed_time_step`, `uses_forcing_file`, etc.) -- Unknown keys not in the config schema -- Empty default-valued fields - -The minified config is still valid for subsequent runs. - -## NetCDF forcing format - -The forcing file must be a NetCDF file with: - -- A variable `ids` containing location ID strings -- A variable `Time` containing epoch timestamps (int64) -- Dimensions `catchment-id` (locations) and `time` (timesteps) -- Data variables with dimensions `[catchment-id, time]` - -All forcing data for a location is preloaded into memory before running that location's models. +See [docs/configuration.md](docs/configuration.md) for all field reference tables and library usage examples. ## Output -Each location produces a CSV at `/outputs/bmi-driver/.csv`: - -```csv -Time Step,Time,Q_OUT,EVAPOTRANS,QINSUR,RAIN_RATE -0,2010-01-01 00:00:00,0.000000000,0.000000000,0.000000000,0.000000000 -1,2010-01-01 01:00:00,0.000012345,0.000000100,0.000001234,0.000000500 -... -``` +Output format is set via `output_format` in the config (`"csv"`, `"netcdf"`, or `"zarr"`). Default is CSV. -Columns are determined by `output_variables` in the formulation params. If `output_variables` is empty, all model output variables are included. +- **CSV**: One file per location at `outputs/bmi-driver/.csv` +- **NetCDF**: Single `outputs/bmi-driver/results.nc` with dimensions `[id, time]` +- **Zarr**: Single `outputs/bmi-driver/results.zarr` store with chunking `[1, n_times]` per variable -## Project structure +## Documentation -``` -src/ -├── main.rs # CLI entry point (bmi-driver binary) -├── lib.rs # Public exports -├── traits.rs # Bmi trait and types -├── error.rs # Error types -├── ffi.rs # C FFI struct definition -├── library.rs # Dynamic library loading (dlopen/dlsym) -├── config.rs # Configuration parsing -├── forcings.rs # NetCDF forcing data -├── runner.rs # Model orchestration and dependency resolution -├── units.rs # Automatic unit conversion -├── aliases.rs # AORC ↔ CSDMS variable name aliases -└── adapters/ - ├── mod.rs # Adapter exports - ├── c.rs # C BMI adapter - ├── fortran.rs # Fortran BMI adapter (feature-gated) - ├── python.rs # Python BMI adapter (feature-gated) - └── sloth.rs # SLOTH dummy model -``` +| Document | Contents | +|----------|----------| +| [docs/configuration.md](docs/configuration.md) | Config field reference tables, library usage | +| [docs/module-types.md](docs/module-types.md) | bmi_c, bmi_fortran, bmi_python, SLOTH details | +| [docs/features.md](docs/features.md) | Unit conversion, variable aliases, dependency resolution, config minification | +| [docs/netcdf-forcing.md](docs/netcdf-forcing.md) | NetCDF forcing file format spec | diff --git a/docs/configuration.md b/docs/configuration.md new file mode 100644 index 0000000..f1c790b --- /dev/null +++ b/docs/configuration.md @@ -0,0 +1,54 @@ +# Configuration Reference + +The config file is `/config/realization.json`. The format is compatible with [ngen](https://github.com/NOAA-OWP/ngen) realization configs -- extra ngen-specific fields are silently ignored. + +## Top-level fields + +| Field | Required | Description | +|-------|----------|-------------| +| `global.formulations` | yes | Array with one `bmi_multi` formulation entry | +| `global.forcing.path` | yes | Path to NetCDF forcing file (relative to data_dir) | +| `time.start_time` | yes | Simulation start, `"YYYY-MM-DD HH:MM:SS"` | +| `time.end_time` | yes | Simulation end, `"YYYY-MM-DD HH:MM:SS"` | +| `time.output_interval` | no | Output interval in seconds (default: 3600) | + +## Formulation fields + +| Field | Required | Description | +|-------|----------|-------------| +| `name` | yes | Must be `"bmi_multi"` for multi-model orchestration | +| `params.modules` | yes | Array of module configs (see below) | +| `params.main_output_variable` | no | Primary output variable name | +| `params.output_variables` | no | Variables to include in output. If empty, all model outputs are written. | + +## Module fields (shared across all adapter types) + +| Field | Required | Description | +|-------|----------|-------------| +| `name` | yes | Adapter type: `"bmi_c"`, `"bmi_fortran"`, `"bmi_python"`, or `"bmi_c++"` | +| `params.model_type_name` | yes | Model identifier (e.g., `"CFE"`, `"NoahOWP"`, `"SLOTH"`) | +| `params.init_config` | yes | Path to model config file. `{{id}}` is replaced with the location ID at runtime. | +| `params.main_output_variable` | yes | The primary output variable of this model | +| `params.library_file` | depends | Path to shared library. Required for C and Fortran adapters. | +| `params.registration_function` | no | Name of the C registration function (default: `"register_bmi"`) | +| `params.variables_names_map` | no | Maps model input names to source variable names (from forcings or upstream models) | +| `params.model_params` | no | Key-value parameters set on the model after initialization via `set_value()` | +| `params.python_type` | no | Python class path for the Python adapter (e.g., `"lstm.bmi_lstm.bmi_LSTM"`) | + +## Library usage + +```rust +use bmi_driver::{ModelRunner, BmiError}; + +fn main() -> Result<(), BmiError> { + let mut runner = ModelRunner::from_config("config/realization.json")?; + runner.initialize("cat-123")?; + runner.run()?; + + let outputs = runner.main_outputs()?; + println!("Completed {} timesteps", outputs.len()); + + runner.finalize()?; + Ok(()) +} +``` diff --git a/docs/features.md b/docs/features.md new file mode 100644 index 0000000..773ab60 --- /dev/null +++ b/docs/features.md @@ -0,0 +1,85 @@ +# Features + +## Unit conversion + +The runner automatically converts units between source variables and model inputs. When a model input is mapped to a source, the runner reads both sides' unit strings and builds a linear conversion (`output = input * scale + offset`). + +### Supported conversions + +| Category | Examples | +|----------|----------| +| Length | m, mm, cm, km, ft, in | +| Pressure | Pa, kPa, hPa, mb, atm, bar | +| Temperature | K, C, F (with offset handling for K/C/F conversions) | +| Rates | mm/s, mm/h, m/s, m s^-1, mm s-1 | +| Mass flux | kg m^-2 s^-1 to/from mm/s (assumes water density 1000 kg/m^3) | +| Dimensionless | `""`, `"1"`, `"-"`, `"m/m"`, `"none"` | + +### Notation variants + +All of these are recognized as equivalent: + +- `mm s^-1`, `mm/s`, `mm s-1` +- `kg m^-2 s^-1`, `kg/m^2/s`, `kg m-2 s-1` +- `W m^-2`, `W/m^2`, `W m-2` + +Use `--units` to print all active conversions for debugging: + +```bash +bmi-driver --units +``` + +If units are unknown or incompatible, the runner falls back to an identity conversion (no change) and prints a warning. + +## Variable aliases + +The runner knows common aliases between AORC forcing field names and CSDMS standard names. When it detects unmapped model inputs that match an available variable under a different name, it suggests adding the mapping. + +| Short name | CSDMS standard name | +|------------|-------------------| +| `precip_rate` | `atmosphere_water__liquid_equivalent_precipitation_rate` | +| `TMP_2maboveground` | `land_surface_air__temperature` | +| `UGRD_10maboveground` | `land_surface_wind__x_component_of_velocity` | +| `VGRD_10maboveground` | `land_surface_wind__y_component_of_velocity` | +| `DLWRF_surface` | `land_surface_radiation__incoming_longwave_flux` | +| `DSWRF_surface` | `land_surface_radiation__incoming_shortwave_flux` | +| `PRES_surface` | `land_surface_air__pressure` | +| `SPFH_2maboveground` | `land_surface_air__specific_humidity` | +| `APCP_surface` | `land_surface_water__precipitation_volume_flux` | + +On the first run, if unmapped inputs are found, the runner prompts: + +``` +Found unmapped model inputs that match available variables: + [1] CFE: "atmosphere_water__liquid_equivalent_precipitation_rate" <- "QINSUR" + +Add these mappings to realization.json? [y/N] +``` + +Answering `y` updates the config file in place. + +## Model dependency resolution + +Models are loaded in dependency order, not config order. The runner uses `variables_names_map` to determine dependencies: + +- The **keys** are the model's input variable names. +- The **values** are source variable names (from forcings or upstream model outputs). + +The runner performs an iterative topological sort: + +1. Start with all forcing variables as available. +2. Find a module whose `variables_names_map` values are all satisfied by available variables. +3. Load that module, add its output variables to the available set. +4. Repeat until all modules are loaded. + +This means modules can be listed in any order in the config. If a circular dependency or missing variable is detected, the runner reports an error. + +## Config minification + +Use `--minify` to strip a realization.json down to only the fields bmi-driver reads: + +```bash +bmi-driver --minify +``` + +This removes ngen-specific fields (`routing`, `forcing.provider`, `allow_exceed_end_time`, `fixed_time_step`, `uses_forcing_file`, etc.), unknown keys not in the config schema, and empty default-valued fields. diff --git a/docs/module-types.md b/docs/module-types.md new file mode 100644 index 0000000..2df9d62 --- /dev/null +++ b/docs/module-types.md @@ -0,0 +1,141 @@ +# Module Types + +## bmi_c -- C shared libraries + +Loads a C shared library via `dlopen` with `RTLD_GLOBAL` (so symbols are shared across libraries). Looks up the registration function by name, calls it to populate a struct of BMI function pointers, then calls through those pointers for all BMI operations. + +```json +{ + "name": "bmi_c", + "params": { + "model_type_name": "CFE", + "library_file": "/path/to/libcfebmi.so", + "init_config": "./config/cat_config/CFE/{{id}}.ini", + "registration_function": "register_bmi_cfe", + "main_output_variable": "Q_OUT", + "variables_names_map": { + "atmosphere_water__liquid_equivalent_precipitation_rate": "QINSUR", + "water_potential_evaporation_flux": "EVAPOTRANS" + } + } +} +``` + +The shared library must export a registration function matching this signature: + +```c +Bmi* register_bmi_cfe(Bmi* model); +``` + +This function fills in the `Bmi` struct's function pointers and returns the pointer. The default function name is `register_bmi` if `registration_function` is omitted. + +## bmi_fortran -- Fortran shared libraries + +**Requires:** `--features fortran` + +Loads Fortran BMI models. The adapter supports two modes: + +1. **Single library** -- the `.so` implements both the registration function and the BMI interface directly. +2. **With middleware** -- a separate C-to-Fortran middleware library translates C-calling-convention calls to Fortran. The runner auto-detects a middleware library if one exists alongside the model library. + +The adapter passes an opaque handle pointer as the first argument to all BMI calls, matching the convention used by Fortran BMI implementations. + +```json +{ + "name": "bmi_fortran", + "params": { + "model_type_name": "NoahOWP", + "library_file": "/path/to/libsurfacebmi.so", + "init_config": "./config/cat_config/NOAH-OWP-M/{{id}}.input", + "main_output_variable": "QINSUR", + "variables_names_map": { + "PRCPNONC": "precip_rate", + "SFCTMP": "TMP_2maboveground", + "SFCPRS": "PRES_surface", + "Q2": "SPFH_2maboveground", + "UU": "UGRD_10maboveground", + "VV": "VGRD_10maboveground", + "LWDN": "DLWRF_surface", + "SOLDN": "DSWRF_surface" + } + } +} +``` + +## bmi_python -- Python BMI classes + +**Requires:** `--features python` + +Loads a Python class that implements the BMI interface. Values are exchanged via numpy arrays. There are two ways to specify the Python class: + +**Option 1: `python_type` (preferred)** -- a dotted module path where the last component is the class name: + +```json +{ + "name": "bmi_python", + "params": { + "model_type_name": "LSTM", + "python_type": "bmi_lstm.bmi_LSTM", + "init_config": "./config/cat_config/LSTM/{{id}}.yml", + "main_output_variable": "land_surface_water__runoff_volume_flux", + "variables_names_map": { + "atmosphere_water__liquid_equivalent_precipitation_rate": "precip_rate", + "land_surface_air__temperature": "TMP_2maboveground" + } + } +} +``` + +This imports `bmi_lstm` and instantiates `bmi_LSTM()`. The package must be installed or on `PYTHONPATH`. + +**Option 2: `library_file` + `registration_function`** -- loads a `.py` file directly: + +```json +{ + "name": "bmi_python", + "params": { + "model_type_name": "MyModel", + "library_file": "/path/to/my_model.py", + "registration_function": "MyBmiClass", + "init_config": "./config/cat_config/MyModel/{{id}}.yml", + "main_output_variable": "output_var" + } +} +``` + +This reads the file, adds its parent directory to `sys.path`, imports the module by filename, and instantiates `MyBmiClass()`. + +## bmi_c++ (SLOTH) -- Built-in dummy model + +The `bmi_c++` adapter name with `model_type_name: "SLOTH"` activates the built-in SLOTH (Simple Lightweight Output Testing Handler) model. SLOTH returns constant values and requires no shared library. It is configured entirely through `model_params`. + +Variables are defined using a special string format: + +``` +"variable_name(count,type,units,location)": value +``` + +| Component | Description | +|-----------|-------------| +| `variable_name` | The BMI output variable name | +| `count` | Array size (typically `1` for scalar) | +| `type` | Data type: `double`, `float`, or `int` | +| `units` | Unit string (e.g., `m`, `1`, `mm/s`) | +| `location` | BMI grid location (e.g., `node`) | +| `value` | Constant value returned for all timesteps | + +```json +{ + "name": "bmi_c++", + "params": { + "model_type_name": "SLOTH", + "init_config": "/dev/null", + "main_output_variable": "z", + "model_params": { + "sloth_ice_fraction_schaake(1,double,m,node)": 0.0, + "sloth_ice_fraction_xinanjiang(1,double,1,node)": 0.0, + "sloth_soil_moisture_profile(1,double,1,node)": 0.0 + } + } +} +``` diff --git a/docs/netcdf-forcing.md b/docs/netcdf-forcing.md new file mode 100644 index 0000000..660dc1f --- /dev/null +++ b/docs/netcdf-forcing.md @@ -0,0 +1,10 @@ +# NetCDF Forcing Format + +The forcing file must be a NetCDF file with: + +- A variable `ids` containing location ID strings +- A variable `Time` containing epoch timestamps (int64) +- Dimensions `catchment-id` (locations) and `time` (timesteps) +- Data variables with dimensions `[catchment-id, time]` + +All forcing data for a location is preloaded into memory before running that location's models. diff --git a/src/adapters/c.rs b/src/adapters/c.rs index f4bea03..f11146f 100644 --- a/src/adapters/c.rs +++ b/src/adapters/c.rs @@ -1,10 +1,9 @@ use std::collections::HashMap; -use std::ffi::CString; use std::os::raw::{c_char, c_void}; use std::path::Path; -use super::{check_initialized, cstr_to_string}; -use crate::error::{BmiError, BmiResult}; +use super::{check_initialized, cstr_to_string, impl_bmi_drop, to_cstring, verify_config_path}; +use crate::error::{function_failed, not_implemented, BmiError, BmiResult}; use crate::ffi::{Bmi as BmiStruct, BMI_MAX_NAME, BMI_SUCCESS}; use crate::library::GlobalLibrary; use crate::traits::{parse_time_units, Bmi, VarType}; @@ -100,16 +99,10 @@ impl BmiC { where F: FnOnce() -> Option i32>, { - let func = f().ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: func_name.into(), - })?; + let func = f().ok_or_else(|| not_implemented(&self.name, func_name))?; let mut buf = vec![0u8; BMI_MAX_NAME]; if unsafe { func(self.ptr(), buf.as_mut_ptr() as *mut c_char) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: func_name.into(), - }); + return Err(function_failed(&self.name, func_name)); } cstr_to_string(&buf) } @@ -118,16 +111,10 @@ impl BmiC { where F: FnOnce() -> Option i32>, { - let func = f().ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: func_name.into(), - })?; + let func = f().ok_or_else(|| not_implemented(&self.name, func_name))?; let mut val = 0i32; if unsafe { func(self.ptr(), &mut val) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: func_name.into(), - }); + return Err(function_failed(&self.name, func_name)); } Ok(val) } @@ -136,16 +123,10 @@ impl BmiC { where F: FnOnce() -> Option i32>, { - let func = f().ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: func_name.into(), - })?; + let func = f().ok_or_else(|| not_implemented(&self.name, func_name))?; let mut val = 0.0f64; if unsafe { func(self.ptr(), &mut val) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: func_name.into(), - }); + return Err(function_failed(&self.name, func_name)); } Ok(val) } @@ -156,19 +137,13 @@ impl BmiC { func: Option i32>, name: &str, ) -> BmiResult { - let func = func.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: func_name.into(), - })?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let func = func.ok_or_else(|| not_implemented(&self.name, func_name))?; + let cname = to_cstring(name)?; let mut buf = vec![0u8; BMI_MAX_NAME]; if unsafe { func(self.ptr(), cname.as_ptr(), buf.as_mut_ptr() as *mut c_char) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: func_name.into(), - }); + return Err(function_failed(&self.name, func_name)); } cstr_to_string(&buf) } @@ -179,17 +154,11 @@ impl BmiC { func: Option i32>, name: &str, ) -> BmiResult { - let func = func.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: func_name.into(), - })?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let func = func.ok_or_else(|| not_implemented(&self.name, func_name))?; + let cname = to_cstring(name)?; let mut val = 0i32; if unsafe { func(self.ptr(), cname.as_ptr(), &mut val) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: func_name.into(), - }); + return Err(function_failed(&self.name, func_name)); } Ok(val) } @@ -200,20 +169,14 @@ impl BmiC { func: Option i32>, count: usize, ) -> BmiResult> { - let func = func.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: func_name.into(), - })?; + let func = func.ok_or_else(|| not_implemented(&self.name, func_name))?; let mut bufs: Vec> = (0..count).map(|_| vec![0u8; BMI_MAX_NAME]).collect(); let mut ptrs: Vec<*mut c_char> = bufs .iter_mut() .map(|b| b.as_mut_ptr() as *mut c_char) .collect(); if unsafe { func(self.ptr(), ptrs.as_mut_ptr()) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: func_name.into(), - }); + return Err(function_failed(&self.name, func_name)); } bufs.iter().map(|b| cstr_to_string(b)).collect() } @@ -242,25 +205,15 @@ impl Bmi for BmiC { model: self.name.clone(), }); } - if !Path::new(config).exists() { - return Err(BmiError::ConfigNotFound { - path: config.into(), - }); - } + verify_config_path(config)?; let func = self .bmi .initialize - .ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "initialize".into(), - })?; - let cconfig = CString::new(config).map_err(|_| BmiError::InvalidUtf8)?; + .ok_or_else(|| not_implemented(&self.name, "initialize"))?; + let cconfig = to_cstring(config)?; if unsafe { func(self.bmi.as_mut(), cconfig.as_ptr()) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "initialize".into(), - }); + return Err(function_failed(&self.name, "initialize")); } self.initialized = true; @@ -274,15 +227,12 @@ impl Bmi for BmiC { fn update(&mut self) -> BmiResult<()> { check_initialized(self.initialized, &self.name)?; - let func = self.bmi.update.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "update".into(), - })?; + let func = self + .bmi + .update + .ok_or_else(|| not_implemented(&self.name, "update"))?; if unsafe { func(self.bmi.as_mut()) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "update".into(), - }); + return Err(function_failed(&self.name, "update")); } Ok(()) } @@ -292,15 +242,9 @@ impl Bmi for BmiC { let func = self .bmi .update_until - .ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "update_until".into(), - })?; + .ok_or_else(|| not_implemented(&self.name, "update_until"))?; if unsafe { func(self.bmi.as_mut(), time) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "update_until".into(), - }); + return Err(function_failed(&self.name, "update_until")); } Ok(()) } @@ -309,17 +253,14 @@ impl Bmi for BmiC { if !self.initialized { return Ok(()); } - let func = self.bmi.finalize.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "finalize".into(), - })?; + let func = self + .bmi + .finalize + .ok_or_else(|| not_implemented(&self.name, "finalize"))?; let result = unsafe { func(self.bmi.as_mut()) }; self.initialized = false; if result != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "finalize".into(), - }); + return Err(function_failed(&self.name, "finalize")); } Ok(()) } @@ -405,73 +346,68 @@ impl Bmi for BmiC { self.call_double("get_time_step", || self.bmi.get_time_step) } + // TODO: get_value_f64/f32/i32 and set_value_f64/f32/i32 could use a generic helper + // to reduce duplication across the three typed variants. Deferred because the unsafe + // FFI calls need careful attention. + fn get_value_f64(&self, name: &str) -> BmiResult> { check_initialized(self.initialized, &self.name)?; - let func = self.bmi.get_value.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "get_value".into(), - })?; + let func = self + .bmi + .get_value + .ok_or_else(|| not_implemented(&self.name, "get_value"))?; let count = self.get_var_nbytes(name)? as usize / std::mem::size_of::(); let mut vals = vec![0.0f64; count]; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { func(self.ptr(), cname.as_ptr(), vals.as_mut_ptr() as *mut c_void) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "get_value".into(), - }); + return Err(function_failed(&self.name, "get_value")); } Ok(vals) } fn get_value_f32(&self, name: &str) -> BmiResult> { check_initialized(self.initialized, &self.name)?; - let func = self.bmi.get_value.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "get_value".into(), - })?; + let func = self + .bmi + .get_value + .ok_or_else(|| not_implemented(&self.name, "get_value"))?; let count = self.get_var_nbytes(name)? as usize / std::mem::size_of::(); let mut vals = vec![0.0f32; count]; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { func(self.ptr(), cname.as_ptr(), vals.as_mut_ptr() as *mut c_void) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "get_value".into(), - }); + return Err(function_failed(&self.name, "get_value")); } Ok(vals) } fn get_value_i32(&self, name: &str) -> BmiResult> { check_initialized(self.initialized, &self.name)?; - let func = self.bmi.get_value.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "get_value".into(), - })?; + let func = self + .bmi + .get_value + .ok_or_else(|| not_implemented(&self.name, "get_value"))?; let count = self.get_var_nbytes(name)? as usize / std::mem::size_of::(); let mut vals = vec![0i32; count]; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { func(self.ptr(), cname.as_ptr(), vals.as_mut_ptr() as *mut c_void) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "get_value".into(), - }); + return Err(function_failed(&self.name, "get_value")); } Ok(vals) } fn set_value_f64(&mut self, name: &str, values: &[f64]) -> BmiResult<()> { check_initialized(self.initialized, &self.name)?; - let func = self.bmi.set_value.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "set_value".into(), - })?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let func = self + .bmi + .set_value + .ok_or_else(|| not_implemented(&self.name, "set_value"))?; + let cname = to_cstring(name)?; if unsafe { func( self.bmi.as_mut(), @@ -480,21 +416,18 @@ impl Bmi for BmiC { ) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "set_value".into(), - }); + return Err(function_failed(&self.name, "set_value")); } Ok(()) } fn set_value_f32(&mut self, name: &str, values: &[f32]) -> BmiResult<()> { check_initialized(self.initialized, &self.name)?; - let func = self.bmi.set_value.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "set_value".into(), - })?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let func = self + .bmi + .set_value + .ok_or_else(|| not_implemented(&self.name, "set_value"))?; + let cname = to_cstring(name)?; if unsafe { func( self.bmi.as_mut(), @@ -503,21 +436,18 @@ impl Bmi for BmiC { ) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "set_value".into(), - }); + return Err(function_failed(&self.name, "set_value")); } Ok(()) } fn set_value_i32(&mut self, name: &str, values: &[i32]) -> BmiResult<()> { check_initialized(self.initialized, &self.name)?; - let func = self.bmi.set_value.ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "set_value".into(), - })?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let func = self + .bmi + .set_value + .ok_or_else(|| not_implemented(&self.name, "set_value"))?; + let cname = to_cstring(name)?; if unsafe { func( self.bmi.as_mut(), @@ -526,10 +456,7 @@ impl Bmi for BmiC { ) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "set_value".into(), - }); + return Err(function_failed(&self.name, "set_value")); } Ok(()) } @@ -539,16 +466,10 @@ impl Bmi for BmiC { let func = self .bmi .get_grid_rank - .ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "get_grid_rank".into(), - })?; + .ok_or_else(|| not_implemented(&self.name, "get_grid_rank"))?; let mut val = 0i32; if unsafe { func(self.ptr(), grid, &mut val) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "get_grid_rank".into(), - }); + return Err(function_failed(&self.name, "get_grid_rank")); } Ok(val) } @@ -558,16 +479,10 @@ impl Bmi for BmiC { let func = self .bmi .get_grid_size - .ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "get_grid_size".into(), - })?; + .ok_or_else(|| not_implemented(&self.name, "get_grid_size"))?; let mut val = 0i32; if unsafe { func(self.ptr(), grid, &mut val) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "get_grid_size".into(), - }); + return Err(function_failed(&self.name, "get_grid_size")); } Ok(val) } @@ -577,25 +492,13 @@ impl Bmi for BmiC { let func = self .bmi .get_grid_type - .ok_or_else(|| BmiError::NotImplemented { - model: self.name.clone(), - func: "get_grid_type".into(), - })?; + .ok_or_else(|| not_implemented(&self.name, "get_grid_type"))?; let mut buf = vec![0u8; BMI_MAX_NAME]; if unsafe { func(self.ptr(), grid, buf.as_mut_ptr() as *mut c_char) } != BMI_SUCCESS { - return Err(BmiError::FunctionFailed { - model: self.name.clone(), - func: "get_grid_type".into(), - }); + return Err(function_failed(&self.name, "get_grid_type")); } cstr_to_string(&buf) } } -impl Drop for BmiC { - fn drop(&mut self) { - if self.initialized { - let _ = self.finalize(); - } - } -} +impl_bmi_drop!(BmiC); diff --git a/src/adapters/fortran.rs b/src/adapters/fortran.rs index fa19cfc..9e78bc2 100644 --- a/src/adapters/fortran.rs +++ b/src/adapters/fortran.rs @@ -1,10 +1,9 @@ use std::collections::HashMap; -use std::ffi::CString; use std::os::raw::{c_char, c_double, c_float, c_int, c_void}; use std::path::Path; -use super::{check_initialized, cstr_to_string}; -use crate::error::{BmiError, BmiResult}; +use super::{check_initialized, cstr_to_string, impl_bmi_drop, to_cstring, verify_config_path}; +use crate::error::{function_failed, BmiError, BmiResult}; use crate::ffi::BMI_MAX_NAME; use crate::library::GlobalLibrary; use crate::traits::{parse_time_units, Bmi, VarType}; @@ -160,10 +159,7 @@ impl BmiFortran { } if handle.is_null() { - return Err(BmiError::FunctionFailed { - model: name, - func: reg_func.into(), - }); + return Err(function_failed(name, reg_func)); } let fns = Self::load_fns(&middleware_lib)?; @@ -205,10 +201,7 @@ impl BmiFortran { } if handle.is_null() { - return Err(BmiError::FunctionFailed { - model: name, - func: reg_func.into(), - }); + return Err(function_failed(name, reg_func)); } let fns = Self::load_fns(&lib)?; @@ -271,12 +264,6 @@ impl BmiFortran { &mut self.handle as *mut *mut c_void as *mut c_void } - fn err(&self, func: &str) -> BmiError { - BmiError::FunctionFailed { - model: self.model_name.clone(), - func: func.into(), - } - } } impl Bmi for BmiFortran { @@ -302,15 +289,11 @@ impl Bmi for BmiFortran { model: self.model_name.clone(), }); } - if !Path::new(config).exists() { - return Err(BmiError::ConfigNotFound { - path: config.into(), - }); - } + verify_config_path(config)?; - let cconfig = CString::new(config).map_err(|_| BmiError::InvalidUtf8)?; + let cconfig = to_cstring(config)?; if unsafe { (self.fns.initialize)(self.handle_ptr_mut(), cconfig.as_ptr()) } != 0 { - return Err(self.err("initialize")); + return Err(function_failed(&self.model_name,"initialize")); } self.initialized = true; @@ -325,7 +308,7 @@ impl Bmi for BmiFortran { fn update(&mut self) -> BmiResult<()> { check_initialized(self.initialized, &self.model_name)?; if unsafe { (self.fns.update)(self.handle_ptr_mut()) } != 0 { - return Err(self.err("update")); + return Err(function_failed(&self.model_name,"update")); } Ok(()) } @@ -334,7 +317,7 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let mut t = time; if unsafe { (self.fns.update_until)(self.handle_ptr_mut(), &mut t) } != 0 { - return Err(self.err("update_until")); + return Err(function_failed(&self.model_name,"update_until")); } Ok(()) } @@ -346,7 +329,7 @@ impl Bmi for BmiFortran { let result = unsafe { (self.fns.finalize)(self.handle_ptr_mut()) }; self.initialized = false; if result != 0 { - return Err(self.err("finalize")); + return Err(function_failed(&self.model_name,"finalize")); } Ok(()) } @@ -357,7 +340,7 @@ impl Bmi for BmiFortran { (self.fns.get_component_name)(self.handle_ptr(), buf.as_mut_ptr() as *mut c_char) } != 0 { - return Err(self.err("get_component_name")); + return Err(function_failed(&self.model_name,"get_component_name")); } cstr_to_string(&buf) } @@ -366,7 +349,7 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let mut count = 0; if unsafe { (self.fns.get_input_item_count)(self.handle_ptr(), &mut count) } != 0 { - return Err(self.err("get_input_item_count")); + return Err(function_failed(&self.model_name,"get_input_item_count")); } Ok(count) } @@ -375,7 +358,7 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let mut count = 0; if unsafe { (self.fns.get_output_item_count)(self.handle_ptr(), &mut count) } != 0 { - return Err(self.err("get_output_item_count")); + return Err(function_failed(&self.model_name,"get_output_item_count")); } Ok(count) } @@ -389,7 +372,7 @@ impl Bmi for BmiFortran { .map(|b| b.as_mut_ptr() as *mut c_char) .collect(); if unsafe { (self.fns.get_input_var_names)(self.handle_ptr(), ptrs.as_mut_ptr()) } != 0 { - return Err(self.err("get_input_var_names")); + return Err(function_failed(&self.model_name,"get_input_var_names")); } bufs.iter().map(|b| cstr_to_string(b)).collect() } @@ -403,24 +386,24 @@ impl Bmi for BmiFortran { .map(|b| b.as_mut_ptr() as *mut c_char) .collect(); if unsafe { (self.fns.get_output_var_names)(self.handle_ptr(), ptrs.as_mut_ptr()) } != 0 { - return Err(self.err("get_output_var_names")); + return Err(function_failed(&self.model_name,"get_output_var_names")); } bufs.iter().map(|b| cstr_to_string(b)).collect() } fn get_var_grid(&self, name: &str) -> BmiResult { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; let mut grid = 0; if unsafe { (self.fns.get_var_grid)(self.handle_ptr(), cname.as_ptr(), &mut grid) } != 0 { - return Err(self.err("get_var_grid")); + return Err(function_failed(&self.model_name,"get_var_grid")); } Ok(grid) } fn get_var_type(&self, name: &str) -> BmiResult { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; let mut buf = vec![0u8; BMI_MAX_NAME]; if unsafe { (self.fns.get_var_type)( @@ -430,14 +413,14 @@ impl Bmi for BmiFortran { ) } != 0 { - return Err(self.err("get_var_type")); + return Err(function_failed(&self.model_name,"get_var_type")); } cstr_to_string(&buf) } fn get_var_units(&self, name: &str) -> BmiResult { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; let mut buf = vec![0u8; BMI_MAX_NAME]; if unsafe { (self.fns.get_var_units)( @@ -447,36 +430,36 @@ impl Bmi for BmiFortran { ) } != 0 { - return Err(self.err("get_var_units")); + return Err(function_failed(&self.model_name,"get_var_units")); } cstr_to_string(&buf) } fn get_var_itemsize(&self, name: &str) -> BmiResult { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; let mut size = 0; if unsafe { (self.fns.get_var_itemsize)(self.handle_ptr(), cname.as_ptr(), &mut size) } != 0 { - return Err(self.err("get_var_itemsize")); + return Err(function_failed(&self.model_name,"get_var_itemsize")); } Ok(size) } fn get_var_nbytes(&self, name: &str) -> BmiResult { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; let mut nbytes = 0; if unsafe { (self.fns.get_var_nbytes)(self.handle_ptr(), cname.as_ptr(), &mut nbytes) } != 0 { - return Err(self.err("get_var_nbytes")); + return Err(function_failed(&self.model_name,"get_var_nbytes")); } Ok(nbytes) } fn get_var_location(&self, name: &str) -> BmiResult { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; let mut buf = vec![0u8; BMI_MAX_NAME]; if unsafe { (self.fns.get_var_location)( @@ -486,7 +469,7 @@ impl Bmi for BmiFortran { ) } != 0 { - return Err(self.err("get_var_location")); + return Err(function_failed(&self.model_name,"get_var_location")); } cstr_to_string(&buf) } @@ -495,7 +478,7 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let mut time = 0.0; if unsafe { (self.fns.get_current_time)(self.handle_ptr(), &mut time) } != 0 { - return Err(self.err("get_current_time")); + return Err(function_failed(&self.model_name,"get_current_time")); } Ok(time) } @@ -504,7 +487,7 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let mut time = 0.0; if unsafe { (self.fns.get_start_time)(self.handle_ptr(), &mut time) } != 0 { - return Err(self.err("get_start_time")); + return Err(function_failed(&self.model_name,"get_start_time")); } Ok(time) } @@ -513,7 +496,7 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let mut time = 0.0; if unsafe { (self.fns.get_end_time)(self.handle_ptr(), &mut time) } != 0 { - return Err(self.err("get_end_time")); + return Err(function_failed(&self.model_name,"get_end_time")); } Ok(time) } @@ -524,7 +507,7 @@ impl Bmi for BmiFortran { if unsafe { (self.fns.get_time_units)(self.handle_ptr(), buf.as_mut_ptr() as *mut c_char) } != 0 { - return Err(self.err("get_time_units")); + return Err(function_failed(&self.model_name,"get_time_units")); } cstr_to_string(&buf) } @@ -533,21 +516,22 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let mut ts = 0.0; if unsafe { (self.fns.get_time_step)(self.handle_ptr(), &mut ts) } != 0 { - return Err(self.err("get_time_step")); + return Err(function_failed(&self.model_name,"get_time_step")); } Ok(ts) } + // TODO: get_value_f64/f32/i32 and set_value_f64/f32/i32 could use a generic helper to reduce duplication. Deferred because the unsafe FFI calls need careful attention. fn get_value_f64(&self, name: &str) -> BmiResult> { check_initialized(self.initialized, &self.model_name)?; let count = self.get_var_nbytes(name)? as usize / 8; let mut vals = vec![0.0f64; count]; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { (self.fns.get_value_double)(self.handle_ptr(), cname.as_ptr(), vals.as_mut_ptr()) } != 0 { - return Err(self.err("get_value_double")); + return Err(function_failed(&self.model_name,"get_value_double")); } Ok(vals) } @@ -556,12 +540,12 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let count = self.get_var_nbytes(name)? as usize / 4; let mut vals = vec![0.0f32; count]; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { (self.fns.get_value_float)(self.handle_ptr(), cname.as_ptr(), vals.as_mut_ptr()) } != 0 { - return Err(self.err("get_value_float")); + return Err(function_failed(&self.model_name,"get_value_float")); } Ok(vals) } @@ -570,18 +554,18 @@ impl Bmi for BmiFortran { check_initialized(self.initialized, &self.model_name)?; let count = self.get_var_nbytes(name)? as usize / 4; let mut vals = vec![0i32; count]; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { (self.fns.get_value_int)(self.handle_ptr(), cname.as_ptr(), vals.as_mut_ptr()) } != 0 { - return Err(self.err("get_value_int")); + return Err(function_failed(&self.model_name,"get_value_int")); } Ok(vals) } fn set_value_f64(&mut self, name: &str, values: &[f64]) -> BmiResult<()> { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { (self.fns.set_value_double)( self.handle_ptr_mut(), @@ -590,14 +574,14 @@ impl Bmi for BmiFortran { ) } != 0 { - return Err(self.err("set_value_double")); + return Err(function_failed(&self.model_name,"set_value_double")); } Ok(()) } fn set_value_f32(&mut self, name: &str, values: &[f32]) -> BmiResult<()> { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { (self.fns.set_value_float)( self.handle_ptr_mut(), @@ -606,14 +590,14 @@ impl Bmi for BmiFortran { ) } != 0 { - return Err(self.err("set_value_float")); + return Err(function_failed(&self.model_name,"set_value_float")); } Ok(()) } fn set_value_i32(&mut self, name: &str, values: &[i32]) -> BmiResult<()> { check_initialized(self.initialized, &self.model_name)?; - let cname = CString::new(name).map_err(|_| BmiError::InvalidUtf8)?; + let cname = to_cstring(name)?; if unsafe { (self.fns.set_value_int)( self.handle_ptr_mut(), @@ -622,7 +606,7 @@ impl Bmi for BmiFortran { ) } != 0 { - return Err(self.err("set_value_int")); + return Err(function_failed(&self.model_name,"set_value_int")); } Ok(()) } @@ -632,7 +616,7 @@ impl Bmi for BmiFortran { let mut g = grid; let mut rank = 0; if unsafe { (self.fns.get_grid_rank)(self.handle_ptr(), &mut g, &mut rank) } != 0 { - return Err(self.err("get_grid_rank")); + return Err(function_failed(&self.model_name,"get_grid_rank")); } Ok(rank) } @@ -642,7 +626,7 @@ impl Bmi for BmiFortran { let mut g = grid; let mut size = 0; if unsafe { (self.fns.get_grid_size)(self.handle_ptr(), &mut g, &mut size) } != 0 { - return Err(self.err("get_grid_size")); + return Err(function_failed(&self.model_name,"get_grid_size")); } Ok(size) } @@ -655,16 +639,10 @@ impl Bmi for BmiFortran { (self.fns.get_grid_type)(self.handle_ptr(), &mut g, buf.as_mut_ptr() as *mut c_char) } != 0 { - return Err(self.err("get_grid_type")); + return Err(function_failed(&self.model_name,"get_grid_type")); } cstr_to_string(&buf) } } -impl Drop for BmiFortran { - fn drop(&mut self) { - if self.initialized { - let _ = self.finalize(); - } - } -} +impl_bmi_drop!(BmiFortran); diff --git a/src/adapters/mod.rs b/src/adapters/mod.rs index 52f3039..1dd7d85 100644 --- a/src/adapters/mod.rs +++ b/src/adapters/mod.rs @@ -18,8 +18,9 @@ pub use python::BmiPython; pub use sloth::BmiSloth; use crate::error::{BmiError, BmiResult}; -use std::ffi::CStr; +use std::ffi::{CStr, CString}; use std::os::raw::c_char; +use std::path::Path; pub fn cstr_to_string(buffer: &[u8]) -> BmiResult { unsafe { CStr::from_ptr(buffer.as_ptr() as *const c_char) } @@ -37,3 +38,29 @@ pub fn check_initialized(initialized: bool, model: &str) -> BmiResult<()> { }) } } + +pub fn to_cstring(s: &str) -> BmiResult { + CString::new(s).map_err(|_| BmiError::InvalidUtf8) +} + +pub fn verify_config_path(config: &str) -> BmiResult<()> { + if !Path::new(config).exists() { + return Err(BmiError::ConfigNotFound { + path: config.into(), + }); + } + Ok(()) +} + +macro_rules! impl_bmi_drop { + ($type:ty) => { + impl Drop for $type { + fn drop(&mut self) { + if self.is_initialized() { + let _ = self.finalize(); + } + } + } + }; +} +pub(crate) use impl_bmi_drop; diff --git a/src/adapters/python.rs b/src/adapters/python.rs index 005d9e5..f868b6e 100644 --- a/src/adapters/python.rs +++ b/src/adapters/python.rs @@ -3,8 +3,8 @@ use std::path::Path; use pyo3::prelude::*; -use super::check_initialized; -use crate::error::{BmiError, BmiResult}; +use super::{check_initialized, impl_bmi_drop, verify_config_path}; +use crate::error::{function_failed, BmiError, BmiResult}; use crate::traits::{parse_time_units, Bmi, VarType}; /// BMI adapter for Python classes implementing the BMI interface. @@ -76,13 +76,13 @@ impl BmiPython { let dot_pos = python_type .rfind('.') - .ok_or_else(|| BmiError::FunctionFailed { - model: name.clone(), - func: format!( + .ok_or_else(|| function_failed( + name.clone(), + format!( "python_type '{}' must be a dotted path like 'package.module.ClassName'", python_type ), - })?; + ))?; let module_path = &python_type[..dot_pos]; let class_name = &python_type[dot_pos + 1..]; @@ -90,10 +90,7 @@ impl BmiPython { let instance = py.import_bound(module_path)?.getattr(class_name)?.call0()?; Ok(instance.unbind()) }) - .map_err(|e| BmiError::FunctionFailed { - model: name.clone(), - func: format!("load_from_type: {e}"), - })?; + .map_err(|e| function_failed(name.clone(), format!("load_from_type: {e}")))?; Ok(Self { model_name: name, @@ -127,10 +124,7 @@ impl BmiPython { let instance = py.import_bound(stem)?.getattr(class_name)?.call0()?; Ok(instance.unbind()) }) - .map_err(|e| BmiError::FunctionFailed { - model: name.clone(), - func: format!("load: {e}"), - })?; + .map_err(|e| function_failed(name.clone(), format!("load: {e}")))?; Ok(Self { model_name: name, @@ -241,11 +235,7 @@ impl Bmi for BmiPython { model: self.model_name.clone(), }); } - if !Path::new(config).exists() { - return Err(BmiError::ConfigNotFound { - path: config.into(), - }); - } + verify_config_path(config)?; Python::with_gil(|py| { self.obj .call_method1(py, "initialize", (config,)) @@ -443,10 +433,4 @@ impl Bmi for BmiPython { } } -impl Drop for BmiPython { - fn drop(&mut self) { - if self.initialized { - let _ = self.finalize(); - } - } -} +impl_bmi_drop!(BmiPython); diff --git a/src/adapters/sloth.rs b/src/adapters/sloth.rs index 181e31e..f2243c1 100644 --- a/src/adapters/sloth.rs +++ b/src/adapters/sloth.rs @@ -1,5 +1,5 @@ use super::check_initialized; -use crate::error::{BmiError, BmiResult}; +use crate::error::{function_failed, BmiResult}; use crate::traits::{Bmi, VarType}; use std::collections::HashMap; @@ -66,10 +66,7 @@ impl BmiSloth { let parts: Vec<&str> = key[paren_start + 1..paren_end].split(',').collect(); if parts.len() != 4 { - return Err(BmiError::FunctionFailed { - model: "SLOTH".into(), - func: format!("Invalid param format: {}", key), - }); + return Err(function_failed("SLOTH", format!("Invalid param format: {}", key))); } let count: usize = parts[0].trim().parse().unwrap_or(1); @@ -84,10 +81,10 @@ impl BmiSloth { _ => VarType::Double, }; - let value: f64 = val.trim().parse().map_err(|_| BmiError::FunctionFailed { - model: "SLOTH".into(), - func: format!("Invalid value '{}' for {}", val, name), - })?; + let value: f64 = val + .trim() + .parse() + .map_err(|_| function_failed("SLOTH", format!("Invalid value '{}' for {}", val, name)))?; Ok(Some(( name, @@ -105,10 +102,7 @@ impl BmiSloth { fn get_var(&self, name: &str) -> BmiResult<&SlothVar> { self.variables .get(name) - .ok_or_else(|| BmiError::FunctionFailed { - model: self.name.clone(), - func: format!("Unknown variable: {}", name), - }) + .ok_or_else(|| function_failed(&self.name, format!("Unknown variable: {}", name))) } } diff --git a/src/lib.rs b/src/lib.rs index 9353452..95a8758 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -27,6 +27,7 @@ pub use config::{ parse_datetime, BmiAdapterType, DownsampleMode, ModuleConfig, OutputFormat, RealizationConfig, UpsampleMode, }; +pub use output::DivideDataStore; pub use error::{BmiError, BmiResult}; pub use forcings::{Forcings, NetCdfForcings}; pub use library::preload_dependencies; diff --git a/src/main.rs b/src/main.rs index 05416af..56603c2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use bmi_driver::{preload_dependencies, BmiError, ModelRunner, OutputFormat}; +use bmi_driver::{preload_dependencies, BmiError, DivideDataStore, ModelRunner, OutputFormat}; use clap::Parser; use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; use std::env; @@ -139,7 +139,7 @@ fn main() -> Result<(), BmiError> { if let (Some(start), Some(end)) = (args.worker_start, args.worker_end) { let output_path = data_dir.join("outputs").join("bmi-driver"); - run_worker(&realization, &locations[start..end], &output_path, start, &args.output_vars) + run_worker(&realization, &locations[start..end], &output_path, start) } else { run_parent(&data_dir, &realization, &locations, jobs, progress) } @@ -149,7 +149,7 @@ fn print_units(realization: &PathBuf, locations: &[String]) -> Result<(), BmiErr let mut runner = ModelRunner::from_config(realization)?; if let Some(loc) = locations.first() { runner.initialize(loc)?; - runner.print_all_unit_info(); + runner.print_unit_conversions(false); runner.finalize()?; } else { eprintln!("No locations found."); @@ -227,15 +227,15 @@ fn run_parent( // Re-initialize with updated config to show new conversions let mut runner2 = ModelRunner::from_config(realization)?; runner2.initialize(loc)?; - runner2.print_active_conversions(); + runner2.print_unit_conversions(true); runner2.finalize()?; } else { eprintln!("Skipping. Running with current config."); - runner.print_active_conversions(); + runner.print_unit_conversions(true); runner.finalize()?; } } else { - runner.print_active_conversions(); + runner.print_unit_conversions(true); runner.finalize()?; } } @@ -514,7 +514,6 @@ fn run_worker( locations: &[String], output_path: &PathBuf, #[cfg_attr(not(feature = "zarr"), allow(unused))] global_start: usize, - #[cfg_attr(not(feature = "zarr"), allow(unused))] cli_output_vars: &Option, ) -> Result<(), BmiError> { let mut runner = ModelRunner::from_config(realization)?; runner.suppress_warnings = true; @@ -529,158 +528,24 @@ fn run_worker( .map(|f| f.params.output_variables.clone()) .unwrap_or_default(); - let report_interval = ((locations.len() as f64) * 0.01).ceil().max(1.0) as usize; - - match output_format { - OutputFormat::Csv => run_worker_csv( - &mut runner, - locations, - output_path, - &output_vars, - start_epoch, - interval, - report_interval, - ), - OutputFormat::Netcdf => run_worker_netcdf( - &mut runner, - locations, - output_path, - &output_vars, - report_interval, - ), - #[cfg(feature = "zarr")] - OutputFormat::Zarr => { - let zarr_vars: Vec = cli_output_vars - .as_deref() - .unwrap_or("") - .split(',') - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .collect(); - run_worker_zarr( - &mut runner, - locations, - output_path, - &output_vars, - &zarr_vars, - global_start, - report_interval, - ) - } - } -} - -fn run_worker_csv( - runner: &mut ModelRunner, - locations: &[String], - output_path: &PathBuf, - output_vars: &[String], - start_epoch: i64, - interval: i64, - report_interval: usize, -) -> Result<(), BmiError> { - for (i, location) in locations.iter().enumerate() { - runner.initialize(location)?; - runner.run()?; - - let columns: Vec<(&str, &Vec)> = if output_vars.is_empty() { - runner - .outputs - .iter() - .map(|(name, vals)| (name.as_str(), vals)) - .collect() - } else { - output_vars - .iter() - .filter_map(|name| runner.outputs(name).ok().map(|vals| (name.as_str(), vals))) - .collect() - }; - - let csv_path = output_path.join(format!("{}.csv", location)); - let mut csv = String::from("Time Step,Time"); - for (name, _) in &columns { - csv.push(','); - csv.push_str(name); - } - csv.push('\n'); - - let n_steps = columns.first().map(|(_, v)| v.len()).unwrap_or(0); - for step in 0..n_steps { - let ts = start_epoch + (step as i64) * interval; - csv.push_str(&format!("{},{}", step, format_epoch(ts))); - for (_, vals) in &columns { - csv.push_str(&format!(",{:.9}", vals[step])); - } - csv.push('\n'); - } - - fs::write(&csv_path, csv).map_err(|e| BmiError::FunctionFailed { - model: "runner".into(), - func: format!("Failed to write {}: {}", csv_path.display(), e), - })?; - - runner.finalize()?; - - if (i + 1) % report_interval == 0 || i + 1 == locations.len() { - println!("{}", i + 1); - } - } - Ok(()) -} - -fn run_worker_netcdf( - runner: &mut ModelRunner, - locations: &[String], - output_path: &PathBuf, - output_vars: &[String], - report_interval: usize, -) -> Result<(), BmiError> { - use bmi_driver::output::netcdf::{LocationResult, NetCdfWriter}; + let mut store = create_output_store( + output_format, + output_path, + &runner, + start_epoch, + interval, + global_start, + locations.first().map(|s| s.as_str()), + )?; - // Determine the worker's tmp file name from the first location - let tmp_name = if let Some(first) = locations.first() { - format!("tmp_{}.nc", first) - } else { - return Ok(()); - }; - let nc_path = output_path.join(&tmp_name); - - let start_time = runner.config.time.start_time.clone(); - let interval = runner.config.time.output_interval; - // Calculate total_steps from config directly — runner.total_steps is 0 - // until initialize() is called, but we need it before the first initialize(). - let start_epoch = bmi_driver::parse_datetime(&start_time)?; - let end_epoch = bmi_driver::parse_datetime(&runner.config.time.end_time)?; - let total_steps = ((end_epoch - start_epoch) / interval) as usize; - - let writer = NetCdfWriter::new(nc_path, &start_time, interval, total_steps)?; + let report_interval = ((locations.len() as f64) * 0.01).ceil().max(1.0) as usize; for (i, location) in locations.iter().enumerate() { runner.initialize(location)?; runner.run()?; - let columns: Vec<(String, Vec)> = if output_vars.is_empty() { - runner - .outputs - .iter() - .map(|(name, vals)| (name.clone(), vals.clone())) - .collect() - } else { - output_vars - .iter() - .filter_map(|name| { - runner - .outputs(name) - .ok() - .map(|vals| (name.clone(), vals.clone())) - }) - .collect() - }; - - writer.write(LocationResult { - location_id: location.clone(), - columns, - })?; + let columns = collect_columns(&runner, &output_vars); + store.write_location(location, &columns)?; runner.finalize()?; @@ -689,98 +554,65 @@ fn run_worker_netcdf( } } - writer.finish()?; + store.finish()?; Ok(()) } -#[cfg(feature = "zarr")] -fn run_worker_zarr( - runner: &mut ModelRunner, - locations: &[String], - output_path: &PathBuf, - output_vars: &[String], - zarr_vars: &[String], - global_start: usize, - report_interval: usize, -) -> Result<(), BmiError> { - let zarr_path = output_path.join("results.zarr"); - - // Use zarr_vars (from parent's --output-vars) if output_vars (from config) is empty - let effective_vars = if output_vars.is_empty() { - zarr_vars +fn collect_columns(runner: &ModelRunner, output_vars: &[String]) -> Vec<(String, Vec)> { + if output_vars.is_empty() { + runner + .outputs + .iter() + .map(|(name, vals)| (name.clone(), vals.clone())) + .collect() } else { output_vars - }; - - for (i, location) in locations.iter().enumerate() { - runner.initialize(location)?; - runner.run()?; - - let columns: Vec<(String, Vec)> = if effective_vars.is_empty() { - runner - .outputs - .iter() - .map(|(name, vals)| (name.clone(), vals.clone())) - .collect() - } else { - effective_vars - .iter() - .filter_map(|name| { - runner - .outputs(name) - .ok() - .map(|vals| (name.clone(), vals.clone())) - }) - .collect() - }; - - bmi_driver::output::zarr::write_location(&zarr_path, global_start + i, &columns)?; - - runner.finalize()?; - - if (i + 1) % report_interval == 0 || i + 1 == locations.len() { - println!("{}", i + 1); - } + .iter() + .filter_map(|name| { + runner + .outputs(name) + .ok() + .map(|vals| (name.clone(), vals.clone())) + }) + .collect() } - Ok(()) } -fn format_epoch(epoch: i64) -> String { - let secs_per_day: i64 = 86400; - let mut remaining = epoch; - let sec = remaining % 60; - remaining /= 60; - let min = remaining % 60; - remaining /= 60; - let hour = remaining % 24; - let mut days = epoch / secs_per_day; - - let leap = |y: i32| (y % 4 == 0 && y % 100 != 0) || y % 400 == 0; - let days_in_month = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]; - - let mut year = 1970i32; - loop { - let yd = if leap(year) { 366 } else { 365 }; - if days < yd { - break; +fn create_output_store( + format: OutputFormat, + output_path: &PathBuf, + runner: &ModelRunner, + start_epoch: i64, + interval: i64, + #[cfg_attr(not(feature = "zarr"), allow(unused))] global_start: usize, + first_location: Option<&str>, +) -> Result, BmiError> { + match format { + OutputFormat::Csv => Ok(Box::new(bmi_driver::output::csv::CsvStore::new( + output_path.clone(), + start_epoch, + interval, + ))), + OutputFormat::Netcdf => { + let first_loc = first_location.unwrap_or("output"); + let nc_path = output_path.join(format!("tmp_{}.nc", first_loc)); + let start_time = &runner.config.time.start_time; + let end_epoch = bmi_driver::parse_datetime(&runner.config.time.end_time)?; + let total_steps = ((end_epoch - start_epoch) / interval) as usize; + Ok(Box::new(bmi_driver::output::netcdf::NetCdfWriter::new( + nc_path, + start_time, + interval, + total_steps, + )?)) } - days -= yd; - year += 1; - } - - let mut month = 0u32; - for m in 0..12 { - let md = days_in_month[m] as i64 + if m == 1 && leap(year) { 1 } else { 0 }; - if days < md { - month = m as u32 + 1; - break; + #[cfg(feature = "zarr")] + OutputFormat::Zarr => { + let zarr_path = output_path.join("results.zarr"); + Ok(Box::new(bmi_driver::output::zarr::ZarrStore::new( + zarr_path, + global_start, + ))) } - days -= md; } - let day = days + 1; - - format!( - "{:04}-{:02}-{:02} {:02}:{:02}:{:02}", - year, month, day, hour, min, sec - ) } diff --git a/src/output/csv.rs b/src/output/csv.rs new file mode 100644 index 0000000..ebdbbe8 --- /dev/null +++ b/src/output/csv.rs @@ -0,0 +1,88 @@ +use std::fs; +use std::path::PathBuf; + +use crate::error::{function_failed, BmiResult}; + +pub struct CsvStore { + output_path: PathBuf, + start_epoch: i64, + interval: i64, +} + +impl CsvStore { + pub fn new(output_path: PathBuf, start_epoch: i64, interval: i64) -> Self { + Self { + output_path, + start_epoch, + interval, + } + } +} + +impl super::DivideDataStore for CsvStore { + fn write_location(&mut self, loc_id: &str, columns: &[(String, Vec)]) -> BmiResult<()> { + let csv_path = self.output_path.join(format!("{}.csv", loc_id)); + + let mut csv = String::from("Time Step,Time"); + for (name, _) in columns { + csv.push(','); + csv.push_str(name); + } + csv.push('\n'); + + let n_steps = columns.first().map(|(_, v)| v.len()).unwrap_or(0); + for step in 0..n_steps { + let ts = self.start_epoch + (step as i64) * self.interval; + csv.push_str(&format!("{},{}", step, format_epoch(ts))); + for (_, vals) in columns { + csv.push_str(&format!(",{:.9}", vals[step])); + } + csv.push('\n'); + } + + fs::write(&csv_path, csv).map_err(|e| { + function_failed("csv_writer", format!("Failed to write {}: {}", csv_path.display(), e)) + })?; + Ok(()) + } +} + +fn format_epoch(epoch: i64) -> String { + let secs_per_day: i64 = 86400; + let mut remaining = epoch; + let sec = remaining % 60; + remaining /= 60; + let min = remaining % 60; + remaining /= 60; + let hour = remaining % 24; + let mut days = epoch / secs_per_day; + + let leap = |y: i32| (y % 4 == 0 && y % 100 != 0) || y % 400 == 0; + let days_in_month = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]; + + let mut year = 1970i32; + loop { + let yd = if leap(year) { 366 } else { 365 }; + if days < yd { + break; + } + days -= yd; + year += 1; + } + + let mut month = 0u32; + for m in 0..12 { + let md = days_in_month[m] as i64 + if m == 1 && leap(year) { 1 } else { 0 }; + if days < md { + month = m as u32 + 1; + break; + } + days -= md; + } + let day = days + 1; + + format!( + "{:04}-{:02}-{:02} {:02}:{:02}:{:02}", + year, month, day, hour, min, sec + ) +} diff --git a/src/output/mod.rs b/src/output/mod.rs index e095f88..da6e84e 100644 --- a/src/output/mod.rs +++ b/src/output/mod.rs @@ -1,3 +1,20 @@ +use crate::error::BmiResult; + +pub mod csv; pub mod netcdf; #[cfg(feature = "zarr")] pub mod zarr; + +/// Unified write interface for per-location timeseries data. +/// +/// Each output format (CSV, NetCDF, Zarr) implements this trait so that +/// the worker loop can write results without knowing the format. +pub trait DivideDataStore { + /// Write one location's output data. Each column is (variable_name, values_per_timestep). + fn write_location(&mut self, loc_id: &str, columns: &[(String, Vec)]) -> BmiResult<()>; + + /// Finalize the store (flush buffers, close files, etc). + fn finish(&mut self) -> BmiResult<()> { + Ok(()) + } +} diff --git a/src/output/netcdf.rs b/src/output/netcdf.rs index 8420f1e..df78190 100644 --- a/src/output/netcdf.rs +++ b/src/output/netcdf.rs @@ -5,13 +5,10 @@ use std::sync::{Arc, Mutex}; use std::thread::{self, JoinHandle}; use std::time::Duration; -use crate::error::{BmiError, BmiResult}; +use crate::error::{function_failed, BmiError, BmiResult}; fn err(msg: String) -> BmiError { - BmiError::FunctionFailed { - model: "netcdf_writer".into(), - func: msg, - } + function_failed("netcdf_writer", msg) } /// Suppress HDF5's default error handler which prints verbose diagnostics @@ -97,7 +94,11 @@ impl NetCdfWriter { }) } - pub fn finish(mut self) -> BmiResult<()> { + pub fn finish_owned(mut self) -> BmiResult<()> { + self.shutdown() + } + + fn shutdown(&mut self) -> BmiResult<()> { let _ = self.sender.send(WriterMessage::Shutdown); if let Some(handle) = self.handle.take() { handle @@ -109,6 +110,19 @@ impl NetCdfWriter { } } +impl super::DivideDataStore for NetCdfWriter { + fn write_location(&mut self, loc_id: &str, columns: &[(String, Vec)]) -> BmiResult<()> { + self.write(LocationResult { + location_id: loc_id.to_string(), + columns: columns.to_vec(), + }) + } + + fn finish(&mut self) -> BmiResult<()> { + self.shutdown() + } +} + fn writer_thread( receiver: Receiver, path: PathBuf, @@ -124,58 +138,39 @@ fn writer_thread( let mut batch: Vec = Vec::new(); let batch_size = 50; + let flush_batch = |file: &mut Option, + var_names: &mut Vec, + location_count: &mut usize, + batch: &mut Vec| + -> BmiResult<()> { + if batch.is_empty() { + return Ok(()); + } + if file.is_none() { + let (f, names) = init_netcdf(&path, start_time, interval, total_steps, &batch[0])?; + *file = Some(f); + *var_names = names; + } + *location_count = + write_batch(file.as_mut().unwrap(), batch, var_names, *location_count)?; + batch.clear(); + Ok(()) + }; + loop { match receiver.recv_timeout(Duration::from_millis(100)) { Ok(WriterMessage::Write(result)) => { batch.push(result); if batch.len() >= batch_size { - if file.is_none() { - let (f, names) = - init_netcdf(&path, start_time, interval, total_steps, &batch[0])?; - file = Some(f); - var_names = names; - } - location_count = - write_batch(file.as_mut().unwrap(), &batch, &var_names, location_count)?; - batch.clear(); + flush_batch(&mut file, &mut var_names, &mut location_count, &mut batch)?; } } - Ok(WriterMessage::Shutdown) => { - if !batch.is_empty() { - if file.is_none() { - let (f, names) = - init_netcdf(&path, start_time, interval, total_steps, &batch[0])?; - file = Some(f); - var_names = names; - } - write_batch(file.as_mut().unwrap(), &batch, &var_names, location_count)?; - } + Ok(WriterMessage::Shutdown) | Err(mpsc::RecvTimeoutError::Disconnected) => { + flush_batch(&mut file, &mut var_names, &mut location_count, &mut batch)?; break; } Err(mpsc::RecvTimeoutError::Timeout) => { - if !batch.is_empty() { - if file.is_none() { - let (f, names) = - init_netcdf(&path, start_time, interval, total_steps, &batch[0])?; - file = Some(f); - var_names = names; - } - location_count = - write_batch(file.as_mut().unwrap(), &batch, &var_names, location_count)?; - batch.clear(); - } - } - Err(mpsc::RecvTimeoutError::Disconnected) => { - if !batch.is_empty() { - if file.is_none() { - let (f, names) = - init_netcdf(&path, start_time, interval, total_steps, &batch[0])?; - file = Some(f); - var_names = names; - } - write_batch(file.as_mut().unwrap(), &batch, &var_names, location_count)?; - } - break; + flush_batch(&mut file, &mut var_names, &mut location_count, &mut batch)?; } } } diff --git a/src/output/zarr.rs b/src/output/zarr.rs index e54a25c..8fca812 100644 --- a/src/output/zarr.rs +++ b/src/output/zarr.rs @@ -6,13 +6,10 @@ use zarrs::filesystem::FilesystemStore; use zarrs::group::GroupBuilder; use zarrs::storage::ReadableWritableListableStorage; -use crate::error::{BmiError, BmiResult}; +use crate::error::{function_failed, BmiError, BmiResult}; fn err(msg: String) -> BmiError { - BmiError::FunctionFailed { - model: "zarr_writer".into(), - func: msg, - } + function_failed("zarr_writer", msg) } /// Create a Zarr V3 store with pre-allocated arrays for all output variables. @@ -129,3 +126,26 @@ pub fn write_location( Ok(()) } + +/// Stateful wrapper around `write_location` that tracks the current location index. +pub struct ZarrStore { + path: std::path::PathBuf, + next_idx: usize, +} + +impl ZarrStore { + pub fn new(path: std::path::PathBuf, start_idx: usize) -> Self { + Self { + path, + next_idx: start_idx, + } + } +} + +impl super::DivideDataStore for ZarrStore { + fn write_location(&mut self, _loc_id: &str, columns: &[(String, Vec)]) -> BmiResult<()> { + write_location(&self.path, self.next_idx, columns)?; + self.next_idx += 1; + Ok(()) + } +} diff --git a/src/runner.rs b/src/runner.rs index 662c51c..faf4c88 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -10,7 +10,9 @@ use crate::aliases; use crate::config::{ parse_datetime, BmiAdapterType, DownsampleMode, ModuleConfig, RealizationConfig, UpsampleMode, }; -use crate::error::{BmiError, BmiResult}; +use crate::error::{function_failed, BmiResult}; + +const TIMESTEP_EPSILON: f64 = 1e-9; use crate::forcings::{Forcings, NetCdfForcings}; use crate::resample; use crate::traits::{Bmi, BmiExt}; @@ -167,16 +169,43 @@ impl ModelRunner { .filter(|d| !self.vars.contains_key(*d)) .cloned() .collect(); - return Err(BmiError::FunctionFailed { - model: "runner".into(), - func: format!("Missing dependencies: {:?}", missing), - }); + return Err(function_failed("runner", format!("Missing dependencies: {:?}", missing))); } } Ok(()) } fn load_model(&mut self, module: &ModuleConfig, loc_id: &str, idx: usize) -> BmiResult<()> { + let model = self.create_adapter(module, loc_id)?; + let timestep_info = self.query_model_timestep(&*model, module); + + for output in model.get_output_var_names()? { + self.source_timesteps + .insert(output.clone(), timestep_info.clone()); + self.vars.insert(output.clone(), VarSource::Model(idx)); + } + + let input_conversions = self.build_unit_conversions(&*model, module); + + self.models.push(ModelInstance { + name: module.params.model_type_name.clone(), + model, + input_map: module.params.variables_names_map.clone(), + main_output: module.params.main_output_variable.clone(), + input_conversions, + timestep_info, + downsample_mode: module.params.downsample_mode, + upsample_mode: module.params.upsample_mode, + }); + Ok(()) + } + + /// Create and initialize the appropriate BMI adapter for a module config. + fn create_adapter( + &self, + module: &ModuleConfig, + loc_id: &str, + ) -> BmiResult> { let is_sloth = module.name == "bmi_c++" && module.params.model_type_name.to_uppercase() == "SLOTH"; @@ -186,10 +215,7 @@ impl ModelRunner { Box::new(sloth) } else { let adapter = BmiAdapterType::from_name(&module.name).ok_or_else(|| { - BmiError::FunctionFailed { - model: module.params.model_type_name.clone(), - func: format!("Unknown adapter: {}", module.name), - } + function_failed(&module.params.model_type_name, format!("Unknown adapter: {}", module.name)) })?; match adapter { @@ -238,12 +264,11 @@ impl ModelRunner { &module.params.registration_function, )?) } else { - return Err(BmiError::FunctionFailed { - model: module.params.model_type_name.clone(), - func: "bmi_python requires either 'python_type' (e.g. \"lstm.bmi_lstm.bmi_LSTM\") \ - or both 'library_file' and 'registration_function'" - .into(), - }); + return Err(function_failed( + &module.params.model_type_name, + "bmi_python requires either 'python_type' (e.g. \"lstm.bmi_lstm.bmi_LSTM\") \ + or both 'library_file' and 'registration_function'", + )); } } } @@ -258,11 +283,14 @@ impl ModelRunner { } } - // Query model timestep + Ok(model) + } + + /// Query a model's timestep and return TimestepInfo. + fn query_model_timestep(&self, model: &dyn Bmi, module: &ModuleConfig) -> TimestepInfo { let model_dt_seconds = match (model.get_time_step(), model.get_time_units()) { (Ok(dt), Ok(units)) => { - let factor = crate::traits::parse_time_units(&units); - let dt_sec = dt * factor; + let dt_sec = dt * crate::traits::parse_time_units(&units); if dt_sec > 0.0 { dt_sec } else { @@ -279,20 +307,19 @@ impl ModelRunner { self.config.time.output_interval as f64 } }; - let model_steps = (self.simulation_span_seconds / model_dt_seconds) as usize; - let timestep_info = TimestepInfo { + TimestepInfo { dt_seconds: model_dt_seconds, - num_steps: model_steps, - }; - - for output in model.get_output_var_names()? { - self.source_timesteps - .insert(output.clone(), timestep_info.clone()); - self.vars.insert(output.clone(), VarSource::Model(idx)); + num_steps: (self.simulation_span_seconds / model_dt_seconds) as usize, } + } - // Build unit conversions for each input mapping - let mut input_conversions = HashMap::new(); + /// Build unit conversions for each input variable mapping. + fn build_unit_conversions( + &self, + model: &dyn Bmi, + module: &ModuleConfig, + ) -> HashMap { + let mut conversions = HashMap::new(); for (model_input, source_var) in &module.params.variables_names_map { let dest_units = model.get_var_units(model_input).unwrap_or_default(); let source_units = self.get_source_units(source_var); @@ -308,29 +335,15 @@ impl ModelRunner { ); } } - input_conversions.insert(model_input.clone(), conv); + conversions.insert(model_input.clone(), conv); } } - - self.models.push(ModelInstance { - name: module.params.model_type_name.clone(), - model, - input_map: module.params.variables_names_map.clone(), - main_output: module.params.main_output_variable.clone(), - input_conversions, - timestep_info, - downsample_mode: module.params.downsample_mode, - upsample_mode: module.params.upsample_mode, - }); - Ok(()) + conversions } pub fn run(&mut self) -> BmiResult<()> { if self.has_run { - return Err(BmiError::FunctionFailed { - model: "runner".into(), - func: "already run".into(), - }); + return Err(function_failed("runner", "already run")); } for i in 0..self.models.len() { @@ -364,7 +377,7 @@ impl ModelRunner { } }; - if (source_ts.dt_seconds - output_dt).abs() < 1e-9 { + if (source_ts.dt_seconds - output_dt).abs() < TIMESTEP_EPSILON { // Same timestep, no resampling needed resampled.insert(name.clone(), vals.clone()); continue; @@ -447,19 +460,16 @@ impl ModelRunner { downsample_mode: DownsampleMode, upsample_mode: UpsampleMode, ) -> BmiResult { - let source_ts = - self.source_timesteps - .get(name) - .ok_or_else(|| BmiError::FunctionFailed { - model: "runner".into(), - func: format!("No timestep info for variable: {}", name), - })?; + let source_ts = self + .source_timesteps + .get(name) + .ok_or_else(|| function_failed("runner", format!("No timestep info for variable: {}", name)))?; let source_dt = source_ts.dt_seconds; match self.vars.get(name) { Some(VarSource::Forcing) => { // Fast path: same timestep - if (source_dt - dest_dt).abs() < 1e-9 { + if (source_dt - dest_dt).abs() < TIMESTEP_EPSILON { let step = (dest_time / source_dt).round() as usize; return self.forcings.get_f64(name, &self.location_id, step); } @@ -473,21 +483,15 @@ impl ModelRunner { ) } Some(VarSource::Model(_)) => { - let source_data = - self.outputs - .get(name) - .ok_or_else(|| BmiError::FunctionFailed { - model: "runner".into(), - func: format!("'{}' not yet computed", name), - })?; + let source_data = self + .outputs + .get(name) + .ok_or_else(|| function_failed("runner", format!("'{}' not yet computed", name)))?; // Fast path: same timestep - if (source_dt - dest_dt).abs() < 1e-9 { + if (source_dt - dest_dt).abs() < TIMESTEP_EPSILON { let step = (dest_time / source_dt).round() as usize; return source_data.get(step).copied().ok_or_else(|| { - BmiError::FunctionFailed { - model: "runner".into(), - func: format!("'{}' not available at step {}", name, step), - } + function_failed("runner", format!("'{}' not available at step {}", name, step)) }); } resample::resample_value( @@ -499,10 +503,7 @@ impl ModelRunner { upsample_mode, ) } - None => Err(BmiError::FunctionFailed { - model: "runner".into(), - func: format!("Unknown variable: {}", name), - }), + None => Err(function_failed("runner", format!("Unknown variable: {}", name))), } } @@ -526,7 +527,7 @@ impl ModelRunner { DownsampleMode::Interpolate => { let lower_val = self.forcings.get_f64(name, &self.location_id, lower_idx)?; let frac = fractional_idx - lower_idx as f64; - if frac.abs() < 1e-12 { + if frac.abs() < TIMESTEP_EPSILON { return Ok(lower_val); } match self @@ -550,10 +551,7 @@ impl ModelRunner { } } if vals.is_empty() { - return Err(BmiError::FunctionFailed { - model: "runner".into(), - func: format!("No forcing data for '{}' in window", name), - }); + return Err(function_failed("runner", format!("No forcing data for '{}' in window", name))); } resample::aggregate(&vals, upsample_mode) } @@ -574,52 +572,40 @@ impl ModelRunner { } } - /// Print only the active (non-identity) unit conversions to stderr. - pub fn print_active_conversions(&self) { - let mut any = false; - for m in &self.models { - for (model_input, source_var) in &m.input_map { - if let Some(conv) = m.input_conversions.get(model_input) { - if conv.is_identity() { - continue; - } - let source_label = self.source_label(source_var); - eprintln!( - " {}: {} ← {} ({}): {}", - m.name, model_input, source_var, source_label, conv - ); - any = true; - } - } + /// Print unit conversions to stderr. + /// If `active_only` is true, only prints non-identity conversions. + /// If false, prints all variable mappings including those without unit info. + pub fn print_unit_conversions(&self, active_only: bool) { + if !active_only { + eprintln!("Unit conversions for this run:"); } - if any { - eprintln!(); - } - } - - /// Print a full summary of all variable mappings and unit conversions. - pub fn print_all_unit_info(&self) { - eprintln!("Unit conversions for this run:"); let mut any = false; for m in &self.models { for (model_input, source_var) in &m.input_map { let source_label = self.source_label(source_var); - if let Some(conv) = m.input_conversions.get(model_input) { + if active_only && conv.is_identity() { + continue; + } eprintln!( " {}: {} ← {} ({}): {}", m.name, model_input, source_var, source_label, conv ); - } else { + } else if !active_only { eprintln!( " {}: {} ← {} ({}): no unit info available", m.name, model_input, source_var, source_label ); + } else { + continue; } any = true; } } - if !any { + if active_only && any { + eprintln!(); + } + if !active_only && !any { eprintln!(" (no variable mappings)"); } } @@ -687,27 +673,18 @@ impl ModelRunner { pub fn main_outputs(&self) -> BmiResult<&Vec> { if !self.has_run { - return Err(BmiError::FunctionFailed { - model: "runner".into(), - func: "call run() first".into(), - }); + return Err(function_failed("runner", "call run() first")); } Ok(&self.final_outputs) } pub fn outputs(&self, name: &str) -> BmiResult<&Vec> { if !self.has_run { - return Err(BmiError::FunctionFailed { - model: "runner".into(), - func: "call run() first".into(), - }); + return Err(function_failed("runner", "call run() first")); } self.outputs .get(name) - .ok_or_else(|| BmiError::FunctionFailed { - model: "runner".into(), - func: format!("Output '{}' not found", name), - }) + .ok_or_else(|| function_failed("runner", format!("Output '{}' not found", name))) } pub fn finalize(&mut self) -> BmiResult<()> { diff --git a/src/types/error.rs b/src/types/error.rs index 84e5aec..2867b0e 100644 --- a/src/types/error.rs +++ b/src/types/error.rs @@ -37,3 +37,19 @@ pub enum BmiError { } pub type BmiResult = Result; + +/// Shorthand for creating a `BmiError::FunctionFailed`. +pub fn function_failed(model: impl Into, msg: impl Into) -> BmiError { + BmiError::FunctionFailed { + model: model.into(), + func: msg.into(), + } +} + +/// Shorthand for creating a `BmiError::NotImplemented`. +pub fn not_implemented(model: impl Into, func: impl Into) -> BmiError { + BmiError::NotImplemented { + model: model.into(), + func: func.into(), + } +} diff --git a/src/variables/resample.rs b/src/variables/resample.rs index bf3a5d3..6ea8de5 100644 --- a/src/variables/resample.rs +++ b/src/variables/resample.rs @@ -1,12 +1,9 @@ use crate::config::{DownsampleMode, UpsampleMode}; -use crate::error::{BmiError, BmiResult}; +use crate::error::{function_failed, BmiResult}; use std::collections::HashMap; -fn err(msg: String) -> BmiError { - BmiError::FunctionFailed { - model: "resample".into(), - func: msg, - } +fn err(msg: String) -> crate::error::BmiError { + function_failed("resample", msg) } /// Resample a single value from a source time series to a destination time. diff --git a/tests/output_formats.rs b/tests/output_formats.rs new file mode 100644 index 0000000..075d57f --- /dev/null +++ b/tests/output_formats.rs @@ -0,0 +1,347 @@ +use bmi_driver::output::DivideDataStore; +use std::path::Path; + +const LOCATIONS: &[&str] = &["cat-100", "cat-200"]; +const VAR_NAMES: &[&str] = &["Q_OUT", "EVAPOTRANS", "RAIN_RATE"]; +const N_STEPS: usize = 25; +const START_TIME: &str = "2010-01-01 00:00:00"; +const INTERVAL: i64 = 3600; + +/// Generate deterministic test data: value = loc_idx * 1000 + var_idx * 100 + step + 0.123456789 +fn generate_data() -> Vec<(String, Vec<(String, Vec)>)> { + LOCATIONS + .iter() + .enumerate() + .map(|(loc_idx, loc_id)| { + let columns: Vec<(String, Vec)> = VAR_NAMES + .iter() + .enumerate() + .map(|(var_idx, var_name)| { + let values: Vec = (0..N_STEPS) + .map(|step| loc_idx as f64 * 1000.0 + var_idx as f64 * 100.0 + step as f64 + 0.123456789) + .collect(); + (var_name.to_string(), values) + }) + .collect(); + (loc_id.to_string(), columns) + }) + .collect() +} + +fn write_all(store: &mut dyn DivideDataStore, data: &[(String, Vec<(String, Vec)>)]) { + for (loc_id, columns) in data { + store.write_location(loc_id, columns).unwrap(); + } + store.finish().unwrap(); +} + +// --- CSV --- + +fn read_csv_data(dir: &Path) -> Vec<(String, Vec<(String, Vec)>)> { + LOCATIONS + .iter() + .map(|loc_id| { + let csv_path = dir.join(format!("{}.csv", loc_id)); + let content = std::fs::read_to_string(&csv_path) + .unwrap_or_else(|e| panic!("Failed to read {}: {}", csv_path.display(), e)); + let mut lines = content.lines(); + let header = lines.next().unwrap(); + // Header: "Time Step,Time,VAR1,VAR2,..." + let var_names: Vec = header.split(',').skip(2).map(|s| s.to_string()).collect(); + let mut columns: Vec> = vec![Vec::new(); var_names.len()]; + + for line in lines { + let fields: Vec<&str> = line.split(',').collect(); + for (i, val_str) in fields.iter().skip(2).enumerate() { + columns[i].push(val_str.parse::().unwrap()); + } + } + + let result: Vec<(String, Vec)> = var_names + .into_iter() + .zip(columns) + .collect(); + (loc_id.to_string(), result) + }) + .collect() +} + +#[test] +fn test_csv_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let data = generate_data(); + + let start_epoch = 1262304000i64; // 2010-01-01 00:00:00 UTC + let mut store = bmi_driver::output::csv::CsvStore::new(dir.path().to_path_buf(), start_epoch, INTERVAL); + write_all(&mut store, &data); + + let read_back = read_csv_data(dir.path()); + assert_eq!(read_back.len(), data.len()); + + for (orig, read) in data.iter().zip(read_back.iter()) { + assert_eq!(orig.0, read.0, "location ID mismatch"); + assert_eq!(orig.1.len(), read.1.len(), "variable count mismatch"); + for (orig_col, read_col) in orig.1.iter().zip(read.1.iter()) { + assert_eq!(orig_col.0, read_col.0, "variable name mismatch"); + assert_eq!(orig_col.1.len(), read_col.1.len(), "value count mismatch"); + for (i, (o, r)) in orig_col.1.iter().zip(read_col.1.iter()).enumerate() { + // CSV writes 9 decimal places + assert!( + (o - r).abs() < 1e-8, + "CSV value mismatch at step {} for {}: {} vs {}", + i, orig_col.0, o, r + ); + } + } + } +} + +// --- NetCDF --- + +fn read_netcdf_data(path: &Path) -> Vec<(String, Vec<(String, Vec)>)> { + let file = netcdf::open(path).unwrap(); + + let n_locs = file.dimension("id").unwrap().len(); + let ids_var = file.variable("id").unwrap(); + + // Discover variable names (everything except time and id) + let var_names: Vec = file + .variables() + .filter(|v| { + let name = v.name(); + name != "time" && name != "id" + }) + .map(|v| v.name()) + .collect(); + + (0..n_locs) + .map(|loc_idx| { + let loc_id = ids_var.get_string(loc_idx).unwrap(); + let columns: Vec<(String, Vec)> = var_names + .iter() + .map(|name| { + let var = file.variable(name).unwrap(); + let f32_vals: Vec = var.get_values((loc_idx, ..)).unwrap(); + let f64_vals: Vec = f32_vals.iter().map(|&v| v as f64).collect(); + (name.clone(), f64_vals) + }) + .collect(); + (loc_id, columns) + }) + .collect() +} + +#[test] +fn test_netcdf_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let data = generate_data(); + + let nc_path = dir.path().join("test.nc"); + let mut store = bmi_driver::output::netcdf::NetCdfWriter::new( + nc_path.clone(), + START_TIME, + INTERVAL, + N_STEPS - 1, // total_steps = N_STEPS - 1 because n_times = total_steps + 1 + ) + .unwrap(); + write_all(&mut store, &data); + + let read_back = read_netcdf_data(&nc_path); + assert_eq!(read_back.len(), data.len()); + + for (orig, read) in data.iter().zip(read_back.iter()) { + assert_eq!(orig.0, read.0, "location ID mismatch"); + assert_eq!(orig.1.len(), read.1.len(), "variable count mismatch"); + for (orig_col, read_col) in orig.1.iter().zip(read.1.iter()) { + assert_eq!(orig_col.0, read_col.0, "variable name mismatch"); + assert_eq!(orig_col.1.len(), read_col.1.len(), "value count mismatch for {}", orig_col.0); + for (i, (o, r)) in orig_col.1.iter().zip(read_col.1.iter()).enumerate() { + // NetCDF stores f32, so tolerance must account for f64→f32→f64 roundtrip + assert!( + (o - r).abs() < 1e-2, + "NetCDF value mismatch at step {} for {}: {} vs {}", + i, orig_col.0, o, r + ); + } + } + } +} + +// --- Zarr --- + +#[cfg(feature = "zarr")] +fn read_zarr_data(path: &Path) -> Vec<(String, Vec<(String, Vec)>)> { + use std::sync::Arc; + use zarrs::filesystem::FilesystemStore; + use zarrs::storage::ReadableWritableListableStorage; + + let store: ReadableWritableListableStorage = + Arc::new(FilesystemStore::new(path).unwrap()); + + let id_array = zarrs::array::Array::open(store.clone(), "/id").unwrap(); + let n_locs = id_array.shape()[0] as usize; + + // Read location IDs + let loc_ids: Vec = (0..n_locs) + .map(|i| { + let chunk: Vec = id_array + .retrieve_chunk::>(&[i as u64]) + .unwrap(); + chunk[0].clone() + }) + .collect(); + + let var_names: Vec = VAR_NAMES.iter().map(|s| s.to_string()).collect(); + + loc_ids + .iter() + .enumerate() + .map(|(loc_idx, loc_id): (usize, &String)| { + let columns: Vec<(String, Vec)> = var_names + .iter() + .map(|name| { + let array = zarrs::array::Array::open(store.clone(), &format!("/{name}")).unwrap(); + let f32_vals: Vec = array + .retrieve_chunk::>(&[loc_idx as u64, 0]) + .unwrap(); + let f64_vals: Vec = f32_vals.iter().map(|&v| v as f64).collect(); + (name.clone(), f64_vals) + }) + .collect(); + (loc_id.clone(), columns) + }) + .collect() +} + +#[cfg(feature = "zarr")] +#[test] +fn test_zarr_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let data = generate_data(); + + let zarr_path = dir.path().join("test.zarr"); + let loc_ids: Vec = LOCATIONS.iter().map(|s| s.to_string()).collect(); + let var_names_owned: Vec = VAR_NAMES.iter().map(|s| s.to_string()).collect(); + + bmi_driver::output::zarr::create_zarr_store( + &zarr_path, + START_TIME, + INTERVAL, + N_STEPS - 1, + &loc_ids, + &var_names_owned, + ) + .unwrap(); + + let mut store = bmi_driver::output::zarr::ZarrStore::new(zarr_path.clone(), 0); + write_all(&mut store, &data); + + let read_back = read_zarr_data(&zarr_path); + assert_eq!(read_back.len(), data.len()); + + for (orig, read) in data.iter().zip(read_back.iter()) { + assert_eq!(orig.0, read.0, "location ID mismatch"); + assert_eq!(orig.1.len(), read.1.len(), "variable count mismatch"); + for (orig_col, read_col) in orig.1.iter().zip(read.1.iter()) { + assert_eq!(orig_col.0, read_col.0, "variable name mismatch"); + assert_eq!(orig_col.1.len(), read_col.1.len(), "value count mismatch for {}", orig_col.0); + for (i, (o, r)) in orig_col.1.iter().zip(read_col.1.iter()).enumerate() { + // Zarr stores f32 + assert!( + (o - r).abs() < 1e-2, + "Zarr value mismatch at step {} for {}: {} vs {}", + i, orig_col.0, o, r + ); + } + } + } +} + +// --- Cross-format equivalence --- + +#[cfg(feature = "zarr")] +#[test] +fn test_all_formats_equivalent() { + let dir = tempfile::tempdir().unwrap(); + let data = generate_data(); + let start_epoch = 1262304000i64; + + // Write CSV + let csv_dir = dir.path().join("csv"); + std::fs::create_dir_all(&csv_dir).unwrap(); + let mut csv_store = bmi_driver::output::csv::CsvStore::new(csv_dir.clone(), start_epoch, INTERVAL); + write_all(&mut csv_store, &data); + + // Write NetCDF + let nc_path = dir.path().join("test.nc"); + let mut nc_store = bmi_driver::output::netcdf::NetCdfWriter::new( + nc_path.clone(), + START_TIME, + INTERVAL, + N_STEPS - 1, + ) + .unwrap(); + write_all(&mut nc_store, &data); + + // Write Zarr + let zarr_path = dir.path().join("test.zarr"); + let loc_ids: Vec = LOCATIONS.iter().map(|s| s.to_string()).collect(); + let var_names_owned: Vec = VAR_NAMES.iter().map(|s| s.to_string()).collect(); + bmi_driver::output::zarr::create_zarr_store( + &zarr_path, + START_TIME, + INTERVAL, + N_STEPS - 1, + &loc_ids, + &var_names_owned, + ) + .unwrap(); + let mut zarr_store = bmi_driver::output::zarr::ZarrStore::new(zarr_path.clone(), 0); + write_all(&mut zarr_store, &data); + + // Read back + let csv_data = read_csv_data(&csv_dir); + let nc_data = read_netcdf_data(&nc_path); + let zarr_data = read_zarr_data(&zarr_path); + + // Compare location IDs + assert_eq!(csv_data.len(), nc_data.len()); + assert_eq!(csv_data.len(), zarr_data.len()); + + for loc_idx in 0..csv_data.len() { + assert_eq!(csv_data[loc_idx].0, nc_data[loc_idx].0, "CSV vs NetCDF loc ID"); + assert_eq!(csv_data[loc_idx].0, zarr_data[loc_idx].0, "CSV vs Zarr loc ID"); + + let csv_cols = &csv_data[loc_idx].1; + let nc_cols = &nc_data[loc_idx].1; + let zarr_cols = &zarr_data[loc_idx].1; + + assert_eq!(csv_cols.len(), nc_cols.len(), "CSV vs NetCDF var count"); + assert_eq!(csv_cols.len(), zarr_cols.len(), "CSV vs Zarr var count"); + + for var_idx in 0..csv_cols.len() { + let csv_vals = &csv_cols[var_idx].1; + let nc_vals = &nc_cols[var_idx].1; + let zarr_vals = &zarr_cols[var_idx].1; + + assert_eq!(csv_vals.len(), nc_vals.len(), "value count mismatch"); + assert_eq!(csv_vals.len(), zarr_vals.len(), "value count mismatch"); + + for step in 0..csv_vals.len() { + // NetCDF and Zarr store f32, CSV stores 9 decimal places + // Compare NetCDF ↔ Zarr (both f32, should be identical) + assert!( + (nc_vals[step] - zarr_vals[step]).abs() < 1e-6, + "NetCDF vs Zarr mismatch at loc={}, var={}, step={}: {} vs {}", + loc_idx, csv_cols[var_idx].0, step, nc_vals[step], zarr_vals[step] + ); + // Compare CSV ↔ NetCDF (CSV has more precision than f32 roundtrip) + assert!( + (csv_vals[step] - nc_vals[step]).abs() < 1e-2, + "CSV vs NetCDF mismatch at loc={}, var={}, step={}: {} vs {}", + loc_idx, csv_cols[var_idx].0, step, csv_vals[step], nc_vals[step] + ); + } + } + } +}