Add functions for computing and plotting transects #144

xylar · 2023-10-28T14:55:01Z

Checklist

Developer's Guide has been updated
API documentation in the Developer's Guide (api.md) has any new or modified class, method and/or functions listed
Documentation has been built locally and changes look as expected
Testing comment in the PR documents testing used to verify the changes

xylar · 2023-10-28T14:57:21Z

Testing

This has been used in testing the branch https://github.com/xylar/polaris/tree/add-isomip-plus-init and produced plots like these:

Sigma coordinate with 10 layers and a thin-film region

Z-star coordinate with 36 layers (and no thin-film region)

cbegeman · 2023-10-31T19:08:55Z

polaris/ocean/viz/transect.py

+    ds_transect = find_transect_levels_and_weights(
+        ds_transect, ds_3d_mesh.layerThickness,
+        ds_3d_mesh.bottomDepth, ds_3d_mesh.maxLevelCell - 1)


Can you explain why triangles (and tripcolor) are preferable to quadrilaterals (and pcolormesh)? I realize that MPAS-Tools is already written for triangles so it would be more expedient to go that route, but I don't understand the motivation for the additional layer of complexity. Quadrilaterals would be nice because you can plot the edges and they are meaningful.

I will need to remind myself about the details but I didn't decide to move to triangles lightly. It was a huge pain in the ass to make the change. So I must have had a good reason that I couldn't use quads anymore that I am not remembering now. I'll look into it and keep you posted. If we can move back to quads, that would be great for a lot of reasons. One is that labeled contours don't work right on triangles.

Looking through the docstrings and such, I don't see why we couldn't use quads. For example,
https://github.com/MPAS-Dev/MPAS-Tools/blob/master/conda_package/mpas_tools/ocean/transects.py#L40-L46
but I don't see why we couldn't use flat or gouraud shading to accomplish the same.

What I recall being tricky was masking out the invalid cells (e.g. below bathymetry). It's straightforward with flat shading -- you just have NaNs for invalid cell values. But for gouraud, you have a problem that you could have an invalid cell surrounded by 4 valid nodes (e.g. because you have a local column of high bathymetry). A way to handle this and still use pcolormesh would be to add mid-point values to each transect segment. This would fix the masking issue for bathymetry and would make the interpolation of field values along the transect slightly more accurate. But it would mean that there would be vertical lines both at intersections of the transect with cell edges and at the midpoint between cell edges. This would make the outlines of quads less meaningful.

We could use the straightforward approach for flat shading and add an midpoint to the transect segments for gouraud shading (presumably discouraging cell outlines in that case).

I can see how this goes at least.

It took me more than a week but I got the transect plotting to work on quads. Showing the edges of the quads still does not have the effect that you were after because the quads do not correspond to MPAS-ocean cells (the cells have to be divided into triangles horizontally and half-levels vertically for interpolation to nodes to work correctly). However, I have stored segments that can be plotted to get MPAS-Ocean layer interfaces and vertical cell boundaries so you can get the effect that you were after.

polaris/ocean/viz/transect.py

The code added here is similar to functionality in MPAS-Tools but it works with data on quads, rather than triangles. Visualization on triangles has proven to have a number of major drawbacks. It is not very intuitive, functionality for matplotlib functions and methods on triangles have fewer features and are less robust than their counterparts for quads, and there does not seem to be support for labeled contours on triangles.

xylar · 2023-11-06T14:15:20Z

@cbegeman, I spent about a week rewriting the transect functionality to use quads. There are a lot of good reasons for us to do this, high among them:

quads are more intuitive
we can plot contours with labels
the matplotlib functionality for plotting fields on quads is a lot more robust, mature, and intuitive than that for triangles
This functionality will be used extensively for the MPAS-Analysis port so it's definitely time well spent.

For now, I have the functionality here in Polaris rather than in MPAS-Tools. This is partly to make it easier to test and make bug fixes. It is also partly to prevent breaking anything that relies on the current MPAS-Tools implementation. Since it's a big change to the API, I don't want to put it into MPAS-Tools until everything that relies on the old API (probably just MPAS-Analysis) has either been updated or declared deprecated.

I will post some new plots with this approach shortly. You can try it out on your case(s) as well. I believe it will fix the issue you pointed out to me on Slack but I haven't tried your branch out myself yet.

cbegeman · 2023-11-06T16:59:10Z

@cbegeman, I spent about a week rewriting the transect functionality to use quads.

Awesome! I'm excited to try this out. I'm having a little bit of trouble because it seems like interp_mpas_to_transect_cells is looking for nVertLevels in ds_transect but my ds_transect does not have that dimension. Can you point me to a branch and task where you have the viz working? I can try to figure it out from there.

xylar · 2023-11-06T17:31:02Z

I'm having a little bit of trouble because it seems like interp_mpas_to_transect_cells is looking for nVertLevels in ds_transect but my ds_transect does not have that dimension.

Two questions:

Did you get your ds_transect from the compute_transect() routine here or from something prior to this update?
Are you trying to call interp_mpas_to_transect_cells() yourself or letting plot_transect() do it for you?

Here is an example of me using compute_transect():
https://github.com/xylar/polaris/blob/add-isomip-plus-init/polaris/ocean/tasks/isomip_plus/init/init.py#L306-L307
and here are 2 different ways I'm using plot_transect():
https://github.com/xylar/polaris/blob/add-isomip-plus-init/polaris/ocean/tasks/isomip_plus/init/init.py#L399-L403
https://github.com/xylar/polaris/blob/add-isomip-plus-init/polaris/ocean/tasks/isomip_plus/init/init.py#L448-L452

xylar · 2023-11-06T17:36:24Z

I'm having a little bit of trouble because it seems like interp_mpas_to_transect_cells is looking for nVertLevels in ds_transect but my ds_transect does not have that dimension.

Taking a look at the code:

polaris/polaris/ocean/viz/transect/vert.py

Lines 202 to 231 in c1efd85

    
           def interp_mpas_to_transect_cells(ds_transect, da): 
        
               """ 
        
               Interpolate a 3D (``nCells`` by ``nVertLevels``) MPAS-Ocean DataArray 
        
               to transect cells 
        
               Parameters 
        
               ---------- 
        
               ds_transect : xarray.Dataset 
        
                   A dataset that defines an MPAS-Ocean transect, the results of calling 
        
                   ``find_transect_levels_and_weights()`` 
        
               da : xarray.DataArray 
        
                   An MPAS-Ocean 3D field with dimensions `nCells`` and ``nVertLevels`` 
        
                   (possibly among others) 
        
               Returns 
        
               ------- 
        
               da_cells : xarray.DataArray 
        
                   The data array interpolated to transect cells with dimensions 
        
                   ``nSegments`` and ``nHalfLevels`` (in addition to whatever 
        
                   dimensions were in ``da`` besides ``nCells`` and ``nVertLevels``) 
        
               """ 
        
               cell_indices = ds_transect.cellIndices 
        
               level_indices = ds_transect.levelIndices 
        
               da_cells = da.isel(nCells=cell_indices, nVertLevels=level_indices) 
        
               da_cells = da_cells.where(ds_transect.validCells) 
        
               return da_cells

it seems like maybe you're passing something for da that doesn't have the expected dimensions (nCells, nVertLevels). I don't see where nVertLevels would be expected in ds_transect and indeed it's also not in mine.

cbegeman · 2023-11-06T18:13:24Z

@xylar Thanks! I forgot that I had vertVelocityTop in there which has dimension nVertLevelsP1. I just changed the dimension and now it works.

I'm a little confused about the vertical lines because I would expect them to change position if they are the boundaries between cells. However, I've tried 3 different line positions so far and I can't get them to change.

xylar · 2023-11-06T18:41:32Z

I'm a little confused about the vertical lines because I would expect them to change position if they are the boundaries between cells.

These vertical lines are from setting cell_boundary_color='black'? I do have to agree that the bars look pretty unexpected. I don't see why there should be 4 so close together.

However, I've tried 3 different line positions so far and I can't get them to change

You're calling compute_transect() again each time you change the line, I guess? Yes, I would expect the locations of cell boundaries to change each time unless they happen to intersect cells the same way each time (e.g. you're moving by exactly an integer number of cells).

cbegeman · 2023-11-06T19:28:14Z

I got the line locations to change by moving diagonally across the domain, so I guess the transects did cut across the cells in similar ways (the lines were perpendicular). Sorry for the false alarm!

xylar · 2023-11-06T20:47:31Z

@cbegeman, thanks. I did a test where I have a transect zigzagging through the ISOMIP+ domain. I do seem to be seeing the right cell boundaries (red) and skipping the internal triangle boundaries (blue dots):

So I think the cell boundaries are working as expected at least in my test.

xylar · 2023-11-07T02:08:49Z

I know what it is. I'm not removing the periodic edges from the mesh and they're contaminating the plot (most likely) and the cell boundaries (for sure)

Polaris' version of the funciton handles periodicity in planar as well as spherical meshes.

cbegeman · 2023-11-07T15:07:11Z

@xylar Should I go ahead and review now?

xylar · 2023-11-07T15:12:04Z

@cbegeman, yes you can try this again. I believe I have handling of periodic boundaries in here correctly now. I added plots to baroclinic channel to test it out:

This seems feasible given the coarse resolution:

polaris/ocean/tasks/baroclinic_channel/viz.py

xylar · 2023-11-07T21:03:21Z

Still working on this...

xylar · 2023-11-07T21:47:13Z

Okay, I think I'm happy with this but open to suggestions:

xylar · 2023-11-07T21:50:18Z

polaris/viz/planar.py

+    if field_name not in ds:
+        raise ValueError(
+            f'{field_name} must be present in ds before plotting.')
+
+    if patches is not None:
+        if patch_mask is None:
+            raise ValueError('You must supply both patches and patch_mask '
+                             'from a previous call to plot_horiz_field()')
+
+    if (transect_x is None) != (transect_y is None):
+        raise ValueError('You must supply both transect_x and transect_y or '
+                         'neither')
+


I just consolidated these at the top because it's a good practice to raise errors early before wasting any time.

cbegeman · 2023-11-08T20:24:31Z

@xylar My understanding is that if you want to plot the interface locations that correspond to the same time slice that the field is derived from, you have to include the mesh stream in the output stream so that those variables are available when you pass the output dataset to compute_transect. For OMEGA, my understanding is that mesh variables will not be included in the output stream. Have you already thought about adding a separate input argument to compute_transect in addition to ds_3d_mesh that is ds.layerThickness.isel(Time=tidx)?

xylar · 2023-11-08T20:33:11Z

I don't really understand yet what Omega thinks of as mesh variables and what are not. Until I understand the boundary, I can't really design around it. For now, I've been assuming anything that is generic to any MPAS mesh can be shared but variables like maxLevelCell, dimensions like nVertLevels and anything else associated with the vertical coordinate can't be separated out.

I'm happy to design the transect for 2 different input files but it's not clear to me that layerThickness is sufficient to accomplish this.

cbegeman · 2023-11-08T20:43:19Z

I don't really understand yet what Omega thinks of as mesh variables and what are not. Until I understand the boundary, I can't really design around it.

Sounds good. We can certainly wait on this.

xylar · 2023-11-08T20:48:40Z

@cbegeman, how would you feel about there being a ds_horiz_mesh parameter and a ds_vert_coord parameter, where the latter has to have:

minLevelCell
maxLevelCell
layerThickness.isel(Time=tidx)
bottomDepth

and can optionally have landIceFraction (or that can be a separate parameter)?

cbegeman · 2023-11-08T21:20:31Z

@xylar I think that makes sense but maybe you're right that we should just wait for OMEGA streams to be worked out so we don't waste time.

Your comment triggered another question though: Can you explain how you're using landIceFraction in compute_transect? It seems like if this is a weight on an ocean field for the interpolation we actually need the grounded land ice fraction which is landIceFraction - landIceFloatingFraction. But your plots look fine so I must be misunderstanding what's going on.

xylar · 2023-11-08T21:26:12Z

I'm actually not using landIceFraction right now. It was used in MPAS-Analysis to make the background light blue where there was ice and brown elsewhere but I didn't try to do that here. So we could remove it for now.

This should allow the horizontal mesh and constituents of the vertical coordinate to come from different data sets as needed.

xylar · 2023-11-09T09:56:43Z

@cbegeman, I added more parameters to compute_transect() and I'm pretty happy with it now. It makes a lot more sense now in baroclinic channel's viz step because layerThickness is used at the final time rather than from the initial condition.

One thing to be careful about is that you now have to subtract 1 from minLevelCell and maxLevelCell as you pass them in.

I removed landIceFraction for now. We can add it back (along with other changes that may be needed) when we want to plot transects for analysis purposes.

xylar · 2023-11-09T16:35:39Z

One weird thing is that, if I plot the transect as originally requested:

it obviously looks quite different from what actually gets plotted:

So it might be better if folks pass the actual coordinates of the transect from ds_transect rather than the ones that were originally requested.

rljacob · 2023-11-09T23:30:07Z

Late but re: separate input for the mesh variables, uxarray has been assuming that the data is in one file and the mesh is in another because we think modeling groups won't put time-invariant things like the mesh info in every output file.

philipwjones · 2023-11-09T23:39:36Z

Some mesh info will be time-dependent - both for wetting/drying and for when we start exercising more of the Lagrangian vertical mesh. For Omega, I suspect that will mean that much of the horz mesh can sit in separate static mesh file, but that some masks and much of the vertical mesh is likely to end up with other time-dependent prognostic/diagnostic vars (in terms of output). We'll know more once we start putting that design doc together.

xylar · 2023-11-11T22:38:11Z

@philipwjones, that's helpful and follows what I was thinking. To support diverse vertical coordinates and processes like glacial isostatic adjustment, bottomDepth, minLevelCell and maxLevelCell (or their Omega equivalents) will all need to be allowed to depend on time.

The solution I have here only assumes (for purposes of transects) that the horizontal mesh is in its own dataset so we should be good.

cbegeman

@xylar Thanks so much for your work on this! I have successfully run the transect routine on a number of test cases. I just have a few minor comments, but feel free to merge when you're happy with it.

docs/developers_guide/ocean/framework.md

polaris/ocean/viz/transect/plot.py

polaris/ocean/tasks/baroclinic_channel/init.py

cbegeman · 2023-11-12T19:32:16Z

polaris/ocean/tasks/baroclinic_channel/viz.py

+        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=9., vmax=13., cmap='cmo.thermal',


Here and below: maybe we want to just define the limits to be the min, max of mpas_field

Yes, probably the best idea.

Remove confusing references to "3D" and indicate that only fields with `nCells` and `nVertLevels` as dimensions can be plotted.

Color the start and end axes of the transects (now off by default). Use the median to compute the x location of the transect. Use vertices to determine the start and end of the transect in y (so the transects doesn't start in the middle of a cell). Use the min and max of temperature to determine `vmin` and `vmax`, rather than hard-coding.

xylar · 2023-11-13T23:44:11Z

I ran the baroclinic channel default test one more time and results were as expected.

xylar · 2023-11-13T23:44:49Z

Thanks for an excellent review, @cbegeman. You had a ton of useful suggestions that made the final result much better.

xylar added enhancement New feature or request ocean Related to ocean tests or analysis labels Oct 28, 2023

xylar force-pushed the add-transect-plots branch 2 times, most recently from a70f1d1 to 8531f70 Compare October 30, 2023 07:44

cbegeman reviewed Oct 31, 2023

View reviewed changes

polaris/ocean/viz/transect.py Outdated Show resolved Hide resolved

xylar added 4 commits November 6, 2023 14:50

Update to mpas_tools 0.27.0

5fb9483

Add functions for computing and plotting transects

9ebaf1e

Switch transect viz to use polaris transects on quads

4ca3362

xylar force-pushed the add-transect-plots branch from 8531f70 to c1efd85 Compare November 6, 2023 14:05

cbegeman mentioned this pull request Nov 6, 2023

Allow patches to be reused in plot_horiz_field() #146

Merged

3 tasks

Handle periodic meshes in mesh_to_triangles()

65474b5

Polaris' version of the funciton handles periodicity in planar as well as spherical meshes.

xylar force-pushed the add-transect-plots branch from a99cb49 to d9b666d Compare November 7, 2023 15:10

cbegeman reviewed Nov 7, 2023

View reviewed changes

polaris/ocean/tasks/baroclinic_channel/viz.py Outdated Show resolved Hide resolved

xylar force-pushed the add-transect-plots branch from d9b666d to 07930eb Compare November 7, 2023 20:38

Add transect plotting to plot_horiz_field()

29f19d3

xylar added 3 commits November 7, 2023 22:46

Add start/end colors on transect axes

186b016

Add transect plots to baroclinic channel

de8d5bf

Add docs

0e2fd48

xylar force-pushed the add-transect-plots branch from 07930eb to 0e2fd48 Compare November 7, 2023 21:46

xylar commented Nov 7, 2023

View reviewed changes

Add more parameters to compute_transect()

dba67c8

This should allow the horizontal mesh and constituents of the vertical coordinate to come from different data sets as needed.

cbegeman approved these changes Nov 12, 2023

View reviewed changes

xylar added 4 commits November 13, 2023 08:40

Clean up docstrings for computing transects

6d58edb

Remove confusing references to "3D" and indicate that only fields with `nCells` and `nVertLevels` as dimensions can be plotted.

Make coloring the start and end of a transect off by default

501360b

Update docs

e269794

xylar merged commit 10afd60 into E3SM-Project:main Nov 13, 2023
4 checks passed

xylar deleted the add-transect-plots branch November 13, 2023 23:44

Add functions for computing and plotting transects #144

Add functions for computing and plotting transects #144

Conversation

xylar commented Oct 28, 2023 • edited Loading

xylar commented Oct 28, 2023 • edited Loading

Testing

Sigma coordinate with 10 layers and a thin-film region

Z-star coordinate with 36 layers (and no thin-film region)

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xylar commented Nov 6, 2023

cbegeman commented Nov 6, 2023 • edited Loading

xylar commented Nov 6, 2023

xylar commented Nov 6, 2023

cbegeman commented Nov 6, 2023

xylar commented Nov 6, 2023 • edited Loading

cbegeman commented Nov 6, 2023

xylar commented Nov 6, 2023

xylar commented Nov 7, 2023

cbegeman commented Nov 7, 2023

xylar commented Nov 7, 2023

xylar commented Nov 7, 2023

xylar commented Nov 7, 2023

Choose a reason for hiding this comment

cbegeman commented Nov 8, 2023

xylar commented Nov 8, 2023

cbegeman commented Nov 8, 2023

xylar commented Nov 8, 2023 • edited Loading

cbegeman commented Nov 8, 2023

xylar commented Nov 8, 2023

xylar commented Nov 9, 2023

xylar commented Nov 9, 2023

rljacob commented Nov 9, 2023

philipwjones commented Nov 9, 2023

xylar commented Nov 11, 2023

cbegeman left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xylar commented Nov 13, 2023

xylar commented Nov 13, 2023

xylar commented Oct 28, 2023 •

edited

Loading

xylar commented Oct 28, 2023 •

edited

Loading

cbegeman commented Nov 6, 2023 •

edited

Loading

xylar commented Nov 6, 2023 •

edited

Loading

xylar commented Nov 8, 2023 •

edited

Loading