Merge pull request #300 from xylar/add-global-combine-topo-step

xylar · web-flow · commit d7cf8ae1b279 · 2025-05-09T11:07:08.000+02:00
Add a step for combining global and Antarctic topography
diff --git a/.github/workflows/build_workflow.yml b/.github/workflows/build_workflow.yml
@@ -106,6 +106,7 @@ jobs:
             --verbose \
             --python=${{ matrix.python-version }}
           source load_polaris_test_mpich.sh
+          python -c "import polaris; import polaris.version; print(polaris.version.__version__)"
 
       - if: ${{ steps.skip_check.outputs.should_skip != 'true' }}
         name: Build Sphinx Docs
diff --git a/docs/Makefile b/docs/Makefile
@@ -30,7 +30,7 @@ help:
 
 
 clean:
-	rm -rf generated
+	rm -rf $(BUILDDIR) developers_guide/generated/ developers_guide/*/generated/ developers_guide/*/*/generated/
 	@$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 
 clean-versioned-html:
@@ -47,6 +47,3 @@ clean-versioned-html:
 # raise warnings to errors
 html-strict:
 	@$(SPHINXBUILD) -b html -nW --keep-going "$(SOURCEDIR)" "$(BUILDDIR)/html" $(SPHINXOPTS) $(O)
-
-clean:
-	rm -rf $(BUILDDIR) developers_guide/generated/ developers_guide/*/generated/
diff --git a/docs/conf.py b/docs/conf.py
@@ -47,6 +47,7 @@
 ]
 
 autosummary_generate = ['developers_guide/api.md',
+                        'developers_guide/e3sm/init/api.md',
                         'developers_guide/mesh/api.md',
                         'developers_guide/ocean/api.md',
                         'developers_guide/seaice/api.md']
diff --git a/docs/developers_guide/e3sm/init/api.md b/docs/developers_guide/e3sm/init/api.md
@@ -0,0 +1,31 @@
+# e3sm/init
+
+```{eval-rst}
+.. currentmodule:: polaris.tasks.e3sm.init
+
+.. autosummary::
+   :toctree: generated/
+
+   add_tasks.add_e3sm_init_tasks
+```
+
+## Tasks
+
+### topo
+
+```{eval-rst}
+.. currentmodule:: polaris.tasks.e3sm.init.topo.combine
+
+.. autosummary::
+   :toctree: generated/
+
+   CombineStep
+   CombineStep.get_subdir
+   CombineStep.setup
+   CombineStep.constrain_resources
+   CombineStep.run
+   CombineTask
+   VizCombinedStep
+   VizCombinedStep.setup
+   VizCombinedStep.run
+```
diff --git a/docs/developers_guide/e3sm/init/index.md b/docs/developers_guide/e3sm/init/index.md
@@ -0,0 +1,13 @@
+(dev-e3sm-init)=
+
+# E3SM initial condition (e3sm/init) component
+
+The `e3sm/init` component defines tasks and steps related to generating ocean
+and sea-ice initial conditions for E3SM.  There are currently no
+component-level config options or shared framework.
+
+```{toctree}
+:titlesonly: true
+
+tasks/index
+```
diff --git a/docs/developers_guide/e3sm/init/tasks/images/bathymetry_500.png b/docs/developers_guide/e3sm/init/tasks/images/bathymetry_500.png
diff --git a/docs/developers_guide/e3sm/init/tasks/index.md b/docs/developers_guide/e3sm/init/tasks/index.md
@@ -0,0 +1,9 @@
+(dev-e3sm-init-tasks)=
+
+# Tasks
+
+```{toctree}
+:titlesonly: true
+
+topo
+```
diff --git a/docs/developers_guide/e3sm/init/tasks/topo.md b/docs/developers_guide/e3sm/init/tasks/topo.md
@@ -0,0 +1,134 @@
+# Topography Tasks
+
+The `polaris.tasks.e3sm.init.topo` module provides tools for working with
+topography data in Polaris. This includes steps for processing, modifying, and
+combining topography datasets to create inputs for E3SM components such as
+MPAS-Ocean. The framework is designed to handle both global and regional
+datasets, supporting various grid types like lat-lon and cubed-sphere grids.
+
+## Combine Steps and Task
+
+```{image} images/bathymetry_500.png
+:align: center
+:width: 500 px
+```
+
+Global bathymetry datasets do not typically include the latest datasets around
+Antarctica needed for ice-sheet and ice-shelf modeling.  For this reason, we
+typically combine a global topography dataset north of the Southern Ocean with
+one for Antarctica.
+
+```{note}
+At the moment, the step for combining these datasets provides fields that are
+masked to locations where the bed topography (bathymetry) is below sea level.
+This avoids the risk of interpolated topography resulting in a bed that is
+above sea level in ocean regions when we perform further interpolation of
+the data to an MPAS mesh.  However, a more general topogrpahy data set will
+likely be needed in the future that accommodates both the ocean and land/river
+components.
+```
+
+The {py:class}`polaris.tasks.e3sm.init.topo.combine.CombineStep` step is a key
+component of the topography framework. It is responsible for combining global
+and Antarctic topography datasets into a single dataset suitable for use in
+E3SM simulations. The step supports blending datasets across specified latitude
+ranges and remapping them to a target grid.
+
+The {py:class}`polaris.tasks.e3sm.init.topo.combine.CombineTask` wraps the
+`CombineStep` into a task that can be used to generate and cache combined
+topography datasets for reuse in other contexts.
+
+The {py:class}`polaris.tasks.e3sm.init.topo.combine.VizCombinedStep` step is
+an optional visualization step that can be added to the workflow to create
+plots of the combined topography dataset. This step is particularly useful for
+debugging or analyzing the combined dataset.
+
+### Key Features
+
+- **Dataset Support**: Supports multiple datasets, including `bedmap3`,
+  `bedmachinev3`, and `gebco2023`.
+- **Grid Types**: Handles both lat-lon and cubed-sphere target grids.
+- **Blending**: Blends global and Antarctic datasets across a configurable
+  latitude range.
+- **Remapping**: Uses tools like `mbtempest`, `ESMF_RegridWeightGen` and
+  `ncremap` for remapping datasets to the target grid.
+- **Output**: Produces combined topography datasets with consistent variables
+  and attributes.
+- **Visualization**: Generates rasterized images of various fields (e.g.,
+  bathymetry, ice draft) using the `datashader` library.
+
+### Configuration Options
+
+The `CombineStep` step is configured through the `[combine_topo]` section in
+the configuration file. Key options include:
+
+- `resolution_latlon`: Target resolution for lat-lon grids (in degrees).
+- `resolution_cubedsphere`: Target resolution for cubed-sphere grids (e.g.,
+  `3000` for NExxx grids).
+- `latmin` and `latmax`: Latitude range for blending datasets.
+- `ntasks` and `min_tasks`: Number of MPI tasks for remapping.
+- `method`: Remapping method (e.g., `bilinear`).
+
+### Workflow
+
+1. **Setup**: The step downloads required datasets and sets up input/output
+   files.
+2. **Modification**: Antarctic and global datasets are modified to include
+   necessary variables and attributes.
+3. **Remapping**: Datasets are remapped to the target grid using SCRIP files
+   and weight generation.
+4. **Blending**: The datasets are blended across the specified latitude range.
+5. **Output**: The combined dataset is saved in NetCDF format.
+8. **Optional Field Plotting**: Each field in the dataset is rasterized and saved as an image with a colorbar.
+
+### Example Usage
+
+Below is an example of how the `CombineStep` can be added to a Polaris
+task:
+
+```python
+from polaris.tasks.e3sm.init.topo.combine import CombineStep
+
+
+component = task.component
+subdir = CombineStep.get_subdir()
+if subdir in component.steps:
+    step = component.steps[subdir]
+else:
+    step = CombineStep(component=component)
+    component.add_step(step)
+task.add_step(step)
+```
+
+To create a `CombineTask` for caching combined datasets:
+
+```python
+from polaris.tasks.e3sm.init.topo.combine import CombineTask
+
+combine_task = CombineTask(component=my_component)
+my_component.add_task(combine_task)
+```
+
+Below is an example of how the `VizCombinedStep` can be added to a Polaris task:
+
+```python
+from polaris.tasks.e3sm.init.topo.combine import VizCombinedStep
+
+viz_step = VizCombinedStep(component=my_component, combine_step=combine_step)
+my_component.add_step(viz_step)
+```
+
+Since there is a single shared step for each pair of Antarctic and global
+datasets, the step should be added only once to the component and the existing
+step (identifiable via its `subdir`) should be used subsequently.
+
+The `VizCombinedStep` is typically added only when visualization is explicitly required, as it is not part of the default workflow.
+
+For more details, refer to the source code of the
+{py:class}`polaris.tasks.e3sm.init.topo.combine.CombineStep` and
+{py:class}`polaris.tasks.e3sm.init.topo.combine.CombineTask` classes.
+
+```{note}
+Since this step is expensive and time-consuming to run, most tasks will
+want to use cached outputs from this step rather than running it in full.
+```
diff --git a/docs/index.md b/docs/index.md
@@ -70,6 +70,7 @@ developers_guide/quick_start
 developers_guide/overview
 developers_guide/command_line
 developers_guide/organization/index
+developers_guide/e3sm/init/index
 developers_guide/mesh/index
 developers_guide/ocean/index
 developers_guide/seaice/index
diff --git a/polaris/cache.py b/polaris/cache.py
@@ -59,14 +59,16 @@ def update_cache(step_paths, date_string=None, dry_run=False):
         else:
             steps[component] = [step]
 
+    out_filename = f'{component.replace("/", "_")}_cached_files.json'
+
     # now, iterate over cores and steps
     for component in steps:
         database_root = config.get('paths', 'database_root')
         cache_root = f'{database_root}/{component}/polaris_cache'
 
-        package = f'polaris.{component}'
+        package = f'polaris.{component.replace("/", ".")}'
         try:
-            with open(f'{component}_cached_files.json') as data_file:
+            with open(out_filename) as data_file:
                 cached_files = json.load(data_file)
         except FileNotFoundError:
             # we don't have a local version of the file yet, let's see if
@@ -86,14 +88,14 @@ def update_cache(step_paths, date_string=None, dry_run=False):
 
             for output in step.outputs:
                 output = os.path.basename(output)
-                out_filename = os.path.join(step_path, output)
+                dest_filename = os.path.join(step_path, output)
                 # remove the component from the file path
-                target = out_filename[len(component) + 1 :]
+                target = dest_filename[len(component) + 1 :]
                 path, ext = os.path.splitext(target)
                 target = f'{path}.{date_string}{ext}'
-                cached_files[out_filename] = target
+                cached_files[dest_filename] = target
 
-                print(out_filename)
+                print(dest_filename)
                 print(f'  ==> {target}')
                 output_path = f'{cache_root}/{target}'
                 print(f'  copy to: {output_path}')
@@ -104,9 +106,8 @@ def update_cache(step_paths, date_string=None, dry_run=False):
                         os.makedirs(directory)
                     except FileExistsError:
                         pass
-                    shutil.copyfile(out_filename, output_path)
+                    shutil.copyfile(dest_filename, output_path)
 
-        out_filename = f'{component}_cached_files.json'
         with open(out_filename, 'w') as data_file:
             json.dump(cached_files, data_file, indent=4)
 
diff --git a/polaris/component.py b/polaris/component.py
@@ -132,7 +132,7 @@ def configure(self, config):
     def _read_cached_files(self):
         """Read in the dictionary of cached files from cached_files.json"""
 
-        package = f'polaris.{self.name}'
+        package = f'polaris.{self.name.replace("/", ".")}'
         filename = 'cached_files.json'
         try:
             pkg_file = imp_res.files(package).joinpath(filename)
diff --git a/polaris/e3sm/__init__.py b/polaris/e3sm/__init__.py
diff --git a/polaris/e3sm/init/__init__.py b/polaris/e3sm/init/__init__.py
diff --git a/polaris/e3sm/init/cached_files.json b/polaris/e3sm/init/cached_files.json
@@ -0,0 +1,5 @@
+{
+    "e3sm/init/topo/combine_bedmap3_gebco2023/ne3000.scrip.nc": "topo/combine_bedmap3_gebco2023/ne3000.scrip.250509.nc",
+    "e3sm/init/topo/combine_bedmap3_gebco2023/bedmap3_gebco2023_ne3000.nc": "topo/combine_bedmap3_gebco2023/bedmap3_gebco2023_ne3000.250509.nc",
+    "e3sm/init/topo/combine_bedmap3_gebco2023/ne3000.g": "topo/combine_bedmap3_gebco2023/ne3000.250509.g"
+}
diff --git a/polaris/list.py b/polaris/list.py
@@ -95,7 +95,7 @@ def list_suites(components=None, verbose=False):
         components = [component.name for component in get_components()]
     print('Suites:')
     for component in components:
-        package = f'polaris.suites.{component}'
+        package = f'polaris.suites.{component.replace("/", ".")}'
         try:
             suites = sorted(
                 [
diff --git a/polaris/setup.py b/polaris/setup.py
@@ -697,7 +697,7 @@ def _get_basic_config(config_file, machine, component_path, component, model):
 
     # add the config options for the component
     config.add_from_package(
-        f'polaris.{component.name}',
+        f'polaris.{component.name.replace("/", ".")}',
         f'{component.name}.cfg',
         exception=False,
     )
diff --git a/polaris/tasks/__init__.py b/polaris/tasks/__init__.py
@@ -3,6 +3,8 @@
 from polaris import Component
 
 # import new components here
+from polaris.tasks.e3sm.init import e3sm_init
+from polaris.tasks.e3sm.init.add_tasks import add_e3sm_init_tasks
 from polaris.tasks.mesh import mesh
 from polaris.tasks.mesh.add_tasks import add_mesh_tasks
 from polaris.tasks.ocean import ocean
@@ -12,6 +14,7 @@
 
 # Add new components alphabetically to this list
 _components: List[Component] = [
+    e3sm_init,
     mesh,
     ocean,
     seaice,
@@ -27,6 +30,7 @@ def get_components():
     global _tasks_added
     if not _tasks_added:
         # add tasks to each component
+        add_e3sm_init_tasks(component=e3sm_init)
         add_mesh_tasks(component=mesh)
         add_ocean_tasks(component=ocean)
         add_seaice_tasks(component=seaice)
diff --git a/polaris/tasks/e3sm/__init__.py b/polaris/tasks/e3sm/__init__.py
diff --git a/polaris/tasks/e3sm/init/__init__.py b/polaris/tasks/e3sm/init/__init__.py
@@ -0,0 +1,4 @@
+from polaris.component import Component
+
+# create a single module-level instance available to other components
+e3sm_init = Component(name='e3sm/init')
diff --git a/polaris/tasks/e3sm/init/add_tasks.py b/polaris/tasks/e3sm/init/add_tasks.py
@@ -0,0 +1,11 @@
+from polaris.tasks.e3sm.init.topo.combine import CombineTask as CombineTopoTask
+
+
+def add_e3sm_init_tasks(component):
+    """
+    Add all tasks to the e3sm/init component
+
+    component : polaris.Component
+        the e3sm/init component that the tasks will be added to
+    """
+    component.add_task(CombineTopoTask(component=component))
diff --git a/polaris/tasks/e3sm/init/topo/__init__.py b/polaris/tasks/e3sm/init/topo/__init__.py
diff --git a/polaris/tasks/e3sm/init/topo/combine/__init__.py b/polaris/tasks/e3sm/init/topo/combine/__init__.py
@@ -0,0 +1,12 @@
+from polaris.tasks.e3sm.init.topo.combine.step import (
+    CombineStep as CombineStep,
+)
+from polaris.tasks.e3sm.init.topo.combine.steps import (
+    get_combine_topo_steps as get_combine_topo_steps,
+)
+from polaris.tasks.e3sm.init.topo.combine.task import (
+    CombineTask as CombineTask,
+)
+from polaris.tasks.e3sm.init.topo.combine.viz import (
+    VizCombinedStep as VizCombinedStep,
+)
diff --git a/polaris/tasks/e3sm/init/topo/combine/combine.cfg b/polaris/tasks/e3sm/init/topo/combine/combine.cfg
@@ -0,0 +1,26 @@
+# config options related combining Antarctic and global topograph datasets
+[combine_topo]
+
+# target grid: lat_lon, cubed_sphere
+target_grid = cubed_sphere
+
+# target resolution (degrees or NExxx)
+resolution_latlon = 0.0125
+resolution_cubedsphere = 3000
+method = bilinear
+
+# threshold for masks below which interpolated variables are not renormalized
+renorm_thresh = 1e-3
+
+# the target and minimum number of MPI tasks to use in remapping
+ntasks = 1280
+min_tasks = 512
+
+# latitudes between which the topography datasets get blended
+latmin = -62.
+latmax = -60.
+
+# the number of decomposition tiles in lat and lon to make global remapping
+# tractable
+lat_tiles = 3
+lon_tiles = 6
diff --git a/polaris/tasks/e3sm/init/topo/combine/step.py b/polaris/tasks/e3sm/init/topo/combine/step.py
diff --git a/polaris/tasks/e3sm/init/topo/combine/steps.py b/polaris/tasks/e3sm/init/topo/combine/steps.py
diff --git a/polaris/tasks/e3sm/init/topo/combine/task.py b/polaris/tasks/e3sm/init/topo/combine/task.py
diff --git a/polaris/tasks/e3sm/init/topo/combine/viz.py b/polaris/tasks/e3sm/init/topo/combine/viz.py

Original file line number	Diff line number	Diff line change
`@@ -47,6 +47,7 @@`
`47`	`47`	`]`
`48`	`48`
`49`	`49`	`autosummary_generate = ['developers_guide/api.md',`
	`50`	`+ 'developers_guide/e3sm/init/api.md',`
`50`	`51`	`'developers_guide/mesh/api.md',`
`51`	`52`	`'developers_guide/ocean/api.md',`
`52`	`53`	`'developers_guide/seaice/api.md']`