Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: cartesian cutting planes through non-cartesian geometries #4847

Open
wants to merge 18 commits into
base: main
Choose a base branch
from

Conversation

chrishavlin
Copy link
Contributor

@chrishavlin chrishavlin commented Mar 6, 2024

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')

spherical_geom_vert_slices

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')

cross_section_example

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
Copy link
Contributor Author

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

@chrishavlin chrishavlin marked this pull request as draft March 6, 2024 19:29
@chrishavlin
Copy link
Contributor Author

pre-commit.ci autofix

@@ -72,12 +72,16 @@ def __init__(self, ds, field_parameters, data_source=None):
self.field_parameters.update(data_source.field_parameters)
self.quantities = DerivedQuantityCollection(self)

def _get_selector_class(self):
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Contributor Author

chrishavlin commented Mar 6, 2024

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
Copy link
Member

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
Copy link
Contributor Author

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.

@chrishavlin chrishavlin changed the title [WIP] ENH: mixed coordinate cutting planes [WIP] ENH: carteisian cutting planes through non-cartesian geometries Mar 19, 2024
@chrishavlin chrishavlin changed the title [WIP] ENH: carteisian cutting planes through non-cartesian geometries [WIP] ENH: cartesian cutting planes through non-cartesian geometries Mar 19, 2024
@neutrinoceros
Copy link
Member

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
Copy link
Contributor Author

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
Copy link
Contributor Author

pre-commit.ci autofix

@chrishavlin
Copy link
Contributor Author

@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):

image

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
Copy link
Contributor Author

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 chrishavlin changed the title [WIP] ENH: cartesian cutting planes through non-cartesian geometries ENH: cartesian cutting planes through non-cartesian geometries Mar 29, 2024
@chrishavlin chrishavlin marked this pull request as ready for review March 29, 2024 20:25
@chrishavlin
Copy link
Contributor Author

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

@chrishavlin chrishavlin added the enhancement Making something better label Mar 29, 2024
@neutrinoceros neutrinoceros self-requested a review April 1, 2024 15:25
@neutrinoceros
Copy link
Member

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants