diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..eaf20eec4e --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +See [AGENTS.md](AGENTS.md) for all repository instructions. diff --git a/docs/developers_guide/e3sm/init/api.md b/docs/developers_guide/e3sm/init/api.md index 33d06f179e..17673e1563 100644 --- a/docs/developers_guide/e3sm/init/api.md +++ b/docs/developers_guide/e3sm/init/api.md @@ -115,3 +115,4 @@ CullTopoTask ``` + diff --git a/docs/developers_guide/e3sm/init/tasks/index.md b/docs/developers_guide/e3sm/init/tasks/index.md index edd507ecb4..3a4cbcdba0 100644 --- a/docs/developers_guide/e3sm/init/tasks/index.md +++ b/docs/developers_guide/e3sm/init/tasks/index.md @@ -21,3 +21,4 @@ topo/combine topo/remap topo/cull ``` + diff --git a/docs/developers_guide/e3sm/init/tasks/topo/cull.md b/docs/developers_guide/e3sm/init/tasks/topo/cull.md index 52a644e707..52209f33e6 100644 --- a/docs/developers_guide/e3sm/init/tasks/topo/cull.md +++ b/docs/developers_guide/e3sm/init/tasks/topo/cull.md @@ -77,7 +77,7 @@ steps, config = get_default_cull_topo_steps( unsmoothed_topo_step=unsmoothed_topo_step, include_viz=False, ) -for step in steps: +for step in steps.values(): component.add_step(step) ``` diff --git a/docs/developers_guide/e3sm/init/tasks/topo/remap.md b/docs/developers_guide/e3sm/init/tasks/topo/remap.md index b5cda4050e..4a01cb5e7a 100644 --- a/docs/developers_guide/e3sm/init/tasks/topo/remap.md +++ b/docs/developers_guide/e3sm/init/tasks/topo/remap.md @@ -65,7 +65,7 @@ steps, config = get_default_remap_topo_steps( smoothing=True, include_viz=True, ) -for step in steps: +for step in steps.values(): component.add_step(step) ``` diff --git a/docs/developers_guide/mesh/api.md b/docs/developers_guide/mesh/api.md index 880b782c0c..f587929550 100644 --- a/docs/developers_guide/mesh/api.md +++ b/docs/developers_guide/mesh/api.md @@ -23,3 +23,52 @@ BaseMeshTask ``` + +### Unified Coastline Tasks + +```{eval-rst} +.. currentmodule:: polaris.tasks.mesh.spherical.unified.coastline + +.. autosummary:: + :toctree: generated/ + + CONVENTIONS + PrepareCoastlineStep + PrepareCoastlineStep.setup + PrepareCoastlineStep.run + build_coastline_datasets + build_coastline_dataset + get_lat_lon_coastline_steps + LatLonCoastlineTask + add_coastline_tasks + VizCoastlineStep + VizCoastlineStep.setup + VizCoastlineStep.run +``` + +### Unified River Tasks + +```{eval-rst} +.. currentmodule:: polaris.tasks.mesh.spherical.unified.river + +.. autosummary:: + :toctree: generated/ + + PrepareRiverSourceStep + PrepareRiverSourceStep.setup + PrepareRiverSourceStep.run + simplify_river_network_feature_collection + PrepareRiverLatLonStep + PrepareRiverLatLonStep.setup + PrepareRiverLatLonStep.run + build_river_network_dataset + PrepareRiverForBaseMeshStep + PrepareRiverForBaseMeshStep.setup + PrepareRiverForBaseMeshStep.run + get_mesh_unified_river_steps + VizRiverStep + VizRiverStep.setup + VizRiverStep.run + UnifiedRiverNetworkTask + add_river_tasks +``` diff --git a/docs/developers_guide/mesh/tasks/base.md b/docs/developers_guide/mesh/tasks/base.md deleted file mode 100644 index 323b506cdd..0000000000 --- a/docs/developers_guide/mesh/tasks/base.md +++ /dev/null @@ -1,10 +0,0 @@ -(dev-mesh-base)= - -# Base Mesh Task - -The the base mesh tasks defined by {py:class}`polaris.tasks.mesh.base.BaseMeshTask` -can be used to generate quasi-uniform (`qu`) and subdivided icosahedral -(`icos`) spherical meshes at quasi-uniform resolutions ranging from 30 to 480 -km. These "base" meshes cover the full sphere (include both continents and -ocean regions). - diff --git a/docs/developers_guide/mesh/tasks/index.md b/docs/developers_guide/mesh/tasks/index.md index c197989aca..8ca6da249c 100644 --- a/docs/developers_guide/mesh/tasks/index.md +++ b/docs/developers_guide/mesh/tasks/index.md @@ -2,8 +2,26 @@ # Tasks +(dev-mesh-base-mesh-task)= + +## Base Mesh Task + +The the base mesh tasks defined by {py:class}`polaris.tasks.mesh.base.BaseMeshTask` +can be used to generate quasi-uniform (`qu`) and subdivided icosahedral +(`icos`) spherical meshes at quasi-uniform resolutions ranging from 30 to 480 +km. These "base" meshes cover the full sphere (include both continents and +ocean regions). + +## Unified Mesh Preparation Tasks + +This section covers the shared task families that prepare reusable coastline +and river products for named global unified meshes. These workflows are +intended to be reused by downstream sizing-field and base-mesh tasks rather +than implemented separately inside each consumer. + ```{toctree} :titlesonly: true -base +unified/coastline +unified/river ``` diff --git a/docs/developers_guide/mesh/tasks/unified/coastline.md b/docs/developers_guide/mesh/tasks/unified/coastline.md new file mode 100644 index 0000000000..b76dc4818c --- /dev/null +++ b/docs/developers_guide/mesh/tasks/unified/coastline.md @@ -0,0 +1,191 @@ +(dev-mesh-unified-coastline)= + +# Coastline Steps and Tasks + +This section of the Developer's guide covers the code for the coastline +extraction workflow. The {ref}`users-mesh-unified-coastline` section describes +the user-facing aspects of the workflow such as the conventions, +algorithms, and config options. + +The `polaris.tasks.mesh.spherical.unified.coastline` module provides a +raster-first coastline workflow on latitude-longitude grids. It reuses the +shared `e3sm/init/topo/combine` workflow (see +{ref}`dev-e3sm-init-topo-combine-tasks`) to create `topography.nc`, derives +convention-specific ocean masks from that dataset, and then computes +signed distance to the nearest raster coastline sample. + +## Available tasks + +The helper +{py:func}`polaris.tasks.mesh.spherical.unified.coastline.add_coastline_tasks` +registers one standalone task for each supported latitude-longitude target +grid: + +- `mesh/spherical/unified/coastline/lat_lon/0.25000_degree/task` +- `mesh/spherical/unified/coastline/lat_lon/0.12500_degree/task` +- `mesh/spherical/unified/coastline/lat_lon/0.06250_degree/task` +- `mesh/spherical/unified/coastline/lat_lon/0.03125_degree/task` + +Each task is an instance of +{py:class}`polaris.tasks.mesh.spherical.unified.coastline.LatLonCoastlineTask`. + +Current testing suggests `0.25000_degree` is useful mainly as a cheaper +inspection tier and is too coarse for scientifically valid coastline products. +The finer `0.12500_degree`, `0.06250_degree`, and `0.03125_degree` tasks are +the supported science-oriented coastline tiers. + +If a developer wants to add more resolutions, they can edit the list of +resolutions in `add_coastline_tasks()`. It would be best to also add +that resolution to the tasks for creating combined topography in +`polaris.tasks.e3sm.init.topo.resolutions.LAT_LON_RESOLUTIONS`, though +that isn't strictly required. + +## Task structure and shared steps + +`LatLonCoastlineTask` first requests the shared latitude-longitude combine +step from the `e3sm/init` component (see {ref}`dev-e3sm-init-topo-combine-tasks`) with +{py:func}`polaris.tasks.e3sm.init.topo.combine.steps.get_lat_lon_topo_steps` +(with `include_viz=False`) and adds that step to the task as `combine_topo`. It then calls +{py:func}`polaris.tasks.mesh.spherical.unified.coastline.get_lat_lon_coastline_steps` +with `include_viz=True`, links the shared `coastline.cfg`, and adds: + +- a shared + {py:class}`polaris.tasks.mesh.spherical.unified.coastline.PrepareCoastlineStep` in + `spherical/unified/coastline/lat_lon//prepare`; and +- a shared + {py:class}`polaris.tasks.mesh.spherical.unified.coastline.VizCoastlineStep` + in `spherical/unified/coastline/lat_lon//prepare/viz`. + +This keeps the expensive combined-topography step shareable across tasks while +still making the coastline products and diagnostics visible from the task work +directory. + +## Implementation map + +`PrepareCoastlineStep.run()` reads `coastline.cfg`, opens `topography.nc`, and +dispatches the numerical work to +{py:func}`polaris.tasks.mesh.spherical.unified.coastline.prepare.build_coastline_datasets`. +That helper operates on `(lat, lon)` NumPy arrays and, for each convention: + +- builds a candidate ocean mask from `base_elevation`, `ice_mask`, and + `grounded_mask`; +- calls + {py:func}`polaris.tasks.mesh.spherical.unified.coastline.prepare._flood_fill_ocean` + to retain only the connected ocean; +- calls + {py:func}`polaris.tasks.mesh.spherical.unified.coastline.prepare._coastline_edges` + to derive the transient raster edge diagnostics used for coastline + sampling; and +- calls + {py:func}`polaris.tasks.mesh.spherical.unified.coastline.prepare._signed_distance_from_mask` + to compute the convention-specific signed-distance field. + +The results are assembled by +{py:func}`polaris.tasks.mesh.spherical.unified.coastline.prepare._build_single_coastline_dataset` +into one raster dataset per convention, with metadata that records the +thresholds, flood-fill seed strategy, and sign convention. The detailed +algorithm is documented in the {ref}`users-mesh-unified-coastline` page so that the +developer guide can stay focused on code organization and extension points. + +The output contract is entirely raster-based. The workflow does not currently +write a persisted vector coastline product or a Natural Earth fallback. + +## Visualization + +`VizCoastlineStep` reads the coastline files and writes: + +- global and Antarctic binary plots of the final `ocean_mask` for each + convention; and +- global and Antarctic signed-distance plots for each convention. + +It also writes `debug_summary.txt`, which records convention-specific counts +such as ocean and land cells and the min/max signed distance. + +## Extension points + +The current implementation is intentionally narrow: lat-lon target grids, +three coastline conventions, and raster outputs. Common developer changes are: + +- adding a new target resolution by updating `add_coastline_tasks()` and, + ideally, the matching shared `e3sm/init/topo/combine` tasks (see + {ref}`dev-e3sm-init-topo-combine-tasks`); +- adding a new convention by extending `CONVENTIONS`, the `candidate_masks` + dictionary in `build_coastline_datasets()`, the visualization step, and the + associated tests; and +- extending the output contract by updating + `_build_single_coastline_dataset()`, the plotting code in `viz.py`, and the + tests that validate dataset contents and behavior. + +## Configuration plumbing + +The coastline workflow is configured through `coastline.cfg`. + +`PrepareCoastlineStep` consumes `[coastline]` options: + +- `resolution_latlon` +- `include_critical_transects` +- `mask_threshold` +- `sea_level_elevation` +- `distance_chunk_size` + +`VizCoastlineStep` consumes `[viz_coastline]` options: + +- `antarctic_max_latitude` +- `dpi` +- `signed_distance_limit` + +The {ref}`users-mesh-unified-coastline` page explains how these options affect the +user-visible behavior of the workflow. + +The reusable coastline-building functions now live in +`polaris.mesh.spherical.coastline`, and +`PrepareCoastlineStep` is responsible only for task-level configuration, +shared critical-transect loading, and writing outputs. Shared default +critical transects are loaded from +`polaris.mesh.spherical.critical_transects` and, when enabled, are rasterized +onto the target lat-lon grid before flood fill. + +## Example usage + +To register the default coastline tasks for the `mesh` component: + +```python +from polaris.tasks.mesh.spherical.unified.coastline import \ + add_coastline_tasks + +add_coastline_tasks(component) +``` + +To reuse the shared coastline steps inside another workflow: + +```python +from polaris.tasks.e3sm.init import e3sm_init +from polaris.tasks.e3sm.init.topo.combine.steps import get_lat_lon_topo_steps +from polaris.tasks.mesh.spherical.unified.coastline import \ + get_lat_lon_coastline_steps + +from polaris.tasks.e3sm.init.topo.combine.step import CombineStep + +combine_steps, _ = get_lat_lon_topo_steps( + component=e3sm_init, + resolution=0.25, + include_viz=False, +) +combine_topo_step = combine_steps[CombineStep.get_name()] + +coastline_steps, config = get_lat_lon_coastline_steps( + component=component, + combine_topo_step=combine_topo_step, + resolution=0.25, + include_viz=True, +) +coastline_step = coastline_steps['coastline'] +``` + +## Test coverage + +Unit tests in `tests/mesh/spherical/unified/test_coastline.py` currently +validate the public coastline dataset contract, convention-specific Antarctic +behavior, exclusion of disconnected inland water, and the use of the +northernmost latitude row for flood-fill seeding. There is not yet a +task-level smoke test for the full standalone workflow. \ No newline at end of file diff --git a/docs/developers_guide/mesh/tasks/unified/river.md b/docs/developers_guide/mesh/tasks/unified/river.md new file mode 100644 index 0000000000..45084d4cb4 --- /dev/null +++ b/docs/developers_guide/mesh/tasks/unified/river.md @@ -0,0 +1,191 @@ +(dev-mesh-unified-river)= + +# River-Network Steps and Tasks + +This section of the Developer's guide covers the code for the unified-mesh +river workflow. The {ref}`users-mesh-unified-river` section describes the +user-facing behavior, configuration options, algorithms, and output products. + +The `polaris.tasks.mesh.spherical.unified.river` package separates the river +workflow into source-level simplification, target-grid rasterization, optional +diagnostics, and downstream base-mesh conditioning. The package makes available +functions that return a dict of shared steps and thin task wrappers so the same +implementation can be reused by standalone inspection tasks and by downstream +unified-mesh tasks. + +## Available tasks + +The helper +{py:func}`polaris.tasks.mesh.spherical.unified.river.add_river_tasks` +registers one standalone task for each mesh name in +`polaris.mesh.spherical.unified.UNIFIED_MESH_NAMES`: + +- {py:class}`polaris.tasks.mesh.spherical.unified.river.UnifiedRiverNetworkTask` at + `mesh/spherical/unified//river/task`. + +{py:func}`polaris.tasks.mesh.add_tasks.add_mesh_tasks` wires these tasks into +the `mesh` component after the generic base-mesh tasks (see +{ref}`dev-mesh-base-mesh-task`) and the shared coastline tasks (see +{ref}`dev-mesh-unified-coastline`). + +## Task structure and shared steps + +`UnifiedRiverNetworkTask` layers shared dependencies in a fixed order: + +1. It requests the shared lat-lon topography step (see + {ref}`dev-e3sm-init-topo-combine-tasks`) from + {py:func}`polaris.tasks.e3sm.init.topo.combine.steps.get_lat_lon_topo_steps` + and exposes it as `combine_topo`. +2. It requests the shared coastline steps (see {ref}`dev-mesh-unified-coastline`) + from + {py:func}`polaris.tasks.mesh.spherical.unified.coastline.get_lat_lon_coastline_steps` + and exposes them as `prepare_coastline` and the optional + `viz_prepare_coastline`. +3. It requests the full set of shared river steps from + {py:func}`polaris.tasks.mesh.spherical.unified.river.get_mesh_unified_river_steps` + and exposes {py:class}`polaris.tasks.mesh.spherical.unified.river.PrepareRiverSourceStep` + as `prepare_source`, + {py:class}`polaris.tasks.mesh.spherical.unified.river.PrepareRiverLatLonStep` + as `prepare_lat_lon`, + {py:class}`polaris.tasks.mesh.spherical.unified.river.VizRiverStep` as + `viz_lat_lon`, and + {py:class}`polaris.tasks.mesh.spherical.unified.river.PrepareRiverForBaseMeshStep` + as `prepare_base_mesh`. + +This ordering keeps the river implementation dependent only on the explicit +products it consumes: simplified HydroRIVERS geometry and the selected +coastline dataset. + +The shared-step factory in `steps.py` is the main extension point for other +task families: + +- {py:func}`polaris.tasks.mesh.spherical.unified.river.get_mesh_unified_river_steps` + builds or reuses the full chain of shared river steps — + {py:class}`polaris.tasks.mesh.spherical.unified.river.PrepareRiverSourceStep`, + {py:class}`polaris.tasks.mesh.spherical.unified.river.PrepareRiverLatLonStep`, + optional {py:class}`polaris.tasks.mesh.spherical.unified.river.VizRiverStep`, + and {py:class}`polaris.tasks.mesh.spherical.unified.river.PrepareRiverForBaseMeshStep` + — returning them as a dict keyed by step name along with the lat-lon mesh + config. + +## Implementation map + +### `source.py` + +`PrepareRiverSourceStep.setup()` registers the HydroRIVERS archive as an input +file through `add_input_file()` using the public URL from `river_network.cfg`. +`PrepareRiverSourceStep.run()` unpacks the archive, converts the shapefile to +GeoJSON, simplifies the source network, and writes: + +- `source_river_network.geojson` +- `simplified_river_network.geojson` +- `retained_outlets.geojson` + +The public helper +{py:func}`polaris.tasks.mesh.spherical.unified.river.simplify_river_network_feature_collection` +contains the source-level retention logic. It builds canonical segments, +validates downstream topology, filters outlet candidates, and traverses +retained basins while preserving main stems and significant tributaries. + +The internal `RiverSegment` dataclass is the canonical representation used by +the simplification helpers. It is intentionally not exported as part of the +public API. + +### `lat_lon.py` + +`PrepareRiverLatLonStep.setup()` links the simplified source products and the +selected coastline NetCDF file. `PrepareRiverLatLonStep.run()` reads those +inputs, calls +{py:func}`polaris.tasks.mesh.spherical.unified.river.build_river_network_dataset`, +and writes: + +- `river_network.nc` +- `river_outlets.geojson` + +`build_river_network_dataset()` is the public target-grid helper. It rasterizes +river channels, snaps ocean outlets against coastline ocean cells, snaps inland +sinks to land cells, and writes a clean mask split between channels, all +outlets, ocean outlets, and inland sinks. + +### `base_mesh.py` + +`PrepareRiverForBaseMeshStep` is produced through `get_mesh_unified_river_steps()` +and is included in `UnifiedRiverNetworkTask` as well as downstream unified +base-mesh consumers. + +Its `run()` method reads the simplified river network and coastline products, +then conditions the retained geometry for direct base-mesh use by: + +- clipping segments inland of the selected coastline by the configured clip + distance; +- simplifying clipped geometry and removing degenerate or too-short pieces; +- dropping outlets that do not remain inland after clipping; and +- regenerating a diagnostic lat-lon mask product from the conditioned river + geometry. + +This step writes `clipped_river_network.geojson`, `clipped_outlets.geojson`, +and `clipped_river_network.nc`. + +### `viz.py` + +`VizRiverStep` is a pure diagnostic consumer of the shared source, coastline, +and lat-lon river products. It writes `river_network_overview.png` and +`debug_summary.txt`, and keeps visualization logic out of the numerical steps. + +## Configuration plumbing + +All shared river steps use mesh-specific configs built through +`polaris.mesh.spherical.unified.get_unified_mesh_config()`. That loader +combines: + +- the generic `unified_mesh.cfg` defaults; +- the shared `river_network.cfg` file; and +- the selected named-mesh config file. + +`PrepareRiverSourceStep` consumes `[river_network]` options. + +`PrepareRiverLatLonStep` consumes `[river_lat_lon]` options and also reads the +selected `unified_mesh` settings such as `mesh_name`, `resolution_latlon`, +and `coastline_convention`. + +`PrepareRiverForBaseMeshStep` consumes `[river_network]` and `unified_mesh` +settings, and currently reads the Antarctic coastline convention from the +`[spherical_mesh]` section when selecting the coastline product for clipping. + +`VizRiverStep` consumes `[viz_river_network].dpi`. + +## Extension points + +Common extension paths for future development are: + +- adding a new named unified mesh by creating a new config file under + `polaris.mesh.spherical.unified`; task registration discovers mesh names + from those config files automatically; +- extending source-level metadata or retention rules in + `simplify_river_network_feature_collection()` and the associated GeoJSON + property builders; +- extending the target-grid output contract in + `build_river_network_dataset()` and the visualization step; and +- adjusting downstream coastline-aware conditioning in + `PrepareRiverForBaseMeshStep` without creating a separate river-processing + code path for base meshes. + +The public API in this package is intentionally narrow: task wrappers, +shared-step factories, step classes, and the two reusable dataset-building +helpers. Most geometry and graph utilities remain private so they can evolve +without breaking downstream callers. + +## Test coverage + +Unit tests in `tests/mesh/spherical/unified/test_river.py` currently cover: + +- source-level outlet filtering, deep main-stem traversal, tributary + retention, and cycle detection; +- HydroRIVERS archive unpacking and shapefile-to-GeoJSON conversion helpers; +- target-grid raster and snapped-outlet contracts; +- coastline-aware base-mesh conditioning helpers and shared-step factories; +- and task registration for all named unified meshes. + +There is not yet a full task-level integration test that runs the end-to-end +river workflow on real HydroRIVERS data inside the documentation build or unit +test suite. \ No newline at end of file diff --git a/docs/users_guide/mesh/tasks/base.md b/docs/users_guide/mesh/tasks/base.md deleted file mode 100644 index 15bb347a40..0000000000 --- a/docs/users_guide/mesh/tasks/base.md +++ /dev/null @@ -1,30 +0,0 @@ -(mesh-base-mesh-task)= - -# Base Mesh Tasks - -The `mesh/spherical/qu/base_mesh/XXXkm/task` and -`mesh/spherical/icos/base_mesh/XXXkm/task` tasks can be used to create -spherical, uniform MPAS meshes at the given resolution that are either -quasi-uniform (`qu`) or subdivided icosahedral (`icos`). - -There are also two variable resolution base meshes: `rrs6to18km` and -`so12to30km`. - -The RRS mesh has a resolution that ranges from 6 km at the -poles to 18 km at the equator, approximately scaling with the Rossby radius. - -```{image} ../../../developers_guide/mesh/framework/images/rrs.png -:align: center -:width: 500 px -``` - -The SO mesh has a quasi-uniform, 30 km background resolution that transitions -to 12 km in region surrounding the Southern Ocean. The boundary of high -resolution region attempts to approximatly follow dynamical contours in -the ocean, since rapid changes in resolution have been shown to steer ocean -currents along isocontours of resolution. - -```{image} ../../../developers_guide/mesh/framework/images/so.png -:align: center -:width: 500 px -``` \ No newline at end of file diff --git a/docs/users_guide/mesh/tasks/index.md b/docs/users_guide/mesh/tasks/index.md index d69275eb53..b27dc1fec7 100644 --- a/docs/users_guide/mesh/tasks/index.md +++ b/docs/users_guide/mesh/tasks/index.md @@ -2,8 +2,47 @@ # Tasks +(mesh-base-mesh-task)= + +## Base Mesh Tasks + +The `mesh/spherical/qu/base_mesh/XXXkm/task` and +`mesh/spherical/icos/base_mesh/XXXkm/task` tasks can be used to create +spherical, uniform MPAS meshes at the given resolution that are either +quasi-uniform (`qu`) or subdivided icosahedral (`icos`). + +There are also two variable resolution base meshes: `rrs6to18km` and +`so12to30km`. + +The RRS mesh has a resolution that ranges from 6 km at the +poles to 18 km at the equator, approximately scaling with the Rossby radius. + +```{image} ../../../developers_guide/mesh/framework/images/rrs.png +:align: center +:width: 500 px +``` + +The SO mesh has a quasi-uniform, 30 km background resolution that transitions +to 12 km in region surrounding the Southern Ocean. The boundary of high +resolution region attempts to approximatly follow dynamical contours in +the ocean, since rapid changes in resolution have been shown to steer ocean +currents along isocontours of resolution. + +```{image} ../../../developers_guide/mesh/framework/images/so.png +:align: center +:width: 500 px +``` + +## Unified Mesh Tasks + +These tasks support the shared preprocessing workflow for named global +unified meshes. They focus on reusable products that downstream sizing-field +and base-mesh steps can consume directly instead of recalculating coastline +and river information independently. + ```{toctree} :titlesonly: true -base +unified/coastline +unified/river ``` diff --git a/docs/users_guide/mesh/tasks/unified/coastline.md b/docs/users_guide/mesh/tasks/unified/coastline.md new file mode 100644 index 0000000000..0c5d46fd96 --- /dev/null +++ b/docs/users_guide/mesh/tasks/unified/coastline.md @@ -0,0 +1,165 @@ +(users-mesh-unified-coastline)= + +# Coastline tasks + +The `mesh/spherical/unified/coastline` tasks build coastline masks and +signed-distance fields from the shared combined-topography products on +latitude-longitude grids. These tasks are intended for inspecting and caching +coastline products that later unified-mesh workflows (such as +{ref}`users-mesh-unified-river`) can reuse. + +Current coastline products are derived only from the combined topography +fields produced by the `e3sm/init/topo/combine` workflow (see +{ref}`e3sm-init-topo-tasks`). Alternate coastline sources such +as Natural Earth are not implemented yet. + +## Available tasks + +Polaris currently provides one standalone coastline task for each supported +latitude-longitude target grid: + +- `mesh/spherical/unified/coastline/lat_lon/0.25000_degree/task` +- `mesh/spherical/unified/coastline/lat_lon/0.12500_degree/task` +- `mesh/spherical/unified/coastline/lat_lon/0.06250_degree/task` +- `mesh/spherical/unified/coastline/lat_lon/0.03125_degree/task` + +Current testing suggests `0.25000_degree` is still useful for lower-cost +inspection but is too coarse for scientifically valid coastline products. +Prefer `0.12500_degree`, `0.06250_degree`, or `0.03125_degree` when coastline +fidelity matters. + +Each task includes: + +- a shared `combine_topo` step from `e3sm/init/topo/combine` (see + {ref}`e3sm-init-topo-tasks`) that creates `topography.nc` on the selected + target grid; +- a shared `prepare` step that writes one coastline NetCDF file for each + supported coastline convention; and +- a shared `viz` step that writes diagnostic PNG images and a text summary. + +## Coastline conventions + +Each run writes three convention-specific NetCDF files: + +- `coastline_calving_front.nc` +- `coastline_grounding_line.nc` +- `coastline_bedrock_zero.nc` + +These correspond to three ways of handling Antarctic floating and grounded +ice when defining the ocean mask: + +- `calving_front`: cells below sea level are ocean only where `ice_mask` + indicates they are ice free; +- `grounding_line`: cells below sea level are ocean where `grounded_mask` + indicates they are not grounded ice; and +- `bedrock_zero`: all cells below sea level are candidate ocean. + +## How the coastline is derived + +For each convention, Polaris builds the coastline fields in four stages: + +1. It classifies each raster cell as candidate ocean or not candidate ocean + using `base_elevation`, `ice_mask`, and `grounded_mask`. +2. It keeps only the connected ocean by labeling four-neighbor connected + components on the latitude-longitude grid, merging components that meet + across the antimeridian, and retaining only components that touch the + northernmost latitude row. +3. It marks coastline transitions on east and north cell edges wherever the + final ocean mask changes from ocean to land. +4. It computes signed distance to the nearest raster coastline sample from + those edge transitions. + +The flood-fill step is what removes disconnected inland basins. A cell can be +below sea level and still be classified as land if it belongs to a component +that never connects to the northern boundary of the grid through candidate +ocean cells. + +The raster flood fill uses four-neighbor connectivity rather than diagonal +connectivity. North-south and east-west neighbors count as connected, but +cells that touch only at a corner do not. Longitudinal periodicity is handled +explicitly, so ocean that crosses the antimeridian remains connected. + +The signed-distance calculation is also raster-based. Polaris places sample +points at the midpoints of coastline edges, converts both those samples and +all grid-cell centers from longitude/latitude to Cartesian coordinates on the +sphere, and builds a `scipy.spatial.cKDTree` from the coastline samples. The +nearest-neighbor search is performed in latitude chunks controlled by +`distance_chunk_size`, then the Cartesian chord distance is converted to a +spherical arc distance in meters. The result is positive in ocean cells and +negative in land cells. + +Because the distance is measured to raster coastline samples rather than to an +exact vector shoreline, the answer is tied to the chosen latitude-longitude +resolution. Finer grids produce a more detailed coastline and a more accurate +approximation to the continuous shoreline. + +## Output fields + +For each convention, the coastline file contains: + +- `ocean_mask` +- `signed_distance` + +These fields have the following interpretation: + +- `ocean_mask`: candidate ocean cells that remain after the flood fill; +- `signed_distance`: signed nearest-coastline distance in meters, negative + over land and positive over ocean. + +## Configuration + +Each task links a shared config file named `coastline.cfg`. The main options +are: + +- `[coastline].resolution_latlon`: the target latitude-longitude + resolution in degrees. The task sets this automatically from the selected + task path. +- `[coastline].include_critical_transects`: whether to apply the + shared `geometric_features` critical land blockages and critical passages + before flood filling. The default is `True`. +- `[coastline].mask_threshold`: threshold for converting remapped + `ice_mask` and `grounded_mask` fields to binary masks. The default is + `0.5`. +- `[coastline].sea_level_elevation`: elevation threshold for + identifying below-sea-level cells. The default is `0.0` m. +- `[coastline].distance_chunk_size`: number of latitude rows + processed at a time when computing signed distance. This changes memory use + and query batching, not the definition of the distance. The default is + `64`. +- `[viz_coastline].antarctic_max_latitude`: northern extent of + Antarctic stereographic plots. The default is `-45.0` degrees. +- `[viz_coastline].dpi`: output resolution for diagnostic plots. The + default is `200`. +- `[viz_coastline].signed_distance_limit`: symmetric colorbar limit + for signed-distance plots. The default is `500000.0` m. + +## Diagnostics + +The visualization step writes global and Antarctic binary plots of the final +`ocean_mask` for each convention, along with matching signed-distance plots. +This keeps the visualization focused on the downstream-relevant land/ocean +classification and coastal-distance product while avoiding the slower +debug-only plots at high target-grid resolutions. + +When critical transects are enabled, the flood fill honors the same shared +critical land blockages and critical passages from `geometric_features` that +are used in E3SM topography culling (see {ref}`e3sm-init-topo-tasks`). This changes connectivity in the final +`ocean_mask` and derived `signed_distance` field without changing the output +schema. + +The file `debug_summary.txt` records convention-specific counts such as the +number of ocean and land cells, along with the minimum and maximum signed +distance. + +## Running a task + +You can set up one of the coastline tasks with standard polaris commands, for +example: + +```bash +polaris setup -t mesh/spherical/unified/coastline/lat_lon/0.12500_degree/task \ + -w coastline_0125 +``` + +After setup, the work directory contains symlinks to the shared `combine_topo`, +`prepare`, and `viz` steps, along with the shared `coastline.cfg` file. \ No newline at end of file diff --git a/docs/users_guide/mesh/tasks/unified/river.md b/docs/users_guide/mesh/tasks/unified/river.md new file mode 100644 index 0000000000..bcacb06602 --- /dev/null +++ b/docs/users_guide/mesh/tasks/unified/river.md @@ -0,0 +1,193 @@ +(users-mesh-unified-river)= + +# River-network tasks + +The unified-mesh river-network tasks simplify HydroRIVERS source data, +reconcile retained outlets with the selected coastline product, and build +target-grid masks that downstream sizing-field and base-mesh workflows (see +{ref}`mesh-base-mesh-task`) can consume directly. + +The standalone tasks expose two layers of the workflow: + +- a source-level task for performing HydroRIVERS conversion, outlet + retention, and basin-aware simplification; and +- a latitude-longitude task for performing coastline-aware outlet snapping, + target-grid masks, and diagnostic plots. + +The downstream base-mesh workflow (see {ref}`mesh-base-mesh-task`) reuses +these same shared steps and then produces coastline-clipped river products for +final mesh generation. + +## Available tasks + +Polaris registers two standalone river tasks for each named unified mesh: + +- `mesh/spherical/unified//river/source/task` +- `mesh/spherical/unified//river/lat_lon/task` + +Supported `mesh_name` values are: + +- `ocn_240km_lnd_240km_riv_240km` +- `ocn_30km_lnd_10km_riv_10km` +- `ocn_rrs_6to18km_lnd_12km_riv_6km` +- `ocn_so_12to30km_lnd_10km_riv_10km` + +Each source task contains the shared `prepare` step that downloads and +simplifies HydroRIVERS. + +Each lat-lon task contains: + +- `prepare_source`, the same shared source-level river step; +- `combine_topo`, the shared `e3sm/init/topo/combine` step (see + {ref}`e3sm-init-topo-tasks`) for the mesh's target latitude-longitude grid; +- `prepare_coastline`, the shared coastline step + (see {ref}`users-mesh-unified-coastline`) for the selected target grid and coastline + convention; +- `viz_prepare_coastline`, an optional coastline diagnostic step that is not + run by default; +- `prepare`, the shared lat-lon river step that writes raster and outlet + products; and +- `viz_river_network`, a diagnostic step that writes an overview plot and + summary text file. + +## How the river network is derived + +The source-level workflow follows a staged simplification strategy: + +1. It downloads the HydroRIVERS archive and converts the shapefile into + `source_river_network.geojson`. +2. It canonicalizes source features into one segment per `hyriv_id` when the + source contains multiple geometries for the same river segment. +3. It filters segments by `drainage_area_threshold`. +4. It validates that the retained `next_down` graph is acyclic. +5. It identifies outlet candidates from segments with `next_down == 0`, then + keeps large, well-separated ocean outlets while retaining explicit inland + sinks. +6. It traverses upstream from each retained outlet to keep main stems and + significant tributaries, preserving basin connectivity and confluence + structure. + +The lat-lon workflow then turns the simplified network into target-grid +products: + +1. It samples the retained river lines densely enough to rasterize them onto + the chosen latitude-longitude grid. +2. It snaps ocean outlets to nearby ocean cells from `coastline.nc` when a + match is available within `outlet_match_tolerance`. +3. It snaps inland sinks to the nearest land cell derived from the coastline + `ocean_mask`. +4. It writes separate masks for channels, all outlets, ocean outlets, and + inland sinks instead of overloading one raster with mixed semantics. + +If an ocean outlet cannot be matched to ocean within tolerance, it is still +snapped to the nearest grid cell, but its `matched_to_ocean` flag remains +`false` and the snapping distance is recorded for diagnostics. + +## Outputs + +The source task writes: + +- `source_river_network.geojson`, the direct conversion of the HydroRIVERS + shapefile; +- `simplified_river_network.geojson`, the retained river segments and their + basin and outlet metadata; and +- `retained_outlets.geojson`, one feature per retained ocean outlet or inland + sink. + +The lat-lon task writes: + +- `river_network.nc`, containing `river_channel_mask`, `river_outlet_mask`, + `river_ocean_outlet_mask`, and `river_inland_sink_mask`; +- `river_outlets.geojson`, containing source and snapped outlet coordinates, + snapped grid indices, snapping distance, and `matched_to_ocean`; and +- `river_network_overview.png` and `debug_summary.txt` from the visualization + step. + +The downstream shared base-mesh river step (see {ref}`mesh-base-mesh-task`) +also writes: + +- `clipped_river_network.geojson`; +- `clipped_outlets.geojson`; +- and `clipped_river_network.nc`. + +These clipped products are intended for the unified base-mesh workflow rather +than the standalone inspection tasks. + +## Configuration + +Each task links the shared `river_network.cfg` file and inherits the selected +mesh's `unified_mesh` settings, including the target-grid resolution and +coastline convention (see {ref}`users-mesh-unified-coastline`). + +The `[river_network]` section controls HydroRIVERS access and source-level +simplification: + +- `hydrorivers_url`: the public HydroRIVERS archive URL. +- `hydrorivers_archive_filename`: local filename used for the downloaded + archive. +- `hydrorivers_shp_directory`: expected directory name after unpacking the + archive. +- `hydrorivers_shp_filename`: HydroRIVERS shapefile basename inside the + unpacked archive. +- `drainage_area_threshold`: minimum drainage area in square meters for + retaining a source segment. +- `outlet_distance_tolerance`: minimum spacing in meters between retained + non-endorheic outlets. +- `tributary_area_ratio`: minimum tributary-to-main-stem drainage-area ratio + for retaining a nearby tributary. +- `base_mesh_clip_distance_km`: inland clip distance used later when + preparing base-mesh river products. +- `base_mesh_simplify_tolerance_km`: geometry simplification tolerance used + for downstream base-mesh products. +- `base_mesh_min_segment_length_km`: minimum retained segment length after + downstream base-mesh clipping. +- `base_mesh_preserve_outlet_stub_km`: optional outlet-stub length preserved + during downstream base-mesh clipping. + +The `[river_lat_lon]` section controls the target-grid products: + +- `outlet_match_tolerance`: maximum snapping distance in meters when matching + an ocean outlet to a coastline ocean cell. +- `channel_subsegment_fraction`: fraction of one target-grid cell used to set + line-sampling density during rasterization. +- `channel_buffer_km`: optional physical buffer around sampled river-channel + points. A value of `0` preserves one-cell-wide rasterization. + +The `[viz_river_network]` section currently contains: + +- `dpi`: output resolution for the river diagnostic plot. + +## Diagnostics and expected behavior + +`river_network_overview.png` overlays the simplified river network and snapped +outlets on the selected coastline background, with a global view and a CONUS +inset. `debug_summary.txt` records counts of retained segments, retained +outlets, matched and unmatched ocean outlets, inland sinks, and raster cells +flagged in each mask. + +The lat-lon NetCDF product also records target-grid metadata such as +`target_grid_resolution_degrees`, `outlet_match_tolerance_m`, +`channel_buffer_m`, `matched_ocean_outlets`, and +`unmatched_ocean_outlets`. + +## Running a task + +To inspect source-level simplification for one named unified mesh: + +```bash +polaris setup -t \ + mesh/spherical/unified/ocn_30km_lnd_10km_riv_10km/river/source/task \ + -w river_source_30km +``` + +To inspect the coastline-aware target-grid products for the same mesh: + +```bash +polaris setup -t \ + mesh/spherical/unified/ocn_30km_lnd_10km_riv_10km/river/lat_lon/task \ + -w river_lat_lon_30km +``` + +The work directory for the lat-lon task contains symlinks to the shared river, +topography, and coastline steps, along with the shared `river_network.cfg` +file. \ No newline at end of file diff --git a/docs/users_guide/tasks.md b/docs/users_guide/tasks.md index abd92238c5..a3412da00c 100644 --- a/docs/users_guide/tasks.md +++ b/docs/users_guide/tasks.md @@ -3,10 +3,10 @@ # Tasks Polaris currently supports tasks for the {ref}`mesh`, {ref}`e3sm-init`, -{ref}`ocean` ([MPAS-Ocean](https://mpas-dev.github.io/ocean/ocean.html)) and -{ref}`seaice` -([MPAS-Seaice](https://mpas-dev.github.io/sea_ice/sea_ice.html)) components. -Land-ice support is planned but has not yet been migrated. +{ref}`ocean` ([MPAS-Ocean](https://mpas-dev.github.io/ocean/ocean.html)), and +{ref}`seaice` ([MPAS-Seaice](https://mpas-dev.github.io/sea_ice/sea_ice.html)) +components. Land-ice support is planned but has not yet been migrated as its +own component. Tasks are grouped under these components and then into common categories for convenience and shared framework. These groupings of tasks have some common purpose or concept. For ocean, this includes "idealized" tasks like diff --git a/polaris/archive.py b/polaris/archive.py index 32b9288e42..03497324ed 100644 --- a/polaris/archive.py +++ b/polaris/archive.py @@ -3,6 +3,20 @@ import zipfile +def extract_zip_subdir(zip_filename, member_prefix, out_dir='.'): + with zipfile.ZipFile(zip_filename) as archive: + member_names = _find_zip_subdir_members(archive, member_prefix) + for member_name in member_names: + member_path = pathlib.PurePosixPath(member_name) + out_filename = pathlib.Path(out_dir).joinpath(*member_path.parts) + out_filename.parent.mkdir(parents=True, exist_ok=True) + with ( + archive.open(member_name) as src, + open(out_filename, 'wb') as dst, + ): + shutil.copyfileobj(src, dst) + + def extract_zip_member(zip_filename, member_name, out_filename): with zipfile.ZipFile(zip_filename) as archive: member_name = find_zip_member(archive, member_name) @@ -27,3 +41,19 @@ def find_zip_member(archive, member_name): f'Could not find ZIP member {member_name!r} in archive ' f'{archive.filename!r}' ) + + +def _find_zip_subdir_members(archive, member_prefix): + prefix = pathlib.PurePosixPath(member_prefix).as_posix().rstrip('/') + prefix = f'{prefix}/' + matches = [ + name + for name in archive.namelist() + if name.startswith(prefix) and not name.endswith('/') + ] + if len(matches) == 0: + raise ValueError( + f'Could not find ZIP subdirectory {member_prefix!r} in archive ' + f'{archive.filename!r}' + ) + return matches diff --git a/polaris/mesh/base/so/__init__.py b/polaris/mesh/base/so/__init__.py index 69934f354c..005f1de2e7 100644 --- a/polaris/mesh/base/so/__init__.py +++ b/polaris/mesh/base/so/__init__.py @@ -1,11 +1,7 @@ import numpy as np -from geometric_features import read_feature_collection -from mpas_tools.mesh.creation.signed_distance import ( - signed_distance_from_geojson, -) -from polaris.constants import get_constant from polaris.mesh import QuasiUniformSphericalMeshStep +from polaris.mesh.base.so.background import build_southern_ocean_background class SOBaseMesh(QuasiUniformSphericalMeshStep): @@ -19,7 +15,8 @@ def setup(self): """ self.add_input_file( - filename='high_res_region.geojson', package=self.__module__ + filename='high_res_region.geojson', + package='polaris.mesh.base.so', ) super().setup() @@ -47,29 +44,17 @@ def build_cell_width_lat_lon(self): dlon = 0.1 dlat = dlon - earth_radius = get_constant('mean_radius') nlon = int(360.0 / dlon) + 1 nlat = int(180.0 / dlat) + 1 lon = np.linspace(-180.0, 180.0, nlon) lat = np.linspace(-90.0, 90.0, nlat) - # start with a uniform max_res km background resolution - cell_width = max_res * np.ones((nlat, nlon)) - - fc = read_feature_collection('high_res_region.geojson') - - so_signed_distance = signed_distance_from_geojson( - fc, lon, lat, earth_radius, max_length=0.25 - ) - - # Equivalent to 20 degrees latitude - trans_width = 1600e3 - trans_start = 500e3 - - weights = 0.5 * ( - 1 + np.tanh((so_signed_distance - trans_start) / trans_width) + cell_width = build_southern_ocean_background( + lat=lat, + lon=lon, + high_res_km=min_res, + low_res_km=max_res, + region_filename='high_res_region.geojson', ) - cell_width = min_res * (1 - weights) + cell_width * weights - return cell_width, lon, lat diff --git a/polaris/mesh/base/so/background.py b/polaris/mesh/base/so/background.py new file mode 100644 index 0000000000..e24bfde43f --- /dev/null +++ b/polaris/mesh/base/so/background.py @@ -0,0 +1,34 @@ +import numpy as np +from geometric_features import read_feature_collection +from mpas_tools.mesh.creation.signed_distance import ( + signed_distance_from_geojson, +) + +from polaris.constants import get_constant + + +def build_southern_ocean_background( + lat, lon, high_res_km, low_res_km, region_filename +): + """ + Build a Southern Ocean background field on a regular lat-lon grid. + """ + earth_radius = get_constant('mean_radius') + fc = read_feature_collection(region_filename) + + so_signed_distance = signed_distance_from_geojson( + fc, lon, lat, earth_radius, max_length=0.25 + ) + + # Equivalent to 20 degrees latitude + transition_width_m = 1600e3 + transition_start_m = 500e3 + + weights = 0.5 * ( + 1 + + np.tanh( + (so_signed_distance - transition_start_m) / transition_width_m + ) + ) + + return high_res_km * (1 - weights) + low_res_km * weights diff --git a/polaris/mesh/spherical/coastline.py b/polaris/mesh/spherical/coastline.py new file mode 100644 index 0000000000..2ff3f4bade --- /dev/null +++ b/polaris/mesh/spherical/coastline.py @@ -0,0 +1,585 @@ +from typing import Any + +import netCDF4 +import numpy as np +import xarray as xr +from scipy import ndimage +from scipy.spatial import cKDTree + +from polaris.constants import get_constant +from polaris.mesh.spherical.critical_transects import CriticalTransects + +CONVENTIONS = ('calving_front', 'grounding_line', 'bedrock_zero') +EARTH_RADIUS = get_constant('mean_radius') + + +def build_coastline_datasets( + ds_topo, + resolution, + mask_threshold=0.5, + sea_level_elevation=0.0, + distance_chunk_size=64, + workers=1, + critical_transects=None, +): + """ + Build coastline datasets from combined topography. + + Parameters + ---------- + ds_topo : xr.Dataset + Combined topography on a lat-lon grid + + resolution : float + The target lat-lon resolution in degrees + + mask_threshold : float, optional + Threshold for converting remapped mask fields to binary masks + + sea_level_elevation : float, optional + Elevation threshold used to classify below-sea-level cells + + distance_chunk_size : int, optional + Number of latitude rows per signed-distance query chunk + + workers : int, optional + Number of workers to use in KD-tree queries + + critical_transects : CriticalTransects, optional + Critical land blockages and passages to rasterize before flood fill + + Returns + ------- + ds_coastlines : dict[str, xr.Dataset] + Coastline datasets keyed by coastline convention + """ + lon = ds_topo.lon.values + lat = ds_topo.lat.values + + base_elevation = ds_topo.base_elevation.transpose('lat', 'lon').values + ice_mask = ( + ds_topo.ice_mask.transpose('lat', 'lon').values >= mask_threshold + ) + grounded_mask = ( + ds_topo.grounded_mask.transpose('lat', 'lon').values >= mask_threshold + ) + below_sea_level = np.isfinite(base_elevation) & ( + base_elevation < sea_level_elevation + ) + + candidate_masks = { + 'calving_front': below_sea_level & np.logical_not(ice_mask), + 'grounding_line': below_sea_level & np.logical_not(grounded_mask), + 'bedrock_zero': below_sea_level, + } + + land_blockages, passages = _rasterize_critical_transects( + critical_transects=critical_transects, + lon=lon, + lat=lat, + ) + + ds_coastlines: dict[str, xr.Dataset] = {} + + for convention in CONVENTIONS: + candidate_ocean = candidate_masks[convention] + candidate_ocean = np.logical_and( + candidate_ocean, np.logical_not(land_blockages) + ) + candidate_ocean = np.logical_or(candidate_ocean, passages) + + ocean_mask = _flood_fill_ocean(candidate_ocean, lat) + edge_east, edge_north = _coastline_edges(ocean_mask) + signed_distance = _signed_distance_from_mask( + ocean_mask=ocean_mask, + edge_east=edge_east, + edge_north=edge_north, + lon=lon, + lat=lat, + chunk_size=distance_chunk_size, + workers=workers, + ) + + ds_coastlines[convention] = _build_single_coastline_dataset( + convention=convention, + ds_topo=ds_topo, + resolution=resolution, + mask_threshold=mask_threshold, + sea_level_elevation=sea_level_elevation, + ocean_mask=ocean_mask.astype(np.int8), + signed_distance=signed_distance.astype(np.float32), + ) + + return ds_coastlines + + +def build_coastline_dataset( + ds_topo, + resolution, + convention, + mask_threshold=0.5, + sea_level_elevation=0.0, + distance_chunk_size=64, + workers=1, + critical_transects=None, +): + """ + Build a coastline dataset for one coastline convention. + """ + ds_coastlines = build_coastline_datasets( + ds_topo=ds_topo, + resolution=resolution, + mask_threshold=mask_threshold, + sea_level_elevation=sea_level_elevation, + distance_chunk_size=distance_chunk_size, + workers=workers, + critical_transects=critical_transects, + ) + return ds_coastlines[convention] + + +def _build_single_coastline_dataset( + convention, + ds_topo, + resolution, + mask_threshold, + sea_level_elevation, + ocean_mask, + signed_distance, +): + """ + Build one convention-specific coastline dataset. + """ + ds_coastline = xr.Dataset( + coords=dict( + lat=ds_topo.lat, + lon=ds_topo.lon, + ) + ) + + dims = ('lat', 'lon') + data_vars: dict[str, np.ndarray[Any, Any]] = { + 'ocean_mask': ocean_mask, + 'signed_distance': signed_distance, + } + for var_name, values in data_vars.items(): + ds_coastline[var_name] = xr.DataArray(values, dims=dims) + + ds_coastline.attrs.update( + dict( + coastline_convention=convention, + target_grid='lat_lon', + target_grid_resolution_degrees=resolution, + coastline_source='combined_topography', + mask_threshold=mask_threshold, + sea_level_elevation=sea_level_elevation, + flood_fill_seed_strategy='candidate_ocean_on_northernmost_row', + sign_convention='negative_over_land_positive_over_ocean', + coastline_edge_definition=( + 'east and north cell-edge transitions on the target grid' + ), + coastline_distance_definition=( + 'spherical nearest-sample distance from raster coastline' + ), + ) + ) + + ds_coastline['ocean_mask'].attrs['long_name'] = ( + 'Exclusive ocean mask after flood fill' + ) + ds_coastline['signed_distance'].attrs['long_name'] = ( + 'Signed distance to the nearest coastline sample' + ) + ds_coastline['signed_distance'].attrs['units'] = 'm' + + return ds_coastline + + +def _flood_fill_ocean(candidate_ocean, lat): + """ + Flood fill candidate ocean regions from the northernmost latitude row. + + Periodicity is enforced in longitude by merging labels on the eastern and + western boundaries before selecting connected components. + """ + structure = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.int8) + labels, label_count = ndimage.label(candidate_ocean, structure=structure) + if label_count == 0: + return np.zeros_like(candidate_ocean, dtype=bool) + + seed_row = int(np.argmax(lat)) + + parent = np.arange(label_count + 1, dtype=np.int64) + + def find(label): + while parent[label] != label: + parent[label] = parent[parent[label]] + label = parent[label] + return label + + def union(left, right): + left_root = find(left) + right_root = find(right) + if left_root != right_root: + parent[right_root] = left_root + + wrap_rows = np.nonzero(candidate_ocean[:, 0] & candidate_ocean[:, -1])[0] + for row in wrap_rows: + union(labels[row, 0], labels[row, -1]) + + seed_labels = np.unique(labels[seed_row, candidate_ocean[seed_row, :]]) + seed_roots = {find(label) for label in seed_labels if label != 0} + if not seed_roots: + return np.zeros_like(candidate_ocean, dtype=bool) + + ocean_mask = np.zeros_like(candidate_ocean, dtype=bool) + for label in range(1, label_count + 1): + if find(label) in seed_roots: + ocean_mask |= labels == label + + return ocean_mask + + +def _coastline_edges(ocean_mask): + """ + Build east and north coastline edge diagnostics from an ocean mask. + """ + edge_east = ocean_mask != np.roll(ocean_mask, -1, axis=1) + edge_north = np.zeros_like(ocean_mask, dtype=bool) + edge_north[:-1, :] = ocean_mask[:-1, :] != ocean_mask[1:, :] + return edge_east, edge_north + + +def _signed_distance_from_mask( + ocean_mask, + edge_east, + edge_north, + lon, + lat, + chunk_size, + workers, +): + """ + Compute signed distance to the nearest raster coastline sample. + """ + sample_points = _coastline_sample_xyz( + edge_east=edge_east, edge_north=edge_north, lon=lon, lat=lat + ) + if sample_points.shape[0] == 0: + return np.where(ocean_mask, np.inf, -np.inf) + + tree = cKDTree(sample_points) + lon_rad = np.deg2rad(lon) + cos_lon = np.cos(lon_rad) + sin_lon = np.sin(lon_rad) + + distance = np.empty(ocean_mask.shape, dtype=np.float64) + for start in range(0, lat.size, chunk_size): + stop = min(start + chunk_size, lat.size) + lat_rad = np.deg2rad(lat[start:stop]) + cos_lat = np.cos(lat_rad)[:, np.newaxis] + sin_lat = np.sin(lat_rad)[:, np.newaxis] + xyz = np.empty((stop - start, lon.size, 3), dtype=np.float64) + xyz[:, :, 0] = cos_lat * cos_lon[np.newaxis, :] + xyz[:, :, 1] = cos_lat * sin_lon[np.newaxis, :] + xyz[:, :, 2] = sin_lat + points = xyz.reshape((-1, 3)) + chord_distance, _ = _tree_query(tree, points, workers) + chord_distance = np.clip(chord_distance, 0.0, 2.0) + angle = 2.0 * np.arcsin(0.5 * chord_distance) + distance[start:stop, :] = angle.reshape((stop - start, lon.size)) + + signed_distance = EARTH_RADIUS * distance + return np.where(ocean_mask, signed_distance, -signed_distance) + + +def _coastline_sample_xyz(edge_east, edge_north, lon, lat): + """ + Convert coastline edge midpoints into Cartesian coordinates. + """ + east_rows, east_cols = np.nonzero(edge_east) + north_rows, north_cols = np.nonzero(edge_north) + + sample_xyz = [] + + if east_rows.size > 0: + east_lon = _angular_midpoint( + lon[east_cols], lon[(east_cols + 1) % lon.size] + ) + east_lat = lat[east_rows] + sample_xyz.append(_lon_lat_to_xyz(east_lon, east_lat)) + + if north_rows.size > 0: + north_lon = lon[north_cols] + north_lat = 0.5 * (lat[north_rows] + lat[north_rows + 1]) + sample_xyz.append(_lon_lat_to_xyz(north_lon, north_lat)) + + if not sample_xyz: + return np.empty((0, 3), dtype=np.float64) + + return np.vstack(sample_xyz) + + +def _angular_midpoint(lon_a, lon_b): + """ + Compute the midpoint between longitudes, respecting antimeridian wrap. + """ + lon_a_rad = np.deg2rad(lon_a) + lon_b_rad = np.deg2rad(lon_b) + x = np.cos(lon_a_rad) + np.cos(lon_b_rad) + y = np.sin(lon_a_rad) + np.sin(lon_b_rad) + return np.rad2deg(np.arctan2(y, x)) + + +def _lon_lat_to_xyz(lon, lat): + """ + Convert lon/lat coordinates in degrees to Cartesian coordinates. + """ + lon_rad = np.deg2rad(lon) + lat_rad = np.deg2rad(lat) + cos_lat = np.cos(lat_rad) + xyz = np.empty((lon_rad.size, 3), dtype=np.float64) + xyz[:, 0] = cos_lat * np.cos(lon_rad) + xyz[:, 1] = cos_lat * np.sin(lon_rad) + xyz[:, 2] = np.sin(lat_rad) + return xyz + + +def _tree_query(tree, points, workers): + """ + Query a KD-tree with SciPy-version-compatible worker support. + """ + try: + return tree.query(points, workers=workers) + except TypeError: + return tree.query(points) + + +def _write_netcdf_with_fill_values(ds, filename, format='NETCDF4'): + """ + Write an xarray Dataset with NetCDF4 fill values where needed. + """ + ds = ds.copy() + fill_values = netCDF4.default_fillvals + encoding = {} + vars_and_coords = list(ds.data_vars.keys()) + list(ds.coords.keys()) + for var_name in vars_and_coords: + ds[var_name].attrs.pop('_FillValue', None) + is_numeric = np.issubdtype(ds[var_name].dtype, np.number) + if is_numeric: + dtype = ds[var_name].dtype + for fill_type, fill_value in fill_values.items(): + if dtype == np.dtype(fill_type): + encoding[var_name] = {'_FillValue': fill_value} + break + else: + encoding[var_name] = {'_FillValue': None} + ds.to_netcdf(filename, encoding=encoding, format=format) + + +def _rasterize_critical_transects( + critical_transects: CriticalTransects | None, + lon, + lat, +): + """ + Rasterize critical land blockages and passages onto a lat-lon grid. + """ + shape = (lat.size, lon.size) + if critical_transects is None: + return ( + np.zeros(shape, dtype=bool), + np.zeros(shape, dtype=bool), + ) + + land_blockages = _rasterize_feature_collection( + feature_collection=critical_transects.land_blockages, + lon=lon, + lat=lat, + ) + passages = _rasterize_feature_collection( + feature_collection=critical_transects.passages, + lon=lon, + lat=lat, + ) + return land_blockages, passages + + +def _rasterize_feature_collection(feature_collection, lon, lat): + """ + Rasterize a transect feature collection to a 4-connected cell mask. + """ + mask = np.zeros((lat.size, lon.size), dtype=bool) + if feature_collection is None: + return mask + + lon_spacing = _representative_spacing(lon) + lat_spacing = _representative_spacing(lat) + + for coordinates in _iter_feature_coordinates(feature_collection): + _rasterize_linestring( + mask=mask, + coordinates=coordinates, + lon=lon, + lat=lat, + lon_spacing=lon_spacing, + lat_spacing=lat_spacing, + ) + + return mask + + +def _iter_feature_coordinates(feature_collection): + """ + Yield the coordinates for each line geometry in a feature collection. + """ + if isinstance(feature_collection, dict): + features = feature_collection.get('features', []) + else: + features = getattr(feature_collection, 'features', []) + + for feature in features: + geometry = feature.get('geometry') + if geometry is None: + continue + + geometry_type = geometry.get('type') + if geometry_type == 'LineString': + yield np.asarray(geometry['coordinates'], dtype=np.float64) + elif geometry_type == 'MultiLineString': + for coordinates in geometry['coordinates']: + yield np.asarray(coordinates, dtype=np.float64) + else: + raise ValueError( + f'Unsupported critical transect geometry type: {geometry_type}' + ) + + +def _rasterize_linestring( + mask, + coordinates, + lon, + lat, + lon_spacing, + lat_spacing, +): + """ + Rasterize one line string by densifying and connecting nearest cells. + """ + if coordinates.shape[0] == 0: + return + + grid_path: list[tuple[int, int]] = [] + for index in range(coordinates.shape[0] - 1): + start = coordinates[index, :2] + stop = coordinates[index + 1, :2] + dense_points = _densify_segment( + start=start, + stop=stop, + lon_spacing=lon_spacing, + lat_spacing=lat_spacing, + ) + for point in dense_points: + row = _nearest_lat_index(point[1], lat) + col = _nearest_lon_index(point[0], lon) + cell = (row, col) + if not grid_path or grid_path[-1] != cell: + grid_path.append(cell) + + if not grid_path: + point = coordinates[0, :2] + grid_path.append( + ( + _nearest_lat_index(point[1], lat), + _nearest_lon_index(point[0], lon), + ) + ) + + start_row, start_col = grid_path[0] + mask[start_row, start_col] = True + for previous, current in zip(grid_path[:-1], grid_path[1:], strict=False): + _mark_grid_connection(mask, previous, current) + + +def _densify_segment(start, stop, lon_spacing, lat_spacing): + """ + Densify a segment enough to avoid gaps when mapping to cell centers. + """ + delta_lon = _wrap_longitude_delta(stop[0] - start[0]) + delta_lat = stop[1] - start[1] + n_lon = 0.0 if lon_spacing == 0.0 else abs(delta_lon) / lon_spacing + n_lat = 0.0 if lat_spacing == 0.0 else abs(delta_lat) / lat_spacing + segments = max(1, int(np.ceil(2.0 * max(n_lon, n_lat)))) + fractions = np.linspace(0.0, 1.0, segments + 1) + dense = np.empty((segments + 1, 2), dtype=np.float64) + dense[:, 0] = start[0] + fractions * delta_lon + dense[:, 1] = start[1] + fractions * delta_lat + return dense + + +def _representative_spacing(values): + """ + Estimate a typical grid spacing from a one-dimensional coordinate. + """ + if values.size < 2: + return 1.0 + + spacing = np.abs(np.diff(values)) + spacing = spacing[spacing > 0.0] + if spacing.size == 0: + return 1.0 + return float(np.median(spacing)) + + +def _nearest_lat_index(value, lat): + """ + Find the nearest latitude index for a point. + """ + return int(np.argmin(np.abs(lat - value))) + + +def _nearest_lon_index(value, lon): + """ + Find the nearest longitude index for a point with periodic wrap. + """ + delta = _wrap_longitude_delta(lon - value) + return int(np.argmin(np.abs(delta))) + + +def _wrap_longitude_delta(delta): + """ + Wrap longitude differences to the shortest path on the sphere. + """ + return np.mod(delta + 180.0, 360.0) - 180.0 + + +def _mark_grid_connection(mask, previous, current): + """ + Mark a 4-connected path between two raster cells. + """ + row, col = previous + stop_row, stop_col = current + n_lon = mask.shape[1] + + while row != stop_row or col != stop_col: + lon_delta = _periodic_index_delta(stop_col - col, n_lon) + if lon_delta != 0: + col = (col + int(np.sign(lon_delta))) % n_lon + mask[row, col] = True + + lat_delta = stop_row - row + if lat_delta != 0: + row += int(np.sign(lat_delta)) + mask[row, col] = True + + +def _periodic_index_delta(delta, size): + """ + Compute the shortest signed periodic index delta. + """ + half_size = size / 2.0 + if delta > half_size: + delta -= size + elif delta < -half_size: + delta += size + return delta diff --git a/polaris/mesh/spherical/unified/__init__.py b/polaris/mesh/spherical/unified/__init__.py index a726e065b2..753071a6fa 100644 --- a/polaris/mesh/spherical/unified/__init__.py +++ b/polaris/mesh/spherical/unified/__init__.py @@ -1,3 +1,18 @@ +from polaris.mesh.spherical.unified.cell_width import ( + UnifiedCellWidthMeshStep as UnifiedCellWidthMeshStep, +) +from polaris.mesh.spherical.unified.configs import ( + RIVER_CONFIG_FILENAME as RIVER_CONFIG_FILENAME, +) +from polaris.mesh.spherical.unified.configs import ( + UNIFIED_MESH_NAMES as UNIFIED_MESH_NAMES, +) +from polaris.mesh.spherical.unified.configs import ( + get_unified_mesh_config as get_unified_mesh_config, +) +from polaris.mesh.spherical.unified.families import ( + get_unified_mesh_family as get_unified_mesh_family, +) from polaris.mesh.spherical.unified.resolutions import ( LAT_LON_TARGET_GRID_RESOLUTIONS as LAT_LON_TARGET_GRID_RESOLUTIONS, ) diff --git a/polaris/mesh/spherical/unified/cell_width.py b/polaris/mesh/spherical/unified/cell_width.py new file mode 100644 index 0000000000..d7b25206c9 --- /dev/null +++ b/polaris/mesh/spherical/unified/cell_width.py @@ -0,0 +1,61 @@ +import os + +import xarray as xr + +from polaris.mesh import QuasiUniformSphericalMeshStep + + +class UnifiedCellWidthMeshStep(QuasiUniformSphericalMeshStep): + """ + A spherical mesh step that consumes an upstream unified sizing field. + """ + + def __init__( + self, + component, + sizing_field_step=None, + name='unified_base_mesh', + subdir=None, + mesh_name='mesh', + ): + super().__init__( + component=component, + name=name, + subdir=subdir, + cell_width=None, + mesh_name=mesh_name, + ) + self.sizing_field_step = sizing_field_step + self.sizing_field_filename = 'sizing_field.nc' + + def setup(self): + """ + Link an upstream sizing field if one has been provided. + """ + if self.sizing_field_step is not None: + self.add_input_file( + filename=self.sizing_field_filename, + work_dir_target=os.path.join( + self.sizing_field_step.path, + self.sizing_field_step.sizing_field_filename, + ), + ) + super().setup() + + def build_cell_width_lat_lon(self): + """ + Read the cell width, lon, and lat directly from sizing_field.nc. + """ + with xr.open_dataset(self.sizing_field_filename) as ds_sizing: + if 'cellWidth' not in ds_sizing: + raise ValueError( + 'Expected variable "cellWidth" in sizing_field.nc.' + ) + if 'lat' not in ds_sizing.coords or 'lon' not in ds_sizing.coords: + raise ValueError( + 'Expected lat/lon coordinates in sizing_field.nc.' + ) + cell_width = ds_sizing.cellWidth.values + lon = ds_sizing.lon.values + lat = ds_sizing.lat.values + return cell_width, lon, lat diff --git a/polaris/mesh/spherical/unified/configs.py b/polaris/mesh/spherical/unified/configs.py new file mode 100644 index 0000000000..ed82ba1084 --- /dev/null +++ b/polaris/mesh/spherical/unified/configs.py @@ -0,0 +1,100 @@ +import importlib.resources as imp_res +import os + +from polaris.config import PolarisConfigParser + +PACKAGE = 'polaris.mesh.spherical.unified' +FAMILY_CONFIG_PACKAGE = 'polaris.mesh.spherical.unified.families' +RIVER_CONFIG_PACKAGE = 'polaris.mesh.spherical.unified.river' +RIVER_CONFIG_FILENAME = 'river_network.cfg' +SIZING_FIELD_CONFIG_PACKAGE = 'polaris.mesh.spherical.unified.sizing_field' +SIZING_FIELD_CONFIG_FILENAME = 'sizing_field.cfg' +DEFAULT_CONFIG_FILENAME = 'unified_mesh.cfg' +DEFAULT_FAMILY_NAME = 'default' + + +def _discover_mesh_names(): + """ + Discover unified-mesh names from config files in the package. + """ + package_files = imp_res.files(PACKAGE) + mesh_names = [] + for resource in package_files.iterdir(): + if not resource.is_file() or not resource.name.endswith('.cfg'): + continue + if resource.name == DEFAULT_CONFIG_FILENAME: + continue + + config = PolarisConfigParser() + config.add_from_package(PACKAGE, resource.name) + config.combine() + combined = config.combined + assert combined is not None + if not combined.has_section('unified_mesh'): + continue + + mesh_name = combined.get('unified_mesh', 'mesh_name') + expected_name = os.path.splitext(resource.name)[0] + if mesh_name != expected_name: + raise ValueError( + f'Unified-mesh config {resource.name!r} declares ' + f'mesh_name={mesh_name!r}, expected {expected_name!r}.' + ) + + mesh_names.append(mesh_name) + + return tuple(sorted(mesh_names)) + + +UNIFIED_MESH_NAMES = _discover_mesh_names() + + +def get_unified_mesh_config(mesh_name, filepath=None): + """ + Load default, generic workflow and per-mesh unified-mesh configs. + """ + config_filename = _get_mesh_config_filename(mesh_name) + family_name = _get_mesh_family_name(config_filename) + config = PolarisConfigParser(filepath=filepath) + config.add_from_package(PACKAGE, DEFAULT_CONFIG_FILENAME) + config.add_from_package(RIVER_CONFIG_PACKAGE, RIVER_CONFIG_FILENAME) + config.add_from_package( + SIZING_FIELD_CONFIG_PACKAGE, SIZING_FIELD_CONFIG_FILENAME + ) + _add_family_config(config=config, family_name=family_name) + config.add_from_package(PACKAGE, config_filename) + config.combine() + return config + + +def _add_family_config(config, family_name): + package_files = imp_res.files(FAMILY_CONFIG_PACKAGE) + family_config_filename = f'{family_name}.cfg' + family_config = package_files.joinpath(family_config_filename) + if family_config.is_file(): + config.add_from_package(FAMILY_CONFIG_PACKAGE, family_config_filename) + + +def _get_mesh_family_name(config_filename): + config = PolarisConfigParser() + config.add_from_package(PACKAGE, config_filename) + config.combine() + combined = config.combined + assert combined is not None + return combined.get( + 'unified_mesh', 'mesh_family', fallback=DEFAULT_FAMILY_NAME + ) + + +def _get_mesh_config_filename(mesh_name): + """ + Get the config filename for one named unified mesh. + """ + if mesh_name not in UNIFIED_MESH_NAMES: + valid_mesh_names = ', '.join(UNIFIED_MESH_NAMES) + raise ValueError( + f'Unexpected unified mesh {mesh_name!r}. Valid mesh names ' + f'are: {valid_mesh_names}' + ) + + return f'{mesh_name}.cfg' diff --git a/polaris/mesh/spherical/unified/families/__init__.py b/polaris/mesh/spherical/unified/families/__init__.py new file mode 100644 index 0000000000..7b1f287ef5 --- /dev/null +++ b/polaris/mesh/spherical/unified/families/__init__.py @@ -0,0 +1,48 @@ +from typing import Protocol + +from polaris.mesh.spherical.unified.families.default import ( + DefaultUnifiedMeshFamily, + build_ocean_background_from_mode, +) +from polaris.mesh.spherical.unified.families.so_region import ( + SORegionUnifiedMeshFamily, +) + + +class UnifiedMeshFamily(Protocol): + name: str + + def setup_sizing_field_step(self, step): ... + + def build_ocean_background(self, ds_coastline, section): ... + + +_FAMILY_LIST: list[UnifiedMeshFamily] = [ + DefaultUnifiedMeshFamily(), + SORegionUnifiedMeshFamily(), +] + +_FAMILIES: dict[str, UnifiedMeshFamily] = { + family.name: family for family in _FAMILY_LIST +} + + +def get_unified_mesh_family(config): + """ + Get the unified-mesh family object for a combined mesh config. + """ + family_name = config.get('unified_mesh', 'mesh_family') + try: + return _FAMILIES[family_name] + except KeyError as exc: + valid_families = ', '.join(sorted(_FAMILIES)) + raise ValueError( + f'Unexpected unified mesh family {family_name!r}. Valid ' + f'families are: {valid_families}' + ) from exc + + +__all__ = [ + 'build_ocean_background_from_mode', + 'get_unified_mesh_family', +] diff --git a/polaris/mesh/spherical/unified/families/default.py b/polaris/mesh/spherical/unified/families/default.py new file mode 100644 index 0000000000..dd7db0fb75 --- /dev/null +++ b/polaris/mesh/spherical/unified/families/default.py @@ -0,0 +1,54 @@ +import mpas_tools.mesh.creation.mesh_definition_tools as mdt +import numpy as np + + +class DefaultUnifiedMeshFamily: + """ + The default unified-mesh family for built-in ocean backgrounds. + """ + + name = 'default' + + def setup_sizing_field_step(self, step): + """ + Add any family-specific sizing-field inputs. + """ + + def build_ocean_background(self, ds_coastline, section): + """ + Build the family ocean background on the shared target grid. + """ + return build_ocean_background_from_mode( + lat=ds_coastline.lat.values, + lon=ds_coastline.lon.values, + mode=section.get('ocean_background_mode'), + min_km=section.getfloat('ocean_background_min_km'), + max_km=section.getfloat('ocean_background_max_km'), + ) + + +def build_ocean_background_from_mode(lat, lon, mode, min_km, max_km): + """ + Build a 2D ocean-background field in km from the default family modes. + """ + if mode == 'constant': + if not np.isclose(min_km, max_km): + raise ValueError( + 'Constant ocean backgrounds require ' + 'ocean_background_min_km and ' + 'ocean_background_max_km to be equal.' + ) + values = np.full((lat.size, lon.size), max_km, dtype=float) + elif mode == 'rrs_latitude': + cell_width_vs_lat = mdt.RRS_CellWidthVsLat( + lat, cellWidthEq=max_km, cellWidthPole=min_km + ) + values = np.outer(cell_width_vs_lat, np.ones(lon.size)) + else: + raise ValueError( + f'Unexpected ocean background mode {mode!r}. Valid options are ' + 'constant and rrs_latitude. Mesh-specific ocean backgrounds ' + 'belong in a unified mesh family implementation.' + ) + + return values.astype(float) diff --git a/polaris/mesh/spherical/unified/families/so_region.py b/polaris/mesh/spherical/unified/families/so_region.py new file mode 100644 index 0000000000..7a2d9ebdd7 --- /dev/null +++ b/polaris/mesh/spherical/unified/families/so_region.py @@ -0,0 +1,34 @@ +SO_REGION_FILENAME = 'high_res_region.geojson' +SO_REGION_PACKAGE = 'polaris.mesh.base.so' + + +class SORegionUnifiedMeshFamily: + """ + The Southern Ocean unified-mesh family. + """ + + name = 'so_region' + + def setup_sizing_field_step(self, step): + """ + Link the shared Southern Ocean refinement region. + """ + step.add_input_file( + filename=SO_REGION_FILENAME, package=SO_REGION_PACKAGE + ) + + def build_ocean_background(self, ds_coastline, section): + """ + Build the Southern Ocean background on the shared target grid. + """ + from polaris.mesh.base.so.background import ( + build_southern_ocean_background, + ) + + return build_southern_ocean_background( + lat=ds_coastline.lat.values, + lon=ds_coastline.lon.values, + high_res_km=section.getfloat('ocean_background_min_km'), + low_res_km=section.getfloat('ocean_background_max_km'), + region_filename=SO_REGION_FILENAME, + ) diff --git a/polaris/mesh/spherical/unified/ocn_240km_lnd_240km_riv_240km.cfg b/polaris/mesh/spherical/unified/ocn_240km_lnd_240km_riv_240km.cfg new file mode 100644 index 0000000000..2392125371 --- /dev/null +++ b/polaris/mesh/spherical/unified/ocn_240km_lnd_240km_riv_240km.cfg @@ -0,0 +1,69 @@ +[unified_mesh] +# name of this unified mesh; must match the config filename +mesh_name = ocn_240km_lnd_240km_riv_240km + +# shared lat-lon target-grid resolution for coastline, river and sizing-field +# masks (degrees) +resolution_latlon = 0.25 + +[river_network] +# minimum drainage area to retain in the simplified network (m^2) +drainage_area_threshold = 2.0e10 + +# inland clip distance used for base-mesh river products (km) +base_mesh_clip_distance_km = 480.0 + +# geometry simplification tolerance used for base-mesh river products (km) +# finer than 1/2 river resolution so river network isn't over-simplified +base_mesh_simplify_tolerance_km = 30.0 + +# minimum retained segment length after base-mesh clipping (km) +# finer than 1/2 river resolution so river network isn't over-simplified +base_mesh_min_segment_length_km = 30.0 + +[river_lat_lon] +# maximum distance allowed when matching an ocean outlet to a coastline cell +outlet_match_tolerance = 50000.0 + +# physical buffer radius around river channels for the lat-lon channel mask; +# 0 preserves single-cell line rasterization +# We use twice the river resolution +channel_buffer_km = 480.0 + +[sizing_field] +# Whether to add a coastline refinement candidate field +enable_coastline_refinement = true + +# Whether to add a river-channel refinement candidate field +enable_river_channel_refinement = true + +# Whether to add a river-outlet refinement candidate field +enable_river_outlet_refinement = true + +# Background ocean resolution mode. Options are: +# constant: one ocean cell width everywhere +# rrs_latitude: RRS latitude-dependent ocean cell width +# Complex ocean backgrounds are delegated to mesh-family implementations. +ocean_background_mode = constant + +# Minimum ocean background cell width (km). For constant backgrounds, set this +# equal to ocean_background_max_km. +ocean_background_min_km = 240.0 + +# Maximum ocean background cell width (km). For rrs_latitude, this is the +# equatorial resolution. For so_region, this is the coarse background. +ocean_background_max_km = 240.0 + + +# Background land cell width (km) +land_background_km = 240.0 + +# Distance (km) over which the land side of the coastline refinement +# transitions back to the land background cell width +coastline_transition_land_km = 480.0 + +# Target cell width (km) along rasterized river channels +river_channel_km = 240.0 + +# Target cell width (km) at river outlets and inland sinks +river_outlet_km = 240.0 diff --git a/polaris/mesh/spherical/unified/ocn_30km_lnd_10km_riv_10km.cfg b/polaris/mesh/spherical/unified/ocn_30km_lnd_10km_riv_10km.cfg new file mode 100644 index 0000000000..417d601872 --- /dev/null +++ b/polaris/mesh/spherical/unified/ocn_30km_lnd_10km_riv_10km.cfg @@ -0,0 +1,66 @@ +[unified_mesh] +# name of this unified mesh; must match the config filename +mesh_name = ocn_30km_lnd_10km_riv_10km + +# shared lat-lon target-grid resolution for coastline, river and sizing-field +# masks (degrees) +resolution_latlon = 0.125 + +[river_network] +# minimum drainage area to retain in the simplified network (m^2) +drainage_area_threshold = 1.0e10 + +# inland clip distance used for base-mesh river products (km) +base_mesh_clip_distance_km = 60.0 + +# geometry simplification tolerance used for base-mesh river products (km) +base_mesh_simplify_tolerance_km = 5.0 + +# minimum retained segment length after base-mesh clipping (km) +base_mesh_min_segment_length_km = 5.0 + +[river_lat_lon] +# maximum distance allowed when matching an ocean outlet to a coastline cell +outlet_match_tolerance = 50000.0 + +# physical buffer radius around river channels for the lat-lon channel mask; +# 0 preserves single-cell line rasterization +# We use twice the river resolution +channel_buffer_km = 20.0 + +[sizing_field] +# Whether to add a coastline refinement candidate field +enable_coastline_refinement = true + +# Whether to add a river-channel refinement candidate field +enable_river_channel_refinement = true + +# Whether to add a river-outlet refinement candidate field +enable_river_outlet_refinement = true + +# Background ocean resolution mode. Options are: +# constant: one ocean cell width everywhere +# rrs_latitude: RRS latitude-dependent ocean cell width +# Complex ocean backgrounds are delegated to mesh-family implementations. +ocean_background_mode = constant + +# Minimum ocean background cell width (km). For constant backgrounds, set this +# equal to ocean_background_max_km. +ocean_background_min_km = 30.0 + +# Maximum ocean background cell width (km). For rrs_latitude, this is the +# equatorial resolution. For so_region, this is the coarse background. +ocean_background_max_km = 30.0 + +# Background land cell width (km) +land_background_km = 10.0 + +# Distance (km) over which the land side of the coastline refinement +# transitions back to the land background cell width +coastline_transition_land_km = 60.0 + +# Target cell width (km) along rasterized river channels +river_channel_km = 10.0 + +# Target cell width (km) at river outlets and inland sinks +river_outlet_km = 10.0 diff --git a/polaris/mesh/spherical/unified/ocn_rrs_6to18km_lnd_12km_riv_6km.cfg b/polaris/mesh/spherical/unified/ocn_rrs_6to18km_lnd_12km_riv_6km.cfg new file mode 100644 index 0000000000..81f2464ab0 --- /dev/null +++ b/polaris/mesh/spherical/unified/ocn_rrs_6to18km_lnd_12km_riv_6km.cfg @@ -0,0 +1,68 @@ +[unified_mesh] +# name of this unified mesh; must match the config filename +mesh_name = ocn_rrs_6to18km_lnd_12km_riv_6km + +# shared lat-lon target-grid resolution for coastline, river and sizing-field +# masks (degrees) +resolution_latlon = 0.03125 + +[river_network] +# minimum drainage area to retain in the simplified network (m^2) +drainage_area_threshold = 1.0e10 + +# inland clip distance used for base-mesh river products (km) +base_mesh_clip_distance_km = 36.0 + +# geometry simplification tolerance used for base-mesh river products (km) +base_mesh_simplify_tolerance_km = 3.0 + +# minimum retained segment length after base-mesh clipping (km) +base_mesh_min_segment_length_km = 3.0 + +[river_lat_lon] +# maximum distance allowed when matching an ocean outlet to a coastline cell +outlet_match_tolerance = 25000.0 + +# physical buffer radius around river channels for the lat-lon channel mask; +# 0 preserves single-cell line rasterization +# We use twice the river resolution +channel_buffer_km = 12.0 + +[sizing_field] +# Whether to add a coastline refinement candidate field +enable_coastline_refinement = true + +# Whether to add a river-channel refinement candidate field +enable_river_channel_refinement = true + +# Whether to add a river-outlet refinement candidate field +enable_river_outlet_refinement = true + +# Background ocean resolution mode. Options are: +# constant: one ocean cell width everywhere +# rrs_latitude: RRS latitude-dependent ocean cell width +# Complex ocean backgrounds are delegated to mesh-family implementations. +ocean_background_mode = rrs_latitude + +# Minimum ocean background cell width (km). For constant backgrounds, set this +# equal to ocean_background_max_km. +ocean_background_min_km = 6.0 + +# Maximum ocean background cell width (km). For rrs_latitude, this is the +# equatorial resolution. For so_region, this is the coarse background. +ocean_background_max_km = 18.0 + +# Background land cell width (km) +land_background_km = 12.0 + +# Distance (km) over which the land side of the coastline refinement +# transitions back to the land background cell width +# about double the coarsest ocean cell size, because we don't want finer land +# or river resolution limiting ocean time steps +coastline_transition_land_km = 36.0 + +# Target cell width (km) along rasterized river channels +river_channel_km = 6.0 + +# Target cell width (km) at river outlets and inland sinks +river_outlet_km = 6.0 diff --git a/polaris/mesh/spherical/unified/ocn_so_12to30km_lnd_10km_riv_10km.cfg b/polaris/mesh/spherical/unified/ocn_so_12to30km_lnd_10km_riv_10km.cfg new file mode 100644 index 0000000000..1ea6ec78e5 --- /dev/null +++ b/polaris/mesh/spherical/unified/ocn_so_12to30km_lnd_10km_riv_10km.cfg @@ -0,0 +1,63 @@ +[unified_mesh] +# name of this unified mesh; must match the config filename +mesh_name = ocn_so_12to30km_lnd_10km_riv_10km + +# Mesh family used to specialize shared unified-mesh behavior +mesh_family = so_region + +# shared lat-lon target-grid resolution for coastline, river and sizing-field +# masks (degrees) +resolution_latlon = 0.0625 + +[river_network] +# minimum drainage area to retain in the simplified network (m^2) +drainage_area_threshold = 1.0e10 + +# inland clip distance used for base-mesh river products (km) +base_mesh_clip_distance_km = 60.0 + +# geometry simplification tolerance used for base-mesh river products (km) +base_mesh_simplify_tolerance_km = 5.0 + +# minimum retained segment length after base-mesh clipping (km) +base_mesh_min_segment_length_km = 5.0 + +[river_lat_lon] +# maximum distance allowed when matching an ocean outlet to a coastline cell +outlet_match_tolerance = 50000.0 + +# physical buffer radius around river channels for the lat-lon channel mask; +# 0 preserves single-cell line rasterization +# We use twice the river resolution +channel_buffer_km = 20.0 + +[sizing_field] +# Whether to add a coastline refinement candidate field +enable_coastline_refinement = true + +# Whether to add a river-channel refinement candidate field +enable_river_channel_refinement = true + +# Whether to add a river-outlet refinement candidate field +enable_river_outlet_refinement = true + +# Minimum ocean background cell width (km). For constant backgrounds, set this +# equal to ocean_background_max_km. +ocean_background_min_km = 12.0 + +# Maximum ocean background cell width (km). For rrs_latitude, this is the +# equatorial resolution. For so_region, this is the coarse background. +ocean_background_max_km = 30.0 + +# Background land cell width (km) +land_background_km = 10.0 + +# Distance (km) over which the land side of the coastline refinement +# transitions back to the land background cell width +coastline_transition_land_km = 60.0 + +# Target cell width (km) along rasterized river channels +river_channel_km = 10.0 + +# Target cell width (km) at river outlets and inland sinks +river_outlet_km = 10.0 diff --git a/polaris/mesh/spherical/unified/river/__init__.py b/polaris/mesh/spherical/unified/river/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/polaris/mesh/spherical/unified/river/river_network.cfg b/polaris/mesh/spherical/unified/river/river_network.cfg new file mode 100644 index 0000000000..6e723c8eab --- /dev/null +++ b/polaris/mesh/spherical/unified/river/river_network.cfg @@ -0,0 +1,51 @@ +[river_network] +# public HydroRIVERS archive URL +hydrorivers_url = https://data.hydrosheds.org/file/HydroRIVERS/HydroRIVERS_v10_shp.zip + +# local filename for the downloaded HydroRIVERS archive +hydrorivers_archive_filename = HydroRIVERS_v10_shp.zip + +# shapefile directory expected after unpacking the archive +hydrorivers_shp_directory = HydroRIVERS_v10_shp + +# HydroRIVERS shapefile basename within the unpacked archive +hydrorivers_shp_filename = HydroRIVERS_v10.shp + +# minimum drainage area to retain in the simplified network (m^2) +drainage_area_threshold = 1.0e10 + +# minimum separation between retained non-endorheic outlets (m) +outlet_distance_tolerance = 10000.0 + +# minimum tributary-to-main-stem drainage-area ratio for retaining a nearby +# tributary at a confluence +tributary_area_ratio = 0.05 + +# inland clip distance for base-mesh river products (km) +base_mesh_clip_distance_km = 20.0 + +# geometry simplification tolerance for base-mesh river products (km) +base_mesh_simplify_tolerance_km = 2.0 + +# minimum retained segment length after base-mesh clipping (km) +base_mesh_min_segment_length_km = 2.0 + +# optional minimum outlet stub length to preserve after clipping (km) +base_mesh_preserve_outlet_stub_km = 0.0 + + +[river_lat_lon] +# maximum distance allowed when matching an ocean outlet to a coastline cell +outlet_match_tolerance = 50000.0 + +# fraction of a grid cell used when sampling lines for rasterization +channel_subsegment_fraction = 0.5 + +# physical buffer radius around river channels for the lat-lon channel mask; +# 0 preserves single-cell line rasterization +channel_buffer_km = 0.0 + + +[viz_river_network] +# output resolution for diagnostic plots +dpi = 200 diff --git a/polaris/mesh/spherical/unified/sizing_field/__init__.py b/polaris/mesh/spherical/unified/sizing_field/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/polaris/mesh/spherical/unified/sizing_field/sizing_field.cfg b/polaris/mesh/spherical/unified/sizing_field/sizing_field.cfg new file mode 100644 index 0000000000..90021e5af5 --- /dev/null +++ b/polaris/mesh/spherical/unified/sizing_field/sizing_field.cfg @@ -0,0 +1,44 @@ +[sizing_field] +# Whether to add a coastline refinement candidate field +enable_coastline_refinement = true + +# Whether to add a river-channel refinement candidate field +enable_river_channel_refinement = true + +# Whether to add a river-outlet refinement candidate field +enable_river_outlet_refinement = true + +# Background ocean resolution mode. Options are: +# constant: one ocean cell width everywhere +# rrs_latitude: RRS latitude-dependent ocean cell width +# Complex ocean backgrounds are delegated to mesh-family implementations. +ocean_background_mode = constant + +# Minimum ocean background cell width (km). For constant backgrounds, set this +# equal to ocean_background_max_km. +ocean_background_min_km = 240.0 + +# Maximum ocean background cell width (km). For rrs_latitude, this is the +# equatorial resolution. For so_region, this is the coarse background. +ocean_background_max_km = 240.0 + +# Background land cell width (km) +land_background_km = 240.0 + +# Distance (km) over which the land side of the coastline refinement +# transitions back to the land background cell width +coastline_transition_land_km = 0.0 + +# Target cell width (km) along rasterized river channels +river_channel_km = 240.0 + +# Target cell width (km) at river outlets and inland sinks +river_outlet_km = 240.0 + + +[viz_sizing_field] +# output resolution for diagnostic plots +dpi = 200 + +# colormap for cell-width diagnostic plots +cell_width_cmap = cmo.deep_r diff --git a/polaris/mesh/spherical/unified/unified_mesh.cfg b/polaris/mesh/spherical/unified/unified_mesh.cfg new file mode 100644 index 0000000000..d0dc925db7 --- /dev/null +++ b/polaris/mesh/spherical/unified/unified_mesh.cfg @@ -0,0 +1,6 @@ +[unified_mesh] +# Antarctic coastline convention shared by all unified-mesh workflow steps +coastline_convention = calving_front + +# Mesh family used to specialize shared unified-mesh behavior +mesh_family = default diff --git a/polaris/tasks/e3sm/init/topo/combine/step.py b/polaris/tasks/e3sm/init/topo/combine/step.py index a0251be220..363ccec502 100644 --- a/polaris/tasks/e3sm/init/topo/combine/step.py +++ b/polaris/tasks/e3sm/init/topo/combine/step.py @@ -76,7 +76,25 @@ def get_subdir(): subdir = f'combine_{CombineStep.ANTARCTIC}_{CombineStep.GLOBAL}' return os.path.join('topo', subdir) - def __init__(self, component, subdir): + @classmethod + def get_name(cls, target_grid, resolution_name): + """ + Get the step name based on the current datasets and target grid. + + Parameters + ---------- + target_grid : {'cubed_sphere', 'lat_lon'} + The type of target grid for the combined topography + + resolution_name : str + The target resolution name for the combined topography + """ + return ( + f'combine_topo_{cls.ANTARCTIC}_{cls.GLOBAL}' + f'_{target_grid}_{resolution_name}' + ) + + def __init__(self, component, subdir, target_grid, resolution_name): """ Create a new step @@ -88,10 +106,14 @@ def __init__(self, component, subdir): subdir : str The subdirectory within the component's work directory + target_grid : {'cubed_sphere', 'lat_lon'} + The type of target grid for the combined topography + + resolution_name : str + The target resolution name for the combined topography + """ - antarctic_dataset = self.ANTARCTIC - global_dataset = self.GLOBAL - name = f'combine_topo_{antarctic_dataset}_{global_dataset}' + name = self.get_name(target_grid, resolution_name) super().__init__( component=component, name=name, diff --git a/polaris/tasks/e3sm/init/topo/combine/steps.py b/polaris/tasks/e3sm/init/topo/combine/steps.py index 4607e59a3b..637d38460f 100644 --- a/polaris/tasks/e3sm/init/topo/combine/steps.py +++ b/polaris/tasks/e3sm/init/topo/combine/steps.py @@ -2,6 +2,7 @@ from polaris.config import PolarisConfigParser from polaris.e3sm.init.topo import format_lat_lon_resolution_name +from polaris.step import Step from polaris.tasks.e3sm.init.topo.combine.step import CombineStep from polaris.tasks.e3sm.init.topo.combine.viz import VizCombinedStep @@ -26,8 +27,9 @@ def get_cubed_sphere_topo_steps(component, resolution, include_viz=False): Returns ------- - steps : list of polaris.Step - The combine topo step and optional visualization step. + steps : dict of str to polaris.Step + The combine topo step and optional visualization step, keyed by step + name. config : polaris.config.PolarisConfigParser The shared config options. @@ -57,8 +59,9 @@ def get_lat_lon_topo_steps(component, resolution, include_viz=False): Returns ------- - steps : list of polaris.Step - The combine topo step and optional visualization step. + steps : dict of str to polaris.Step + The combine topo step and optional visualization step, keyed by step + name. config : polaris.config.PolarisConfigParser The shared config options. @@ -91,8 +94,9 @@ def _get_target_topo_steps(component, target_grid, resolution, include_viz): Returns ------- - steps : list of polaris.Step - The combine topo step and optional visualization step. + steps : dict of str to polaris.Step + The combine topo step and optional visualization step, keyed by step + name. config : polaris.config.PolarisConfigParser The shared config options. @@ -122,14 +126,16 @@ def _get_target_topo_steps(component, target_grid, resolution, include_viz): else: config.set('combine_topo', 'resolution_latlon', f'{resolution}') - steps = [] + steps: dict[str, Step] = {} combine_step = component.get_or_create_shared_step( step_cls=CombineStep, subdir=subdir, config=config, + target_grid=target_grid, + resolution_name=resolution_name, ) combine_step._set_res_and_outputs(update=False) - steps.append(combine_step) + steps[combine_step.name] = combine_step if include_viz: viz_subdir = os.path.join(combine_step.subdir, 'viz') @@ -140,7 +146,7 @@ def _get_target_topo_steps(component, target_grid, resolution, include_viz): config_filename='combine_topo.cfg', combine_step=combine_step, ) - steps.append(viz_step) + steps[viz_step.name] = viz_step return steps, config diff --git a/polaris/tasks/e3sm/init/topo/combine/task.py b/polaris/tasks/e3sm/init/topo/combine/task.py index 1daefd2e6a..0a276fb007 100644 --- a/polaris/tasks/e3sm/init/topo/combine/task.py +++ b/polaris/tasks/e3sm/init/topo/combine/task.py @@ -49,7 +49,7 @@ def __init__(self, component, resolution): include_viz=True, ) self.set_shared_config(config, link='combine_topo.cfg') - for step in steps: + for step in steps.values(): self.add_step(step) @@ -88,5 +88,5 @@ def __init__(self, component, resolution): include_viz=True, ) self.set_shared_config(config, link='combine_topo.cfg') - for step in steps: + for step in steps.values(): self.add_step(step) diff --git a/polaris/tasks/e3sm/init/topo/cull/steps.py b/polaris/tasks/e3sm/init/topo/cull/steps.py index 13b93fba1a..109bc724e8 100644 --- a/polaris/tasks/e3sm/init/topo/cull/steps.py +++ b/polaris/tasks/e3sm/init/topo/cull/steps.py @@ -33,8 +33,8 @@ def get_default_cull_topo_steps( Returns ------- - steps : list of polaris.Step - Steps for masking and then culling topography + steps : dict of str to polaris.Step + Steps for masking and then culling topography, keyed by step name """ mesh_name = base_mesh_step.mesh_name @@ -48,7 +48,7 @@ def get_default_cull_topo_steps( ) config = PolarisConfigParser(filepath=filepath) config.add_from_package('polaris.tasks.e3sm.init.topo.cull', 'cull.cfg') - steps: list[Step] = [] + steps: dict[str, Step] = {} step_name = 'cull_mask' subdir = os.path.join(mesh_name, 'topo', 'cull', 'mask') @@ -61,7 +61,7 @@ def get_default_cull_topo_steps( unsmoothed_topo_step=unsmoothed_topo_step, name=step_name, ) - steps.append(cull_mask_step) + steps[cull_mask_step.name] = cull_mask_step step_name = 'cull_mesh' subdir = os.path.join(mesh_name, 'topo', 'cull', 'mesh') @@ -74,6 +74,6 @@ def get_default_cull_topo_steps( cull_mask_step=cull_mask_step, name=step_name, ) - steps.append(cull_mesh_step) + steps[cull_mesh_step.name] = cull_mesh_step return steps, config diff --git a/polaris/tasks/e3sm/init/topo/cull/task.py b/polaris/tasks/e3sm/init/topo/cull/task.py index 9ca8dac49d..3926a5fea9 100644 --- a/polaris/tasks/e3sm/init/topo/cull/task.py +++ b/polaris/tasks/e3sm/init/topo/cull/task.py @@ -75,7 +75,7 @@ def __init__( include_viz=include_viz, ) self.set_shared_config(config, link='cull_topo.cfg') - for step in steps: + for step in steps.values(): self.add_step(step) self.combine_topo_step = combine_topo_step diff --git a/polaris/tasks/e3sm/init/topo/cull/tasks.py b/polaris/tasks/e3sm/init/topo/cull/tasks.py index f5b474984b..4794ff179c 100644 --- a/polaris/tasks/e3sm/init/topo/cull/tasks.py +++ b/polaris/tasks/e3sm/init/topo/cull/tasks.py @@ -6,6 +6,7 @@ from polaris.tasks.e3sm.init.topo.combine import ( get_cubed_sphere_topo_steps, ) +from polaris.tasks.e3sm.init.topo.combine.step import CombineStep from polaris.tasks.e3sm.init.topo.cull.task import CullTopoTask from polaris.tasks.e3sm.init.topo.remap.steps import ( get_default_remap_topo_steps, @@ -26,7 +27,9 @@ def add_cull_topo_tasks(component): combine_topo_steps, _ = get_cubed_sphere_topo_steps( component=component, resolution=resolution ) - combine_steps[low_res] = combine_topo_steps[0] + combine_steps[low_res] = combine_topo_steps[ + CombineStep.get_name('cubed_sphere', f'ne{resolution}') + ] base_mesh_steps = get_base_mesh_steps() @@ -42,8 +45,8 @@ def add_cull_topo_tasks(component): smoothing=True, include_viz=False, ) - remap_mask_step = remap_topo_steps[0] - unsmoothed_topo_step = remap_topo_steps[1] + remap_mask_step = remap_topo_steps['mask_topo'] + unsmoothed_topo_step = remap_topo_steps['remap_unsmoothed'] task = CullTopoTask( component=component, diff --git a/polaris/tasks/e3sm/init/topo/remap/steps.py b/polaris/tasks/e3sm/init/topo/remap/steps.py index 0866dd54cd..3c43154131 100644 --- a/polaris/tasks/e3sm/init/topo/remap/steps.py +++ b/polaris/tasks/e3sm/init/topo/remap/steps.py @@ -50,9 +50,10 @@ def get_default_remap_topo_steps( Returns ------- - steps : list of polaris.Step + steps : dict of str to polaris.Step Steps for remapping topography (without smoothing and optionally with - smoothing) as well as, if requested, for visualizing the results. + smoothing) as well as, if requested, for visualizing the results, + keyed by step name. """ mesh_name = base_mesh_step.mesh_name @@ -70,7 +71,7 @@ def get_default_remap_topo_steps( config.add_from_package( 'polaris.tasks.e3sm.init.topo.remap', 'remap_low_res.cfg' ) - steps: list[Step] = [] + steps: dict[str, Step] = {} step_name = 'mask_topo' subdir = os.path.join(mesh_name, 'topo', 'remap', 'mask') @@ -82,7 +83,7 @@ def get_default_remap_topo_steps( combine_topo_step=combine_topo_step, name=step_name, ) - steps.append(mask_step) + steps[mask_step.name] = mask_step if smoothing: suffixes = ['unsmoothed', 'smoothed'] @@ -108,7 +109,7 @@ def get_default_remap_topo_steps( if suffix == 'unsmoothed': unsmoothed_topo = remap_step - steps.append(remap_step) + steps[remap_step.name] = remap_step if include_viz: step_name = f'viz_remapped_{suffix}' @@ -121,6 +122,6 @@ def get_default_remap_topo_steps( remap_step=remap_step, name=step_name, ) - steps.append(viz_step) + steps[viz_step.name] = viz_step return steps, config diff --git a/polaris/tasks/e3sm/init/topo/remap/task.py b/polaris/tasks/e3sm/init/topo/remap/task.py index 0bb6656050..1a8396c664 100644 --- a/polaris/tasks/e3sm/init/topo/remap/task.py +++ b/polaris/tasks/e3sm/init/topo/remap/task.py @@ -72,7 +72,7 @@ def __init__( include_viz=include_viz, ) self.set_shared_config(config, link='remap_topo.cfg') - for step in steps: + for step in steps.values(): self.add_step(step) self.combine_topo_step = combine_topo_step diff --git a/polaris/tasks/e3sm/init/topo/remap/tasks.py b/polaris/tasks/e3sm/init/topo/remap/tasks.py index a7da62ef25..dd4a8e92b0 100644 --- a/polaris/tasks/e3sm/init/topo/remap/tasks.py +++ b/polaris/tasks/e3sm/init/topo/remap/tasks.py @@ -6,6 +6,7 @@ from polaris.tasks.e3sm.init.topo.combine import ( get_cubed_sphere_topo_steps, ) +from polaris.tasks.e3sm.init.topo.combine.step import CombineStep from polaris.tasks.e3sm.init.topo.remap import RemapTopoTask @@ -23,7 +24,9 @@ def add_remap_topo_tasks(component): combine_topo_steps, _ = get_cubed_sphere_topo_steps( component=component, resolution=resolution ) - combine_steps[low_res] = combine_topo_steps[0] + combine_steps[low_res] = combine_topo_steps[ + CombineStep.get_name('cubed_sphere', f'ne{resolution}') + ] base_mesh_steps = get_base_mesh_steps() diff --git a/polaris/tasks/mesh/add_tasks.py b/polaris/tasks/mesh/add_tasks.py index 10fc531510..f497bbcef0 100644 --- a/polaris/tasks/mesh/add_tasks.py +++ b/polaris/tasks/mesh/add_tasks.py @@ -1,4 +1,11 @@ from polaris.tasks.mesh.base import add_base_mesh_tasks +from polaris.tasks.mesh.spherical.unified.coastline import ( + add_coastline_tasks, +) +from polaris.tasks.mesh.spherical.unified.river import add_river_tasks +from polaris.tasks.mesh.spherical.unified.sizing_field import ( + add_sizing_field_tasks, +) def add_mesh_tasks(component): @@ -10,3 +17,6 @@ def add_mesh_tasks(component): """ # add tasks alphabetically add_base_mesh_tasks(component=component) + add_coastline_tasks(component=component) + add_river_tasks(component=component) + add_sizing_field_tasks(component=component) diff --git a/polaris/tasks/mesh/spherical/__init__.py b/polaris/tasks/mesh/spherical/__init__.py new file mode 100644 index 0000000000..2a2735b62d --- /dev/null +++ b/polaris/tasks/mesh/spherical/__init__.py @@ -0,0 +1 @@ +"""Tasks under the spherical mesh component.""" diff --git a/polaris/tasks/mesh/spherical/unified/__init__.py b/polaris/tasks/mesh/spherical/unified/__init__.py new file mode 100644 index 0000000000..05712bad95 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/__init__.py @@ -0,0 +1,19 @@ +from polaris.tasks.mesh.spherical.unified.sizing_field import ( + BuildSizingFieldStep, + SizingFieldTask, + VizSizingFieldStep, + add_sizing_field_tasks, + get_lat_lon_sizing_field_steps, + get_sizing_field_config, + sizing_field_dataset, +) + +__all__ = [ + 'BuildSizingFieldStep', + 'SizingFieldTask', + 'VizSizingFieldStep', + 'add_sizing_field_tasks', + 'sizing_field_dataset', + 'get_lat_lon_sizing_field_steps', + 'get_sizing_field_config', +] diff --git a/polaris/tasks/mesh/spherical/unified/coastline/__init__.py b/polaris/tasks/mesh/spherical/unified/coastline/__init__.py new file mode 100644 index 0000000000..3441f183b7 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/coastline/__init__.py @@ -0,0 +1,24 @@ +from polaris.mesh.spherical.coastline import ( + CONVENTIONS as CONVENTIONS, +) +from polaris.mesh.spherical.coastline import ( + build_coastline_dataset as build_coastline_dataset, +) +from polaris.mesh.spherical.coastline import ( + build_coastline_datasets as build_coastline_datasets, +) +from polaris.tasks.mesh.spherical.unified.coastline.prepare import ( + PrepareCoastlineStep as PrepareCoastlineStep, +) +from polaris.tasks.mesh.spherical.unified.coastline.steps import ( + get_lat_lon_coastline_steps as get_lat_lon_coastline_steps, +) +from polaris.tasks.mesh.spherical.unified.coastline.task import ( + LatLonCoastlineTask as LatLonCoastlineTask, +) +from polaris.tasks.mesh.spherical.unified.coastline.tasks import ( + add_coastline_tasks as add_coastline_tasks, +) +from polaris.tasks.mesh.spherical.unified.coastline.viz import ( + VizCoastlineStep as VizCoastlineStep, +) diff --git a/polaris/tasks/mesh/spherical/unified/coastline/coastline.cfg b/polaris/tasks/mesh/spherical/unified/coastline/coastline.cfg new file mode 100644 index 0000000000..fd3eb0f7af --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/coastline/coastline.cfg @@ -0,0 +1,27 @@ +[coastline] +# target resolution (degrees) +resolution_latlon = 0.25 + +# whether to apply shared critical transects from geometric_features before +# flood filling the candidate ocean mask +include_critical_transects = True + +# threshold used to convert remapped ice and grounded masks to binary masks +mask_threshold = 0.5 + +# sea level used in the bedrock threshold +sea_level_elevation = 0.0 + +# number of latitude rows to process in each distance-query chunk +distance_chunk_size = 64 + + +[viz_coastline] +# rows south of this latitude are shown in the Antarctic stereographic plots +antarctic_max_latitude = -45.0 + +# output resolution for the diagnostic plots +dpi = 200 + +# symmetric absolute limit in meters for signed-distance plots +signed_distance_limit = 500000.0 diff --git a/polaris/tasks/mesh/spherical/unified/coastline/prepare.py b/polaris/tasks/mesh/spherical/unified/coastline/prepare.py new file mode 100644 index 0000000000..1767261419 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/coastline/prepare.py @@ -0,0 +1,107 @@ +import os + +import xarray as xr + +from polaris.mesh.spherical.coastline import ( + CONVENTIONS, + _write_netcdf_with_fill_values, + build_coastline_dataset, + build_coastline_datasets, +) +from polaris.mesh.spherical.critical_transects import ( + load_default_critical_transects, +) +from polaris.step import Step + +__all__ = [ + 'CONVENTIONS', + 'PrepareCoastlineStep', + 'build_coastline_dataset', + 'build_coastline_datasets', +] + + +class PrepareCoastlineStep(Step): + """ + Prepare coastline masks and signed-distance fields on a lat-lon grid. + """ + + def __init__(self, component, combine_step, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + combine_step : polaris.tasks.e3sm.init.topo.combine.CombineStep + The step with combined topography on the target lat-lon grid + + subdir : str + The subdirectory within the component's work directory + """ + super().__init__( + component=component, + name='coastline', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.combine_step = combine_step + self.output_filenames = { + convention: f'coastline_{convention}.nc' + for convention in CONVENTIONS + } + + def setup(self): + """ + Set up the step in the work directory, including linking inputs. + """ + combine_step = self.combine_step + self.add_input_file( + filename='topography.nc', + work_dir_target=os.path.join( + combine_step.path, combine_step.combined_filename + ), + ) + for filename in self.output_filenames.values(): + self.add_output_file(filename=filename) + + def run(self): + """ + Run this step. + """ + section = self.config['coastline'] + resolution = section.getfloat('resolution_latlon') + include_critical_transects = section.getboolean( + 'include_critical_transects' + ) + mask_threshold = section.getfloat('mask_threshold') + sea_level_elevation = section.getfloat('sea_level_elevation') + distance_chunk_size = section.getint('distance_chunk_size') + + critical_transects = None + if include_critical_transects: + critical_transects = load_default_critical_transects() + + ds_topo = xr.open_dataset('topography.nc') + ds_coastlines = build_coastline_datasets( + ds_topo=ds_topo, + resolution=resolution, + mask_threshold=mask_threshold, + sea_level_elevation=sea_level_elevation, + distance_chunk_size=distance_chunk_size, + workers=self.cpus_per_task, + critical_transects=critical_transects, + ) + for convention, ds_coastline in ds_coastlines.items(): + ds_coastline.attrs['source_topography'] = ( + self.combine_step.combined_filename + ) + ds_coastline.attrs['source_topography_step'] = ( + self.combine_step.subdir + ) + _write_netcdf_with_fill_values( + ds_coastline, self.output_filenames[convention] + ) diff --git a/polaris/tasks/mesh/spherical/unified/coastline/steps.py b/polaris/tasks/mesh/spherical/unified/coastline/steps.py new file mode 100644 index 0000000000..fe516ca374 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/coastline/steps.py @@ -0,0 +1,104 @@ +import os + +from polaris.config import PolarisConfigParser +from polaris.e3sm.init.topo import format_lat_lon_resolution_name +from polaris.step import Step +from polaris.tasks.mesh.spherical.unified.coastline.prepare import ( + PrepareCoastlineStep, +) +from polaris.tasks.mesh.spherical.unified.coastline.viz import ( + VizCoastlineStep, +) + + +def get_lat_lon_coastline_steps( + component, combine_topo_step, resolution, include_viz=False +): + """ + Get shared coastline-preparation steps for a lat-lon target grid. + + Parameters + ---------- + component : polaris.Component + The component the steps belong to + + combine_topo_step : polaris.tasks.e3sm.init.topo.combine.CombineStep + The shared combined-topography step on the target lat-lon grid + + resolution : float + The latitude-longitude resolution in degrees + + include_viz : bool, optional + Whether to include a visualization step + + Returns + ------- + steps : dict of str to polaris.Step + The coastline step and optional visualization step, keyed by step name + + config : polaris.config.PolarisConfigParser + The shared config options + """ + resolution_name = format_lat_lon_resolution_name(resolution) + config_filename = 'coastline.cfg' + filepath = os.path.join( + component.name, + 'spherical', + 'unified', + 'coastline', + 'lat_lon', + resolution_name, + config_filename, + ) + config = _get_lat_lon_coastline_config( + component=component, + filepath=filepath, + config_filename=config_filename, + resolution=resolution, + ) + + steps: dict[str, Step] = {} + subdir = os.path.join( + 'spherical', + 'unified', + 'coastline', + 'lat_lon', + resolution_name, + 'prepare', + ) + coastline_step = component.get_or_create_shared_step( + step_cls=PrepareCoastlineStep, + subdir=subdir, + config=config, + config_filename=config_filename, + combine_step=combine_topo_step, + ) + steps[coastline_step.name] = coastline_step + + if include_viz: + viz_subdir = os.path.join(subdir, 'viz') + viz_step = component.get_or_create_shared_step( + step_cls=VizCoastlineStep, + subdir=viz_subdir, + config=config, + config_filename=config_filename, + coastline_step=coastline_step, + ) + steps[viz_step.name] = viz_step + + return steps, config + + +def _get_lat_lon_coastline_config( + component, filepath, config_filename, resolution +): + if filepath in component.configs: + return component.configs[filepath] + + config = PolarisConfigParser(filepath=filepath) + config.add_from_package( + 'polaris.tasks.mesh.spherical.unified.coastline', + config_filename, + ) + config.set('coastline', 'resolution_latlon', f'{resolution}') + return config diff --git a/polaris/tasks/mesh/spherical/unified/coastline/task.py b/polaris/tasks/mesh/spherical/unified/coastline/task.py new file mode 100644 index 0000000000..5082af524f --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/coastline/task.py @@ -0,0 +1,68 @@ +import os + +from polaris.e3sm.init.topo import format_lat_lon_resolution_name +from polaris.task import Task +from polaris.tasks.e3sm.init import e3sm_init +from polaris.tasks.e3sm.init.topo.combine.step import CombineStep +from polaris.tasks.e3sm.init.topo.combine.steps import ( + get_lat_lon_topo_steps, +) +from polaris.tasks.mesh.spherical.unified.coastline.steps import ( + get_lat_lon_coastline_steps, +) + + +class LatLonCoastlineTask(Task): + """ + A task for preparing coastline products on a latitude-longitude grid. + """ + + def __init__(self, component, resolution): + """ + Create a new task. + + Parameters + ---------- + component : polaris.Component + The component the task belongs to + + resolution : float + The latitude-longitude resolution in degrees + """ + resolution_name = format_lat_lon_resolution_name(resolution) + subdir = os.path.join( + 'spherical', + 'unified', + 'coastline', + 'lat_lon', + resolution_name, + 'task', + ) + super().__init__( + component=component, + name=f'coastline_lat_lon_{resolution_name}_task', + subdir=subdir, + ) + + combine_steps, _ = get_lat_lon_topo_steps( + component=e3sm_init, + resolution=resolution, + include_viz=False, + ) + self.combine_topo_step = combine_steps[ + CombineStep.get_name( + target_grid='lat_lon', resolution_name=resolution_name + ) + ] + self.add_step(self.combine_topo_step, symlink='combine_topo') + + steps, config = get_lat_lon_coastline_steps( + component=component, + combine_topo_step=self.combine_topo_step, + resolution=resolution, + include_viz=True, + ) + self.set_shared_config(config, link='coastline.cfg') + for step in steps.values(): + symlink = os.path.basename(step.subdir) + self.add_step(step, symlink=symlink) diff --git a/polaris/tasks/mesh/spherical/unified/coastline/tasks.py b/polaris/tasks/mesh/spherical/unified/coastline/tasks.py new file mode 100644 index 0000000000..04d950ab82 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/coastline/tasks.py @@ -0,0 +1,19 @@ +from polaris.mesh.spherical.unified import LAT_LON_TARGET_GRID_RESOLUTIONS +from polaris.tasks.mesh.spherical.unified.coastline.task import ( + LatLonCoastlineTask, +) + + +def add_coastline_tasks(component): + """ + Add standalone coastline tasks for each supported lat-lon target grid. + + Parameters + ---------- + component : polaris.Component + The mesh component that the tasks belong to + """ + for resolution in LAT_LON_TARGET_GRID_RESOLUTIONS: + component.add_task( + LatLonCoastlineTask(component=component, resolution=resolution) + ) diff --git a/polaris/tasks/mesh/spherical/unified/coastline/viz.py b/polaris/tasks/mesh/spherical/unified/coastline/viz.py new file mode 100644 index 0000000000..e022a80f45 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/coastline/viz.py @@ -0,0 +1,438 @@ +import os + +import cartopy.crs as ccrs +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr +from matplotlib.colors import TwoSlopeNorm +from matplotlib.lines import Line2D +from numpy.typing import NDArray +from pyremap.descriptor.utility import interp_extrap_corner + +from polaris.mesh.spherical.critical_transects import ( + load_default_critical_transects, +) +from polaris.step import Step +from polaris.tasks.mesh.spherical.unified.coastline.prepare import ( + CONVENTIONS, +) +from polaris.viz import use_mplstyle + + +class VizCoastlineStep(Step): + """ + A step for visualizing coastline diagnostics. + """ + + def __init__(self, component, coastline_step, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + coastline_step : polaris.tasks.mesh.spherical.unified.coastline. + PrepareCoastlineStep + The coastline step to visualize + + subdir : str + The subdirectory within the component's work directory + """ + super().__init__( + component=component, + name='viz_coastline', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.coastline_step = coastline_step + + def setup(self): + """ + Set up the step in the work directory, including linking inputs. + """ + for convention in CONVENTIONS: + filename = self.coastline_step.output_filenames[convention] + self.add_input_file( + filename=filename, + work_dir_target=os.path.join( + self.coastline_step.path, filename + ), + ) + + def run(self): + """ + Run this step. + """ + use_mplstyle() + + viz_section = self.config['viz_coastline'] + prepare_section = self.config['coastline'] + dpi = viz_section.getint('dpi') + antarctic_max_latitude = viz_section.getfloat('antarctic_max_latitude') + signed_distance_limit = viz_section.getfloat('signed_distance_limit') + plot_info = _get_plot_info_from_file( + self.coastline_step.output_filenames[CONVENTIONS[0]] + ) + line_overlays = None + if prepare_section.getboolean('include_critical_transects'): + try: + line_overlays = _get_critical_transect_overlays() + except Exception: # pragma: no cover - diagnostic fallback + self.logger.warning( + 'Could not load critical transects for coastline ' + 'overlays.', + exc_info=True, + ) + + with open('debug_summary.txt', 'w') as summary: + for convention in CONVENTIONS: + filename = self.coastline_step.output_filenames[convention] + with xr.open_dataset(filename) as ds_coastline: + ocean_mask = ds_coastline.ocean_mask.values > 0 + signed_distance = ds_coastline.signed_distance.values + + self._plot_binary_field( + plot_info=plot_info, + field=ocean_mask, + title=f'{convention} ocean mask', + out_prefix=f'{convention}_ocean_mask', + dpi=dpi, + antarctic_max_latitude=antarctic_max_latitude, + line_overlays=line_overlays, + ) + self._plot_signed_distance( + plot_info=plot_info, + signed_distance=signed_distance, + title=f'{convention} signed distance', + out_prefix=f'{convention}_signed_distance', + distance_limit=signed_distance_limit, + dpi=dpi, + antarctic_max_latitude=antarctic_max_latitude, + ) + self._write_convention_summary( + summary=summary, + convention=convention, + ocean_mask=ocean_mask, + signed_distance=signed_distance, + ) + + def _plot_binary_field( + self, + plot_info, + field, + title, + out_prefix, + dpi, + antarctic_max_latitude, + line_overlays=None, + ): + """ + Plot a binary diagnostic field for both global and Antarctic views. + """ + field = np.asarray(field, dtype=float) + self._plot_scalar_field( + plot_info=plot_info, + field=field, + title=title, + out_filename=f'{out_prefix}_global.png', + dpi=dpi, + antarctic_max_latitude=antarctic_max_latitude, + projection='global', + cmap='Greys', + vmin=0.0, + vmax=1.0, + line_overlays=line_overlays, + ) + self._plot_scalar_field( + plot_info=plot_info, + field=field, + title=title, + out_filename=f'{out_prefix}_antarctic.png', + dpi=dpi, + antarctic_max_latitude=antarctic_max_latitude, + projection='antarctic', + cmap='Greys', + vmin=0.0, + vmax=1.0, + line_overlays=line_overlays, + ) + + def _plot_signed_distance( + self, + plot_info, + signed_distance, + title, + out_prefix, + distance_limit, + dpi, + antarctic_max_latitude, + ): + """ + Plot signed distance for both global and Antarctic views. + """ + field = 1.0e-3 * signed_distance + limit_km = 1.0e-3 * distance_limit + norm = TwoSlopeNorm(vmin=-limit_km, vcenter=0.0, vmax=limit_km) + + self._plot_scalar_field( + plot_info=plot_info, + field=field, + title=f'{title} (km)', + out_filename=f'{out_prefix}_global.png', + dpi=dpi, + antarctic_max_latitude=antarctic_max_latitude, + projection='global', + cmap='RdBu_r', + norm=norm, + ) + self._plot_scalar_field( + plot_info=plot_info, + field=field, + title=f'{title} (km)', + out_filename=f'{out_prefix}_antarctic.png', + dpi=dpi, + antarctic_max_latitude=antarctic_max_latitude, + projection='antarctic', + cmap='RdBu_r', + norm=norm, + ) + + def _plot_scalar_field( + self, + plot_info, + field, + title, + out_filename, + dpi, + antarctic_max_latitude, + projection, + cmap, + vmin=None, + vmax=None, + norm=None, + line_overlays=None, + ): + """ + Plot a scalar field with pcolormesh. + """ + ref_projection = ccrs.PlateCarree() + ax_projection = _get_projection(projection) + figsize = (20, 8) if projection == 'global' else (12, 12) + + ordered_field = field[:, plot_info['center_order']] + + fig = plt.figure(figsize=figsize, dpi=dpi) + ax = plt.axes(projection=ax_projection) + _configure_axes( + ax=ax, + projection=projection, + antarctic_max_latitude=antarctic_max_latitude, + ) + mesh = ax.pcolormesh( + plot_info['lon_corner'], + plot_info['lat_corner'], + ordered_field, + cmap=cmap, + vmin=vmin, + vmax=vmax, + norm=norm, + transform=ref_projection, + ) + if line_overlays is not None: + _plot_line_overlays( + ax=ax, + projection=projection, + line_overlays=line_overlays, + ) + fig.colorbar(mesh, ax=ax, shrink=0.7) + ax.set_title(title) + fig.savefig(out_filename, bbox_inches='tight', pad_inches=0.1) + plt.close(fig) + + @staticmethod + def _write_convention_summary( + summary, + convention, + ocean_mask, + signed_distance, + ): + """ + Write convention-specific ocean-mask counts to the summary file. + """ + total_cells = ocean_mask.size + summary.write(f'[{convention}]\n') + summary.write(f'total_cells: {total_cells}\n') + summary.write(f'connected_ocean_cells: {int(ocean_mask.sum())}\n') + summary.write(f'land_cells: {int(total_cells - ocean_mask.sum())}\n') + summary.write( + f'min_signed_distance_m: {float(np.nanmin(signed_distance))}\n' + ) + summary.write( + f'max_signed_distance_m: {float(np.nanmax(signed_distance))}\n' + ) + summary.write('\n') + + +def _get_plot_info(lon, lat): + """ + Get longitude and latitude arrays for plotting. + """ + plot_lon = np.where(lon > 180.0, lon - 360.0, lon) + center_order = np.argsort(plot_lon) + plot_lon = plot_lon[center_order] + lon_corner = interp_extrap_corner(plot_lon) + lat_corner = interp_extrap_corner(lat) + + return dict( + center_order=center_order, + lon=plot_lon, + lon_corner=lon_corner, + lat=lat, + lat_corner=lat_corner, + ) + + +def _get_plot_info_from_file(filename): + """ + Read grid coordinates from a coastline file and build plot metadata. + """ + with xr.open_dataset(filename) as ds_coastline: + return _get_plot_info( + lon=ds_coastline.lon.values, + lat=ds_coastline.lat.values, + ) + + +def _get_projection(projection): + """ + Get the requested Cartopy projection. + """ + if projection == 'global': + return ccrs.PlateCarree() + if projection == 'antarctic': + return ccrs.SouthPolarStereo() + raise ValueError(f'Unexpected projection: {projection}') + + +def _get_critical_transect_overlays(): + """ + Build line overlays for critical passages and land blockages. + """ + critical_transects = load_default_critical_transects() + return [ + dict( + label='Critical passages', + color='tab:blue', + segments=_get_feature_segments(critical_transects.passages), + ), + dict( + label='Critical blockages', + color='tab:red', + segments=_get_feature_segments(critical_transects.land_blockages), + ), + ] + + +def _configure_axes(ax, projection, antarctic_max_latitude): + """ + Configure common map settings for coastline diagnostics. + """ + ref_projection = ccrs.PlateCarree() + if projection == 'global': + ax.set_global() + elif projection == 'antarctic': + ax.set_extent( + [-180.0, 180.0, -90.0, antarctic_max_latitude], + crs=ref_projection, + ) + else: + raise ValueError(f'Unexpected projection: {projection}') + + ax.gridlines(color='0.8', linestyle=':', linewidth=0.5) + + +def _plot_line_overlays(ax, projection, line_overlays): + """ + Plot optional line overlays and add a legend when present. + """ + ref_projection = ccrs.PlateCarree() + line_width = 1.0 if projection == 'antarctic' else 1.3 + legend_handles = [] + + for overlay in line_overlays: + segments = overlay['segments'] + if len(segments) == 0: + continue + + for lon, lat in segments: + ax.plot( + lon, + lat, + color=overlay['color'], + linewidth=line_width, + transform=ref_projection, + zorder=3, + ) + + legend_handles.append( + Line2D( + [0.0], + [0.0], + color=overlay['color'], + linewidth=2.0, + label=overlay['label'], + ) + ) + + if legend_handles: + location = 'lower left' if projection == 'global' else 'upper left' + ax.legend( + handles=legend_handles, + loc=location, + framealpha=0.9, + ) + + +def _get_feature_segments(feature_collection): + """ + Convert feature-collection line geometries into plot-ready segments. + """ + segments: list[tuple[NDArray[np.float64], NDArray[np.float64]]] = [] + if feature_collection is None: + return segments + + if isinstance(feature_collection, dict): + features = feature_collection.get('features', []) + else: + features = getattr(feature_collection, 'features', []) + + for feature in features: + geometry = feature.get('geometry') + if geometry is None: + continue + + geometry_type = geometry.get('type') + if geometry_type == 'LineString': + parts = [geometry['coordinates']] + elif geometry_type == 'MultiLineString': + parts = geometry['coordinates'] + else: + raise ValueError( + f'Unsupported critical transect geometry type: {geometry_type}' + ) + + for coordinates in parts: + coordinates = np.asarray(coordinates, dtype=np.float64) + lon = np.where( + coordinates[:, 0] > 180.0, + coordinates[:, 0] - 360.0, + coordinates[:, 0], + ) + lat = coordinates[:, 1] + if lon.size > 1: + segments.append((lon, lat)) + + return segments diff --git a/polaris/tasks/mesh/spherical/unified/river/__init__.py b/polaris/tasks/mesh/spherical/unified/river/__init__.py new file mode 100644 index 0000000000..13d8578cde --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/__init__.py @@ -0,0 +1,33 @@ +from polaris.tasks.mesh.spherical.unified.river.base_mesh import ( + PrepareRiverForBaseMeshStep, +) +from polaris.tasks.mesh.spherical.unified.river.lat_lon import ( + PrepareRiverLatLonStep, + build_river_network_dataset, +) +from polaris.tasks.mesh.spherical.unified.river.source import ( + PrepareRiverSourceStep, + simplify_river_network_feature_collection, +) +from polaris.tasks.mesh.spherical.unified.river.steps import ( + get_mesh_unified_river_steps, +) +from polaris.tasks.mesh.spherical.unified.river.task import ( + UnifiedRiverNetworkTask, +) +from polaris.tasks.mesh.spherical.unified.river.tasks import ( + add_river_tasks, +) +from polaris.tasks.mesh.spherical.unified.river.viz import VizRiverStep + +__all__ = [ + 'PrepareRiverLatLonStep', + 'PrepareRiverForBaseMeshStep', + 'PrepareRiverSourceStep', + 'VizRiverStep', + 'UnifiedRiverNetworkTask', + 'add_river_tasks', + 'build_river_network_dataset', + 'get_mesh_unified_river_steps', + 'simplify_river_network_feature_collection', +] diff --git a/polaris/tasks/mesh/spherical/unified/river/base_mesh.py b/polaris/tasks/mesh/spherical/unified/river/base_mesh.py new file mode 100644 index 0000000000..8028508f46 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/base_mesh.py @@ -0,0 +1,566 @@ +import os +from typing import Any + +import numpy as np +import xarray as xr +from mpas_tools.mesh.interpolation import interp_bilin +from shapely.geometry import LineString + +from polaris.step import Step +from polaris.tasks.mesh.spherical.unified.river.lat_lon import ( + build_river_network_dataset, +) +from polaris.tasks.mesh.spherical.unified.river.source import ( + RiverSegment, + _haversine_distance, + _read_geojson, + _write_geojson, + read_river_segments_from_feature_collection, + river_segments_to_feature_collection, +) + + +class PrepareRiverForBaseMeshStep(Step): + """ + Prepare clipped river products for base-mesh consumers. + """ + + def __init__(self, component, prepare_step, coastline_step, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + prepare_step : polaris.tasks.mesh.spherical.unified.river.source. + PrepareRiverSourceStep + The shared source river preparation step + + coastline_step : polaris.tasks.mesh.spherical.unified.coastline. + PrepareCoastlineStep + The shared coastline preparation step + + subdir : str + The subdirectory within the component's work directory + """ + super().__init__( + component=component, + name='river_base_mesh', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.prepare_step = prepare_step + self.coastline_step = coastline_step + self.clipped_filename = 'clipped_river_network.geojson' + self.clipped_outlets_filename = 'clipped_outlets.geojson' + self.masks_filename = 'clipped_river_network.nc' + + def setup(self): + """ + Link simplified river and coastline products and declare outputs. + """ + convention = self.config.get( + 'spherical_mesh', 'antarctic_boundary_convention' + ) + self.add_input_file( + filename='simplified_river_network.geojson', + work_dir_target=os.path.join( + self.prepare_step.path, self.prepare_step.simplified_filename + ), + ) + self.add_input_file( + filename='retained_outlets.geojson', + work_dir_target=os.path.join( + self.prepare_step.path, self.prepare_step.outlets_filename + ), + ) + self.add_input_file( + filename='coastline.nc', + work_dir_target=os.path.join( + self.coastline_step.path, + self.coastline_step.output_filenames[convention], + ), + ) + self.add_output_file(filename=self.clipped_filename) + self.add_output_file(filename=self.clipped_outlets_filename) + self.add_output_file(filename=self.masks_filename) + + def run(self): + """ + Clip and simplify the retained river network for base-mesh use. + """ + section = self.config['river_network'] + river_fc = _read_geojson('simplified_river_network.geojson') + outlet_fc = _read_geojson('retained_outlets.geojson') + with xr.open_dataset('coastline.nc') as ds_coastline: + segments = read_river_segments_from_feature_collection(river_fc) + clipped_segments = condition_base_mesh_river_segments( + segments=segments, + ds_coastline=ds_coastline, + clip_distance_m=1.0e3 + * section.getfloat('base_mesh_clip_distance_km'), + simplify_tolerance_deg=_km_to_equatorial_degrees( + section.getfloat('base_mesh_simplify_tolerance_km') + ), + min_segment_length_m=1.0e3 + * section.getfloat('base_mesh_min_segment_length_km'), + preserve_outlet_stub_m=1.0e3 + * section.getfloat('base_mesh_preserve_outlet_stub_km'), + ) + clipped_fc = river_segments_to_feature_collection(clipped_segments) + clipped_outlets_fc = clip_outlet_feature_collection( + outlet_feature_collection=outlet_fc, + ds_coastline=ds_coastline, + clip_distance_m=1.0e3 + * section.getfloat('base_mesh_clip_distance_km'), + ) + ds_river, _ = build_river_network_dataset( + river_feature_collection=clipped_fc, + outlet_feature_collection=clipped_outlets_fc, + ds_coastline=ds_coastline, + resolution=self.config.getfloat( + 'unified_mesh', 'resolution_latlon' + ), + outlet_match_tolerance=self.config['river_lat_lon'].getfloat( + 'outlet_match_tolerance' + ), + channel_subsegment_fraction=self.config[ + 'river_lat_lon' + ].getfloat('channel_subsegment_fraction'), + channel_buffer_km=self.config['river_lat_lon'].getfloat( + 'channel_buffer_km' + ), + ) + + metadata = dict( + mesh_name=self.config.get('unified_mesh', 'mesh_name'), + resolution_latlon=self.config.getfloat( + 'unified_mesh', 'resolution_latlon' + ), + source_river_step=self.prepare_step.subdir, + source_coastline_step=self.coastline_step.subdir, + clip_distance_m=1.0e3 + * section.getfloat('base_mesh_clip_distance_km'), + simplify_tolerance_km=section.getfloat( + 'base_mesh_simplify_tolerance_km' + ), + min_segment_length_km=section.getfloat( + 'base_mesh_min_segment_length_km' + ), + preserve_outlet_stub_km=section.getfloat( + 'base_mesh_preserve_outlet_stub_km' + ), + ) + clipped_fc['metadata'] = metadata + clipped_outlets_fc['metadata'] = dict(metadata) + ds_river.attrs.update(metadata) + _write_geojson(clipped_fc, self.clipped_filename) + _write_geojson(clipped_outlets_fc, self.clipped_outlets_filename) + ds_river.to_netcdf(self.masks_filename) + + +def condition_base_mesh_river_segments( + segments, + ds_coastline, + clip_distance_m, + simplify_tolerance_deg, + min_segment_length_m, + preserve_outlet_stub_m=0.0, +): + """ + Clip, simplify, and clean river segments for base-mesh use. + + Parameters + ---------- + segments : list of RiverSegment + River segments to condition + + ds_coastline : xarray.Dataset + Coastline dataset with ``signed_distance``, ``lon``, and ``lat`` + variables + + clip_distance_m : float + Distance inland from the coastline at which segments are clipped, + in meters; segment portions seaward of this boundary are removed + + simplify_tolerance_deg : float + Douglas-Peucker simplification tolerance applied after clipping, + in degrees + + min_segment_length_m : float + Minimum retained segment length after simplification, in meters + + preserve_outlet_stub_m : float, optional + If positive, a stub of at least this length is kept for outlet + segments that would otherwise be fully clipped, preserving the + outlet location + + Returns + ------- + list of RiverSegment + Conditioned segments sorted by outlet group, drainage area, and + HydroRIVERS identifier + """ + clipped_segments: list[RiverSegment] = [] + lon = ds_coastline.lon.values.astype(float) + lat = ds_coastline.lat.values.astype(float) + signed_distance = ds_coastline.signed_distance.values.astype(float) + threshold_m = -float(clip_distance_m) + + all_coords = [ + np.asarray(seg.geometry.coords, dtype=float) for seg in segments + ] + seg_sizes = [len(c) for c in all_coords] + if seg_sizes: + all_distances = _interpolate_signed_distance( + coords=np.vstack(all_coords), + lon=lon, + lat=lat, + signed_distance=signed_distance, + ) + per_segment_distances = np.split( + all_distances, np.cumsum(seg_sizes)[:-1] + ) + else: + per_segment_distances = [] + + for segment, coords, point_signed_distance in zip( + segments, all_coords, per_segment_distances, strict=True + ): + clipped_geometries = _clip_line_string_with_distances( + coords=coords, + point_signed_distance=point_signed_distance, + threshold_m=threshold_m, + preserve_outlet_stub_m=preserve_outlet_stub_m, + ) + for geometry in clipped_geometries: + simplified = geometry.simplify( + simplify_tolerance_deg, preserve_topology=False + ) + cleaned = _clean_conditioned_geometry( + simplified, min_segment_length_m + ) + if cleaned is None: + continue + clipped_segments.append( + _conditioned_segment_from_geometry(segment, cleaned) + ) + + return sorted( + clipped_segments, + key=lambda segment: ( + segment.outlet_hyriv_id or -1, + -segment.drainage_area, + segment.hyriv_id, + tuple(segment.geometry.coords[-1]), + ), + ) + + +def clip_outlet_feature_collection( + outlet_feature_collection, + ds_coastline, + clip_distance_m, +) -> dict[str, Any]: + """ + Retain only outlets that remain inland after base-mesh clipping. + + Parameters + ---------- + outlet_feature_collection : dict + A GeoJSON feature collection of outlet points + + ds_coastline : xarray.Dataset + Coastline dataset with ``signed_distance``, ``lon``, and ``lat`` + variables + + clip_distance_m : float + Distance inland from the coastline beyond which outlets are + removed, in meters + + Returns + ------- + dict + A GeoJSON feature collection containing only the outlets whose + signed distance remains within the clipping threshold + """ + lon = ds_coastline.lon.values.astype(float) + lat = ds_coastline.lat.values.astype(float) + signed_distance = ds_coastline.signed_distance.values.astype(float) + threshold_m = -float(clip_distance_m) + all_features = outlet_feature_collection['features'] + + if not all_features: + return dict(type='FeatureCollection', features=[]) + + coords = np.array( + [feature['geometry']['coordinates'] for feature in all_features], + dtype=float, + ) + distances = _interpolate_signed_distance( + coords=coords, + lon=lon, + lat=lat, + signed_distance=signed_distance, + ) + features = [ + feature + for feature, dist in zip(all_features, distances, strict=True) + if dist <= threshold_m + ] + return dict(type='FeatureCollection', features=features) + + +def _conditioned_segment_from_geometry(segment, geometry): + outlet_type = segment.outlet_type + outlet_hyriv_id = segment.outlet_hyriv_id + if geometry.coords[-1] != segment.geometry.coords[-1]: + outlet_type = None + outlet_hyriv_id = None + + return RiverSegment( + geometry=geometry, + hyriv_id=segment.hyriv_id, + main_riv=segment.main_riv, + ord_stra=segment.ord_stra, + drainage_area=segment.drainage_area, + next_down=segment.next_down, + endorheic=segment.endorheic, + river_name=segment.river_name, + outlet_type=outlet_type, + outlet_hyriv_id=outlet_hyriv_id, + ) + + +def _clean_conditioned_geometry(geometry, min_segment_length_m): + """ + Drop degenerate or too-short conditioned geometries. + """ + if not isinstance(geometry, LineString): + return None + + coords = np.asarray(geometry.coords, dtype=float) + if coords.shape[0] < 2: + return None + + deduped = [coords[0]] + for point in coords[1:]: + if not np.allclose(deduped[-1], point): + deduped.append(point) + + if len(deduped) < 2: + return None + + cleaned = LineString(np.asarray(deduped)) + if _line_string_length_m(cleaned) < min_segment_length_m: + return None + + return cleaned + + +def _clip_line_string_with_distances( + coords, + point_signed_distance, + threshold_m, + preserve_outlet_stub_m, +): + clipped_coords = _clip_coords_by_threshold( + coords=coords, + point_signed_distance=point_signed_distance, + threshold_m=threshold_m, + ) + if len(clipped_coords) == 0 and preserve_outlet_stub_m > 0.0: + stub = _build_outlet_stub( + coords=coords, + point_signed_distance=point_signed_distance, + threshold_m=threshold_m, + preserve_outlet_stub_m=preserve_outlet_stub_m, + ) + if stub is not None: + clipped_coords = [stub] + + return [LineString(line) for line in clipped_coords if len(line) >= 2] + + +def _build_outlet_stub( + coords, + point_signed_distance, + threshold_m, + preserve_outlet_stub_m, +): + for start_index in range(coords.shape[0] - 1, 0, -1): + start_point = coords[start_index - 1] + end_point = coords[start_index] + start_distance = float(point_signed_distance[start_index - 1]) + end_distance = float(point_signed_distance[start_index]) + if start_distance <= threshold_m or end_distance <= threshold_m: + continue + + crossing = _interpolate_threshold_crossing( + start_point=start_point, + end_point=end_point, + start_distance=start_distance, + end_distance=end_distance, + threshold_m=threshold_m, + ) + stub = LineString([crossing, end_point]) + if _line_string_length_m(stub) >= preserve_outlet_stub_m: + return np.vstack([crossing, end_point]) + + return None + + +def _clip_coords_by_threshold(coords, point_signed_distance, threshold_m): + clipped_lines: list[np.ndarray] = [] + current_points: list[np.ndarray] = [] + + for start_index in range(coords.shape[0] - 1): + start_point = coords[start_index] + end_point = coords[start_index + 1] + start_distance = float(point_signed_distance[start_index]) + end_distance = float(point_signed_distance[start_index + 1]) + start_valid = start_distance <= threshold_m + end_valid = end_distance <= threshold_m + + if start_valid and not current_points: + _append_point_if_distinct(current_points, start_point) + + if start_valid and end_valid: + _append_point_if_distinct(current_points, end_point) + continue + + if start_valid and not end_valid: + _append_point_if_distinct( + current_points, + _interpolate_threshold_crossing( + start_point=start_point, + end_point=end_point, + start_distance=start_distance, + end_distance=end_distance, + threshold_m=threshold_m, + ), + ) + if len(current_points) >= 2: + clipped_lines.append(np.vstack(current_points)) + current_points = [] + continue + + if not start_valid and end_valid: + _append_point_if_distinct( + current_points, + _interpolate_threshold_crossing( + start_point=start_point, + end_point=end_point, + start_distance=start_distance, + end_distance=end_distance, + threshold_m=threshold_m, + ), + ) + _append_point_if_distinct(current_points, end_point) + + if len(current_points) >= 2: + clipped_lines.append(np.vstack(current_points)) + + return clipped_lines + + +def _line_string_length_m(geometry): + coords = np.asarray(geometry.coords, dtype=float) + if coords.shape[0] < 2: + return 0.0 + + segment_lengths = _haversine_distance( + coords[:-1, 0], + coords[:-1, 1], + coords[1:, 0], + coords[1:, 1], + ) + return float(np.sum(segment_lengths)) + + +def _km_to_equatorial_degrees(distance_km): + return float(distance_km) / 111.0 + + +def _interpolate_signed_distance(coords, lon, lat, signed_distance): + sample_lon = np.asarray(coords[:, 0], dtype=float) + interp_lon, interp_field, sample_lon = ( + _prepare_periodic_longitude_interpolation( + lon=lon, + field=signed_distance, + sample_lon=sample_lon, + ) + ) + sample_lat = np.asarray(coords[:, 1], dtype=float) + return interp_bilin( + x=interp_lon, + y=lat, + field=interp_field, + xCell=sample_lon, + yCell=sample_lat, + ) + + +def _prepare_periodic_longitude_interpolation(lon, field, sample_lon): + lon = np.asarray(lon, dtype=float) + field = np.asarray(field) + sample_lon = np.asarray(sample_lon, dtype=float) + + if lon.ndim != 1 or lon.size < 2: + return lon, field, sample_lon + + lon_spacing = np.diff(lon) + if not np.all(lon_spacing > 0.0): + return lon, field, sample_lon + + spacing = float(np.median(lon_spacing)) + includes_periodic_endpoint = np.isclose( + lon[-1] - lon[0], 360.0, atol=spacing * 0.5 + ) + omits_periodic_endpoint = np.isclose( + lon[-1] - lon[0] + spacing, 360.0, atol=spacing * 0.5 + ) + + if not includes_periodic_endpoint and not omits_periodic_endpoint: + return lon, field, sample_lon + + original_sample_lon = sample_lon.copy() + sample_lon = lon[0] + np.mod(sample_lon - lon[0], 360.0) + + if includes_periodic_endpoint: + wrap_mask = np.isclose(sample_lon, lon[0]) & ( + original_sample_lon > lon[-1] + ) + sample_lon[wrap_mask] = lon[-1] + return lon, field, sample_lon + + lon = np.append(lon, lon[0] + 360.0) + field = np.concatenate([field, field[:, :1]], axis=1) + return lon, field, sample_lon + + +def _interpolate_threshold_crossing( + start_point, + end_point, + start_distance, + end_distance, + threshold_m, +): + if np.isclose(end_distance, start_distance): + fraction = 0.5 + else: + fraction = (threshold_m - start_distance) / ( + end_distance - start_distance + ) + + fraction = float(np.clip(fraction, 0.0, 1.0)) + return start_point + fraction * (end_point - start_point) + + +def _append_point_if_distinct(points, point): + point = np.asarray(point, dtype=np.float64) + if len(points) == 0 or not np.allclose(points[-1], point): + points.append(point) diff --git a/polaris/tasks/mesh/spherical/unified/river/lat_lon.py b/polaris/tasks/mesh/spherical/unified/river/lat_lon.py new file mode 100644 index 0000000000..dcdc59b32a --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/lat_lon.py @@ -0,0 +1,592 @@ +import os + +import numpy as np +import xarray as xr +from scipy.spatial import cKDTree +from shapely.geometry import Point, mapping, shape + +from polaris.step import Step +from polaris.tasks.mesh.spherical.unified.river.source import ( + EARTH_RADIUS, + _haversine_distance, + _read_geojson, + _write_geojson, + read_river_segments_from_feature_collection, +) + + +class PrepareRiverLatLonStep(Step): + """ + Rasterize a simplified river network onto a shared lat-lon target grid. + """ + + def __init__(self, component, prepare_step, coastline_step, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + prepare_step : polaris.tasks.mesh.spherical.unified.river.source. + PrepareRiverSourceStep + The shared source river preparation step + + coastline_step : polaris.tasks.mesh.spherical.unified.coastline. + PrepareCoastlineStep + The shared coastline preparation step + + subdir : str + The subdirectory within the component's work directory + """ + super().__init__( + component=component, + name='river_lat_lon', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.prepare_step = prepare_step + self.coastline_step = coastline_step + self.masks_filename = 'river_network.nc' + self.outlets_filename = 'river_outlets.geojson' + + def setup(self): + """ + Link shared source and coastline inputs and declare outputs. + """ + convention = self.config.get('unified_mesh', 'coastline_convention') + self.add_input_file( + filename='simplified_river_network.geojson', + work_dir_target=os.path.join( + self.prepare_step.path, self.prepare_step.simplified_filename + ), + ) + self.add_input_file( + filename='retained_outlets.geojson', + work_dir_target=os.path.join( + self.prepare_step.path, self.prepare_step.outlets_filename + ), + ) + self.add_input_file( + filename='coastline.nc', + work_dir_target=os.path.join( + self.coastline_step.path, + self.coastline_step.output_filenames[convention], + ), + ) + self.add_output_file(filename=self.masks_filename) + self.add_output_file(filename=self.outlets_filename) + + def run(self): + """ + Build target-grid river masks and snapped outlet diagnostics. + """ + section = self.config['river_lat_lon'] + river_fc = _read_geojson('simplified_river_network.geojson') + outlet_fc = _read_geojson('retained_outlets.geojson') + ds_coastline = xr.open_dataset('coastline.nc') + ds_river, snapped_outlets = build_river_network_dataset( + river_feature_collection=river_fc, + outlet_feature_collection=outlet_fc, + ds_coastline=ds_coastline, + resolution=self.config.getfloat( + 'unified_mesh', 'resolution_latlon' + ), + outlet_match_tolerance=section.getfloat('outlet_match_tolerance'), + channel_subsegment_fraction=section.getfloat( + 'channel_subsegment_fraction' + ), + channel_buffer_km=section.getfloat('channel_buffer_km'), + ) + ds_river.attrs['source_river_step'] = self.prepare_step.subdir + ds_river.attrs['source_coastline_step'] = self.coastline_step.subdir + ds_river.attrs['mesh_name'] = self.config.get( + 'unified_mesh', 'mesh_name' + ) + ds_river.attrs['coastline_convention'] = self.config.get( + 'unified_mesh', 'coastline_convention' + ) + ds_river.to_netcdf(self.masks_filename) + _write_geojson(snapped_outlets, self.outlets_filename) + + +def build_river_network_dataset( + river_feature_collection, + outlet_feature_collection, + ds_coastline, + resolution, + outlet_match_tolerance, + channel_subsegment_fraction=0.5, + channel_buffer_km=0.0, +): + """ + Rasterize a simplified river network onto a regular lat-lon grid. + + Parameters + ---------- + river_feature_collection : dict + A GeoJSON feature collection of simplified river line segments + + outlet_feature_collection : dict + A GeoJSON feature collection of outlet points + + ds_coastline : xarray.Dataset + Coastline dataset with ``lat``, ``lon``, and ``ocean_mask`` + variables + + resolution : float + The target grid resolution in degrees + + outlet_match_tolerance : float + Maximum distance in meters within which an ocean outlet is + snapped to the nearest ocean cell; outlets beyond this distance + fall back to the nearest grid cell + + channel_subsegment_fraction : float, optional + Fraction of the target resolution used as the maximum + inter-sample spacing when densifying river channels for + rasterization + + channel_buffer_km : float, optional + Physical buffer radius in kilometres around each sampled channel + point within which additional grid cells are marked + + Returns + ------- + ds_river : xarray.Dataset + Dataset with ``river_channel_mask``, ``river_outlet_mask``, + ``river_ocean_outlet_mask``, and ``river_inland_sink_mask`` + on the target lat-lon grid + + snapped_outlets : dict + A GeoJSON feature collection of snapped outlet points with + diagnostic snapping properties + """ + river_segments = read_river_segments_from_feature_collection( + river_feature_collection + ) + lat = ds_coastline.lat.values + lon = ds_coastline.lon.values + shape_2d = (lat.size, lon.size) + + river_channel_mask = np.zeros(shape_2d, dtype=np.int8) + river_outlet_mask = np.zeros(shape_2d, dtype=np.int8) + river_ocean_outlet_mask = np.zeros(shape_2d, dtype=np.int8) + river_inland_sink_mask = np.zeros(shape_2d, dtype=np.int8) + channel_buffer_m = channel_buffer_km * 1.0e3 + + for segment in river_segments: + sample_points = _sample_line( + segment.geometry, + resolution=resolution, + subsegment_fraction=channel_subsegment_fraction, + ) + for sample_lon, sample_lat in sample_points: + _mark_channel_sample( + mask=river_channel_mask, + sample_lon=sample_lon, + sample_lat=sample_lat, + lon=lon, + lat=lat, + buffer_m=channel_buffer_m, + ) + + snapped_features = [] + ocean_outlets = 0 + unmatched_ocean_outlets = 0 + ocean_mask = _mask_from_dataset(ds_coastline, 'ocean_mask') + land_mask = _get_land_mask(ds_coastline) + ocean_kdtree, ocean_indices_arr = _build_cell_kdtree(ocean_mask, lon, lat) + land_kdtree, land_indices_arr = _build_cell_kdtree(land_mask, lon, lat) + for feature in outlet_feature_collection['features']: + properties = dict(feature['properties']) + source_point = shape(feature['geometry']) + source_lon = float(source_point.x) + source_lat = float(source_point.y) + outlet_type = properties['outlet_type'] + if outlet_type == 'ocean': + ocean_outlets += 1 + ( + lat_index, + lon_index, + snapped_lon, + snapped_lat, + matched_to_ocean, + snapping_distance, + ) = _match_ocean_outlet( + source_lon=source_lon, + source_lat=source_lat, + lon=lon, + lat=lat, + ocean_kdtree=ocean_kdtree, + ocean_indices_arr=ocean_indices_arr, + outlet_match_tolerance=outlet_match_tolerance, + ) + if not matched_to_ocean: + unmatched_ocean_outlets += 1 + river_ocean_outlet_mask[lat_index, lon_index] = 1 + else: + ( + lat_index, + lon_index, + snapped_lon, + snapped_lat, + snapping_distance, + ) = _match_land_point( + source_lon=source_lon, + source_lat=source_lat, + lon=lon, + lat=lat, + land_kdtree=land_kdtree, + land_indices_arr=land_indices_arr, + ) + matched_to_ocean = False + river_inland_sink_mask[lat_index, lon_index] = 1 + + river_outlet_mask[lat_index, lon_index] = 1 + properties.update( + source_lon=source_lon, + source_lat=source_lat, + snapped_lon=snapped_lon, + snapped_lat=snapped_lat, + snapped_lat_index=int(lat_index), + snapped_lon_index=int(lon_index), + snapping_distance_m=snapping_distance, + matched_to_ocean=matched_to_ocean, + ) + snapped_features.append( + dict( + type='Feature', + properties=properties, + geometry=mapping(Point(snapped_lon, snapped_lat)), + ) + ) + + ds_river = xr.Dataset( + coords=dict(lat=ds_coastline.lat, lon=ds_coastline.lon) + ) + dims = ('lat', 'lon') + ds_river['river_channel_mask'] = xr.DataArray( + river_channel_mask, dims=dims + ) + ds_river['river_outlet_mask'] = xr.DataArray(river_outlet_mask, dims=dims) + ds_river['river_ocean_outlet_mask'] = xr.DataArray( + river_ocean_outlet_mask, dims=dims + ) + ds_river['river_inland_sink_mask'] = xr.DataArray( + river_inland_sink_mask, dims=dims + ) + ds_river.attrs.update( + dict( + target_grid='lat_lon', + target_grid_resolution_degrees=resolution, + outlet_match_tolerance_m=outlet_match_tolerance, + channel_buffer_m=channel_buffer_m, + unmatched_ocean_outlets=unmatched_ocean_outlets, + matched_ocean_outlets=ocean_outlets - unmatched_ocean_outlets, + ) + ) + ds_river['river_channel_mask'].attrs['long_name'] = ( + 'Lat-lon mask for retained river channels' + ) + ds_river['river_outlet_mask'].attrs['long_name'] = ( + 'Lat-lon mask for all retained river outlets and inland sinks' + ) + ds_river['river_ocean_outlet_mask'].attrs['long_name'] = ( + 'Lat-lon mask for retained ocean-draining river outlets' + ) + ds_river['river_inland_sink_mask'].attrs['long_name'] = ( + 'Lat-lon mask for retained inland sinks' + ) + + snapped_outlets = dict( + type='FeatureCollection', + features=snapped_features, + metadata=dict( + mesh_name=ds_river.attrs.get('mesh_name'), + target_grid='lat_lon', + target_grid_resolution_degrees=resolution, + outlet_match_tolerance_m=outlet_match_tolerance, + channel_buffer_m=channel_buffer_m, + unmatched_ocean_outlets=unmatched_ocean_outlets, + ), + ) + return ds_river, snapped_outlets + + +def _mask_from_dataset(ds_coastline, variable_name): + """ + Read an integer mask from a coastline dataset if available. + """ + if variable_name not in ds_coastline: + return None + return ds_coastline[variable_name].values.astype(bool) + + +def _get_land_mask(ds_coastline): + """ + Derive the land mask from ocean_mask. + """ + return np.logical_not(ds_coastline.ocean_mask.values.astype(bool)) + + +def _build_cell_kdtree(mask, lon, lat): + """ + Build a 3-D spherical KD-tree from grid cells where mask is True. + + Converts lon/lat to unit-sphere Cartesian coordinates so that the + nearest neighbour in Euclidean 3-D space equals the nearest neighbour + by great-circle distance. Returns (None, None) when the mask is empty. + """ + if mask is None or not np.any(mask): + return None, None + indices = np.argwhere(mask) + cell_lon_rad = np.deg2rad(lon[indices[:, 1]]) + cell_lat_rad = np.deg2rad(lat[indices[:, 0]]) + xyz = np.column_stack( + [ + np.cos(cell_lat_rad) * np.cos(cell_lon_rad), + np.cos(cell_lat_rad) * np.sin(cell_lon_rad), + np.sin(cell_lat_rad), + ] + ) + return cKDTree(xyz), indices + + +def _sample_line(geometry, resolution, subsegment_fraction): + """ + Sample a line densely enough to rasterize it onto a regular grid. + """ + coords = np.asarray(geometry.coords) + sample_points = [tuple(coords[0])] + max_step = resolution * subsegment_fraction + for start, end in zip(coords[:-1], coords[1:], strict=True): + lon0, lat0 = start + lon1, lat1 = end + delta_lon = _wrapped_longitude_difference(lon1 - lon0) + delta_lat = lat1 - lat0 + segment_extent = max(abs(delta_lon), abs(delta_lat)) + n_steps = max(1, int(np.ceil(segment_extent / max_step))) + for index in range(1, n_steps + 1): + fraction = index / n_steps + sample_points.append( + ( + _wrap_longitude(lon0 + fraction * delta_lon), + lat0 + fraction * delta_lat, + ) + ) + return sample_points + + +def _nearest_grid_index(sample_lon, sample_lat, lon, lat): + """ + Find the nearest lat-lon grid-cell center. + """ + lon_index = _nearest_periodic_index(sample_lon, lon) + lat_index = _nearest_bounded_index(sample_lat, lat) + return lat_index, lon_index + + +def _mark_channel_sample(mask, sample_lon, sample_lat, lon, lat, buffer_m): + """ + Mark grid cells within a physical buffer of a sampled river point. + """ + lat_index, lon_index = _nearest_grid_index( + sample_lon=sample_lon, + sample_lat=sample_lat, + lon=lon, + lat=lat, + ) + mask[lat_index, lon_index] = 1 + if buffer_m <= 0.0: + return + + angular_buffer = buffer_m / EARTH_RADIUS + lat_delta = np.rad2deg(angular_buffer) + cos_lat = max(np.cos(np.deg2rad(sample_lat)), 1.0e-6) + lon_delta = min(180.0, np.rad2deg(angular_buffer / cos_lat)) + lat_indices = _coordinate_window(sample_lat, lat, lat_delta) + lon_indices = _coordinate_window(sample_lon, lon, lon_delta, periodic=True) + + candidate_lat = lat[lat_indices][:, np.newaxis] + candidate_lon = lon[lon_indices][np.newaxis, :] + distances = _haversine_distance( + sample_lon, sample_lat, candidate_lon, candidate_lat + ) + buffered_mask = distances <= buffer_m + mask[np.ix_(lat_indices, lon_indices)] |= buffered_mask.astype(np.int8) + + +def _nearest_bounded_index(value, coord): + """ + Find the nearest index in a regular coordinate array. + """ + if coord.size == 1: + return 0 + + step = coord[1] - coord[0] + raw_index = int(np.rint((value - coord[0]) / step)) + return int(np.clip(raw_index, 0, coord.size - 1)) + + +def _nearest_periodic_index(value, coord): + """ + Find the nearest index in a regular periodic coordinate array. + """ + if coord.size == 1: + return 0 + + step = coord[1] - coord[0] + raw_index = int(np.rint((value - coord[0]) / step)) + return raw_index % coord.size + + +def _coordinate_window(value, coord, delta, periodic=False): + """ + Find candidate coordinate indices within one regular-grid window. + """ + if coord.size == 1: + return np.array([0], dtype=int) + + spacing = abs(coord[1] - coord[0]) + half_width = int(np.ceil(delta / spacing)) + 1 + if periodic: + center = _nearest_periodic_index(value, coord) + indices = np.arange( + center - half_width, center + half_width + 1, dtype=int + ) + return np.unique(indices % coord.size) + + center = _nearest_bounded_index(value, coord) + start = max(0, center - half_width) + stop = min(coord.size, center + half_width + 1) + return np.arange(start, stop, dtype=int) + + +def _match_ocean_outlet( + source_lon, + source_lat, + lon, + lat, + ocean_kdtree, + ocean_indices_arr, + outlet_match_tolerance, +): + """ + Match an ocean outlet to the nearest ocean cell. + """ + if ocean_kdtree is not None: + lon_rad = np.deg2rad(source_lon) + lat_rad = np.deg2rad(source_lat) + source_xyz = [ + np.cos(lat_rad) * np.cos(lon_rad), + np.cos(lat_rad) * np.sin(lon_rad), + np.sin(lat_rad), + ] + _, tree_idx = ocean_kdtree.query(source_xyz) + lat_index = int(ocean_indices_arr[tree_idx, 0]) + lon_index = int(ocean_indices_arr[tree_idx, 1]) + snapped_lon = float(lon[lon_index]) + snapped_lat = float(lat[lat_index]) + snapping_distance = float( + _haversine_distance( + source_lon, source_lat, snapped_lon, snapped_lat + ) + ) + matched_to_ocean = snapping_distance <= outlet_match_tolerance + if matched_to_ocean: + return ( + lat_index, + lon_index, + snapped_lon, + snapped_lat, + True, + snapping_distance, + ) + + lat_index, lon_index = _nearest_grid_index( + sample_lon=source_lon, + sample_lat=source_lat, + lon=lon, + lat=lat, + ) + snapped_lon = float(lon[lon_index]) + snapped_lat = float(lat[lat_index]) + snapping_distance = float( + _haversine_distance(source_lon, source_lat, snapped_lon, snapped_lat) + ) + return ( + lat_index, + lon_index, + snapped_lon, + snapped_lat, + False, + snapping_distance, + ) + + +def _match_land_point( + source_lon, + source_lat, + lon, + lat, + land_kdtree, + land_indices_arr, +): + """ + Match an inland sink to the nearest land cell. + """ + if land_kdtree is not None: + lon_rad = np.deg2rad(source_lon) + lat_rad = np.deg2rad(source_lat) + source_xyz = [ + np.cos(lat_rad) * np.cos(lon_rad), + np.cos(lat_rad) * np.sin(lon_rad), + np.sin(lat_rad), + ] + _, tree_idx = land_kdtree.query(source_xyz) + lat_index = int(land_indices_arr[tree_idx, 0]) + lon_index = int(land_indices_arr[tree_idx, 1]) + snapped_lon = float(lon[lon_index]) + snapped_lat = float(lat[lat_index]) + snapping_distance = float( + _haversine_distance( + source_lon, source_lat, snapped_lon, snapped_lat + ) + ) + return ( + lat_index, + lon_index, + snapped_lon, + snapped_lat, + snapping_distance, + ) + + lat_index, lon_index = _nearest_grid_index( + sample_lon=source_lon, + sample_lat=source_lat, + lon=lon, + lat=lat, + ) + snapped_lon = float(lon[lon_index]) + snapped_lat = float(lat[lat_index]) + snapping_distance = float( + _haversine_distance(source_lon, source_lat, snapped_lon, snapped_lat) + ) + return lat_index, lon_index, snapped_lon, snapped_lat, snapping_distance + + +def _wrapped_longitude_difference(delta_lon): + """ + Wrap a longitude difference into the [-180, 180) interval. + """ + return (delta_lon + 180.0) % 360.0 - 180.0 + + +def _wrap_longitude(lon): + """ + Wrap a longitude into the [-180, 180) interval. + """ + return _wrapped_longitude_difference(lon) diff --git a/polaris/tasks/mesh/spherical/unified/river/source.py b/polaris/tasks/mesh/spherical/unified/river/source.py new file mode 100644 index 0000000000..4c780a3776 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/source.py @@ -0,0 +1,741 @@ +import json +import os +from collections import defaultdict +from dataclasses import dataclass, replace + +import numpy as np +import shapefile +from shapely import line_merge, unary_union +from shapely.geometry import LineString, Point, mapping, shape + +from polaris.archive import extract_zip_subdir +from polaris.constants import get_constant +from polaris.step import Step + +EARTH_RADIUS = get_constant('mean_radius') +KM2_TO_M2 = 1.0e6 + + +@dataclass(frozen=True) +class RiverSegment: + """ + Canonical representation of one river-network segment. + + Attributes + ---------- + geometry : shapely.geometry.LineString + The segment geometry in longitude-latitude coordinates + + hyriv_id : int + The HydroRIVERS unique segment identifier + + main_riv : int + The HydroRIVERS main-river identifier + + ord_stra : int + The Strahler stream order + + drainage_area : float + The upstream drainage area in square meters + + next_down : int + The HydroRIVERS identifier of the next downstream segment + (0 for outlets) + + endorheic : int + 1 if the segment drains to an inland sink, 0 otherwise + + river_name : str or None, optional + The river name, if available + + outlet_type : str or None, optional + The outlet type (``'ocean'`` or ``'inland_sink'``), set when the + segment has been assigned to a basin + + outlet_hyriv_id : int or None, optional + The HydroRIVERS identifier of the outlet for the basin this + segment belongs to + """ + + geometry: LineString + hyriv_id: int + main_riv: int + ord_stra: int + drainage_area: float + next_down: int + endorheic: int + river_name: str | None = None + outlet_type: str | None = None + outlet_hyriv_id: int | None = None + + @property + def endpoint(self) -> tuple[float, float]: + """ + Return the downstream endpoint of the segment. + """ + lon, lat = self.geometry.coords[-1] + return float(lon), float(lat) + + +class PrepareRiverSourceStep(Step): + """ + Prepare a simplified, source-grid-independent river network. + """ + + def __init__(self, component, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + subdir : str + The subdirectory within the component's work directory + """ + super().__init__( + component=component, + name='river_source', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.raw_source_filename = 'source_river_network.geojson' + self.simplified_filename = 'simplified_river_network.geojson' + self.outlets_filename = 'retained_outlets.geojson' + + def setup(self): + """ + Add the HydroRIVERS shapefile archive input and declare outputs. + """ + section = self.config['river_network'] + archive_filename = section.get('hydrorivers_archive_filename') + source_url = section.get('hydrorivers_url') + self.add_input_file( + filename=archive_filename, + target=archive_filename, + url=source_url, + database='river_network', + ) + self.add_output_file(filename=self.raw_source_filename) + self.add_output_file(filename=self.simplified_filename) + self.add_output_file(filename=self.outlets_filename) + + def run(self): + """ + Download, unpack, convert, and simplify HydroRIVERS source data. + """ + section = self.config['river_network'] + archive_filename = section.get('hydrorivers_archive_filename') + shp_directory = section.get('hydrorivers_shp_directory') + shp_filename = section.get('hydrorivers_shp_filename') + _unpack_hydrorivers_archive(archive_filename, shp_directory) + _convert_hydrorivers_shapefile_to_geojson( + shp_filename=os.path.join(shp_directory, shp_filename), + output_filename=self.raw_source_filename, + ) + + feature_collection = _read_geojson(self.raw_source_filename) + simplified_fc, outlets_fc = simplify_river_network_feature_collection( + feature_collection=feature_collection, + drainage_area_threshold=section.getfloat( + 'drainage_area_threshold' + ), + outlet_distance_tolerance=section.getfloat( + 'outlet_distance_tolerance' + ), + tributary_area_ratio=section.getfloat('tributary_area_ratio'), + ) + simplified_fc['metadata'] = dict( + mesh_name=self.config.get('unified_mesh', 'mesh_name'), + resolution_latlon=self.config.getfloat( + 'unified_mesh', 'resolution_latlon' + ), + hydrorivers_url=section.get('hydrorivers_url'), + hydrorivers_archive_filename=archive_filename, + hydrorivers_shp_directory=shp_directory, + hydrorivers_shp_filename=shp_filename, + drainage_area_threshold=section.getfloat( + 'drainage_area_threshold' + ), + outlet_distance_tolerance=section.getfloat( + 'outlet_distance_tolerance' + ), + tributary_area_ratio=section.getfloat('tributary_area_ratio'), + ) + outlets_fc['metadata'] = dict(simplified_fc['metadata']) + _write_geojson(simplified_fc, self.simplified_filename) + _write_geojson(outlets_fc, self.outlets_filename) + + +def _unpack_hydrorivers_archive(archive_filename, data_directory): + """ + Unpack the HydroRIVERS archive if needed. + """ + if os.path.isdir(data_directory): + return + + extract_zip_subdir(archive_filename, data_directory) + + if not os.path.isdir(data_directory): + raise OSError( + 'Unpacked HydroRIVERS archive but did not find the expected ' + f'data directory {data_directory!r}.' + ) + + +def _convert_hydrorivers_shapefile_to_geojson(shp_filename, output_filename): + """ + Convert the HydroRIVERS shapefile to GeoJSON. + """ + reader = shapefile.Reader(shp_filename) + features = [] + for shape_record in reader.iterShapeRecords(): + features.append( + dict( + type='Feature', + properties=shape_record.record.as_dict(), + geometry=shape_record.shape.__geo_interface__, + ) + ) + _write_geojson( + dict(type='FeatureCollection', features=features), output_filename + ) + + +def simplify_river_network_feature_collection( + feature_collection, + drainage_area_threshold, + outlet_distance_tolerance, + tributary_area_ratio=0.05, +): + """ + Simplify a HydroRIVERS-style feature collection. + + Parameters + ---------- + feature_collection : dict + A GeoJSON feature collection + + drainage_area_threshold : float + Minimum retained drainage area in square meters + + outlet_distance_tolerance : float + Minimum retained spacing between non-endorheic outlets in meters + + tributary_area_ratio : float, optional + The minimum tributary-to-main-stem drainage-area ratio for retaining + a nearby tributary at a confluence + + Returns + ------- + simplified_fc : dict + A GeoJSON feature collection for the simplified river network + + outlets_fc : dict + A GeoJSON feature collection for retained outlet points + """ + segments = read_river_segments_from_feature_collection(feature_collection) + segments = [ + segment + for segment in segments + if segment.drainage_area >= drainage_area_threshold + ] + if len(segments) == 0: + raise ValueError( + 'No river-network segments remain after applying the drainage ' + 'area threshold.' + ) + + segment_by_id = {segment.hyriv_id: segment for segment in segments} + _validate_acyclic_downstream_graph(segment_by_id) + + upstream_map: dict[int, list[RiverSegment]] = defaultdict(list) + for segment in segments: + if segment.next_down != 0: + upstream_map[segment.next_down].append(segment) + + outlet_candidates = [ + segment for segment in segments if segment.next_down == 0 + ] + retained_outlets = _filter_outlets( + outlet_candidates, outlet_distance_tolerance + ) + + retained_segments: dict[int, RiverSegment] = {} + for outlet in retained_outlets: + if outlet.endorheic == 1: + outlet_type = 'inland_sink' + else: + outlet_type = 'ocean' + _retain_basin_segments( + segment=replace( + outlet, + outlet_type=outlet_type, + outlet_hyriv_id=outlet.hyriv_id, + ), + upstream_map=upstream_map, + segment_by_id=segment_by_id, + retained_segments=retained_segments, + basin_segments=[], + distance_tolerance=outlet_distance_tolerance, + tributary_area_ratio=tributary_area_ratio, + ) + + simplified_segments = sorted( + retained_segments.values(), + key=lambda segment: ( + segment.outlet_hyriv_id or -1, + -segment.drainage_area, + segment.hyriv_id, + ), + ) + outlet_segments = sorted( + [segment for segment in retained_outlets], + key=lambda segment: (-segment.drainage_area, segment.hyriv_id), + ) + + return ( + river_segments_to_feature_collection(simplified_segments), + outlet_segments_to_feature_collection(outlet_segments), + ) + + +def read_river_segments_from_feature_collection( + feature_collection, +) -> list[RiverSegment]: + """ + Convert a GeoJSON feature collection to canonical river segments. + + Parameters + ---------- + feature_collection : dict + A GeoJSON feature collection with HydroRIVERS-style properties + + Returns + ------- + list of RiverSegment + Canonical river segments; geometries with duplicate HydroRIVERS + identifiers are merged into a single segment + """ + merged_segments: dict[int, RiverSegment] = {} + for feature in feature_collection['features']: + segment = _segment_from_feature(feature) + existing = merged_segments.get(segment.hyriv_id) + if existing is None: + merged_segments[segment.hyriv_id] = segment + continue + + merged_geometry = line_merge( + unary_union([existing.geometry, segment.geometry]) + ) + if not isinstance(merged_geometry, LineString): + merged_geometry = LineString( + np.vstack( + [ + np.asarray(existing.geometry.coords), + np.asarray(segment.geometry.coords), + ] + ) + ) + merged_segments[segment.hyriv_id] = replace( + existing, geometry=merged_geometry + ) + + return list(merged_segments.values()) + + +def river_segments_to_feature_collection(segments): + """ + Convert canonical river segments to a GeoJSON feature collection. + + Parameters + ---------- + segments : list of RiverSegment + Canonical river segments + + Returns + ------- + dict + A GeoJSON feature collection with one feature per segment, + preserving all HydroRIVERS properties + """ + features = [] + for segment in segments: + properties = dict( + hyriv_id=segment.hyriv_id, + main_riv=segment.main_riv, + ord_stra=segment.ord_stra, + drainage_area=segment.drainage_area, + upland_skm=segment.drainage_area / KM2_TO_M2, + next_down=segment.next_down, + endorheic=segment.endorheic, + outlet_type=segment.outlet_type, + outlet_hyriv_id=segment.outlet_hyriv_id, + ) + if segment.river_name is not None: + properties['river_name'] = segment.river_name + features.append( + dict( + type='Feature', + properties=properties, + geometry=mapping(segment.geometry), + ) + ) + return dict(type='FeatureCollection', features=features) + + +def outlet_segments_to_feature_collection(segments): + """ + Convert retained outlet segments to a GeoJSON feature collection. + + Parameters + ---------- + segments : list of RiverSegment + Retained outlet segments + + Returns + ------- + dict + A GeoJSON feature collection of outlet point features + """ + features = [] + for segment in segments: + lon, lat = segment.endpoint + if segment.endorheic == 1: + outlet_type = 'inland_sink' + else: + outlet_type = 'ocean' + features.append( + dict( + type='Feature', + properties=dict( + hyriv_id=segment.hyriv_id, + main_riv=segment.main_riv, + drainage_area=segment.drainage_area, + upland_skm=segment.drainage_area / KM2_TO_M2, + endorheic=segment.endorheic, + outlet_type=outlet_type, + ), + geometry=mapping(Point(lon, lat)), + ) + ) + return dict(type='FeatureCollection', features=features) + + +def _retain_basin_segments( + segment, + upstream_map, + segment_by_id, + retained_segments, + basin_segments, + distance_tolerance, + tributary_area_ratio, +): + """ + Iteratively retain a main stem and major tributaries for one basin. + """ + basin_visited: set[int] = set() + stack = [segment] + + while len(stack) > 0: + current = stack.pop() + if current.hyriv_id not in segment_by_id: + continue + if current.hyriv_id in basin_visited: + continue + + basin_visited.add(current.hyriv_id) + retained_segments[current.hyriv_id] = current + basin_segments.append(current) + + keep_segments = _select_upstream_segments( + segment=current, + upstream_map=upstream_map, + basin_segments=basin_segments, + distance_tolerance=distance_tolerance, + tributary_area_ratio=tributary_area_ratio, + ) + + for upstream in reversed(keep_segments): + annotated = replace( + upstream, + outlet_type=current.outlet_type, + outlet_hyriv_id=current.outlet_hyriv_id, + ) + stack.append(annotated) + + +def _select_upstream_segments( + segment, + upstream_map, + basin_segments, + distance_tolerance, + tributary_area_ratio, +): + """ + Select retained upstream segments for one confluence. + """ + upstream_segments = sorted( + upstream_map.get(segment.hyriv_id, []), + key=lambda upstream: ( + upstream.drainage_area, + upstream.ord_stra, + upstream.hyriv_id, + ), + reverse=True, + ) + if len(upstream_segments) == 0: + return [] + + keep_segments = [upstream_segments[0]] + primary = upstream_segments[0] + for upstream in upstream_segments[1:]: + keep = upstream.drainage_area >= ( + tributary_area_ratio * primary.drainage_area + ) + if not keep: + keep = _distance_to_basin(upstream, basin_segments) >= ( + distance_tolerance + ) + if keep: + keep_segments.append(upstream) + + return keep_segments + + +def _validate_acyclic_downstream_graph(segment_by_id): + """ + Validate that retained segments do not contain NEXT_DOWN cycles. + """ + checked: set[int] = set() + + for hyriv_id in segment_by_id: + if hyriv_id in checked: + continue + + path: list[int] = [] + path_index: dict[int, int] = {} + current_id = hyriv_id + + while current_id in segment_by_id: + if current_id in checked: + break + if current_id in path_index: + cycle = path[path_index[current_id] :] + [current_id] + cycle_text = ' -> '.join(str(item) for item in cycle) + raise ValueError( + 'Cycle detected in retained river-network NEXT_DOWN ' + f'chain: {cycle_text}' + ) + + path_index[current_id] = len(path) + path.append(current_id) + + next_down = segment_by_id[current_id].next_down + if next_down == 0: + break + current_id = next_down + + checked.update(path) + + +def _distance_to_basin(segment, basin_segments): + """ + Compute the minimum distance from a segment to already retained segments. + """ + if len(basin_segments) == 0: + return np.inf + return min( + _minimum_line_distance(segment.geometry, other.geometry) + for other in basin_segments + ) + + +def _minimum_line_distance(line_a, line_b): + """ + Approximate the minimum distance between two lines on the sphere. + """ + coords_a = np.asarray(line_a.coords) + coords_b = np.asarray(line_b.coords) + min_distance = np.inf + for lon_a, lat_a in coords_a: + distances = _haversine_distance( + lon_a, + lat_a, + coords_b[:, 0], + coords_b[:, 1], + ) + min_distance = min(min_distance, float(np.min(distances))) + for lon_b, lat_b in coords_b: + distances = _haversine_distance( + lon_b, + lat_b, + coords_a[:, 0], + coords_a[:, 1], + ) + min_distance = min(min_distance, float(np.min(distances))) + return min_distance + + +def _filter_outlets(outlet_candidates, outlet_distance_tolerance): + """ + Retain large, well-separated outlets while preserving inland sinks. + """ + retained_outlets: list[RiverSegment] = [] + for candidate in sorted( + outlet_candidates, + key=lambda outlet: ( + outlet.drainage_area, + outlet.ord_stra, + outlet.hyriv_id, + ), + reverse=True, + ): + keep = True + lon, lat = candidate.endpoint + for retained in retained_outlets: + if candidate.endorheic == 1 and retained.endorheic == 1: + continue + retained_lon, retained_lat = retained.endpoint + distance = _haversine_distance( + lon, lat, retained_lon, retained_lat + ) + if float(distance) < outlet_distance_tolerance: + keep = False + break + if keep: + retained_outlets.append(candidate) + + return retained_outlets + + +def _segment_from_feature(feature): + """ + Parse one GeoJSON feature into a canonical river segment. + """ + properties = feature['properties'] + geometry = _line_string_from_geojson(feature['geometry']) + hyriv_id = _get_int_property(properties, ('hyriv_id', 'HYRIV_ID')) + main_riv = _get_int_property(properties, ('main_riv', 'MAIN_RIV')) + ord_stra = _get_int_property(properties, ('ord_stra', 'ORD_STRA')) + next_down = _get_int_property(properties, ('next_down', 'NEXT_DOWN')) + endorheic = _get_int_property(properties, ('endorheic', 'ENDORHEIC')) + drainage_area = _get_float_property( + properties, + ('drainage_area',), + default=None, + ) + if drainage_area is None: + drainage_area = _get_float_property( + properties, ('upland_skm', 'UPLAND_SKM') + ) + drainage_area *= KM2_TO_M2 + river_name = _get_str_property( + properties, ('river_name', 'RIVER_NAME', 'RIV_NAME', 'NAME') + ) + + return RiverSegment( + geometry=geometry, + hyriv_id=hyriv_id, + main_riv=main_riv, + ord_stra=ord_stra, + drainage_area=drainage_area, + next_down=next_down, + endorheic=endorheic, + river_name=river_name, + ) + + +def _line_string_from_geojson(geometry): + """ + Convert a GeoJSON line geometry to a single LineString. + """ + geom = shape(geometry) + if isinstance(geom, LineString): + return geom + + merged = line_merge(geom) + if isinstance(merged, LineString): + return merged + + raise ValueError( + f'Unsupported river geometry type {geom.geom_type!r}; expected a ' + 'LineString-compatible geometry.' + ) + + +def _get_int_property(properties, keys): + """ + Read an integer property using one of several candidate keys. + """ + value = _get_property(properties, keys) + return int(value) + + +def _get_float_property(properties, keys, default=np.nan): + """ + Read a float property using one of several candidate keys. + """ + value = _get_property(properties, keys, default=default) + if value is None: + return None + return float(value) + + +def _get_str_property(properties, keys): + """ + Read a string property if available. + """ + value = _get_property(properties, keys, default=None) + if value in (None, ''): + return None + return str(value) + + +def _get_property(properties, keys, default=np.nan): + """ + Read one of several candidate properties from a mapping. + """ + for key in keys: + if key in properties: + return properties[key] + if default is np.nan: + raise ValueError( + f'Could not find any of the required properties {keys!r}.' + ) + return default + + +def _haversine_distance(lon_a, lat_a, lon_b, lat_b): + """ + Compute great-circle distance in meters. + """ + lon_a = np.deg2rad(lon_a) + lat_a = np.deg2rad(lat_a) + lon_b = np.deg2rad(lon_b) + lat_b = np.deg2rad(lat_b) + delta_lon = lon_b - lon_a + delta_lat = lat_b - lat_a + haversine = ( + np.sin(delta_lat / 2.0) ** 2 + + np.cos(lat_a) * np.cos(lat_b) * np.sin(delta_lon / 2.0) ** 2 + ) + return 2.0 * EARTH_RADIUS * np.arcsin(np.sqrt(haversine)) + + +def _read_geojson(filename): + """ + Read a GeoJSON file into a dictionary. + """ + with open(filename, 'r', encoding='utf-8') as infile: + return json.load(infile) + + +def _write_geojson(feature_collection, filename): + """ + Write a GeoJSON feature collection. + """ + with open(filename, 'w', encoding='utf-8') as outfile: + json.dump(feature_collection, outfile, indent=2, sort_keys=True) diff --git a/polaris/tasks/mesh/spherical/unified/river/steps.py b/polaris/tasks/mesh/spherical/unified/river/steps.py new file mode 100644 index 0000000000..fc68023c17 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/steps.py @@ -0,0 +1,149 @@ +import os + +from polaris.mesh.spherical.unified import ( + RIVER_CONFIG_FILENAME, + get_unified_mesh_config, +) +from polaris.step import Step +from polaris.tasks.mesh.spherical.unified.river.base_mesh import ( + PrepareRiverForBaseMeshStep, +) +from polaris.tasks.mesh.spherical.unified.river.lat_lon import ( + PrepareRiverLatLonStep, +) +from polaris.tasks.mesh.spherical.unified.river.source import ( + PrepareRiverSourceStep, +) +from polaris.tasks.mesh.spherical.unified.river.viz import VizRiverStep + + +def get_mesh_unified_river_steps( + component, + coastline_step, + mesh_name, + include_viz=False, +): + """ + Get shared river-network steps for one mesh (source, lat-lon, base-mesh). + + Parameters + ---------- + component : polaris.Component + The component the steps belong to + + coastline_step : polaris.Step + The shared coastline-preparation step for the lat-lon grid + + mesh_name : str + The name of the unified mesh + + include_viz : bool, optional + Whether to include a visualization step + + Returns + ------- + steps : dict of str to polaris.Step + The river steps keyed by step name: source prepare, lat-lon prepare, + optional viz, and base-mesh prepare + + config : polaris.config.PolarisConfigParser + The shared lat-lon river config options + """ + config_filename = RIVER_CONFIG_FILENAME + + source_filepath = os.path.join( + component.name, + 'spherical', + 'unified', + mesh_name, + 'river', + 'source', + config_filename, + ) + source_config = _get_mesh_river_config( + component=component, mesh_name=mesh_name, filepath=source_filepath + ) + source_subdir = os.path.join( + 'spherical', 'unified', mesh_name, 'river', 'source', 'prepare' + ) + prepare_step = component.get_or_create_shared_step( + step_cls=PrepareRiverSourceStep, + subdir=source_subdir, + config=source_config, + config_filename=config_filename, + ) + + lat_lon_filepath = os.path.join( + component.name, + 'spherical', + 'unified', + mesh_name, + 'river', + 'lat_lon', + config_filename, + ) + lat_lon_config = _get_mesh_river_config( + component=component, mesh_name=mesh_name, filepath=lat_lon_filepath + ) + + lat_lon_subdir = os.path.join( + 'spherical', 'unified', mesh_name, 'river', 'lat_lon', 'prepare' + ) + prepare_lat_lon_step = component.get_or_create_shared_step( + step_cls=PrepareRiverLatLonStep, + subdir=lat_lon_subdir, + config=lat_lon_config, + config_filename=config_filename, + prepare_step=prepare_step, + coastline_step=coastline_step, + ) + steps: dict[str, Step] = {} + steps[prepare_step.name] = prepare_step + steps[prepare_lat_lon_step.name] = prepare_lat_lon_step + + if include_viz: + viz_subdir = os.path.join(lat_lon_subdir, 'viz') + viz_step = component.get_or_create_shared_step( + step_cls=VizRiverStep, + subdir=viz_subdir, + config=lat_lon_config, + config_filename=config_filename, + prepare_step=prepare_step, + river_step=prepare_lat_lon_step, + ) + steps[viz_step.name] = viz_step + + base_mesh_filepath = os.path.join( + component.name, + 'spherical', + 'unified', + mesh_name, + 'river', + 'base_mesh', + config_filename, + ) + base_mesh_config = _get_mesh_river_config( + component=component, mesh_name=mesh_name, filepath=base_mesh_filepath + ) + + base_mesh_subdir = os.path.join( + 'spherical', 'unified', mesh_name, 'river', 'base_mesh', 'prepare' + ) + prepare_base_mesh_step = component.get_or_create_shared_step( + step_cls=PrepareRiverForBaseMeshStep, + subdir=base_mesh_subdir, + config=base_mesh_config, + config_filename=config_filename, + prepare_step=prepare_step, + coastline_step=coastline_step, + ) + steps[prepare_base_mesh_step.name] = prepare_base_mesh_step + + return steps, lat_lon_config + + +def _get_mesh_river_config(component, mesh_name, filepath): + if filepath in component.configs: + return component.configs[filepath] + + return get_unified_mesh_config(mesh_name=mesh_name, filepath=filepath) diff --git a/polaris/tasks/mesh/spherical/unified/river/task.py b/polaris/tasks/mesh/spherical/unified/river/task.py new file mode 100644 index 0000000000..6c09d66958 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/task.py @@ -0,0 +1,92 @@ +import os + +from polaris.e3sm.init.topo import format_lat_lon_resolution_name +from polaris.mesh.spherical.unified import ( + RIVER_CONFIG_FILENAME, + get_unified_mesh_config, +) +from polaris.task import Task +from polaris.tasks.e3sm.init import e3sm_init +from polaris.tasks.e3sm.init.topo.combine.step import CombineStep +from polaris.tasks.e3sm.init.topo.combine.steps import ( + get_lat_lon_topo_steps, +) +from polaris.tasks.mesh.spherical.unified.coastline.steps import ( + get_lat_lon_coastline_steps, +) +from polaris.tasks.mesh.spherical.unified.river.steps import ( + get_mesh_unified_river_steps, +) + + +class UnifiedRiverNetworkTask(Task): + """ + A standalone task for preparing all river-network products for one mesh. + """ + + def __init__(self, component, mesh_name): + """ + Create a new task. + + Parameters + ---------- + component : polaris.Component + The component the task belongs to + + mesh_name : str + The name of the unified mesh + """ + config = get_unified_mesh_config(mesh_name=mesh_name) + resolution = config.getfloat('unified_mesh', 'resolution_latlon') + subdir = os.path.join( + 'spherical', 'unified', mesh_name, 'river', 'task' + ) + super().__init__( + component=component, + name=f'river_network_{mesh_name}_task', + subdir=subdir, + ) + + combine_steps, _ = get_lat_lon_topo_steps( + component=e3sm_init, + resolution=resolution, + include_viz=False, + ) + resolution_name = format_lat_lon_resolution_name(resolution) + self.combine_topo_step = combine_steps[ + CombineStep.get_name( + target_grid='lat_lon', resolution_name=resolution_name + ) + ] + self.add_step(self.combine_topo_step, symlink='combine_topo') + + coastline_steps, _ = get_lat_lon_coastline_steps( + component=component, + combine_topo_step=self.combine_topo_step, + resolution=resolution, + include_viz=True, + ) + self.coastline_step = coastline_steps['coastline'] + self.add_step(self.coastline_step, symlink='prepare_coastline') + self.coastline_viz_step = coastline_steps['viz_coastline'] + self.add_step( + self.coastline_viz_step, + symlink='viz_prepare_coastline', + run_by_default=False, + ) + + river_steps, config = get_mesh_unified_river_steps( + component=component, + coastline_step=self.coastline_step, + mesh_name=mesh_name, + include_viz=True, + ) + self.set_shared_config(config, link=RIVER_CONFIG_FILENAME) + self.prepare_source_step = river_steps['river_source'] + self.lat_lon_step = river_steps['river_lat_lon'] + self.lat_lon_viz_step = river_steps['viz_river_network'] + self.base_mesh_step = river_steps['river_base_mesh'] + self.add_step(self.prepare_source_step, symlink='prepare_source') + self.add_step(self.lat_lon_step, symlink='prepare_lat_lon') + self.add_step(self.lat_lon_viz_step, symlink='viz_lat_lon') + self.add_step(self.base_mesh_step, symlink='prepare_base_mesh') diff --git a/polaris/tasks/mesh/spherical/unified/river/tasks.py b/polaris/tasks/mesh/spherical/unified/river/tasks.py new file mode 100644 index 0000000000..da47499a9c --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/tasks.py @@ -0,0 +1,19 @@ +from polaris.mesh.spherical.unified import UNIFIED_MESH_NAMES +from polaris.tasks.mesh.spherical.unified.river.task import ( + UnifiedRiverNetworkTask, +) + + +def add_river_tasks(component): + """ + Add standalone river-network tasks. + + Parameters + ---------- + component : polaris.Component + The mesh component that the tasks belong to + """ + for mesh_name in UNIFIED_MESH_NAMES: + component.add_task( + UnifiedRiverNetworkTask(component=component, mesh_name=mesh_name) + ) diff --git a/polaris/tasks/mesh/spherical/unified/river/viz.py b/polaris/tasks/mesh/spherical/unified/river/viz.py new file mode 100644 index 0000000000..83e46608d6 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/river/viz.py @@ -0,0 +1,488 @@ +import os + +import cartopy.crs as ccrs +import matplotlib.colors as mcolors +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr +from matplotlib.lines import Line2D +from shapely.geometry import LineString, MultiLineString, box, shape + +from polaris.step import Step +from polaris.tasks.mesh.spherical.unified.river.source import _read_geojson +from polaris.viz import use_mplstyle + +CONUS_EXTENT = (-128.0, -65.0, 22.0, 52.0) + +COAST_COLOR = '#222222' +LAND_COLOR = '#c8ced6' +OCEAN_COLOR = '#e7f0f7' +RIVER_COLOR = '#0a6ba8' +MATCHED_COLOR = '#c43c39' +UNMATCHED_COLOR = '#7a2cb8' +INLAND_SINK_COLOR = '#c97800' + + +class VizRiverStep(Step): + """ + A step for visualizing river-network diagnostics. + """ + + def __init__(self, component, prepare_step, river_step, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + prepare_step : polaris.tasks.mesh.spherical.unified.river.source. + PrepareRiverSourceStep + The shared source-level river step + + river_step : polaris.tasks.mesh.spherical.unified.river.lat_lon. + PrepareRiverLatLonStep + The shared lat-lon river step to visualize + + subdir : str + The subdirectory within the component's work directory + """ + super().__init__( + component=component, + name='viz_river_network', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.prepare_step = prepare_step + self.river_step = river_step + self.output_filenames = [ + 'river_network_overview.png', + 'debug_summary.txt', + ] + + def setup(self): + """ + Set up the step in the work directory, including linking inputs. + """ + convention = self.river_step.config.get( + 'unified_mesh', 'coastline_convention' + ) + self.add_input_file( + filename='simplified_river_network.geojson', + work_dir_target=os.path.join( + self.prepare_step.path, self.prepare_step.simplified_filename + ), + ) + self.add_input_file( + filename='retained_outlets.geojson', + work_dir_target=os.path.join( + self.prepare_step.path, self.prepare_step.outlets_filename + ), + ) + self.add_input_file( + filename='river_network.nc', + work_dir_target=os.path.join( + self.river_step.path, self.river_step.masks_filename + ), + ) + self.add_input_file( + filename='river_outlets.geojson', + work_dir_target=os.path.join( + self.river_step.path, self.river_step.outlets_filename + ), + ) + self.add_input_file( + filename='coastline.nc', + work_dir_target=os.path.join( + self.river_step.coastline_step.path, + self.river_step.coastline_step.output_filenames[convention], + ), + ) + for filename in self.output_filenames: + self.add_output_file(filename=filename) + + def run(self): + """ + Run this step. + """ + use_mplstyle() + + dpi = self.config['viz_river_network'].getint('dpi') + + simplified_fc = _read_geojson('simplified_river_network.geojson') + retained_outlets_fc = _read_geojson('retained_outlets.geojson') + snapped_outlets_fc = _read_geojson('river_outlets.geojson') + + with xr.open_dataset('river_network.nc') as ds_river: + lon = ds_river.lon.values + lat = ds_river.lat.values + river_channel_mask = ds_river.river_channel_mask.values > 0 + river_outlet_mask = ds_river.river_outlet_mask.values > 0 + river_ocean_outlet_mask = ( + ds_river.river_ocean_outlet_mask.values > 0 + ) + river_inland_sink_mask = ds_river.river_inland_sink_mask.values > 0 + matched_ocean_outlets = int( + ds_river.attrs['matched_ocean_outlets'] + ) + unmatched_ocean_outlets = int( + ds_river.attrs['unmatched_ocean_outlets'] + ) + + with xr.open_dataset('coastline.nc') as ds_coastline: + ocean_mask = ds_coastline.ocean_mask.values > 0 + land_mask = _get_land_mask(ds_coastline) + + self._plot_network_overview( + lon=lon, + lat=lat, + land_mask=land_mask, + ocean_mask=ocean_mask, + simplified_fc=simplified_fc, + snapped_outlets_fc=snapped_outlets_fc, + dpi=dpi, + out_filename='river_network_overview.png', + ) + + _write_summary( + filename='debug_summary.txt', + simplified_fc=simplified_fc, + retained_outlets_fc=retained_outlets_fc, + snapped_outlets_fc=snapped_outlets_fc, + river_channel_mask=river_channel_mask, + river_outlet_mask=river_outlet_mask, + river_ocean_outlet_mask=river_ocean_outlet_mask, + river_inland_sink_mask=river_inland_sink_mask, + matched_ocean_outlets=matched_ocean_outlets, + unmatched_ocean_outlets=unmatched_ocean_outlets, + ) + + def _plot_network_overview( + self, + lon, + lat, + land_mask, + ocean_mask, + simplified_fc, + snapped_outlets_fc, + dpi, + out_filename, + ): + """ + Plot the simplified river network and snapped outlet classes. + """ + fig, ax, inset_ax = _setup_axes_with_inset(dpi=dpi) + for current_ax in (ax, inset_ax): + _plot_context( + ax=current_ax, + lon=lon, + lat=lat, + land_mask=land_mask, + ocean_mask=ocean_mask, + ) + _plot_lines( + ax=current_ax, + feature_collection=simplified_fc, + color=RIVER_COLOR, + lw=0.9, + ) + _plot_snapped_points( + ax=current_ax, + feature_collection=snapped_outlets_fc, + size=38, + ) + _add_figure_legend(fig=fig, handles=_get_overview_legend_handles()) + ax.set_title('Simplified river network and snapped outlets') + fig.savefig(out_filename, bbox_inches='tight') + plt.close(fig) + + +def _setup_axes_with_inset(dpi): + """ + Create a global Robinson plot with a CONUS inset. + """ + fig = plt.figure(figsize=(16.0, 9.0), dpi=dpi) + ax = fig.add_axes([0.03, 0.12, 0.68, 0.78], projection=ccrs.Robinson()) + ax.set_global() + ax.coastlines(linewidth=0.45, color=COAST_COLOR) + + inset_ax = fig.add_axes( + [0.73, 0.17, 0.24, 0.28], projection=ccrs.PlateCarree() + ) + inset_ax.set_extent(CONUS_EXTENT, crs=ccrs.PlateCarree()) + inset_ax.coastlines(linewidth=0.45, color=COAST_COLOR) + inset_ax.set_title('CONUS inset', fontsize=11, pad=4.0) + + _plot_inset_outline(ax=ax, extent=CONUS_EXTENT) + return fig, ax, inset_ax + + +def _plot_context(ax, lon, lat, land_mask, ocean_mask): + """ + Plot muted land and ocean context layers. + """ + _plot_ocean_mask(ax=ax, lon=lon, lat=lat, ocean_mask=ocean_mask) + _plot_land_mask(ax=ax, lon=lon, lat=lat, land_mask=land_mask) + + +def _plot_inset_outline(ax, extent): + """ + Highlight the inset region on the global map. + """ + ax.add_geometries( + [box(extent[0], extent[2], extent[1], extent[3])], + crs=ccrs.PlateCarree(), + facecolor='none', + edgecolor='#4a5568', + linewidth=1.0, + linestyle='--', + zorder=6, + ) + + +def _plot_land_mask(ax, lon, lat, land_mask): + """ + Plot the land mask as a muted background. + """ + if land_mask is None: + return + background = np.where(land_mask, 1.0, np.nan) + ax.imshow( + background, + origin='lower', + extent=_get_extent(lon=lon, lat=lat), + transform=ccrs.PlateCarree(), + cmap=mcolors.ListedColormap([LAND_COLOR]), + interpolation='nearest', + alpha=0.8, + zorder=0, + ) + + +def _get_land_mask(ds_coastline): + """ + Derive the land mask from ocean_mask. + """ + return ds_coastline.ocean_mask.values <= 0 + + +def _plot_ocean_mask(ax, lon, lat, ocean_mask): + """ + Plot the ocean mask as a muted context layer. + """ + ocean = np.where(ocean_mask, 1.0, np.nan) + ax.imshow( + ocean, + origin='lower', + extent=_get_extent(lon=lon, lat=lat), + transform=ccrs.PlateCarree(), + cmap=mcolors.ListedColormap([OCEAN_COLOR]), + interpolation='nearest', + alpha=0.85, + zorder=-1, + ) + + +def _plot_lines(ax, feature_collection, color, lw): + """ + Plot river lines from a GeoJSON feature collection. + """ + for feature in feature_collection['features']: + geom = shape(feature['geometry']) + if isinstance(geom, LineString): + line_geometries = [geom] + elif isinstance(geom, MultiLineString): + line_geometries = list(geom.geoms) + else: + continue + for line in line_geometries: + coords = np.asarray(line.coords) + ax.plot( + coords[:, 0], + coords[:, 1], + color=color, + linewidth=lw, + alpha=0.95, + transform=ccrs.PlateCarree(), + zorder=2, + ) + + +def _plot_snapped_points(ax, feature_collection, size): + """ + Plot snapped outlet points with colors by outlet type and match status. + """ + for feature in feature_collection['features']: + props = feature['properties'] + geom = shape(feature['geometry']) + color, marker = _get_outlet_style( + outlet_type=props['outlet_type'], + matched=props['matched_to_ocean'], + ) + ax.scatter( + [float(geom.x)], + [float(geom.y)], + color=color, + marker=marker, + s=size, + edgecolors='white', + linewidths=0.4, + transform=ccrs.PlateCarree(), + zorder=5, + ) + + +def _get_outlet_style(outlet_type, matched): + """ + Get marker styling for a snapped outlet. + """ + if outlet_type == 'inland_sink': + return INLAND_SINK_COLOR, '^' + if matched: + return MATCHED_COLOR, 'o' + return UNMATCHED_COLOR, 's' + + +def _get_overview_legend_handles(): + """ + Build legend handles for the overview figure. + """ + return [ + Line2D( + [0], + [0], + color=RIVER_COLOR, + linewidth=1.4, + label='Simplified river network', + ), + Line2D( + [0], + [0], + marker='o', + linestyle='None', + color=MATCHED_COLOR, + markerfacecolor=MATCHED_COLOR, + markersize=8, + label='Matched ocean outlet', + ), + Line2D( + [0], + [0], + marker='s', + linestyle='None', + color=UNMATCHED_COLOR, + markerfacecolor=UNMATCHED_COLOR, + markersize=8, + label='Unmatched ocean outlet', + ), + Line2D( + [0], + [0], + marker='^', + linestyle='None', + color=INLAND_SINK_COLOR, + markerfacecolor=INLAND_SINK_COLOR, + markersize=8, + label='Inland sink outlet', + ), + ] + + +def _add_figure_legend(fig, handles): + """ + Add a shared legend below the main map. + """ + fig.legend( + handles=handles, + loc='lower center', + bbox_to_anchor=(0.5, 0.02), + ncol=2, + frameon=True, + ) + + +def _get_extent(lon, lat): + """ + Get an image extent from regular lat-lon cell centers. + """ + lon_edges = _compute_edges(lon) + lat_edges = _compute_edges(lat) + return [lon_edges[0], lon_edges[-1], lat_edges[0], lat_edges[-1]] + + +def _compute_edges(values): + """ + Compute cell edges from regular cell-center coordinates. + """ + values = np.asarray(values) + if values.size == 1: + delta = 1.0 + return np.array([values[0] - 0.5 * delta, values[0] + 0.5 * delta]) + midpoints = 0.5 * (values[:-1] + values[1:]) + first = values[0] - 0.5 * (values[1] - values[0]) + last = values[-1] + 0.5 * (values[-1] - values[-2]) + return np.concatenate(([first], midpoints, [last])) + + +def _write_summary( + filename, + simplified_fc, + retained_outlets_fc, + snapped_outlets_fc, + river_channel_mask, + river_outlet_mask, + river_ocean_outlet_mask, + river_inland_sink_mask, + matched_ocean_outlets, + unmatched_ocean_outlets, +): + """ + Write a compact text summary of river-network diagnostics. + """ + drainage_areas = [ + feature['properties']['drainage_area'] + for feature in simplified_fc['features'] + ] + snapping_distances = [ + feature['properties']['snapping_distance_m'] + for feature in snapped_outlets_fc['features'] + ] + with open(filename, 'w', encoding='utf-8') as summary: + summary.write('River Network Diagnostics\n') + summary.write('=========================\n\n') + summary.write( + f'Simplified segments: {len(simplified_fc["features"])}\n' + ) + summary.write( + f'Retained outlets: {len(retained_outlets_fc["features"])}\n' + ) + summary.write( + f'Rasterized channel cells: {int(np.sum(river_channel_mask))}\n' + ) + summary.write( + f'Rasterized outlet cells: {int(np.sum(river_outlet_mask))}\n' + ) + summary.write( + 'Rasterized ocean-outlet cells: ' + f'{int(np.sum(river_ocean_outlet_mask))}\n' + ) + summary.write( + 'Rasterized inland-sink cells: ' + f'{int(np.sum(river_inland_sink_mask))}\n' + ) + summary.write(f'Matched ocean outlets: {matched_ocean_outlets}\n') + summary.write(f'Unmatched ocean outlets: {unmatched_ocean_outlets}\n') + if len(drainage_areas) > 0: + summary.write( + 'Drainage area range (km^2): ' + f'{min(drainage_areas) / 1.0e6:.1f} to ' + f'{max(drainage_areas) / 1.0e6:.1f}\n' + ) + if len(snapping_distances) > 0: + summary.write( + 'Snapping distance range (km): ' + f'{min(snapping_distances) / 1.0e3:.1f} to ' + f'{max(snapping_distances) / 1.0e3:.1f}\n' + ) diff --git a/polaris/tasks/mesh/spherical/unified/sizing_field/__init__.py b/polaris/tasks/mesh/spherical/unified/sizing_field/__init__.py new file mode 100644 index 0000000000..c6aa8e10c4 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/sizing_field/__init__.py @@ -0,0 +1,29 @@ +from polaris.tasks.mesh.spherical.unified.sizing_field.build import ( + BuildSizingFieldStep, + sizing_field_dataset, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.configs import ( + get_sizing_field_config, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.steps import ( + get_lat_lon_sizing_field_steps, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.task import ( + SizingFieldTask, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.tasks import ( + add_sizing_field_tasks, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.viz import ( + VizSizingFieldStep, +) + +__all__ = [ + 'BuildSizingFieldStep', + 'SizingFieldTask', + 'VizSizingFieldStep', + 'add_sizing_field_tasks', + 'sizing_field_dataset', + 'get_lat_lon_sizing_field_steps', + 'get_sizing_field_config', +] diff --git a/polaris/tasks/mesh/spherical/unified/sizing_field/build.py b/polaris/tasks/mesh/spherical/unified/sizing_field/build.py new file mode 100644 index 0000000000..cbd41db276 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/sizing_field/build.py @@ -0,0 +1,434 @@ +import os + +import numpy as np +import xarray as xr + +from polaris.mesh.spherical.unified import get_unified_mesh_family +from polaris.step import Step + + +class BuildSizingFieldStep(Step): + """ + Build a unified sizing field on a shared lat-lon target grid. + + The step reads the prepared coastline and river-network products for one + mesh, delegates ocean-background construction to the mesh family, and + writes a multi-variable sizing-field dataset to ``sizing_field.nc``. + + Attributes + ---------- + coastline_step : polaris.Step + The shared coastline-preparation step whose outputs are consumed + + river_step : polaris.Step + The shared lat-lon river-network step whose masks are consumed + + sizing_field_filename : str + Name of the output sizing-field NetCDF file + """ + + def __init__(self, component, coastline_step, river_step, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + coastline_step : polaris.Step + The shared coastline-preparation step for the lat-lon grid + + river_step : polaris.Step + The shared lat-lon river-network step + + subdir : str + The subdirectory within the component work directory where the + step runs + """ + super().__init__( + component=component, + name='sizing_field', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.coastline_step = coastline_step + self.river_step = river_step + self.sizing_field_filename = 'sizing_field.nc' + + def setup(self): + """ + Link coastline and river products, configure the mesh family, and + declare outputs. + + This method links the coastline convention file and river-network mask + file as inputs, calls the mesh-family hook to register any additional + family-specific inputs (such as SO-region masks), and declares the + output sizing-field file. + """ + convention = self.config.get('unified_mesh', 'coastline_convention') + self.add_input_file( + filename='coastline.nc', + work_dir_target=os.path.join( + self.coastline_step.path, + self.coastline_step.output_filenames[convention], + ), + ) + self.add_input_file( + filename='river_network.nc', + work_dir_target=os.path.join( + self.river_step.path, self.river_step.masks_filename + ), + ) + self._get_mesh_family().setup_sizing_field_step(self) + self.add_output_file(filename=self.sizing_field_filename) + + def run(self): + """ + Build the sizing-field dataset and write it to NetCDF. + + Opens the coastline and river-network files, constructs the + ocean-background field via the mesh family, calls + :py:func:`sizing_field_dataset` to compose the full sizing field, and + writes the result to ``sizing_field.nc``. + """ + section = self.config['sizing_field'] + unified_section = self.config['unified_mesh'] + + with xr.open_dataset('coastline.nc') as ds_coastline: + with xr.open_dataset('river_network.nc') as ds_river: + ocean_background = self._get_ocean_background( + ds_coastline=ds_coastline, section=section + ) + ds_sizing = sizing_field_dataset( + ds_coastline=ds_coastline, + ds_river=ds_river, + resolution=unified_section.getfloat('resolution_latlon'), + mesh_name=unified_section.get('mesh_name'), + ocean_background=ocean_background, + land_background_km=section.getfloat('land_background_km'), + enable_coastline_refinement=section.getboolean( + 'enable_coastline_refinement' + ), + coastline_transition_land_km=section.getfloat( + 'coastline_transition_land_km' + ), + enable_river_channel_refinement=section.getboolean( + 'enable_river_channel_refinement' + ), + river_channel_km=section.getfloat('river_channel_km'), + enable_river_outlet_refinement=section.getboolean( + 'enable_river_outlet_refinement' + ), + river_outlet_km=section.getfloat('river_outlet_km'), + ) + + ds_sizing.attrs['source_coastline_step'] = self.coastline_step.subdir + ds_sizing.attrs['source_river_step'] = self.river_step.subdir + ds_sizing.to_netcdf(self.sizing_field_filename) + + def _get_ocean_background(self, ds_coastline, section): + """ + Build the ocean background field for the shared target grid. + """ + return self._get_mesh_family().build_ocean_background( + ds_coastline=ds_coastline, section=section + ) + + def _get_mesh_family(self): + """ + Get the unified-mesh family object for this step's shared config. + """ + return get_unified_mesh_family(self.config) + + +def sizing_field_dataset( + ds_coastline, + ds_river, + resolution, + mesh_name, + ocean_background, + land_background_km=240.0, + enable_coastline_refinement=True, + coastline_transition_land_km=0.0, + enable_river_channel_refinement=True, + river_channel_km=240.0, + enable_river_outlet_refinement=True, + river_outlet_km=240.0, +): + """ + Build the unified sizing-field dataset from coastline and river products. + + Parameters + ---------- + ds_coastline : xarray.Dataset + Coastline dataset with ``ocean_mask``, ``signed_distance``, ``lat``, + and ``lon`` variables on a shared lat-lon target grid + + ds_river : xarray.Dataset + River-network dataset with ``river_channel_mask`` and + ``river_outlet_mask`` on the same lat-lon grid as ``ds_coastline`` + + resolution : float + The latitude-longitude grid resolution in degrees; stored as a + dataset attribute + + mesh_name : str + The name of the unified mesh; stored as a dataset attribute + + ocean_background : array-like + 2-D array of shape ``(n_lat, n_lon)`` giving the target cell width in + km for open-ocean cells + + land_background_km : float, optional + Target cell width in km for land cells that are not modified by any + refinement control + + enable_coastline_refinement : bool, optional + Whether to apply the coastline-proximity refinement + + coastline_transition_land_km : float, optional + Width in km of the linear-transition zone on the land side of the + coastline; 0 means apply the ocean background only at the coastline + raster edge + + enable_river_channel_refinement : bool, optional + Whether to refine cells on the river-channel mask + + river_channel_km : float, optional + Target cell width in km on river-channel mask cells + + enable_river_outlet_refinement : bool, optional + Whether to refine cells on the river-outlet mask + + river_outlet_km : float, optional + Target cell width in km on river-outlet mask cells + + Returns + ------- + ds_sizing : xarray.Dataset + Dataset on the shared lat-lon grid containing ``cellWidth`` (the + final cell-width field in km) and diagnostic variables that record + each intermediate sizing stage and the active control index + """ + _validate_shared_grid(ds_coastline=ds_coastline, ds_river=ds_river) + + lat = ds_coastline.lat.values + lon = ds_coastline.lon.values + dims = ('lat', 'lon') + ocean_background = np.asarray(ocean_background, dtype=float) + expected_shape = (lat.size, lon.size) + if ocean_background.shape != expected_shape: + raise ValueError( + 'Ocean background must have shape ' + f'{expected_shape}, got {ocean_background.shape}.' + ) + + ocean_mask = ds_coastline.ocean_mask.values.astype(bool) + signed_distance = ds_coastline.signed_distance.values.astype(float) + river_channel_mask = ds_river.river_channel_mask.values.astype(bool) + river_outlet_mask = ds_river.river_outlet_mask.values.astype(bool) + land_background = np.full(ocean_background.shape, land_background_km) + background = np.where(ocean_mask, ocean_background, land_background) + + river_channel_candidate = _build_mask_candidate( + background=land_background, + mask=river_channel_mask, + enabled=enable_river_channel_refinement, + target_km=river_channel_km, + ) + river_outlet_candidate = _build_mask_candidate( + background=land_background, + mask=river_outlet_mask, + enabled=enable_river_outlet_refinement, + target_km=river_outlet_km, + ) + + land_river_stack = np.stack( + [ + land_background, + river_channel_candidate, + river_outlet_candidate, + ], + axis=0, + ) + land_river_control = np.argmin(land_river_stack, axis=0).astype(np.int8) + land_river_cell_width = np.take_along_axis( + land_river_stack, land_river_control[np.newaxis, ...], axis=0 + )[0] + land_river_control = np.array([0, 2, 3], dtype=np.int8)[land_river_control] + + composed_background = np.where( + ocean_mask, ocean_background, land_river_cell_width + ) + coastline_candidate = _build_coastline_candidate( + background=composed_background, + ocean_background=ocean_background, + ocean_mask=ocean_mask, + signed_distance=signed_distance, + enabled=enable_coastline_refinement, + land_transition_km=coastline_transition_land_km, + ) + final_cell_width = coastline_candidate + + active_control = np.where(ocean_mask, 0, land_river_control) + active_control = active_control.astype(np.int8) + coastline_active = ~np.isclose(coastline_candidate, composed_background) + active_control[coastline_active] = 1 + coastal_transition_delta = final_cell_width - composed_background + + ds_sizing = xr.Dataset( + coords=dict(lat=ds_coastline.lat, lon=ds_coastline.lon) + ) + data_vars = { + 'cellWidth': final_cell_width.astype(np.float32), + 'background_cell_width': background.astype(np.float32), + 'ocean_background_cell_width': ocean_background.astype(np.float32), + 'land_river_cell_width': land_river_cell_width.astype(np.float32), + 'pre_coastline_cell_width': composed_background.astype(np.float32), + 'coastline_cell_width': coastline_candidate.astype(np.float32), + 'coastal_transition_delta': coastal_transition_delta.astype( + np.float32 + ), + 'river_channel_cell_width': river_channel_candidate.astype(np.float32), + 'river_outlet_cell_width': river_outlet_candidate.astype(np.float32), + 'active_control': active_control, + } + for var_name, values in data_vars.items(): + ds_sizing[var_name] = xr.DataArray(values, dims=dims) + + ds_sizing.attrs.update( + dict( + mesh_name=mesh_name, + target_grid='lat_lon', + target_grid_resolution_degrees=resolution, + active_control_meanings=( + '0=background 1=coastline 2=river_channel 3=river_outlet' + ), + ) + ) + _add_mask_candidate_attrs( + ds_sizing=ds_sizing, + prefix='river_channel', + mask=river_channel_mask, + candidate=river_channel_candidate, + background=land_background, + ) + _add_mask_candidate_attrs( + ds_sizing=ds_sizing, + prefix='river_outlet', + mask=river_outlet_mask, + candidate=river_outlet_candidate, + background=land_background, + ) + ds_sizing.cellWidth.attrs['units'] = 'km' + ds_sizing.background_cell_width.attrs['units'] = 'km' + ds_sizing.ocean_background_cell_width.attrs['units'] = 'km' + ds_sizing.land_river_cell_width.attrs['units'] = 'km' + ds_sizing.pre_coastline_cell_width.attrs['units'] = 'km' + ds_sizing.coastline_cell_width.attrs['units'] = 'km' + ds_sizing.coastal_transition_delta.attrs['units'] = 'km' + ds_sizing.river_channel_cell_width.attrs['units'] = 'km' + ds_sizing.river_outlet_cell_width.attrs['units'] = 'km' + ds_sizing.active_control.attrs['long_name'] = ( + 'Index of the active sizing control' + ) + + return ds_sizing + + +def _validate_shared_grid(ds_coastline, ds_river): + """ + Validate that coastline and river products share the same target grid. + """ + for coord_name in ['lat', 'lon']: + if ( + coord_name not in ds_coastline.coords + or coord_name not in ds_river.coords + ): + raise ValueError(f'Missing required coordinate {coord_name!r}.') + + if not np.array_equal( + ds_coastline[coord_name].values, ds_river[coord_name].values + ): + raise ValueError( + 'Coastline and river products must share the same lat-lon ' + f'grid. Mismatch found in coordinate {coord_name!r}.' + ) + + +def _add_mask_candidate_attrs(ds_sizing, prefix, mask, candidate, background): + """ + Add provenance counts for a mask-based sizing candidate. + """ + ds_sizing.attrs[f'{prefix}_mask_count'] = int(np.count_nonzero(mask)) + ds_sizing.attrs[f'{prefix}_finer_than_background_count'] = int( + np.count_nonzero(mask & (candidate < background)) + ) + ds_sizing.attrs[f'{prefix}_equal_to_background_count'] = int( + np.count_nonzero(mask & np.isclose(candidate, background)) + ) + ds_sizing.attrs[f'{prefix}_coarser_than_background_count'] = int( + np.count_nonzero(mask & (candidate > background)) + ) + + +def _build_mask_candidate(background, mask, enabled, target_km): + """ + Build a mask-based candidate field in km. + """ + if not enabled: + return background.copy() + + return np.where(mask, target_km, background).astype(float) + + +def _build_coastline_candidate( + background, + ocean_background, + ocean_mask, + signed_distance, + enabled, + land_transition_km, +): + """ + Build the coastline candidate field in km. + """ + candidate = background.copy() + if not enabled: + return candidate + + land_transition_m = land_transition_km * 1.0e3 + land_side = np.logical_not(ocean_mask) & (signed_distance <= 0.0) + + _apply_transition( + candidate=candidate, + target=ocean_background.astype(float), + background=background, + signed_distance=np.abs(signed_distance), + mask=land_side, + transition_m=land_transition_m, + ) + return candidate + + +def _apply_transition( + candidate, target, background, signed_distance, mask, transition_m +): + """ + Apply a linear transition from the coastline target back to background. + """ + if transition_m < 0.0: + raise ValueError('Transition widths must be nonnegative.') + + if transition_m == 0.0: + zero_mask = mask & np.isclose(signed_distance, 0.0) + candidate[zero_mask] = target[zero_mask] + return + + transition_mask = mask & (signed_distance <= transition_m) + fraction = signed_distance[transition_mask] / transition_m + candidate[transition_mask] = target[transition_mask] + fraction * ( + background[transition_mask] - target[transition_mask] + ) diff --git a/polaris/tasks/mesh/spherical/unified/sizing_field/configs.py b/polaris/tasks/mesh/spherical/unified/sizing_field/configs.py new file mode 100644 index 0000000000..823883ec26 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/sizing_field/configs.py @@ -0,0 +1,40 @@ +from polaris.mesh.spherical.unified.configs import ( + UNIFIED_MESH_NAMES, + get_unified_mesh_config, +) + + +def get_sizing_field_config(mesh_name, filepath=None): + """ + Load one unified mesh config with sizing-field defaults. + + Parameters + ---------- + mesh_name : str + The name of the unified mesh; must be a member of + :py:data:`polaris.mesh.spherical.unified.UNIFIED_MESH_NAMES` + + filepath : str, optional + The path to the config file; passed directly to + :py:func:`polaris.mesh.spherical.unified.get_unified_mesh_config` + + Returns + ------- + config : polaris.config.PolarisConfigParser + The merged config that combines the generic ``unified_mesh.cfg`` + defaults, the shared sizing-field defaults, and the named-mesh + config file + + Raises + ------ + ValueError + If ``mesh_name`` is not a recognized unified mesh name + """ + if mesh_name not in UNIFIED_MESH_NAMES: + valid_mesh_names = ', '.join(UNIFIED_MESH_NAMES) + raise ValueError( + f'Unexpected unified mesh {mesh_name!r}. Valid mesh names ' + f'are: {valid_mesh_names}' + ) + + return get_unified_mesh_config(mesh_name=mesh_name, filepath=filepath) diff --git a/polaris/tasks/mesh/spherical/unified/sizing_field/steps.py b/polaris/tasks/mesh/spherical/unified/sizing_field/steps.py new file mode 100644 index 0000000000..6234f70973 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/sizing_field/steps.py @@ -0,0 +1,101 @@ +import os + +from polaris.step import Step +from polaris.tasks.mesh.spherical.unified.sizing_field.build import ( + BuildSizingFieldStep, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.configs import ( + get_sizing_field_config, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.viz import ( + VizSizingFieldStep, +) + + +def get_lat_lon_sizing_field_steps( + component, + coastline_step, + river_step, + mesh_name, + include_viz=False, +): + """ + Get shared sizing-field steps for one named mesh. + + Parameters + ---------- + component : polaris.Component + The component the steps belong to + + coastline_step : polaris.Step + The shared coastline-preparation step for the lat-lon grid + + river_step : polaris.Step + The shared lat-lon river-network step whose masks are consumed by + the sizing-field build step + + mesh_name : str + The name of the unified mesh + + include_viz : bool, optional + Whether to include a visualization step + + Returns + ------- + steps : dict of str to polaris.Step + The sizing-field steps keyed by step name: the build step and, + if ``include_viz`` is ``True``, the visualization step + + config : polaris.config.PolarisConfigParser + The shared sizing-field config options + """ + config_filename = 'sizing_field.cfg' + filepath = os.path.join( + component.name, + 'spherical', + 'unified', + mesh_name, + 'sizing_field', + config_filename, + ) + config = _get_lat_lon_sizing_field_config( + component=component, mesh_name=mesh_name, filepath=filepath + ) + + subdir = os.path.join( + 'spherical', + 'unified', + mesh_name, + 'sizing_field', + 'build', + ) + build_step = component.get_or_create_shared_step( + step_cls=BuildSizingFieldStep, + subdir=subdir, + config=config, + config_filename=config_filename, + coastline_step=coastline_step, + river_step=river_step, + ) + steps: dict[str, Step] = {} + steps[build_step.name] = build_step + + if include_viz: + viz_subdir = os.path.join(subdir, 'viz') + viz_step = component.get_or_create_shared_step( + step_cls=VizSizingFieldStep, + subdir=viz_subdir, + config=config, + config_filename=config_filename, + sizing_step=build_step, + ) + steps[viz_step.name] = viz_step + + return steps, config + + +def _get_lat_lon_sizing_field_config(component, mesh_name, filepath): + if filepath in component.configs: + return component.configs[filepath] + + return get_sizing_field_config(mesh_name=mesh_name, filepath=filepath) diff --git a/polaris/tasks/mesh/spherical/unified/sizing_field/task.py b/polaris/tasks/mesh/spherical/unified/sizing_field/task.py new file mode 100644 index 0000000000..64c79779c8 --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/sizing_field/task.py @@ -0,0 +1,118 @@ +import os + +from polaris.e3sm.init.topo.resolutions import format_lat_lon_resolution_name +from polaris.mesh.spherical.unified import get_unified_mesh_config +from polaris.task import Task +from polaris.tasks.e3sm.init import e3sm_init +from polaris.tasks.e3sm.init.topo.combine.step import CombineStep +from polaris.tasks.e3sm.init.topo.combine.steps import ( + get_lat_lon_topo_steps, +) +from polaris.tasks.mesh.spherical.unified.coastline.steps import ( + get_lat_lon_coastline_steps, +) +from polaris.tasks.mesh.spherical.unified.river.steps import ( + get_mesh_unified_river_steps, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.steps import ( + get_lat_lon_sizing_field_steps, +) + + +class SizingFieldTask(Task): + """ + A standalone task for building one named unified sizing field. + + The task layers shared dependencies in a fixed order: + + 1. The shared lat-lon combined-topography step from the ``e3sm/init`` + component (see :ref:`dev-e3sm-init-topo-combine-tasks`), exposed as + ``combine_topo``. + 2. The shared coastline-preparation steps from + :py:func:`polaris.tasks.mesh.spherical.unified.coastline.\ +get_lat_lon_coastline_steps`, exposed as ``prepare_coastline`` and the + optional ``viz_prepare_coastline`` (not run by default). + 3. The shared river-network steps from + :py:func:`polaris.tasks.mesh.spherical.unified.river.\ +get_mesh_unified_river_steps`, with the source-prepare step exposed as + ``prepare_river_source`` and the lat-lon step as + ``prepare_river_lat_lon``. + 4. The shared sizing-field steps from + :py:func:`polaris.tasks.mesh.spherical.unified.sizing_field.\ +get_lat_lon_sizing_field_steps`, exposed by their step names. + + Parameters + ---------- + component : polaris.Component + The component the task belongs to + + mesh_name : str + The name of the unified mesh + """ + + def __init__(self, component, mesh_name): + config = get_unified_mesh_config(mesh_name=mesh_name) + resolution = config.getfloat('unified_mesh', 'resolution_latlon') + subdir = os.path.join( + 'spherical', + 'unified', + mesh_name, + 'sizing_field', + 'task', + ) + super().__init__( + component=component, + name=f'sizing_field_{mesh_name}_task', + subdir=subdir, + ) + + combine_steps, _ = get_lat_lon_topo_steps( + component=e3sm_init, + resolution=resolution, + include_viz=False, + ) + resolution_name = format_lat_lon_resolution_name(resolution) + self.combine_topo_step = combine_steps[ + CombineStep.get_name( + target_grid='lat_lon', resolution_name=resolution_name + ) + ] + self.add_step(self.combine_topo_step, symlink='combine_topo') + + coastline_steps, _ = get_lat_lon_coastline_steps( + component=component, + combine_topo_step=self.combine_topo_step, + resolution=resolution, + include_viz=True, + ) + self.coastline_step = coastline_steps['coastline'] + self.add_step(self.coastline_step, symlink='prepare_coastline') + self.coastline_viz_step = coastline_steps['viz_coastline'] + self.add_step( + self.coastline_viz_step, + symlink='viz_prepare_coastline', + run_by_default=False, + ) + + river_steps, _ = get_mesh_unified_river_steps( + component=component, + coastline_step=self.coastline_step, + mesh_name=mesh_name, + include_viz=False, + ) + self.prepare_source_step = river_steps['river_source'] + self.add_step(self.prepare_source_step, symlink='prepare_river_source') + self.river_step = river_steps['river_lat_lon'] + self.add_step(self.river_step, symlink='prepare_river_lat_lon') + + sizing_steps, config = get_lat_lon_sizing_field_steps( + component=component, + coastline_step=self.coastline_step, + river_step=self.river_step, + mesh_name=mesh_name, + include_viz=True, + ) + self.set_shared_config(config, link='sizing_field.cfg') + for step in sizing_steps.values(): + symlink = os.path.basename(step.subdir) + self.add_step(step, symlink=symlink) diff --git a/polaris/tasks/mesh/spherical/unified/sizing_field/tasks.py b/polaris/tasks/mesh/spherical/unified/sizing_field/tasks.py new file mode 100644 index 0000000000..bd13162f7f --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/sizing_field/tasks.py @@ -0,0 +1,22 @@ +from polaris.mesh.spherical.unified import UNIFIED_MESH_NAMES +from polaris.tasks.mesh.spherical.unified.sizing_field.task import ( + SizingFieldTask, +) + + +def add_sizing_field_tasks(component): + """ + Add standalone sizing-field tasks for all supported unified meshes. + + One :py:class:`SizingFieldTask` is registered for each mesh name in + :py:data:`polaris.mesh.spherical.unified.UNIFIED_MESH_NAMES`. + + Parameters + ---------- + component : polaris.Component + The component to register the tasks with + """ + for mesh_name in UNIFIED_MESH_NAMES: + component.add_task( + SizingFieldTask(component=component, mesh_name=mesh_name) + ) diff --git a/polaris/tasks/mesh/spherical/unified/sizing_field/viz.py b/polaris/tasks/mesh/spherical/unified/sizing_field/viz.py new file mode 100644 index 0000000000..82925fce1b --- /dev/null +++ b/polaris/tasks/mesh/spherical/unified/sizing_field/viz.py @@ -0,0 +1,334 @@ +import os + +import cartopy.crs as ccrs +import cmocean # noqa: F401 +import matplotlib.colors as mcolors +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +from polaris.step import Step +from polaris.viz import use_mplstyle + + +class VizSizingFieldStep(Step): + """ + Visualize sizing-field diagnostics produced by + :py:class:`BuildSizingFieldStep`. + + The step writes a global overview plot, an active-control map, and a + plain-text summary of min/max cell widths and provenance counts. + + Attributes + ---------- + sizing_step : polaris.Step + The shared sizing-field build step whose output is visualized + + output_filenames : list of str + The names of the output files written by this step + """ + + def __init__(self, component, sizing_step, subdir): + """ + Create a new step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to + + sizing_step : polaris.Step + The shared sizing-field build step + + subdir : str + The subdirectory within the component work directory where the + step runs + """ + super().__init__( + component=component, + name='viz_sizing_field', + subdir=subdir, + cpus_per_task=1, + min_cpus_per_task=1, + ) + self.sizing_step = sizing_step + self.output_filenames = [ + 'sizing_field_overview.png', + 'active_control.png', + 'debug_summary.txt', + ] + + def setup(self): + """ + Link the sizing-field dataset and declare outputs. + """ + self.add_input_file( + filename='sizing_field.nc', + work_dir_target=os.path.join( + self.sizing_step.path, self.sizing_step.sizing_field_filename + ), + ) + for filename in self.output_filenames: + self.add_output_file(filename=filename) + + def run(self): + """ + Create simple global diagnostic plots and a text summary. + """ + use_mplstyle() + dpi = self.config['viz_sizing_field'].getint('dpi') + cmap = self.config['viz_sizing_field'].get('cell_width_cmap') + + with xr.open_dataset('sizing_field.nc') as ds: + lon = ds.lon.values + lat = ds.lat.values + cell_width_fields = [ + ('cellWidth', 'Final cell width (km)'), + ('ocean_background_cell_width', 'Ocean background (km)'), + ('land_river_cell_width', 'Land/river composite (km)'), + ] + delta_field = ds.coastal_transition_delta.values + + fig = plt.figure( + figsize=(15.0, 9.2), dpi=dpi, constrained_layout=True + ) + grid = fig.add_gridspec(4, 2, height_ratios=[1.0, 1.0, 0.12, 0.14]) + axes = np.array( + [ + fig.add_subplot(grid[0, 0], projection=ccrs.PlateCarree()), + fig.add_subplot(grid[0, 1], projection=ccrs.PlateCarree()), + fig.add_subplot(grid[1, 0], projection=ccrs.PlateCarree()), + fig.add_subplot(grid[1, 1], projection=ccrs.PlateCarree()), + ] + ) + delta_cax = fig.add_subplot(grid[2, 1]) + shared_cax = fig.add_subplot(grid[3, :]) + norm, ticks = _get_cell_width_norm_and_ticks( + fields=[ + ds[var_name].values for var_name, _ in cell_width_fields + ] + ) + image = None + for ax, (var_name, title) in zip( + axes[:3], cell_width_fields, strict=True + ): + image = self._plot_scalar_field( + ax=ax, + lon=lon, + lat=lat, + field=ds[var_name].values, + title=title, + cmap=cmap, + norm=norm, + add_colorbar=False, + ) + delta_norm, delta_ticks = _get_difference_norm_and_ticks( + field=delta_field + ) + delta_image = self._plot_scalar_field( + ax=axes[3], + lon=lon, + lat=lat, + field=delta_field, + title='Coastal transition delta (km)', + cmap='cmo.balance', + norm=delta_norm, + add_colorbar=False, + ) + assert image is not None + colorbar = fig.colorbar( + image, + cax=shared_cax, + orientation='horizontal', + label='Cell width (km)', + ) + if ticks is not None: + colorbar.set_ticks(np.asarray(ticks).tolist()) + delta_colorbar = fig.colorbar( + delta_image, + cax=delta_cax, + orientation='horizontal', + label='Final minus pre-coastline cell width (km)', + ) + if delta_ticks is not None: + delta_colorbar.set_ticks(delta_ticks) + fig.savefig('sizing_field_overview.png', bbox_inches='tight') + plt.close(fig) + + fig = plt.figure(figsize=(11.0, 5.0), dpi=dpi) + ax = plt.axes(projection=ccrs.PlateCarree()) + control_cmap = mcolors.ListedColormap( + ['#3b528b', '#5ec962', '#fdae61', '#d7191c'] + ) + control_norm = mcolors.BoundaryNorm( + np.arange(-0.5, 4.5, 1.0), control_cmap.N + ) + image = self._plot_scalar_field( + ax=ax, + lon=lon, + lat=lat, + field=ds.active_control.values, + title='Active sizing control', + cmap=control_cmap, + norm=control_norm, + add_colorbar=False, + ) + colorbar = fig.colorbar( + image, ax=ax, ticks=np.arange(4), shrink=0.75 + ) + colorbar.ax.set_yticklabels( + ['background', 'coastline', 'river channel', 'river outlet'] + ) + fig.savefig('active_control.png', bbox_inches='tight') + plt.close(fig) + + self._write_summary(ds) + + def _plot_scalar_field( + self, + ax, + lon, + lat, + field, + title, + cmap, + norm=None, + vmin=None, + vmax=None, + add_colorbar=True, + ): + """ + Plot one scalar field on a simple global Plate Carree map. + """ + ax.set_global() + image = ax.imshow( + field, + origin='lower', + extent=[lon.min(), lon.max(), lat.min(), lat.max()], + transform=ccrs.PlateCarree(), + cmap=cmap, + norm=norm, + vmin=vmin, + vmax=vmax, + interpolation='nearest', + ) + ax.set_title(title) + if add_colorbar: + plt.colorbar(image, ax=ax, shrink=0.75) + return image + + @staticmethod + def _write_summary(ds): + """ + Write a compact summary of the sizing-field dataset. + """ + with open('debug_summary.txt', 'w') as summary: + summary.write(f'mesh_name: {ds.attrs["mesh_name"]}\n') + summary.write( + 'target_grid_resolution_degrees: ' + f'{ds.attrs["target_grid_resolution_degrees"]}\n' + ) + summary.write( + f'cellWidth_min_km: {float(np.min(ds.cellWidth.values)):.6f}\n' + ) + summary.write( + f'cellWidth_max_km: {float(np.max(ds.cellWidth.values)):.6f}\n' + ) + background = ds.background_cell_width.values + count = int(np.count_nonzero(ds.cellWidth.values != background)) + summary.write( + f'cellWidth_differs_from_background_count: {count}\n' + ) + for var_name in [ + 'background_cell_width', + 'ocean_background_cell_width', + 'land_river_cell_width', + 'pre_coastline_cell_width', + 'coastline_cell_width', + 'coastal_transition_delta', + 'river_channel_cell_width', + 'river_outlet_cell_width', + ]: + values = ds[var_name].values + summary.write( + f'{var_name}_min_km: {float(np.min(values)):.6f}\n' + ) + summary.write( + f'{var_name}_max_km: {float(np.max(values)):.6f}\n' + ) + if var_name not in [ + 'background_cell_width', + 'ocean_background_cell_width', + 'coastal_transition_delta', + ]: + count = int(np.count_nonzero(values < background)) + summary.write( + f'{var_name}_finer_than_background_count: {count}\n' + ) + for prefix in ['river_channel', 'river_outlet']: + _write_attr_count_summary( + summary=summary, attrs=ds.attrs, prefix=prefix + ) + for control_value, label in enumerate( + ['background', 'coastline', 'river_channel', 'river_outlet'] + ): + count = int( + np.count_nonzero(ds.active_control.values == control_value) + ) + summary.write(f'{label}_count: {count}\n') + + +def _get_cell_width_norm_and_ticks(fields): + """ + Get a shared linear norm and optional readable ticks for cell-width plots. + """ + finite_values = np.concatenate( + [field[np.isfinite(field)].ravel() for field in fields] + ) + vmin = float(np.min(finite_values)) + vmax = float(np.max(finite_values)) + + if np.isclose(vmin, vmax): + padding = max(abs(vmin) * 0.01, 1.0) + return mcolors.Normalize(vmin=vmin - padding, vmax=vmax + padding), [ + vmin + ] + + unique_values = np.unique(finite_values) + if unique_values.size <= 6: + ticks = unique_values + else: + ticks = None + return mcolors.Normalize(vmin=vmin, vmax=vmax), ticks + + +def _get_difference_norm_and_ticks(field): + """ + Get a zero-centered norm and optional readable ticks for difference plots. + """ + finite_values = field[np.isfinite(field)] + max_abs = float(np.max(np.abs(finite_values))) + if np.isclose(max_abs, 0.0): + max_abs = 1.0 + ticks = [0.0] + else: + ticks = [-max_abs, 0.0, max_abs] + + return mcolors.TwoSlopeNorm( + vmin=-max_abs, vcenter=0.0, vmax=max_abs + ), ticks + + +def _write_attr_count_summary(summary, attrs, prefix): + """ + Write optional sizing-candidate provenance counts. + """ + for suffix in [ + 'mask_count', + 'finer_than_background_count', + 'equal_to_background_count', + 'coarser_than_background_count', + ]: + attr_name = f'{prefix}_{suffix}' + if attr_name in attrs: + summary.write(f'{attr_name}: {attrs[attr_name]}\n') diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py index b063f5d2ed..520627e4fb 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py @@ -1,10 +1,15 @@ import os from polaris.config import PolarisConfigParser +from polaris.e3sm.init.topo import format_lat_lon_resolution_name +from polaris.step import Step from polaris.tasks.e3sm.init import e3sm_init from polaris.tasks.e3sm.init.topo.combine import ( get_lat_lon_topo_steps, ) +from polaris.tasks.e3sm.init.topo.combine.step import ( + CombineStep as TopoCombineStep, +) from .combine import CombineStep from .extrapolate import ExtrapolateStep @@ -24,7 +29,8 @@ def get_woa23_topography_step(): steps, _ = get_lat_lon_topo_steps( component=e3sm_init, resolution=0.25, include_viz=False ) - return steps[0] + woa23_res_name = format_lat_lon_resolution_name(0.25) + return steps[TopoCombineStep.get_name('lat_lon', woa23_res_name)] def get_woa23_steps(component, combine_topo_step): @@ -42,8 +48,9 @@ def get_woa23_steps(component, combine_topo_step): Returns ------- - steps : list of polaris.Step - Shared steps for combining and extrapolating WOA23 data. + steps : dict of {str: polaris.Step} + Shared steps for combining and extrapolating WOA23 data, keyed by + step name. config : polaris.config.PolarisConfigParser The shared config options for the task and its steps. @@ -86,4 +93,9 @@ def get_woa23_steps(component, combine_topo_step): combine_topo_step=combine_topo_step, ) - return [combine_step, extrapolate_step, viz_step], config + steps: dict[str, Step] = { + combine_step.name: combine_step, + extrapolate_step.name: extrapolate_step, + viz_step.name: viz_step, + } + return steps, config diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py index b1d5c311cb..ebb9da149c 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py @@ -32,7 +32,7 @@ def __init__(self, component): self.set_shared_config(config) self.add_step(self.combine_topo_step, symlink='combine_topo') - for step in steps: + for step in steps.values(): self.add_step(step, run_by_default=step.name != 'viz') def configure(self): diff --git a/tests/mesh/spherical/unified/test_coastline.py b/tests/mesh/spherical/unified/test_coastline.py new file mode 100644 index 0000000000..ae2524081e --- /dev/null +++ b/tests/mesh/spherical/unified/test_coastline.py @@ -0,0 +1,464 @@ +import configparser +from types import SimpleNamespace + +import numpy as np +import xarray as xr +from scipy import ndimage + +from polaris.component import Component +from polaris.mesh.spherical.coastline import ( + CONVENTIONS, + _rasterize_feature_collection, + build_coastline_dataset, + build_coastline_datasets, +) +from polaris.mesh.spherical.critical_transects import CriticalTransects +from polaris.task import Task +from polaris.tasks.mesh.spherical.unified.coastline import ( + get_lat_lon_coastline_steps, +) +from polaris.tasks.mesh.spherical.unified.coastline import ( + prepare as prepare_module, +) +from polaris.tasks.mesh.spherical.unified.coastline.prepare import ( + PrepareCoastlineStep, +) + + +def test_coastline_contract_and_variants(): + ds_topo = _make_topography_dataset( + base_elevation=np.array( + [ + [-10.0, -10.0, -10.0, -10.0], + [-10.0, 20.0, 20.0, -10.0], + [-10.0, -10.0, -10.0, -10.0], + [-10.0, -20.0, -20.0, -10.0], + ] + ), + ice_mask=np.array( + [ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 1.0, 0.0], + ] + ), + grounded_mask=np.array( + [ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] + ), + ) + + ds_coastlines = build_coastline_datasets( + ds_topo=ds_topo, + resolution=1.0, + mask_threshold=0.5, + distance_chunk_size=2, + ) + ds_coastline = ds_coastlines['grounding_line'] + calving = ds_coastlines['calving_front'].ocean_mask.values + grounding = ds_coastline.ocean_mask.values + bedrock = ds_coastlines['bedrock_zero'].ocean_mask.values + + expected_vars = {'ocean_mask', 'signed_distance'} + assert set(ds_coastline.data_vars) == expected_vars + assert tuple(ds_coastlines.keys()) == ( + 'calving_front', + 'grounding_line', + 'bedrock_zero', + ) + assert ds_coastline.attrs['coastline_convention'] == 'grounding_line' + assert ( + ds_coastline.attrs['flood_fill_seed_strategy'] + == 'candidate_ocean_on_northernmost_row' + ) + + assert calving[3, 1] == 0 + assert grounding[3, 1] == 1 + assert grounding[3, 2] == 0 + assert bedrock[3, 2] == 1 + + assert 'coastline_mask' not in ds_coastline + assert 'coastline_edge_east' not in ds_coastline + assert 'coastline_edge_north' not in ds_coastline + assert 'candidate_ocean_mask' not in ds_coastline + assert 'land_mask' not in ds_coastline + + signed_distance = ds_coastline.signed_distance.values + assert np.isfinite(signed_distance).all() + assert signed_distance[0, 0] > 0.0 + assert signed_distance[1, 1] < 0.0 + + +def test_coastline_excludes_disconnected_inland_water(): + ds_topo = _make_topography_dataset( + base_elevation=np.array( + [ + [-10.0, -10.0, -10.0, -10.0], + [20.0, 20.0, 20.0, 20.0], + [20.0, -10.0, -10.0, 20.0], + [20.0, 20.0, 20.0, 20.0], + ] + ) + ) + + ds_coastline = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='bedrock_zero', + mask_threshold=0.5, + distance_chunk_size=2, + ) + ocean_mask = ds_coastline.ocean_mask.values + + assert ocean_mask[2, 1] == 0 + assert ocean_mask[2, 2] == 0 + + +def test_coastline_uses_northernmost_latitude_for_seed_row(): + ds_topo = _make_topography_dataset( + lat=np.array([-60.0, -20.0, 20.0, 60.0]), + base_elevation=np.array( + [ + [-10.0, -10.0, -10.0, -10.0], + [-10.0, -10.0, -10.0, -10.0], + [-10.0, -10.0, -10.0, -10.0], + [-10.0, -10.0, -10.0, -10.0], + ] + ), + ice_mask=np.array( + [ + [1.0, 1.0, 1.0, 1.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ] + ), + ) + + ds_coastline = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='calving_front', + mask_threshold=0.5, + distance_chunk_size=2, + ) + ocean_mask = ds_coastline.ocean_mask.values + + assert ocean_mask.sum() == ocean_mask.size - ds_topo.lon.size + assert np.all(ocean_mask[0, :] == 0) + assert np.all(ocean_mask[-1, :] == 1) + + +def test_coastline_none_transects_matches_legacy_behavior(): + ds_topo = _make_topography_dataset( + base_elevation=np.array( + [ + [-10.0, -10.0, -10.0, -10.0], + [20.0, 20.0, 20.0, 20.0], + [-10.0, -10.0, -10.0, -10.0], + [-10.0, -10.0, -10.0, -10.0], + ] + ) + ) + + legacy = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='bedrock_zero', + distance_chunk_size=2, + ) + explicit_none = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='bedrock_zero', + distance_chunk_size=2, + critical_transects=None, + ) + + np.testing.assert_array_equal( + legacy.ocean_mask.values, explicit_none.ocean_mask.values + ) + np.testing.assert_allclose( + legacy.signed_distance.values, + explicit_none.signed_distance.values, + ) + + +def test_critical_land_blockage_closes_a_narrow_ocean_connection(): + base_elevation = -10.0 * np.ones((5, 5)) + base_elevation[1, :] = 20.0 + base_elevation[1, 2] = -10.0 + ds_topo = _make_topography_dataset(base_elevation=base_elevation) + + baseline = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='bedrock_zero', + distance_chunk_size=2, + ) + critical_transects = CriticalTransects( + land_blockages=_feature_collection([(0.0, 40.0), (0.0, 20.0)]), + passages=None, + ) + blocked = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='bedrock_zero', + distance_chunk_size=2, + critical_transects=critical_transects, + ) + + assert baseline.ocean_mask.values[-1, 2] == 1 + assert blocked.ocean_mask.values[-1, 2] == 0 + + +def test_critical_passage_connects_otherwise_disconnected_ocean(): + base_elevation = -10.0 * np.ones((5, 5)) + base_elevation[1, :] = 20.0 + ds_topo = _make_topography_dataset(base_elevation=base_elevation) + + baseline = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='bedrock_zero', + distance_chunk_size=2, + ) + critical_transects = CriticalTransects( + land_blockages=None, + passages=_feature_collection([(0.0, 40.0), (0.0, 20.0)]), + ) + connected = build_coastline_dataset( + ds_topo=ds_topo, + resolution=1.0, + convention='bedrock_zero', + distance_chunk_size=2, + critical_transects=critical_transects, + ) + + assert baseline.ocean_mask.values[-1, 2] == 0 + assert connected.ocean_mask.values[-1, 2] == 1 + + +def test_diagonal_transect_rasterization_is_four_connected(): + lon = np.array([-135.0, -45.0, 45.0, 135.0]) + lat = np.array([60.0, 20.0, -20.0, -60.0]) + mask = _rasterize_feature_collection( + feature_collection=_feature_collection( + [(-120.0, 60.0), (20.0, -60.0)] + ), + lon=lon, + lat=lat, + ) + + labels, count = ndimage.label( + mask, + structure=np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.int8), + ) + assert labels.any() + assert count == 1 + assert mask.sum() > 2 + + +def test_antimeridian_transect_rasterization_uses_periodic_longitude(): + lon = np.array([-135.0, -45.0, 45.0, 135.0]) + lat = np.array([60.0, 20.0, -20.0, -60.0]) + mask = _rasterize_feature_collection( + feature_collection=_feature_collection( + [(170.0, 20.0), (-170.0, 20.0)] + ), + lon=lon, + lat=lat, + ) + + assert mask[1, 0] + assert mask[1, -1] + + +def test_coastline_step_configures_critical_transects(monkeypatch): + component = Component('mesh') + combine_step = SimpleNamespace( + combined_filename='combined_topography.nc', + subdir='combine', + path='combine', + ) + dummy_dataset = xr.Dataset() + + calls = [] + sentinel = object() + + def fake_build(**kwargs): + calls.append(kwargs['critical_transects']) + return {convention: dummy_dataset.copy() for convention in CONVENTIONS} + + monkeypatch.setattr(prepare_module, 'build_coastline_datasets', fake_build) + monkeypatch.setattr( + prepare_module, + '_write_netcdf_with_fill_values', + lambda ds, filename: None, + ) + monkeypatch.setattr( + prepare_module.xr, + 'open_dataset', + lambda filename: xr.Dataset(), + ) + + enabled_calls = {'count': 0} + + def fake_loader(): + enabled_calls['count'] += 1 + return sentinel + + monkeypatch.setattr( + prepare_module, + 'load_default_critical_transects', + fake_loader, + ) + + disabled_step = PrepareCoastlineStep( + component=component, + combine_step=combine_step, + subdir='prepare_disabled', + ) + disabled_step.config = _make_prepare_config(False) + disabled_step.run() + + enabled_step = PrepareCoastlineStep( + component=component, + combine_step=combine_step, + subdir='prepare_enabled', + ) + enabled_step.config = _make_prepare_config(True) + enabled_step.run() + + assert calls == [None, sentinel] + assert enabled_calls['count'] == 1 + + +def test_get_lat_lon_coastline_steps_reuses_shared_config_for_viz(): + component = Component('mesh') + combine_step = SimpleNamespace( + combined_filename='combined_topography.nc', + subdir='combine', + path='combine', + ) + + steps_without_viz, config_without_viz = get_lat_lon_coastline_steps( + component=component, + combine_topo_step=combine_step, + resolution=0.25, + include_viz=False, + ) + steps_with_viz, config_with_viz = get_lat_lon_coastline_steps( + component=component, + combine_topo_step=combine_step, + resolution=0.25, + include_viz=True, + ) + + assert len(steps_without_viz) == 1 + assert len(steps_with_viz) == 2 + assert steps_without_viz['coastline'] is steps_with_viz['coastline'] + assert config_without_viz is config_with_viz + + +def test_task_can_include_shared_coastline_viz_without_default_run(): + component = Component('mesh') + combine_step = SimpleNamespace( + combined_filename='combined_topography.nc', + subdir='combine', + path='combine', + ) + coastline_steps, _ = get_lat_lon_coastline_steps( + component=component, + combine_topo_step=combine_step, + resolution=0.25, + include_viz=True, + ) + + task = Task(component=component, name='coastline_consumer') + task.add_step(coastline_steps['coastline'], symlink='coastline') + task.add_step( + coastline_steps['viz_coastline'], + symlink='viz_coastline', + run_by_default=False, + ) + + assert 'coastline' in task.steps + assert 'viz_coastline' in task.steps + assert 'coastline' in task.steps_to_run + assert 'viz_coastline' not in task.steps_to_run + + +def _make_topography_dataset( + base_elevation, + lon=None, + lat=None, + ice_mask=None, + grounded_mask=None, +): + if lon is None: + n_lon = base_elevation.shape[1] + if n_lon == 5: + lon = np.array([-144.0, -72.0, 0.0, 72.0, 144.0]) + else: + lon = np.array([-135.0, -45.0, 45.0, 135.0]) + if lat is None: + n_lat = base_elevation.shape[0] + if n_lat == 5: + lat = np.array([60.0, 30.0, 0.0, -30.0, -60.0]) + else: + lat = np.array([60.0, 20.0, -20.0, -60.0]) + + if ice_mask is None: + ice_mask = np.zeros_like(base_elevation) + if grounded_mask is None: + grounded_mask = np.zeros_like(base_elevation) + + return xr.Dataset( + data_vars=dict( + base_elevation=(('lat', 'lon'), base_elevation), + ice_mask=(('lat', 'lon'), ice_mask), + grounded_mask=(('lat', 'lon'), grounded_mask), + ), + coords=dict( + lat=xr.DataArray(np.asarray(lat), dims=('lat',)), + lon=xr.DataArray(np.asarray(lon), dims=('lon',)), + ), + ) + + +def _feature_collection(*lines): + return { + 'type': 'FeatureCollection', + 'features': [ + { + 'type': 'Feature', + 'properties': {}, + 'geometry': { + 'type': 'LineString', + 'coordinates': list(line), + }, + } + for line in lines + ], + } + + +def _make_prepare_config(include_critical_transects): + config = configparser.ConfigParser() + config.add_section('coastline') + config.set('coastline', 'resolution_latlon', '1.0') + config.set( + 'coastline', + 'include_critical_transects', + str(include_critical_transects), + ) + config.set('coastline', 'mask_threshold', '0.5') + config.set('coastline', 'sea_level_elevation', '0.0') + config.set('coastline', 'distance_chunk_size', '2') + return config diff --git a/tests/mesh/spherical/unified/test_river.py b/tests/mesh/spherical/unified/test_river.py new file mode 100644 index 0000000000..a77bb98663 --- /dev/null +++ b/tests/mesh/spherical/unified/test_river.py @@ -0,0 +1,820 @@ +import os +import zipfile +from types import SimpleNamespace + +import numpy as np +import pytest +import shapefile +import xarray as xr + +from polaris.component import Component +from polaris.mesh.spherical.unified import ( + UNIFIED_MESH_NAMES, + get_unified_mesh_config, +) +from polaris.tasks.mesh.spherical.unified.river import ( + add_river_tasks, + build_river_network_dataset, + get_mesh_unified_river_steps, + simplify_river_network_feature_collection, +) +from polaris.tasks.mesh.spherical.unified.river.base_mesh import ( + clip_outlet_feature_collection, + condition_base_mesh_river_segments, +) +from polaris.tasks.mesh.spherical.unified.river.source import ( + _convert_hydrorivers_shapefile_to_geojson, + _unpack_hydrorivers_archive, + read_river_segments_from_feature_collection, +) + + +def test_simplify_river_network_filters_outlets_and_minor_tributaries(): + feature_collection = dict( + type='FeatureCollection', + features=[ + _line_feature( + hyriv_id=10, + coords=[(-1.0, 0.0), (0.0, 0.0)], + next_down=0, + drainage_area=100.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=11, + coords=[(-2.0, 0.0), (-1.0, 0.0)], + next_down=10, + drainage_area=80.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=12, + coords=[(-1.5, 1.0), (-1.0, 0.0)], + next_down=10, + drainage_area=35.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=13, + coords=[(-1.6, 0.2), (-1.2, 0.0)], + next_down=10, + drainage_area=10.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=20, + coords=[(0.0, 0.05), (0.05, 0.05)], + next_down=0, + drainage_area=90.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=30, + coords=[(50.0, 10.0), (50.1, 10.0)], + next_down=0, + drainage_area=25.0e6, + endorheic=1, + ), + _line_feature( + hyriv_id=31, + coords=[(50.0, 10.05), (50.1, 10.05)], + next_down=0, + drainage_area=20.0e6, + endorheic=1, + ), + ], + ) + + simplified_fc, outlets_fc = simplify_river_network_feature_collection( + feature_collection=feature_collection, + drainage_area_threshold=5.0e6, + outlet_distance_tolerance=30.0e3, + tributary_area_ratio=0.4, + ) + + simplified_ids = { + feature['properties']['hyriv_id'] + for feature in simplified_fc['features'] + } + outlet_ids = { + feature['properties']['hyriv_id'] for feature in outlets_fc['features'] + } + + assert simplified_ids == {10, 11, 12, 30, 31} + assert outlet_ids == {10, 30, 31} + + +def test_simplify_river_network_handles_deep_main_stem(): + n_segments = 1500 + features = [] + for hyriv_id in range(1, n_segments + 1): + next_down = hyriv_id - 1 if hyriv_id > 1 else 0 + features.append( + _line_feature( + hyriv_id=hyriv_id, + coords=[ + (-float(hyriv_id), 0.0), + (-float(hyriv_id) + 1.0, 0.0), + ], + next_down=next_down, + drainage_area=(n_segments - hyriv_id + 1) * 1.0e6, + endorheic=0, + ) + ) + + simplified_fc, outlets_fc = simplify_river_network_feature_collection( + feature_collection=dict(type='FeatureCollection', features=features), + drainage_area_threshold=1.0, + outlet_distance_tolerance=1.0, + tributary_area_ratio=0.05, + ) + + simplified_ids = { + feature['properties']['hyriv_id'] + for feature in simplified_fc['features'] + } + outlet_ids = { + feature['properties']['hyriv_id'] for feature in outlets_fc['features'] + } + + assert len(simplified_ids) == n_segments + assert outlet_ids == {1} + + +def test_simplify_river_network_rejects_next_down_cycles(): + feature_collection = dict( + type='FeatureCollection', + features=[ + _line_feature( + hyriv_id=10, + coords=[(0.0, 0.0), (1.0, 0.0)], + next_down=11, + drainage_area=50.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=11, + coords=[(1.0, 0.0), (2.0, 0.0)], + next_down=10, + drainage_area=60.0e6, + endorheic=0, + ), + ], + ) + + with pytest.raises(ValueError, match='Cycle detected'): + simplify_river_network_feature_collection( + feature_collection=feature_collection, + drainage_area_threshold=1.0, + outlet_distance_tolerance=1.0, + tributary_area_ratio=0.05, + ) + + +def test_simplify_river_network_preserves_branch_traversal_order(): + feature_collection = dict( + type='FeatureCollection', + features=[ + _line_feature( + hyriv_id=10, + coords=[(0.0, 0.0), (1.0, 0.0)], + next_down=0, + drainage_area=100.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=11, + coords=[(-1.0, 0.0), (0.0, 0.0)], + next_down=10, + drainage_area=80.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=12, + coords=[(0.0, 1.0), (0.0, 0.0)], + next_down=10, + drainage_area=60.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=21, + coords=[(-2.0, 0.02), (-1.0, 0.0)], + next_down=11, + drainage_area=70.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=22, + coords=[(0.0, 2.0), (0.0, 1.0)], + next_down=12, + drainage_area=55.0e6, + endorheic=0, + ), + _line_feature( + hyriv_id=23, + coords=[(-2.0, 0.03), (0.0, 1.0)], + next_down=12, + drainage_area=20.0e6, + endorheic=0, + ), + ], + ) + + simplified_fc, _ = simplify_river_network_feature_collection( + feature_collection=feature_collection, + drainage_area_threshold=1.0, + outlet_distance_tolerance=20.0e3, + tributary_area_ratio=0.4, + ) + + simplified_ids = { + feature['properties']['hyriv_id'] + for feature in simplified_fc['features'] + } + + assert simplified_ids == {10, 11, 12, 21, 22} + + +def test_build_river_network_dataset_contract_and_snapped_outlets(): + river_fc = dict( + type='FeatureCollection', + features=[ + _line_feature( + hyriv_id=10, + coords=[(-90.0, 0.0), (-45.0, 0.0)], + next_down=0, + drainage_area=100.0e6, + endorheic=0, + outlet_type='ocean', + outlet_hyriv_id=10, + ), + _line_feature( + hyriv_id=30, + coords=[(90.0, 0.0), (45.0, 0.0)], + next_down=0, + drainage_area=20.0e6, + endorheic=1, + outlet_type='inland_sink', + outlet_hyriv_id=30, + ), + ], + ) + outlet_fc = dict( + type='FeatureCollection', + features=[ + _point_feature( + hyriv_id=10, + coords=(-46.0, 1.0), + drainage_area=100.0e6, + endorheic=0, + outlet_type='ocean', + ), + _point_feature( + hyriv_id=30, + coords=(46.0, 1.0), + drainage_area=20.0e6, + endorheic=1, + outlet_type='inland_sink', + ), + ], + ) + ds_coastline = xr.Dataset( + data_vars=dict( + ocean_mask=( + ('lat', 'lon'), + np.array( + [ + [0, 0, 0, 0, 0], + [0, 1, 0, 0, 0], + [0, 0, 0, 0, 0], + ], + dtype=np.int8, + ), + ), + ), + coords=dict( + lat=xr.DataArray(np.array([10.0, 0.0, -10.0]), dims=('lat',)), + lon=xr.DataArray( + np.array([-90.0, -45.0, 0.0, 45.0, 90.0]), dims=('lon',) + ), + ), + ) + + ds_river, snapped_outlets = build_river_network_dataset( + river_feature_collection=river_fc, + outlet_feature_collection=outlet_fc, + ds_coastline=ds_coastline, + resolution=45.0, + outlet_match_tolerance=200.0e3, + channel_subsegment_fraction=0.5, + ) + + expected_vars = { + 'river_channel_mask', + 'river_outlet_mask', + 'river_ocean_outlet_mask', + 'river_inland_sink_mask', + } + assert expected_vars.issubset(set(ds_river.data_vars)) + assert ds_river.attrs['matched_ocean_outlets'] == 1 + assert ds_river.attrs['unmatched_ocean_outlets'] == 0 + + assert ds_river.river_channel_mask.sum() > 0 + assert ds_river.river_ocean_outlet_mask.sel(lat=0.0, lon=-45.0) == 1 + assert ds_river.river_inland_sink_mask.sel(lat=0.0, lon=45.0) == 1 + + snapped_by_id = { + feature['properties']['hyriv_id']: feature + for feature in snapped_outlets['features'] + } + assert snapped_by_id[10]['properties']['snapped_lon'] == -45.0 + assert snapped_by_id[10]['properties']['snapped_lat'] == 0.0 + assert snapped_by_id[30]['properties']['snapped_lon'] == 45.0 + assert snapped_by_id[30]['properties']['snapped_lat'] == 0.0 + assert snapped_by_id[10]['properties']['matched_to_ocean'] + + +def test_build_river_network_dataset_applies_physical_channel_buffer(): + river_fc = dict( + type='FeatureCollection', + features=[ + _line_feature( + hyriv_id=10, + coords=[(0.0, 59.0), (0.0, 61.0)], + next_down=0, + drainage_area=100.0e6, + endorheic=0, + ), + ], + ) + ds_coastline = xr.Dataset( + data_vars=dict( + ocean_mask=( + ('lat', 'lon'), + np.zeros((3, 5), dtype=np.int8), + ), + ), + coords=dict( + lat=xr.DataArray(np.array([61.0, 60.0, 59.0]), dims=('lat',)), + lon=xr.DataArray( + np.array([-2.0, -1.0, 0.0, 1.0, 2.0]), dims=('lon',) + ), + ), + ) + + ds_river, _ = build_river_network_dataset( + river_feature_collection=river_fc, + outlet_feature_collection=dict(type='FeatureCollection', features=[]), + ds_coastline=ds_coastline, + resolution=1.0, + outlet_match_tolerance=200.0e3, + channel_subsegment_fraction=0.5, + channel_buffer_km=80.0, + ) + + assert ds_river.attrs['channel_buffer_m'] == 80.0e3 + assert ds_river.river_channel_mask.sel(lat=60.0, lon=0.0) == 1 + assert ds_river.river_channel_mask.sel(lat=60.0, lon=-1.0) == 1 + assert ds_river.river_channel_mask.sel(lat=60.0, lon=1.0) == 1 + assert ds_river.river_channel_mask.sel(lat=60.0, lon=-2.0) == 0 + assert ds_river.river_channel_mask.sel(lat=60.0, lon=2.0) == 0 + + +def test_build_river_network_dataset_derives_land_mask_from_ocean_mask(): + river_fc = dict(type='FeatureCollection', features=[]) + outlet_fc = dict( + type='FeatureCollection', + features=[ + _point_feature( + hyriv_id=30, + coords=(46.0, 1.0), + drainage_area=20.0e6, + endorheic=1, + outlet_type='inland_sink', + ), + ], + ) + ds_coastline = xr.Dataset( + data_vars=dict( + ocean_mask=( + ('lat', 'lon'), + np.array( + [ + [1, 1, 1, 1, 1], + [1, 1, 1, 0, 1], + [1, 1, 1, 1, 1], + ], + dtype=np.int8, + ), + ), + ), + coords=dict( + lat=xr.DataArray(np.array([10.0, 0.0, -10.0]), dims=('lat',)), + lon=xr.DataArray( + np.array([-90.0, -45.0, 0.0, 45.0, 90.0]), dims=('lon',) + ), + ), + ) + + ds_river, snapped_outlets = build_river_network_dataset( + river_feature_collection=river_fc, + outlet_feature_collection=outlet_fc, + ds_coastline=ds_coastline, + resolution=45.0, + outlet_match_tolerance=200.0e3, + channel_subsegment_fraction=0.5, + ) + + assert ds_river.river_inland_sink_mask.sel(lat=0.0, lon=45.0) == 1 + snapped_feature = snapped_outlets['features'][0] + assert snapped_feature['properties']['snapped_lon'] == 45.0 + assert snapped_feature['properties']['snapped_lat'] == 0.0 + + +def test_condition_base_mesh_river_segments_clips_then_simplifies(): + ds_coastline = xr.Dataset( + data_vars=dict( + ocean_mask=(('lat', 'lon'), np.zeros((2, 3), dtype=np.int8)), + signed_distance=( + ('lat', 'lon'), + np.array( + [ + [-2.0e5, -2.0e5, 1.0e5], + [-2.0e5, -2.0e5, 1.0e5], + ] + ), + ), + ), + coords=dict( + lat=xr.DataArray(np.array([-5.0, 5.0]), dims=('lat',)), + lon=xr.DataArray(np.array([0.0, 10.0, 20.0]), dims=('lon',)), + ), + ) + river_fc = dict( + type='FeatureCollection', + features=[ + _line_feature( + hyriv_id=10, + coords=[(0.0, 0.0), (8.0, 0.0), (16.0, 0.0)], + next_down=0, + drainage_area=100.0e6, + endorheic=0, + outlet_type='ocean', + outlet_hyriv_id=10, + ) + ], + ) + + segments = condition_base_mesh_river_segments( + segments=read_river_segments_from_feature_collection(river_fc), + ds_coastline=ds_coastline, + clip_distance_m=50.0e3, + simplify_tolerance_deg=0.5, + min_segment_length_m=100.0, + ) + + assert len(segments) == 1 + coords = np.asarray(segments[0].geometry.coords) + assert coords.shape[0] == 2 + assert coords[0, 0] == 0.0 + assert 8.0 < coords[-1, 0] < 16.0 + assert segments[0].outlet_type is None + assert segments[0].outlet_hyriv_id is None + + +def test_condition_base_mesh_river_segments_drops_short_fragments(): + ds_coastline = xr.Dataset( + data_vars=dict( + ocean_mask=(('lat', 'lon'), np.zeros((2, 3), dtype=np.int8)), + signed_distance=( + ('lat', 'lon'), + np.array( + [ + [1.0e5, -1.0e4, 1.0e5], + [1.0e5, -1.0e4, 1.0e5], + ] + ), + ), + ), + coords=dict( + lat=xr.DataArray(np.array([-5.0, 5.0]), dims=('lat',)), + lon=xr.DataArray(np.array([0.0, 1.0, 2.0]), dims=('lon',)), + ), + ) + river_fc = dict( + type='FeatureCollection', + features=[ + _line_feature( + hyriv_id=10, + coords=[(0.0, 0.0), (1.0, 0.0), (2.0, 0.0)], + next_down=0, + drainage_area=100.0e6, + endorheic=0, + ) + ], + ) + + segments = condition_base_mesh_river_segments( + segments=read_river_segments_from_feature_collection(river_fc), + ds_coastline=ds_coastline, + clip_distance_m=20.0e3, + simplify_tolerance_deg=0.0, + min_segment_length_m=200.0e3, + ) + + assert segments == [] + + +def test_clip_outlet_feature_collection_removes_ocean_outlets(): + ds_coastline = xr.Dataset( + data_vars=dict( + ocean_mask=(('lat', 'lon'), np.zeros((2, 2), dtype=np.int8)), + signed_distance=( + ('lat', 'lon'), + np.array( + [ + [-2.0e5, 1.0e5], + [-2.0e5, 1.0e5], + ] + ), + ), + ), + coords=dict( + lat=xr.DataArray(np.array([-5.0, 5.0]), dims=('lat',)), + lon=xr.DataArray(np.array([0.0, 10.0]), dims=('lon',)), + ), + ) + outlet_fc = dict( + type='FeatureCollection', + features=[ + _point_feature( + hyriv_id=10, + coords=(9.0, 0.0), + drainage_area=100.0e6, + endorheic=0, + outlet_type='ocean', + ), + _point_feature( + hyriv_id=20, + coords=(1.0, 0.0), + drainage_area=50.0e6, + endorheic=1, + outlet_type='inland_sink', + ), + ], + ) + + clipped = clip_outlet_feature_collection( + outlet_feature_collection=outlet_fc, + ds_coastline=ds_coastline, + clip_distance_m=50.0e3, + ) + + clipped_ids = { + feature['properties']['hyriv_id'] for feature in clipped['features'] + } + assert clipped_ids == {20} + + +def test_build_river_network_dataset_marks_distant_ocean_outlet_unmatched(): + river_fc = dict(type='FeatureCollection', features=[]) + outlet_fc = dict( + type='FeatureCollection', + features=[ + _point_feature( + hyriv_id=10, + coords=(-46.0, 1.0), + drainage_area=100.0e6, + endorheic=0, + outlet_type='ocean', + ), + ], + ) + ds_coastline = xr.Dataset( + data_vars=dict( + ocean_mask=( + ('lat', 'lon'), + np.array( + [ + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 0], + ], + dtype=np.int8, + ), + ), + ), + coords=dict( + lat=xr.DataArray(np.array([10.0, 0.0, -10.0]), dims=('lat',)), + lon=xr.DataArray( + np.array([-90.0, -45.0, 0.0, 45.0, 90.0]), dims=('lon',) + ), + ), + ) + + ds_river, snapped_outlets = build_river_network_dataset( + river_feature_collection=river_fc, + outlet_feature_collection=outlet_fc, + ds_coastline=ds_coastline, + resolution=45.0, + outlet_match_tolerance=200.0e3, + channel_subsegment_fraction=0.5, + ) + + assert ds_river.attrs['matched_ocean_outlets'] == 0 + assert ds_river.attrs['unmatched_ocean_outlets'] == 1 + + snapped_feature = snapped_outlets['features'][0] + assert not snapped_feature['properties']['matched_to_ocean'] + assert snapped_feature['properties']['snapped_lon'] == -45.0 + assert snapped_feature['properties']['snapped_lat'] == 0.0 + + +def test_unpack_hydrorivers_archive(tmp_path): + archive_filename = tmp_path / 'HydroRIVERS_v10_shp.zip' + shp_directory = tmp_path / 'HydroRIVERS_v10_shp' + + with zipfile.ZipFile(archive_filename, 'w') as archive: + archive.writestr('HydroRIVERS_v10_shp/HydroRIVERS_v10.shp', 'x') + + cwd = os.getcwd() + os.chdir(tmp_path) + try: + _unpack_hydrorivers_archive( + archive_filename=str(archive_filename), + data_directory='HydroRIVERS_v10_shp', + ) + finally: + os.chdir(cwd) + + assert shp_directory.is_dir() + assert (shp_directory / 'HydroRIVERS_v10.shp').exists() + + +def test_convert_hydrorivers_shapefile_to_geojson(tmp_path): + base = tmp_path / 'HydroRIVERS_v10' + writer = shapefile.Writer(str(base)) + writer.field('HYRIV_ID', 'N', 10, 0) + writer.field('MAIN_RIV', 'N', 10, 0) + writer.field('ORD_STRA', 'N', 10, 0) + writer.field('UPLAND_SKM', 'F', 18, 3) + writer.field('NEXT_DOWN', 'N', 10, 0) + writer.field('ENDORHEIC', 'N', 10, 0) + writer.line([[[-1.0, 0.0], [0.0, 0.0]]]) + writer.record(1, 1, 1, 100.0, 0, 0) + writer.close() + + output_filename = tmp_path / 'source_river_network.geojson' + _convert_hydrorivers_shapefile_to_geojson( + shp_filename=f'{base}.shp', + output_filename=str(output_filename), + ) + + feature_collection = output_filename.read_text(encoding='utf-8') + assert 'HYRIV_ID' in feature_collection + assert 'LineString' in feature_collection + + +def test_mesh_river_step_factories_use_mesh_subdirs(): + component = Component(name='mesh') + mesh_name = 'ocn_30km_lnd_10km_riv_10km' + coastline_step = _coastline_step(mesh_name) + unified_steps, unified_config = get_mesh_unified_river_steps( + component=component, + coastline_step=coastline_step, + mesh_name=mesh_name, + include_viz=False, + ) + source_step = unified_steps['river_source'] + lat_lon_step = unified_steps['river_lat_lon'] + base_mesh_step = unified_steps['river_base_mesh'] + + assert ( + source_step.subdir + == f'spherical/unified/{mesh_name}/river/source/prepare' + ) + assert ( + lat_lon_step.subdir + == f'spherical/unified/{mesh_name}/river/lat_lon/prepare' + ) + assert unified_config.get('unified_mesh', 'mesh_name') == mesh_name + assert not unified_config.has_option( + 'river_lat_lon', 'coastline_convention' + ) + assert not unified_config.has_option('river_lat_lon', 'resolution_latlon') + + assert ( + base_mesh_step.subdir + == f'spherical/unified/{mesh_name}/river/base_mesh/prepare' + ) + + +def test_mesh_river_step_factories_reuse_shared_configs(): + component = Component(name='mesh') + mesh_name = 'ocn_30km_lnd_10km_riv_10km' + + coastline_step = _coastline_step(mesh_name) + unified_steps_first, unified_config_first = get_mesh_unified_river_steps( + component=component, + coastline_step=coastline_step, + mesh_name=mesh_name, + include_viz=False, + ) + unified_steps_second, unified_config_second = get_mesh_unified_river_steps( + component=component, + coastline_step=coastline_step, + mesh_name=mesh_name, + include_viz=True, + ) + + # source and lat-lon steps are shared across both calls + assert ( + unified_steps_first['river_source'] + is unified_steps_second['river_source'] + ) + assert ( + unified_steps_first['river_lat_lon'] + is unified_steps_second['river_lat_lon'] + ) + assert unified_config_first is unified_config_second + # without viz: {river_source, river_lat_lon, river_base_mesh} + # with viz: {river_source, river_lat_lon, viz_river_network, + # river_base_mesh} + assert len(unified_steps_first) == 3 + assert len(unified_steps_second) == 4 + assert ( + unified_steps_first['river_base_mesh'] + is unified_steps_second['river_base_mesh'] + ) + + +def test_add_river_tasks_registers_mesh_tasks(): + component = Component(name='mesh') + add_river_tasks(component=component) + + assert len(component.tasks) == len(UNIFIED_MESH_NAMES) + + for mesh_name in UNIFIED_MESH_NAMES: + subdir = f'spherical/unified/{mesh_name}/river/task' + assert subdir in component.tasks + task = component.tasks[subdir] + assert task.name == f'river_network_{mesh_name}_task' + + +def _line_feature( + hyriv_id, + coords, + next_down, + drainage_area, + endorheic, + outlet_type=None, + outlet_hyriv_id=None, +): + return dict( + type='Feature', + properties=dict( + hyriv_id=hyriv_id, + main_riv=hyriv_id, + ord_stra=1, + drainage_area=drainage_area, + next_down=next_down, + endorheic=endorheic, + outlet_type=outlet_type, + outlet_hyriv_id=outlet_hyriv_id, + ), + geometry=dict(type='LineString', coordinates=coords), + ) + + +def _point_feature( + hyriv_id, + coords, + drainage_area, + endorheic, + outlet_type, +): + return dict( + type='Feature', + properties=dict( + hyriv_id=hyriv_id, + main_riv=hyriv_id, + drainage_area=drainage_area, + endorheic=endorheic, + outlet_type=outlet_type, + ), + geometry=dict(type='Point', coordinates=coords), + ) + + +def _coastline_step(mesh_name): + config = get_unified_mesh_config(mesh_name=mesh_name) + resolution = config.getfloat('unified_mesh', 'resolution_latlon') + convention = config.get('unified_mesh', 'coastline_convention') + return SimpleNamespace( + subdir=( + 'spherical/unified/coastline/lat_lon/' + f'{resolution:.5f}_degree/prepare' + ), + path='coastline', + output_filenames={convention: 'coastline.nc'}, + ) diff --git a/tests/mesh/spherical/unified/test_sizing_field.py b/tests/mesh/spherical/unified/test_sizing_field.py new file mode 100644 index 0000000000..f1e330b245 --- /dev/null +++ b/tests/mesh/spherical/unified/test_sizing_field.py @@ -0,0 +1,568 @@ +import importlib.resources as imp_res +import os +from types import SimpleNamespace + +import numpy as np +import xarray as xr + +from polaris.component import Component +from polaris.mesh.spherical.unified import ( + UNIFIED_MESH_NAMES, + UnifiedCellWidthMeshStep, + get_unified_mesh_family, +) +from polaris.mesh.spherical.unified.families.default import ( + build_ocean_background_from_mode, +) +from polaris.mesh.spherical.unified.families.so_region import ( + SO_REGION_FILENAME, + SO_REGION_PACKAGE, +) +from polaris.tasks.mesh.spherical.unified.sizing_field import ( + BuildSizingFieldStep, + get_lat_lon_sizing_field_steps, + get_sizing_field_config, + sizing_field_dataset, +) +from polaris.tasks.mesh.spherical.unified.sizing_field.tasks import ( + add_sizing_field_tasks, +) + + +def test_sizing_field_mesh_outputs_and_active_control(): + ds_coastline = _make_coastline_dataset( + ocean_mask=np.array([[1, 1, 0], [1, 0, 0]], dtype=np.int8), + signed_distance=np.array( + [[1000.0, 2000.0, -500.0], [3000.0, -4000.0, -1000.0]] + ), + lat=np.array([-60.0, 0.0]), + ) + ds_river = _make_river_dataset( + river_channel_mask=np.array([[0, 0, 0], [0, 1, 0]], dtype=np.int8), + river_outlet_mask=np.array([[0, 1, 0], [0, 0, 0]], dtype=np.int8), + lat=np.array([-60.0, 0.0]), + ) + + ds_uniform = sizing_field_dataset( + ds_coastline=ds_coastline, + ds_river=ds_river, + resolution=0.25, + mesh_name='ocn_240km_lnd_240km_riv_240km', + ocean_background=_constant_ocean_background( + ds_coastline=ds_coastline, value=240.0 + ), + land_background_km=240.0, + river_channel_km=240.0, + river_outlet_km=240.0, + ) + assert ds_uniform.attrs['mesh_name'] == 'ocn_240km_lnd_240km_riv_240km' + assert 'profile_name' not in ds_uniform.attrs + np.testing.assert_allclose(ds_uniform.cellWidth.values, 240.0) + assert np.all(ds_uniform.active_control.values == 0) + + ds_split = sizing_field_dataset( + ds_coastline=ds_coastline, + ds_river=ds_river, + resolution=0.125, + mesh_name='ocn_30km_lnd_10km_riv_10km', + ocean_background=_constant_ocean_background( + ds_coastline=ds_coastline, value=30.0 + ), + land_background_km=10.0, + river_channel_km=10.0, + river_outlet_km=10.0, + ) + expected_split = np.array([[30.0, 30.0, 10.0], [30.0, 10.0, 10.0]]) + np.testing.assert_allclose(ds_split.cellWidth.values, expected_split) + assert ds_split.active_control.values[0, 1] == 0 + assert ds_split.active_control.values[1, 1] == 0 + assert ds_split.attrs['river_channel_mask_count'] == 1 + assert ds_split.attrs['river_channel_finer_than_background_count'] == 0 + assert ds_split.attrs['river_channel_equal_to_background_count'] == 1 + assert ds_split.attrs['river_outlet_mask_count'] == 1 + assert ds_split.attrs['river_outlet_finer_than_background_count'] == 0 + assert ds_split.attrs['river_outlet_equal_to_background_count'] == 1 + + ds_rrs = sizing_field_dataset( + ds_coastline=ds_coastline, + ds_river=ds_river, + resolution=0.03125, + mesh_name='ocn_rrs_6to18km_lnd_12km_riv_6km', + ocean_background=build_ocean_background_from_mode( + lat=ds_coastline.lat.values, + lon=ds_coastline.lon.values, + mode='rrs_latitude', + min_km=6.0, + max_km=18.0, + ), + land_background_km=12.0, + river_channel_km=6.0, + river_outlet_km=6.0, + ) + ocean_values = ds_rrs.cellWidth.values[:, 0] + assert ocean_values[0] < ocean_values[1] + assert np.isclose(ds_rrs.cellWidth.values[0, 1], ocean_values[0]) + assert np.isclose(ds_rrs.cellWidth.values[1, 1], 6.0) + assert ds_rrs.active_control.values[0, 1] == 0 + assert ds_rrs.active_control.values[1, 1] == 2 + + +def test_sizing_field_coastline_transition_on_land(): + ds_coastline = _make_coastline_dataset( + ocean_mask=np.array([[1, 1, 0, 0]], dtype=np.int8), + signed_distance=np.array([[0.0, 1000.0, -1000.0, -3000.0]]), + lat=np.array([0.0]), + lon=np.array([-30.0, -10.0, 10.0, 30.0]), + ) + ds_river = _make_river_dataset( + river_channel_mask=np.array([[0, 0, 0, 0]], dtype=np.int8), + river_outlet_mask=np.array([[0, 0, 0, 0]], dtype=np.int8), + lat=np.array([0.0]), + lon=np.array([-30.0, -10.0, 10.0, 30.0]), + ) + + ds_sizing = sizing_field_dataset( + ds_coastline=ds_coastline, + ds_river=ds_river, + resolution=0.125, + mesh_name='transition_test', + ocean_background=_constant_ocean_background( + ds_coastline=ds_coastline, value=30.0 + ), + land_background_km=10.0, + coastline_transition_land_km=2.0, + enable_river_channel_refinement=False, + enable_river_outlet_refinement=False, + ) + + np.testing.assert_allclose( + ds_sizing.coastline_cell_width.values[0], + np.array([30.0, 30.0, 20.0, 10.0]), + ) + np.testing.assert_allclose( + ds_sizing.cellWidth.values[0], np.array([30.0, 30.0, 20.0, 10.0]) + ) + np.testing.assert_array_equal( + ds_sizing.active_control.values[0], + np.array([0, 0, 1, 0], dtype=np.int8), + ) + + +def test_sizing_field_rivers_composed_before_coastline_transition(): + ds_coastline = _make_coastline_dataset( + ocean_mask=np.array([[1, 1, 1, 0]], dtype=np.int8), + signed_distance=np.array([[0.0, 1000.0, 3000.0, -1000.0]]), + lat=np.array([0.0]), + lon=np.array([-30.0, -10.0, 10.0, 30.0]), + ) + ds_river = _make_river_dataset( + river_channel_mask=np.array([[0, 0, 0, 0]], dtype=np.int8), + river_outlet_mask=np.array([[1, 0, 1, 0]], dtype=np.int8), + lat=np.array([0.0]), + lon=np.array([-30.0, -10.0, 10.0, 30.0]), + ) + + ds_sizing = sizing_field_dataset( + ds_coastline=ds_coastline, + ds_river=ds_river, + resolution=0.125, + mesh_name='river_transition_test', + ocean_background=_constant_ocean_background( + ds_coastline=ds_coastline, value=30.0 + ), + land_background_km=10.0, + coastline_transition_land_km=0.0, + river_channel_km=5.0, + river_outlet_km=5.0, + ) + + np.testing.assert_allclose( + ds_sizing.river_outlet_cell_width.values[0], + np.array([5.0, 10.0, 5.0, 10.0]), + ) + np.testing.assert_allclose( + ds_sizing.ocean_background_cell_width.values[0], + np.array([30.0, 30.0, 30.0, 30.0]), + ) + np.testing.assert_allclose( + ds_sizing.land_river_cell_width.values[0], + np.array([5.0, 10.0, 5.0, 10.0]), + ) + np.testing.assert_allclose( + ds_sizing.pre_coastline_cell_width.values[0], + np.array([30.0, 30.0, 30.0, 10.0]), + ) + np.testing.assert_allclose( + ds_sizing.cellWidth.values[0], np.array([30.0, 30.0, 30.0, 10.0]) + ) + np.testing.assert_allclose( + ds_sizing.coastal_transition_delta.values[0], + np.array([0.0, 0.0, 0.0, 0.0]), + ) + np.testing.assert_array_equal( + ds_sizing.active_control.values[0], + np.array([0, 0, 0, 0], dtype=np.int8), + ) + + +def test_sizing_field_coastline_overrides_finer_land_and_river_controls(): + ds_coastline = _make_coastline_dataset( + ocean_mask=np.array([[1, 0, 0]], dtype=np.int8), + signed_distance=np.array([[0.0, -1000.0, -3000.0]]), + lat=np.array([0.0]), + lon=np.array([-30.0, 0.0, 30.0]), + ) + ds_river = _make_river_dataset( + river_channel_mask=np.array([[0, 1, 0]], dtype=np.int8), + river_outlet_mask=np.array([[0, 0, 0]], dtype=np.int8), + lat=np.array([0.0]), + lon=np.array([-30.0, 0.0, 30.0]), + ) + + ds_sizing = sizing_field_dataset( + ds_coastline=ds_coastline, + ds_river=ds_river, + resolution=0.125, + mesh_name='coastal_override_test', + ocean_background=_constant_ocean_background( + ds_coastline=ds_coastline, value=30.0 + ), + land_background_km=10.0, + coastline_transition_land_km=2.0, + river_channel_km=5.0, + river_outlet_km=5.0, + ) + + np.testing.assert_allclose( + ds_sizing.land_river_cell_width.values[0], np.array([10.0, 5.0, 10.0]) + ) + np.testing.assert_allclose( + ds_sizing.coastline_cell_width.values[0], np.array([30.0, 17.5, 10.0]) + ) + np.testing.assert_allclose( + ds_sizing.cellWidth.values[0], np.array([30.0, 17.5, 10.0]) + ) + np.testing.assert_allclose( + ds_sizing.coastal_transition_delta.values[0], + np.array([0.0, 12.5, 0.0]), + ) + np.testing.assert_array_equal( + ds_sizing.active_control.values[0], np.array([0, 1, 0], dtype=np.int8) + ) + + +def test_unified_cell_width_mesh_step_reads_sizing_field(tmp_path): + step = UnifiedCellWidthMeshStep( + component=Component(name='mesh'), + subdir='spherical/unified/base_mesh/test', + ) + step.work_dir = str(tmp_path) + + ds = xr.Dataset( + data_vars=dict( + cellWidth=(('lat', 'lon'), np.array([[1.0, 2.0], [3.0, 4.0]])) + ), + coords=dict( + lat=xr.DataArray(np.array([-45.0, 45.0]), dims=('lat',)), + lon=xr.DataArray(np.array([-90.0, 90.0]), dims=('lon',)), + ), + ) + ds.to_netcdf(tmp_path / 'sizing_field.nc') + + cwd = os.getcwd() + try: + os.chdir(tmp_path) + cell_width, lon, lat = step.build_cell_width_lat_lon() + finally: + os.chdir(cwd) + + np.testing.assert_allclose(cell_width, ds.cellWidth.values) + np.testing.assert_allclose(lon, ds.lon.values) + np.testing.assert_allclose(lat, ds.lat.values) + + +def test_generic_ocean_background_modes(): + lat = np.array([-60.0, 0.0]) + lon = np.array([-120.0, 0.0, 120.0]) + + constant = build_ocean_background_from_mode( + lat=lat, + lon=lon, + mode='constant', + min_km=30.0, + max_km=30.0, + ) + np.testing.assert_allclose(constant, 30.0) + + rrs = build_ocean_background_from_mode( + lat=lat, + lon=lon, + mode='rrs_latitude', + min_km=6.0, + max_km=18.0, + ) + assert rrs.shape == (2, 3) + assert rrs[0, 0] < rrs[1, 0] + + +def test_constant_ocean_background_requires_equal_min_max(): + lat = np.array([-60.0, 0.0]) + lon = np.array([-120.0, 0.0, 120.0]) + + try: + build_ocean_background_from_mode( + lat=lat, + lon=lon, + mode='constant', + min_km=6.0, + max_km=18.0, + ) + except ValueError as exc: + message = str(exc) + else: + raise AssertionError('Expected constant mode to raise ValueError.') + + assert 'ocean_background_min_km' in message + assert 'ocean_background_max_km' in message + + +def test_generic_ocean_background_rejects_ec_latitude(): + lat = np.array([-60.0, 0.0]) + lon = np.array([-120.0, 0.0, 120.0]) + + try: + build_ocean_background_from_mode( + lat=lat, + lon=lon, + mode='ec_latitude', + min_km=6.0, + max_km=18.0, + ) + except ValueError as exc: + message = str(exc) + else: + raise AssertionError('Expected ec_latitude to raise ValueError.') + + assert 'ec_latitude' in message + assert 'Mesh-specific ocean backgrounds' in message + + +def test_add_sizing_field_tasks_registers_named_meshes(): + component = Component(name='mesh') + add_sizing_field_tasks(component=component) + + assert 'ocn_so_12to30km_lnd_10km_riv_10km' in UNIFIED_MESH_NAMES + assert len(component.tasks) == len(UNIFIED_MESH_NAMES) + + for mesh_name in UNIFIED_MESH_NAMES: + subdir = f'spherical/unified/{mesh_name}/sizing_field/task' + assert subdir in component.tasks + task = component.tasks[subdir] + assert task.name == f'sizing_field_{mesh_name}_task' + + +def test_sizing_field_step_factory_uses_mesh_subdir(): + component = Component(name='mesh') + mesh_name = 'ocn_30km_lnd_10km_riv_10km' + coastline_step = SimpleNamespace( + subdir='spherical/unified/coastline/lat_lon/0.12500_degree/prepare', + path='coastline', + output_filenames={'calving_front': 'coastline.nc'}, + ) + river_step = SimpleNamespace( + subdir='spherical/unified/ocn_30km_lnd_10km_riv_10km/river/lat_lon/prepare', + path='river', + masks_filename='river_network.nc', + ) + + steps, config = get_lat_lon_sizing_field_steps( + component=component, + coastline_step=coastline_step, + river_step=river_step, + mesh_name=mesh_name, + include_viz=True, + ) + + assert ( + steps['sizing_field'].subdir + == f'spherical/unified/{mesh_name}/sizing_field/build' + ) + assert ( + steps['viz_sizing_field'].subdir + == f'{steps["sizing_field"].subdir}/viz' + ) + assert config.get('unified_mesh', 'mesh_name') == mesh_name + assert config.has_option('unified_mesh', 'coastline_convention') + + +def test_sizing_field_step_factory_uses_mesh_family(): + component = Component(name='mesh') + mesh_name = 'ocn_so_12to30km_lnd_10km_riv_10km' + coastline_step = SimpleNamespace( + subdir='spherical/unified/coastline/lat_lon/0.06250_degree/prepare', + path='coastline', + output_filenames={'calving_front': 'coastline.nc'}, + ) + river_step = SimpleNamespace( + subdir='spherical/unified/ocn_so_12to30km_lnd_10km_riv_10km/' + 'river/lat_lon/prepare', + path='river', + masks_filename='river_network.nc', + ) + + steps, config = get_lat_lon_sizing_field_steps( + component=component, + coastline_step=coastline_step, + river_step=river_step, + mesh_name=mesh_name, + include_viz=False, + ) + + assert type(steps['sizing_field']) is BuildSizingFieldStep + assert config.get('unified_mesh', 'mesh_family') == 'so_region' + assert get_unified_mesh_family(config).name == 'so_region' + + +def test_sizing_field_step_factory_reuses_shared_config_for_viz(): + component = Component(name='mesh') + mesh_name = 'ocn_30km_lnd_10km_riv_10km' + coastline_step = SimpleNamespace( + subdir='spherical/unified/coastline/lat_lon/0.12500_degree/prepare', + path='coastline', + output_filenames={'calving_front': 'coastline.nc'}, + ) + river_step = SimpleNamespace( + subdir='spherical/unified/ocn_30km_lnd_10km_riv_10km/river/lat_lon/prepare', + path='river', + masks_filename='river_network.nc', + ) + + build_steps, _ = get_lat_lon_sizing_field_steps( + component=component, + coastline_step=coastline_step, + river_step=river_step, + mesh_name=mesh_name, + include_viz=False, + ) + steps, config = get_lat_lon_sizing_field_steps( + component=component, + coastline_step=coastline_step, + river_step=river_step, + mesh_name=mesh_name, + include_viz=True, + ) + + assert steps['sizing_field'] is build_steps['sizing_field'] + assert len(steps) == 2 + assert config is component.configs[config.filepath] + + +def test_so_mesh_family_links_shared_region_and_builds_field( + tmp_path, +): + component = Component(name='mesh') + coastline_step = SimpleNamespace( + subdir='spherical/unified/coastline/lat_lon/0.06250_degree/prepare', + path='coastline', + output_filenames={'calving_front': 'coastline.nc'}, + ) + river_step = SimpleNamespace( + subdir='spherical/unified/ocn_so_12to30km_lnd_10km_riv_10km/' + 'river/lat_lon/prepare', + path='river', + masks_filename='river_network.nc', + ) + config = get_sizing_field_config( + mesh_name='ocn_so_12to30km_lnd_10km_riv_10km', + filepath='mesh/spherical/unified/' + 'ocn_so_12to30km_lnd_10km_riv_10km/sizing_field/sizing_field.cfg', + ) + step = BuildSizingFieldStep( + component=component, + coastline_step=coastline_step, + river_step=river_step, + subdir='spherical/unified/ocn_so_12to30km_lnd_10km_riv_10km/' + 'sizing_field/build', + ) + step.set_shared_config(config, link='sizing_field.cfg') + step.setup() + + assert get_unified_mesh_family(config).name == 'so_region' + assert any( + input_file['filename'] == SO_REGION_FILENAME + and input_file['package'] == SO_REGION_PACKAGE + for input_file in step.input_data + ) + + ds_coastline = _make_coastline_dataset( + ocean_mask=np.ones((3, 4), dtype=np.int8), + signed_distance=np.zeros((3, 4), dtype=float), + lat=np.array([-80.0, -60.0, -20.0]), + lon=np.array([-180.0, -60.0, 60.0, 180.0]), + ) + + region_resource = imp_res.files(SO_REGION_PACKAGE).joinpath( + SO_REGION_FILENAME + ) + with imp_res.as_file(region_resource) as region_path: + os.symlink(region_path, tmp_path / SO_REGION_FILENAME) + cwd = os.getcwd() + try: + os.chdir(tmp_path) + ocean_background = step._get_ocean_background( + ds_coastline=ds_coastline, + section=config['sizing_field'], + ) + finally: + os.chdir(cwd) + + assert ocean_background.shape == (3, 4) + assert np.min(ocean_background) < np.max(ocean_background) + assert np.mean(ocean_background[0, :]) < np.mean(ocean_background[-1, :]) + + +def _make_coastline_dataset(ocean_mask, signed_distance, lat=None, lon=None): + if lat is None: + lat = np.array([-30.0, 30.0]) + if lon is None: + lon = np.array([-120.0, 0.0, 120.0]) + + return xr.Dataset( + data_vars=dict( + ocean_mask=(('lat', 'lon'), ocean_mask), + signed_distance=(('lat', 'lon'), signed_distance), + ), + coords=dict( + lat=xr.DataArray(np.asarray(lat), dims=('lat',)), + lon=xr.DataArray(np.asarray(lon), dims=('lon',)), + ), + ) + + +def _make_river_dataset( + river_channel_mask, river_outlet_mask, lat=None, lon=None +): + if lat is None: + lat = np.array([-30.0, 30.0]) + if lon is None: + lon = np.array([-120.0, 0.0, 120.0]) + + zeros = np.zeros_like(river_channel_mask) + return xr.Dataset( + data_vars=dict( + river_channel_mask=(('lat', 'lon'), river_channel_mask), + river_outlet_mask=(('lat', 'lon'), river_outlet_mask), + river_ocean_outlet_mask=(('lat', 'lon'), zeros), + river_inland_sink_mask=(('lat', 'lon'), zeros), + ), + coords=dict( + lat=xr.DataArray(np.asarray(lat), dims=('lat',)), + lon=xr.DataArray(np.asarray(lon), dims=('lon',)), + ), + ) + + +def _constant_ocean_background(ds_coastline, value): + lat_size = ds_coastline.sizes['lat'] + lon_size = ds_coastline.sizes['lon'] + return np.full((lat_size, lon_size), value, dtype=float) diff --git a/tests/e3sm/init/topo/test_combine_topo.py b/tests/test_archive.py similarity index 62% rename from tests/e3sm/init/topo/test_combine_topo.py rename to tests/test_archive.py index c17965b29b..b8c8476614 100644 --- a/tests/e3sm/init/topo/test_combine_topo.py +++ b/tests/test_archive.py @@ -7,7 +7,7 @@ def _load_archive_module(): module_path = ( - Path(__file__).resolve().parents[4] / 'polaris' / 'archive.py' + Path(__file__).resolve().parents[1] / 'polaris' / 'archive.py' ) spec = importlib.util.spec_from_file_location( 'polaris_archive', module_path @@ -46,3 +46,22 @@ def test_find_zip_member_raises_on_ambiguous_basename(tmp_path): with zipfile.ZipFile(zip_path) as archive: with pytest.raises(ValueError, match='Multiple ZIP members'): archive_module.find_zip_member(archive, 'GEBCO_2023.nc') + + +def test_extract_zip_subdir_extracts_matching_members(tmp_path): + archive_module = _load_archive_module() + zip_path = tmp_path / 'HydroRIVERS_v10_shp.zip' + + with zipfile.ZipFile(zip_path, 'w') as archive: + archive.writestr('HydroRIVERS_TechDoc_v10.pdf', b'pdf') + archive.writestr('HydroRIVERS_v10_shp/HydroRIVERS_v10.shp', b'shp') + archive.writestr('HydroRIVERS_v10_shp/HydroRIVERS_v10.dbf', b'dbf') + + archive_module.extract_zip_subdir( + str(zip_path), 'HydroRIVERS_v10_shp', out_dir=str(tmp_path) + ) + + shp_dir = tmp_path / 'HydroRIVERS_v10_shp' + assert (shp_dir / 'HydroRIVERS_v10.shp').read_bytes() == b'shp' + assert (shp_dir / 'HydroRIVERS_v10.dbf').read_bytes() == b'dbf' + assert not (tmp_path / 'HydroRIVERS_TechDoc_v10.pdf').exists()