neuropsychology/NeuroKit

# -*- coding: utf-8 -*-
from warnings import warn

import numpy as np

from ..misc import NeuroKitWarning, check_random_state, listify


def signal_simulate(
    duration=10,
    sampling_rate=1000,
    frequency=1,
    amplitude=0.5,
    noise=0,
    silent=False,
    random_state=None,
):
    """**Simulate a continuous signal**

    Parameters
    ----------
    duration : float
        Desired length of duration (s).
    sampling_rate : int
        The desired sampling rate (in Hz, i.e., samples/second).
    frequency : float or list
        Oscillatory frequency of the signal (in Hz, i.e., oscillations per second).
    amplitude : float or list
        Amplitude of the oscillations.
    noise : float
        Noise level (amplitude of the laplace noise).
    silent : bool
        If ``False`` (default), might print warnings if impossible frequencies are queried.
    random_state : None, int, numpy.random.RandomState or numpy.random.Generator
        Seed for the random number generator. See for ``misc.check_random_state`` for further information.

    Returns
    -------
    array
        The simulated signal.

    Examples
    --------
    .. ipython:: python

      import pandas as pd
      import neurokit2 as nk

      @savefig p_signal_simulate1.png scale=100%
      pd.DataFrame({
          "1Hz": nk.signal_simulate(duration=5, frequency=1),
          "2Hz": nk.signal_simulate(duration=5, frequency=2),
          "Multi": nk.signal_simulate(duration=5, frequency=[0.5, 3], amplitude=[0.5, 0.2])
      }).plot()
      @suppress
      plt.close()

    """
    n_samples = int(np.rint(duration * sampling_rate))
    period = 1 / sampling_rate
    seconds = np.arange(n_samples) * period

    signal = np.zeros(seconds.size)
    params = listify(frequency=frequency, amplitude=amplitude)

    for i in range(len(params["frequency"])):

        freq = params["frequency"][i]
        amp = params["amplitude"][i]
        # Apply a very conservative Nyquist criterion in order to ensure
        # sufficiently sampled signals.
        nyquist = sampling_rate * 0.1
        if freq > nyquist:
            if not silent:
                warn(
                    f"Skipping requested frequency"
                    f" of {freq} Hz since it cannot be resolved at the"
                    f" sampling rate of {sampling_rate} Hz. Please increase"
                    f" sampling rate to {freq * 10} Hz or choose frequencies"
                    f" smaller than or equal to {nyquist} Hz.",
                    category=NeuroKitWarning,
                )
            continue
        # Also make sure that at leat one period of the frequency can be
        # captured over the duration of the signal.
        if (1 / freq) > duration:
            if not silent:
                warn(
                    f"Skipping requested frequency"
                    f" of {freq} Hz since its period of {1 / freq} seconds"
                    f" exceeds the signal duration of {duration} seconds."
                    f" Please choose frequencies larger than"
                    f" {1 / duration} Hz or increase the duration of the"
                    f" signal above {1 / freq} seconds.",
                    category=NeuroKitWarning,
                )
            continue

        signal += _signal_simulate_sinusoidal(x=seconds, frequency=freq, amplitude=amp)
        # Add random noise
        if noise > 0:
            rng = check_random_state(random_state)
            signal += rng.laplace(0, noise, len(signal))

    return signal


# =============================================================================
# Simple Sinusoidal Model
# =============================================================================
def _signal_simulate_sinusoidal(x, frequency=100, amplitude=0.5):

    signal = amplitude * np.sin(2 * np.pi * frequency * x)

    return signal