Remove expensive equilibrium deepcopying

freegsnke/GSstaticsolver.py

@@ -769,7 +769,7 @@ def optimize_currents(
             currents = np.copy(self.dummy_current)
             currents[i] = 1.0 * delta_current[i]
             currents = full_current_vec + constrain.rebuild_full_current_vec(currents)
-            self.eq2 = deepcopy(eq)
+            self.eq2 = eq.create_auxiliary_equilibrium()
             self.eq2.tokamak.set_all_coil_currents(currents)
             self.forward_solve(
                 eq=self.eq2,

freegsnke/build_machine.py

@@ -587,6 +587,21 @@ def build_active_coil_dict(active_coils):
     return coils_dict
 
 
+def copy_tokamak(tokamak: Machine):
+    new_tokamak = tokamak.copy()
+
+    new_tokamak.coils_dict = tokamak.coils_dict.copy()
+    new_tokamak.coils_list = tokamak.coils_list[::]
+    new_tokamak.n_active_coils = tokamak.n_active_coils
+    new_tokamak.n_passive_coils = tokamak.n_passive_coils
+    new_tokamak.n_coils = tokamak.n_coils
+
+    # add probe object attribute to tokamak (not strictly required)
+    new_tokamak.probes = tokamak.probes
+
+    return new_tokamak
+
+
 if __name__ == "__main__":
     for coil_name in active_coils:
         print([pol for pol in active_coils[coil_name]])

freegsnke/equilibrium_update.py

@@ -28,6 +28,7 @@
 from scipy import interpolate
 
 from . import limiter_func, virtual_circuits
+from .build_machine import copy_tokamak
 
 
 class Equilibrium(freegs4e.equilibrium.Equilibrium):
@@ -62,6 +63,58 @@ def __init__(self, *args, **kwargs):
             float
         )
 
+    def create_auxiliary_equilibrium(self):
+        """Creates the auxiliary equilibrium object.
+
+        The auxiliary object returned from this method is essentially
+        a copy of the equilibrium object (self) however it is manually
+        setup and so won't contain all attributes on self (especially custom
+        attributes). It is NOT _guaranteed_ to be the same as a deepcopy, or even
+        a shallow copy.
+        """
+        # __new__ stops __init__ being called.
+        # This is necessary because the __init__ method does expensive
+        # calculations which we can just copy the results of
+        equilibrium = Equilibrium.__new__(Equilibrium)
+
+        # attributes that FreeGS4e sets
+        equilibrium.tokamak = copy_tokamak(self.tokamak)
+        equilibrium.Rmin = self.Rmin
+        equilibrium.Rmax = self.Rmax
+        equilibrium.Zmin = self.Zmin
+        equilibrium.Zmax = self.Zmax
+        equilibrium.nx = self.nx
+        equilibrium.ny = self.ny
+        equilibrium.dR = self.dR
+        equilibrium.dZ = self.dZ
+        equilibrium._applyBoundary = self._applyBoundary
+        equilibrium._pgreen = self._pgreen
+        equilibrium._vgreen = self._vgreen
+        equilibrium._current = self._current
+        equilibrium.order = self.order
+        equilibrium._solver = self._solver
+
+        # attributes the FreeGSNKE sets
+        equilibrium.solved = self.solved
+        equilibrium.psi_func_interp = self.psi_func_interp
+        equilibrium.nxh = self.nxh
+        equilibrium.nyh = self.nyh
+        equilibrium.Rnxh = self.Rnxh
+        equilibrium.Znyh = self.Znyh
+        equilibrium.limiter_handler = self.limiter_handler  # should be safe not to copy
+
+        # attributes that actually need to be copied
+        equilibrium.R_1D = np.copy(self.R_1D)
+        equilibrium.Z_1D = np.copy(self.Z_1D)
+        equilibrium.R = np.copy(self.R)
+        equilibrium.Z = np.copy(self.Z)
+        equilibrium.tokamak_psi = np.copy(self.tokamak_psi)
+        equilibrium.plasma_psi = np.copy(self.plasma_psi)
+        equilibrium.mask_inside_limiter = np.copy(self.mask_inside_limiter)
+        equilibrium.mask_outside_limiter = np.copy(self.mask_outside_limiter)
+
+        return equilibrium
+
     def adjust_psi_plasma(
         self,
     ):

freegsnke/nonlinear_solve.py

@@ -197,9 +197,9 @@ def __init__(
         print("-----")
 
         # set internal copy of the equilibrium and profile
-        self.eq1 = deepcopy(eq)
+        self.eq1 = eq.create_auxiliary_equilibrium()
         self.profiles1 = deepcopy(profiles)
-        self.eq2 = deepcopy(eq)
+        self.eq2 = eq.create_auxiliary_equilibrium()
         self.profiles2 = deepcopy(profiles)
         self.Iy = self.limiter_handler.Iy_from_jtor(profiles.jtor).copy()
         self.nIy = np.linalg.norm(self.Iy)
@@ -2016,7 +2016,7 @@ def initialize_from_ICs(
         # set internal copy of the equilibrium and profile
         # note that at this stage, the equilibrium may have vessel currents.
         # These can not be reproduced exactly if modes are truncated.
-        self.eq1 = deepcopy(eq)
+        self.eq1 = eq.create_auxiliary_equilibrium()
         self.profiles1 = deepcopy(profiles)
         # The pair self.eq1 and self.profiles1 is the pair that is advanced at each timestep.
         # Their properties evolve according to the dynamics.
@@ -2043,7 +2043,7 @@ def initialize_from_ICs(
         # self.eq2 and self.profiles2 are used as auxiliary objects when solving for the dynamics
         # They are used for all intermediate calculations, so
         # they should not be used to extract properties of the evolving equilibrium
-        self.eq2 = deepcopy(self.eq1)
+        self.eq2 = self.eq1.create_auxiliary_equilibrium()
         self.profiles2 = deepcopy(self.profiles1)
 
         # self.Iy is the istantaneous 1d vector representing the plasma current distribution
@@ -2125,7 +2125,8 @@ def step_complete_assign(self, working_relative_tol_GS, from_linear=False):
 
         if from_linear:
             self.profiles1 = deepcopy(self.profiles2)
-            self.eq1 = deepcopy(self.eq2)
+            self.eq1 = self.eq2
+            self.eq2 = self.eq1.create_auxiliary_equilibrium()
         else:
             self.eq1.plasma_psi = np.copy(self.trial_plasma_psi)
             self.profiles1.Ip = self.trial_currents[-1] * self.plasma_norm_factor

freegsnke/passive_structure.py

@@ -15,9 +15,9 @@
 it under the terms of the GNU Lesser General Public License as published by
 the Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.
-  
+
 You should have received a copy of the GNU Lesser General Public License
-along with FreeGSNKE.  If not, see <http://www.gnu.org/licenses/>.   
+along with FreeGSNKE.  If not, see <http://www.gnu.org/licenses/>.
 """
 
 import freegs4e
@@ -83,6 +83,36 @@ def __init__(
 
         self.greens = {}
 
+    def copy(self):
+        # dont instantiate the new object, it will be slow
+        new_obj = type(self).__new__(type(self))
+
+        new_obj.turns = self.turns
+        new_obj.control = self.turns
+        new_obj.current = self.current
+        new_obj.refine_mode = self.refine_mode
+
+        # ASSUMING the shape will never be modified in-place
+        new_obj.area = self.area
+        new_obj.R = self.R
+        new_obj.Z = self.Z
+        new_obj.Len = self.Len
+        new_obj.Rpolygon = self.Rpolygon
+        new_obj.Zpolygon = self.Zpolygon
+        new_obj.vertices = self.vertices
+        new_obj.polygon = self.polygon
+        new_obj.n_refine = self.n_refine
+        new_obj.filaments = self.filaments
+
+        # This performs a shallow copy of the greens dictionary.
+        # This implicitly assumes that the dictionary might be modified
+        # e.g. self.greens["psi"] = new_array (this would be fine)
+        # but its values WON't be modified in place
+        # e.g. self.greens["psi"][:] = new_array (this would cause problems)
+        new_obj.greens = self.greens.copy()
+
+        return new_obj
+
     def create_RZ_key(self, R, Z):
         """
         Produces tuple (Rmin,Rmax,Zmin,Zmax,nx,ny) to access correct dictionary entry of greens function.

freegsnke/virtual_circuits.py

@@ -646,7 +646,7 @@ def calculate_VC(
 
         # make copies of the newly solved equilibrium and profile objects
         # these are used for all GS solves below
-        self._eq2 = deepcopy(eq)
+        self._eq2 = eq.create_auxiliary_equilibrium()
         self._profiles2 = deepcopy(profiles)
 
         # for each coil, prepare by inferring delta(I_j) corresponding to a change delta(I_y)
@@ -656,7 +656,6 @@ def calculate_VC(
                 print(
                     f"{j}th coil ({coils[j]}) using initial current shift {starting_dI[j]}."
                 )
-            # self._eq2 = deepcopy(eq)
             self.prepare_build_dIydI_j(j, coils, target_dIy, starting_dI[j])
 
         if verbose:
@@ -666,7 +665,6 @@ def calculate_VC(
         # for each coil, build the Jacobian using the value of delta(I_j) inferred earlier
         # by self.prepare_build_dIydI_j.
         for j in np.arange(len(coils)):
-            # self._eq2 = deepcopy(eq)
             # each shape matrix row is derivative of targets wrt the final coil current change
             shape_matrix[:, j] = self.build_dIydI_j(
                 j, coils, targets, targets_options, non_standard_targets, verbose
@@ -779,7 +777,7 @@ def apply_VC(
         )
 
         # store copies of the eq and profile objects
-        eq_new = deepcopy(eq)
+        eq_new = eq.create_auxiliary_equilibrium()
         profiles_new = deepcopy(profiles)
 
         # assign currents to the required coils in the eq object

requirements-freegs4e.txt

@@ -1 +1 @@
-freegs4e~=0.11
+freegs4e~=0.12