diff --git a/src/simsopt/field/coil.py b/src/simsopt/field/coil.py index 99772cdac..a5a4136bd 100644 --- a/src/simsopt/field/coil.py +++ b/src/simsopt/field/coil.py @@ -10,7 +10,7 @@ __all__ = ['Coil', 'Current', 'coils_via_symmetries', 'load_coils_from_makegrid_file', 'apply_symmetries_to_currents', 'apply_symmetries_to_curves', - 'coils_to_makegrid', 'coils_to_focus', 'CoilSet'] + 'coils_to_makegrid', 'coils_to_focus', 'CoilSet', 'ReducedCoilSet'] class Coil(sopp.Coil, Optimizable): @@ -391,7 +391,7 @@ def __init__(self, base_coils=None, coils=None, surface=None): super().__init__(depends_on=base_coils) @classmethod - def for_surface(cls, surf, coil_current=1e5, coils_per_period=5, nfp=None, current_constraint="fix_all", **kwargs ): + def for_surface(cls, surf, coil_current=1e5, coils_per_period=5, nfp=None, current_constraint="fix_all", **kwargs): """ Create a CoilSet for a given surface. The coils are created using :obj:`create_equally_spaced_curves` with the given parameters. @@ -461,7 +461,7 @@ def for_spec_equil(cls, spec, coils_per_period=5, current_constraint="fix_all", nfp = spec.nfp use_stellsym = spec.stellsym total_current = spec.poloidal_current_amperes - total_coil_number = coils_per_period * nfp * (1 + int(use_stellsym)) #only coils for half-period + total_coil_number = coils_per_period * nfp * (1 + int(use_stellsym)) # only coils for half-period if spec.freebound: surface = spec.computational_boundary else: @@ -532,7 +532,7 @@ def get_dof_orders(self): """ import re dof_names = self.dof_names - orders = np.zeros(self.dof_size) # dofs which are not Fourier coefficients are treated as zeroth' order. + orders = np.zeros(self.dof_size) # dofs which are not Fourier coefficients are treated as zeroth' order. # test if coils are CurveXYZFourier: if type(self.coils[0].curve) is not CurveXYZFourier: raise ValueError("Coils must be of type CurveXYZFourier") @@ -545,7 +545,6 @@ def get_dof_orders(self): orders[dof_names.index(name)] = order return orders - def flux_penalty(self, target=None): """ Return the penalty function for the quadratic flux penalty on @@ -591,7 +590,6 @@ def cs_distance_penalty(self, DISTANCE_THRESHOLD): curves = [coil.curve for coil in self.coils] return CurveSurfaceDistance(curves, self.surface, DISTANCE_THRESHOLD) - def lp_curvature_penalty(self, CURVATURE_THRESHOLD, p=2): """ Return a penalty function on the curvature of the coils that is @@ -648,7 +646,6 @@ def to_vtk(self, filename, add_biotsavart=True, close=False): pointData = None self.surface.to_vtk(filename+'_surface', extra_data=pointData) - def plot(self, **kwargs): """ Plot the coil's curve. This method is just shorthand for calling @@ -662,9 +659,10 @@ def plot(self, **kwargs): #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) + class ReducedCoilSet(CoilSet): """ - An Optimizable that replaces a CoilSet but has as degrees of + An Optimizable that replaces a CoilSet and has as degrees of freedom linear combinations of the CoilSet degrees of freedom Single-stage optimization is often accompanied with an explosion @@ -684,7 +682,21 @@ class ReducedCoilSet(CoilSet): 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): + def __init__(self, coilset=None, nsv=None, s_diag=None, u_matrix=None, vh_matrix=None): + """ + create a ReducedCoilSet from a CoilSet and the three elements of an SVD of the Jacobian of a mapping. + """ + # create standards if initialized with None + if coilset is None: # need 'is' (idenity comparison), overloaded __eq__ fails here. + coilset = CoilSet() + if nsv is None: + nsv = coilset.dof_size + if s_diag is None: + s_diag = np.ones(nsv) + if u_matrix is None: + u_matrix = np.eye(nsv) + if vh_matrix is None: + vh_matrix = np.eye(nsv) # Reproduce CoilSet's attributes self.coilset = coilset self.base_coils = coilset.base_coils @@ -693,31 +705,136 @@ def __init__(self, coilset, u_matrix, s_diag, vh_matrix): self.bs = coilset.bs self.bs.set_points(self._surface.gamma().reshape((-1, 3))) # Add the SVD matrices + self.nsv = nsv 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) + Optimizable.__init__(self, x0=np.zeros_like(self._s_diag[:nsv]), names=self._make_names(), external_dof_setter=ReducedCoilSet.set_dofs) + @classmethod + def from_function(cls, coilset, target_function, nsv='nonzero', threshold=1e-08): + """ + Reduce an existing CoilSet to a ReducedCoilSet using the SVD of the mapping + given by target_function. + Args: + coilset: The CoilSet to reduce + target_function: a function that takes a CoilSet and maps to a different space that is relevant for the optimization problem. + The SVD will be calculated on the Jacobian of f: coilDOFs -> target + For example a function that calculates the fourier coefficients of the + normal field on the surface. + nsv: The number of singular values to keep, 'nonzero' or 'greaterthan' defaults to 'nonzero'. + If 'nonzero', all nonzero singular values are kept. + If 'greaterthan', all singular values greater than threshold are kept. + + + """ + from simsopt._core.finite_difference import FiniteDifference + from simsopt._core import make_optimizable + function = make_optimizable(target_function, coilset) + fd = FiniteDifference(function.J) + jaccers = fd.jac() + u_matrix, s_diag, vh_matrix = np.linalg.svd(jaccers) + if nsv == 'nonzero': + nsv = len(s_diag) + else: + assert isinstance(nsv, int), "nsv must be an integer or 'nonzero'" + return cls(coilset, nsv, s_diag, u_matrix, vh_matrix) @property def surface(self): - return self.surface + return self._surface + + @property + def rsv(self): + """ + right-singular vectors of the svd + """ + return self._vh_matrix.tolist() + + @property + def lsv(self): + """ + left-singular vectors of the svd + """ + return self._u_matrix.T.tolist() @surface.setter def surface(self, *args, **kwargs): + """ + settiing a new surface requires re-calculating + the SVD matrices. + """ raise ValueError("Not supported yet") def _make_names(self): - names = [f"DOFvector{n}" for n in range(len(self.s_diag))] + names = [f"sv{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): + if len(x) != self.nsv: raise ValueError("Wrong number of DOFs") - self.x = x - self.coilset.x = self._coil_x0 + self._vh_matrix @ (x * self.s_diag) - - + padx = np.pad(x, (0, len(self._coil_x0)-self.nsv), mode='constant', constant_values=0) + pads= np.pad(self._s_diag[:self.nsv], (0, len(self._coil_x0) - self.nsv), mode='constant', constant_values=0) + self.coilset.x = self._coil_x0 + (padx * pads) @ self._vh_matrix # multiply x by singular value so that low singular values have less effect. [could also put in trust region... is this best? or divide by?] + def plot_singular_vector(self, n, eps=1e-4, show_delta_B=True, engine='mayavi', show=False, **kwargs): + """ + Plot the n-th singular vector. + args: + n: the index of the singular vector + eps: the magnitude of the displacement. + show_delta_B: whether to show the change in the magnetic field + engine: the plotting engine to use. Defaults to 'mayavi' + """ + from simsopt.geo import plot, SurfaceRZFourier + from mayavi import mlab + if n > self.nsv: + raise ValueError("n must be smaller than the number of singular values") + singular_vector = np.array(self.rsv[n]) + + plotsurf = SurfaceRZFourier.from_other_surface(self.surface, range=SurfaceRZFourier.RANGE_FULL_TORUS) + if show_delta_B: + bs = self.bs + initial_points = self.bs.get_points_cart_ref() + bs.set_points(plotsurf.gamma().reshape((-1, 3))) + + current_x = np.copy(self.x) + startpositions = [np.copy(coil.curve.gamma()) for coil in self.coilset.coils] + if show_delta_B: + startB = np.copy(np.sum(bs.B().reshape((plotsurf.quadpoints_phi.size, plotsurf.quadpoints_theta.size, 3)) * plotsurf.unitnormal()*-1, axis=2)) + startB = np.concatenate((startB, startB[:1,:]), axis=0) + startB = np.concatenate((startB, startB[:,:1]), axis=1) + + # Perturb the coils by the singular vector + self.coilset.x = self.coilset.x + singular_vector*eps + newpositions = [np.copy(coil.curve.gamma()) for coil in self.coilset.coils] + if show_delta_B: + changedB = np.copy(np.sum(bs.B().reshape((plotsurf.quadpoints_phi.size, plotsurf.quadpoints_theta.size, 3)) * plotsurf.unitnormal()*-1, axis=2)) + # close the plot + changedB = np.concatenate((changedB, changedB[:1,:]), axis=0) + changedB = np.concatenate((changedB, changedB[:,:1]), axis=1) + # plot the displacement vectors + for newcoilpos, startcoilpos in zip(newpositions, startpositions): + diffs = (0.05/eps)* (startcoilpos - newcoilpos) + x = startcoilpos[:, 0] + y = startcoilpos[:, 1] + z = startcoilpos[:, 2] + # enlarge the difference vectors for better visibility + dx = diffs[:, 0] + dy = diffs[:, 1] + dz = diffs[:, 2] + mlab.quiver3d(x, y, z, dx, dy, dz, line_width=4, **kwargs) + + if show_delta_B: + plot([plotsurf,], engine='mayavi', wireframe=False, close=True, scalars=changedB-startB, colormap='Reds', show=False, **kwargs) + else: + plot([plotsurf,], engine='mayavi', wireframe=False, close=True, colormap='Reds', show=False, **kwargs) + # plot the original coils again + self.x = current_x + plot(self.coilset.coils, close=True, engine='mayavi', color=(1,1,1), show=show, **kwargs) + # set bs set points back + if show_delta_B: + bs.set_points(initial_points) + return changedB, startB \ No newline at end of file diff --git a/src/simsopt/field/normal_field.py b/src/simsopt/field/normal_field.py index 85f9cd055..20ed84508 100644 --- a/src/simsopt/field/normal_field.py +++ b/src/simsopt/field/normal_field.py @@ -477,11 +477,12 @@ 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) + 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): """ A SPEC NormalField generated by a CoilSet. @@ -500,25 +501,25 @@ class CoilNormalField(NormalField): This property is cached, and recomputed only when the parents' DOFS (the coils) change. """ - def __init__(self, coilset: 'CoilSet'=None): + def __init__(self, coilset: 'CoilSet' = None): self._vns = None self._vnc = None # Set coilset and boundary: if not given create standard ones. if coilset is not None: - self.coilset = coilset + self._coilset = coilset self.computational_boundary = coilset.surface else: from simsopt.field import CoilSet surface = SurfaceRZFourier() - self.coilset = CoilSet.for_surface(surface) - self.computational_boundary = self.coilset.surface + self._coilset = CoilSet.for_surface(surface) + self.computational_boundary = self._coilset.surface self.nfp = self.computational_boundary.nfp self.stellsym = self.computational_boundary.stellsym self.mpol = self.computational_boundary.mpol self.ntor = self.computational_boundary.ntor - Optimizable.__init__(self, depends_on=[self.coilset]) # call the Optimizable constructor, skip the NormalField constructor + Optimizable.__init__(self, depends_on=[self._coilset]) # call the Optimizable constructor, skip the NormalField constructor @classmethod def from_spec_object(cls, spec, coils_per_period=6, optimize_coils=False, TARGET_LENGTH=1000): @@ -580,6 +581,38 @@ def vnc(self): @vnc.setter def vnc(self): raise AttributeError('you cannot set vnc, the coils do this!') + + @property + def coilset(self): + return self._coilset + + @coilset.setter # TODO: change DOFS and replace parent correctly!! + def coilset(self, coilset): + from simsopt.field import CoilSet, ReducedCoilSet + assert isinstance(coilset, (CoilSet, ReducedCoilSet)) + self.remove_parent(self._coilset) + self._coilset = coilset + self.append_parent(coilset) + self.recompute_bell() + + def reduce_coilset(self, nsv='nonzero'): + """ + Replace the coilset with a ReducedCoilSet keeping the first nsv singular values. + """ + from simsopt.field import ReducedCoilSet + thiscoilset = self.coilset + if type(self.coilset) is ReducedCoilSet: + thiscoilset = self.coilset.coilset + + def target_function(coilset): + cnf = CoilNormalField(coilset) + output = cnf.vns.ravel()[coilset.surface.ntor+1:] + if not coilset.surface.stellsym: + np.append(output, cnf.vnc.ravel()[coilset.surface.ntor:]) + return np.ravel(output) + + reduced_coilset = ReducedCoilSet.from_function(thiscoilset, target_function, nsv=nsv) + self.coilset = reduced_coilset def recompute_bell(self, parent=None): # Should parent be CoilSet? self._vnc = None @@ -662,7 +695,7 @@ def fun(dofs): dofs = JF.x res = minimize(fun, dofs, jac=True, method='L-BFGS-B', - options={'maxiter': MAXITER, 'maxcor': 300, 'iprint': 5}, tol=1e-15) + options={'maxiter': MAXITER, 'maxcor': 300, 'iprint': 5}, tol=1e-15) print(f'the maximum difference between coil Vns and target Vns is: {np.max(np.abs(self.vns-targetvns))}') 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))}') diff --git a/src/simsopt/mhd/spec.py b/src/simsopt/mhd/spec.py index 998b16b0a..c52378c60 100644 --- a/src/simsopt/mhd/spec.py +++ b/src/simsopt/mhd/spec.py @@ -367,7 +367,7 @@ def boundary(self, boundary): self._boundary = boundary self.append_parent(boundary) return - else: #in freeboundary case, plasma boundary is is not a parent + else: # in freeboundary case, plasma boundary is is not a parent self._boundary = boundary @property @@ -392,7 +392,6 @@ def normal_field(self, normal_field): self.append_parent(normal_field) return - @property def computational_boundary(self): """ @@ -941,8 +940,8 @@ def run(self, update_guess: bool = True): si.zac[0:mn] = self.axis['zac'] # Set initial guess - if self.initial_guess is not None: # note: self.initial_guess is None for workers!! only leaders read and do_stuff with initial_guess - self._set_spec_initial_guess() # workers get the info through a broadcast. this line fails if workers get a guess set + if self.initial_guess is not None: # note: self.initial_guess is None for workers!! only leaders read and do_stuff with initial_guess + self._set_spec_initial_guess() # workers get the info through a broadcast. this line fails if workers get a guess set # write the boundary which is a guess in freeboundary if self.freebound: @@ -1050,7 +1049,7 @@ def run(self, update_guess: bool = True): logger.debug('About to call check_inputs') spec.allglobal.check_inputs() logger.debug('About to call broadcast_inputs') - self.mpi.comm_groups.Barrier() # wait for + self.mpi.comm_groups.Barrier() # wait for spec.allglobal.broadcast_inputs() logger.debug('About to call preset') spec.preset() @@ -1133,7 +1132,7 @@ def run(self, update_guess: bool = True): # Enforce SPEC to use initial guess, but only group leaders have the necessary arrays if self.mpi.proc0_groups: - self.initial_guess = new_guess #set this here? or make whole above monster proc0-exclusive + self.initial_guess = new_guess # set this here? or make whole above monster proc0-exclusive self.inputlist.linitialize = 0 # Group leaders handle deletion of files: @@ -1184,8 +1183,6 @@ def replace_normal_field_with_coils(self, filename=None, TARGET_LENGTH=3000): coil_normal_field = CoilNormalField.from_spec_object(self, optimize_coils=True, TARGET_LENGTH=TARGET_LENGTH) self.normal_field = coil_normal_field - - def array_translator(self, array=None, style='spec'): """ Returns a SpecFourierArray object to help transforming between