Spaces:
Paused
Paused
File size: 3,360 Bytes
5e0b58b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """Runtime checkpoint sidecar tests."""
from __future__ import annotations
import random
from dataclasses import dataclass
from pathlib import Path
from src.storage.checkpoint import (
capture_runtime_checkpoint,
read_runtime_checkpoint,
write_runtime_checkpoint,
)
@dataclass(slots=True)
class Event:
source_id: int
target_id: int
weight: float
delivery_tick: int
@dataclass(slots=True)
class Neuron:
a: float
b: float
c: float
d: float
v: float
u: float
energy: float
spike_cost: float
threshold_adaptation: float
spike_counter: int
last_spike_tick: int
last_external_current: float
last_synaptic_current: float
firing_rate_estimate: float = 0.0
_spike_count_window: int = 0
_last_update_tick: int = 0
pre_trace: float = 0.0
post_trace: float = 0.0
@dataclass(slots=True)
class Synapse:
target_id: int
weight: float
delay: int
eligibility: float
last_pre_spike: int
@dataclass(slots=True)
class Network:
rng: random.Random
current_tick: int
total_spikes: int
total_events_processed: int
pending_currents: dict[int, float]
input_cells: set[int]
output_cells: set[int]
event_slots: list[list[Event]]
neurons: dict[int, Neuron]
synapses: dict[int, list[Synapse]]
def test_runtime_checkpoint_roundtrip(tmp_path: Path) -> None:
"""Exact neuron, synapse, RNG and queue state survives JSON roundtrip."""
network = Network(
rng=random.Random(42),
current_tick=10,
total_spikes=7,
total_events_processed=5,
pending_currents={1: 3.5},
input_cells={1},
output_cells={2},
event_slots=[[Event(1, 2, 0.4000000000000001, 11)], []],
neurons={
1: Neuron(
a=0.020000000000000004,
b=0.20000000000000004,
c=-65.0,
d=8.000000000000002,
v=-64.123456789,
u=-12.987654321,
energy=0.9987654321,
spike_cost=0.0010000000000000002,
threshold_adaptation=0.00123456789,
spike_counter=3,
last_spike_tick=9,
last_external_current=1.25,
last_synaptic_current=0.75,
)
},
synapses={
1: [
Synapse(
target_id=2,
weight=0.3333333333333333,
delay=2,
eligibility=0.1234567890123456,
last_pre_spike=8,
)
]
},
)
checkpoint = capture_runtime_checkpoint(network) # type: ignore[arg-type]
path = tmp_path / "runtime.json"
write_runtime_checkpoint(path, checkpoint)
loaded = read_runtime_checkpoint(path)
assert loaded == checkpoint
assert loaded.neuron_states[0].a == network.neurons[1].a
assert loaded.neuron_states[0].membrane_v == network.neurons[1].v
assert loaded.neuron_states[0].energy == network.neurons[1].energy
assert loaded.synapse_states[0].weight == network.synapses[1][0].weight
restored_rng = random.Random()
restored_rng.setstate(
(loaded.rng.version, tuple(loaded.rng.state), loaded.rng.gauss_next)
)
assert restored_rng.random() == network.rng.random()
|