Form Compiler#

NeuralMag features a form compiler that translates arbitrary functionals and linear weak forms into efficient, vectorized backend code. As an input it takes SymPy expressions written with a small set of special symbols for discretized functions and integration measures.

When you need this page

You normally don’t have to read any of this to use NeuralMag — the form compiler is an implementation detail of the built-in field terms. The only common reason to learn it is to add a new field term by writing its energy as a SymPy expression. Everything else is handled automatically.

Function spaces#

NeuralMag supports two function spaces for discretized fields:

a piecewise polynomial space with degrees of freedom on the nodes of the cuboid mesh, and
a piecewise constant space with values defined on the cells.

The basis functions \(\phi_{ijk}\) for the nodal space are defined per cell as the product of 1D functions in each principal direction:

\[\phi_{ijk}(x, y, z) = \phi_i^x(x)\,\phi_j^y(y)\,\phi_k^z(z),\]

where \(i,j,k\) are the node indices in each direction. The 1D basis functions \(\phi^x_i\) are first-order Lagrange polynomials satisfying the nodal condition

\[\phi^x_{i}(x_j) = \delta_{ij},\]

i.e.

\[\begin{split}\phi^x_{i}(x) = \begin{cases} (x - x_{i-1}) / (x_i - x_{i-1}) & \text{if } x_{i-1} \le x \le x_i, \\ (x - x_{i+1}) / (x_i - x_{i+1}) & \text{if } x_i \le x \le x_{i+1}, \\ 0 & \text{else}. \end{cases}\end{split}\]

The piecewise-constant basis functions \(\theta_{ijk}\) are defined on the cells bounded by the nodes:

\[\theta_{ijk}(x, y, z) = \theta^x_i(x)\,\theta^y_j(y)\,\theta^z_k(z),\]

with

\[\begin{split}\theta^x_i(x) = \begin{cases} 1 & \text{if } x_i \le x \le x_{i+1}, \\ 0 & \text{else}. \end{cases}\end{split}\]

NeuralMag supports nodal and cell-based functions in arbitrary dimensions and even allows mixing of nodal and cell discretization in different principal directions. For the symbolic representation of discretized fields, the form compiler exposes Variable(). For instance, a scalar field that is nodal in all 3 directions is declared as

u = Variable("u", "nnn")

The first argument is the variable name; the second specifies the space in each direction ("n" for nodal, "c" for cell). A piecewise constant 3D field would be "ccc"; a mixed space such as "ncn" yields the basis function

\[\kappa_{ijk}(x, y, z) = \phi_i^x(x)\,\theta_j^y(y)\,\phi_k^z(z).\]

Defining a functional#

A field contribution is described by an energy density: a function e_expr(m, dim, options) that combines Variable() fields and the integration measure dV into a single SymPy form. This is exactly the method a FieldTerm subclass implements. The built-in ExchangeField, for instance, is just

from neuralmag.common.engine import N, Variable, dV

def e_expr(m, dim, options):
    A = Variable("material__A", "c" * dim)   # cell-constant exchange constant
    return A * (m.diff(N.x).dot(m.diff(N.x))
                + m.diff(N.y).dot(m.diff(N.y))
                + m.diff(N.z).dot(m.diff(N.z))) * dV(dim)

Here m is the (vector, nodal) magnetization, A is a piecewise-constant material parameter on the cells, and dV(dim) is the volume measure. The form compiler turns this single expression into the field, energy-density and total-energy kernels through the three stages below.

The kernel-once pipeline#

Rather than symbolically unrolling Gauss quadrature and the test-function basis (one expanded term per quadrature point and per element corner), NeuralMag extracts the pointwise integrand once and lets the backend contract it against small reference-element tables. The whole pipeline lives in neuralmag.common.engine and the backend CodeBlock; you can run it by hand on any e_expr.

1. Extract the kernel. extract_kernel() builds m as a nodal quadrature field, evaluates e_expr, strips the measure, and rewrites every field value and gradient at the quadrature point into independent symbols (qv for a value component, qg for a gradient component):

kernel, meta, args = extract_kernel(e_expr, dim, options)

For the exchange energy the resulting kernel is the bare integrand

\[A\,\rho \sum_{l,d} \Bigl(\tfrac{\partial m_l}{\partial x_d}\Bigr)^2 ,\]

where \(\rho\) is the geometry density (see Domains and Material Density). meta maps each interpolated field to its (spaces, shape) (here {"m": ("nnn", (3,))}) and args carries the measure metadata. No differentiation has happened yet — extract_kernel works the same on a hand-written linear form.

2. Vary to get the effective field. The effective field is the variation (Gâteaux derivative) of the energy with respect to m. variation() differentiates the kernel with respect to the value and gradient symbols of a chosen field:

Kv, Kg, test = variation(kernel, "m", meta, dim)

It returns the coefficients of the test value (Kv[l]) and test gradient (Kg[l][d]). For exchange the energy depends only on \(\nabla m\), so

Kv = [0, 0, 0]
Kg[l][d] = 2 * A * rho * (d m_l / d x_d)     # coefficient of the test gradient

i.e. the familiar weak form \(\int 2A\,\nabla\vec v : \nabla\vec m\,\dx\).

3. Emit backend code. The backend CodeBlock assembles the linear form from the kernel coefficients using the discretization’s gather / interpolate / contract / scatter primitives and the reference-element tables from neuralmag.common.quadrature. The emitted functions are branch-free tensor code. The total-energy reduction generated for the exchange term, for example, reads (lightly trimmed):

def E(dx, m, material__A, rho):
    # gather: element-local corner values of m
    m_M = torch.stack([m[:-1,:-1,:-1], ..., m[1:,1:,1:]], dim=-2)
    # interpolate: gradient at the quadrature points via the reference table
    m_g = torch.einsum('qad,...al->...qld', _t0_DPHI, m_M) * (1.0 / dx[:3])
    vol = dx[0]*dx[1]*dx[2]
    # the pointwise kernel A·ρ·Σ (∂m_l/∂x_d)^2, evaluated at each quad point
    kern = material__A[..., None] * rho[..., None] * (
        m_g[..., 0, 0]**2 + ... + m_g[..., 2, 2]**2)
    # contract with the Gauss weights and sum over the mesh
    return (vol * torch.einsum('q,...q->...', _t0_W, kern)).sum()

The effective-field function h is assembled the same way from Kv/Kg and then divided by a mass-lumped node weight; the energy density e is the kernel projected onto nodes. FieldTerm.register() compiles these h, e and E functions through the active backend and binds them onto the state. You normally never call the three stages yourself — implementing e_expr is enough.

Choosing the quadrature order#

Because the tensor-product P1 basis is multilinear, the per-axis polynomial degree of a kernel in the quadrature symbols is known at compile time (axis_degrees). The form compiler uses it to pick the minimal Gauss order that integrates each form exactly, and to detect affine (degree \(\le 1\)) forms whose quadrature folds into a constant element matrix — which is why the exchange field above collapses to a single stencil contraction. A non-polynomial kernel (e.g. a field in a denominator) has no exact order and falls back to the configured config.fem["n_gauss"].

What about bilinear forms?#

NeuralMag does not generate code for bilinear forms in the sense of assembling explicit system matrices: with iterative solvers you only need the action of the operator on a vector, and the linear-form machinery already provides exactly that. Effective fields, projection forms (see Discretization), and the Gâteaux derivatives used by every field term are all expressed through linear forms.