Merging in master.

src/simsopt/field/__init__.py

@@ -0,0 +1,21 @@
+from .biotsavart import *
+from .boozermagneticfield import *
+from .coil import *
+from .magneticfield import *
+from .magneticfieldclasses import *
+from .mgrid import *
+from .normal_field import *
+from .tracing import *
+from .magnetic_axis_helpers import *
+
+__all__ = (
+    biotsavart.__all__
+    + boozermagneticfield.__all__
+    + coil.__all__
+    + magneticfield.__all__
+    + magneticfieldclasses.__all__
+    + mgrid.__all__
+    + normal_field.__all__
+    + tracing.__all__
+    + magnetic_axis_helpers.__all__
+)

src/simsopt/field/biotsavart.py

@@ -1,9 +1,8 @@
-import json
 import numpy as np
-from monty.json import MSONable, MontyDecoder, MontyEncoder
 
 import simsoptpp as sopp
 from .magneticfield import MagneticField
+from .._core.json import GSONDecoder
 
 __all__ = ['BiotSavart']
 
@@ -210,16 +209,20 @@ def A_vjp(self, v):
         res_current = [np.sum(v * dA_by_dcoilcurrents[i]) for i in range(len(dA_by_dcoilcurrents))]
         return sum([coils[i].vjp(res_gamma[i], res_gammadash[i], np.asarray([res_current[i]])) for i in range(len(coils))])
 
-    def as_dict(self) -> dict:
-        d = MSONable.as_dict(self)
+    def as_dict(self, serial_objs_dict) -> dict:
+        d = super().as_dict(serial_objs_dict=serial_objs_dict)
         d["points"] = self.get_points_cart()
         return d
 
     @classmethod
-    def from_dict(cls, d):
-        decoder = MontyDecoder()
-        coils = decoder.process_decoded(d["coils"])
+    def from_dict(cls, d, serial_objs_dict, recon_objs):
+        decoder = GSONDecoder()
+        xyz = decoder.process_decoded(d["points"],
+                                      serial_objs_dict=serial_objs_dict,
+                                      recon_objs=recon_objs)
+        coils = decoder.process_decoded(d["coils"],
+                                        serial_objs_dict=serial_objs_dict,
+                                        recon_objs=recon_objs)
         bs = cls(coils)
-        xyz = decoder.process_decoded(d["points"])
         bs.set_points_cart(xyz)
         return bs

src/simsopt/field/boozermagneticfield.py

@@ -7,6 +7,8 @@
 
 logger = logging.getLogger(__name__)
 
+__all__ = ["BoozerMagneticField","BoozerAnalytic","BoozerRadialInterpolant","InterpolatedBoozerField"]
+
 try:
     from mpi4py import MPI
 except ImportError as e:

src/simsopt/field/coil.py

@@ -1,14 +1,16 @@
+from math import pi
+import numpy as np
+
 from simsopt._core.optimizable import Optimizable
 from simsopt._core.derivative import Derivative
 from simsopt.geo.curvexyzfourier import CurveXYZFourier
-from simsopt.geo.curve import RotatedCurve, Curve
+from simsopt.geo.curve import RotatedCurve
 import simsoptpp as sopp
-from math import pi
-import numpy as np
-from monty.json import MontyDecoder, MSONable
 
-__all__ = ['Coil', 'Current', 'coils_via_symmetries',
-           'apply_symmetries_to_currents', 'apply_symmetries_to_curves']
+
+__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']
 
 
 class Coil(sopp.Coil, Optimizable):
@@ -22,7 +24,7 @@ def __init__(self, curve, current):
         self._curve = curve
         self._current = current
         sopp.Coil.__init__(self, curve, current)
-        Optimizable.__init__(self, x0=np.asarray([]), depends_on=[curve, current])
+        Optimizable.__init__(self, depends_on=[curve, current])
 
     def vjp(self, v_gamma, v_gammadash, v_current):
         return self.curve.dgamma_by_dcoeff_vjp(v_gamma) \
@@ -38,16 +40,6 @@ def plot(self, **kwargs):
         """
         return self.curve.plot(**kwargs)
 
-    def as_dict(self) -> dict:
-        return MSONable.as_dict(self)
-
-    @classmethod
-    def from_dict(cls, d):
-        decoder = MontyDecoder()
-        current = decoder.process_decoded(d["current"])
-        curve = decoder.process_decoded(d["curve"])
-        return cls(curve, current)
-
 
 class CurrentBase(Optimizable):
 
@@ -90,24 +82,21 @@ class Current(sopp.Current, CurrentBase):
     of coils that are constrained to use the same current.
     """
 
-    def __init__(self, current, **kwargs):
+    def __init__(self, current, dofs=None, **kwargs):
         sopp.Current.__init__(self, current)
-        CurrentBase.__init__(self, external_dof_setter=sopp.Current.set_dofs,
-                             x0=self.get_dofs(), **kwargs)
+        if dofs is None:
+            CurrentBase.__init__(self, external_dof_setter=sopp.Current.set_dofs,
+                                 x0=self.get_dofs(), **kwargs)
+        else:
+            CurrentBase.__init__(self, external_dof_setter=sopp.Current.set_dofs,
+                                 dofs=dofs, **kwargs)
 
     def vjp(self, v_current):
         return Derivative({self: v_current})
 
-    def as_dict(self) -> dict:
-        d = {}
-        d["@module"] = self.__class__.__module__
-        d["@class"] = self.__class__.__name__
-        d["current"] = self.get_value()
-        return d
-
-    @classmethod
-    def from_dict(cls, d):
-        return cls(d["current"])
+    @property
+    def current(self):
+        return self.get_value()
 
 
 class ScaledCurrent(sopp.CurrentBase, CurrentBase):
@@ -120,24 +109,14 @@ def __init__(self, current_to_scale, scale, **kwargs):
         self.current_to_scale = current_to_scale
         self.scale = scale
         sopp.CurrentBase.__init__(self)
-        CurrentBase.__init__(self, x0=np.asarray([]),
-                             depends_on=[current_to_scale], **kwargs)
+        CurrentBase.__init__(self, depends_on=[current_to_scale], **kwargs)
 
     def vjp(self, v_current):
         return self.scale * self.current_to_scale.vjp(v_current)
 
     def get_value(self):
         return self.scale * self.current_to_scale.get_value()
 
-    def as_dict(self) -> dict:
-        return MSONable.as_dict(self)
-
-    @classmethod
-    def from_dict(cls, d):
-        decoder = MontyDecoder()
-        current = decoder.process_decoded(d["current_to_scale"])
-        return cls(current, d["scale"])
-
 
 class CurrentSum(sopp.CurrentBase, CurrentBase):
     """
@@ -148,24 +127,14 @@ def __init__(self, current_a, current_b):
         self.current_a = current_a
         self.current_b = current_b
         sopp.CurrentBase.__init__(self)
-        CurrentBase.__init__(self, x0=np.asarray([]), depends_on=[current_a, current_b])
+        CurrentBase.__init__(self, depends_on=[current_a, current_b])
 
     def vjp(self, v_current):
         return self.current_a.vjp(v_current) + self.current_b.vjp(v_current)
 
     def get_value(self):
         return self.current_a.get_value() + self.current_b.get_value()
 
-    def as_dict(self) -> dict:
-        return MSONable.as_dict(self)
-
-    @classmethod
-    def from_dict(cls, d):
-        decoder = MontyDecoder()
-        current_a = decoder.process_decoded(d["current_a"])
-        current_b = decoder.process_decoded(d["current_b"])
-        return cls(current_a, current_b)
-
 
 def apply_symmetries_to_curves(base_curves, nfp, stellsym):
     """
@@ -215,3 +184,144 @@ def coils_via_symmetries(curves, currents, nfp, stellsym):
     currents = apply_symmetries_to_currents(currents, nfp, stellsym)
     coils = [Coil(curv, curr) for (curv, curr) in zip(curves, currents)]
     return coils
+
+
+def load_coils_from_makegrid_file(filename, order, ppp=20):
+    """
+    This function loads a file in MAKEGRID input format containing the Cartesian coordinates 
+    and the currents for several coils and returns an array with the corresponding coils. 
+    The format is described at
+    https://princetonuniversity.github.io/STELLOPT/MAKEGRID
+
+    Args:
+        filename: file to load.
+        order: maximum mode number in the Fourier expansion.
+        ppp: points-per-period: number of quadrature points per period.
+
+    Returns:
+        A list of ``Coil`` objects with the Fourier coefficients and currents given by the file.
+    """
+    with open(filename, 'r') as f:
+        all_coils_values = f.read().splitlines()[3:] 
+
+    currents = []
+    flag = True
+    for j in range(len(all_coils_values)-1):
+        vals = all_coils_values[j].split()
+        if flag:
+            currents.append(float(vals[3]))
+            flag = False
+        if len(vals) > 4:
+            flag = True
+
+    curves = CurveXYZFourier.load_curves_from_makegrid_file(filename, order=order, ppp=ppp)
+    coils = [Coil(curves[i], Current(currents[i])) for i in range(len(curves))]
+
+    return coils
+
+
+def coils_to_makegrid(filename, curves, currents, groups=None, nfp=1, stellsym=False):
+    """
+    Export a list of Curve objects together with currents in MAKEGRID input format, so they can 
+    be used by MAKEGRID and FOCUS. The format is introduced at
+    https://princetonuniversity.github.io/STELLOPT/MAKEGRID
+    Note that this function does not generate files with MAKEGRID's *output* format.
+
+    Args:
+        filename: Name of the file to write.
+        curves: A python list of Curve objects.
+        currents: Coil current of each curve.
+        groups: Coil current group. Coils in the same group will be assembled together. Defaults to None.
+        nfp: The number of field periodicity. Defaults to 1.
+        stellsym: Whether or not following stellarator symmetry. Defaults to False.
+    """
+
+    assert len(curves) == len(currents)
+    coils = coils_via_symmetries(curves, currents, nfp, stellsym)
+    ncoils = len(coils)
+    if groups is None:
+        groups = np.arange(ncoils) + 1
+    else:
+        assert len(groups) == ncoils
+        # should be careful. SIMSOPT flips the current, but actually should change coil order
+    with open(filename, "w") as wfile:
+        wfile.write("periods {:3d} \n".format(nfp)) 
+        wfile.write("begin filament \n")
+        wfile.write("mirror NIL \n")
+        for icoil in range(ncoils):
+            x = coils[icoil].curve.gamma()[:, 0]
+            y = coils[icoil].curve.gamma()[:, 1]
+            z = coils[icoil].curve.gamma()[:, 2]
+            for iseg in range(len(x)):  # the last point matches the first one;
+                wfile.write(
+                    "{:23.15E} {:23.15E} {:23.15E} {:23.15E}\n".format(
+                        x[iseg], y[iseg], z[iseg], coils[icoil].current.get_value()
+                    )
+                )
+            wfile.write(
+                "{:23.15E} {:23.15E} {:23.15E} {:23.15E} {:} {:10} \n".format(
+                    x[0], y[0], z[0], 0.0, groups[icoil], coils[icoil].curve.name
+                )
+            )
+        wfile.write("end \n")
+    return
+
+
+def coils_to_focus(filename, curves, currents, nfp=1, stellsym=False, Ifree=False, Lfree=False):
+    """
+    Export a list of Curve objects together with currents in FOCUS format, so they can 
+    be used by FOCUS. The format is introduced at
+    https://princetonuniversity.github.io/FOCUS/rdcoils.pdf
+    This routine only works with curves of type CurveXYZFourier,
+    not other curve types.
+
+    Args:
+        filename: Name of the file to write.
+        curves: A python list of CurveXYZFourier objects.
+        currents: Coil current of each curve.
+        nfp: The number of field periodicity. Defaults to 1.      
+        stellsym: Whether or not following stellarator symmetry. Defaults to False.
+        Ifree: Flag specifying whether the coil current is free. Defaults to False.
+        Lfree: Flag specifying whether the coil geometry is free. Defaults to False.
+    """
+    from simsopt.geo import CurveLength
+
+    assert len(curves) == len(currents)
+    ncoils = len(curves)
+    if stellsym:
+        symm = 2  # both periodic and symmetric
+    elif nfp > 1 and not stellsym:
+        symm = 1  # only periodicity
+    else:
+        symm = 0  # no periodicity or symmetry
+    if nfp > 1:
+        print('Please note: FOCUS sets Nfp in the plasma file.')
+    with open(filename, 'w') as f:
+        f.write('# Total number of coils \n')
+        f.write('  {:d} \n'.format(ncoils))
+        for i in range(ncoils):
+            assert isinstance(curves[i], CurveXYZFourier)
+            nf = curves[i].order
+            xyz = curves[i].full_x.reshape((3, -1))
+            xc = xyz[0, ::2]
+            xs = np.concatenate(([0.], xyz[0, 1::2]))
+            yc = xyz[1, ::2]
+            ys = np.concatenate(([0.], xyz[1, 1::2]))
+            zc = xyz[2, ::2]
+            zs = np.concatenate(([0.], xyz[2, 1::2]))
+            length = CurveLength(curves[i]).J()
+            nseg = len(curves[i].quadpoints)
+            f.write('#------------{:d}----------- \n'.format(i+1))
+            f.write('# coil_type  symm  coil_name \n')
+            f.write('  {:d}   {:d}  {:} \n'.format(1, symm, curves[i].name))
+            f.write('# Nseg current Ifree Length Lfree target_length \n')
+            f.write('  {:d} {:23.15E} {:d} {:23.15E} {:d} {:23.15E} \n'.format(nseg, currents[i].get_value(), Ifree, length, Lfree, length))
+            f.write('# NFcoil \n')
+            f.write('  {:d} \n'.format(nf))
+            f.write('# Fourier harmonics for coils ( xc; xs; yc; ys; zc; zs) \n')
+            for r in [xc, xs, yc, ys, zc, zs]:  # 6 lines
+                for k in range(nf+1):
+                    f.write('{:23.15E} '.format(r[k]))
+                f.write('\n')
+        f.write('\n')
+    return