Fix area -> aspect ratio in AspectRatio, tidy up boozerQA_ls_mpi.py a…

…nd test_mpi_optimizable

examples/2_Intermediate/boozerQA_ls_mpi.py

@@ -1,26 +1,6 @@
-#!/usr/bin/env python3
-from simsopt.geo import curves_to_vtk, MajorRadius, CurveLength, CurveCurveDistance, NonQuasiSymmetricRatio, Iotas,\
-        BoozerResidual, LpCurveCurvature, MeanSquaredCurvature, ArclengthVariation
-from simsopt._core import load
-from simsopt.objectives import MPIObjective, MPIOptimizable
-from simsopt.field import BiotSavart
-from simsopt.objectives import QuadraticPenalty
-from scipy.optimize import minimize
-from simsopt.util import proc0_print
-import numpy as np
-import os
-try:
-    from mpi4py import MPI
-    comm = MPI.COMM_WORLD
-    rank = comm.rank
-    size = comm.size
-
-except ImportError:
-    comm = None
-    size = 1
-    rank = 0
+#!/usr/bin/env python
 
-"""
+r"""
 This example optimizes the NCSX coils and currents for QA on potentially multiple surfaces using the BoozerLS approach.  
 For a single surface, the objective is:
 
@@ -52,9 +32,33 @@
 and nsurfaces=Nranks=2, then the proper call is:
     mpirun -np 2 ./boozerQA_ls_mpi.py
 
-More details on this work can be found at or doi:10.1063/5.0129716 arxiv:2210.03248.
+More details on this work can be found at
+A Giuliani et al, "Direct stellarator coil optimization for nested magnetic surfaces with precise
+quasi-symmetry", Physics of Plasmas 30, 042511 (2023) doi:10.1063/5.0129716
+or arxiv:2210.03248.
 """
 
+import os
+import numpy as np
+from scipy.optimize import minimize
+from simsopt.geo import curves_to_vtk, MajorRadius, CurveLength, CurveCurveDistance, NonQuasiSymmetricRatio, Iotas,\
+        BoozerResidual, LpCurveCurvature, MeanSquaredCurvature, ArclengthVariation
+from simsopt._core import load
+from simsopt.objectives import MPIObjective, MPIOptimizable
+from simsopt.field import BiotSavart
+from simsopt.objectives import QuadraticPenalty
+from simsopt.util import proc0_print, in_github_actions
+try:
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    rank = comm.rank
+    size = comm.size
+
+except ImportError:
+    comm = None
+    size = 1
+    rank = 0
+
 # Directory for output
 IN_DIR = "./inputs/input_ncsx/"
 OUT_DIR = "./output/"
@@ -67,7 +71,7 @@
 
 # you can optimize for QA on up to 10 surfaces, by changing nsurfaces below.
 nsurfaces = 2
-assert nsurfaces <=10
+assert nsurfaces <= 10
 
 surfaces = surfaces[:nsurfaces]
 boozer_surfaces = boozer_surfaces[:nsurfaces]
@@ -205,8 +209,7 @@ def callback(x):
 ################################################################################
 """)
 # Number of iterations to perform:
-ci = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true']
-MAXITER = 50 if ci else 1e3
+MAXITER = 50 if in_github_actions else 1e3
 
 res = minimize(fun, dofs, jac=True, method='BFGS', options={'maxiter': MAXITER}, tol=1e-15, callback=callback)
 if comm is None or rank == 0:
@@ -215,5 +218,5 @@ def callback(x):
     if comm is None or rank == 0:
         surface.to_vtk(OUT_DIR + f"surf_opt_{idx}")
 
-proc0_print("End of 2_Intermediate/boozerQA_ls.py")
-proc0_print("====================================")
+proc0_print("End of 2_Intermediate/boozerQA_ls_mpi.py")
+proc0_print("========================================")

src/simsopt/geo/surfaceobjectives.py

@@ -18,7 +18,7 @@
 
 class AspectRatio(Optimizable):
     """
-    Wrapper class for surface area label.
+    Wrapper class for surface aspect ratio.
     """
 
     def __init__(self, surface, range=None, nphi=None, ntheta=None):
@@ -46,7 +46,7 @@ def __init__(self, surface, range=None, nphi=None, ntheta=None):
 
     def J(self):
         """
-        Compute the area of a surface.
+        Compute the aspect ratio of a surface.
         """
         return self.surface.aspect_ratio()
 

tests/geo/surface_test_helpers.py

@@ -79,7 +79,7 @@ def get_boozer_surface(label="Volume", nphi=None, ntheta=None, boozer_type='exac
     Returns a boozer surface that will be used in unit tests.
     """
 
-    assert label == "Volume" or label == "ToroidalFlux" or label == "Area" or label == "AspectRatio"
+    assert label in ["Volume", "ToroidalFlux", "Area", "AspectRatio"]
 
     if boozer_type == 'exact':
         assert weight_inv_modB == False

tests/objectives/test_utilities.py

@@ -63,10 +63,8 @@ def test_quadratic_penalty(self):
         with self.assertRaises(Exception):
             self.subtest_quadratic_penalty(curve, J.J()+0.1, 'NotInList')
 
+    @unittest.skipIf(MPI is None, "mpi4py not found")
     def test_mpi_objective(self):
-        if MPI is None:
-            print("skip test_mpi_objective")
-            return
         comm = MPI.COMM_WORLD
 
         c = self.create_curve()
@@ -87,13 +85,11 @@ def test_mpi_objective(self):
             assert abs(Jmpi1.J() - sum(J.J() for J in Js)/n) < 1e-14
             assert np.sum(np.abs(Jmpi1.dJ() - sum(J.dJ() for J in Js)/n)) < 1e-14
 
+    @unittest.skipIf(MPI is None, "mpi4py not found")
     def test_mpi_optimizable(self):
         """
         This test checks that the `x` attribute of the surfaces is correctly communicated across the ranks.
         """
-        if MPI is None:
-            print("skip test_mpi_optimizable")
-            return
 
         comm = MPI.COMM_WORLD
         for size in [1, 2, 3, 4, 5]: