Clean-up from code review

Co-authored-by: Xylar Asay-Davis <[email protected]>
E3SM-Project · Apr 12, 2023 · af73520 · af73520
1 parent e3593e8
commit af73520
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/polaris/ocean/tests/baroclinic_channel/viz.py b/polaris/ocean/tests/baroclinic_channel/viz.py
@@ -47,5 +47,5 @@ def run(self):
         plot_horiz_field(config, ds, ds_mesh, 'normalVelocity',
                          'final_normalVelocity.png',
                          t_index=t_index,
-                         vmin=-1 * max_velocity, vmax=max_velocity,
+                         vmin=-max_velocity, vmax=max_velocity,
                          cmap='cmo.balance', show_patch_edges=True)
diff --git a/polaris/viz/__init__.py b/polaris/viz/__init__.py
@@ -83,11 +83,11 @@ def plot_horiz_field(config, ds, ds_mesh, field_name, out_file_name,
     if 'Time' in field.dims and t_index is None:
         t_index = 0
     if t_index is not None:
-        field = field[t_index, :]
+        field = field.isel(Time=t_index)
     if 'nVertLevels' in field.dims and z_index is None:
         z_index = 0
     if z_index is not None:
-        field = field[:, z_index]
+        field = field.isel(nVertLevels=z_index)
 
     if 'nCells' in field.dims:
         ocean_mask = ds_mesh.maxLevelCell - 1 >= 0