quantum-playground 0.3.0__py3-none-any.whl

qplay/__init__.py

@@ -0,0 +1,29 @@
+"""qplay: a small physics and maths playground.
+
+Modules
+-------
+integrators  time-stepping schemes (Euler, RK4, velocity Verlet, adaptive wrapper)
+mechanics    oscillators, pendulums, projectiles, N-body gravity
+waves        1D wave and heat equations, Fourier series, spectra
+em           point-charge fields, Laplace solver, charged-particle motion
+quantum      Schroedinger solvers (1D and 2D), tunnelling, double slit, hydrogen orbitals, spin-1/2,
+             oscillator eigenstates and coherent states
+statmech     2D Ising model and Onsager's exact solution
+relativity   Lorentz transformations and special-relativity helpers
+kinetic      hard-sphere gas, Maxwell-Boltzmann distribution, ideal-gas law
+atomic       angular momentum, spin-orbit coupling, Zeeman effect, hydrogen fine structure
+nuclear      radioactive decay, Bateman decay chains, Monte Carlo decay
+circuits     RC charging and discharging, series RLC resonance
+chaos        Lorenz system, logistic map, Mandelbrot/Julia sets
+symbolic     SymPy helpers: Euler-Lagrange equations, Lagrangian -> simulation
+hamiltonian  Legendre transform, Hamilton's equations, Hamiltonian -> simulation
+dynamics     phase portraits and fixed-point classification for 2D systems
+complexplane domain colouring and the argument principle
+linalg       Gram-Schmidt, power iteration, how a matrix maps the unit circle
+numerics     root finding, quadrature, finite differences, Monte Carlo
+plotting     figure saving and GIF helpers
+demos        ready-made figures, run with ``python -m qplay run <name>``
+animations   GIFs, run with ``python -m qplay animate <name>``
+"""
+
+__version__ = "0.3.0"

qplay/__main__.py

@@ -0,0 +1,58 @@
+"""Command line.
+
+    qplay list                         show demos and animations
+    qplay run <name|all> [--out DIR]   render demo figures (PNG)
+    qplay animate <name|all> [--frames N] [--fps F] [--out DIR]   render animations (GIF)
+"""
+
+from __future__ import annotations
+
+import argparse
+
+from . import __version__, animations, demos
+
+
+def _summary(fn):
+    doc = (fn.__doc__ or "").strip().splitlines()
+    return doc[0] if doc else ""
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="qplay", description="Physics and maths playground")
+    parser.add_argument("--version", action="version", version=__version__)
+    sub = parser.add_subparsers(dest="command", required=True)
+    sub.add_parser("list", help="list demos and animations")
+    run = sub.add_parser("run", help="render demos to PNG files")
+    run.add_argument("names", nargs="+", help="demo names, or 'all'")
+    run.add_argument("--out", default="figures", help="output directory (default: figures)")
+    anim = sub.add_parser("animate", help="render animations to GIF files")
+    anim.add_argument("names", nargs="+", help="animation names, or 'all'")
+    anim.add_argument("--out", default="figures", help="output directory (default: figures)")
+    anim.add_argument("--frames", type=int, default=120, help="number of frames (default: 120)")
+    anim.add_argument("--fps", type=int, default=24, help="frames per second (default: 24)")
+    args = parser.parse_args(argv)
+
+    if args.command == "list":
+        print("Demos (qplay run NAME):")
+        for name in sorted(demos.DEMOS):
+            print(f"  {name:22s}{_summary(demos.DEMOS[name])}")
+        print("Animations (qplay animate NAME):")
+        for name in sorted(animations.ANIMATIONS):
+            print(f"  {name:22s}{_summary(animations.ANIMATIONS[name])}")
+        return 0
+
+    registry = demos.DEMOS if args.command == "run" else animations.ANIMATIONS
+    names = sorted(registry) if args.names == ["all"] else args.names
+    for name in names:
+        try:
+            if args.command == "run":
+                print(demos.run(name, args.out))
+            else:
+                print(animations.render(name, args.out, n_frames=args.frames, fps=args.fps))
+        except KeyError as err:
+            parser.error(str(err.args[0]))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

qplay/animations.py

@@ -0,0 +1,172 @@
+"""Animated GIFs. Each animation takes ``n_frames`` and returns ``(fig, update, frames)``.
+
+Register one with ``@animation("name")``; render with ``python -m qplay animate name``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import PowerNorm
+
+from . import kinetic, mechanics, plotting, quantum, statmech
+
+ANIMATIONS: dict = {}
+
+
+def animation(name):
+    def register(fn):
+        ANIMATIONS[name] = fn
+        return fn
+
+    return register
+
+
+def render(name, outdir="figures", n_frames=120, fps=24):
+    """Render one animation to ``<outdir>/<name>.gif`` and return the path."""
+    if name not in ANIMATIONS:
+        raise KeyError(f"unknown animation {name!r}; available: {', '.join(sorted(ANIMATIONS))}")
+    fig, update, frames = ANIMATIONS[name](n_frames)
+    return plotting.save_gif(fig, update, frames, Path(outdir) / f"{name}.gif", fps=fps)
+
+
+@animation("double-pendulum")
+def double_pendulum_anim(n_frames):
+    """A chaotic double pendulum with a fading trail."""
+    sub = 10  # integrate on a finer grid so the trail is smooth
+    t, y = mechanics.double_pendulum([2.2, 2.6, 0.0, 0.0], t_max=15, n=n_frames * sub)
+    x1, y1, x2, y2 = mechanics.double_pendulum_xy(y)
+    fig, ax = plt.subplots(figsize=(4, 4))
+    ax.set(xlim=(-2.2, 2.2), ylim=(-2.2, 2.2), aspect="equal", xticks=[], yticks=[])
+    ax.grid(False)
+    (trail,) = ax.plot([], [], color="C1", lw=1, alpha=0.6)
+    (rods,) = ax.plot([], [], "o-", color="k", lw=2, ms=7, mfc="C0")
+    label = ax.text(0.03, 0.95, "", transform=ax.transAxes, va="top")
+
+    def update(frame):
+        i = frame * sub
+        rods.set_data([0, x1[i], x2[i]], [0, y1[i], y2[i]])
+        start = max(0, i - 25 * sub)
+        trail.set_data(x2[start:i + 1], y2[start:i + 1])
+        label.set_text(f"t = {t[i]:.1f} s")
+
+    return fig, update, range(n_frames)
+
+
+@animation("wavepacket")
+def wavepacket_anim(n_frames):
+    """A Gaussian wavepacket tunnelling through and reflecting off a barrier."""
+    x = np.linspace(-60, 60, 2048, endpoint=False)
+    V = np.where(np.abs(x) < 1.0, 2.0, 0.0)
+    psi0 = quantum.gaussian_wavepacket(x, x0=-20, sigma=3, k0=1.5)
+    steps = 3000
+    times, psi = quantum.split_step(psi0, x, V + quantum.absorbing_layer(x, 10.0),
+                                    dt=0.01, steps=steps, save_every=max(1, steps // n_frames))
+    fig, ax = plt.subplots(figsize=(6, 3.2))
+    ax.fill_between(x, V / 2.0 * 0.15, color="0.85")
+    (line,) = ax.plot([], [], color="C0")
+    (real,) = ax.plot([], [], color="C1", lw=0.6, alpha=0.6)
+    ax.set(xlim=(-45, 45), ylim=(-0.4, 0.4), xlabel="x")
+    label = ax.text(0.02, 0.92, "", transform=ax.transAxes)
+
+    def update(i):
+        line.set_data(x, np.abs(psi[i]) ** 2)
+        real.set_data(x, psi[i].real * 0.6)
+        label.set_text(rf"$|\psi|^2$ and Re $\psi$,  t = {times[i]:.1f}")
+
+    return fig, update, range(min(n_frames, len(psi)))
+
+
+@animation("figure-eight")
+def figure_eight_anim(n_frames):
+    """Three equal masses chasing each other around a figure-eight."""
+    fe = mechanics.FIGURE_EIGHT
+    _, pos, _ = mechanics.nbody(fe["masses"], fe["pos0"], fe["vel0"], fe["period"], n=n_frames)
+    fig, ax = plt.subplots(figsize=(5, 2.8))
+    ax.plot(pos[:, 0, 0], pos[:, 0, 1], color="0.85", lw=1)
+    ax.set(xlim=(-1.2, 1.2), ylim=(-0.45, 0.45), aspect="equal", xticks=[], yticks=[])
+    ax.grid(False)
+    bodies = ax.scatter(pos[0, :, 0], pos[0, :, 1], s=60, c=["C0", "C1", "C2"], zorder=3)
+
+    def update(i):
+        bodies.set_offsets(pos[i])
+
+    return fig, update, range(n_frames)
+
+
+@animation("double-slit")
+def double_slit_anim(n_frames):
+    """A 2D wavepacket passing through two slits and interfering."""
+    steps = 1600
+    run = quantum.double_slit_experiment(save_every=max(1, steps // n_frames))
+    x, y, psi = run["x"], run["y"], run["psi"]
+    amp = np.abs(psi)
+    # square-root colour scale: the transmitted ~13% stays visible next to the bright incident packet
+    norm = PowerNorm(gamma=0.5, vmin=0, vmax=np.percentile(amp, 99.9))
+    fig, ax = plt.subplots(figsize=(4.6, 4.25))
+    fig.subplots_adjust(left=0.02, right=0.98, bottom=0.02, top=0.98)
+    fig.patch.set_facecolor("#0B0B0D")  # dark margins, so the GIF sits cleanly on dark pages
+    im = ax.imshow(amp[0], origin="lower", extent=[x[0], x[-1], y[0], y[-1]], cmap="magma", norm=norm)
+    ax.contour(x, y, run["V"], [1], colors="c", linewidths=0.8)
+    ax.set(xlim=(-15, 50), ylim=(-30, 30), xticks=[], yticks=[])
+    ax.grid(False)
+    label = ax.text(0.02, 0.94, "", transform=ax.transAxes, color="w")
+
+    def update(i):
+        im.set_data(amp[i])
+        label.set_text(rf"$|\psi|$,  t = {run['times'][i]:.1f}")
+
+    return fig, update, range(min(n_frames, len(psi)))
+
+
+@animation("ising")
+def ising_anim(n_frames):
+    """Domains coarsening in the 2D Ising model at the critical temperature."""
+    snaps: list = []
+    statmech.ising(L=128, T=statmech.T_CRITICAL, sweeps=n_frames, start="hot", seed=3, snapshots=snaps)
+    fig, ax = plt.subplots(figsize=(4, 4))
+    im = ax.imshow(snaps[0], cmap="gray", interpolation="nearest", vmin=-1, vmax=1)
+    ax.set(xticks=[], yticks=[])
+    ax.grid(False)
+    title = ax.set_title("")
+
+    def update(i):
+        im.set_data(snaps[i])
+        title.set_text(rf"Ising at $T_c$, sweep {i + 1}")
+
+    return fig, update, range(n_frames)
+
+
+@animation("gas")
+def gas_anim(n_frames):
+    """Discs in a box: every one starts at the same speed and the speeds spread into Maxwell-Boltzmann."""
+    gas = kinetic.gas_box(n=400, dim=2, radius=0.009, t_max=8, save_every=2, seed=4)
+    stride = max(1, (len(gas.t) - 1) // max(1, n_frames - 1))
+    idx = np.arange(0, len(gas.t), stride)[:n_frames]
+    speeds = np.linalg.norm(gas.velocities, axis=-1)
+    kT = gas.temperature()
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(7.6, 3.8), gridspec_kw={"width_ratios": [1, 1]})
+    ax1.set(xlim=(0, 1), ylim=(0, 1), aspect="equal", xticks=[], yticks=[])
+    ax1.grid(False)
+    dots = ax1.scatter(*gas.positions[0].T, s=9, c=speeds[0], cmap="plasma", vmin=0, vmax=2.5)
+    bins = np.linspace(0, 3, 31)
+    vs = np.linspace(0, 3, 200)
+    ax2.plot(vs, kinetic.maxwell_boltzmann_speed(vs, kT, dim=2), "k", lw=1.2, label="Maxwell-Boltzmann (2D)")
+    bars = ax2.bar(bins[:-1], np.zeros(30), width=bins[1] - bins[0], align="edge", color="C0", alpha=0.6)
+    ax2.set(xlim=(0, 3), ylim=(0, 2.2), xlabel="speed", yticks=[])
+    ax2.legend(loc="upper right", fontsize=8)
+    label = ax2.set_title("")
+
+    def update(frame):
+        k = idx[frame]
+        dots.set_offsets(gas.positions[k])
+        dots.set_array(speeds[k])
+        window = speeds[max(0, k - 40):k + 1].ravel()       # smooth the histogram over recent steps
+        hist, _ = np.histogram(window, bins=bins, density=True)
+        for bar, h in zip(bars, hist):
+            bar.set_height(h)
+        label.set_text(f"t = {gas.t[k]:.2f}")
+
+    return fig, update, range(len(idx))

qplay/atomic.py

@@ -0,0 +1,86 @@
+r"""Atomic structure: angular momentum matrices, spin-orbit coupling, the Zeeman effect and
+hydrogen fine structure.
+
+For one electron with orbital angular momentum $l$ and spin $s = \tfrac12$ in a magnetic field
+$B$ along $z$ (energies in units where $\hbar = 1$):
+
+$$H = \xi\,\mathbf L\cdot\mathbf S + \mu_B B\,(L_z + g_s S_z).$$
+
+At $B = 0$ the levels are labelled by $j = l\pm\tfrac12$ with
+$E = \tfrac{\xi}{2}\,[\,j(j+1) - l(l+1) - \tfrac34\,]$. A weak field splits each into
+$E = g_J m_j \mu_B B$ with the Landé factor $g_J$; a strong field gives the Paschen-Back limit
+$E \to \mu_B B\,(m_l + g_s m_s)$.
+
+Hydrogen fine structure (Dirac, to order $\alpha^2$):
+$$E_{nj} = -\frac{R_\infty hc}{n^2}\left[1 + \frac{\alpha^2}{n^2}\left(\frac{n}{j+\tfrac12} - \frac34\right)\right].$$
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy import constants
+
+ALPHA = constants.alpha
+RYDBERG_EV = constants.physical_constants["Rydberg constant times hc in eV"][0]
+MU_B_EV_PER_T = constants.physical_constants["Bohr magneton in eV/T"][0]
+ELECTRON_G = -constants.physical_constants["electron g factor"][0]   # about 2.0023
+
+
+def angular_momentum(j):
+    r"""Matrices ``(Jx, Jy, Jz)`` for spin ``j`` in the basis $m = j, j-1, \dots, -j$ ($\hbar = 1$).
+
+    Built from $J_\pm|j,m\rangle = \sqrt{j(j+1) - m(m\pm1)}\,|j,m\pm1\rangle$.
+    """
+    m = np.arange(j, -j - 1, -1)
+    jp = np.diag(np.sqrt(j * (j + 1) - m[1:] * (m[1:] + 1)), 1).astype(complex)  # raising
+    jm = jp.conj().T
+    return (jp + jm) / 2, (jp - jm) / 2j, np.diag(m).astype(complex)
+
+
+def basis_labels(l, s=0.5):
+    r"""``(m_l, m_s)`` for each basis vector of :func:`spin_orbit_zeeman_hamiltonian`."""
+    return [(ml, ms) for ml in np.arange(l, -l - 1, -1) for ms in np.arange(s, -s - 1, -1)]
+
+
+def spin_orbit_zeeman_hamiltonian(l, xi, B, mu_b=1.0, gs=2.0, s=0.5):
+    r"""$H = \xi\,\mathbf L\cdot\mathbf S + \mu_B B\,(L_z + g_s S_z)$ in the $|m_l, m_s\rangle$ basis.
+
+    ``xi`` and ``mu_b * B`` share whatever energy unit you choose.
+    """
+    Lx, Ly, Lz = angular_momentum(l)
+    Sx, Sy, Sz = angular_momentum(s)
+    il, is_ = np.eye(Lz.shape[0]), np.eye(Sz.shape[0])
+    LS = sum(np.kron(a, b) for a, b in ((Lx, Sx), (Ly, Sy), (Lz, Sz)))
+    return xi * LS + mu_b * B * (np.kron(Lz, is_) + gs * np.kron(il, Sz))
+
+
+def zeeman_levels(l, xi, B_values, mu_b=1.0, gs=2.0):
+    """Sorted energy levels for each field in ``B_values``; shape ``(len(B_values), 2(2l+1))``."""
+    return np.array([np.linalg.eigvalsh(spin_orbit_zeeman_hamiltonian(l, xi, B, mu_b, gs))
+                     for B in np.atleast_1d(B_values)])
+
+
+def lande_g(l, j, s=0.5, gs=2.0):
+    r"""$g_J = 1 + (g_s - 1)\,\dfrac{j(j+1) + s(s+1) - l(l+1)}{2j(j+1)}$."""
+    return 1 + (gs - 1) * (j * (j + 1) + s * (s + 1) - l * (l + 1)) / (2 * j * (j + 1))
+
+
+def spin_orbit_energy(l, j, xi, s=0.5):
+    r"""$\tfrac{\xi}{2}\,[\,j(j+1) - l(l+1) - s(s+1)\,]$, the eigenvalue of $\xi\,\mathbf L\cdot\mathbf S$."""
+    return xi / 2 * (j * (j + 1) - l * (l + 1) - s * (s + 1))
+
+
+def hydrogen_fine_structure_eV(n, j):
+    r"""Hydrogen level $E_{nj}$ in eV including fine structure to order $\alpha^2$ (see module docstring)."""
+    return -RYDBERG_EV / n**2 * (1 + ALPHA**2 / n**2 * (n / (j + 0.5) - 0.75))
+
+
+def dirac_energy_eV(n, j):
+    r"""Exact Dirac energy minus $m_ec^2$:
+    $E = m_ec^2\left[\left(1 + \dfrac{\alpha^2}{(n - \delta)^2}\right)^{-1/2} - 1\right]$,
+    $\delta = j + \tfrac12 - \sqrt{(j+\tfrac12)^2 - \alpha^2}$."""
+    mc2 = constants.physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
+    k = j + 0.5
+    delta = k - np.sqrt(k**2 - ALPHA**2)
+    x = (ALPHA / (n - delta)) ** 2
+    return mc2 * (1 / np.sqrt(1 + x) - 1)

qplay/chaos.py

@@ -0,0 +1,83 @@
+r"""Chaos and fractals: Lorenz system, logistic map, Mandelbrot and Julia sets.
+
+Logistic map: $x_{n+1} = r x_n (1 - x_n)$.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .integrators import solve
+
+
+def lorenz(state0=(1.0, 1.0, 1.0), sigma=10.0, rho=28.0, beta=8 / 3, t_max=40.0, n=4000):
+    """Lorenz system. Returns ``(t, y)`` with columns ``x, y, z``."""
+    def rhs(_, s):
+        x, y, z = s
+        return [sigma * (y - x), x * (rho - z) - y, x * y - beta * z]
+
+    return solve(rhs, state0, (0, t_max), n, rtol=1e-9, atol=1e-11)
+
+
+def logistic_map(r, x0=0.5, n=100):
+    """First ``n`` iterates (starting with ``x0``) of the logistic map."""
+    x = np.empty(n)
+    x[0] = x0
+    for i in range(n - 1):
+        x[i + 1] = r * x[i] * (1 - x[i])
+    return x
+
+
+def bifurcation(r_values, n_transient=500, n_keep=100, x0=0.5):
+    """Attractor of the logistic map for each ``r``. Returns flat arrays ``(r, x)`` for scatter plots."""
+    r_values = np.asarray(r_values, float)
+    x = np.full_like(r_values, x0)
+    for _ in range(n_transient):
+        x = r_values * x * (1 - x)
+    rs, xs = [], []
+    for _ in range(n_keep):
+        x = r_values * x * (1 - x)
+        rs.append(r_values.copy())
+        xs.append(x.copy())
+    return np.concatenate(rs), np.concatenate(xs)
+
+
+def lyapunov_logistic(r, n=20000, x0=0.4, n_transient=200):
+    r"""Lyapunov exponent $\lambda=\langle\ln|r(1-2x)|\rangle$; positive means chaos."""
+    x = x0
+    for _ in range(n_transient):
+        x = r * x * (1 - x)
+    total = 0.0
+    for _ in range(n):
+        x = r * x * (1 - x)
+        total += np.log(abs(r * (1 - 2 * x)) + 1e-300)
+    return total / n
+
+
+def escape_time(c, max_iter=100, z0=0.0):
+    """Iterations of $z\\to z^2+c$ before $|z|>2$ (``max_iter`` if it never escapes)."""
+    c = np.asarray(c, complex)
+    z = np.broadcast_to(np.asarray(z0, complex), c.shape).copy()
+    count = np.full(c.shape, max_iter, dtype=int)
+    alive = np.ones(c.shape, dtype=bool)
+    for i in range(max_iter):
+        z[alive] = z[alive] ** 2 + c[alive]
+        escaped = alive & (np.abs(z) > 2)
+        count[escaped] = i
+        alive &= ~escaped
+    return count
+
+
+def _grid(xlim, ylim, width, height):
+    X, Y = np.meshgrid(np.linspace(*xlim, width), np.linspace(*ylim, height))
+    return X + 1j * Y
+
+
+def mandelbrot(xlim=(-2.2, 0.8), ylim=(-1.2, 1.2), width=600, height=400, max_iter=100):
+    return escape_time(_grid(xlim, ylim, width, height), max_iter)
+
+
+def julia(c, xlim=(-1.6, 1.6), ylim=(-1.2, 1.2), width=600, height=400, max_iter=100):
+    """Julia set for $z\\to z^2 + c$: the grid supplies the starting points $z_0$."""
+    z0 = _grid(xlim, ylim, width, height)
+    return escape_time(np.full(z0.shape, c, dtype=complex), max_iter, z0=z0)

qplay/circuits.py

@@ -0,0 +1,75 @@
+r"""Electric circuits: capacitor charge and discharge, and the series RLC circuit.
+
+Series RLC with source $V(t)$ (charge $q$ on the capacitor, current $I = \dot q$):
+
+$$L\ddot q + R\dot q + \frac{q}{C} = V(t).$$
+
+With $L = 0$ this is the RC circuit, time constant $\tau = RC$. The undamped resonant angular
+frequency is $\omega_0 = 1/\sqrt{LC}$, the quality factor $Q = \frac1R\sqrt{L/C}$, and for a drive
+$V_0\cos\omega t$ the steady current amplitude is $V_0/|Z|$ with
+$Z = R + i\left(\omega L - \frac{1}{\omega C}\right)$.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .integrators import solve
+
+
+def _source(V):
+    return V if callable(V) else (lambda t, V=float(V): V)
+
+
+def rc_discharge(t, V0, R, C):
+    r"""Capacitor voltage while discharging: $V = V_0 e^{-t/RC}$."""
+    return V0 * np.exp(-np.asarray(t, float) / (R * C))
+
+
+def rc_charge(t, V0, R, C):
+    r"""Capacitor voltage while charging from a supply $V_0$: $V = V_0(1 - e^{-t/RC})$."""
+    return V0 * -np.expm1(-np.asarray(t, float) / (R * C))
+
+
+def rc_circuit(R, C, V=0.0, q0=0.0, t_max=1.0, n=500):
+    r"""Integrate $R\dot q + q/C = V(t)$. ``V`` is a number or a function of $t$.
+
+    Returns ``(t, q, I)``.
+    """
+    src = _source(V)
+    t, y = solve(lambda t, q: [(src(t) - q[0] / C) / R], [q0], (0, t_max), n)
+    q = y[:, 0]
+    return t, q, (np.array([src(ti) for ti in t]) - q / C) / R
+
+
+def rlc_series(R, L, C, V=0.0, q0=0.0, i0=0.0, t_max=1.0, n=1000):
+    r"""Integrate $L\ddot q + R\dot q + q/C = V(t)$. Returns ``(t, q, I)``."""
+    src = _source(V)
+
+    def rhs(t, y):
+        q, i = y
+        return [i, (src(t) - R * i - q / C) / L]
+
+    t, y = solve(rhs, [q0, i0], (0, t_max), n)
+    return t, y[:, 0], y[:, 1]
+
+
+def resonant_frequency(L, C):
+    r"""$\omega_0 = 1/\sqrt{LC}$ in rad/s."""
+    return 1 / np.sqrt(L * C)
+
+
+def quality_factor(R, L, C):
+    r"""$Q = \dfrac{1}{R}\sqrt{\dfrac{L}{C}}$; the resonance full width at half power is $\omega_0/Q = R/L$."""
+    return np.sqrt(L / C) / R
+
+
+def impedance(omega, R, L, C):
+    r"""Complex impedance $Z = R + i(\omega L - 1/\omega C)$ of the series circuit."""
+    omega = np.asarray(omega, float)
+    return R + 1j * (omega * L - 1 / (omega * C))
+
+
+def current_amplitude(omega, V0, R, L, C):
+    r"""Steady-state current amplitude $I_0 = V_0/|Z|$ for a drive $V_0\cos\omega t$."""
+    return V0 / np.abs(impedance(omega, R, L, C))

qplay/complexplane.py

@@ -0,0 +1,41 @@
+r"""Complex functions: domain colouring and the argument principle.
+
+Argument principle: $\frac{1}{2\pi i}\oint_C \frac{f'(z)}{f(z)}\,dz = N_{\text{zeros}} - N_{\text{poles}}$ inside $C$,
+which equals the number of times $f(C)$ winds around $0$.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from matplotlib.colors import hsv_to_rgb
+
+
+def complex_grid(xlim=(-2, 2), ylim=(-2, 2), n=500):
+    X, Y = np.meshgrid(np.linspace(*xlim, n), np.linspace(*ylim, n))
+    return X + 1j * Y
+
+
+def domain_coloring(f, xlim=(-2, 2), ylim=(-2, 2), n=500):
+    """RGB image of ``f``: hue shows $\\arg f$ (red = positive real), brightness bands show $\\log_2|f|$.
+
+    Returns an array of shape ``(n, n, 3)`` with row 0 at ``ylim[0]`` (use ``origin="lower"``).
+    """
+    with np.errstate(all="ignore"):
+        w = f(complex_grid(xlim, ylim, n))
+        hue = (np.angle(w) / (2 * np.pi)) % 1.0
+        bands = np.log2(np.abs(w)) % 1.0
+    bad = ~np.isfinite(w)
+    value = np.where(bad, 1.0, 0.55 + 0.45 * np.nan_to_num(bands))
+    sat = np.where(bad, 0.0, 0.85)
+    return hsv_to_rgb(np.stack([np.nan_to_num(hue), sat, value], axis=-1))
+
+
+def winding_number(f, center=0.0, radius=1.0, n=4000):
+    """How many times ``f`` maps the circle ``|z - center| = radius`` around the origin.
+
+    By the argument principle this is (zeros - poles) inside the circle, counted with multiplicity.
+    The circle must not pass through a zero or pole.
+    """
+    theta = np.linspace(0, 2 * np.pi, n + 1)
+    w = f(center + radius * np.exp(1j * theta))
+    return int(round(np.sum(np.diff(np.unwrap(np.angle(w)))) / (2 * np.pi)))