Flow solver

Conjugate heat transfer with no mesh in the loop — Brinkman-penalised lattice Boltzmann and a single-domain energy solve on the scene’s own field.

cadjoint.flow solves forced convection over a part (air pushed along a duct, heat conducted out of the part and carried away by the air) on a fixed lattice, and differentiates the pressure drop and the temperature with respect to the geometry with no mesh generation in the loop.

Example

import jax
jax.config.update("jax_enable_x64", True)

from cadjoint import extract_parameters, functionalize
from cadjoint.flow import FlowConfig, FlowGrid, SteadyOptions, sample_solid_fraction, solve
from cadjoint.geometry import Vector
from cadjoint.sdf.primitives import Box

fin = Box(Vector([0.3, 0.4, 0.6], free=True, name="size"))
free, fixed, _ = extract_parameters(fin)
evaluate = functionalize(fin)

grid = FlowGrid(shape=(16, 32, 16), origin=(-1.0, -2.0, -1.0), size=(2.0, 4.0, 2.0))
chi = sample_solid_fraction(evaluate(free, fixed), grid)
print(chi.shape, float(chi.max()), grid.suggested_epsilon())   # (16, 32, 16) 1.0 0.0625

config = FlowConfig(
    shape=grid.shape, inlet_speed=0.02, reynolds=30.0,
    steady=SteadyOptions(tol=1e-10, adjoint_solver="fixed_point",
                         adjoint_tol=1e-10, adjoint_max_steps=4000),
)
print(round(config.omega, 4), round(config.mach, 3))   # 1.8797 0.035

result = solve(chi, config, cell_volume=grid.cell_volume)
print(float(result.pressure_drop), float(result.heat_transfer))

Brinkman penalisation

The scene’s SDF is sampled on a fixed lattice and mapped to a solid fraction chi(x) in [0, 1]. chi enters the momentum equation as a Brinkman drag −alpha_max · chi · u, which stops the fluid where the design is solid. The grid never moves, so the chain from a sketch point to a pressure drop is one JAX expression:

sketch point → SDF → chi = profile(−f(x) / ε) → alpha = alpha_max · chi
             → D3Q19 lattice-Boltzmann fixed point → pressure drop, heat proxy

Grid and solid fraction

FlowGrid is the lattice in world coordinates, with the duct along +Y; sample_solid_fraction evaluates a scene SDF on it. The profile that maps distance to occupancy has compact support, because a sigmoid never reaches zero and, multiplied by alpha_max, its tail blocks every open channel. The default is a C² quintic, which also keeps a finite-difference check second order.

Solver settings

FlowConfig holds the lattice-unit settings (inlet speed, Reynolds number, drag at chi = 1) and derives the BGK relaxation rate omega from them. It rejects an omega above 1.95 at construction: BGK loses stability as omega approaches 2, and the penalised solid lowers that limit.

FlowResult carries the converged populations, density and velocity fields and two scalars: the inlet-to-outlet pressure_drop, and heat_transfer, an ∫ chi |u| proxy for a cooling number without an energy solve. FlowStudy (below) reports a temperature and is the recommended objective. convergence returns the residual history.

Gradients

Nothing in the chain contours, meshes or decides membership, so jax.grad runs through the SDF, the profile and the fixed point:

def objective(free):
    chi = sample_solid_fraction(evaluate(free, fixed), grid)
    return solve(chi, config, cell_volume=grid.cell_volume).pressure_drop

value, grad = jax.value_and_grad(objective)(free)
print(float(value), grad["size"])
# 0.0039086 [6.884e-04 2.700e-03 8.27e-07]

Widening the fin across the duct (size[0]) raises the pressure drop at 6.9e-4 per unit; lengthening it along the flow (size[1]) costs about four times as much, and its height (size[2]) almost nothing, because the box already spans most of the duct’s height. A central difference on size[0] with h = 1e-3 gives 6.897e-4, 0.2 % from the adjoint, which is the difference scheme’s truncation. On the starter sink’s sketch handles the research note measures 3e-8.

The gradient comes from the implicit function theorem rather than from taping the pseudo-time march: one adjoint solve at the converged state, with memory independent of the number of steps. On the starter sink that is 457 MB against the tape’s 21 GB, and the two agree to 5.6e-11 once the tape’s march has converged.

adjoint_solver Behaviour
"gmres" fastest when it converges; holds restart Krylov vectors
"fixed_point" Richardson iteration on the same system; converges exactly when the forward march does; three states of memory

Disagreement between the two means one solve had not converged. The default adjoint_max_steps of 400 is too few on this duct for either solver (with it, the example reports 8.1e-4 rather than 6.9e-4), so the configuration above raises it. Check a new configuration against a finite difference once.

Conjugate heat transfer

A FlowStudy declares the coupled problem the way a ThermalStudy declares conduction (in the scene, captured by capture_studies, with describe() and its own boundary conditions) and solves both the flow and the temperature. On the fin box from above:

from cadjoint.studies import Nodes
from cadjoint.flow import FlowStudy, HeatSource, Inlet, Outlet, Walls

cooling = FlowStudy(
    name="duct-cooling",
    resolution=(20, 26, 12),        # cubic cells: 1.40/20 = 1.82/26 = 0.84/12
    bounds=(-0.70, -0.91, -0.50), size=(1.40, 1.82, 0.84),
    reynolds=25.0, conductivity_ratio=200.0,
    bcs=[
        Inlet(velocity=0.02, temperature=0.0),
        Outlet(),
        Walls(),                       # no-slip always; adiabatic by default
        HeatSource(Nodes.box([-0.14, -0.18, -0.42], [0.14, 0.18, -0.18]), power=1.0),
    ],
)
result = cooling.solve(evaluate(free, fixed))
print(float(result.peak_temperature), float(result.pressure_drop))

scenes/duct_sink.py declares the same study around a three-fin sink. python scenes/duct_sink.py runs in well under a minute and prints:

peak temperature     0.42973      thermal resistance   0.42973
mean temperature     0.41265      pressure drop        0.010414
outlet bulk air      0.27618      energy imbalance    -2.566e-11

The cells must be cubic. The solve is entirely in lattice units (streaming moves one cell per axis, the energy stencil is one cell wide), so size only decides where the SDF is sampled, and a lattice with unequal spacings solves the duct stretched by their ratio.

A scene declaring a flow study opens, renders and compiles in the playground, and the Studies window reads and edits it like any other study.

Energy equation

The energy equation is solved once over the whole lattice, with the conductivity interpolated by chi between the air’s and the solid’s:

div(k grad T)  =  (rho cp)_f  u · grad T  −  q

Inside the metal u is zero (removed by the Brinkman drag), so the equation reduces to conduction; in the channels k is the air’s and advection dominates. No interface condition is imposed: one unknown field admits no temperature jump, and a conservative finite-volume flux is shared by the cells on either side of a face, so temperature and heat flux are continuous by construction, with no interface to locate or mesh.

The coupling is one-way, which is exact in this model: the momentum solve takes chi, an inlet velocity and a viscosity, none of which depends on temperature. Two-way coupling would add buoyancy, which is negligible at small Richardson number; result.richardson reports it (2.1e-4 on scenes/duct_sink.py).

Volumetric selections

HeatSource and HeldTemperature take the same Nodes.box / sphere / halfspace / cylinder selections as a mesh study, composed with &, | and ~. A lattice has no boundary surface, so a selection picks every cell centre satisfying the criterion rather than boundary nodes. Nodes.side and Nodes.predicate are rejected with the alternative named. A region thinner than a cell contains no centres and is rejected as well.

Warnings

result.warnings() reports:

  • a lattice whose cells are not cubes, which solves a stretched duct;
  • a Richardson number above 0.1;
  • a Reynolds number past the turbulent transition;
  • an under-resolved thermal boundary layer;
  • an energy balance that does not close.

Gradient through both solves

d(peak temperature)/d(geometry) runs through the momentum fixed point’s adjoint and the energy solve’s transposed linear solve, which are independent:

from cadjoint.flow.precision import double_precision

def objective(parameters):
    return cooling.solve(evaluate(parameters, fixed)).peak_temperature

with double_precision():
    gradient = jax.grad(objective)(free)

Against a central difference on a box in a duct, the difference converges to the adjoint at second order, down to 4.6e-8 relative. On a new configuration, check the convergence rate, not one step size: where a perturbation puts the surface exactly on a plane of cell centres, the objective has a kink and the difference converges only at first order. It still converges to the adjoint, but a single ratio at a single h cannot distinguish that from a wrong gradient.

Precision

The solve needs float64 and enables it for the duration of the call. A scene therefore sets no jax flags: setting jax_enable_x64 at module scope would make the whole process float64, and the WGSL backend has no 64-bit type, so the scene would not open in the viewer.

The scope covers a forward solve but not a gradient, because jax.grad runs its backward pass after the call returns. Differentiating through a study needs x64 around the gradient, as with the FEM solvers: double_precision() above, or precision="double" on an Optimization. Without it the backward pass fails with lax.dynamic_update_slice requires arguments to have the same dtypes, got float32, float64.

Verification

Check Result
pure conduction vs. the analytic parabola 3.9e-3 → 6.1e-5 over 8 → 64 cells, exactly 2nd order
the same problem on the FEM ThermalStudy agrees to 2.4e-4 of the peak at 32 cells (all of it the lattice’s truncation; the FEM side is exact to 1e-14)
metal–air interface vs. a two-layer series resistance 1.4e-13 at a conductivity ratio of 8000
advection–diffusion vs. the analytic column 3.9e-16 (exponential scheme); upwind is off by 0.198
square-duct Nusselt number vs. the textbook 2.976 2.927 at 22 cells across (−1.65%), 2.857 at 14 (−4.0%)
global energy balance 2.6e-11 of the injected power

Shape optimization

scenes/duct_fairing.py optimizes one blunt body immersed in a duct, with the pressure drop of the converged flow as the objective. Run from the repository root:

from cadjoint.optimize import Optimization
from scenes.duct_fairing import fairing, streamlining_cost

streamline = Optimization(
    name="streamline",
    objective=streamlining_cost,   # dp / dp_ref + 6 (V/V_ref - 1)^2
    of=fairing,                    # the box's free half-extents
    steps=16, learning_rate=0.005, method="adam",
    precision="double",
)

The volume term is required because pressure drop alone is minimized by removing the obstacle. In a duct the pressure drop is dominated by blockage, so at constant volume the descent trades frontal area for length. At the starting box d(dp)/d(half-width) = +5.709e-2 across the flow against +1.242e-2 along it.

Sixteen steps take about two and a half minutes on a laptop CPU:

aspect (y/x)  0.750 -> 1.011      frontal area  0.16000 -> 0.13277
volume        0.04800 -> 0.04892  pressure drop 4.597e-3 -> 2.684e-3

The volume drifts 1.9%, so the scene also solves the start box scaled isotropically to the final volume. It costs 4.706e-3 against 2.684e-3 for the optimized shape at the same volume: a 43% reduction from proportions alone.

main() verifies the gradient on every run: it asserts that the relative error against a central difference falls by about 100× when the step size falls by 10×, the second-order convergence that the "smootherstep" solid-fraction profile preserves.

A heat-sink objective needs both temperature and pressure drop: on scenes/duct_sink.py the sink costs 11.7× the empty duct’s pressure drop while holding the die 0.430 above inlet air, and thickening a fin moves the two in opposite directions. Without the pressure term an optimizer converges to a solid block. One FlowStudy reports both.

Plugin

The solver ships as flow_brinkman, a plugin of kind flow_solver. The other solver plugins take a mesh; this one takes one array:

from cadjoint.plugins import plugin_for_kind

flow = plugin_for_kind("flow_solver")
print(flow.name, sorted(flow.capabilities.differentiable_inputs))
# flow_brinkman ['chi', 'inlet_velocity']

apply and vector_jacobian_product through the runtime are bit-identical to the in-process solve and gradient, because both sides run the same custom_vjp. It needs no conda and no extra: cadjoint.flow is pure JAX.

Limitations

  • The model is laminar. A forced-convection heat sink often runs above the duct transition at Re ≈ 2300, where this solver reports the laminar answer and so overstates the temperature.
  • It runs on CPU at 12–16 million lattice updates a second, about 2 % of the memory roofline, because stream() is nineteen jnp.rolls where a fused kernel makes one pass. Iterations to convergence grow linearly with the grid’s linear size (about 119·N + 1400), so a 128³ solve takes 48 minutes on this CPU and four minutes on an A100 as written.
  • A Nusselt number within a couple of percent needs about 22 cells across a channel; the examples here use 12 to 16.
  • Adjoint memory: GMRES at 128³ holds 12.8 GB of Krylov vectors.
  • The design space is a few primitive dimensions.
  • Read §8.9 of the research note before quoting an absolute number; it lists which results speed, turbulence and a wider design space would change.

API