Merge pull request #449 from hiddenSymmetries/master

Coil_forces

.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<2.0.0" 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 f90wrap
 
     - name: Install booz_xform
       if: contains(matrix.packages, 'vmec') || contains(matrix.packages, 'all')
@@ -125,9 +125,6 @@ jobs:
         pip install -e SPEC/Utilities/pythontools
         python -c "import py_spec; print('success')"
 
-    - name: Install f90wrap
-      run: pip install -U git+https://github.com/zhucaoxiang/f90wrap
-
     - name: Build SPEC python wrapper.
       if: contains(matrix.packages, 'spec') || contains(matrix.packages, 'all')
       run: |
@@ -166,8 +163,9 @@ jobs:
     - name: Install simsopt package
       if: contains(matrix.packages, 'spec') || contains(matrix.packages, 'all')
       run: |
+        pip install -v git+https://github.com/zhisong/pyoculus
         pip install -v .
-        pip install mpi4py py_spec pyoculus h5py
+        pip install mpi4py py_spec h5py
 
     - name: Install simsopt package
       if: contains(matrix.packages, 'none')
@@ -235,13 +233,18 @@ jobs:
         mpiexec -n 2 coverage run -m unittest discover -k "mpi" -s tests -v
         mpiexec -n 3 --oversubscribe coverage run -m unittest discover -k "mpi" -s tests -v
 
+    - name: ls to see coverage files that were produced
+      if: contains(matrix.test-type, 'unit')
+      run: ls -al
+
     - name: Upload uncombined coverage to github
       if: contains(matrix.test-type, 'unit')
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       with:
-        name: gh-actions-parallel-coverage
+        name: gh-actions-parallel-coverage-${{ matrix.packages }}-${{ matrix.python-version }}
         path: .coverage.*
         if-no-files-found: error
+        include-hidden-files: true
 
   coverage:
     runs-on: ubuntu-latest
@@ -263,10 +266,10 @@ jobs:
       run: pip install .
 
     - name: Download artifact
-      uses: actions/download-artifact@v3
+      uses: actions/download-artifact@v4
       with:
-        name: gh-actions-parallel-coverage
         path: ./
+        merge-multiple: true
 
     - name: Display coverage files before combine 
       run: ls -a
@@ -278,7 +281,7 @@ jobs:
         coverage xml
 
     - name: Upload coverage to github
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       with:
         name: tox-gh-actions-coverage
         path: coverage.xml

.github/workflows/tests.yml

@@ -93,7 +93,10 @@ jobs:
     - name: Install python dependencies
       run: |
         sudo apt-get install graphviz graphviz-dev
-        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
+        pip install "numpy<2.0.0" cmake scikit-build f90nml ninja wheel setuptools sympy qsc pyevtk matplotlib plotly networkx pygraphviz mpi4py py_spec h5py ground bentley_ottmann f90wrap
+
+    - name: Install pyoculus
+      run: pip install -v git+https://github.com/zhisong/pyoculus
 
     - name: Install booz_xform
       run: pip install -v git+https://github.com/hiddenSymmetries/booz_xform
@@ -128,9 +131,6 @@ jobs:
         pip install -e SPEC/Utilities/pythontools
         python -c "import py_spec; print('success')"
 
-    - name: Install f90wrap
-      run: pip install -U git+https://github.com/zhucaoxiang/f90wrap
-
     - name: Build SPEC python wrapper.
       run: |
         cd SPEC
@@ -200,6 +200,10 @@ jobs:
         mpiexec -n 2 coverage run -m unittest discover -k "mpi" -s tests -v
         mpiexec -n 3 --oversubscribe coverage run -m unittest discover -k "mpi" -s tests -v
 
+    - name: ls to see coverage files that were produced
+      if: contains(matrix.test-type, 'unit')
+      run: ls -al
+
     - name: Combine coverage reports
       if: contains(matrix.test-type, 'unit')
       run: | 
@@ -209,7 +213,7 @@ jobs:
 
     - name: Upload coverage to github
       if: contains(matrix.test-type, 'unit')
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       with:
         name: tox-gh-actions-coverage
         path: coverage.xml

.github/workflows/wheel.yml

@@ -68,7 +68,7 @@ jobs:
     # alternatively, to publish when a GitHub Release is created, use the following rule:
     # if: github.event_name == 'release' && github.event.action == 'published'
     steps:
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4.1.7
         with:
           name: artifact
           path: dist

ci/Dockerfile.ubuntu

@@ -55,7 +55,7 @@ RUN git clone --depth 1 https://github.com/hiddenSymmetries/VMEC2000.git /src/VM
     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 pyoculus
+RUN /venv/bin/pip install git+https://github.com/zhisong/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

@@ -72,7 +72,7 @@ Stage: devel
     cd .. && \
     rm -rf VMEC2000
 
-    /venv/bin/pip install pyoculus
+    /venv/bin/pip install git+https://github.com/zhisong/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

src/simsopt/solve/mpi.py

@@ -23,6 +23,7 @@
     MPI = None
 
 from .._core.optimizable import Optimizable
+from .._core.util import Struct
 from ..util.mpi import MpiPartition
 from .._core.finite_difference import MPIFiniteDifference
 from ..objectives.least_squares import LeastSquaresProblem
@@ -209,8 +210,13 @@ def _f_proc0(x):
                 x0 = np.copy(prob.x)
                 logger.info("Using finite difference method implemented in "
                             "SIMSOPT for evaluating gradient")
-                result = least_squares(_f_proc0, x0, jac=fd.jac, verbose=2,
-                                       **kwargs)
+                try:
+                    result = least_squares(_f_proc0, x0, jac=fd.jac, verbose=2,
+                                           **kwargs)
+                except:
+                    print("Failure on proc0_world")
+                    result = Struct()
+                    result.x = x0
 
     else:
         leaders_action = lambda mpi, data: None
@@ -230,8 +236,9 @@ def _f_proc0(x):
     if mpi.proc0_world:
         x = result.x
 
-        objective_file.close()
-        if save_residuals:
+        if objective_file is not None:
+            objective_file.close()
+        if save_residuals and residuals_file is not None:
             residuals_file.close()
 
     datalog_started = False

tests/solve/test_mpi.py

@@ -9,6 +9,7 @@
     MPI = None
 
 from simsopt._core.optimizable import Optimizable
+from simsopt._core import ObjectiveFailure
 from simsopt.objectives.least_squares import LeastSquaresProblem
 if MPI is not None:
     from simsopt.util.mpi import MpiPartition
@@ -87,6 +88,12 @@ def f1(self):
     return_fn_map = {'f0': f0, 'f1': f1}
 
 
+class FailingOptimizable(Optimizable):
+    def residuals(self):
+        raise ObjectiveFailure("foo")
+        return self.x - np.array([10, 9, 8, 7])
+
+
 @unittest.skipIf(MPI is None, "Requires mpi4py")
 class MPISolveTests(unittest.TestCase):
 
@@ -137,3 +144,18 @@ def test_parallel_optimization_with_grad(self):
                             self.assertAlmostEqual(prob.x[0], 1)
                             self.assertAlmostEqual(prob.x[1], 1)
 
+    def test_objective_failure_with_mpi(self):
+        """
+        If the objective function fails on the first evaluation, make sure the code does not hang.
+        """
+        with ScratchDir("."):
+            for ngroups in range(1, 4):
+                mpi = MpiPartition(ngroups)
+
+                opt = FailingOptimizable(x0=np.array([5, 6, 7, 8.0]))
+
+                prob = LeastSquaresProblem.from_tuples(
+                    [(opt.residuals, 0, 1)]
+                )
+
+                least_squares_mpi_solve(prob, mpi, grad=True)