Allow selecting a faster second-order static GS operator

examples/example02 - static_forward_solve_MASTU.ipynb

@@ -131,7 +131,9 @@
     "\n",
     "We can now load FreeGSNKE's Grad-Shafranov static solver. The equilibrium is used to inform the solver of the computational domain and of the tokamak properties. The solver below can be used for both inverse and forward solve modes.\n",
     "\n",
-    "Note: It's not necessary to instantiate a new solver when aiming to use it on new or different equilibria, as long as the integration domain, mesh grid, and tokamak are consistent across solves. "
+    "Note: It's not necessary to instantiate a new solver when aiming to use it on new or different equilibria, as long as the integration domain, mesh grid, and tokamak are consistent across solves.\n",
+    "\n",
+    "The default `gs_operator_order=4` uses FreeGS4E's fourth-order finite-difference Grad-Shafranov operator. For exploratory calculations where a small discretisation-accuracy trade-off is acceptable, `gs_operator_order=2` selects FreeGS4E's second-order operator. With the direct sparse solver used here, this reduces matrix construction and LU factorisation costs; repeated back-solves have similar cost. The operator order is fixed when the solver is instantiated."
    ]
   },
   {
@@ -141,7 +143,11 @@
    "outputs": [],
    "source": [
     "from freegsnke import GSstaticsolver\n",
-    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    "
+    "\n",
+    "GSStaticSolver = GSstaticsolver.NKGSsolver(\n",
+    "    eq,\n",
+    "    gs_operator_order=4,  # use 2 for lower direct-solver setup cost\n",
+    ")"
    ]
   },
   {

freegsnke/GSstaticsolver.py

@@ -73,14 +73,15 @@ def __init__(
         l2_reg=1e-6,
         collinearity_reg=1e-6,
         seed=42,
+        gs_operator_order=4,
     ):
         """
         Initialise the Grad–Shafranov nonlinear solver.
 
         The constructor prepares all numerical operators required for
         nonlinear GS solving, including:
 
-            • Linear GS multigrid solver
+            • Direct sparse linear GS solver
             • Green's function boundary response operator
             • Newton–Krylov nonlinear solver backend
             • Random generator for Krylov direction perturbations
@@ -110,6 +111,12 @@ def __init__(
                 • Krylov perturbation generation
                 • Directional exploration in nonlinear solve
 
+        gs_operator_order : {2, 4}, optional (default=4)
+            Finite-difference order of the linear Grad-Shafranov operator.
+            Fourth order is more accurate; second order reduces sparse matrix
+            construction and factorisation costs when that accuracy trade-off
+            is acceptable.
+
         Attributes
         ----------
         self.R, self.Z : ndarray
@@ -159,6 +166,14 @@ def __init__(
         dZ = Z[0, 1] - Z[0, 0]
         self.dRdZ = dR * dZ
 
+        if gs_operator_order == 2:
+            gs_operator = freegs4e.gradshafranov.GSsparse
+        elif gs_operator_order == 4:
+            gs_operator = freegs4e.gradshafranov.GSsparse4thOrder
+        else:
+            raise ValueError("gs_operator_order must be either 2 or 4")
+        self.gs_operator_order = gs_operator_order
+
         # nonlinear solver backend
         self.nksolver = nk_solver.nksolver(
             problem_dimension=self.nx * self.ny,
@@ -170,9 +185,7 @@ def __init__(
         self.linear_GS_solver = freegs4e.multigrid.createVcycle(
             nx,
             ny,
-            freegs4e.gradshafranov.GSsparse4thOrder(
-                eq.R[0, 0], eq.R[-1, 0], eq.Z[0, 0], eq.Z[0, -1]
-            ),
+            gs_operator(eq.R[0, 0], eq.R[-1, 0], eq.Z[0, 0], eq.Z[0, -1]),
             nlevels=1,
             ncycle=1,
             niter=2,

freegsnke/tests/test_static_solver.py

@@ -175,3 +175,34 @@ def test_static_solve(create_machine):
     assert np.allclose(
         eq.psi(), test_psi, atol=(np.max(test_psi) - np.min(test_psi)) * 0.003
     ), "Psi map differs significantly from the test map"
+
+
+def test_second_order_static_solve(create_machine):
+    """The opt-in second-order GS operator produces a consistent equilibrium."""
+    eq, profiles, _ = create_machine
+
+    from freegsnke import GSstaticsolver
+
+    eq.tokamak.set_coil_current("P6", 0)
+    eq.tokamak["P6"].control = False
+    eq.tokamak["Solenoid"].control = False
+    eq.tokamak.set_coil_current("Solenoid", 15000)
+    eq.tokamak.setControlCurrents(np.load(STATIC_CURRENT_BASELINE))
+
+    solver = GSstaticsolver.NKGSsolver(eq, gs_operator_order=2)
+    solver.forward_solve(eq, profiles, 1e-8, suppress=True)
+
+    reference_psi = np.load(STATIC_PSI_BASELINE)
+    tolerance = np.ptp(reference_psi) * 0.003
+    assert solver.gs_operator_order == 2
+    assert np.allclose(eq.psi(), reference_psi, atol=tolerance)
+
+
+def test_static_solver_rejects_invalid_operator_order(create_machine):
+    """Only the two finite-difference operators supplied by FreeGS4E are valid."""
+    eq, _, _ = create_machine
+
+    from freegsnke import GSstaticsolver
+
+    with pytest.raises(ValueError, match="gs_operator_order"):
+        GSstaticsolver.NKGSsolver(eq, gs_operator_order=3)