gwsounds.py

import marimo

__generated_with = "0.15.2"
app = marimo.App(width="full")


@app.cell
def _():
    import gwsurrogate
    import sxs
    import numpy as np
    from scipy.io import wavfile
    from pathlib import Path


    # Useful constants
    π = np.pi
    c = float(299_792_458)   # meters/second
    GM_sun = 1.32712440041e20  # meters³/second²
    au = float(149_597_870_700)  # meters
    pc = 1*au / (π / 648_000)  # meters
    Mpc = 1_000_000*pc  # meters
    sample_rate = 44100  # Hz  (CD-quality audio; certainly overkill)
    return GM_sun, Path, c, gwsurrogate, np, sample_rate, sxs, wavfile, π


@app.cell
def _(gwsurrogate):
    ## There are two good options for waveform models.  The first allows for precessing waveforms; the second provides longer waveforms but only nonprecessing.
    #model = "NRSur7dq4"
    model = "NRHybSur2dq15"

    # Load the surrogate model (downloading if necessary)
    sur = gwsurrogate.LoadSurrogate(model)
    return (sur,)


@app.cell
def _():
    # These two parameters basically control how long the signal is
    f_min = 60  # Hz
    Mtot = 1  # solar masses
    return Mtot, f_min


@app.cell
def _():
    # Choose physical parameters for the surrogate model.
    # # Choose precessing parameters (spin components in x and y directions) only if using NRSur7dq4
    # q = 4                           # mass ratio, mA/mB >= 1.
    # chiA = [-0.2, 0.4, 0.1]         # Dimensionless spin of heavier BH
    # chiB = [-0.5, 0.2, -0.4]        # Dimensionless of lighter BH
    # Components like the following are nonprecessing and can be used for either surrogate model:
    q = 4                           # mass ratio, mA/mB >= 1.
    chiA = [0., 0., 0.4]         # Dimensionless spin of heavier BH
    chiB = [0., 0., 0.]        # Dimensionless of lighter BH
    return chiA, chiB, q


@app.cell
def _(GM_sun, Mtot, c, f_min, q, sample_rate):
    # Some calculations needed to scale things properly

    M1 = Mtot * q/(1+q)*GM_sun / c**3  # seconds
    M2 = Mtot * 1/(1+q)*GM_sun / c**3  # seconds
    M = M1 + M2  # seconds
    # dL = 100*Mpc / c  # seconds
    # z = 0.094

    # True sampling rate
    Δtₒ = 1/sample_rate  # seconds

    # Sampling rate in units of total mass M (for surrogate evaluation)
    dt = Δtₒ / M

    # Initial frequency, f_low=0 returns the full surrogate
    f_low = f_min * M  # dimensionless frequency
    return M, dt, f_low, Δto


@app.cell
def _(M, chiA, chiB, dt, f_low, np, q, sur, sxs):
    # Now, we actually evaluate the signal

    # h is dictionary of spin-weighted spherical harmonic modes
    # t is the corresponding time array in units of M
    # dyn stands for dynamics, do dyn.keys() to see contents
    t, h, _ = sur(q, chiA, chiB, dt=dt, f_low=f_low)

    ℓm = h.keys()
    ℓ_min = min(ℓ for ℓ, m in ℓm)
    ℓ_max = max(ℓ for ℓ, m in ℓm)

    data = np.array([
        h[(ℓ, m)] if (ℓ, m) in h else np.zeros_like(h[(2,2)])  # zero-pad missing modes
        for ℓ in range(ℓ_min, ℓ_max + 1)
        for m in range(-ℓ, ℓ + 1)
    ]).T

    w = sxs.WaveformModes(
        data, t*M,
        time_axis=0, modes_axis=1, ell_min=ℓ_min, ell_max=ℓ_max,
        data_type="h", spin_weight=-2,
    )

    ω0 = np.linalg.norm(w.angular_velocity, axis=1)[1]

    # Evaluate the waveform in a particular direction
    θ, ϕ = 0.1, 0.2
    h = w.evaluate(θ, ϕ)
    return h, ω0


@app.cell
def _(Mtot, Path, h, np, sample_rate, wavfile, Δto, π, ω0):
    # Convert to audio and export as a standard WAV file.

    audio = np.asarray(h.real, dtype=np.float64)
    audio = np.ravel(audio)

    # Fast turn-on window over a few initial orbits using ω0; eliminates startup clicking sound
    n_orbits_window = 3.0
    T_orbit0 = 2 * π / ω0
    taper_duration = n_orbits_window * T_orbit0
    n_taper = min(audio.size, max(1, int(np.ceil(taper_duration / Δtₒ))))
    turn_on = np.ones(audio.size, dtype=np.float64)
    turn_on[:n_taper] = np.sin(np.linspace(0.0, 0.5 * π, n_taper))**2
    audio *= turn_on

    if audio.size == 0:
        raise ValueError("`h` is empty; cannot write audio file.")

    audio_f32 = audio.astype(np.float32)
    output_path = Path(__file__).with_name(f"bbh_waveform_{Mtot}Msun.wav")
    wavfile.write(str(output_path), sample_rate, audio_f32)
    print(f"Wrote {output_path} at {sample_rate} Hz (32-bit float WAV)")
    return


if __name__ == "__main__":
    app.run()