ENH: cartesian cutting planes through non-cartesian geometries

.gitignore

@@ -39,6 +39,7 @@ yt/utilities/lib/bitarray.c
 yt/utilities/lib/bounded_priority_queue.c
 yt/utilities/lib/bounding_volume_hierarchy.c
 yt/utilities/lib/contour_finding.c
+yt/utilities/lib/coordinate_utilities.c
 yt/utilities/lib/cykdtree/kdtree.cpp
 yt/utilities/lib/cykdtree/utils.cpp
 yt/utilities/lib/cyoctree.c

nose_ignores.txt

@@ -1,6 +1,8 @@
 --ignore-file=test_add_field\.py
 --ignore-file=test_ambiguous_fields\.py
 --ignore-file=test_callable_grids\.py
+--ignore-file=test_cartesian_cutting_plane\.py
+--ignore-file=test_cartesian_cutting_plane_selector\.py
 --ignore-file=test_callbacks_geographic\.py
 --ignore-file=test_commons\.py
 --ignore-file=test_cython_fortran_utils\.py

tests/tests.yaml

@@ -172,6 +172,8 @@ other_tests:
      - "--ignore-file=test_add_field\\.py"
      - "--ignore-file=test_ambiguous_fields\\.py"
      - "--ignore-file=test_callable_grids\\.py"
+     - "--ignore-file=test_cartesian_cutting_plane\\.py"
+     - "--ignore-file=test_cartesian_cutting_plane_selector\\.py"
      - "--ignore-file=test_callbacks_geographic\\.py"
      - "--ignore-file=test_commons\\.py"
      - "--ignore-file=test_cython_fortran_utils\\.py"

yt/data_objects/selection_objects/data_selection_objects.py

@@ -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):
+        s_module = getattr(self, "_selector_module", yt.geometry.selection_routines)
+        sclass = getattr(s_module, f"{self._type_name}_selector", None)
+        return sclass
+
     @property
     def selector(self):
         if self._selector is not None:
             return self._selector
-        s_module = getattr(self, "_selector_module", yt.geometry.selection_routines)
-        sclass = getattr(s_module, f"{self._type_name}_selector", None)
+        sclass = self._get_selector_class()
         if sclass is None:
             raise YTDataSelectorNotImplemented(self._type_name)
 

yt/data_objects/selection_objects/slices.py

@@ -15,6 +15,8 @@
     validate_object,
     validate_width_tuple,
 )
+from yt.geometry import selection_routines
+from yt.geometry.geometry_enum import Geometry
 from yt.utilities.exceptions import YTNotInsideNotebook
 from yt.utilities.minimal_representation import MinimalSliceData
 from yt.utilities.orientation import Orientation
@@ -169,7 +171,6 @@ class YTCuttingPlane(YTSelectionContainer2D):
     data_source: optional
         Draw the selection from the provided data source rather than
         all data associated with the dataset
-
     Notes
     -----
 
@@ -233,33 +234,42 @@ def __init__(
     def normal(self):
         return self._norm_vec
 
+    def _current_chunk_xyz(self):
+        x = self._current_chunk.fcoords[:, 0]
+        y = self._current_chunk.fcoords[:, 1]
+        z = self._current_chunk.fcoords[:, 2]
+        return x, y, z
+
     def _generate_container_field(self, field):
         if self._current_chunk is None:
             self.index._identify_base_chunk(self)
         if field == "px":
-            x = self._current_chunk.fcoords[:, 0] - self.center[0]
-            y = self._current_chunk.fcoords[:, 1] - self.center[1]
-            z = self._current_chunk.fcoords[:, 2] - self.center[2]
+            x, y, z = self._current_chunk_xyz()
+            x = x - self.center[0]
+            y = y - self.center[1]
+            z = z - self.center[2]
             tr = np.zeros(x.size, dtype="float64")
             tr = self.ds.arr(tr, "code_length")
             tr += x * self._x_vec[0]
             tr += y * self._x_vec[1]
             tr += z * self._x_vec[2]
             return tr
         elif field == "py":
-            x = self._current_chunk.fcoords[:, 0] - self.center[0]
-            y = self._current_chunk.fcoords[:, 1] - self.center[1]
-            z = self._current_chunk.fcoords[:, 2] - self.center[2]
+            x, y, z = self._current_chunk_xyz()
+            x = x - self.center[0]
+            y = y - self.center[1]
+            z = z - self.center[2]
             tr = np.zeros(x.size, dtype="float64")
             tr = self.ds.arr(tr, "code_length")
             tr += x * self._y_vec[0]
             tr += y * self._y_vec[1]
             tr += z * self._y_vec[2]
             return tr
         elif field == "pz":
-            x = self._current_chunk.fcoords[:, 0] - self.center[0]
-            y = self._current_chunk.fcoords[:, 1] - self.center[1]
-            z = self._current_chunk.fcoords[:, 2] - self.center[2]
+            x, y, z = self._current_chunk_xyz()
+            x = x - self.center[0]
+            y = y - self.center[1]
+            z = z - self.center[2]
             tr = np.zeros(x.size, dtype="float64")
             tr = self.ds.arr(tr, "code_length")
             tr += x * self._norm_vec[0]
@@ -353,6 +363,7 @@ def to_frb(self, width, resolution, height=None, periodic=False):
         >>> frb = cutting.to_frb((1.0, "pc"), 1024)
         >>> write_image(np.log10(frb["gas", "density"]), "density_1pc.png")
         """
+
         if is_sequence(width):
             validate_width_tuple(width)
             width = self.ds.quan(width[0], width[1])
@@ -363,8 +374,150 @@ def to_frb(self, width, resolution, height=None, periodic=False):
             height = self.ds.quan(height[0], height[1])
         if not is_sequence(resolution):
             resolution = (resolution, resolution)
+
         from yt.visualization.fixed_resolution import FixedResolutionBuffer
 
         bounds = (-width / 2.0, width / 2.0, -height / 2.0, height / 2.0)
         frb = FixedResolutionBuffer(self, bounds, resolution, periodic=periodic)
         return frb
+
+
+class YTCartesianCuttingPlane(YTCuttingPlane):
+    """
+    A YTCartesianCuttingPlane (ds.cartesian_cutting) is similar to YTCuttingPlane (ds.cutting)
+    but the cutting plane is always defined in cartesian coordinates, allowing arbitrary slices
+    through datasets defined in non-cartesian geometries.
+
+    Parameters
+    ----------
+    normal : array_like
+        The vector that defines the desired plane in cartesian coordinates.
+    center : array_like
+        The center of the cutting plane, where the normal vector is anchored, in
+        cartesian coordinates.
+    north_vector: array_like, optional
+        An optional vector to describe the north-facing direction in the resulting
+        plane, in cartesian coordinates.
+    ds: ~yt.data_objects.static_output.Dataset, optional
+        An optional dataset to use rather than self.ds
+    field_parameters : dictionary
+         A dictionary of field parameters than can be accessed by derived
+         fields.
+    data_source: optional
+        Draw the selection from the provided data source rather than
+        all data associated with the dataset
+    edge_tol: float
+        Optional edge tolerance (default 1e-12). This controls the fuzziness
+        of element-plane intersection to account for floating point errors
+        in coordinate transformations. If your slice is missing elements,
+        try increasing this number a bit.
+    """
+
+    _type_name = "cartesian_cutting"
+    _con_args = ("normal", "center")
+    _tds_attrs = ("_inv_mat",)
+    _tds_fields = ("x", "y", "z", "dx")
+    _container_fields = ("px", "py", "pz", "pdx", "pdy", "pdz")
+    _supported_geometries = (Geometry.SPHERICAL,)
+
+    def __init__(
+        self,
+        normal,
+        center,
+        north_vector=None,
+        ds=None,
+        field_parameters=None,
+        data_source=None,
+        edge_tol=1e-12,
+    ):
+        super().__init__(
+            normal,
+            center,
+            north_vector=north_vector,
+            ds=ds,
+            field_parameters=field_parameters,
+            data_source=data_source,
+        )
+        self._ds_geom = self.ds.geometry
+        self._validate_geometry()
+        self.edge_tol = edge_tol
+
+    def _validate_geometry(self):
+        if self._ds_geom not in self._supported_geometries:
+            if self._ds_geom is Geometry.CARTESIAN:
+                msg = (
+                    "YTCuttingPlaneMixedCoords is not supported for cartesian "
+                    "coordinates: use YTCuttingPlane instead (i.e., ds.cutting)."
+                )
+                raise NotImplementedError(msg)
+            else:
+                self._raise_unsupported_geometry()
+
+    def _raise_unsupported_geometry(self):
+        msg = (
+            "YTCuttingPlaneMixedCoords only supports the following "
+            f"geometries: {self._supported_geometries}. The current"
+            f" geometry is {self._ds_geom}."
+        )
+        raise NotImplementedError(msg)
+
+    @property
+    def _index_fields(self):
+        # note: using the default axis order here because the index fields
+        # will are accessed by-chunk and passed down to the pixelizer
+        # with an expected ordering matching the default ordering.
+        ax_order = self.ds.coordinates._default_axis_order
+        fields = [("index", fld) for fld in ax_order]
+        fields += [("index", f"d{fld}") for fld in ax_order]
+        return fields
+
+    @property
+    def _cartesian_to_native(self):
+        if self._ds_geom is Geometry.SPHERICAL:
+            from yt.utilities.lib.coordinate_utilities import cartesian_to_spherical
+
+            return cartesian_to_spherical
+        self._raise_unsupported_geometry()
+
+    @property
+    def _native_to_cartesian(self):
+        if self._ds_geom is Geometry.SPHERICAL:
+            from yt.utilities.lib.coordinate_utilities import spherical_to_cartesian
+
+            return spherical_to_cartesian
+        self._raise_unsupported_geometry()
+
+    def _plane_coords(self, in_plane_x, in_plane_y):
+        # calculates the 3d coordinates of points on the plane in the
+        # native coordinate system of the dataset.
+
+        # actual x, y, z locations of each point in the plane
+        c = self.center.d
+        x_global = in_plane_x * self._x_vec[0] + in_plane_y * self._y_vec[0] + c[0]
+        y_global = in_plane_x * self._x_vec[1] + in_plane_y * self._y_vec[1] + c[1]
+        z_global = in_plane_x * self._x_vec[2] + in_plane_y * self._y_vec[2] + c[2]
+
+        # now transform to the native coordinates
+        return self._cartesian_to_native(x_global, y_global, z_global)
+
+    def to_pw(self, fields=None, center="center", width=None, axes_unit=None):
+        msg = (
+            "to_pw is not implemented for mixed coordinate slices. You can create"
+            " plots manually using to_frb() to generate a fixed resolution array."
+        )
+        raise NotImplementedError(msg)
+
+    def _get_selector_class(self):
+        s_module = getattr(self, "_selector_module", selection_routines)
+        if self.ds.geometry is Geometry.SPHERICAL:
+            type_name = self._type_name + "_spherical"
+        else:
+            self._raise_unsupported_geometry()
+        sclass = getattr(s_module, f"{type_name}_selector", None)
+        return sclass
+
+    def _current_chunk_xyz(self):
+        x = self._current_chunk.fcoords[:, 0]
+        y = self._current_chunk.fcoords[:, 1]
+        z = self._current_chunk.fcoords[:, 2]
+        return self._native_to_cartesian(x, y, z)

yt/data_objects/tests/test_cartesian_cutting_plane.py

@@ -0,0 +1,114 @@
+import itertools
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+import unyt
+
+import yt
+from yt.geometry.coordinates.spherical_coordinates import spherical_to_cartesian
+from yt.testing import fake_amr_ds
+
+
+def test_cartesian_cutting_plane():
+    ds = fake_amr_ds(geometry="spherical")
+    normal = np.array([0.0, 0.0, 1.0])
+    plane_center = np.array([0.0, 0.0, 0.5])
+    slc = ds.cartesian_cutting(normal, plane_center)
+    frb = slc.to_frb(2.0, 800)
+    bvals = frb[("index", "r")]
+    mask = frb.get_mask(("index", "r"))
+    # note: the min value of r on the plane will be the z value of the
+    # plane center. how close it is to the correct answer will depend
+    # on the size of the elements.
+    assert np.allclose(bvals[mask].min().d, plane_center[2], atol=0.02)
+
+
+def _get_spherical_uniform_grid(shp, bbox, axis_order):
+
+    data = {"density": np.random.random(shp)}
+
+    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 = yt.load_uniform_grid(
+        data,
+        shp,
+        bbox=bbox,
+        geometry="spherical",
+        axis_order=axis_order,
+        length_unit="m",
+    )
+
+    ds.add_field(
+        name=("index", "z_val"), function=_z, sampling_type="cell", take_log=False
+    )
+    return ds
+
+
+@pytest.fixture
+def spherical_ds():
+
+    shp = (32, 32, 32)
+    bbox = np.array([[0.0, 1.0], [0, np.pi], [0, 2 * np.pi]])
+    ax_order = ("r", "theta", "phi")
+    return _get_spherical_uniform_grid(shp, bbox, ax_order)
+
+
+def test_cartesian_cutting_plane_fixed_z(spherical_ds):
+    ds = spherical_ds
+    normal = np.array([0.0, 0.0, 1.0])
+    center = np.array([0.0, 0.0, 0.5])
+    slc = ds.cartesian_cutting(normal, center)
+    zvals = slc["index", "z_val"].to("code_length").d
+    assert np.allclose(zvals, ds.quan(0.5, "code_length").d, atol=0.05)
+
+
+@pytest.mark.mpl_image_compare
+def test_vertical_slice_at_sphere_edge(spherical_ds):
+    ds = spherical_ds
+    normal = np.array([0.0, 1.0, 0.0])
+    center = np.array([0.0, 0.75, 0.0])
+    slc = ds.cartesian_cutting(normal, center)
+    frb = slc.to_frb(2.0, 50)
+    vals = frb["index", "z_val"].to("code_length")
+    vals[~frb.get_mask(("index", "z_val"))] = np.nan
+
+    f, axs = plt.subplots(1)
+    axs.imshow(vals, origin="lower", extent=frb.bounds)
+    return f
+
+
+def test_cartesian_cutting_plane_with_axis_ordering():
+    # check that slicing works with any axis order
+    shp = (32, 32, 32)
+    axes = ["r", "theta", "phi"]
+    bbox_ranges = {"r": [0.0, 1.0], "theta": [0, np.pi], "phi": [0, 2 * np.pi]}
+
+    # set the attributes for the plane, including a north vector found
+    # for an arbitrary point on the plane.
+    normal = np.array([1.0, 1.0, 1.0])
+    center = np.array([0.0, 0.0, 0.0])
+    x, y = 1.0, 1.0
+    z = -x * normal[0] - y * normal[1]
+    north_pt = np.array([x, y, z])
+    assert np.dot(normal, north_pt) == 0.0  # just to be sure...
+
+    frb_vals = []
+    for axis_order in itertools.permutations(axes):
+        bbox = np.zeros((3, 2))
+        for i, ax in enumerate(axis_order):
+            bbox[i, :] = bbox_ranges[ax]
+        ds = _get_spherical_uniform_grid(shp, bbox, tuple(axis_order))
+        slc = ds.cartesian_cutting(normal, center, north_vector=north_pt)
+        frb = slc.to_frb(2.0, 50)
+        vals = frb["index", "z_val"].to("code_length")
+        vals[~frb.get_mask(("index", "z_val"))] = np.nan
+        frb_vals.append(vals.d)
+
+    for frb_z in frb_vals[1:]:
+        np.allclose(frb_z, frb_vals[0])