HexBlock.rotate updates child spatial locators

armi/reactor/blocks.py

@@ -21,6 +21,7 @@
 from typing import Optional, Type, Tuple, ClassVar
 import collections
 import copy
+import functools
 import math
 
 import numpy as np
@@ -1559,26 +1560,49 @@ class during axial expansion.
         """
         self.p.axialExpTargetComponent = targetComponent.name
 
-    def getPinCoordinates(self):
+    def getPinLocations(self) -> list[grids.IndexLocation]:
+        """Produce all the index locations for pins in the block.
+
+        Returns
+        -------
+        list[grids.IndexLocation]
+            Integer locations where pins can be found in the block.
+
+        Notes
+        -----
+        Only components with ``Flags.CLAD`` are considered to define a pin's location.
+
+        See Also
+        --------
+        :meth:`getPinCoordinates` - companion for this method.
+        """
+        items = []
+        for clad in self.getChildrenWithFlags(Flags.CLAD):
+            if isinstance(clad.spatialLocator, grids.MultiIndexLocation):
+                items.extend(clad.spatialLocator)
+            else:
+                items.append(clad.spatialLocator)
+        return items
+
+    def getPinCoordinates(self) -> np.ndarray:
         """
         Compute the local centroid coordinates of any pins in this block.
 
         The pins must have a CLAD-flagged component for this to work.
 
         Returns
         -------
-        localCoordinates : list
-            list of (x,y,z) pairs representing each pin in the order they are listed as children
+        localCoords : numpy.ndarray
+            ``(N, 3)`` array of coordinates for pins locations. ``localCoords[i]`` contains a triplet of
+            the x, y, z location for pin ``i``. Ordered according to how they are listed as children
+
+        See Also
+        --------
+        :meth:`getPinLocations` - companion for this method
         """
-        coords = []
-        for clad in self.getChildrenWithFlags(Flags.CLAD):
-            if isinstance(clad.spatialLocator, grids.MultiIndexLocation):
-                coords.extend(
-                    [locator.getLocalCoordinates() for locator in clad.spatialLocator]
-                )
-            else:
-                coords.append(clad.spatialLocator.getLocalCoordinates())
-        return coords
+        indices = self.getPinLocations()
+        coords = [location.getLocalCoordinates() for location in indices]
+        return np.array(coords)
 
     def getTotalEnergyGenerationConstants(self):
         """
@@ -2017,7 +2041,7 @@ def setPinPowers(self, powers, powerKeySuffix=""):
             else:
                 self.p.linPowByPin = self.p[powerKey]
 
-    def rotate(self, rad):
+    def rotate(self, rad: float):
         """
         Rotates a block's spatially varying parameters by a specified angle in the
         counter-clockwise direction.
@@ -2028,6 +2052,22 @@ def rotate(self, rad):
 
         The pin indexing, as stored on the ``pinLocation`` parameter, is also updated.
 
+        .. impl:: A hexagonal block shall be rotatable by 60 degree increments.
+            :id: I_ARMI_ROTATE_HEX
+            :implements: R_ARMI_ROTATE_HEX
+
+        .. impl:: Rotating a hex block updates the orientation parameter.
+            :id: I_ARMI_ROTATE_HEX_ORIENTATION
+            :implements: R_ARMI_ROTATE_HEX_PARAMS
+
+        .. imp:: Rotating a hex block updates parameters on the boundary of the hexagon.
+            :id: I_ARMI_ROTATE_HEX_BOUNDARY
+            :tests: R_ARMI_ROTATE_HEX_PARAMS
+
+        .. impl:: Rotating a hex block updates the spatial coordinates on contained objects.
+            :id: I_ARMI_ROTATE_HEX_PIN
+            :tests: R_ARMI_ROTATE_HEX
+
         Parameters
         ----------
         rad: float, required
@@ -2037,10 +2077,44 @@ def rotate(self, rad):
 
         """
         rotNum = round((rad % (2 * math.pi)) / math.radians(60))
-        self._rotatePins(rotNum)
+        self._rotateChildLocations(rad, rotNum)
+        self.p.orientation[2] += rotNum * 60
         self._rotateBoundaryParameters(rotNum)
         self._rotateDisplacement(rad)
 
+    def _rotateChildLocations(self, radians: float, rotNum: int):
+        """Update spatial locators for children."""
+        if self.spatialGrid is None:
+            return
+
+        locationRotator = functools.partial(
+            self.spatialGrid.rotateIndex, rotations=rotNum
+        )
+        rotationMatrix = np.array(
+            [
+                [math.cos(radians), -math.sin(radians)],
+                [math.sin(radians), math.cos(radians)],
+            ]
+        )
+        for c in self:
+            if isinstance(c.spatialLocator, grids.MultiIndexLocation):
+                newLocations = list(map(locationRotator, c.spatialLocator))
+                c.spatialLocator = grids.MultiIndexLocation(self.spatialGrid)
+                c.spatialLocator.extend(newLocations)
+            elif isinstance(c.spatialLocator, grids.CoordinateLocation):
+                oldCoords = c.spatialLocator.getLocalCoordinates()
+                newXY = rotationMatrix.dot(oldCoords[:2])
+                newLocation = grids.CoordinateLocation(
+                    newXY[0], newXY[1], oldCoords[2], self.spatialGrid
+                )
+                c.spatialLocator = newLocation
+            elif isinstance(c.spatialLocator, grids.IndexLocation):
+                c.spatialLocator = locationRotator(c.spatialLocator)
+            elif c.spatialLocator is not None:
+                msg = f"{c} on {self} has an invalid spatial locator for rotation: {c.spatialLocator}"
+                runLog.error(msg)
+                raise TypeError(msg)
+
     def _rotateBoundaryParameters(self, rotNum: int):
         """Rotate any parameters defined on the corners or edge of bounding hexagon.
 
@@ -2090,98 +2164,6 @@ def _rotateDisplacement(self, rad: float):
             self.p.displacementX = dispx * math.cos(rad) - dispy * math.sin(rad)
             self.p.displacementY = dispx * math.sin(rad) + dispy * math.cos(rad)
 
-    def _rotatePins(self, rotNum, justCompute=False):
-        """
-        Rotate the pins of a block, which means rotating the indexing of pins. Note that this does
-        not rotate all block quantities, just the pins.
-
-        Parameters
-        ----------
-        rotNum : int, required
-            An integer from 0 to 5, indicating the number of counterclockwise 60-degree rotations
-            from the CURRENT orientation. Degrees of counter-clockwise rotation = 60*rot
-
-        justCompute : boolean, optional
-            If True, rotateIndexLookup will be returned but NOT assigned to the object parameter
-            self.p.pinLocation. If False, rotateIndexLookup will be returned AND assigned to the
-            object variable self.p.pinLocation.  Useful for figuring out which rotation is best
-            to minimize burnup, etc.
-
-        Returns
-        -------
-        rotateIndexLookup : dict of ints
-            This is an index lookup (or mapping) between pin ids and pin locations. The pin
-            indexing is 1-D (not ring,pos or GEODST). The "ARMI pin ordering" is used for location,
-            which is counter-clockwise from 1 o'clock. Pin ids are always consecutively
-            ordered starting at 1, while pin locations are not once a rotation has been
-            applied.
-
-        Notes
-        -----
-        Changing (x,y) positions of pins does NOT constitute rotation, because the indexing of pin
-        atom densities must be re-ordered.  Re-order indexing of pin-level quantities, NOT (x,y)
-        locations of pins.  Otherwise, subchannel input will be in wrong order.
-
-        How rotations works is like this. There are pins with unique pin numbers in each block.
-        These pin numbers will not change no matter what happens to a block, so if you have pin 1,
-        you always have pin 1. However, these pins are all in pinLocations, and these are what
-        change with rotations. At BOL, a pin's pinLocation is equal to its pin number, but after
-        a rotation, this will no longer be so.
-
-        So, all params that don't care about exactly where in space the pin is (such as depletion)
-        can just use the pin number, but anything that needs to know the spatial location (such as
-        fluxRecon, which interpolates the flux spatially, or subchannel codes, which needs to know where the
-        power is) need to map through the pinLocation parameters.
-
-        This method rotates the pins by changing the pinLocation parameter.
-
-        See Also
-        --------
-        armi.reactor.blocks.HexBlock.rotate
-            Rotates the entire block (pins, ducts, and spatial quantities).
-
-        Examples
-        --------
-            rotateIndexLookup[i_after_rotation-1] = i_before_rotation-1
-        """
-        if not 0 <= rotNum <= 5:
-            raise ValueError(
-                "Cannot rotate {0} to rotNum {1}. Must be 0-5. ".format(self, rotNum)
-            )
-
-        numPins = self.getNumPins()
-        hexRings = hexagon.numRingsToHoldNumCells(numPins)
-        fullNumPins = hexagon.totalPositionsUpToRing(hexRings)
-        rotateIndexLookup = dict(
-            zip(range(1, fullNumPins + 1), range(1, fullNumPins + 1))
-        )
-
-        # Look up the current orientation and add this to it. The math below just rotates
-        # from the reference point so we need a total rotation.
-        rotNum = int((self.getRotationNum() + rotNum) % 6)
-
-        # non-trivial rotation requested
-        # start at 2 because pin 1 never changes (it's in the center!)
-        for pinNum in range(2, fullNumPins + 1):
-            if rotNum == 0:
-                # Rotation to reference orientation. Pin locations are pin IDs.
-                pass
-            else:
-                newPinLocation = hexagon.getIndexOfRotatedCell(pinNum, rotNum)
-                # Assign "before" and "after" pin indices to the index lookup
-                rotateIndexLookup[pinNum] = newPinLocation
-
-        # Because the above math creates indices based on the absolute rotation number,
-        # the old values of pinLocation (if they've been set in the past) can be overwritten
-        # with new numbers
-        if not justCompute:
-            self.setRotationNum(rotNum)
-            self.p["pinLocation"] = [
-                rotateIndexLookup[pinNum] for pinNum in range(1, fullNumPins + 1)
-            ]
-
-        return rotateIndexLookup
-
     def verifyBlockDims(self):
         """Perform some checks on this type of block before it is assembled."""
         try:
@@ -2286,13 +2268,13 @@ def getPinToDuctGap(self, cold=False):
 
         return pinToDuctGap
 
-    def getRotationNum(self):
+    def getRotationNum(self) -> int:
         """Get index 0 through 5 indicating number of rotations counterclockwise around the z-axis."""
         return (
             np.rint(self.p.orientation[2] / 360.0 * 6) % 6
         )  # assume rotation only in Z
 
-    def setRotationNum(self, rotNum):
+    def setRotationNum(self, rotNum: int):
         """
         Set orientation based on a number 0 through 5 indicating number of rotations
         counterclockwise around the z-axis.

armi/reactor/grids/hexagonal.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from collections import deque
 from math import sqrt
 from typing import Tuple, List, Optional
 
@@ -23,7 +24,7 @@
     BOUNDARY_60_DEGREES,
     BOUNDARY_CENTER,
 )
-from armi.reactor.grids.locations import IJKType, IJType
+from armi.reactor.grids.locations import IJKType, IJType, IndexLocation
 from armi.reactor.grids.structuredGrid import StructuredGrid
 from armi.utils import hexagon
 
@@ -221,7 +222,8 @@ def indicesToRingPos(i: int, j: int) -> Tuple[int, int]:
         positionBase = 1 + edge * (ring - 1)
         return ring, positionBase + offset
 
-    def getMinimumRings(self, n: int) -> int:
+    @staticmethod
+    def getMinimumRings(n: int) -> int:
         """
         Return the minimum number of rings needed to fit ``n`` objects.
 
@@ -232,7 +234,8 @@ def getMinimumRings(self, n: int) -> int:
         """
         return hexagon.numRingsToHoldNumCells(n)
 
-    def getPositionsInRing(self, ring: int) -> int:
+    @staticmethod
+    def getPositionsInRing(ring: int) -> int:
         """Return the number of positions within a ring."""
         return hexagon.numPositionsInRing(ring)
 
@@ -572,3 +575,70 @@ def generateSortedHexLocationList(self, nLocs: int):
         )
 
         return locList[:nLocs]
+
+    def rotateIndex(self, loc: IndexLocation, rotations: int) -> IndexLocation:
+        """Find the new location of an index after some number of CCW rotations.
+
+        Parameters
+        ----------
+        loc : IndexLocation
+            Starting index
+        rotations : int
+            Number of counter clockwise rotations
+
+        Returns
+        -------
+        IndexLocation
+            Index in the grid after rotation
+
+        Notes
+        -----
+        Rotation uses a three-dimensional index in what can be known elsewhere
+        by the confusing name of "cubic" coordinate system for a hexagon. Cubic stems
+        from the notion of using three dimensions, ``(q, r, s)`` to describe a point in the
+        hexagonal grid. The conversion from the indexing used in the ARMI framework follows::
+
+            q = i
+            r = j
+            # s = - q - r = - (q + r)
+            s = -(i + j)
+
+        The motivation for the cubic notation is rotation is far simpler: a clockwise
+        rotation by 60 degrees results in a shifting and negating of the coordinates. So
+        the first rotation of ``(q, r, s)`` would produce a new coordinate
+        ``(-r, -s, -q)``. Another rotation would produce ``(s, q, r)``, and so on.
+
+        Raises
+        ------
+        TypeError
+            If ``loc.grid`` is populated and not consistent with this grid. For example,
+            it doesn't make sense to rotate an index from a Cartesian grid in a hexagonal coordinate
+            system, nor hexagonal grid with different orientation (flats up vs. corners up)
+        """
+        if self._roughlyEqual(loc.grid) or loc.grid is None:
+            i, j, k = loc[:3]
+            buffer = deque((i, j, -(i + j)))
+            buffer.rotate(-rotations)
+            newI = buffer[0]
+            newJ = buffer[1]
+            if rotations % 2:
+                newI *= -1
+                newJ *= -1
+            return IndexLocation(newI, newJ, k, loc.grid)
+        raise TypeError(
+            f"Refusing to rotate an index {loc} from a grid {loc.grid} that "
+            f"is not consistent with {self}"
+        )
+
+    def _roughlyEqual(self, other) -> bool:
+        """Check that two hex grids are nearly identical.
+
+        Would the same ``(i, j, k)`` index in ``self`` be the same location in ``other``?
+        """
+        if other is self:
+            return True
+        return (
+            isinstance(other, HexGrid)
+            and other.pitch == self.pitch
+            and other.cornersUp == self.cornersUp
+        )

armi/reactor/grids/locations.py

@@ -97,7 +97,7 @@ def __hash__(self) -> Hashable:
         """
         return hash((self.i, self.j, self.k))
 
-    def __eq__(self, other: Union[Tuple[int, int, int], "LocationBase"]) -> bool:
+    def __eq__(self, other: Union[IJKType, "LocationBase"]) -> bool:
         if isinstance(other, tuple):
             return (self.i, self.j, self.k) == other
         if isinstance(other, LocationBase):