physicskit 0.1.0__py3-none-any.whl

physicskit/quantum/core/eigensolvers.py

@@ -0,0 +1,311 @@
+r"""Matrix Numerov eigensolver for the time-independent Schrodinger equation.
+
+Solves
+
+.. math::
+
+    -\frac{\hbar^2}{2m}\,\psi''(x) + V(x)\,\psi(x) = E\,\psi(x)
+
+on a uniform grid with :math:`\psi = 0` at the domain boundary, using the
+generalized-eigenvalue formulation of the Numerov method (Pillai, Goglio &
+Walker, *Am. J. Phys.* **80**, 1017 (2012)).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from numba import njit
+from scipy.linalg import eigh
+from scipy.sparse import diags
+
+from .._compat import trapz
+
+__all__ = [
+    "EigenResult",
+    "NumerovSolver",
+    "infinite_well",
+    "finite_well",
+    "asymmetric_step_well",
+    "linear_gravitational_well",
+    "double_well",
+    "harmonic_well",
+]
+
+
+@njit(cache=True)
+def _tridiag_arrays(V: np.ndarray, h: float, hbar: float, m: float):
+    n = V.shape[0]
+    a_main = np.empty(n)
+    a_off = np.empty(n - 1)
+    b_main = np.empty(n)
+    b_off = np.empty(n - 1)
+
+    kin = hbar**2 / (2.0 * m * h * h)
+    for i in range(n):
+        a_main[i] = 2.0 * kin + (10.0 / 12.0) * V[i]
+        b_main[i] = 10.0 / 12.0
+    for i in range(n - 1):
+        a_off[i] = -kin + (1.0 / 12.0) * 0.5 * (V[i] + V[i + 1])
+        b_off[i] = 1.0 / 12.0
+    return a_main, a_off, b_main, b_off
+
+
+@dataclass
+class EigenResult:
+    """Container for the output of :meth:`NumerovSolver.solve`."""
+
+    x: np.ndarray
+    """numpy.ndarray: The spatial grid the wavefunctions are sampled on."""
+
+    energies: np.ndarray
+    r"""numpy.ndarray: Eigenenergies :math:`E_n`, ascending, shape ``(n_states,)``."""
+
+    wavefunctions: np.ndarray
+    r"""numpy.ndarray: Normalized eigenfunctions :math:`\psi_n(x)`, shape
+    ``(n_states, len(x))``, with :math:`\int \lvert\psi_n\rvert^2\,dx = 1`
+    and :math:`\psi_n(x[0]) = \psi_n(x[-1]) = 0`."""
+
+
+class NumerovSolver:
+    r"""Bound-state solver for an arbitrary 1D potential :math:`V(x)` on a fixed grid.
+
+    Parameters
+    ----------
+    x : array_like
+        Uniform spatial grid. The wavefunction is forced to vanish at
+        ``x[0]`` and ``x[-1]`` (hard walls), so the domain should extend
+        well into the classically forbidden region for soft potentials.
+    V : callable or array_like
+        Potential energy, either a function ``V(x)`` or a precomputed array
+        matching ``x``.
+    hbar : float, default=1.0
+        Value of :math:`\hbar` to use.
+    m : float, default=1.0
+        Particle mass.
+
+    Notes
+    -----
+    The Numerov finite-difference relation for :math:`\psi'' = f\psi` is
+    recast as the generalized eigenvalue problem
+
+    .. math::
+
+        \hat A \boldsymbol\psi = E \hat B \boldsymbol\psi, \qquad
+        \hat A = -\frac{\hbar^2}{2m h^2}\hat T + \hat B\,\mathrm{diag}(V),
+
+    where :math:`\hat T` and :math:`\hat B` are the tridiagonal matrices
+    :math:`\mathrm{tridiag}(1,-2,1)` and :math:`\mathrm{tridiag}(1,10,1)/12`
+    and :math:`h` is the grid spacing. Validated against the analytic
+    infinite-well and harmonic-oscillator spectra to four decimal places.
+    """
+
+    def __init__(self, x: np.ndarray, V, hbar: float = 1.0, m: float = 1.0):
+        self.x = np.asarray(x, dtype=float)
+        h = self.x[1] - self.x[0]
+        if not np.allclose(np.diff(self.x), h, rtol=1e-8):
+            raise ValueError("NumerovSolver requires a uniform grid.")
+        self.h = h
+        self.hbar = hbar
+        self.m = m
+        self.V_full = np.asarray(V(self.x) if callable(V) else V, dtype=float)
+
+    def solve(self, n_states: int = 6) -> EigenResult:
+        """Compute the lowest bound states.
+
+        Parameters
+        ----------
+        n_states : int, default=6
+            Number of lowest-energy eigenstates to return. Must be smaller
+            than the number of interior grid points.
+
+        Returns
+        -------
+        EigenResult
+            The energies and normalized wavefunctions.
+
+        Raises
+        ------
+        ValueError
+            If ``n_states`` is not smaller than the number of interior grid
+            points.
+        """
+        # Interior points only; psi vanishes at the two boundary points.
+        V_int = self.V_full[1:-1]
+        n = V_int.shape[0]
+        if n_states >= n:
+            raise ValueError("n_states must be smaller than the number of interior grid points.")
+
+        a_main, a_off, b_main, b_off = _tridiag_arrays(V_int, self.h, self.hbar, self.m)
+
+        A = diags([a_off, a_main, a_off], offsets=[-1, 0, 1]).toarray()
+        B = diags([b_off, b_main, b_off], offsets=[-1, 0, 1]).toarray()
+
+        eigvals, eigvecs = eigh(A, B)
+
+        energies = eigvals[:n_states]
+        psis = np.zeros((n_states, self.x.shape[0]))
+        for k in range(n_states):
+            psi_int = eigvecs[:, k]
+            psi = np.zeros(self.x.shape[0])
+            psi[1:-1] = psi_int
+            norm = np.sqrt(trapz(psi**2, self.x))
+            psis[k] = psi / norm
+
+        return EigenResult(x=self.x, energies=energies, wavefunctions=psis)
+
+
+# --- Convenience potential builders ---------------------------------------
+
+
+def infinite_well():
+    r"""The infinite square well: :math:`V(x) = 0` inside the domain.
+
+    The grid boundaries passed to :class:`NumerovSolver` act as the
+    infinite walls, since the wavefunction is forced to vanish there.
+
+    Returns
+    -------
+    callable
+        A function ``V(x)`` returning zeros the same shape as ``x``.
+    """
+
+    def V(x):
+        return np.zeros_like(x)
+
+    return V
+
+
+def finite_well(V0: float, width: float):
+    r"""A finite square well,
+
+    .. math::
+
+        V(x) = \begin{cases} 0 & \lvert x\rvert \le \text{width}/2 \\
+                              V_0 & \text{otherwise} \end{cases}.
+
+    Parameters
+    ----------
+    V0 : float
+        Height of the barrier outside the well.
+    width : float
+        Full width of the well (zero-potential region).
+
+    Returns
+    -------
+    callable
+        A function ``V(x)``.
+    """
+
+    def V(x):
+        return np.where(np.abs(x) <= width / 2, 0.0, V0)
+
+    return V
+
+
+def asymmetric_step_well(width: float, V_left: float, V_right: float):
+    r"""A well of unequal left/right wall heights,
+
+    .. math::
+
+        V(x) = \begin{cases}
+            V_\text{left}  & x < -\text{width}/2 \\
+            0              & \lvert x\rvert \le \text{width}/2 \\
+            V_\text{right} & x > \text{width}/2
+        \end{cases}.
+
+    Parameters
+    ----------
+    width : float
+        Full width of the zero-potential region.
+    V_left : float
+        Wall height on the left side.
+    V_right : float
+        Wall height on the right side.
+
+    Returns
+    -------
+    callable
+        A function ``V(x)``.
+    """
+
+    def V(x):
+        out = np.zeros_like(x)
+        out[x < -width / 2] = V_left
+        out[x > width / 2] = V_right
+        return out
+
+    return V
+
+
+def linear_gravitational_well(alpha: float):
+    r"""The V-shaped 'quantum bouncer' well, :math:`V(x) = \alpha \lvert x\rvert`.
+
+    A particle in this potential is the quantum analogue of a ball bouncing
+    under uniform gravity (:math:`\alpha = mg` on the half-line); its
+    eigenstates are Airy functions (see :func:`physicskit.quantum.chapters.potentials.airy_bouncer_energies`).
+
+    Parameters
+    ----------
+    alpha : float
+        Slope of the linear potential (e.g. :math:`mg` for gravity).
+
+    Returns
+    -------
+    callable
+        A function ``V(x)``.
+    """
+
+    def V(x):
+        return alpha * np.abs(x)
+
+    return V
+
+
+def double_well(lam: float, a: float):
+    r"""The symmetric quartic double well, :math:`V(x) = \lambda (x^2 - a^2)^2`.
+
+    Has minima at :math:`x = \pm a`, each of depth 0, separated by a barrier
+    of height :math:`\lambda a^4` at :math:`x=0`.
+
+    Parameters
+    ----------
+    lam : float
+        Overall strength :math:`\lambda` of the quartic potential.
+    a : float
+        Half-separation of the two minima.
+
+    Returns
+    -------
+    callable
+        A function ``V(x)``.
+    """
+
+    def V(x):
+        return lam * (x**2 - a**2) ** 2
+
+    return V
+
+
+def harmonic_well(m: float = 1.0, omega: float = 1.0):
+    r"""The quantum harmonic oscillator potential,
+    :math:`V(x) = \tfrac{1}{2} m \omega^2 x^2`.
+
+    Parameters
+    ----------
+    m : float, default=1.0
+        Particle mass.
+    omega : float, default=1.0
+        Angular frequency :math:`\omega`.
+
+    Returns
+    -------
+    callable
+        A function ``V(x)``.
+    """
+
+    def V(x):
+        return 0.5 * m * omega**2 * x**2
+
+    return V

physicskit/quantum/core/operators.py

@@ -0,0 +1,296 @@
+r"""Dirac algebra helpers: Pauli matrices, spin/ladder operators, commutators.
+
+Conventions follow Cohen-Tannoudji, Diu & Laloe, with :math:`\hbar = 1`
+unless a function explicitly takes ``hbar`` as a parameter.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+__all__ = [
+    "sigma_x",
+    "sigma_y",
+    "sigma_z",
+    "identity2",
+    "PAULI",
+    "spin_operator",
+    "annihilation_operator",
+    "creation_operator",
+    "number_operator",
+    "position_operator",
+    "momentum_operator",
+    "commutator",
+    "anticommutator",
+    "expectation",
+    "is_hermitian",
+]
+
+# --- Pauli matrices -------------------------------------------------------
+
+sigma_x = np.array([[0, 1], [1, 0]], dtype=complex)
+r"""The Pauli matrix :math:`\sigma_x = \begin{pmatrix} 0 & 1 \\ 1 & 0 \end{pmatrix}`."""
+
+sigma_y = np.array([[0, -1j], [1j, 0]], dtype=complex)
+r"""The Pauli matrix :math:`\sigma_y = \begin{pmatrix} 0 & -i \\ i & 0 \end{pmatrix}`."""
+
+sigma_z = np.array([[1, 0], [0, -1]], dtype=complex)
+r"""The Pauli matrix :math:`\sigma_z = \begin{pmatrix} 1 & 0 \\ 0 & -1 \end{pmatrix}`."""
+
+identity2 = np.eye(2, dtype=complex)
+"""The :math:`2\\times 2` identity matrix."""
+
+PAULI = {"x": sigma_x, "y": sigma_y, "z": sigma_z, "i": identity2}
+"""Dict mapping ``'x'``, ``'y'``, ``'z'``, ``'i'`` to the corresponding matrix."""
+
+
+def spin_operator(axis: str, s: float = 0.5, hbar: float = 1.0) -> np.ndarray:
+    r"""Angular-momentum operator :math:`\hat S_\text{axis}` for spin quantum number ``s``.
+
+    For :math:`s=1/2` this reduces to :math:`(\hbar/2)\,\sigma_\text{axis}`. For
+    general :math:`s`, the :math:`(2s+1)`-dimensional matrix representation is
+    built from the ladder operators
+
+    .. math::
+
+        \hat S_\pm \lvert s, m\rangle
+            = \hbar\sqrt{s(s+1) - m(m\pm 1)}\, \lvert s, m\pm 1\rangle,
+
+    with :math:`\hat S_x = (\hat S_+ + \hat S_-)/2` and
+    :math:`\hat S_y = (\hat S_+ - \hat S_-)/2i`.
+
+    Parameters
+    ----------
+    axis : {'x', 'y', 'z', '+', '-'}
+        Which component (or ladder operator) to return.
+    s : float, default=0.5
+        Spin quantum number (may be integer or half-integer).
+    hbar : float, default=1.0
+        Value of :math:`\hbar` to use.
+
+    Returns
+    -------
+    numpy.ndarray
+        The :math:`(2s+1)\times(2s+1)` operator matrix, in the
+        :math:`\lvert s, m\rangle` basis ordered :math:`m = s, s-1, \dots, -s`.
+
+    Raises
+    ------
+    ValueError
+        If ``axis`` is not one of ``'x'``, ``'y'``, ``'z'``, ``'+'``, ``'-'``.
+    """
+    dim = int(round(2 * s + 1))
+    m_values = np.array([s - k for k in range(dim)])  # m = s, s-1, ..., -s
+
+    if axis == "z":
+        return hbar * np.diag(m_values).astype(complex)
+
+    s_plus = np.zeros((dim, dim), dtype=complex)
+    for k in range(dim - 1):
+        m = m_values[k + 1]
+        coeff = hbar * np.sqrt(s * (s + 1) - m * (m + 1))
+        s_plus[k, k + 1] = coeff
+    s_minus = s_plus.conj().T
+
+    if axis == "+":
+        return s_plus
+    if axis == "-":
+        return s_minus
+    if axis == "x":
+        return 0.5 * (s_plus + s_minus)
+    if axis == "y":
+        return (s_plus - s_minus) / (2j)
+    raise ValueError(f"Unknown axis {axis!r}; expected 'x', 'y', 'z', '+' or '-'.")
+
+
+# --- Harmonic-oscillator ladder operators (truncated Fock basis) ---------
+
+
+def annihilation_operator(n_max: int) -> np.ndarray:
+    r"""Lowering operator :math:`\hat a` in a truncated Fock basis.
+
+    Acts as :math:`\hat a \lvert n\rangle = \sqrt{n}\,\lvert n-1\rangle` on the
+    truncated basis :math:`\{\lvert 0\rangle, \dots, \lvert n_\text{max}-1\rangle\}`.
+
+    Parameters
+    ----------
+    n_max : int
+        Dimension of the truncated Fock space.
+
+    Returns
+    -------
+    numpy.ndarray
+        The :math:`n_\text{max}\times n_\text{max}` matrix representation of
+        :math:`\hat a`.
+    """
+    a = np.zeros((n_max, n_max), dtype=complex)
+    for n in range(1, n_max):
+        a[n - 1, n] = np.sqrt(n)
+    return a
+
+
+def creation_operator(n_max: int) -> np.ndarray:
+    r"""Raising operator :math:`\hat a^\dagger`, the Hermitian conjugate of
+    :func:`annihilation_operator`.
+
+    Parameters
+    ----------
+    n_max : int
+        Dimension of the truncated Fock space.
+
+    Returns
+    -------
+    numpy.ndarray
+        The :math:`n_\text{max}\times n_\text{max}` matrix representation of
+        :math:`\hat a^\dagger`.
+    """
+    return annihilation_operator(n_max).conj().T
+
+
+def number_operator(n_max: int) -> np.ndarray:
+    r"""Number operator :math:`\hat N = \hat a^\dagger \hat a`, diagonal with
+    eigenvalues :math:`0, 1, \dots, n_\text{max}-1`.
+
+    Parameters
+    ----------
+    n_max : int
+        Dimension of the truncated Fock space.
+
+    Returns
+    -------
+    numpy.ndarray
+        The :math:`n_\text{max}\times n_\text{max}` matrix representation of
+        :math:`\hat N`.
+    """
+    a = annihilation_operator(n_max)
+    return a.conj().T @ a
+
+
+def position_operator(n_max: int, m: float = 1.0, omega: float = 1.0, hbar: float = 1.0) -> np.ndarray:
+    r"""Harmonic-oscillator position operator in the truncated Fock basis,
+
+    .. math::
+
+        \hat x = \sqrt{\frac{\hbar}{2 m \omega}} \left(\hat a + \hat a^\dagger\right).
+
+    Parameters
+    ----------
+    n_max : int
+        Dimension of the truncated Fock space.
+    m : float, default=1.0
+        Particle mass.
+    omega : float, default=1.0
+        Angular frequency of the oscillator.
+    hbar : float, default=1.0
+        Value of :math:`\hbar` to use.
+
+    Returns
+    -------
+    numpy.ndarray
+        The :math:`n_\text{max}\times n_\text{max}` matrix representation of
+        :math:`\hat x`.
+    """
+    a = annihilation_operator(n_max)
+    return np.sqrt(hbar / (2 * m * omega)) * (a + a.conj().T)
+
+
+def momentum_operator(n_max: int, m: float = 1.0, omega: float = 1.0, hbar: float = 1.0) -> np.ndarray:
+    r"""Harmonic-oscillator momentum operator in the truncated Fock basis,
+
+    .. math::
+
+        \hat p = i\sqrt{\frac{\hbar m \omega}{2}} \left(\hat a^\dagger - \hat a\right).
+
+    Parameters
+    ----------
+    n_max : int
+        Dimension of the truncated Fock space.
+    m : float, default=1.0
+        Particle mass.
+    omega : float, default=1.0
+        Angular frequency of the oscillator.
+    hbar : float, default=1.0
+        Value of :math:`\hbar` to use.
+
+    Returns
+    -------
+    numpy.ndarray
+        The :math:`n_\text{max}\times n_\text{max}` matrix representation of
+        :math:`\hat p`.
+    """
+    a = annihilation_operator(n_max)
+    return 1j * np.sqrt(hbar * m * omega / 2) * (a.conj().T - a)
+
+
+# --- Algebra utilities ------------------------------------------------------
+
+
+def commutator(A: np.ndarray, B: np.ndarray) -> np.ndarray:
+    r"""The commutator :math:`[\hat A, \hat B] = \hat A \hat B - \hat B \hat A`.
+
+    Parameters
+    ----------
+    A, B : numpy.ndarray
+        Square matrices of matching shape.
+
+    Returns
+    -------
+    numpy.ndarray
+        :math:`\hat A \hat B - \hat B \hat A`.
+    """
+    return A @ B - B @ A
+
+
+def anticommutator(A: np.ndarray, B: np.ndarray) -> np.ndarray:
+    r"""The anticommutator :math:`\{\hat A, \hat B\} = \hat A \hat B + \hat B \hat A`.
+
+    Parameters
+    ----------
+    A, B : numpy.ndarray
+        Square matrices of matching shape.
+
+    Returns
+    -------
+    numpy.ndarray
+        :math:`\hat A \hat B + \hat B \hat A`.
+    """
+    return A @ B + B @ A
+
+
+def expectation(operator: np.ndarray, state: np.ndarray) -> complex:
+    r"""Expectation value :math:`\langle\psi\rvert \hat A \lvert\psi\rangle` of
+    an operator in a discrete-basis state vector.
+
+    Parameters
+    ----------
+    operator : numpy.ndarray
+        The operator matrix :math:`\hat A`.
+    state : numpy.ndarray
+        The (not necessarily normalized) state vector :math:`\lvert\psi\rangle`;
+        flattened if not already 1D.
+
+    Returns
+    -------
+    complex
+        :math:`\langle\psi\rvert \hat A \lvert\psi\rangle`.
+    """
+    state = np.asarray(state).reshape(-1)
+    return complex(np.vdot(state, operator @ state))
+
+
+def is_hermitian(A: np.ndarray, atol: float = 1e-10) -> bool:
+    r"""Check whether a matrix is Hermitian, :math:`\hat A = \hat A^\dagger`.
+
+    Parameters
+    ----------
+    A : numpy.ndarray
+        The matrix to test.
+    atol : float, default=1e-10
+        Absolute tolerance passed to :func:`numpy.allclose`.
+
+    Returns
+    -------
+    bool
+        ``True`` if ``A`` equals its conjugate transpose within ``atol``.
+    """
+    return np.allclose(A, A.conj().T, atol=atol)