Merge pull request #438 from hiddenSymmetries/cbs/coilset_reducedcoil…

…set_coilnormalfield

Single-stage optimization with spec: CoilSet, ReducedCoilSet and CoilNormalField (PR 2 of 2)

src/simsopt/field/__init__.py

@@ -1,6 +1,7 @@
 from .biotsavart import *
 from .boozermagneticfield import *
 from .coil import *
+from .coilset import *
 from .magneticfield import *
 from .magneticfieldclasses import *
 from .mgrid import *
@@ -13,6 +14,7 @@
     biotsavart.__all__
     + boozermagneticfield.__all__
     + coil.__all__
+    + coilset.__all__
     + magneticfield.__all__
     + magneticfieldclasses.__all__
     + mgrid.__all__

src/simsopt/field/coil.py

@@ -10,7 +10,8 @@
 
 __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']
+           'coils_to_makegrid', 'coils_to_focus'
+           ]
 
 
 class Coil(sopp.Coil, Optimizable):

src/simsopt/field/normal_field.py

@@ -4,6 +4,7 @@
 
 from .._core.optimizable import DOFs, Optimizable
 from simsopt.geo import SurfaceRZFourier
+from .coilset import CoilSet
 
 logger = logging.getLogger(__name__)
 
@@ -13,7 +14,7 @@
     py_spec = None
     logger.debug(str(e))
 
-__all__ = ['NormalField']
+__all__ = ['NormalField', 'CoilNormalField']
 
 
 class NormalField(Optimizable):
@@ -189,7 +190,7 @@ def get_dofs(self):
                     dofs[ii] = self.vnc[mm, input_ntor+nn]
         return dofs
 
-    def get_index_in_array(self, m, n, mpol=None, ntor=None, even=False):
+    def get_index_in_array(self, m, n, mpol=None, ntor=None):
         """
         Returns position of mode (m,n) in vns or vnc array
 
@@ -212,8 +213,6 @@ def get_index_in_array(self, m, n, mpol=None, ntor=None, even=False):
             raise ValueError('n out of bound')
         if m == 0 and n < 0:
             raise ValueError('n has to be positive if m==0')
-        if not even and m == 0 and n == 0:
-            raise ValueError('m=n=0 not supported for odd series')
 
         return [m, n+ntor]
 
@@ -428,7 +427,9 @@ def get_vnc_asarray(self, mpol=None, ntor=None):
         elif ntor > self.ntor:
             raise ValueError('ntor out of bound')
 
-        vnc = self.vns
+        vnc = self.vnc
+        if vnc is None:
+            vnc = np.zeros((mpol, 2*ntor+1))
 
         return vnc[0:mpol, self.ntor-ntor:self.ntor+ntor+1]
 
@@ -509,8 +510,238 @@ def get_real_space_field(self):
         provided field on the computational boundary. 
         The returned array will be of size specified by the surfaces'  quadpoints and located on the quadpoints. 
         """
-        vns, vnc = self.get_vns_vnc_asarray(mpol=self.mpol, ntor=self.ntor)
+        vns, vnc = self.get_vns_vnc_asarray()
         BdotN_unnormalized = self.surface.inverse_fourier_transform_scalar(vns, vnc, normalization=(2*np.pi)**2, stellsym=self.stellsym)
         normal_field_real_space = -1 * BdotN_unnormalized / np.linalg.norm(self.surface.normal(), axis=-1)
         return normal_field_real_space
 
+class CoilNormalField(NormalField):
+    """
+    A SPEC NormalField generated by a CoilSet. 
+    
+    The CoilNormalField provides the same interface as the 
+    NormalField, but its degrees of freedom are inherited
+    from its CoilSet parent. 
+
+    Args:
+        coilset: The CoilSet object from which to inherit the degrees of freedom        
+
+    Properties:
+        surface: The computational boundary of the SPEC simulation, 
+        that is managed by the CoilSet. 
+        vns/vnc: fourier harmonics of the normal field. 
+        This property is cached, and recomputed only when the parents' DOFS (the
+        coils) change.
+    """
+    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
+        else:  
+            self._coilset = CoilSet()
+
+        self.nfp = self.surface.nfp
+        self.stellsym = self.surface.stellsym
+        self.mpol = self.surface.mpol
+        self.ntor = self.surface.ntor   
+        Optimizable.__init__(self, depends_on=[self._coilset])  # call the Optimizable constructor, skip the NormalField constructor
+
+    @property
+    def surface(self): 
+        return self._coilset.surface
+
+    @surface.setter
+    def surface(self, boundary):
+        self._coilset.surface = boundary
+
+    @classmethod
+    def from_spec_object(cls, *args, **kwargs):
+        raise ValueError('Generate the coils separately and pass them to the CoilNormalField constructor')
+
+    @classmethod
+    def from_spec(cls, *args, **kwargs):
+        raise ValueError('Generate the coils separately and pass them to the CoilNormalField constructor')
+
+    @property
+    def vns(self):
+        if self._vns is None:
+            bnormal = np.sum(self.coilset.bs.B().reshape((self.surface.quadpoints_phi.size, self.surface.quadpoints_theta.size, 3)) * self.surface.normal()*-1, axis=2)
+            Vns, Vnc = self.surface.fourier_transform_scalar(bnormal[:, :], normalization=(2*np.pi)**2, stellsym=self.stellsym)
+            self._vns = Vns
+            self._vnc = Vnc
+        return self._vns
+
+    @vns.setter
+    def vns(self, *args, **kwargs):
+        raise AttributeError('you cannot set vns, the coils do this!')
+
+    @property
+    def vnc(self):
+        if self._vnc is None:
+            bnormal = np.sum(self.coilset.bs.B().reshape((self.surface.quadpoints_phi.size, self.surface.quadpoints_theta.size, 3)) * self.surface.normal()*-1, axis=2)
+            Vns, Vnc = self.surface.fourier_transform_scalar(bnormal[:, :], normalization=(2*np.pi)**2, stellsym=self.stellsym)
+            self._vns = Vns
+            self._vnc = Vnc
+        return self._vnc
+
+    @vnc.setter
+    def vnc(self, *args, **kwargs):
+        raise AttributeError('you cannot set vnc, the coils do this!')
+
+    @property
+    def coilset(self):
+        return self._coilset
+
+    @coilset.setter  
+    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 Re:w
+        ducedCoilSet keeping the first nsv singular values.
+
+        Note: Should this be done by proc0 and broadcast? Is SVD deterministic?
+        """
+        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) # dummy CoilNormalField to evaluate the vnc and vnc
+            output = cnf.vns.ravel()[coilset.surface.ntor+1:] #remove leading zeros
+            if not coilset.surface.stellsym:
+                output = np.append(output, cnf.vnc.ravel()[coilset.surface.ntor:]) #remove leading zeros
+            return np.ravel(output)
+
+        reduced_coilset = thiscoilset.reduce(target_function, nsv=nsv)
+        logger.info(f'CoilNormalField replaced Coilset with ReducedCoilsSet with {reduced_coilset.nsv} singular values')
+        logger.debug('first right-singular vector: ')
+        logger.debug(reduced_coilset.rsv[0])
+        logger.debug('singular values: ')
+        logger.debug(reduced_coilset._s_diag)
+        self.coilset = reduced_coilset  # using setter; replaces parent
+
+    def recompute_bell(self, parent=None):  # Should parent be CoilSet?
+        self._vnc = None
+        self._vns = None
+
+    def get_vns(self, m, n):
+        self.check_mn(m, n)
+        index = self.get_index_in_array(m, n)
+        return self.vns[index[0], index[1]]  # calls cache'd getter
+
+    def get_vnc(self, m, n):
+        self.check_mn(m, n)
+        index = self.get_index_in_array(m, n)
+        return self.vnc[index[0], index[1]]  # calls cache'd getter
+
+    def set_vns(self, *args, **kwargs):
+        raise AttributeError('you cannot set vns, the coils do this!')
+
+    def set_vnc(self, *args, **kwargs):
+        raise AttributeError('you cannot set vnc, the coils do this!')
+
+    def get_vns_asarray(self):
+        return self.vns
+
+    def set_vns_asarray(self, *args, **kwargs):
+        raise AttributeError('you cannot set vns, the coils do this!')
+
+    def get_vnc_asarray(self):
+        return self.vnc
+
+    def set_vnc_asarray(self, *args, **kwargs):
+        raise AttributeError('you cannot set vnc, the coils do this!')
+
+    def get_vns_vnc_asarray(self):
+        return self.vns, self.vnc
+
+    def set_vns_vnc_asarray(self, *args, **kwargs):
+        raise AttributeError('you cannot set fourier components, the coils do this!')
+
+    def change_resolution(self, *args, **kwargs):
+        raise ValueError('CoilNormalField has no resolution, change parameters in its coilset')
+
+    def fixed_range(self, *args, **kwargs):
+        raise ValueError('no sense in fixing anything in a CoilNormalField')
+
+    def get_index_in_array(self, m, n, mpol=None, ntor=None):
+        """
+        Returns position of mode (m,n) in vns or vnc array
+
+        Args:
+        - m: poloidal mode number
+        - n: toroidal mode number (normalized by Nfp)
+        - mpol: resolution of dofs array. If None (by default), use self.mpol
+        - ntor: resolution of dofs array. If None (by default), use self.ntor
+        - even: set to True to get vnc. Default is False
+        """
+
+        if mpol is None:
+            mpol = self.mpol
+        if ntor is None:
+            ntor = self.ntor
+
+        if m < 0 or m > mpol:
+            raise ValueError('m out of bound')
+        if abs(n) > ntor:
+            raise ValueError('n out of bound')
+        if m == 0 and n < 0:
+            raise ValueError('n has to be positive if m==0')
+
+        return [m, n+ntor]
+
+    def get_dofs(self, *args, **kwargs):
+        raise ValueError('CoilNormalField does not have its own degrees of freedom')
+
+    def get_index_in_dofs(self, *args, **kwargs):
+        raise ValueError('CoilNormalField does not have its own degrees of freedom')
+
+    def optimize_coils(self, targetvns, targetvnc=None, TARGET_LENGTH=1000, MAXITER=300):
+        r"""
+        Simple convenience function to 
+        optimize the coils to match the target vns and vnc using a FOCUS-style algorithm.
+
+        Uses the simplest FOCUS optimization consisting of only
+        the quadratic flux penalty and a length penalty. 
+
+        Args: 
+            targetvns: The target odd fourier modes of :math:`\mathbf{B}\cdot\mathbf{\vec{n}}`. 2D array of size
+                (mpol+1)x(2ntor+1). 
+            targetvnc: The target even fourier modes of :math:`\mathbf{B}\cdot\mathbf{\vec{n}}`. 2D array of size
+                (mpol+1)x(2ntor+1). Ignored if stellsym if True. 
+            TARGET_LENGTH: The target length of the coils. Default is 1000. 
+            MAXITER: The maximum number of iterations. Default is 1000.
+        returns:
+            res: the result object from scipy.optimize.minimize
+        """
+        from scipy.optimize import minimize
+        if targetvnc is None:
+            targetvnc = np.zeros_like(targetvns)
+        BdotN_unnormalized = self.surface.inverse_fourier_transform_scalar(targetvns, targetvnc, normalization=(2*np.pi)**2, stellsym=self.stellsym)
+        target = -1 * BdotN_unnormalized / np.linalg.norm(self.surface.normal(), axis=-1)
+        JF = self.coilset.flux_penalty(target=target)\
+            + self.coilset.length_penalty(TOTAL_LENGTH=TARGET_LENGTH, f='max')
+
+        def fun(dofs):
+            JF.x = dofs
+            return JF.J(), JF.dJ()
+
+        dofs = JF.x
+
+        res = minimize(fun, dofs, jac=True, method='L-BFGS-B',
+                       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))}')
+        return res
+

src/simsopt/geo/surface.py

@@ -841,6 +841,22 @@ def interpolate_on_arclength_grid(self, function, theta_evaluate):
             function_interpolated[iphi, :] = f(theta_evaluate[iphi, :])
 
         return function_interpolated
+
+    @property
+    def deduced_range(self):
+        """
+        The quadpoints of a surface can be anything, but are often set to 
+        'full torus', 'field period' or 'half period'. 
+        Since this is not stored in the object, but often useful to know
+        this function deduces the range from the quadpoints
+        """
+        if np.isclose(self.quadpoints_phi[-1], 1-1/len(self.quadpoints_phi), atol=1e-10):
+            return Surface.RANGE_FULL_TORUS
+        elif self.quadpoints_phi[0] == 0:
+            return Surface.RANGE_FIELD_PERIOD
+        else:
+            return Surface.RANGE_HALF_PERIOD
+