physicskit 0.1.0__py3-none-any.whl

physicskit/astro/stellar_dynamo.py

@@ -0,0 +1,628 @@
+r"""Stellar convection and the alpha-omega mean-field magnetic dynamo.
+
+The Sun's large-scale magnetic field -- the one that reverses polarity
+every ~11 years and paints the "butterfly diagram" of sunspots migrating
+from mid-latitudes toward the equator -- is not a fossil field left over
+from formation. It is continuously regenerated by the interplay of two
+processes going on simultaneously in the convection zone: small-scale,
+helical turbulent convection (the "alpha effect", parametrizing the net
+effect of many individually chaotic convective eddies twisted by the
+Sun's rotation and stratification) and the large-scale rotational shear
+between the Sun's differentially rotating layers (the "Omega effect",
+which stretches a poloidal field into a much stronger toroidal one).
+Together these sustain the field against Ohmic decay -- this is the
+alpha-omega dynamo mechanism first proposed by Parker (1955), and it is
+the standard textbook explanation for the solar cycle.
+
+A crucial and easy-to-miss subtlety (Zeldovich's antidynamo theorem) is
+that a strictly two-dimensional, fully resolved magnetic field confined
+to a plane *cannot* be sustained by 2D fluid motion alone -- genuine
+dynamo action is an intrinsically three-dimensional effect, arising from
+the helical (non-mirror-symmetric) twisting that only a 3D, rotating,
+stratified flow can produce. This module therefore does not attempt to
+build one dynamically self-consistent 2D "dynamo simulation" (which
+would either not be a real dynamo, or would silently smuggle in 3D
+physics through the back door). Instead, exactly as the textbook
+treatment splits the problem, it implements two separate,
+self-contained pieces:
+
+- A resolved 2D Boussinesq convection simulation
+  (:func:`simulate_stellar_convection`) showing the turbulent
+  convective rolls and plumes whose unresolved, helical, three-dimensional
+  fine structure is -- in the real Sun -- the physical origin of the
+  alpha effect. This piece is a qualitative picture of the turbulence,
+  not a dynamo; by Zeldovich's theorem it could never sustain a resolved
+  2D field, and it is not asked to.
+- The linearized 1D alpha-omega **mean-field** dynamo equations
+  (:func:`simulate_alpha_omega_dynamo`), in which the alpha effect and
+  the rotational shear enter as prescribed coefficients rather than
+  being derived from the convection simulation above. This piece *does*
+  robustly sustain and grow a large-scale field, and reproduces the
+  real solar dynamo's most famous piece of phenomenology: a migrating
+  dynamo wave, whose space-time diagram is the butterfly diagram.
+
+The two pieces are deliberately not coupled to each other; each is
+correct and complete on its own terms.
+
+Part 1: 2D convective rolls
+----------------------------
+
+A standard doubly periodic, pseudo-spectral Boussinesq system for
+vorticity :math:`\omega`, streamfunction :math:`\psi`, and temperature
+perturbation :math:`T'` about an unstably stratified periodic background
+:math:`T_{eq}(y) = -\beta\cos(2\pi y/L_y)`:
+
+.. math::
+
+    \partial_t \omega + \mathbf{u}\cdot\nabla\omega
+        = \nu\nabla^2\omega + g\alpha_{th}\,\partial_x T' \\
+    \partial_t T' + \mathbf{u}\cdot\nabla T'
+        = \kappa\nabla^2 T' - u_y\,\frac{dT_{eq}}{dy}
+          + \frac{T_{eq}(y) - T'}{\tau_{relax}} \\
+    \nabla^2\psi = -\omega, \qquad \mathbf{u} = (\partial_y\psi, -\partial_x\psi)
+
+The Newtonian-cooling term :math:`(T_{eq}-T')/\tau_{relax}` (together
+with the advection of the background gradient, :math:`-u_y\,dT_{eq}/dy`)
+is the standard trick for sustaining an unstable stratification in a
+periodic box with no physical walls -- without it, convection would mix
+the background profile flat and decay back to a quiescent state.
+
+Part 2: The alpha-omega mean-field dynamo
+-------------------------------------------
+
+A linearized 1D reduction (Parker 1955; standard in mean-field dynamo
+theory) for the poloidal-flux proxy :math:`A(x,t)` and toroidal field
+:math:`B(x,t)`, periodic in :math:`x` (a local-Cartesian stand-in for
+stellar latitude or radius):
+
+.. math::
+
+    \partial_t A = \alpha B + \eta\,\partial_x^2 A \\
+    \partial_t B = G\,\partial_x A + \eta\,\partial_x^2 B
+
+with constant alpha-effect coefficient :math:`\alpha`, shear rate
+(Omega effect) :math:`G`, and magnetic diffusivity :math:`\eta`.
+Substituting a plane-wave mode :math:`(A,B)\propto e^{ikx-i\omega_f t+\sigma t}`
+turns this linear system into the 2x2 eigenvalue problem
+
+.. math::
+
+    \begin{pmatrix}\sigma-i\omega_f+\eta k^2 & -\alpha \\ -iGk & \sigma-i\omega_f+\eta k^2\end{pmatrix}
+    \begin{pmatrix}A\\B\end{pmatrix} = 0,
+
+whose characteristic equation, with :math:`s=\sigma-i\omega_f+\eta k^2`,
+is :math:`s^2 = i\alpha G k`. Since :math:`i\alpha Gk` is purely
+imaginary, its principal square root always has real and imaginary parts
+of equal magnitude, :math:`\sqrt{|\alpha Gk|/2}` -- giving the closed-form
+growth rate and wave frequency implemented in
+:func:`alpha_omega_growth_rate` and :func:`alpha_omega_wave_frequency`,
+and checked against the full simulation in this module's test suite.
+Because the system is linear, each Fourier mode's amplitude can be
+advanced *exactly* (no time-discretization error at all) via the 2x2
+matrix exponential of its own coefficient matrix -- see
+:func:`simulate_alpha_omega_dynamo`.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+__all__ = [
+    "convective_roll_rhs",
+    "simulate_stellar_convection",
+    "kinetic_energy",
+    "alpha_omega_growth_rate",
+    "alpha_omega_wave_frequency",
+    "simulate_alpha_omega_dynamo",
+    "dominant_mode_growth_rate",
+]
+
+
+# ---------------------------------------------------------------------------
+# Part 1: 2D convective rolls (Boussinesq vorticity-streamfunction-temperature)
+# ---------------------------------------------------------------------------
+
+
+def _spectral_grid_2d(nx: int, ny: int, Lx: float, Ly: float):
+    """Build the doubly periodic real-space and wavenumber grids for Part 1.
+
+    Fields are laid out with shape ``(ny, nx)`` -- rows indexed by ``y``,
+    columns by ``x`` -- the usual image-like convention (matches
+    ``imshow``'s default axes and this module's documented snapshot
+    shape). Self-contained (not imported from :mod:`physicskit.fluids`),
+    per this package's no-cross-subpackage-imports convention, but built
+    the same way as :func:`physicskit.fluids.core.grid.spectral_grid`:
+    ``np.fft.fftfreq``-based wavenumbers, and the ``k=0`` mode of ``K2``
+    set to 1 so the Poisson solve is safe to divide by.
+    """
+    x = np.linspace(0.0, Lx, nx, endpoint=False)
+    y = np.linspace(0.0, Ly, ny, endpoint=False)
+    X, Y = np.meshgrid(x, y, indexing="xy")
+    kx = 2.0 * np.pi * np.fft.fftfreq(nx, d=Lx / nx)
+    ky = 2.0 * np.pi * np.fft.fftfreq(ny, d=Ly / ny)
+    KX, KY = np.meshgrid(kx, ky, indexing="xy")
+    K2 = KX**2 + KY**2
+    K2 = K2.copy()
+    K2[0, 0] = 1.0
+    return X, Y, KX, KY, K2
+
+
+def _poisson_solve_streamfunction(omega: NDArray[np.float64], K2: NDArray[np.float64]) -> NDArray[np.float64]:
+    """Solve :math:`\\nabla^2\\psi=-\\omega` for the streamfunction via FFT (see :func:`_spectral_grid_2d`)."""
+    omega_hat = np.fft.fft2(omega)
+    psi_hat = omega_hat / K2
+    psi_hat[0, 0] = 0.0
+    return np.real(np.fft.ifft2(psi_hat))
+
+
+def _velocity_from_streamfunction(psi: NDArray[np.float64], KX: NDArray[np.float64], KY: NDArray[np.float64]):
+    """Recover :math:`\\mathbf{u}=(\\partial_y\\psi,-\\partial_x\\psi)` from the streamfunction spectrally."""
+    psi_hat = np.fft.fft2(psi)
+    u = np.real(np.fft.ifft2(1j * KY * psi_hat))
+    v = np.real(np.fft.ifft2(-1j * KX * psi_hat))
+    return u, v
+
+
+def convective_roll_rhs(
+    omega: NDArray[np.float64],
+    T: NDArray[np.float64],
+    Y: NDArray[np.float64],
+    KX: NDArray[np.float64],
+    KY: NDArray[np.float64],
+    K2: NDArray[np.float64],
+    nu: float,
+    kappa: float,
+    g_alpha: float,
+    beta: float,
+    Ly: float,
+    tau_relax: float,
+):
+    """Right-hand side of the 2D Boussinesq convective-roll system, evaluated pseudo-spectrally.
+
+    Mirrors the structure of the Rayleigh-Taylor Boussinesq right-hand
+    side used elsewhere in this codebase for a buoyancy-vorticity system
+    (vorticity forced by the buoyancy field's :math:`x`-derivative,
+    :math:`g\\,\\partial_x b`, since gravity and the background
+    stratification both point along :math:`y`), plus the Newtonian-cooling
+    forcing that sustains the unstable background profile
+    :math:`T_{eq}(y)=-\\beta\\cos(2\\pi y/L_y)` against being mixed away by
+    the convection it drives (see the module docstring).
+
+    Parameters
+    ----------
+    omega, T : ndarray of float, shape (ny, nx)
+        Vorticity and temperature-perturbation fields.
+    Y : ndarray of float, shape (ny, nx)
+        The ``y`` coordinate grid, from :func:`_spectral_grid_2d`, used to
+        evaluate the background profile :math:`T_{eq}(y)` and its
+        gradient pointwise.
+    KX, KY, K2 : ndarray of float, shape (ny, nx)
+        Wavenumber grids from :func:`_spectral_grid_2d`.
+    nu : float
+        Kinematic viscosity (vorticity diffusivity).
+    kappa : float
+        Thermal diffusivity.
+    g_alpha : float
+        Combined buoyancy coefficient :math:`g\\,\\alpha_{th}`.
+    beta : float
+        Amplitude of the background stratification :math:`T_{eq}(y)`.
+    Ly : float
+        Domain size along :math:`y` (the period of :math:`T_{eq}`).
+    tau_relax : float
+        Newtonian-cooling relaxation time back to :math:`T_{eq}(y)`.
+
+    Returns
+    -------
+    domega_dt, dT_dt : ndarray of float, shape (ny, nx)
+        Time derivatives of vorticity and temperature perturbation.
+    """
+    psi = _poisson_solve_streamfunction(omega, K2)
+    u, v = _velocity_from_streamfunction(psi, KX, KY)
+
+    omega_hat = np.fft.fft2(omega)
+    domega_dx = np.real(np.fft.ifft2(1j * KX * omega_hat))
+    domega_dy = np.real(np.fft.ifft2(1j * KY * omega_hat))
+    diffusion_omega = np.real(np.fft.ifft2(-nu * K2 * omega_hat))
+
+    T_hat = np.fft.fft2(T)
+    dT_dx = np.real(np.fft.ifft2(1j * KX * T_hat))
+    dT_dy = np.real(np.fft.ifft2(1j * KY * T_hat))
+    diffusion_T = np.real(np.fft.ifft2(-kappa * K2 * T_hat))
+
+    Teq = -beta * np.cos(2.0 * np.pi * Y / Ly)
+    dTeq_dy = beta * (2.0 * np.pi / Ly) * np.sin(2.0 * np.pi * Y / Ly)
+
+    domega_dt = -(u * domega_dx + v * domega_dy) + diffusion_omega + g_alpha * dT_dx
+    dT_dt = -(u * dT_dx + v * dT_dy) + diffusion_T - v * dTeq_dy + (Teq - T) / tau_relax
+    return domega_dt, dT_dt
+
+
+def simulate_stellar_convection(
+    nx: int,
+    ny: int,
+    Lx: float,
+    Ly: float,
+    nu: float = 0.08,
+    kappa: float = 0.08,
+    g_alpha: float = 2.0,
+    beta: float = 1.0,
+    tau_relax: float = 0.5,
+    dt: float = 0.005,
+    n_steps: int = 5000,
+    save_every: int = 50,
+    seed: int | None = None,
+):
+    """Time-step 2D Boussinesq convective rolls with pseudo-spectral RK4.
+
+    Starts from a small random vorticity and temperature perturbation
+    (root-mean-square amplitude ``1e-3``) and integrates
+    :func:`convective_roll_rhs` with classical 4th-order Runge-Kutta,
+    exactly the scheme used for the Boussinesq Rayleigh-Taylor system
+    elsewhere in this codebase. The default parameters were tuned (see
+    this module's test suite) so that the periodic background
+    stratification :math:`T_{eq}(y)=-\\beta\\cos(2\\pi y/L_y)`, sustained
+    against mixing by the ``tau_relax`` Newtonian-cooling forcing, robustly
+    grows the seed noise into sustained convective rolls -- kinetic energy
+    measurably larger at the end of the run than at the start -- without
+    blowing up.
+
+    Parameters
+    ----------
+    nx, ny : int
+        Number of grid points along ``x`` and ``y``.
+    Lx, Ly : float
+        Physical domain size along ``x`` and ``y``.
+    nu : float, default 0.08
+        Kinematic viscosity.
+    kappa : float, default 0.08
+        Thermal diffusivity.
+    g_alpha : float, default 2.0
+        Combined buoyancy coefficient :math:`g\\,\\alpha_{th}`.
+    beta : float, default 1.0
+        Amplitude of the background stratification.
+    tau_relax : float, default 0.5
+        Newtonian-cooling relaxation time.
+    dt : float, default 0.005
+        Time step.
+    n_steps : int, default 5000
+        Total number of RK4 steps.
+    save_every : int, default 50
+        Save a snapshot every this many steps.
+    seed : int, optional
+        Random seed for the initial perturbation.
+
+    Returns
+    -------
+    times : ndarray of float, shape (n_saved,)
+    omega_snapshots : ndarray of float, shape (n_saved, ny, nx)
+    T_snapshots : ndarray of float, shape (n_saved, ny, nx)
+
+    See Also
+    --------
+    convective_roll_rhs : The right-hand side stepped here.
+    kinetic_energy : Diagnostic used to confirm convection is sustained.
+
+    Examples
+    --------
+    >>> times, omega_snaps, T_snaps = simulate_stellar_convection(
+    ...     32, 32, 2 * np.pi, 2 * np.pi, n_steps=20, save_every=10, seed=0)
+    >>> omega_snaps.shape[1:]
+    (32, 32)
+    >>> bool(np.all(np.isfinite(omega_snaps)))
+    True
+    """
+    X, Y, KX, KY, K2 = _spectral_grid_2d(nx, ny, Lx, Ly)
+
+    rng = np.random.default_rng(seed)
+    omega = 1.0e-3 * rng.standard_normal((ny, nx))
+    T = 1.0e-3 * rng.standard_normal((ny, nx))
+
+    times = [0.0]
+    omega_snapshots = [omega.copy()]
+    T_snapshots = [T.copy()]
+
+    t = 0.0
+    for step in range(1, n_steps + 1):
+        k1o, k1T = convective_roll_rhs(omega, T, Y, KX, KY, K2, nu, kappa, g_alpha, beta, Ly, tau_relax)
+        k2o, k2T = convective_roll_rhs(omega + dt / 2 * k1o, T + dt / 2 * k1T, Y, KX, KY, K2, nu, kappa, g_alpha, beta, Ly, tau_relax)
+        k3o, k3T = convective_roll_rhs(omega + dt / 2 * k2o, T + dt / 2 * k2T, Y, KX, KY, K2, nu, kappa, g_alpha, beta, Ly, tau_relax)
+        k4o, k4T = convective_roll_rhs(omega + dt * k3o, T + dt * k3T, Y, KX, KY, K2, nu, kappa, g_alpha, beta, Ly, tau_relax)
+        omega = omega + dt / 6 * (k1o + 2 * k2o + 2 * k3o + k4o)
+        T = T + dt / 6 * (k1T + 2 * k2T + 2 * k3T + k4T)
+        t += dt
+        if step % save_every == 0:
+            times.append(t)
+            omega_snapshots.append(omega.copy())
+            T_snapshots.append(T.copy())
+
+    return np.array(times), np.array(omega_snapshots), np.array(T_snapshots)
+
+
+def kinetic_energy(omega_snapshot: NDArray[np.float64], KX: NDArray[np.float64], KY: NDArray[np.float64], K2: NDArray[np.float64]) -> float:
+    """Total kinetic energy of the flow implied by a vorticity snapshot.
+
+    Recovers the streamfunction and velocity field from `omega_snapshot`
+    via the same spectral Poisson solve used during time-stepping, then
+    sums :math:`\\frac{1}{2}(u^2+v^2)` over the grid. Used as the
+    diagnostic that convection in :func:`simulate_stellar_convection`
+    actually spins up (grows) rather than decaying back to quiescence.
+
+    Parameters
+    ----------
+    omega_snapshot : ndarray of float, shape (ny, nx)
+        A single vorticity snapshot, e.g. one entry of the
+        `omega_snapshots` array returned by :func:`simulate_stellar_convection`.
+    KX, KY, K2 : ndarray of float, shape (ny, nx)
+        Wavenumber grids from :func:`_spectral_grid_2d`.
+
+    Returns
+    -------
+    float
+        Total kinetic energy :math:`\\sum \\frac{1}{2}(u^2+v^2)`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> _, _, KX, KY, K2 = _spectral_grid_2d(32, 32, 2 * np.pi, 2 * np.pi)
+    >>> KE = kinetic_energy(np.zeros((32, 32)), KX, KY, K2)
+    >>> KE == 0.0
+    True
+    """
+    psi = _poisson_solve_streamfunction(omega_snapshot, K2)
+    u, v = _velocity_from_streamfunction(psi, KX, KY)
+    return float(0.5 * np.sum(u**2 + v**2))
+
+
+# ---------------------------------------------------------------------------
+# Part 2: alpha-omega mean-field dynamo
+# ---------------------------------------------------------------------------
+
+
+def alpha_omega_growth_rate(alpha: float, shear: float, k, eta: float):
+    """Closed-form linear growth rate of the alpha-omega dynamo wave.
+
+    From the dispersion relation :math:`s^2=i\\alpha G k` (see the module
+    docstring), the growing branch's real part is
+    :math:`\\sigma=\\sqrt{|\\alpha G k|/2}-\\eta k^2` -- positive
+    (dynamo action beats Ohmic decay) once :math:`|\\alpha Gk|` is large
+    enough relative to :math:`\\eta^2k^4`, i.e. at long enough wavelength
+    for fixed diffusivity.
+
+    Parameters
+    ----------
+    alpha : float
+        Alpha-effect coefficient.
+    shear : float
+        Rotational shear rate :math:`G` (the Omega effect).
+    k : float or ndarray
+        Wavenumber(s).
+    eta : float
+        Magnetic diffusivity.
+
+    Returns
+    -------
+    float or ndarray
+        Growth rate :math:`\\sigma`, same shape as `k`.
+
+    See Also
+    --------
+    alpha_omega_wave_frequency : The companion imaginary part (wave frequency magnitude).
+    simulate_alpha_omega_dynamo : Full simulation this closed form is checked against.
+
+    Examples
+    --------
+    >>> round(alpha_omega_growth_rate(alpha=1.0, shear=5.0, k=1.0, eta=0.05), 4)
+    1.5311
+    """
+    k = np.asarray(k, dtype=float)
+    sigma = np.sqrt(np.abs(alpha * shear * k) / 2.0) - eta * k**2
+    return sigma if sigma.ndim else float(sigma)
+
+
+def alpha_omega_wave_frequency(alpha: float, shear: float, k):
+    """Closed-form magnitude of the alpha-omega dynamo wave's angular frequency.
+
+    :math:`|\\omega_f| = \\sqrt{|\\alpha Gk|/2}` (see the module docstring
+    for the derivation). Only the magnitude is returned; the propagation
+    *direction* depends on the sign of :math:`\\alpha Gk` and this
+    module's specific stepping convention -- empirically, with
+    :func:`simulate_alpha_omega_dynamo`'s time-stepping and the
+    :math:`e^{ikx}` Fourier convention used throughout, the wave
+    (for ``alpha * shear > 0``, i.e. a positive-``k`` mode with
+    :math:`\\alpha Gk>0``) migrates toward **decreasing** :math:`x`, i.e.
+    it travels in the ``-x`` direction (verified against the simulation
+    in this module's test suite and visible in
+    :func:`physicskit.astro.visualizers.plot_dynamo_butterfly_diagram`'s
+    diagonal stripes).
+
+    Parameters
+    ----------
+    alpha : float
+        Alpha-effect coefficient.
+    shear : float
+        Rotational shear rate :math:`G`.
+    k : float or ndarray
+        Wavenumber(s).
+
+    Returns
+    -------
+    float or ndarray
+        :math:`|\\omega_f|`, same shape as `k`.
+
+    Examples
+    --------
+    >>> round(alpha_omega_wave_frequency(alpha=1.0, shear=5.0, k=1.0), 4)
+    1.5811
+    """
+    k = np.asarray(k, dtype=float)
+    freq = np.sqrt(np.abs(alpha * shear * k) / 2.0)
+    return freq if freq.ndim else float(freq)
+
+
+def simulate_alpha_omega_dynamo(
+    A0: NDArray[np.float64],
+    B0: NDArray[np.float64],
+    alpha: float,
+    shear: float,
+    eta: float,
+    Lx: float,
+    dt: float,
+    n_steps: int,
+    save_every: int = 1,
+):
+    """Time-step the linear 1D alpha-omega mean-field dynamo equations exactly, per Fourier mode.
+
+    Because :math:`\\partial_t A=\\alpha B+\\eta\\partial_x^2 A`,
+    :math:`\\partial_t B=G\\partial_x A+\\eta\\partial_x^2 B` is linear,
+    each Fourier mode's amplitude pair :math:`(\\hat A(k,t),\\hat B(k,t))`
+    evolves independently under a 2x2 constant-coefficient linear ODE,
+    :math:`d/dt\\,(\\hat A,\\hat B)^T = M(k)\\,(\\hat A,\\hat B)^T` with
+
+    .. math::
+
+        M(k) = \\begin{pmatrix}-\\eta k^2 & \\alpha \\\\ iGk & -\\eta k^2\\end{pmatrix},
+
+    so it can be advanced *exactly* over a step `dt` via the matrix
+    exponential :math:`\\exp(M(k)\\,dt)`, with no time-discretization error
+    whatsoever. Writing :math:`M(k)=-\\eta k^2 I + N(k)` with
+    :math:`N(k)=\\begin{pmatrix}0&\\alpha\\\\iGk&0\\end{pmatrix}`, one has
+    :math:`N(k)^2 = i\\alpha Gk\\,I`, so with :math:`s_0=\\sqrt{i\\alpha Gk}`
+    (principal branch; :math:`s_0=0` at :math:`k=0`),
+
+    .. math::
+
+        \\exp(N(k)\\,dt) = \\cosh(s_0 dt)\\,I + \\frac{\\sinh(s_0 dt)}{s_0}\\,N(k)
+
+    (the usual generalized Euler formula for a matrix whose square is a
+    scalar multiple of the identity), and
+    :math:`\\exp(M(k)\\,dt) = e^{-\\eta k^2 dt}\\exp(N(k)\\,dt)`. This is
+    applied to every Fourier mode at once (vectorized over `k`, not
+    looped), then inverse-transformed back to real space.
+
+    Parameters
+    ----------
+    A0, B0 : ndarray of float, shape (nx,)
+        Initial poloidal-flux proxy and toroidal field on a periodic grid
+        of ``nx = len(A0)`` points over ``[0, Lx)``.
+    alpha : float
+        Alpha-effect coefficient.
+    shear : float
+        Rotational shear rate :math:`G`.
+    eta : float
+        Magnetic diffusivity.
+    Lx : float
+        Physical domain size.
+    dt : float
+        Time step.
+    n_steps : int
+        Total number of steps.
+    save_every : int, default 1
+        Save a snapshot every this many steps.
+
+    Returns
+    -------
+    times : ndarray of float, shape (n_saved,)
+    A_snapshots : ndarray of float, shape (n_saved, nx)
+    B_snapshots : ndarray of float, shape (n_saved, nx)
+
+    See Also
+    --------
+    alpha_omega_growth_rate : Closed-form growth rate this stepping is checked against.
+    dominant_mode_growth_rate : Extracts the numerically observed growth rate from the output here.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> x = np.linspace(0, 2 * np.pi, 32, endpoint=False)
+    >>> A0 = 1e-3 * np.cos(x)
+    >>> B0 = np.zeros_like(A0)
+    >>> times, A_snaps, B_snaps = simulate_alpha_omega_dynamo(A0, B0, alpha=1.0, shear=5.0, eta=0.05, Lx=2 * np.pi, dt=0.01, n_steps=20, save_every=5)
+    >>> A_snaps.shape
+    (5, 32)
+    >>> bool(np.all(np.isfinite(A_snaps)))
+    True
+    """
+    A0 = np.asarray(A0, dtype=float)
+    B0 = np.asarray(B0, dtype=float)
+    nx = A0.shape[0]
+    k = 2.0 * np.pi * np.fft.fftfreq(nx, d=Lx / nx)
+
+    c = 1j * alpha * shear * k
+    s0 = np.sqrt(c)
+    dt_s0 = s0 * dt
+    cosh_term = np.cosh(dt_s0)
+    sinh_over_s0 = np.where(s0 == 0, dt, np.divide(np.sinh(dt_s0), s0, out=np.zeros_like(dt_s0), where=(s0 != 0)))
+    decay = np.exp(-eta * k**2 * dt)
+
+    A_hat = np.fft.fft(A0)
+    B_hat = np.fft.fft(B0)
+
+    times = [0.0]
+    A_snapshots = [A0.copy()]
+    B_snapshots = [B0.copy()]
+
+    t = 0.0
+    for step in range(1, n_steps + 1):
+        new_A_hat = decay * (cosh_term * A_hat + sinh_over_s0 * alpha * B_hat)
+        new_B_hat = decay * (cosh_term * B_hat + sinh_over_s0 * (1j * k * shear) * A_hat)
+        A_hat, B_hat = new_A_hat, new_B_hat
+        t += dt
+        if step % save_every == 0:
+            times.append(t)
+            A_snapshots.append(np.real(np.fft.ifft(A_hat)))
+            B_snapshots.append(np.real(np.fft.ifft(B_hat)))
+
+    return np.array(times), np.array(A_snapshots), np.array(B_snapshots)
+
+
+def dominant_mode_growth_rate(A_snapshots: NDArray[np.float64], dt_save: float, window_frac: float = 0.5) -> float:
+    """Numerically observed exponential growth rate of the total field energy.
+
+    Fits a line to :math:`\\log\\left(\\sum_x A(x,t)^2\\right)` versus time
+    (least squares) and returns half the slope -- since energy
+    :math:`\\propto A^2` grows as :math:`e^{2\\sigma t}` for a field
+    growing as :math:`e^{\\sigma t}`. Only the *last* `window_frac` of the
+    snapshots are used for the fit: an alpha-omega dynamo mode has two
+    eigenbranches (:math:`\\sigma\\pm` in the module docstring's
+    derivation), and unless the initial condition happens to be a pure
+    eigenmode of the faster-growing branch (e.g. `B0` all zero, as in
+    this module's typical test setup), the early evolution is a
+    transient mix of both branches that will bias a whole-history fit;
+    restricting to the late-time window isolates the asymptotically
+    dominant branch, which is what :func:`alpha_omega_growth_rate`
+    predicts. Used to validate :func:`simulate_alpha_omega_dynamo`
+    against that closed form.
+
+    Parameters
+    ----------
+    A_snapshots : ndarray of float, shape (n_saved, nx)
+        Field snapshots, e.g. `A_snapshots` or `B_snapshots` from
+        :func:`simulate_alpha_omega_dynamo`.
+    dt_save : float
+        Time interval between consecutive snapshots.
+    window_frac : float, default 0.5
+        Fraction of the snapshot history (counted from the end) used for
+        the fit.
+
+    Returns
+    -------
+    float
+        The fitted growth rate :math:`\\sigma`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> t = np.arange(10) * 0.1
+    >>> sigma_true = 2.0
+    >>> A_snaps = np.exp(sigma_true * t)[:, None] * np.ones((10, 4))
+    >>> round(dominant_mode_growth_rate(A_snaps, dt_save=0.1), 6)
+    2.0
+    """
+    energy = np.sum(A_snapshots**2, axis=1)
+    n = energy.shape[0]
+    times = np.arange(n) * dt_save
+    start = max(0, n - max(2, int(np.ceil(n * window_frac))))
+    log_energy = np.log(energy[start:] + 1e-300)
+    slope, _ = np.polyfit(times[start:], log_energy, 1)
+    return float(slope / 2.0)