[WIP] Naive off-axis slice attempt for spherical-ish geometry

yt/data_objects/selection_objects/slices.py

@@ -15,6 +15,7 @@
     validate_object,
     validate_width_tuple,
 )
+from yt.geometry.api import Geometry
 from yt.utilities.exceptions import YTNotInsideNotebook
 from yt.utilities.minimal_representation import MinimalSliceData
 from yt.utilities.orientation import Orientation
@@ -50,6 +51,11 @@ class YTSlice(YTSelectionContainer2D):
     data_source: optional
         Draw the selection from the provided data source rather than
         all data associated with the data_set
+    offset: array_like, optional
+        Apply an offset to the slicing operation.  Only operable in
+        spherical geometries, to facilitate rotation. This will be added to
+        coordinates as they are compared against the selection routine
+        (and will be supplied to any plot object as well.)
 
     Examples
     --------
@@ -62,11 +68,18 @@ class YTSlice(YTSelectionContainer2D):
 
     _top_node = "/Slices"
     _type_name = "slice"
-    _con_args = ("axis", "coord")
+    _con_args = ("axis", "coord", "offset")
     _container_fields = ("px", "py", "pz", "pdx", "pdy", "pdz")
 
     def __init__(
-        self, axis, coord, center=None, ds=None, field_parameters=None, data_source=None
+        self,
+        axis,
+        coord,
+        center=None,
+        offset=None,
+        ds=None,
+        field_parameters=None,
+        data_source=None,
     ):
         validate_axis(ds, axis)
         validate_float(coord)
@@ -79,6 +92,14 @@ def __init__(
         YTSelectionContainer2D.__init__(self, axis, ds, field_parameters, data_source)
         self._set_center(center)
         self.coord = coord
+        if (
+            ds.geometry
+            in (Geometry.SPHERICAL, Geometry.GEOGRAPHIC, Geometry.INTERNAL_GEOGRAPHIC)
+            and offset is not None
+        ):
+            self.offset = offset
+        else:
+            self.offset = [0.0, 0.0, 0.0]
 
     def _generate_container_field(self, field):
         xax = self.ds.coordinates.x_axis[self.axis]

yt/geometry/_selection_routines/cutting_plane_selector.pxi

@@ -1,3 +1,5 @@
+from libc.math cimport sin, cos, atan2, sqrt, acos
+
 cdef class CuttingPlaneSelector(SelectorObject):
     cdef public np.float64_t norm_vec[3]
     cdef public np.float64_t d
@@ -114,4 +116,27 @@ cdef class CuttingPlaneSelector(SelectorObject):
     def _get_state_attnames(self):
         return ("d", "norm_vec")
 
+cdef class SphericalCuttingPlaneSelector(CuttingPlaneSelector):
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    @cython.cdivision(True)
+    cdef void transform_rtp_to_xyz(self, np.float64_t in_pos[3], np.float64_t out_pos[3]) noexcept nogil:
+        # See `spherical_coordinates.py` for more details
+        # r => in_pos[0] theta => in_pos[1] phi => in_pos[2]
+        out_pos[0] = in_pos[0] * cos(in_pos[2]) * sin(in_pos[1])
+        out_pos[1] = in_pos[0] * sin(in_pos[2]) * sin(in_pos[1])
+        out_pos[2] = in_pos[0] * cos(in_pos[1])
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    @cython.cdivision(True)
+    cdef void transform_xyz_to_rtp(self, np.float64_t in_pos[3], np.float64_t out_pos[3]) noexcept nogil:
+        # we have to have 012 be xyz and 012 be rtp
+        out_pos[0] = sqrt(in_pos[0]*in_pos[0] + in_pos[1]*in_pos[1] + in_pos[2]*in_pos[2])
+        out_pos[1] = acos(in_pos[2] / out_pos[0])
+        out_pos[2] = atan2(in_pos[1], in_pos[0])
+
 cutting_selector = CuttingPlaneSelector
+
+spherical_cutting_selector = SphericalCuttingPlaneSelector

yt/geometry/_selection_routines/slice_selector.pxi

@@ -3,6 +3,8 @@ cdef class SliceSelector(SelectorObject):
     cdef public np.float64_t coord
     cdef public int ax, ay
     cdef public int reduced_dimensionality
+    cdef np.float64_t _offset[3]
+    cdef object offset
 
     def __init__(self, dobj):
         self.axis = dobj.axis
@@ -15,10 +17,20 @@ cdef class SliceSelector(SelectorObject):
             self.reduced_dimensionality = 1
         else:
             self.reduced_dimensionality = 0
+        self.offset = getattr(dobj, "offset", [0, 0, 0])
 
         self.ax = (self.axis+1) % 3
         self.ay = (self.axis+2) % 3
 
+    @property
+    def offset(self):
+        return tuple(self._offset[i] for i in range(3))
+
+    @offset.setter
+    def offset(self, value):
+        for i in range(3):
+            self._offset[i] = value[i]
+
     @cython.boundscheck(False)
     @cython.wraparound(False)
     @cython.cdivision(True)
@@ -44,7 +56,7 @@ cdef class SliceSelector(SelectorObject):
             for i in range(3):
                 if i == self.axis:
                     icoord = <np.uint64_t>(
-                        (self.coord - gobj.LeftEdge.d[i])/gobj.dds[i])
+                        (self.coord - (gobj.LeftEdge.d[i] + self._offset[i]))/gobj.dds[i])
                     # clip coordinate to avoid seg fault below if we're
                     # exactly at a grid boundary
                     ind[i][0] = iclip(
@@ -69,8 +81,8 @@ cdef class SliceSelector(SelectorObject):
     cdef int select_cell(self, np.float64_t pos[3], np.float64_t dds[3]) noexcept nogil:
         if self.reduced_dimensionality == 1:
             return 1
-        if pos[self.axis] + 0.5*dds[self.axis] > self.coord \
-           and pos[self.axis] - 0.5*dds[self.axis] - grid_eps <= self.coord:
+        if (pos[self.axis] + self._offset[self.axis]) + 0.5*dds[self.axis] > self.coord \
+           and (pos[self.axis] + self._offset[self.axis]) - 0.5*dds[self.axis] - grid_eps <= self.coord:
             return 1
         return 0
 
@@ -85,7 +97,7 @@ cdef class SliceSelector(SelectorObject):
         if self.reduced_dimensionality == 1:
             return 1
         cdef np.float64_t dist = self.periodic_difference(
-            pos[self.axis], self.coord, self.axis)
+            pos[self.axis] + self._offset[self.axis], self.coord, self.axis)
         if dist*dist < radius*radius:
             return 1
         return 0
@@ -97,7 +109,7 @@ cdef class SliceSelector(SelectorObject):
                                np.float64_t right_edge[3]) noexcept nogil:
         if self.reduced_dimensionality == 1:
             return 1
-        if left_edge[self.axis] - grid_eps <= self.coord < right_edge[self.axis]:
+        if (left_edge[self.axis] + self._offset[self.axis]) - grid_eps <= self.coord < (right_edge[self.axis] + self._offset[self.axis]):
             return 1
         return 0
 
@@ -108,15 +120,16 @@ cdef class SliceSelector(SelectorObject):
                                np.float64_t right_edge[3]) noexcept nogil:
         if self.reduced_dimensionality == 1:
             return 2
-        if left_edge[self.axis] - grid_eps <= self.coord < right_edge[self.axis]:
+        if (left_edge[self.axis] + self._offset[self.axis]) - grid_eps <= self.coord < (right_edge[self.axis] + self._offset[self.axis]):
             return 2 # a box with non-zero volume can't be inside a plane
         return 0
 
     def _hash_vals(self):
         return (("axis", self.axis),
-                ("coord", self.coord))
+                ("coord", self.coord),
+                ("offset", self.offset))
 
     def _get_state_attnames(self):
-        return ("axis", "coord", "ax", "ay", "reduced_dimensionality")
+        return ("axis", "coord", "ax", "ay", "reduced_dimensionality", "offset")
 
 slice_selector = SliceSelector

yt/geometry/coordinates/spherical_coordinates.py

@@ -128,6 +128,7 @@ def _ortho_pixelize(
     ):
         # use Aitoff projection
         # http://paulbourke.net/geometry/transformationprojection/
+        # We should be supplying an offset here, but it needs checking.
         bounds = tuple(_.ndview for _ in self._aitoff_bounds)
         buff, mask = pixelize_aitoff(
             azimuth=data_source["py"],
@@ -157,6 +158,7 @@ def _cyl_pixelize(self, data_source, field, bounds, size, antialias, dimension):
                 data_source[field],
                 bounds,
                 return_mask=True,
+                offset=data_source.offset,
             )
         elif name == "phi":
             # Note that we feed in buff.T here
@@ -169,6 +171,7 @@ def _cyl_pixelize(self, data_source, field, bounds, size, antialias, dimension):
                 data_source[field],
                 bounds,
                 return_mask=True,
+                offset=data_source.offset,
             ).T
         else:
             raise RuntimeError

yt/utilities/lib/pixelization_routines.pyx

@@ -558,6 +558,7 @@ def pixelize_cylinder(np.float64_t[:,:] buff,
                       extents,
                       *,
                       int return_mask=0,
+                      offset=None,
 ):
 
     cdef np.float64_t x, y, dx, dy, r0, theta0
@@ -566,7 +567,11 @@ def pixelize_cylinder(np.float64_t[:,:] buff,
     cdef np.float64_t r_inc, theta_inc
     cdef np.float64_t costheta, sintheta
     cdef int i, i1, pi, pj
-
+    cdef np.float64_t voff[3]
+    if offset is None:
+        voff[0] = voff[1] = voff[2] = 0.0
+    else:
+        voff[0], voff[1], voff[2] = offset
     cdef int imin, imax
     imin = np.asarray(radius).argmin()
     imax = np.asarray(radius).argmax()
@@ -612,8 +617,8 @@ def pixelize_cylinder(np.float64_t[:,:] buff,
     r_inc = 0.5 * fmin(dx, dy)
 
     for i in range(radius.shape[0]):
-        r0 = radius[i]
-        theta0 = theta[i]
+        r0 = radius[i] + voff[0]
+        theta0 = theta[i] + voff[1]
         dr_i = dradius[i]
         dtheta_i = dtheta[i]
         # Skip out early if we're offsides, for zoomed in plots