From b3615953ee7c0a013db28e1a8a54ef08dd9ec23e Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Wed, 24 Sep 2025 09:37:55 +0100 Subject: [PATCH 1/3] Manually create auxiliary equilibrium --- freegsnke/build_machine.py | 15 ++++++++++ freegsnke/equilibrium_update.py | 53 +++++++++++++++++++++++++++++++++ freegsnke/nonlinear_solve.py | 3 +- freegsnke/passive_structure.py | 32 ++++++++++++++++++++ 4 files changed, 102 insertions(+), 1 deletion(-) diff --git a/freegsnke/build_machine.py b/freegsnke/build_machine.py index 4731c52e..7200235c 100644 --- a/freegsnke/build_machine.py +++ b/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 + 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]]) diff --git a/freegsnke/equilibrium_update.py b/freegsnke/equilibrium_update.py index 47a7553c..6b02603a 100644 --- a/freegsnke/equilibrium_update.py +++ b/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, ): diff --git a/freegsnke/nonlinear_solve.py b/freegsnke/nonlinear_solve.py index afcf04ac..e827cfc1 100644 --- a/freegsnke/nonlinear_solve.py +++ b/freegsnke/nonlinear_solve.py @@ -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 diff --git a/freegsnke/passive_structure.py b/freegsnke/passive_structure.py index 4727ab95..49260bbf 100644 --- a/freegsnke/passive_structure.py +++ b/freegsnke/passive_structure.py @@ -20,6 +20,8 @@ along with FreeGSNKE. If not, see . """ +from copy import deepcopy + import freegs4e import matplotlib.pyplot as plt import numpy as np @@ -83,6 +85,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. From 57ca4256448903ce8bc20303c6b2844cdf194757 Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Mon, 6 Oct 2025 11:16:01 +0100 Subject: [PATCH 2/3] Use create_auxiliary_equilibrium instead of deepcopying --- freegsnke/GSstaticsolver.py | 2 +- freegsnke/build_machine.py | 4 ++-- freegsnke/nonlinear_solve.py | 8 ++++---- freegsnke/passive_structure.py | 6 ++---- freegsnke/virtual_circuits.py | 6 ++---- 5 files changed, 11 insertions(+), 15 deletions(-) diff --git a/freegsnke/GSstaticsolver.py b/freegsnke/GSstaticsolver.py index 07f20d9b..e37372fd 100644 --- a/freegsnke/GSstaticsolver.py +++ b/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, diff --git a/freegsnke/build_machine.py b/freegsnke/build_machine.py index 7200235c..d742ef02 100644 --- a/freegsnke/build_machine.py +++ b/freegsnke/build_machine.py @@ -590,8 +590,8 @@ def build_active_coil_dict(active_coils): def copy_tokamak(tokamak: Machine): new_tokamak = tokamak.copy() - new_tokamak.coils_dict = tokamak.coils_dict - new_tokamak.coils_list = tokamak.coils_list + 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 diff --git a/freegsnke/nonlinear_solve.py b/freegsnke/nonlinear_solve.py index e827cfc1..835edd55 100644 --- a/freegsnke/nonlinear_solve.py +++ b/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 diff --git a/freegsnke/passive_structure.py b/freegsnke/passive_structure.py index 49260bbf..0a3788f2 100644 --- a/freegsnke/passive_structure.py +++ b/freegsnke/passive_structure.py @@ -15,13 +15,11 @@ 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 . +along with FreeGSNKE. If not, see . """ -from copy import deepcopy - import freegs4e import matplotlib.pyplot as plt import numpy as np diff --git a/freegsnke/virtual_circuits.py b/freegsnke/virtual_circuits.py index 5479ba23..184e007a 100644 --- a/freegsnke/virtual_circuits.py +++ b/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 From da1bdfd2e32e28ecce7d6e82bcb74d053005e1dc Mon Sep 17 00:00:00 2001 From: George Holt Date: Wed, 8 Oct 2025 16:14:57 +0100 Subject: [PATCH 3/3] Bump freegs4e requirements -> ~=0.12 --- requirements-freegs4e.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-freegs4e.txt b/requirements-freegs4e.txt index 2b01e26e..db1fd7c1 100644 --- a/requirements-freegs4e.txt +++ b/requirements-freegs4e.txt @@ -1 +1 @@ -freegs4e~=0.11 +freegs4e~=0.12