Dipole and passive coil arrays branch #483

akaptano · 2025-04-21T19:22:58Z

Hi all, this has been a long while coming, but I have my branch cleaned up that does the following:

Implements dipole array optimization (including merging in Jake Halpbern's functionality he nicely shared with me)
Implements passive array optimization
Implements Jax-based net force, net torque, pointwise force, pointwise torque objectives for optimization. Note that MixedLpCurveForce and similar functions starting with 'Mixed' should be used if the number of coils is greater than 20 or 30. The reason is that the tree dependencies in SIMSOPT grow exponentially fast and slow down the code to a halt if any of the objective functions create a BiotSavart class object inside of them (as the default LpCurveForce and other functions do!). As far as I know, this big issue is not seen elsewhere because it requires a fair number of coils to see it, and it also requires that there is an objective that itself creates a BiotSavart object.
Adds a tutorial file to the SIMSOPT documentation on minimal, advanced, and passive dipole coil array optimization
Adds a fair number of new example files, including the two files for the tutorial in 3_Advanced/ including dipole_array_tutorial.py and dipole_array_tutorial_advanced.py. There are also files for running advanced scripts and coil force Pareto scans in 3_Advanced/dipole_coil_optimization, 3_Advanced/passive_coil_optimization, and 3_Advanced/coil_force_optimization.

The code is caught up and merged with several other branches, including master, Siena's coil forces branch, somewhat Jake Halpbern's dipole optimization branch, etc.

… currents don't flip either. A_matrix is correctly symmetrized WITHOUT multiplying it by any factors since A_matrix @ I agrees with A_matrix_full @ I_all. Moreover, A_matrix_full @ I_all produces correctly symmetrized Paraview plots as far as I can tell. So as it stands I think the terms in the least squares is working, still need to debug what is going on with the derivatives, and confusing to write out the whole derivative matrix for all the symmetrized coils since the dA/dkappa denominator depends on the symmetrized kappa.

…res and getting the alpha,delta angles are correct for any discrete symmetries. However, still working on getting jacobian calculations working. It seems like the derivatives depend on the parameters being symmetrized, and I think I should account for that like dA_dalphaj = dA_dalphaj_prime * dalphaj_prime_dalphaj but this is only working so far for nfp 2 and not for the QH and NCSX examples. Next steps are continue to debug to try and get the derivatives operating correctly.

…ore careful thought since both delta_prime and alpha_prime can change when a perturbation to delta is made. Had some more weird bugs during unit testing but it looks like the unit tests are only failing when a coil is accidentally placed so that it intersects with one of the symmetry planes. Going to add a warning about this and check more carefully tomorrow.

…x this. The reason is that there was still an extra ambiguity in the orientations from alpha in -pi < alpha < pi, while alpha actually should be bounded by pi/2 on both sides. Restricting the range seems to fix all the issues and jacobian calculations appear to be working now. More debugging to do and then big code clean up needed.

… look pretty good but still have some small disagreements with the finite differences so hard to say if there is a bug. Found that I was not checking the discrete symmetrization of the coils quite right, which could be important since coils intersecting the symmetry planes will do lots of weird things. Think I fixed this.

…obian calculation. alphas need to be between -pi/2 and pi/2 the entire optimization and same for deltas between -pi and pi. Otherwise, you can get a minus sign on the gradient because of an orientation flip of the coil. Adding these as bounds in the optimization seems to work so far. Compared with finite differences for a while and looks good. Then added some STLSQ optimization to QH and started trying to get high-performance results. Some decent initial results and trying to figure out how to make poincare plots with it now. Added a function in c++ for initializing all the symmetrized coil angles but its no faster than numpy yet.

…es in good shape. Did a fix to test if there are coil coil intersections in the grid.

…STLSQ so far but requires tuning the threshold nicely. Added some code to save the solution after each iteration of STLSQ. Added some code in the QH example to save the qfm surface and now running vmec in a separate script that can be called with MPI. This seems to run for a very long time so still working on getting a proper vmec and poincare section during post processing.

…anged the aspect ratio of the coils so that each coil has roughly twice as much current as before. Probably dont want to push this any further.

…but now seeing some issues with optimizing and not sure where this is coming from.

…aused by eps too small so the coils can reach the BFGS bounds. walked it back a little and seems to help so far. Tweaked the examples a bit more. Added check in the grid to look for possible PSC-TFcoil intersections.

…d tune this a bit more.

…he PSCs can wiggle. Added new example using GPMO. Minor changes to the PSC grid class although reinstated aspect ratio 0.1 for the more complex runs.

…pportunity to do joint optimization with the TF coils, but havent quite figured this out yet.

…t alone working well). The issue is that when the PSC coils start again, they seem to have the wrong TF coils, or not been optimized consistently wwhen the TF coils were going.

…Downside is that it is very slow. Going to take another look at trying to get the gradients working.

…joint optimizations. Even if I fix all the PSCs and dont do any optimization on them, letting the TF coils further optimize does not seem to help. Issues could be that the finite differences are not good enough, that the PSCs are a small effect, or something. But seems like it should work because spent time getting the calculation with finite differences looking sensible.

…oils. Need to check that the PSC currents are being updated correctly and only the orientations are being optimized.

… updates the jacobian. Not sure yet how to do this right while keeping the whole psc_array updated. Need to figure out where exactly in the code the coil dofs get changed. Then can update all the coil dofs at once, after the full jacobian computed.

…loser but gradients are zero from the derivatives with respect to the currents and not sure why.

…e direct coil optimization in stage_two_optimization_planar_coils.py and a pared-down file QA_psc.py. The results are often a bit worse using the direct coil optimization -- it seems more prone to get stuck in a local minima but maybe this is a result of numerical error accumulation anyways. However the discrepancies are small and seems to be consistent with the psc solutions. For instance, I can initialize a filament optimization correctly from the psc solution. Next step is to unfix the currents in the PSCs and do the self-consistent optimization as the currents change with the orientations.

… optimization.

…pscs. first issue is that I need a PSC curve class, since otherwise passing the derivatives through the coils doesnt work correctly. It needs to be passing a derivative associated with the curve, not with the coil or the current class. Secondly, think I found a bug related to the derivatives associated with the symmetrized curves. Not sure yet if I have gotten it exactly right.

…on (only has one argument -- oops). Debugged a bit more. Got the unit tests all working again. Still not agreeing with FDs.

… Still have discrepancy in the QA psc example, but this may be because of the way its comparing with finite differences. Need to test the full chain of derivatives through the SquaredFlux.

… flux calculation in unit tests.

…correctly compute the rotated coils contribution. Unit tests confusingly pass a basic test with ncoils = 8, but fail with a different number of coils. Similar confusing behavior when I change epsilon value of the finite differences. Seems like curve and current variables are getting changed accidentally somewhere. Will debug the basis test with different number of coils first.

… jacobian calculation sometimes have a big disagreement (looks like an error in the finite differencesbut overall it shows the convergence as epsilon gets smallers. This is just using default CurvePlanarFourier objects with no PSC usage whatsoever so far. Also absolute errors (instead of relative errors) are returned, which can be confusing when there are big errors on small numbers.

mishapadidar · 2025-05-07T20:12:27Z

src/simsopt/field/force.py

+    usually using curves_to_vtk.
+
+    Args:
+        coils (list): List of coils to optimize.


why "coils to optimize"?

mishapadidar · 2025-05-07T20:15:56Z

src/simsopt/field/force.py

+
+    Args:
+        coil (Coil): Coil to optimize.
+        regularization (Regularization): Regularization object.


is regularization an object or a function? Perhaps say that this must be regularization_circ or regularization_rect

mishapadidar · 2025-05-07T20:16:46Z

src/simsopt/field/force.py

+
+
+# @jit
+def lp_force_pure_deprecated(gamma, gammadash, gammadashdash, quadpoints, current, regularization, B_mutual, p, threshold, downsample):


@mbkumar how do we want to handle the deprecation?

src/simsopt/field/tracing.py

src/simsopt/geo/__init__.py

mishapadidar · 2025-05-07T20:19:48Z

src/simsopt/geo/curve.py

        return dkappadash_by_dcoeff

+    def center(self, gamma, gammadash):
+        # Compute the centroid of the curve


full docstring please. this is a good function to have.

mishapadidar · 2025-05-07T20:23:03Z

src/simsopt/geo/curve.py

+    # @jit
+    def frenet_frame_pure(self, dofs):
+        r"""
+        This function returns the Frenet frame, :math:`(\mathbf{t}, \mathbf{n}, \mathbf{b})`,


say this is for jax/autodiff plz

mishapadidar · 2025-05-07T20:23:15Z

src/simsopt/geo/curve.py

+        b = jnp.cross(t, n, axis=1)
+        return t, n, b
+
+    def frenet_frame(self):


mishapadidar · 2025-05-07T20:27:08Z

src/simsopt/geo/curve.py



-def curves_to_vtk(curves, filename, close=False, extra_data=None):
+def curves_to_vtk(curves, filename, close=False, I=None, extra_point_data=None,


this might be a breaking change (extra_data to extra_point_data). perhaps stick NetForces and NetTorques in there. simplify

mishapadidar · 2025-05-07T20:27:41Z

src/simsopt/geo/curve.py



-def create_equally_spaced_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None):
+def setup_uniform_grid(s, s_inner, s_outer, Nx, Ny, Nz, Nmin_factor=2.5):


what does this do? Why do we need?

mishapadidar · 2025-05-07T20:31:37Z

src/simsopt/geo/curve.py

+    return curves, all_curves
+
+
+def create_equally_spaced_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None, jax_flag=False):


could we use a more descriptive name for jax_flag?

mishapadidar · 2025-05-07T20:35:11Z

src/simsopt/geo/curve.py



-def create_equally_spaced_planar_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None):
+def create_equally_spaced_planar_curves(


mishapadidar · 2025-05-07T20:35:25Z

src/simsopt/geo/curve.py


-def create_equally_spaced_planar_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None):
+def create_equally_spaced_planar_curves(
+        ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None, jax_flag=False):


change jax_flag to more descrptive.

src/simsopt/geo/curveobjectives.py

mishapadidar · 2025-05-07T20:47:06Z

src/simsopt/geo/curveplanarellipticalcylindrical.py

+import numpy as np
+from .curve import JaxCurve
+
+__all__ = ['CurvePlanarEllipticalCylindrical', 'create_equally_spaced_cylindrical_curves',


Can we figure out what the CurvePlanarEllipticalCylindrical class is doing? Is this redundant with CurvePlanarFourier?

mishapadidar · 2025-05-07T20:48:40Z

src/simsopt/geo/curveplanarellipticalcylindrical.py

+    return final_gamma
+
+
+class CurvePlanarEllipticalCylindrical(JaxCurve):


@jhalpern30 can you chime in on what this class is doing and why it is useful?

I didn't realize that had made it into this PR. This was a hyper-specific class we used for designing the hybrid dipole experiment. I didn't write the code (Rohan Lopez did Spring 2024), but am decently familiar with it since I cleaned up/reformatted a good bit of it a few months back. Basically, its a planar, elliptical curve that can be rotated/shifted in cylindrical coordinates, we used it to make the TFs tilt about a specific pivot point.

I wouldn't include this - I think all the functionality could be done with CurvePlanarFourier, with the only advantage being analytically representing an ellipse instead of the Fourier series in CurvePlanarFourier. This is a byproduct of two first year grad students not knowing the full feature list within simsopt :)

I think it's somewhat useful still. Keeping for now

mishapadidar · 2025-05-07T20:49:15Z

src/simsopt/geo/curveplanarellipticalcylindrical.py

+        return rtpz_name + angle_name
+
+
+def create_equally_spaced_cylindrical_curves(ncurves, nfp, stellsym, R0, a, b, numquadpoints=32):


can we reuse the existing create_equally_spaced_curves ?

mishapadidar · 2025-05-07T20:50:16Z

src/simsopt/geo/curveplanarfourier.py

+                                     + center[2])))
+
+
+class JaxCurvePlanarFourier(JaxCurve):


If this is not used in PSCArray etc, lets move this to a later PR?

src/simsopt/geo/surfaceobjectives.py

src/simsopt/mhd/virtual_casing.py

src/simsopt/util/permanent_magnet_helper_functions.py

src/simsoptpp/curvexyzfourier.cpp

tests/field/test_biotsavart.py

tests/field/test_coil.py

mishapadidar

Lots of good changes. The plan is to split this into 4 PRs: PSC/Dipole Arrays with applicable tests/examples, Forces/Torques with applicable tests/examples, remaining tests examples that depend on both, miscellaneous changes that don't fit in to the previous ones.
Please add documentation.

…ra singleton dimension.

…e sure the unit tests and examples run well again.

…. Still have some issues from newly merged code slightly changing test outputs.

akaptano and others added 30 commits May 31, 2024 16:43

Sped up some calculations and working on getting the different exampl…

016eb0c

…es in good shape. Did a fix to test if there are coil coil intersections in the grid.

Made some changes to NCSX example to get it a bit more compelling. Ch…

d8307f6

…anged the aspect ratio of the coils so that each coil has roughly twice as much current as before. Probably dont want to push this any further.

Kept working on the examples. Cleaned up the grid class a little bit …

d31de9f

…but now seeing some issues with optimizing and not sure where this is coming from.

Had some success with smaller grids and bigger coils. Going to try an…

799e7b1

…d tune this a bit more.

Got the QH example there. Added new example that restricts how much t…

76527a7

…he PSCs can wiggle. Added new example using GPMO. Minor changes to the PSC grid class although reinstated aspect ratio 0.1 for the more complex runs.

Played around a bit more with a QA example and realized there is an o…

f108114

…pportunity to do joint optimization with the TF coils, but havent quite figured this out yet.

another small fix to try and get the join optimization consistent (le…

e28a90c

…t alone working well). The issue is that when the PSC coils start again, they seem to have the wrong TF coils, or not been optimized consistently wwhen the TF coils were going.

After some more debugging, joint optimization appears to be working. …

27d3452

…Downside is that it is very slow. Going to take another look at trying to get the gradients working.

Got something interesting working where I directly optimize the PSC c…

3f5a94c

…oils. Need to check that the PSC currents are being updated correctly and only the orientations are being optimized.

In the middle of debugging the psc optimization. Seem to be getting c…

e08647b

…loser but gradients are zero from the derivatives with respect to the currents and not sure why.

Slowly working up to getting the new jacobian working during filament…

03eafba

… optimization.

Few more changes to get things running at least.

fc63408

Made some small fixes, such as correctly using arcsin in c++ and pyth…

1d5f9f4

…on (only has one argument -- oops). Debugged a bit more. Got the unit tests all working again. Still not agreeing with FDs.

Got dpsi and dI derivatives seemingly working well in the unit tests.…

8e8c234

… Still have discrepancy in the QA psc example, but this may be because of the way its comparing with finite differences. Need to test the full chain of derivatives through the SquaredFlux.

Merge branch 'master' into psc_arrays

38fe23a

Still trying to check all the partial derivatives through the squared…

f280093

… flux calculation in unit tests.