flybrain/services/flysim/crates/flybrain-core/tests/plasticity.rs
acamilo 660c3cf00d
Some checks failed
ci / node 22 (test + typecheck) (push) Has been cancelled
ci / rust stable (cargo test --workspace --release) (push) Has been cancelled
ci / infra/tests/lint.sh (push) Has been cancelled
ci / playwright apps/stage (allowed to fail) (push) Has been cancelled
flybrain v0.4.0: public tree (history retained privately)
2026-09-21 15:09:46 +00:00

455 lines
16 KiB
Rust

//! Unit tests ported from `packages/brain/tests/plasticity.test.ts`.
//!
//! The golden scenarios pin the numbers; these pin the *rules*: what makes an edge plastic, which
//! pairings count, where the clamps sit, and that a rejected checkpoint mutates nothing.
mod common;
use std::sync::Arc;
use common::{toy_dataset, visual_toy_dataset};
use flybrain_core::lif::{LifConfig, LifNetwork};
use flybrain_core::plasticity::{
plasticity_version, PlasticityConfig, PlasticityState, RewardModulatedStdp, PLASTICITY_VERSION,
};
use flybrain_core::pool::WorkerPool;
/// One causal pre(0)->post(1) pairing 10 ms apart, as in the original unit tests.
fn causal(plasticity: &mut RewardModulatedStdp, ms: f64) {
plasticity.observe(&[1], 1, &[ms - 10.0, -1e6, -1e6, -1e6], ms);
}
fn rule(config: PlasticityConfig) -> RewardModulatedStdp {
RewardModulatedStdp::new(&toy_dataset(), config)
}
fn network(config: LifConfig, plasticity: PlasticityConfig) -> LifNetwork {
LifNetwork::new(Arc::new(visual_toy_dataset()), config, plasticity).expect("a network")
}
#[test]
fn causal_eligibility_is_necessary_and_learning_stays_in_anatomical_scope() {
let mut plasticity = rule(PlasticityConfig::default());
assert_eq!(plasticity.edges, vec![0]);
// No eligibility yet: reward alone changes nothing.
plasticity.reinforce(1.0, 100.0);
assert_eq!(plasticity.gain(0), 1.0);
// Eligibility but zero reward: `reinforce` returns before the restoring term too.
causal(&mut plasticity, 110.0);
plasticity.reinforce(0.0, 120.0);
assert_eq!(plasticity.gain(0), 1.0);
plasticity.reinforce(1.0, 130.0);
assert!(plasticity.gain(0) > 1.0);
// Only the selected edge moved; the rest of the connectome is immutable.
for edge in 1..5 {
assert_eq!(plasticity.gain(edge), 1.0, "edge {edge}");
}
let before = plasticity.export_state();
plasticity.reinforce(0.0, 140.0);
assert_eq!(plasticity.export_state(), before, "zero reward is a no-op");
}
#[test]
fn eligibility_decays_over_five_seconds_and_anti_causal_pairing_depresses() {
let mut immediate = rule(PlasticityConfig::default());
let mut delayed = rule(PlasticityConfig::default());
causal(&mut immediate, 100.0);
causal(&mut delayed, 100.0);
immediate.reinforce(1.0, 100.0);
delayed.reinforce(1.0, 5100.0);
assert!(
delayed.gain(0) > 1.0 && delayed.gain(0) < immediate.gain(0),
"a 5 s old trace must still potentiate, but less"
);
// Post before pre: the anti-causal branch depresses.
let mut plasticity = rule(PlasticityConfig::default());
plasticity.observe(&[0], 1, &[-1e6, 90.0, -1e6, -1e6], 100.0);
plasticity.reinforce(1.0, 100.0);
assert!(plasticity.gain(0) < 1.0);
// 20,000 alternating reinforcements never leave the clamp, so excitation never changes sign.
for index in 0..20_000 {
let ms = 200.0 + f64::from(index);
causal(&mut plasticity, ms);
plasticity.reinforce(if index < 10_000 { 10.0 } else { -10.0 }, ms);
assert!(
plasticity.gain(0) >= 0.899999 && plasticity.gain(0) <= 1.100001,
"gain left the clamp at index {index}: {}",
plasticity.gain(0)
);
}
}
#[test]
fn simultaneous_spikes_never_pair_in_either_direction() {
// The window is 0 < dt <= 100 ms, so dt = 0 is excluded, and dt = 101 is out of range.
let mut plasticity = rule(PlasticityConfig::default());
plasticity.observe(&[1], 1, &[100.0, -1e6, -1e6, -1e6], 100.0);
plasticity.reinforce(1.0, 100.0);
assert_eq!(plasticity.gain(0), 1.0, "dt = 0 must not pair");
let mut plasticity = rule(PlasticityConfig::default());
plasticity.observe(&[1], 1, &[100.0, -1e6, -1e6, -1e6], 201.0);
plasticity.reinforce(1.0, 201.0);
assert_eq!(plasticity.gain(0), 1.0, "dt = 101 is outside the window");
let mut plasticity = rule(PlasticityConfig::default());
plasticity.observe(&[1], 1, &[100.0, -1e6, -1e6, -1e6], 200.0);
plasticity.reinforce(1.0, 200.0);
assert!(plasticity.gain(0) > 1.0, "dt = 100 is inside the window");
}
#[test]
fn traces_are_clipped_after_every_update_not_at_reinforcement_time() {
let mut plasticity = rule(PlasticityConfig::default());
// 0.1 per pair at dt -> 0 would exceed 1 after eleven pairs without the per-update clip.
for index in 0..40 {
let ms = 100.0 + f64::from(index);
plasticity.observe(&[1], 1, &[ms - 1.0, -1e6, -1e6, -1e6], ms);
}
assert!(plasticity.traces.iter().all(|trace| trace.abs() <= 1.0));
assert!(
plasticity.traces[0] > 0.9,
"the trace should be near saturation"
);
}
#[test]
fn pre_post_roles_and_budget_choose_which_edges_are_plastic() {
// Toy connectome edges: 0: 0->1 (10), 1: 0->2 (-5), 2: 0->3 (20), 3: 1->3 (10), 4: 2->1 (5).
let with = |pre: &str, post: &str, budget: usize| {
rule(PlasticityConfig {
pre_role: pre.to_string(),
post_role: post.to_string(),
budget,
..PlasticityConfig::default()
})
.edges
.clone()
};
assert_eq!(with("kenyon", "mbon", 16_384), vec![0]);
assert_eq!(with("kenyon", "motor", 16_384), vec![2]);
assert_eq!(with("mbon", "motor", 16_384), vec![3]);
assert_eq!(with("mbon", "mbon", 16_384), vec![4]);
assert_eq!(with("kenyon", "missing", 16_384), Vec::<u32>::new());
// Budget keeps the strongest candidates, edge index breaking ties, re-sorted by edge.
let mut wide = toy_dataset();
wide.meta.roles.insert("kenyon".to_string(), vec![0, 1, 2]);
wide.meta.roles.insert("mbon".to_string(), vec![1, 2, 3]);
assert_eq!(
RewardModulatedStdp::new(&wide, PlasticityConfig::default()).edges,
vec![0, 2, 3, 4]
);
let budget_two = PlasticityConfig {
budget: 2,
..PlasticityConfig::default()
};
assert_eq!(
RewardModulatedStdp::new(&wide, budget_two).edges,
vec![0, 2],
"the two strongest are edges 2 (20) and 0 (10), re-sorted by edge"
);
// The topology hash separates all three selections even where the version string does not.
let motor = PlasticityConfig {
post_role: "motor".to_string(),
..PlasticityConfig::default()
};
let topologies: std::collections::HashSet<u32> = [
rule(PlasticityConfig::default()).topology(),
rule(motor).topology(),
RewardModulatedStdp::new(&wide, PlasticityConfig::default()).topology(),
]
.into_iter()
.collect();
assert_eq!(topologies.len(), 3);
}
#[test]
fn statistics_keep_the_historical_site_field_names() {
let mut plasticity = rule(PlasticityConfig::default());
let before = plasticity.statistics();
assert_eq!(before.version, PLASTICITY_VERSION);
assert!(before.enabled);
assert_eq!(before.synapses, 1);
assert_eq!(before.mushroom, 1);
assert_eq!(before.output, 0);
assert_eq!(before.updates, 0.0);
assert_eq!(before.changed, 0);
assert_eq!(before.mean_change, 0.0);
assert_eq!(before.max_change, 0.0);
assert_eq!(before.signal, 0.0);
causal(&mut plasticity, 100.0);
plasticity.reinforce(1.0, 100.0);
let after = plasticity.statistics();
assert_eq!(after.mushroom, plasticity.edges.len() as u64);
assert_eq!(after.output, 0);
assert_eq!(after.updates, 1.0);
assert_eq!(after.changed, 1);
assert!(after.max_change > 0.0);
assert_eq!(after.signal, flybrain_core::jsmath::tanh(1.0));
}
#[test]
fn clamps_follow_the_configured_gain_bounds() {
let wide = || PlasticityConfig {
min_gain: 0.5,
max_gain: 3.0,
learning_rate: 1.0,
..PlasticityConfig::default()
};
let mut plasticity = rule(wide());
for index in 0..50 {
let ms = 100.0 + f64::from(index);
causal(&mut plasticity, ms);
plasticity.reinforce(10.0, ms);
}
assert_eq!(plasticity.gain(0), 3.0);
let state = plasticity.export_state();
rule(wide())
.import_state(Some(&state))
.expect("the same bounds accept their own state");
// The default radius (0.100001) rejects a gain of 3.
let default_version = PlasticityState {
version: PLASTICITY_VERSION.to_string(),
..state
};
assert_eq!(
rule(PlasticityConfig::default())
.import_state(Some(&default_version))
.unwrap_err()
.message(),
"Invalid plasticity values"
);
}
#[test]
fn disabled_plasticity_ignores_observation_and_reinforcement() {
let mut plasticity = rule(PlasticityConfig::default());
plasticity.enabled = false;
causal(&mut plasticity, 100.0);
plasticity.reinforce(1.0, 100.0);
assert_eq!(plasticity.gain(0), 1.0);
assert_eq!(plasticity.statistics().updates, 0.0);
plasticity.enabled = true;
causal(&mut plasticity, 200.0);
plasticity.reinforce(1.0, 200.0);
assert!(plasticity.gain(0) > 1.0);
plasticity.clear_eligibility(200.0);
assert!(plasticity.traces.iter().all(|trace| *trace == 0.0));
assert!(plasticity.touched.iter().all(|touched| *touched == 200.0));
assert_eq!(plasticity.statistics().signal, 0.0);
// clearEligibility does not touch gains.
assert!(plasticity.gain(0) > 1.0);
}
#[test]
fn a_default_version_state_is_incompatible_with_a_non_default_configuration() {
let state = rule(PlasticityConfig::default()).export_state();
let faster = PlasticityConfig {
learning_rate: 0.004,
..PlasticityConfig::default()
};
assert_eq!(
rule(faster)
.import_state(Some(&state))
.unwrap_err()
.message(),
"Incompatible plasticity topology/version"
);
// Same rule constants, different site: the topology hash rejects it even though versions match.
let site = || PlasticityConfig {
post_role: "motor".to_string(),
..PlasticityConfig::default()
};
assert_eq!(rule(site()).version, PLASTICITY_VERSION);
assert_eq!(
rule(site())
.import_state(Some(&state))
.unwrap_err()
.message(),
"Incompatible plasticity topology/version"
);
// And through the network, which delegates the check.
let slower = PlasticityConfig {
trace_ms: 4000.0,
..PlasticityConfig::default()
};
let mut brain = network(LifConfig::default(), slower);
let default_state = network(LifConfig::default(), PlasticityConfig::default()).export_state();
assert_eq!(
brain.import_state(&default_state).unwrap_err().message(),
"Incompatible plasticity topology/version"
);
}
#[test]
fn invalid_imports_reject_before_any_neural_or_plastic_mutation() {
let mut brain = network(LifConfig::default(), PlasticityConfig::default());
let original = brain.export_state();
let mut bad_ms = original.clone();
bad_ms.ms = f64::NAN;
assert_eq!(
brain.import_state(&bad_ms).unwrap_err().message(),
"Invalid neural checkpoint values"
);
assert_eq!(brain.export_state(), original);
/// An expected message paired with the corruption that should produce it.
type Case = (&'static str, Box<dyn Fn(&mut PlasticityState)>);
let patches: Vec<Case> = vec![
(
"Incompatible plasticity topology/version",
Box::new(|state: &mut PlasticityState| state.version = "wrong".to_string()),
),
(
"Incompatible plasticity topology/version",
Box::new(|state: &mut PlasticityState| state.topology = 0),
),
(
"Invalid plasticity values",
Box::new(|state: &mut PlasticityState| state.gains = vec![f32::NAN]),
),
(
"Invalid plasticity values",
Box::new(|state: &mut PlasticityState| state.traces = vec![2.0]),
),
(
"Invalid plasticity values",
Box::new(|state: &mut PlasticityState| state.touched = vec![-1.0]),
),
(
"Invalid plasticity metadata",
Box::new(|state: &mut PlasticityState| state.updates = -1.0),
),
(
"Invalid plasticity metadata",
Box::new(|state: &mut PlasticityState| state.signal = 2.0),
),
];
for (message, patch) in patches {
let mut state = original.plasticity.clone();
patch(&mut state);
assert_eq!(
brain
.plasticity
.import_state(Some(&state))
.unwrap_err()
.message(),
message
);
assert_eq!(brain.export_state(), original, "mutated on {message}");
}
}
#[test]
fn exact_neural_continuation_survives_a_long_run_and_a_checkpoint() {
// 2^26 ms exercises the Float64 pairing timestamps: 1-ms differences must still resolve.
let mut source = network(LifConfig::default(), PlasticityConfig::default());
source.ms = f64::from(1u32 << 26);
source.step(15);
let ms = source.ms;
causal(&mut source.plasticity, ms);
source.plasticity.reinforce(1.0, ms);
let mut restored = network(LifConfig::default(), PlasticityConfig::default());
restored
.import_state(&source.export_state())
.expect("import");
assert_eq!(restored.export_state(), source.export_state());
source.step(25);
restored.step(25);
let source_ms = source.ms;
let restored_ms = restored.ms;
source.plasticity.reinforce(0.5, source_ms);
restored.plasticity.reinforce(0.5, restored_ms);
assert_eq!(restored.export_state(), source.export_state());
}
#[test]
fn a_warm_up_without_a_framebuffer_keeps_visual_neurons_finite() {
// The regression test for the prototype's uninitialized warm-up bug.
let mut brain = network(LifConfig::default(), PlasticityConfig::default());
brain.step(100);
assert!(brain.membrane.iter().all(|value| value.is_finite()));
assert_eq!(brain.visual_drive(), &[0.0]);
network(LifConfig::default(), PlasticityConfig::default())
.import_state(&brain.export_state())
.expect("a zero-drive checkpoint must load");
}
#[test]
fn the_default_configuration_keeps_the_historical_version() {
assert_eq!(
plasticity_version(&PlasticityConfig::default()),
PLASTICITY_VERSION
);
assert_eq!(
rule(PlasticityConfig::default()).export_state().version,
PLASTICITY_VERSION
);
}
/// `observe` sharded by slot range is the sequential walk, including where a shard is empty.
///
/// `golden_real` proves this on the real connectome's 16,384 slots as part of its whole-state
/// thread sweep. This one is the degenerate end: the toy fixture has a single plastic slot, so
/// every worker but one gets an empty range, and a pool with more workers than slots has to be a
/// no-op for the rest rather than a panic or a double write.
#[test]
fn sharded_observe_matches_the_sequential_walk_even_with_empty_shards() {
let spikes = [1u32, 0, 1];
let last_spike = [90.0, 95.0, -1e6, -1e6];
let mut reference = rule(PlasticityConfig::default());
reference.observe(&spikes, spikes.len(), &last_spike, 100.0);
reference.observe(&spikes, spikes.len(), &last_spike, 101.0);
assert!(
reference.traces.iter().any(|trace| *trace != 0.0),
"the fixture has to actually pair, or this test proves nothing"
);
for workers in [1usize, 2, 4, 16] {
let pool = WorkerPool::new(workers, "test").expect("a pool");
let mut sharded = rule(PlasticityConfig::default());
let bounds = sharded.slot_shards(workers);
assert_eq!(bounds.len(), workers + 1);
assert_eq!(bounds[0], 0);
assert_eq!(
*bounds.last().expect("a last bound") as usize,
sharded.gains.len(),
"the shards have to cover every slot"
);
for ms in [100.0, 101.0] {
sharded.observe_sharded(
&spikes,
spikes.len(),
&last_spike,
ms,
Some((&pool, &bounds)),
);
}
assert_eq!(sharded.traces, reference.traces, "traces, {workers} workers");
assert_eq!(
sharded.touched, reference.touched,
"touched, {workers} workers"
);
}
}