Numerical APIs¶

fno.Math provides backend-neutral array algebra, while fno.Field describes OpenFOAM fields and finite-volume operations. NumPy arrays remain in NumPy, JAX arrays and tracers remain in JAX, and symbolic field expressions are evaluated by OpenFOAM at the declared solver boundary.

Execution lanes¶

Yes means that a fully symbolic function executes inside OpenFOAM. Resident means that the operation executes in ModelHost for NumPy or JAX inputs. A mixed function uses the resident lane.

Math reference¶

All Math functions return the native result type of their operands: NumPy input produces NumPy output, JAX input remains in JAX, and a symbolic FieldExpression produces another expression. Arguments named axis, axes, shape and subscripts follow NumPy conventions on the resident path.

fno.Math.outer(a, b) maps two vector fields to a tensor field, and fno.Math.trace(A) maps a tensor field to a scalar field. fno.Context.delta_time() reads the active OpenFOAM time-step size at the invocation boundary.

Native graphs map matmul(A, B), einsum("...ii->...", A) and einsum("...ij,...ij->...", A, B) to OpenFOAM contractions. Positive dyadic powers, such as x**1.5, and scalar where(x > value, a, b) also remain native.

A dyadic exponent is a positive integer divided by a power of two. Other powers, einsum patterns and arbitrary Python control flow use the resident path.

Mathematical operators¶

The same mathematical function can operate on ordinary arrays for local checks or on symbolic FieldExpression values during closure declaration.

# Define the Smagorinsky function

def smagorinsky_function(velocity_grad, delta, C_k=0.094, C_e=1.048):
    strain_rate = fno.Math.symm(velocity_grad)
    strain_rate_trace = fno.Math.trace(strain_rate)
    dev_strain_rate = fno.Math.dev(strain_rate)

    coefficient_a = C_e / delta
    coefficient_b = (2.0 / 3.0) * strain_rate_trace
    coefficient_c = (
        2.0
        * C_k
        * delta
        * fno.Math.ddot(dev_strain_rate, strain_rate)
    )

    discriminant = coefficient_b**2 + 4.0 * coefficient_a * coefficient_c
    sqrt_k = (
        -coefficient_b
        + fno.Math.sqrt(fno.Math.maximum(discriminant, 0.0))
    ) / (2.0 * coefficient_a)
    sqrt_k = fno.Math.maximum(sqrt_k, 0.0)

    return C_k * delta * sqrt_k

Numerical safeguards belong to the function or trained model that defines the closure. FoamNordic does not silently clip a resident model output after inference.

OpenFOAM field expressions¶

Stored fields and derived fields use the same input declaration.

# Define OpenFOAM closure inputs

inputs = {
    "velocity_grad": fno.Field.grad("U"),
    "convection": fno.Field.div("phi", "U"),
    "thermal_diffusion": fno.Field.laplacian("alphaEff", "T"),
    "vorticity": fno.Field.curl("U"),
    "delta": fno.Field.delta(),
}

Field expressions compose without manually constructing OpenFOAM expression strings.

# Define the strain-rate magnitude

velocity_grad = fno.Field.grad("U")
strain = fno.Math.symm(velocity_grad)
strain_rate = fno.Math.sqrt(2.0 * fno.Math.ddot(strain, strain))

Field expressions support matching-field addition and subtraction. Multiplication accepts a scalar field or finite dimensionless scalar on either side; division accepts a scalar field or finite dimensionless scalar as its denominator. sqrt requires a scalar field. Field.laplacian(diffusivity, field) accepts a scalar field expression as its diffusivity.

Field and Context reference¶

Field.fields() is output metadata rather than a finite-volume expression. It expands against the solver object registry and fails if the pattern matches no fields. Patch expressions identify boundary data explicitly and remain subject to the selected adapter contract.

fno.Field.transport(value, *, diffusivity, source, implicit_sink,
                    lower_bound=0.0)

Finite-volume operators use the copied system/fvSchemes. Exact schemes and usable section defaults are preserved. Convenient mode adds a fallback only when both are absent; strict mode rejects the missing scheme.

Fallbacks are Gauss linear for gradient, curl and divergence, and Gauss linear corrected for Laplacians. The source case stays unchanged.

Native algebra¶

A symbolic Operator.function can run directly in OpenFOAM when every operation is supported and it uses complete fields. Functions that use patches or unsupported Python operations run through ModelHost instead. Native closures and transforms support Observe and profile_every without exchanging fields.

# Define the Smagorinsky closure

closure = fno.Closure(
    name="nutFjord",
    operator=fno.Operator.function(smagorinsky_function),
    inputs={
        "velocity_grad": fno.Field.grad("U"),
        "delta": fno.Field.delta(),
    },
    outputs={
        "nut": fno.Field("nut"),
    },
)

Functions containing Python control flow, NumPy, JAX or model evaluation remain on the resident path. Their OpenFOAM input expressions are still evaluated before packing, so derived fields do not move to Python merely because the model is resident.

Native transport¶

Field.transport() marks one implicit scalar finite-volume solve inside a mathematical Operator.function. The graph is traced during declaration and executed by kEqnFjord without a ModelHost in the corrector loop.

# Define the k-equation function

def keqn_function(k, velocity_grad, nu, delta, C_k=0.094, C_e=1.048):
    k_positive = fno.Math.maximum(k, 1.0e-10)
    sqrt_k = fno.Math.sqrt(k_positive)
    strain = fno.Math.dev(2.0 * fno.Math.symm(velocity_grad))

    nut = C_k * delta * sqrt_k
    production = nut * fno.Math.ddot(velocity_grad, strain)
    dissipation = C_e * sqrt_k / delta

    k_new = fno.Field.transport(
        value=k,
        diffusivity=nu + nut,
        source=production,
        implicit_sink=dissipation,
        lower_bound=1.0e-10,
    )

    return {
        "k": k_new,
        "nut": C_k * delta * fno.Math.sqrt(k_new),
    }


# Define the k-equation closure

closure = fno.Closure(
    name="kEqnFjord",
    operator=fno.Operator.function(keqn_function),
    inputs={
        "k": fno.Field("k"),
        "velocity_grad": fno.Field.grad("U"),
        "nu": fno.Field("nu"),
        "delta": fno.Field.delta(),
    },
    outputs={
        "k": fno.Field("k"),
        "nut": fno.Field("nut"),
    },
)

The current native transport contract permits one scalar solve and fixes the kEqnFjord outputs as k and nut. Use the value returned by Field.transport() for post-solve calculations; mixing it with the original k in one output is rejected. Observe samples the requested OpenFOAM fields after completed corrections, and profile_every measures native algebra, matrix assembly, solving and field publication.

Random keys¶

FoamNordic uses immutable, backend-neutral keys for stochastic closures and transforms.

key = fno.Random.key(42, scope="global")


def perturb(velocity, *, key):
    scale_key, noise_key = fno.Random.split(key, 2)
    scale = fno.Random.uniform(scale_key, low=0.995, high=1.005)
    noise = fno.Random.normal(noise_key, shape=velocity.shape, std=1.0e-6)
    return {"velocity": velocity * scale + noise}

Reusing a key repeats a draw; split keys for independent draws. Distributions are uniform, normal, integers and bernoulli.

scope="global" shares an invocation key across ranks; scope="rank" folds in the negotiated rank. to_jax() creates a JAX key. Reproducibility holds within each backend, not bit-for-bit across NumPy and JAX.