Merge branch 'master' into Vmec_without_vmec

docs/source/example_coils.rst

@@ -49,7 +49,7 @@ some classes that will be used::
   import numpy as np
   from scipy.optimize import minimize
   from simsopt.geo.surfacerzfourier import SurfaceRZFourier
-  from simsopt.objectives.fluxobjective import SquaredFlux, CoilOptObjective
+  from simsopt.objectives.fluxobjective import SquaredFlux
   from simsopt.geo.curve import curves_to_vtk, create_equally_spaced_curves
   from simsopt.field.biotsavart import BiotSavart
   from simsopt.field.coil import Current, coils_via_symmetries
@@ -189,7 +189,14 @@ The result is 0.19 Tesla. We now define the objective function::
   Jf = SquaredFlux(s, bs)
   Jls = [CurveLength(c) for c in base_curves]
   Jdist = MinimumDistance(curves, MIN_DIST)
-  objective = CoilOptObjective(Jf, Jls, ALPHA, Jdist, BETA)
+  # Scale and add terms to form the total objective function:
+  objective = Jf + ALPHA * sum(Jls) + BETA * Jdist
+
+In the last line, we have used the fact that the Optimizable objects
+representing the individual terms in the objective can be scaled by a
+constant and added.  (This feature applies to Optimizable objects that
+have a function ``J()`` returning the objective and, if gradients are
+used, a function ``dJ()`` returning the gradient.)
 
 You can check the degrees of freedom that will be varied in the
 optimization by printing the ``dof_names`` property of the objective::

examples/1_Simple/logger_example.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 import logging
-from simsopt import initialize_logging
+from simsopt.util.log import initialize_logging
 
 """
 Example file for transparently logging both MPI and serial jobs
@@ -20,6 +20,7 @@
     comm = MPI.COMM_WORLD
 except:
     comm = None
+    print("MPI not found")
 
 if comm is not None:
     initialize_logging(mpi=True, filename='mpi.log')

examples/2_Intermediate/stage_two_optimization.py

@@ -20,7 +20,7 @@
 import numpy as np
 from scipy.optimize import minimize
 from simsopt.geo.surfacerzfourier import SurfaceRZFourier
-from simsopt.objectives.fluxobjective import SquaredFlux, CoilOptObjective
+from simsopt.objectives.fluxobjective import SquaredFlux
 from simsopt.geo.curve import curves_to_vtk, create_equally_spaced_curves
 from simsopt.field.biotsavart import BiotSavart
 from simsopt.field.coil import Current, coils_via_symmetries
@@ -91,12 +91,16 @@
 Jls = [CurveLength(c) for c in base_curves]
 Jdist = MinimumDistance(curves, MIN_DIST)
 
-JF = CoilOptObjective(Jf, Jls, ALPHA, Jdist, BETA)
-
+# Form the total objective function. To do this, we can exploit the
+# fact that Optimizable objects with J() and dJ() functions can be
+# multiplied by scalars and added:
+JF = Jf + ALPHA * sum(Jls) + BETA * Jdist
 
 # We don't have a general interface in SIMSOPT for optimisation problems that
 # are not in least-squares form, so we write a little wrapper function that we
 # pass directly to scipy.optimize.minimize
+
+
 def fun(dofs):
     JF.x = dofs
     J = JF.J()

examples/2_Intermediate/stage_two_optimization_stochastic.py

@@ -19,16 +19,16 @@
 the latter is independent for each coil.
 """
 
+import os
+from pathlib import Path
+import numpy as np
 from simsopt.geo.surfacerzfourier import SurfaceRZFourier
-from simsopt.objectives.fluxobjective import SquaredFlux, CoilOptObjective
+from simsopt.objectives.fluxobjective import SquaredFlux
 from simsopt.geo.curve import RotatedCurve, curves_to_vtk, create_equally_spaced_curves
 from simsopt.field.biotsavart import BiotSavart
 from simsopt.field.coil import Current, Coil, coils_via_symmetries
 from simsopt.geo.curveobjectives import CurveLength, MinimumDistance
 from simsopt.geo.curveperturbed import GaussianSampler, CurvePerturbed, PerturbationSample
-import numpy as np
-import os
-from pathlib import Path
 
 # Number of unique coil shapes, i.e. the number of coils per half field period:
 # (Since the configuration has nfp = 2, multiply by 4 to get the total number of coils.)
@@ -124,19 +124,23 @@
 Jls = [CurveLength(c) for c in base_curves]
 Jdist = MinimumDistance(curves, MIN_DIST)
 
-JF = CoilOptObjective(Jfs, Jls, ALPHA, Jdist, BETA)
-
+# Form the total objective function. To do this, we can exploit the
+# fact that Optimizable objects with J() and dJ() functions can be
+# multiplied by scalars and added:
+JF = (1.0 / N_SAMPLES) * sum(Jfs) + ALPHA * sum(Jls) + BETA * Jdist
 
 # We don't have a general interface in SIMSOPT for optimisation problems that
 # are not in least-squares form, so we write a little wrapper function that we
 # pass directly to scipy.optimize.minimize
+
+
 def fun(dofs):
     JF.x = dofs
     J = JF.J()
     grad = JF.dJ()
     cl_string = ", ".join([f"{J.J():.3f}" for J in Jls])
     mean_AbsB = np.mean(bs.AbsB())
-    jf = sum(J.J() for J in JF.Jfluxs)/len(JF.Jfluxs)
+    jf = sum(J.J() for J in Jfs) / N_SAMPLES
     print(f"J={J:.3e}, Jflux={jf:.3e}, sqrt(Jflux)/Mean(|B|)={np.sqrt(jf)/mean_AbsB:.3e}, CoilLengths=[{cl_string}], ||∇J||={np.linalg.norm(grad):.3e}")
     return J, grad
 

src/simsopt/_core/derivative.py

@@ -1,7 +1,6 @@
 import numpy as np
 import numbers
 import collections
-from .graph_optimizable import Optimizable
 
 
 class OptimizableDefaultDict(collections.defaultdict):
@@ -14,6 +13,7 @@ def __init__(self, d):
         super().__init__(None, d)
 
     def __missing__(self, key):
+        from .graph_optimizable import Optimizable  # Import here to avoid circular import
         assert isinstance(key, Optimizable)
         self[key] = value = np.zeros((key.local_full_dof_size, ))
         return value
@@ -155,10 +155,14 @@ def __rmul__(self, other):
             x[k] *= other
         return Derivative(x)
 
-    def __call__(self, optim: Optimizable):
+    def __call__(self, optim):
         """
         Get the derivative with respect to all DOFs that ``optim`` depends on.
+
+        Args:
+            optim: An Optimizable object
         """
+        from .graph_optimizable import Optimizable  # Import here to avoid circular import
         assert isinstance(optim, Optimizable)
         deps = optim.ancestors + [optim]
         derivs = []