physicskit 0.1.0__py3-none-any.whl

physicskit/quantum/core/solvers.py

@@ -0,0 +1,271 @@
+r"""Split-operator FFT propagators for the time-dependent Schrodinger equation.
+
+Both solvers implement the symmetric (Strang) splitting
+
+.. math::
+
+    \psi(t+dt) = e^{-i \hat V dt / 2\hbar}\, e^{-i \hat T dt/\hbar}\,
+                 e^{-i \hat V dt / 2\hbar}\, \psi(t),
+
+which is unitary to machine precision regardless of :math:`dt` (probability
+is exactly conserved up to FFT/round-off error), and second-order accurate
+in :math:`dt`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import numpy as np
+
+from .._compat import trapz
+
+__all__ = ["SplitOperatorSolver1D", "SplitOperatorSolver2D"]
+
+
+class SplitOperatorSolver1D:
+    r"""Propagate a 1D wavepacket under potential :math:`V(x)` or :math:`V(x,t)`.
+
+    Parameters
+    ----------
+    x : numpy.ndarray
+        Uniform spatial grid.
+    V : callable or array_like
+        Potential energy: a function ``V(x)`` (static) or ``V(x, t)``
+        (time-dependent, detected by inspecting the function's signature),
+        or a precomputed array matching ``x``.
+    hbar : float, default=1.0
+        Value of :math:`\hbar` to use.
+    m : float, default=1.0
+        Particle mass.
+    dt : float, default=1e-3
+        Time step.
+    """
+
+    def __init__(self, x: np.ndarray, V, hbar: float = 1.0, m: float = 1.0, dt: float = 1e-3):
+        self.x = np.asarray(x, dtype=float)
+        self.N = self.x.shape[0]
+        self.dx = self.x[1] - self.x[0]
+        self.hbar = hbar
+        self.m = m
+        self.dt = dt
+        self._V = V
+        self._time_dependent = callable(V) and _accepts_time(V, n_spatial_args=1)
+
+        self.k = 2 * np.pi * np.fft.fftfreq(self.N, d=self.dx)
+        self._kin_phase = np.exp(-1j * hbar * self.k**2 * dt / (2 * m))
+
+        if not self._time_dependent:
+            V0 = V(self.x) if callable(V) else np.asarray(V, dtype=float)
+            self._half_pot_phase = np.exp(-1j * V0 * dt / (2 * hbar))
+        else:
+            self._half_pot_phase = None
+
+    def _half_potential_phase(self, t: float) -> np.ndarray:
+        if self._half_pot_phase is not None:
+            return self._half_pot_phase
+        Vt = self._V(self.x, t)
+        return np.exp(-1j * Vt * self.dt / (2 * self.hbar))
+
+    def step(self, psi: np.ndarray, t: float = 0.0) -> np.ndarray:
+        """Advance the wavefunction by one time step ``dt``.
+
+        Parameters
+        ----------
+        psi : numpy.ndarray
+            Complex wavefunction at time ``t``, sampled on ``self.x``.
+        t : float, default=0.0
+            Current time (only used if the potential is time-dependent).
+
+        Returns
+        -------
+        numpy.ndarray
+            The wavefunction at time ``t + dt``.
+        """
+        psi = psi * self._half_potential_phase(t)
+        psi = np.fft.ifft(self._kin_phase * np.fft.fft(psi))
+        psi = psi * self._half_potential_phase(t + self.dt)
+        return psi
+
+    def propagate(self, psi0: np.ndarray, n_steps: int, t0: float = 0.0, save_every: int = 1) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+        """Propagate an initial state and record its history.
+
+        Parameters
+        ----------
+        psi0 : numpy.ndarray
+            Initial wavefunction (renormalized internally).
+        n_steps : int
+            Total number of time steps to take.
+        t0 : float, default=0.0
+            Initial time.
+        save_every : int, default=1
+            Save a snapshot every ``save_every`` steps (plus the initial state).
+
+        Returns
+        -------
+        frames : numpy.ndarray
+            Stacked wavefunction snapshots, shape ``(n_saved, len(x))``.
+        times : numpy.ndarray
+            The time of each snapshot, shape ``(n_saved,)``.
+        """
+        psi = np.asarray(psi0, dtype=complex).copy()
+        psi /= np.sqrt(trapz(np.abs(psi) ** 2, self.x))
+
+        frames = [psi.copy()]
+        times = [t0]
+        t = t0
+        for step_idx in range(1, n_steps + 1):
+            psi = self.step(psi, t)
+            t = t0 + step_idx * self.dt
+            if step_idx % save_every == 0:
+                frames.append(psi.copy())
+                times.append(t)
+        return np.array(frames), np.array(times)
+
+    def norm(self, psi: np.ndarray) -> float:
+        r"""Total probability :math:`\int \lvert\psi(x)\rvert^2\,dx`.
+
+        Parameters
+        ----------
+        psi : numpy.ndarray
+            Wavefunction sampled on ``self.x``.
+
+        Returns
+        -------
+        float
+            The normalization integral (should stay at 1.0 under unitary
+            evolution).
+        """
+        return float(trapz(np.abs(psi) ** 2, self.x))
+
+
+class SplitOperatorSolver2D:
+    r"""Propagate a 2D wavepacket under potential :math:`V(x,y)` or :math:`V(x,y,t)`.
+
+    Parameters
+    ----------
+    x, y : numpy.ndarray
+        Uniform spatial grids along each axis.
+    V : callable or array_like
+        Potential energy: a function ``V(x, y)`` (static) or ``V(x, y, t)``
+        (time-dependent, detected by inspecting the function's signature),
+        or a precomputed array matching the ``(x, y)`` meshgrid.
+    hbar : float, default=1.0
+        Value of :math:`\hbar` to use.
+    m : float, default=1.0
+        Particle mass.
+    dt : float, default=1e-3
+        Time step.
+    """
+
+    def __init__(self, x: np.ndarray, y: np.ndarray, V, hbar: float = 1.0, m: float = 1.0, dt: float = 1e-3):
+        self.x = np.asarray(x, dtype=float)
+        self.y = np.asarray(y, dtype=float)
+        self.dx = self.x[1] - self.x[0]
+        self.dy = self.y[1] - self.y[0]
+        self.hbar = hbar
+        self.m = m
+        self.dt = dt
+        self._V = V
+        self._time_dependent = callable(V) and _accepts_time(V, n_spatial_args=2)
+
+        self.X, self.Y = np.meshgrid(self.x, self.y, indexing="ij")
+        kx = 2 * np.pi * np.fft.fftfreq(self.x.shape[0], d=self.dx)
+        ky = 2 * np.pi * np.fft.fftfreq(self.y.shape[0], d=self.dy)
+        KX, KY = np.meshgrid(kx, ky, indexing="ij")
+        self._kin_phase = np.exp(-1j * hbar * (KX**2 + KY**2) * dt / (2 * m))
+
+        if not self._time_dependent:
+            V0 = V(self.X, self.Y) if callable(V) else np.asarray(V, dtype=float)
+            self._half_pot_phase = np.exp(-1j * V0 * dt / (2 * hbar))
+        else:
+            self._half_pot_phase = None
+
+    def _half_potential_phase(self, t: float) -> np.ndarray:
+        if self._half_pot_phase is not None:
+            return self._half_pot_phase
+        Vt = self._V(self.X, self.Y, t)
+        return np.exp(-1j * Vt * self.dt / (2 * self.hbar))
+
+    def step(self, psi: np.ndarray, t: float = 0.0) -> np.ndarray:
+        """Advance the wavefunction by one time step ``dt``.
+
+        Parameters
+        ----------
+        psi : numpy.ndarray
+            Complex wavefunction at time ``t``, sampled on the ``(x, y)`` grid.
+        t : float, default=0.0
+            Current time (only used if the potential is time-dependent).
+
+        Returns
+        -------
+        numpy.ndarray
+            The wavefunction at time ``t + dt``.
+        """
+        psi = psi * self._half_potential_phase(t)
+        psi = np.fft.ifft2(self._kin_phase * np.fft.fft2(psi))
+        psi = psi * self._half_potential_phase(t + self.dt)
+        return psi
+
+    def propagate(self, psi0: np.ndarray, n_steps: int, t0: float = 0.0, save_every: int = 1):
+        """Propagate an initial state and record its history.
+
+        Parameters
+        ----------
+        psi0 : numpy.ndarray
+            Initial wavefunction (renormalized internally).
+        n_steps : int
+            Total number of time steps to take.
+        t0 : float, default=0.0
+            Initial time.
+        save_every : int, default=1
+            Save a snapshot every ``save_every`` steps (plus the initial state).
+
+        Returns
+        -------
+        frames : numpy.ndarray
+            Stacked wavefunction snapshots, shape ``(n_saved, *X.shape)``.
+        times : numpy.ndarray
+            The time of each snapshot, shape ``(n_saved,)``.
+        """
+        psi = np.asarray(psi0, dtype=complex).copy()
+        psi /= np.sqrt(self.norm(psi))
+
+        frames = [psi.copy()]
+        times = [t0]
+        t = t0
+        for step_idx in range(1, n_steps + 1):
+            psi = self.step(psi, t)
+            t = t0 + step_idx * self.dt
+            if step_idx % save_every == 0:
+                frames.append(psi.copy())
+                times.append(t)
+        return np.array(frames), np.array(times)
+
+    def norm(self, psi: np.ndarray) -> float:
+        r"""Total probability :math:`\iint \lvert\psi(x,y)\rvert^2\,dx\,dy`.
+
+        Parameters
+        ----------
+        psi : numpy.ndarray
+            Wavefunction sampled on the ``(x, y)`` grid.
+
+        Returns
+        -------
+        float
+            The normalization integral (should stay at 1.0 under unitary
+            evolution).
+        """
+        density = np.abs(psi) ** 2
+        return float(trapz(trapz(density, self.y, axis=1), self.x))
+
+
+def _accepts_time(func: Callable, n_spatial_args: int) -> bool:
+    """True if ``func`` takes more than ``n_spatial_args`` positional
+    parameters, i.e. it is V(x, ..., t) rather than V(x, ...)."""
+    try:
+        import inspect
+
+        return len(inspect.signature(func).parameters) > n_spatial_args
+    except (TypeError, ValueError):
+        return False

physicskit/quantum/tests/conftest.py

@@ -0,0 +1,7 @@
+"""Test-session setup: force Matplotlib's non-interactive Agg backend
+before anything imports pyplot, so the suite (and CI, which has no
+display) can exercise the plotting/animation code headlessly."""
+
+import matplotlib
+
+matplotlib.use("Agg")

physicskit/quantum/tests/test_animations.py

@@ -0,0 +1,225 @@
+"""Every dynamical demo produces a genuine frame-by-frame animation (not a
+single static snapshot), and the underlying physics is spot-checked against
+known behavior (spreading, tunneling, splitting, entanglement growth, ...).
+"""
+
+import numpy as np
+import pytest
+from matplotlib.animation import FuncAnimation, PillowWriter
+
+from physicskit.quantum._compat import trapz
+from physicskit.quantum.chapters.entanglement import IsingEntangler
+from physicskit.quantum.chapters.harmonic_spin import HarmonicOscillator
+from physicskit.quantum.chapters.hydrogen_am import HydrogenOrbital, orbital_superposition_density
+from physicskit.quantum.chapters.potentials import DoubleWellSimulator, FiniteSquareWell
+from physicskit.quantum.chapters.spin import RabiProblem, SternGerlach
+from physicskit.quantum.chapters.wave_packets import GaussianDispersion, propagate_double_slit
+from physicskit.quantum.visualizers.bloch_sphere import animate_bloch_sphere, state_to_bloch_trajectory
+from physicskit.quantum.visualizers.entanglement import animate_entanglement_growth
+from physicskit.quantum.visualizers.orbitals import animate_orbital_beating
+from physicskit.quantum.visualizers.wavefunctions import animate_density, animate_density_2d
+
+
+def _save_and_check(anim: FuncAnimation, tmp_path, name: str):
+    out = tmp_path / name
+    anim.save(out, writer=PillowWriter(fps=10))
+    assert out.exists() and out.stat().st_size > 0
+
+
+# 1. Double-slit: genuine 2D time propagation through a two-gap wall.
+@pytest.mark.slow
+def test_double_slit_propagation_animates(tmp_path):
+    x = np.linspace(-15, 15, 96)
+    y = np.linspace(-12, 12, 96)
+    X, Y, frames, times = propagate_double_slit(
+        x,
+        y,
+        x0=-8.0,
+        k0=8.0,
+        sigma_x=0.7,
+        sigma_y=4.0,
+        wall_x=0.0,
+        n_steps=160,
+        save_every=20,
+        dt=1e-4,
+    )
+    assert frames.shape[0] == len(times) >= 5
+    norm0 = float(np.sum(np.abs(frames[0]) ** 2)) * (x[1] - x[0]) * (y[1] - y[0])
+    norm_end = float(np.sum(np.abs(frames[-1]) ** 2)) * (x[1] - x[0]) * (y[1] - y[0])
+    assert abs(norm0 - norm_end) < 1e-2  # probability conserved through the slits
+
+    anim = animate_density_2d(X, Y, frames, times=times)
+    _save_and_check(anim, tmp_path, "double_slit.gif")
+
+
+# 2. Quantum tunnelling: a wavepacket incident on a barrier splits into
+# reflected + transmitted pieces.
+@pytest.mark.slow
+def test_barrier_tunnelling_animates(tmp_path):
+    barrier = FiniteSquareWell(V0=-6.0, width=1.0)
+    x, frames, times = barrier.wavepacket_scattering(
+        x_extent=25.0,
+        n_points=1024,
+        x0=-10.0,
+        sigma0=1.5,
+        k0=None,
+        dt=1e-3,
+        n_steps=5000,
+        save_every=500,
+    )
+    assert frames.shape[0] == len(times) >= 5
+
+    dx = x[1] - x[0]
+    final_density = np.abs(frames[-1]) ** 2
+    transmitted = float(np.sum(final_density[x > 0.5]) * dx)
+    reflected = float(np.sum(final_density[x < -0.5]) * dx)
+    assert transmitted > 0.01  # some probability tunnels through
+    assert reflected > 0.01  # some probability reflects
+    assert abs((transmitted + reflected) - 1.0) < 0.2
+
+    anim = animate_density(x, frames, times=times)
+    _save_and_check(anim, tmp_path, "barrier_tunnelling.gif")
+
+
+# 3. Hydrogen orbitals: a two-eigenstate superposition beats in density,
+# unlike either stationary orbital alone.
+def test_hydrogen_orbital_beating_animates():
+    orb_a = HydrogenOrbital(2, 0, 0)
+    orb_b = HydrogenOrbital(3, 1, 0)
+    period = 2 * np.pi / (orb_b.energy - orb_a.energy)
+    times = np.linspace(0, abs(period), 4)
+
+    fig = animate_orbital_beating(orb_a, orb_b, times, n_points=12)
+    assert len(fig.frames) == len(times)
+    assert fig.data[0].value.size > 0
+    assert any(m["buttons"] for m in fig.layout.updatemenus)  # a Play button exists
+
+    r = np.linspace(0.1, 15, 20)
+    theta = np.full_like(r, np.pi / 3)
+    phi = np.zeros_like(r)
+    density_0 = orbital_superposition_density(orb_a, orb_b, r, theta, phi, 0.0)
+    density_half = orbital_superposition_density(orb_a, orb_b, r, theta, phi, abs(period) / 2)
+    assert np.max(np.abs(density_0 - density_half)) > 1e-6  # genuinely time-dependent
+
+
+# 4. Harmonic-oscillator eigenstate superposition: plumbing onto animate_density.
+@pytest.mark.slow
+def test_harmonic_superposition_animates(tmp_path):
+    ho = HarmonicOscillator()
+    x = np.linspace(-10, 10, 400)
+    times = np.linspace(0, 2 * np.pi / ho.omega, 12)
+    frames = ho.superposition_trajectory([0, 1, 2], [1.0, 1.0j, 0.5], x, times)
+
+    for psi in frames:
+        assert abs(trapz(np.abs(psi) ** 2, x) - 1.0) < 1e-6
+
+    anim = animate_density(x, frames, times=times)
+    _save_and_check(anim, tmp_path, "harmonic_superposition.gif")
+
+
+# 5. Free wave-packet dispersion: plumbing onto animate_density.
+@pytest.mark.slow
+def test_wave_packet_dispersion_animates(tmp_path):
+    gd = GaussianDispersion(x0=0.0, sigma0=1.0, k0=3.0)
+    x = np.linspace(-30, 30, 600)
+    times = np.linspace(0, 6.0, 10)
+    frames = gd.trajectory(x, times)
+
+    widths = [gd.width(t) for t in times]
+    assert widths[-1] > widths[0]  # the packet spreads
+
+    anim = animate_density(x, frames, times=times)
+    _save_and_check(anim, tmp_path, "wave_packet_dispersion.gif")
+
+
+# 6. Scattering off a finite well: resonance/reflection dynamics, sharing
+# FiniteSquareWell.wavepacket_scattering with the barrier case above.
+@pytest.mark.slow
+def test_well_scattering_animates(tmp_path):
+    well = FiniteSquareWell(V0=20.0, width=2.0)
+    x, frames, times = well.wavepacket_scattering(
+        x_extent=25.0,
+        n_points=1024,
+        x0=-10.0,
+        sigma0=1.5,
+        k0=None,
+        dt=1e-3,
+        n_steps=5000,
+        save_every=500,
+    )
+    assert frames.shape[0] == len(times) >= 5
+
+    anim = animate_density(x, frames, times=times)
+    _save_and_check(anim, tmp_path, "well_scattering.gif")
+
+
+# 7. Stern-Gerlach: the joint spatial density splits into two lobes.
+@pytest.mark.slow
+def test_stern_gerlach_splitting_animates(tmp_path):
+    sg = SternGerlach(mu=2.0, grad_B=1.0, sigma0=1.0, sigma_x=2.0, k0=4.0)
+    x = np.linspace(-15, 15, 60)
+    y = np.linspace(-15, 15, 80)
+    times = np.array([0.0, 1.0, 2.5])
+    stack = sg.joint_density_stack(x, y, times)
+    assert stack.shape == (len(times), len(x), len(y))
+
+    def y_spread(density_xy):
+        marginal_y = density_xy.sum(axis=0)
+        marginal_y /= marginal_y.sum()
+        mean_y = np.sum(y * marginal_y)
+        return np.sqrt(np.sum((y - mean_y) ** 2 * marginal_y))
+
+    assert y_spread(stack[-1]) > y_spread(stack[0])  # splits apart over time
+
+    anim = animate_density_2d(x, y, stack, times=times, phase_colored=False)
+    _save_and_check(anim, tmp_path, "stern_gerlach.gif")
+
+
+# 8. Rabi oscillations on the Bloch sphere.
+@pytest.mark.slow
+def test_rabi_bloch_sphere_animates(tmp_path):
+    rabi = RabiProblem(omega0=1.0, omega_d=1.0, Omega=0.5)  # on resonance
+    t_pi = np.pi / rabi.Omega  # a resonant pi-pulse fully inverts the population
+    assert rabi.excited_state_population(np.array([t_pi]))[0] > 0.999
+
+    times = np.linspace(0, t_pi, 20)
+    states = rabi.state_trajectory(times)
+    assert np.allclose(np.sum(np.abs(states) ** 2, axis=1), 1.0, atol=1e-10)
+
+    trajectory = state_to_bloch_trajectory(states)
+    assert np.allclose(np.linalg.norm(trajectory, axis=1), 1.0, atol=1e-8)
+
+    anim = animate_bloch_sphere(trajectory, times=times)
+    _save_and_check(anim, tmp_path, "rabi_bloch.gif")
+
+
+# 9. Entangling Ising coupling: concurrence grows from 0 toward 1.
+@pytest.mark.slow
+def test_entanglement_growth_animates(tmp_path):
+    ising = IsingEntangler(J=1.0)
+    assert abs(ising.concurrence(ising.initial_state())) < 1e-12  # starts unentangled
+
+    times = np.linspace(0, np.pi / (2 * ising.J), 30)
+    concurrence = ising.concurrence_trajectory(times)
+    assert concurrence[0] < 1e-6
+    assert concurrence.max() > 0.99  # reaches (near-)maximal entanglement
+
+    anim = animate_entanglement_growth(times, concurrence)
+    _save_and_check(anim, tmp_path, "entanglement_growth.gif")
+
+
+# 10. Double-well tunneling oscillation: almost pure plumbing onto animate_density.
+@pytest.mark.slow
+def test_double_well_tunneling_animates(tmp_path):
+    dw = DoubleWellSimulator(lam=0.3, a=1.5)
+    result = dw.solve(n_states=2)
+    times = np.linspace(0, result.tunneling_period, 16)
+    frames = dw.tunneling_wavefunction(result, times, side="left")
+    assert frames.shape == (len(times), len(dw.x))
+
+    left_prob = dw.left_well_probability(result, times, side="left")
+    assert left_prob[0] > 0.9  # starts localized on the left
+    assert left_prob[np.argmin(np.abs(times - result.tunneling_period / 2))] < 0.5  # tunnels to the right
+
+    anim = animate_density(dw.x, frames, times=times)
+    _save_and_check(anim, tmp_path, "double_well.gif")

physicskit/quantum/tests/test_compat.py

@@ -0,0 +1,43 @@
+"""Test for physicskit.quantum._compat's scipy<1.15 fallback branch.
+
+The dev/CI environment installs scipy>=1.15 (where scipy.special
+already has sph_harm_y natively), so the ``except ImportError`` branch
+that builds a fallback from the older, deprecated ``sph_harm`` is dead
+code under the installed version -- but it is real code protecting
+users on scipy>=1.10,<1.15 (the package's declared minimum), so it's
+worth exercising directly rather than pragma-excluding it.
+
+This simulates the old-scipy environment by temporarily removing
+``scipy.special.sph_harm_y`` and installing a fake ``sph_harm`` with
+the old (m, n, phi, theta) signature/convention, then reloading the
+compat module so it re-runs its import-time fallback logic, and
+restores everything afterward regardless of outcome.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+import scipy.special as sp
+
+import physicskit.quantum._compat as compat
+
+
+def test_sph_harm_y_fallback_matches_native_implementation_when_scipy_is_old():
+    real_sph_harm_y = sp.sph_harm_y
+
+    def fake_old_sph_harm(m, n, phi, theta):
+        # The old, deprecated scipy.special.sph_harm's argument order and
+        # convention: (m, n, azimuthal, polar).
+        return real_sph_harm_y(n, m, theta, phi)
+
+    del sp.sph_harm_y
+    sp.sph_harm = fake_old_sph_harm
+    try:
+        importlib.reload(compat)
+        n, m, theta, phi = 2, 1, 0.3, 0.7
+        assert compat.sph_harm_y(n, m, theta, phi) == real_sph_harm_y(n, m, theta, phi)
+    finally:
+        del sp.sph_harm
+        sp.sph_harm_y = real_sph_harm_y
+        importlib.reload(compat)

physicskit/quantum/tests/test_eigensolvers.py

@@ -0,0 +1,24 @@
+"""Error-path tests for NumerovSolver that test_physics_checks.py's
+accuracy checks never exercise: rejecting a non-uniform grid and
+rejecting a requested state count that isn't smaller than the number
+of interior grid points."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from physicskit.quantum.core.eigensolvers import NumerovSolver, infinite_well
+
+
+def test_numerov_solver_rejects_non_uniform_grid():
+    x = np.array([0.0, 0.1, 0.3, 0.6, 1.0])  # unevenly spaced
+    with pytest.raises(ValueError):
+        NumerovSolver(x, infinite_well())
+
+
+def test_numerov_solver_rejects_n_states_not_smaller_than_interior_points():
+    x = np.linspace(0, 1, 6)  # 4 interior points
+    solver = NumerovSolver(x, infinite_well())
+    with pytest.raises(ValueError):
+        solver.solve(n_states=4)