functionalize

functionalize

Compile parameterized SDF trees to pure JAX functions.

The SDF tree is cadjoint’s intermediate representation: primitives, booleans, transforms and patterns, each with named parameter slots. This module lowers it to JAX, and how it does that decides how big the compiled program is.

Traced naively, the tree flattens: a pattern of eight ribs emits eight copies of the rib, a subtree used by two booleans is emitted twice, and every parameter value is folded in as a literal — so a slider edit produces a program that is structurally identical to the one before it and yet compiles from scratch. Three things keep that from happening here.

• Sharing. A node is built once per object, not once per occurrence, so a tool cut from several bodies is one function.
• Outlining. A node that is evaluated more than once — a pattern’s child, a shared subtree — is wrapped in :func:jax.jit, which StableHLO keeps as a func.func and a func.call per instance. The WGSL backend already emits one shader function per func.func, so the same move shrinks the shader.
• Parameters as arguments. :func:functionalize_parametric hands the parameter dicts to the jitted function rather than closing over them, so the lowered HLO is byte-identical across value edits and the persistent compilation cache hits.

Everything is opt-out through :func:cadjoint.sdf._lowering.scalar_lowering, which the shader backend holds while it traces: WGSL has no type wider than a mat4, so the array-shaped forms inside the primitives cannot be lowered there and the flat emission is still the right one.

Functions

Name	Description
functionalize	Compile an SDF to a pure function with free and fixed parameters.
functionalize_parametric	Compile an SDF to a function whose parameters are arguments.
functionalize_scene	Compile a geometry SDF tree to pure (sdf, material_fn) closures.
functionalize_scene_parametric	Compile a scene so its parameters are arguments, not literals.

functionalize

functionalize.functionalize(sdf)

Compile an SDF to a pure function with free and fixed parameters.

Returns a curried function with signature

sdf_fn(free_params, fixed_params) -> (point -> distance)

Nodes are built once per object and repeated nodes are outlined, so a pattern of N instances and a subtree used by two parents each cost one copy in the emitted program rather than N and two.

Parameters

Name	Type	Description	Default
sdf		The SDF to compile	required

Returns

Name	Type	Description
Callable	Callable	Curried function sdf_fn(free_params, fixed_params) -> (point -> distance) mapping parameter dicts to a callable point: Array (3,) -> distance: Array ().

Example

radius = Scalar(value=1.0, free=True, name='radius')
sphere = Sphere(radius=radius)
sdf_fn = functionalize(sphere)
distance = sdf_fn({'sphere_0.radius': 2.0}, {})(jnp.array([0., 0., 0.]))

functionalize_parametric

functionalize.functionalize_parametric(sdf)

Compile an SDF to a function whose parameters are arguments.

:func:functionalize closes over the parameter values, so jax.jit of the result folds every one of them into the program as a literal: two designs that differ only in a slider value lower to two different HLO modules, compile separately, and miss each other in the persistent compilation cache. This form takes them as arrays instead, so the lowered text is byte-identical for every value of every parameter and one compiled executable serves them all.

Parameters

Name	Type	Description	Default
sdf		The SDF to compile.	required

Returns

Name	Type	Description
	Callable	fn(free_params, fixed_params, point) -> distance, safe to wrap in :func:jax.jit directly.

Example

free, fixed, _ = extract_parameters(scene)
evaluate = jax.jit(functionalize_parametric(scene))
distance = evaluate(free, fixed, jnp.zeros(3))

functionalize_scene

functionalize.functionalize_scene(geometry)

Compile a geometry SDF tree to pure (sdf, material_fn) closures.

Returns a curried function::

scene_fn = functionalize_scene(geometry)
sdf, material_fn = scene_fn(free_params, fixed_params)
distance = sdf(point)
mat_dict = material_fn(point)

Both sdf and material_fn are pure JAX functions. This helper is for geometry and material evaluation; the forward image renderer deliberately has a separate, non-differentiable API.

The node counter matches extract_parameters exactly (same DFS order, same counter increments for every Fluent node), so path-keyed fixed params align correctly — including for a node reached twice, which is numbered twice there and built only once here.

Parameters

Name	Type	Description	Default
geometry		Root SDF node of the geometry tree.	required

Returns

Name	Type	Description
	Callable	(free_params, fixed_params) -> (sdf_fn, material_fn)

functionalize_scene_parametric

functionalize.functionalize_scene_parametric(geometry)

Compile a scene so its parameters are arguments, not literals.

The :func:functionalize_parametric of :func:functionalize_scene: one lowered program serves every value of every parameter, which is what makes a shader’s uniform block possible and what makes the on-disk compilation cache hit across edits and processes.

Parameters

Name	Type	Description	Default
geometry		Root SDF node of the geometry tree.	required

Returns

Name	Type	Description
	Callable	fn(free_params, fixed_params, point) -> (distance, material dict).