Node Authoring Guide#

This guide covers everything required to add a new SimulationNode to MADDENING. Every requirement here is derived from the Documentation Architecture and enforced by CI.

The Contract#

A MADDENING node is a pure function wrapped in a descriptor class:

  • initial_state(params) -> dict — returns the initial state arrays

  • update(state, boundary_inputs, dt) -> new_state — returns a new state dict

  • update() must be JAX-traceable: no Python-level side effects, no data-dependent control flow, no print statements. Use jnp.where instead of if/else.

  • State is immutable — return a new dict, don’t mutate in place

  • Parameters live in self.params, not in state

Directory Structure#

src/maddening/nodes/your_node.py        # Node implementation
tests/nodes/test_your_node.py           # Unit tests
tests/verification/test_your_node_*.py  # Verification benchmark(s)
docs/algorithm_guide/nodes/your_node.md # Algorithm documentation

Step-by-Step#

1. Implement the Node#

"""YourNode -- one-line description."""

import jax.numpy as jnp

from maddening.core.node import SimulationNode
from maddening.core.compliance.metadata import (
    NodeMeta, StabilityLevel, ValidatedRegime, Reference,
)
from maddening.core.compliance.stability import stability


@stability(StabilityLevel.EXPERIMENTAL)
class YourNode(SimulationNode):
    """NumPy-style docstring.

    Parameters
    ----------
    name : str
        Unique node name.
    timestep : float
        Simulation timestep in seconds.
    ...

    Boundary inputs
    ---------------
    ...
    """

    meta = NodeMeta(
        algorithm_id="MADD-NODE-XXX",      # Get next available ID
        algorithm_version="1.0.0",
        stability=StabilityLevel.EXPERIMENTAL,
        description="One-line description",
        governing_equations=r"...",          # LaTeX
        discretization="...",
        assumptions=(
            "...",
        ),
        limitations=(
            "...",
        ),
        validated_regimes=(
            ValidatedRegime("param", min_val, max_val, "units"),
        ),
        references=(
            Reference("AuthorYear", "Description"),
        ),
        hazard_hints=(
            "...",
        ),
    )

    def __init__(self, name, timestep, **kwargs):
        params = {"key": value, ...}
        state_spec = {"field": shape_tuple, ...}
        super().__init__(name, timestep, params, state_spec)

    def initial_state(self):
        return {"field": jnp.zeros(self.state_spec["field"])}

    def update(self, state, boundary_inputs, dt):
        # Pure JAX operations only
        ...
        return new_state

2. Required Metadata (NodeMeta)#

Every node must have a meta ClassVar. Required fields:

Field

Required

Description

algorithm_id

Yes

MADD-NODE-XXX format (get next available from existing nodes)

algorithm_version

Yes

Semantic version of this algorithm implementation

stability

Yes

StabilityLevel.EXPERIMENTAL for new nodes

description

Yes

One-line description

assumptions

Yes

Tuple of strings — every physical/mathematical assumption

limitations

Yes

Tuple of strings — every known failure mode

hazard_hints

Yes

Tuple of strings — qualitative risks for ISO 14971 input

Recommended fields:

Field

Description

governing_equations

LaTeX string of the governing equations

discretization

Description of the numerical method

validated_regimes

Tuple of ValidatedRegime — quantitative parameter bounds

references

Tuple of Reference — BibTeX keys from docs/bibliography.bib

implementation_map

Dict mapping equation terms to Python qualified names

Scope distinction: validated_regimes is for quantitative parameter-bound risks (“CFL must be < 0.5”). hazard_hints is for qualitative non-parameter-bound risks (“wall bounce-back assumes rigid walls”). A given risk goes in exactly one, never both.

3. Write the Algorithm Guide#

Copy docs/algorithm_guide/nodes/_template.md and fill in every section:

docs/algorithm_guide/nodes/your_node.md

Mandatory sections: Summary, Governing Equations, Discretization, Implementation Mapping, Assumptions and Simplifications, Validated Physical Regimes, Known Limitations and Failure Modes, Stability Conditions, State Variables, Parameters, Boundary Inputs, References, Verification Evidence, Changelog.

Implementation Mapping: Trace every equation term to a specific Python/JAX function. No silent omissions. Terms handled by JAX primitives (e.g., jnp.fft.rfftn()) must be documented as such.

References: Use Pandoc-style [@Key] citations where Key matches an entry in docs/bibliography.bib. Include YAML frontmatter:

---
bibliography: ../../bibliography.bib
---

Each reference also gets a human-readable inline description for GitHub/VS Code readability.

4. Add Bibliography Entries#

Add BibTeX entries for any cited references to docs/bibliography.bib:

@article{AuthorYear,
  author  = {Last, First and Last2, First2},
  title   = {Title of the Paper},
  journal = {Journal Name},
  year    = {2024},
  volume  = {1},
  pages   = {1--10},
  doi     = {10.xxxx/yyyy},
}

CI validates that every [@Key] citation in algorithm guides resolves to an entry in the bib file (scripts/check_citations.py).

5. Write Tests#

Unit tests (tests/nodes/test_your_node.py) — mandatory:

  • Test each public method

  • Test normal operation and edge cases

  • Verify JAX-traceability: jax.jit(node.update)(state, {}, dt) must work

  • Verify jax.grad compatibility if the node supports differentiation

  • Verify jax.vmap compatibility if the node supports batching

Verification benchmark (tests/verification/test_your_node_*.py) — mandatory for physics nodes:

  • Compare against analytical solution or published reference data

  • Register with @verification_benchmark:

from maddening.core.validation import verification_benchmark

@verification_benchmark(
    benchmark_id="MADD-VER-XXX",
    description="Your benchmark description",
    node_class="YourNode",
    reference="AuthorYear",
)
def test_your_analytical_comparison():
    ...

Integration test — mandatory:

  • Test the node within a GraphManager (add node, connect edges, run steps)

6. Apply the @stability Decorator#

from maddening.core.compliance.stability import stability

@stability(StabilityLevel.EXPERIMENTAL)
class YourNode(SimulationNode):
    ...

New nodes start as EXPERIMENTAL. Promote to STABLE when:

  • At least one verification benchmark passes

  • Algorithm guide is complete

  • API has been stable for at least one minor release

7. Document Known Limitations#

If your node has known failure modes or limitations, add entries to docs/validation/known_anomalies.yaml:

- anomaly_id: "MADD-ANO-XXX"
  title: "YourNode: brief description of limitation"
  description: "Full description..."
  severity: "major"            # critical | major | minor
  safety_relevance: "context_dependent"
  safety_relevance_rationale: "..."
  affected_components: ["YourNode"]
  affected_versions: ["0.2.0"]
  status: "open"
  workaround: "..."

Run python -m maddening.compliance check-anomalies docs/validation/known_anomalies.yaml to validate.

8. Run CI Checks#

Before committing, verify everything passes:

# Activate the virtual environment
source ../venvs/.maddening/bin/activate

# Run the full test suite
PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/ -v --tb=short --ignore=tests/viz

# Run compliance checks
python scripts/check_anomalies.py
python scripts/check_impl_mapping.py
python scripts/check_citations.py

Boundary Input Specification#

Override boundary_input_spec() to declare what boundary inputs your node expects. This enables validation, documentation generation, and correct initialization for additive inputs.

from maddening.core.node import BoundaryInputSpec

def boundary_input_spec(self):
    return {
        "left_temperature": BoundaryInputSpec(
            shape=(), description="Dirichlet BC at left end",
        ),
        "heat_source": BoundaryInputSpec(
            shape=(self.params["n_cells"],),
            description="Volumetric heat source",
            coupling_type="additive",  # multiple edges sum
        ),
    }

Each entry maps a boundary input name to a BoundaryInputSpec with:

  • shape: array shape (empty tuple for scalar)

  • dtype: JAX dtype (default jnp.float32)

  • default: default value if not supplied

  • coupling_type: "replacive" (last edge wins, default) or "additive" (edges sum)

  • description: human-readable description

Exposing Flux Quantities#

Override compute_boundary_fluxes() to expose derived quantities (forces, heat fluxes) that other nodes can consume via edges. Flux fields are NOT part of state — they are computed on-the-fly.

def compute_boundary_fluxes(self, state, boundary_inputs, dt):
    T = state["temperature"]
    dx = self.params["length"] / self.params["n_cells"]
    alpha = self.params["thermal_diffusivity"]
    return {
        "left_heat_flux": -alpha * (T[1] - T[0]) / dx,
        "right_heat_flux": -alpha * (T[-1] - T[-2]) / dx,
    }

Requirements:

  • Must be a pure JAX-traceable function (same rules as update)

  • Return a dict of JAX arrays

  • Keys become available as source_field on edges

  • Called automatically after each node update during edge resolution

Additive vs Replacive Inputs#

By default, if multiple edges write to the same boundary input, the last one wins (“replacive”). For inputs that should accumulate (e.g., forces from multiple sources), mark them as "additive":

# In boundary_input_spec():
"force": BoundaryInputSpec(shape=(2,), coupling_type="additive")

# When adding edges:
gm.add_edge("spring1", "body", "spring_force", "force", additive=True)
gm.add_edge("spring2", "body", "spring_force", "force", additive=True)
# body receives the SUM of both spring forces

The first additive edge sets the initial value; subsequent additive edges accumulate via addition.

Coupling Patterns#

Value coupling (most common)#

One node’s state field feeds another’s boundary input:

from maddening.core.coupling_helpers import add_value_coupling
add_value_coupling(gm, "ball", "spring", "position", "anchor_position")

Flux coupling#

One node’s flux output feeds another’s boundary input:

from maddening.core.coupling_helpers import add_flux_coupling
add_flux_coupling(gm, "rod_a", "rod_b", "right_heat_flux", "heat_source")

Dirichlet-Neumann coupling#

The classic partitioned approach — one node gets a value BC, the other gets a flux BC:

from maddening.core.coupling_helpers import add_dirichlet_neumann_pair
add_dirichlet_neumann_pair(
    gm,
    dirichlet_node="rod_a",  # receives temperature (value)
    neumann_node="rod_b",    # receives heat flux
    value_field="temperature",
    flux_field="right_heat_flux",
    value_input="right_temperature",
    flux_input="heat_source",
    value_transform=lambda T: T[0],
)

Robin coupling#

Combines value and flux for better convergence:

from maddening.core.coupling_helpers import add_robin_coupling
add_robin_coupling(
    gm, "rod_a", "rod_b",
    value_field_a="temperature", flux_field_a="right_heat_flux",
    value_field_b="temperature", flux_field_b="left_heat_flux",
    input_a="right_temperature", input_b="left_temperature",
    alpha=0.5,  # mixing: 0=pure Neumann, 1=pure Dirichlet
)

Transform Registration#

Edge transforms are Python callables applied to data as it flows along edges. For local development, inline lambdas work fine. For USD serialization (and eventually IEC 62304 traceability), transforms must be registered with a unique name.

Registering a transform#

from maddening.core.transforms import register_transform

@register_transform("extract_right_boundary")
def extract_right_boundary(T):
    """Extract the rightmost cell of a temperature array."""
    return T[-1]

Once registered, use either the callable or its string name in edges:

gm.add_edge("rod_a", "rod_b", "temperature", "left_temperature",
            transform="extract_right_boundary")
# OR equivalently:
gm.add_edge("rod_a", "rod_b", "temperature", "left_temperature",
            transform=extract_right_boundary)

When to register#

Scenario

Registration required?

Local development / testing

No (lambdas OK)

USD graph serialization

Yes (unregistered raises UnregisteredTransformError)

Production simulation scenarios

Recommended (enables traceability)

Built-in transforms#

These are pre-registered and always available:

Name

Description

extract_first

arr[0]

extract_last

arr[-1]

extract_second

arr[1]

extract_second_last

arr[-2]

negate

-x

identity

pass-through

scale(factor)

factor * x (auto-registered as "scale_{factor}")

CI validation#

The scripts/check_transforms.py CI script scans production code for string transform references and verifies they resolve in the registry. Run it alongside the other compliance checks:

python scripts/check_transforms.py

New Node Checklist#

  • [ ] SimulationNode subclass with initial_state() and update()

  • [ ] update() is JAX-traceable (jit, grad, vmap compatible)

  • [ ] boundary_input_spec() overridden (declares expected boundary inputs)

  • [ ] compute_boundary_fluxes() overridden if node exposes flux quantities

  • [ ] Additive inputs marked with coupling_type="additive" in spec

  • [ ] @stability(StabilityLevel.EXPERIMENTAL) decorator applied

  • [ ] NodeMeta metadata attached (algorithm ID, stability, assumptions, limitations, hazard_hints)

  • [ ] NumPy-style docstring with Parameters, Boundary inputs

  • [ ] Algorithm guide document in docs/algorithm_guide/nodes/ following _template.md

  • [ ] All template sections filled in (no empty tables or placeholder text)

  • [ ] Implementation Mapping traces every equation term to code

  • [ ] [@Key] citations reference entries in docs/bibliography.bib

  • [ ] YAML frontmatter with bibliography: ../../bibliography.bib

  • [ ] Assumptions and simplifications listed

  • [ ] Validated physical regimes documented

  • [ ] Known limitations and failure modes documented

  • [ ] At least one registered @verification_benchmark

  • [ ] Unit tests covering normal operation, edge cases, JAX-traceability

  • [ ] Integration test within a GraphManager

  • [ ] Known limitations entered in docs/validation/known_anomalies.yaml

  • [ ] Entry in docs/bibliography.bib for primary reference

  • [ ] Edge transforms registered via @register_transform if used in production scenarios

  • [ ] All CI checks pass: check_anomalies.py, check_impl_mapping.py, check_citations.py, check_transforms.py