Utility function for auditory oddball stimulus generation.

oddball.py

"""Utilities for generating an auditory oddball stimulus."""
import os
from pathlib import Path
from typing import Iterable, Optional, Tuple, Type, Union

import librosa
import numpy as np
import soundfile as sf

# Exceptions.
LENGTH_MISMATCH: str = (
    "Length of `oddball_distances` and `distance_weights` must be the same."
)

# General-purpose path-like input type representation.
PathLike: Type = Union[str, bytes, os.PathLike, Path]


def normalize_distance_parameters(
    distances: Union[int, Iterable[int]],
    weights: Union[float, Iterable[float]],
) -> Tuple[np.ndarray, np.ndarray]:
    """Validate oddball distances and weights.

    Parameters
    ----------
    distances : Union[int, Iterable[int]]
        Distance between oddball clicks, or a list of possible distances (s)
    weights : Union[float, Iterable[float]]
        Probability of each possible oddball distance

    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        Normalized oddball distances and weights

    Raises
    ------
    ValueError
        If `oddball_distances` and `distance_weights` are not of the same length.
    """
    # Standardize distances input format.
    distances = np.array(
        [distances]
        if isinstance(distances, (int, float))
        else list(distances),
        dtype=int,
    )
    weights = np.array(
        [weights] if isinstance(weights, (float, int)) else weights,
        dtype=float,
    )
    # Validate input.
    if len(distances) != len(weights):
        raise ValueError(LENGTH_MISMATCH)
    # Normalize weights.
    weights /= weights.sum()
    return distances, weights


def generate_oddball_indices(
    distances: np.ndarray, weights: np.ndarray, duration: float
) -> np.ndarray:
    """Generate indices of oddball clicks.

    Parameters
    ----------
    distances : np.ndarray
        Distance between oddball clicks (s)
    weights : np.ndarray
        Probability of each possible oddball distance
    duration : float
        Duration of the created stimulus (s)

    Returns
    -------
    np.ndarray
        Indices of oddball clicks
    """
    # Generate oddball distances.
    distances = np.random.choice(distances, int(duration), p=weights)
    oddball_indices = np.cumsum(distances)
    oddball_indices = oddball_indices[oddball_indices < duration]
    return oddball_indices


def infer_stimulus_timings(
    click_frequency: float, duration: float, oddball_indices: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
    """Generate click times.

    Parameters
    ----------
    click_frequency : float
        Click frequency (Hz)
    duration : float
        Duration of the created stimulus (s)
    oddball_indices : np.ndarray
        Indices of oddball clicks

    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        Click times, oddball times
    """
    n_clicks = int(click_frequency * duration)
    click_times = np.linspace(0, duration, n_clicks, endpoint=False)
    oddball_times = click_times[oddball_indices]
    click_times = np.setdiff1d(click_times, oddball_times)
    return click_times, oddball_times


def create_oddball_audio(
    click_times: np.ndarray,
    oddball_times: np.ndarray,
    click_tone: float,
    oddball_tone: float,
    click_duration: float,
    oddball_duration: float,
    sampling_rate: float,
) -> np.ndarray:
    """Generate audio for complete oddball configuration.

    Parameters
    ----------
    click_times : np.ndarray
        Click times
    oddball_times : np.ndarray
        Oddball times
    click_tone : float
        Non-oddball click tone (Hz)
    oddball_tone : float
        Oddball click tone (Hz)
    click_duration : float
        Non-oddball click duration (s)
    oddball_duration : float
        Oddball click duration (s)
    sampling_rate : float
        Sampling rate (Hz)

    Returns
    -------
    np.ndarray
        Generated stimulus data
    """
    # Create click train.
    click_train = librosa.clicks(
        times=click_times,
        click_freq=click_tone,
        click_duration=click_duration,
        sr=sampling_rate,
    )
    # Create oddball train.
    oddball_train = librosa.clicks(
        times=oddball_times,
        click_freq=oddball_tone,
        click_duration=oddball_duration,
        length=len(click_train),
        sr=sampling_rate,
    )
    # Combine click and oddball trains.
    return click_train + oddball_train


def generate_oddball_stimulus(
    click_frequency: float = 1.0,
    click_duration: float = 0.4,
    click_tone: float = 1000.0,
    oddball_tone: float = 1200.0,
    oddball_duration: float = 0.4,
    oddball_distances: Union[int, Iterable[int]] = 10,
    distance_weights: Union[float, Iterable[float]] = 1.0,
    duration: float = 600.0,
    sampling_rate: int = 44_100,
    output_path: Optional[PathLike] = None,
) -> np.ndarray:
    """Generate an auditory oddball stimulus.

    Parameters
    ----------
    click_frequency : float
        Click frequency (Hz)
    click_duration : float
        Click duration (s)
    click_tone : float
        Non-oddball click tone (Hz)
    oddball_tone : float
        Oddball click tone (Hz)
    oddball_distances : Union[int, Iterable[int]]
        Distance between oddball clicks, or a list of possible distances (s)
    distance_weights : Union[float, Iterable[float]]
        Probability of each possible oddball distance
    duration : float
        Duration of the created stimulus (s)
    sampling_rate : float
        Audio data sampling rate (Hz)
    output_path : Optional[PathLike]
        Output path

    Returns
    -------
    np.ndarray
        Generated stimulus data
    """
    oddball_distances, distance_weights = normalize_distance_parameters(
        oddball_distances, distance_weights
    )
    oddball_indices = generate_oddball_indices(
        oddball_distances, distance_weights, duration
    )
    click_times, oddball_times = infer_stimulus_timings(
        click_frequency, duration, oddball_indices
    )
    stimulus = create_oddball_audio(
        click_times,
        oddball_times,
        click_tone,
        oddball_tone,
        click_duration,
        oddball_duration,
        sampling_rate,
    )
    if output_path is not None:
        sf.write(output_path, stimulus, sampling_rate)
    return stimulus