Cbs/fieldine integrator

[smiet]

Revamp of fieldline integration. Some tests still need to be written, but this should simplify the user experience in generating Poincare plots.

What’s new

Integrator base class with subclasses:

• SimsoptFieldlineIntegrator (interface to the sopp integration routines
• ScipyFieldlineIntegrator (integration of the magnetic field using scipy.solve_ivp methods).

These subclasses provide methods to return Poincare sections, three dimensional field line trajectories, as well as the ability to integrate from any starting point over a given toroidal distance.

The ScipyFieldlineIntegrator provides methods to integrate $R$ and $Z$ using $\phi$ as the 'time' variable in the ODE, providing crisper Poincare sections and integrates about as fast as the sopp methods (which evolves the three coordinates independently).

The Integrator subclasses also provide an interface to the PoincarePlotter class, which provides an easy interface to create Poincare plots of the field.

The PoincarePlotter uses caching and the simsopt dependency graph (parent Integrator with parent MagneticField) to trigger recomputes when necessary. Manipulation of the res_phi_hits and res_tys arrays is a thing of the past!

give it a spin!:

from simsopt.configs import get_data
from simsopt.field import SimsoptFieldlineIntegrator, PoincarePlotter
from simsopt.geo import plot
import numpy as np

base_coils, base_currents, ma, nfp, bs = get_data('ncsx')

axis_rz = ma.gamma()[0][::2]
# startpoints linspace from axis out: 
start_points_RZ = np.linspace(axis_rz, axis_rz+np.array([0.1, 0]), 10)


integrator = SimsoptFieldlineIntegrator(bs, stellsym=True, nfp=nfp, R0=axis_rz[0], tol=1e-9)
pplotter = PoincarePlotter(integrator, start_points_RZ, n_transits=200, phis=4)

fig, axs = pplotter.plot_poincare_all()


pplotter.plot_fieldline_trajectories_3d(engine='mayavi', show=False, opacity=0.2, tube_radius=0.001)
plot(bs.coils, engine='mayavi', show=False)
pplotter.plot_poincare_in_3d(engine='mayavi', show=True, scale_factor=0.01)

[codecov[bot]]

Codecov Report

:x: Patch coverage is 89.14186% with 62 lines in your changes missing coverage. Please review. :white_check_mark: Project coverage is 92.45%. Comparing base (0da1f07) to head (bdc0ab5). :warning: Report is 13 commits behind head on master.

Files with missing lines	Patch %	Lines
src/simsopt/field/tracing.py	89.14%	62 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #555      +/-   ##
==========================================
- Coverage   92.56%   92.45%   -0.12%     
==========================================
  Files          83       83              
  Lines       16229    16798     +569     
==========================================
+ Hits        15023    15530     +507     
- Misses       1206     1268      +62     

Flag	Coverage Δ
unittests	92.45% <89.14%> (-0.12%)	:arrow_down:

Flags with carried forward coverage won't be shown. Click here to find out more.

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[smiet]

@mishapadidar, @landreman, this PR is ready for review after fiddling a bit with some final tests.

I know the codecov is a bit low, but that is because the runners cannot test the plotting routines using Mayavi (I am using mayvi myself, so I know these routines work).

[smiet]

@andrewgiuliani Thank you for your review, I have added tests to cover all the missed lines that I could. The ones that are not hit are:

• the mayavi plotting routines (because we do not have a runner with mayavi installed, but they will run when that is the case)
• The lines that gather data from several processes (though these should be hit in the examples, which are run with multiprocessing, but maybe coverage is unaware)
• a few others for which I could not find an elegant way to hit, like the exception to not make proc0print fail on non-mpi but mpi-capable systems.

I hope that we can accept the few misses, the diff is only a few percent below the current coverage.

There are some other failing tests because SPEC install on pytnon 3.9 is currently broken.

@landreman @mishapadidar, could I ask you for a review if you have a bit of time to spare? A few people already using it have told me this is very pleasant to work with, and I would like to make these features available in the not-too-distant-future.

I know that there are many more features to add, but I will deal with these in later PRs, otherwise this one will bloat even larger than it currently is.

Future plans include:

• [ ] make InterpolatedField aware of its underlying MagneticField DoF's and trigger re-computation on first access after DoF change.
• [ ] include a SimsoptParticleIntegrator (requires me to first familiarize myself with how particle tracing works... or I could delegate...)
• [ ] deal with BoozerMagneticField in integration and in particle integration.
• [ ] methods for Integrator to generate Surface and Curve objects from the integration results.

[smiet]

Thanks @mishapadidar for your points so far. You raise a few points for discussion:

Implement compute_fieldlines with the scipy integration?

I would like to move away completely from the functional paradigm for Poincare plotting, and suggest users to only use class-based Integrators. I left the compute_fieldlines method for reverse compatibility, but want to deprecate that method of performing integration. @landreman @andrewgiuliani, what do you think, update the compute_fieldlines function or move towards deprecating it (and leave it be for now)?

Limiting the scope of each integrator

Scipy and Simsopt fieldline integrators perform different computations, and use different methods, but should we limit each integrator to only one way of solving the equations? Should we merge the two integrators and let it switch between whether sopp or scipy is used based on the requested data from the PoincarePlotter?

I would hope to leave this work to a later PR, and I am not equipped to modify the sopp integrator to implement the features I want

Derivatives of integration

This is in the pipeline. I am trying to get this PR out to cement the class interfaces, and am developing a topological optimisation toolset in parallel that will do this. This will involve giving the ScipyFieldlineIntegrator methods to differentiate integrations, evaluate action integrals, etc. and defining new Optimizables with proper J and dJ functions.