pysmo.tools.noise

def generate_noise(
    model: NoiseModel,
    npts: int,
    delta: pd.Timedelta = SeismogramDefaults.delta,
    begin_time: pd.Timestamp = SeismogramDefaults.begin_time,
    return_velocity: bool = False,
    seed: int | None = None,
) -> MiniSeismogram:
    """Generate a random seismogram from a noise model.

    The amplitude spectrum is prescribed by the noise model and random phases
    are drawn uniformly from `[-π, π]`. The combined spectrum is transformed
    back to the time domain via an inverse FFT. Internally the computation is
    performed on the next power-of-two length greater than or equal to `npts`
    to ensure an efficient FFT; the central `npts` samples are then extracted
    from the result to avoid edge artefacts near the start and end of the
    generated buffer.

    Args:
        model: Noise model used to compute seismic noise.
        npts: Number of samples in the output seismogram.
        delta: Sampling interval of the generated noise.
        begin_time: Begin time of the output seismogram.
        return_velocity: If `True`, integrate the acceleration spectrum
            (division by `iω` in the frequency domain) to return ground
            velocity instead.
        seed: Random seed for reproducibility (e.g. in tests).

    Raises:
        ValueError: If `npts` is not a positive integer.

    Returns:
        Seismogram containing the generated noise. Data represent ground
        acceleration (arbitrary units matching the noise model's PSD) unless
        `return_velocity=True`, in which case they represent ground velocity.

    Examples:
        ```python
        >>> import pandas as pd
        >>> from pysmo import MiniSeismogram
        >>> from pysmo.tools.noise import peterson, generate_noise
        >>> model = peterson(0.0)
        >>> noise = generate_noise(
        ...     model=model, npts=64, delta=pd.Timedelta(seconds=1.0), seed=42
        ... )
        >>> isinstance(noise, MiniSeismogram)
        True
        >>> len(noise.data)
        64
        >>> noise.data[:3]
        array([-4.20418403e-09, -1.25152943e-08, -8.42501771e-09])
        >>>
        ```
    """
    if npts < 1:
        raise ValueError(f"npts={npts} must be a positive integer.")

    dt = delta.total_seconds()

    # Next power of 2 of the number of points (and at least 2, so the FFT is
    # always well defined), for an efficient FFT.
    NPTS = max(int(2 ** np.ceil(np.log2(npts))), 2)

    # Frequencies from DC to Nyquist for a real signal of length NPTS.
    freqs = np.fft.rfftfreq(NPTS, d=dt)

    # Period corresponding to each frequency; DC has no period and is handled
    # separately below (amplitude forced to zero there).
    periods = np.empty_like(freqs)
    periods[0] = np.inf
    periods[1:] = 1 / freqs[1:]

    T = model.T.total_seconds()
    if periods[1:].min() < T.min() or periods[1:].max() > T.max():
        warnings.warn(
            "Requested frequencies extend beyond the noise model's tabulated "
            "period range; extrapolating with the nearest tabulated value.",
            stacklevel=2,
        )

    # Interpolate psd in log-period space: Peterson's model is a piecewise
    # power law, i.e. linear in log10(period), not in period itself.
    Pxx = np.interp(np.log10(periods), np.log10(T), model.psd)

    # Recreate the amplitude spectrum, with the DC term set to 0 (mean=0).
    # Pxx is a one-sided PSD (Peterson's convention, matching
    # pysmo.tools.signal.psd's scipy.signal.welch default), which doubles the
    # power of every interior bin to account for the folded negative
    # frequencies. Since only the positive-frequency half is synthesised here,
    # the target power must be halved to compensate, otherwise a one-sided
    # PSD estimate of the generated noise comes out 3 dB too high.
    amplitude = np.sqrt(10 ** (Pxx / 10) * NPTS / (2 * dt))
    amplitude[0] = 0.0

    # Phase is randomly generated; DC and Nyquist are forced real, as
    # required for a Hermitian-symmetric spectrum representing a real signal.
    rng = np.random.default_rng(seed)
    phase = (rng.random(freqs.size) - 0.5) * np.pi * 2
    phase[0] = 0.0
    phase[-1] = 0.0

    spectrum = amplitude * (np.cos(phase) + 1j * np.sin(phase))

    if return_velocity:
        # Integrate acceleration to velocity in the frequency domain
        # (division by iω), rather than the time domain: a cumulative
        # time-domain integrator is only marginally stable and introduces
        # unbounded low-frequency drift.
        omega = 2 * np.pi * freqs
        velocity_spectrum = np.zeros_like(spectrum)
        velocity_spectrum[1:] = spectrum[1:] / (1j * omega[1:])
        spectrum = velocity_spectrum

    result = np.fft.irfft(spectrum, n=NPTS)

    start = (NPTS - npts) // 2
    end = start + npts
    return MiniSeismogram(begin_time=begin_time, delta=delta, data=result[start:end])

def peterson(noise_level: float) -> NoiseModel:
    """Generate a noise model by interpolating between Peterson's[^1]
    New Low Noise Model (NLNM) and New High Noise Model (NHNM).

    [^1]: Peterson, Jon R. Observations and Modeling of Seismic Background
        Noise. Report, 93–322, 1993, https://doi.org/10.3133/ofr93322. USGS
        Publications Warehouse.

    Args:
        noise_level: Determines the noise level of the generated noise model.
            A noise level of 0 returns the NLNM, 1 returns the NHNM, and
            anything > 0 and < 1 returns an interpolated model that lies
            between the NLNM and NHNM.

    Returns:
        Noise model.

    Examples:
        ```python
        >>> from pysmo.tools.noise import peterson, NLNM, NHNM
        >>> peterson(0.0) == NLNM
        True
        >>> peterson(1.0) == NHNM
        True
        >>> model = peterson(0.5)
        >>> model.psd[0]  # midpoint of NLNM and NHNM at the shortest period
        np.float64(-129.75)
        >>>
        ```
    """
    # check for valid input
    if not 0 <= noise_level <= 1:
        raise ValueError(
            f"Parameter noise_level={noise_level} is not within 0-1 range."
        )

    # calculate noise model
    if noise_level == 0:
        return NLNM
    elif noise_level == 1:
        return NHNM
    else:
        T_common = np.unique(
            np.concatenate((NLNM.T.total_seconds(), NHNM.T.total_seconds()))
        )
        # Peterson's model is a piecewise power law, i.e. linear in
        # log10(period), not in period itself, so interpolate in log-period
        # space to stay close to the true curve between tabulated points.
        log_T_common = np.log10(T_common)
        NLNM_interp = np.interp(
            log_T_common, np.log10(NLNM.T.total_seconds()), NLNM.psd
        )
        NHNM_interp = np.interp(
            log_T_common, np.log10(NHNM.T.total_seconds()), NHNM.psd
        )
        dB = NLNM_interp + (NHNM_interp - NLNM_interp) * noise_level
        return NoiseModel(psd=dB, T=pd.to_timedelta(T_common, unit="s"))