Spaces:
Paused
Paused
File size: 4,501 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 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """Trace A/B divergence tick-by-tick between K and N."""
from __future__ import annotations
import json
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from src.config.loader import ConfigDict
from src.homeostasis.engine import HomeostasisEngine
from src.learning.learning_engine import LearningEngine
from src.storage.checkpoint import capture_runtime_checkpoint, write_runtime_checkpoint
from src.storage.core_restore import restore_full
from src.storage.runtime import StorageRuntimeConfig, StorageSession
from tests._restore_helpers import (
K,
N,
build_absolute_schedule,
create_network,
make_config,
run_absolute_schedule,
)
OUT_DIR = Path(__file__).resolve().parent.parent / "tmp" / "trace_diag"
OUT_DIR.mkdir(parents=True, exist_ok=True)
def _write_artifacts(network, homeo, learn, tmp_path, schedule, config):
rt = StorageRuntimeConfig(
snapshot_path=tmp_path / "base.b5d",
journal_path=tmp_path / "base.b5d.journal",
commit_interval_ticks=1,
)
with StorageSession(network, rt):
pass
from tests._restore_helpers import capture_learning_state
learn_state = capture_learning_state(learn)
checkpoint = capture_runtime_checkpoint(
network,
homeostasis_rates=homeo._rates_hz,
learning_states=learn_state["states"] if learn_state else None,
pending_rewards=learn_state["pending_rewards"] if learn_state else None,
)
cp_path = tmp_path / "runtime.json"
write_runtime_checkpoint(cp_path, checkpoint)
cfg_path = tmp_path / "config.json"
cfg_path.write_text(json.dumps(config, sort_keys=True), encoding="utf-8")
sched_path = tmp_path / "schedule.json"
sched_path.write_text(json.dumps(schedule, sort_keys=True), encoding="utf-8")
return {
"snapshot": rt.snapshot_path,
"journal": rt.journal_path,
"checkpoint": cp_path,
"config": cfg_path,
"schedule": sched_path,
}
def _neuron_0_state(network, homeo):
n = network.neurons[0]
return {
"tick": network.current_tick,
"v": n.v,
"u": n.u,
"threshold_adaptation": n.threshold_adaptation,
"last_spike_tick": n.last_spike_tick,
"homeo_rate": homeo._rates_hz.get(0, 0.0) if homeo else None,
}
def main():
config = make_config()
schedule = build_absolute_schedule(config, N)
# Path A: run to N, recording neuron 0 each tick after K
net_a = create_network(config)
homeo_a = HomeostasisEngine(net_a, config)
homeo_a.attach()
learn_a = LearningEngine(net_a, config)
learn_a.attach()
run_absolute_schedule(net_a, schedule, K)
trace_a = []
while net_a.current_tick < N:
net_a.step()
trace_a.append(_neuron_0_state(net_a, homeo_a))
# Path B: run to K, restore, then step-by-step to N
net_b = create_network(config)
homeo_b = HomeostasisEngine(net_b, config)
homeo_b.attach()
learn_b = LearningEngine(net_b, config)
learn_b.attach()
run_absolute_schedule(net_b, schedule, K)
tmp_path = OUT_DIR / "path_b"
tmp_path.mkdir(exist_ok=True)
artifacts = _write_artifacts(net_b, homeo_b, learn_b, tmp_path, schedule, config)
bundle = restore_full(
snapshot_path=artifacts["snapshot"],
journal_path=artifacts["journal"],
checkpoint_path=artifacts["checkpoint"],
config=ConfigDict(config),
recovered_path=tmp_path / "recovered.b5d",
create_homeostasis_engine=True,
create_learning_engine=True,
)
trace_b = []
while bundle.network.current_tick < N:
bundle.network.step()
trace_b.append(_neuron_0_state(bundle.network, bundle.homeostasis_engine))
# Compare tick-by-tick
diff_tick = None
for i, (sa, sb) in enumerate(zip(trace_a, trace_b)):
if sa != sb:
diff_tick = i
break
out = OUT_DIR / "trace_comparison.json"
out.write_text(
json.dumps(
{"diff_tick_index": diff_tick, "A": trace_a, "B": trace_b}, indent=2
),
encoding="utf-8",
)
print(f"Trace written to {out}")
if diff_tick is not None:
tick = trace_a[diff_tick]["tick"]
print(f"First divergence at tick {tick}")
print(f" A: {trace_a[diff_tick]}")
print(f" B: {trace_b[diff_tick]}")
else:
print("No divergence found between K and N")
if __name__ == "__main__":
main()
|