diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..0203922 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,38 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.12", "3.13"] + + defaults: + run: + # Login shell so the micromamba environment stays activated across steps. + shell: bash -el {0} + + steps: + - uses: actions/checkout@v4 + + # ecflow is not on PyPI; it is installed from conda-forge via the + # environment file. river-route and the rest come along with it. + - name: Create environment + uses: mamba-org/setup-micromamba@v2 + with: + environment-file: environment.yml + create-args: >- + python=${{ matrix.python-version }} + cache-environment: true + + - name: Install package with dev dependencies + run: pip install -e ".[dev]" + + - name: Run tests + run: pytest -q --cov=geoglows_ecflow.resources --cov-report=term-missing diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..943ace6 --- /dev/null +++ b/PLAN.md @@ -0,0 +1,116 @@ +# Workflow Simplification Plan + +A living plan for simplifying the `geoglows_ecflow` workflow. Branched off the +`rapid-to-river-route` work (PR #27), so all references below assume the +river-route codebase, not the RAPID `main`. + +## Goals + +1. **Reduce complexity** — eliminate duplication, simplify functions, separate + inputs from logic. +2. **Make implicit explicit** — name constants and configuration, split mixed + functions, document non-obvious behavior. +3. **Improve maintainability** — make different configurations easy to run and + add a unit-test safety net. + +## Scope + +- **In scope:** `geoglows_ecflow/resources/*.py`, + `geoglows_ecflow/workflow/builders/builder.py`, + `geoglows_ecflow/workflow/parts/*`, and the `.ecf` task scripts. +- **Mostly frozen:** `geoglows_ecflow/workflow/comfies/*` — vendored ECMWF + framework code (Apache 2.0, ~4,700 lines). Left untouched except for the + minimal Python-3.12+ compatibility fixes noted below (the suite could not be + imported at all without them). + +## Decisions + +- **Tests:** unit tests on the pure functions, plus in-memory + suite-definition smoke tests (these were originally deferred but pulled + forward to guard the builder refactor). +- **HRES member:** high-resolution is always ensemble member **52**, kept as a + single named constant `HRES_ENSEMBLE_MEMBER` rather than a configurable value. + +--- + +## Phase 1 — Pure cleanup (no behavior change) — DONE + +- [x] `builder.py`: fix `Task("dimmy")` typo. +- [x] `helper_functions.py`: fix `create_logger` so file logging attaches. +- [x] `builder.py`: remove duplicate imports and the duplicated `nodes` import. +- [x] `builder.py`: remove read-but-unused config vars. +- [x] `builder.py`: fix stale docstring ("GLOFAS suite" → GEOGloWS). +- [x] `generate_esri_table.py`: fix `int or str` type hint. +- [x] Standardize the `argparse(nargs=1)` + `args.x[0]` antipattern. + +## Phase 2 — Test harness + CI — DONE + +- [x] Add `pytest` (+ `pytest-cov`) as a dev dependency. +- [x] Create `tests/` with fixtures. +- [x] Unit tests for the pure functions (ensemble parsing, VPU listing, date + parsing, state-init lookback, forecast preprocess, return-period / + exceedance, init-flow time index, zarr round-trip). +- [x] GitHub Actions workflow (`.github/workflows/tests.yml`). Now builds a + conda env from `environment.yml` (ecflow has no PyPI wheel) on a + Python **3.12–3.13** matrix and runs the whole suite. + +## Phase 3 — Centralize duplication — DONE + +- [x] Shared zarr-writing helper (`resources/zarr_io.py`). +- [x] `helper_functions.load_forecast_run` loader. +- [x] `RETURN_PERIODS` defined once and used for the ladders. +- [x] Standardized `logging` setup (`configure_logging`). + +## Phase 4 — Make implicit explicit — MOSTLY DONE + +- [x] `HRES_ENSEMBLE_MEMBER = 52` — used in `netcdf_to_zarr.py` and the + `range(1, HRES_ENSEMBLE_MEMBER + 1)` ensemble loop in `builder.py`. +- [ ] `nco_calc.ecf` `grep -v ..._52.nc` (×3) — **decision pending**: wire an + ecflow `%HRES_MEMBER%` variable vs. leave `52` + a comment. Highest risk: + wrong wiring silently changes which member is excluded from the mean. +- [x] `EMOS_BASE != "12"` gate (×3) → `is_00z_cycle()` helper. +- [x] Magic numbers named: thickness ladder (`THICKNESS_THRESHOLDS`), + stream-order (`MIN_STREAM_ORDER`), 10-day window (`FORECAST_WINDOW_DAYS`), + `MEM` values (`ENS_TASK_MEM_MB` / `ARCHIVE_QINIT_MEM_MB`). +- [x] Timer offsets (`hours=7`/`hours=9`/`"14:15"`) — **resolved by deletion**: + they lived only in the broken `rd`/research-mode branch, which has been + removed (see follow-ups). Nothing to extract; revisit if `rd` returns. +- [x] Consolidate `self.config.get(...)` reads in `builder.py` into one + documented block. + +## Follow-ups (later tasks) + +- ecFlow-**server** tests (building the def against a live server / `--dry`); + the in-memory structural tests are done, this is the heavier version. +- README refresh (carried over from PR #27 review). +- **comfies is incompatible with ecflow 5.17+.** Its node wrappers set + `Variable.parent` (`ooflow.py:1925`), which ecflow 5.17 made a read-only + built-in, so building any suite raises `AttributeError`. Worked around by + pinning `ecflow<5.17` in `environment.yml`; the real fix is to rename + comfies' parent-tracking attribute so it no longer collides. +- **Research (`rd`) mode removed.** `mode='rd'` was the only path with + `follow_osuite=False` and crashed unconditionally at `barrier_hh.ymd` + (`barrier_hh` is a `NominalTime`, which has no `ymd`) — broken since the + original 2024-07-25 authoring, so never usable. Removed the `rd` choice, the + `follow_osuite`/`in_production`/`in_test` flags, the non-`follow_osuite` + branch (the `+7h`/`+9h`/`14:15` run timers), and the dead crash line. `prod` + and `test` are the remaining modes. If research mode is wanted again, it + should be reintroduced correctly (with the intended barrier-repeat wiring). + +--- + +## Current status + +**Branch `workflow-simplification`** (fork `JakeGimenes`), open as a PR against +`rapid-to-river-route`. Phases 1–3 complete; Phase 4 complete except the +`nco_calc.ecf` decision. The timer constants are resolved by deletion (the +`rd`/research mode that owned them has been removed). **33 pytest tests pass** — +the resources tests run anywhere; the suite-definition tests require `ecflow` +(conda-forge). + +The vendored `comfies` framework got the minimum Python-3.12+ compatibility +fixes needed to import it at all (`imp` → `importlib`, `pkg_resources` → +`packaging`); everything else in `comfies/*` is unchanged. + +**Remaining actionable work:** the `nco_calc.ecf` `HRES_MEMBER` decision and the +README refresh. diff --git a/README.md b/README.md index a974f57..f9c8722 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# ECFLOW RAPID workflow for GEOGloWS +# ECFLOW workflow for GEOGloWS streamflow forecasting ![GEOGloWS VPUCode Coverage](images/geoglows_vpucode_coverage.png) *Coverage of GEOGloWS VPUCode basins. Source: [Riley Hales](mailto:rchales@byu.edu).* @@ -18,8 +18,7 @@ pip install -e . ## Non-Python Dependencies -- rapid>=20210423 -- ecflow>=5.11.3 +- ecflow>=5.11.3,<5.17 - nco>=5.1.8 - ksh>=2020.0.0 @@ -33,7 +32,7 @@ pip install -e . mars_bond_id='251' staticdata = '/path/to/assets' workroot = f'/path/to/workroot' - mode = 'test' # suite mode ('rd':research, 'test':test, 'prod':production) + mode = 'test' # suite mode ('test':test, 'prod':production) expver = 'geoglows' exparch = '/path/to/archive' iniexparch = '/path/to/init_archive' @@ -81,18 +80,10 @@ pip install -e . ) # -------------------------------------------- - # Configuration of EFAS software packages + # Configuration of GEOGloWS software packages # which are installed together with the suite. # -------------------------------------------- packages = dict( - model = dict( - srcdir = 'git+https://github.com/c-h-david/rapid.git@20210423', - ), - - petsc = dict( - srcdir = srcroot + 'petsc_reqs', - ), - scripts = dict( srcdir = srcroot + 'scripts', ), @@ -123,19 +114,26 @@ ecflow_start.sh -d /path/to/ecflow_home ## Local run example -```Python -import subprocess -from geoglows_ecflow import geoglows_forecast_job, client +Generate the suite definition (via CLI or Python): -# Start server -subprocess.run(['bash', '/path/to/local_server_start.sh']) +```bash +gdeploy --config /path/to/config.cfg +``` -# Create definition -geoglows_forecast_job.create("/path/to/config.cfg") +```python +from geoglows_ecflow.workflow.create import main +main("/path/to/config.cfg") +``` + +Start a local ecflow server, then load and begin the suite: + +```bash +bash /path/to/local_ecflow_start.sh +``` -# Add definition to server -client.add_definition("/path/to/definition.def", ":") +```python +from geoglows_ecflow import client -# Begin definition -client.begin("definition_name") +client.add_definition("/path/to/deploy_dir/suite.def", "localhost:2500") +client.begin("suite_name", "localhost:2500") ``` diff --git a/environment.yml b/environment.yml index 4f996d3..7710e78 100644 --- a/environment.yml +++ b/environment.yml @@ -3,7 +3,11 @@ channels: - conda-forge - defaults dependencies: - - ecflow + - python>=3.12,<3.14 + - pip + # comfies' node wrappers set Variable.parent, which ecflow 5.17 made + # read-only; pin below that until comfies is updated. + - ecflow<5.17 - nco - pyyaml - numpy @@ -17,4 +21,4 @@ dependencies: - boto3 - requests - pip: - - basininflow @ git+https://pypi.org/project/basininflow/ + - river-route>=2.1.1 diff --git a/geoglows_ecflow/resources/RAPIDpy/__init__.py b/geoglows_ecflow/resources/RAPIDpy/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/geoglows_ecflow/resources/RAPIDpy/dataset.py b/geoglows_ecflow/resources/RAPIDpy/dataset.py deleted file mode 100644 index 736e199..0000000 --- a/geoglows_ecflow/resources/RAPIDpy/dataset.py +++ /dev/null @@ -1,1236 +0,0 @@ -# -*- coding: utf-8 -*- -""" - dataset.py - RAPIDpy - - Created by Alan D Snow, 2016. - License: BSD-3-Clause -""" -import datetime -from csv import writer as csv_writer - -import numpy as np -import pandas as pd -from netCDF4 import Dataset, num2date -from numpy.ma import is_masked -from pytz import utc - -from .helper_functions import log, open_csv - - -# ----------------------------------------------------------------------------- -# Helper Function -# ----------------------------------------------------------------------------- -def compare_qout_files(dataset1_path, dataset2_path): - """ - This function compares the output of RAPID Qout and tells you where - they are different. - """ - qout_same = False - - d1 = RAPIDDataset(dataset1_path) - d2 = RAPIDDataset(dataset2_path) - - if len(d1.get_river_id_array()) != len(d2.get_river_id_array()): - log("Length of COMID/rivid input not the same.", - "ERROR") - - if not (d1.get_river_id_array() == d2.get_river_id_array()).all(): - log("COMID/rivid order is different in each dataset." - " Reordering data for comparison.", - "WARNING") - - d2_reordered_river_index_list = [] - for rivid in d1.get_river_id_array(): - reordered_index = np.where(d2.get_river_id_array() == rivid)[0][0] - d2_reordered_river_index_list.append(reordered_index) - d2_reordered_qout = d2.get_qout_index(d2_reordered_river_index_list) - else: - d2_reordered_qout = d2.get_qout() - - # get where the files are different - d1_qout = d1.get_qout() - where_diff = np.where(d1_qout != d2_reordered_qout) - un_where_diff = np.unique(where_diff[0]) - - # if different, check to see how different - if un_where_diff.any(): - decimal_test = 7 - while decimal_test > 0: - try: - np.testing.assert_almost_equal(d1_qout, - d2_reordered_qout, - decimal=decimal_test) - log("ALMOST EQUAL to {0} decimal places.".format(decimal_test), - "INFO") - qout_same = True - decimal_test = -1 - except AssertionError as ex: - if decimal_test <= 1: - log(ex, "WARNING") - decimal_test -= 1 - - log("Number of different timeseries: {0}".format(len(un_where_diff)), - "INFO") - log("COMID idexes where different: {0}".format(un_where_diff), - "INFO") - log("COMID idexes where different: {0}".format(un_where_diff), - "INFO") - index = un_where_diff[0] - log("Dataset 1 example. COMID index: " - "{0}".format(d1.get_qout_index(index)), - "INFO") - log("Dataset 2 example. COMID index: " - "{0}".format(d2_reordered_qout[index, :]), - "INFO") - - else: - qout_same = True - log("Output Qout data is the same.", - "INFO") - - d1.close() - d2.close() - return qout_same - - -# ------------------------------------------------------------------------------ -# Main Dataset Manager Class -# ------------------------------------------------------------------------------ -class RAPIDDataset(object): - """ - This class is designed to access data from the RAPID Qout - NetCDF file. - - Attributes - ---------- - filename: str - Path to the RAPID Qout NetCDF file. - river_id_dimension: str, optional - Name of the river ID dimension. Default is to search through - a pre-defined list. - river_id_variable: str, optional - Name of the river ID variable. Default is to search through - a pre-defined list. - streamflow_variable: str, optional - Name of the streamflow varaible. Default is to search through - a pre-defined list. - datetime_simulation_start: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the simulation start time. - simulation_time_step_seconds: int, optional - This is the time step of the simulation output in seconds. - out_tzinfo: tzinfo, optional - Time zone to output data as. The dates will be converted from UTC - to the time zone input. Default is UTC. - - - Example:: - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - #USE FUNCTIONS TO ACCESS DATA HERE - - """ - - # pylint: disable=too-many-instance-attributes - def __init__(self, filename, - river_id_dimension="", - river_id_variable="", - streamflow_variable="", - datetime_simulation_start=None, - simulation_time_step_seconds=None, - out_tzinfo=None): - """ - Initialize the class with variables given by the user - """ - self.qout_nc = Dataset(filename, mode='r') - - # determine river ID dimension - self.river_id_dimension = river_id_dimension - if not river_id_dimension: - if 'rivid' in self.qout_nc.dimensions: - self.river_id_dimension = 'rivid' - elif 'COMID' in self.qout_nc.dimensions: - self.river_id_dimension = 'COMID' - elif 'station' in self.qout_nc.dimensions: - self.river_id_dimension = 'station' - elif 'DrainLnID' in self.qout_nc.dimensions: - self.river_id_dimension = 'DrainLnID' - elif 'FEATUREID' in self.qout_nc.dimensions: - self.river_id_dimension = 'FEATUREID' - else: - raise IndexError('Could not find river ID dimension.') - elif river_id_dimension not in self.qout_nc.dimensions: - raise IndexError('Could not find river ID dimension:' - ' {0}.'.format(river_id_dimension)) - - self.size_river_id = len(self.qout_nc - .dimensions[self.river_id_dimension]) - - variable_keys = self.qout_nc.variables.keys() - - # determine streamflow variable - self.q_var_name = streamflow_variable - if not streamflow_variable: - if 'Qout' in variable_keys: - self.q_var_name = 'Qout' - elif 'streamflow' in variable_keys: - self.q_var_name = 'streamflow' - elif 'm3_riv' in variable_keys: - self.q_var_name = 'm3_riv' - else: - raise IndexError('ERROR: Could not find flow variable.' - ' Looked for Qout, streamflow, and m3_riv.') - elif streamflow_variable not in variable_keys: - raise IndexError('Could not find flow variable.' - ' Looked for {0}.'.format(streamflow_variable)) - - self.size_q_var = len(self.qout_nc.variables[self.q_var_name]) - - # determine time dimension - if 'time' in self.qout_nc.dimensions: - self.size_time = len(self.qout_nc.dimensions['time']) - elif 'Time' in self.qout_nc.dimensions: - self.size_time = len(self.qout_nc.dimensions['Time']) - else: - raise IndexError('Could not find time dimension.') - - # determine river ID variable - self.river_id_variable = river_id_variable - if not river_id_variable: - if 'rivid' in variable_keys: - self.river_id_variable = 'rivid' - elif 'COMID' in variable_keys: - self.river_id_variable = 'COMID' - elif 'station_id' in variable_keys: - self.river_id_variable = 'station_id' - elif 'DrainLnID' in variable_keys: - self.river_id_variable = 'DrainLnID' - elif 'FEATUREID' in variable_keys: - self.river_id_variable = 'FEATUREID' - else: - log('Could not find river ID variable' - ' in {0}.'.format(variable_keys), - "WARNING") - elif river_id_variable not in variable_keys: - log('Could not find river ID variable:' - ' {0}.'.format(river_id_variable), - "WARNING") - - self.out_tzinfo = out_tzinfo - self.datetime_simulation_start = datetime_simulation_start - self.simulation_time_step_seconds = simulation_time_step_seconds - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - self.close() - - def close(self): - """Close the dataset.""" - self.qout_nc.close() - - def _is_legacy_time_valid(self): - """ - This determines whether or not legacy time is set correctly. - - Returns - ------- - boolean: - True if the legacy time is setup correctly, otherwise false. - """ - return self.datetime_simulation_start is not None and \ - self.simulation_time_step_seconds is not None - - def is_time_variable_valid(self): - """ - This function returns whether or not the time variable - is valid. - - Returns - ------- - boolean - True if the time variable is valid, otherwise false. - - - Example:: - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - if qout_nc.is_time_variable_valid(): - #DO WORK HERE - - """ - # pylint: disable=len-as-condition - time_var_valid = False - if 'time' in self.qout_nc.variables.keys(): - if len(self.qout_nc.dimensions['time']) > 0: - if not is_masked(self.qout_nc.variables['time'][:]): - try: - timestep = ( - datetime.datetime.utcfromtimestamp(self.qout_nc.variables['time'][1]) - - datetime.datetime.utcfromtimestamp(self.qout_nc.variables['time'][0]) - # datetime.datetime.fromtimestamp(self.qout_nc.variables['time'][1], datetime.UTC) - - # datetime.datetime.fromtimestamp(self.qout_nc.variables['time'][0], datetime.UTC) - ).total_seconds() - if timestep > 0: - time_var_valid = True - except ValueError: - pass - - return time_var_valid - - def raise_time_valid(self): - """Raise ValueError if time not valid""" - if not (self.is_time_variable_valid() or self._is_legacy_time_valid()): - raise IndexError("Valid time variable not found. Valid time" - " variable required in Qout file to proceed ...") - - def get_time_array(self, - datetime_simulation_start=None, - simulation_time_step_seconds=None, - return_datetime=False, - time_index_array=None): - """ - This method extracts or generates an array of time. - The new version of RAPID output has the time array stored. - However, the old version requires the user to know when the - simulation began and the time step of the output. - - Parameters - ---------- - datetime_simulation_start: :obj:`datetime.datetime`, optional - The start datetime of the simulation. Only required if the time - variable is not included in the file. - simulation_time_step_seconds: int, optional - The time step of the file in seconds. Only required if the time - variable is not included in the file. - return_datetime: bool, optional - If true, it converts the data to a list of datetime objects. - Default is False. - time_index_array: list or :obj:`numpy.array`, optional - This is used to extract the datetime values by index from the main - list. This can be from the *get_time_index_range* function. - - Returns - ------- - list: - An array of integers representing seconds since Jan 1, 1970 UTC - or datetime objects if *return_datetime* is set to True. - - These examples demonstrates how to retrieve or generate a time array - to go along with your RAPID streamflow series. - - - CF-Compliant Qout File Example: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - #retrieve integer timestamp array - time_array = qout_nc.get_time_array() - - #or, to get datetime array - time_datetime = qout_nc.get_time_array(return_datetime=True) - - - Legacy Qout File Example: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - with RAPIDDataset(path_to_rapid_qout, - datetime_simulation_start=datetime(1980, 1, 1), - simulation_time_step_seconds=3 * 3600)\ - as qout_nc: - - #retrieve integer timestamp array - time_array = qout_nc.get_time_array() - - #or, to get datetime array - time_datetime = qout_nc.get_time_array(return_datetime=True) - - """ - # Original Qout file - if datetime_simulation_start is not None: - self.datetime_simulation_start = datetime_simulation_start - if simulation_time_step_seconds is not None: - self.simulation_time_step_seconds = simulation_time_step_seconds - - epoch = datetime.datetime(1970, 1, 1, tzinfo=utc) - time_units = "seconds since {0}".format(epoch) - - # CF-1.6 compliant file - if self.is_time_variable_valid(): - time_array = self.qout_nc.variables['time'][:] - if self.qout_nc.variables['time'].units: - time_units = self.qout_nc.variables['time'].units - - # Original Qout file - elif self._is_legacy_time_valid(): - initial_time_seconds = ((self.datetime_simulation_start - .replace(tzinfo=utc) - epoch) - .total_seconds() + - self.simulation_time_step_seconds) - final_time_seconds = (initial_time_seconds + - self.size_time * - self.simulation_time_step_seconds) - time_array = np.arange(initial_time_seconds, - final_time_seconds, - self.simulation_time_step_seconds) - else: - raise ValueError("This file does not contain the time" - " variable. To get time array, add" - " datetime_simulation_start and" - " simulation_time_step_seconds") - - if time_index_array is not None: - time_array = time_array[time_index_array] - - if return_datetime: - try: - # only_use_cftime_datetime=True is default after cftime PR #135 - time_array = num2date(time_array, time_units, only_use_cftime_datetimes=False) - except: - time_array = num2date(time_array, time_units) - - if self.out_tzinfo is not None: - for i in range(len(time_array)): - # convert time to output timezone - time_array[i] = utc.localize(time_array[i]) \ - .astimezone(self.out_tzinfo) \ - .replace(tzinfo=None) - - return time_array - - def get_time_index_range(self, - date_search_start=None, - date_search_end=None, - time_index_start=None, - time_index_end=None, - time_index=None): - """ - Generates a time index range based on time bounds given. - This is useful for subset data extraction. - - Parameters - ---------- - date_search_start: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the minimum date for - starting. - date_search_end: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the maximum date - for ending. - time_index_start: int, optional - This is the index of the start of the time array subset. - Useful for the old file version. - time_index_end: int, optional - This is the index of the end of the time array subset. - Useful for the old file version. - time_index: int, optional - This is the index of time to return in the case that your - code only wants one index. Used internally. - - Returns - ------- - :obj:`numpy.array`: - This is an array of time indices used to extract a subset of data. - - - CF-Compliant Qout File Example: - - .. code:: python - - from datetime import datetime - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - time_index_range = qout_nc.get_time_index_range( - date_search_start=datetime(1980, 1, 1), - date_search_end=datetime(1980, 12, 11)) - - - Legacy Qout File Example: - - .. code:: python - - from datetime import datetime - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - with RAPIDDataset(path_to_rapid_qout, - datetime_simulation_start=datetime(1980, 1, 1), - simulation_time_step_seconds=3600) as qout_nc: - - time_index_range = qout_nc.get_time_index_range( - date_search_start=datetime(1980, 1, 1), - date_search_end=datetime(1980, 12, 11)) - - """ - # get the range of time based on datetime range - time_range = None - if ((self.is_time_variable_valid() or self._is_legacy_time_valid()) and - (date_search_start is not None or - date_search_end is not None)): - - log("Determining time range ({0} to {1})" - "...".format(date_search_start, date_search_end), - "INFO") - time_array = self.get_time_array() - if date_search_start is not None: - date_search_start_utc = date_search_start - if self.out_tzinfo is not None: - date_search_start_utc = self.out_tzinfo \ - .localize(date_search_start) \ - .astimezone(utc) \ - .replace(tzinfo=None) - seconds_start = (date_search_start_utc - - datetime.datetime(1970, 1, 1)).total_seconds() - time_range = np.where(time_array >= seconds_start)[0] - - if date_search_end is not None: - date_search_end_utc = date_search_end - if self.out_tzinfo is not None: - date_search_end_utc = self.out_tzinfo \ - .localize(date_search_end) \ - .astimezone(utc) \ - .replace(tzinfo=None) - - seconds_end = (date_search_end_utc - - datetime.datetime(1970, 1, 1)).total_seconds() - if time_range is not None: - time_range = np.intersect1d(time_range, - np.where(time_array <= - seconds_end)[0]) - else: - time_range = np.where(time_array <= seconds_end)[0] - - # get the range of time based on time index range - elif time_index_start is not None or time_index_end is not None: - if time_index_start is None: - time_index_start = 0 - if time_index_end is None: - time_index_end = self.size_time - time_range = range(time_index_start, time_index_end) - - # get only one time step - elif time_index is not None: - time_range = [time_index] - # return all - else: - time_range = range(self.size_time) - - return time_range - - def get_river_id_array(self): - """ - This method returns the river ID array for this file. - - Returns - ------- - :obj:`numpy.array`: - An array of the river ID's - - - Example:: - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - river_ids = qout_nc.get_river_id_array() - - """ - return self.qout_nc.variables[self.river_id_variable][:] - - def get_river_index(self, river_id): - """ - This method retrieves the river index in the netCDF - dataset corresponding to the river ID. - - Parameters - ---------- - river_id: int - The ID of the river segment. - - Returns - ------- - int: - The index of the river ID's in the file. - - - Example:: - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - river_id = 53458 - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - river_index = qout_nc.get_river_index(river_id) - - """ - try: - return np.where(self.get_river_id_array() == river_id)[0][0] - except IndexError: - raise IndexError("ERROR: River ID {0} not found in dataset " - "...".format(river_id)) - - def get_subset_riverid_index_list(self, river_id_list): - """ - Gets the subset riverid_list from the netcdf file - Optional returns include the list of valid river ids in the dataset - as well as a list of missing rive rids - - Parameters - ---------- - river_id_list: list or :obj:`numpy.array` - Array of river ID's for the river segments you want the index of. - - Returns - ------- - :obj:`numpy.array` - A sorted array of the river index in the NetCDF file that - were found. - :obj:`numpy.array` - A sorted array of the river IDs that were found. - list - An array of the missing river ids. - - """ - netcdf_river_indices_list = [] - valid_river_ids = [] - missing_river_ids = [] - for river_id in river_id_list: - # get where streamids are in netcdf file - try: - netcdf_river_indices_list \ - .append(self.get_river_index(river_id)) - valid_river_ids.append(river_id) - except IndexError: - log("ReachID {0} not found in netCDF dataset." - " Skipping ...".format(river_id), - "WARNING") - missing_river_ids.append(river_id) - - np_valid_river_indices_list = np.array(netcdf_river_indices_list) - np_valid_river_ids = np.array(valid_river_ids) - sorted_indexes = np.argsort(np_valid_river_indices_list) - - return (np_valid_river_indices_list[sorted_indexes], - np_valid_river_ids[sorted_indexes], - np.array(missing_river_ids)) - - def get_qout(self, - river_id_array=None, - date_search_start=None, - date_search_end=None, - time_index_start=None, - time_index_end=None, - time_index=None, - time_index_array=None, - daily=False, - pd_filter=None, - filter_mode="mean", - as_dataframe=False): - """ - This method extracts streamflow data by a single river ID - or by a river ID array. It has options to extract by date - or by date index. - - Parameters - ---------- - river_id_array: :obj:`numpy.array` or list or int, optional - A single river ID or an array of river IDs. - date_search_start: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the minimum date - for starting. - date_search_end: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the maximum date - for ending. - time_index_start: int, optional - This is the index of the start of the time array subset. - Useful for the old file version. - time_index_end: int, optional - This is the index of the end of the time array subset. - Useful for the old file version. - time_index: int, optional - This is the index of time to return in the case that your - code only wants one index. Used internally. - time_index_array: list or :obj:`numpy.array`, optional - This is used to extract the vales only for particular dates. - This can be from the *get_time_index_range* function. - daily: bool, optional - If true, this will convert qout to daily average. - pd_filter: str, optional - This is a valid pandas resample frequency filter. - filter_mode: str, optional - You can get the daily average "mean" or the maximum "max". - Default is "mean". - as_dataframe: bool, optional - Return as a pandas dataframe object. Default is False. - - - Returns - ------- - qout_array: :obj:`numpy.array` - This is a 1D or 2D array or a single value depending on your - input search. - - - This example demonstrates how to retrieve the streamflow associated - with the reach you are interested in:: - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - river_id = 500 - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - streamflow_array = qout_nc.get_qout(river_id) - - This example demonstrates how to retrieve the streamflow within a date - range associated with the reach you are interested in:: - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - river_id = 500 - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - streamflow_array = qout_nc.get_qout( - river_id, - date_search_start=datetime(1985,1,1), - date_search_end=datetime(1985,2,4)) - - """ - # get indices of where the streamflow data is - riverid_index_list_subset = None - if river_id_array is not None: - if not hasattr(river_id_array, "__len__"): - river_id_array = [river_id_array] - riverid_index_list_subset = \ - self.get_subset_riverid_index_list(river_id_array)[0] - - return self.get_qout_index(riverid_index_list_subset, - date_search_start, - date_search_end, - time_index_start, - time_index_end, - time_index, - time_index_array, - daily, - pd_filter, - filter_mode, - as_dataframe) - - def get_qout_index(self, - river_index_array=None, - date_search_start=None, - date_search_end=None, - time_index_start=None, - time_index_end=None, - time_index=None, - time_index_array=None, - daily=False, - pd_filter=None, - filter_mode="mean", - as_dataframe=False): - """ - This method extracts streamflow data by river index. - It allows for extracting single or multiple river streamflow arrays - It has options to extract by date or by date index. - - See: :meth:`RAPIDpy.RAPIDDataset.get_qout` - """ - if river_index_array is not None: - if hasattr(river_index_array, "__len__"): - if len(river_index_array) == 1: - river_index_array = river_index_array[0] - - if time_index_array is None: - time_index_array = self.get_time_index_range(date_search_start, - date_search_end, - time_index_start, - time_index_end, - time_index) - - qout_variable = self.qout_nc.variables[self.q_var_name] - qout_dimensions = qout_variable.dimensions - if qout_dimensions[0].lower() == 'time' and \ - qout_dimensions[1].lower() == self.river_id_dimension.lower(): - if time_index_array is not None and river_index_array is not None: - streamflow_array = qout_variable[time_index_array, - river_index_array].transpose() - elif time_index_array is not None: - streamflow_array = qout_variable[time_index_array, :] \ - .transpose() - elif river_index_array is not None: - streamflow_array = qout_variable[:, river_index_array] \ - .transpose() - else: - streamflow_array = qout_variable[:].transpose() - elif qout_dimensions[1].lower() == 'time' and \ - qout_dimensions[0].lower() == self.river_id_dimension.lower(): - if time_index_array is not None and river_index_array is not None: - streamflow_array = qout_variable[river_index_array, - time_index_array] - elif time_index_array is not None: - streamflow_array = qout_variable[:, time_index_array] - elif river_index_array is not None: - streamflow_array = qout_variable[river_index_array, :] - else: - streamflow_array = qout_variable[:] - else: - raise Exception("Invalid RAPID Qout file dimensions ...") - - if daily: - pd_filter = "D" - - if pd_filter is not None or as_dataframe: - time_array = self.get_time_array(return_datetime=True, - time_index_array=time_index_array) - qout_df = pd.DataFrame(streamflow_array.T, index=time_array) - - if pd_filter is not None: - qout_df = qout_df.resample(pd_filter) - if filter_mode == "mean": - qout_df = qout_df.mean() - elif filter_mode == "max": - qout_df = qout_df.max() - else: - raise Exception("Invalid filter_mode ...") - - if as_dataframe: - return qout_df - - try: - streamflow_array = qout_df.to_numpy().T - except: - streamflow_array = qout_df.as_matrix().T - - if streamflow_array.ndim > 0 and streamflow_array.shape[0] == 1: - streamflow_array = streamflow_array[0] - - return streamflow_array - - def write_flows_to_csv(self, path_to_output_file, - river_index=None, - river_id=None, - date_search_start=None, - date_search_end=None, - daily=False, - filter_mode="mean"): - """ - Write out RAPID output to CSV file. - - .. note:: Need either *reach_id* or *reach_index* parameter, - but either can be used. - - Parameters - ---------- - path_to_output_file: str - Path to the output csv file. - river_index: :obj:`datetime.datetime`, optional - This is the index of the river in the file you want the - streamflow for. - river_id: :obj:`datetime.datetime`, optional - This is the river ID that you want the streamflow for. - date_search_start: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the minimum date - for starting. - date_search_end: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the maximum date - for ending. - daily: bool, optional - If True and the file is CF-Compliant, write out daily flows. - filter_mode: str, optional - You can get the daily average "mean" or the maximum "max". - Default is "mean". - - - Example writing entire time series to file: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - river_id = 3624735 - path_to_rapid_qout = '/path/to/Qout.nc' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - #for writing entire time series to file - qout_nc.write_flows_to_csv('/timeseries/Qout_3624735.csv', - river_id=river_id, - ) - - - #if file is CF compliant, you can write out daily average - - #NOTE: Getting the river index is not necessary - #this is just an example of how to use this - river_index = qout_nc.get_river_index(river_id) - qout_nc.write_flows_to_csv('/timeseries/Qout_daily.csv', - river_index=river_index, - daily=True, - ) - - Example writing entire time series as daily average to file: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - river_id = 3624735 - path_to_rapid_qout = '/path/to/Qout.nc' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - #NOTE: Getting the river index is not necessary - #this is just an example of how to use this - river_index = qout_nc.get_river_index(river_id) - - #if file is CF compliant, you can write out daily average - qout_nc.write_flows_to_csv('/timeseries/Qout_daily.csv', - river_index=river_index, - daily=True, - ) - - Example writing entire time series as daily average to file: - - .. code:: python - - from datetime import datetime - from RAPIDpy import RAPIDDataset - - river_id = 3624735 - path_to_rapid_qout = '/path/to/Qout.nc' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - # if file is CF compliant, you can filter by date - qout_nc.write_flows_to_csv( - '/timeseries/Qout_daily_date_filter.csv', - river_id=river_id, - daily=True, - date_search_start=datetime(2002, 8, 31), - date_search_end=datetime(2002, 9, 15), - filter_mode="max" - ) - """ - if river_id is not None: - river_index = self.get_river_index(river_id) - elif river_id is None and river_index is None: - raise ValueError("Need reach id or reach index ...") - - # analyze and write - if self.is_time_variable_valid() or self._is_legacy_time_valid(): - qout_df = self.get_qout_index(river_index, - date_search_start=date_search_start, - date_search_end=date_search_end, - daily=daily, - filter_mode=filter_mode, - as_dataframe=True) - - qout_df.to_csv(path_to_output_file, header=False) - - else: - log("Valid time variable not found. Printing values only ...", "WARNING") - qout_arr = self.get_qout_index(river_index) - with open_csv(path_to_output_file, 'w') as outcsv: - writer = csv_writer(outcsv) - for index in range(len(qout_arr)): - writer.writerow([index, "{0:.5f}".format(qout_arr[index])]) - - def write_flows_to_gssha_time_series_xys(self, - path_to_output_file, - series_name, - series_id, - river_index=None, - river_id=None, - date_search_start=None, - date_search_end=None, - daily=False, - filter_mode="mean"): - """ - Write out RAPID output to GSSHA WMS time series xys file. - - Parameters - ---------- - path_to_output_file: str - Path to the output xys file. - series_name: str - The name for the series. - series_id: int - The ID to give the series. - river_index: :obj:`datetime.datetime`, optional - This is the index of the river in the file you want the - streamflow for. - river_id: :obj:`datetime.datetime`, optional - This is the river ID that you want the streamflow for. - date_search_start: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the minimum date for - starting. - date_search_end: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the maximum date for - ending. - daily: bool, optional - If True and the file is CF-Compliant, write out daily flows. - filter_mode: str, optional - You can get the daily average "mean" or the maximum "max". - Defauls is "mean". - - - Example writing entire time series to file: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - river_id = 3624735 - path_to_rapid_qout = '/path/to/Qout.nc' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - qout_nc.write_flows_to_gssha_time_series_xys( - '/timeseries/Qout_{0}.xys'.format(river_id), - series_name="RAPID_TO_GSSHA_{0}".format(river_id), - series_id=34, - river_id=river_id) - - - Example writing entire time series as daily average to file: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - river_id = 3624735 - path_to_rapid_qout = '/path/to/Qout.nc' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - # NOTE: Getting the river index is not necessary - # this is just an example of how to use this - river_index = qout_nc.get_river_index(river_id) - - # if file is CF compliant, you can write out daily average - qout_nc.write_flows_to_gssha_time_series_xys( - '/timeseries/Qout_daily.xys', - series_name="RAPID_TO_GSSHA_{0}".format(river_id), - series_id=34, - river_index=river_index, - daily=True) - - - Example writing subset of time series as daily maximum to file: - - .. code:: python - - from datetime import datetime - from RAPIDpy import RAPIDDataset - - river_id = 3624735 - path_to_rapid_qout = '/path/to/Qout.nc' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - # NOTE: Getting the river index is not necessary - # this is just an example of how to use this - river_index = qout_nc.get_river_index(river_id) - - # if file is CF compliant, you can filter by date and - # get daily values - qout_nc.write_flows_to_gssha_time_series_xys( - '/timeseries/Qout_daily_date_filter.xys', - series_name="RAPID_TO_GSSHA_{0}".format(river_id), - series_id=34, - river_index=river_index, - date_search_start=datetime(2002, 8, 31), - date_search_end=datetime(2002, 9, 15), - daily=True, - filter_mode="max") - - """ - if river_id is not None: - river_index = self.get_river_index(river_id) - elif river_id is None and river_index is None: - raise ValueError(" Need reach id or reach index ...") - - self.raise_time_valid() - - # analyze and write - qout_df = self.get_qout_index(river_index, - date_search_start=date_search_start, - date_search_end=date_search_end, - daily=daily, - filter_mode=filter_mode, - as_dataframe=True) - - with open_csv(path_to_output_file, 'w') as out_ts: - out_ts.write("XYS {0} {1} \"{2}\"\r\n".format(series_id, - len(qout_df.index), - series_name)) - for index, pd_row in qout_df.iterrows(): - date_str = index.strftime("%m/%d/%Y %I:%M:%S %p") - out_ts.write("\"{0}\" {1:.5f}\n".format(date_str, - pd_row[0])) - - def write_flows_to_gssha_time_series_ihg(self, - path_to_output_file, - connection_list_file, - date_search_start=None, - date_search_end=None, - daily=False, - filter_mode="mean"): - # pylint: disable=line-too-long - """ - Write out RAPID output to GSSHA time series ihg file - - .. note:: See: http://www.gsshawiki.com/Surface_Water_Routing:Introducing_Dischage/Constituent_Hydrographs - - .. note:: GSSHA project card is CHAN_POINT_INPUT - - Parameters - ---------- - path_to_output_file: str - Path to the output xys file. - connection_list_file: str - CSV file with link_id, node_id, baseflow, and rapid_rivid header - and rows with data. - date_search_start: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the minimum date - for starting. - date_search_end: :obj:`datetime.datetime`, optional - This is a datetime object with the date of the maximum date - for ending. - daily: bool, optional - If True and the file is CF-Compliant, write out daily flows. - filter_mode: str, optional - You can get the daily average "mean" or the maximum "max". - Defauls is "mean". - - - Example connection list file:: - - link_id, node_id, baseflow, rapid_rivid - 599, 1, 0.0, 80968 - 603, 1, 0.0, 80967 - - - Example writing entire time series to file: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - connection_list_file = '/path/to/connection_list_file.csv' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - #for writing entire time series to file - qout_nc.write_flows_to_gssha_time_series_ihg( - '/timeseries/Qout_3624735.ihg', - connection_list_file) - - - Example writing entire time series as daily average to file: - - .. code:: python - - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - connection_list_file = '/path/to/connection_list_file.csv' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - # if file is CF compliant, you can write out daily average - qout_nc.write_flows_to_gssha_time_series_ihg( - '/timeseries/Qout_3624735.ihg', - connection_list_file, - daily=True) - - - Example writing subset of time series as daily maximum to file: - - .. code:: python - - from datetime import datetime - from RAPIDpy import RAPIDDataset - - path_to_rapid_qout = '/path/to/Qout.nc' - connection_list_file = '/path/to/connection_list_file.csv' - - with RAPIDDataset(path_to_rapid_qout) as qout_nc: - # if file is CF compliant, you can filter by - # date and get daily values - qout_nc.write_flows_to_gssha_time_series_ihg( - '/timeseries/Qout_daily_date_filter.ihg', - connection_list_file, - date_search_start=datetime(2002, 8, 31), - date_search_end=datetime(2002, 9, 15), - daily=True, - filter_mode="max") - """ # noqa - self.raise_time_valid() - - # analyze and write - with open_csv(path_to_output_file, 'w') as out_ts: - # HEADER SECTION EXAMPLE: - # NUMPT 3 - # POINT 1 599 0.0 - # POINT 1 603 0.0 - # POINT 1 605 0.0 - - connection_list = np.loadtxt(connection_list_file, - skiprows=1, ndmin=1, - delimiter=',', - usecols=(0, 1, 2, 3), - dtype={'names': ('link_id', - 'node_id', - 'baseflow', - 'rapid_rivid'), - 'formats': ('i8', 'i8', - 'f4', 'i8') - }, - ) - - out_ts.write("NUMPT {0}\n".format(connection_list.size)) - - river_idx_list = [] - for connection in connection_list: - out_ts.write("POINT {0} {1} {2}\n" - "".format(connection['node_id'], - connection['link_id'], - connection['baseflow'], - ), - ) - river_idx_list.append( - self.get_river_index(connection['rapid_rivid']) - ) - - # INFLOW SECTION EXAMPLE: - # NRPDS 54 - # INPUT 2002 01 01 00 00 15.551210 12.765090 0.000000 - # INPUT 2002 01 02 00 00 15.480830 12.765090 0.000000 - # INPUT 2002 01 03 00 00 16.078910 12.765090 0.000000 - # ... - qout_df = self.get_qout_index( - river_idx_list, - date_search_start=date_search_start, - date_search_end=date_search_end, - daily=daily, - filter_mode=filter_mode, - as_dataframe=True) - - out_ts.write("NRPDS {0}\n".format(len(qout_df.index))) - - for index, pd_row in qout_df.iterrows(): - date_str = index.strftime("%Y %m %d %H %M") - qout_str = " ".join(["{0:.5f}".format(pd_row[column]) - for column in qout_df]) - out_ts.write("INPUT {0} {1}\n".format(date_str, qout_str)) diff --git a/geoglows_ecflow/resources/RAPIDpy/helper_functions.py b/geoglows_ecflow/resources/RAPIDpy/helper_functions.py deleted file mode 100644 index 0798eaa..0000000 --- a/geoglows_ecflow/resources/RAPIDpy/helper_functions.py +++ /dev/null @@ -1,145 +0,0 @@ -# -*- coding: utf-8 -*- -""" - helper_functions.py - RAPIDpy - - Created by Alan D Snow, 2015. -""" -import csv -import netCDF4 as nc -from os import remove -from sys import version_info - -from numpy.testing import assert_almost_equal -from numpy import array as np_array -from numpy import float32 as np_float32 - - -# ----------------------------------------------------------------------------- -# HELPER FUNCTIONS -# ----------------------------------------------------------------------------- -# pylint: disable=line-too-long -def open_csv(csv_file, mode='r'): - """ - Get mode depending on Python version - Based on: http://stackoverflow.com/questions/29840849/writing-a-csv-file-in-python-that-works-for-both-python-2-7-and-python-3-3-in - """ # noqa - if version_info[0] == 2: # Not named on 2.6 - access = '{0}b'.format(mode) - kwargs = {} - else: - access = '{0}t'.format(mode) - kwargs = {'newline': ''} - - return open(csv_file, access, **kwargs) - - -def log(message, severity="INFO", print_debug=True): - """Logs, prints, or raises a message. - - Arguments: - message -- message to report - severity -- string of one of these values: - CRITICAL|ERROR|WARNING|INFO|DEBUG - """ - - print_me = ['WARNING', 'INFO', 'DEBUG'] - if severity in print_me: - if severity == 'DEBUG': - if print_debug: - print("{0}: {1}".format(severity, message)) - else: - print("{0}: {1}".format(severity, message)) - else: - raise Exception("{0}: {1}".format(severity, message)) - - -def csv_to_list(csv_file, delimiter=','): - """ - Reads in a CSV file and returns the contents as list, - where every row is stored as a sublist, and each element - in the sublist represents 1 cell in the table. - """ - with open_csv(csv_file) as csv_con: - if len(delimiter) > 1: - dialect = csv.Sniffer().sniff(csv_con.read(1024), - delimiters=delimiter) - csv_con.seek(0) - reader = csv.reader(csv_con, dialect) - else: - reader = csv.reader(csv_con, delimiter=delimiter) - return list(reader) - -def netcdf_to_list(netcdf_file): - """ - Reads in a netCDF file and returns the contents as list - """ - with nc.Dataset(netcdf_file, 'r') as nc_con: - netcdf_list = [[i] for i in nc_con['Qout'][0][:]] - return netcdf_list - - -def compare_csv_decimal_files(file1, file2, header=True, timeseries=False): - """ - This function compares two csv files - """ - # CHECK NUM LINES - with open_csv(file1) as fh1, \ - open_csv(file2) as fh2: - assert sum(1 for _ in fh1) == sum(1 for _ in fh2) - - with open_csv(file1) as fh1, \ - open_csv(file2) as fh2: - csv1 = csv.reader(fh1) - csv2 = csv.reader(fh2) - - if header: - assert next(csv1) == next(csv2) # header - - while True: - try: - row1 = next(csv1) - row2 = next(csv2) - compare_start_index = 0 - if timeseries: - assert row1[0] == row2[0] # check dates - compare_start_index = 1 - - assert_almost_equal( - np_array(row1[compare_start_index:], dtype=np_float32), - np_array(row2[compare_start_index:], dtype=np_float32), - decimal=2) - except StopIteration: - break - return True - - -def compare_csv_timeseries_files(file1, file2, header=True): - """ - This function compares two csv files - """ - return compare_csv_decimal_files(file1, file2, header, True) - - -def remove_files(*args): - """ - This function removes all files input as arguments - """ - for arg in args: - try: - remove(arg) - except OSError: - pass - - -def add_latlon_metadata(lat_var, lon_var): - """Adds latitude and longitude metadata""" - lat_var.long_name = 'latitude' - lat_var.standard_name = 'latitude' - lat_var.units = 'degrees_north' - lat_var.axis = 'Y' - - lon_var.long_name = 'longitude' - lon_var.standard_name = 'longitude' - lon_var.units = 'degrees_east' - lon_var.axis = 'X' \ No newline at end of file diff --git a/geoglows_ecflow/resources/RAPIDpy/postprocess_merge.py b/geoglows_ecflow/resources/RAPIDpy/postprocess_merge.py deleted file mode 100644 index c579ea8..0000000 --- a/geoglows_ecflow/resources/RAPIDpy/postprocess_merge.py +++ /dev/null @@ -1,539 +0,0 @@ -# -*- coding: utf-8 -*- -""" -merge.py -RAPIDpy - -Created by Tim Whitaker, 2015. -Modified by Alan D Snow, 2015-2016 - -Copies data from RAPID netCDF output to a CF-compliant netCDF file. -Code originated from Tim Whitaker at University of Texas. The code was -modified by Alan Snow at US Army ERDC. - -Remarks: - A new netCDF file is created with data from RAPID [1] simulation model - output. The result follows CF conventions [2] with additional metadata - prescribed by the NODC timeSeries Orthogonal template [3] for time series - at discrete point feature locations. - - This script was created for the National Flood Interoperability Experiment, - and so metadata in the result reflects that. - -Requires: - netcdf4-python - https://github.com/Unidata/netcdf4-python - -Inputs: - Lookup CSV table with COMID, Lat, Lon, and Elev_m columns. Columns must - be in that order and these must be the first four columns. The order of - COMIDs in the table must match the order of features in the netCDF file. - - -/////////////////////////////////////////////////// -netcdf result_2014100520141101 { -dimensions: - Time = UNLIMITED ; // (224 currently) - COMID = 61818 ; -variables: - float Qout(Time, COMID) ; -/////////////////////////////////////////////////// - -Outputs: - CF-compliant netCDF file of RAPID results, named with original filename - with "_CF" appended to the filename. File is written to 'output' folder. - - Input netCDF file is archived or deleted, based on 'archive' config - parameter. - -Usage: - import the script, e.g., import ConvertRAPIDOutputToCF as cf. - - -References: - [1] http://rapid-hub.org/ - [2] http://cfconventions.org/ - [3] http://www.nodc.noaa.gov/data/formats/netcdf/v1.1/ -""" -import os -from datetime import datetime - -import numpy as np -from netCDF4 import Dataset -from pytz import utc - -# local -from .dataset import RAPIDDataset -from .helper_functions import (add_latlon_metadata, csv_to_list, - netcdf_to_list, remove_files, log) - - -class ConvertRAPIDOutputToCF(object): - """ - Class to convert RAPID output to be CF compliant. You can also use this to - combine consecutive RAPID output files into one file. - - Parameters - ---------- - rapid_output_file: str or list - Path to a single RAPID Qout file or a list of RAPID Qout files. - start_datetime: :obj:`datetime.datetime` - Datetime object with the time of the start of the simulation. - time_step: int or list - Time step of simulation in seconds if single Qout file or a list of - time steps corresponding to each Qout file in the *rapid_output_file*. - qinit_file: str, optional - Path to the Qinit file for the simulation. If used, it will use the - values in the file for the flow at simulation time zero. - comid_lat_lon_z_file: str, optional - Path to comid_lat_lon_z file. If included, the spatial information - will be added to the output NetCDF file. - rapid_connect_file: str, optional - Path to RAPID connect file. This is required if *qinit_file* is added. - project_name: str, optional - Name of your project in the output file. Default is - "Default RAPID Project". - output_id_dim_name: str, optional - Name of the output river ID dimension name. Default is 'rivid'. - output_flow_var_name: str, optional - Name of streamflow variable in output file, typically - 'Qout' or 'm3_riv'. Default is 'Qout'. - print_debug: bool, optional - If True, the debug output will be printed to the console. - Default is False. - - - .. warning:: This code replaces the first file with the combined output and - deletes the second file. BACK UP YOUR FILES!!!! - - - Example: - - .. code:: python - - import datetime - from RAPIDpy.postprocess import ConvertRAPIDOutputToCF - - file1 = "/path/to/Qout_1980to1981.nc" - file2 = "/path/to/Qout_1981to1982.nc" - - cv = ConvertRAPIDOutputToCF( - rapid_output_file=[file1, file2], - start_datetime=datetime.datetime(2005,1,1), - time_step=[3*3600, 3*3600], - project_name="NLDAS(VIC)-RAPID historical flows by US Army ERDC") - cv.convert() - - """ - - # pylint: disable= - def __init__(self, - rapid_output_file, - start_datetime, - time_step, - qinit_file="", - comid_lat_lon_z_file="", - rapid_connect_file="", - project_name="Default RAPID Project", - output_id_dim_name='rivid', - output_flow_var_name='Qout', - print_debug=False): - if not isinstance(rapid_output_file, list): - self.rapid_output_file_list = [rapid_output_file] - else: - self.rapid_output_file_list = rapid_output_file - self.start_datetime = start_datetime.replace(tzinfo=utc) - - if not isinstance(time_step, list): - self.time_step_array = [time_step] - else: - self.time_step_array = time_step - - self.qinit_file = qinit_file - self.comid_lat_lon_z_file = comid_lat_lon_z_file - self.rapid_connect_file = rapid_connect_file - self.project_name = project_name - self.output_id_dim_name = output_id_dim_name - self.output_flow_var_name = output_flow_var_name - self.print_debug = print_debug - self.cf_compliant_file = '%s_CF.nc' % os.path.splitext( - self.rapid_output_file_list[0])[0] - self.cf_nc = None - self.raw_nc_list = [] - - def _validate_raw_nc(self): - """Checks that raw netCDF file has the right dimensions and variables. - - Returns - ------- - int: - Length of rivid dimension. - int: - Length of time dimension. - - Remarks: Raises exception if file doesn't validate. - """ - self.raw_nc_list = [] - # add one for the first flow value RAPID - # does not include - total_time_len = 1 - id_len_list = [] - for rapid_output_file in self.rapid_output_file_list: - qout_nc = RAPIDDataset(rapid_output_file) - id_len_list.append(qout_nc.size_river_id) - total_time_len += qout_nc.size_time - self.raw_nc_list.append(qout_nc) - - # make sure river id lists are the same - for id_len_undex in range(1, len(id_len_list)): - if id_len_list[id_len_undex] != id_len_list[0]: - raise Exception("River ID size is different in " - "one of the files ...") - - for raw_nc_index in range(1, len(self.raw_nc_list)): - if not (self.raw_nc_list[raw_nc_index].get_river_id_array() == - self.raw_nc_list[0].get_river_id_array()).all(): - raise Exception("River IDs are different in " - "files ...") - - return id_len_list[0], total_time_len - - def _initialize_output(self, time_len, id_len): - """Creates netCDF file with CF dimensions and variables, but no data. - - Arguments - --------- - time_len: int - Length of time dimension (number of time steps). - id_len: int - Length of Id dimension (number of time series). - - """ - log('Initializing new file %s' % self.cf_compliant_file, 'INFO') - - self.cf_nc = Dataset(self.cf_compliant_file, 'w', - format='NETCDF3_CLASSIC') - - # Create global attributes - log(' globals', 'DEBUG', self.print_debug) - self.cf_nc.featureType = 'timeSeries' - self.cf_nc.Metadata_Conventions = 'Unidata Dataset Discovery v1.0' - self.cf_nc.Conventions = 'CF-1.6' - self.cf_nc.cdm_data_type = 'Station' - self.cf_nc.nodc_template_version = ( - 'NODC_NetCDF_TimeSeries_Orthogonal_Template_v1.1') - self.cf_nc.standard_name_vocabulary = \ - ('NetCDF Climate and Forecast (CF) ' - 'Metadata Convention Standard Name ' - 'Table v28') - self.cf_nc.title = 'RAPID Result' - self.cf_nc.summary = \ - ("Results of RAPID river routing simulation. Each river " - "reach (i.e., feature) is represented by a point " - "feature at its midpoint, and is identified by the " - "reach's unique NHDPlus COMID identifier.") - self.cf_nc.time_coverage_resolution = 'point' - self.cf_nc.geospatial_lat_min = 0.0 - self.cf_nc.geospatial_lat_max = 0.0 - self.cf_nc.geospatial_lat_units = 'degrees_north' - self.cf_nc.geospatial_lat_resolution = 'midpoint of stream feature' - self.cf_nc.geospatial_lon_min = 0.0 - self.cf_nc.geospatial_lon_max = 0.0 - self.cf_nc.geospatial_lon_units = 'degrees_east' - self.cf_nc.geospatial_lon_resolution = 'midpoint of stream feature' - self.cf_nc.geospatial_vertical_min = 0.0 - self.cf_nc.geospatial_vertical_max = 0.0 - self.cf_nc.geospatial_vertical_units = 'm' - self.cf_nc.geospatial_vertical_resolution = \ - 'midpoint of stream feature' - self.cf_nc.geospatial_vertical_positive = 'up' - self.cf_nc.project = self.project_name - self.cf_nc.processing_level = 'Raw simulation result' - self.cf_nc.keywords_vocabulary = \ - ('NASA/Global Change Master Directory ' - '(GCMD) Earth Science Keywords. Version ' - '8.0.0.0.0') - self.cf_nc.keywords = 'DISCHARGE/FLOW' - self.cf_nc.comment = \ - 'Result time step(s) (seconds): ' + str(self.time_step_array) - - timestamp = datetime.utcnow().isoformat() + 'Z' - self.cf_nc.date_created = timestamp - self.cf_nc.history = \ - (timestamp + '; added time, lat, lon, z, crs variables; ' - 'added metadata to conform to NODC_NetCDF_TimeSeries_' - 'Orthogonal_Template_v1.1') - - # Create dimensions - log(' dimming', 'DEBUG', self.print_debug) - self.cf_nc.createDimension('time', time_len) - self.cf_nc.createDimension(self.output_id_dim_name, id_len) - - # Create variables - log(' time_series_var', 'DEBUG', self.print_debug) - time_series_var = \ - self.cf_nc.createVariable(self.output_id_dim_name, 'i4', - (self.output_id_dim_name,)) - time_series_var.long_name = ( - 'Unique NHDPlus COMID identifier for each river reach feature') - time_series_var.cf_role = 'timeseries_id' - - log(' time_var', 'DEBUG', self.print_debug) - time_var = self.cf_nc.createVariable('time', 'i4', ('time',)) - time_var.long_name = 'time' - time_var.standard_name = 'time' - time_var.units = 'seconds since 1970-01-01 00:00:00 0:00' - time_var.axis = 'T' - - # only add if user adds - if self.comid_lat_lon_z_file and \ - os.path.exists(self.comid_lat_lon_z_file): - log(' lat_var', 'DEBUG', self.print_debug) - lat_var = self.cf_nc.createVariable('lat', 'f8', - (self.output_id_dim_name,), - fill_value=-9999.0) - - log(' lon_var', 'DEBUG', self.print_debug) - lon_var = self.cf_nc.createVariable('lon', 'f8', - (self.output_id_dim_name,), - fill_value=-9999.0) - - add_latlon_metadata(lat_var, lon_var) - - log(' z_var', 'DEBUG', self.print_debug) - z_var = self.cf_nc.createVariable('z', 'f8', - (self.output_id_dim_name,), - fill_value=-9999.0) - z_var.long_name = ('Elevation referenced to the North American ' - 'Vertical Datum of 1988 (NAVD88)') - z_var.standard_name = 'surface_altitude' - z_var.units = 'm' - z_var.axis = 'Z' - z_var.positive = 'up' - - log(' crs_var', 'DEBUG', self.print_debug) - crs_var = self.cf_nc.createVariable('crs', 'i4') - crs_var.grid_mapping_name = 'latitude_longitude' - crs_var.epsg_code = 'EPSG:4326' # WGS 84 - crs_var.semi_major_axis = 6378137.0 - crs_var.inverse_flattening = 298.257223563 - - def _write_comid_lat_lon_z(self): - """Add latitude, longitude, and z values for each netCDF feature - - Remarks: - Lookup table is a CSV file with COMID, Lat, Lon, - and Elev_m columns. - Columns must be in that order and these must be the first - four columns. - """ - # only add if user adds - if self.comid_lat_lon_z_file and \ - os.path.exists(self.comid_lat_lon_z_file): - # get list of COMIDS - lookup_table = csv_to_list(self.comid_lat_lon_z_file) - lookup_comids = np.array([int(float(row[0])) for row in - lookup_table[1:]]) - - # Get relevant arrays while we update them - nc_comids = self.cf_nc.variables[self.output_id_dim_name][:] - lats = self.cf_nc.variables['lat'][:] - lons = self.cf_nc.variables['lon'][:] - zs = self.cf_nc.variables['z'][:] - - min_lat = None - max_lat = None - min_lon = None - max_lon = None - z_min = None - z_max = None - - # Process each row in the lookup table - for nc_index, nc_comid in enumerate(nc_comids): - try: - lookup_index = \ - np.where(lookup_comids == nc_comid)[0][0] + 1 - except IndexError: - log('rivid %s missing in comid_lat_lon_z file' % nc_comid, 'ERROR') - - lat = float(lookup_table[lookup_index][1]) - lats[nc_index] = lat - if min_lat is None or lat < min_lat: - min_lat = lat - if max_lat is None or lat > max_lat: - max_lat = lat - - lon = float(lookup_table[lookup_index][2]) - lons[nc_index] = lon - if min_lon is None or lon < min_lon: - min_lon = lon - if max_lon is None or lon > max_lon: - max_lon = lon - - z = float(lookup_table[lookup_index][3]) - zs[nc_index] = z - if z_min is None or z < z_min: - z_min = z - if z_max is None or z > z_max: - z_max = z - - # Overwrite netCDF variable values - self.cf_nc.variables['lat'][:] = lats - self.cf_nc.variables['lon'][:] = lons - self.cf_nc.variables['z'][:] = zs - - # Update metadata - if min_lat is not None: - self.cf_nc.geospatial_lat_min = min_lat - if max_lat is not None: - self.cf_nc.geospatial_lat_max = max_lat - if min_lon is not None: - self.cf_nc.geospatial_lon_min = min_lon - if max_lon is not None: - self.cf_nc.geospatial_lon_max = max_lon - if z_min is not None: - self.cf_nc.geospatial_vertical_min = z_min - if z_max is not None: - self.cf_nc.geospatial_vertical_max = z_max - else: - log('No comid_lat_lon_z file. Not adding values ...', 'INFO') - - def _generate_time_values(self): - """ - Generates time values for out nc file - """ - # Populate time values - log('writing times', 'INFO') - d1970 = datetime(1970, 1, 1, tzinfo=utc) - time_array = [[int((self.start_datetime - d1970).total_seconds())]] - - datetime_nc_start_simulation = self.start_datetime - for raw_nc_index, raw_nc in enumerate(self.raw_nc_list): - raw_nc_time = raw_nc.get_time_array( - datetime_simulation_start=datetime_nc_start_simulation, - simulation_time_step_seconds=self.time_step_array[ - raw_nc_index]) - - time_array.append(raw_nc_time) - datetime_nc_start_simulation = datetime.utcfromtimestamp(raw_nc_time[-1]) - - self.cf_nc.variables['time'][:] = np.concatenate(time_array) - end_date = datetime.utcfromtimestamp(self.cf_nc.variables['time'][-1]) - self.cf_nc.time_coverage_start = self.start_datetime.isoformat() + 'Z' - self.cf_nc.time_coverage_end = end_date.isoformat() + 'Z' - - def _copy_streamflow_values(self): - """ - Copies streamflow values from raw output to CF file - """ - log('Creating streamflow variable', 'INFO') - q_var = self.cf_nc.createVariable( - self.output_flow_var_name, 'f4', (self.output_id_dim_name, 'time')) - q_var.long_name = 'Discharge' - q_var.units = 'm^3/s' - q_var.coordinates = 'time lat lon z' - q_var.grid_mapping = 'crs' - q_var.source = ('Generated by the Routing Application for Parallel ' - 'computatIon of Discharge (RAPID) river routing ' - 'model.') - q_var.references = 'http://rapid-hub.org/' - q_var.comment = ('lat, lon, and z values taken at midpoint of river ' - 'reach feature') - - log('Copying streamflow values', 'INFO') - master_begin_time_step_index = 1 - # to reduce RAM, copy by chunks - max_2d_dimension = 1000000000 # ~8GB Max - for raw_nc in self.raw_nc_list: - max_time_step_size = min(raw_nc.size_time, - max(1, int(float(max_2d_dimension) / - float(raw_nc.size_river_id)))) - raw_nc_begin_time_step_index = 0 - for raw_nc_time_index in \ - range(0, raw_nc.size_time, max_time_step_size): - time_interval_size = \ - max(1, min(raw_nc.size_time - raw_nc_time_index, - max_time_step_size)) - - raw_nc_end_time_step_index = raw_nc_begin_time_step_index + time_interval_size - master_end_time_step_index = master_begin_time_step_index + time_interval_size - - q_var[:, - master_begin_time_step_index:master_end_time_step_index] \ - = raw_nc.get_qout( - time_index_start=raw_nc_begin_time_step_index, - time_index_end=raw_nc_end_time_step_index) - - master_begin_time_step_index = master_end_time_step_index - raw_nc_begin_time_step_index = raw_nc_end_time_step_index - - log('Adding initial streamflow values', 'INFO') - # add initial flow to RAPID output file - if self.qinit_file and self.rapid_connect_file: - lookup_table = csv_to_list(self.rapid_connect_file) - lookup_comids = np.array([int(float(row[0])) for row - in lookup_table]) - - if self.qinit_file.endswith(".csv"): - init_flow_table = csv_to_list(self.qinit_file) - else: - init_flow_table = netcdf_to_list(self.qinit_file) - - for index, comid in enumerate( - self.cf_nc.variables[self.output_id_dim_name][:]): - try: - lookup_index = np.where(lookup_comids == comid)[0][0] - except IndexError: - log('COMID %s misssing in rapid_connect file' % comid,'ERROR') - q_var[index, 0] = float(init_flow_table[lookup_index][0]) - else: - for index, comid in enumerate( - self.cf_nc.variables[self.output_id_dim_name][:]): - q_var[index, 0] = 0 - - def convert(self): - """ - Copies data from RAPID netCDF output to a CF-compliant netCDF file. - """ - try: - log('Processing %s ...' % self.rapid_output_file_list[0]) - time_start_conversion = datetime.utcnow() - - # Validate the raw netCDF file - log('validating input netCDF file', 'INFO') - id_len, time_len = self._validate_raw_nc() - - # Initialize the output file (create dimensions and variables) - log('initializing output', 'INFO') - self._initialize_output(time_len, id_len) - - self._generate_time_values() - - # copy river ids over - self.cf_nc.variables[self.output_id_dim_name][:] = \ - self.raw_nc_list[0].get_river_id_array() - - # Populate comid, lat, lon, z - log('writing comid lat lon z') - lookup_start = datetime.now() - self._write_comid_lat_lon_z() - duration = str((datetime.now() - lookup_start).total_seconds()) - log('Lookup Duration (s): ' + duration) - - # Create a variable for streamflow. This is big, and slows down - # previous steps if we do it earlier. - self._copy_streamflow_values() - - # close files - for raw_nc in self.raw_nc_list: - raw_nc.close() - self.cf_nc.close() - - # delete original RAPID output - remove_files(*self.rapid_output_file_list) - - # rename nc compliant file to original name - os.rename(self.cf_compliant_file, self.rapid_output_file_list[0]) - log('Time to process %s' % (datetime.utcnow() - time_start_conversion)) - except Exception: - # delete cf RAPID output - remove_files(self.cf_compliant_file) - raise diff --git a/geoglows_ecflow/resources/RAPIDpy/rapid.py b/geoglows_ecflow/resources/RAPIDpy/rapid.py deleted file mode 100644 index e5d01f7..0000000 --- a/geoglows_ecflow/resources/RAPIDpy/rapid.py +++ /dev/null @@ -1,1224 +0,0 @@ -# -*- coding: utf-8 -*- -""" - rapid.py - RAPIDpy - - Created by Alan D Snow, 2015. - License: BSD-3-Clause -""" -import datetime -import os -from calendar import isleap -from csv import writer as csvwriter -from multiprocessing import cpu_count -from subprocess import Popen, PIPE -from time import gmtime - -import netCDF4 as nc -import numpy as np -import xarray -from dateutil.parser import parse -from requests import get - -from .dataset import RAPIDDataset -from .helper_functions import csv_to_list, log, open_csv -from .postprocess_merge import ConvertRAPIDOutputToCF - - -# ----------------------------------------------------------------------------- -# Main RAPID Manager Class -# ----------------------------------------------------------------------------- -class RAPID(object): - """ - This class is designed to prepare the rapid_namelist file and run - the RAPID program. There are also other utilities added. - - Attributes - ---------- - rapid_executable_location: str, optional - Path to the RAPID executable location. - num_processors: int, optional - Number of procesors to use. Default is 1. - Overridden if *use_all_processors* is True. - use_all_processors: bool, optional - If set to True, the RAPID program will use all available processors. - Default is False. - cygwin_bin_location: str, optional - If using Windows, this is the path to the Cygwin 'bin' directory. - mpiexec_command: str, optional - This is the mpi execute commmand. Default is "mpiexec". - ksp_type: str, optional - This is the solver type. Default is "richardson". - **kwargs: str, optional - Keyword arguments matching the input parameters in the RAPID namelist. - - - Linux Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - rapid_executable_location='~/work/rapid/run/rapid' - use_all_processors=True, - ZS_TauR=24 * 3600, - ZS_dtR=15 * 60, - ZS_TauM=365 * 24 * 3600, - ZS_dtM=24 * 3600 - ) - - - Windows with Cygwin Example: - - .. code:: python - - from RAPIDpy import RAPID - - cygwin_exe = 'C:/cygwin64/home/username/work/rapid/run/rapid' - rapid_manager = RAPID( - rapid_executable_location=cygwin_exe, - cygwin_bin_location='C:/cygwin64/bin', - use_all_processors=True, - ZS_TauR=24 * 3600, - ZS_dtR=15 * 60, - ZS_TauM=365 * 24 * 3600, - ZS_dtM=24 * 3600 - ) - - """ - - # pylint: disable=too-many-instance-attributes - def __init__(self, - rapid_executable_location="", - num_processors=1, - use_all_processors=False, - cygwin_bin_location="", - mpiexec_command="mpiexec", - ksp_type="richardson", - **kwargs): - """ - Initialize the class with variables given by the user - """ - if os.name == "nt" and not \ - (cygwin_bin_location or os.path.exists(cygwin_bin_location)) \ - and rapid_executable_location: - raise Exception("Required to have cygwin_bin_location set " - "if using windows!") - - self._rapid_executable_location = rapid_executable_location - self._cygwin_bin_location = cygwin_bin_location - self._cygwin_bash_exe_location = \ - os.path.join(cygwin_bin_location, "bash.exe") - self._mpiexec_command = mpiexec_command - self._ksp_type = ksp_type - - # use all processors makes precedent over num_processors arg - if use_all_processors is True: - self._num_processors = cpu_count() - elif num_processors > cpu_count(): - log("Num processors requested exceeded max. Set to max ...", - "WARNING") - self._num_processors = cpu_count() - else: - self._num_processors = num_processors - - # --------------------------------------------------------------------- - # Runtime options - # --------------------------------------------------------------------- - self.BS_opt_Qinit = False - # !.false. --> no read initial flow .true. --> read initial flow - self.BS_opt_Qfinal = False - # !.false. --> no write final flow .true. --> write final flow - self.BS_opt_dam = False - # !.false. --> no dam model used .true. --> dam model used - self.BS_opt_for = False - # !.false. --> no forcing .true. --> forcing - self.BS_opt_influence = False - # !.false. --> no output influence .true. --> output influence - self.IS_opt_routing = 1 - # !1 --> matrix-based Muskingum 2 --> traditional Muskingum - # !3 --> Transbnd. matrix-based - self.IS_opt_run = 1 - # !1 --> regular run 2 --> parameter optimization - self.IS_opt_phi = 1 - # !1 --> phi1 2 --> phi2 - # --------------------------------------------------------------------- - # Temporal information - # --------------------------------------------------------------------- - # NOTE: ALL TIME IN SECONDS! - # ALWAYS USED - self.ZS_TauR = 0 - # duration of routing procedure (time step of runoff data) - self.ZS_dtR = 0 - # internal routing time step - # ONLY FOR REGULAR RUN - self.ZS_TauM = 0 - # total simulation time - self.ZS_dtM = 0 - # input time step - # ONLY FOR OPTIMIZATION RUN - self.ZS_TauO = 0 - # total optimization time - self.ZS_dtO = 0 - # observation time step - # FORCING MODE (replace some values with observations) - self.ZS_dtF = 0 - # time step of forcing data - # --------------------------------------------------------------------- - # Domain in which input data is available - # --------------------------------------------------------------------- - self.IS_riv_tot = 0 - # number of river reaches in rapid connect file - self.rapid_connect_file = '' - # path to rapid_connect file - self.IS_max_up = 0 - # maximum number of ustream segments - self.Vlat_file = '' - # path to runoff file - # --------------------------------------------------------------------- - # Domain in which model runs - # --------------------------------------------------------------------- - self.IS_riv_bas = 0 - # number of river reaches in subbasin - self.riv_bas_id_file = '' - # subbasin reach id file - # --------------------------------------------------------------------- - # Initial instantaneous flow file - # --------------------------------------------------------------------- - self.Qinit_file = '' - # initial flow file (same order as rapid_connect) - # --------------------------------------------------------------------- - # Final instantaneous flow file - # --------------------------------------------------------------------- - self.Qfinal_file = '' - # path to output final flow file - # --------------------------------------------------------------------- - # Available dam data - # --------------------------------------------------------------------- - self.IS_dam_tot = 0 - # number of dams - self.dam_tot_id_file = '' - # ids of dam location - # --------------------------------------------------------------------- - # Dam data used - # --------------------------------------------------------------------- - self.IS_dam_use = 0 - # number in subset of dam data to use - self.dam_use_id_file = '' - # ids of subset of dams - # --------------------------------------------------------------------- - # Available forcing data - # --------------------------------------------------------------------- - self.IS_for_tot = 0 - self.for_tot_id_file = '' - self.Qfor_file = '' - # --------------------------------------------------------------------- - # Forcing data used as model runs - # --------------------------------------------------------------------- - self.IS_for_use = 0 - self.for_use_id_file = '' - # --------------------------------------------------------------------- - # File where max (min) of absolute values of b (QoutR) are stored - # --------------------------------------------------------------------- - self.babsmax_file = '' - self.QoutRabsmin_file = '' - self.QoutRabsmax_file = '' - # --------------------------------------------------------------------- - # Regular model run - # --------------------------------------------------------------------- - self.k_file = '' - self.x_file = '' - self.Qout_file = '' - # --------------------------------------------------------------------- - # Optimization - # --------------------------------------------------------------------- - self.ZS_phifac = 0 - # --------------------------------------------------------------------- - # Routing parameters - # --------------------------------------------------------------------- - self.kfac_file = '' - self.xfac_file = '' - self.ZS_knorm_init = 0 - self.ZS_xnorm_init = 0 - # --------------------------------------------------------------------- - # Gage observations - # --------------------------------------------------------------------- - self.IS_obs_tot = 0 - self.obs_tot_id_file = '' - self.Qobs_file = '' - self.Qobsbarrec_file = '' - self.IS_obs_use = 0 - self.obs_use_id_file = '' - self.IS_strt_opt = 0 - - self.update_parameters(**kwargs) - - def _get_cygwin_path(self, windows_path): - """ - Convert windows path to cygpath - """ - conv_cmd = [os.path.join(self._cygwin_bin_location, "cygpath.exe"), - "-u", windows_path] - process = Popen(conv_cmd, - stdout=PIPE, stderr=PIPE, shell=False) - out, err = process.communicate() - if err: - print(err) - raise Exception(err) - - return out.strip() - - def _create_symlink_cygwin(self, initial_path, final_path): - """ - Use cygqin to generate symbolic link - """ - symlink_cmd = [os.path.join(self._cygwin_bin_location, "ln.exe"), - "-s", self._get_cygwin_path(initial_path), - self._get_cygwin_path(final_path)] - process = Popen(symlink_cmd, - stdout=PIPE, stderr=PIPE, shell=False) - out, err = process.communicate() - if err: - print(err) - raise Exception(err) - - return out.strip() - - def _dos2unix_cygwin(self, file_path): - """ - Use cygwin to convert file to unix format - """ - dos2unix_cmd = \ - [os.path.join(self._cygwin_bin_location, "dos2unix.exe"), - self._get_cygwin_path(file_path)] - process = Popen(dos2unix_cmd, - stdout=PIPE, stderr=PIPE, shell=False) - process.communicate() - - def update_parameters(self, **kwargs): - """ - You can add or update rapid namelist parameters by using the name of - the variable in the rapid namelist file (this is case sensitive). - - Parameters - ---------- - **kwargs: str, optional - Keyword arguments matching the input parameters - in the RAPID namelist. - - - Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - rapid_executable_location='~/work/rapid/run/rapid' - use_all_processors=True, - ZS_TauR=24 * 3600, - ZS_dtR=15 * 60, - ZS_TauM=365 * 24 * 3600, - ZS_dtM=24 * 3600 - ) - - rapid_manager.update_parameters( - rapid_connect_file='../rapid-io/input/rapid_connect.csv', - Vlat_file='../rapid-io/input/m3_riv.nc', - riv_bas_id_file='../rapid-io/input/riv_bas_id.csv', - k_file='../rapid-io/input/k.csv', - x_file='../rapid-io/input/x.csv', - Qout_file='../rapid-io/output/Qout.nc', - ) - - """ - # set arguments based off of user input - for key, value in list(kwargs.items()): - if key in dir(self) and not key.startswith('_'): - setattr(self, key, value) - else: - log("Invalid RAPID parameter %s." % key, - "ERROR") - - def update_reach_number_data(self): - """ - Update the reach number data for the namelist based on input files. - - .. warning:: You need to make sure you set *rapid_connect_file* - and *riv_bas_id_file* before running this function. - - - Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - rapid_connect_file='../rapid-io/input/rapid_connect.csv', - riv_bas_id_file='../rapid-io/input/riv_bas_id.csv', - ) - - rapid_manager.update_reach_number_data() - - - Example with forcing data: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - rapid_connect_file='../rapid-io/input/rapid_connect.csv', - riv_bas_id_file='../rapid-io/input/riv_bas_id.csv', - Qfor_file='../rapid-io/input/qfor_file.csv', - for_tot_id_file='../rapid-io/input/for_tot_id_file.csv', - for_use_id_file='../rapid-io/input/for_use_id_file.csv', - ZS_dtF=3*60*60, - BS_opt_for=True - ) - - rapid_manager.update_reach_number_data() - - """ - if not self.rapid_connect_file: - log("Missing rapid_connect_file. " - "Please set before running this function ...", - "ERROR") - - if not self.riv_bas_id_file: - log("Missing riv_bas_id_file. " - "Please set before running this function ...", - "ERROR") - - # get rapid connect info - rapid_connect_table = np.loadtxt(self.rapid_connect_file, - ndmin=2, delimiter=",", dtype=int) - - self.IS_riv_tot = int(rapid_connect_table.shape[0]) - self.IS_max_up = int(rapid_connect_table[:, 2].max()) - - # get riv_bas_id info - riv_bas_id_table = np.loadtxt(self.riv_bas_id_file, - ndmin=1, delimiter=",", - usecols=(0,), dtype=int) - self.IS_riv_bas = int(riv_bas_id_table.size) - - # add the forcing files - if not self.for_tot_id_file: - self.IS_for_tot = 0 - log("Missing for_tot_id_file. Skipping ...", - "WARNING") - else: - # get riv_bas_id info - for_tot_id_table = np.loadtxt(self.for_tot_id_file, - ndmin=1, delimiter=",", - usecols=(0,), dtype=int) - self.IS_for_tot = int(for_tot_id_table.size) - - if not self.for_use_id_file: - self.IS_for_use = 0 - log("Missing for_use_id_file. Skipping ...", - "WARNING") - else: - # get riv_bas_id info - for_use_id_table = np.loadtxt(self.for_use_id_file, - ndmin=1, delimiter=",", - usecols=(0,), dtype=int) - self.IS_for_use = int(for_use_id_table.size) - - def update_simulation_runtime(self): - """ - Updates the total simulation duration from - the m3 file (Vlat_file) and the time step (ZS_TauR). - - .. warning:: You need to set the m3 file (Vlat_file) and the - time step (ZS_TauR) before runnning this function. - - - Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - Vlat_file='../rapid-io/input/m3_riv.csv', - ZS_TauR=3*3600, - ) - - rapid_manager.update_simulation_runtime() - """ - if not self.Vlat_file or not os.path.exists(self.Vlat_file): - log("Need Vlat_file to proceed ...", - "ERROR") - - if self.ZS_TauR <= 0: - log("Missing routing time step ...", - "ERROR") - - try: - self.ZS_TauR = int(self.ZS_TauR) - except ValueError: - log("Invalid routing time step: {0} ...".format(self.ZS_TauR), - "ERROR") - - with RAPIDDataset(self.Vlat_file) as m3_nc: - self.ZS_TauM = m3_nc.size_time * self.ZS_TauR - self.ZS_TauO = m3_nc.size_time * self.ZS_TauR - - def generate_namelist_file(self, rapid_namelist_file): - """ - Generate rapid_namelist file. - - Parameters - ---------- - rapid_namelist_file: str - Path of namelist file to generate from - parameters added to the RAPID manager. - """ - log("Generating RAPID namelist file ...", - "INFO") - try: - os.remove(rapid_namelist_file) - except OSError: - pass - - with open(rapid_namelist_file, 'w') as new_file: - new_file.write('&NL_namelist\n') - for attr, value in sorted(list(self.__dict__.items())): - if not attr.startswith('_'): - if attr.startswith('BS'): - new_file.write("{0} = .{1}.\n" - .format(attr, str(value).lower())) - elif isinstance(value, int): - new_file.write("%s = %s\n" % (attr, value)) - else: - if value: - if os.name == "nt": - # if windows generate file with cygpath - value = self._get_cygwin_path(value) - new_file.write("%s = \'%s\'\n" % (attr, value)) - new_file.write("/\n") - - def update_namelist_file(self, rapid_namelist_file, - new_namelist_file=None): - """ - Update existing namelist file with new parameters - - Parameters - ---------- - rapid_namelist_file: str - Path of namelist file to use in the simulation. It will be - updated with any parameters added to the RAPID manager. - new_namelist_file: str, optional - Path to output the updated namelist file. - """ - if os.path.exists(rapid_namelist_file) and rapid_namelist_file: - log("Adding missing inputs from RAPID input file ...", - "INFO") - with open(rapid_namelist_file, 'r') as old_file: - for line in old_file: - line = line.strip() - if not line[:1].isalpha() or not line: - continue - line_split = line.split("=") - attr = line_split[0].strip() - value = None - if len(line_split) > 1: - value = line_split[1].strip() \ - .replace("'", "").replace('"', "") - # convert integers to integers - try: - value = int(value) - except ValueError: - pass - # remove dots from beginning & end of value - if attr.startswith('BS'): - value = value.replace(".", "") - # add attribute if exists - if attr in dir(self) and not attr.startswith('_'): - # set attribute if not set already - if not getattr(self, attr): - setattr(self, attr, value) - else: - log("Invalid argument {0}. Skipping ...".format(attr), - "INFO") - - if new_namelist_file is None: - new_namelist_file = rapid_namelist_file - - self.generate_namelist_file(new_namelist_file) - else: - log("RAPID namelist file to update not found.", - "ERROR") - - def make_output_cf_compliant(self, - simulation_start_datetime, - comid_lat_lon_z_file="", - project_name="Normal RAPID project"): - """ - This function converts the RAPID output to be CF compliant. - This will require a *comid_lat_lon_z.csv* file - (See: :func:`~RAPIDpy.gis.centroid.FlowlineToPoint` to - generate the file). - - .. note:: It prepends time an initial flow to your simulation from the - *qinit_file*. If no qinit file is given, an initial value - of zero is added. - - .. warning:: This will delete your original Qout file. - - Parameters - ---------- - simulation_start_datetime: datetime - Datetime object with the start date of the simulation. - comid_lat_lon_z_file: str, optional - Path to the *comid_lat_lon_z.csv* file. If none given, - spatial information will be skipped. - project_name: str, optional - Name of project to add to the RAPID output file. - - - Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - rapid_executable_location='~/work/rapid/run/rapid' - use_all_processors=True, - ZS_TauR=24*3600, - ZS_dtR=15*60, - ZS_TauM=365*24*3600, - ZS_dtM=24*3600 - rapid_connect_file='../rapid-io/input/rapid_connect.csv', - Vlat_file='../rapid-io/input/m3_riv.nc', - riv_bas_id_file='../rapid-io/input/riv_bas_id.csv', - k_file='../rapid-io/input/k.csv', - x_file='../rapid-io/input/x.csv', - Qout_file='../rapid-io/output/Qout.nc' - ) - - rapid_manager.run() - - rapid_manager.make_output_cf_compliant( - simulation_start_datetime=datetime.datetime(1980, 1, 1), - comid_lat_lon_z_file='../rapid-io/input/comid_lat_lon_z.csv', - project_name="ERA Interim Historical flows by US Army ERDC" - ) - - """ - with RAPIDDataset(self.Qout_file) as qout_nc: - if qout_nc.is_time_variable_valid(): - log("RAPID Qout file already CF compliant ...", - "INFO") - return - - crv = ConvertRAPIDOutputToCF( - rapid_output_file=self.Qout_file, - start_datetime=simulation_start_datetime, - time_step=self.ZS_TauR, - qinit_file=self.Qinit_file, - comid_lat_lon_z_file=comid_lat_lon_z_file, - rapid_connect_file=self.rapid_connect_file, - project_name=project_name, - output_id_dim_name='rivid', - output_flow_var_name='Qout', - print_debug=False - ) - crv.convert() - - def run(self, rapid_namelist_file=""): - """ - Run RAPID program and generate file based on inputs - This will generate your rapid_namelist file and run RAPID from wherever - you call this script (your working directory). - - Parameters - ---------- - rapid_namelist_file: str, optional - Path of namelist file to use in the simulation. - It will be updated with any parameters added to the RAPID manager. - - - Linux Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - rapid_executable_location='~/work/rapid/src/rapid' - use_all_processors=True, - ) - - rapid_manager.update_parameters( - rapid_connect_file='../rapid-io/input/rapid_connect.csv', - Vlat_file='../rapid-io/input/m3_riv.nc', - riv_bas_id_file='../rapid-io/input/riv_bas_id.csv', - k_file='../rapid-io/input/k.csv', - x_file='../rapid-io/input/x.csv', - Qout_file='../rapid-io/output/Qout.nc', - ) - - rapid_manager.update_reach_number_data() - rapid_manager.update_simulation_runtime() - rapid_manager.run( - rapid_namelist_file='../rapid-io/input/rapid_namelist') - - - Linux Reservoir Forcing Flows Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - rapid_executable_location='~/work/rapid/src/rapid', - num_processors=4, - IS_for_tot=4, - IS_for_use=4, - for_tot_id_file='../rapid-io/input/dam_id.csv', - for_use_id_file='../rapid-io/input/dam_id.csv', - Qfor_file='../rapid-io/input/qout_dams.csv', - ZS_dtF=86400, - BS_opt_for=True, - ) - - rapid_manager.run( - rapid_namelist_file='../rapid-io/input/rapid_namelist_regular') - - Windows with Cygwin Example: - - .. code:: python - - from RAPIDpy import RAPID - from os import path - - rapid_exe_path = 'C:/cygwin64/home/username/rapid/run/rapid', - rapid_manager = RAPID( - rapid_executable_location=rapid_exe_path, - cygwin_bin_location='C:/cygwin64/bin', - use_all_processors=True, - ZS_TauR=24*3600, - ZS_dtR=15*60, - ZS_TauM=365*24*3600, - ZS_dtM=24*3600 - ) - - rapid_input = 'C:/cygwin64/home/username/rapid-io/input' - rapid_output = 'C:/cygwin64/home/username/rapid-io/output' - rapid_manager.update_parameters( - rapid_connect_file=path.join(rapid_input, 'rapid_connect.csv'), - Vlat_file=path.join(rapid_input, 'm3_riv.nc'), - riv_bas_id_file=path.join(rapid_input, 'riv_bas_id.csv'), - k_file=path.join(rapid_input, 'k.csv'), - x_file=path.join(rapid_input, 'x.csv'), - Qout_file=path.join(rapid_output, 'Qout.nc'), - ) - - rapid_manager.update_reach_number_data() - rapid_manager.update_simulation_runtime() - rapid_manager.run() - """ - if not self._rapid_executable_location: - log("Missing rapid_executable_location. " - "Please set before running this function ...", - "ERROR") - - time_start = datetime.datetime.utcnow() - temp_rapid_namelist_file = os.path.join(os.getcwd(), "rapid_namelist") - - if not rapid_namelist_file or not os.path.exists(rapid_namelist_file): - # generate input file if it does not exist - self.generate_namelist_file(temp_rapid_namelist_file) - with open(temp_rapid_namelist_file, 'r') as file_: - log(file_.read(), "INFO") - else: - # update existing file - self.update_namelist_file(rapid_namelist_file, - temp_rapid_namelist_file) - - local_rapid_executable_location = \ - os.path.join(os.path.dirname(temp_rapid_namelist_file), - "rapid_exe_symlink") - - def rapid_cleanup(*args): - """ - Cleans up the rapid files generated by the process - """ - for arg in args: - # remove files - try: - os.remove(arg) - except OSError: - pass - - # create link to RAPID if needed - temp_link_to_rapid = "" - # pylint: disable=no-member - if self._rapid_executable_location != \ - local_rapid_executable_location: - rapid_cleanup(local_rapid_executable_location) - if os.name == "nt": - self._create_symlink_cygwin(self._rapid_executable_location, - local_rapid_executable_location) - else: - os.symlink(self._rapid_executable_location, - local_rapid_executable_location) - temp_link_to_rapid = local_rapid_executable_location - - # run RAPID - log("Running RAPID ...", - "INFO") - if os.name == "nt": - local_rapid_executable_location = \ - self._get_cygwin_path(local_rapid_executable_location) - - # htcondor will not allow mpiexec for single processor jobs - # this was added for that purpose - run_rapid_command = [local_rapid_executable_location, - "-ksp_type", self._ksp_type] - - if self._num_processors > 1: - run_rapid_command = [self._mpiexec_command, - "-n", str(self._num_processors)] \ - + run_rapid_command - - process = Popen(run_rapid_command, - stdout=PIPE, stderr=PIPE, shell=False) - out, err = process.communicate() - if err: - rapid_cleanup(temp_link_to_rapid, temp_rapid_namelist_file) - raise Exception(err) - else: - log('RAPID output:', - "INFO") - for line in out.split(b'\n'): - print(line) - rapid_cleanup(temp_link_to_rapid, temp_rapid_namelist_file) - log("Time to run RAPID: %s" % (datetime.datetime.utcnow() - time_start), - "INFO") - - def generate_qinit_from_past_qout(self, qinit_file, time_index=-1, - out_datetime=None): - """ - Generate qinit from a RAPID qout file - - Parameters - ---------- - qinit_file: str - Path to output qinit_file. - time_index: int, optional - Index of simulation to generate initial flow file. - Default is the last index. - out_datetime: :obj:`datetime.datetime`, optional - Datetime object containing time of initialization. - - - Example: - - .. code:: python - - from RAPIDpy import RAPID - - rapid_manager = RAPID( - Qout_file='/output_mississippi-nfie/Qout_k2v1_2005to2009.nc', - rapid_connect_file='/input_mississippi_nfie/rapid_connect.csv' - ) - - rapid_manager.generate_qinit_from_past_qout( - qinit_file='/input_mississippi_nfie/Qinit_2008_flood.csv', - time_index=10162 - ) - - """ - if not self.Qout_file or not os.path.exists(self.Qout_file): - log('Missing Qout_file. ' - 'Please set before running this function ...', - "ERROR") - - if not self.rapid_connect_file or not self.rapid_connect_file: - log('Missing rapid_connect file. ' - 'Please set before running this function ...', - "ERROR") - - log("Generating qinit file from qout file ...", - "INFO") - # get information from dataset - with xarray.open_dataset(self.Qout_file) as qds: - rivid_array = qds.rivid.values - if out_datetime is None: - streamflow_values = qds.isel(time=time_index).Qout.values - else: - streamflow_values = qds.sel(time=str(out_datetime)).Qout.values - - log("Reordering data ...", - "INFO") - - stream_id_array = np.loadtxt(self.rapid_connect_file, - ndmin=1, delimiter=",", - usecols=(0,), dtype=int) - init_flows_array = np.zeros(stream_id_array.size) - for riv_bas_index, riv_bas_id in enumerate(rivid_array): - try: - data_index = np.where(stream_id_array == riv_bas_id)[0][0] - init_flows_array[data_index] = streamflow_values[riv_bas_index] - except IndexError: - log('riv bas id {0} not found in connectivity list.' - .format(riv_bas_id), - "WARNING") - - log("Writing to file ...", - "INFO") - if qinit_file.endswith(".csv"): - with open_csv(qinit_file, 'w') as qinit_out: - for init_flow in init_flows_array: - qinit_out.write('{0}\n'.format(init_flow)) - else: - with nc.Dataset(qinit_file, "w", format="NETCDF3_CLASSIC") as qinit_out: - qinit_out.createDimension('Time', 1) - qinit_out.createDimension('rivid', stream_id_array.size) - var_Qout = qinit_out.createVariable('Qout', 'f8', ('Time', 'rivid',)) - var_Qout[:] = init_flows_array - - self.Qinit_file = qinit_file - self.BS_opt_Qinit = True - log("Initialization Complete!", - "INFO") - - def generate_seasonal_intitialization( - self, - qinit_file, - datetime_start_initialization=datetime.datetime.utcnow() - ): - """This creates a seasonal qinit file from a RAPID qout file. This - requires a simulation Qout file with a longer time period of record and - to be CF compliant. It takes the average of the current date +- 3 days - and goes back as far as possible. - - Parameters - ---------- - qinit_file: str - Path to output qinit_file. - datetime_start_initialization: :obj:`datetime.datetime`, optional - Datetime object with date of simulation to go back through the - years and get a running average to generate streamflow - initialization. Default is utcnow. - - - Example: - - .. code:: python - - from RAPIDpy.rapid import RAPID - - rapid_manager = RAPID( - Qout_file='/output_mississippi-nfie/Qout_2000to2015.nc', - rapid_connect_file='/input_mississippi_nfie/rapid_connect.csv' - ) - - rapid_manager.generate_seasonal_intitialization( - qinit_file='/input_mississippi_nfie/Qinit_seasonal_avg.csv' - ) - """ - if not self.Qout_file or not os.path.exists(self.Qout_file): - log("Missing Qout_file. " - "Please set before running this function ...", - "ERROR") - - if not self.rapid_connect_file or not self.rapid_connect_file: - log("Missing rapid_connect file. " - "Please set before running this function ...", - "ERROR") - - day_of_year = datetime_start_initialization.timetuple().tm_yday - min_day = day_of_year - 3 - max_day = day_of_year + 3 - - with RAPIDDataset(self.Qout_file) as qout_hist_nc: - if not qout_hist_nc.is_time_variable_valid(): - log("File must be CF 1.6 compliant " - "with valid time variable ...", - "ERROR") - - log("Generating seasonal average qinit file from qout file ...", - "INFO") - - log("Determining dates with streamflows of interest ...", - "INFO") - - time_indices = [] - for idx, ttt in enumerate(qout_hist_nc.get_time_array()): - var_time = gmtime(ttt) - compare_yday = var_time.tm_yday - # move day back one past because of leap year adds - # a day after feb 29 (day 60) - if isleap(var_time.tm_year) and compare_yday > 60: - compare_yday -= 1 - # check if date within range of season - if min_day <= compare_yday < max_day: - time_indices.append(idx) - - if not time_indices: - log("No time steps found within range ...", - "ERROR") - - log("Extracting data ...", - "INFO") - - streamflow_array = \ - qout_hist_nc.get_qout(time_index_array=time_indices) - - log("Reordering data...", - "INFO") - stream_id_array = np.loadtxt(self.rapid_connect_file, - ndmin=1, delimiter=",", - usecols=(0,), dtype=int) - init_flows_array = np.zeros(stream_id_array.size) - for riv_bas_index, riv_bas_id in enumerate( - qout_hist_nc.get_river_id_array()): - try: - data_index = np.where(stream_id_array == riv_bas_id)[0][0] - init_flows_array[data_index] = \ - np.mean(streamflow_array[riv_bas_index]) - except IndexError: - log('riv_bas_id {0} not found in connectivity list.' - .format(riv_bas_id), - "WARNING") - - log("Writing to file ...", - "INFO") - if qinit_file.endswith(".csv"): - with open_csv(qinit_file, 'w') as qinit_out: - for init_flow in init_flows_array: - qinit_out.write('{}\n'.format(init_flow)) - else: - with nc.Dataset(qinit_file, "w", format="NETCDF3_CLASSIC") as qinit_out: - qinit_out.createDimension('Time', 1) - qinit_out.createDimension('rivid', stream_id_array.size) - var_Qout = qinit_out.createVariable('Qout', 'f8', ('Time', 'rivid',)) - var_Qout[:] = init_flows_array - - log("Initialization Complete!", - "INFO") - - def generate_usgs_avg_daily_flows_opt(self, - reach_id_gage_id_file, - start_datetime, - end_datetime, - out_streamflow_file, - out_stream_id_file): - """ - Generate daily streamflow file and stream id file required for - calibration or for substituting flows based on USGS gage ids - associated with stream ids. - - Parameters - ---------- - reach_id_gage_id_file: str - Path to reach_id_gage_id file. - start_datetime: datetime - A datetime object with the start date to download data. - end_datetime: datetime - A datetime object with the end date to download data. - out_streamflow_file: str - The path to output the streamflow file for RAPID. - out_stream_id_file: str - The path to output the stream ID file associated with the - streamflow file for RAPID. - - - Example *reach_id_gage_id_file*:: - - COMID, USGS_GAGE_ID - 2000, 503944 - ... - - .. warning:: Overuse will get you blocked from downloading data from - USGS. - - .. warning:: This code does not clean the data in any way. Thus, you - are likely to run into issues if you simply use the raw - data. - - .. warning:: The code skips gages that do not have data - for the entire time period. - - - Simple Example: - - .. code:: python - - import datetime - from os.path import join - from RAPIDpy import RAPID - - main_path = "/home/username/data" - - rapid_manager = RAPID() - rapid_manager.generate_usgs_avg_daily_flows_opt( - reach_id_gage_id_file=join(main_path, "usgsgage_id_comid.csv"), - start_datetime=datetime.datetime(2000,1,1), - end_datetime=datetime.datetime(2014,12,31), - out_streamflow_file=join(main_path,"streamflow_2000_2014.csv"), - out_stream_id_file=join(main_path,"streamid_2000_2014.csv") - ) - - - Complex Example: - - .. code:: python - - import datetime - from os.path import join - from RAPIDpy import RAPID - - main_path = "/home/username/data" - - rapid_manager = RAPID( - rapid_executable_location='~/work/rapid/run/rapid' - use_all_processors=True, - ZS_TauR=24*3600, - ZS_dtR=15*60, - ZS_TauM=365*24*3600, - ZS_dtM=24*3600 - ) - - rapid_manager.update_parameters( - rapid_connect_file='../rapid-io/input/rapid_connect.csv', - Vlat_file='../rapid-io/input/m3_riv.nc', - riv_bas_id_file='../rapid-io/input/riv_bas_id.csv', - k_file='../rapid-io/input/k.csv', - x_file='../rapid-io/input/x.csv', - Qout_file='../rapid-io/output/Qout.nc', - ) - - rapid_manager.update_reach_number_data() - rapid_manager.update_simulation_runtime() - rapid_manager.generate_usgs_avg_daily_flows_opt( - reach_id_gage_id_file=join(main_path, "usgsgage_id_comid.csv"), - start_datetime=datetime.datetime(2000,1,1), - end_datetime=datetime.datetime(2014,12,31), - out_streamflow_file=join(main_path,"streamflow_2000_2014.csv"), - out_stream_id_file=join(main_path,"streamid_2000_2014.csv") - ) - rapid_manager.run() - - """ - log("Generating avg streamflow file and stream id file " - "required for calibration ...", - "INFO") - log("Generating avg streamflow file and stream id file " - "required for calibration ...", - "INFO") - reach_id_gage_id_list = csv_to_list(reach_id_gage_id_file) - gage_data_matrix = [] - valid_comid_list = [] - - # add extra day as it includes the start date - # (e.g. 7-5 is 2 days, but have data for 5,6,7, so +1) - num_days_needed = (end_datetime - start_datetime).days + 1 - - gage_id_list = [] - for row in reach_id_gage_id_list[1:]: - station_id = row[1] - if len(row[1]) == 7: - station_id = '0' + row[1] - gage_id_list.append(station_id) - - num_gage_id_list = np.array(gage_id_list, dtype=np.int32) - log("Querying Server for Data ...", - "INFO") - - query_params = { - 'format': 'json', - 'sites': ",".join(gage_id_list), - 'startDT': start_datetime.strftime("%Y-%m-%d"), - 'endDT': end_datetime.strftime("%Y-%m-%d"), - 'parameterCd': '00060', # streamflow - 'statCd': '00003' # average - } - response = get("http://waterservices.usgs.gov/nwis/dv", - params=query_params) - - if not response.ok: - log("USGS query error ...", - "WARNING") - return - - requested_data = None - try: - requested_data = response.json()['value']['timeSeries'] - except IndexError: - pass - - if requested_data is not None: - for time_series in enumerate(requested_data): - usgs_station_full_name = time_series[1]['name'] - usgs_station_id = usgs_station_full_name.split(":")[1] - gage_data = [] - for time_step in time_series[1]['values'][0]['value']: - local_datetime = parse(time_step['dateTime']) - if local_datetime > end_datetime: - break - - if local_datetime >= start_datetime: - if not time_step['value']: - log("MISSING DATA for USGS Station {0} {1} {2}" - .format(usgs_station_id, - local_datetime, - time_step['value']), - "WARNING") - gage_data.append( - float(time_step['value']) / 35.3146667) - - try: - # get where streamids associated with USGS station ID - streamid_index = \ - np.where(num_gage_id_list == - int(float(usgs_station_id)))[0][0] + 1 - except (IndexError, ValueError): - log("USGS Station {0} not found in list ..." - .format(usgs_station_id), - "WARNING") - raise - - if len(gage_data) == num_days_needed: - gage_data_matrix.append(gage_data) - valid_comid_list.append( - reach_id_gage_id_list[streamid_index][0]) - else: - log("StreamID {0} USGS Station {1} MISSING {2} " - "DATA VALUES".format( - reach_id_gage_id_list[streamid_index][0], - usgs_station_id, - num_days_needed - len(gage_data)), - "WARNING") - - if gage_data_matrix and valid_comid_list: - log("Writing Output ...", - "INFO") - np_array = np.array(gage_data_matrix).transpose() - with open_csv(out_streamflow_file, 'w') as gage_data: - wgd = csvwriter(gage_data) - for row in np_array: - wgd.writerow(row) - - with open_csv(out_stream_id_file, 'w') as comid_data: - wcd = csvwriter(comid_data) - for row in valid_comid_list: - wcd.writerow([int(float(row))]) - - # set parameters for RAPID run - self.IS_obs_tot = len(valid_comid_list) - self.obs_tot_id_file = out_stream_id_file - self.Qobs_file = out_streamflow_file - self.IS_obs_use = len(valid_comid_list) - self.obs_use_id_file = out_stream_id_file - else: - log("No valid data returned ...", - "WARNING") diff --git a/geoglows_ecflow/resources/archive_to_aws.py b/geoglows_ecflow/resources/archive_to_aws.py index 6bc9309..fcaae55 100644 --- a/geoglows_ecflow/resources/archive_to_aws.py +++ b/geoglows_ecflow/resources/archive_to_aws.py @@ -1,18 +1,19 @@ import argparse import glob -import json import os import boto3 import yaml +from geoglows_ecflow.resources.helper_functions import load_forecast_run + def upload_to_s3(workspace: str, aws_config_file: str): """ Uploads GEOGloWS forecast output to AWS. Args: - workspace (str): Path to rapid_run.json base directory. + workspace (str): Path to forecast_run.json base directory. aws_config_file (str): Path to AWS config file. """ with open(aws_config_file, "r") as f: @@ -22,10 +23,9 @@ def upload_to_s3(workspace: str, aws_config_file: str): forecast_bucket_uri = config["bucket_forecast_archive"] mapstyletable_bucket_uri = config["bucket_maptable_archive"] - with open(os.path.join(workspace, "rapid_run.json"), "r") as f: - data = json.load(f) - date = data["date"] - rapid_output_path = data["output_dir"] + data = load_forecast_run(workspace) + date = data["date"] + output_dir = data["output_dir"] # Create an S3 client s3 = boto3.client( @@ -35,7 +35,7 @@ def upload_to_s3(workspace: str, aws_config_file: str): ) for forecast_nc in sorted( - glob.glob(os.path.join(rapid_output_path, f"Qout_*.nc")) + glob.glob(os.path.join(output_dir, f"Qout_*.nc")) ): s3.upload_file( forecast_nc, @@ -60,17 +60,15 @@ def upload_to_s3(workspace: str, aws_config_file: str): argparser = argparse.ArgumentParser() argparser.add_argument( "workspace", - nargs=1, help="Path to suite home directory", ) argparser.add_argument( "aws_config_file", - nargs=1, help="Path to AWS config file", ) args = argparser.parse_args() - workspace = args.workspace[0] - aws_config_file = args.aws_config_file[0] + workspace = args.workspace + aws_config_file = args.aws_config_file upload_to_s3(workspace, aws_config_file) diff --git a/geoglows_ecflow/resources/combine_esri_tables.py b/geoglows_ecflow/resources/combine_esri_tables.py index 632cacd..bb99e74 100644 --- a/geoglows_ecflow/resources/combine_esri_tables.py +++ b/geoglows_ecflow/resources/combine_esri_tables.py @@ -1,19 +1,20 @@ import argparse import logging -import sys import glob import os import pandas as pd +from geoglows_ecflow.resources.helper_functions import configure_logging + def combine_esri_tables(workspace: str): """Combines the map_style_tables from each VPU into 1 CSV file per time step with rows from all VPUs Args: - workspace (str): Path to rapid_run.json base directory. + workspace (str): Path to forecast_run.json base directory. """ # get path to tables from workspace @@ -56,10 +57,6 @@ def combine_esri_tables(workspace: str): args = argparser.parse_args() workspace = args.workspace[0] - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - stream=sys.stdout, - ) + configure_logging() combine_esri_tables(workspace) diff --git a/geoglows_ecflow/resources/compute_init_flows.py b/geoglows_ecflow/resources/compute_init_flows.py index 82c815c..b10cc04 100644 --- a/geoglows_ecflow/resources/compute_init_flows.py +++ b/geoglows_ecflow/resources/compute_init_flows.py @@ -1,91 +1,36 @@ import argparse import os -import json -import netCDF4 as nc import pandas as pd import xarray as xr -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "workspace", - nargs=1, - help="Path to rapid_run.json base directory.", - ) - parser.add_argument( - "vpu", - nargs=1, - help="vpu number to process.", - ) - - args = parser.parse_args() - workspace = args.workspace[0] - vpu = args.vpu[0] - - with open(os.path.join(workspace, "rapid_run.json"), "r") as f: - data = json.load(f) - ymd = data["date"] - - output_file = os.path.join(workspace, 'input', vpu, f'Qinit_{ymd}.nc') +from geoglows_ecflow.resources.helper_functions import load_forecast_run - with xr.open_dataset(os.path.join(workspace, "output", f"nces_avg_{vpu}.nc")) as average_flows: - qinit_values = average_flows.Qout[7, :] - river_ids = average_flows.rivid[:] - init_date = pd.to_datetime(average_flows.time[7].values).strftime("%Y-%m-%d %X") - with nc.Dataset(output_file, "w", format="NETCDF3_CLASSIC") as qinit_nc: - # create dimensions - qinit_nc.createDimension("time", 1) - qinit_nc.createDimension("rivid", river_ids.shape[0]) +# Time index 7 of the ensemble-mean Q corresponds to t+24h on the ENS +# 3h-resolution grid (the HRES member is excluded from the average upstream by +# nco_calc.ecf). The next cycle reads this file as river-route's +# channel_state_init_file via run_river_route_forecast._find_state_init. +INIT_TIME_INDEX = 7 - qout_var = qinit_nc.createVariable("Qout", "f8", ("time", "rivid")) - qout_var[:] = qinit_values - qout_var.long_name = "instantaneous river water discharge downstream of each river reach" - qout_var.units = "m3 s-1" - qout_var.coordinates = "lon lat" - qout_var.grid_mapping = "crs" - qout_var.cell_methods = "time: point" - # rivid - rivid_var = qinit_nc.createVariable("rivid", "i4", ("rivid",)) - rivid_var[:] = river_ids - rivid_var.long_name = "unique identifier for each river reach" - rivid_var.units = "1" - rivid_var.cf_role = "timeseries_id" +def main(workspace: str, vpu: str) -> None: + ymd = load_forecast_run(workspace)["date"] - # time - time_var = qinit_nc.createVariable("time", "i4", ("time",)) - time_var[:] = 0 - time_var.long_name = "time" - time_var.standard_name = "time" - time_var.units = f'seconds since {init_date}' # Must be seconds - time_var.axis = "T" - time_var.calendar = "gregorian" + avg_path = os.path.join(workspace, "output", f"nces_avg_{vpu}.nc") + out_path = os.path.join(workspace, "input", vpu, f"Qinit_{ymd}.parquet") - # longitude - lon_var = qinit_nc.createVariable("lon", "f8", ("rivid",)) - lon_var[:] = 0 - lon_var.long_name = "longitude of a point related to each river reach" - lon_var.standard_name = "longitude" - lon_var.units = "degrees_east" - lon_var.axis = "X" + with xr.open_dataset(avg_path) as ds: + # river-route's channel_state_init_file is a single-column parquet + # whose row order must match river_id ordering in params.parquet. + pd.DataFrame({"Q": ds["Q"].isel(time=INIT_TIME_INDEX).values}).to_parquet( + out_path + ) - # latitude - lat_var = qinit_nc.createVariable("lat", "f8", ("rivid",)) - lat_var[:] = 0 - lat_var.long_name = "latitude of a point related to each river reach" - lat_var.standard_name = "latitude" - lat_var.units = "degrees_north" - lat_var.axis = "Y" - # crs - crs_var = qinit_nc.createVariable("crs", "i4") - crs_var.grid_mapping_name = "latitude_longitude" - crs_var.epsg_code = "EPSG:4326" # WGS 84 - crs_var.semi_major_axis = 6378137.0 - crs_var.inverse_flattening = 298.257223563 - - # add global attributes - qinit_nc.Conventions = "CF-1.6" - qinit_nc.featureType = "timeSeries" +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("workspace", help="Path to forecast_run.json base directory.") + parser.add_argument("vpu", help="VPU number to process.") + args = parser.parse_args() + main(args.workspace, args.vpu) diff --git a/geoglows_ecflow/resources/concat_forecast_warnings.py b/geoglows_ecflow/resources/concat_forecast_warnings.py index c4fa6d6..3070272 100644 --- a/geoglows_ecflow/resources/concat_forecast_warnings.py +++ b/geoglows_ecflow/resources/concat_forecast_warnings.py @@ -1,10 +1,11 @@ import os import glob -import sys import logging import pandas as pd import argparse +from geoglows_ecflow.resources.helper_functions import configure_logging + def concat_warnings(workdir: str) -> None: """ @@ -46,15 +47,11 @@ def concat_warnings(workdir: str) -> None: help="Path to the daily workspace directory, named in YYYYMMDDHH " "format, containing (1) *.runoff.nc IFS forecast files, " "(2) an output directory of routed discharge netcdfs, " - "(3) symlinks to the rapid inputs and return periods directories", + "(3) symlinks to the per-VPU inputs and return periods directories", ) args = parser.parse_args() workspace = args.workspace[0] - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - stream=sys.stdout, - ) + configure_logging() concat_warnings(workspace) diff --git a/geoglows_ecflow/resources/day_one_forecast.py b/geoglows_ecflow/resources/day_one_forecast.py index 70791f4..2ed7cda 100644 --- a/geoglows_ecflow/resources/day_one_forecast.py +++ b/geoglows_ecflow/resources/day_one_forecast.py @@ -2,33 +2,21 @@ import glob import logging import os -import sys import argparse import numpy as np import pandas as pd import xarray as xr import netCDF4 as nc -import dask -from numcodecs import Blosc +from geoglows_ecflow.resources.helper_functions import ( + RETURN_PERIODS, + configure_logging, +) +from geoglows_ecflow.resources.zarr_io import write_dataset_to_zarr -def merge_forecast_qout_files(rapid_output: str, vpu: str | int): - # list the forecast files - prediction_files = sorted( - glob.glob(os.path.join(rapid_output, f"Qout_{vpu}_*.nc")) - ) - - # merge them into a single file joined by ensemble number - ensemble_index_list = [] - qout_datasets = [] - for forecast_nc in prediction_files: - ensemble_index_list.append( - int(os.path.basename(forecast_nc)[:-3].split("_")[-1]) - ) - qout_datasets.append(xr.open_dataset(forecast_nc).Qout) - return xr.concat( - qout_datasets, pd.Index(ensemble_index_list, name="ensemble") - ) +# Only streams of at least this Strahler order are checked against return +# periods for the warnings summary (smaller headwater streams are skipped). +MIN_STREAM_ORDER = 3 def check_for_return_period_flow( @@ -36,59 +24,38 @@ def check_for_return_period_flow( ): max_flow = max(forecasted_flows_df["means"]) - # temporary dates - date_r5 = "" - date_r10 = "" - date_r25 = "" - date_r50 = "" - date_r100 = "" - - # retrieve return period flow levels from dataframe - r2 = float(rp_data["rp2"].values[0]) - r5 = float(rp_data["rp5"].values[0]) - r10 = float(rp_data["rp10"].values[0]) - r25 = float(rp_data["rp25"].values[0]) - r50 = float(rp_data["rp50"].values[0]) - r100 = float(rp_data["rp100"].values[0]) - - # then compare the timeseries to the return period thresholds - if max_flow >= r2: - date_r2 = get_time_of_first_exceedance(forecasted_flows_df, r2) + # retrieve return period flow levels from the dataframe + thresholds = { + rp: float(rp_data[f"rp{rp}"].values[0]) for rp in RETURN_PERIODS + } + # if the flow is not larger than the smallest return period, return the # dataframe without appending anything - else: + if max_flow < thresholds[RETURN_PERIODS[0]]: return largeflows_df - # check the rest of the return period flow levels - if max_flow >= r5: - date_r5 = get_time_of_first_exceedance(forecasted_flows_df, r5) - if max_flow >= r10: - date_r10 = get_time_of_first_exceedance(forecasted_flows_df, r10) - if max_flow >= r25: - date_r25 = get_time_of_first_exceedance(forecasted_flows_df, r25) - if max_flow >= r50: - date_r50 = get_time_of_first_exceedance(forecasted_flows_df, r50) - if max_flow >= r100: - date_r100 = get_time_of_first_exceedance(forecasted_flows_df, r100) - - new_row = pd.DataFrame( - { - "comid": rp_data.index[0], - "stream_order": stream_order, - "max_forecasted_flow": round(max_flow, 2), - "date_exceeds_return_period_2": date_r2, - "date_exceeds_return_period_5": date_r5, - "date_exceeds_return_period_10": date_r10, - "date_exceeds_return_period_25": date_r25, - "date_exceeds_return_period_50": date_r50, - "date_exceeds_return_period_100": date_r100, - }, - index=[0], - ) + # compare the timeseries to each return period threshold, ascending, so the + # progressive masking inside get_time_of_first_exceedance is preserved + exceedance_dates = {} + for rp in RETURN_PERIODS: + if max_flow >= thresholds[rp]: + exceedance_dates[rp] = get_time_of_first_exceedance( + forecasted_flows_df, thresholds[rp] + ) + else: + exceedance_dates[rp] = "" - largeflows_df = pd.concat([largeflows_df, new_row], ignore_index=True) + row = { + "comid": rp_data.index[0], + "stream_order": stream_order, + "max_forecasted_flow": round(max_flow, 2), + } + for rp in RETURN_PERIODS: + row[f"date_exceeds_return_period_{rp}"] = exceedance_dates[rp] - return largeflows_df + new_row = pd.DataFrame(row, index=[0]) + + return pd.concat([largeflows_df, new_row], ignore_index=True) def get_time_of_first_exceedance(forecasted_flows_df, flow): @@ -100,8 +67,8 @@ def get_time_of_first_exceedance(forecasted_flows_df, flow): def postprocess_vpu( vpu, - rapid_input, - rapid_output, + input_dir, + output_dir, return_periods_dir, forecast_records, ): @@ -126,13 +93,13 @@ def postprocess_vpu( logging.info(" merging forecasts") merged_forecasts = xr.open_dataset( - os.path.join(rapid_output, f"nces_avg_{vpu}.nc") + os.path.join(output_dir, f"nces_avg_{vpu}.nc") ) # collect the times and comids from the forecasts logging.info(" reading info from forecasts") times = pd.to_datetime(pd.Series(merged_forecasts.time)) - comids = pd.Series(merged_forecasts.rivid) + comids = pd.Series(merged_forecasts.river_id) tomorrow = times[0] + pd.Timedelta(days=1) year = times[0].strftime("%Y") @@ -145,10 +112,11 @@ def postprocess_vpu( # read the list of large streams logging.info(" creating dataframe of large streams") - streams_file_path = os.path.join(rapid_input, "master_table.parquet") + streams_file_path = os.path.join(input_dir, "master_table.parquet") streams_df = pd.read_parquet(streams_file_path) large_vpu_streams_df = streams_df[ - (streams_df["VPUCode"] == int(vpu)) & ((streams_df["strmOrder"] >= 3)) + (streams_df["VPUCode"] == int(vpu)) + & (streams_df["strmOrder"] >= MIN_STREAM_ORDER) ] # get the list of comids @@ -161,7 +129,7 @@ def postprocess_vpu( # now process the mean flows for each river in the vpu for comid in comids: # compute the timeseries of average flows - means = merged_forecasts.sel(rivid=comid).Qout.values.flatten() + means = merged_forecasts.sel(river_id=comid).Q.values.flatten() # put it in a dataframe with the times series forecasted_flows = ( @@ -191,7 +159,7 @@ def postprocess_vpu( logging.info(" updating the forecast records file") try: update_forecast_records( - vpu, forecast_records, rapid_output, year, first_day_flows, times + vpu, forecast_records, output_dir, year, first_day_flows, times ) except Exception as e: logging.info(" unexpected error updating the forecast records") @@ -208,7 +176,7 @@ def postprocess_vpu( .replace({"": np.nan}) ) largeflows.to_parquet( - os.path.join(rapid_output, f"forecastwarnings_{vpu}.parquet") + os.path.join(output_dir, f"forecastwarnings_{vpu}.parquet") ) return @@ -231,31 +199,25 @@ def update_forecast_records( reference = nc.Dataset(reference) # make a new record file record = nc.Dataset(record_path, "w") - # copy the right dimensions and variables + # copy the right dimensions and variables. lat/lon are deliberately + # not carried into the record: they're dropped during the zarr + # conversion below, and river-route's native output doesn't + # necessarily include them. record.createDimension("time", None) - record.createDimension("rivid", reference.dimensions["rivid"].size) + record.createDimension("river_id", reference.dimensions["river_id"].size) record.createVariable( "time", reference.variables["time"].dtype, dimensions=("time",) ) record.createVariable( - "lat", reference.variables["lat"].dtype, dimensions=("rivid",) - ) - record.createVariable( - "lon", reference.variables["lon"].dtype, dimensions=("rivid",) + "river_id", reference.variables["river_id"].dtype, dimensions=("river_id",) ) record.createVariable( - "rivid", reference.variables["rivid"].dtype, dimensions=("rivid",) - ) - record.createVariable( - "Qout", - reference.variables["Qout"].dtype, - dimensions=("time", "rivid"), + "Q", + reference.variables["Q"].dtype, + dimensions=("time", "river_id"), fill_value=np.nan, ) - # and also prepopulate the lat, lon, and rivid fields - record.variables["rivid"][:] = reference.variables["rivid"][:] - record.variables["lat"][:] = reference.variables["lat"][:] - record.variables["lon"][:] = reference.variables["lon"][:] + record.variables["river_id"][:] = reference.variables["river_id"][:] # set the time variable attributes record.variables["time"].setncattr( @@ -291,7 +253,7 @@ def update_forecast_records( end_time_index = start_time_index + len(first_day_flows[0]) # convert all those saved flows to a np array and write to the netcdf first_day_flows = np.asarray(first_day_flows) - record_netcdf.variables["Qout"][ + record_netcdf.variables["Q"][ start_time_index:end_time_index, : ] = first_day_flows.T @@ -313,48 +275,16 @@ def netcdf_forecast_record_to_zarr(record_path) -> None: zarr_path = record_path.replace(".nc", ".zarr") record_nc = xr.open_dataset(record_path) - with dask.config.set(**{ - 'array.slicing.split_large_chunks': False, - # set the max chunk size to 5MB - 'array.chunk-size': '40MB', - # use the threads scheduler - 'scheduler': 'threads', - # set the maximum memory target usage to 90% of total memory - 'distributed.worker.memory.target': 0.80, - # do not allow spilling to disk - 'distributed.worker.memory.spill': False, - # specify the amount of resources to allocate to dask workers - 'distributed.worker.resources': { - 'memory': 3e9, # 1e9=1GB, this is the amount per worker - 'cpu': os.cpu_count(), # num CPU per worker - } - }): - #set compressing information - logging.info("Configuring compression") - - #if we get rid of dask, we can get rid of the compressor - #the compressor throws an error for version 3 so specify version 2 - compressor = Blosc(cname="zstd", clevel=3, shuffle=Blosc.BITSHUFFLE) - encoding = {'Qout': {"compressor": compressor}} - - logging.info("Writing to zarr") - ( - record_nc - .drop_vars(["lat", "lon"]) - .chunk({ - "time": -1, - "rivid": "auto" - }) - .to_zarr( - zarr_path, - consolidated=True, - encoding=encoding, - mode = 'w', - zarr_version=2 - ) - ) - - record_nc.close() + logging.info("Writing to zarr") + write_dataset_to_zarr( + record_nc, + zarr_path, + {"time": -1, "river_id": "auto"}, + drop_vars=["lat", "lon"], + mode="w", + zarr_version=2, + ) + record_nc.close() logging.info("Done") @@ -362,34 +292,31 @@ def netcdf_forecast_record_to_zarr(record_path) -> None: parser = argparse.ArgumentParser() parser.add_argument( "workspace", - nargs=1, help="path to the daily workspace directory", ) parser.add_argument( "vpu", - nargs=1, help="VPU number", ) parser.add_argument( "output_dir", - nargs=1, help="path to the forecast records output directory", ) args = parser.parse_args() - workspace = args.workspace[0] - vpu = args.vpu[0] - rapid_input = os.path.join(workspace, "input") - rapid_output = os.path.join(workspace, "output") + workspace = args.workspace + vpu = args.vpu + input_dir = os.path.join(workspace, "input") + output_dir = os.path.join(workspace, "output") returnperiods = os.path.join(workspace, "return_periods_dir") - forecast_records = args.output_dir[0] - rapid_output = os.path.join(workspace, "output") + forecast_records = args.output_dir + output_dir = os.path.join(workspace, "output") # start logging - logging.basicConfig(stream=sys.stdout, level=logging.INFO) + configure_logging() postprocess_vpu( - vpu, rapid_input, rapid_output, returnperiods, forecast_records + vpu, input_dir, output_dir, returnperiods, forecast_records ) logging.info("Finished at " + datetime.datetime.now().strftime("%c")) diff --git a/geoglows_ecflow/resources/generate_esri_table.py b/geoglows_ecflow/resources/generate_esri_table.py index ab9d190..c9f628b 100644 --- a/geoglows_ecflow/resources/generate_esri_table.py +++ b/geoglows_ecflow/resources/generate_esri_table.py @@ -1,40 +1,53 @@ import argparse import logging import os -import sys import netCDF4 as nc import pandas as pd import xarray as xr +from geoglows_ecflow.resources.helper_functions import ( + RETURN_PERIODS, + configure_logging, +) + +# Only the first 10 days of the forecast are summarized in the style table. +FORECAST_WINDOW_DAYS = 10 + +# Mean-flow thresholds (m^3/s) that drive the map line-thickness ladder. Flows +# below the first threshold get thickness 1; each threshold crossed bumps the +# thickness by one (levels 2..6). +THICKNESS_THRESHOLDS = [20, 250, 1500, 10000, 30000] + def postprocess_vpu_forecast_directory( - rapid_output: str, + output_dir: str, returnperiods: str, - vpu: int or str, + vpu: int | str, ): # creates file name for the csv file date_string = os.path.basename( - os.path.dirname(rapid_output) + os.path.dirname(output_dir) ) # should be a date in YYYYMMDDHH format style_table_file_name = f"mapstyletable_{vpu}_{date_string}.parquet" - if os.path.exists(os.path.join(rapid_output, style_table_file_name)): + if os.path.exists(os.path.join(output_dir, style_table_file_name)): logging.info(f"Style table already exists: {style_table_file_name}") return logging.info(f"Creating style table: {style_table_file_name}") - nces_output_filename = os.path.join(rapid_output, f"nces_avg_{vpu}.nc") + nces_output_filename = os.path.join(output_dir, f"nces_avg_{vpu}.nc") # read the date and COMID lists from one of the netcdfs with xr.open_dataset(nces_output_filename) as ds: - comids = ds["rivid"][:].values + comids = ds["river_id"][:].values dates = pd.to_datetime(ds["time"][:].values) - mean_flows = ds["Qout"][:].values.round(2) + mean_flows = ds["Q"][:].values.round(2) mean_flow_df = pd.DataFrame(mean_flows, columns=comids, index=dates) # limit both dataframes to the first 10 days mean_flow_df = mean_flow_df[ - mean_flow_df.index <= mean_flow_df.index[0] + pd.Timedelta(days=10) + mean_flow_df.index + <= mean_flow_df.index[0] + pd.Timedelta(days=FORECAST_WINDOW_DAYS) ] # creating pandas dataframe with return periods @@ -43,32 +56,21 @@ def postprocess_vpu_forecast_directory( with nc.Dataset(rp_path, "r") as rp_ncfile: rp_df = pd.DataFrame( { - "return_2": rp_ncfile.variables["rp2"][:], - "return_5": rp_ncfile.variables["rp5"][:], - "return_10": rp_ncfile.variables["rp10"][:], - "return_25": rp_ncfile.variables["rp25"][:], - "return_50": rp_ncfile.variables["rp50"][:], - "return_100": rp_ncfile.variables["rp100"][:], + f"return_{rp}": rp_ncfile.variables[f"rp{rp}"][:] + for rp in RETURN_PERIODS }, - index=rp_ncfile.variables["rivid"][:], + index=rp_ncfile.variables["river_id"][:], ) mean_thickness_df = pd.DataFrame(columns=comids, index=dates, dtype=int) mean_thickness_df[:] = 1 - mean_thickness_df[mean_flow_df >= 20] = 2 - mean_thickness_df[mean_flow_df >= 250] = 3 - mean_thickness_df[mean_flow_df >= 1500] = 4 - mean_thickness_df[mean_flow_df >= 10000] = 5 - mean_thickness_df[mean_flow_df >= 30000] = 6 + for level, threshold in enumerate(THICKNESS_THRESHOLDS, start=2): + mean_thickness_df[mean_flow_df >= threshold] = level mean_ret_per_df = pd.DataFrame(columns=comids, index=dates, dtype=int) mean_ret_per_df[:] = 0 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_2"], axis=1)] = 2 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_5"], axis=1)] = 5 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_10"], axis=1)] = 10 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_25"], axis=1)] = 25 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_50"], axis=1)] = 50 - mean_ret_per_df[mean_flow_df.gt(rp_df["return_100"], axis=1)] = 100 + for rp in RETURN_PERIODS: + mean_ret_per_df[mean_flow_df.gt(rp_df[f"return_{rp}"], axis=1)] = rp mean_flow_df = mean_flow_df.stack().to_frame().rename(columns={0: "mean"}) mean_thickness_df = ( @@ -84,7 +86,7 @@ def postprocess_vpu_forecast_directory( df, left_index=True, right_index=True ) - maptable_outdir = os.path.join(rapid_output, "map_style_tables") + maptable_outdir = os.path.join(output_dir, "map_style_tables") if not os.path.exists(maptable_outdir): os.makedirs(maptable_outdir) @@ -109,22 +111,17 @@ def postprocess_vpu_forecast_directory( help="Path to the daily workspace directory, named in YYYYMMDDHH " "format, containing (1) *.runoff.nc IFS forecast files, " "(2) an output directory of routed discharge netcdfs, " - "(3) symlinks to the rapid inputs and return periods directories", + "(3) symlinks to the per-VPU inputs and return periods directories", ) parser.add_argument("vpu", nargs=1, help="id number of vpu to process") args = parser.parse_args() workspace = args.workspace[0] - rapid_output = os.path.join(workspace, "output") + output_dir = os.path.join(workspace, "output") returnperiods = os.path.join(workspace, "return_periods_dir") vpu = args.vpu[0] - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - stream=sys.stdout, - ) + configure_logging() - params = [rapid_output, returnperiods, vpu] + params = [output_dir, returnperiods, vpu] postprocess_vpu_forecast_directory(*params) diff --git a/geoglows_ecflow/resources/helper_functions.py b/geoglows_ecflow/resources/helper_functions.py index a9841bd..1ebdcab 100644 --- a/geoglows_ecflow/resources/helper_functions.py +++ b/geoglows_ecflow/resources/helper_functions.py @@ -2,11 +2,47 @@ # See: spt_compute (https://github.com/erdc/spt_compute) # Updated by Michael Souffront, 2023 +import json import os import sys import re import logging as log -from glob import glob + +# Return periods (years) used throughout the forecast post-processing. Defined +# once here and imported by the modules that build return-period ladders. +RETURN_PERIODS = [2, 5, 10, 25, 50, 100] + +# The high-resolution (HRES) forecast is always ensemble member 52. Members +# 1-51 are the ensemble perturbations. Defined once here and imported wherever +# the workflow needs to single out (or exclude) the HRES member. +HRES_ENSEMBLE_MEMBER = 52 + +# Shared logging format so every resource module logs identically. +LOG_FORMAT = "%(asctime)s %(levelname)s %(message)s" +LOG_DATEFMT = "%Y-%m-%d %H:%M:%S" + + +def configure_logging(level: str = "INFO") -> None: + """Configure root logging with the shared format, writing to stdout.""" + log.basicConfig( + level=level, + format=LOG_FORMAT, + datefmt=LOG_DATEFMT, + stream=sys.stdout, + ) + + +def load_forecast_run(workspace: str) -> dict: + """Load and parse /forecast_run.json. + + Args: + workspace (str): Directory containing forecast_run.json. + + Returns: + dict: The parsed forecast-run manifest. + """ + with open(os.path.join(workspace, "forecast_run.json"), "r") as f: + return json.load(f) def create_logger( @@ -22,13 +58,13 @@ def create_logger( else: handler = log.StreamHandler(sys.stdout) - handler.setLevel(level) - handler.setFormatter( - log.Formatter("%(asctime)s - %(levelname)s - %(message)s") - ) + handler.setLevel(level) + handler.setFormatter( + log.Formatter("%(asctime)s - %(levelname)s - %(message)s") + ) - # Add the handler to the logger - logger.addHandler(handler) + # Add the handler to the logger + logger.addHandler(handler) return logger @@ -67,7 +103,7 @@ def get_valid_vpucode_list(input_directory: str) -> list[str]: Get a list of vpucodes from the input directory. Args: - input_directory (str): Path to the rapid input directory. + input_directory (str): Path to the per-VPU input directory. Returns: list[str]: List of valid directories (vpucodes). @@ -83,28 +119,6 @@ def get_valid_vpucode_list(input_directory: str) -> list[str]: return valid_input_directories -def find_current_rapid_output( - forecast_directory: str, vpu: str | int -) -> list | None: - """Finds output from RAPID for a specific VPU. - - Args: - forecast_directory (str): Path to forecast directory. - vpu (str | int): VPU code. - - Returns: - list | None: List of paths to RAPID output files or None if not found. - """ - if os.path.exists(forecast_directory): - basin_files = glob( - os.path.join(forecast_directory, f"Qout_{vpu}_*.nc") - ) - if len(basin_files) > 0: - return basin_files - # there are none found - return None - - def get_ensemble_number_from_forecast(forecast_name: str) -> int: """Gets the datetimestep from forecast. diff --git a/geoglows_ecflow/resources/netcdf_to_zarr.py b/geoglows_ecflow/resources/netcdf_to_zarr.py index 78e0bce..934ef8f 100644 --- a/geoglows_ecflow/resources/netcdf_to_zarr.py +++ b/geoglows_ecflow/resources/netcdf_to_zarr.py @@ -1,21 +1,20 @@ import argparse import glob -import json import logging import os import shutil -import sys -import dask import numpy as np import xarray as xr -from numcodecs import Blosc -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(message)s", - stream=sys.stdout, +from geoglows_ecflow.resources.helper_functions import ( + HRES_ENSEMBLE_MEMBER, + configure_logging, + load_forecast_run, ) +from geoglows_ecflow.resources.zarr_io import write_dataset_to_zarr + +configure_logging() def netcdf_forecasts_to_zarr(workspace: str) -> None: @@ -23,71 +22,49 @@ def netcdf_forecasts_to_zarr(workspace: str) -> None: Converts the netcdf forecast files to zarr. Args: - workspace (str): Path to rapid_run.json base directory. + workspace (str): Path to forecast_run.json base directory. """ - with open(os.path.join(workspace, "rapid_run.json"), "r") as f: - data = json.load(f) - rapid_output = data["output_dir"] - date = data["date"] + data = load_forecast_run(workspace) + output_dir = data["output_dir"] + date = data["date"] vpu_nums = sorted( - set([os.path.basename(x).split("_")[1] for x in glob.glob(os.path.join(rapid_output, f"Qout_*_52.nc"))]) + set([os.path.basename(x).split("_")[1] for x in glob.glob(os.path.join(output_dir, f"Qout_*_{HRES_ENSEMBLE_MEMBER}.nc"))]) ) - qout_1_51_files = sorted([os.path.join(rapid_output, f"Qout_{vpu}.nc") for vpu in vpu_nums]) - qout_52_files = sorted(glob.glob(os.path.join(rapid_output, f"Qout_*_52.nc"))) - zarr_file_path = os.path.join(rapid_output, f"Qout_{date}.zarr") + qout_1_51_files = sorted([os.path.join(output_dir, f"Qout_{vpu}.nc") for vpu in vpu_nums]) + qout_52_files = sorted(glob.glob(os.path.join(output_dir, f"Qout_*_{HRES_ENSEMBLE_MEMBER}.nc"))) + zarr_file_path = os.path.join(output_dir, f"Qout_{date}.zarr") if os.path.exists(zarr_file_path): shutil.rmtree(zarr_file_path) - with dask.config.set(**{ - 'array.slicing.split_large_chunks': False, - # set the max chunk size to 5MB - 'array.chunk-size': '40MB', - # use the threads scheduler - 'scheduler': 'threads', - # set the maximum memory target usage to 90% of total memory - 'distributed.worker.memory.target': 0.80, - # do not allow spilling to disk - 'distributed.worker.memory.spill': False, - # specify the amount of resources to allocate to dask workers - 'distributed.worker.resources': { - 'memory': 3e9, # 1e9=1GB, this is the amount per worker - 'cpu': os.cpu_count(), # num CPU per worker - } - }): - logging.info("Opening ensembles 1-51 datasets") - with xr.open_mfdataset(qout_1_51_files, combine="nested", concat_dim="rivid") as ds151: + logging.info("Opening ensembles 1-51 datasets") + with xr.open_mfdataset( + qout_1_51_files, combine="nested", concat_dim="river_id" + ) as ds151: + logging.info("Assigning the ensemble coordinate variable") + ds151 = ds151.assign_coords( + ensemble=np.arange(1, HRES_ENSEMBLE_MEMBER) + ) + logging.info("Opening ensemble 52 dataset") + with xr.open_mfdataset( + qout_52_files, combine="nested", concat_dim="river_id" + ) as ds52: logging.info("Assigning the ensemble coordinate variable") - ds151 = ds151.assign_coords(ensemble=np.arange(1, 52)) - logging.info("Opening ensemble 52 dataset") - with xr.open_mfdataset(qout_52_files, combine="nested", concat_dim="rivid") as ds52: - logging.info("Assigning the ensemble coordinate variable") - ds52 = ds52.assign_coords(ensemble=52) + ds52 = ds52.assign_coords(ensemble=HRES_ENSEMBLE_MEMBER) - logging.info("Concatenating 1-51 and 52 datasets") - ds = xr.concat([ds151, ds52], dim="ensemble") + logging.info("Concatenating 1-51 and 52 datasets") + ds = xr.concat([ds151, ds52], dim="ensemble") - logging.info("Configuring compression") - compressor = Blosc(cname="zstd", clevel=3, shuffle=Blosc.BITSHUFFLE) - encoding = {'Qout': {"compressor": compressor}} - logging.info("Writing to zarr") - ( - ds - .drop_vars(["crs", "lat", "lon", "time_bnds", "Qout_err"]) - .chunk({ - "time": -1, - "rivid": "auto", - "ensemble": -1 - }) - .to_zarr( - zarr_file_path, - consolidated=True, - encoding=encoding, - ) - ) - logging.info("Done") + logging.info("Writing to zarr") + write_dataset_to_zarr( + ds, + zarr_file_path, + {"time": -1, "river_id": "auto", "ensemble": -1}, + drop_vars=["crs", "lat", "lon", "time_bnds"], + ) + logging.info("Done") if __name__ == "__main__": @@ -95,8 +72,7 @@ def netcdf_forecasts_to_zarr(workspace: str) -> None: parser.add_argument( "workspace", - nargs=1, help="Path to the suite home directory.", ) args = parser.parse_args() - netcdf_forecasts_to_zarr(workspace=args.workspace[0]) + netcdf_forecasts_to_zarr(workspace=args.workspace) diff --git a/geoglows_ecflow/resources/prep_rapid_forecast.py b/geoglows_ecflow/resources/prep_rapid_forecast.py deleted file mode 100644 index b0747b2..0000000 --- a/geoglows_ecflow/resources/prep_rapid_forecast.py +++ /dev/null @@ -1,120 +0,0 @@ -import os -import sys -import logging -import json -import argparse -from glob import glob -from geoglows_ecflow.resources.helper_functions import ( - get_valid_vpucode_list, - get_ensemble_number_from_forecast, -) - -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - stream=sys.stdout, -) - - -def rapid_forecast_preprocess( - workspace: str, - rapid_input: str, - rapid_output: str, - runoff_dir: str, - initialize_flows: bool = True, -) -> list[tuple]: - """Creates a dict of jobs to run. - - Args: - workspace (str): Path where the rapid_run.json will be created. - rapid_input (str): Path to the rapid input. - rapid_output (str): Path to the rapid output. - runoff_dir (str): Path to runoff base directory containing ensemble - runoff files. E.g. '/path/to/runoff_dir' - initialize_flows (bool, optional): Whether to initialize flows. - - Returns: - dict[dict]: dict of jobs to run. - """ - # Create master dict - master_dict = { - "input_dir": rapid_input, - "output_dir": rapid_output, - "runoff_dir": runoff_dir, - "date": os.path.basename(runoff_dir), - } - - # Get list of rapid vpu input directories - rapid_vpu_input_dirs = get_valid_vpucode_list(rapid_input) - - # Get list of ensemble runoff files - ensemble_runoff_list = glob(os.path.join(runoff_dir, "*.runoff.*nc")) - - # Make the largest files first - ensemble_runoff_list.sort( - key=lambda x: int(os.path.basename(x).split(".")[0]), reverse=True - ) # key=os.path.getsize - - # submit jobs to downsize ecmwf files to vpu - for vpu in rapid_vpu_input_dirs: - logging.info(f"Adding rapid input directory {vpu}") - - # get vpu-specific input directory - master_vpu_input_dir = os.path.join(rapid_input, vpu) - - # create output directory if not exist - if not os.path.exists(rapid_output): - os.makedirs(rapid_output) - - # create jobs - for runoff in ensemble_runoff_list: - ensemble_number = get_ensemble_number_from_forecast(runoff) - - # get output file names - outflow_file_name = f"Qout_{vpu}_{ensemble_number}.nc" - - # get output full path - master_rapid_outflow_file = os.path.join( - rapid_output, outflow_file_name - ) - - # add job to master dict - master_dict[f"job_{vpu}_{ensemble_number}"] = { - "runoff": runoff, - "vpu": vpu, - "ensemble": ensemble_number, - "input_dir": master_vpu_input_dir, - "output_file": master_rapid_outflow_file, - "init_flows": initialize_flows, - } - - with open(os.path.join(workspace, "rapid_run.json"), "w") as f: - json.dump(master_dict, f) - - logging.info("Completed creating job config json file") - - return master_dict - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "workspace", - nargs=1, - help="Path to workspace directory", - ) - - args = parser.parse_args() - workspace = args.workspace[0] - - rapid_input = os.path.join(workspace, "input") - rapid_output = os.path.join(workspace, "output") - runoff_dir = workspace - - rapid_forecast_preprocess( - workspace=workspace, - rapid_input=rapid_input, - rapid_output=rapid_output, - runoff_dir=runoff_dir, - ) diff --git a/geoglows_ecflow/resources/prep_river_route_forecast.py b/geoglows_ecflow/resources/prep_river_route_forecast.py new file mode 100644 index 0000000..d1e9b6b --- /dev/null +++ b/geoglows_ecflow/resources/prep_river_route_forecast.py @@ -0,0 +1,104 @@ +import argparse +import json +import logging +import os +import sys +from glob import glob + +from geoglows_ecflow.resources.helper_functions import ( + configure_logging, + get_ensemble_number_from_forecast, + get_valid_vpucode_list, +) + + +configure_logging() + + +def forecast_preprocess( + workspace: str, + input_dir: str, + output_dir: str, + runoff_dir: str, + initialize_flows: bool = True, +) -> dict: + """Build the per-(vpu, ens) job manifest that the cycle's ens_member tasks + consume. + + Walks input_dir for VPU subdirs (3-digit names) and runoff_dir for ensemble + runoff netCDFs, then writes /forecast_run.json with one + job__ entry per (VPU, ensemble member) pair. Each entry is read + verbatim by run_river_route_forecast and compute_init_flows. + + Args: + workspace: Where to write forecast_run.json. + input_dir: Per-VPU static-data root (each VPU subdir holds + params.parquet, weights.nc, and Qinit_.parquet files). + output_dir: Where ens_member tasks will write Qout__.nc. + runoff_dir: Directory holding ensemble runoff netCDF files + (`.runoff.nc` etc.). The directory's basename is taken as + the cycle date. + initialize_flows: Whether each job should look for a prior-cycle + Qinit; pass False for cold starts. + + Returns: + The full master dict serialized to forecast_run.json. + """ + os.makedirs(output_dir, exist_ok=True) + + master = { + "input_dir": input_dir, + "output_dir": output_dir, + "runoff_dir": runoff_dir, + "date": os.path.basename(runoff_dir), + } + + vpus = get_valid_vpucode_list(input_dir) + + # Largest ensemble number first so HRES (mem 52) starts early; helps + # spread load across workers since HRES has the longest single-task time. + runoff_files = sorted( + glob(os.path.join(runoff_dir, "*.runoff.*nc")), + key=get_ensemble_number_from_forecast, + reverse=True, + ) + + for vpu in vpus: + logging.info(f"Adding VPU input directory {vpu}") + vpu_input_dir = os.path.join(input_dir, vpu) + for runoff in runoff_files: + ens = get_ensemble_number_from_forecast(runoff) + master[f"job_{vpu}_{ens}"] = { + "runoff": runoff, + "vpu": vpu, + "ensemble": ens, + "input_dir": vpu_input_dir, + "output_file": os.path.join(output_dir, f"Qout_{vpu}_{ens}.nc"), + "init_flows": initialize_flows, + } + + with open(os.path.join(workspace, "forecast_run.json"), "w") as f: + json.dump(master, f) + + logging.info("Wrote forecast_run.json") + return master + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Build the per-(vpu, ens) job manifest for one forecast cycle." + ) + parser.add_argument("workspace", help="Path to workspace directory") + args = parser.parse_args(argv) + + forecast_preprocess( + workspace=args.workspace, + input_dir=os.path.join(args.workspace, "input"), + output_dir=os.path.join(args.workspace, "output"), + runoff_dir=args.workspace, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/geoglows_ecflow/resources/run_rapid_forecast.py b/geoglows_ecflow/resources/run_rapid_forecast.py deleted file mode 100644 index 88e684c..0000000 --- a/geoglows_ecflow/resources/run_rapid_forecast.py +++ /dev/null @@ -1,271 +0,0 @@ -import os -import argparse -import datetime -import json -from glob import glob -from .RAPIDpy.rapid import RAPID -from geoglows_ecflow.resources.helper_functions import ( - create_logger, - get_ensemble_number_from_forecast, - case_insensitive_file_search, -) -from shutil import move -from basininflow.inflow import create_inflow_file - - -def rapid_forecast_exec( - workspace: str, - job_id: str, - rapid_executable_location: str, - mp_execute_directory: str, - subprocess_forecast_log_dir: str, -) -> None: - """Runs GEOGloWS RAPID forecast. - - Args: - workspace (str): Path to rapid_run.json. - job_id (str): Job ID. - rapid_executable_location (str): Path to RAPID executable. - mp_execute_directory (str): Path intermediate directory for RAPID. - subprocess_forecast_log_dir (str): Path to RAPID log directory. - """ - if not os.path.exists(mp_execute_directory): - os.mkdir(mp_execute_directory) - if not os.path.exists(subprocess_forecast_log_dir): - os.mkdir(subprocess_forecast_log_dir) - - with open(os.path.join(workspace, "rapid_run.json"), "r") as f: - data = json.load(f) - date = data["date"] - - # Get job - job = data.get(job_id, {}) - - # Check if job is empty - if not job: - raise ValueError(f"Job {job_id} not found.") - - runoff = job["runoff"] - vpucode = job["vpu"] - rapid_vpu_input_dir = job["input_dir"] - master_rapid_outflow_file = job["output_file"] - initialize_flows = job["init_flows"] - - rapid_logger = create_logger( - "rapid_run_logger", - "DEBUG", - os.path.join(subprocess_forecast_log_dir, f"{job_id}.log"), - ) - - rapid_logger.info(f"Preparing {job_id} directories.") - execute_directory = os.path.join(mp_execute_directory, job_id) - output_base_dir = os.path.dirname(master_rapid_outflow_file) - try: - if not os.path.exists(execute_directory): - os.mkdir(execute_directory) - except OSError as e: - raise OSError(f"Failed to create {execute_directory}: {e}") - - try: - if not os.path.exists(output_base_dir): - os.makedirs(output_base_dir) - except OSError as e: - raise OSError(f"Failed to create {output_base_dir}: {e}") - - time_start_all = datetime.datetime.utcnow() - rapid_logger.info(f"Creating inflow files for {job_id}.") - - os.chdir(execute_directory) - ens_number = get_ensemble_number_from_forecast(runoff) - - # prepare ECMWF file for RAPID - rapid_logger.info( - f"Running RAPID downscaling for vpu: {vpucode}, " - f"ensemble: {ens_number}." - ) - - # set up RAPID manager - try: - rapid_connect_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"rapid_connect\.csv" - ) - riv_bas_id_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"riv_bas_id.*?\.csv" - ) - comid_lat_lon_z_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"comid_lat_lon_z.*?\.csv" - ) - weight_table = case_insensitive_file_search( - rapid_vpu_input_dir, r"weight_ifs_48r1.*?\.csv" - ) - k_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"k\.csv" - ) - x_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"x\.csv" - ) - except Exception as e: - rapid_logger.critical(f"input file not found: {e}") - raise - - rapid_manager = RAPID( - rapid_executable_location=rapid_executable_location, - rapid_connect_file=rapid_connect_file, - riv_bas_id_file=riv_bas_id_file, - k_file=k_file, - x_file=x_file, - ZS_dtM=3 * 60 * 60, # Assume 3hr time step - ) - - # check for forcing flows - try: - Qfor_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"qfor\.csv" - ) - for_tot_id_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"for_tot_id\.csv" - ) - for_use_id_file = case_insensitive_file_search( - rapid_vpu_input_dir, r"for_use_id\.csv" - ) - - rapid_manager.update_parameters( - Qfor_file=Qfor_file, - for_tot_id_file=for_tot_id_file, - for_use_id_file=for_use_id_file, - ZS_dtF=3 * 60 * 60, # forcing time interval - BS_opt_for=True, - ) - except Exception: - rapid_logger.info("Forcing files not found. Skipping forcing ...") - pass - - rapid_manager.update_reach_number_data() - - outflow_file_name = os.path.join( - execute_directory, f"Qout_{vpucode}_{ens_number}.nc" - ) - - # Get qinit file - qinit_file = "" - BS_opt_Qinit = False - if initialize_flows: - # Look for qinit files for the past 3 days; - # Try seasonal average file if not - for day in [24, 48, 72]: - past_date = ( - datetime.datetime.strptime(date, "%Y%m%d%H") - - datetime.timedelta(hours=day) - ).strftime("%Y%m%d%H") - qinit_file = os.path.join( - rapid_vpu_input_dir, f"Qinit_{past_date}.nc" - ) - BS_opt_Qinit = qinit_file and os.path.exists(qinit_file) - if BS_opt_Qinit: - break - - if not BS_opt_Qinit: - print( - "Qinit file not found. " - "Trying to initialize from Seasonal Averages ..." - ) - try: - qinit_file = glob( - os.path.join(rapid_vpu_input_dir, "seasonal_qinit*.nc") - )[0] - BS_opt_Qinit = qinit_file and os.path.exists(qinit_file) - except Exception: - print("Failed to initialize from Seasonal Averages.") - print( - f"WARNING: {qinit_file} not found. " - "Not initializing ..." - ) - qinit_file = "" - - # Create inflow directory - inflow_dir = os.path.join(workspace, "inflows") - if not os.path.exists(inflow_dir): - os.mkdir(inflow_dir) - - # Create inflow - create_inflow_file( - lsm_data=runoff, - input_dir=rapid_vpu_input_dir, - inflow_dir=inflow_dir, - weight_table=weight_table, - comid_lat_lon_z=comid_lat_lon_z_file, - cumulative=True, - file_label=ens_number, - ) - - # Get forecast chronometry - interval = 3 if ens_number < 52 else 1 - duration = 360 if ens_number < 52 else 240 - - # Get inflow file path - inflow_file_path = case_insensitive_file_search( - inflow_dir, rf"m3_{vpucode}.*_{ens_number}\.nc" - ) - - try: - rapid_manager.update_parameters( - ZS_TauR=interval * 60 * 60, - ZS_dtR=15 * 60, - ZS_TauM=duration * 60 * 60, - ZS_dtM=interval * 60 * 60, - ZS_dtF=interval * 60 * 60, - Vlat_file=inflow_file_path, - Qout_file=outflow_file_name, - Qinit_file=qinit_file, - BS_opt_Qinit=BS_opt_Qinit, - ) - - # run RAPID - rapid_manager.run() - except Exception as e: - rapid_logger.critical(f"Failed to run RAPID: {e}.") - raise - - time_stop_all = datetime.datetime.utcnow() - delta_time = time_stop_all - time_start_all - rapid_logger.info(f"Total time to compute: {delta_time}") - - node_rapid_outflow_file = os.path.join( - execute_directory, os.path.basename(master_rapid_outflow_file) - ) - - move(node_rapid_outflow_file, master_rapid_outflow_file) - - -if __name__ == "__main__": - argparser = argparse.ArgumentParser() - argparser.add_argument( - "workspace", - nargs=1, - help="Path to suite home directory", - ) - argparser.add_argument( - "job_id", - nargs=1, - help="Job ID", - ) - argparser.add_argument( - "rapid_executable_location", - nargs=1, - help="Path to RAPID executable", - ) - - args = argparser.parse_args() - workspace = args.workspace[0] - job_id = args.job_id[0] - rapid_executable_location = args.rapid_executable_location[0] - mp_execute_directory = os.path.join(workspace, "execute") - subprocess_forecast_log_dir = os.path.join(workspace, "subprocess") - - rapid_forecast_exec( - workspace, - job_id, - rapid_executable_location, - mp_execute_directory, - subprocess_forecast_log_dir, - ) diff --git a/geoglows_ecflow/resources/run_river_route_forecast.py b/geoglows_ecflow/resources/run_river_route_forecast.py new file mode 100644 index 0000000..6c8424a --- /dev/null +++ b/geoglows_ecflow/resources/run_river_route_forecast.py @@ -0,0 +1,90 @@ +import argparse +import datetime +import os +import sys +from glob import glob + +import river_route as rr + +from geoglows_ecflow.resources.helper_functions import ( + create_logger, + get_ensemble_number_from_forecast, + load_forecast_run, +) + + +def _find_state_init(vpu_input_dir: str, date: str) -> str | None: + """Find prior-cycle Qinit at 24/48/72h lookback, then seasonal fallback.""" + # Forecasts run at 00 and 12 UTC; the 24/48/72h lookback tolerates one or + # two missed cycles before falling through to a seasonal climatology. + # Qinit_.parquet is written by the previous cycle's + # compute_init_flows step (ensemble mean). + base = datetime.datetime.strptime(date, "%Y%m%d%H") + for hrs in (24, 48, 72): + past = (base - datetime.timedelta(hours=hrs)).strftime("%Y%m%d%H") + cand = os.path.join(vpu_input_dir, f"Qinit_{past}.parquet") + if os.path.exists(cand): + return cand + seasonal = sorted(glob(os.path.join(vpu_input_dir, "seasonal_qinit*.parquet"))) + return seasonal[0] if seasonal else None + + +def river_route_forecast_exec(workspace: str, job_id: str, log_dir: str) -> None: + data = load_forecast_run(workspace) + + job = data.get(job_id) + if not job: + raise ValueError(f"Job '{job_id}' not found in forecast_run.json") + + date = data["date"] + runoff = job["runoff"] + vpu = job["vpu"] + vpu_input_dir = job["input_dir"] + discharge_file = job["output_file"] + initialize_flows = job["init_flows"] + + logger = create_logger( + "river_route_logger", "INFO", os.path.join(log_dir, f"{job_id}.log") + ) + + ens = get_ensemble_number_from_forecast(runoff) + state_init = _find_state_init(vpu_input_dir, date) if initialize_flows else None + state_final = os.path.join(vpu_input_dir, f"Qfinal_{date}_{ens}.parquet") + + logger.info(f"Routing vpu={vpu} ens={ens} runoff={runoff}") + logger.info(f" state_init={state_init or '(none)'}") + logger.info(f" state_final={state_final}") + logger.info(f" discharge={discharge_file}") + + rr.RapidMuskingum( + params_file=os.path.join(vpu_input_dir, "params.parquet"), + grid_runoff_files=[runoff], + grid_weights_file=os.path.join(vpu_input_dir, "weights.nc"), + discharge_files=[discharge_file], + channel_state_init_file=state_init, + channel_state_final_file=state_final, + # ECMWF runoff is accumulated since forecast start; the default + # ("incremental") would silently produce wrong day-1 values. + grid_accumulation_type="cumulative", + runoff_processing_mode="ensemble", + var_x="lon", + var_y="lat", + progress_bar=False, + log_level="INFO", + ).route() + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Run river-route forecast for one (vpu, ens) job.") + parser.add_argument("workspace", help="Path to directory containing forecast_run.json") + parser.add_argument("job_id", help="Job key inside forecast_run.json (e.g. job__)") + args = parser.parse_args(argv) + + log_dir = os.path.join(args.workspace, "subprocess") + os.makedirs(log_dir, exist_ok=True) + river_route_forecast_exec(args.workspace, args.job_id, log_dir) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/geoglows_ecflow/resources/zarr_io.py b/geoglows_ecflow/resources/zarr_io.py new file mode 100644 index 0000000..c86b1af --- /dev/null +++ b/geoglows_ecflow/resources/zarr_io.py @@ -0,0 +1,67 @@ +"""Shared helpers for writing xarray datasets to zarr. + +The dask configuration and Blosc/zstd compression settings were previously +copy-pasted into both ``netcdf_to_zarr`` and ``day_one_forecast``. They live +here so the heavy ``dask``/``numcodecs`` imports stay out of the lightweight +``helper_functions`` module. +""" + +import os + +import dask +from numcodecs import Blosc + +# Dask settings shared by every zarr write in the workflow. +DASK_ZARR_CONFIG = { + "array.slicing.split_large_chunks": False, + # cap the max chunk size + "array.chunk-size": "40MB", + # use the threads scheduler + "scheduler": "threads", + # set the maximum memory target usage to 80% of total memory + "distributed.worker.memory.target": 0.80, + # do not allow spilling to disk + "distributed.worker.memory.spill": False, + # resources to allocate to dask workers + "distributed.worker.resources": { + "memory": 3e9, # 1e9=1GB, per worker + "cpu": os.cpu_count(), # num CPU per worker + }, +} + + +def write_dataset_to_zarr( + ds, + zarr_path: str, + chunks: dict, + drop_vars=(), + **to_zarr_kwargs, +) -> None: + """Write an open xarray dataset to a consolidated, Blosc-compressed zarr. + + The caller is responsible for opening (and closing) ``ds``; this helper + only applies the shared dask config, drops the requested vars, rechunks, + and writes. Per-call differences (chunk dims, dropped vars, ``mode``, + ``zarr_version``) are passed in rather than hardcoded. + + Args: + ds: An open xarray Dataset. + zarr_path (str): Destination zarr store path. + chunks (dict): Chunk spec passed to ``ds.chunk``. + drop_vars: Variable names to drop (ignored if absent). + **to_zarr_kwargs: Extra keyword args forwarded to ``to_zarr`` + (e.g. ``mode``, ``zarr_version``). + """ + with dask.config.set(**DASK_ZARR_CONFIG): + compressor = Blosc(cname="zstd", clevel=3, shuffle=Blosc.BITSHUFFLE) + encoding = {"Q": {"compressor": compressor}} + ( + ds.drop_vars(list(drop_vars), errors="ignore") + .chunk(chunks) + .to_zarr( + zarr_path, + consolidated=True, + encoding=encoding, + **to_zarr_kwargs, + ) + ) diff --git a/geoglows_ecflow/workflow/builders/builder.py b/geoglows_ecflow/workflow/builders/builder.py index 608c3f7..630f6f8 100644 --- a/geoglows_ecflow/workflow/builders/builder.py +++ b/geoglows_ecflow/workflow/builders/builder.py @@ -1,24 +1,35 @@ from geoglows_ecflow.workflow.builders.base import GEOGLOWSBaseBuilder from geoglows_ecflow.workflow.comfies.ooflow import Trigger, Defuser from geoglows_ecflow.workflow.comfies.ooflow import all_complete, Event, complete -from geoglows_ecflow.workflow.comfies.ooflow import complete, Limit, InLimit, Variable +from geoglows_ecflow.workflow.comfies.ooflow import Limit, InLimit, Variable from geoglows_ecflow.workflow.comfies.ooflow import RepeatDate, Defstatus -from geoglows_ecflow.workflow.parts.nodes import Family, Task +from geoglows_ecflow.workflow.parts.nodes import Family, Task, NominalTime from geoglows_ecflow.workflow.parts.times import ( t2t, - Timer, CronDateRefresh, CronDataAvailability, ) -from geoglows_ecflow.workflow.comfies.dateandtime import Date, TimeDelta, CalSeq +from geoglows_ecflow.workflow.comfies.dateandtime import Date, CalSeq from geoglows_ecflow.workflow.parts.admin import AdminFamily from geoglows_ecflow.workflow.parts.epilogs import DummyEpilog from geoglows_ecflow.workflow.parts.repeats import calseq_repeat from geoglows_ecflow.workflow.parts.packages import PackageInstallers from geoglows_ecflow.workflow.comfies.partition import partition +from geoglows_ecflow.resources.helper_functions import HRES_ENSEMBLE_MEMBER +# Scheduler memory reservations (MB) for the heavier tasks. +ENS_TASK_MEM_MB = 6000 +ARCHIVE_QINIT_MEM_MB = 4000 -from geoglows_ecflow.workflow.parts.nodes import Family, Task, NominalTime + +def is_00z_cycle(nominal_time): + """Whether a nominal-time family is the 00Z cycle. + + The 00Z cycle runs the full ensemble pipeline; the 12Z cycle runs a + reduced tree. The cycle is identified by the EMOS_BASE variable the + NominalTime family carries ("00" or "12"). + """ + return nominal_time.get_variable("EMOS_BASE").value() != "12" class Builder(GEOGLOWSBaseBuilder): @@ -28,47 +39,41 @@ class Builder(GEOGLOWSBaseBuilder): ecflow_module = "geoglows_ecflow.workflow.parts.nodes" scripts = [ - "geoglows_ecflow/workflow/scripts/rapid", + "geoglows_ecflow/workflow/scripts/routing", "geoglows_ecflow/workflow/scripts/common", ] includes = [ - "geoglows_ecflow/workflow/scripts/rapid", + "geoglows_ecflow/workflow/scripts/routing", "geoglows_ecflow/workflow/scripts/common", ] def build(self): """ - Create parts and wire them together into an GLOFAS suite. + Create parts and wire them together into a GEOGloWS suite. Naming conventions: n_* -- ecFlow Node object e_* -- ecFlow Event object """ super(Builder, self).build() cfg = self.config - # get suite configuration parameters from the deployment config file - suite_name = self.config.get("name") - mode = self.config.get("mode", choices=["prod", "test", "rd"]) - first_date = self.config.get("first_date", type=int) - last_date = self.config.get("last_date", type=int, default="20300101") - first_barrier = self.config.get( - "first_barrier", type=int, default=first_date - ) - archive_path = self.config.get("exparch") - suite_dir = self.config.get("workroot") - - with_flood_hazard = self.config.get("with_floodhazard", default=False) - wb_days = self.config.get("wb_days", type=int, default=10) - - # initially empty suite, provided by parent - # class will be filled up with content here. - mars_nworkers = self.config.get("mars_workers", type=int, default=1) - ens_members = self.config.get("ens_members", type=int, default=51) - ens_range = self.config.get("ens_range", type=int, default=30) - suite = self.suite - par_jobvars = self.jobvars.dest("parallel", fallback="PARENT") - # Selectable Trigger suites + # All tunable parameters read from the deployment config file are + # gathered here so what the suite exposes is visible at a glance. + # (exparch/workroot are consumed by the task scripts via templating, + # so they are intentionally not read here.) + suite_name = cfg.get("name") + # Validate the run mode (consumed by suite.h via templating); the + # builder no longer branches on it, so the return is discarded. + cfg.get("mode", choices=["prod", "test"]) + first_date = cfg.get("first_date", type=int) + last_date = cfg.get("last_date", type=int, default="20300101") + first_barrier = cfg.get("first_barrier", type=int, default=first_date) + mars_nworkers = cfg.get("mars_workers", type=int, default=1) + ens_members = cfg.get("ens_members", type=int, default=51) + vpu_list = cfg.get("vpu_list", type=list, default=[]) + + # Operational suites this suite triggers off (normalized to lead "/"). o_suite = cfg.get("o_suite", default="/o") mc_suite = cfg.get("mc_suite", default="/mc") if o_suite[0] != "/": @@ -76,8 +81,7 @@ def build(self): if mc_suite[0] != "/": mc_suite = f"/{mc_suite}" - # these flags are not user-configurable but - # depend on other flags + suite = self.suite # admin family n_admin = Family("admin") @@ -99,17 +103,6 @@ def build(self): with_webpush = False with_diss = False - follow_osuite = False - in_production = False - in_test = False - - if mode == "prod": - in_production = True - if mode == "test": - in_test = True - - if in_production or in_test: - follow_osuite = True # make family n_make = Family("make") @@ -117,13 +110,6 @@ def build(self): packages=["scripts"] ) - n_build_petsc = Task("build_petsc") - if "cc" in self.config.get( - "jobs.destinations.default.host", default="lxc" - ): - n_build_petsc.add_defstatus(complete) - n_build_rapid = Task("build_rapid") - n_build_rapid.trigger = n_build_petsc.complete n_build_venv = Task("build_venv") n_packages.trigger = n_build_venv.complete n_statics = Task("install_static_data") @@ -138,8 +124,6 @@ def build(self): n_make.add( Variable("SMSTRIES", 1), n_build_venv, - n_build_petsc, - n_build_rapid, n_packages, n_statics, n_initialize, @@ -147,8 +131,6 @@ def build(self): n_make.add_inlimit("make") - vpu_list = self.config.get("vpu_list", type=list, default=[]) - n_make.add(Variable("YMD", first_date)) suite.add(n_make, n_admin) @@ -189,47 +171,38 @@ def build(self): (barrier_00, main_00, lag_00), ): cycle = str(main_hh.time.hh) - tnom = main_hh.time - - if follow_osuite: - self.defs.add_extern(f"{mc_suite}/main:YMD") - self.defs.add_extern(f"{mc_suite}/main/{cycle}/fc0015d/fc") - self.defs.add_extern(f"{o_suite}/main:YMD") - self.defs.add_extern(f"{o_suite}/main/{cycle}/fc/model") - n_run_hr = Family("run_hr").add( - Trigger( - f"({o_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " - f"and {o_suite}/main/{cycle}/fc/model == complete) " - f"or ({o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" - ) - ) - n_run_hr.add( - Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) - ) - n_run_en = Family("run_en").add( - Trigger( - f"({mc_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " - f"and {mc_suite}/main/{cycle}/fc0015d/fc == complete) " - f"or ({mc_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" - ) - ) - n_run_en.add( - Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) - ) - n_barrier_epilog = Family("last").add( - Trigger( - f"{o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD" - ), - Task("sleep").add(Trigger("0==1"), Defuser("1==1")), + self.defs.add_extern(f"{mc_suite}/main:YMD") + self.defs.add_extern(f"{mc_suite}/main/{cycle}/fc0015d/fc") + self.defs.add_extern(f"{o_suite}/main:YMD") + self.defs.add_extern(f"{o_suite}/main/{cycle}/fc/model") + n_run_hr = Family("run_hr").add( + Trigger( + f"({o_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " + f"and {o_suite}/main/{cycle}/fc/model == complete) " + f"or ({o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" + ) + ) + n_run_hr.add( + Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) + ) + n_run_en = Family("run_en").add( + Trigger( + f"({mc_suite}/main:YMD == /{suite_name}/barrier/daily:YMD " + f"and {mc_suite}/main/{cycle}/fc0015d/fc == complete) " + f"or ({mc_suite}/main:YMD > /{suite_name}/barrier/daily:YMD)" ) + ) + n_run_en.add( + Task("dummy").add(Trigger("0==1")).add(Defuser("1==1")) + ) - else: - n_run_hr = Family("run_hr") - n_run_en = Family("run_en") - n_run_hr.add(Task("dummy"), Timer(tnom + TimeDelta(hours=7))) - n_run_en.add(Task("dimmy"), Timer(tnom + TimeDelta(hours=9))) - n_barrier_epilog = DummyEpilog(done=Timer("14:15")) + n_barrier_epilog = Family("last").add( + Trigger( + f"{o_suite}/main:YMD > /{suite_name}/barrier/daily:YMD" + ), + Task("sleep").add(Trigger("0==1"), Defuser("1==1")), + ) barrier_hh.add(n_run_hr, n_run_en) n_barrier_daily.add(barrier_hh) @@ -245,8 +218,7 @@ def build(self): n_hr.add(n_ret_hr) n_ens = Family("ens") - if follow_osuite: - n_ens.trigger = n_run_en.complete.across("YMD") + n_ens.trigger = n_run_en.complete.across("YMD") n_ens.trigger &= n_initialize.complete n_ens.add(Variable("CONTEXT", "ens")) n_ret_ens = Family("retrieve") @@ -267,14 +239,14 @@ def build(self): n_prep_ens.trigger &= n_ret_hr.complete n_ens_ens = Family("ens_members") n_ens_ens.trigger = n_prep_ens.complete - n_ens_ens.add_variable("MEM", 6000) + n_ens_ens.add_variable("MEM", ENS_TASK_MEM_MB) n_ens.add(n_ret_ens) - if main_hh.get_variable("EMOS_BASE").value() != "12": + if is_00z_cycle(main_hh): n_ens.add(n_prep_ens, n_ens_ens) for vpu in vpu_list: # Create the ensemble tasks - for mem in reversed(range(1, 53)): + for mem in reversed(range(1, HRES_ENSEMBLE_MEMBER + 1)): n_member = Family(f"{vpu}_{mem:02d}").add( Task("ens_member"), Variable("JOB_ID", f"job_{vpu}_{mem}"), @@ -309,7 +281,7 @@ def build(self): n_forecast_warnings.trigger = n_vpus.complete n_archive_qinit = Task("archive_qinit") - n_archive_qinit.add_variable('MEM', 4000) + n_archive_qinit.add_variable('MEM', ARCHIVE_QINIT_MEM_MB) n_archive_qinit.add_variable('NCPUS', 12) n_archive_qinit.trigger = n_vpus.complete @@ -326,15 +298,6 @@ def build(self): n_diss_ip.trigger = n_ret_ens.complete & n_ret_hr.complete n_diss.add(n_diss_ip) - n_diss_fc = Family("diss_fc") - n_diss_fc.add(Variable("CONTEXT", "rapid")) - n_diss_fc.defuser = e_no_diss - n_diss_fc.add(Task("diss")) - n_diss_fc.trigger = n_nc_to_zarr.complete & n_plain_table.complete & n_forecast_warnings.complete - - if main_hh.get_variable("EMOS_BASE").value() != "12": - n_diss.add(n_diss_fc) - n_web = Family("web_push") n_web.trigger = n_nc_to_zarr.complete & n_vpus.complete n_web_prod = Family("prod") @@ -345,10 +308,7 @@ def build(self): n_web_test.defuser = e_no_web_test n_web.add(n_web_prod, n_web_test) - if not follow_osuite: - barrier_ymd = barrier_hh.ymd - - if main_hh.get_variable("EMOS_BASE").value() != "12": + if is_00z_cycle(main_hh): main_hh.add( n_initialize, n_hr, @@ -373,7 +333,7 @@ def build(self): n_lag_arch_init.defuser = e_no_ecfs_archive n_lag_arch_fc = Task("arch_fc") n_lag_arch_fc.defuser = e_no_ecfs_archive - if main_hh.get_variable("EMOS_BASE").value() != "12": + if is_00z_cycle(main_hh): lag_hh.add(n_lag_arch_init, n_lag_arch_fc) lag_hh.trigger = main_hh.complete.across("YMD") n_daily_lag.add(lag_hh) diff --git a/geoglows_ecflow/workflow/comfies/config.py b/geoglows_ecflow/workflow/comfies/config.py index 1075676..ce48dbc 100644 --- a/geoglows_ecflow/workflow/comfies/config.py +++ b/geoglows_ecflow/workflow/comfies/config.py @@ -14,7 +14,8 @@ import os import sys -import imp +import importlib.util +from importlib.machinery import SourceFileLoader import collections.abc as collections import datetime @@ -54,6 +55,26 @@ class ConfigItemTypeError(ConfigError): # ---------------------------------------- +def _load_source(name, path): + """Load a Python source file as a module and return it. + + An explicit ``SourceFileLoader`` is used so the file is parsed as Python + source regardless of its extension (config files use ``.cfg``). + """ + loader = SourceFileLoader(name, path) + spec = importlib.util.spec_from_file_location(name, path, loader=loader) + module = importlib.util.module_from_spec(spec) + # Register before executing so the file's own relative imports resolve; + # drop it again if execution fails, leaving no half-built module behind. + sys.modules[name] = module + try: + loader.exec_module(module) + except BaseException: + sys.modules.pop(name, None) + raise + return module + + class ConfigSource(object): """ Abstract Base Class. @@ -105,7 +126,7 @@ def _load(self, path): try: old_dont_write_bytecode = sys.dont_write_bytecode sys.dont_write_bytecode = True - data = imp.load_source('_sdeploy_config_'+path, path).__dict__ + data = _load_source('_sdeploy_config_'+path, path).__dict__ sys.dont_write_bytecode = old_dont_write_bytecode except IOError as e: msg = path + ": " + e.strerror diff --git a/geoglows_ecflow/workflow/comfies/ooflow.py b/geoglows_ecflow/workflow/comfies/ooflow.py index 0a3c79b..f35e85a 100644 --- a/geoglows_ecflow/workflow/comfies/ooflow.py +++ b/geoglows_ecflow/workflow/comfies/ooflow.py @@ -21,7 +21,7 @@ import os from functools import wraps -from pkg_resources import parse_version +from packaging.version import parse as parse_version import ecflow from .py2 import reduce try: diff --git a/geoglows_ecflow/workflow/comfies/sdeploy.py b/geoglows_ecflow/workflow/comfies/sdeploy.py index c92227d..1193d19 100644 --- a/geoglows_ecflow/workflow/comfies/sdeploy.py +++ b/geoglows_ecflow/workflow/comfies/sdeploy.py @@ -28,8 +28,9 @@ import subprocess import contextlib import logging as log -import imp import importlib +import importlib.util +import importlib.machinery import argparse from geoglows_ecflow.workflow.comfies.memoize import memoize @@ -60,7 +61,7 @@ TemplateError, ) from geoglows_ecflow.workflow.comfies.sjob import SshHost -from pkg_resources import parse_version +from packaging.version import parse as parse_version from geoglows_ecflow.workflow.comfies.version import __version__ from geoglows_ecflow.workflow.comfies.py2 import basestring import ecflow @@ -368,12 +369,20 @@ def module_exists(name, path): """ Test if module exists in path. If module doesn't exist return False. + + Walk the dotted name one component at a time, descending into each + package's directory, and return the final source path (truthy) or False. """ for x in name.split("."): - try: - file, path, descr = imp.find_module(x, [path]) - except ImportError: + spec = importlib.machinery.PathFinder.find_spec(x, [path]) + if spec is None: return False + if spec.submodule_search_locations: + # x is a package; search the next component inside its directory + path = list(spec.submodule_search_locations)[0] + else: + # x is a module; its origin is the source file + path = spec.origin return path @@ -772,9 +781,9 @@ class BaseBuilder(object): def __init__(self, config): globals()["ecflow"] = importlib.import_module(self.ecflow_module) - if parse_version(self.comfies_minimum_version) > parse_version( - __version__ - ): + if self.comfies_minimum_version != "any" and parse_version( + self.comfies_minimum_version + ) > parse_version(__version__): raise ComfiesVersionError( "This suite needs version {}" " or later of comfies package".format( diff --git a/geoglows_ecflow/workflow/parts/start.py b/geoglows_ecflow/workflow/parts/start.py deleted file mode 100644 index 255bede..0000000 --- a/geoglows_ecflow/workflow/parts/start.py +++ /dev/null @@ -1,39 +0,0 @@ -from geoglows_ecflow.workflow.parts.nodes import Family, Task, Trigger -from geoglows_ecflow.workflow.parts.packages import PackageInstallers -from geoglows_ecflow.workflow.comfies.ooflow import Variable, Label, Defuser, complete -from geoglows_ecflow.workflow.comfies.ooflow import NullEvent, Null - -""" -Nodes executed once to initialise newly created suite. -""" - - -class MakeFamily(Family): - - """ - sync_install tasks run on the computational - cluster and install EFAS software packages. - """ - - def __init__(self): - - super(MakeFamily, self).__init__('make') - - # Create a family of package installers - # (each task installs an individual EFAS package) - - n_packages = PackageInstallers( - packages = [ - 'scripts', - 'rapidpy', - 'geoglows_ecflow', - 'basininflow' - ] - ) - - self.add( - Variable('SMSTRIES', 1), - n_packages - ) - - self.add_inlimit('make') diff --git a/geoglows_ecflow/workflow/scripts/common/conda.h b/geoglows_ecflow/workflow/scripts/common/conda.h index 2d469e2..812b477 100644 --- a/geoglows_ecflow/workflow/scripts/common/conda.h +++ b/geoglows_ecflow/workflow/scripts/common/conda.h @@ -6,10 +6,10 @@ set +eu #conda config --append envs_dirs $suite_libdir/virtualenv _CONDA_SET_GEOTIFF_CSV="" -GDAL_DATA=$suite_libdir/virtualenvs/rapid/share/gdal -GDAL_DRIVER_PATH=$suite_libdir/virtualenvs/rapid/lib/gdalplugins +GDAL_DATA=$suite_libdir/virtualenvs/routing/share/gdal +GDAL_DRIVER_PATH=$suite_libdir/virtualenvs/routing/lib/gdalplugins GEOTIFF_CSV='' -PROJ_LIB=$suite_libdir/virtualenvs/rapid/share/proj +PROJ_LIB=$suite_libdir/virtualenvs/routing/share/proj conda config --get -conda activate $suite_libdir/virtualenvs/rapid +conda activate $suite_libdir/virtualenvs/routing set -eu \ No newline at end of file diff --git a/geoglows_ecflow/workflow/scripts/common/install_package.ecf b/geoglows_ecflow/workflow/scripts/common/install_package.ecf index e4a1b1b..a13c3f9 100644 --- a/geoglows_ecflow/workflow/scripts/common/install_package.ecf +++ b/geoglows_ecflow/workflow/scripts/common/install_package.ecf @@ -3,7 +3,7 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce %includeonce # Where is the source code of the package. diff --git a/geoglows_ecflow/workflow/scripts/common/petsc.h b/geoglows_ecflow/workflow/scripts/common/petsc.h deleted file mode 100644 index 9157943..0000000 --- a/geoglows_ecflow/workflow/scripts/common/petsc.h +++ /dev/null @@ -1 +0,0 @@ -# Dummy PETSC \ No newline at end of file diff --git a/geoglows_ecflow/workflow/scripts/rapid/build_petsc.ecf b/geoglows_ecflow/workflow/scripts/rapid/build_petsc.ecf deleted file mode 100644 index 01f12a9..0000000 --- a/geoglows_ecflow/workflow/scripts/rapid/build_petsc.ecf +++ /dev/null @@ -1,62 +0,0 @@ -#!/bin/bash -l -%manual - - This task builds PETSC on systems that do not have PETSC available or where we require - a specific version - - Also installs Mpich and fblaslapack - - OPERATORS - - Leave Aborted - Inform analyst during weekdays. - - - ANALYSTS - - Should only run on a suite installation. - Contact person responsible during working hours. - -%end - -%includeonce -%includeonce -%includeonce - -srcdir= -stage_dir=$suite_libdir/petsc - -mkdir -p $stage_dir -cd $stage_dir -if [ ${srcdir%%+*} = "git" ]; then - - # Get sources - echo "detecting git" - giturl=${srcdir%%@*} - giturl=${giturl##*+} - gitbranch=${srcdir##*@} - git clone $giturl --branch $gitbranch --depth 1 $stage_dir - stage_info="files fetched from $giturl@$gitbranch" -else - rsync -aip --chmod=ug+rw,o+r --exclude '.svn/*' $srcdir/* $stage_dir > stage_package_filelist -fi - -cat stage_package_filelist -stage_package_nfiles=$(cat stage_package_filelist | grep '^.[^d]' | wc -l) -stage_info="$stage_package_nfiles files fetched from $srcdir" - -for tar in petsc-3.13.0.tar.gz mpich-3.3.2.tar.gz pkg-fblaslapack.tar.gz;do - tar -xf $tar -done - -PETSC_DIR=$stage_dir/petsc-3.13.0 -cd petsc-3.13.0 - -#/rapid_install_prereqs.sh -python3 './configure' 'PETSC_DIR='$PWD 'PETSC_ARCH=linux-gcc-c' '--download-fblaslapack' '--download-mpich='../mpich-3.3.2.tar.gz '--with-cc=gcc' '--with-fc=gfortran' '--with-clanguage=c' '--with-debugging=0' - -#make $PETSC_DIR PETSC_ARCH=linux-gcc-c all -#make $PETSC_DIR PETSC_ARCH=linux-gcc-c check -make all -make check - diff --git a/geoglows_ecflow/workflow/scripts/rapid/build_rapid.ecf b/geoglows_ecflow/workflow/scripts/rapid/build_rapid.ecf deleted file mode 100644 index 5e26b0e..0000000 --- a/geoglows_ecflow/workflow/scripts/rapid/build_rapid.ecf +++ /dev/null @@ -1,62 +0,0 @@ -#!/bin/bash -l -%manual - - This task builds the rapid binary - - OPERATORS - - Leave Aborted - Inform analyst during weekdays. - - - ANALYSTS - - Should only run on a suite installation. - Contact person responsible during working hours. - -%end - -%includeonce -%includeonce - - -%includeonce -%includeonce -%includeonce - -srcdir= - -blddir=$suite_dir/build/rapid/ -PETSC_INSTALL_DIR=$suite_dir/lib/ -blddir_petsc=$suite_dir/build/petsc/ -rm -rf $blddir -mkdir -p $blddir $PETSC_INSTALL_DIR $blddir_petsc - -if [ ${srcdir%%+*} = "git" ]; then - - # Get sources - echo "detecting git" - giturl=${srcdir%%@*} - giturl=${giturl##*+} - gitbranch=${srcdir##*@} - git clone $giturl --branch $gitbranch --depth 1 $blddir - stage_info="files fetched from $giturl@$gitbranch" -else - rsync -aip --chmod=ug+rw,o+r --exclude '.svn/*' --exclude '.git/*' $srcdir/* $blddir/ -fi - - -cd $blddir - -%include -cd src -make rapid - -mkdir -p $suite_dir/bin -cp $blddir/src/rapid $suite_dir/bin/rapid -chmod 755 $suite_dir/bin/rapid - -cd ../tst -gfortran -o tst_run_conv_Qinit tst_run_conv_Qinit.f90 -I $TACC_NETCDF_INC -L $TACC_NETCDF_LIB -lnetcdff -cp tst_run_conv_Qinit $suite_dir/bin/ -chmod 755 $suite_dir/bin/tst_run_conv_Qinit \ No newline at end of file diff --git a/geoglows_ecflow/workflow/scripts/rapid/rapid.h b/geoglows_ecflow/workflow/scripts/rapid/rapid.h deleted file mode 100644 index 0a3e420..0000000 --- a/geoglows_ecflow/workflow/scripts/rapid/rapid.h +++ /dev/null @@ -1,11 +0,0 @@ -# Rapid.h - -%includeonce - -export TACC_NETCDF_LIB=$NETCDF4_DIR/lib -export TACC_NETCDF_INC=$NETCDF4_DIR/include -export PETSC_DIR=$suite_libdir/petsc/petsc-3.13.0 -export PETSC_ARCH='linux-gcc-c' -export LD_LIBRARY_PATH=$TACC_NETCDF_LIB -export PATH=$PATH:/$PETSC_DIR/$PETSC_ARCH/bin -export PATH=$suite_dir/bin:$PATH diff --git a/geoglows_ecflow/workflow/scripts/rapid/rapid_venv.h b/geoglows_ecflow/workflow/scripts/rapid/rapid_venv.h deleted file mode 100644 index ed33fb7..0000000 --- a/geoglows_ecflow/workflow/scripts/rapid/rapid_venv.h +++ /dev/null @@ -1,8 +0,0 @@ -%includeonce -%includeonce -%includeonce - -set +u -export PYTHONPATH=$suite_libdir/virtualenvs/rapid/lib/python3.10/site-packages/:$suite_dir/lib/python:$PYTHONPATH -export LD_LIBRARY_PATH=$suite_libdir/virtualenvs/rapid/lib -set -u diff --git a/geoglows_ecflow/workflow/scripts/rapid/arch_fc.ecf b/geoglows_ecflow/workflow/scripts/routing/arch_fc.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/arch_fc.ecf rename to geoglows_ecflow/workflow/scripts/routing/arch_fc.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/arch_init.ecf b/geoglows_ecflow/workflow/scripts/routing/arch_init.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/arch_init.ecf rename to geoglows_ecflow/workflow/scripts/routing/arch_init.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/archive_qinit.ecf b/geoglows_ecflow/workflow/scripts/routing/archive_qinit.ecf similarity index 60% rename from geoglows_ecflow/workflow/scripts/rapid/archive_qinit.ecf rename to geoglows_ecflow/workflow/scripts/routing/archive_qinit.ecf index d642046..d40f939 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/archive_qinit.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/archive_qinit.ecf @@ -1,19 +1,19 @@ #!/bin/bash -l %manual - Create initial conditions for next run the rapid workflow + Archive Qinit files for the next forecast cycle %end %includeonce %includeonce -ls -al $ens_rapid_input -cd $ens_rapid_input +ls -al $ens_input +cd $ens_input find . -type f -name "Qinit_${ens_basetime}.*" -print > Qinit_files if [[ $(cat Qinit_files | wc -l) -ne $(find . -mindepth 1 -maxdepth 1 -type d | wc -l) ]];then echo "Fail, less Qinit files than expected" false fi -tar cvzf $ens_fcdir/init_files_$ens_basetime.tar.gz -C $ens_rapid_input --files-from=Qinit_files +tar cvzf $ens_fcdir/init_files_$ens_basetime.tar.gz -C $ens_input --files-from=Qinit_files diff --git a/geoglows_ecflow/workflow/scripts/rapid/archive_to_aws.ecf b/geoglows_ecflow/workflow/scripts/routing/archive_to_aws.ecf similarity index 96% rename from geoglows_ecflow/workflow/scripts/rapid/archive_to_aws.ecf rename to geoglows_ecflow/workflow/scripts/routing/archive_to_aws.ecf index 79a975a..2bb9b8b 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/archive_to_aws.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/archive_to_aws.ecf @@ -6,7 +6,7 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce #python3 -m geoglows_ecflow.resources.archive_to_aws $ens_fcdir $suite_staticdata/aws_config.yml diff --git a/geoglows_ecflow/workflow/scripts/rapid/bc.ecf b/geoglows_ecflow/workflow/scripts/routing/bc.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/bc.ecf rename to geoglows_ecflow/workflow/scripts/routing/bc.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/build_venv.ecf b/geoglows_ecflow/workflow/scripts/routing/build_venv.ecf similarity index 71% rename from geoglows_ecflow/workflow/scripts/rapid/build_venv.ecf rename to geoglows_ecflow/workflow/scripts/routing/build_venv.ecf index 0b576ed..08d449b 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/build_venv.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/build_venv.ecf @@ -1,7 +1,7 @@ #!/bin/bash -l %manual - Build Virtual Env for Rapid Workflow + Build Virtual Env for Routing Workflow OPERATORS @@ -22,20 +22,20 @@ %includeonce _CONDA_SET_GEOTIFF_CSV="" -GDAL_DATA=$suite_libdir/virtualenvs/rapid/share/gdal -GDAL_DRIVER_PATH=$suite_libdir/virtualenvs/rapid/lib/gdalplugins +GDAL_DATA=$suite_libdir/virtualenvs/routing/share/gdal +GDAL_DRIVER_PATH=$suite_libdir/virtualenvs/routing/lib/gdalplugins GEOTIFF_CSV='' -PROJ_LIB=$suite_libdir/virtualenvs/rapid/share/proj +PROJ_LIB=$suite_libdir/virtualenvs/routing/share/proj export _CONDA_SET_GDAL_DATA="" export _CONDA_SET_GDAL_DRIVER_PATH="" export _CONDA_SET_PROJ_LIB="" export LC_ALL=C -rm -rf $suite_libdir/virtualenvs/rapid +rm -rf $suite_libdir/virtualenvs/routing module load conda/22.11.1-2 cat > venv.yaml << EOF -name: rapid +name: routing channels: - conda-forge - defaults @@ -55,16 +55,18 @@ dependencies: - netcdf4==1.6.5 - zarr==2.16.1 - boto3==1.28.65 - - pangaea==0.0.4 - nco==5.1.9 - awscli + - pip + - pip: + - river-route>=2.1.1 EOF conda config --set auto_activate_base True -conda create python=3.10 -c conda-forge -p $suite_libdir/virtualenvs/rapid -conda activate $suite_libdir/virtualenvs/rapid +conda create python=3.10 -c conda-forge -p $suite_libdir/virtualenvs/routing +conda activate $suite_libdir/virtualenvs/routing set +u conda install -c conda-forge -y mamba -mamba env update -f venv.yaml -p $suite_libdir/virtualenvs/rapid +mamba env update -f venv.yaml -p $suite_libdir/virtualenvs/routing set -u %include diff --git a/geoglows_ecflow/workflow/scripts/rapid/check_mars.ecf b/geoglows_ecflow/workflow/scripts/routing/check_mars.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/check_mars.ecf rename to geoglows_ecflow/workflow/scripts/routing/check_mars.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/clean.ecf b/geoglows_ecflow/workflow/scripts/routing/clean.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/clean.ecf rename to geoglows_ecflow/workflow/scripts/routing/clean.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/combine_forecast_warnings.ecf b/geoglows_ecflow/workflow/scripts/routing/combine_forecast_warnings.ecf similarity index 87% rename from geoglows_ecflow/workflow/scripts/rapid/combine_forecast_warnings.ecf rename to geoglows_ecflow/workflow/scripts/routing/combine_forecast_warnings.ecf index 04c1c5f..926b8ea 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/combine_forecast_warnings.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/combine_forecast_warnings.ecf @@ -7,7 +7,7 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce python3 -m geoglows_ecflow.resources.concat_forecast_warnings $ens_fcdir diff --git a/geoglows_ecflow/workflow/scripts/rapid/combine_plain_table.ecf b/geoglows_ecflow/workflow/scripts/routing/combine_plain_table.ecf similarity index 79% rename from geoglows_ecflow/workflow/scripts/rapid/combine_plain_table.ecf rename to geoglows_ecflow/workflow/scripts/routing/combine_plain_table.ecf index 327db62..eab5d56 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/combine_plain_table.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/combine_plain_table.ecf @@ -7,20 +7,20 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce python3 -m geoglows_ecflow.resources.combine_esri_tables $ens_fcdir # Do not zip. it is long and unnecessary -#cd $ens_rapid_output/map_style_tables/ +#cd $ens_output/map_style_tables/ #for file in *.csv; do # tar -rvf ${ens_fcdir}/mapstyletable_${ens_basetime}.tar $file #done #gzip -f ${ens_fcdir}/mapstyletable_${ens_basetime}.tar # ## create VPU archive (see diss.ecf) -#cd $ens_rapid_output -#for path in $(find $ens_rapid_input/ -mindepth 1 -maxdepth 1 -type d); do +#cd $ens_output +#for path in $(find $ens_input/ -mindepth 1 -maxdepth 1 -type d); do # for file in Qout_*.nc; do # basename $path | xargs -I {} tar -rvf $ens_fcdir/Qout_{}_${ens_basetime}.tar $file; # done diff --git a/geoglows_ecflow/workflow/scripts/rapid/comp_init.ecf b/geoglows_ecflow/workflow/scripts/routing/comp_init.ecf similarity index 66% rename from geoglows_ecflow/workflow/scripts/rapid/comp_init.ecf rename to geoglows_ecflow/workflow/scripts/routing/comp_init.ecf index 300767b..60d1574 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/comp_init.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/comp_init.ecf @@ -1,14 +1,14 @@ #!/bin/bash -l %manual - Create initial conditions for next run the rapid workflow + Compute initial conditions for the next forecast cycle %end %includeonce %includeonce %includeonce -%includeonce +%includeonce python3 -m geoglows_ecflow.resources.compute_init_flows $ens_fcdir %VPU% diff --git a/geoglows_ecflow/workflow/scripts/rapid/day_one.ecf b/geoglows_ecflow/workflow/scripts/routing/day_one.ecf similarity index 87% rename from geoglows_ecflow/workflow/scripts/rapid/day_one.ecf rename to geoglows_ecflow/workflow/scripts/routing/day_one.ecf index 221d22f..6c5931d 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/day_one.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/day_one.ecf @@ -7,6 +7,6 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce python3 -m geoglows_ecflow.resources.day_one_forecast $ens_fcdir %VPU% $ens_workdir/forecast_records diff --git a/geoglows_ecflow/workflow/scripts/rapid/diss.ecf b/geoglows_ecflow/workflow/scripts/routing/diss.ecf similarity index 72% rename from geoglows_ecflow/workflow/scripts/rapid/diss.ecf rename to geoglows_ecflow/workflow/scripts/routing/diss.ecf index a5af964..a4a0a72 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/diss.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/diss.ecf @@ -38,12 +38,4 @@ case $context in ecpds -echost aux -destination CEMS_Flood_Glofas -force -lifetime 7d -source Runoff.${ens_ymd}.${ens_base}.exp${ens_mars_expver}.Fgrid.netcdf.tar.gz -target /tcyc/Runoff.${ens_ymd}.${ens_base}.exp${ens_mars_expver}.Fgrid.netcdf.tar.gz ecpds -echost aux -destination CEMS_Flood_Rapid -force -lifetime 7d -source Runoff.${ens_ymd}.${ens_base}.exp${ens_mars_expver}.Fgrid.netcdf.tar.gz -target Runoff.${ens_ymd}.${ens_base}.exp${ens_mars_expver}.Fgrid.netcdf.tar.gz ;; - -# rapid) -# for file in $ens_fcdir/*.tar.gz;do - #ecpds -echost aux -destination CEMS_Flood_Glofas -force -lifetime 7d -source $file -target /tcyc/rapid_${ens_ymd}_${ens_base}/${file##*/} - #ecpds -echost aux -destination CEMS_Flood_Rapid -force -lifetime 7d -source $file -target rapid_${ens_ymd}_${ens_base}/${file##*/} -# done - -# ;; esac diff --git a/geoglows_ecflow/workflow/scripts/rapid/ens.h b/geoglows_ecflow/workflow/scripts/routing/ens.h similarity index 93% rename from geoglows_ecflow/workflow/scripts/rapid/ens.h rename to geoglows_ecflow/workflow/scripts/routing/ens.h index 67e814c..ab04b57 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/ens.h +++ b/geoglows_ecflow/workflow/scripts/routing/ens.h @@ -36,8 +36,8 @@ ens_pworkdir=$suite_workdir ens_inputdir=$ens_workdir/grib/$ens_basetime ens_fcdir=$ens_workdir/fc/$ens_basetime -ens_rapid_input=$ens_fcdir/input -ens_rapid_output=$ens_fcdir/output +ens_input=$ens_fcdir/input +ens_output=$ens_fcdir/output ens_member=%MEMBER:0% ens_nmembers=%MEMBERS:51% diff --git a/geoglows_ecflow/workflow/scripts/rapid/ens_member.ecf b/geoglows_ecflow/workflow/scripts/routing/ens_member.ecf similarity index 59% rename from geoglows_ecflow/workflow/scripts/rapid/ens_member.ecf rename to geoglows_ecflow/workflow/scripts/routing/ens_member.ecf index 7c6fc88..00d36c6 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/ens_member.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/ens_member.ecf @@ -1,7 +1,7 @@ #!/bin/bash -l %manual - Ensemble Member for Rapid Workflow + Ensemble Member for Routing Workflow %end @@ -9,12 +9,9 @@ %includeonce %includeonce %includeonce -%includeonce -%includeonce +%includeonce -export LD_LIBRARY_PATH=$TACC_NETCDF_LIB - mkdir -p $sim_fcdir/execute mkdir -p $sim_fcdir/subprocess @@ -25,4 +22,4 @@ if [[ ! -f $sim_fcdir/${ens_member#0}.runoff.nc ]];then false fi -python3 -m geoglows_ecflow.resources.run_rapid_forecast $sim_fcdir %JOB_ID% $suite_dir/bin/rapid +python3 -m geoglows_ecflow.resources.run_river_route_forecast $sim_fcdir %JOB_ID% diff --git a/geoglows_ecflow/workflow/scripts/rapid/initialize.ecf b/geoglows_ecflow/workflow/scripts/routing/initialize.ecf similarity index 68% rename from geoglows_ecflow/workflow/scripts/rapid/initialize.ecf rename to geoglows_ecflow/workflow/scripts/routing/initialize.ecf index 024c2b4..0067b4c 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/initialize.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/initialize.ecf @@ -23,8 +23,6 @@ %includeonce %includeonce %includeonce -%includeonce -%includeonce ecp $suite_iniarchdir/$ens_pyear/$ens_pmonth/$ens_pymd/init_files_${ens_basetime}.tar.gz ./ tar xvf init_files_${ens_basetime}.tar.gz @@ -32,7 +30,4 @@ tar xvf init_files_${ens_basetime}.tar.gz for dir in $(find . -mindepth 1 -maxdepth 1 -regextype posix-extended -type d -regex '.*/[0-9]{3}'); do mkdir -p $suite_staticdata/input/$dir/ cp $dir/Qinit_${ens_basetime}.* $suite_staticdata/input/$dir/ - if [[ ! -f $suite_staticdata/input/$dir/Qinit_${ens_basetime}.nc ]] && [[ -f $suite_staticdata/input/$dir/Qinit_${ens_basetime}.csv ]]; then - tst_run_conv_Qinit $suite_staticdata/input/$dir/Qinit_${ens_basetime}.csv $suite_staticdata/input/$dir/Qinit_${ens_basetime}.nc - fi done diff --git a/geoglows_ecflow/workflow/scripts/rapid/install_package.ecf b/geoglows_ecflow/workflow/scripts/routing/install_package.ecf similarity index 99% rename from geoglows_ecflow/workflow/scripts/rapid/install_package.ecf rename to geoglows_ecflow/workflow/scripts/routing/install_package.ecf index efee3fc..61ca8b7 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/install_package.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/install_package.ecf @@ -94,7 +94,7 @@ done if [[ -f _setup_ ]]; then . '_setup_' fi -%includeonce +%includeonce dest_dir=$suite_dir/lib/%PACKAGE% make uninstall PREFIX=$dest_dir LIBDIR=$suite_libdir diff --git a/geoglows_ecflow/workflow/scripts/rapid/install_static_data.ecf b/geoglows_ecflow/workflow/scripts/routing/install_static_data.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/install_static_data.ecf rename to geoglows_ecflow/workflow/scripts/routing/install_static_data.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/nc_to_zarr.ecf b/geoglows_ecflow/workflow/scripts/routing/nc_to_zarr.ecf similarity index 89% rename from geoglows_ecflow/workflow/scripts/rapid/nc_to_zarr.ecf rename to geoglows_ecflow/workflow/scripts/routing/nc_to_zarr.ecf index f3e39cd..544fd6d 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/nc_to_zarr.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/nc_to_zarr.ecf @@ -7,7 +7,7 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce postprocess_rapid_output -d $ens_fcdir/output python3 -m geoglows_ecflow.resources.netcdf_to_zarr $ens_fcdir diff --git a/geoglows_ecflow/workflow/scripts/rapid/nco_calc.ecf b/geoglows_ecflow/workflow/scripts/routing/nco_calc.ecf similarity index 93% rename from geoglows_ecflow/workflow/scripts/rapid/nco_calc.ecf rename to geoglows_ecflow/workflow/scripts/routing/nco_calc.ecf index 8e893f8..b266af8 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/nco_calc.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/nco_calc.ecf @@ -9,8 +9,8 @@ %includeonce %includeonce -echo "Looking for rapid outputs in directory: $ens_rapid_output" -cd $ens_rapid_output +echo "Looking for routing outputs in directory: $ens_output" +cd $ens_output vpu_number=%VPU% echo "Your VPU number is ${vpu_number}" diff --git a/geoglows_ecflow/workflow/scripts/rapid/plain_table.ecf b/geoglows_ecflow/workflow/scripts/routing/plain_table.ecf similarity index 90% rename from geoglows_ecflow/workflow/scripts/rapid/plain_table.ecf rename to geoglows_ecflow/workflow/scripts/routing/plain_table.ecf index 7f0577a..8fe93b8 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/plain_table.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/plain_table.ecf @@ -7,7 +7,7 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce ln -snf $suite_staticdata/return_periods_dir $ens_fcdir/return_periods_dir diff --git a/geoglows_ecflow/workflow/scripts/rapid/prep_task.ecf b/geoglows_ecflow/workflow/scripts/routing/prep_task.ecf similarity index 55% rename from geoglows_ecflow/workflow/scripts/rapid/prep_task.ecf rename to geoglows_ecflow/workflow/scripts/routing/prep_task.ecf index 2cfa7c5..50ed4ca 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/prep_task.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/prep_task.ecf @@ -1,7 +1,7 @@ #!/bin/bash -l %manual - Prepatory Task for the rapid workflow + Preparatory Task for the routing workflow @@ -11,14 +11,13 @@ %includeonce %includeonce %includeonce -%includeonce +%includeonce mkdir -p $sim_fcdir -mkdir -p $sim_fcdir/inflows ln -snf $suite_staticdata/input $sim_fcdir/input -python3 -m geoglows_ecflow.resources.prep_rapid_forecast $sim_fcdir +python3 -m geoglows_ecflow.resources.prep_river_route_forecast $sim_fcdir %include diff --git a/geoglows_ecflow/workflow/scripts/rapid/remove_conda.h b/geoglows_ecflow/workflow/scripts/routing/remove_conda.h similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/remove_conda.h rename to geoglows_ecflow/workflow/scripts/routing/remove_conda.h diff --git a/geoglows_ecflow/workflow/scripts/rapid/retrieve_ens.ecf b/geoglows_ecflow/workflow/scripts/routing/retrieve_ens.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/retrieve_ens.ecf rename to geoglows_ecflow/workflow/scripts/routing/retrieve_ens.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/retrieve_hres.ecf b/geoglows_ecflow/workflow/scripts/routing/retrieve_hres.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/retrieve_hres.ecf rename to geoglows_ecflow/workflow/scripts/routing/retrieve_hres.ecf diff --git a/geoglows_ecflow/workflow/scripts/routing/routing_venv.h b/geoglows_ecflow/workflow/scripts/routing/routing_venv.h new file mode 100644 index 0000000..a060259 --- /dev/null +++ b/geoglows_ecflow/workflow/scripts/routing/routing_venv.h @@ -0,0 +1,8 @@ +%includeonce +%includeonce +%includeonce + +set +u +export PYTHONPATH=$suite_libdir/virtualenvs/routing/lib/python3.10/site-packages/:$suite_dir/lib/python:$PYTHONPATH +export LD_LIBRARY_PATH=$suite_libdir/virtualenvs/routing/lib +set -u diff --git a/geoglows_ecflow/workflow/scripts/rapid/run_hr.ecf b/geoglows_ecflow/workflow/scripts/routing/run_hr.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/run_hr.ecf rename to geoglows_ecflow/workflow/scripts/routing/run_hr.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/sim.h b/geoglows_ecflow/workflow/scripts/routing/sim.h similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/sim.h rename to geoglows_ecflow/workflow/scripts/routing/sim.h diff --git a/geoglows_ecflow/workflow/scripts/rapid/sim_%CONTEXT%.h b/geoglows_ecflow/workflow/scripts/routing/sim_%CONTEXT%.h similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/sim_%CONTEXT%.h rename to geoglows_ecflow/workflow/scripts/routing/sim_%CONTEXT%.h diff --git a/geoglows_ecflow/workflow/scripts/rapid/sim_ens.h b/geoglows_ecflow/workflow/scripts/routing/sim_ens.h similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/sim_ens.h rename to geoglows_ecflow/workflow/scripts/routing/sim_ens.h diff --git a/geoglows_ecflow/workflow/scripts/rapid/sim_hres.h b/geoglows_ecflow/workflow/scripts/routing/sim_hres.h similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/sim_hres.h rename to geoglows_ecflow/workflow/scripts/routing/sim_hres.h diff --git a/geoglows_ecflow/workflow/scripts/rapid/sleep.ecf b/geoglows_ecflow/workflow/scripts/routing/sleep.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/sleep.ecf rename to geoglows_ecflow/workflow/scripts/routing/sleep.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/suite.h b/geoglows_ecflow/workflow/scripts/routing/suite.h similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/suite.h rename to geoglows_ecflow/workflow/scripts/routing/suite.h diff --git a/geoglows_ecflow/workflow/scripts/rapid/toggles.ecf b/geoglows_ecflow/workflow/scripts/routing/toggles.ecf similarity index 100% rename from geoglows_ecflow/workflow/scripts/rapid/toggles.ecf rename to geoglows_ecflow/workflow/scripts/routing/toggles.ecf diff --git a/geoglows_ecflow/workflow/scripts/rapid/web_push.ecf b/geoglows_ecflow/workflow/scripts/routing/web_push.ecf similarity index 97% rename from geoglows_ecflow/workflow/scripts/rapid/web_push.ecf rename to geoglows_ecflow/workflow/scripts/routing/web_push.ecf index 79ac1bb..74c8c59 100644 --- a/geoglows_ecflow/workflow/scripts/rapid/web_push.ecf +++ b/geoglows_ecflow/workflow/scripts/routing/web_push.ecf @@ -43,7 +43,7 @@ if [[ $suite_mode == "prod" ]] || [[ $suite_mode == "test" ]];then for host in $hosts;do - srcdir=$ens_rapid_output + srcdir=$ens_output destdir=$destdir_root/forecasts ssh efas@$host "umask 002; mkdir -p $destdir" diff --git a/pyproject.toml b/pyproject.toml index a78fb81..f8f1677 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "geoglows_ecflow" version = "3.1.1" -description = "ECFLOW RAPID workflow for GEOGloWS" +description = "ECFLOW workflow for GEOGloWS streamflow forecasting" authors = [ { name = "Michael Souffront", email = "msouffront@aquaveo.com" }, { name = "Riley Hales", email = "rchales@byu.edu" }, @@ -17,12 +17,22 @@ dependencies = [ "fastparquet>=2023.8.0", "zarr>=2.16.1", "boto3>=1.28.65", - "basininflow>=0.13.0", + "river-route>=2.1.1", + "packaging>=21.0", ] -requires-python = ">=3.10" +requires-python = ">=3.12,<3.14" readme = "README.md" license = { text = "BSD-3-Clause" } +[project.optional-dependencies] +dev = [ + "pytest>=7.0", + "pytest-cov>=4.0", +] + +[tool.pytest.ini_options] +testpaths = ["tests"] + [tool.coverage.report] show_missing = true diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..711c7b1 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,60 @@ +"""Shared pytest fixtures and import shims for the resources test suite. + +These tests pin the *current* behavior of the pure helper functions used by +the river-route forecast workflow. The functions under test never touch the +heavy optional dependency ``river_route``, so when it is not installed we +inject a Dummy module into ``sys.modules`` before any test imports +``run_river_route_forecast``. When the real package *is* installed (e.g. on the +Ubuntu workflow box), it is used as-is. +""" + +import sys +import types + +import pandas as pd +import pytest + +# ``run_river_route_forecast`` does ``import river_route as rr`` at module top, +# but ``_find_state_init`` (the function under test) never touches it. Prefer +# the real package when installed; only fall back to a Dummy module so the +# import succeeds where the dependency is absent. +if "river_route" not in sys.modules: + try: + import river_route # noqa: F401 (use the real package when present) + except ModuleNotFoundError: + sys.modules["river_route"] = types.ModuleType("river_route") + + +@pytest.fixture +def make_flows(): + """Return a factory building a forecasted-flows DataFrame. + + The frame mirrors what ``day_one_forecast`` passes around: a ``means`` + column of forecast values and a parallel ``times`` column. + """ + + def _make(means, times=None): + if times is None: + times = [f"t{i}" for i in range(len(means))] + return pd.DataFrame({"means": list(means), "times": list(times)}) + + return _make + + +@pytest.fixture +def rp_table(): + """A single-row return-period table indexed by comid. + + Columns match what ``check_for_return_period_flow`` reads: rp2..rp100. + """ + return pd.DataFrame( + { + "rp2": [2.0], + "rp5": [5.0], + "rp10": [10.0], + "rp25": [25.0], + "rp50": [50.0], + "rp100": [100.0], + }, + index=[12345], + ) diff --git a/tests/test_compute_init_flows.py b/tests/test_compute_init_flows.py new file mode 100644 index 0000000..afe8bb6 --- /dev/null +++ b/tests/test_compute_init_flows.py @@ -0,0 +1,41 @@ +"""Test that compute_init_flows selects the configured time index.""" + +import json +import os + +import numpy as np +import pandas as pd +import xarray as xr + +from geoglows_ecflow.resources import compute_init_flows + + +def test_init_time_index_constant(): + # t+24h on the ENS 3h grid; documented contract relied on downstream. + assert compute_init_flows.INIT_TIME_INDEX == 7 + + +def test_qinit_parquet_holds_q_at_init_time_index(tmp_path): + vpu = "101" + ymd = "2023010100" + + (tmp_path / "forecast_run.json").write_text(json.dumps({"date": ymd})) + output_dir = tmp_path / "output" + output_dir.mkdir() + (tmp_path / "input" / vpu).mkdir(parents=True) + + # Q[time, river_id] with distinct values so the selected slice is unique. + data = np.arange(10 * 3).reshape(10, 3) + ds = xr.Dataset( + {"Q": (("time", "river_id"), data)}, + coords={"time": range(10), "river_id": [1, 2, 3]}, + ) + ds.to_netcdf(output_dir / f"nces_avg_{vpu}.nc") + + compute_init_flows.main(str(tmp_path), vpu) + + out = pd.read_parquet( + os.path.join(str(tmp_path), "input", vpu, f"Qinit_{ymd}.parquet") + ) + expected = data[compute_init_flows.INIT_TIME_INDEX] + assert out["Q"].tolist() == expected.tolist() diff --git a/tests/test_day_one_forecast.py b/tests/test_day_one_forecast.py new file mode 100644 index 0000000..79493e3 --- /dev/null +++ b/tests/test_day_one_forecast.py @@ -0,0 +1,49 @@ +"""Tests for the return-period exceedance logic in day_one_forecast.""" + +import pandas as pd + +from geoglows_ecflow.resources.day_one_forecast import ( + check_for_return_period_flow, + get_time_of_first_exceedance, +) + + +def test_first_exceedance_returns_earliest_time_at_or_above_flow(make_flows): + flows = make_flows([1, 5, 10, 2], times=["a", "b", "c", "d"]) + + # Rows below 4 (1 and 2) are dropped; first remaining time is "b" (5). + assert get_time_of_first_exceedance(flows, 4) == "b" + + +def test_below_smallest_return_period_returns_input_unchanged( + make_flows, rp_table +): + largeflows = pd.DataFrame() + flows = make_flows([0.5, 1.0]) # max 1.0 < rp2 (2.0) + + result = check_for_return_period_flow(largeflows, flows, 3, rp_table) + + assert result is largeflows + assert len(result) == 0 + + +def test_appends_row_with_exceeded_thresholds(make_flows, rp_table): + largeflows = pd.DataFrame() + flows = make_flows([1, 3, 6, 8], times=["a", "b", "c", "d"]) # max 8 + + result = check_for_return_period_flow(largeflows, flows, 3, rp_table) + + assert len(result) == 1 + row = result.iloc[0] + assert row["comid"] == 12345 + assert row["stream_order"] == 3 + assert row["max_forecasted_flow"] == 8 + # rp2 (2) and rp5 (5) are exceeded; the input frame is mutated between + # calls, so the first time >= rp5 is "c", not "b". + assert row["date_exceeds_return_period_2"] == "b" + assert row["date_exceeds_return_period_5"] == "c" + # rp10 (10) and above are not reached (max is 8). + assert row["date_exceeds_return_period_10"] == "" + assert row["date_exceeds_return_period_25"] == "" + assert row["date_exceeds_return_period_50"] == "" + assert row["date_exceeds_return_period_100"] == "" diff --git a/tests/test_generate_esri_table.py b/tests/test_generate_esri_table.py new file mode 100644 index 0000000..ace68f5 --- /dev/null +++ b/tests/test_generate_esri_table.py @@ -0,0 +1,82 @@ +"""Return-period ladder tests for generate_esri_table.""" + +import netCDF4 as nc +import numpy as np +import pandas as pd +import xarray as xr + +from geoglows_ecflow.resources.generate_esri_table import ( + postprocess_vpu_forecast_directory, +) + +VPU = "101" +DATE = "20230101.00" + + +def _write_inputs(tmp_path): + """Create the nces-average and return-period netCDFs the function reads.""" + base = tmp_path / DATE # basename(dirname(output_dir)) becomes the date + output_dir = base / "output" + returnperiods = base / "rp" + output_dir.mkdir(parents=True) + returnperiods.mkdir(parents=True) + + # river 1 flows at 300 (exceeds rp5 but not rp10); river 2 stays at 5. + times = pd.date_range("2023-01-01", periods=2, freq="3h") + flows = np.array([[300.0, 5.0], [300.0, 5.0]]) + ds = xr.Dataset( + {"Q": (("time", "river_id"), flows)}, + coords={"time": times, "river_id": [1, 2]}, + ) + ds.to_netcdf(output_dir / f"nces_avg_{VPU}.nc") + + rp = nc.Dataset(str(returnperiods / f"returnperiods_{VPU}.nc"), "w") + rp.createDimension("river_id", 2) + rp.createVariable("river_id", "i4", ("river_id",))[:] = [1, 2] + thresholds = {"rp2": 10, "rp5": 100, "rp10": 500, + "rp25": 1000, "rp50": 2000, "rp100": 5000} + for name, value in thresholds.items(): + rp.createVariable(name, "f8", ("river_id",))[:] = [value, value] + rp.close() + + return output_dir, returnperiods + + +def test_return_period_ladder_assigns_expected_levels(tmp_path): + output_dir, returnperiods = _write_inputs(tmp_path) + + postprocess_vpu_forecast_directory( + str(output_dir), str(returnperiods), VPU + ) + + table = pd.read_parquet( + output_dir + / "map_style_tables" + / f"mapstyletable_{VPU}_{DATE}.parquet" + ) + ret_per_by_comid = table.groupby("comid")["ret_per"].unique() + + # 300 exceeds rp5 (100) but not rp10 (500) -> level 5. + assert ret_per_by_comid[1].tolist() == [5] + # 5 is below rp2 (10) -> level 0. + assert ret_per_by_comid[2].tolist() == [0] + + +def test_thickness_ladder_assigns_expected_levels(tmp_path): + output_dir, returnperiods = _write_inputs(tmp_path) + + postprocess_vpu_forecast_directory( + str(output_dir), str(returnperiods), VPU + ) + + table = pd.read_parquet( + output_dir + / "map_style_tables" + / f"mapstyletable_{VPU}_{DATE}.parquet" + ) + thickness_by_comid = table.groupby("comid")["thickness"].unique() + + # 300 crosses the 250 threshold but not 1500 -> thickness 3. + assert thickness_by_comid[1].tolist() == [3] + # 5 is below the first (20) threshold -> thickness 1. + assert thickness_by_comid[2].tolist() == [1] diff --git a/tests/test_helper_functions.py b/tests/test_helper_functions.py new file mode 100644 index 0000000..f67aba6 --- /dev/null +++ b/tests/test_helper_functions.py @@ -0,0 +1,58 @@ +"""Unit tests pinning current behavior of helper_functions pure functions.""" + +import json +import os + +import pytest + +from geoglows_ecflow.resources.helper_functions import ( + get_date_from_forecast_dir, + get_ensemble_number_from_forecast, + get_valid_vpucode_list, + load_forecast_run, +) + + +@pytest.mark.parametrize( + "forecast_name, expected", + [ + ("1.runoff.nc", 1), + ("52.runoff.nc", 52), + ("/abs/path/to/7.runoff.nc", 7), + # Legacy RAPID-era name: the ".205.runoff.grib.runoff.netcdf" suffix + # routes to split(".")[2] for the ensemble number. + ("20230101.00.52.205.runoff.grib.runoff.netcdf", 52), + ], +) +def test_get_ensemble_number_from_forecast(forecast_name, expected): + assert get_ensemble_number_from_forecast(forecast_name) == expected + + +def test_get_valid_vpucode_list_keeps_three_digit_dirs(tmp_path): + (tmp_path / "101").mkdir() + (tmp_path / "102").mkdir() + # Non-three-digit dir and a file are both skipped. + (tmp_path / "abc").mkdir() + (tmp_path / "103.txt").write_text("not a dir") + + result = get_valid_vpucode_list(str(tmp_path)) + + assert sorted(result) == ["101", "102"] + + +def test_get_date_from_forecast_dir_parses_timestamp(): + forecast_dir = os.path.join("some", "root", "20230101.00") + + assert get_date_from_forecast_dir(forecast_dir) == "20230101.00" + + +def test_get_date_from_forecast_dir_raises_without_match(): + with pytest.raises(AttributeError): + get_date_from_forecast_dir("no-date-here") + + +def test_load_forecast_run_round_trips_manifest(tmp_path): + manifest = {"date": "2023010100", "output_dir": "/out", "job_101_52": {}} + (tmp_path / "forecast_run.json").write_text(json.dumps(manifest)) + + assert load_forecast_run(str(tmp_path)) == manifest diff --git a/tests/test_prep_river_route_forecast.py b/tests/test_prep_river_route_forecast.py new file mode 100644 index 0000000..7b6a4f3 --- /dev/null +++ b/tests/test_prep_river_route_forecast.py @@ -0,0 +1,95 @@ +"""Tests for the forecast job-manifest builder (forecast_preprocess).""" + +import json +import os + +from geoglows_ecflow.resources.prep_river_route_forecast import ( + forecast_preprocess, +) + + +def _build_cycle(tmp_path): + """Create a minimal input/runoff layout and return its paths.""" + input_dir = tmp_path / "input" + for vpu in ("101", "102"): + (input_dir / vpu).mkdir(parents=True) + output_dir = tmp_path / "output" + # The runoff dir's basename is taken verbatim as the cycle date. + runoff_dir = tmp_path / "2023010100" + runoff_dir.mkdir() + for mem in (1, 2, 52): + (runoff_dir / f"{mem}.runoff.nc").write_text("") + return input_dir, output_dir, runoff_dir + + +def test_manifest_top_level_shape(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + assert master["date"] == "2023010100" + assert master["input_dir"] == str(input_dir) + assert master["output_dir"] == str(output_dir) + assert master["runoff_dir"] == str(runoff_dir) + + +def test_one_job_per_vpu_ensemble_pair(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + job_keys = {k for k in master if k.startswith("job_")} + assert job_keys == { + "job_101_1", "job_101_2", "job_101_52", + "job_102_1", "job_102_2", "job_102_52", + } + + +def test_jobs_ordered_hres_first(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + ens_order_101 = [ + int(k.rsplit("_", 1)[1]) for k in master if k.startswith("job_101_") + ] + # Largest ensemble number (HRES = 52) is scheduled first. + assert ens_order_101 == [52, 2, 1] + + +def test_job_entry_fields(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), + str(input_dir), + str(output_dir), + str(runoff_dir), + initialize_flows=False, + ) + + job = master["job_101_52"] + assert job["vpu"] == "101" + assert job["ensemble"] == 52 + assert job["input_dir"] == str(input_dir / "101") + assert job["output_file"] == str(output_dir / "Qout_101_52.nc") + assert job["init_flows"] is False + assert job["runoff"].endswith("52.runoff.nc") + + +def test_manifest_written_to_disk(tmp_path): + input_dir, output_dir, runoff_dir = _build_cycle(tmp_path) + + master = forecast_preprocess( + str(tmp_path), str(input_dir), str(output_dir), str(runoff_dir) + ) + + written_path = os.path.join(str(tmp_path), "forecast_run.json") + with open(written_path) as f: + assert json.load(f) == master diff --git a/tests/test_run_river_route_forecast.py b/tests/test_run_river_route_forecast.py new file mode 100644 index 0000000..7652b14 --- /dev/null +++ b/tests/test_run_river_route_forecast.py @@ -0,0 +1,41 @@ +"""Tests for _find_state_init prior-cycle Qinit lookup. + +``run_river_route_forecast`` imports ``river_route`` at module scope; the +conftest Dummy shim lets this import succeed without the real dependency. +""" + +from geoglows_ecflow.resources.run_river_route_forecast import _find_state_init + +DATE = "2023010100" # base cycle: 2023-01-01 00 UTC + + +def _touch(directory, name): + path = directory / name + path.write_text("") + return str(path) + + +def test_returns_24h_lookback_when_present(tmp_path): + # 24h before the base cycle. + expected = _touch(tmp_path, "Qinit_2022123100.parquet") + + assert _find_state_init(str(tmp_path), DATE) == expected + + +def test_falls_through_to_72h_lookback(tmp_path): + # 24h and 48h missing; only the 72h-prior file exists. + expected = _touch(tmp_path, "Qinit_2022122900.parquet") + + assert _find_state_init(str(tmp_path), DATE) == expected + + +def test_seasonal_fallback_when_no_recent_qinit(tmp_path): + _touch(tmp_path, "seasonal_qinit_b.parquet") + expected = _touch(tmp_path, "seasonal_qinit_a.parquet") + + # Seasonal candidates are sorted; the lexicographically first is chosen. + assert _find_state_init(str(tmp_path), DATE) == expected + + +def test_returns_none_when_nothing_found(tmp_path): + assert _find_state_init(str(tmp_path), DATE) is None diff --git a/tests/test_suite_definition.py b/tests/test_suite_definition.py new file mode 100644 index 0000000..dcc3e66 --- /dev/null +++ b/tests/test_suite_definition.py @@ -0,0 +1,145 @@ +"""Structural smoke tests for the GEOGloWS suite definition. + +The suite definition is built entirely in memory here (no ecflow server, no +files written to disk) and the resulting node tree is asserted directly. These +pin the structure the routing workflow depends on so that changes to the +builder are caught. + +Requires the ``ecflow`` Python bindings; the module is skipped where they are +not installed. +""" + +import pytest + +pytest.importorskip("ecflow") + +from geoglows_ecflow.workflow.comfies.config import Config # noqa: E402 +from geoglows_ecflow.workflow.builders.builder import Builder # noqa: E402 + +SUITE = "geoglows_test" + +# Families that are only attached on the 00Z cycle (the full forecast +# pipeline). The 12Z cycle builds a reduced tree without them. +GATED_FAMILIES = { + "initialize", + "vpu_list", + "nc_to_zarr", + "combine_plain_table", + "combine_forecast_warnings", + "archive_qinit", + "archive_to_aws", + "web_push", +} + + +class _DictConfigSource: + """Minimal in-memory config source, so no ``.cfg`` file is needed.""" + + def __init__(self, data): + self.name = SUITE + self.origin = "" + self.data = data + + +def _make_config(mode="prod", vpu_list=("101",)): + """A minimal in-memory ``Config`` sufficient to construct the builder.""" + data = { + "name": SUITE, + "mode": mode, + "first_date": "20230101", + "last_date": "20230102", + "first_barrier": "20230101", + "exparch": "/tmp/arch", + "workroot": "/tmp/work", + "vpu_list": list(vpu_list), + "ens_members": 51, + "mars_workers": 1, + "target": {"root": "/tmp/target"}, + "jobs": { + "root": "/tmp/jobs", + "destinations": {"default": {"name": "localhost"}}, + }, + } + return Config(_DictConfigSource(data)) + + +def build_defs(mode="prod", vpu_list=("101",)): + """Build the suite definition in memory for the given run mode. + + The returned object is the ecflow ``Defs``. Building runs the builder's + own structural validation (``defs.check()``), so a malformed tree raises + here rather than producing a bad definition. + """ + builder = Builder(_make_config(mode=mode, vpu_list=vpu_list)) + builder() # runs build() and defs.check() + return builder.defs + + +def _member_family_names(defs, cycle): + """Names of the ensemble-member families on a cycle, or None if absent.""" + node = defs.find_abs_node(f"/{SUITE}/main/{cycle}/ens/ens_members") + if node is None: + return None + return sorted(n.name() for n in node.nodes) + + +def _cycle_children(defs, cycle): + node = defs.find_abs_node(f"/{SUITE}/main/{cycle}") + return [n.name() for n in node.nodes] + + +@pytest.fixture(scope="module") +def prod_defs(): + """A production-mode suite definition, built once and treated read-only.""" + return build_defs(mode="prod") + + +def test_00_cycle_builds_one_family_per_ensemble_member(prod_defs): + expected = sorted(f"101_{m:02d}" for m in range(1, 53)) + assert _member_family_names(prod_defs, "00") == expected + + +def test_12_cycle_has_no_ensemble_members(prod_defs): + assert _member_family_names(prod_defs, "12") is None + + +def test_00_cycle_builds_the_full_pipeline(prod_defs): + assert GATED_FAMILIES <= set(_cycle_children(prod_defs, "00")) + + +def test_12_cycle_skips_the_gated_families(prod_defs): + assert _cycle_children(prod_defs, "12") == ["hres", "ens", "diss"] + + +def test_nominal_time_families_carry_their_cycle_in_emos_base(prod_defs): + assert prod_defs.find_abs_node( + f"/{SUITE}/main/00" + ).find_variable("EMOS_BASE").value() == "00" + assert prod_defs.find_abs_node( + f"/{SUITE}/main/12" + ).find_variable("EMOS_BASE").value() == "12" + + +def test_ensemble_and_archive_memory_limits(prod_defs): + ens_members = prod_defs.find_abs_node( + f"/{SUITE}/main/00/ens/ens_members" + ) + archive_qinit = prod_defs.find_abs_node( + f"/{SUITE}/main/00/archive_qinit" + ) + assert ens_members.find_variable("MEM").value() == "6000" + assert archive_qinit.find_variable("MEM").value() == "4000" + + +def test_vpu_count_scales_the_member_families(): + defs = build_defs(mode="prod", vpu_list=("101", "102")) + assert len(_member_family_names(defs, "00")) == 2 * 52 + + +def test_default_minimum_comfies_version_does_not_block_construction(): + # A builder that leaves comfies_minimum_version at its "any" default must + # still construct: the version gate treats "any" as "no minimum". + class _AnyVersionBuilder(Builder): + comfies_minimum_version = "any" + + _AnyVersionBuilder(_make_config()) diff --git a/tests/test_zarr_io.py b/tests/test_zarr_io.py new file mode 100644 index 0000000..1d4c282 --- /dev/null +++ b/tests/test_zarr_io.py @@ -0,0 +1,35 @@ +"""Tests for the shared zarr-writing helper.""" + +import os + +import numpy as np +import xarray as xr + +from geoglows_ecflow.resources.zarr_io import write_dataset_to_zarr + + +def test_write_dataset_to_zarr_round_trips_and_drops_vars(tmp_path): + data = np.arange(5 * 3).reshape(5, 3).astype(float) + ds = xr.Dataset( + { + "Q": (("time", "river_id"), data), + "lat": (("river_id",), [10.0, 11.0, 12.0]), + }, + coords={"time": range(5), "river_id": [1, 2, 3]}, + ) + zarr_path = str(tmp_path / "out.zarr") + + write_dataset_to_zarr( + ds, + zarr_path, + {"time": -1, "river_id": "auto"}, + drop_vars=["lat"], + mode="w", + zarr_version=2, + ) + + # consolidated=True writes a .zmetadata file at the store root. + assert os.path.exists(os.path.join(zarr_path, ".zmetadata")) + with xr.open_zarr(zarr_path) as result: + assert "lat" not in result.variables + assert result["Q"].values.tolist() == data.tolist()