deepSTRF.metrics package

Submodules

deepSTRF.metrics.losses module

Loss functions for deepSTRF: gradient-bearing, NaN-aware, reduction over neurons.

See docs/_source/md/metrics_paradigm.md for shape conventions, NaN handling, and reduction semantics.

deepSTRF.metrics.losses.mse_loss(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Boolean-masked mean squared error, reduced over the neuron axis.

Parameters:

• pred (torch.Tensor) – Prediction tensor of shape (B, N, 1, T).

• gt (torch.Tensor) – Ground-truth tensor of shape (B, N, 1, T) (PSTH or single-trial target) or (B, N, R, T) with R > 1 (raw responses), in which case it is collapsed to PSTH via nanmean(dim=2, keepdim=True) before the loss is computed. gt may contain NaN; positions where the resulting PSTH is NaN are dropped from the per-neuron mean.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to the post-collapse gt shape (B, N, 1, T). If None, defaults to ~gt.isnan(). If provided, REPLACES (does not augment) the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis. 'none' returns the per-neuron vector; 'mean' / 'sum' reduce it via nanmean / nansum.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

deepSTRF.metrics.losses.poisson_loss(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean', log_input: bool = False, validate_input: bool = False, eps: float = 1e-08) → Tensor

Negative Poisson log-likelihood (without the log(gt!) constant).

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T). Interpreted as the rate λ when log_input=False and as the log-rate η = log(λ) when log_input=True (see Notes).

• gt (torch.Tensor) – Ground-truth target. A pre-computed PSTH (B, N, 1, T) or a raw responses tensor (B, N, R, T) with R > 1; in the latter case it is collapsed to PSTH via nanmean(dim=2, keepdim=True) first. May contain NaN.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to (B, N, 1, T). If None, defaults to ~gt.isnan(). If provided, REPLACES (does not augment) the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis ('none' keeps the (N,) vector).

• log_input (bool, default False) – Selects the prediction parameterisation (see Notes).

• validate_input (bool, default False) – When log_input=False, raise on negative pred at masked-in positions instead of silently clamping inside the log. Costs a per-step CPU sync.

• eps (float, default 1e-8) – Floor applied to pred inside the log when log_input=False.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

Notes

Two parameterisations of the prediction are supported, matching the canonical-link logic of generalised linear models:

• log_input=False (default): the loss is pred − gt · log(pred + eps) per element. pred must be non-negative for the log to be meaningful; the implementation silently clamps pred to ≥ eps inside the log to avoid NaN propagation (the linear term keeps its sign). Pass validate_input=True for a loud failure on negative predictions.

• log_input=True: the loss becomes exp(pred) − gt · pred, which is well-defined for any real-valued pred. This is the standard trick for pairing a Poisson NLL with an unbounded readout (Linear, sign-permitting ParametricSigmoid, etc.). See metrics_paradigm.md §6.2 for the GLM-canonical-link derivation.

The log(gt!) Stirling term is not added — for non-integer gt (e.g. trial-averaged PSTH binned counts) it is meaningless, and for integer gt it is constant in pred so it does not affect optimisation. Users who want the full likelihood for AIC/BIC can add it themselves.

deepSTRF.metrics.performance module

Performance metrics for deepSTRF: NaN-aware, eval-only, reduction over neurons.

See docs/_source/md/metrics_paradigm.md for the design rationale, shape conventions, NaN handling, and per-metric formulas.

deepSTRF.metrics.performance.coherence(pred: Tensor, gt: Tensor, dt_ms: float, reduction: str = 'mean') → Tensor

Magnitude-squared coherence per neuron (mean over frequency bins).

Eval-only: uses scipy.signal.coherence, no gradient.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T).

• gt (torch.Tensor) – Ground-truth PSTH of shape (B, N, 1, T).

• dt_ms (float) – Time-bin width in milliseconds; sets the sampling rate fs = 1 / (dt_ms · 1e-3) passed to scipy.signal.coherence.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

Raises:

ValueError – If any element of pred or gt is NaN (this metric does not tolerate NaN; pre-flatten to a NaN-free subset first).

deepSTRF.metrics.performance.compute_CCmax(responses: Tensor, max_iters: int = 126) → Tensor

CCmax (Hsu / Spearman-Brown) per (B,) cell.

Internal helper kept importable for the audio loaders (legacy CCmax / TTRC pipeline). Prefer compute_neuron_quality() on the dataset itself.

Parameters:

• responses (torch.Tensor) – Responses of shape (B, R, T). NaN marks invalid repeats / time bins (dropped per (b,)).

• max_iters (int, default 126) – Cap on random half-splits per cell.

Returns:

Shape (B,). 1.0 for cells with R = 1; NaN for cells with ρ_half ≤ 0.

Return type:

torch.Tensor

deepSTRF.metrics.performance.compute_TTRC(responses: Tensor) → Tensor

Trial-to-trial response correlation per (B,) cell.

Internal helper kept importable for the audio loaders (legacy CCmax / TTRC pipeline). Prefer compute_neuron_quality() on the dataset itself.

Parameters:

responses (torch.Tensor) – Responses of shape (B, R, T). NaN-aware.

Returns:

Shape (B,). 1.0 for cells with R = 1; NaN for cells with no valid trial pair.

Return type:

torch.Tensor

deepSTRF.metrics.performance.corrcoef(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Pearson correlation per neuron over flattened valid (B, T) positions.

See metrics_paradigm.md §6.3.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T).

• gt (torch.Tensor) – Pre-computed PSTH (B, N, 1, T) or raw responses (B, N, R, T) with R > 1 (collapsed to PSTH via nanmean(dim=2, keepdim=True) first). May contain NaN.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to (B, N, 1, T). If None, defaults to ~gt.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. Neurons with fewer than 2 valid positions or zero variance yield NaN.

Return type:

torch.Tensor

deepSTRF.metrics.performance.fve(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Fraction of variance explained (R²) per neuron, over flattened valid positions.

FVE_n = 1 - SS_res / SS_tot where SS_tot is the variance of gt only. Negative when the prediction is worse than predicting the mean.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T).

• gt (torch.Tensor) – Pre-computed PSTH (B, N, 1, T) or raw responses (B, N, R, T) with R > 1 (collapsed to PSTH via nanmean(dim=2, keepdim=True) first). May contain NaN.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to (B, N, 1, T). If None, defaults to ~gt.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. Neurons with fewer than 2 valid positions or zero target variance yield NaN.

Return type:

torch.Tensor

deepSTRF.metrics.performance.noise_power(responses: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Sahani-Linden noise power per neuron (NP = TP - SP).

Parameters:

• responses (torch.Tensor) – Raw responses of shape (B, N, R, T). NaN marks missing repeats / time bins.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to responses. If None, defaults to ~responses.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

deepSTRF.metrics.performance.normalized_corrcoef(pred: Tensor, responses: Tensor, method: str = 'schoppe', mask: Tensor | None = None, reduction: str = 'mean', ccmax_iters: int = 126) → Tensor

Noise-corrected correlation coefficient per neuron.

See metrics_paradigm.md §6.4.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T) (the model-output convention; the R-axis must be 1).

• responses (torch.Tensor) – Raw responses (B, N, R, T). The PSTH is formed internally via nanmean(dim=2, keepdim=True); the noise ceiling uses the full repeat axis. May contain NaN.

• method ({'schoppe', 'hsu'}, default 'schoppe') – 'schoppe' divides by sqrt(var(pred) · SP) (Schoppe et al. 2016); 'hsu' divides the raw Pearson by CCmax (Hsu/Spearman-Brown).

• mask (torch.Tensor, optional) – Bool tensor broadcastable to the PSTH (B, N, 1, T). If None, defaults to ~psth.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

• ccmax_iters (int, default 126) – Cap on random half-splits per (stim, neuron) for the 'hsu' CCmax estimate.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. In the single-trial degenerate case (R = 1 everywhere) the noise correction is undefined and the raw corrcoef(pred, psth) is returned.

Return type:

torch.Tensor

Raises:

ValueError – If method is not 'schoppe' or 'hsu', or if the input shapes are inconsistent.

deepSTRF.metrics.performance.signal_power(responses: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Sahani-Linden signal power per neuron.

Computed per-stimulus (with the per-stim valid repeat count) then length-weighted across stimuli. Cells without any qualifying stim — needs ≥ 2 valid repeats and ≥ 2 valid time bins — return NaN.

Parameters:

• responses (torch.Tensor) – Raw responses of shape (B, N, R, T). NaN marks missing repeats / time bins.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to responses. If None, defaults to ~responses.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

deepSTRF.metrics.performance.snr(responses: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Signal-to-noise ratio per neuron (SNR = SP / NP).

Parameters:

• responses (torch.Tensor) – Raw responses of shape (B, N, R, T). NaN marks missing repeats / time bins.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to responses. If None, defaults to ~responses.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. Noiseless cells (NP ≈ 0) yield +inf; the caller decides whether to filter.

Return type:

torch.Tensor

Module contents

deepSTRF.metrics.coherence(pred: Tensor, gt: Tensor, dt_ms: float, reduction: str = 'mean') → Tensor

Magnitude-squared coherence per neuron (mean over frequency bins).

Eval-only: uses scipy.signal.coherence, no gradient.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T).

• gt (torch.Tensor) – Ground-truth PSTH of shape (B, N, 1, T).

• dt_ms (float) – Time-bin width in milliseconds; sets the sampling rate fs = 1 / (dt_ms · 1e-3) passed to scipy.signal.coherence.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

Raises:

ValueError – If any element of pred or gt is NaN (this metric does not tolerate NaN; pre-flatten to a NaN-free subset first).

deepSTRF.metrics.corrcoef(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Pearson correlation per neuron over flattened valid (B, T) positions.

See metrics_paradigm.md §6.3.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T).

• gt (torch.Tensor) – Pre-computed PSTH (B, N, 1, T) or raw responses (B, N, R, T) with R > 1 (collapsed to PSTH via nanmean(dim=2, keepdim=True) first). May contain NaN.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to (B, N, 1, T). If None, defaults to ~gt.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. Neurons with fewer than 2 valid positions or zero variance yield NaN.

Return type:

torch.Tensor

deepSTRF.metrics.fve(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Fraction of variance explained (R²) per neuron, over flattened valid positions.

FVE_n = 1 - SS_res / SS_tot where SS_tot is the variance of gt only. Negative when the prediction is worse than predicting the mean.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T).

• gt (torch.Tensor) – Pre-computed PSTH (B, N, 1, T) or raw responses (B, N, R, T) with R > 1 (collapsed to PSTH via nanmean(dim=2, keepdim=True) first). May contain NaN.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to (B, N, 1, T). If None, defaults to ~gt.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. Neurons with fewer than 2 valid positions or zero target variance yield NaN.

Return type:

torch.Tensor

deepSTRF.metrics.mse_loss(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Boolean-masked mean squared error, reduced over the neuron axis.

Parameters:

• pred (torch.Tensor) – Prediction tensor of shape (B, N, 1, T).

• gt (torch.Tensor) – Ground-truth tensor of shape (B, N, 1, T) (PSTH or single-trial target) or (B, N, R, T) with R > 1 (raw responses), in which case it is collapsed to PSTH via nanmean(dim=2, keepdim=True) before the loss is computed. gt may contain NaN; positions where the resulting PSTH is NaN are dropped from the per-neuron mean.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to the post-collapse gt shape (B, N, 1, T). If None, defaults to ~gt.isnan(). If provided, REPLACES (does not augment) the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis. 'none' returns the per-neuron vector; 'mean' / 'sum' reduce it via nanmean / nansum.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

deepSTRF.metrics.noise_power(responses: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Sahani-Linden noise power per neuron (NP = TP - SP).

Parameters:

• responses (torch.Tensor) – Raw responses of shape (B, N, R, T). NaN marks missing repeats / time bins.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to responses. If None, defaults to ~responses.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

deepSTRF.metrics.normalized_corrcoef(pred: Tensor, responses: Tensor, method: str = 'schoppe', mask: Tensor | None = None, reduction: str = 'mean', ccmax_iters: int = 126) → Tensor

Noise-corrected correlation coefficient per neuron.

See metrics_paradigm.md §6.4.

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T) (the model-output convention; the R-axis must be 1).

• responses (torch.Tensor) – Raw responses (B, N, R, T). The PSTH is formed internally via nanmean(dim=2, keepdim=True); the noise ceiling uses the full repeat axis. May contain NaN.

• method ({'schoppe', 'hsu'}, default 'schoppe') – 'schoppe' divides by sqrt(var(pred) · SP) (Schoppe et al. 2016); 'hsu' divides the raw Pearson by CCmax (Hsu/Spearman-Brown).

• mask (torch.Tensor, optional) – Bool tensor broadcastable to the PSTH (B, N, 1, T). If None, defaults to ~psth.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

• ccmax_iters (int, default 126) – Cap on random half-splits per (stim, neuron) for the 'hsu' CCmax estimate.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. In the single-trial degenerate case (R = 1 everywhere) the noise correction is undefined and the raw corrcoef(pred, psth) is returned.

Return type:

torch.Tensor

Raises:

ValueError – If method is not 'schoppe' or 'hsu', or if the input shapes are inconsistent.

deepSTRF.metrics.poisson_loss(pred: Tensor, gt: Tensor, mask: Tensor | None = None, reduction: str = 'mean', log_input: bool = False, validate_input: bool = False, eps: float = 1e-08) → Tensor

Negative Poisson log-likelihood (without the log(gt!) constant).

Parameters:

• pred (torch.Tensor) – Prediction of shape (B, N, 1, T). Interpreted as the rate λ when log_input=False and as the log-rate η = log(λ) when log_input=True (see Notes).

• gt (torch.Tensor) – Ground-truth target. A pre-computed PSTH (B, N, 1, T) or a raw responses tensor (B, N, R, T) with R > 1; in the latter case it is collapsed to PSTH via nanmean(dim=2, keepdim=True) first. May contain NaN.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to (B, N, 1, T). If None, defaults to ~gt.isnan(). If provided, REPLACES (does not augment) the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis ('none' keeps the (N,) vector).

• log_input (bool, default False) – Selects the prediction parameterisation (see Notes).

• validate_input (bool, default False) – When log_input=False, raise on negative pred at masked-in positions instead of silently clamping inside the log. Costs a per-step CPU sync.

• eps (float, default 1e-8) – Floor applied to pred inside the log when log_input=False.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

Notes

Two parameterisations of the prediction are supported, matching the canonical-link logic of generalised linear models:

• log_input=False (default): the loss is pred − gt · log(pred + eps) per element. pred must be non-negative for the log to be meaningful; the implementation silently clamps pred to ≥ eps inside the log to avoid NaN propagation (the linear term keeps its sign). Pass validate_input=True for a loud failure on negative predictions.

• log_input=True: the loss becomes exp(pred) − gt · pred, which is well-defined for any real-valued pred. This is the standard trick for pairing a Poisson NLL with an unbounded readout (Linear, sign-permitting ParametricSigmoid, etc.). See metrics_paradigm.md §6.2 for the GLM-canonical-link derivation.

The log(gt!) Stirling term is not added — for non-integer gt (e.g. trial-averaged PSTH binned counts) it is meaningless, and for integer gt it is constant in pred so it does not affect optimisation. Users who want the full likelihood for AIC/BIC can add it themselves.

deepSTRF.metrics.signal_power(responses: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Sahani-Linden signal power per neuron.

Computed per-stimulus (with the per-stim valid repeat count) then length-weighted across stimuli. Cells without any qualifying stim — needs ≥ 2 valid repeats and ≥ 2 valid time bins — return NaN.

Parameters:

• responses (torch.Tensor) – Raw responses of shape (B, N, R, T). NaN marks missing repeats / time bins.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to responses. If None, defaults to ~responses.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar.

Return type:

torch.Tensor

deepSTRF.metrics.snr(responses: Tensor, mask: Tensor | None = None, reduction: str = 'mean') → Tensor

Signal-to-noise ratio per neuron (SNR = SP / NP).

Parameters:

• responses (torch.Tensor) – Raw responses of shape (B, N, R, T). NaN marks missing repeats / time bins.

• mask (torch.Tensor, optional) – Bool tensor broadcastable to responses. If None, defaults to ~responses.isnan(); if provided, REPLACES the NaN-derived mask.

• reduction ({'none', 'mean', 'sum'}, default 'mean') – Reduction over the neuron axis.

Returns:

Shape (N,) if reduction='none', otherwise a scalar. Noiseless cells (NP ≈ 0) yield +inf; the caller decides whether to filter.

Return type:

torch.Tensor