Merge pull request #144 from xylar/add-transect-plots

Add functions for computing and plotting transects
E3SM-Project · Nov 13, 2023 · 10afd60 · 10afd60
2 parents e545833 + e269794
commit 10afd60
Show file tree

Hide file tree

Showing 11 changed files with 1,730 additions and 16 deletions.
diff --git a/deploy/default.cfg b/deploy/default.cfg
@@ -24,7 +24,7 @@ geometric_features = 1.2.0
 jigsaw = 0.9.14
 jigsawpy = 0.3.3
 mache = 1.16.0
-mpas_tools = 0.25.0
+mpas_tools = 0.27.0
 otps = 2021.10
 parallelio = 2.6.0
 

diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md
@@ -314,3 +314,22 @@
    vertical.zlevel.compute_z_level_resting_thickness
    vertical.zstar.init_z_star_vertical_coord
 ```
+
+### Visualization
+
+```{eval-rst}
+.. currentmodule:: polaris.ocean.viz
+
+.. autosummary::
+   :toctree: generated/
+
+   compute_transect
+   plot_transect
+   transect.horiz.find_spherical_transect_cells_and_weights
+   transect.horiz.find_planar_transect_cells_and_weights
+   transect.horiz.make_triangle_tree
+   transect.horiz.mesh_to_triangles
+   transect.vert.find_transect_levels_and_weights
+   transect.vert.interp_mpas_to_transect_cells
+   transect.vert.interp_mpas_to_transect_nodes
+
diff --git a/docs/developers_guide/ocean/framework.md b/docs/developers_guide/ocean/framework.md
@@ -422,3 +422,15 @@ density, which is horizontally constant and increases with depth.
 The {py:func}`polaris.ocean.rpe.compute_rpe()` is used to compute the RPE as
 a function of time in a series of one or more output files.  The RPE is stored
 in `rpe.csv` and also returned as a numpy array for plotting and analysis.
+
+## Visualization
+
+The `polaris.ocean.viz` module provides functions for making plots that are
+specific to the ocean component.
+
+The `polaris.ocean.viz.transect` modules includes functions for computing
+({py:func}`polaris.ocean.viz.compute_transect()`) and plotting
+({py:func}`polaris.ocean.viz.plot_transect()`) transects through the ocean
+from a sequence of x-y or latitude-longitude coordinates.  Currently, only
+transects on xarray data arrays with dimensions `nCells` by `nVertLevels` are
+supported.
diff --git a/polaris/ocean/tasks/baroclinic_channel/init.py b/polaris/ocean/tasks/baroclinic_channel/init.py
@@ -8,6 +8,7 @@
 from polaris import Step
 from polaris.mesh.planar import compute_planar_hex_nx_ny
 from polaris.ocean.vertical import init_vertical_coord
+from polaris.ocean.viz import compute_transect, plot_transect
 from polaris.viz import plot_horiz_field
 
 
@@ -162,8 +163,34 @@ def run(self):
         write_netcdf(ds, 'initial_state.nc')
 
         cell_mask = ds.maxLevelCell >= 1
-        plot_horiz_field(ds, ds_mesh, 'temperature',
-                         'initial_temperature.png', cell_mask=cell_mask)
+
         plot_horiz_field(ds, ds_mesh, 'normalVelocity',
                          'initial_normal_velocity.png', cmap='cmo.balance',
                          show_patch_edges=True, cell_mask=cell_mask)
+
+        y_min = ds_mesh.yVertex.min().values
+        y_max = ds_mesh.yVertex.max().values
+        x_mid = ds_mesh.xCell.median().values
+
+        y = xr.DataArray(data=np.linspace(y_min, y_max, 2), dims=('nPoints',))
+        x = x_mid * xr.ones_like(y)
+
+        ds_transect = compute_transect(
+            x=x, y=y, ds_horiz_mesh=ds_mesh,
+            layer_thickness=ds.layerThickness.isel(Time=0),
+            bottom_depth=ds.bottomDepth, min_level_cell=ds.minLevelCell - 1,
+            max_level_cell=ds.maxLevelCell - 1, spherical=False)
+
+        field_name = 'temperature'
+        vmin = ds[field_name].min().values
+        vmax = ds[field_name].max().values
+        plot_transect(ds_transect=ds_transect,
+                      mpas_field=ds[field_name].isel(Time=0),
+                      title=f'{field_name} at x={1e-3 * x_mid:.1f} km',
+                      out_filename=f'initial_{field_name}_section.png',
+                      vmin=vmin, vmax=vmax, cmap='cmo.thermal',
+                      colorbar_label=r'$^\circ$C', color_start_and_end=True)
+
+        plot_horiz_field(ds, ds_mesh, 'temperature', 'initial_temperature.png',
+                         vmin=vmin, vmax=vmax, cmap='cmo.thermal',
+                         cell_mask=cell_mask, transect_x=x, transect_y=y)
diff --git a/polaris/ocean/tasks/baroclinic_channel/viz.py b/polaris/ocean/tasks/baroclinic_channel/viz.py
@@ -3,6 +3,7 @@
 import xarray as xr
 
 from polaris import Step
+from polaris.ocean.viz import compute_transect, plot_transect
 from polaris.viz import plot_horiz_field
 
 
@@ -42,13 +43,40 @@ def run(self):
         ds = xr.load_dataset('output.nc')
         t_index = ds.sizes['Time'] - 1
         cell_mask = ds_init.maxLevelCell >= 1
-        plot_horiz_field(ds, ds_mesh, 'temperature',
-                         'final_temperature.png', t_index=t_index,
-                         cell_mask=cell_mask)
         max_velocity = np.max(np.abs(ds.normalVelocity.values))
         plot_horiz_field(ds, ds_mesh, 'normalVelocity',
                          'final_normalVelocity.png',
                          t_index=t_index,
                          vmin=-max_velocity, vmax=max_velocity,
                          cmap='cmo.balance', show_patch_edges=True,
                          cell_mask=cell_mask)
+
+        y_min = ds_mesh.yVertex.min().values
+        y_max = ds_mesh.yVertex.max().values
+        x_mid = ds_mesh.xCell.median().values
+
+        y = xr.DataArray(data=np.linspace(y_min, y_max, 2), dims=('nPoints',))
+        x = x_mid * xr.ones_like(y)
+
+        ds_transect = compute_transect(
+            x=x, y=y, ds_horiz_mesh=ds_mesh,
+            layer_thickness=ds.layerThickness.isel(Time=t_index),
+            bottom_depth=ds_init.bottomDepth,
+            min_level_cell=ds_init.minLevelCell - 1,
+            max_level_cell=ds_init.maxLevelCell - 1,
+            spherical=False)
+
+        field_name = 'temperature'
+        vmin = ds[field_name].min().values
+        vmax = ds[field_name].max().values
+        mpas_field = ds[field_name].isel(Time=t_index)
+        plot_transect(ds_transect=ds_transect, mpas_field=mpas_field,
+                      title=f'{field_name} at x={1e-3 * x_mid:.1f} km',
+                      out_filename=f'final_{field_name}_section.png',
+                      vmin=vmin, vmax=vmax, cmap='cmo.thermal',
+                      colorbar_label=r'$^\circ$C', color_start_and_end=True)
+
+        plot_horiz_field(ds, ds_mesh, 'temperature', 'final_temperature.png',
+                         t_index=t_index, vmin=vmin, vmax=vmax,
+                         cmap='cmo.thermal', cell_mask=cell_mask, transect_x=x,
+                         transect_y=y)
diff --git a/polaris/ocean/viz/__init__.py b/polaris/ocean/viz/__init__.py
@@ -0,0 +1,2 @@
+from polaris.ocean.viz.transect.plot import plot_transect
+from polaris.ocean.viz.transect.vert import compute_transect
diff --git a/polaris/ocean/viz/transect/__init__.py b/polaris/ocean/viz/transect/__init__.py
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		from polaris.ocean.viz.transect.plot import plot_transect
		from polaris.ocean.viz.transect.vert import compute_transect