van_rossum_loss

Contents

van_rossum_loss#

class braintools.fit.losses.van_rossum_loss(prediction, target, tau, dt=None, *, reduction='mean', squared=False)#

Compute a differentiable van Rossum distance between two dense spike rasters.

The difference of the two rasters is filtered by the causal exponential kernel \(K(t) = \frac{1}{\tau} e^{-t/\tau} H(t)\) with a scan over time,

\[d_k = e^{-\Delta t/\tau}\, d_{k-1} + \frac{\hat s_k - s_k}{\tau},\]

which is the exact filtered difference just after step \(k\). Between grid points it decays exponentially, so its squared integral over \([0, T\Delta t]\) is exactly

\[D^2 = \frac{\tau}{2}\left(1 - e^{-2\Delta t/\tau}\right) \sum_k d_k^2,\]

and the distance of each neuron is \(D = \sqrt{D^2}\) [1].

Parameters:

  • prediction (Any) – Raster of shape (T,) or (T, N), time on axis 0. Values may be real, e.g. spikes from a surrogate-gradient function such as braintools.surrogate.sigmoid().

  • target (Any) – Raster of the same shape.

  • tau (Any) – Kernel time constant: a Quantity, or a plain number in milliseconds. May be traced (and differentiated).

  • dt (Any) – Raster time step. A plain number is in milliseconds (the spike-time convention of braintools.analysis.spike and brainstate, same as tau), not seconds as in psd_log_distance(); pass a Quantity (e.g. 0.1 * u.ms) to avoid ambiguity. Defaults to brainstate.environ.get_dt().

  • reduction (str) – How the per-neuron distances are combined: their mean, their sum, or no reduction (an (N,) array; () for a 1-D raster).

  • squared (bool) – Return \(D^2\) instead of \(D\) per neuron (smooth at \(D = 0\)).

Returns:

The reduced distance, in \(\mathrm{ms}^{-1/2}\) (D² in \(\mathrm{ms}^{-1}\)), the convention of van_rossum_distance() with times in ms.

Return type:

Array

Raises:

  • ValueError – If a raster is not 1-D or 2-D, the shapes differ, dt is missing, or reduction is unknown.

  • TypeError – If a raster carries units, or tau/dt are not times.

See also

braintools.analysis.spike.van_rossum_distance

Event-based closed form.

rate_loss

Notes

For spikes on the grid (spike k at time \(k\,\Delta t\)) the result equals van_rossum_distance(times_1, times_2, tau, t_max=T*dt) up to float round-off: the only discretisation is that spike times are snapped to the grid. Works under jax.jit and jax.vmap and is differentiable with respect to the raster values and tau; at \(D = 0\) the gradient of \(D\) is taken to be 0.

References

Examples

>>> import brainunit as u
>>> import jax.numpy as jnp
>>> from braintools.analysis.spike import van_rossum_distance
>>> from braintools.fit import losses
>>> a = jnp.zeros(100).at[jnp.array([10, 30, 50])].set(1.)
>>> b = jnp.zeros(100).at[jnp.array([12, 32, 52])].set(1.)
>>> d = losses.van_rossum_loss(a, b, tau=5. * u.ms, dt=0.1 * u.ms)
>>> round(float(d), 4)
0.1512
>>> ref = van_rossum_distance(jnp.array([1., 3., 5.]), jnp.array([1.2, 3.2, 5.2]),
...                           tau=5., t_max=10.)
>>> round(float(ref), 4)
0.1512
Contents