fix: preserve plasma flux interpolator updates

Author: nicamo

freegsnke/equilibrium_update.py

@@ -65,6 +65,11 @@ def __init__(self, *args, **kwargs):
             float
         )
 
+    def _updatePlasmaPsi(self, plasma_psi):
+        """Update plasma flux while retaining the checked FreeGSNKE interpolator."""
+        super()._updatePlasmaPsi(plasma_psi)
+        self.psi_func_interp = self.__dict__.pop("psi_func")
+
     def update_machine_description(
         self,
         active_coils_data=None,

freegsnke/tests/test_static_solver.py

@@ -175,6 +175,9 @@ 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"
+    assert np.allclose(
+        eq.psi_func(eq.R, eq.Z, grid=False), eq.plasma_psi
+    ), "Plasma-flux interpolator is stale after the solve"
 
 
 def test_second_order_static_solve(create_machine):