Merge pull request #310 from MigMadeira/coils2simsgeo

Method for loading makegrid files

src/simsopt/field/coil.py

@@ -8,7 +8,7 @@
 import simsoptpp as sopp
 
 
-__all__ = ['Coil', 'Current', 'coils_via_symmetries',
+__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']
 
@@ -186,6 +186,40 @@ def coils_via_symmetries(curves, currents, nfp, stellsym):
     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 

src/simsopt/geo/curvexyzfourier.py

@@ -3,11 +3,13 @@
 
 import numpy as np
 import jax.numpy as jnp
+from scipy.fft import rfft
 
 from .curve import Curve, JaxCurve
 from .._core.json import GSONDecoder
 import simsoptpp as sopp
 
+
 __all__ = ['CurveXYZFourier', 'JaxCurveXYZFourier']
 
 
@@ -103,6 +105,90 @@ def load_curves_from_file(filename, order=None, ppp=20, delimiter=','):
             coils[ic].local_x = np.concatenate(dofs)
         return coils
 
+    @staticmethod
+    def load_curves_from_makegrid_file(filename: str, order: int, ppp=20):
+        """
+        This function loads a Makegrid input file containing the Cartesian
+        coordinates for several coils and finds the corresponding Fourier
+        coefficients through an fft. The format is described at
+        https://princetonuniversity.github.io/STELLOPT/MAKEGRID
+
+        Args:
+            filename: file to load.
+            order: maximum mode number in the Fourier series. 
+            ppp: points-per-period: number of quadrature points per period.
+
+        Returns:
+            A list of ``CurveXYZFourier`` objects.
+        """
+
+        with open(filename, 'r') as f:
+            file_lines = f.read().splitlines()[3:] 
+
+        curve_data = []
+        single_curve_data = []
+        for j_line in range(len(file_lines)):
+            vals = file_lines[j_line].split()
+            n_vals = len(vals)
+            if n_vals == 4:
+                float_vals = [float(val) for val in vals[:3]]
+                single_curve_data.append(float_vals)
+            elif n_vals == 6:
+                # This must be the last line of the coil
+                curve_data.append(single_curve_data)
+                single_curve_data = []
+            elif n_vals == 1:
+                # Presumably the line that is just "end"
+                break
+            else:
+                raise RuntimeError("Should not get here")
+
+        coil_data = []
+
+        # Compute the Fourier coefficients for each coil
+        for curve in curve_data:
+            xArr, yArr, zArr = np.transpose(curve)
+
+            curves_Fourier = []
+
+            # Compute the Fourier coefficients
+            for x in [xArr, yArr, zArr]:
+                assert len(x) >= 2*order  # the order of the fft is limited by the number of samples
+                xf = rfft(x) / len(x)
+
+                fft_0 = [xf[0].real]  # find the 0 order coefficient
+                fft_cos = 2 * xf[1:order + 1].real  # find the cosine coefficients
+                fft_sin = -2 * xf[:order + 1].imag  # find the sine coefficients
+
+                combined_fft = np.concatenate([fft_sin, fft_0, fft_cos])
+                curves_Fourier.append(combined_fft)
+
+            coil_data.append(np.concatenate(curves_Fourier))
+
+        coil_data = np.asarray(coil_data)
+        coil_data = coil_data.reshape(6 * len(curve_data), order + 1)  # There are 6 * order coefficients per coil
+        coil_data = np.transpose(coil_data)
+
+        assert coil_data.shape[1] % 6 == 0
+        assert order <= coil_data.shape[0]-1
+
+        num_coils = coil_data.shape[1] // 6
+        coils = [CurveXYZFourier(order*ppp, order) for i in range(num_coils)]
+        for ic in range(num_coils):
+            dofs = coils[ic].dofs_matrix
+            dofs[0][0] = coil_data[0, 6*ic + 1]
+            dofs[1][0] = coil_data[0, 6*ic + 3]
+            dofs[2][0] = coil_data[0, 6*ic + 5]
+            for io in range(0, min(order, coil_data.shape[0] - 1)):
+                dofs[0][2*io+1] = coil_data[io+1, 6*ic + 0]
+                dofs[0][2*io+2] = coil_data[io+1, 6*ic + 1]
+                dofs[1][2*io+1] = coil_data[io+1, 6*ic + 2]
+                dofs[1][2*io+2] = coil_data[io+1, 6*ic + 3]
+                dofs[2][2*io+1] = coil_data[io+1, 6*ic + 4]
+                dofs[2][2*io+2] = coil_data[io+1, 6*ic + 5]
+            coils[ic].local_x = np.concatenate(dofs)
+        return coils
+
 
 def jaxfouriercurve_pure(dofs, quadpoints, order):
     k = len(dofs)//3

tests/field/test_coil.py

@@ -1,14 +1,14 @@
 import unittest
 import json
-
+import os
 import numpy as np
 
 from simsopt.geo.curvexyzfourier import CurveXYZFourier, JaxCurveXYZFourier
 from simsopt.geo.curverzfourier import CurveRZFourier
 from simsopt.geo.curvehelical import CurveHelical
 from simsopt.geo.curve import RotatedCurve
 from simsopt.field.coil import Coil, Current, ScaledCurrent, CurrentSum
-from simsopt.field.coil import coils_to_makegrid, coils_to_focus
+from simsopt.field.coil import coils_to_makegrid, coils_to_focus, load_coils_from_makegrid_file
 from simsopt.field.biotsavart import BiotSavart
 from simsopt._core.json import GSONEncoder, GSONDecoder, SIMSON
 from simsopt.configs import get_ncsx_data
@@ -149,3 +149,46 @@ def test_focus(self):
         stellsym = True
         curves, currents, ma = get_ncsx_data()        
         coils_to_makegrid('coils.test', curves, currents, nfp=3, stellsym=True)
+
+    def test_load_coils_from_makegrid_file(self):     
+        order = 25
+        ppp = 10
+
+        curves, currents, ma = get_ncsx_data(Nt_coils=order, ppp=ppp)  
+        coils_to_makegrid("coils.file_to_load", curves, currents, nfp=1)
+
+        loaded_coils = load_coils_from_makegrid_file("coils.file_to_load", order, ppp)
+
+        gamma = [curve.gamma() for curve in curves]
+        loaded_gamma = [coil.curve.gamma() for coil in loaded_coils]
+        loaded_currents = [coil.current for coil in loaded_coils]
+        coils = [Coil(curve, current) for curve, current in zip(curves, currents)]
+
+        for j_coil in range(len(coils)):
+            np.testing.assert_allclose(
+                currents[j_coil].get_value(),
+                loaded_currents[j_coil].get_value()
+            )
+            np.testing.assert_allclose(curves[j_coil].x, loaded_coils[j_coil].curve.x)
+
+        np.random.seed(1)
+
+        bs = BiotSavart(coils)
+        loaded_bs = BiotSavart(loaded_coils)
+
+        points = np.asarray(17 * [[0.9, 0.4, -0.85]])
+        points += 0.01 * (np.random.rand(*points.shape) - 0.5)
+        bs.set_points(points)
+        loaded_bs.set_points(points)
+
+        B = bs.B()
+        loaded_B = loaded_bs.B()
+
+        np.testing.assert_allclose(B, loaded_B)
+        np.testing.assert_allclose(gamma, loaded_gamma)
+
+        os.remove("coils.file_to_load")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/geo/test_curve.py

@@ -1,6 +1,8 @@
 import logging
 import unittest
 import json
+import os
+
 
 import numpy as np
 
@@ -10,7 +12,10 @@
 from simsopt.geo.curvehelical import CurveHelical
 from simsopt.geo.curve import RotatedCurve, curves_to_vtk
 from simsopt.geo import parameters
-from simsopt.configs.zoo import get_ncsx_data
+from simsopt.configs.zoo import get_ncsx_data, get_w7x_data  
+from simsopt.field.coil import coils_to_makegrid
+from simsopt.geo import CurveLength, CurveCurveDistance
+
 
 try:
     import pyevtk
@@ -464,6 +469,45 @@ def test_serialization(self):
                 with self.subTest(curvetype=curvetype, rotated=rotated):
                     self.subtest_serialization(curvetype, rotated)
 
+    def test_load_curves_from_makegrid_file(self):
+        get_config_functions = [get_ncsx_data, get_w7x_data]
+        order = 10
+        ppp = 4
+
+        for get_config_function in get_config_functions:
+            curves, currents, ma = get_config_function(Nt_coils=order, ppp=ppp)  
+
+            # write coils to MAKEGRID file
+            coils_to_makegrid("coils.file_to_load", curves, currents, nfp=1)
+            loaded_curves = CurveXYZFourier.load_curves_from_makegrid_file("coils.file_to_load", order, ppp)
+
+            assert len(curves) == len(loaded_curves)
+
+            for j in range(len(curves)):
+                np.testing.assert_allclose(curves[j].x, loaded_curves[j].x)
+
+            gamma = [curve.gamma() for curve in curves]
+            loaded_gamma = [curve.gamma() for curve in loaded_curves]
+
+            np.testing.assert_allclose(gamma, loaded_gamma)
+
+            kappa = [np.max(curve.kappa()) for curve in curves]
+            loaded_kappa = [np.max(curve.kappa()) for curve in loaded_curves]
+
+            np.testing.assert_allclose(kappa, loaded_kappa)
+
+            length = [CurveLength(c).J() for c in curves]
+            loaded_length = [CurveLength(c).J() for c in loaded_curves]
+
+            np.testing.assert_allclose(length, loaded_length)
+
+            ccdist = CurveCurveDistance(curves, 0).J()
+            loaded_ccdist = CurveCurveDistance(loaded_curves, 0).J()
+
+            np.testing.assert_allclose(ccdist, loaded_ccdist)
+
+            os.remove("coils.file_to_load")
+
 
 if __name__ == "__main__":
     unittest.main()

tests/geo/test_surface_rzfourier.py

@@ -416,10 +416,10 @@ def test_from_pyQSC(self):
         full_torus = SurfaceRZFourier.from_pyQSC(qsc, r=0.1, ntheta=100, mpol=6, ntor=6)
         full_period = SurfaceRZFourier(mpol=full_torus.mpol, ntor=full_torus.ntor, stellsym=full_torus.stellsym, nfp=full_torus.nfp, quadpoints_phi=phis, quadpoints_theta=thetas)
         full_period.x = full_torus.x
-        
+
         np.testing.assert_allclose(full_torus.rc, full_period.rc)
         np.testing.assert_allclose(full_torus.zs, full_period.zs)
-        
+
         np.random.seed(1)
         # non stell sym for code coverage
         qsc = Qsc(ma.rc, np.insert(ma.zs, 0, 0), rs=np.random.rand(5)*1e-7, zc=np.random.rand(5)*1e-7, nfp=3, etabar=-0.408)