Merge pull request #100 from FusionComputingLab/share-profile-geometry

Share immutable profile geometry

Author: kpentland

freegsnke/jtor_update.py

@@ -98,7 +98,7 @@ def copy(self, obj=None):
         -----
         Copy semantics:
         - Immutable / scalar attributes are copied directly
-        - NumPy arrays and simple containers are shallow-copied where safe
+        - Grid geometry and limiter masks are shared through limiter_handler
         - Selected complex objects are deep-copied (e.g. via copy.deepcopy)
         - Some large system objects are shared by reference (e.g. limiter_handler)
 
@@ -122,17 +122,13 @@ def copy(self, obj=None):
         copy_into(self, obj, "dRdZ")
         copy_into(self, obj, "nx")
 
-        copy_into(self, obj, "dR_dZ", mutable=True)
-        copy_into(self, obj, "grid_points", mutable=True)
-        copy_into(self, obj, "eqRidx", mutable=True)
-        copy_into(self, obj, "eqZidx", mutable=True)
-        copy_into(self, obj, "idx_grid_points", mutable=True)
-        copy_into(self, obj, "R0Z0", mutable=True)
-        copy_into(self, obj, "mask_inside_limiter", mutable=True)
-        copy_into(self, obj, "mask_outside_limiter", mutable=True)
-        copy_into(self, obj, "limiter_mask_out", mutable=True)
-        copy_into(self, obj, "limiter_mask_for_plotting", mutable=True)
-        copy_into(self, obj, "edge_mask", mutable=True)
+        copy_into(self, obj, "dR_dZ")
+        copy_into(self, obj, "eqRidx")
+        copy_into(self, obj, "eqZidx")
+        copy_into(self, obj, "R0Z0")
+        copy_into(self, obj, "mask_inside_limiter")
+        copy_into(self, obj, "mask_outside_limiter")
+        copy_into(self, obj, "limiter_mask_out")
         obj.inputs = self.inputs[::]  # shallow copy suffices
 
         # *Should* not be necessary to copy this
@@ -187,46 +183,25 @@ def set_masks(self, eq):
         eq : FreeGSNKE Equilibrium object
             Equilibrium defining the computational domain, including 1D and 2D
             coordinate grids and limiter geometry.
+
+        Notes
+        -----
+        Fixed grid arrays are shared with ``eq.limiter_handler`` and treated as
+        read-only by profile calculations.
         """
         self.dR = eq.R_1D[1] - eq.R_1D[0]
         self.dZ = eq.Z_1D[1] - eq.Z_1D[0]
-        self.dR_dZ = np.array([self.dR, self.dZ])
-        self.R0Z0 = np.array([eq.R_1D[0], eq.Z_1D[0]])
         self.dRdZ = self.dR * self.dZ
-        self.grid_points = np.concatenate(
-            (eq.R[:, :, np.newaxis], eq.Z[:, :, np.newaxis]), axis=-1
-        )
         self.nx, self.ny = np.shape(eq.R)
-        self.eqRidx = np.tile(np.arange(self.nx)[:, np.newaxis], (1, self.ny))
-        self.eqZidx = np.tile(np.arange(self.ny)[:, np.newaxis], (1, self.nx)).T
-        self.idx_grid_points = np.concatenate(
-            (self.eqRidx[:, :, np.newaxis], self.eqZidx[:, :, np.newaxis]), axis=-1
-        ).reshape(-1, 2)
 
         self.limiter_handler = eq.limiter_handler
-
-        # self.core_mask_limiter = eq.limiter_handler.core_mask_limiter
-
-        self.mask_inside_limiter = eq.limiter_handler.mask_inside_limiter
-
-        mask_outside_limiter = np.logical_not(eq.limiter_handler.mask_inside_limiter)
-        # Note the factor 2 is not a typo: used in critical.inside_mask
-        self.mask_outside_limiter = (2 * mask_outside_limiter).astype(float)
-
-        self.limiter_mask_out = eq.limiter_handler.limiter_mask_out
-
-        self.limiter_mask_for_plotting = (
-            eq.limiter_handler.mask_inside_limiter
-            + eq.limiter_handler.make_layer_mask(
-                eq.limiter_handler.mask_inside_limiter, layer_size=1
-            )
-        ) > 0
-
-        # set mask of the edge domain pixels
-        self.edge_mask = np.zeros_like(eq.R)
-        self.edge_mask[0, :] = self.edge_mask[:, 0] = self.edge_mask[-1, :] = (
-            self.edge_mask[:, -1]
-        ) = 1
+        self.dR_dZ = self.limiter_handler.dR_dZ
+        self.R0Z0 = self.limiter_handler.R0Z0
+        self.eqRidx = self.limiter_handler.eqRidx
+        self.eqZidx = self.limiter_handler.eqZidx
+        self.mask_inside_limiter = self.limiter_handler.mask_inside_limiter
+        self.mask_outside_limiter = self.limiter_handler.mask_outside_limiter
+        self.limiter_mask_out = self.limiter_handler.limiter_mask_out
 
     def select_refinement(self, eq, refine_jtor, nnx, nny):
         """
@@ -408,9 +383,9 @@ def diverted_critical(
             )
         )[np.newaxis]
         # build diverted_core_mask
-        diverted_core_mask = path.contains_points(self.idx_grid_points).reshape(
-            (self.nx, self.ny)
-        )
+        diverted_core_mask = path.contains_points(
+            self.limiter_handler.idx_grid_points
+        ).reshape((self.nx, self.ny))
 
         return opt, xpt, diverted_core_mask, psi_bndry
 

freegsnke/limiter_func.py

@@ -32,7 +32,8 @@ class Limiter_handler:
     objects to determine whether grid points lie inside the allowable plasma
     region and to support limiter-dependent calculations.
 
-    Each profile object typically instantiates its own Limiter_handler.
+    An equilibrium owns one handler; profile objects share its fixed grid geometry
+    and limiter masks.
 
     Notes
     -----
@@ -61,6 +62,8 @@ def __init__(self, eq, limiter):
 
         self.dR = self.eqR[1, 0] - self.eqR[0, 0]
         self.dZ = self.eqZ[0, 1] - self.eqZ[0, 0]
+        self.dR_dZ = np.array([self.dR, self.dZ])
+        self.R0Z0 = np.array([self.eqR_1D[0], self.eqZ_1D[0]])
         self.dRdZ = self.dR * self.dZ
         self.nx, self.ny = np.shape(eq.R)
         self.nxny = self.nx * self.ny
@@ -70,9 +73,22 @@ def __init__(self, eq, limiter):
 
         self.validate_limiter_inside_domain()
         self.build_mask_inside_limiter()
+        self.mask_outside_limiter = (
+            2 * np.logical_not(self.mask_inside_limiter)
+        ).astype(float)
         self.limiter_points()
         self.plasma_pts = self.extract_plasma_pts(eq.R, eq.Z, self.mask_inside_limiter)
         self.idxs_mask = self.extract_index_mask(self.mask_inside_limiter)
+        self._idx_grid_points = None
+
+    @property
+    def idx_grid_points(self):
+        """Grid indices for contour fallback, built only when first required."""
+        if self._idx_grid_points is None:
+            self._idx_grid_points = np.column_stack(
+                (self.eqRidx.reshape(-1), self.eqZidx.reshape(-1))
+            )
+        return self._idx_grid_points
 
     def validate_limiter_inside_domain(self):
         """Raise a clear error if the limiter is not inside the solution domain.

freegsnke/tests/test_jtor_update.py

@@ -101,6 +101,32 @@ def test_profiles_BetapIp(create_machine):
     ), "The profiles object does not have the xpt, opt and jtor attributes"
 
 
+def test_profile_geometry_is_shared_and_fallback_indices_are_lazy(create_machine):
+    eq, _ = create_machine
+    profiles = jtor.ConstrainPaxisIp(eq, 8.1e3, 6.2e5, 0.5)
+    profiles.inputs = []
+    copied_profiles = jtor.Jtor_universal.copy(profiles)
+    handler = eq.limiter_handler
+
+    for name in (
+        "dR_dZ",
+        "R0Z0",
+        "eqRidx",
+        "eqZidx",
+        "mask_inside_limiter",
+        "mask_outside_limiter",
+        "limiter_mask_out",
+    ):
+        assert getattr(profiles, name) is getattr(handler, name)
+        assert getattr(copied_profiles, name) is getattr(handler, name)
+
+    assert handler._idx_grid_points is None
+    assert handler.idx_grid_points.shape == (eq.nx * eq.ny, 2)
+    expected_indices = np.indices((eq.nx, eq.ny)).reshape(2, -1).T
+    assert np.array_equal(handler.idx_grid_points, expected_indices)
+    assert handler.idx_grid_points is handler._idx_grid_points
+
+
 def test_contour_fallback_coordinates_and_current_sign(create_machine):
     """The contour fallback is invariant under simultaneous psi and Ip reversal."""
     eq, _ = create_machine