Merge pull request #451 from daringli/coilgroup_names

Specify coil group names in load_coils_from_makegrid_file

src/simsopt/field/coil.py

@@ -186,7 +186,7 @@ def coils_via_symmetries(curves, currents, nfp, stellsym):
     return coils
 
 
-def load_coils_from_makegrid_file(filename, order, ppp=20):
+def load_coils_from_makegrid_file(filename, order, ppp=20, group_names=None):
     """
     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. 
@@ -197,10 +197,16 @@ def load_coils_from_makegrid_file(filename, order, ppp=20):
         filename: file to load.
         order: maximum mode number in the Fourier expansion.
         ppp: points-per-period: number of quadrature points per period.
+        group_names: List of coil group names (str). Only get coils in coil groups that are in the list.
 
     Returns:
         A list of ``Coil`` objects with the Fourier coefficients and currents given by the file.
     """
+
+    if isinstance(group_names,str):
+        # Handle case of a single string
+        group_names = [group_names]
+
     with open(filename, 'r') as f:
         all_coils_values = f.read().splitlines()[3:] 
 
@@ -209,12 +215,19 @@ def load_coils_from_makegrid_file(filename, order, ppp=20):
     for j in range(len(all_coils_values)-1):
         vals = all_coils_values[j].split()
         if flag:
-            currents.append(float(vals[3]))
+            curr = float(vals[3])
             flag = False
         if len(vals) > 4:
             flag = True
-
-    curves = CurveXYZFourier.load_curves_from_makegrid_file(filename, order=order, ppp=ppp)
+            if group_names is None:
+                currents.append(curr)
+            else:
+                this_group_name = vals[5]
+                if this_group_name in group_names:
+                    currents.append(curr)  
+
+
+    curves = CurveXYZFourier.load_curves_from_makegrid_file(filename, order=order, ppp=ppp, group_names=group_names)    
     coils = [Coil(curves[i], Current(currents[i])) for i in range(len(curves))]
 
     return coils

src/simsopt/geo/curvexyzfourier.py

@@ -105,7 +105,7 @@ def load_curves_from_file(filename, order=None, ppp=20, delimiter=','):
         return coils
 
     @staticmethod
-    def load_curves_from_makegrid_file(filename: str, order: int, ppp=20):
+    def load_curves_from_makegrid_file(filename: str, order: int, ppp=20, group_names = None):
         """
         This function loads a Makegrid input file containing the Cartesian
         coordinates for several coils and finds the corresponding Fourier
@@ -116,6 +116,7 @@ def load_curves_from_makegrid_file(filename: str, order: int, ppp=20):
             filename: file to load.
             order: maximum mode number in the Fourier series. 
             ppp: points-per-period: number of quadrature points per period.
+            group_names: List of coil group names (str). If not 'None', only get coils in coil groups that are in the list.
 
         Returns:
             A list of ``CurveXYZFourier`` objects.
@@ -134,7 +135,12 @@ def load_curves_from_makegrid_file(filename: str, order: int, ppp=20):
                 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)
+                if group_names is None:
+                    curve_data.append(single_curve_data)
+                else:
+                    this_group_name = vals[5]
+                    if this_group_name in group_names:
+                        curve_data.append(single_curve_data)
                 single_curve_data = []
             elif n_vals == 1:
                 # Presumably the line that is just "end"

tests/field/test_coil.py

@@ -14,6 +14,9 @@
 from simsopt._core.json import GSONEncoder, GSONDecoder, SIMSON
 from simsopt.configs import get_ncsx_data
 
+import os
+from . import TEST_DIR
+
 
 def get_curve(curvetype, rotated, x=np.asarray([0.5])):
     np.random.seed(2)
@@ -187,6 +190,31 @@ def test_load_coils_from_makegrid_file(self):
         np.testing.assert_allclose(B, loaded_B)
         np.testing.assert_allclose(gamma, loaded_gamma)
 
+    def test_load_coils_from_makegrid_file_group(self):
+        order = 25
+        ppp = 10
+
+        # Coil group_names is a list of strings
+        filecoils = os.path.join(TEST_DIR, "coils.M16N08")   
+        coils = load_coils_from_makegrid_file(filecoils, order, ppp, group_names = ["245th-coil","100th-coil"])
+        all_coils = load_coils_from_makegrid_file(filecoils, order, ppp)
+        #     NOTE: coils will be returned in order they appear in the file, not in order of listed groups.
+        #     So group_names = ["245th-coil","100th-coil"] gives the array [<coil nr 100>, <coil nr 245>]
+        compare_coils = [all_coils[99],all_coils[244]]
+        gamma = [coil.curve.gamma() for coil in coils]
+        compare_gamma = [coil.curve.gamma() for coil in compare_coils]
+        np.testing.assert_allclose(gamma, compare_gamma)
+
+        # Coil group_names is a single string
+        coils = load_coils_from_makegrid_file(filecoils, order, ppp, group_names = "256th-coil")
+        all_coils = load_coils_from_makegrid_file(filecoils, order, ppp)
+        compare_coils = [all_coils[255]]
+        gamma = [coil.curve.gamma() for coil in coils]
+        compare_gamma = [coil.curve.gamma() for coil in compare_coils]
+        np.testing.assert_allclose(gamma, compare_gamma)
+
+
+
     def test_equally_spaced_planar_curves(self):
         ncoils = 4
         nfp = 4