braintools 0.0.2__py2.py3-none-any.whl

braintools/metric/_lfp.py

@@ -0,0 +1,139 @@
+# 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 braincore as bc
+import jax
+from jax import numpy as jnp
+
+__all__ = [
+  'unitary_LFP',
+]
+
+
+def unitary_LFP(
+    times: bc.typing.ArrayLike,
+    spikes: bc.typing.ArrayLike,
+    spike_type: str,
+    xmax: bc.typing.ArrayLike = 0.2,
+    ymax: bc.typing.ArrayLike = 0.2,
+    va: bc.typing.ArrayLike = 200.,
+    lambda_: bc.typing.ArrayLike = 0.2,
+    sig_i: bc.typing.ArrayLike = 2.1,
+    sig_e: bc.typing.ArrayLike = 2.1 * 1.5,
+    location: str = 'soma layer',
+    seed: bc.typing.SeedOrKey = None
+) -> jax.Array:
+  """
+  A kernel-based method to calculate unitary local field potentials (uLFP)
+  from a network of spiking neurons [1]_.
+
+  .. note::
+     This method calculates LFP only from the neuronal spikes. It does not consider
+     the subthreshold synaptic events, or the dendritic voltage-dependent ion channels.
+
+  Examples
+  --------
+
+  If you have spike data of excitatory and inhibtiory neurons, you can get the LFP
+  by the following methods:
+
+  >>> import braincore as bc
+  >>> import jax
+  >>> import braintools as bt
+  >>> n_time = 1000
+  >>> n_exc = 100
+  >>> n_inh = 25
+  >>> times = jax.numpy.arange(n_time) * 0.1
+  >>> exc_sps = bc.random.random((n_time, n_exc)) < 0.3
+  >>> inh_sps = bc.random.random((n_time, n_inh)) < 0.4
+  >>> lfp = bt.metric.unitary_LFP(times, exc_sps, 'exc')
+  >>> lfp += bt.metric.unitary_LFP(times, inh_sps, 'inh')
+
+  Parameters
+  ----------
+  times: ndarray
+    The times of the recording points.
+  spikes: ndarray
+    The spikes of excitatory neurons recorded by brainpy monitors.
+  spike_type: str
+    The neuron type of the spike trains. It can be "exc" or "inh".
+  location: str
+    The location of the spikes recorded. It can be "soma layer", "deep layer",
+    "superficial layer" and "surface".
+  xmax: float
+    Size of the array (in mm).
+  ymax: float
+    Size of the array (in mm).
+  va: int, float
+    The axon velocity (mm/sec).
+  lambda_: float
+    The space constant (mm).
+  sig_i: float
+    The std-dev of inhibition (in ms)
+  sig_e: float
+    The std-dev for excitation (in ms).
+  seed: int
+    The random seed.
+
+  References
+  ----------
+  .. [1] Telenczuk, Bartosz, Maria Telenczuk, and Alain Destexhe. "A kernel-based
+         method to calculate local field potentials from networks of spiking
+         neurons." Journal of Neuroscience Methods 344 (2020): 108871.
+
+  """
+  if spike_type not in ['exc', 'inh']:
+    raise ValueError('"spike_type" should be "exc or ""inh". ')
+  if spikes.ndim != 2:
+    raise ValueError('"E_spikes" should be a matrix with shape of (num_time, num_neuron). '
+                     f'But we got {spikes.shape}')
+  if times.shape[0] != spikes.shape[0]:
+    raise ValueError('times and spikes should be consistent at the firs axis. '
+                     f'Bug we got {times.shape[0]} != {spikes.shape}.')
+
+  # Distributing cells in a 2D grid
+  rng = bc.random.RandomState(seed)
+  num_neuron = spikes.shape[1]
+  pos_xs, pos_ys = rng.rand(2, num_neuron).value * jnp.array([[xmax], [ymax]])
+  pos_xs, pos_ys = jnp.asarray(pos_xs), jnp.asarray(pos_ys)
+
+  # distance/coordinates
+  xe, ye = xmax / 2, ymax / 2  # coordinates of electrode
+  dist = jnp.sqrt((pos_xs - xe) ** 2 + (pos_ys - ye) ** 2)  # distance to electrode in mm
+
+  # amplitude
+  if location == 'soma layer':
+    amp_e, amp_i = 0.48, 3.  # exc/inh uLFP amplitude (soma layer)
+  elif location == 'deep layer':
+    amp_e, amp_i = -0.16, -0.2  # exc/inh uLFP amplitude (deep layer)
+  elif location == 'superficial layer':
+    amp_e, amp_i = 0.24, -1.2  # exc/inh uLFP amplitude (superficial layer)
+  elif location == 'surface layer':
+    amp_e, amp_i = -0.08, 0.3  # exc/inh uLFP amplitude (surface)
+  else:
+    raise NotImplementedError
+  A = jnp.exp(-dist / lambda_) * (amp_e if spike_type == 'exc' else amp_i)
+
+  # delay
+  delay = 10.4 + dist / va  # delay to peak (in ms)
+
+  # LFP Calculation
+  iis, ids = jnp.where(spikes)
+  tts = times[iis] + delay[ids]
+  exc_amp = A[ids]
+  tau = (2 * sig_e * sig_e) if spike_type == 'exc' else (2 * sig_i * sig_i)
+  return bc.transform.for_loop(lambda t: jnp.sum(exc_amp * jnp.exp(-(t - tts) ** 2 / tau)), times)

braintools/metric/_ranking.py

@@ -0,0 +1,163 @@
+# 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.
+# ==============================================================================
+
+"""
+Ranking losses.
+
+A ranking loss is a differentiable function that expresses the cost of a ranking
+induced by item scores compared to a ranking induced from relevance labels.
+
+Ranking losses are designed to operate on the last dimension of its inputs. The
+leading dimensions are considered batch dimensions.
+
+Standalone usage:
+
+>>> import braintools as bt
+>>> scores = jnp.array([2., 1., 3.])
+>>> labels = jnp.array([1., 0., 0.])
+>>> loss = bt.metric.ranking_softmax_loss(scores, labels)
+>>> print(f"{loss:.3f}")
+1.408
+
+Usage with a batch of data and a mask to indicate valid items.
+
+>>> scores = jnp.array([[2., 1., 0.], [1., 0.5, 1.5]])
+>>> labels = jnp.array([[1., 0., 0.], [0., 0., 1.]])
+>>> where = jnp.array([[True, True, False], [True, True, True]])
+>>> loss = bt.metric.ranking_softmax_loss(scores, labels, where=where)
+>>> print(f"{loss:.3f}")
+0.497
+"""
+
+from typing import Callable, Optional
+
+import jax
+import jax.numpy as jnp
+import braincore as bc
+
+__all__ = [
+  'ranking_softmax_loss',
+]
+
+
+def _safe_reduce(
+    a: bc.typing.ArrayLike,
+    where: Optional[bc.typing.ArrayLike] = None,
+    reduce_fn: Optional[Callable[..., bc.typing.ArrayLike]] = None,
+) -> bc.typing.ArrayLike:
+  """Reduces the values of given array while preventing NaN in the output.
+
+  For :func:`jax.numpy.mean` reduction, this additionally prevents ``NaN`` in
+  the output if all elements are masked. This can happen with pairwise losses
+  where there are no valid pairs because all labels are the same. In this
+  situation, 0 is returned instead.
+
+  When there is no ``reduce_fn``, this will set elements of ``a`` to 0 according
+  to the ``where`` mask.
+
+  Args:
+    a: The :class:`jax.Array` to reduce.
+    where: An optional :class:`jax.Array` indicating which elements to include
+      in the reduction.
+    reduce_fn: The function used to reduce. If None, no reduction is performed.
+
+  Returns:
+    The result of reducing the values of ``a`` using given ``reduce_fn``.
+  """
+  # Reduce values if there is a reduce_fn, otherwise keep the values as-is.
+  output = reduce_fn(a, where=where) if reduce_fn is not None else a
+
+  if reduce_fn is jnp.mean:
+    # For mean reduction, we have to check whether the input contains any NaN
+    # values, to ensure that masked mean reduction does not hide them (see
+    # below).
+    is_input_valid = jnp.logical_not(jnp.any(jnp.isnan(a)))
+
+    # The standard jnp.mean implementation returns NaN if `where` is False
+    # everywhere. This can happen in our case, e.g. pairwise losses with no
+    # valid pairs. Instead, we prefer that the loss returns 0 in these cases.
+    # Note that this only hides those NaN values if the input did not contain
+    # any NaN values. Otherwise it just returns the output as-is.
+    output = jnp.where(jnp.isnan(output) & is_input_valid, 0.0, output)
+
+  if reduce_fn is None and where is not None:
+    # When there is no reduce_fn (i.e. we are returning an unreduced
+    # loss/metric), set the values of `a` to 0 for invalid (masked) items.
+    # This makes sure that manual sum reduction on an unreduced loss works as
+    # expected:
+    # `jnp.sum(loss_fn(reduce_fn=None)) == loss_fn(reduce_fn=jnp.sum)`
+    output = jnp.where(where, output, 0.0)
+
+  return output
+
+
+def ranking_softmax_loss(
+    logits: bc.typing.ArrayLike,
+    labels: bc.typing.ArrayLike,
+    *,
+    where: Optional[bc.typing.ArrayLike] = None,
+    weights: Optional[bc.typing.ArrayLike] = None,
+    reduce_fn: Optional[Callable[..., bc.typing.ArrayLike]] = jnp.mean
+) -> bc.typing.ArrayLike:
+  r"""Ranking softmax loss.
+
+  Definition:
+
+  .. math::
+      \ell(s, y) = -\sum_i y_i \log \frac{\exp(s_i)}{\sum_j \exp(s_j)}
+
+  Args:
+    logits: A ``[..., list_size]``-:class:`~jax.Array`, indicating the score of
+      each item.
+    labels: A ``[..., list_size]``-:class:`~jax.Array`, indicating the relevance
+      label for each item.
+    where: An optional ``[..., list_size]``-:class:`~jax.Array`, indicating
+      which items are valid for computing the loss. Items for which this is
+      False will be ignored when computing the loss.
+    weights: An optional ``[..., list_size]``-:class:`~jax.Array`, indicating
+      the weight for each item.
+    reduce_fn: An optional function that reduces the loss values. Can be
+      :func:`jax.numpy.sum` or :func:`jax.numpy.mean`. If ``None``, no reduction
+      is performed.
+
+  Returns:
+    The ranking softmax loss.
+  """
+  assert bc.math.is_float(logits), "logits must be a float type."
+  labels = labels.astype(logits.dtype)
+
+  # Applies mask so that masked elements do not count towards the loss.
+  if where is not None:
+    labels = jnp.where(where, labels, jnp.zeros_like(labels))
+    logits = jnp.where(where, logits, -jnp.ones_like(logits) * jnp.inf)
+
+  # Apply weights to labels.
+  if weights is not None:
+    labels *= weights
+
+  # Scales labels and logits to match the cross entropy loss.
+  logits_log_softmax = jax.nn.log_softmax(logits, axis=-1)
+
+  # Computes per-element cross entropy.
+  softmax_cross_entropy = labels * logits_log_softmax
+
+  # Reduces softmax cross-entropy loss.
+  loss = -jnp.sum(softmax_cross_entropy, axis=-1, where=where)
+
+  # Setup mask to ignore lists with only invalid items in reduce_fn.
+  if where is not None:
+    where = jnp.any(where, axis=-1)
+
+  return _safe_reduce(loss, where=where, reduce_fn=reduce_fn)

braintools/metric/_ranking_test.py

@@ -0,0 +1,242 @@
+# 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 functools
+import math
+from absl.testing import parameterized
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from braintools.metric import _ranking
+
+# Export symbols from math for conciser test value definitions.
+exp = math.exp
+log = math.log
+logloss = lambda x: log(1.0 + exp(-x))
+sigmoid = lambda x: 1.0 / (1.0 + exp(-x))
+
+
+class RankingLossesTest(parameterized.TestCase):
+
+  @parameterized.parameters([
+      {
+          "loss_fn": _ranking.ranking_softmax_loss,
+          "expected_value": -(
+              log(exp(2.0) / (exp(0.0) + exp(3.0) + exp(1.0) + exp(2.0)))
+              + log(exp(1.0) / (exp(0.0) + exp(3.0) + exp(1.0) + exp(2.0)))
+          ),
+      },
+  ])
+  def test_computes_loss_value(self, loss_fn, expected_value):
+    scores = jnp.asarray([0.0, 3.0, 1.0, 2.0])
+    labels = jnp.asarray([0.0, 0.0, 1.0, 1.0])
+
+    loss = loss_fn(scores, labels, reduce_fn=jnp.sum)
+
+    np.testing.assert_allclose(jnp.asarray(expected_value), loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      {
+          "loss_fn": _ranking.ranking_softmax_loss,
+          "expected_value": -(
+              (-2.1e26 - (0.0 + -2.1e26 + 3.4e37 + 42.0))
+              + (3.4e37 - (0.0 + -2.1e26 + 3.4e37 + 42.0))
+          ),
+      },
+  ])
+  def test_computes_loss_with_extreme_inputs(self, loss_fn, expected_value):
+    scores = jnp.asarray([0.0, -2.1e26, 3.4e37, 42.0])
+    labels = jnp.asarray([0.0, 1.0, 1.0, 0.0])
+
+    loss = loss_fn(scores, labels, reduce_fn=jnp.sum)
+
+    np.testing.assert_allclose(jnp.asarray(expected_value), loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      {"loss_fn": _ranking.ranking_softmax_loss, "expected_value": 0.0},
+  ])
+  def test_computes_loss_for_zero_labels(self, loss_fn, expected_value):
+    scores = jnp.asarray([0.0, 3.0, 1.0, 2.0])
+    labels = jnp.asarray([0.0, 0.0, 0.0, 0.0])
+
+    loss = loss_fn(scores, labels, reduce_fn=jnp.sum)
+
+    np.testing.assert_allclose(jnp.asarray(expected_value), loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      {
+          "loss_fn": _ranking.ranking_softmax_loss,
+          "expected_value": -(
+              2.0 * log(exp(2.0) / (exp(0.0) + exp(3.0) + exp(1.0) + exp(2.0)))
+              + log(exp(1.0) / (exp(0.0) + exp(3.0) + exp(1.0) + exp(2.0)))
+          ),
+      },
+  ])
+  def test_computes_weighted_loss_value(self, loss_fn, expected_value):
+    scores = jnp.asarray([0.0, 3.0, 2.0, 1.0])
+    labels = jnp.asarray([0.0, 0.0, 1.0, 1.0])
+    weights = jnp.asarray([1.0, 1.0, 2.0, 1.0])
+
+    loss = loss_fn(scores, labels, weights=weights, reduce_fn=jnp.sum)
+
+    np.testing.assert_allclose(jnp.asarray(expected_value), loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      {
+          "loss_fn": _ranking.ranking_softmax_loss,
+          "expected_value": [
+              -(
+                  log(exp(2.0) / (exp(0.0) + exp(3.0) + exp(1.0) + exp(2.0)))
+                  + log(exp(1.0) / (exp(0.0) + exp(3.0) + exp(1.0) + exp(2.0)))
+              ),
+              -(
+                  2.0
+                  * log(exp(3.0) / (exp(3.0) + exp(1.0) + exp(4.0) + exp(2.0)))
+                  + log(exp(4.0) / (exp(3.0) + exp(1.0) + exp(4.0) + exp(2.0)))
+              ),
+          ],
+      },
+  ])
+  def test_computes_loss_value_with_vmap(self, loss_fn, expected_value):
+    scores = jnp.asarray([[0.0, 3.0, 1.0, 2.0], [3.0, 1.0, 4.0, 2.0]])
+    labels = jnp.asarray([[0.0, 0.0, 1.0, 1.0], [2.0, 0.0, 1.0, 0.0]])
+
+    loss_fn = functools.partial(loss_fn, reduce_fn=jnp.sum)
+    vmap_loss_fn = jax.vmap(loss_fn, in_axes=(0, 0), out_axes=0)
+    loss = vmap_loss_fn(scores, labels)
+
+    np.testing.assert_allclose(jnp.asarray(expected_value), loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      {
+          "loss_fn": _ranking.ranking_softmax_loss,
+          "expected_value": [
+              -log(exp(2.0) / (exp(2.0) + exp(1.0) + exp(3.0))),
+              -log(exp(1.5) / (exp(1.0) + exp(0.5) + exp(1.5))),
+          ],
+          "normalizer": 2.0,
+      },
+  ])
+  def test_computes_reduced_loss(self, loss_fn, expected_value, normalizer):
+    scores = jnp.array([[2.0, 1.0, 3.0], [1.0, 0.5, 1.5]])
+    labels = jnp.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+    expected_value = jnp.asarray(expected_value)
+
+    mean_loss = loss_fn(scores, labels, reduce_fn=jnp.mean)
+    sum_loss = loss_fn(scores, labels, reduce_fn=jnp.sum)
+
+    np.testing.assert_allclose(
+        mean_loss, jnp.sum(expected_value) / normalizer, rtol=1e-3
+    )
+    np.testing.assert_allclose(sum_loss, jnp.sum(expected_value), rtol=1e-3)
+
+  @parameterized.parameters([
+      {"loss_fn": _ranking.ranking_softmax_loss, "expected_shape": (2,)},
+  ])
+  def test_computes_unreduced_loss(self, loss_fn, expected_shape):
+    scores = jnp.array([[2.0, 1.0, 3.0], [1.0, 0.5, 1.5]])
+    labels = jnp.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+
+    none_loss = loss_fn(scores, labels, reduce_fn=None)
+    sum_loss = loss_fn(scores, labels, reduce_fn=jnp.sum)
+
+    self.assertEqual(none_loss.shape, expected_shape)
+    self.assertEqual(jnp.sum(none_loss), sum_loss)
+
+  @parameterized.parameters([
+      _ranking.ranking_softmax_loss,
+  ])
+  def test_computes_loss_value_with_where(self, loss_fn):
+    scores = jnp.asarray([0.0, 3.0, 1.0, 2.0])
+    labels = jnp.asarray([0.0, 0.0, 2.0, 1.0])
+    where = jnp.asarray([True, True, True, False])
+    expected_scores = jnp.asarray([0.0, 3.0, 1.0])
+    expected_labels = jnp.asarray([0.0, 0.0, 2.0])
+
+    loss = loss_fn(scores, labels, where=where)
+    expected_loss = loss_fn(expected_scores, expected_labels)
+
+    np.testing.assert_allclose(expected_loss, loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      _ranking.ranking_softmax_loss,
+  ])
+  def test_computes_loss_value_with_all_masked(self, loss_fn):
+    scores = jnp.asarray([0.0, 3.0, 1.0, 2.0])
+    labels = jnp.asarray([0.0, 0.0, 1.0, 1.0])
+    where = jnp.asarray([False, False, False, False])
+
+    loss = loss_fn(scores, labels, where=where)
+
+    np.testing.assert_allclose(jnp.asarray(0.0), loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      _ranking.ranking_softmax_loss,
+  ])
+  def test_computes_loss_with_arbitrary_batch_dimensions(self, loss_fn):
+    scores = jnp.asarray([2.0, 3.0, 1.0])
+    labels = jnp.asarray([0.0, 0.0, 1.0])
+    where = jnp.asarray([False, True, True])
+    original_loss = loss_fn(scores, labels, where=where)
+
+    scores = jnp.asarray([[[[2.0, 3.0, 1.0]]]])
+    labels = jnp.asarray([[[[0.0, 0.0, 1.0]]]])
+    where = jnp.asarray([[[[False, True, True]]]])
+    batched_loss = loss_fn(scores, labels, where=where)
+
+    np.testing.assert_allclose(original_loss, batched_loss, rtol=1e-3)
+
+  @parameterized.parameters([
+      _ranking.ranking_softmax_loss,
+  ])
+  def test_grad_does_not_return_nan_for_zero_labels(self, loss_fn):
+    scores = jnp.asarray([0.0, 3.0, 1.0, 2.0])
+    labels = jnp.asarray([0.0, 0.0, 0.0, 0.0])
+
+    grads = jax.grad(loss_fn)(scores, labels, reduce_fn=jnp.mean)
+
+    np.testing.assert_array_equal(
+        jnp.isnan(grads), jnp.zeros_like(jnp.isnan(grads))
+    )
+
+  @parameterized.parameters([
+      _ranking.ranking_softmax_loss,
+  ])
+  def test_grad_does_not_return_nan_with_all_masked(self, loss_fn):
+    scores = jnp.asarray([0.0, 3.0, 1.0, 2.0])
+    labels = jnp.asarray([0.0, 0.0, 1.0, 1.0])
+    where = jnp.asarray([False, False, False, False])
+
+    grads = jax.grad(loss_fn)(scores, labels, where=where, reduce_fn=jnp.mean)
+
+    np.testing.assert_array_equal(
+        jnp.isnan(grads), jnp.zeros_like(jnp.isnan(grads))
+    )
+
+  @parameterized.parameters([
+      _ranking.ranking_softmax_loss,
+  ])
+  def test_ignores_lists_containing_only_invalid_items(self, loss_fn):
+    scores = jnp.asarray([[0.0, 3.0, 1.0, 2.0], [3.0, 1.0, 4.0, 2.0]])
+    labels = jnp.asarray([[0.0, 0.0, 1.0, 1.0], [2.0, 0.0, 1.0, 0.0]])
+    mask = jnp.asarray([[1, 1, 1, 1], [0, 0, 0, 0]], dtype=jnp.bool_)
+
+    output = loss_fn(scores, labels, where=mask)
+    expected = loss_fn(scores[0, :], labels[0, :])
+
+    np.testing.assert_allclose(output, expected, rtol=1e-3)
+