updated

Author: Garvity

physics/hamiltonian.py

@@ -1,38 +1,40 @@
 """
-Hamiltonian functions for classical and quantum  mechanics.
+Hamiltonian functions for classical and quantum mechanics.
 
 This module provides two educational, minimal implementations:
 
-- classical_hamiltonian(mass, momentum, potential_energy): Computes H = T + V for a particle/system, where
-    T = p^2 / (2 m) is the kinetic energy expressed in terms of momentum p, and V is
-  the potential energy (can be a scalar or an array broadcastable to p).
+- ham_c(mass, momentum, potential_energy):
+    Computes H = T + V, where T = p^2/(2m), and V is the potential (scalar or array).
 
-- quantum_hamiltonian_1d(mass, hbar, potential_energy, dx): Builds the 1D Hamiltonian matrix for a
-    particle in a potential V using second-order central finite differences for the
-  kinetic energy operator: T = - (hbar^2 / 2m) d^2/dx^2 with Dirichlet boundaries.
+- ham_1d(mass, hbar, potential_energy, dx):
+    Builds the 1D Hamiltonian using second-order central differences for the kinetic
+    operator: T = - (hbar^2 / 2m) d^2/dx^2 with Dirichlet boundaries.
 
-These functions are intended for learners to quickly prototype and simulate basic
-physical systems.
+These functions help learners quickly prototype and simulate basic physical systems.
 
 References
 ----------
-- Classical Hamiltonian mechanics: https://en.wikipedia.org/wiki/Hamiltonian_mechanics
-- Discrete 1D Schrödinger operator: https://en.wikipedia.org/wiki/Finite_difference_method
+- Classical Hamiltonian mechanics:
+    https://en.wikipedia.org/wiki/Hamiltonian_mechanics
+- Discrete 1D Schrödinger operator:
+    https://en.wikipedia.org/wiki/Finite_difference_method
 """
 
-from __future__ import annotations
-
 from typing import Any
 
 import numpy as np
 
 
-def classical_hamiltonian(mass: float, momentum: Any, potential_energy: Any) -> Any:
+def ham_c(
+    mass: float,
+    momentum: Any,
+    potential_energy: Any,
+) -> Any:
     """
     Classical Hamiltonian H = T + V with T = p^2 / (2 m).
 
     The function supports scalars or array-like inputs for momentum ``p`` and
-    potential energy ``v``; NumPy broadcasting rules apply. If inputs are scalars,
+    potential energy ``v``. NumPy broadcasting rules apply. If inputs are scalars,
     a float is returned; otherwise a NumPy array is returned.
 
     Parameters
@@ -47,19 +49,19 @@ def classical_hamiltonian(mass: float, momentum: Any, potential_energy: Any) ->
     Returns
     -------
     float | np.ndarray
-    The Hamiltonian value(s) H = p^2/(2m) + V.
+        The Hamiltonian value(s) H = p^2/(2m) + V.
 
     Examples
     --------
     Free particle with p = 3 kg·m/s and m = 2 kg (v = 0):
-    >>> classical_hamiltonian(2.0, 3.0, 0.0)
+    >>> ham_c(2.0, 3.0, 0.0)
     2.25
 
     Harmonic oscillator snapshot with vectorized p and v:
     >>> mass = 1.0
     >>> momentum = np.array([0.0, 1.0, 2.0])
-    >>> potential_energy = np.array([0.5, 0.5, 0.5])  # e.g., 1/2 k x^2 at three positions
-    >>> classical_hamiltonian(mass, momentum, potential_energy).tolist()
+    >>> potential_energy = np.array([0.5, 0.5, 0.5])  # e.g., 1/2 k x^2 at positions
+    >>> ham_c(mass, momentum, potential_energy).tolist()
     [0.5, 1.0, 2.5]
     """
     if mass <= 0:
@@ -77,8 +79,12 @@ def classical_hamiltonian(mass: float, momentum: Any, potential_energy: Any) ->
     return h
 
 
-def quantum_hamiltonian_1d(
-    mass: float, hbar: float, potential_energy: Any, dx: float
+
+def ham_1d(
+    mass: float,
+    hbar: float,
+    potential_energy: Any,
+    dx: float,
 ) -> np.ndarray:
     """
     Construct the 1D quantum Hamiltonian matrix using finite differences.
@@ -89,7 +95,7 @@ def quantum_hamiltonian_1d(
     H = - (hbar^2 / 2m) d^2/dx^2 + V
 
     On a uniform grid with spacing ``dx`` and N sites, the Laplacian is
-    approximated by the tridiagonal matrix with main diagonal ``-2`` and
+    approximated by a tridiagonal matrix with main diagonal ``-2`` and
     off-diagonals ``+1``. The resulting kinetic term has main diagonal
     ``(hbar^2)/(m*dx^2)`` and off-diagonals ``-(hbar^2)/(2*m*dx^2)``.
 
@@ -111,9 +117,12 @@ def quantum_hamiltonian_1d(
 
     Examples
     --------
-    Free particle (v=0) on a small grid: main diagonal = 1/dx^2, off = -1/(2*dx^2) in units m=hbar=1.
+    Free particle (v=0) on a small grid: main diagonal = 1/dx^2, off = -1/(2*dx^2)
+    in units m=hbar=1.
     >>> n, dx = 5, 0.1
-    >>> h = quantum_hamiltonian_1d(mass=1.0, hbar=1.0, potential_energy=np.zeros(n), dx=dx)
+    >>> h = ham_1d(
+    ...     mass=1.0, hbar=1.0, potential_energy=np.zeros(n), dx=dx
+    ... )
     >>> float(h[0, 0])
     99.99999999999999
     >>> float(h[0, 1])
@@ -122,7 +131,7 @@ def quantum_hamiltonian_1d(
     Add a harmonic-like site potential to the diagonal:
     >>> x = dx * (np.arange(n) - (n-1)/2)
     >>> potential_energy = 0.5 * x**2  # k=m=omega=1 for illustration
-    >>> h2 = quantum_hamiltonian_1d(1.0, 1.0, potential_energy, dx)
+    >>> h2 = ham_1d(1.0, 1.0, potential_energy, dx)
     >>> np.allclose(np.diag(h2) - np.diag(h), potential_energy)
     True
     """
@@ -158,3 +167,7 @@ def quantum_hamiltonian_1d(
     import doctest
 
     doctest.testmod(verbose=True)
+
+# Backward-compatible aliases
+classical_hamiltonian = ham_c
+quantum_hamiltonian_1d = ham_1d