Skip to content

🧪 3. Test a Node

In a run started with mujou.run, dt follows the real time that passes, and the order in which the nodes update differs every time, so the results differ slightly from run to run. This page tests the node of 2. Write a Node with mujou.testing, which runs nodes with the same result every time.

The tools

Import it with import mujou.testing; import mujou alone does not.

Name Role
mujou.testing.run(root, ...) Runs the tree with the timer, inputs and state of mujou.run (no workers) and returns a ManualRun. No node updates until step updates it. It runs on a new mujou.time.VirtualTimer by default
ManualRun.step(dt, *nodes) Moves the timer dt seconds forward, then updates the given nodes once each, in the given order
ManualRun.timer The run's VirtualTimer

ManualRun is a Run, so write, read, state_dict and the rest work as with mujou.run.

A test

class TestLeaky:
    def test_follows_the_exact_solution(self) -> None:
        tau, h, steps = 0.1, 0.001, 100
        node = Leaky(tau)
        one = torch.tensor(1.0, dtype=torch.float64)
        expected = 0.0
        with mujou.testing.run(node, inputs={node.prev: one}) as run:
            for _ in range(steps):
                run.step(h)  # 1 ms later, then every node once (period is ignored)
                expected += (1 - math.exp(-h / tau)) * (1.0 - expected)
            y, t = run.read_with_time(node.next)
        assert t is not None and t.ns == steps * 1_000_000  # sent at the last step
        torch.testing.assert_close(y, torch.tensor(expected, dtype=torch.float64))

uv run pytest examples/test_leaky.py
  1. mujou.testing.run(node, inputs=...) runs the node on a new VirtualTimer at time 0. inputs are the values of the input ports from the start
  2. Each run.step(h) moves the time h seconds forward and updates every node once, so the dt of the node is exactly h
  3. The reference is the exact solution, computed with the same coefficient. Compare it with torch.testing.assert_close, or with torch.equal when the reference does the same operations in the same order (below)
  4. The time the node sent the output is steps milliseconds: no real time enters a manual run

The order of the updates

class TestLeakyChain:
    def test_follows_the_discrete_algorithm(self) -> None:
        tau, h, steps = 0.01, 0.001, 50
        u = torch.tensor(1.0, dtype=torch.float64)
        # The reference: the discrete algorithm, a then b at every step.
        a, b = SCALAR.zeros(), SCALAR.zeros()
        for _ in range(steps):
            a += (1 - math.exp(-h / tau)) * (u - a)
            b += (1 - math.exp(-h / tau)) * (a - b)  # b reads the new a

        net = LeakyChain(tau)
        with mujou.testing.run(net, inputs={net.prev: u}) as run:
            for _ in range(steps):
                run.step(h, net.a, net.b)  # a, then b, which reads a's new value
            y = run.read(net.next)
        assert torch.equal(y, b)  # the same operations: equal to the last bit
  • step(dt) without nodes updates every node once, in the order of the tree (the order the attributes were assigned)
  • The order changes the result. With step(h, net.a, net.b), b reads the value a sent in this step; with step(h, net.b, net.a), the one from the step before
  • A node's dt is the time since its last update, the first from the start of the run. A node passed twice in one step updates once: the second time is skipped, since its time has not moved
  • The period of a node is not used in a manual run
  • A node that raises makes step raise NetworkError at once

Save and restore

A run that takes the state midway and goes on in a new instance matches a run that went on without stopping, to the last bit:

def run_steps(n: int, state: dict[str, Any] | None = None) -> dict[str, Any]:
    """Run a new ``LeakyChain`` for ``n`` steps of 1 ms and return its state."""
    net = LeakyChain(0.01)
    u = torch.tensor(1.0, dtype=torch.float64)
    with mujou.testing.run(net, inputs={net.prev: u}, state=state) as run:
        for _ in range(n):
            run.step(0.001)
    return run.state_dict()


class TestRestore:
    def test_a_restored_run_goes_on_as_one_that_did_not_stop(self) -> None:
        straight = run_steps(20)
        # Ten steps, then ten more in a new instance restored from the first half,
        # on a VirtualTimer of its own that starts at the time of the state.
        resumed = run_steps(10, state=run_steps(10))
        x = resumed["nodes"]["b"]["state"]["x"]
        assert torch.equal(x, straight["nodes"]["b"]["state"]["x"])

A run restored from a state runs on a VirtualTimer of its own, which starts at the time of the state.

Exceptions

When a node raises, step raises NetworkError, whose __cause__ is the node's exception. Every later call of the run raises it again, including the shutdown at the end of the with block, so put pytest.raises around the whole block, in the same with statement:

with pytest.raises(NetworkError) as info, mujou.testing.run(net) as run:
    run.step(0.001)  # a node raises ValueError here
assert isinstance(info.value.__cause__, ValueError)

Other conditions for determinism

  • Give a seed to the nodes that draw random numbers, and draw them from self.generator (see Threads and Devices)
  • Move the time only with step. Nothing is sent unless step is called, so wait_for and subscribe without a timeout wait forever on the thread that steps
  • To put several runs on one time axis, pass the same mujou.time.VirtualTimer() to each: a step of one moves the others too

Next