From f51a558152ac4694c635ad1daf1750ad0b529e07 Mon Sep 17 00:00:00 2001 From: Chris Date: Mon, 8 Jan 2024 15:53:45 +0100 Subject: [PATCH] added CoilSet.get_dof_orders, develop on ReducedCoilSet --- src/simsopt/field/coil.py | 95 +++++++++++++++++++++++++++++-- src/simsopt/field/normal_field.py | 7 ++- 2 files changed, 97 insertions(+), 5 deletions(-) diff --git a/src/simsopt/field/coil.py b/src/simsopt/field/coil.py index 39d6ab493..732c28e29 100644 --- a/src/simsopt/field/coil.py +++ b/src/simsopt/field/coil.py @@ -436,7 +436,6 @@ def from_mgrid_file(cls, makegrid_file, surface): coils = load_coils_from_makegrid_file(makegrid_file, order=6, ppp=20) return cls(base_coils=coils, coils=coils, surface=surface) - @classmethod def for_spec_equil(cls, spec, coils_per_period=5, current_constraint="fix_all", **kwargs): """ @@ -520,6 +519,33 @@ def _cicrclecoils_around_surface(surf, nfp=None, coils_per_period=4, order=6, R0 # take the magnitude of the first-order poloidal Fourier coefficient R1 = np.sqrt(surf.to_RZFourier().get_rc(1, 0)**2 + surf.to_RZFourier().get_zs(1, 0)**2) * factor return create_equally_spaced_curves(coils_per_period, nfp, stellsym=use_stellsym, R0=R0, R1=R1, order=order) + + def get_dof_orders(self): + """ + get the Fourier order of the degrees of freedom corresponding to the Fourier + coefficients describing the coils. + Useful for reducing the trust region for the higher order coefficients. + in optimization. + + Returns: + An array with the fourier order of each degree of freedom. + """ + import re + dof_names = self.dof_names + orders = np.ones(self.dof_size) + # test if coils are CurveXYZFourier: + if type(self.coils[0].curve) is not CurveXYZFourier: + raise ValueError("Coils must be of type CurveXYZFourier") + # run through names, extract order using regular expression + for name in dof_names: + if name.startswith("Curve"): + match = re.search(r'\((\d+)\)', name) + if match: + order = int(match.group(1)) + if order > 0: # leave zeroth order alone + orders[dof_names.index(name)] = order + return orders + def flux_penalty(self, target=None): """ @@ -529,12 +555,13 @@ def flux_penalty(self, target=None): from simsopt.objectives import SquaredFlux return SquaredFlux(self.surface, self.bs, target=target) - def length_penalty(self, TOTAL_LENGTH): + def length_penalty(self, TOTAL_LENGTH, f): """ Return a QuadraticPenalty on the total length of the coils if it is larger than TARGET_LENGTH (do not penalize if shorter) Args: TOTAL_LENGTH: The threshold length above which the penalty is applied + f: type of penalty, "min", "max" or "identity" """ from simsopt.objectives import QuadraticPenalty from simsopt.geo import CurveLength @@ -543,7 +570,7 @@ def length_penalty(self, TOTAL_LENGTH): # summing optimizables makes new optimizables lenth_optimizable = sum(CurveLength(coil.curve) for coil in self.base_coils)*coil_multiplication_factor # return the penalty function - return QuadraticPenalty(lenth_optimizable, TOTAL_LENGTH, "max") + return QuadraticPenalty(lenth_optimizable, TOTAL_LENGTH, f) def cc_distance_penalty(self, DISTANCE_THRESHOLD): """ @@ -634,4 +661,64 @@ def plot(self, **kwargs): #plot the coils, which are already a list: plot(self.coils, **kwargs, show=False) #generate a full torus for plotting, plot it: - SurfaceRZFourier.from_other_surface(self.surface.to_RZFourier(), range=SurfaceRZFourier.RANGE_FULL_TORUS).plot(**kwargs, show=True) \ No newline at end of file + SurfaceRZFourier.from_other_surface(self.surface.to_RZFourier(), range=SurfaceRZFourier.RANGE_FULL_TORUS).plot(**kwargs, show=True) + +class ReducedCoilSet(CoilSet): + """ + An Optimizable that replaces a CoilSet but has as degrees of + freedom linear combinations of the CoilSet degrees of freedom + + Single-stage optimization is often accompanied with an explosion + in the numbers of degrees of freedom needed to accurately represent + the stellarator, as each coilneeds 6 DOFs per coil per Fourier order. + This makes solving the optimization problem harder in two ways: + 1: the computational cost of finite-difference gradient evaluation is + linear in the number of DOFs, and 2: the optimization problem can become + rank-deficient: changes to your DOFs cancel each other, and the minmum + is no longer unique. + + This class solves these problems by allowing physics-based dimensionality + reduction of the optimization space using singluar value-decomposition + of a fast-evalable map. + + A ReducedCoilSet consists of a Coilset, and the three elements of an SVD of + the Jacobian defined by a mapping from the coils' DOFs to a physically relevant, + fast-evaluating target. + """ + def __init__(self, coilset, u_matrix, s_diag, vh_matrix): + # Reproduce CoilSet's attributes + self.coilset = coilset + self.base_coils = coilset.base_coils + self.coils = coilset.coils + self._surface = coilset.surface + self.bs = coilset.bs + self.bs.set_points(self._surface.gamma().reshape((-1, 3))) + # Add the SVD matrices + self._u_matrix = u_matrix + self._vh_matrix = vh_matrix + self._coil_x0 = self.coilset.x + self._s_diag = s_diag + Optimizable.__init__(self, x0=np.zeros_like(self._s_diag), names=self._make_names(), external_dof_setter=ReducedCoilSet.set_dofs) + + + @property + def surface(self): + return self.surface + + @surface.setter + def surface(self, *args, **kwargs): + raise ValueError("Not supported yet") + + def _make_names(self): + names = [f"DOFvector{n}" for n in range(len(self.s_diag))] + return names + + def set_dofs(self, x): + #check if correct! + if len(x) != len(self.s_diag): + raise ValueError("Wrong number of DOFs") + self.x = x + self.coilset.x = self._coil_x0 + self._vh_matrix @ (x * self.s_diag) + + + diff --git a/src/simsopt/field/normal_field.py b/src/simsopt/field/normal_field.py index 68044d29c..85f9cd055 100644 --- a/src/simsopt/field/normal_field.py +++ b/src/simsopt/field/normal_field.py @@ -476,6 +476,11 @@ def set_vns_vnc_asarray(self, vns, vnc, mpol=None, ntor=None): self.set_vns_asarray(vns, mpol, ntor) + def get_real_space_field(self): + vns, vnc = self.get_vns_vnc_asarray(mpol = self.mpol, ntor=self.ntor) + BdotN_unnormalized = self.computational_boundary.inverse_fourier_transform_field(vns, vnc, normalization=(2*np.pi)**2, stellsym=self.stellsym) + normal_field_real_space = -1 * BdotN_unnormalized / np.linalg.norm(self.computational_boundary.normal(), axis=-1) + return normal_field_real_space class CoilNormalField(NormalField): """ @@ -662,4 +667,4 @@ def fun(dofs): print(f'The root mean squared difference between the Vns produced by the coils and the target is: {np.sqrt(np.mean((self.vns-targetvns)**2))}') - \ No newline at end of file +