psd_log_distance

Contents

psd_log_distance#

class braintools.fit.losses.psd_log_distance(prediction, target, dt, nperseg=None, band=None, *, noverlap=None, eps=None)#

Compute the mean squared difference of the log10 power spectra of two signals.

\[L = \frac{1}{|F| C} \sum_{f \in F} \sum_{c=1}^{C} \left[\log_{10}(\hat P_c(f) + \epsilon) - \log_{10}(P_c(f) + \epsilon)\right]^2,\]

where \(\hat P\) and \(P\) are Welch power spectral densities from power_spectral_density() (Hann window), \(F\) is the set of frequencies (optionally restricted to band) and \(C\) the number of channels.

Parameters:
  • prediction (Any) – Signal of shape (n_time,) or (n_time, n_channels).

  • target (Any) – Signal of the same shape. If one input carries units, both must, with the same dimension; the target is converted to the prediction’s unit.

  • dt (Any) – Sampling interval, static. A plain number is in seconds (the continuous-signal convention of braintools.analysis.spectral), not milliseconds as in rate_loss() and van_rossum_loss(); pass a Quantity (e.g. 0.1 * u.ms) to avoid ambiguity.

  • nperseg (int | None) – Welch segment length (static). Default n_time // 8.

  • band (Tuple[Any, Any] | None) – (f_low, f_high); only frequencies with f_low <= f <= f_high enter the mean. Plain numbers are in Hz.

  • noverlap (int | None) – Overlap between segments (static). Default nperseg // 2.

  • eps (float | None) – Floor added to the PSDs before the logarithm. Defaults to the smallest positive normal number of the dtype.

Returns:

Scalar loss in squared decades.

Return type:

Array

Raises:
  • ValueError – If the shapes differ, or band contains no frequency of the Welch grid.

  • TypeError – If only one input carries units, or dt/band have the wrong dimension.

Notes

The frequency grid rfftfreq(nperseg, dt) is computed on the host, so the band selection is a static index and the loss works under jax.jit and jax.vmap. It is differentiable with respect to the signals. An overall gain changes every log-PSD by the same amount, so the loss is sensitive to amplitude as well as to spectral shape [1].

References

Examples

>>> import brainunit as u
>>> import jax.numpy as jnp
>>> from braintools.fit import losses
>>> t = jnp.arange(1000) * 1e-3
>>> x = jnp.sin(2 * jnp.pi * 10 * t)
>>> float(losses.psd_log_distance(x, x, dt=1. * u.ms, nperseg=200))
0.0
>>> # Doubling the amplitude raises the PSD by log10(4) everywhere.
>>> round(float(losses.psd_log_distance(2 * x, x, 1e-3, 200, band=(5., 50.))), 4)
0.3625