Generalize topo/combine to work on lat-lon grids

docs/developers_guide/e3sm/init/api.md

@@ -26,7 +26,10 @@
    CombineStep.setup
    CombineStep.constrain_resources
    CombineStep.run
-   CombineTask
+   get_cubed_sphere_topo_steps
+   get_lat_lon_topo_steps
+   CubedSphereCombineTask
+   LatLonCombineTask
    VizCombinedStep
    VizCombinedStep.setup
    VizCombinedStep.run

docs/developers_guide/e3sm/init/tasks/topo/combine.md

@@ -18,32 +18,37 @@ 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.CubedSphereCombineTask`
+and
+{py:class}`polaris.tasks.e3sm.init.topo.combine.LatLonCombineTask`
+wrap the `CombineStep` into tasks 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.
 
-## High-Resolution and Low-Resolution Versions
+## Target Grids and Resolutions
 
-There are two versions of the combine steps and task:
+The combine framework is now organized explicitly around the target grid and
+resolution rather than around a special “low-resolution” mode. Current tasks
+include:
 
-1. **Standard (High-Resolution) Version**: This version maps to a
-   high-resolution (ne3000, ~1 km) cubed-sphere grid by default, producing
-   topogrpahy that is suitable for remapping to standard and high-resolution
-   MPAS meshes (~60 km and finer).
+1. Cubed-sphere topography on `ne3000`
+2. Cubed-sphere topography on `ne120`
+3. Latitude-longitude topography on `0.2500_degree`
 
-2. **Low-Resolution Version**: This version uses a coarser ne120 (~25 km) grid
-   for faster remapping to coarse-resolution MPAS meshes (e.g., Icos240). It is
-   designed to reduce computational cost while still providing adequate accuracy
-   for low-resolution simulations used for regression testing rather than
-   science.
+These appear in task paths such as:
 
-The low-resolution version can be selected by setting the `low_res` parameter
-to `True` when creating the `CombineStep` or `CombineTask`.
+- `e3sm/init/topo/combine_bedmap3_gebco2023/cubed_sphere/ne3000/task`
+- `e3sm/init/topo/combine_bedmap3_gebco2023/cubed_sphere/ne120/task`
+- `e3sm/init/topo/combine_bedmap3_gebco2023/lat_lon/0.2500_degree/task`
+
+This structure makes it easier for downstream consumers such as ocean
+hydrography preprocessing to reuse combined topography products without
+embedding product-specific names into `e3sm/init`.
 
 ## Key Features
 
@@ -73,9 +78,6 @@ the configuration file. Key options include:
 - `lat_tiles` and `lon_tiles`: Number of tiles to split the global dataset for parallel remapping.
 - `renorm_thresh`: Threshold for renormalizing Antarctic variables during blending.
 
-For the low-resolution version, additional configuration options are provided
-in the `combine_low_res.cfg` file.
-
 ## Workflow
 
 1. **Setup**: The step downloads required datasets and sets up input/output
@@ -98,21 +100,30 @@ task:
 from polaris.tasks.e3sm.init.topo.combine import CombineStep
 
 component = task.component
-subdir = CombineStep.get_subdir(low_res=False)
+subdir = CombineStep.get_subdir()
 if subdir in component.steps:
     step = component.steps[subdir]
 else:
-    step = CombineStep(component=component, low_res=False)
+    step = CombineStep(component=component, subdir=subdir)
     component.add_step(step)
 task.add_step(step)
 ```
 
-To create a `CombineTask` for caching combined datasets:
+To create a cubed-sphere combine task for caching combined datasets:
+
+```python
+from polaris.tasks.e3sm.init.topo.combine import CubedSphereCombineTask
+
+combine_task = CubedSphereCombineTask(component=my_component, resolution=3000)
+my_component.add_task(combine_task)
+```
+
+To create a latitude-longitude combine task:
 
 ```python
-from polaris.tasks.e3sm.init.topo.combine import CombineTask
+from polaris.tasks.e3sm.init.topo.combine import LatLonCombineTask
 
-combine_task = CombineTask(component=my_component, low_res=False)
+combine_task = LatLonCombineTask(component=my_component, resolution=0.25)
 my_component.add_task(combine_task)
 ```
 
@@ -133,7 +144,8 @@ The `VizCombinedStep` is typically added only when visualization is explicitly r
 
 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.
+{py:class}`polaris.tasks.e3sm.init.topo.combine.CubedSphereCombineTask` and
+{py:class}`polaris.tasks.e3sm.init.topo.combine.LatLonCombineTask` classes.
 
 ```{note}
 Since this step is expensive and time-consuming to run, most tasks will

polaris/tasks/e3sm/init/add_tasks.py

@@ -1,4 +1,7 @@
-from polaris.tasks.e3sm.init.topo.combine import CombineTask as CombineTopoTask
+from polaris.tasks.e3sm.init.topo.combine import (
+    CubedSphereCombineTask,
+    LatLonCombineTask,
+)
 from polaris.tasks.e3sm.init.topo.cull import add_cull_topo_tasks
 from polaris.tasks.e3sm.init.topo.remap import add_remap_topo_tasks
 
@@ -10,9 +13,15 @@ def add_e3sm_init_tasks(component):
     component : polaris.Component
         the e3sm/init component that the tasks will be added to
     """
-    for low_res in [False, True]:
+    for cubed_sphere_res in [3000, 120]:
         component.add_task(
-            CombineTopoTask(component=component, low_res=low_res)
+            CubedSphereCombineTask(
+                component=component, resolution=cubed_sphere_res
+            )
+        )
+    for lat_lon_res in [0.0625, 0.25, 1.0]:
+        component.add_task(
+            LatLonCombineTask(component=component, resolution=lat_lon_res)
         )
 
     add_remap_topo_tasks(component=component)

polaris/tasks/e3sm/init/topo/combine/__init__.py

@@ -2,10 +2,16 @@
     CombineStep as CombineStep,
 )
 from polaris.tasks.e3sm.init.topo.combine.steps import (
-    get_combine_topo_steps as get_combine_topo_steps,
+    get_cubed_sphere_topo_steps as get_cubed_sphere_topo_steps,
+)
+from polaris.tasks.e3sm.init.topo.combine.steps import (
+    get_lat_lon_topo_steps as get_lat_lon_topo_steps,
+)
+from polaris.tasks.e3sm.init.topo.combine.task import (
+    CubedSphereCombineTask as CubedSphereCombineTask,
 )
 from polaris.tasks.e3sm.init.topo.combine.task import (
-    CombineTask as CombineTask,
+    LatLonCombineTask as LatLonCombineTask,
 )
 from polaris.tasks.e3sm.init.topo.combine.viz import (
     VizCombinedStep as VizCombinedStep,

polaris/tasks/e3sm/init/topo/combine/combine.cfg

@@ -26,3 +26,39 @@ latmax = -60.
 # tractable
 lat_tiles = 3
 lon_tiles = 6
+
+
+[viz_combine_topo_base_elevation]
+colormap_name = cmo.topo
+norm_type = linear
+norm_args = {'vmin': -9000, 'vmax': 9000}
+under_color = black
+over_color = yellow
+
+[viz_combine_topo_ice_thickness]
+colormap_name = cmo.ice_r
+norm_type = linear
+norm_args = {'vmin': 0, 'vmax': 4000}
+under_color = cyan
+over_color = black
+
+[viz_combine_topo_ice_draft]
+colormap_name = cmo.topo
+norm_type = linear
+norm_args = {'vmin': -2000, 'vmax': 2000}
+under_color = black
+over_color = yellow
+
+[viz_combine_topo_ice_mask]
+colormap_name = cmo.amp_r
+norm_type = linear
+norm_args = {'vmin': 0, 'vmax': 1}
+under_color = red
+over_color = yellow
+
+[viz_combine_topo_grounded_mask]
+colormap_name = cmo.amp_r
+norm_type = linear
+norm_args = {'vmin': 0, 'vmax': 1}
+under_color = red
+over_color = yellow

polaris/tasks/e3sm/init/topo/combine/step.py

@@ -56,21 +56,14 @@ class CombineStep(Step):
     }
 
     @staticmethod
-    def get_subdir(low_res):
+    def get_subdir():
         """
-        Get the subdirectory for the step based on the datasets
-        Parameters
-        ----------
-        low_res : bool, optional
-            Whether to use the low resolution configuration options
+        Get the base subdirectory for the step based on the datasets.
         """
-        suffix = '_low_res' if low_res else ''
-        subdir = (
-            f'combine_{CombineStep.ANTARCTIC}_{CombineStep.GLOBAL}{suffix}'
-        )
+        subdir = f'combine_{CombineStep.ANTARCTIC}_{CombineStep.GLOBAL}'
         return os.path.join('topo', subdir)
 
-    def __init__(self, component, subdir, low_res=False):
+    def __init__(self, component, subdir):
         """
         Create a new step
 
@@ -82,13 +75,10 @@ def __init__(self, component, subdir, low_res=False):
         subdir : str
             The subdirectory within the component's work directory
 
-        low_res : bool, optional
-            Whether to use the low resolution configuration options
         """
         antarctic_dataset = self.ANTARCTIC
         global_dataset = self.GLOBAL
-        suffix = '_low_res' if low_res else ''
-        name = f'combine_topo_{antarctic_dataset}_{global_dataset}{suffix}'
+        name = f'combine_topo_{antarctic_dataset}_{global_dataset}'
         super().__init__(
             component=component,
             name=name,
@@ -266,11 +256,15 @@ def _set_res_and_outputs(self, update):
                 f'{res_name}.nc',
             ]
         )
-        self.exodus_filename = f'{self.resolution_name}.g'
+        if target_grid == 'cubed_sphere':
+            self.exodus_filename = f'{self.resolution_name}.g'
+        else:
+            self.exodus_filename = None
 
         self.add_output_file(filename=self.dst_scrip_filename)
         self.add_output_file(filename=self.combined_filename)
-        self.add_output_file(filename=self.exodus_filename)
+        if self.exodus_filename is not None:
+            self.add_output_file(filename=self.exodus_filename)
 
         if update:
             # We need to set absolute paths
@@ -282,10 +276,10 @@ def _set_res_and_outputs(self, update):
 
     def _modify_gebco(self, in_filename, out_filename):
         """
-        Modify GEBCO to include lon/lat bounds located at grid edges
+        Modify GEBCO to include lon/lat bounds and an ocean mask
         """
         logger = self.logger
-        logger.info('Adding bounds to GEBCO lat/lon')
+        logger.info('Adding bounds and an ocean mask to GEBCO')
 
         # Modify GEBCO
         gebco = xr.open_dataset(in_filename)
@@ -299,6 +293,7 @@ def _modify_gebco(self, in_filename, out_filename):
         gebco['lon_bnds'] = lon_bnds.transpose('lon', 'bnds')
         gebco.lat.attrs['bounds'] = 'lat_bnds'
         gebco.lon.attrs['bounds'] = 'lon_bnds'
+        gebco['ocean_mask'] = (gebco.elevation < 0.0).astype(float)
 
         # Write modified GEBCO to netCDF
         _write_netcdf_with_fill_values(gebco, out_filename)
@@ -451,10 +446,18 @@ def _create_global_tile(self, global_filename, lon_tile, lat_tile):
 
         # Select tile from dataset
         tile = ds.isel(lat=lat_indices, lon=lon_indices)
+        lat = tile.lat.values.copy()
+        lat_attrs = tile.lat.attrs.copy()
+        # xarray may expose coordinate views from ``isel()`` as read-only.
+        # Reassign corrected pole coordinates instead of mutating in place,
+        # while preserving CF metadata that ESMF uses to detect CFGRID files.
         if lat_tile == 0:
-            tile.lat.values[0] = -90.0  # Correct south pole
+            lat[0] = -90.0  # Correct south pole
         if lat_tile == lat_tiles - 1:
-            tile.lat.values[-1] = 90.0  # Correct north pole
+            lat[-1] = 90.0  # Correct north pole
+        tile = tile.assign_coords(
+            lat=xr.DataArray(lat, dims=tile.lat.dims, attrs=lat_attrs)
+        )
 
         # Write tile to netCDF
         _write_netcdf_with_fill_values(tile, out_filename)
@@ -702,14 +705,14 @@ def _remap_global(self, global_filename, out_filename):
 
                 # Add tile to remapped global topography
                 logger.info(f'    adding {remapped_filename}')
-                elevation = xr.open_dataset(remapped_filename).elevation
-                elevation = elevation.where(elevation.notnull(), 0.0)
-                if 'elevation' in global_remapped:
-                    global_remapped['elevation'] = (
-                        global_remapped.elevation + elevation
-                    )
-                else:
-                    global_remapped['elevation'] = elevation
+                ds_tile = xr.open_dataset(remapped_filename)
+                for field in ['elevation', 'ocean_mask']:
+                    da = ds_tile[field]
+                    da = da.where(da.notnull(), 0.0)
+                    if field in global_remapped:
+                        global_remapped[field] = global_remapped[field] + da
+                    else:
+                        global_remapped[field] = da
 
         # Write tile to netCDF
         logger.info(f'    writing {out_filename}')
@@ -781,6 +784,10 @@ def _combine_datasets(
         global_elevation = global_elevation.where(
             global_elevation.notnull(), 0.0
         )
+        global_ocean_mask = ds_global.ocean_mask
+        global_ocean_mask = global_ocean_mask.where(
+            global_ocean_mask.notnull(), 0.0
+        )
 
         # Load and mask Antarctic dataset
         ds_antarctic = xr.open_dataset(antarctic_filename)
@@ -795,6 +802,7 @@ def _combine_datasets(
             'ice_draft',
             'ice_mask',
             'grounded_mask',
+            'ocean_mask',
         ]
 
         for var in vars:
@@ -822,13 +830,20 @@ def _combine_datasets(
         # Add masks
         for field in ['ice_mask', 'grounded_mask']:
             combined[field] = ds_antarctic[field]
+        antarctic_ocean_mask = ds_antarctic.ocean_mask.where(
+            ds_antarctic.ocean_mask.notnull(), global_ocean_mask
+        )
+        combined['ocean_mask'] = (
+            alpha * global_ocean_mask + (1.0 - alpha) * antarctic_ocean_mask
+        )
 
         # Add fill values
         fill_vals = {
             'ice_draft': 0.0,
             'ice_thickness': 0.0,
             'ice_mask': 0.0,
             'grounded_mask': 0.0,
+            'ocean_mask': 0.0,
         }
         for field, fill_val in fill_vals.items():
             valid = combined[field].notnull()