From 619920c15d7dc4cc94d36de6cf467fbbc0e29df5 Mon Sep 17 00:00:00 2001 From: Matt Date: Sat, 25 Jun 2022 08:22:17 +0200 Subject: [PATCH 1/2] Tweaks to docs --- docs/source/index.rst | 11 +++-------- src/simsopt/configs/zoo.py | 1 + 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/docs/source/index.rst b/docs/source/index.rst index f6172a575..aa54a9198 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -87,15 +87,10 @@ optimization. Others include `STELLOPT .. toctree:: :maxdepth: 3 - :caption: API + :caption: API reference - simsopt - -.. toctree:: - :maxdepth: 3 - :caption: Developer's API - - simsopt + Puiblic functions and classes + For developers Indices and tables ================== diff --git a/src/simsopt/configs/zoo.py b/src/simsopt/configs/zoo.py index 9f72a5348..bab28688e 100644 --- a/src/simsopt/configs/zoo.py +++ b/src/simsopt/configs/zoo.py @@ -12,6 +12,7 @@ def get_ncsx_data(Nt_coils=25, Nt_ma=10, ppp=10): """ Get a configuration that corresponds to the modular coils of the NCSX experiment (circular coils are not included). + Args: Nt_coils: order of the curves representing the coils. Nt_ma: order of the curve representing the magnetic axis. From 533b405385332c81f3bb4c0c73a77651f411a45f Mon Sep 17 00:00:00 2001 From: Matt Date: Sat, 25 Jun 2022 15:25:18 +0200 Subject: [PATCH 2/2] Improvements to docs --- docs/source/example_coils.rst | 281 ++++++++++-------- docs/source/example_quasisymmetry.rst | 13 +- docs/source/example_vmec_only.rst | 10 +- docs/source/fields.rst | 2 +- docs/source/geo.rst | 30 +- docs/source/index.rst | 4 +- docs/source/optimizable.rst | 12 +- docs/source/overview.rst | 14 +- docs/source/tracing.rst | 2 +- .../stage_two_optimization_minimal.py | 146 +++++++++ examples/run_serial_examples | 1 + src/simsopt/field/__init__.py | 3 +- src/simsopt/field/magneticfield.py | 2 + src/simsopt/geo/__init__.py | 2 + src/simsopt/geo/curve.py | 2 +- src/simsopt/geo/surface.py | 2 + 16 files changed, 350 insertions(+), 176 deletions(-) create mode 100755 examples/1_Simple/stage_two_optimization_minimal.py diff --git a/docs/source/example_coils.rst b/docs/source/example_coils.rst index 7ad7bc057..96ed357bc 100644 --- a/docs/source/example_coils.rst +++ b/docs/source/example_coils.rst @@ -1,9 +1,9 @@ Coil optimization ================= -Here we work through an example of "stage 2" coil optimization. In -this approach, we consider that a surface of constant magnetic flux (plasma -boundary shape) has already been found in what is commonly known as +Here we work through several examples of "stage 2" coil optimization. In +this approach, a target plasma shape (i.e. boundary flux surface shape) +has already been found in a "stage 1" optimization, and the present task is to optimize the shapes of coils to produce this target field. For this stage-2 problem, no MHD codes are used (example, VMEC or SPEC), so @@ -12,55 +12,48 @@ there is no need to have them installed. We first describe the basic principles of "stage 2" coil optimization and how they are used in SIMSOPT, namely, how to set up the optimization problem in the minimal example ``examples/2_Intermediate/stage_two_optimization.py``. -We then add increasingly more terms to the objective function in order to illustrate -how to perform various extensions of the minimal example, namely: -- Take into account coil perturbations via systematic and statistical -errors using a stochastic optimization method similar to the example in -``examples/2_Intermediate/stage_two_optimization_stochastic.py``. -- Take into account the presence of the plasma and its respective magnetic field -using the virtual casing principle, similar to the example -in ``examples/2_Intermediate/stage_two_optimization_finite_beta.py``. -- Take into account finite coil width using a multifilament approach -similar to the example in ``examples/3_Advanced/stage_two_optimization_finite_build.py``. - -.. _simple_stage2: - -Simple objective function ---------------------------- - -The first form of the objective function :math:`J` (or cost function) -that will be minimized in order to find coils that yield the desired magnetic field is -the one presented in ``examples/2_Intermediate/stage_two_optimization.py`` and given by:: - - J = (1/2) \int |B dot n|^2 ds - + LENGTH_WEIGHT * \sum_j L_j - + DISTANCE_WEIGHT * MininumDistancePenalty(DISTANCE_THRESHOLD) - + CURVATURE_WEIGHT * CurvaturePenalty(CURVATURE_THRESHOLD) - + MSC_WEIGHT * MeanSquaredCurvaturePenalty(MSC_THRESHOLD) +We then add several advanced features to the problem: +- Including more terms in the objective to regularize the coils, + which may help with engineering feasibility. +- Taking into account coil perturbations via systematic and statistical + errors using a stochastic optimization method similar to the example in + ``examples/2_Intermediate/stage_two_optimization_stochastic.py``. +- Taking into account the presence of plasma current + using the virtual casing principle, similar to the example + in ``examples/2_Intermediate/stage_two_optimization_finite_beta.py``. +- Taking into account finite coil width using a multi-filament approach + similar to the example in ``examples/3_Advanced/stage_two_optimization_finite_build.py``. + + +.. _minimal_stage2: + +Minimal example +--------------- + +A minimal objective function (cost function) for coil optimization is + +.. math:: + + J = \frac{1}{2} \int |\vec{B} \cdot \vec{n}|^2 ds + + (\mathtt{LENGTH\_WEIGHT}) \left(L_* - \sum_j L_j\right)^2. + +This example can also be found in +``examples/1_Simple/stage_two_optimization_minimal.py``. The first term in the right-hand-side term is the "quadratic flux", the area integral over the target plasma surface of the square of the magnetic field normal to the surface. If the coils exactly produce a flux surface of the target shape, this term will vanish. Next, :math:`L_j` -is the length of the :math:`j`-th coil. The scalar regularization -parameter `LENGTH_WEIGHT` is chosen to balance a trade-off: large +is the length of the :math:`j`-th coil. +The number :math:`L_*` +is chosen to balance a trade-off: small values will give smooth coils at the cost of inaccuracies in producing -the target field; small values of `LENGTH_WEIGHT` will give a more +the target field; large values of :math:`L_*` will give a more accurate match to the target field at the cost of complicated coils. - -Next, `MininumDistancePenalty` is a penalty that prevents coils from -becoming too close. For its precise definition, see -:obj:`~simsopt.geo.CurveCurveDistance`. The constant -`DISTANCE_WEIGHT` is selected to balance this minimum distance penalty -against the other objectives. Analytic derivatives are used for the -optimization. - -Finally, the two remaining terms `CurvaturePenalty` and `MeanSquaredCurvaturePenalty` -are regularization terms that prevent the coils from becoming to wiggly. -These try to make the coils as smooth as possible to accomodate -possible engineering constraints. +The weight parameter ``LENGTH_WEIGHT`` provides some additional control over the regularization, +though results should not be very sensitive to it, and a value of 1.0 is fine. In this first section we consider vacuum fields only, so the magnetic field due to current in the plasma does not need to be subtracted in the @@ -73,6 +66,13 @@ simsopt using OpenMP and vectorization, but MPI is not used, so the ``mpi4py`` python package is not needed. This example can be run in a few seconds on a laptop. +To solve this optimization problem in simsopt, we first import the necessary classes:: + + from simsopt.geo import SurfaceRZFourier, create_equally_spaced_curves, \ + CurveLength, plot, curves_to_vtk + from simsopt.field import Current, coils_via_symmetries, BiotSavart + from simsopt.objectives import SquaredFlux, QuadraticPenalty + The target plasma surface is given in the VMEC input file ``tests/test_files/input.LandremanPaul2021_QA``. We load the surface using: @@ -85,19 +85,20 @@ We load the surface using: You can adjust the directory in ``"filename"`` as appropriate for your system. The surface could also be read in from a VMEC wout file using -:obj:`simsopt.geo.surfacerzfourier.SurfaceRZFourier.from_wout()`. -Note that VMEC does not need to be installed to initialize a surface from a -VMEC input or output file. As surface objects carry a grid of -"quadrature points" at which the position vector is evaluated -we may want these points to cover different -ranges of the toroidal angle. For this problem with stellarator -symmetry and field-period symmetry, we need only consider half of a -field period in order to evaluate integrals over the entire +:obj:`simsopt.geo.SurfaceRZFourier.from_wout()`. Note that VMEC does +not need to be installed to initialize a surface from a VMEC input or +output file. As surface objects carry a grid of "quadrature points" at +which the position vector is evaluated we may want these points to +cover different ranges of the toroidal angle. For this problem with +stellarator symmetry and field-period symmetry, we need only consider +half of a field period in order to evaluate integrals over the entire surface. For this reason, the ``range`` parameter of the surface is -set to ``"half period"`` here. Possible options are ``full torus``, ``field period`` and ``half period``. +set to ``"half period"`` here. Possible options in general are ``"full +torus"``, ``"field period"`` and ``"half period"``, but ``"half +period"`` is significantly more efficient than the other options for +this problem. -After importing the necessary classes (see ``examples/2_Intermediate/stage_two_optimization.py``), -we set the initial condition for the coils, which will be equally spaced circles. +We next set the initial condition for the coils, which will be equally spaced circles. We consider here a case with four unique coil shapes, each of which is repeated four times due to stellarator symmetry and two-field-period symmetry, giving a total of 16 coils. The four unique coil shapes are called the "base coils". Each copy of a coil also carries the same current, @@ -121,8 +122,10 @@ The initial conditions are set as follows:: order = 5 base_curves = create_equally_spaced_curves(ncoils, s.nfp, stellsym=True, R0=R0, R1=R1, order=order) - base_currents = [Current(1e5) for i in range(ncoils)] + base_currents = [Current(1.0) * 1e5 for i in range(ncoils)] +In the last line, the degrees of freedom for the current are set to be of order 1, +but the actual currents are ``1e5`` larger so the magnetic field will be on the order of 1 Tesla. One detail of optimizing coils for a vacuum configuration is that the optimizer can "cheat" by making all the currents go to zero, which makes the quadratic flux vanish. To close this loophole, we can fix @@ -130,7 +133,7 @@ the current of the first base coil:: base_currents[0].fix_all() -(A ``Current`` object only has one degree of freedom, hence we can use +(A :obj:`~simsopt.field.Current` object only has one degree of freedom, hence we can use ``fix_all()``.) If you wish, you can fix the currents in all the coils to force them to have the same value. Now the full set of 16 coils can be obtained using stellarator symmetry and field-period @@ -187,55 +190,31 @@ and instruct it to evaluate the field on the surface:: (The surface position vector ``gamma()`` returns an array of size ``(nphi, ntheta, 3)``, which we reshaped here to -``(nphi*ntheta, 3)`` for the -:obj:`~simsopt.field.BiotSavart` object.) +``(nphi*ntheta, 3)`` as required for the +:obj:`~simsopt.field.BiotSavart` object.) To check the size of the field normal to the target surface before optimization we can run:: B_dot_n = np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2) - print('Initial max B dot n:', np.max(B_dot_n)) + print('Initial max|B dot n|:', np.max(np.abs(B_dot_n))) -The result is 0.19 Tesla. +The result is 0.23 Tesla. We now define the objective function by stating what are the weights used and the corresponding terms:: - # Weight on the curve lengths in the objective function. We use the `Weight` - # class here to later easily adjust the scalar value and rerun the optimization - # without having to rebuild the objective. - LENGTH_WEIGHT = Weight(1e-6) - - # Threshold and weight for the coil-to-coil distance penalty in the objective function: - CC_THRESHOLD = 0.1 - CC_WEIGHT = 1000 - - # Threshold and weight for the coil-to-surface distance penalty in the objective function: - CS_THRESHOLD = 0.3 - CS_WEIGHT = 10 - - # Threshold and weight for the curvature penalty in the objective function: - CURVATURE_THRESHOLD = 5. - CURVATURE_WEIGHT = 1e-6 + # Weight on the curve lengths in the objective function. + LENGTH_WEIGHT = 1.0 - # Threshold and weight for the mean squared curvature penalty in the objective function: - MSC_THRESHOLD = 5 - MSC_WEIGHT = 1e-6 + # Target value for the sum of the lengths of the base coils: + LENGTH_TARGET = 18.0 # Define the individual terms objective function: Jf = SquaredFlux(s, bs) Jls = [CurveLength(c) for c in base_curves] - Jccdist = CurveCurveDistance(curves, CC_THRESHOLD, num_basecurves=ncoils) - Jcsdist = CurveSurfaceDistance(curves, s, CS_THRESHOLD) - Jcs = [LpCurveCurvature(c, 2, CURVATURE_THRESHOLD) for c in base_curves] - Jmscs = [MeanSquaredCurvature(c) for c in base_curves] # Form the total objective function. - JF = Jf \ - + LENGTH_WEIGHT * sum(Jls) \ - + CC_WEIGHT * Jccdist \ - + CS_WEIGHT * Jcsdist \ - + CURVATURE_WEIGHT * sum(Jcs) \ - + MSC_WEIGHT * sum(QuadraticPenalty(J, MSC_THRESHOLD) for J in Jmscs) + objective = Jf + LENGTH_WEIGHT * QuadraticPenalty(sum(Jls), LENGTH_TARGET) 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 @@ -270,7 +249,7 @@ from scipy `_:: res = minimize(fun, objective.x, jac=True, method='L-BFGS-B', - options={'maxiter': 200, 'iprint': 5}, tol=1e-15) + options={'maxiter': 300, 'iprint': 5}, tol=1e-15) The optimization takes a few seconds, and the output will look like @@ -293,10 +272,10 @@ The optimization takes a few seconds, and the output will look like At iterate 10 f= 1.13772D-04 |proj g|= 6.27872D-04 ... - At iterate 195 f= 1.81723D-05 |proj g|= 4.18583D-06 - - At iterate 200 f= 1.81655D-05 |proj g|= 6.31030D-06 + At iterate 295 f= 6.20353D-07 |proj g|= 9.16379D-06 + At iterate 300 f= 6.19516D-07 |proj g|= 1.66448D-05 + * * * Tit = total number of iterations @@ -310,18 +289,18 @@ The optimization takes a few seconds, and the output will look like * * * N Tit Tnf Tnint Skip Nact Projg F - 135 200 234 1 0 0 6.310D-06 1.817D-05 - F = 1.8165520700970273E-005 - + 135 300 409 2 0 0 1.664D-05 6.195D-07 + F = 6.1951581434132075E-007 + STOP: TOTAL NO. of ITERATIONS REACHED LIMIT -You can adjust parameters such as the tolerance and number of +(You may obtain somewhat different values). You can adjust parameters such as the tolerance and number of iterations. Let us check the final :math:`\vec{B}\cdot\vec{n}` on the surface:: B_dot_n = np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2) - print('Final max B dot n:', np.max(B_dot_n)) + print('Final max|B dot n|:', np.max(np.abs(B_dot_n))) -The final value is 0.0017 Tesla, reduced two orders of magnitude from +The final value is 0.0015 Tesla, reduced two orders of magnitude from the initial state. As with the initial conditions, you can plot the optimized coil shapes directly from simsopt using @@ -341,6 +320,7 @@ syntax:: .. image:: coils_final.png :width: 500 +(Your coils may look somewhat different). The optimized value of the current in coil ``j`` can be obtained using ``coils[j].current.get_value()``. The optimized Fourier coefficients for coil ``j`` can be obtained from ``coils[j].curve.x``, where the @@ -348,19 +328,61 @@ meaning of each array element can be seen from ``coils[j].curve.dof_names``. The position vector for coil ``j`` in Cartesian coordinates can be obtained from ``coils[j].curve.gamma()``. +At the end of the optimization, the results can be saved as follows:: + + bs.save("biot_savart_opt.json") + +This line saves the optimized coil shapes and currents, along with the +:obj:`~simsopt.field.BiotSavart` object that evaluates the magnetic +field produced by the coils. The resulting json file can be loaded in +later in a separate script to analyze the results (making Poincare +plots, etc.) + + +Further coil regularization terms +--------------------------------- + +In the previous example we have already seen several objective function terms available for coil regularization: + +- :obj:`~simsopt.geo.CurveLength`: The length of a coil +- :obj:`~simsopt.objectives.QuadraticPenalty`: Useful for making quantities equal to, greater than, or less than a target value. + +However, better coils can be obtained if additional terms are included in the objective function +to control the coil curvature and coil-to-coil separation. +A large number of other terms are available in :obj:`simsopt.geo` to include in the objective function, +many of which are illustrated in +``examples/2_Intermediate/stage_two_optimization.py``. Other available objective terms include + +- :obj:`~simsopt.geo.MeanSquaredCurvature`: Reduces the coil curvature evenly along the curve. +- :obj:`~simsopt.geo.LpCurveCurvature`: Penalizes values of the coil curvature that exceed a threshold. +- :obj:`~simsopt.geo.LpCurveTorsion`: Penalizes values of the coil torsion that exceed a threshold. Coil optimization does not tend to work well when this term is included. +- :obj:`~simsopt.geo.CurveCurveDistance`: Useful for ensuring the minimum coil-to-coil distance is at least a specified target value. +- :obj:`~simsopt.geo.CurveSurfaceDistance`: Useful for ensuring the minimum coil-to-plasma distance is at least a specified target value. +- :obj:`~simsopt.geo.ArclengthVariation`: Ensures the curves are parameterized using (approximately) a uniform-arclength parameter. + +You can click on any of the links above in this section to see the precise definitions of these objective terms. + +Another useful class for forming objective functions is +:obj:`simsopt.objectives.Weight`, which is a wrapper for a float. +This object can be used for the weight in front of a term in the +objective function. Although a standard float works for this purpose +as well, the advantage of using :obj:`~simsopt.objectives.Weight` is +that when the value of the weight is changed, the objective function +object is automatically updated with the new value. + Stochastic Optimization ---------------------------- +----------------------- In this example we solve a stochastic version of -the :ref:`first example here `. As before, +the :ref:`first example here `. As before, the goal is to find coils that generate a specific target normal field on a given surface. As we are still considering a vacuum field the target is just zero. The target equilibrium is the `precise QA configuration `_. The complete script can be found in ``examples/2_Intermediate/stage_two_optimization_stochastic.py``. -The objective function similar to :ref:`the first example ` -with small modifications:: +The objective function is similar to ``examples/2_Intermediate/stage_two_optimization.py`` +but with a few modifications:: J = (1/2) Mean(\int |B dot n|^2 ds) + LENGTH_WEIGHT * (sum CurveLength) @@ -369,20 +391,22 @@ with small modifications:: + MSC_WEIGHT * MeanSquaredCurvaturePenalty(MSC_THRESHOLD) + ARCLENGTH_WEIGHT * ArclengthVariation -The first term is altered to be given by the Mean of the flux instead of the -flux itself. In here, the Mean is approximated by a sample average over perturbed coils. +The first term is altered to be given by the mean (expected value) of the flux over a distribution +of possible perturbed coils instead of the +flux itself. The mean is approximated by a sample average over perturbed coils. The coil perturbations for each coil are the sum of a 'systematic error' and a 'statistical error'. The former satisfies rotational and stellarator symmetry, the latter is independent for each coil. -An extra term term is also added, namely the variation of the arclength along a curve -named `ArclengthVariation`. The idea is to avoid ill-posedness of curve objectives due to -non-uniqueness of the underlying parametrization. -For more info see :obj:`~simsopt.geo.ArclengthVariation`. +The last term in the objective, :obj:`~simsopt.geo.ArclengthVariation`, +tends to be useful to include when doing stochastic optimization. Without this term, +which tends to make the curve parameterization have a uniform arclength, +the optimizer can exploit the curve parameterization rather than actually changing the shape of the curve. +This is possible since the Gaussian process model for curve perturbations is based on the curve parameter rather than arclength. We now define the objective function by stating what are the weights used and the corresponding terms. Besides the terms in -:ref:`the first example `, we additionally define:: +the previous examples, we additionally define:: # Weight for the arclength variation penalty in the objective function: ARCLENGTH_WEIGHT = 1e-2 @@ -437,18 +461,17 @@ added to each of the final coils, and is independent between coils. Finite Beta Optimization ---------------------------- - -In this example, we solve a finite beta version of -the :ref:`first example here `. -By finite beta, it is understood that the effect of -the plasma is also taken into accout when calculating -the normal field on a given surface. Therefore, the -target quantity :math:`B_{external}\cdot \mathbf n` is no longer zero +------------------------ + +In this example, we solve a variant of +the :ref:`first example here ` +in which there is current inside the plasma, which contributes significantly to +the total magnetic field.Therefore, the +target quantity :math:`\vec{B}_{external}\cdot \vec{n}` is no longer zero and a virtual casing calculation is used to find its value. The complete script can be found in ``examples/2_Intermediate/stage_two_finite_beta.py``. -We use an objective function similar to :ref:`the first example ` +We use an objective function similar to :ref:`the first example ` with small modifications:: J = (1/2) \int |(B_{BiotSavart} - B_{External}) dot n|^2 ds @@ -501,32 +524,32 @@ from the initial coil length is penalized; curvature and distance are not targeted here. Finite Build Optimization ---------------------------- +------------------------- In this final example, we perform a stage 2 optimization with finite build coils. The script for this case can be found in ``examples/3_Advanced/stage_two_optimization_finite_build.py``. -In particular, we use a multifilament approach that follows:: +In particular, we use a multifilament approach that follows - Optimization of finite-build stellarator coils, - Singh, Luquant, et al. Journal of Plasma Physics 86.4 (2020). +.. code-block:: + + "Optimization of finite-build stellarator coils," + Singh, Luquant, et al. Journal of Plasma Physics 86, 4 (2020) to approximate a finite build coil in order to have finite thickness. The target equilibrium is the `precise QA configuration `_. -Besides the degrees of freedom listed in :ref:`first example here `, +Besides the degrees of freedom listed in :ref:`first example here `, in this case, we have additional degrees of freedom related to the rotation of the coil pack. The objective function is given by:: J = (1/2) \int |(B_{BiotSavart}) dot n|^2 ds + LENGTH_PEN * (sum CurveLength) - + DIST_PEN * PairwiseDistancePenalty + + DIST_PEN * CurveCurveDistance -In here, the `PairwiseDistancePenalty` is the same term as the -`MininumDistancePenalty` of the :ref:`first example ` -that is used to prevents coils from +Here, :obj:`~simsopt.geo.CurveCurveDistance` is used to prevents coils from becoming too close. The constant -`DIST_PEN` is selected to balance this minimum distance penalty +``DIST_PEN`` is selected to balance this minimum distance penalty against the other objectives. To initialize the finite build optimization, we use the definitions below:: diff --git a/docs/source/example_quasisymmetry.rst b/docs/source/example_quasisymmetry.rst index c5c0956e0..1f18b4224 100644 --- a/docs/source/example_quasisymmetry.rst +++ b/docs/source/example_quasisymmetry.rst @@ -41,8 +41,7 @@ imports of the classes and functions we will need:: import numpy as np from simsopt.util import MpiPartition - from simsopt.mhd import Vmec - from simsopt.mhd import QuasisymmetryRatioResidual + from simsopt.mhd import Vmec, QuasisymmetryRatioResidual from simsopt.objectives import LeastSquaresProblem from simsopt.solve import least_squares_mpi_solve @@ -235,8 +234,7 @@ As usual, we begin with the necessary imports:: import numpy as np from simsopt.util import MpiPartition - from simsopt.mhd import Vmec - from simsopt.mhd import QuasisymmetryRatioResidual + from simsopt.mhd import Vmec, QuasisymmetryRatioResidual from simsopt.objectives import LeastSquaresProblem from simsopt.solve import least_squares_mpi_solve @@ -274,7 +272,7 @@ axisymmetric:: It can be seen here that we are seeking a configuration with aspect ratio 6, and average iota of 0.42, slightly above the resonance at 2 / -5 = 0.4. The function :func:`simsopt.mhd.vmec.Vmec.mean_iota()` used +5 = 0.4. The function :func:`simsopt.mhd.Vmec.mean_iota()` used here returns :math:`\int_0^1 \iota\, ds` where :math:`s` is the toroidal flux normalized by its value at the VMEC boundary. @@ -336,7 +334,7 @@ Bmn objective Here we show an alternative method of quasisymmetry optimization using a different objective function, -:obj:`simsopt.mhd.boozer.Quasisymmetry`, based on the +:obj:`simsopt.mhd.Quasisymmetry`, based on the symmetry-breaking Fourier mode aplitudes :math:`B_{m,n}` in Boozer coordinates. This example can also be found in the ``examples/2_Intermediate`` directory as @@ -345,8 +343,7 @@ coordinates. This example can also be found in the In this case, the imports needed are:: from simsopt.util import MpiPartition - from simsopt.mhd import Vmec - from simsopt.mhd import Boozer, Quasisymmetry + from simsopt.mhd import Vmec, Boozer, Quasisymmetry from simsopt.objectives import LeastSquaresProblem from simsopt.solve import least_squares_mpi_solve diff --git a/docs/source/example_vmec_only.rst b/docs/source/example_vmec_only.rst index 047aae1e5..b919cc5fa 100644 --- a/docs/source/example_vmec_only.rst +++ b/docs/source/example_vmec_only.rst @@ -60,8 +60,8 @@ For simplicity, MPI parallelization will not be used for now. To start, we must import several classes:: from simsopt.mhd import Spec - from simsopt.objectives.least_squares import LeastSquaresProblem - from simsopt.solve.serial import least_squares_serial_solve + from simsopt.objectives import LeastSquaresProblem + from simsopt.solve import least_squares_serial_solve Then we create the equilibrium object, starting from an input file:: @@ -141,15 +141,15 @@ computation, or both at the same time. To introduce MPI we first initialize an :obj:`simsopt.util.MpiPartition` object and choose the number of worker groups. The instance is then passed as an argument to the Vmec object and to the -:meth:`simsopt.solver.mpi_solve.least_squares_mpi_solve` function. +:meth:`simsopt.solve.least_squares_mpi_solve` function. For more details about MPI, see :doc:`mpi`. The complete example is then as follows:: - from simsopt.util.mpi import MpiPartition + from simsopt.util import MpiPartition from simsopt.mhd import Vmec from simsopt.objectives import LeastSquaresProblem - from simsopt.solve.mpi import least_squares_mpi_solve + from simsopt.solve import least_squares_mpi_solve # In the next line, we can adjust how many groups the pool of MPI # processes is split into. diff --git a/docs/source/fields.rst b/docs/source/fields.rst index 41e0ff8b4..c0b94698f 100644 --- a/docs/source/fields.rst +++ b/docs/source/fields.rst @@ -194,7 +194,7 @@ operator creates the latter.) Example:: - from simsopt.field.magneticfieldclasses import ToroidalField, CircularCoil + from simsopt.field import ToroidalField, CircularCoil field1 = CircularCoil(I=1.e7, r0=1.) field2 = ToroidalField(R0=1., B0=1.) diff --git a/docs/source/geo.rst b/docs/source/geo.rst index 04b26376b..d926afc04 100644 --- a/docs/source/geo.rst +++ b/docs/source/geo.rst @@ -53,15 +53,15 @@ the derivative of quantities such as curvature (via ``Curve.dkappa_by_dcoeff()``) or torsion (via ``Curve.dtorsion_by_dcoeff()``). -A number of quantities are implemented in -:obj:`simsopt.geo.curveobjectives` and are computed on a +A number of objective functions are available in +:obj:`simsopt.geo` and are computed on a :obj:`simsopt.geo.Curve`: -- ``CurveLength``: computes the length of the ``Curve``. -- ``LpCurveCurvature``: computes a penalty based on the :math:`L_p` norm of the curvature on a curve. -- ``LpCurveTorsion``: computes a penalty based on the :math:`L_p` norm of the torsion on a curve. -- ``CurveCurveDistance``: computes a penalty term on the minimum distance between a set of curves. -- ``CurveSurfaceDistance``: computes a penalty term on the minimum distance between a set of curves and a surface. +- :obj:`~simsopt.geo.CurveLength`: computes the length of the ``Curve``. +- :obj:`~simsopt.geo.LpCurveCurvature`: computes a penalty based on the :math:`L_p` norm of the curvature on a curve. +- :obj:`~simsopt.geo.LpCurveTorsion`: computes a penalty based on the :math:`L_p` norm of the torsion on a curve. +- :obj:`~simsopt.geo.CurveCurveDistance`: computes a penalty term on the minimum distance between a set of curves. +- :obj:`~simsopt.geo.CurveSurfaceDistance`: computes a penalty term on the minimum distance between a set of curves and a surface. The value of the quantity and its derivative with respect to the curve dofs can be obtained by calling e.g., ``CurveLength.J()`` and @@ -152,10 +152,10 @@ the :obj:`~simsopt.geo.Curve` class: - ``Surface.second_fund_form()``: returns a ``(n_phi, n_theta, 3)`` array containing :math:`[\hat{\textbf{n}}(\phi_i, \theta_j) \cdot \partial^2_{\phi,\phi} \Gamma(\phi_i, \theta_j), \hat{\textbf{n}}(\phi_i, \theta_j) \cdot \partial^2_{\phi,\theta} \Gamma(\phi_i, \theta_j), \hat{\textbf{n}}(\phi_i, \theta_j) \cdot \partial^2_{\theta,\theta} \Gamma(\phi_i, \theta_j)]` for :math:`i\in\{1, \ldots, n_\phi\}, j\in\{1, \ldots, n_\theta\}` where :math:`\hat{\textbf{n}}` is the unit normal. - ``Surface.surface_curvatures()``: returns a ``(n_phi, n_theta, 4)`` array containing :math:`[H(\phi_i, \theta_j),K(\phi_i, \theta_j),\kappa_1(\phi_i, \theta_j),\kappa_2(\phi_i, \theta_j)]` for :math:`i\in\{1, \ldots, n_\phi\}, j\in\{1, \ldots, n_\theta\}` where :math:`H` is the mean curvature, :math:`K` is the Gaussian curvature, and :math:`\kappa_{1,2}` are the principal curvatures with :math:`\kappa_1>\kappa_2`. -A number of quantities are implemented in :obj:`simsopt.geo.surfaceobjectives` and are computed on a :obj:`simsopt.geo.Surface`: +A number of objective functions related to surfaces are available in :obj:`simsopt.geo`: -- ``ToroidalFlux``: computes the flux through a toroidal cross section of a ``Surface``. -- ``PrincipalCurvature``: computes a metric which penalizes large values of the principal curvatures of a given ``Surface``. +- :obj:`~simsopt.geo.ToroidalFlux`: computes the flux through a toroidal cross section of a ``Surface``. +- :obj:`~simsopt.geo.PrincipalCurvature`: computes a metric which penalizes large values of the principal curvatures of a given ``Surface``. The value of the quantity and its derivative with respect to the surface dofs can be obtained by calling e.g., ``ToroidalFlux.J()`` and ``ToroidalFlux.dJ_dsurfacecoefficients()``. @@ -185,12 +185,12 @@ plotting engine by passing the ``engine`` keyword argument, e.g. if can use the ``close`` argument to control whether segments are drawn between the last quadrature point and the first. For these and other options, see the API documentation for -:func:`simsopt.geo.curve.Curve.plot()` and -:func:`simsopt.geo.surface.Surface.plot()`. +:func:`simsopt.geo.Curve.plot()` and +:func:`simsopt.geo.Surface.plot()`. If you have multiple curve and/or surface objects, a convenient way to plot them together on the same axes is the function -:func:`simsopt.geo.plotting.plot()`, which accepts a list of objects as +:func:`simsopt.geo.plot()`, which accepts a list of objects as its argument. Any keywords passed to this function are passed to the ``.plot()`` methods of the individual objects, so you may wish to pass keywords such as ``engine`` or ``close``. Alternatively, you can also @@ -201,6 +201,6 @@ It is also possible to export curve and surface objects in VTK format, so they can be viewed in Paraview. This functionality requires the python package ``pyevtk``, which can be installed via ``pip install pyevtk``. A list of curve objects can be exported using the function -:func:`simsopt.geo.curve.curves_to_vtk()`. To export a VTK file for a +:func:`simsopt.geo.curves_to_vtk()`. To export a VTK file for a surface, call the ``.to_vtk(filename)`` function of the object. See -:func:`simsopt.geo.surface.Surface.to_vtk()` for more details. +:func:`simsopt.geo.Surface.to_vtk()` for more details. diff --git a/docs/source/index.rst b/docs/source/index.rst index aa54a9198..3a82431d7 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -89,8 +89,8 @@ optimization. Others include `STELLOPT :maxdepth: 3 :caption: API reference - Puiblic functions and classes - For developers + Public functions and classes + Full listing (for developers) Indices and tables ================== diff --git a/docs/source/optimizable.rst b/docs/source/optimizable.rst index 7e6b7acdd..c2aa22bc1 100644 --- a/docs/source/optimizable.rst +++ b/docs/source/optimizable.rst @@ -25,7 +25,7 @@ algorithms. Users can create their own optimizable objects in two ways. One method is to create a standard python function, and apply the -:obj:`simsopt.make_optimizable()` function to it, as explained +:meth:`simsopt.make_optimizable()` function to it, as explained below. Or, you can directly subclass :obj:`simsopt._core.Optimizable`. @@ -533,7 +533,7 @@ Specifying least-squares objective functions A common use case is to minimize a nonlinear least-squares objective function, which consists of a sum of several terms. In this case the -:obj:`simsopt.objectives.least_squares.LeastSquaresProblem` +:obj:`simsopt.objectives.LeastSquaresProblem` class can be used. Suppose we want to solve a least-squares optimization problem in which an :obj:`~simsopt._core.Optimizable` object ``obj`` has @@ -562,18 +562,18 @@ The corresponding objective funtion is then ``weight1 * scalar functions and by 1D numpy array-valued functions. Note that the function handles that are specified should be members of an :obj:`~simsopt._core.Optimizable` object. As -:obj:`~simsopt.objectives.least_squares.LeastSquaresProblem` is +:obj:`~simsopt.objectives.LeastSquaresProblem` is a subclass of :obj:`~simsopt._core.Optimizable`, the free dofs of all the objects that go into the objective function are available in the global state vector ``prob.x``. The overall scalar objective function is available from -:func:`simsopt.objectives.least_squares.LeastSquaresProblem.objective`. +:func:`simsopt.objectives.LeastSquaresProblem.objective`. The vector of residuals before scaling by the ``weight`` factors ``obj.func() - goal`` is available from -:func:`simsopt.objectives.least_squares.LeastSquaresProblem.unweighted_residuals`. +:func:`simsopt.objectives.LeastSquaresProblem.unweighted_residuals`. The vector of residuals after scaling by the ``weight`` factors, ``sqrt(weight) * (obj.func() - goal)``, is available from -:func:`simsopt.objectives.least_squares.LeastSquaresProblem.residuals`. +:func:`simsopt.objectives.LeastSquaresProblem.residuals`. Least-squares problems can also be defined in an alternative way:: diff --git a/docs/source/overview.rst b/docs/source/overview.rst index 332ea2f8e..f625cdf51 100644 --- a/docs/source/overview.rst +++ b/docs/source/overview.rst @@ -25,9 +25,9 @@ Simsopt does not use input data files to define optimization problems, in contrast to ``STELLOPT``. Rather, problems are specified using a python driver script, in which objects are defined and configured. However, objects related to specific physics codes may use -their own input files. In particular, a :obj:`simsopt.mhd.vmec.Vmec` object +their own input files. In particular, a :obj:`simsopt.mhd.Vmec` object can be initialized using a standard VMEC ``input.*`` input file, and a -:obj:`simsopt.mhd.spec.Spec` object can be initialized using a standard +:obj:`simsopt.mhd.Spec` object can be initialized using a standard SPEC ``*.sp`` input file. @@ -72,18 +72,18 @@ Some typical objects are a MHD equilibrium represented by the VMEC or SPEC code, or some electromagnetic coils. To define objective functions, a variety of additional objects can be defined that depend on the MHD equilibrium or coils, such as a -:obj:`simsopt.mhd.boozer.Boozer` object for Boozer-coordinate -transformation, a :obj:`simsopt.mhd.spec.Residue` object to represent +:obj:`simsopt.mhd.Boozer` object for Boozer-coordinate +transformation, a :obj:`simsopt.mhd.Residue` object to represent Greene's residue of a magnetic island, or a -:obj:`simsopt.geo.objectives.LpCurveCurvature` penalty on coil +:obj:`simsopt.geo.LpCurveCurvature` penalty on coil curvature. More details about setting degrees of freedom and defining objective functions can be found on the :doc:`optimizable` page. For the solution step, two functions are provided presently, -:meth:`simsopt.solve.serial.least_squares_serial_solve` and -:meth:`simsopt.solve.mpi.least_squares_mpi_solve`. The first +:meth:`simsopt.solve.least_squares_serial_solve` and +:meth:`simsopt.solve.least_squares_mpi_solve`. The first is simpler, while the second allows MPI-parallelized finite differences to be used in the optimization. diff --git a/docs/source/tracing.rst b/docs/source/tracing.rst index f9604ab01..020985e23 100644 --- a/docs/source/tracing.rst +++ b/docs/source/tracing.rst @@ -18,7 +18,7 @@ Particle motion in cylindrical coordinates ------------------------------------------ The main function to use in this case is -:obj:`simsopt.field.tracing.trace_particles` (click the link for more +:obj:`simsopt.field.trace_particles` (click the link for more information on the input and output parameters) and it is able to use two different sets of equations depending on the input parameter ``mode``: diff --git a/examples/1_Simple/stage_two_optimization_minimal.py b/examples/1_Simple/stage_two_optimization_minimal.py new file mode 100755 index 000000000..924c129c8 --- /dev/null +++ b/examples/1_Simple/stage_two_optimization_minimal.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python +r"""In this example we solve a FOCUS like Stage II coil optimisation problem: the +goal is to find coils that generate a specific target normal field on a given +surface. In this particular case we consider a vacuum field, so the target is +just zero. + +The objective is given by + + J = (1/2) \int |B dot n|^2 ds + + LENGTH_WEIGHT * (LENGTH_TARGET - sum CurveLengths)^2 + +If LENGTH_TARGET is increased, the coils +become more complicated, but the target normal field may be +achieved more accurately. + +The target equilibrium is the QA configuration of arXiv:2108.03711. + +Better coils can be obtained by including more terms in the objective +function to control the coil curvature and coil-to-coil distance. For +an example of this more sophisticated coil optimization see +examples/2_Intermediate/stage_two_optimization.py. + +""" + +import os +from pathlib import Path +import numpy as np +from scipy.optimize import minimize + +from simsopt.geo import SurfaceRZFourier, create_equally_spaced_curves, \ + CurveLength, plot, curves_to_vtk +from simsopt.field import Current, coils_via_symmetries, BiotSavart +from simsopt.objectives import SquaredFlux, QuadraticPenalty + +# 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.) +ncoils = 4 + +# Major radius for the initial circular coils: +R0 = 1.0 + +# Minor radius for the initial circular coils: +R1 = 0.5 + +# Number of Fourier modes describing each Cartesian component of each coil: +order = 5 + +# Weight on the curve lengths in the objective function. +LENGTH_WEIGHT = 1.0 + +# Target value for the sum of the lengths of the base coils: +LENGTH_TARGET = 18.0 + +# Number of iterations to perform: +ci = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true'] +MAXITER = 50 if ci else 300 + +# File for the desired boundary magnetic surface: +TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve() +filename = TEST_DIR / 'input.LandremanPaul2021_QA' + +# Directory for output +OUT_DIR = "./output/" +os.makedirs(OUT_DIR, exist_ok=True) + +####################################################### +# End of input parameters. +####################################################### + +# Initialize the boundary magnetic surface: +nphi = 32 +ntheta = 32 +s = SurfaceRZFourier.from_vmec_input(filename, range="half period", nphi=nphi, ntheta=ntheta) + +# Create the initial coils: +base_curves = create_equally_spaced_curves(ncoils, s.nfp, stellsym=True, R0=R0, R1=R1, order=order) +base_currents = [Current(1.0) * 1e5 for i in range(ncoils)] +# Since the target field is zero, one possible solution is just to set all +# currents to 0. To avoid the minimizer finding that solution, we fix one +# of the currents: +base_currents[0].fix_all() + +coils = coils_via_symmetries(base_curves, base_currents, s.nfp, True) +bs = BiotSavart(coils) +bs.set_points(s.gamma().reshape((-1, 3))) + +curves = [c.curve for c in coils] +curves_to_vtk(curves, OUT_DIR + "curves_init", close=True) +pointData = {"B_N": np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)[:, :, None]} +s.to_vtk(OUT_DIR + "surf_init", extra_data=pointData) + +# Define the individual terms objective function: +Jf = SquaredFlux(s, bs) +Jls = [CurveLength(c) for c in base_curves] + +# 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 + LENGTH_WEIGHT * QuadraticPenalty(sum(Jls), LENGTH_TARGET) + +B_dot_n = np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2) +print('Initial max|B dot n|:', np.max(np.abs(B_dot_n))) +print('Names of the dofs:', JF.dof_names) + +# 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 + return JF.J(), JF.dJ() + + +print(""" +################################################################################ +### Perform a Taylor test ###################################################### +################################################################################ +""") +f = fun +dofs = JF.x +np.random.seed(1) +h = np.random.uniform(size=dofs.shape) +J0, dJ0 = f(dofs) +dJh = sum(dJ0 * h) +for eps in [1e-3, 1e-4, 1e-5, 1e-6, 1e-7]: + J1, _ = f(dofs + eps*h) + J2, _ = f(dofs - eps*h) + print("err", (J1-J2)/(2*eps) - dJh) + +print(""" +################################################################################ +### Run the optimisation ####################################################### +################################################################################ +""") +res = minimize(fun, dofs, jac=True, method='L-BFGS-B', + options={'maxiter': MAXITER, 'maxcor': 300, 'iprint': 5}, tol=1e-15) +curves_to_vtk(curves, OUT_DIR + f"curves_opt", close=True) +pointData = {"B_N": np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)[:, :, None]} +s.to_vtk(OUT_DIR + "surf_opt", extra_data=pointData) + +B_dot_n = np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2) +print('Final max|B dot n|:', np.max(np.abs(B_dot_n))) + +# Save the optimized coil shapes and currents so they can be loaded into other scripts for analysis: +bs.save(OUT_DIR + "biot_savart_opt.json") diff --git a/examples/run_serial_examples b/examples/run_serial_examples index c65e4753f..91d49472d 100755 --- a/examples/run_serial_examples +++ b/examples/run_serial_examples @@ -11,6 +11,7 @@ set -ex ./1_Simple/tracing_fieldlines_QA.py ./1_Simple/tracing_particle.py ./1_Simple/qfm.py +./1_Simple/stage_two_optimization_minimal.py ./2_Intermediate/QSC.py ./2_Intermediate/boozer.py ./2_Intermediate/stage_two_optimization.py diff --git a/src/simsopt/field/__init__.py b/src/simsopt/field/__init__.py index d8f189fec..2183f2aa1 100644 --- a/src/simsopt/field/__init__.py +++ b/src/simsopt/field/__init__.py @@ -1,8 +1,9 @@ from .biotsavart import * from .coil import * +from .magneticfield import * from .magneticfieldclasses import * from .boozermagneticfield import * from .tracing import * -__all__ = (biotsavart.__all__ + coil.__all__ + magneticfieldclasses.__all__ + +__all__ = (biotsavart.__all__ + coil.__all__ + magneticfield.__all__ + magneticfieldclasses.__all__ + boozermagneticfield.__all__ + tracing.__all__) diff --git a/src/simsopt/field/magneticfield.py b/src/simsopt/field/magneticfield.py index 9d8995319..062967377 100644 --- a/src/simsopt/field/magneticfield.py +++ b/src/simsopt/field/magneticfield.py @@ -5,6 +5,8 @@ from .._core.optimizable import Optimizable from .._core.derivative import Derivative +__all__ = ['MagneticField', 'MagneticFieldSum', 'MagneticFieldMultiply'] + class MagneticField(sopp.MagneticField, Optimizable): diff --git a/src/simsopt/geo/__init__.py b/src/simsopt/geo/__init__.py index 818eb611d..bbc47fc74 100644 --- a/src/simsopt/geo/__init__.py +++ b/src/simsopt/geo/__init__.py @@ -14,6 +14,7 @@ from .boozersurface import * from .qfmsurface import * +from .surface import * from .surfacegarabedian import * from .surfacehenneberg import * from .surfaceobjectives import * @@ -26,6 +27,7 @@ curveperturbed.__all__ + curveobjectives.__all__ + finitebuild.__all__ + plotting.__all__ + boozersurface.__all__ + qfmsurface.__all__ + + surface.__all__ + surfacegarabedian.__all__ + surfacehenneberg.__all__ + surfacerzfourier.__all__ + surfacexyzfourier.__all__ + surfacexyztensorfourier.__all__ + surfaceobjectives.__all__) diff --git a/src/simsopt/geo/curve.py b/src/simsopt/geo/curve.py index 166461eec..90f982670 100644 --- a/src/simsopt/geo/curve.py +++ b/src/simsopt/geo/curve.py @@ -13,7 +13,7 @@ from .jit import jit from .plotting import fix_matplotlib_3d -__all__ = ['RotatedCurve', 'curves_to_vtk', 'create_equally_spaced_curves'] +__all__ = ['Curve', 'RotatedCurve', 'curves_to_vtk', 'create_equally_spaced_curves'] @jit diff --git a/src/simsopt/geo/surface.py b/src/simsopt/geo/surface.py index 280eb4085..2468bb692 100644 --- a/src/simsopt/geo/surface.py +++ b/src/simsopt/geo/surface.py @@ -13,6 +13,8 @@ from .._core.dev import SimsoptRequires from .plotting import fix_matplotlib_3d +__all__ = ['Surface', 'signed_distance_from_surface', 'SurfaceClassifier', 'SurfaceScaled', 'best_nphi_over_ntheta'] + class Surface(Optimizable): r"""