diff --git a/yt/data_objects/selection_objects/slices.py b/yt/data_objects/selection_objects/slices.py index 69547f1c295..40dbf616ced 100644 --- a/yt/data_objects/selection_objects/slices.py +++ b/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] diff --git a/yt/geometry/_selection_routines/slice_selector.pxi b/yt/geometry/_selection_routines/slice_selector.pxi index 0e09ae889c2..77fdb1099f3 100644 --- a/yt/geometry/_selection_routines/slice_selector.pxi +++ b/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 = ( - (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 diff --git a/yt/geometry/coordinates/spherical_coordinates.py b/yt/geometry/coordinates/spherical_coordinates.py index 165287393a2..007a9e6b10f 100644 --- a/yt/geometry/coordinates/spherical_coordinates.py +++ b/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 diff --git a/yt/utilities/lib/pixelization_routines.pyx b/yt/utilities/lib/pixelization_routines.pyx index 72197abdb3c..cfb88e19c72 100644 --- a/yt/utilities/lib/pixelization_routines.pyx +++ b/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