|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved. |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | +# ============================================================================== |
| 16 | +"""Tests for ``brainpy.dynold.synplast.short_term_plasticity`` (STD / STP). |
| 17 | +
|
| 18 | +These STP components are attached as the ``stp=`` slot of a ``TwoEndConn`` |
| 19 | +synapse; ``register_master`` allocates their state from the master's |
| 20 | +pre-synaptic group. The regressions here pin P11-M1: the discrete |
| 21 | +Tsodyks-Markram jumps must act on the value *at spike arrival* (the decayed |
| 22 | +local), not the pre-decay state held over from the previous step. |
| 23 | +""" |
| 24 | + |
| 25 | +import unittest |
| 26 | + |
| 27 | +import numpy as np |
| 28 | + |
| 29 | +import brainpy as bp |
| 30 | +import brainpy.math as bm |
| 31 | +from brainpy.context import share |
| 32 | + |
| 33 | + |
| 34 | +def _make_std(num=4, tau=200., U=0.07): |
| 35 | + """Build an STD bound to a real master synapse.""" |
| 36 | + pre = bp.neurons.LIF(num) |
| 37 | + post = bp.neurons.LIF(num) |
| 38 | + syn = bp.synapses.Exponential(pre, post, bp.connect.One2One(), |
| 39 | + stp=bp.synplast.STD(tau=tau, U=U), |
| 40 | + comp_method='dense') |
| 41 | + return syn.stp |
| 42 | + |
| 43 | + |
| 44 | +def _make_stp(num=4, U=0.15, tau_f=1500., tau_d=200.): |
| 45 | + pre = bp.neurons.LIF(num) |
| 46 | + post = bp.neurons.LIF(num) |
| 47 | + syn = bp.synapses.Exponential(pre, post, bp.connect.One2One(), |
| 48 | + stp=bp.synplast.STP(U=U, tau_f=tau_f, tau_d=tau_d), |
| 49 | + comp_method='dense') |
| 50 | + return syn.stp |
| 51 | + |
| 52 | + |
| 53 | +class TestSTD(unittest.TestCase): |
| 54 | + def setUp(self): |
| 55 | + bm.random.seed(0) |
| 56 | + bm.set_dt(0.1) |
| 57 | + |
| 58 | + def test_first_spike_from_rest(self): |
| 59 | + std = _make_std(3, tau=200., U=0.07) |
| 60 | + share.save(t=0.0, dt=bm.dt) |
| 61 | + std.update(bm.ones(3, dtype=bool)) |
| 62 | + # from rest x=1 -> x^+ = 1 - U = 0.93 (decay over one dt is negligible) |
| 63 | + np.testing.assert_allclose(bm.as_jax(std.x.value), np.full(3, 1 - 0.07), atol=2e-3) |
| 64 | + |
| 65 | + def test_jump_uses_decayed_state(self): |
| 66 | + # P11-M1: depress, let x recover for one step (no spike), then spike. The |
| 67 | + # depression must scale with the *decayed* x (= x^- at spike arrival), |
| 68 | + # i.e. x^+ = x_dec - U*x_dec, NOT x_dec - U*x_prev. |
| 69 | + U, tau, dt = 0.5, 50., bm.dt |
| 70 | + std = _make_std(1, tau=tau, U=U) |
| 71 | + share.save(t=0.0, dt=dt) |
| 72 | + std.update(bm.ones(1, dtype=bool)) # x drops to ~1-U |
| 73 | + x_prev = float(bm.as_jax(std.x.value)[0]) |
| 74 | + share.save(t=float(dt), dt=dt) |
| 75 | + std.update(bm.ones(1, dtype=bool)) # recover one dt, then spike |
| 76 | + x_after = float(bm.as_jax(std.x.value)[0]) |
| 77 | + |
| 78 | + # decayed value at spike arrival |
| 79 | + x_dec = x_prev + (1 - x_prev) / tau * float(dt) |
| 80 | + expected_correct = x_dec - U * x_dec |
| 81 | + expected_buggy = x_dec - U * x_prev |
| 82 | + self.assertAlmostEqual(x_after, expected_correct, places=5) |
| 83 | + # the two differ enough (recovery over dt) that the buggy form is rejected |
| 84 | + self.assertNotAlmostEqual(expected_correct, expected_buggy, places=7) |
| 85 | + |
| 86 | + |
| 87 | +class TestSTP(unittest.TestCase): |
| 88 | + def setUp(self): |
| 89 | + bm.random.seed(0) |
| 90 | + bm.set_dt(0.1) |
| 91 | + |
| 92 | + def test_jump_uses_decayed_state(self): |
| 93 | + # P11-M1: u^+ = u^- + U(1-u^-) and x^+ = x^- - u^+ x^- must use the |
| 94 | + # decayed (current-time) locals, not the previous-step Variables. |
| 95 | + U, tau_f, tau_d, dt = 0.5, 100., 50., bm.dt |
| 96 | + stp = _make_stp(1, U=U, tau_f=tau_f, tau_d=tau_d) |
| 97 | + share.save(t=0.0, dt=dt) |
| 98 | + stp.update(bm.ones(1, dtype=bool)) |
| 99 | + u_prev = float(bm.as_jax(stp.u.value)[0]) |
| 100 | + x_prev = float(bm.as_jax(stp.x.value)[0]) |
| 101 | + share.save(t=float(dt), dt=dt) |
| 102 | + stp.update(bm.ones(1, dtype=bool)) |
| 103 | + u_after = float(bm.as_jax(stp.u.value)[0]) |
| 104 | + x_after = float(bm.as_jax(stp.x.value)[0]) |
| 105 | + |
| 106 | + # decayed locals at spike arrival (exp_auto integrates exactly here) |
| 107 | + u_dec = u_prev + (U - u_prev / tau_f) * float(dt) |
| 108 | + x_dec = x_prev + (1 - x_prev) / tau_d * float(dt) |
| 109 | + u_correct = u_dec + U * (1 - u_dec) |
| 110 | + x_correct = x_dec - u_correct * x_dec |
| 111 | + self.assertAlmostEqual(u_after, u_correct, places=4) |
| 112 | + self.assertAlmostEqual(x_after, x_correct, places=4) |
| 113 | + # buggy variants (using the pre-decay Variables) must be distinguishable |
| 114 | + u_buggy = u_dec + U * (1 - u_prev) |
| 115 | + self.assertNotAlmostEqual(u_correct, u_buggy, places=6) |
| 116 | + |
| 117 | + |
| 118 | +if __name__ == '__main__': |
| 119 | + unittest.main() |
0 commit comments