Add vertical crossection support

[rajeeja]

#1296

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[philipc2]

Hey @rajeeja

I'm not sure if the implementation here is what we are looking for with vertical cross sections. Let me start with an example.

Below is temperature variable that varies in the vertical direction along nIsoLevelsT

<xarray.UxDataArray 't_isobaric' (n_face: 637604, nIsoLevelsT: 5)> Size: 13MB
[3188020 values with dtype=float32]
Dimensions without coordinates: n_face, nIsoLevelsT
Attributes:
    units:      K
    long_name:  Temperature interpolated to isobaric surfaces defined in t_is...

Plotting one of the levels looks like the following:

Let's say we want to compute a vertical cross section along a line of constant latitude. We need to essentially sample points along that line to create our raster. The resulting plot should look something like this:

Notice this is very similar to us simply doing:

uxds['t_isobaric'].cross_section.constant_latitude(lat).isel(Time=0, nIsoLevelsT=0).plot()

So the result above is similar to the bottom most plot in the stacked vertical cross section. However, instead of plotting the true Grid geometry, we are plotting the resulting raster

I'm happy to look at this more tomorrow when we meet. Here's a great reference too

https://www.ncl.ucar.edu/Applications/transect.shtml

[philipc2]

Here's the code I used to generate the vertical cross section plot.

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt 

lat = 40.0
samples = 150

subset = uxds['t_isobaric'].cross_section.constant_latitude(lat).isel(Time=0)

# Points to Sample
lons = np.linspace(-180, 180, samples)
lats = np.ones_like(lons) * lat

# find which face(s) each point hits
faces_along_lat = subset.uxgrid.get_faces_containing_point(
    np.column_stack((lons, lats)), return_counts=False
)

# collapse to a 1D index: use -1 where no face hits
face_idx = np.array(
    [row[0] if row else -1 for row in faces_along_lat],
    dtype=int
)

# pre‐allocate and fill
n_levels = subset.sizes['nIsoLevelsT']
data = np.full((samples, n_levels), fill_value=np.nan, dtype=subset.dtype)

valid = face_idx >= 0
data[valid, :] = subset.values[face_idx[valid], :]


# wrap in DataArray so you keep coords
da = xr.DataArray(
    data,
    dims=('sample', 'level'),
    coords={
        'sample': np.arange(samples),
        'lon'  : ('sample', lons),
        'lat'  : ('sample', lats),
    },
)

fig, ax = plt.subplots(figsize=(10, 4))
da.plot(x='lon', y='level', ax=ax, cmap='Blues')

ax.set_xlabel("Longitude")
ax.set_ylabel("Isobaric level index")
ax.set_title("Vertical Cross Section at 40 degrees latitude")
ax.set_xlim(-150, -50)

plt.tight_layout()

[rajeeja]

Thanks, I will look into this and try to modify the code tonight, let's talk about it tomorrow when we meet.

[philipc2]

Another example that is closer to the NCL showcase with a countour plot.

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

lat = 40.0
samples = 1000

# extract the cross‐section
subset = uxds['t_isobaric'].cross_section.constant_latitude(lat).isel(Time=0)

# sample points along the line of latitude
lons = np.linspace(-180, 180, samples)
lats = np.full(samples, lat)

# find which face each point hits
faces_along_lat = subset.uxgrid.get_faces_containing_point(
    np.column_stack((lons, lats)), return_counts=False
)
face_idx = np.array([row[0] if row else -1 for row in faces_along_lat], dtype=int)

# pull out the data
n_levels = subset.sizes['nIsoLevelsT']
data = np.full((samples, n_levels), np.nan, subset.dtype)
valid = face_idx >= 0
data[valid, :] = subset.values[face_idx[valid], :]

# grab the real isobaric values
iso_levels = uxds['t_iso_levels'].values  # shape (n_levels,)

da = xr.DataArray(
    data,
    dims=('sample', 'level'),
    coords={
        'lon': ('sample', lons),
        't_iso_levels': ('level', iso_levels),
    },
    name='temperature'
)

Lon, Press = np.meshgrid(lons, iso_levels)

fig, ax = plt.subplots(figsize=(10, 4))
cf = ax.contourf(
    Lon,
    Press,
    data.T,
    levels=20,
    cmap='viridis'
)
cb = fig.colorbar(cf, ax=ax, label='Temperature (K)')

ax.set_xlabel("Longitude (°E)")
ax.set_ylabel("Pressure Level (hPa)")
ax.set_title(f"Vertical Cross‐Section at {lat}°N")
ax.set_xlim(-150, -50)

plt.gca().invert_yaxis()  
plt.tight_layout()
plt.show()

[philipc2]

I'd lean towards not duplicating these methods here, since the core idea of a constant_latitude or constant_longitude cross-section is independent of any leading/vertical dimensions.

Instead, maybe we could adding the following parameters to enable us to sample the resulting spatial cross section along the transect, with some number of samples like in the example I shared above.

# default cross section, returns a UxDataArray with an associated Grid
uxds['t2m'].cross_section.constant_latitude(lat=45.0)

# one example, we can brainstorm other ideas 
uxds['t2m'].cross_section.constant_latitude(lat=45.0, n_samples=100)

[erogluorhan]

Implementation for vertical cross-sections looks great, but I think I am a bit worried about the API direction this PR proposes:

• Changing the API of the user functions uxda.cross_section.constant_latitude() and uxda.cross_section.constant_longitude() significantly.
  • The semantic seems to be changing. We have created and advertised these functions as them creating brand new subset grids so far during tutorials, etc. There is knowledge based on that and there might be many workflows created using them in that way. Now changing the semantics not only needs to be justified and not be right away, but also would break many UXarray workflows.
    • Related to this, return types being converted to xarray
• I think there should still be a difference between how we implemented (horizontal) cross-sections so far (i.e. in the way described above that generates a subset grid) and how vertical cross-sections are meant to be (structured constant lat/lon data through sampling for mainly visual inspection purposes)
  • I don't know if aiming to implement them all together in the same API at the cost of significant changes to the existing API is a good idea. Even though all of them are called "cross-sections", there is actually a usage difference between the existing ones and the vertical cross-section I believe, e.g. one can run stats on the output of horizontal cross-sections we used so far, but not on the output of what this PR implements I believe (at least no stats on the actual grid topology).

[philipc2]

@erogluorhan

Thanks for the feedback. We've discussed a few of these points during our co-working

• Existing cross_section functionality will be preserved and left unchanged
• To enable visualuzations of vertical cross sections, sampling along the transect needs to be done, which can use our get_faces_containing_point(), similar to how we rasterize for MPl/Cartopy
• This added functionality will be added to the existing cross_section methods through additional parameters, but we are still thinking of the best approach for it.

[erogluorhan]

• This added functionality will be added to the existing cross_section methods through additional parameters, but we are still thinking of the best approach for it.

While this could help much for the sake of clear code and API, the difference of purposes between those existing and newer functions makes me hesitant about this, i.e. accessing two different purposes through the same API with parameters might not be the best option.

Also, the same functions returning different objects (either UxDataArray or xarray.DataArray) depending on an optional parameter could create a lot of confusions with the usage.

[erogluorhan]

• This added functionality will be added to the existing cross_section methods through additional parameters, but we are still thinking of the best approach for it.

While this could help much for the sake of clear code and API, the difference of purposes between those existing and newer functions makes me hesitant about this, i.e. accessing two different purposes through the same API with parameters might not be the best option.

Also, the same functions returning different objects (either UxDataArray or xarray.DataArray) depending on an optional parameter could create a lot of confusions with the usage.

How about keeping the existing functions exactly the same but adding a new one for the vertical only, which would look something like this:

uxda.cross_sections.vertical(... , constant_dim=lat, dim_value=0, other_dim_range=(min,max,step), ...)

• This would work for both lat & lon through the same API
  • A similar way would even be done for the existing cross-sections, i.e. instead of having separate constant_dim and constant_dim_interval functions, we'd have a single API that'd have params for dim selections and interval vs single value selections?
• step in other_dim_range instead of number of samples to give better control for resolution of samples and have consistency with other API such as zonal_mean etc.
  • I might be missing some user insights about the term "number of samples" though

[philipc2]

How about keepag the existing functions exactly the same but adding a new one for the vertical only, which would look something like this

My main concern with this is that our existing cross_section operators already preserve vertical dimensions. Having a separate "vertical" method doesn't make much sense to me.

For vertical cross section analysis (but not exclusively, also for temporal or other cross sections), it necessary to sample along the line of intersection as opposed to directly return a brand new grid. This allows us to then draw contour or other plots along the sample spatial dimension such as the one below.

I'm still in favor of including parameters to dictate the output. With the examples above, there would be two types of cross sections.

Unstructured (Current Implementation)

• A cross section is performed along some transect (i.e. constant latitude)
• A UxDataArray with a valid Grid is returned, with data remaining on the unstructured grid elements

Structured/Sampled

• Data is sampled along the transect, such that the output is no longer tied to any unstructured grid elements
• Output would be a Xarray.DataArray, since we no longer would attach a grid and analysis can be performed without implementing new functionality

[philipc2]

Also, the same functions returning different objects (either UxDataArray or xarray.DataArray) depending on an optional parameter could create a lot of confusions with the usage.

I'm not too concerned about this, as long as:

• Additional parameters do not impact existing workflows
• Proper Typing and Docstrings
• Well-written user guide and example notebooks

Once an unstructured grid is sampled, such as onto a structured grid, we should no longer return UXarray objects, it should fall back to Xarray so that we do not need to reimplement things.

[rajeeja]

The Notebook needs to be rethought completely to make a distinction on how to call for horizontal and vertical cross sections with the current implementation it does get a little tricky, but I agree with Philip that it will be simpler and more maintainable if we use a common interface for both.

[philipc2]

Pinging @brianpm for any feedback on these, especially regarding any expectations from the research/science applications for these types of workflows.

Using the implementation in #1321, creating vertical (and also temporal) cross section plots is seamless since we can utilize existing Xarray functionality.

Here's the data variable I'm using (taken from @rytam2's E3SM dataset) in our docs (which can be found here for download)

<xarray.UxDataArray 'T' (time: 11, lev: 72, n_face: 21600)> Size: 68MB
dask.array<concatenate, shape=(11, 72, 21600), dtype=float32, chunksize=(1, 72, 21600), chunktype=numpy.ndarray>
Coordinates:
  * lev      (lev) float64 576B 0.1238 0.1828 0.2699 ... 986.2 993.8 998.5
  * time     (time) object 88B 0001-02-01 00:00:00 ... 0005-02-01 00:00:00
Dimensions without coordinates: n_face
Attributes:
    mdims:          1
    units:          K
    long_name:      Temperature
    standard_name:  air_temperature
    cell_methods:   time: mean

# Compute a interpolated cross section by sampling the line of constant latitude
lat_cs = uxds['T'].compute().cross_section.constant_latitude(lat=45.0, interpolate=True)

# vertical cross section
lat_cs.isel(time=0).plot(cmap='inferno')

# vertical cross section contour 
lat_cs.isel(time=0).plot.contourf(levels=10, cmap='inferno')

# temporal cross section
lat_cs.isel(lev=0, time=[0, 1, 2, 3, 4]).plot(cmap='inferno')

[erogluorhan]

My main concern with this is that our existing cross_section operators already preserve vertical dimensions. Having a separate "vertical" method doesn't make much sense to me.

This looks like a difference in perspective and it's understandable. We apparently need input from the actual users for the right perspective. To me though, the cross-section is designed to subset data on different references, and the existing functions are only either zonal or meridional with all those vertical levels still present (at least with the design so far), and despite the current understanding, they are not "horizontal". Vertical levels being preserved does not make the either horizontal or vertical.

For vertical cross section analysis (but not exclusively, also for temporal or other cross sections), it necessary to sample along the line of intersection as opposed to directly return a brand new grid. This allows us to then draw contour or other plots along the sample spatial dimension such as the one below.

And, because of this distinction, I don't think they should all be part of the same API. They are different in nature than the existing functions.

I'm still in favor of including parameters to dictate the output. With the examples above, there would be two types of cross sections.

Unstructured (Current Implementation)

• A cross section is performed along some transect (i.e. constant latitude)
• A UxDataArray with a valid Grid is returned, with data remaining on the unstructured grid elements

Structured/Sampled

• Data is sampled along the transect, such that the output is no longer tied to any unstructured grid elements
• Output would be a Xarray.DataArray, since we no longer would attach a grid and analysis can be performed without implementing new functionality

This cross-sections definitions make total sense to me, but I believe designing all of them in the same API with the help of optional parameters, which most of the time is easy to be overseen by the user, creates a lot of confusion at the user's end.

[erogluorhan]

Also, the same functions returning different objects (either UxDataArray or xarray.DataArray) depending on an optional parameter could create a lot of confusions with the usage.

I'm not too concerned about this, as long as:

• Additional parameters do not impact existing workflows
• Proper Typing and Docstrings
• Well-written user guide and example notebooks

Due to my comments above, I am still concerned about this way forward, considering the potential confusions for the end user of these functions.

Once an unstructured grid is sampled, such as onto a structured grid, we should no longer return UXarray objects, it should fall back to Xarray so that we do not need to reimplement things.

I agree, but not from the same API through optional params

[erogluorhan]

Due to my comments above, I am still very concerned about this way forward, considering the potential confusions for the end user of these functions.

Let's all try to collect as much user input on this possible then reunite to discuss.

[philipc2]

Thanks for the feedback @erogluorhan

This implementation from MetPy could serve as a good design reference.

[philipc2]

This cross-sections definitions make total sense to me, but I believe designing all of them in the same API with the help of optional parameters, which most of the time is easy to be overseen by the user, creates a lot of confusion at the user's end.

While I understand the concern here, I don't believe that introducing additional parameters to extend functionality would lead to more confusion than introduction brand-new methods, which would be very similar to existing ones.

This looks like a difference in perspective and it's understandable. We apparently need input from the actual users for the right perspective. To me though, the cross-section is designed to subset data on different references, and the existing functions are only either zonal or meridional with all those vertical levels still present (at least with the design so far), and despite the current understanding, they are not "horizontal". Vertical levels being preserved does not make the either horizontal or vertical.

From what I've seen in the NCL and MetPy examples, generally speaking a cross-section is defined by taking some transect defined by two arbitrary lat-lon points and obtaining the data that falls between the path formed by connecting them.

The ideal of vertical or temporal cross sections is then done after the fact by selecting the dimension you'd want to compare along the intersection, such as my example above for the Level and Time dimensions.

Vertical levels being preserved does not make the either horizontal or vertical.

That's correct, but preserving vertical (or other) levels does enable us to perform vertical cross-section analysis, since we would have reduced our data from the grid geometry onto the samples along the transect.

The great part of using Xarray for our implementation is that dimensions are preserved, and the user can dictate which ones they want to keep or drop, or which ones they'd like to select for plotting. I don't believe explicitly handling these through our implementation is necessary.

[philipc2]

And, because of this distinction, I don't think they should all be part of the same API. They are different in nature than the existing functions.

Assuming that we do not make an explicit distinction between vertical, temporal, or other types of applications of cross sections, I'm not sure what the best way to split up the grid-based and interpolated cross sections would be.

For example, we currently have:

• cross_section.constant_latitude(lat)
• cross_section.constant_longitude(lat)

We could add additional methods to the existing cross_section accessor, such as:

• cross_section.interpolated_constant_latitude(lat)
• cross_section.interpolated_constant_longitude(lat)

Or, as I've put together in #1321, we could dictate the output by using an interpolate parameter

• cross_section.constant_latitude(lat, interpolate=True)
• cross_section.constant_longitude(lat, interpolate=True)

[erogluorhan]

Hey @falkojudt , we were able to come up with a vertical cross section implementation here but are having a hard time deciding how to design its function signatures, i.e. to add into existing functions or create new ones etc (Please see th conversation in this PR). From a meteorological use perspective, please let us know if you'd have any suggestions.

[clyne]

Just wanted to chime in with a few things to consider.

• There is a subtle distinction between 3D to 2D dimensionality reduction via indexing (slicing in Xarray parlance) and interpolation to a plane.

• For structured data, an axis-aligned plane can be constructed via indexing. For an arbitrarily oriented plane, the plane needs to be defined in space and the data that intersect it resampled onto that plane. The sampling is typically performed along a regular grid with some sensible coordinate system within the plane.

• For unstructured data, interpolation onto a plane is generally the only option for going from 3D to 2D. The caveat with climate & weather unstructured grids is that they are structured along the vertical axis. So slicing still makes sense along Z.

• Xarray supports slicing via isel() and interpolation via interp(). These methods are well understood by Xarray users. It makes sense to leverage these when developing capability for UXarray. However, convenience functions would be useful for specialized cases such as a vertical cross section with a plane with constant lat or lon

[erogluorhan]

Assuming that we do not make an explicit distinction between vertical, temporal, or other types of applications of cross sections, I'm not sure what the best way to split up the grid-based and interpolated cross sections would be.

For example, we currently have:

• cross_section.constant_latitude(lat)
• cross_section.constant_longitude(lat)

I think having these functions designed as the main cross-section functionality at the time is giving us the actual hardship now re design. Maybe these current functions were/are more suitable for .subset or some kind of xarray-like slicing API rather than "cross-sections" (See @clyne 's previous comment).

Also, we had a good conversation at today's Raijin monthly meeting and there were similar thoughts there; for instance, cross-sections should never return a new grid topology and rather only return some form of xarray output, and those @clyne mentioned. (@brianpm @clyne @lantaosun please chime in if you think more on that)

Maybe we need to consider moving away from providing these current functions under cross_section? e.g. subset.constant_latitude(lat) etc?

We could add additional methods to the existing cross_section accessor, such as:

• cross_section.interpolated_constant_latitude(lat)
• cross_section.interpolated_constant_longitude(lat)

Or, as I've put together in #1321, we could dictate the output by using an interpolate parameter

• cross_section.constant_latitude(lat, interpolate=True)
• cross_section.constant_longitude(lat, interpolate=True)

If we decide to perform the above suggestion re current functionality, I think we'd have a couple options:

• Either, there might not be any need for adding additional methods here, e.g. interpolated_constant_latitude(lat) etc. Rather we'd implement cross_section.constant_longitude(lat) as interpolate=True by default?
• Or, @clyne's suggestion, we might implement these interpolated cross-sections through an interpolate interface?

All that said, I still think that combining the existing and new ones in the same API through optional parameters (e.g. interpolate=True) might not be the best idea.

[philipc2]

I agree with @erogluorhan that the current implementation of cross_sections would be a good fit within the subset family of methods.

[philipc2]

Xarray supports slicing via isel() and interpolation via interp(). These methods are well understood by Xarray users. It makes sense to leverage these when developing capability for UXarray. However, convenience functions would be useful for specialized cases such as a vertical cross section with a plane with constant lat or lon

Or, @clyne's suggestion, we might implement these interpolated cross-sections through an interpolate interface?

I like this idea.

# Subset the original data to include elements that intersect a line of constant latitude
data_along_lat = uxds['T'].subset.constant_latitude(lat=45.0)

# Interpolate the result onto a structured Xarary dataset
data_along_lat_interpolated = data_along_lat.interp(lon=np.linspace(-180, 180, 10)

Right now, interp() doesn't work on grid dimensions (n_node, n_edge or n_face), but we could consider extending the functionality to perform the sampling that we put together as part of this implementation.

[falkojudt]

Here are my thoughts:

• The MetPy example that Philipp mentioned is probably closest to what we need:
  https://unidata.github.io/MetPy/latest/api/generated/metpy.interpolate.cross_section.html

• One important feature is the ability to define cross sections in arbitrary directions—not just strictly zonal or meridional—as shown in the MetPy example.

• Regarding the vertical dimension: while slicing along model levels is straightforward, we’re typically more interested in height levels (e.g., meters above ground/sea level) or pressure levels. This would likely require interpolation. (One issue with using model levels directly—especially when the vertical coordinate is terrain-following—is that the resulting cross section can appear misleading. For example, over mountainous terrain, a model-level slice might look flat, whereas a cross section in height or pressure space would clearly show the underlying topography.)

[philipc2]

One important feature is the ability to define cross sections in arbitrary directions—not just strictly zonal or meridional—as shown in the MetPy example.

This is something we can implement.

Regarding the vertical dimension: while slicing along model levels is straightforward, we’re typically more interested in height levels (e.g., meters above ground/sea level) or pressure levels. This would likely require interpolation. (One issue with using model levels directly—especially when the vertical coordinate is terrain-following—is that the resulting cross section can appear misleading. For example, over mountainous terrain, a model-level slice might look flat, whereas a cross section in height or pressure space would clearly show the underlying topography.)

Is this something that can already be done easily with Xarray?

[falkojudt]

Is this something that can already be done easily with Xarray?

I am not sure, hopefully someone else would know.

[rajeeja]

Thanks for the great feedback and discussion here. We are now not implement vertical crossection as a separate method. Subsetting will handle the grid informed crossection that will return the data on the original grid elements, while, the crossection method will sample data along intersection. Closing this infavor of #1321