# /// 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"""
        # Two climates under the same Sun

        The simplest climate model balances sunlight absorbed against heat
        radiated to space, for the planet as a whole:

        $$C\,\frac{dT}{dt} = \frac{S}{4}\bigl(1 - \alpha(T)\bigr) - \varepsilon\,\sigma T^4$$

        $S$ is the solar constant, $\alpha$ the fraction of sunlight reflected
        (the albedo), and $\varepsilon\sigma T^4$ the infrared emission, with an
        effective emissivity $\varepsilon$ that stands in for the greenhouse
        effect. Ice reflects more sunlight than ocean and land, so the albedo
        rises as the planet cools. That one feedback, studied by Budyko and by
        Sellers in 1969, gives this model two stable climates at today's
        sunlight: the one we live in, and a frozen one.
        """
    )
    return


@app.cell
def _():
    import marimo as mo
    import numpy as np
    import matplotlib.pyplot as plt
    return mo, np, plt


@app.cell
def _(np):
    SIGMA = 5.670e-8  # Stefan-Boltzmann constant, W m^-2 K^-4
    S_TODAY = 1361.0  # solar constant, W m^-2
    EMISSIVITY = 0.61  # effective, chosen so today's climate sits near 288 K
    HEAT_CAPACITY = 4.0e8  # J m^-2 K^-1, about a 100 m ocean mixed layer
    YEAR = 3.156e7

    def albedo(T, ice=0.7, warm=0.3):
        """Albedo falls smoothly from an ice-covered to an ice-free planet near 265 K."""
        return ice + (warm - ice) * (1 + np.tanh((T - 265.0) / 10.0)) / 2

    def net_heating(T, solar):
        """Absorbed minus emitted flux, W m^-2, at a solar constant `solar` * S_TODAY."""
        return solar * S_TODAY / 4 * (1 - albedo(T)) - EMISSIVITY * SIGMA * T**4

    def equilibria(solar):
        """Equilibrium temperatures and whether each one is stable."""
        T = np.linspace(150, 360, 4200)
        f = net_heating(T, solar)
        crossing = np.flatnonzero(np.sign(f[:-1]) != np.sign(f[1:]))
        roots = T[crossing] - f[crossing] * (T[1] - T[0]) / (f[crossing + 1] - f[crossing])
        stable = f[crossing + 1] < f[crossing]
        return roots, stable

    return HEAT_CAPACITY, S_TODAY, YEAR, albedo, equilibria, net_heating


@app.cell
def _(mo):
    solar = mo.ui.slider(0.70, 1.60, step=0.01, value=1.00, label="Sunlight (fraction of today's)")
    solar
    return (solar,)


@app.cell
def _(S_TODAY, albedo, equilibria, mo, np, plt, solar):
    T = np.linspace(180, 340, 400)
    roots, stable = equilibria(solar.value)

    fig, ax = plt.subplots(figsize=(7, 3.8))
    ax.plot(T, solar.value * S_TODAY / 4 * (1 - albedo(T)), color="#9a6700", lw=2, label="absorbed sunlight")
    ax.plot(T, 0.61 * 5.670e-8 * T**4, color="#0969da", lw=2, label="emitted infrared")
    for _root, _ok in zip(roots, stable):
        ax.plot(
            _root,
            0.61 * 5.670e-8 * _root**4,
            "o",
            ms=8,
            color="#1f2328" if _ok else "white",
            markeredgecolor="#1f2328",
        )
    ax.set(xlabel="global mean temperature (K)", ylabel="flux (W m⁻²)")
    ax.legend(frameon=False, loc="upper left")
    fig.tight_layout()

    described = ", ".join(
        f"{_root:.0f} K ({_root - 273.15:+.0f} °C, {'stable' if _ok else 'unstable'})"
        for _root, _ok in zip(roots, stable)
    )
    mo.vstack([fig, mo.md(f"Equilibria at {solar.value:.2f} × today's sunlight: {described}.")])
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        Where the curves cross, the planet neither warms nor cools. A filled
        point is stable: nudge the temperature and it returns. The open point is
        unstable and separates the two basins. Lower the sunlight until the warm
        crossing disappears; the planet then has nowhere to go but the frozen
        state.

        ## The hysteresis loop

        Follow the equilibria as the Sun slowly dims and then brightens again.
        A warm planet stays warm until the warm branch ends, then drops to a
        snowball. Once frozen it stays frozen until the cold branch ends, at a
        much brighter Sun. This is the Snowball Earth puzzle: how the planet
        froze over during the Neoproterozoic, and how it ever thawed. The
        accepted answer is volcanic carbon dioxide accumulating for millions of
        years under the ice, which this model represents only through
        $\varepsilon$.
        """
    )
    return


@app.cell
def _(equilibria, np, plt, solar):
    fractions = np.linspace(0.7, 1.6, 361)
    points = [(s, r, ok) for s in fractions for r, ok in zip(*equilibria(s))]
    s_all, r_all, ok_all = map(np.array, zip(*points))
    warm_end = min(s for s, r, ok in points if ok and r > 260)
    cold_end = max(s for s, r, ok in points if ok and r < 260)

    fig2, ax2 = plt.subplots(figsize=(7, 3.8))
    ax2.scatter(s_all[ok_all], r_all[ok_all], s=3, color="#1f2328", label="stable")
    ax2.scatter(s_all[~ok_all], r_all[~ok_all], s=3, color="#8c959f", label="unstable")
    ax2.annotate("", xy=(warm_end, 225), xytext=(warm_end, 270), arrowprops={"arrowstyle": "->", "color": "#0969da"})
    ax2.annotate("", xy=(cold_end, 300), xytext=(cold_end, 255), arrowprops={"arrowstyle": "->", "color": "#cf222e"})
    ax2.axvline(solar.value, color="#9a6700", lw=0.8)
    ax2.set(xlabel="sunlight (fraction of today's)", ylabel="equilibrium temperature (K)")
    ax2.legend(frameon=False, loc="upper left")
    fig2.tight_layout()
    fig2
    return cold_end, warm_end


@app.cell
def _(cold_end, mo, warm_end):
    mo.md(
        f"""
        In this model the warm climate ends at **{warm_end:.2f}** times today's
        sunlight and the frozen one only at **{cold_end:.2f}** times. Between
        the two, which climate the planet has depends on where it came from.
        The Sun was about 70 percent as bright four billion years ago, when
        geology shows liquid water: the faint young Sun paradox. A stronger
        greenhouse effect then is the usual resolution. Lower the emissivity
        in your fork to try it.
        """
    )
    return


@app.cell
def _(mo):
    start_T = mo.ui.slider(200, 320, step=1, value=262, label="Starting temperature (K)")
    start_T
    return (start_T,)


@app.cell
def _(HEAT_CAPACITY, YEAR, net_heating, np, plt, solar, start_T):
    dt = 0.05 * YEAR
    path = [float(start_T.value)]
    for _ in range(1200):
        path.append(path[-1] + dt * net_heating(path[-1], solar.value) / HEAT_CAPACITY)
    years = np.arange(len(path)) * dt / YEAR

    fig3, ax3 = plt.subplots(figsize=(7, 3))
    ax3.plot(years, path, color="#0969da", lw=2)
    ax3.set(xlabel="years", ylabel="temperature (K)")
    ax3.set_title(f"From {start_T.value} K at {solar.value:.2f} × today's sunlight: {path[-1]:.0f} K after {years[-1]:.0f} years")
    fig3.tight_layout()
    fig3
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        ---
        **Model and limits.** Zero-dimensional energy balance: one temperature
        for the whole planet, no latitude, no ocean circulation, no carbon
        cycle. Albedo 0.7 for an ice-covered and 0.3 for an ice-free planet,
        joined smoothly around 265 K; effective emissivity 0.61; heat capacity
        of a 100 m ocean layer, which sets only the time scale. The thresholds
        depend on these choices, the existence of two branches does not.
        **References.** M. I. Budyko, *Tellus* 21 (1969) 611;
        W. D. Sellers, *J. Appl. Meteor.* 8 (1969) 392;
        P. F. Hoffman et al., *Science* 281 (1998) 1342.
        Fork it to add latitude bands, as Budyko and Sellers did.
        """
    )
    return


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