Reduce static GS boundary Green memory

freegsnke/GSstaticsolver.py

@@ -132,7 +132,8 @@ def __init__(
             Multigrid solver for linearised GS equation.
 
         self.greenfunc
-            Boundary response Green's function matrix.
+            Boundary response Green's function matrix for source points inside
+            the limiter, where the plasma current is confined.
 
         self.nksolver
             Newton–Krylov nonlinear solver backend.
@@ -152,9 +153,6 @@ def __init__(
         self.R = R
         self.Z = Z
 
-        R_1D = R[:, 0]
-        Z_1D = Z[0, :]
-
         # number of grid points
         nx, ny = np.shape(R)
         self.nx = nx
@@ -203,22 +201,12 @@ def __init__(
         )
         self.bndry_indices = bndry_indices
 
-        # Compute Green's function mapping:
-        #
-        #   Jtor(R',Z') → ψ_boundary(R,Z)
-        greenfunc = Greens(
-            R[np.newaxis, :, :],
-            Z[np.newaxis, :, :],
-            R_1D[bndry_indices[:, 0]][:, np.newaxis, np.newaxis],
-            Z_1D[bndry_indices[:, 1]][:, np.newaxis, np.newaxis],
+        # Plasma current is confined inside the limiter, so only those Green
+        # columns contribute to the free-boundary condition.
+        self.plasma_source_mask = np.asarray(
+            eq.limiter_handler.mask_inside_limiter, dtype=bool
         )
-
-        # remove singular self-interaction terms
-        zeros = np.ones_like(greenfunc)
-        zeros[
-            np.arange(len(bndry_indices)), bndry_indices[:, 0], bndry_indices[:, 1]
-        ] = 0
-        self.greenfunc = greenfunc * zeros * self.dRdZ
+        self.greenfunc = self._build_boundary_green(self.plasma_source_mask)
 
         # Precompute geometric RHS coefficient
         # Comes from GS equation:
@@ -228,6 +216,35 @@ def __init__(
         # random generator used for NK search direction exploration
         self.rng = np.random.default_rng(seed=seed)
 
+    def _build_boundary_green(self, source_mask):
+        """Build the boundary Green matrix for a selected set of source points."""
+        source_indices = np.flatnonzero(source_mask)
+        boundary_indices = np.ravel_multi_index(
+            (self.bndry_indices[:, 0], self.bndry_indices[:, 1]),
+            (self.nx, self.ny),
+        )
+        flat_R = self.R.reshape(-1)
+        flat_Z = self.Z.reshape(-1)
+        greenfunc = Greens(
+            flat_R[source_indices][np.newaxis, :],
+            flat_Z[source_indices][np.newaxis, :],
+            flat_R[boundary_indices][:, np.newaxis],
+            flat_Z[boundary_indices][:, np.newaxis],
+        )
+
+        # Remove singular self-interactions when the selected sources include
+        # points on the computational boundary.
+        positions = np.searchsorted(source_indices, boundary_indices)
+        valid = positions < len(source_indices)
+        matches = np.zeros_like(valid)
+        matches[valid] = source_indices[positions[valid]] == boundary_indices[valid]
+        greenfunc[np.flatnonzero(matches), positions[matches]] = 0.0
+        return np.ascontiguousarray(greenfunc * self.dRdZ)
+
+    def _boundary_flux_from_jtor(self, jtor):
+        """Return boundary flux from plasma current inside the limiter."""
+        return self.greenfunc @ jtor[self.plasma_source_mask]
+
     def freeboundary(self, plasma_psi, tokamak_psi, profiles):
         """
         Apply free-boundary Grad–Shafranov boundary conditions and compute
@@ -306,15 +323,11 @@ def freeboundary(self, plasma_psi, tokamak_psi, profiles):
         #
         # psi_boundary = ∫ G(R,Z; R',Z') Jtor(R',Z') dR'dZ'
         #
-        # Implemented using tensor contraction:
-        #
-        # Contract:
-        #   greenfunc axis (1,2) with jtor axis (0,1)
-        #
-        # Result is flattened boundary flux vector.
+        # Implemented as a matrix-vector product over source points inside the
+        # limiter, outside which the plasma current is identically zero.
         # ------------------------------------------------------------
         self.psi_boundary = np.zeros_like(self.R)
-        psi_bnd = np.tensordot(self.greenfunc, self.jtor, axes=([1, 2], [0, 1]))
+        psi_bnd = self._boundary_flux_from_jtor(self.jtor)
 
         # ------------------------------------------------------------
         # Map flattened Green's solution back to boundary grid

freegsnke/tests/test_static_solver.py

@@ -206,3 +206,34 @@ def test_static_solver_rejects_invalid_operator_order(create_machine):
 
     with pytest.raises(ValueError, match="gs_operator_order"):
         GSstaticsolver.NKGSsolver(eq, gs_operator_order=3)
+
+
+def test_limiter_reduced_boundary_green_is_exact(create_machine):
+    """Limiter reduction preserves the boundary flux for confined current."""
+    eq, _, _ = create_machine
+
+    from freegsnke import GSstaticsolver
+
+    solver = GSstaticsolver.NKGSsolver(eq)
+    boundary_indices = solver.bndry_indices
+    full_green = freegs4e.gradshafranov.Greens(
+        eq.R[np.newaxis, :, :],
+        eq.Z[np.newaxis, :, :],
+        eq.R[:, 0][boundary_indices[:, 0]][:, np.newaxis, np.newaxis],
+        eq.Z[0, :][boundary_indices[:, 1]][:, np.newaxis, np.newaxis],
+    )
+    full_green[
+        np.arange(len(boundary_indices)),
+        boundary_indices[:, 0],
+        boundary_indices[:, 1],
+    ] = 0.0
+    full_green *= solver.dRdZ
+
+    jtor = np.zeros_like(eq.R)
+    jtor[solver.plasma_source_mask] = np.random.default_rng(0).random(
+        np.count_nonzero(solver.plasma_source_mask)
+    )
+    expected = np.tensordot(full_green, jtor, axes=([1, 2], [0, 1]))
+    actual = solver._boundary_flux_from_jtor(jtor)
+
+    np.testing.assert_allclose(actual, expected, rtol=2e-14, atol=1e-14)