Code Generation

An equilibrium object carries nothing but the parameters of a field. It cannot be evaluated, and it is not meant to be: the quantities a simulation needs are the magnetic field, the frame along it and a tower of derivatives, none of which the equilibrium states. They are derived from it, symbolically, and compiled. This page describes how, what that costs, and what to do when the result is not what was expected.

What the Equilibrium Supplies

An equilibrium is required to define very little. Everything on this list is a method taking the coordinates and the equilibrium, written generically enough to accept symbolic arguments:

periodic is the one exception to the shape above: it takes the equilibrium alone, because the answer is a property of the chart rather than of a point in it. It is also the only entry with no default. A chart that has not answered raises a MethodError when a field is built from it, which is deliberate — an all-false default would let a periodic chart report no periodicity silently, and the bounds cannot be read for it, since a bounded range does not imply that a coordinate wraps.

The name is GeometricBase.periodic, which already means one Bool per component there. It is not periodicity: GeometricEquations gives that generic an (xmin, xmax) tuple, so answering it with a Bool vector would give one name two shapes.

A field built on CartesianEquilibrium inherits the identity chart, J = 1, the euclidean metric, orientation = +1 and no periodicity in any coordinate, so in that case only the vector potential is left. That is why ThetaPinchEquilibrium is a few dozen lines.

The struct itself holds the parameters — B₀ for the θ-pinch, R₀, B₀ and q₀ for the tokamaks — and the methods read them off it. They are not baked into the generated code; see The Parameters Are Arguments, Not Literals.

The Symbolic Trace

FieldFunctions(equ) begins by declaring symbolic variables for the time and the three coordinates and calling the methods above with them. Because those methods are written generically over AbstractVector, they return symbolic expressions rather than numbers, and what comes back is the vector potential and the chart as algebraic expressions in $\xi$.

From there the derivation is the standard one of differential geometry, and it is carried out with Symbolics.jl:

1. The tangent map $DF$ is the derivative of the chart; its inverse $\bar{DF}$ is the derivative of the inverse chart, pulled back onto $\xi$.
2. The metric is inverted symbolically, and both are differentiated twice.
3. The magnetic two-form is the exterior derivative of the vector potential, and the magnetic field is its Hodge dual. The Hodge star is orientation-dependent and is given the signed determinant $\det DF = \mathrm{orientation} \cdot J$, not the volume element.
4. The magnitude, the unit vector and the perpendicular frame follow, the frame being built from the first coordinate basis vector whose cross product with $b$ the trace does not show to be zero. The choice is made once, so the frame reverses where $b$ is parallel to that vector; The Frame along the Magnetic Field says where.
5. The electric field is minus the gradient of the potential.
6. Everything is differentiated as far as the tables in Interface say.

The full derivation is in Fields.

Each resulting expression is then handed to Symbolics.build_function, which turns it into Julia code and compiles it. The functions are RuntimeGeneratedFunctions: ordinary compiled Julia functions that were built at run time rather than parsed from a file. They are stored in the FieldFunctions object, which is the reason a field is a value that can be passed around rather than a set of names in a module.

The Parameters Are Arguments, Not Literals

One detail of the trace decides a great deal. It does not run against the equilibrium you passed in, but against a copy of it whose parameters are symbolic. AxisymmetricTokamakCylindrical is traced in terms of $R_0$, $B_0$ and $q_0$ rather than 6.2, 5.3 and 2.0, so the generated code takes them as an argument. The equilibrium struct remains where the values live, and the field passes them in on every call.

The consequence is that the generated code depends on the equilibrium's type, not on its parameters. Two fields of the same type share it exactly:

using ElectromagneticFields

a = FieldFunctions(AxisymmetricTokamakCylindricalEquilibrium(1.0, 1.0, 2.0))
b = FieldFunctions(AxisymmetricTokamakCylindricalEquilibrium(6.2, 5.3, 1.7))

shared = typeof(functions(a).B♭.f) === typeof(functions(b).B♭.f)
shared, B♭(a, 0.0, [1.05, 0.25, 0.5]), B♭(b, 0.0, [1.05, 0.25, 0.5])

(true, [-0.11904761904761904, 0.02380952380952383, 1.0], [-0.742296918767507, -15.291316526610647, 32.86])

Same code, different answers. This is why the precompilation described below works for parameter values nobody anticipated, and it is why parameter_names matters: anything the A₁, φ or metric methods read that is not listed as a parameter gets frozen into the code as a literal. The default — every field of the struct except name — is what the equilibria here want.

Common Subexpression Elimination

Symbolic differentiation produces expressions that share a great deal of structure and state none of it. The second derivatives of the Solov'ev flux function contain the same logarithm dozens of times over, and written out literally they are both unreadable and slow.

build_function is therefore asked to eliminate common subexpressions: each repeated subexpression is computed once into a local and referred to afterwards. This is on by default and is value-preserving — subexpressions are named, never rewritten — and it can be turned off:

field = FieldFunctions(equ; cse = false)

The only reason to do so is to read the generated code against a paper. The result is considerably slower.

Reading the Generated Code

To see what was generated for one quantity, take the symbolic expressions and ask build_function for an expression instead of a function:

using ElectromagneticFields: generate_field_expressions
using Symbolics

exprs = generate_field_expressions(AxisymmetricTokamakCylindricalEquilibrium(),
                                   ZeroPerturbation())
t, ξ, p = exprs.arguments

Symbolics.build_function(exprs.B♭, t, ξ, p; expression = Val{true}, cse = true)

(:(function (t, ˍ₋arg2, ˍ₋arg3)
      #= /home/runner/.julia/packages/Symbolics/JBOvO/src/build_function.jl:441 =# @inbounds begin
              #= /home/runner/.julia/packages/Symbolics/JBOvO/src/build_function.jl:441 =#
              begin
                  begin
                      var"##cse#1" = -1
                      var"##cse#2" = (*)((*)(var"##cse#1", ˍ₋arg2[2]), ˍ₋arg3[2])
                      var"##cse#3" = (*)(ˍ₋arg2[1], ˍ₋arg3[3])
                      var"##cse#4" = (/)(var"##cse#2", var"##cse#3")
                      var"##cse#5" = (*)(var"##cse#1", ˍ₋arg3[1])
                      var"##cse#6" = (+)(ˍ₋arg2[1], var"##cse#5")
                      var"##cse#7" = (*)(var"##cse#6", ˍ₋arg3[2])
                      var"##cse#8" = (/)(var"##cse#7", var"##cse#3")
                      var"##cse#9" = 1//2
                      var"##cse#10" = -1//2
                      var"##cse#11" = (*)((*)(var"##cse#10", ˍ₋arg3[1]), ˍ₋arg3[2])
                      var"##cse#12" = (/)(var"##cse#11", ˍ₋arg2[1])
                      var"##cse#13" = (*)((*)(var"##cse#1", ˍ₋arg3[1]), ˍ₋arg3[2])
                      var"##cse#14" = 2
                      var"##cse#15" = (*)(var"##cse#14", ˍ₋arg2[1])
                      var"##cse#16" = (/)(var"##cse#13", var"##cse#15")
                      var"##cse#17" = (+)(var"##cse#12", var"##cse#16")
                      var"##cse#18" = (*)(var"##cse#9", var"##cse#17")
                      var"##cse#19" = (*)(var"##cse#1", var"##cse#12")
                      var"##cse#20" = (*)(var"##cse#1", var"##cse#16")
                      var"##cse#21" = (+)(var"##cse#19", var"##cse#20")
                      var"##cse#22" = (*)(var"##cse#10", var"##cse#21")
                      var"##cse#23" = (+)(var"##cse#18", var"##cse#22")
                      var"##cse#24" = (*)((*)(var"##cse#1", ˍ₋arg2[1]), var"##cse#23")
                      begin
                          #= /home/runner/.julia/packages/SymbolicUtils/fK7cM/src/code.jl:1337 =#
                          (SymbolicUtils.Code.create_array)(typeof(ˍ₋arg2), nothing, Val{1}(), Val{(3,)}(), var"##cse#4", var"##cse#8", var"##cse#24")
                      end
                  end
              end
          end
  end), :(function (ˍ₋out, t, ˍ₋arg2, ˍ₋arg3)
      #= /home/runner/.julia/packages/Symbolics/JBOvO/src/build_function.jl:443 =# @inbounds begin
              #= /home/runner/.julia/packages/Symbolics/JBOvO/src/build_function.jl:443 =#
              begin
                  begin
                      var"##cse#1" = -1
                      var"##cse#2" = (*)((*)(var"##cse#1", ˍ₋arg2[2]), ˍ₋arg3[2])
                      var"##cse#3" = (*)(ˍ₋arg2[1], ˍ₋arg3[3])
                      var"##cse#4" = (/)(var"##cse#2", var"##cse#3")
                      var"##cse#5" = (*)(var"##cse#1", ˍ₋arg3[1])
                      var"##cse#6" = (+)(ˍ₋arg2[1], var"##cse#5")
                      var"##cse#7" = (*)(var"##cse#6", ˍ₋arg3[2])
                      var"##cse#8" = (/)(var"##cse#7", var"##cse#3")
                      var"##cse#9" = 1//2
                      var"##cse#10" = -1//2
                      var"##cse#11" = (*)((*)(var"##cse#10", ˍ₋arg3[1]), ˍ₋arg3[2])
                      var"##cse#12" = (/)(var"##cse#11", ˍ₋arg2[1])
                      var"##cse#13" = (*)((*)(var"##cse#1", ˍ₋arg3[1]), ˍ₋arg3[2])
                      var"##cse#14" = 2
                      var"##cse#15" = (*)(var"##cse#14", ˍ₋arg2[1])
                      var"##cse#16" = (/)(var"##cse#13", var"##cse#15")
                      var"##cse#17" = (+)(var"##cse#12", var"##cse#16")
                      var"##cse#18" = (*)(var"##cse#9", var"##cse#17")
                      var"##cse#19" = (*)(var"##cse#1", var"##cse#12")
                      var"##cse#20" = (*)(var"##cse#1", var"##cse#16")
                      var"##cse#21" = (+)(var"##cse#19", var"##cse#20")
                      var"##cse#22" = (*)(var"##cse#10", var"##cse#21")
                      var"##cse#23" = (+)(var"##cse#18", var"##cse#22")
                      var"##cse#24" = (*)((*)(var"##cse#1", ˍ₋arg2[1]), var"##cse#23")
                      #= /home/runner/.julia/packages/SymbolicUtils/fK7cM/src/code.jl:1272 =# @inbounds begin
                              #= /home/runner/.julia/packages/SymbolicUtils/fK7cM/src/code.jl:1268 =#
                              ˍ₋out[1] = var"##cse#4"
                              ˍ₋out[2] = var"##cse#8"
                              ˍ₋out[3] = var"##cse#24"
                              #= /home/runner/.julia/packages/SymbolicUtils/fK7cM/src/code.jl:1270 =#
                              ˍ₋out
                          end
                  end
              end
          end
  end))

exprs.arguments are the symbols the code is written in: the time, the three coordinates, and the parameters — which appear in the result as R₀, B₀ and q₀ rather than as numbers, since the code takes them as its third argument. The rest of exprs is a NamedTuple keyed exactly like functions(field), so any quantity in the tables of Interface can be inspected this way. Passing cse = false gives the unshared form, which is the one to read against a paper.

What It Costs

The trace, the code generation and the first compilation all happen inside the call to FieldFunctions. Evaluating the result afterwards is cheap — every accessor is allocation-free and type-stable, and a whole tensor is computed in one pass with its subexpressions shared, which is usually faster than computing the components one at a time.

Building it is the part worth understanding. Measured on the toroidal tokamak in a cold session:

The striking thing is where the time is not. Compiling the generated code is 5% of the bill. Almost all of it is Julia compiling Symbolics' own machinery for the expression types a trace produces — and that is paid per expression shape, not per field. Because the parameters are arguments, the shape is a property of the equilibrium's type, so building the same equilibrium again with different parameters costs nothing measurable.

That is what makes the cost recoverable, and this package recovers all of it for the equilibria it ships. Two things combine:

• the generated functions of a type already built are cached, keyed by what actually determines them — the equilibrium and perturbation types, the parameter names and shapes, cse and cache_module — but not by the parameter values, which the code does not contain;
• a PrecompileTools workload builds every shipped equilibrium and traces a field of it during precompilation, so the constructors, the compiled specializations and the cache itself all land in the package image.

A fresh session therefore finds them already built. Measured against the same package with the workload removed, one cold process per figure:

at the price of this package's own precompilation, 1.4 s → 20 s, paid once per version.

The workload names every shipped Solov'ev preset, not one per type. Those seven span three types, so a field is a cache hit for all but the first of each — but a preset is its own function and has to be compiled: leaving the FRC and NSTX ones out puts the constructor figure above at 0.11 s instead of 0.06 s, and naming them costs 0.2 MB of package image and 0.2 s of precompilation. The three AxisymmetricTokamak*ITER presets are not named; they wrap a constructor the workload already builds, and compiling them costs 10 µs.

An equilibrium you define yourself is traced the first time and cached thereafter, so it costs a fraction of a second once per session — or nothing at all, if you precompile it in your own package as described below.

One consequence of the cache is worth knowing. It is a Dict, so FieldFunctions(equ) is not type-inferable at its call site: the object it returns is concretely typed and every accessor on it is type-stable and allocation-free, but the construction itself is not. Fields are normally built in a constructor or a setup step and then handed to something else, which is a function barrier already, so this costs nothing in practice — but build a field and evaluate it in the same function body and the evaluation will be inferred as Any.

Precompiling a Field in Your Own Package

The remaining fraction of a second per field can be removed too. A package that always uses the same field can build it during its own precompilation, so loading it costs nothing at all.

Two things are needed. A generated function keeps its body in a cache inside some module, and that cache only survives precompilation when the module is the one being precompiled; @precompilable_fields prepares the calling module, and cache_module tells FieldFunctions to use it. A PrecompileTools.@compile_workload then caches the compiled code of the accessors as well.

module MyModel

using ElectromagneticFields
using PrecompileTools

@precompilable_fields

const FIELD = FieldFunctions(AxisymmetricTokamakCylindricalEquilibrium(6.2, 5.3, 2.0);
                             cache_module = @__MODULE__)

sample(t, ξ) = B♭(FIELD, t, ξ)

@setup_workload begin
    ξ = [6.5, 0.5, 0.25]
    @compile_workload begin
        sample(0.0, ξ)
    end
end

end

Loading MyModel in a fresh session then costs only what loading ElectromagneticFields costs, and the first evaluation of the field takes microseconds: no trace, no code generation, no compilation. Omitting cache_module is not an error — the field is simply rebuilt on every load.

Adding a Field

Adding a field to the package, or defining one in your own code, means adding methods to the generics in the first table. For a field in cartesian coordinates, the vector potential is enough:

using ElectromagneticFields
import ElectromagneticFields: A₁, A₂, A₃, get_functions, X, Y, Z

struct MyPinch{T <: Number} <: ElectromagneticFields.CartesianEquilibrium
    name::String
    B₀::T
    MyPinch{T}(B₀::T) where {T <: Number} = new("MyPinch", B₀)
end

MyPinch(B₀::T) where {T <: Number} = MyPinch{T}(B₀)

A₁(x::AbstractVector, equ::MyPinch) = -equ.B₀ * Y(x, equ) / 2
A₂(x::AbstractVector, equ::MyPinch) = +equ.B₀ * X(x, equ) / 2
A₃(x::AbstractVector, equ::MyPinch) = zero(eltype(x))

get_functions(::MyPinch) = (X = X, Y = Y, Z = Z)

field = FieldFunctions(MyPinch(2.0))
B♭(field, 0.0, [0.5, 0.5, 0.5])

3-element StaticArraysCore.SVector{3, Float64} with indices SOneTo(3):
 0.0
 0.0
 2.0

Three rules govern what those methods may contain.

Follow the constructor convention. The inner constructor takes the parameters in order and sets name itself, and there is an outer one supplying defaults — which is what lets the trace build the copy with symbolic parameters. MyPinch{T}(B₀::T), not MyPinch(B₀::T): the trace calls MyPinch{Num}(...). A type that cannot be written this way needs a symbolic_copy method.

Write them generically. The argument is annotated AbstractVector and the element type is left open, because the generator calls them with a vector of symbolic variables. Annotating AbstractVector{Float64}, or calling a function that only accepts numbers, fails at trace time. zero(eltype(x)) rather than 0.0 is the same rule applied to constants.

Every function used must be differentiable by Symbolics. The elementary functions are, as are NaNMath.log and the two-argument atan, both of which this package relies on. A branch on the value of a coordinate is not: if x[1] > 0 cannot be traced, and the expression has to be written without it.

A new chart needs its geometry to be consistent. Defining x¹ … ξ³, the metric, J, orientation and periodic is not enough on its own; all but periodic have to fit together. The test suite asserts $J = \sqrt{\det(DF^T DF)}$ and $\det DF = \mathrm{orientation} \cdot J$ for every equilibrium, and the identities in Interface are the same checks to run on a chart of your own. Handedness in particular has no visible symptom when it is wrong other than a magnetic field pointing the wrong way — see orientation.

A perturbation is defined the same way, subtyping AnalyticPerturbation (or CartesianPerturbation) and defining whichever of A₁, A₂, A₃ and φ it contributes. It must live on the same chart as the equilibrium it perturbs, and the generator asserts that the two agree on J and g. The two are added together symbolically before any code is generated, so FieldFunctions(equ, pert) is a field like any other and costs nothing extra to evaluate.

When It Goes Wrong

A MethodError during FieldFunctions almost always means one of the methods above is not generic enough — a concrete type annotation, or a call into something that does not accept symbolic arguments. The stack trace names the method.

A field that evaluates to NaN on a chart with a singular point is usually the chart rather than the generator: 1/cos²θ gauges and log R potentials are genuinely undefined where they are undefined. NaNMath variants are used where a domain error would otherwise propagate.

A magnetic field pointing the wrong way is an orientation error. Compare against a finite-difference curl of the vector potential in cartesian coordinates, which is what test_curl in the test suite does.

A disagreement with the reference values can be checked against scripts/verify_against_symengine.jl, which evaluates every quantity for all twenty equilibria and compares them with the recorded values in scripts/reference-symengine.tsv.