Sofia: a simulation-expert agent with its own DSL for physics

The Sofia DSL (Domain Specific Language) provides the right level of abstraction for designing physics simulations: mathematics. The Sofia agent draws on a base of domain-specific knowledge, deep understanding of the Sofia DSL, and previous experience to rapidly produce bespoke simulations in the Sofia DSL.

Sofia DSL and code generation

The Sofia DSL describes a simulation at the level of the equations, building on the popular Python library Sympy, so the interface feels familiar. Fields are continuous. A conservation law is written as du/dt = F(u, ∇u, …) using operators such as diff, laplacian and divergence. Derivative stencils, boundary conditions, and time-stepping schemes are lowered automatically into discrete forms. Every level of lowering is directly accessible, meaning that discrete methods such as lattice-Boltzmann can be expressed in the same language. Boundaries, staggered grids, and geometries from CAD are first-class citizens. Sofia's code generation layers deterministically turn symbolic expressions into concrete array operations, apply staggering, boundary conditions and optimisations such as CSE. The code generation emits per-cell updates as NVIDIA Warp kernels, with the control flow, JIT compilation, reductions and gradients handled by JAX. All of this is only made possible by the interoperability of Warp and JAX.

Because the physics is written at the level of the equations, a solver program is short and easy for a modelling expert to read. An incompressible Navier–Stokes solver with a penalised body and a far-field fringe is twelve lines of equations in Sofia. Written directly as Warp kernels, the same solver needs the derivative stencils, the staggered indexing, the boundary conditions and ghost cells, the projection step and the Runge–Kutta staging all spelled out by hand, which runs to several hundred lines of kernel code. This difference has the most impact when the physics changes. Adding a term to the momentum equation or swapping the advection scheme is a one-line change in Sofia. In hand-written kernels it requires a rewrite of every kernel that touches that field, plus a new round of debugging and validation.

# Incompressible Navier–Stokes in Sofia: penalised body, far-field fringe, projection
G = Grid(ndim=2)
x, y = G.coordinates
u, v, p = (Field(n, domain=G) for n in ("u", "v", "p"))

def rhs(a, a_free):
    return (-u * diff_upwind(a, x, variable=u) - v * diff_upwind(a, y, variable=v)
            + nu * laplacian(a)            # viscous
            - chi * a / eta                # no-slip by Brinkman penalisation
            - sig * (a - a_free))          # far-field fringe

step = ir.RKn(F={u: rhs(u, U0), v: rhs(v, 0.0)}, order=3, dt=dt,
              post=ir.ProjectDivFree((u, v), p))
sim = backend.jit(build_simulation(code, specs, kernel_backend="warp"))

Sofia agent

Sofia ships with an agent harness. The agent knows which physics to apply for a given simulation task and which numerical schemes best suit it. It knows how to build and run a Sofia simulation, how to read the result, and what common failure modes to look for. It runs the solver and validates things such as conservation, boundary behaviour and grid convergence. What it hands back is not a black-box solver with opaque physics, but a validated solver written in the Sofia DSL, short enough for the engineer who asked for it to read, review and change.

Figure 1. The Sofia pipeline. The agent iterates on a solver written in the Sofia DSL until it passes validation, then delivers a solver package. Sofia generates Warp kernels for the cell updates, geometry handling, and meshing, with JAX code for control flow, targeting NVIDIA GPUs or a CPU from the same script.

How NVIDIA Warp accelerates Sofia

Sofia does not have its own GPU runtime: it lowers operations to NVIDIA Warp kernels and JAX functions, and orchestrates them with JAX. The geometry, meshing, and particle simulation pipelines of Sofia build on the components offered by the Warp framework. The Warp/JAX interoperability means that features such as auto-differentiation through a full solve work out of the box.

Engineering problems usually arrive with CAD models. Warp's mesh data structures and methods such as closest-point and signed-distance queries let Sofia handle complex geometries without leaving the GPU, enabling GPU-native implementations of techniques required for modern physics simulation such as wall modelling in computational fluid dynamics.

Code complexity reduction from Sofia

Below is part of a Warp kernel for one part of a simulation, generated by Sofia. This is the type and length of code that would need to be written, maintained, and verified by experts (or coding agents) building a simulation. Below it is the Sofia code that generated it.

@warp.kernel
def kernel_impl_3907048(
    v_s_1: warp.array2d[warp.float64],
    v_s_0: warp.array2d[warp.float64],
    zeta_s_2: warp.array2d[warp.float64],
    zeta_s_1: warp.array2d[warp.float64],
    zb: warp.array2d[warp.float64],
    zeta_s_0: warp.array2d[warp.float64],
    v_s_2: warp.array2d[warp.float64],
    u_s_2: warp.array2d[warp.float64],
    u_s_0: warp.array2d[warp.float64],
    u_s_1: warp.array2d[warp.float64],
    u_3906986: warp.array2d[warp.float64],
):
    i, j = warp.tid()
    x0 = zeta_s_0[i, j]
    x1 = zeta_s_1[i, j]
    x2 = zeta_s_2[i, j]
    x3 = zeta_s_0[((1) + (i)) % 204, j]
    x4 = zeta_s_1[((1) + (i)) % 204, j]
    x5 = zeta_s_2[((1) + (i)) % 204, j]
    x6 = u_s_0[i, j]
    x7 = u_s_1[i, j]
    x8 = u_s_2[i, j]
    x9 = u_s_0[((1) + (i)) % 204, j]
    # ... x10-x27: 18 more stencil reads, face-averaged v, upwind coefficients
    u_3906986[i, j] = (
        ((warp.float64(0.998096209169431)) * (x6))
        + ((warp.float64(0.000475947707642355)) * (x10))
        + ((warp.float64(0.000475947707642355)) * (x11))
        + ((warp.float64(0.000475947707642355)) * (x12))
        # ... 17 more terms: RK3 weights (1/6, 4/6, 1/6) x viscosity and pressure gradient
        + (
            (warp.float64(-0.0951895415284710))
            * (x22)
            * (
                (
                    (
                        ((warp.float64(-1)) * (x25))
                        + ((warp.float64(0.200000000000000)) * (x11))
                        + (
                            (warp.float64(-0.0500000000000000))
                            * (u_s_0[i, ((2) + (j)) % 44])
                        )
                    )
                    if (x22 < warp.float64(-0.00000100000000000000))
                    else (
                        (
                            (x25)
                            + (
                                (warp.float64(0.0500000000000000))
                                * (u_s_0[i, ((-2) + (j)) % 44])
                            )
                            + ((warp.float64(-0.200000000000000)) * (x12))
                        )
                        if (x22 > warp.float64(0.00000100000000000000))
                        else (
                            ((warp.float64(0.0500000000000000)) * (x11))
                            + ((warp.float64(-0.0500000000000000)) * (x12))
                        
                )
            )
        )
        # ... 227 more lines: the other five upwind switches, bed friction

)
                    )

This was a small fraction of the code generated from the Sofia code, which defines the full simulation:

G = Grid(ndim=2, boundaries=((GH - 0.5, -(GH - 0.5)), (GH - 0.5, -(GH - 0.5))))
x, y = G.coordinates
zeta = Field("zeta", domain=G)
u = Field("u", domain=G, staggered=Staggered.XFACE)
v = Field("v", domain=G, staggered=Staggered.YFACE)
zb = Field("zb", domain=G)
h = zeta - zb

def momentum(a, dx_):
    return (-u * diff_upwind(a, x, variable=u) - v * diff_upwind(a, y, variable=v)
            - G_ACC * diff(zeta, dx_)
            + NU * laplacian(a)
            - CF * a * sp.sqrt(u**2 + v**2) / sp.Max(h, H_MIN))

F = {zeta: -diff(h * u, x, allow_expand=False) - diff(h * v, y, allow_expand=False),
        u: momentum(u, x),
        v: momentum(v, y)}

start_x, end_x = sp.Eq(x, G.DomainStart), sp.Eq(x, G.DomainEnd)
start_y, end_y = sp.Eq(y, G.DomainStart), sp.Eq(y, G.DomainEnd)
bcs = ir.CodeBlock(
    apply_boundary_conditions(((sp.Eq(diff(zeta, x), 0), start_x),
                                (sp.Eq(zeta, Z_OUT), end_x),
                                (sp.Eq(diff(zeta, y), 0), start_y),
                                (sp.Eq(diff(zeta, y), 0), end_y)),
                                field=zeta, as_eval_assignment=True),
    apply_boundary_conditions(((sp.Eq(u, Q_IN / sp.Max(h, H_MIN)), start_x),
                                (sp.Eq(diff(u, x), 0), end_x),
                                (sp.Eq(diff(u, y), 0), start_y),
                                (sp.Eq(diff(u, y), 0), end_y)),
                                field=u, as_eval_assignment=True),
    apply_boundary_conditions(((sp.Eq(diff(v, x), 0), start_x),
                                (sp.Eq(diff(v, x), 0), end_x),
                                (sp.Eq(v, 0), start_y),
                                (sp.Eq(v, 0), end_y)),
                                field=v, as_eval_assignment=True))

code = ir.CodeBlock(
    ir.Inputs((zeta, u, v, zb)),
    bcs,                                   # fill the ghosts before the first stage
    ir.For(length=nsteps, body=ir.RKn(F=F, order=3, dt=DT, post=bcs)),
    ir.Outputs((zeta, u, v)))
shape = (NX + 2 * GH, NY + 2 * GH)
specs = {G: dict(shape=shape, extent=((0.0, NX * DX), (0.0, NY * DX)))}
sim = build_simulation(code, specs, kernel_backend="warp")