Splineboris (Cleaned)

[SiBuijs]

Description

This is a pull request from the splineboris_cleaned branch, which itself is a squashed commit from the splineboris branch (see PR Splineboris#750). The reason for this was to remove large data files from the commit history.

Added the SplineBoris and SplineBorisSequence to the elements.

SplineBoris (xtrack/beam_elements/elements.py)

The SplineBoris is based on a magnetic field description from bpmeth. We describe the field components (Bx, By, Bs) and the derivatives of Bx and By w.r.t. x as fourth order polynomials. These coefficients are ultimately passed to a field-evaluation function that is constructed using bpmeth. SplineBoris accepts the coefficients of these polynomials, as well as the start and end s-coordinate of the element and the number of steps to track particles through the magnetic field. The tracking algorithm is the Boris integrator (see xtrack/beam_elements/elements_src/splineboris.h and xtrack/beam_elements/elements_src/track_splineboris.h), which is written in C, but adopted from the existing Python implementation from BorisSpatialIntegrator. The magnetic field evaluation function is generated from xtrack/beam_elements/elements_src/_generate_bpmeth_to_C.py. The resulting file is xtrack/beam_elements/elements_src/_spline_B_field_eval.h.

SplineBorisSequence (xtrack/beam_elements/elements.py)

The SplineBorisSequence element allows for defining a sequence of SplineBoris elements to describe more complicated field structures. It requires a pandas dataframe of aforementioned polynomial coefficients indexed by their start and end s-coordinates. The SplineBorisSequence constructs a line consisting of individual SplineBoris elements, which can used as a usual line.

FieldFitter (see xtrack/_temp/field_fitter.py)

The required pandas dataframe can be generated using the FieldFitter class, which requires a path to a file of raw data which can either be a csv file or pandas dataframe and is set to df_raw_data. At each $s$-point in the data, it fits $B_x(x)$ and $B_y(x)$ as a polynomial of deg+1. It then computes successive derivatives of these polynomials and evaluates them at $x=x_0$, the $x$-coordinate of the longitudinal axis. The values of these derivatives are stored in df_on_axis_raw.

After df_on_axis_raw is filled, the fitter checks the order-of-magnitude of the derivatives w.r.t. the main field. This tolerance can be set with field_tol, which is $10^-3$ by default. More specifically, it checks if:

$$\frac{1}{n!} \mathrm{max}\bigg| \frac{\partial^n B_i}{\partial x^n} \bigg| \mathrm{max}|x|^n < \mathrm{field_tol} \mathrm{max}(B_x,\ B_y,\ Bs)$$

Here, $n$ is the order of the derivative and $i \in (x,\ y)$. If the above is true, it is assumed that the particular combination of $B_i$ and derivative order $n$ has negligible effect on the dynamics and will not be fitted. This serves as a noise-filter.

Then, the FieldFitter finds extrema in the remaining fields/derivatives and defines these extrema as region boundaries. Then, it cuts the regions between the extrema in regions of min_region_size points. The aim of this method is to provide good accuracy around the extrema, as well as in between them.

To each of these regions, a polynomial of degree 4 is fitted. The resulting polynomial coefficients are stored in df_fit_pars, indexed by their magnetic field, derivative, s_start, s_end, idx_start and idx_end. This df_fit_pars can be read directly by SplineBorisSequence to build a sequence of individual SplineBoris elements with the correct set of parameters.

Parameter Ordering

An important point of note: The field evaluation function (_spline_B_field_eval.h) requires the parameters to be ordered as: The coefficients of the Bs polynomial, the coefficients of Bx and subsequent derivatives, the coefficients of By and subsequent derivatives. This parameter ordering is defined inside SplineBoris under the ParamFormat class, which contains several methods to check parameter ordering, as well as the definition of the polynomial used to describe the field. The FieldFitter etc. derive their parameter ordering and naming convention from ParamFormat.

Tests:

• test_splineboris_homogeneous_analytic: Tests the SplineBoris for a homogeneous magnetic field at various angles against an analytical solution.
• test_splineboris_homogeneous_rbend: Tests the SplineBoris for a similar field as in the test above, but against the RBend in Xsuite.
• test_uniform_solenoid: Tests the SplineBoris for a uniform solenoid against the UniformSolenoid in Xsuite.
• test_splineboris_solenoid_vs_variable_solenoid: Tests the SplineBoris(Sequence) against the VariableSolenoid in Xsuite.
• test_splineboris_undulator_vs_boris_spatial: Tests the SplineBoris(Sequence) against the BorisSpatialIntegrator in Xsuite on the same field map.
• test_splineboris_rotated_undulator_vs_boris_spatial: Rotates the field map from the previous test by 90 degrees and checks that the fit-parameters are also rotated by 90 degrees (that is, $B_x \rightarrow B_y$ and $B_y \rightarrow -B_x$).
• test_splineboris_radiation: Tests mean and quantum radiation in a homogeneous magnetic field against an analytical solution.
• test_splineboris_variable_solenoid_radiation: Similar to test_splineboris_solenoid_vs_variable_solenoid, but tests radiation.
• test_splineboris_spin_uniform_solenoid: Tests the SplineBoris spin tracking against the UniformSolenoid in Xsuite.
• test_splineboris_spin_quadrupole: Tests the SplineBoris spin tracking in an ideal quadrupole against the Quadrupole in Xsuite.

Examples:

• 000a_sls_no_undulators: Reference to later examples, twissing the SLS
• 000b_sls_no_undulators_closed_spin: Reference to later examples, SLS with spin tracking
• 001_fieldfitter_basic_usage: Simple example applying the FieldFitter and plotting the results
• 002_homogeneous_rbend: Similar to the RBend test, demonstrates the SplineBoris against the RBend for several "roll" angles
• 003a_undulator_open_spin_tracking: Open spin-tracking through the undulator
• 003b_undulator_open_spin_tracking_radiation: Open spin and radiation tracking through the undulator
• 004a_sls_with_undulators: Inserts the undulators in the SLS and performs a closed orbit Twiss
• 004b_sls_with_undulators_closed_spin: Same as 004a but with spin
• 004c_sls_with_undulators_closed_spin_radiation: Same as 004a but with spin and radiation
• 004d_sls_offset_undulators: Same as 004a, but with the undulators offset
• 004e_sls_offset_undulators_closed_spin: Same as 004d, but with spin tracking
• 005_solenoid: Similar to the test_splineboris_solenoid_vs_variable_solenoid test, but with a few plots to show the tracks

Small Fix

In test_data, there used to be two directories, named sls and sls_2.0, which contained madx files with different names, but same content. Removed sls_2.0 and changed all references to the file therein to test_data/sls. These changes occur in examples/tests involving the radiation integrals. The relevant tests pass.

In addition: Added the calculation of the damping partition numbers to the radiation integrals in twiss.py:
$$J_x = 1 - I_{4x}/I_2$$
$$J_y = 1 - I_{4y}/I_2$$
$$J_\epsilon = 2 + I_4/I_2$$

Closes # .

Checklist

Mandatory:

• I have added tests to cover my changes
• All the tests are passing, including my new ones
• I described my changes in this PR description

Optional:

• The code I wrote follows good style practices (see PEP 8 and PEP 20).
• I have updated the docs in relation to my changes, if applicable
• I have tested also GPU contexts

[giadarol]

check that this uses the file in test_data folder

[SiBuijs]

Checked, was incorrect, corrected it.

[giadarol]

As above

[SiBuijs]

Same, was incorrect, corrected it.

[giadarol]

This path refers to your computer

[SiBuijs]

Corrected!

[giadarol]

Still taking a path, let's use a DataFrame

[giadarol]

Let's not Taylor the fitter to a specific format

[SiBuijs]

Removed the "path acceptance" from the FieldFitter. Now only accepts DataFrames. The user has to convert their data to a DataFrame of the correct format.

[giadarol]

Give directly the field component derivaritives at edges and average

[SiBuijs]

Done. We now only store:

$$f(s_n)$$ $$f'(s_n)$$ $$f(s_{n+1})$$ $$f'(s_{n+1})$$ $$\frac{1}{s_{n+1} - s_n} \int_{s_n}^{s_{n+1}} f(s) \ ds$$

Field evaluator still requires the usual polynomial coefficients (or we would have to cut open bpmeth). This makes the fitter much cleaner. Conversion from the above to normal polynomial coefficients happens only at construction of SplineBoris.

[giadarol]

note down the version of BPMETH you used (as it is still evolving quite a bit, it might be useful for future debugging to be able to go back

[SiBuijs]

Done. Version seems to be stuck at 0.0.0, so specified the commit hash. Let me know if that should be done differently.

[giadarol]

If this is part of the Xtrack source code, the understore should go away

[SiBuijs]

Done!

[giadarol]

Add a method to evaluate the field (you can import "_spline_B_field_eval_python.py‎")
Rename it to have no underscore.

Keep the underscore only for files that are not executed by the user

[SiBuijs]

Done, do we need a plotter as well?

[SiBuijs]

Done. Also added a demonstration of the field evaluation in examples/splineboris/005_solenoid.py, which plots the field against the analytical solution (and matches very well).