ENH: cartesian cutting planes through non-cartesian geometries

[chrishavlin]

This PR introduces a "Cartesian Cutting Plane", which allows you to define a cutting plane in cartesian coordinates and extract a fixed resolution buffer from the intersection with a non-cartesian 3D dataset. This PR only enables functionality for data in spherical coordinates, but it should be straightforward to apply to other 3D non-cartesian coordinates. Visualization is limited to constructing fixed resolution buffers at present.

Examples

See the included new example notebook for more examples, here is a direct link to the rendered file: link.

But here's a couple:

import numpy as np 
import matplotlib.pyplot as plt 
import yt 
import unyt
from yt.geometry.coordinates.spherical_coordinates import spherical_to_cartesian
from yt.testing import fake_amr_ds 


# data set up: load some data in spherical coordinates, add a field
# for calculating the cartesian z position of cells.

def _z(field, data):
    r = data['index', 'r']
    theta = data['index', 'theta']
    phi = data['index', 'phi']
    _, _, z = spherical_to_cartesian(r, theta, phi)
    return unyt.unyt_array(z, r.units)

ds = fake_amr_ds(geometry='spherical')

ds.add_field(name=('index', "z_val"),
                 function=_z, 
                 sampling_type='cell',                  
                 take_log=False)

# construct our cutting plane, same syntax as `ds.cutting`. 
# the following plane is parallel to the x-z plane 
# and offset to y = 0.6
normal = np.array([0., 1., 0.])
center = np.array([0.0, 0.6, 0.0])
slc = ds.cartesian_cutting(normal, center)

# extract a fixed resolution buffer and make some plots
# (widths here are define as in-plane distances).  
frb = slc.to_frb(2.0, 400)

flds = [('index', 'z_val'), ('stream', 'Density')]
fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(8,3))

for ifld, fld in enumerate(flds):
    vals = frb[fld]
    vals[~frb.get_mask(fld)] = np.nan
    im = axs[ifld].imshow(vals,
               extent=frb.bounds, 
               origin='lower', 
               cmap='plasma',               
              )
    c = plt.colorbar(im)
    c.set_label(fld)
plt.tight_layout()    
fig.savefig('spherical_geom_vert_slices.png')

note that the curvature in the z values in the left plot arise because we're plotting the z value of cell centers.

One of my favorites applications is that this PR allows you to easily create cross sections!

# create a sub-region to sample from 
left_edge = ds.arr([0.5, 30*np.pi/180, 20*np.pi/180], 'code_length')
right_edge = ds.arr([0.9, 60*np.pi/180, 40*np.pi/180], 'code_length')
c = (left_edge + right_edge) / 2.0

region = ds.region(c, left_edge, right_edge)

# set the start and end of the cross section.
pt1 = np.array([0.9, left_edge[1].d, c[2].d])
pt2 = np.array([0.9, right_edge[1].d, c[2].d])

# construct a plane passing through those start and end points
pts = np.column_stack([pt1, pt2, c.d])
x, y, z = spherical_to_cartesian(pts[0,:], pts[1,:], pts[2,:])

normal = -np.cross((x[0], y[0], z[0]), (x[1], y[1], z[1]) )
center = np.array([x[2], y[2], z[2]])
north_vector = center # orient to center so that it appears "up" in the plot

slc = ds.cartesian_cutting(normal, center, north_vector=north_vector, data_source=region)
frb = slc.to_frb(0.5, 400)

vals = frb['stream', 'density']
vals[~frb.get_mask(('stream', 'density'))] = np.nan

fig, axs = plt.subplots(1)
im = axs.imshow(vals, 
           extent=frb.bounds, 
           origin='lower', 
           cmap='plasma',           
          )
axs.axes.axes.set_axis_off()
c = plt.colorbar(im)
fig.savefig('cross_section_example.png')

See the notebook for some more cases.

Some notes

I've gone through several iterations of this the past couple months, many dead ends and many commits (most of which I squashed before PR creation). So here are some extra notes with comments on possible alternate approaches:

design choices

• (this point no longer relevant, keeping the separate data object) The new selection object (YTCuttingPlaneMixedCoords YTCartesianCuttingPlane) is accessed from ds.cutting_mixed ds.carteisan_cutting. @matthewturk and I chatted a bit about whether there's much utility in the current ds.cutting for non-cartesian geometries -- it may make more sense to have this PR simply change how ds.cutting works with non-cartesian geometries rather than introduce a new data object. But I thought initial review would be easier with a separate selection object. If others agree I can refactor to remove ds.cartesian_cutting and tie in the new behavior to ds.cutting.
• I found myself re-writing the spherical-to-cartesian transformation in multiple places so at some point refactored so that I could define it in a single location.
• I included a public function yt.utilities.lib.coordinate_utilities.cartesian_bboxes() that is not explicitly used in the rest of the PR (it calls functions that are used by the rest of the PR). The main reason is so that I can use it over in yt_idv for pre-computing bounding boxes in our spherical-volume rendering attempts over there.

limitations and future PRs

• I haven't managed to get this working with yt's PlotWindow interface due to all the validations for axis names and coordinate centers, so at present the YTCartesianCuttingPlane only implements a to_frb and not a to_pw, so some manual plotting is required. There's enough to review here (and IMO enough useful functionality) that I'd rather leave any PlotWindow integration to a future PR...
• enabling this for geographic coordinates should be trivial, 3-d cylindrical would be a bit more effort but do-able.

Related PRs

This PR probably replaces #4752 (the present PR actually borrowed one early commit from there, which I've kept in the commit history here. I ended up pretty much re-writing the code changes related to that commit, so while I could scrub it from the commit history here, I like having it as it reflects the development history more accurately. but if it's confusing, I can rebase here).

This PR partially addresses #4750 but until the work here fully integrates with yt's plotting interface I think that PR should stay open.

[chrishavlin]

mainly a WIP because I found a small bug with an edge case while writing up examples, but feel free to jump in.

[chrishavlin]

pre-commit.ci autofix

[chrishavlin]

note to reviewers: this change was needed so that the selector that is chosen can depend on the geometry. this method gets over-ridden in the new selector object.

[matthewturk]

I don't remember if this was something I experimented with or not, but it seems like a good idea. Although, I wonder if we could just wrap it in the selector property itself?

[chrishavlin]

one final note to self so i don't forget: the edge case currently failing is in bbox intersection with a spherical shell and the plane falling on one side of all the edges of a large spherical element (kinda like an extreme version of the cusp problem).

[neutrinoceros]

I love it ! One quick comment now is that I'm not sure about the name: I don't get it at first, and when I do I find myself expecting more than it really is. In other words, "mixed coordinate cutting planes" is confusing to me. Would "cartesian cutting planes" be an accurate description ?

[chrishavlin]

One quick comment now is that I'm not sure about the name: I don't get it at first, and when I do I find myself expecting more than it really is. In other words, "mixed coordinate cutting planes" is confusing to me. Would "cartesian cutting planes" be an accurate description ?

Yup, I'm not crazy about the name either actually, it's just where I ended up when this got to the point where it was ready to initially share :)

"cartesian cutting planes" is accurate. Maybe ds.cartesian_cutting_plane would certainly be clearer if a little verbose.

Also, I did go back and forth on the best way to expose the new functionality given that ds.cutting and ds.slice already exist. Initially, I hooked it into the existing ds.cutting with a new keyword argument: ds.cutting(..., slice_on_index: bool = True) but the code got a little harder to follow so I split it into it's own dedicated selection object while I got it working. But at this point, I could go back to that design, though I'd want to come up with a better keyword variable name: maybe a plane_geometry str argument? i.e., ds.cutting(...., plane_geometry='cartesian') would give you a cartesian cutting plane through a non-cartesian geometry.

[neutrinoceros]

Something like ds.cutting(...., plane_geometry='cartesian') would immediately suggest a whole new realm of possibilities that we don't (and likely won't) support. So, I think having a separate new method is fine, and actually our best option.

[chrishavlin]

renamed to YTCartesianCuttingPlane and ds.cartesian_cutting. also fixed the large-element edge case that was causing grid selection to fail in certain cases.

[chrishavlin]

pre-commit.ci autofix

[chrishavlin]

@mathewturk asked in passing whether the intersections match what you'd get from calculating plane-sphere intersections... and then I went on a deep dive over in this notebook to check. The notebooks calculates the area of intersection for the cartesian cutting plane from a fixed res buff of the ('index', 'ones') field and tests a number of different parameters that might affect accuracy. It does a pretty good job. The biggest control is the resolution of the fixed res buffer, here's a plot of the calculated and theoretical area of intersection for a plane passing through the origin (so that the area of intersection should be pi given dataset r bounds of 0,1):

so as long as you're using a frb with resolution above ~100 in each dimension it'll do a decent job. The notebook has some other tests in it to verify things are working as expected (varying the plane position, varying the dataset bounds).

[chrishavlin]

oh and the last push adds an image test, so that test suite will fail... i'll open a PR to the image answer repo later on.

[chrishavlin]

feeling pretty good about where this is at! so just switched out of draft mode.

[neutrinoceros]

Requesting a review from myself so I don't forget about it, but I can't make any promises as to when I'll be able to actually review this massive PR. Hopefully sometimes this month.