allowing varible offsets for polygon.offset (#3120)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

openmc/model/surface_composite.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from copy import copy
@@ -1316,25 +1317,43 @@ class Polygon(CompositeSurface):
             surfsets.append(surf_ops)
         return surfsets
 
-    def offset(self, distance):
+    def offset(self, distance: float | Sequence[float] | np.ndarray) -> Polygon:
         """Offset this polygon by a set distance
 
         Parameters
         ----------
-        distance : float
+        distance : float or sequence of float or np.ndarray
             The distance to offset the polygon by. Positive is outward
-            (expanding) and negative is inward (shrinking).
+            (expanding) and negative is inward (shrinking). If a float is
+            provided, the same offset is applied to all vertices. If a list or
+            tuple is provided, each vertex gets a different offset. If an
+            iterable or numpy array is provided, each vertex gets a different
+            offset.
 
         Returns
         -------
         offset_polygon : openmc.model.Polygon
         """
+
+        if isinstance(distance, float):
+            distance = np.full(len(self.points), distance)
+        elif isinstance(distance, Sequence):
+            distance = np.array(distance)
+        elif not isinstance(distance, np.ndarray):
+            raise TypeError("Distance must be a float or sequence of float.")
+
+        if len(distance) != len(self.points):
+            raise ValueError(
+                f"Length of distance {len(distance)} array must "
+                f"match number of polygon points {len(self.points)}"
+            )
+
         normals = np.insert(self._normals, 0, self._normals[-1, :], axis=0)
         cos2theta = np.sum(normals[1:, :]*normals[:-1, :], axis=-1, keepdims=True)
         costheta = np.cos(np.arccos(cos2theta) / 2)
         nvec = (normals[1:, :] + normals[:-1, :])
         unit_nvec = nvec / np.linalg.norm(nvec, axis=-1, keepdims=True)
-        disp_vec = distance / costheta * unit_nvec
+        disp_vec = distance[:, np.newaxis] / costheta * unit_nvec
 
         return type(self)(self.points + disp_vec, basis=self.basis)
 

tests/unit_tests/test_surface_composite.py

@@ -347,10 +347,13 @@ def test_polygon():
         assert any([points_in[i] in reg for reg in star_poly.regions])
         assert points_in[i] not in +star_poly
         assert (0, 0, 0) not in -star_poly
-        if basis != 'rz':
-            offset_star = star_poly.offset(.6)
-            assert (0, 0, 0) in -offset_star
-            assert any([(0, 0, 0) in reg for reg in offset_star.regions])
+        if basis != "rz":
+            for offsets in [0.6, np.array([0.6] * 10), [0.6] * 10]:
+                offset_star = star_poly.offset(offsets)
+                assert (0, 0, 0) in -offset_star
+                assert any([(0, 0, 0) in reg for reg in offset_star.regions])
+            with pytest.raises(ValueError):
+                star_poly.offset([0.6, 0.6])
 
     # check invalid Polygon input points
     # duplicate points not just at start and end