# /// script
# requires-python = ">=3.12"
# dependencies = ["marimo>=0.25.1", "numpy", "matplotlib"]
# ///

import marimo

__generated_with = "0.25.1"
app = marimo.App(width="medium")


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        # Listen to two black holes merge

        On 14 September 2015 the two LIGO detectors recorded a signal that
        rose in pitch and loudness for a fifth of a second and stopped: two
        black holes, about 36 and 29 times the mass of the Sun, spiralling
        into each other 1.3 billion light years away. It was the first direct
        detection of gravitational waves.

        Long before the merger, an orbiting pair loses energy to gravitational
        waves at a rate set by one combination of the masses, the chirp mass
        $\mathcal{M} = (m_1 m_2)^{3/5} / (m_1 + m_2)^{1/5}$. To leading order
        the wave's frequency $\tau$ seconds before the merger is

        $$f(\tau) = \frac{1}{\pi}\left(\frac{5}{256\,\tau}\right)^{3/8}
        \left(\frac{G\mathcal{M}}{c^3}\right)^{-5/8}.$$

        Choose the masses and the distance, look at the wave, and play it.
        """
    )
    return


@app.cell
def _():
    import io
    import wave

    import marimo as mo
    import numpy as np
    import matplotlib.pyplot as plt
    return io, mo, np, plt, wave


@app.cell
def _(mo):
    m1 = mo.ui.slider(1.4, 80, step=0.1, value=36, label="Mass 1 (Suns)")
    m2 = mo.ui.slider(1.4, 80, step=0.1, value=29, label="Mass 2 (Suns)")
    distance = mo.ui.slider(40, 3000, step=10, value=410, label="Distance (megaparsecs)")
    mo.vstack([m1, m2, distance])
    return distance, m1, m2


@app.cell
def _(np):
    G, C = 6.674e-11, 2.998e8
    M_SUN, MPC = 1.989e30, 3.086e22

    def chirp(m1, m2, distance_mpc, f_low=20.0, rate=4096):
        """Leading-order inspiral from f_low up to the innermost stable orbit.

        Returns time (s, zero at the end), strain h(t) for a face-on source,
        frequency f(t), and the stopping frequency.
        """
        total = (m1 + m2) * M_SUN
        mc = (m1 * m2) ** 0.6 / (m1 + m2) ** 0.2 * M_SUN
        tc = G * mc / C**3
        f_isco = C**3 / (6**1.5 * np.pi * G * total)
        tau_low = 5 / 256 * tc ** (-5 / 3) * (np.pi * f_low) ** (-8 / 3)
        tau_end = 5 / 256 * tc ** (-5 / 3) * (np.pi * f_isco) ** (-8 / 3)
        tau = np.arange(tau_low, tau_end, -1 / rate)
        f = (5 / (256 * tau)) ** 0.375 * tc ** (-0.625) / np.pi
        phase = -2 * (5 * tc) ** (-0.625) * tau**0.625
        amplitude = 4 / (distance_mpc * MPC) * (G * mc / C**2) ** (5 / 3) * (np.pi * f / C) ** (2 / 3)
        return -(tau - tau_end), amplitude * np.cos(phase), f, f_isco

    return (chirp,)


@app.cell
def _(chirp, distance, m1, m2, mo, np, plt):
    t, h, f, f_isco = chirp(m1.value, m2.value, distance.value)
    cycles = np.sum(f) / 4096
    last = t > -min(0.5, -t.min())

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(7.2, 5), sharex=True)
    ax1.plot(t[last], h[last] * 1e21, color="#0072B2", lw=1)
    ax1.set_ylabel("strain h (10⁻²¹)")
    ax2.plot(t[last], f[last], color="#D55E00", lw=2)
    ax2.set(xlabel="seconds before the end of the inspiral", ylabel="frequency (Hz)")
    fig.tight_layout()

    mo.vstack(
        [
            fig,
            mo.md(
                f"From 20 Hz, where LIGO starts to hear it, the pair has **{-t.min():.2f} s** and "
                f"**{cycles:,.0f} wave cycles** left. The inspiral ends near **{f_isco:.0f} Hz**, "
                f"where the orbit can no longer stay stable. The peak strain here is "
                f"**{np.abs(h).max():.1e}**: a 4 km detector arm changes length by "
                f"{np.abs(h).max() * 4000:.1e} m, about 1/{1.7e-15 / (np.abs(h).max() * 4000):,.0f} "
                "of the width of a proton."
            ),
        ]
    )
    return h, t


@app.cell
def _(mo):
    speed = mo.ui.dropdown({"real time": 1, "4 × faster, two octaves higher": 4}, value="real time", label="Playback")
    speed
    return (speed,)


@app.cell
def _(h, io, mo, np, speed, wave):
    rate = 4096 * speed.value
    pcm = (h / np.abs(h).max() * 0.9 * 32767).astype(np.int16)
    pcm = np.concatenate([np.zeros(rate // 4, np.int16), pcm, np.zeros(rate // 4, np.int16)])
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as out:
        out.setnchannels(1)
        out.setsampwidth(2)
        out.setframerate(rate)
        out.writeframes(pcm.tobytes())
    mo.vstack(
        [
            mo.audio(buffer.getvalue()),
            mo.md(
                "The sound is the strain itself, scaled to full volume. Heavy pairs chirp "
                "low, below what laptop speakers play well; try the faster playback or "
                "two neutron stars (1.4 and 1.4 Suns), whose chirp lasts much longer."
            ),
        ]
    )
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        ---
        **Model and limits.** Leading (quadrupole, Newtonian) order in the
        orbital speed, circular orbits, a face-on source, and no detector
        response. The waveform stops at the innermost stable circular orbit;
        the merger and ringdown that follow need numerical relativity, and the
        real GW150914 signal peaked near 150 Hz. Masses here are as measured
        at the source, without the redshift correction.
        **References.** B. P. Abbott et al. (LIGO and Virgo), *Phys. Rev. Lett.*
        116 (2016) 061102; P. C. Peters and J. Mathews, *Phys. Rev.* 131 (1963)
        435; M. Maggiore, *Gravitational Waves*, vol. 1 (2008), ch. 4.
        Fork it to add the next post-Newtonian term or a detector's noise curve.
        """
    )
    return


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