Merge pull request #101 from xylar/fix-max-level-cell-in-horiz-viz

Get maxLevelCell from ds, not ds_mesh in viz
E3SM-Project · Jul 27, 2023 · 76c6700 · 76c6700
2 parents 90a8a9a + ff9e229
commit 76c6700
Show file tree

Hide file tree

Showing 7 changed files with 13 additions and 11 deletions.
diff --git a/polaris/ocean/tests/baroclinic_channel/init.py b/polaris/ocean/tests/baroclinic_channel/init.py
@@ -90,8 +90,6 @@ def run(self):
 
         init_vertical_coord(config, ds)
 
-        ds_mesh['maxLevelCell'] = ds.maxLevelCell
-
         x_min = x_cell.min().values
         x_max = x_cell.max().values
         y_min = y_cell.min().values

diff --git a/polaris/ocean/tests/baroclinic_channel/viz.py b/polaris/ocean/tests/baroclinic_channel/viz.py
@@ -24,6 +24,9 @@ def __init__(self, test_case):
         self.add_input_file(
             filename='mesh.nc',
             target='../init/culled_mesh.nc')
+        self.add_input_file(
+            filename='init.nc',
+            target='../init/initial_state.nc')
         self.add_input_file(
             filename='output.nc',
             target='../forward/output.nc')
@@ -33,7 +36,9 @@ def run(self):
         Run this step of the test case
         """
         ds_mesh = xr.load_dataset('mesh.nc')
+        ds_init = xr.load_dataset('init.nc')
         ds = xr.load_dataset('output.nc')
+        ds['maxLevelCell'] = ds_init.maxLevelCell
         t_index = ds.sizes['Time'] - 1
         plot_horiz_field(ds, ds_mesh, 'temperature',
                          'final_temperature.png', t_index=t_index)

diff --git a/polaris/ocean/tests/inertial_gravity_wave/init.py b/polaris/ocean/tests/inertial_gravity_wave/init.py
@@ -79,7 +79,6 @@ def run(self):
         ds['fEdge'] = f0 * xr.ones_like(ds_mesh.xEdge)
         ds['fVertex'] = f0 * xr.ones_like(ds_mesh.xVertex)
 
-        ds_mesh['maxLevelCell'] = ds.maxLevelCell
         exact_solution = ExactSolution(ds, config)
 
         ssh = exact_solution.ssh(0.0)

diff --git a/polaris/ocean/tests/inertial_gravity_wave/viz.py b/polaris/ocean/tests/inertial_gravity_wave/viz.py
@@ -40,7 +40,7 @@ def __init__(self, test_case, resolutions):
         for resolution in resolutions:
             self.add_input_file(
                 filename=f'mesh_{resolution}km.nc',
-                target=f'../{resolution}km/initial_state/culled_mesh.nc')
+                target=f'../{resolution}km/init/culled_mesh.nc')
             self.add_input_file(
                 filename=f'init_{resolution}km.nc',
                 target=f'../{resolution}km/init/initial_state.nc')
@@ -69,6 +69,7 @@ def run(self):
             ds_mesh = xr.open_dataset(f'mesh_{res}km.nc')
             ds_init = xr.open_dataset(f'init_{res}km.nc')
             ds = xr.open_dataset(f'output_{res}km.nc')
+            ds['maxLevelCell'] = ds_init.maxLevelCell
             exact = ExactSolution(ds_init, config)
 
             t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(),

diff --git a/polaris/ocean/tests/manufactured_solution/init.py b/polaris/ocean/tests/manufactured_solution/init.py
@@ -79,8 +79,6 @@ def run(self):
         ds['fEdge'] = coriolis_parameter * xr.ones_like(ds_mesh.xEdge)
         ds['fVertex'] = coriolis_parameter * xr.ones_like(ds_mesh.xVertex)
 
-        ds_mesh['maxLevelCell'] = ds.maxLevelCell
-
         # Evaluate the exact solution at time=0
         exact_solution = ExactSolution(config, ds)
         ssh = exact_solution.ssh(0.0)

diff --git a/polaris/ocean/tests/manufactured_solution/viz.py b/polaris/ocean/tests/manufactured_solution/viz.py
@@ -40,7 +40,7 @@ def __init__(self, test_case, resolutions):
         for resolution in resolutions:
             self.add_input_file(
                 filename=f'mesh_{resolution}km.nc',
-                target=f'../{resolution}km/initial_state/culled_mesh.nc')
+                target=f'../{resolution}km/init/culled_mesh.nc')
             self.add_input_file(
                 filename=f'init_{resolution}km.nc',
                 target=f'../{resolution}km/init/initial_state.nc')
@@ -68,6 +68,7 @@ def run(self):
             ds_mesh = xr.open_dataset(f'mesh_{res}km.nc')
             ds_init = xr.open_dataset(f'init_{res}km.nc')
             ds = xr.open_dataset(f'output_{res}km.nc')
+            ds['maxLevelCell'] = ds_init.maxLevelCell
             exact = ExactSolution(config, ds_init)
 
             t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(),

diff --git a/polaris/viz/__init__.py b/polaris/viz/__init__.py
@@ -31,7 +31,7 @@ def plot_horiz_field(ds, ds_mesh, field_name, out_file_name=None,  # noqa: C901
         A data set containing fieldName
 
     ds_mesh : xarray.Dataset
-        A data set containing mesh variables
+        A data set containing horizontal mesh variables
 
     field_name: str
         The name of the variable to plot, which must be present in ds
@@ -83,9 +83,9 @@ def plot_horiz_field(ds, ds_mesh, field_name, out_file_name=None,  # noqa: C901
     if title is None:
         title = field_name
 
-    if 'maxLevelCell' not in ds_mesh:
+    if 'maxLevelCell' not in ds:
         raise ValueError(
-            'maxLevelCell must be added to ds_mesh before plotting.')
+            'maxLevelCell must be added to ds before plotting.')
     if field_name not in ds:
         raise ValueError(
             f'{field_name} must be present in ds before plotting.')
@@ -102,7 +102,7 @@ def plot_horiz_field(ds, ds_mesh, field_name, out_file_name=None,  # noqa: C901
         field = field.isel(nVertLevels=z_index)
 
     if 'nCells' in field.dims:
-        ocean_mask = ds_mesh.maxLevelCell - 1 >= 0
+        ocean_mask = ds.maxLevelCell - 1 >= 0
         ocean_patches, ocean_mask = _compute_cell_patches(ds_mesh, ocean_mask)
     elif 'nEdges' in field.dims:
         ocean_mask = np.ones_like(field, dtype='bool')