braintools 0.0.2__py2.py3-none-any.whl

braintools/metric/_correlation.py

@@ -0,0 +1,302 @@
+# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+from typing import Union
+
+import braincore as bc
+import numpy as onp
+from jax import vmap, lax, numpy as jnp
+
+__all__ = [
+  'cross_correlation',
+  'voltage_fluctuation',
+  'matrix_correlation',
+  'weighted_correlation',
+  'functional_connectivity',
+  'functional_connectivity_dynamics',
+]
+
+
+def cross_correlation(spikes: bc.typing.ArrayLike,
+                      bin: Union[int, float],
+                      dt: Union[int, float] =None,
+                      method: str = 'loop'):
+  r"""Calculate cross correlation index between neurons.
+
+  The coherence [1]_ between two neurons i and j is measured by their
+  cross-correlation of spike trains at zero time lag within a time bin
+  of :math:`\Delta t = \tau`. More specifically, suppose that a long
+  time interval T is divided into small bins of :math:`\Delta t` and
+  that two spike trains are given by :math:`X(l)=` 0 or 1, :math:`Y(l)=` 0
+  or 1, :math:`l=1,2, \ldots, K(T / K=\tau)`. Thus, we define a coherence
+  measure for the pair as:
+
+  .. math::
+
+      \kappa_{i j}(\tau)=\frac{\sum_{l=1}^{K} X(l) Y(l)}
+      {\sqrt{\sum_{l=1}^{K} X(l) \sum_{l=1}^{K} Y(l)}}
+
+  The population coherence measure :math:`\kappa(\tau)` is defined by the
+  average of :math:`\kappa_{i j}(\tau)` over many pairs of neurons in the
+  network.
+
+  .. note::
+     To JIT compile this function, users should make ``bin``, ``dt``, ``numpy`` static.
+     For example, ``partial(brainpy.measure.cross_correlation, bin=10, numpy=False)``.
+
+  Parameters
+  ----------
+  spikes : ndarray
+      The history of spike states of the neuron group.
+  bin : float, int
+      The time bin to normalize spike states.
+  dt : float, optional
+      The time precision.
+  method: str
+    The method to calculate all pairs of cross correlation.
+    Supports two kinds of methods: `loop` and `vmap`.
+    `vmap` method needs much more memory.
+
+  Returns
+  -------
+  cc_index : float
+      The cross correlation value which represents the synchronization index.
+
+  References
+  ----------
+  .. [1] Wang, Xiao-Jing, and György Buzsáki. "Gamma oscillation by synaptic
+         inhibition in a hippocampal interneuronal network model." Journal of
+         neuroscience 16.20 (1996): 6402-6413.
+  """
+  dt = bc.environ.get_dt() if dt is None else dt
+  bin_size = int(bin / dt)
+  num_hist, num_neu = spikes.shape
+  num_bin = int(onp.ceil(num_hist / bin_size))
+  if num_bin * bin_size != num_hist:
+    spikes = jnp.append(spikes, jnp.zeros((num_bin * bin_size - num_hist, num_neu)), axis=0)
+  states = spikes.T.reshape((num_neu, num_bin, bin_size))
+  states = jnp.asarray(jnp.sum(states, axis=2) > 0., dtype=jnp.float_)
+  indices = jnp.tril_indices(num_neu, k=-1)
+
+  if method == 'loop':
+    def _f(i, j):
+      sqrt_ij = jnp.sqrt(jnp.sum(states[i]) * jnp.sum(states[j]))
+      return lax.cond(sqrt_ij == 0.,
+                      lambda _: 0.,
+                      lambda _: jnp.sum(states[i] * states[j]) / sqrt_ij,
+                      None)
+
+    res = bc.transform.for_loop(_f, *indices)
+
+  elif method == 'vmap':
+    @vmap
+    def _cc(i, j):
+      sqrt_ij = jnp.sqrt(jnp.sum(states[i]) * jnp.sum(states[j]))
+      return lax.cond(sqrt_ij == 0.,
+                      lambda _: 0.,
+                      lambda _: jnp.sum(states[i] * states[j]) / sqrt_ij,
+                      None)
+
+    res = _cc(*indices)
+  else:
+    raise ValueError(f'Do not support {method}. We only support "loop" or "vmap".')
+
+  return jnp.mean(jnp.asarray(res))
+
+
+def _f_signal(signal):
+  return jnp.mean(signal * signal) - jnp.mean(signal) ** 2
+
+
+def voltage_fluctuation(potentials, method='loop'):
+  r"""Calculate neuronal synchronization via voltage variance.
+
+  The method comes from [1]_ [2]_ [3]_.
+
+  First, average over the membrane potential :math:`V`
+
+  .. math::
+
+      V(t) = \frac{1}{N} \sum_{i=1}^{N} V_i(t)
+
+  The variance of the time fluctuations of :math:`V(t)` is
+
+  .. math::
+
+      \sigma_V^2 = \left\langle \left[ V(t) \right]^2 \right\rangle_t -
+      \left[ \left\langle V(t) \right\rangle_t \right]^2
+
+  where :math:`\left\langle \ldots \right\rangle_t = (1 / T_m) \int_0^{T_m} dt \, \ldots`
+  denotes time-averaging over a large time, :math:`\tau_m`. After normalization
+  of :math:`\sigma_V` to the average over the population of the single cell
+  membrane potentials
+
+  .. math::
+
+      \sigma_{V_i}^2 = \left\langle\left[ V_i(t) \right]^2 \right\rangle_t -
+      \left[ \left\langle V_i(t) \right\rangle_t \right]^2
+
+  one defines a synchrony measure, :math:`\chi (N)`, for the activity of a system
+  of :math:`N` neurons by:
+
+  .. math::
+
+      \chi^2 \left( N \right) = \frac{\sigma_V^2}{ \frac{1}{N} \sum_{i=1}^N
+      \sigma_{V_i}^2}
+
+  .. [1] Golomb, D. and Rinzel J. (1993) Dynamics of globally coupled
+         inhibitory neurons with heterogeneity. Phys. Rev. E 48:4810-4814.
+  .. [2] Golomb D. and Rinzel J. (1994) Clustering in globally coupled
+         inhibitory neurons. Physica D 72:259-282.
+  .. [3] David Golomb (2007) Neuronal synchrony measures. Scholarpedia, 2(1):1347.
+
+  Args:
+    potentials: The membrane potential matrix of the neuron group.
+    method: The method to calculate all pairs of cross correlation.
+       Supports two kinds of methods: `loop` and `vmap`.
+      `vmap` method will consume much more memory.
+
+  Returns:
+    sync_index: The synchronization index.
+  """
+
+  avg = jnp.mean(potentials, axis=1)
+  avg_var = jnp.mean(avg * avg) - jnp.mean(avg) ** 2
+
+  if method == 'loop':
+    _var = bc.transform.for_loop(_f_signal, jnp.moveaxis(potentials, 0, 1))
+  elif method == 'vmap':
+    _var = vmap(_f_signal, in_axes=1)(potentials)
+  else:
+    raise ValueError(f'Do not support {method}. We only support "loop" or "vmap".')
+  
+  var_mean = jnp.mean(_var)
+  r = jnp.where(var_mean == 0., 1., avg_var / var_mean)
+  return r
+
+
+def matrix_correlation(x, y):
+  """Pearson correlation of the lower triagonal of two matrices.
+
+    The triangular matrix is offset by k = 1 in order to ignore the diagonal line
+
+  Parameters
+  ----------
+  x: ndarray
+    First matrix.
+  y: ndarray
+    Second matrix
+  numpy: bool
+    Whether we use numpy array as the functional output.
+    If ``False``, this function can be JIT compiled.
+
+  Returns
+  -------
+  coef: ndarray
+    Correlation coefficient
+  """
+  if x.ndim != 2:
+    raise ValueError(f'Only support 2d array, but we got a array '
+                     f'with the shape of {x.shape}')
+  if y.ndim != 2:
+    raise ValueError(f'Only support 2d array, but we got a array '
+                     f'with the shape of {y.shape}')
+  x = x[jnp.triu_indices_from(x, k=1)]
+  y = y[jnp.triu_indices_from(y, k=1)]
+  cc = jnp.corrcoef(x, y)[0, 1]
+  return cc
+
+
+def functional_connectivity(activities):
+  """Functional connectivity matrix of timeseries activities.
+
+  Parameters
+  ----------
+  activities: ndarray
+    The multidimensional array with the shape of ``(num_time, num_sample)``.
+  numpy: bool
+    Whether we use numpy array as the functional output.
+    If ``False``, this function can be JIT compiled.
+
+  Returns
+  -------
+  connectivity_matrix: ndarray
+    ``num_sample x num_sample`` functional connectivity matrix.
+  """
+  if activities.ndim != 2:
+    raise ValueError('Only support 2d array with shape of "(num_time, num_sample)". '
+                     f'But we got a array with the shape of {activities.shape}')
+  fc = jnp.corrcoef(activities.T)
+  return jnp.nan_to_num(fc)
+
+
+def functional_connectivity_dynamics(activities, window_size=30, step_size=5):
+  """Computes functional connectivity dynamics (FCD) matrix.
+
+  Parameters
+  ----------
+  activities: ndarray
+    The time series with shape of ``(num_time, num_sample)``.
+  window_size: int
+    Size of each rolling window in time steps, defaults to 30.
+  step_size: int
+    Step size between each rolling window, defaults to 5.
+
+  Returns
+  -------
+  fcd_matrix: ndarray
+    FCD matrix.
+  """
+  pass
+
+
+def weighted_correlation(x, y, w):
+  """Weighted Pearson correlation of two data series.
+
+  Parameters
+  ----------
+  x: ndarray
+    The data series 1.
+  y: ndarray
+    The data series 2.
+  w: ndarray
+    Weight vector, must have same length as x and y.
+
+  Returns
+  -------
+  corr: ndarray
+    Weighted correlation coefficient.
+  """
+
+  def _weighted_mean(x, w):
+    """Weighted Mean"""
+    return jnp.sum(x * w) / jnp.sum(w)
+
+  def _weighted_cov(x, y, w):
+    """Weighted Covariance"""
+    return jnp.sum(w * (x - _weighted_mean(x, w)) * (y - _weighted_mean(y, w))) / jnp.sum(w)
+
+  if x.ndim != 1:
+    raise ValueError(f'Only support 1d array, but we got a array '
+                     f'with the shape of {x.shape}')
+  if y.ndim != 1:
+    raise ValueError(f'Only support 1d array, but we got a array '
+                     f'with the shape of {y.shape}')
+  if w.ndim != 1:
+    raise ValueError(f'Only support 1d array, but we got a array '
+                     f'with the shape of {w.shape}')
+  return _weighted_cov(x, y, w) / jnp.sqrt(_weighted_cov(x, x, w) * _weighted_cov(y, y, w))

braintools/metric/_correlation_test.py

@@ -0,0 +1,124 @@
+# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+
+import unittest
+from functools import partial
+import braintools as bt
+
+from jax import jit
+import jax.numpy as jnp
+import braincore as bc
+
+
+class TestCrossCorrelation(unittest.TestCase):
+  def test_c(self):
+    bc.random.seed()
+    spikes = jnp.asarray([[1, 0, 1, 0, 1, 0, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1, 0, 0]]).T
+    cc1 = bt.metric.cross_correlation(spikes, 1., dt=1.)
+    f_cc = jit(partial(bt.metric.cross_correlation, bin=1, dt=1.))
+    cc2 = f_cc(spikes)
+    print(cc1, cc2)
+    self.assertTrue(cc1 == cc2)
+    bc.util.clear_buffer_memory()
+
+  def test_cc(self):
+    bc.random.seed()
+    spikes = jnp.ones((1000, 10))
+    cc1 = bt.metric.cross_correlation(spikes, 1.)
+    self.assertTrue(cc1 == 1.)
+
+    spikes = jnp.zeros((1000, 10))
+    cc2 = bt.metric.cross_correlation(spikes, 1.)
+    self.assertTrue(cc2 == 0.)
+
+    bc.util.clear_buffer_memory()
+
+  def test_cc2(self):
+    bc.random.seed()
+    spikes = bc.random.randint(0, 2, (1000, 10))
+    print(bt.metric.cross_correlation(spikes, 1.))
+    print(bt.metric.cross_correlation(spikes, 0.5))
+    bc.util.clear_buffer_memory()
+
+  def test_cc3(self):
+    bc.random.seed()
+    spikes = bc.random.random((1000, 100)) < 0.8
+    print(bt.metric.cross_correlation(spikes, 1.))
+    print(bt.metric.cross_correlation(spikes, 0.5))
+    bc.util.clear_buffer_memory()
+
+  def test_cc4(self):
+    bc.random.seed()
+    spikes = bc.random.random((1000, 100)) < 0.2
+    print(bt.metric.cross_correlation(spikes, 1.))
+    print(bt.metric.cross_correlation(spikes, 0.5))
+    bc.util.clear_buffer_memory()
+
+  def test_cc5(self):
+    bc.random.seed()
+    spikes = bc.random.random((1000, 100)) < 0.05
+    print(bt.metric.cross_correlation(spikes, 1.))
+    print(bt.metric.cross_correlation(spikes, 0.5))
+    bc.util.clear_buffer_memory()
+
+
+class TestVoltageFluctuation(unittest.TestCase):
+  def test_vf1(self):
+    bc.random.seed()
+    voltages = bc.random.normal(0, 10, size=(100, 10))
+    print(bt.metric.voltage_fluctuation(voltages))
+
+    with bc.environ.context(precision=64):
+      voltages = jnp.ones((100, 10))
+      r1 = bt.metric.voltage_fluctuation(voltages)
+  
+      jit_f = jit(partial(bt.metric.voltage_fluctuation))
+      jit_f = jit(lambda a: bt.metric.voltage_fluctuation(a))
+      r2 = jit_f(voltages)
+      print(r1, r2)  # TODO: JIT results are different?
+      # self.assertTrue(r1 == r2)
+
+    bc.util.clear_buffer_memory()
+
+
+class TestFunctionalConnectivity(unittest.TestCase):
+  def test_cf1(self):
+    bc.random.seed()
+    act = bc.random.random((10000, 3))
+    r1 = bt.metric.functional_connectivity(act)
+
+    jit_f = jit(partial(bt.metric.functional_connectivity))
+    r2 = jit_f(act)
+
+    self.assertTrue(jnp.allclose(r1, r2))
+    bc.util.clear_buffer_memory()
+
+
+class TestMatrixCorrelation(unittest.TestCase):
+  def test_mc(self):
+    bc.random.seed()
+    A = bc.random.random((100, 100))
+    B = bc.random.random((100, 100))
+    r1 = (bt.metric.matrix_correlation(A, B))
+
+    jit_f = jit(partial(bt.metric.matrix_correlation))
+    r2 = jit_f(A, B)
+    self.assertTrue(jnp.allclose(r1, r2))
+    bc.util.clear_buffer_memory()
+
+

braintools/metric/_fenchel_young.py

@@ -0,0 +1,67 @@
+# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Fenchel-Young losses."""
+
+from typing import Any, Protocol
+
+import jax.numpy as jnp
+
+__all__ = [
+    "make_fenchel_young_loss",
+]
+
+
+class MaxFun(Protocol):
+
+  def __call__(self, scores, *args, **kwargs: Any):
+    ...
+
+
+def make_fenchel_young_loss(
+    max_fun: MaxFun
+):
+  """Creates a Fenchel-Young loss from a max function.
+
+  WARNING: The resulting loss accepts an arbitrary number of leading dimensions
+  with the fy_loss operating over the last dimension. The jaxopt version of this
+  function would instead flatten any vector in a single big 1D vector.
+
+  Examples:
+    Given a max function, e.g., the log sum exp, you can construct a
+    Fenchel-Young loss easily as follows:
+
+    >>> from jax.scipy.special import logsumexp
+    >>> fy_loss = make_fy_loss(max_fun=logsumexp)
+
+  Reference:
+    Blondel et al. `Learning with Fenchel-Young Losses
+    <https://arxiv.org/pdf/1901.02324.pdf>`_, 2020
+
+  Args:
+    max_fun: the max function on which the Fenchel-Young loss is built.
+
+  Returns:
+    A Fenchel-Young loss function with the same signature.
+  """
+
+  vdot_last_dim = jnp.vectorize(jnp.vdot, signature="(n),(n)->()")
+  max_fun_last_dim = jnp.vectorize(max_fun, signature="(n)->()")
+
+  def fenchel_young_loss(scores, targets, *args, **kwargs):
+    max_value = max_fun_last_dim(scores, *args, **kwargs)
+    return max_value - vdot_last_dim(targets, scores)
+
+  return fenchel_young_loss

braintools/metric/_fenchel_young_test.py

@@ -0,0 +1,55 @@
+# Copyright 2024 BrainPy Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========================================================================
+
+import jax
+import jax.numpy as jnp
+from absl.testing import parameterized
+from jax.scipy.special import logsumexp
+
+from braintools.metric import _classification
+from braintools.metric import _fenchel_young
+
+
+def one_hot_argmax(inputs: jnp.ndarray) -> jnp.ndarray:
+  """An argmax one-hot function for arbitrary shapes."""
+  inputs_flat = jnp.reshape(inputs, (-1))
+  flat_one_hot = jax.nn.one_hot(jnp.argmax(inputs_flat), inputs_flat.shape[0])
+  return jnp.reshape(flat_one_hot, inputs.shape)
+
+
+class FenchelYoungTest(parameterized.TestCase):
+
+  def test_fenchel_young_reg(self):
+    # Checks the behavior of the Fenchel-Young loss.
+    fy_loss = (_fenchel_young.make_fenchel_young_loss(logsumexp))
+    rng = jax.random.PRNGKey(0)
+    rngs = jax.random.split(rng, 2)
+    theta_true = jax.random.uniform(rngs[0], (8, 5))
+    y_true = jax.vmap(jax.nn.softmax)(theta_true)
+    theta_random = jax.random.uniform(rngs[1], (8, 5))
+    y_random = jax.vmap(jax.nn.softmax)(theta_random)
+    grad_random = jax.vmap(jax.grad(fy_loss))(theta_random, y_true)
+    # Checks that the gradient of the loss takes the correct form.
+    self.assertTrue(jnp.allclose(grad_random, y_random - y_true, rtol=1e-4))
+    y_one_hot = jax.vmap(one_hot_argmax)(theta_true)
+    int_one_hot = jnp.where(y_one_hot == 1.)[1]
+    loss_one_hot = jax.vmap(fy_loss)(theta_random, y_one_hot)
+    log_loss = jax.vmap(_classification.softmax_cross_entropy_with_integer_labels)(theta_random, int_one_hot)
+    # Checks that the FY loss associated to logsumexp is correct.
+    self.assertTrue(jnp.allclose(loss_one_hot, log_loss, rtol=1e-4))
+    # Checks that vmapping or not is equivalent.
+    loss_one_hot_no_vmap = fy_loss(theta_random, y_one_hot)
+    self.assertTrue(jnp.allclose(loss_one_hot, loss_one_hot_no_vmap, rtol=1e-4))
+

braintools/metric/_firings.py

@@ -0,0 +1,85 @@
+# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+from typing import Union
+import numpy as onp
+import jax.numpy as jnp
+import braincore as bc
+
+
+__all__ = [
+  'raster_plot',
+  'firing_rate',
+]
+
+
+def raster_plot(sp_matrix: bc.typing.ArrayLike,
+                times: bc.typing.ArrayLike):
+  """Get spike raster plot which displays the spiking activity
+  of a group of neurons over time.
+
+  Parameters
+  ----------
+  sp_matrix : bnp.ndarray
+      The matrix which record spiking activities.
+  times : bnp.ndarray
+      The time steps.
+
+  Returns
+  -------
+  raster_plot : tuple
+      Include (neuron index, spike time).
+  """
+  times = onp.asarray(times)
+  elements = onp.where(sp_matrix > 0.)
+  index = elements[1]
+  time = times[elements[0]]
+  return index, time
+
+
+def firing_rate(spikes: bc.typing.ArrayLike,
+                width: Union[int, float],
+                dt: Union[int, float]=None):
+  r"""Calculate the mean firing rate over in a neuron group.
+
+  This method is adopted from Brian2.
+
+  The firing rate in trial :math:`k` is the spike count :math:`n_{k}^{sp}`
+  in an interval of duration :math:`T` divided by :math:`T`:
+
+  .. math::
+
+      v_k = {n_k^{sp} \over T}
+
+  Parameters
+  ----------
+  spikes : ndarray
+    The spike matrix which record spiking activities.
+  width : int, float
+    The width of the ``window`` in millisecond.
+  dt : float, optional
+    The sample rate.
+
+  Returns
+  -------
+  rate : ndarray
+      The population rate in Hz, smoothed with the given window.
+  """
+  dt = bc.environ.get_dt() if (dt is None) else dt
+  width1 = int(width / 2 / dt) * 2 + 1
+  window = jnp.ones(width1) * 1000 / width
+  return jnp.convolve(jnp.mean(spikes, axis=1), window, mode='same')

braintools/metric/_firings_test.py

@@ -0,0 +1,41 @@
+# Copyright 2024- BrainPy Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+
+import unittest
+import braincore as bc
+import jax.numpy as jnp
+import braintools as bt
+
+
+class TestFiringRate(unittest.TestCase):
+  def test_fr1(self):
+    spikes = jnp.ones((1000, 10))
+    print(bt.metric.firing_rate(spikes, 1.))
+
+  def test_fr2(self):
+    bc.random.seed()
+    spikes = bc.random.random((1000, 10)) < 0.2
+    print(bt.metric.firing_rate(spikes, 1.))
+    print(bt.metric.firing_rate(spikes, 10.))
+
+  def test_fr3(self):
+    bc.random.seed()
+    spikes = bc.random.random((1000, 10)) < 0.02
+    print(bt.metric.firing_rate(spikes, 1.))
+    print(bt.metric.firing_rate(spikes, 5.))
+