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.

    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 via the
            trapezoidal rule to return ground velocity instead.
        seed: Random seed for reproducibility (e.g. in tests).

    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.
    """
    # Sampling frequency
    Fs = 1 / delta.total_seconds()

    # Nyquist frequency
    Fnyq = 0.5 / delta.total_seconds()

    # get next power of 2 of the nunmber of points and calculate frequencies from
    # Fs/NPTS to Fnyq (we skip a frequency of 0 for now to avoid dividing by 0)
    NPTS = int(2 ** np.ceil(np.log2(npts)))
    freqs = np.linspace(Fs / NPTS, Fnyq, NPTS - 1)

    # interpolate psd and recreate amplitude spectrum with the first
    # term=0 (i.e. mean=0).
    Pxx = np.interp(1 / freqs, model.T.total_seconds(), model.psd)
    spectrum = np.append(
        np.array([0]), np.sqrt(10 ** (Pxx / 10) * NPTS / delta.total_seconds() * 2)
    )

    # phase is randomly generated
    rng = np.random.default_rng(seed)
    phase = (rng.random(NPTS) - 0.5) * np.pi * 2

    # combine amplitude with phase and perform ifft
    NewX = spectrum * (np.cos(phase) + 1j * np.sin(phase))
    acceleration = np.fft.irfft(NewX)

    start = int((NPTS - npts) / 2)
    end = start + npts
    if return_velocity:
        velocity = cumulative_trapezoid(acceleration, dx=delta.total_seconds())
        velocity = velocity[start:end]
        return MiniSeismogram(begin_time=begin_time, delta=delta, data=velocity)
    acceleration = acceleration[start:end]
    return MiniSeismogram(begin_time=begin_time, delta=delta, data=acceleration)

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.
    """
    # 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()))
        )
        NLNM_interp = np.interp(T_common, NLNM.T.total_seconds(), NLNM.psd)
        NHNM_interp = np.interp(T_common, 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"))