Merge pull request #195 from hiddenSymmetries/ml/profiles

Radial profiles and bootstrap current

Author: landreman

docs/source/optimizable.rst

@@ -184,7 +184,10 @@ also manipulate the fixed/free status of dofs using the functions
   >>> c.x
 
   array([ 1. ,  0.1,  0. , -2. ,  0. ,  0.3,  3. , -0.5,  0.4])
-  
+
+
+.. _dependecies:
+
 Dependencies
 ------------
 

docs/source/simsopt.mhd.rst

@@ -4,6 +4,14 @@ simsopt.mhd package
 Submodules
 ----------
 
+simsopt.mhd.bootstrap module
+----------------------------
+
+.. automodule:: simsopt.mhd.bootstrap
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 simsopt.mhd.boozer module
 -------------------------
 
@@ -12,6 +20,14 @@ simsopt.mhd.boozer module
    :undoc-members:
    :show-inheritance:
 
+simsopt.mhd.profiles module
+---------------------------
+
+.. automodule:: simsopt.mhd.profiles
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 simsopt.mhd.spec module
 -----------------------
 

src/simsopt/geo/surface.py

@@ -604,3 +604,31 @@ def to_vtk(self, filename, h=0.01):
         self.dist.evaluate_batch(RPhiZ, vals)
         vals = vals.reshape(R.shape)
         gridToVTK(filename, X, Y, Z, pointData={"levelset": vals})
+
+
+class SurfaceScaled(Optimizable):
+    """
+    Allows you to take any Surface class and scale the dofs. This is
+    useful for stage-1 optimization.
+    """
+
+    def __init__(self, surf, scale_factors):
+        self.surf = surf
+        self.scale_factors = scale_factors
+        super().__init__(x0=surf.x / scale_factors, names=surf.local_dof_names)
+
+    def recompute_bell(self, parent=None):
+        self.surf.local_full_x = self.local_full_x * self.scale_factors
+
+    def to_RZFourier(self):
+        return self.surf.to_RZFourier()
+
+    def update_fixed(self):
+        """
+        Copy the fixed status from self.surf to self.
+        """
+        for j, is_free in enumerate(self.surf.local_dofs_free_status):
+            if is_free:
+                self.unfix(j)
+            else:
+                self.fix(j)

src/simsopt/geo/surfacerzfourier.py

@@ -76,6 +76,7 @@ def __init__(self, nfp=1, stellsym=True, mpol=1, ntor=0,
         Surface.__init__(self, x0=self.get_dofs(),
                          external_dof_setter=SurfaceRZFourier.set_dofs_impl,
                          names=self._make_names())
+        self._make_mn()
 
     def get_dofs(self):
         """
@@ -88,11 +89,18 @@ def set_dofs(self, dofs):
 
     def _make_names(self):
         """
-        Form a list of names of the `rc`, `zs`, `rs`, or `zc` array elements.
+        Form a list of names of the ``rc``, ``zs``, ``rs``, or ``zc``
+        array elements.  The order of these four arrays here must
+        match the order in ``set_dofs_impl()`` and ``get_dofs()`` in
+        ``src/simsoptpp/surfacerzfourier.h``.
         """
-        names = self._make_names_helper('rc', True) + self._make_names_helper('zs', False)
-        if not self.stellsym:
-            names += self._make_names_helper('rs', False) + self._make_names_helper('zc', True)
+        if self.stellsym:
+            names = self._make_names_helper('rc', True) + self._make_names_helper('zs', False)
+        else:
+            names = self._make_names_helper('rc', True) \
+                + self._make_names_helper('rs', False) \
+                + self._make_names_helper('zc', True) \
+                + self._make_names_helper('zs', False)
         return names
 
     def _make_names_helper(self, prefix, include0):
@@ -106,6 +114,23 @@ def _make_names_helper(self, prefix, include0):
             names += [prefix + '(' + str(m) + ',' + str(n) + ')' for n in range(-self.ntor, self.ntor + 1)]
         return names
 
+    def _make_mn(self):
+        """
+        Make the list of m and n values.
+        """
+        m1d = np.arange(self.mpol + 1)
+        n1d = np.arange(-self.ntor, self.ntor + 1)
+        n2d, m2d = np.meshgrid(n1d, m1d)
+        m0 = m2d.flatten()[self.ntor:]
+        n0 = n2d.flatten()[self.ntor:]
+        m = np.concatenate((m0, m0[1:]))
+        n = np.concatenate((n0, n0[1:]))
+        if not self.stellsym:
+            m = np.concatenate((m, m))
+            n = np.concatenate((n, n))
+        self.m = m
+        self.n = n
+
     @classmethod
     def from_wout(cls,
                   filename: str,
@@ -362,6 +387,7 @@ def change_resolution(self, mpol, ntor):
                 if not self.stellsym:
                     self.rs[m, n + ntor] = old_rs[m, n + old_ntor]
                     self.zc[m, n + ntor] = old_zc[m, n + old_ntor]
+        self._make_mn()
 
         # Update the dofs object
         self._dofs = DOFs(self.get_dofs(), self._make_names())