From 7a06830c467c0ca0d3a51dee2fc3a13b3a6232fe Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 26 Mar 2026 21:03:49 +0100 Subject: [PATCH 01/13] Generalize e3sm/init topography combine tasks Organize combined-topography tasks by target grid and resolution rather than by downstream consumer, and add a reusable 0.25-degree lat-lon task for cached products. This renames the task helpers and task classes, updates remap and cull callers, and refreshes the e3sm/init developer docs to match the new layout. --- docs/developers_guide/e3sm/init/api.md | 5 +- .../e3sm/init/tasks/topo/combine.md | 62 ++++++---- polaris/tasks/e3sm/init/add_tasks.py | 10 +- .../tasks/e3sm/init/topo/combine/__init__.py | 10 +- polaris/tasks/e3sm/init/topo/combine/step.py | 20 +--- polaris/tasks/e3sm/init/topo/combine/steps.py | 113 ++++++++++++++---- polaris/tasks/e3sm/init/topo/combine/task.py | 74 +++++++++--- polaris/tasks/e3sm/init/topo/cull/tasks.py | 9 +- polaris/tasks/e3sm/init/topo/remap/tasks.py | 9 +- 9 files changed, 223 insertions(+), 89 deletions(-) diff --git a/docs/developers_guide/e3sm/init/api.md b/docs/developers_guide/e3sm/init/api.md index b72bf26cca..daf22ca34d 100644 --- a/docs/developers_guide/e3sm/init/api.md +++ b/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 diff --git a/docs/developers_guide/e3sm/init/tasks/topo/combine.md b/docs/developers_guide/e3sm/init/tasks/topo/combine.md index 57eba97862..e66ca372f3 100644 --- a/docs/developers_guide/e3sm/init/tasks/topo/combine.md +++ b/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 diff --git a/polaris/tasks/e3sm/init/add_tasks.py b/polaris/tasks/e3sm/init/add_tasks.py index 4f5c008b76..52dc5e4146 100644 --- a/polaris/tasks/e3sm/init/add_tasks.py +++ b/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,10 +13,11 @@ 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 resolution in [3000, 120]: component.add_task( - CombineTopoTask(component=component, low_res=low_res) + CubedSphereCombineTask(component=component, resolution=resolution) ) + component.add_task(LatLonCombineTask(component=component, resolution=0.25)) add_remap_topo_tasks(component=component) diff --git a/polaris/tasks/e3sm/init/topo/combine/__init__.py b/polaris/tasks/e3sm/init/topo/combine/__init__.py index 3d3f21a179..1496c36bf4 100644 --- a/polaris/tasks/e3sm/init/topo/combine/__init__.py +++ b/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, diff --git a/polaris/tasks/e3sm/init/topo/combine/step.py b/polaris/tasks/e3sm/init/topo/combine/step.py index d9f5ac095c..c95efb60c9 100644 --- a/polaris/tasks/e3sm/init/topo/combine/step.py +++ b/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, diff --git a/polaris/tasks/e3sm/init/topo/combine/steps.py b/polaris/tasks/e3sm/init/topo/combine/steps.py index 846f1eeab6..f59bde733e 100644 --- a/polaris/tasks/e3sm/init/topo/combine/steps.py +++ b/polaris/tasks/e3sm/init/topo/combine/steps.py @@ -5,64 +5,131 @@ from polaris.tasks.e3sm.init.topo.combine.viz import VizCombinedStep -def get_combine_topo_steps(component, include_viz=False, low_res=False): +def get_cubed_sphere_topo_steps(component, resolution, include_viz=False): """ - Get the shared combine topo step for the given component, adding it if - it doesn't exist + Get shared combined-topography steps for a cubed-sphere target grid. Parameters ---------- component : polaris.Component - The component that the steps will be added to + The component that the steps will be added to. + + resolution : int + The cubed-sphere resolution, such as 3000 or 120. + + include_viz : bool, optional + Whether to include the visualization step or not. + + Returns + ------- + steps : list of polaris.Step + The combine topo step and optional visualization step. + + config : polaris.config.PolarisConfigParser + The shared config options. + """ + return _get_target_topo_steps( + component=component, + target_grid='cubed_sphere', + resolution=resolution, + include_viz=include_viz, + ) + + +def get_lat_lon_topo_steps(component, resolution, include_viz=False): + """ + Get shared combined-topography steps for a latitude-longitude target grid. + + Parameters + ---------- + component : polaris.Component + The component that the steps will be added to. + + resolution : float + The latitude-longitude resolution in degrees. include_viz : bool, optional - Whether to include the visualization step or not - Default is False. + Whether to include the visualization step or not. - low_res : bool, optional - Whether to use low resolution config options + Returns + ------- + steps : list of polaris.Step + The combine topo step and optional visualization step. + + config : polaris.config.PolarisConfigParser + The shared config options. + """ + return _get_target_topo_steps( + component=component, + target_grid='lat_lon', + resolution=resolution, + include_viz=include_viz, + ) + + +def _get_target_topo_steps(component, target_grid, resolution, include_viz): + """ + Get shared combined-topography steps for a target grid and resolution. + + Parameters + ---------- + component : polaris.Component + The component that the steps will be added to. + + target_grid : {'cubed_sphere', 'lat_lon'} + The type of target grid. + + resolution : int or float + The target resolution as a cubed-sphere face count or degrees. + + include_viz : bool + Whether to include the visualization step or not. Returns ------- steps : list of polaris.Step - The combine topo step and optional visualization step + The combine topo step and optional visualization step. config : polaris.config.PolarisConfigParser - The shared config options + The shared config options. """ + if target_grid == 'cubed_sphere': + resolution_name = f'ne{int(resolution)}' + elif target_grid == 'lat_lon': + resolution_name = f'{float(resolution):.4f}_degree' + else: + raise ValueError(f'Unexpected target grid: {target_grid}') - subdir = CombineStep.get_subdir(low_res=low_res) + subdir = os.path.join( + CombineStep.get_subdir(), target_grid, resolution_name + ) - # add default config options for combining topo -- since these are - # shared steps, the config options need to be defined separately from any - # task this may be added to config_filename = 'combine_topo.cfg' filepath = os.path.join(component.name, subdir, config_filename) config = PolarisConfigParser(filepath=filepath) config.add_from_package( 'polaris.tasks.e3sm.init.topo.combine', 'combine.cfg' ) - if low_res: - config.add_from_package( - 'polaris.tasks.e3sm.init.topo.combine', 'combine_low_res.cfg' - ) + config.set('combine_topo', 'target_grid', target_grid) + if target_grid == 'cubed_sphere': + config.set('combine_topo', 'resolution_cubedsphere', f'{resolution}') + else: + config.set('combine_topo', 'resolution_latlon', f'{resolution}') steps = [] - # no config_filename is needed here since the shared config file is - # in this steps work directory combine_step = component.get_or_create_shared_step( step_cls=CombineStep, subdir=subdir, config=config, - low_res=low_res, ) + combine_step._set_res_and_outputs(update=False) steps.append(combine_step) if include_viz: - subdir = os.path.join(combine_step.subdir, 'viz') + viz_subdir = os.path.join(combine_step.subdir, 'viz') viz_step = component.get_or_create_shared_step( step_cls=VizCombinedStep, - subdir=subdir, + subdir=viz_subdir, config=config, config_filename='combine_topo.cfg', combine_step=combine_step, diff --git a/polaris/tasks/e3sm/init/topo/combine/task.py b/polaris/tasks/e3sm/init/topo/combine/task.py index c5156f5027..12f83b6e5c 100644 --- a/polaris/tasks/e3sm/init/topo/combine/task.py +++ b/polaris/tasks/e3sm/init/topo/combine/task.py @@ -2,43 +2,89 @@ from polaris.task import Task from polaris.tasks.e3sm.init.topo.combine.step import CombineStep -from polaris.tasks.e3sm.init.topo.combine.steps import get_combine_topo_steps +from polaris.tasks.e3sm.init.topo.combine.steps import ( + get_cubed_sphere_topo_steps, + get_lat_lon_topo_steps, +) -class CombineTask(Task): +class CubedSphereCombineTask(Task): """ - A task for creating the combined topography dataset, used to create the - files to be cached for use in all other contexts + A task for creating combined topography on a cubed-sphere target grid. """ - def __init__(self, component, low_res): + def __init__(self, component, resolution): """ - Create a new task + Create a new task. Parameters ---------- component : polaris.Component - The component the task belongs to + The component the task belongs to. - low_res : bool - Whether to use low resolution config options + resolution : int + The cubed-sphere resolution, such as 3000 or 120. """ antarctic_dataset = CombineStep.ANTARCTIC global_dataset = CombineStep.GLOBAL - suffix = '_low_res' if low_res else '' name = ( - f'combine_topo_{antarctic_dataset}_{global_dataset}{suffix}_task' + f'combine_topo_{antarctic_dataset}_{global_dataset}_' + f'cubed_sphere_ne{resolution}_task' + ) + subdir = os.path.join( + CombineStep.get_subdir(), + 'cubed_sphere', + f'ne{resolution}', + 'task', ) - subdir = os.path.join(CombineStep.get_subdir(low_res=low_res), 'task') super().__init__( component=component, name=name, subdir=subdir, ) - steps, config = get_combine_topo_steps( + steps, config = get_cubed_sphere_topo_steps( component=component, + resolution=resolution, include_viz=True, - low_res=low_res, + ) + self.set_shared_config(config, link='combine_topo.cfg') + for step in steps: + self.add_step(step) + + +class LatLonCombineTask(Task): + """ + A task for creating combined topography 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 = f'{resolution:.4f}_degree' + subdir = os.path.join( + CombineStep.get_subdir(), + 'lat_lon', + resolution_name, + 'task', + ) + super().__init__( + component=component, + name=f'combine_topo_lat_lon_{resolution_name}_task', + subdir=subdir, + ) + steps, config = get_lat_lon_topo_steps( + component=component, + resolution=resolution, + include_viz=False, ) self.set_shared_config(config, link='combine_topo.cfg') for step in steps: diff --git a/polaris/tasks/e3sm/init/topo/cull/tasks.py b/polaris/tasks/e3sm/init/topo/cull/tasks.py index e021ca61bb..2008e5b4b1 100644 --- a/polaris/tasks/e3sm/init/topo/cull/tasks.py +++ b/polaris/tasks/e3sm/init/topo/cull/tasks.py @@ -1,5 +1,7 @@ from polaris.mesh.base import get_base_mesh_steps -from polaris.tasks.e3sm.init.topo.combine.steps import get_combine_topo_steps +from polaris.tasks.e3sm.init.topo.combine.steps import ( + get_cubed_sphere_topo_steps, +) from polaris.tasks.e3sm.init.topo.cull.task import CullTopoTask from polaris.tasks.e3sm.init.topo.remap.steps import ( get_default_remap_topo_steps, @@ -16,8 +18,9 @@ def add_cull_topo_tasks(component): combine_steps = {} for low_res in [False, True]: - combine_topo_steps, _ = get_combine_topo_steps( - component=component, low_res=low_res + resolution = 120 if low_res else 3000 + combine_topo_steps, _ = get_cubed_sphere_topo_steps( + component=component, resolution=resolution ) combine_steps[low_res] = combine_topo_steps[0] diff --git a/polaris/tasks/e3sm/init/topo/remap/tasks.py b/polaris/tasks/e3sm/init/topo/remap/tasks.py index 913c8b1de9..58fb83b199 100644 --- a/polaris/tasks/e3sm/init/topo/remap/tasks.py +++ b/polaris/tasks/e3sm/init/topo/remap/tasks.py @@ -1,5 +1,7 @@ from polaris.mesh.base import get_base_mesh_steps -from polaris.tasks.e3sm.init.topo.combine.steps import get_combine_topo_steps +from polaris.tasks.e3sm.init.topo.combine.steps import ( + get_cubed_sphere_topo_steps, +) from polaris.tasks.e3sm.init.topo.remap import RemapTopoTask @@ -13,8 +15,9 @@ def add_remap_topo_tasks(component): combine_steps = {} for low_res in [False, True]: - combine_topo_steps, _ = get_combine_topo_steps( - component=component, low_res=low_res + resolution = 120 if low_res else 3000 + combine_topo_steps, _ = get_cubed_sphere_topo_steps( + component=component, resolution=resolution ) combine_steps[low_res] = combine_topo_steps[0] From 1732bce29ba59db9765a544e8978b47e026b7850 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sun, 5 Apr 2026 07:38:28 -0500 Subject: [PATCH 02/13] Fix lat coordinate in tiles in combine topo step --- polaris/tasks/e3sm/init/topo/combine/step.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/polaris/tasks/e3sm/init/topo/combine/step.py b/polaris/tasks/e3sm/init/topo/combine/step.py index c95efb60c9..e3c7ba6b97 100644 --- a/polaris/tasks/e3sm/init/topo/combine/step.py +++ b/polaris/tasks/e3sm/init/topo/combine/step.py @@ -441,10 +441,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) From 018c552f476d002f14bc07dd1adca17c9a54cb49 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sun, 5 Apr 2026 10:44:02 -0500 Subject: [PATCH 03/13] Don't add exodus output file for non-cubed sphere topo --- polaris/tasks/e3sm/init/topo/combine/step.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/polaris/tasks/e3sm/init/topo/combine/step.py b/polaris/tasks/e3sm/init/topo/combine/step.py index e3c7ba6b97..9240c8204f 100644 --- a/polaris/tasks/e3sm/init/topo/combine/step.py +++ b/polaris/tasks/e3sm/init/topo/combine/step.py @@ -256,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 From 7129e4b3869db37608db26087ab94b307c02f4dd Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sun, 5 Apr 2026 13:30:44 -0500 Subject: [PATCH 04/13] Add ocean_mask to combine topo --- polaris/tasks/e3sm/init/topo/combine/step.py | 33 ++++++++++++++------ 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/polaris/tasks/e3sm/init/topo/combine/step.py b/polaris/tasks/e3sm/init/topo/combine/step.py index 9240c8204f..1d783276e6 100644 --- a/polaris/tasks/e3sm/init/topo/combine/step.py +++ b/polaris/tasks/e3sm/init/topo/combine/step.py @@ -276,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) @@ -293,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) @@ -704,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}') @@ -783,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) @@ -797,6 +802,7 @@ def _combine_datasets( 'ice_draft', 'ice_mask', 'grounded_mask', + 'ocean_mask', ] for var in vars: @@ -824,6 +830,12 @@ 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 = { @@ -831,6 +843,7 @@ def _combine_datasets( '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() From effdbb5b4f76868719662ebd43a32153565b0dcd Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 11 Apr 2026 10:48:03 +0200 Subject: [PATCH 05/13] Add more lat-lon resolutions to topo/combine --- polaris/tasks/e3sm/init/add_tasks.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/polaris/tasks/e3sm/init/add_tasks.py b/polaris/tasks/e3sm/init/add_tasks.py index 52dc5e4146..e7ebb5a390 100644 --- a/polaris/tasks/e3sm/init/add_tasks.py +++ b/polaris/tasks/e3sm/init/add_tasks.py @@ -13,11 +13,16 @@ def add_e3sm_init_tasks(component): component : polaris.Component the e3sm/init component that the tasks will be added to """ - for resolution in [3000, 120]: + for cubed_sphere_res in [3000, 120]: component.add_task( - CubedSphereCombineTask(component=component, resolution=resolution) + 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) ) - component.add_task(LatLonCombineTask(component=component, resolution=0.25)) add_remap_topo_tasks(component=component) From 6a625ced063dc98a6bfd0dcb5132313152d507f0 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 11 Apr 2026 11:04:24 +0200 Subject: [PATCH 06/13] Add viz for lat-lon in topo/combine --- .../tasks/e3sm/init/topo/combine/combine.cfg | 36 +++++++++ polaris/tasks/e3sm/init/topo/combine/task.py | 2 +- polaris/tasks/e3sm/init/topo/combine/viz.py | 79 +++++++++++++++---- 3 files changed, 100 insertions(+), 17 deletions(-) diff --git a/polaris/tasks/e3sm/init/topo/combine/combine.cfg b/polaris/tasks/e3sm/init/topo/combine/combine.cfg index b86a38de50..4fe3f28cdb 100644 --- a/polaris/tasks/e3sm/init/topo/combine/combine.cfg +++ b/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 diff --git a/polaris/tasks/e3sm/init/topo/combine/task.py b/polaris/tasks/e3sm/init/topo/combine/task.py index 12f83b6e5c..a83e903837 100644 --- a/polaris/tasks/e3sm/init/topo/combine/task.py +++ b/polaris/tasks/e3sm/init/topo/combine/task.py @@ -84,7 +84,7 @@ def __init__(self, component, resolution): steps, config = get_lat_lon_topo_steps( component=component, resolution=resolution, - include_viz=False, + include_viz=True, ) self.set_shared_config(config, link='combine_topo.cfg') for step in steps: diff --git a/polaris/tasks/e3sm/init/topo/combine/viz.py b/polaris/tasks/e3sm/init/topo/combine/viz.py index 3243fcdd3d..cc1777857b 100644 --- a/polaris/tasks/e3sm/init/topo/combine/viz.py +++ b/polaris/tasks/e3sm/init/topo/combine/viz.py @@ -1,5 +1,6 @@ import os +import cmocean # noqa: F401 import matplotlib.pyplot as plt import numpy as np import pandas as pd @@ -8,6 +9,7 @@ from mpl_toolkits.axes_grid1 import make_axes_locatable from polaris.step import Step +from polaris.viz import plot_global_lat_lon_field class VizCombinedStep(Step): @@ -60,39 +62,82 @@ def setup(self): work_dir_target=os.path.join(combine_step.path, topo_filename), ) - self.add_input_file( - filename='cubed_sphere.g', - work_dir_target=os.path.join(combine_step.path, exodus_filename), - ) + if exodus_filename is not None: + self.add_input_file( + filename='cubed_sphere.g', + work_dir_target=os.path.join( + combine_step.path, exodus_filename + ), + ) def run(self): """ Run this step """ - colormaps = { - 'base_elevation': 'cmo.deep_r', - 'ice_thickness': 'cmo.ice_r', - 'ice_draft': 'cmo.deep_r', - 'ice_mask': 'cmo.amp_r', - 'grounded_mask': 'cmo.amp_r', + colormap_sections = { + 'base_elevation': 'viz_combine_topo_base_elevation', + 'ice_thickness': 'viz_combine_topo_ice_thickness', + 'ice_draft': 'viz_combine_topo_ice_draft', + 'ice_mask': 'viz_combine_topo_ice_mask', + 'grounded_mask': 'viz_combine_topo_grounded_mask', } ds_data = xr.open_dataset('topography.nc') + target_grid = self.config.get('combine_topo', 'target_grid') - # Use one field to define the valid mask (they all share indexing) - valid_mask = np.isfinite(ds_data['base_elevation'].values) + if target_grid == 'cubed_sphere': + self._plot_cubed_sphere_fields(ds_data) + elif target_grid == 'lat_lon': + self._plot_lat_lon_fields(ds_data, colormap_sections) + else: + raise ValueError(f'Unexpected target grid: {target_grid}') - # Build mesh only once + def _plot_cubed_sphere_fields(self, ds_data): + """ + Plot fields on the cubed-sphere target grid. + """ + valid_mask = np.isfinite(ds_data['base_elevation'].values) vertices, tris = self._load_trimesh_geometry( 'cubed_sphere.g', valid_mask ) - # Plot each field + colormaps = { + 'base_elevation': 'cmo.deep_r', + 'ice_thickness': 'cmo.ice_r', + 'ice_draft': 'cmo.deep_r', + 'ice_mask': 'cmo.amp_r', + 'grounded_mask': 'cmo.amp_r', + } + for field, colormap in colormaps.items(): self.logger.info(f'Plotting field: {field}') data = ds_data[field].values[valid_mask] - self._plot_field(vertices, tris, data, field, colormap) + self._plot_cubed_sphere_field( + vertices, tris, data, field, colormap + ) + + def _plot_lat_lon_fields(self, ds_data, colormap_sections): + """ + Plot fields on the latitude-longitude target grid. + """ + lon = ds_data.lon.values + lat = ds_data.lat.values + + for field, section in colormap_sections.items(): + self.logger.info(f'Plotting field: {field}') + data_array = ds_data[field].transpose('lat', 'lon').values + plot_global_lat_lon_field( + lon=lon, + lat=lat, + data_array=data_array, + out_filename=f'{field}.png', + config=self.config, + colormap_section=section, + title=f'{self.combine_step.resolution_name} {field}', + plot_land=False, + colorbar_label=field, + ) @staticmethod def _load_trimesh_geometry(exodus_path, valid_mask): @@ -125,7 +170,9 @@ def _load_trimesh_geometry(exodus_path, valid_mask): return vertices, tris - def _plot_field(self, vertices, tris, field_data, field_name, cmap): + def _plot_cubed_sphere_field( + self, vertices, tris, field_data, field_name, cmap + ): """ Rasterize and save a trisurf-style field image using Datashader. """ From f1f705c57f7b657aa98e6503b6833561a74b08c1 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 26 Mar 2026 21:04:12 +0100 Subject: [PATCH 07/13] Add global_ocean WOA23 hydrography task Introduce the ocean/global_ocean/hydrography/woa23 task with shared combine and extrapolate steps, draft user and developer documentation, and API entries. The task consumes the reusable e3sm/init lat-lon combined-topography product so future hydrography products can share the same cached input workflow. --- docs/developers_guide/ocean/api.md | 33 +- .../ocean/tasks/global_ocean.md | 69 ++++ docs/developers_guide/ocean/tasks/index.md | 1 + docs/users_guide/ocean/tasks/global_ocean.md | 83 +++++ docs/users_guide/ocean/tasks/index.md | 1 + polaris/tasks/ocean/add_tasks.py | 2 + polaris/tasks/ocean/global_ocean/__init__.py | 15 + .../global_ocean/hydrography/__init__.py | 3 + .../hydrography/woa23/__init__.py | 5 + .../global_ocean/hydrography/woa23/combine.py | 170 +++++++++ .../hydrography/woa23/extrapolate.py | 329 ++++++++++++++++++ .../global_ocean/hydrography/woa23/steps.py | 78 +++++ .../global_ocean/hydrography/woa23/task.py | 50 +++ .../global_ocean/hydrography/woa23/woa23.cfg | 6 + 14 files changed, 844 insertions(+), 1 deletion(-) create mode 100644 docs/developers_guide/ocean/tasks/global_ocean.md create mode 100644 docs/users_guide/ocean/tasks/global_ocean.md create mode 100644 polaris/tasks/ocean/global_ocean/__init__.py create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/__init__.py create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index 4c1fe455d0..560018e984 100644 --- a/docs/developers_guide/ocean/api.md +++ b/docs/developers_guide/ocean/api.md @@ -198,6 +198,38 @@ viz.Viz.run ``` +### global_ocean + +```{eval-rst} +.. currentmodule:: polaris.tasks.ocean.global_ocean + +.. autosummary:: + :toctree: generated/ + + add_global_ocean_tasks +``` + +### global_ocean.hydrography.woa23 + +```{eval-rst} +.. currentmodule:: polaris.tasks.ocean.global_ocean.hydrography.woa23 + +.. autosummary:: + :toctree: generated/ + + Woa23 + get_woa23_topography_step + get_woa23_steps + + CombineStep + CombineStep.setup + CombineStep.run + + ExtrapolateStep + ExtrapolateStep.setup + ExtrapolateStep.run +``` + ### ice_shelf_2d ```{eval-rst} @@ -617,4 +649,3 @@ vertical.ztilde.pressure_and_spec_vol_from_state_at_geom_height vertical.ztilde.geom_height_from_pseudo_height ``` - diff --git a/docs/developers_guide/ocean/tasks/global_ocean.md b/docs/developers_guide/ocean/tasks/global_ocean.md new file mode 100644 index 0000000000..998564fc85 --- /dev/null +++ b/docs/developers_guide/ocean/tasks/global_ocean.md @@ -0,0 +1,69 @@ +(dev-ocean-global-ocean)= + +# global_ocean + +The `global_ocean` tasks in `polaris.tasks.ocean.global_ocean` are intended +for preprocessing and initialization workflows that are upstream of any +particular MPAS mesh. The first task category under this framework is +`hydrography/woa23`, which builds a reusable hydrography product from the +World Ocean Atlas 2023 on its native 0.25-degree latitude-longitude grid. + +(dev-ocean-global-ocean-framework)= + +## framework + +The shared config options for the WOA23 hydrography task are described in +{ref}`ocean-global-ocean` in the User's Guide. + +The implementation is intentionally organized around reusable Polaris steps +rather than around the legacy Compass `utility/extrap_woa` multiprocessing +workflow. One notable design choice is that the task reuses the combined +topography product from `e3sm/init` rather than taking a raw topography +filename as a task-specific input. + +### cached topography dependency + +The helper +{py:func}`polaris.tasks.ocean.global_ocean.hydrography.woa23.get_woa23_topography_step` +creates a shared `e3sm/init` {py:class}`polaris.tasks.e3sm.init.topo.combine.step.CombineStep` +configured for a 0.25-degree lat-lon target grid. The +{py:class}`polaris.tasks.ocean.global_ocean.hydrography.woa23.task.Woa23` +task adds this step with a symlink `combine_topo` and prefers to use a cached +version of its outputs when matching entries are available in the +`e3sm/init` cache database. + +This keeps the expensive topography blending logic in one place and makes the +ocean hydrography preprocessing task consistent with the broader Polaris +approach to shared, cacheable preprocessing steps. + +(dev-ocean-global-ocean-woa23)= + +## hydrography/woa23 + +The {py:class}`polaris.tasks.ocean.global_ocean.hydrography.woa23.task.Woa23` +task is the Polaris port of the WOA preprocessing part of the legacy Compass +workflow. + +### combine + +The class +{py:class}`polaris.tasks.ocean.global_ocean.hydrography.woa23.combine.CombineStep` +combines January and annual WOA23 temperature and salinity climatologies into +a single dataset. January values are used where they exist, and annual values +fill deeper levels where the monthly product is not available. + +This step also converts in-situ temperature to potential temperature with +`gsw`, producing `woa_combined.nc`. + +### extrapolate + +The class +{py:class}`polaris.tasks.ocean.global_ocean.hydrography.woa23.extrapolate.ExtrapolateStep` +uses the cached combined-topography product on the WOA grid together with +`woa_combined.nc` to build a 3D ocean mask and then fill missing WOA values in +two stages: + +1. Horizontal then vertical extrapolation within the ocean mask +2. Horizontal then vertical extrapolation into land and grounded-ice regions + +The final output is `woa23_decav_0.25_jan_extrap.nc`. diff --git a/docs/developers_guide/ocean/tasks/index.md b/docs/developers_guide/ocean/tasks/index.md index fc689f1c67..851685ff50 100644 --- a/docs/developers_guide/ocean/tasks/index.md +++ b/docs/developers_guide/ocean/tasks/index.md @@ -13,6 +13,7 @@ cosine_bell customizable_viz external_gravity_wave geostrophic +global_ocean divergent_2d ice_shelf_2d inertial_gravity_wave diff --git a/docs/users_guide/ocean/tasks/global_ocean.md b/docs/users_guide/ocean/tasks/global_ocean.md new file mode 100644 index 0000000000..0b2d8d8175 --- /dev/null +++ b/docs/users_guide/ocean/tasks/global_ocean.md @@ -0,0 +1,83 @@ +(ocean-global-ocean)= + +# global_ocean + +This category contains ocean preprocessing tasks that are upstream of any +particular MPAS mesh. The first task builds a reusable World Ocean Atlas 2023 +(WOA23) hydrography product on the native 0.25-degree latitude-longitude grid. + +## supported models + +This task is model-independent and does not require either MPAS-Ocean or +Omega to be built. + +(ocean-global-ocean-woa23)= + +## woa23 + +This task is the Polaris port of the legacy Compass +`utility/extrap_woa` workflow. It combines January and annual WOA23 +climatologies, uses a cached `e3sm/init` combined-topography product on the +WOA grid to define the ocean mask used during preprocessing, and then fills +missing temperature and salinity values through staged horizontal and vertical +extrapolation. + +The task can be set up with: + +```bash +polaris setup -t ocean/global_ocean/hydrography/woa23 ... +``` + +### description + +The task is organized into three inspectable steps: + +1. `combine_topo` reuses a cached `e3sm/init` combined-topography step + configured for the WOA23 0.25-degree latitude-longitude grid. +2. `combine` creates `woa_combined.nc` by combining January and annual WOA23 + fields and converting temperature to potential temperature. +3. `extrapolate` creates the final + `woa23_decav_0.25_jan_extrap.nc` product. + +This layout is intended to match Polaris shared-step conventions so the WOA23 +preprocessing pipeline can later be reused by mesh-dependent +`global_ocean` initialization tasks. + +### mesh + +N/A. This task operates on the native WOA23 latitude-longitude grid rather +than an MPAS mesh. + +### vertical grid + +N/A. The task preserves the standard WOA23 depth levels. + +### initial conditions + +The source fields come from the WOA23 January and annual climatologies in the +Polaris input database. + +### forcing + +N/A. + +### time step and run duration + +N/A. + +### config options + +```cfg +# Options related to generating a reusable WOA23 hydrography product +[woa23] + +# the minimum weight sum needed to mark a new cell valid in horizontal +# extrapolation +extrap_threshold = 0.01 +``` + +### cores + +The local `combine` and `extrapolate` steps run serially. The +`combine_topo` step is intended to use the cached `e3sm/init` output because +regenerating the combined topography product is substantially more expensive. diff --git a/docs/users_guide/ocean/tasks/index.md b/docs/users_guide/ocean/tasks/index.md index 8d6befc217..f2f70af621 100644 --- a/docs/users_guide/ocean/tasks/index.md +++ b/docs/users_guide/ocean/tasks/index.md @@ -13,6 +13,7 @@ cosine_bell customizable_viz external_gravity_wave geostrophic +global_ocean divergent_2d ice_shelf_2d inertial_gravity_wave diff --git a/polaris/tasks/ocean/add_tasks.py b/polaris/tasks/ocean/add_tasks.py index 26d35beaa5..8b5d9d67e4 100644 --- a/polaris/tasks/ocean/add_tasks.py +++ b/polaris/tasks/ocean/add_tasks.py @@ -7,6 +7,7 @@ add_external_gravity_wave_tasks as add_external_gravity_wave_tasks, ) from polaris.tasks.ocean.geostrophic import add_geostrophic_tasks +from polaris.tasks.ocean.global_ocean import add_global_ocean_tasks from polaris.tasks.ocean.ice_shelf_2d import add_ice_shelf_2d_tasks from polaris.tasks.ocean.inertial_gravity_wave import ( add_inertial_gravity_wave_tasks as add_inertial_gravity_wave_tasks, @@ -53,5 +54,6 @@ def add_ocean_tasks(component): add_cosine_bell_tasks(component=component) add_external_gravity_wave_tasks(component=component) add_geostrophic_tasks(component=component) + add_global_ocean_tasks(component=component) add_isomip_plus_tasks(component=component, mesh_type='spherical') add_sphere_transport_tasks(component=component) diff --git a/polaris/tasks/ocean/global_ocean/__init__.py b/polaris/tasks/ocean/global_ocean/__init__.py new file mode 100644 index 0000000000..16c2caac12 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/__init__.py @@ -0,0 +1,15 @@ +from polaris.tasks.ocean.global_ocean.hydrography.woa23 import ( + Woa23 as Woa23, +) + + +def add_global_ocean_tasks(component): + """ + Add tasks for global-ocean preprocessing and initialization. + + Parameters + ---------- + component : polaris.tasks.ocean.Ocean + The ocean component to which the tasks will be added. + """ + component.add_task(Woa23(component=component)) diff --git a/polaris/tasks/ocean/global_ocean/hydrography/__init__.py b/polaris/tasks/ocean/global_ocean/hydrography/__init__.py new file mode 100644 index 0000000000..23238e0f13 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/__init__.py @@ -0,0 +1,3 @@ +from polaris.tasks.ocean.global_ocean.hydrography.woa23 import ( + Woa23 as Woa23, +) diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py new file mode 100644 index 0000000000..da90111410 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py @@ -0,0 +1,5 @@ +from .combine import CombineStep as CombineStep +from .extrapolate import ExtrapolateStep as ExtrapolateStep +from .steps import get_woa23_steps as get_woa23_steps +from .steps import get_woa23_topography_step as get_woa23_topography_step +from .task import Woa23 as Woa23 diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py new file mode 100644 index 0000000000..b4ec69f1c8 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py @@ -0,0 +1,170 @@ +import gsw +import numpy as np +import xarray as xr +from mpas_tools.io import write_netcdf + +from polaris import Step + + +class CombineStep(Step): + """ + A step for combining January and annual WOA23 climatologies. + """ + + def __init__(self, component, subdir): + """ + Create the step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to. + + subdir : str + The subdirectory for the step. + """ + super().__init__( + component=component, + name='combine', + subdir=subdir, + ntasks=1, + min_tasks=1, + ) + self.add_output_file(filename='woa_combined.nc') + + def setup(self): + """ + Set up input files for the step. + """ + super().setup() + + base_url = ( + 'https://www.ncei.noaa.gov/thredds-ocean/fileServer/woa23/DATA' + ) + directories = { + 'temp': { + 'ann': 'temperature/netcdf/decav91C0/0.25', + 'jan': 'temperature/netcdf/decav91C0/0.25', + }, + 'salin': { + 'ann': 'salinity/netcdf/decav91C0/0.25', + 'jan': 'salinity/netcdf/decav91C0/0.25', + }, + } + filenames = { + 'temp': { + 'ann': 'woa23_decav91C0_t00_04.nc', + 'jan': 'woa23_decav91C0_t01_04.nc', + }, + 'salin': { + 'ann': 'woa23_decav91C0_s00_04.nc', + 'jan': 'woa23_decav91C0_s01_04.nc', + }, + } + + for field in ['temp', 'salin']: + for season in ['jan', 'ann']: + woa_filename = filenames[field][season] + woa_dir = directories[field][season] + self.add_input_file( + filename=f'woa_{field}_{season}.nc', + target=woa_filename, + database='initial_condition_database', + url=f'{base_url}/{woa_dir}/{woa_filename}', + ) + + def run(self): + """ + Combine January and annual climatologies and convert temperature to + potential temperature. + """ + logger = self.logger + logger.info('Combining January and annual WOA23 climatologies') + + with xr.open_dataset('woa_temp_ann.nc', decode_times=False) as ds_temp: + ds_out = xr.Dataset() + for var in ['lon', 'lat', 'depth']: + ds_out[var] = ds_temp[var] + ds_out[f'{var}_bnds'] = ds_temp[f'{var}_bnds'] + + var_map = {'temp': 't_an', 'salin': 's_an'} + for field, var_name in var_map.items(): + with xr.open_dataset( + f'woa_{field}_ann.nc', decode_times=False + ) as ds_ann: + ds_ann = ds_ann.isel(time=0, drop=True) + with xr.open_dataset( + f'woa_{field}_jan.nc', decode_times=False + ) as ds_jan: + ds_jan = ds_jan.isel(time=0, drop=True) + slices = [] + for depth_index in range(ds_ann.sizes['depth']): + if depth_index < ds_jan.sizes['depth']: + ds = ds_jan + else: + ds = ds_ann + slices.append(ds[var_name].isel(depth=depth_index)) + + ds_out[var_name] = xr.concat(slices, dim='depth') + ds_out[var_name].attrs = ds_ann[var_name].attrs + + ds_out = self._temp_to_potential_temp(ds_out) + write_netcdf(ds_out, 'woa_combined.nc') + logger.info('Wrote woa_combined.nc') + + @staticmethod + def _temp_to_potential_temp(ds): + """ + Convert WOA in-situ temperature to potential temperature. + + Parameters + ---------- + ds : xarray.Dataset + A combined WOA dataset with in-situ temperature and salinity. + + Returns + ------- + ds : xarray.Dataset + The dataset with temperature replaced by potential temperature. + """ + dims = ds.t_an.dims + slices = [] + for depth_index in range(ds.sizes['depth']): + temp_slice = ds.t_an.isel(depth=depth_index) + in_situ_temp = temp_slice.values + practical_salinity = ds.s_an.isel(depth=depth_index).values + lat = ds.lat.broadcast_like(temp_slice).values + lon = ds.lon.broadcast_like(temp_slice).values + z = -ds.depth.isel(depth=depth_index).values + pressure = gsw.p_from_z(z, lat) + + mask = np.isfinite(in_situ_temp) + potential_temp = np.full(in_situ_temp.shape, np.nan) + absolute_salinity = gsw.SA_from_SP( + practical_salinity[mask], + pressure[mask], + lon[mask], + lat[mask], + ) + potential_temp[mask] = gsw.pt_from_t( + absolute_salinity, + in_situ_temp[mask], + pressure[mask], + p_ref=0.0, + ) + slices.append( + xr.DataArray( + data=potential_temp, + dims=temp_slice.dims, + attrs=temp_slice.attrs, + ) + ) + + ds['pt_an'] = xr.concat(slices, dim='depth').transpose(*dims) + ds.pt_an.attrs['standard_name'] = 'sea_water_potential_temperature' + ds.pt_an.attrs['long_name'] = ( + 'Objectively analyzed mean fields for ' + 'sea_water_potential_temperature at standard depth levels.' + ) + + return ds.drop_vars('t_an') diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py new file mode 100644 index 0000000000..8119bff071 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py @@ -0,0 +1,329 @@ +import numpy as np +import xarray as xr +from mpas_tools.io import write_netcdf +from scipy.signal import convolve2d + +from polaris import Step + + +class ExtrapolateStep(Step): + """ + A step for extrapolating WOA23 into missing ocean, land and ice regions. + """ + + output_filename = 'woa23_decav_0.25_jan_extrap.nc' + + def __init__(self, component, subdir, combine_step, combine_topo_step): + """ + Create the step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to. + + subdir : str + The subdirectory for the step. + + combine_step : polaris.Step + The step that produces the combined WOA23 dataset. + + combine_topo_step : polaris.Step + The cached ``e3sm/init`` step that produces combined topography on + the WOA23 grid. + """ + super().__init__( + component=component, + name='extrapolate', + subdir=subdir, + ntasks=1, + min_tasks=1, + ) + self.combine_step = combine_step + self.combine_topo_step = combine_topo_step + self.add_output_file(filename=self.output_filename) + + def setup(self): + """ + Set up input files for the step. + """ + super().setup() + self.add_input_file( + filename='woa.nc', + work_dir_target=f'{self.combine_step.path}/woa_combined.nc', + ) + self.add_input_file( + filename='topography.nc', + work_dir_target=( + f'{self.combine_topo_step.path}/' + f'{self.combine_topo_step.combined_filename}' + ), + ) + + def run(self): + """ + Extrapolate WOA23 horizontally and vertically in two stages. + """ + logger = self.logger + logger.info('Building a 3D ocean mask on the WOA23 grid') + self._make_3d_ocean_mask() + + logger.info('Horizontally extrapolating within the ocean mask') + self._extrap_horiz( + in_filename='woa.nc', + out_filename='woa_extrap_ocean_horiz.nc', + use_ocean_mask=True, + ) + + logger.info('Vertically extrapolating within the ocean mask') + self._extrap_vert( + in_filename='woa_extrap_ocean_horiz.nc', + out_filename='woa_extrap_ocean.nc', + use_ocean_mask=True, + ) + + logger.info('Horizontally extrapolating into land and grounded ice') + self._extrap_horiz( + in_filename='woa_extrap_ocean.nc', + out_filename='woa_extrap_horiz.nc', + use_ocean_mask=False, + ) + + logger.info('Vertically extrapolating into land and grounded ice') + self._extrap_vert( + in_filename='woa_extrap_horiz.nc', + out_filename=self.output_filename, + use_ocean_mask=False, + ) + logger.info(f'Wrote {self.output_filename}') + + @staticmethod + def _make_3d_ocean_mask(): + """ + Build a three-dimensional mask of valid ocean cells on the WOA grid. + """ + with xr.open_dataset('topography.nc') as ds_topo: + with xr.open_dataset('woa.nc', decode_times=False) as ds_woa: + ds_out = xr.Dataset() + for var in ['lon', 'lat', 'depth']: + ds_out[var] = ds_woa[var] + ds_out[f'{var}_bnds'] = ds_woa[f'{var}_bnds'] + + z_top = -ds_woa.depth_bnds.isel(nbounds=0) + ocean_mask_3d = ( + (ds_topo.base_elevation <= z_top) + & (ds_topo.ocean_mask >= 0.5) + ).transpose('depth', 'lat', 'lon') + ds_out['ocean_mask'] = ocean_mask_3d.astype(np.int8) + + write_netcdf(ds_out, 'ocean_mask.nc') + + def _extrap_horiz(self, in_filename, out_filename, use_ocean_mask): + """ + Extrapolate horizontally on each depth level. + + Parameters + ---------- + in_filename : str + The input file to read. + + out_filename : str + The output file to write. + + use_ocean_mask : bool + Whether to restrict filling to the remapped ocean mask. + """ + with xr.open_dataset(in_filename, decode_times=False) as ds: + ds_out = ds.load() + + ocean_mask = None + if use_ocean_mask: + with xr.open_dataset('ocean_mask.nc') as ds_mask: + ocean_mask = ds_mask.ocean_mask.values.astype(bool) + + ndepth = ds_out.sizes['depth'] + for depth_index in range(ndepth): + self.logger.info( + f' Horizontal fill for depth {depth_index + 1}/{ndepth}' + ) + level_mask = None + if ocean_mask is not None: + level_mask = ocean_mask[depth_index, :, :] + ds_level = ds_out.isel(depth=depth_index) + ds_level = self._extrap_level( + ds_level=ds_level, + ocean_mask=level_mask, + threshold=self.config.getfloat('woa23', 'extrap_threshold'), + ) + for field_name in ['pt_an', 's_an']: + ds_out[field_name][depth_index, :, :] = ds_level[field_name] + + write_netcdf(ds_out, out_filename) + + def _extrap_vert(self, in_filename, out_filename, use_ocean_mask): + """ + Extrapolate vertically from shallower depths into deeper missing + values. + + Parameters + ---------- + in_filename : str + The input file to read. + + out_filename : str + The output file to write. + + use_ocean_mask : bool + Whether to restrict filling to the remapped ocean mask. + """ + with xr.open_dataset(in_filename, decode_times=False) as ds: + ds_out = ds.load() + + ocean_mask = None + if use_ocean_mask: + with xr.open_dataset('ocean_mask.nc') as ds_mask: + ocean_mask = ds_mask.ocean_mask.values.astype(bool) + + ndepth = ds_out.sizes['depth'] + for field_name in ['pt_an', 's_an']: + field = ds_out[field_name].values + for depth_index in range(1, ndepth): + mask = np.isnan(field[depth_index, :, :]) + if ocean_mask is not None: + mask = np.logical_and(mask, ocean_mask[depth_index, :, :]) + field[depth_index, :, :][mask] = field[depth_index - 1, :, :][ + mask + ] + + write_netcdf(ds_out, out_filename) + + @staticmethod + def _extrap_level(ds_level, ocean_mask, threshold): + """ + Extrapolate a single depth level horizontally. + + Parameters + ---------- + ds_level : xarray.Dataset + A single-depth WOA dataset. + + ocean_mask : numpy.ndarray or None + A Boolean mask for valid ocean cells at this depth. + + threshold : float + The minimum valid weight sum needed to mark a new cell valid. + + Returns + ------- + ds_level : xarray.Dataset + The filled depth level. + """ + kernel = ExtrapolateStep._get_kernel() + field = ds_level.pt_an.values + + valid = np.isfinite(field) + orig_mask = valid.copy() + if ocean_mask is not None: + invalid_after_fill = np.logical_not( + np.logical_or(valid, ocean_mask) + ) + else: + invalid_after_fill = None + + fields = { + field_name: ds_level[field_name].values.copy() + for field_name in ['pt_an', 's_an'] + } + + nlon = field.shape[1] + lon_with_halo = np.array( + [nlon - 2, nlon - 1] + list(range(nlon)) + [0, 1] + ) + lon_no_halo = list(range(2, nlon + 2)) + + prev_fill_count = 0 + while True: + valid_weight_sum = _extrap_with_halo( + field=valid.astype(float), + kernel=kernel, + valid=valid, + lon_with_halo=lon_with_halo, + lon_no_halo=lon_no_halo, + ) + if invalid_after_fill is not None: + valid_weight_sum[invalid_after_fill] = 0.0 + + new_valid = valid_weight_sum > threshold + fill_mask = np.logical_and(new_valid, np.logical_not(orig_mask)) + fill_count = int(np.count_nonzero(fill_mask)) + if fill_count == prev_fill_count: + break + + for _field_name, field_values in fields.items(): + field_extrap = _extrap_with_halo( + field=field_values, + kernel=kernel, + valid=valid, + lon_with_halo=lon_with_halo, + lon_no_halo=lon_no_halo, + ) + field_values[fill_mask] = ( + field_extrap[fill_mask] / valid_weight_sum[fill_mask] + ) + + valid = new_valid + prev_fill_count = fill_count + + for field_name, field_values in fields.items(): + if invalid_after_fill is not None: + field_values[invalid_after_fill] = np.nan + ds_level[field_name] = (ds_level[field_name].dims, field_values) + + return ds_level + + @staticmethod + def _get_kernel(): + """ + Build the small averaging kernel used for horizontal extrapolation. + + Returns + ------- + kernel : numpy.ndarray + A two-dimensional Gaussian-like averaging kernel. + """ + coordinates = np.arange(-1, 2) + x, y = np.meshgrid(coordinates, coordinates) + return np.exp(-0.5 * (x**2 + y**2)) + + +def _extrap_with_halo(field, kernel, valid, lon_with_halo, lon_no_halo): + """ + Extrapolate a two-dimensional field using a periodic halo in longitude. + + Parameters + ---------- + field : numpy.ndarray + The field to extrapolate. + + kernel : numpy.ndarray + The horizontal averaging kernel. + + valid : numpy.ndarray + A Boolean mask of valid values. + + lon_with_halo : numpy.ndarray + Longitude indices including the periodic halo. + + lon_no_halo : list of int + Longitude indices excluding the periodic halo. + + Returns + ------- + field_extrap : numpy.ndarray + The convolved field without the periodic halo. + """ + masked_field = field.copy() + masked_field[np.logical_not(valid)] = 0.0 + field_with_halo = masked_field[:, lon_with_halo] + field_extrap = convolve2d(field_with_halo, kernel, mode='same') + return field_extrap[:, lon_no_halo] diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py new file mode 100644 index 0000000000..1b421e3b26 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py @@ -0,0 +1,78 @@ +import os + +from polaris.config import PolarisConfigParser +from polaris.tasks.e3sm.init import e3sm_init +from polaris.tasks.e3sm.init.topo.combine import ( + get_lat_lon_topo_steps, +) + +from .combine import CombineStep +from .extrapolate import ExtrapolateStep + + +def get_woa23_topography_step(): + """ + Get the cached combined-topography step for the WOA23 target grid. + + Returns + ------- + combine_topo_step : polaris.tasks.e3sm.init.topo.combine.step.CombineStep + A shared step from ``e3sm/init`` configured for the WOA23 0.25-degree + latitude-longitude grid. + """ + steps, _ = get_lat_lon_topo_steps( + component=e3sm_init, resolution=0.25, include_viz=False + ) + return steps[0] + + +def get_woa23_steps(component, combine_topo_step): + """ + Get the shared steps for building the reusable WOA23 hydrography product. + + Parameters + ---------- + component : polaris.tasks.ocean.Ocean + The ocean component the steps belong to. + + combine_topo_step : polaris.tasks.e3sm.init.topo.combine.step.CombineStep + The cached ``e3sm/init`` step that produces topography on the WOA23 + grid. + + Returns + ------- + steps : list of polaris.Step + Shared steps for combining and extrapolating WOA23 data. + + config : polaris.config.PolarisConfigParser + The shared config options for the task and its steps. + """ + subdir = 'global_ocean/hydrography/woa23' + config_filename = 'woa23.cfg' + config = PolarisConfigParser( + filepath=os.path.join(component.name, subdir, config_filename) + ) + config.add_from_package( + 'polaris.tasks.ocean.global_ocean.hydrography.woa23', + config_filename, + ) + + combine_subdir = os.path.join(subdir, 'combine') + combine_step = component.get_or_create_shared_step( + step_cls=CombineStep, + subdir=combine_subdir, + config=config, + config_filename=config_filename, + ) + + extrapolate_subdir = os.path.join(subdir, 'extrapolate') + extrapolate_step = component.get_or_create_shared_step( + step_cls=ExtrapolateStep, + subdir=extrapolate_subdir, + config=config, + config_filename=config_filename, + combine_step=combine_step, + combine_topo_step=combine_topo_step, + ) + + return [combine_step, extrapolate_step], config diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py new file mode 100644 index 0000000000..2d6b67a526 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py @@ -0,0 +1,50 @@ +import os + +from polaris import Task +from polaris.tasks.ocean.global_ocean.hydrography.woa23.steps import ( + get_woa23_steps, + get_woa23_topography_step, +) + + +class Woa23(Task): + """ + A task for building a reusable WOA23 hydrography product. + """ + + def __init__(self, component): + """ + Create the task. + + Parameters + ---------- + component : polaris.tasks.ocean.Ocean + The ocean component the task belongs to. + """ + subdir = 'global_ocean/hydrography/woa23' + super().__init__(component=component, name='woa23', subdir=subdir) + + self.combine_topo_step = get_woa23_topography_step() + steps, config = get_woa23_steps( + component=component, + combine_topo_step=self.combine_topo_step, + ) + self.set_shared_config(config) + + self.add_step(self.combine_topo_step, symlink='combine_topo') + for step in steps: + self.add_step(step) + + def configure(self): + """ + Use the cached combined-topography product from ``e3sm/init``. + """ + super().configure() + cache_keys = [ + os.path.join(self.combine_topo_step.path, output) + for output in self.combine_topo_step.outputs + ] + cached_files = self.combine_topo_step.component.cached_files + self.combine_topo_step.cached = all( + filename in cached_files for filename in cache_keys + ) diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg new file mode 100644 index 0000000000..48ebcbeadb --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg @@ -0,0 +1,6 @@ +# Options related to generating a reusable WOA23 hydrography product +[woa23] + +# the minimum weight sum needed to mark a new cell valid in horizontal +# extrapolation +extrap_threshold = 0.01 From 80c923ab6dddf7cbd7546b7f5372a77391e83bce Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sun, 5 Apr 2026 08:37:38 -0500 Subject: [PATCH 08/13] Set parallel system for all components This is necessary if shared steps belong to a different component than the one for the task they belong to, but they use a parallel command. This has to happen at both setup and runtime --- polaris/parallel.py | 86 +++++++++++++++++++++++++++++++++++++++++++ polaris/run/serial.py | 7 ++-- polaris/setup.py | 3 +- 3 files changed, 91 insertions(+), 5 deletions(-) create mode 100644 polaris/parallel.py diff --git a/polaris/parallel.py b/polaris/parallel.py new file mode 100644 index 0000000000..159b29e329 --- /dev/null +++ b/polaris/parallel.py @@ -0,0 +1,86 @@ +from polaris.config import PolarisConfigParser + + +def set_parallel_systems(tasks, config: PolarisConfigParser): + """ + Set the active parallel system on every component referenced by the task + and step graph. + + Parameters + ---------- + tasks : dict of polaris.Task + Tasks to scan for referenced components + + config : polaris.config.PolarisConfigParser + The config to use in constructing the parallel systems + """ + seen_components: set[int] = set() + seen_steps: set[int] = set() + + for task in tasks.values(): + _set_parallel_system_for_component( + task.component, config, seen_components + ) + for step in task.steps.values(): + _set_parallel_systems_for_step( + step, config, seen_components, seen_steps + ) + + +def _set_parallel_systems_for_step( + step, config: PolarisConfigParser, seen_components, seen_steps +): + """ + Set the active parallel system on a step's component and recursively on + the components of any step dependencies. + + Parameters + ---------- + step : polaris.Step + The step to scan + + config : polaris.config.PolarisConfigParser + The config to use in constructing the parallel systems + + seen_components : set of int + The ids of components that have already been initialized + + seen_steps : set of int + The ids of steps that have already been visited + """ + step_id = id(step) + if step_id in seen_steps: + return + seen_steps.add(step_id) + + _set_parallel_system_for_component(step.component, config, seen_components) + + for dependency in step.dependencies.values(): + _set_parallel_systems_for_step( + dependency, config, seen_components, seen_steps + ) + + +def _set_parallel_system_for_component( + component, config: PolarisConfigParser, seen_components +): + """ + Set the active parallel system for a component once. + + Parameters + ---------- + component : polaris.Component + The component to initialize + + config : polaris.config.PolarisConfigParser + The config to use in constructing the parallel system + + seen_components : set of int + The ids of components that have already been initialized + """ + component_id = id(component) + if component_id in seen_components: + return + + seen_components.add(component_id) + component.set_parallel_system(config) diff --git a/polaris/run/serial.py b/polaris/run/serial.py index 8ee30dd72b..f410097406 100644 --- a/polaris/run/serial.py +++ b/polaris/run/serial.py @@ -12,6 +12,7 @@ from polaris import Task from polaris.logging import log_function_call, log_method_call +from polaris.parallel import set_parallel_systems from polaris.run import ( complete_step_run, load_dependencies, @@ -68,7 +69,7 @@ def run_tasks( task = next(iter(suite['tasks'].values())) component = task.component common_config = setup_config(task.base_work_dir, f'{component.name}.cfg') - component.set_parallel_system(common_config) + set_parallel_systems(suite['tasks'], common_config) available_resources = component.get_available_resources() # start logging to stdout/stderr @@ -174,7 +175,7 @@ def run_single_step(step_is_subprocess=False, quiet=False): config = setup_config(step.base_work_dir, step.config.filepath) task.config = config - step.component.set_parallel_system(config) + set_parallel_systems({task.path: task}, config) available_resources = step.component.get_available_resources() mpas_tools.io.default_format = config.get('io', 'format') @@ -373,8 +374,6 @@ def _log_and_run_task( config = setup_config(task.base_work_dir, task.config.filepath) task.config = config - task.component.set_parallel_system(config) - available_resources = task.component.get_available_resources() mpas_tools.io.default_format = config.get('io', 'format') mpas_tools.io.default_engine = config.get('io', 'engine') diff --git a/polaris/setup.py b/polaris/setup.py index b67a510829..5efa6a22c4 100644 --- a/polaris/setup.py +++ b/polaris/setup.py @@ -14,6 +14,7 @@ from polaris.io import symlink from polaris.job import write_job_script from polaris.machines import discover_machine +from polaris.parallel import set_parallel_systems from polaris.tasks import get_components @@ -157,7 +158,7 @@ def setup_tasks( ) component.configure(basic_config, list(tasks.values())) - component.set_parallel_system(basic_config) + set_parallel_systems(tasks, basic_config) provenance.write( work_dir, From 0c917c1379ebac08877ebffc1396b14c384553c7 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sun, 5 Apr 2026 13:30:44 -0500 Subject: [PATCH 09/13] Use ocean_mask from combine topo --- .../global_ocean/hydrography/woa23/extrapolate.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py index 8119bff071..9867e47d31 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py @@ -110,14 +110,21 @@ def _make_3d_ocean_mask(): ds_out[f'{var}_bnds'] = ds_woa[f'{var}_bnds'] z_top = -ds_woa.depth_bnds.isel(nbounds=0) + ocean_mask = ExtrapolateStep._get_ocean_mask(ds_topo) ocean_mask_3d = ( - (ds_topo.base_elevation <= z_top) - & (ds_topo.ocean_mask >= 0.5) + (ds_topo.base_elevation <= z_top) & (ocean_mask >= 0.5) ).transpose('depth', 'lat', 'lon') ds_out['ocean_mask'] = ocean_mask_3d.astype(np.int8) write_netcdf(ds_out, 'ocean_mask.nc') + @staticmethod + def _get_ocean_mask(ds_topo): + """ + Return the 2D ocean mask from combined topography. + """ + return ds_topo.ocean_mask + def _extrap_horiz(self, in_filename, out_filename, use_ocean_mask): """ Extrapolate horizontally on each depth level. From f54c7affd197f09826be615d5d1e0a62e4811a6b Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 6 Apr 2026 02:14:00 -0500 Subject: [PATCH 10/13] Add optional viz step to woa23 extrap --- .../hydrography/woa23/__init__.py | 1 + .../global_ocean/hydrography/woa23/steps.py | 13 +- .../global_ocean/hydrography/woa23/task.py | 2 +- .../global_ocean/hydrography/woa23/viz.py | 386 ++++++++++++++++++ .../global_ocean/hydrography/woa23/woa23.cfg | 28 ++ 5 files changed, 428 insertions(+), 2 deletions(-) create mode 100644 polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py index da90111410..315f4bdaf2 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/__init__.py @@ -3,3 +3,4 @@ from .steps import get_woa23_steps as get_woa23_steps from .steps import get_woa23_topography_step as get_woa23_topography_step from .task import Woa23 as Woa23 +from .viz import Woa23VizStep as Woa23VizStep diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py index 1b421e3b26..b063f5d2ed 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/steps.py @@ -8,6 +8,7 @@ from .combine import CombineStep from .extrapolate import ExtrapolateStep +from .viz import Woa23VizStep def get_woa23_topography_step(): @@ -75,4 +76,14 @@ def get_woa23_steps(component, combine_topo_step): combine_topo_step=combine_topo_step, ) - return [combine_step, extrapolate_step], config + viz_subdir = os.path.join(subdir, 'viz') + viz_step = component.get_or_create_shared_step( + step_cls=Woa23VizStep, + subdir=viz_subdir, + config=config, + config_filename=config_filename, + extrapolate_step=extrapolate_step, + combine_topo_step=combine_topo_step, + ) + + return [combine_step, extrapolate_step, viz_step], config diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py index 2d6b67a526..b1d5c311cb 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/task.py @@ -33,7 +33,7 @@ def __init__(self, component): self.add_step(self.combine_topo_step, symlink='combine_topo') for step in steps: - self.add_step(step) + self.add_step(step, run_by_default=step.name != 'viz') def configure(self): """ diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py new file mode 100644 index 0000000000..a5c40e8655 --- /dev/null +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py @@ -0,0 +1,386 @@ +import cmocean # noqa: F401 +import matplotlib.pyplot as plt +import numpy as np +import pyproj +import xarray as xr + +from polaris import Step +from polaris.viz import plot_global_lat_lon_field, use_mplstyle +from polaris.viz.spherical import _setup_colormap + + +class Woa23VizStep(Step): + """ + A step for visualizing extrapolated WOA23 hydrography. + """ + + def __init__(self, component, subdir, extrapolate_step, combine_topo_step): + """ + Create the step. + + Parameters + ---------- + component : polaris.Component + The component the step belongs to. + + subdir : str + The subdirectory for the step. + + extrapolate_step : polaris.Step + The step that extrapolates WOA23 hydrography. + + combine_topo_step : polaris.Step + The cached ``e3sm/init`` step that produces combined topography on + the WOA23 grid. + """ + super().__init__( + component=component, + name='viz', + subdir=subdir, + ntasks=1, + min_tasks=1, + ) + self.extrapolate_step = extrapolate_step + self.combine_topo_step = combine_topo_step + + def setup(self): + """ + Set up input and output files for the step. + """ + super().setup() + self.add_input_file( + filename='woa.nc', + work_dir_target=( + f'{self.extrapolate_step.path}/' + f'{self.extrapolate_step.output_filename}' + ), + ) + self.add_input_file( + filename='topography.nc', + work_dir_target=( + f'{self.combine_topo_step.path}/' + f'{self.combine_topo_step.combined_filename}' + ), + ) + + self.outputs = [] + for depth in self.config.getlist( + 'woa23', 'horizontal_plot_depths', dtype=float + ): + depth_tag = self._depth_tag(depth) + self.add_output_file(filename=f'pt_an_depth_{depth_tag}.png') + self.add_output_file(filename=f's_an_depth_{depth_tag}.png') + + self.add_output_file(filename='filchner_section.png') + self.add_output_file(filename='ross_section.png') + + def run(self): + """ + Plot horizontal fields at selected depths and two Antarctic transects. + """ + use_mplstyle() + + with xr.open_dataset('woa.nc', decode_times=False) as ds_woa: + ds_woa = ds_woa.load() + with xr.open_dataset('topography.nc') as ds_topo: + ds_topo = ds_topo.load() + + self._plot_horizontal_maps(ds_woa=ds_woa, ds_topo=ds_topo) + self._plot_configured_transects(ds_woa=ds_woa, ds_topo=ds_topo) + + def _plot_horizontal_maps(self, ds_woa, ds_topo): + logger = self.logger + fields = [ + ( + 'pt_an', + 'Potential temperature', + r'Potential temperature ($^{\circ}$C)', + 'woa23_viz_temperature', + ), + ( + 's_an', + 'Salinity', + 'Salinity (PSU)', + 'woa23_viz_salinity', + ), + ] + + requested_depths = self.config.getlist( + 'woa23', 'horizontal_plot_depths', dtype=float + ) + depth_values = ds_woa.depth.values + + for requested_depth in requested_depths: + depth_index = int( + np.abs(depth_values - requested_depth).argmin().item() + ) + actual_depth = float(depth_values[depth_index]) + depth_tag = self._depth_tag(requested_depth) + water_mask = self._get_level_water_mask( + ds_topo=ds_topo, depth=actual_depth + ) + + logger.info( + f'Plotting horizontal maps for requested depth ' + f'{requested_depth:g} m using WOA level {actual_depth:g} m' + ) + for ( + field_name, + field_title, + colorbar_label, + cmap_section, + ) in fields: + data = ds_woa[field_name].isel(depth=depth_index).values + data = np.where(water_mask, data, np.nan) + title = self._horizontal_title( + field_title=field_title, + requested_depth=requested_depth, + actual_depth=actual_depth, + ) + out_filename = f'{field_name}_depth_{depth_tag}.png' + plot_global_lat_lon_field( + lon=ds_woa.lon.values, + lat=ds_woa.lat.values, + data_array=data, + out_filename=out_filename, + config=self.config, + colormap_section=cmap_section, + title=title, + colorbar_label=colorbar_label, + ) + + def _plot_configured_transects(self, ds_woa, ds_topo): + transects = [ + ( + 'Filchner Trough', + 'filchner', + 'filchner_section.png', + ), + ( + 'Ross Ice Shelf cavity', + 'ross', + 'ross_section.png', + ), + ] + + for transect_name, prefix, out_filename in transects: + start_lon = self.config.getfloat('woa23', f'{prefix}_start_lon') + start_lat = self.config.getfloat('woa23', f'{prefix}_start_lat') + end_lon = self.config.getfloat('woa23', f'{prefix}_end_lon') + end_lat = self.config.getfloat('woa23', f'{prefix}_end_lat') + self.logger.info(f'Plotting {transect_name} section') + self._plot_transect( + ds_woa=ds_woa, + ds_topo=ds_topo, + title=transect_name, + start_lon=start_lon, + start_lat=start_lat, + end_lon=end_lon, + end_lat=end_lat, + out_filename=out_filename, + ) + + def _plot_transect( + self, + ds_woa, + ds_topo, + title, + start_lon, + start_lat, + end_lon, + end_lat, + out_filename, + ): + distance, lon, lat = self._build_transect( + start_lon=start_lon, + start_lat=start_lat, + end_lon=end_lon, + end_lat=end_lat, + ) + interp_lon = self._normalize_longitudes(lon, ds_woa.lon.values) + coords = { + 'lat': xr.DataArray(lat, dims=('nPoints',)), + 'lon': xr.DataArray(interp_lon, dims=('nPoints',)), + } + + ds_section = ds_woa[['pt_an', 's_an']].interp(**coords) + ds_topo_section = ds_topo[ + ['base_elevation', 'ice_draft', 'ocean_mask'] + ].interp(**coords) + + top_depth = np.maximum(-ds_topo_section.ice_draft.values, 0.0) + bottom_depth = np.maximum(-ds_topo_section.base_elevation.values, 0.0) + ocean_mask = ds_topo_section.ocean_mask.values > 0.5 + valid_column = np.logical_and(ocean_mask, bottom_depth > top_depth) + + depth = ds_woa.depth.values + depth_bounds = self._depth_bounds(ds_woa.depth_bnds.values) + distance_bounds = self._distance_bounds(distance) + water_mask = ( + (depth[:, np.newaxis] >= top_depth[np.newaxis, :]) + & (depth[:, np.newaxis] <= bottom_depth[np.newaxis, :]) + & valid_column[np.newaxis, :] + ) + + fields = [ + ( + 'pt_an', + 'Potential temperature', + r'Potential temperature ($^{\circ}$C)', + 'woa23_viz_temperature', + ), + ( + 's_an', + 'Salinity', + 'Salinity (PSU)', + 'woa23_viz_salinity', + ), + ] + + fig, axes = plt.subplots( + 2, + 1, + figsize=(12, 8), + sharex=True, + constrained_layout=True, + ) + + max_depth = float(depth_bounds[-1]) + for ax, (field_name, field_title, colorbar_label, cmap_section) in zip( + axes, fields, strict=True + ): + plot_data = ds_section[field_name].values + plot_data = np.where(water_mask, plot_data, np.nan) + colormap, norm, ticks = _setup_colormap(self.config, cmap_section) + image = ax.pcolormesh( + distance_bounds, + depth_bounds, + plot_data, + cmap=colormap, + norm=norm, + shading='auto', + ) + + # Grounded ice and land are shown as a light gray background, + # floating ice by the cavity roof and the bed by dark gray fill. + ax.fill_between( + distance, + 0.0, + max_depth, + where=np.logical_not(valid_column), + color='0.88', + linewidth=0.0, + ) + ax.fill_between( + distance, + 0.0, + top_depth, + where=top_depth > 0.0, + color='lightsteelblue', + linewidth=0.0, + ) + ax.fill_between( + distance, + bottom_depth, + max_depth, + where=valid_column, + color='dimgray', + linewidth=0.0, + ) + ax.plot(distance, top_depth, color='k', linewidth=1.0) + ax.plot(distance, bottom_depth, color='k', linewidth=1.0) + ax.set_ylabel('Depth (m)') + ax.set_title(field_title) + ax.set_ylim(max_depth, 0.0) + colorbar = fig.colorbar(image, ax=ax, pad=0.01) + colorbar.set_label(colorbar_label) + if ticks is not None: + colorbar.set_ticks(ticks) + colorbar.set_ticklabels([f'{tick}' for tick in ticks]) + + axes[-1].set_xlabel('Distance along transect (km)') + fig.suptitle( + f'{title}: ({start_lon:.2f}, {start_lat:.2f}) to ' + f'({end_lon:.2f}, {end_lat:.2f})' + ) + fig.savefig(out_filename, bbox_inches='tight', pad_inches=0.1) + plt.close(fig) + + @staticmethod + def _depth_tag(depth): + depth_str = f'{depth:07.2f}' + depth_str = depth_str.replace('.', 'p').replace('-', 'm') + return f'{depth_str}m' + + @staticmethod + def _horizontal_title(field_title, requested_depth, actual_depth): + if np.isclose(requested_depth, actual_depth): + return f'{field_title} at {actual_depth:g} m' + return ( + f'{field_title} near {requested_depth:g} m ' + f'(WOA level {actual_depth:g} m)' + ) + + @staticmethod + def _get_level_water_mask(ds_topo, depth): + top_depth = np.maximum(-ds_topo.ice_draft.values, 0.0) + bottom_depth = np.maximum(-ds_topo.base_elevation.values, 0.0) + ocean_mask = ds_topo.ocean_mask.values > 0.5 + return (top_depth <= depth) & (bottom_depth >= depth) & ocean_mask + + @staticmethod + def _build_transect(start_lon, start_lat, end_lon, end_lat): + geod = pyproj.Geod(ellps='WGS84') + _, _, total_distance = geod.inv(start_lon, start_lat, end_lon, end_lat) + total_distance_km = total_distance * 1.0e-3 + npoints = int(np.clip(total_distance_km / 10.0, 200, 1000)) + if npoints < 2: + npoints = 2 + + if npoints > 2: + intermediate = geod.npts( + start_lon, + start_lat, + end_lon, + end_lat, + npoints - 2, + ) + lon = np.array( + [start_lon] + [point[0] for point in intermediate] + [end_lon] + ) + lat = np.array( + [start_lat] + [point[1] for point in intermediate] + [end_lat] + ) + else: + lon = np.array([start_lon, end_lon]) + lat = np.array([start_lat, end_lat]) + + _, _, segment_distance = geod.inv(lon[:-1], lat[:-1], lon[1:], lat[1:]) + distance = np.zeros(npoints) + distance[1:] = np.cumsum(segment_distance) * 1.0e-3 + return distance, lon, lat + + @staticmethod + def _normalize_longitudes(lon, lon_coord): + lon_min = float(np.nanmin(lon_coord)) + lon_max = float(np.nanmax(lon_coord)) + if lon_min >= 0.0 and lon_max > 180.0: + return np.mod(lon, 360.0) + return ((lon + 180.0) % 360.0) - 180.0 + + @staticmethod + def _depth_bounds(depth_bounds): + lower = depth_bounds[:, 0] + upper = depth_bounds[:, 1] + return np.concatenate(([lower[0]], upper)) + + @staticmethod + def _distance_bounds(distance): + if distance.size == 1: + return np.array([0.0, distance[0]]) + + bounds = np.zeros(distance.size + 1) + bounds[1:-1] = 0.5 * (distance[:-1] + distance[1:]) + bounds[0] = 0.0 + bounds[-1] = distance[-1] + return bounds diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg index 48ebcbeadb..c5bcaf955b 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg @@ -4,3 +4,31 @@ # the minimum weight sum needed to mark a new cell valid in horizontal # extrapolation extrap_threshold = 0.01 + +# target depths for horizontal plots of the extrapolated product (m) +horizontal_plot_depths = 0.0, 200.0, 400.0, 600.0, 800.0 + +# endpoints of a transect through Filchner Trough and into the Filchner +# ice-shelf cavity +filchner_start_lon = -46.0 +filchner_start_lat = -74.5 +filchner_end_lon = -38.5 +filchner_end_lat = -81.2 + +# endpoints of a transect through the Ross Ice Shelf cavity +ross_start_lon = 176.0 +ross_start_lat = -76.5 +ross_end_lon = -171.0 +ross_end_lat = -84.0 + + +[woa23_viz_temperature] +colormap_name = cmo.thermal +norm_type = linear +norm_args = {'vmin': -2.5, 'vmax': 30.0} + + +[woa23_viz_salinity] +colormap_name = cmo.haline +norm_type = linear +norm_args = {'vmin': 32.0, 'vmax': 37.0} From c53b36319ec308197a22613b89bf84b9f844cb3d Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 11 Apr 2026 17:56:02 +0200 Subject: [PATCH 11/13] Several fixes to WOA viz --- .../global_ocean/hydrography/woa23/viz.py | 144 ++++++++++++------ .../global_ocean/hydrography/woa23/woa23.cfg | 7 +- 2 files changed, 106 insertions(+), 45 deletions(-) diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py index a5c40e8655..cf76b5e25e 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py @@ -69,10 +69,16 @@ def setup(self): ): depth_tag = self._depth_tag(depth) self.add_output_file(filename=f'pt_an_depth_{depth_tag}.png') + self.add_output_file( + filename=f'pt_an_depth_{depth_tag}_filled.png' + ) self.add_output_file(filename=f's_an_depth_{depth_tag}.png') + self.add_output_file(filename=f's_an_depth_{depth_tag}_filled.png') self.add_output_file(filename='filchner_section.png') + self.add_output_file(filename='filchner_section_filled.png') self.add_output_file(filename='ross_section.png') + self.add_output_file(filename='ross_section_filled.png') def run(self): """ @@ -131,21 +137,37 @@ def _plot_horizontal_maps(self, ds_woa, ds_topo): cmap_section, ) in fields: data = ds_woa[field_name].isel(depth=depth_index).values - data = np.where(water_mask, data, np.nan) - title = self._horizontal_title( + ocean_title = self._horizontal_title( field_title=field_title, requested_depth=requested_depth, actual_depth=actual_depth, + filled=False, + ) + plot_global_lat_lon_field( + lon=ds_woa.lon.values, + lat=ds_woa.lat.values, + data_array=np.where(water_mask, data, np.nan), + out_filename=f'{field_name}_depth_{depth_tag}.png', + config=self.config, + colormap_section=cmap_section, + title=ocean_title, + colorbar_label=colorbar_label, + ) + + filled_title = self._horizontal_title( + field_title=field_title, + requested_depth=requested_depth, + actual_depth=actual_depth, + filled=True, ) - out_filename = f'{field_name}_depth_{depth_tag}.png' plot_global_lat_lon_field( lon=ds_woa.lon.values, lat=ds_woa.lat.values, data_array=data, - out_filename=out_filename, + out_filename=f'{field_name}_depth_{depth_tag}_filled.png', config=self.config, colormap_section=cmap_section, - title=title, + title=filled_title, colorbar_label=colorbar_label, ) @@ -155,15 +177,18 @@ def _plot_configured_transects(self, ds_woa, ds_topo): 'Filchner Trough', 'filchner', 'filchner_section.png', + 'filchner_section_filled.png', ), ( 'Ross Ice Shelf cavity', 'ross', 'ross_section.png', + 'ross_section_filled.png', ), ] - for transect_name, prefix, out_filename in transects: + max_depth = self.config.getfloat('woa23', 'section_max_depth') + for transect_name, prefix, out_filename, filled_filename in transects: start_lon = self.config.getfloat('woa23', f'{prefix}_start_lon') start_lat = self.config.getfloat('woa23', f'{prefix}_start_lat') end_lon = self.config.getfloat('woa23', f'{prefix}_end_lon') @@ -178,6 +203,20 @@ def _plot_configured_transects(self, ds_woa, ds_topo): end_lon=end_lon, end_lat=end_lat, out_filename=out_filename, + max_depth=max_depth, + filled=False, + ) + self._plot_transect( + ds_woa=ds_woa, + ds_topo=ds_topo, + title=transect_name, + start_lon=start_lon, + start_lat=start_lat, + end_lon=end_lon, + end_lat=end_lat, + out_filename=filled_filename, + max_depth=max_depth, + filled=True, ) def _plot_transect( @@ -190,6 +229,8 @@ def _plot_transect( end_lon, end_lat, out_filename, + max_depth, + filled, ): distance, lon, lat = self._build_transect( start_lon=start_lon, @@ -245,12 +286,15 @@ def _plot_transect( constrained_layout=True, ) - max_depth = float(depth_bounds[-1]) + max_depth = min(max_depth, float(depth_bounds[-1])) + top_depth_plot = np.minimum(top_depth, max_depth) + bottom_depth_plot = np.minimum(bottom_depth, max_depth) for ax, (field_name, field_title, colorbar_label, cmap_section) in zip( axes, fields, strict=True ): plot_data = ds_section[field_name].values - plot_data = np.where(water_mask, plot_data, np.nan) + if not filled: + plot_data = np.where(water_mask, plot_data, np.nan) colormap, norm, ticks = _setup_colormap(self.config, cmap_section) image = ax.pcolormesh( distance_bounds, @@ -261,34 +305,36 @@ def _plot_transect( shading='auto', ) - # Grounded ice and land are shown as a light gray background, - # floating ice by the cavity roof and the bed by dark gray fill. - ax.fill_between( - distance, - 0.0, - max_depth, - where=np.logical_not(valid_column), - color='0.88', - linewidth=0.0, - ) - ax.fill_between( - distance, - 0.0, - top_depth, - where=top_depth > 0.0, - color='lightsteelblue', - linewidth=0.0, - ) - ax.fill_between( - distance, - bottom_depth, - max_depth, - where=valid_column, - color='dimgray', - linewidth=0.0, - ) - ax.plot(distance, top_depth, color='k', linewidth=1.0) - ax.plot(distance, bottom_depth, color='k', linewidth=1.0) + if not filled: + # Grounded ice and land are shown as a light gray + # background, floating ice by the cavity roof and the bed + # by dark gray fill. + ax.fill_between( + distance, + 0.0, + max_depth, + where=np.logical_not(valid_column), + color='0.88', + linewidth=0.0, + ) + ax.fill_between( + distance, + 0.0, + top_depth_plot, + where=top_depth > 0.0, + color='lightsteelblue', + linewidth=0.0, + ) + ax.fill_between( + distance, + bottom_depth_plot, + max_depth, + where=valid_column, + color='dimgray', + linewidth=0.0, + ) + ax.plot(distance, top_depth_plot, color='k', linewidth=1.0) + ax.plot(distance, bottom_depth_plot, color='k', linewidth=1.0) ax.set_ylabel('Depth (m)') ax.set_title(field_title) ax.set_ylim(max_depth, 0.0) @@ -300,7 +346,8 @@ def _plot_transect( axes[-1].set_xlabel('Distance along transect (km)') fig.suptitle( - f'{title}: ({start_lon:.2f}, {start_lat:.2f}) to ' + f'{self._transect_title(title=title, filled=filled)}: ' + f'({start_lon:.2f}, {start_lat:.2f}) to ' f'({end_lon:.2f}, {end_lat:.2f})' ) fig.savefig(out_filename, bbox_inches='tight', pad_inches=0.1) @@ -313,13 +360,24 @@ def _depth_tag(depth): return f'{depth_str}m' @staticmethod - def _horizontal_title(field_title, requested_depth, actual_depth): + def _horizontal_title(field_title, requested_depth, actual_depth, filled): if np.isclose(requested_depth, actual_depth): - return f'{field_title} at {actual_depth:g} m' - return ( - f'{field_title} near {requested_depth:g} m ' - f'(WOA level {actual_depth:g} m)' - ) + title = f'{field_title} at {actual_depth:g} m' + else: + title = ( + f'{field_title} near {requested_depth:g} m ' + f'(WOA level {actual_depth:g} m)' + ) + + if filled: + return f'{title} (filled field)' + return f'{title} (ocean only)' + + @staticmethod + def _transect_title(title, filled): + if filled: + return f'{title} (filled field)' + return f'{title} (ocean only)' @staticmethod def _get_level_water_mask(ds_topo, depth): diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg index c5bcaf955b..014248e275 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg @@ -8,16 +8,19 @@ extrap_threshold = 0.01 # target depths for horizontal plots of the extrapolated product (m) horizontal_plot_depths = 0.0, 200.0, 400.0, 600.0, 800.0 +# maximum depth to include in section plots (m) +section_max_depth = 2000.0 + # endpoints of a transect through Filchner Trough and into the Filchner # ice-shelf cavity filchner_start_lon = -46.0 -filchner_start_lat = -74.5 +filchner_start_lat = -71.5 filchner_end_lon = -38.5 filchner_end_lat = -81.2 # endpoints of a transect through the Ross Ice Shelf cavity ross_start_lon = 176.0 -ross_start_lat = -76.5 +ross_start_lat = -72.0 ross_end_lon = -171.0 ross_end_lat = -84.0 From f7016af964535fb8f6a40a87d173c389c71b136e Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 13 Apr 2026 21:22:28 +0200 Subject: [PATCH 12/13] Fix the transects in the viz step --- .../global_ocean/hydrography/woa23/viz.py | 42 +++++++++++++------ .../global_ocean/hydrography/woa23/woa23.cfg | 14 +++++++ 2 files changed, 43 insertions(+), 13 deletions(-) diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py index cf76b5e25e..56d009007d 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py @@ -238,16 +238,25 @@ def _plot_transect( end_lon=end_lon, end_lat=end_lat, ) - interp_lon = self._normalize_longitudes(lon, ds_woa.lon.values) + ds_woa = self._add_periodic_lon(ds_woa) + ds_topo = self._add_periodic_lon(ds_topo) coords = { 'lat': xr.DataArray(lat, dims=('nPoints',)), - 'lon': xr.DataArray(interp_lon, dims=('nPoints',)), + 'lon': xr.DataArray(lon, dims=('nPoints',)), } ds_section = ds_woa[['pt_an', 's_an']].interp(**coords) - ds_topo_section = ds_topo[ - ['base_elevation', 'ice_draft', 'ocean_mask'] - ].interp(**coords) + ds_topo_section = ds_topo[['base_elevation', 'ice_draft']].interp( + **coords + ) + ds_topo_section['ocean_mask'] = ( + ds_topo[['ocean_mask']] + .interp( + method='nearest', + **coords, + ) + .ocean_mask + ) top_depth = np.maximum(-ds_topo_section.ice_draft.values, 0.0) bottom_depth = np.maximum(-ds_topo_section.base_elevation.values, 0.0) @@ -268,13 +277,13 @@ def _plot_transect( 'pt_an', 'Potential temperature', r'Potential temperature ($^{\circ}$C)', - 'woa23_viz_temperature', + 'woa23_viz_section_temperature', ), ( 's_an', 'Salinity', 'Salinity (PSU)', - 'woa23_viz_salinity', + 'woa23_viz_section_salinity', ), ] @@ -419,12 +428,19 @@ def _build_transect(start_lon, start_lat, end_lon, end_lat): return distance, lon, lat @staticmethod - def _normalize_longitudes(lon, lon_coord): - lon_min = float(np.nanmin(lon_coord)) - lon_max = float(np.nanmax(lon_coord)) - if lon_min >= 0.0 and lon_max > 180.0: - return np.mod(lon, 360.0) - return ((lon + 180.0) % 360.0) - 180.0 + def _add_periodic_lon(ds): + lon = ds.lon + lon_sections = [lon - 360.0, lon, lon + 360.0] + lon_periodic = xr.concat(lon_sections, dim='lon') + + periodic = [] + for offset in [-360.0, 0.0, 360.0]: + shifted = ds.assign_coords(lon=ds.lon + offset) + periodic.append(shifted) + + ds_periodic = xr.concat(periodic, dim='lon', data_vars='all') + ds_periodic = ds_periodic.assign_coords(lon=lon_periodic) + return ds_periodic.sortby('lon') @staticmethod def _depth_bounds(depth_bounds): diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg index 014248e275..ad1b34a2a8 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/woa23.cfg @@ -35,3 +35,17 @@ norm_args = {'vmin': -2.5, 'vmax': 30.0} colormap_name = cmo.haline norm_type = linear norm_args = {'vmin': 32.0, 'vmax': 37.0} + + +[woa23_viz_section_temperature] +colormap_name = cmo.thermal +norm_type = linear +norm_args = {'vmin': -2.5, 'vmax': 1.5} +colorbar_ticks = -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5 + + +[woa23_viz_section_salinity] +colormap_name = cmo.haline +norm_type = linear +norm_args = {'vmin': 34.0, 'vmax': 35.0} +colorbar_ticks = 34.0, 34.2, 34.4, 34.6, 34.8, 35.0 From 59b0dc6abd8b1664a0fbf7804ede46a6a24cae25 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 14 Apr 2026 14:12:47 +0200 Subject: [PATCH 13/13] Switch WOA23 steps to produce CT/SA We want that, not PT and S, so we are focused on Omega more than MPAS-Ocean. --- .../ocean/tasks/global_ocean.md | 5 +- docs/users_guide/ocean/tasks/global_ocean.md | 3 +- .../global_ocean/hydrography/woa23/combine.py | 56 +++++++++----- .../hydrography/woa23/extrapolate.py | 8 +- .../global_ocean/hydrography/woa23/viz.py | 36 ++++----- .../hydrography/woa23/test_combine.py | 76 +++++++++++++++++++ 6 files changed, 140 insertions(+), 44 deletions(-) create mode 100644 tests/ocean/global_ocean/hydrography/woa23/test_combine.py diff --git a/docs/developers_guide/ocean/tasks/global_ocean.md b/docs/developers_guide/ocean/tasks/global_ocean.md index 998564fc85..8afcf3f8e1 100644 --- a/docs/developers_guide/ocean/tasks/global_ocean.md +++ b/docs/developers_guide/ocean/tasks/global_ocean.md @@ -52,8 +52,9 @@ combines January and annual WOA23 temperature and salinity climatologies into a single dataset. January values are used where they exist, and annual values fill deeper levels where the monthly product is not available. -This step also converts in-situ temperature to potential temperature with -`gsw`, producing `woa_combined.nc`. +WOA23 supplies in-situ temperature and practical salinity, so this step uses +`gsw` to derive conservative temperature and absolute salinity for the +canonical `woa_combined.nc` product. ### extrapolate diff --git a/docs/users_guide/ocean/tasks/global_ocean.md b/docs/users_guide/ocean/tasks/global_ocean.md index 0b2d8d8175..c319197553 100644 --- a/docs/users_guide/ocean/tasks/global_ocean.md +++ b/docs/users_guide/ocean/tasks/global_ocean.md @@ -35,7 +35,8 @@ The task is organized into three inspectable steps: 1. `combine_topo` reuses a cached `e3sm/init` combined-topography step configured for the WOA23 0.25-degree latitude-longitude grid. 2. `combine` creates `woa_combined.nc` by combining January and annual WOA23 - fields and converting temperature to potential temperature. + in-situ temperature and practical-salinity fields, then deriving + conservative temperature and absolute salinity. 3. `extrapolate` creates the final `woa23_decav_0.25_jan_extrap.nc` product. diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py index b4ec69f1c8..bf10a8695f 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/combine.py @@ -75,8 +75,8 @@ def setup(self): def run(self): """ - Combine January and annual climatologies and convert temperature to - potential temperature. + Combine January and annual climatologies and derive conservative + temperature and absolute salinity. """ logger = self.logger logger.info('Combining January and annual WOA23 climatologies') @@ -108,14 +108,15 @@ def run(self): ds_out[var_name] = xr.concat(slices, dim='depth') ds_out[var_name].attrs = ds_ann[var_name].attrs - ds_out = self._temp_to_potential_temp(ds_out) + ds_out = self._to_canonical_teos10(ds_out) write_netcdf(ds_out, 'woa_combined.nc') logger.info('Wrote woa_combined.nc') @staticmethod - def _temp_to_potential_temp(ds): + def _to_canonical_teos10(ds): """ - Convert WOA in-situ temperature to potential temperature. + Convert WOA in-situ temperature and practical salinity to canonical + conservative temperature and absolute salinity. Parameters ---------- @@ -125,10 +126,11 @@ def _temp_to_potential_temp(ds): Returns ------- ds : xarray.Dataset - The dataset with temperature replaced by potential temperature. + The dataset with conservative temperature and absolute salinity. """ dims = ds.t_an.dims - slices = [] + ct_slices = [] + sa_slices = [] for depth_index in range(ds.sizes['depth']): temp_slice = ds.t_an.isel(depth=depth_index) in_situ_temp = temp_slice.values @@ -138,33 +140,47 @@ def _temp_to_potential_temp(ds): z = -ds.depth.isel(depth=depth_index).values pressure = gsw.p_from_z(z, lat) - mask = np.isfinite(in_situ_temp) - potential_temp = np.full(in_situ_temp.shape, np.nan) - absolute_salinity = gsw.SA_from_SP( + mask = np.isfinite(in_situ_temp) & np.isfinite(practical_salinity) + conservative_temp = np.full(in_situ_temp.shape, np.nan) + absolute_salinity = np.full(practical_salinity.shape, np.nan) + absolute_salinity[mask] = gsw.SA_from_SP( practical_salinity[mask], pressure[mask], lon[mask], lat[mask], ) - potential_temp[mask] = gsw.pt_from_t( - absolute_salinity, + conservative_temp[mask] = gsw.CT_from_t( + absolute_salinity[mask], in_situ_temp[mask], pressure[mask], - p_ref=0.0, ) - slices.append( + ct_slices.append( xr.DataArray( - data=potential_temp, + data=conservative_temp, dims=temp_slice.dims, attrs=temp_slice.attrs, ) ) + sa_slices.append( + xr.DataArray( + data=absolute_salinity, + dims=temp_slice.dims, + attrs=ds.s_an.attrs, + ) + ) - ds['pt_an'] = xr.concat(slices, dim='depth').transpose(*dims) - ds.pt_an.attrs['standard_name'] = 'sea_water_potential_temperature' - ds.pt_an.attrs['long_name'] = ( + ds['ct_an'] = xr.concat(ct_slices, dim='depth').transpose(*dims) + ds.ct_an.attrs['standard_name'] = 'sea_water_conservative_temperature' + ds.ct_an.attrs['long_name'] = ( + 'Objectively analyzed mean fields for ' + 'sea_water_conservative_temperature at standard depth levels.' + ) + ds['sa_an'] = xr.concat(sa_slices, dim='depth').transpose(*dims) + ds.sa_an.attrs['standard_name'] = 'sea_water_absolute_salinity' + ds.sa_an.attrs['long_name'] = ( 'Objectively analyzed mean fields for ' - 'sea_water_potential_temperature at standard depth levels.' + 'sea_water_absolute_salinity at standard depth levels.' ) + ds.sa_an.attrs['units'] = 'g kg-1' - return ds.drop_vars('t_an') + return ds.drop_vars(['t_an', 's_an']) diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py index 9867e47d31..92d83e58cc 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/extrapolate.py @@ -162,7 +162,7 @@ def _extrap_horiz(self, in_filename, out_filename, use_ocean_mask): ocean_mask=level_mask, threshold=self.config.getfloat('woa23', 'extrap_threshold'), ) - for field_name in ['pt_an', 's_an']: + for field_name in ['ct_an', 'sa_an']: ds_out[field_name][depth_index, :, :] = ds_level[field_name] write_netcdf(ds_out, out_filename) @@ -192,7 +192,7 @@ def _extrap_vert(self, in_filename, out_filename, use_ocean_mask): ocean_mask = ds_mask.ocean_mask.values.astype(bool) ndepth = ds_out.sizes['depth'] - for field_name in ['pt_an', 's_an']: + for field_name in ['ct_an', 'sa_an']: field = ds_out[field_name].values for depth_index in range(1, ndepth): mask = np.isnan(field[depth_index, :, :]) @@ -226,7 +226,7 @@ def _extrap_level(ds_level, ocean_mask, threshold): The filled depth level. """ kernel = ExtrapolateStep._get_kernel() - field = ds_level.pt_an.values + field = ds_level.ct_an.values valid = np.isfinite(field) orig_mask = valid.copy() @@ -239,7 +239,7 @@ def _extrap_level(ds_level, ocean_mask, threshold): fields = { field_name: ds_level[field_name].values.copy() - for field_name in ['pt_an', 's_an'] + for field_name in ['ct_an', 'sa_an'] } nlon = field.shape[1] diff --git a/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py index 56d009007d..485ff716c7 100644 --- a/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py +++ b/polaris/tasks/ocean/global_ocean/hydrography/woa23/viz.py @@ -68,12 +68,14 @@ def setup(self): 'woa23', 'horizontal_plot_depths', dtype=float ): depth_tag = self._depth_tag(depth) - self.add_output_file(filename=f'pt_an_depth_{depth_tag}.png') + self.add_output_file(filename=f'ct_an_depth_{depth_tag}.png') self.add_output_file( - filename=f'pt_an_depth_{depth_tag}_filled.png' + filename=f'ct_an_depth_{depth_tag}_filled.png' + ) + self.add_output_file(filename=f'sa_an_depth_{depth_tag}.png') + self.add_output_file( + filename=f'sa_an_depth_{depth_tag}_filled.png' ) - self.add_output_file(filename=f's_an_depth_{depth_tag}.png') - self.add_output_file(filename=f's_an_depth_{depth_tag}_filled.png') self.add_output_file(filename='filchner_section.png') self.add_output_file(filename='filchner_section_filled.png') @@ -98,15 +100,15 @@ def _plot_horizontal_maps(self, ds_woa, ds_topo): logger = self.logger fields = [ ( - 'pt_an', - 'Potential temperature', - r'Potential temperature ($^{\circ}$C)', + 'ct_an', + 'Conservative temperature', + r'Conservative temperature ($^{\circ}$C)', 'woa23_viz_temperature', ), ( - 's_an', - 'Salinity', - 'Salinity (PSU)', + 'sa_an', + 'Absolute salinity', + r'Absolute salinity (g kg$^{-1}$)', 'woa23_viz_salinity', ), ] @@ -245,7 +247,7 @@ def _plot_transect( 'lon': xr.DataArray(lon, dims=('nPoints',)), } - ds_section = ds_woa[['pt_an', 's_an']].interp(**coords) + ds_section = ds_woa[['ct_an', 'sa_an']].interp(**coords) ds_topo_section = ds_topo[['base_elevation', 'ice_draft']].interp( **coords ) @@ -274,15 +276,15 @@ def _plot_transect( fields = [ ( - 'pt_an', - 'Potential temperature', - r'Potential temperature ($^{\circ}$C)', + 'ct_an', + 'Conservative temperature', + r'Conservative temperature ($^{\circ}$C)', 'woa23_viz_section_temperature', ), ( - 's_an', - 'Salinity', - 'Salinity (PSU)', + 'sa_an', + 'Absolute salinity', + r'Absolute salinity (g kg$^{-1}$)', 'woa23_viz_section_salinity', ), ] diff --git a/tests/ocean/global_ocean/hydrography/woa23/test_combine.py b/tests/ocean/global_ocean/hydrography/woa23/test_combine.py new file mode 100644 index 0000000000..6c4af0a7be --- /dev/null +++ b/tests/ocean/global_ocean/hydrography/woa23/test_combine.py @@ -0,0 +1,76 @@ +import gsw +import numpy as np +import pytest +import xarray as xr + +from polaris.tasks.ocean.global_ocean.hydrography.woa23.combine import ( + CombineStep, +) + + +def test_to_canonical_teos10(): + ds = xr.Dataset( + data_vars={ + 't_an': ( + ('depth', 'lat', 'lon'), + np.array( + [ + [[1.0, 2.0], [3.0, np.nan]], + [[0.5, 1.5], [2.5, 3.5]], + ] + ), + ), + 's_an': ( + ('depth', 'lat', 'lon'), + np.array( + [ + [[34.5, 34.7], [34.9, 35.1]], + [[34.6, 34.8], [35.0, np.nan]], + ] + ), + ), + }, + coords={ + 'depth': ('depth', np.array([0.0, 1000.0])), + 'lat': ('lat', np.array([-45.0, 10.0])), + 'lon': ('lon', np.array([20.0, 160.0])), + }, + ) + + ds_out = CombineStep._to_canonical_teos10(ds) + + assert 't_an' not in ds_out + assert 's_an' not in ds_out + assert ds_out.ct_an.attrs['standard_name'] == ( + 'sea_water_conservative_temperature' + ) + assert ds_out.sa_an.attrs['standard_name'] == ( + 'sea_water_absolute_salinity' + ) + assert ds_out.sa_an.attrs['units'] == 'g kg-1' + + expected_ct = np.full(ds_out.ct_an.shape, np.nan) + expected_sa = np.full(ds_out.sa_an.shape, np.nan) + for depth_index, depth in enumerate(ds.depth.values): + temp = ds.t_an.isel(depth=depth_index).values + practical_salinity = ds.s_an.isel(depth=depth_index).values + level = ds.t_an.isel(depth=depth_index) + lat = ds.lat.broadcast_like(level).values + lon = ds.lon.broadcast_like(level).values + pressure = gsw.p_from_z(-depth, lat) + mask = np.isfinite(temp) & np.isfinite(practical_salinity) + + expected_sa[depth_index, :, :][mask] = gsw.SA_from_SP( + practical_salinity[mask], + pressure[mask], + lon[mask], + lat[mask], + ) + expected_ct[depth_index, :, :][mask] = gsw.CT_from_t( + expected_sa[depth_index, :, :][mask], + temp[mask], + pressure[mask], + ) + + assert ds_out.ct_an.values == pytest.approx(expected_ct, nan_ok=True) + assert ds_out.sa_an.values == pytest.approx(expected_sa, nan_ok=True)