merged master changes

.github/workflows/conda.yml

@@ -15,7 +15,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        platform: [ubuntu-latest,  macos-11]
+        platform: [ubuntu-latest,  macos-12]
         python-version: ["3.9"]
 
     runs-on: ${{ matrix.platform }}

.github/workflows/extensive_test.yml

@@ -98,7 +98,7 @@ jobs:
     - name: Install python dependencies
       run: |
            sudo apt-get install graphviz graphviz-dev
-           pip install wheel numpy scipy f90nml h5py scikit-build cmake qsc sympy pyevtk matplotlib ninja plotly networkx pygraphviz ground bentley_ottmann
+           pip install wheel "numpy<2.0.0" scipy f90nml h5py scikit-build cmake qsc sympy pyevtk matplotlib ninja plotly networkx pygraphviz ground bentley_ottmann
 
     - name: Install booz_xform
       if: contains(matrix.packages, 'vmec') || contains(matrix.packages, 'all')

.github/workflows/non_simd_tests.yml

@@ -84,7 +84,7 @@ jobs:
     - name: Install python dependencies
       run: |
         sudo apt-get install graphviz graphviz-dev
-        pip install wheel numpy scipy f90nml h5py scikit-build cmake qsc sympy pyevtk matplotlib ninja plotly networkx pygraphviz
+        pip install wheel "numpy<2.0.0" scipy f90nml h5py scikit-build cmake qsc sympy pyevtk matplotlib ninja plotly networkx pygraphviz ground bentley_ottmann
 
     - name: Install booz_xform
       run: pip install -v git+https://github.com/hiddenSymmetries/booz_xform

.github/workflows/singularity.yml

@@ -12,6 +12,7 @@ on:
 
   # Do the builds on all pull requests (to test them)
   pull_request: []
+  workflow_dispatch:
 
 jobs:
   build-containers:

.github/workflows/tests.yml

@@ -93,7 +93,7 @@ jobs:
     - name: Install python dependencies
       run: |
         sudo apt-get install graphviz graphviz-dev
-        pip install numpy cmake scikit-build f90nml ninja wheel setuptools sympy qsc pyevtk matplotlib plotly networkx pygraphviz mpi4py  py_spec pyoculus h5py
+        pip install "numpy<2.0.0" cmake scikit-build f90nml ninja wheel setuptools sympy qsc pyevtk matplotlib plotly networkx pygraphviz mpi4py py_spec pyoculus h5py ground bentley_ottmann
 
     - name: Install booz_xform
       run: pip install -v git+https://github.com/hiddenSymmetries/booz_xform

.github/workflows/wheel.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [macos-11, ubuntu-20.04]
+        os: [macos-12, ubuntu-20.04]
 
     steps:
       - uses: actions/checkout@v4

CMakeLists.txt

@@ -119,7 +119,7 @@ set(XTENSOR_USE_TBB 0)
 
 pybind11_add_module(${PROJECT_NAME}
     src/simsoptpp/python.cpp src/simsoptpp/python_surfaces.cpp src/simsoptpp/python_curves.cpp
-    src/simsoptpp/boozerresidual.cpp
+    src/simsoptpp/boozerresidual_py.cpp
     src/simsoptpp/python_magneticfield.cpp src/simsoptpp/python_tracing.cpp src/simsoptpp/python_distance.cpp
     src/simsoptpp/biot_savart_py.cpp
     src/simsoptpp/biot_savart_vjp_py.cpp

LICENSE

@@ -19,3 +19,4 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
+

ci/Dockerfile.ubuntu

@@ -5,7 +5,7 @@ FROM --platform=linux/amd64 ubuntu:22.04 as intermediate
 RUN apt update && \
     DEBIAN_FRONTEND=noninteractive  apt-get install -y  build-essential gfortran git m4 wget cmake ninja-build \
     libopenblas-dev libfftw3-dev libhdf5-dev libhdf5-serial-dev libnetcdf-dev libnetcdff-dev libgl1-mesa-dev \
-    python3-dev python3-pip python3-venv
+    python3-dev python3-pip python3-venv libxrender1
 
 # Install mpich manually
 WORKDIR /src
@@ -40,20 +40,22 @@ RUN python3 -m pip install wheel
 RUN python3 -m venv /venv/
 
 RUN /venv/bin/pip install -U pip
-RUN /venv/bin/python -m pip install numpy scipy jax jaxlib f90nml mpi4py jupyter notebook ipython qsc sympy scikit-build ninja "pybind11[global]" cmake
-RUN /venv/bin/pip install git+https://github.com/zhucaoxiang/f90wrap
+RUN /venv/bin/python -m pip install numpy scipy jax jaxlib f90nml mpi4py jupyter notebook ipython qsc sympy scikit-build ninja "pybind11[global]" cmake f90wrap h5py
 ENV PATH="/venv/bin:${PATH}"
 
 RUN git clone --depth 1 https://github.com/PrincetonUniversity/SPEC.git /src/SPEC && \
     cd /src/SPEC   &&  \
-    /venv/bin/pip install -v . 2>&1 | tee spec_build.log
+    /venv/bin/pip install -v . 2>&1 | tee spec_build.log && \
+    cd Utilities/pythontools && \
+    /venv/bin/pip install -r requirements.txt && \
+    /venv/bin/pip install -v .
 
 RUN git clone --depth 1 https://github.com/hiddenSymmetries/VMEC2000.git /src/VMEC && \
     cd /src/VMEC && \
     cp cmake/machines/ubuntu.json cmake_config_file.json && \
     /venv/bin/pip install -v . 2>&1 | tee vmec_build.log
 
-RUN /venv/bin/pip install h5py pyoculus py_spec
+RUN /venv/bin/pip install pyoculus
 RUN /venv/bin/pip install vtk==9.2.6 pyqt5 matplotlib pyevtk plotly
 RUN /venv/bin/pip install mayavi
 RUN /venv/bin/pip install  git+https://github.com/hiddenSymmetries/booz_xform

ci/singularity.def

@@ -8,7 +8,7 @@ Stage: devel
     libfabric-dev libfabric-bin  \
     libhwloc-dev libpmix-dev libevent-dev \
     libopenblas-dev libfftw3-dev libhdf5-dev libhdf5-serial-dev libnetcdf-dev libnetcdff-dev libgl1-mesa-dev \
-    python3-dev python3-pip python3-venv
+    python3-dev python3-pip python3-venv libxrender1
 
 
 %post
@@ -44,8 +44,7 @@ Stage: devel
     python3 -m venv /venv/
     . /venv/bin/activate
     /venv/bin/pip install -U pip
-    /venv/bin/pip install numpy scipy jax jaxlib f90nml jupyter notebook ipython qsc sympy scikit-build ninja "pybind11[global]" cmake
-    /venv/bin/pip install git+https://github.com/zhucaoxiang/f90wrap
+    /venv/bin/pip install numpy scipy jax jaxlib f90nml jupyter notebook ipython qsc sympy scikit-build ninja "pybind11[global]" cmake f90wrap h5py
     PATH="/usr/local/bin:/venv/bin:${PATH}"
     LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH}"
 
@@ -60,7 +59,10 @@ Stage: devel
     cd SPEC   &&  \
     /venv/bin/python setup.py bdist_wheel && \
     /venv/bin/pip install -v dist/*.whl &&  \
-    cd .. && \
+    cd Utilities/pythontools && \
+    /venv/bin/pip install -r requirements.txt && \
+    /venv/bin/pip install . && \
+    cd ../../.. && \
     rm -rf SPEC
 
     git clone --depth 1 https://github.com/hiddenSymmetries/VMEC2000.git && \
@@ -70,7 +72,7 @@ Stage: devel
     cd .. && \
     rm -rf VMEC2000
 
-    /venv/bin/pip install h5py pyoculus py_spec
+    /venv/bin/pip install pyoculus
     /venv/bin/pip install vtk==9.2.6 pyqt5 matplotlib pyevtk plotly
     /venv/bin/pip install mayavi
     /venv/bin/pip install  git+https://github.com/hiddenSymmetries/booz_xform

docs/source/containers.rst

@@ -130,7 +130,7 @@ issues associated with Docker containers. Shifter allows to you use
 the simsopt Docker image files hosted on Docker Hub.  Detailed
 instructions for using Shifter can be found at the `NERSC page on the
 simsopt wiki
-<https://github.com/hiddenSymmetries/simsopt/wiki/NERSC-Cori>`_.
+<https://github.com/hiddenSymmetries/simsopt/wiki/NERSC-Perlmutter>`_.
 
 
 .. _singularity_doc:

docs/source/simsopt.geo.rst

@@ -84,6 +84,14 @@ simsopt.geo.curvecws module
    :undoc-members:
    :show-inheritance:
 
+simsopt.geo.curvexyzfouriersymmetries module
+--------------------------------------------
+
+.. automodule:: simsopt.geo.curvexyzfouriersymmetries
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 simsopt.geo.finitebuild module
 ------------------------------
 

examples/2_Intermediate/boozer.py

@@ -50,6 +50,8 @@
 print(f"After LBFGS:   iota={res['iota']:.3f}, tf={tf.J():.3f}, area={s.area():.3f}, ||residual||={np.linalg.norm(boozer_surface_residual(s, res['iota'], res['G'], bs, derivatives=0)):.3e}")
 if "DISPLAY" in os.environ:
     s.plot()
+
+boozer_surface.need_to_run_code = True
 # now drive the residual down using a specialised least squares algorithm
 res = boozer_surface.minimize_boozer_penalty_constraints_ls(tol=1e-10, maxiter=100, constraint_weight=100., iota=res['iota'], G=res['G'], method='manual')
 print(f"After Lev-Mar: iota={res['iota']:.3f}, tf={tf.J():.3f}, area={s.area():.3f}, ||residual||={np.linalg.norm(boozer_surface_residual(s, res['iota'], res['G'], bs, derivatives=0)):.3e}")

examples/2_Intermediate/boozerQA_ls_mpi.py

@@ -1,51 +1,89 @@
-#!/usr/bin/env python3
-from simsopt.geo import SurfaceXYZTensorFourier, SurfaceRZFourier, BoozerSurface, curves_to_vtk, boozer_surface_residual, \
-    ToroidalFlux, Volume, MajorRadius, CurveLength, CurveCurveDistance, NonQuasiSymmetricRatio, Iotas, BoozerResidual, \
-    LpCurveCurvature, MeanSquaredCurvature, ArclengthVariation
+#!/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:
+
+    J = ( \int_S B_nonQA**2 dS )/(\int_S B_QA dS)
+        + \int \int BoozerResidual(varphi, theta)^2 d\varphi d\theta
+        + 0.5*(iota - iota_0)**2
+        + 0.5*(major_radius - target_major_radius)**2
+        + 0.5*max(\sum_{coils} CurveLength - CurveLengthTarget, 0)**2
+        + other coil regularization penalties on curvature, mean squared curvature,
+          coil-to-coil distance, and coil arclength.
+
+We first load surfaces in the NCSX equilibrium, then optimize for QA on them. 
+The first term the objective minimizes the deviation from quasi-axisymmetry on the loaded surfaces.
+The second term minimizes the Boozer residual on the surfaces, which pushes the optimizer to heal
+islands and generalized chaos.  It is typically a good idea to weight this penalty term to be ~ the
+same order of magnitude of the first term in the objective.
+
+The objective also includes penalty terms on the rotational transform, major radius, total coil length, curvature,
+mean squared curvature, and coil arclength.  The rotational transform and major radius penalty ensures 
+that the surfaces' rotational transform and aspect ratio do not stray too far from the value in the initial configuration.
+There are also a few coil regularization penalties to prevent the coils from becoming too complex.  The BFGS 
+optimizer is used, and quasisymmetry is improved substantially on the surface.  Surface solves using the BoozerLS 
+approach can be costly, so this script supports distributing the solves across multiple MPI ranks.
+
+You can change the value of the variable `nsurfaces` below to optimize for nested surfaces and QS on up to 10 surfaces.  
+The BoozerSurface solves can be distributed to Nranks ranks using:
+    mpirun -np Nranks ./boozerQA_ls_mpi.py
+where `nsurfaces` is the number of surfaces your optimizing on.  For example, if you want one surface solve per rank,
+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
+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, coils_via_symmetries
-from simsopt.configs import get_ncsx_data
+from simsopt.field import BiotSavart
 from simsopt.objectives import QuadraticPenalty
-from scipy.optimize import minimize
-import numpy as np
-import os
+from simsopt.util import proc0_print, in_github_actions
 try:
     from mpi4py import MPI
     comm = MPI.COMM_WORLD
     rank = comm.rank
     size = comm.size
 
-    def pprint(*args, **kwargs):
-        if comm.rank == 0:  # only print on rank 0
-            print(*args, **kwargs)
-
 except ImportError:
     comm = None
     size = 1
-    pprint = print
+    rank = 0
 
 # Directory for output
 IN_DIR = "./inputs/input_ncsx/"
 OUT_DIR = "./output/"
 os.makedirs(OUT_DIR, exist_ok=True)
 
-pprint("Running 2_Intermediate/boozerQA.py")
-pprint("================================")
+proc0_print("Running 2_Intermediate/boozerQA_ls_mpi.py")
+proc0_print("=========================================")
+
+base_curves, base_currents, coils, curves, surfaces, boozer_surfaces, ress = load(IN_DIR + "ncsx_init.json")
 
-base_curves, base_currents, coils, curves, surfaces, boozer_surfaces, ress = load(IN_DIR + f"ncsx_init.json")
+# you can optimize for QA on up to 10 surfaces, by changing nsurfaces below.
 nsurfaces = 2
+assert nsurfaces <= 10
+
 surfaces = surfaces[:nsurfaces]
 boozer_surfaces = boozer_surfaces[:nsurfaces]
 ress = ress[:nsurfaces]
-
 for boozer_surface, res in zip(boozer_surfaces, ress):
-    boozer_surface.run_code('ls', res['iota'], res['G'], verbose=False)
+    boozer_surface.run_code(res['iota'], res['G'])
 
 mpi_surfaces = MPIOptimizable(surfaces, ["x"], comm)
 mpi_boozer_surfaces = MPIOptimizable(boozer_surfaces, ["res", "need_to_run_code"], comm)
 
 mrs = [MajorRadius(boozer_surface) for boozer_surface in boozer_surfaces]
+mrs_equality = [len(mrs)*QuadraticPenalty(mr, mr.J(), 'identity') if idx == len(mrs)-1 else 0*QuadraticPenalty(mr, mr.J(), 'identity') for idx, mr in enumerate(mrs)]
 iotas = [Iotas(boozer_surface) for boozer_surface in boozer_surfaces]
 nonQSs = [NonQuasiSymmetricRatio(boozer_surface, BiotSavart(coils)) for boozer_surface in boozer_surfaces]
 brs = [BoozerResidual(boozer_surface, BiotSavart(coils)) for boozer_surface in boozer_surfaces]
@@ -68,7 +106,7 @@ def pprint(*args, **kwargs):
 Jiotas = QuadraticPenalty(mean_iota, IOTAS_TARGET, 'identity')
 JnonQSRatio = MPIObjective(nonQSs, comm, needs_splitting=True)
 JBoozerResidual = MPIObjective(brs, comm, needs_splitting=True)
-Jmajor_radius = MPIObjective([len(mrs)*QuadraticPenalty(mr, mr.J(), 'identity') if idx == 0 else 0*QuadraticPenalty(mr, mr.J(), 'identity') for idx, mr in enumerate(mrs)], comm, needs_splitting=True)
+Jmajor_radius = MPIObjective(mrs_equality, comm, needs_splitting=True)
 
 ls = [CurveLength(c) for c in base_curves]
 Jls = QuadraticPenalty(sum(ls), float(sum(ls).J()), 'max')
@@ -86,15 +124,16 @@ def pprint(*args, **kwargs):
 # let's fix the coil current
 base_currents[0].fix_all()
 
-boozer_surface.surface.to_vtk(OUT_DIR + f"surf_init_{rank}")
-if comm is None or comm.rank == 0:
-    curves_to_vtk(curves, OUT_DIR + f"curves_init")
+if comm is None or rank == 0:
+    curves_to_vtk(curves, OUT_DIR + "curves_init")
+for idx, surface in enumerate(mpi_surfaces):
+    if comm is None or rank == 0:
+        surface.to_vtk(OUT_DIR + f"surf_init_{idx}")
 
 # dictionary used to save the last accepted surface dofs in the line search, in case Newton's method fails
 prevs = {'sdofs': [surface.x.copy() for surface in mpi_surfaces], 'iota': [boozer_surface.res['iota'] for boozer_surface in mpi_boozer_surfaces],
          'G': [boozer_surface.res['G'] for boozer_surface in mpi_boozer_surfaces], 'J': JF.J(), 'dJ': JF.dJ().copy(), 'it': 0}
 
-
 def fun(dofs):
     # initialize to last accepted surface values
     for idx, surface in enumerate(mpi_surfaces):
@@ -103,15 +142,20 @@ def fun(dofs):
         boozer_surface.res['iota'] = prevs['iota'][idx]
         boozer_surface.res['G'] = prevs['G'][idx]
 
-    alldofs = MPI.COMM_WORLD.allgather(dofs)
-    assert np.all(np.norm(alldofs[0]-d) == 0 for d in alldofs)
-
+    # this check makes sure that all ranks have exactly the same dofs
+    if comm is not None:
+        alldofs = comm.allgather(dofs)
+        assert np.all(np.all(alldofs[0]-d == 0) for d in alldofs)
+
     JF.x = dofs
     J = JF.J()
     grad = JF.dJ()
 
-    success = np.all([boozer_surface.res['success'] for boozer_surface in mpi_boozer_surfaces])
-    if not success:
+    # check to make sure that all the surface solves succeeded
+    success1 = np.all([boozer_surface.res['success'] for boozer_surface in mpi_boozer_surfaces])
+    # check to make sure that the surfaces are not self-intersecting
+    success2 = np.all([not surface.is_self_intersecting() for surface in mpi_surfaces])
+    if not (success1 and success2):
         J = prevs['J']
         grad = -prevs['dJ']
         for idx, boozer_surface in enumerate(mpi_boozer_surfaces):
@@ -139,38 +183,40 @@ def callback(x):
     outstr += f"{'nonQS ratio':{width}}" + ", ".join([f'{np.sqrt(nonqs.J()):.6e}' for nonqs in nonQSs]) + "\n"
     outstr += f"{'Boozer Residual':{width}}" + ", ".join([f'{br.J():.6e}' for br in brs]) + "\n"
     outstr += f"{'<ι>':{width}} {mean_iota.J():.6f} \n"
-    outstr += f"{'ι on surfaces':{width}}" + ", ".join([f"{boozer_surface.res['iota']:.6f}" for boozer_surface in boozer_surfaces]) + "\n"
-    outstr += f"{'major radius on surfaces':{width}}" + ', '.join([f'{surface.major_radius():.6f}' for surface in surfaces]) + "\n"
-    outstr += f"{'minor radius on surfaces':{width}}" + ', '.join([f'{surface.minor_radius():.6f}' for surface in surfaces]) + "\n"
-    outstr += f"{'aspect ratio radius on surfaces':{width}}" + ', '.join([f'{surface.aspect_ratio():.6f}' for surface in surfaces]) + "\n"
-    outstr += f"{'volume':{width}}" + ', '.join([f'{surface.volume():.6f}' for surface in surfaces]) + "\n"
-    outstr += f"{'surfaces are self-intersecting':{width}}" + ', '.join([f'{surface.is_self_intersecting()}' for surface in surfaces]) + "\n"
+    outstr += f"{'ι on surfaces':{width}}" + ", ".join([f"{boozer_surface.res['iota']:.6f}" for boozer_surface in mpi_boozer_surfaces]) + "\n"
+    outstr += f"{'major radius on surfaces':{width}}" + ', '.join([f'{surface.major_radius():.6f}' for surface in mpi_surfaces]) + "\n"
+    outstr += f"{'minor radius on surfaces':{width}}" + ', '.join([f'{surface.minor_radius():.6f}' for surface in mpi_surfaces]) + "\n"
+    outstr += f"{'aspect ratio radius on surfaces':{width}}" + ', '.join([f'{surface.aspect_ratio():.6f}' for surface in mpi_surfaces]) + "\n"
+    outstr += f"{'volume':{width}}" + ', '.join([f'{surface.volume():.6f}' for surface in mpi_surfaces]) + "\n"
+    outstr += f"{'surfaces are self-intersecting':{width}}" + ', '.join([f'{surface.is_self_intersecting()}' for surface in mpi_surfaces]) + "\n"
     outstr += f"{'shortest coil to coil distance':{width}} {Jccdist.shortest_distance():.3f} \n"
     outstr += f"{'coil lengths':{width}}" + ', '.join([f'{J.J():5.6f}' for J in ls]) + "\n"
     outstr += f"{'coil length sum':{width}} {sum(J.J() for J in ls):.3f} \n"
     outstr += f"{'max κ':{width}}" + ', '.join([f'{np.max(c.kappa()):.6f}' for c in base_curves]) + "\n"
     outstr += f"{'∫ κ^2 dl / ∫ dl':{width}}" + ', '.join([f'{Jmsc.J():.6f}' for Jmsc in msc_list]) + "\n"
     outstr += "\n\n"
 
-    pprint(outstr)
+    proc0_print(outstr)
     prevs['it'] += 1
 
 
 dofs = JF.x
 callback(dofs)
 
-pprint("""
+proc0_print("""
 ################################################################################
 ### Run the optimization #######################################################
 ################################################################################
 """)
 # 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)
-curves_to_vtk(curves, OUT_DIR + f"curves_opt")
-boozer_surface.surface.to_vtk(OUT_DIR + "surf_opt")
-
-pprint("End of 2_Intermediate/boozerQA_ls.py")
-pprint("================================")
+if comm is None or rank == 0:
+    curves_to_vtk(curves, OUT_DIR + "curves_opt")
+for idx, surface in enumerate(mpi_surfaces):
+    if comm is None or rank == 0:
+        surface.to_vtk(OUT_DIR + f"surf_opt_{idx}")
+
+proc0_print("End of 2_Intermediate/boozerQA_ls_mpi.py")
+proc0_print("========================================")