braintools.metric module

`braintools.metric` module#

Metrics and Loss Functions for Neural Networks and Neuroscience.

This module provides comprehensive metrics and loss functions for both machine learning and neuroscience applications, including classification, regression, ranking, spike train analysis, and local field potential (LFP) analysis.

Key Features:

Classification Losses: Binary and multi-class cross-entropy, hinge loss, focal loss
Regression Losses: MSE, MAE, Huber loss, cosine similarity
Ranking Losses: Softmax ranking loss for learning to rank
Spike Train Metrics: Firing rate, synchrony, distance measures
LFP Analysis: Power spectral density, coherence, phase-amplitude coupling
Correlation Analysis: Cross-correlation, functional connectivity
Pairwise Metrics: Cosine similarity for pairwise comparisons

Quick Start - Classification:

import jax.numpy as jnp
from braintools.metric import softmax_cross_entropy, sigmoid_focal_loss

# Multi-class classification
logits = jnp.array([[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]])
labels = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
loss = softmax_cross_entropy(logits, labels)

# Binary classification with focal loss
predictions = jnp.array([0.9, 0.1, 0.7])
targets = jnp.array([1.0, 0.0, 1.0])
focal_loss = sigmoid_focal_loss(predictions, targets)

Quick Start - Regression:

import jax.numpy as jnp
from braintools.metric import squared_error, huber_loss

predictions = jnp.array([1.5, 2.3, 3.1])
targets = jnp.array([1.0, 2.0, 3.0])

# Mean squared error
mse = squared_error(predictions, targets)

# Huber loss (robust to outliers)
huber = huber_loss(predictions, targets, delta=1.0)

Quick Start - Spike Train Analysis:

import brainunit as u
import brainstate
import jax.numpy as jnp
from braintools.metric import (
    firing_rate, victor_purpura_distance,
    spike_train_synchrony
)

# Population firing rate from a spike matrix ``(num_time, num_neurons)``,
# smoothed with a sliding window of the given ``width``.
spikes = (brainstate.random.rand(1000, 50) < 0.1).astype(float)
rate = firing_rate(spikes, width=10.0 * u.ms, dt=1.0 * u.ms)

# Victor-Purpura distance between two spike-time trains (unitless times).
train1 = jnp.array([0.1, 0.3, 0.5])
train2 = jnp.array([0.12, 0.31, 0.52])
distance = victor_purpura_distance(train1, train2, cost_factor=1.0)

# Spike-train synchrony over a sliding window.
spike_matrix = (brainstate.random.rand(1000, 10) < 0.1).astype(float)
synchrony = spike_train_synchrony(spike_matrix, window_size=5.0 * u.ms, dt=1.0 * u.ms)

Classification Losses:

import jax.numpy as jnp
from braintools.metric import (
    sigmoid_binary_cross_entropy,
    softmax_cross_entropy_with_integer_labels,
    hinge_loss,
    multiclass_hinge_loss,
    kl_divergence,
    sigmoid_focal_loss
)

# Binary cross-entropy
logits = jnp.array([2.0, -1.0, 0.5])
labels = jnp.array([1.0, 0.0, 1.0])
bce = sigmoid_binary_cross_entropy(logits, labels)

# Multi-class with integer labels
logits = jnp.array([[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]])
labels = jnp.array([0, 1])  # Class indices
ce = softmax_cross_entropy_with_integer_labels(logits, labels)

# Hinge loss for SVM-style classification
predictions = jnp.array([0.9, -0.5, 0.3])
targets = jnp.array([1.0, -1.0, 1.0])
hinge = hinge_loss(predictions, targets)

# KL divergence
p = jnp.array([0.5, 0.3, 0.2])
q = jnp.array([0.4, 0.4, 0.2])
kl = kl_divergence(p, q)

# Focal loss for imbalanced datasets
predictions = jnp.array([0.9, 0.1, 0.6])
targets = jnp.array([1.0, 0.0, 1.0])
focal = sigmoid_focal_loss(predictions, targets, alpha=0.25, gamma=2.0)

Regression Losses:

import jax.numpy as jnp
from braintools.metric import (
    squared_error,
    absolute_error,
    l1_loss,
    l2_loss,
    huber_loss,
    log_cosh,
    cosine_similarity,
    cosine_distance
)

predictions = jnp.array([1.5, 2.3, 3.1, 4.2])
targets = jnp.array([1.0, 2.0, 3.0, 4.0])

# Various regression losses
mse = squared_error(predictions, targets)
mae = absolute_error(predictions, targets)
l1 = l1_loss(predictions, targets)
l2 = l2_loss(predictions, targets)
huber = huber_loss(predictions, targets, delta=1.0)
log_cosh_loss = log_cosh(predictions, targets)

# Cosine similarity/distance
x = jnp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
y = jnp.array([[1.0, 2.0, 3.0], [7.0, 8.0, 9.0]])
similarity = cosine_similarity(x, y)
distance = cosine_distance(x, y)

Spike Train Analysis:

import brainunit as u
import brainstate
import jax.numpy as jnp
from braintools.metric import (
    raster_plot,
    firing_rate,
    victor_purpura_distance,
    van_rossum_distance,
    spike_train_synchrony,
    burst_synchrony_index,
    phase_locking_value,
    spike_time_tiling_coefficient,
    correlation_index
)

# Raster plot data extraction (spike matrix is ``(num_time, num_neurons)``).
spike_matrix = jnp.array([[1, 0], [0, 1], [1, 1], [0, 0]])
times = jnp.arange(4) * 0.1 * u.second
neuron_ids, spike_times = raster_plot(spike_matrix, times)

# Smoothed population firing rate.
spikes = (brainstate.random.rand(1000, 20) < 0.1).astype(float)
rate = firing_rate(spikes, width=100.0 * u.ms, dt=1.0 * u.ms)

# Distance metrics between two spike-time trains (unitless times).
train1 = jnp.array([0.1, 0.3, 0.5])
train2 = jnp.array([0.12, 0.31, 0.52])
vp_dist = victor_purpura_distance(train1, train2, cost_factor=1.0)
vr_dist = van_rossum_distance(train1, train2, tau=0.01)

# Population synchrony measures (all operate on a spike matrix).
sm = (brainstate.random.rand(1000, 10) < 0.1).astype(float)
synchrony = spike_train_synchrony(sm, window_size=5.0 * u.ms, dt=1.0 * u.ms)
burst_sync = burst_synchrony_index(sm, dt=1.0 * u.ms)
plv = phase_locking_value(sm, reference_freq=40.0 * u.Hz, dt=1.0 * u.ms)
sttc = spike_time_tiling_coefficient(sm, dt=1.0 * u.ms, tau=0.005)
corr_idx = correlation_index(sm, window_size=5.0 * u.ms, dt=1.0 * u.ms)

LFP Analysis:

import brainunit as u
import brainstate
import jax.numpy as jnp
from braintools.metric import (
    unitary_LFP,
    power_spectral_density,
    coherence_analysis,
    phase_amplitude_coupling,
    theta_gamma_coupling,
    current_source_density,
    spectral_entropy,
    lfp_phase_coherence
)

# All spectral functions take the sampling step ``dt`` (not a rate ``fs``).
dt = 1.0 * u.ms
t = jnp.arange(2000) * 0.001  # seconds
lfp_signal = jnp.sin(2 * jnp.pi * 10 * t)  # 10 Hz oscillation

# Unitary LFP from a spike matrix ``(num_time, num_neurons)``; pass 'exc'/'inh'.
times = jnp.arange(2000) * 0.001  # ms time base (unitless)
spikes = (brainstate.random.rand(100, 2000) < 0.02).astype(float)
ulfp = unitary_LFP(times, spikes, 'exc', seed=42)

# Power spectral density (Welch).
freqs, psd = power_spectral_density(lfp_signal, dt=dt)

# Magnitude-squared coherence between two signals.
signal2 = jnp.sin(2 * jnp.pi * 10 * t + 0.1)
freqs, coherence = coherence_analysis(lfp_signal, signal2, dt=dt)

# Phase-amplitude coupling (Tort modulation index).
modulation_index, bin_centers, mean_amplitudes = phase_amplitude_coupling(
    lfp_signal, dt=dt,
    phase_freq_range=(4, 8),       # theta band
    amplitude_freq_range=(30, 80)  # gamma band
)

# Theta-gamma coupling (scalar modulation index).
tgc = theta_gamma_coupling(lfp_signal, dt=dt)

# Current source density along the laminar (channel) axis.
lfp_laminar = brainstate.random.randn(16, 2000)  # 16 channels x time
csd = current_source_density(lfp_laminar, electrode_spacing=100.0 * u.um, axis=0)

# Spectral entropy.
entropy = spectral_entropy(lfp_signal, dt=dt)

# Phase coherence across a multi-channel signal ``(num_time, num_channels)``.
multi_channel = brainstate.random.randn(2000, 4)
phase_coh = lfp_phase_coherence(multi_channel, dt=dt, freq_band=(8, 12))

Correlation Analysis:

import jax.numpy as jnp
import brainstate
from braintools.metric import (
    cross_correlation,
    voltage_fluctuation,
    matrix_correlation,
    weighted_correlation,
    functional_connectivity,
    functional_connectivity_dynamics
)

# Cross-correlation synchrony index from a spike matrix ``(num_time, num_neurons)``.
spikes = (brainstate.random.rand(1000, 10) < 0.1).astype(float)
cc = cross_correlation(spikes, bin=10, dt=1)

# Voltage-fluctuation synchrony index ``(num_time, num_neurons)``.
voltages = brainstate.random.randn(1000, 100)
vf = voltage_fluctuation(voltages)

# Correlation between the upper triangles of two square matrices.
a = brainstate.random.randn(20, 20)
b = brainstate.random.randn(20, 20)
corr = matrix_correlation(a, b)

# Weighted Pearson correlation between two 1-D series.
x = jnp.array([1., 2., 3., 4., 5.])
y = jnp.array([2., 4., 5., 4., 5.])
weights = jnp.array([1., 1., 2., 2., 1.])
w_corr = weighted_correlation(x, y, weights)

# Functional connectivity matrix from ``(num_time, num_signals)`` data.
time_series = brainstate.random.randn(1000, 10)  # 1000 samples, 10 regions
fc = functional_connectivity(time_series)

# Dynamic functional connectivity across sliding windows.
fc_dynamics = functional_connectivity_dynamics(
    time_series,
    window_size=100,
    step_size=50
)

Advanced: Fenchel-Young Losses:

import jax.numpy as jnp
from jax.scipy.special import logsumexp
from braintools.metric import make_fenchel_young_loss

# ``max_fun`` is applied over the last axis via ``jnp.vectorize`` with
# signature ``(n)->()``, so it must map a 1-D score vector to a scalar.
# ``logsumexp`` recovers the softmax cross-entropy Fenchel-Young loss.
loss_fn = make_fenchel_young_loss(max_fun=logsumexp)
scores = jnp.array([[2.0, 1.0, 3.0], [1.5, 2.5, 1.0]])
targets = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
loss = loss_fn(scores, targets)

Ranking Losses:

import jax.numpy as jnp
from braintools.metric import ranking_softmax_loss

# Ranking loss for learning to rank
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]])
loss = ranking_softmax_loss(scores, labels)

Utilities:

import jax.numpy as jnp
from braintools.metric import smooth_labels

# Label smoothing for regularization
labels = jnp.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
smoothed = smooth_labels(labels, alpha=0.1)

Comprehensive metric collection covering spiking activity, statistical analysis, and supervised learning objectives for neural modeling.

Classification Losses#

Objective functions for training classifiers on neural or behavioral labels.

`sigmoid_binary_cross_entropy`	Compute element-wise sigmoid cross entropy given logits and labels.
`hinge_loss`	Compute the hinge loss for binary classification.
`perceptron_loss`	Compute the binary perceptron loss.
`softmax_cross_entropy`	Compute the softmax cross entropy between logits and labels.
`softmax_cross_entropy_with_integer_labels`	Compute softmax cross entropy between logits and integer labels.
`multiclass_hinge_loss`	Compute multiclass hinge loss for classification.
`multiclass_perceptron_loss`	Compute multiclass perceptron loss for classification.
`poly_loss_cross_entropy`	Compute PolyLoss cross entropy between logits and labels.
`kl_divergence`	Compute the Kullback-Leibler divergence (relative entropy) loss.
`kl_divergence_with_log_targets`	Compute KL divergence when both predictions and targets are in log-space.
`convex_kl_divergence`	Compute a convex version of the Kullback-Leibler divergence loss.
`ctc_loss`	Compute Connectionist Temporal Classification (CTC) loss.
`ctc_loss_with_forward_probs`	Compute CTC loss and forward probabilities for sequence alignment.
`sigmoid_focal_loss`	Compute sigmoid focal loss for addressing class imbalance.
`nll_loss`	Compute negative log likelihood loss for classification.

Correlation#

Tools for measuring synchrony, functional connectivity, and aggregated correlations between neural signals.

`cross_correlation`	Calculate cross-correlation index between neurons.
`voltage_fluctuation`	Calculate neuronal synchronization via voltage variance analysis.
`matrix_correlation`	Compute Pearson correlation of upper triangular elements of two matrices.
`weighted_correlation`	Compute weighted Pearson correlation between two data series.
`functional_connectivity`	Compute functional connectivity matrix from time series data.
`functional_connectivity_dynamics`	Compute functional connectivity dynamics (FCD) matrix.

Fenchel-Young Loss#

Generalized convex losses derived from Fenchel-Young duality for structured prediction problems.

make_fenchel_young_loss

Create a Fenchel-Young loss function from a max function.

Spike Firing#

Descriptive statistics that summarize firing rates, timing variability, and spiking reliability.

`raster_plot`	Extract spike times and neuron indices for raster plot visualization.
`firing_rate`	Calculate the smoothed population firing rate from spike data.
`victor_purpura_distance`	Calculate Victor-Purpura distance between two spike trains.
`van_rossum_distance`	Calculate van Rossum distance between two spike trains.
`spike_train_synchrony`	Calculate spike train synchrony using the SPIKE-synchronization measure.
`burst_synchrony_index`	Calculate burst synchrony index based on co-occurring burst events.
`phase_locking_value`	Calculate phase-locking value (PLV) for spike synchronization.
`spike_time_tiling_coefficient`	Calculate Spike Time Tiling Coefficient (STTC).
`correlation_index`	Calculate correlation index for spike train synchrony.

Local Field Potential#

Metrics tailored to local field potential (LFP) analysis such as spectral characteristics and connectivity.

`unitary_LFP`	Calculate unitary local field potentials (uLFP) from spike train data.
`power_spectral_density`	Estimate the one-sided power spectral density (PSD) using Welch's method.
`coherence_analysis`	Compute the magnitude-squared coherence between two LFP signals (Welch).
`phase_amplitude_coupling`	Compute phase-amplitude coupling (PAC) via the Tort modulation index.
`theta_gamma_coupling`	Compute theta-gamma coupling strength using standard frequency bands.
`current_source_density`	Compute current source density (CSD) from laminar LFP recordings.
`spectral_entropy`	Compute the normalized spectral entropy of an LFP signal.
`lfp_phase_coherence`	Compute pairwise phase coherence between LFP channels in a frequency band.

Ranking Losses#

Losses for ordered prediction tasks including pairwise and list-wise ranking setups.

ranking_softmax_loss

Compute ranking softmax loss for learning-to-rank applications.

Regression Losses#

Continuous-valued error metrics for fitting neural or behavioral signals.

`squared_error`	Compute element-wise squared error between predictions and targets.
`absolute_error`	Compute element-wise absolute error between predictions and targets.
`l1_loss`	Measure the mean absolute error (MAE) between each element in the logits \(x\) and the targets \(y\).
`l2_loss`	Calculate the L2 loss for a set of predictions.
`l2_norm`	Compute the L2 norm of the difference between predictions and targets.
`safe_norm`	Compute vector or matrix norm with gradient-safe lower bound.
`huber_loss`	Compute Huber loss combining L1 and L2 properties for robust regression.
`log_cosh`	Calculate the log-cosh loss for a set of predictions.
`cosine_similarity`	Compute cosine similarity between predicted and target vectors.
`cosine_distance`	Compute the cosine distance between targets and predictions.

Pairwise Metrics#

Pairwise distance and similarity matrices between sets of samples.

`pairwise_cosine_similarity`	Compute the pairwise cosine similarity matrix between samples in `X` and `Y`.
`pairwise_cosine_distance`	Compute the pairwise cosine distance matrix between samples in `X` and `Y`.

Smoothing Losses#

Regularizers that promote smooth trajectories or label distributions over time.

smooth_labels

Apply label smoothing regularization to one-hot encoded labels.

braintools.metric module

Contents

`braintools.metric` module#

Classification Losses#

Correlation#

Fenchel-Young Loss#

Spike Firing#

Local Field Potential#

Ranking Losses#

Regression Losses#

Pairwise Metrics#

Smoothing Losses#

Modeling

Infrastructure

Compilation

braintools.metric module

Contents

braintools.metric module#

Classification Losses#

Correlation#

Fenchel-Young Loss#

Spike Firing#

Local Field Potential#

Ranking Losses#

Regression Losses#

Pairwise Metrics#

Smoothing Losses#

`braintools.metric` module#