From 9983ee1a7eb9b95bd45123a9d44a021f789f0410 Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Sat, 9 Nov 2024 00:52:40 +0100 Subject: [PATCH] allowing varible offsets for polygon.offset (#3120) Co-authored-by: Paul Romano --- openmc/model/surface_composite.py | 27 ++++++++++++++++++---- tests/unit_tests/test_surface_composite.py | 11 +++++---- 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/openmc/model/surface_composite.py b/openmc/model/surface_composite.py index df290329647..7a5ab8f1002 100644 --- a/openmc/model/surface_composite.py +++ b/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 @@ def _decompose_polygon_into_convex_sets(self): 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) diff --git a/tests/unit_tests/test_surface_composite.py b/tests/unit_tests/test_surface_composite.py index 963bbe00d19..015ac667e45 100644 --- a/tests/unit_tests/test_surface_composite.py +++ b/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