physicskit 0.1.0__py3-none-any.whl

physicskit/semiclassical/tests/test_path_integral.py

@@ -0,0 +1,194 @@
+"""Physics-correctness tests for physicskit.semiclassical.core.path_integral."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from physicskit.semiclassical.core.path_integral import (
+    build_phasor_diagram,
+    discretized_action,
+    feynman_phasor_partial_sums,
+    free_particle_classical_action,
+    free_particle_classical_path,
+    harmonic_oscillator_classical_action,
+    harmonic_oscillator_classical_path,
+    sample_random_paths,
+)
+
+
+def test_discretized_action_matches_free_particle_closed_form_exactly():
+    # The free-particle classical path is a straight line, so the midpoint-rule
+    # discretization is exact (constant velocity, zero potential) for any n_slices.
+    x0, xf, T, m = 0.0, 2.0, 1.0, 1.0
+    _, x = free_particle_classical_path(x0, xf, T, m, n_slices=200)
+    S_discrete = discretized_action(x, dt=T / 200, m=m, V=lambda x: 0.0 * x)
+    S_exact = free_particle_classical_action(x0, xf, T, m)
+    assert S_discrete == pytest.approx(S_exact, rel=1e-10)
+
+
+def test_discretized_action_converges_to_harmonic_oscillator_closed_form():
+    # Unlike the free particle, the classical path here is curved, so the midpoint
+    # rule has O(dt^2) discretization error. At n_slices=200 the achieved relative
+    # error is ~2.5e-4 (checked empirically), well inside the 1% tolerance used here.
+    x0, xf, T, m, omega = 0.3, -0.7, 2.3, 1.0, 1.1
+    n_slices = 200
+    _, x = harmonic_oscillator_classical_path(x0, xf, T, m, omega, n_slices)
+    V = lambda x: 0.5 * m * omega**2 * x**2
+    S_discrete = discretized_action(x, dt=T / n_slices, m=m, V=V)
+    S_exact = harmonic_oscillator_classical_action(x0, xf, T, m, omega)
+    assert S_discrete == pytest.approx(S_exact, rel=1e-2)
+
+
+def test_free_particle_classical_path_minimizes_action_among_random_paths():
+    # For a free particle the classical (straight-line) path is a true minimum of
+    # the action, not merely a stationary point, so every randomly sampled path
+    # must have action >= the classical action (within numerical tolerance).
+    x0, xf, T, m, n_slices = 0.0, 1.5, 1.0, 1.0, 40
+    S_cl = free_particle_classical_action(x0, xf, T, m)
+    paths = sample_random_paths(x0, xf, T, n_slices, n_paths=50, sigma=0.4, seed=1)
+    dt = T / n_slices
+    actions = np.array([discretized_action(p, dt, m, lambda x: 0.0 * x) for p in paths])
+    assert np.all(actions >= S_cl - 1e-8)
+
+
+def test_harmonic_oscillator_classical_action_closer_than_random_perturbation():
+    # Unlike the free particle, the harmonic-oscillator classical path is only a
+    # stationary point of the action (not necessarily a minimum), so we only check
+    # that it is closer to the exact classical action than a randomly perturbed
+    # path -- a weaker but still meaningful sanity check.
+    x0, xf, T, m, omega, n_slices = 0.2, 0.8, 1.0, 1.0, 1.3, 60
+    S_cl = harmonic_oscillator_classical_action(x0, xf, T, m, omega)
+    _, x_cl = harmonic_oscillator_classical_path(x0, xf, T, m, omega, n_slices)
+    V = lambda x: 0.5 * m * omega**2 * x**2
+    dt = T / n_slices
+    S_cl_discrete = discretized_action(x_cl, dt, m, V)
+
+    random_paths = sample_random_paths(x0, xf, T, n_slices, n_paths=1, sigma=0.5, seed=3)
+    S_random = discretized_action(random_paths[0], dt, m, V)
+
+    assert abs(S_cl_discrete - S_cl) < abs(S_random - S_cl)
+
+
+@pytest.mark.parametrize(
+    "target,func,kwargs",
+    [
+        ("path", harmonic_oscillator_classical_path, dict(x0=0.0, xf=1.0, m=1.0, omega=1.0, n_slices=10)),
+        ("action", harmonic_oscillator_classical_action, dict(x0=0.0, xf=1.0, m=1.0, omega=1.0)),
+    ],
+)
+def test_harmonic_oscillator_raises_at_conjugate_point(target, func, kwargs):
+    # omega*T = pi exactly is a conjugate point: the classical path is non-unique.
+    T = np.pi / kwargs["omega"]
+    with pytest.raises(ValueError):
+        func(T=T, **kwargs)
+
+
+def test_feynman_phasor_partial_sums_matches_hand_computation():
+    actions = np.array([0.3, -1.2, 5.0, 0.0])
+    hbar = 0.5
+    classical_action = 0.1
+    paths = np.zeros((actions.shape[0], 3))  # only shape[0] is used by the function
+
+    partial = feynman_phasor_partial_sums(paths, actions, hbar, classical_action=classical_action)
+
+    # Independent hand computation: sort by |action - classical_action| ascending,
+    # then cumulatively sum exp(i*action/hbar) in that order, starting from 0.
+    order = np.argsort(np.abs(actions - classical_action))
+    expected = np.concatenate([[0.0 + 0.0j], np.cumsum(np.exp(1j * actions[order] / hbar))])
+
+    assert np.allclose(partial, expected)
+    assert partial.shape == (actions.shape[0] + 1,)
+    assert partial[0] == 0.0
+
+
+def test_phasor_sum_concentrates_near_classical_path_as_hbar_shrinks():
+    # The core physics claim of this module: shrinking hbar makes the propagator's
+    # support concentrate ever more tightly around the classical path. Quantify
+    # this by comparing, at two very different hbar values, what fraction of the
+    # final (all-paths) partial-sum magnitude is already captured by just the
+    # closest 10% of paths (sorted by proximity to the classical action).
+    def frac_from_closest_10_percent(hbar, seed):
+        _, _, _, _, partial = build_phasor_diagram(
+            x0=0.0,
+            xf=1.0,
+            T=1.0,
+            m=1.0,
+            hbar=hbar,
+            potential="free",
+            n_slices=40,
+            n_paths=400,
+            sigma=0.5,
+            seed=seed,
+        )
+        n_total = partial.shape[0] - 1
+        idx = max(1, round(0.1 * n_total))
+        return abs(partial[idx]) / abs(partial[-1])
+
+    # hbar=20 is much larger than the typical |S - S_cl| spread of this ensemble
+    # (tens at most), so phases barely wind and the closest 10% contributes about
+    # the "trivial" 10% share. hbar=1.0 is much smaller than that spread, so the
+    # closest paths reinforce while distant ones increasingly cancel.
+    large_hbar, small_hbar = 20.0, 1.0
+    seed = 42
+
+    frac_large = frac_from_closest_10_percent(large_hbar, seed)
+    frac_small = frac_from_closest_10_percent(small_hbar, seed)
+
+    assert frac_small > frac_large
+
+
+def test_free_particle_functions_reject_nonpositive_t_and_bad_n_slices():
+    with pytest.raises(ValueError):
+        free_particle_classical_path(x0=0.0, xf=1.0, T=0.0, m=1.0, n_slices=4)
+    with pytest.raises(ValueError):
+        free_particle_classical_path(x0=0.0, xf=1.0, T=1.0, m=1.0, n_slices=0)
+    with pytest.raises(ValueError):
+        free_particle_classical_action(x0=0.0, xf=1.0, T=0.0, m=1.0)
+
+
+def test_harmonic_oscillator_functions_reject_nonpositive_t_and_bad_n_slices():
+    with pytest.raises(ValueError):
+        harmonic_oscillator_classical_path(x0=0.0, xf=1.0, T=0.0, m=1.0, omega=1.0, n_slices=4)
+    with pytest.raises(ValueError):
+        harmonic_oscillator_classical_path(x0=0.0, xf=1.0, T=1.0, m=1.0, omega=1.0, n_slices=0)
+    with pytest.raises(ValueError):
+        harmonic_oscillator_classical_action(x0=0.0, xf=1.0, T=0.0, m=1.0, omega=1.0)
+
+
+def test_sample_random_paths_rejects_nonpositive_t_and_bad_counts():
+    with pytest.raises(ValueError):
+        sample_random_paths(x0=0.0, xf=1.0, T=0.0, n_slices=4, n_paths=2, sigma=0.1, seed=0)
+    with pytest.raises(ValueError):
+        sample_random_paths(x0=0.0, xf=1.0, T=1.0, n_slices=0, n_paths=2, sigma=0.1, seed=0)
+    with pytest.raises(ValueError):
+        sample_random_paths(x0=0.0, xf=1.0, T=1.0, n_slices=4, n_paths=0, sigma=0.1, seed=0)
+
+
+def test_feynman_phasor_partial_sums_rejects_empty_actions_and_zero_hbar():
+    with pytest.raises(ValueError):
+        feynman_phasor_partial_sums(paths=np.empty((0, 1)), actions=np.array([]), hbar=0.5)
+    with pytest.raises(ValueError):
+        feynman_phasor_partial_sums(paths=np.zeros((2, 1)), actions=np.array([0.1, 0.2]), hbar=0.0)
+
+
+def test_build_phasor_diagram_harmonic_potential_matches_classical_path_and_action():
+    omega = 1.5
+    paths, actions, x_cl, S_cl, partial = build_phasor_diagram(
+        x0=0.0, xf=1.0, T=1.0, m=1.0, hbar=0.1, potential="harmonic", omega=omega, n_slices=10, n_paths=5, sigma=0.1, seed=0
+    )
+    expected_path = harmonic_oscillator_classical_path(x0=0.0, xf=1.0, T=1.0, m=1.0, omega=omega, n_slices=10)[1]
+    expected_action = harmonic_oscillator_classical_action(x0=0.0, xf=1.0, T=1.0, m=1.0, omega=omega)
+    assert np.allclose(x_cl, expected_path)
+    assert S_cl == pytest.approx(expected_action)
+    assert paths.shape[0] == actions.shape[0] == 6
+    assert partial.shape == (7,)
+
+
+def test_build_phasor_diagram_rejects_nonpositive_t_missing_omega_and_bad_potential():
+    with pytest.raises(ValueError):
+        build_phasor_diagram(x0=0.0, xf=1.0, T=0.0, m=1.0, hbar=0.1, potential="free", n_slices=10, n_paths=5, sigma=0.1, seed=0)
+    with pytest.raises(ValueError):
+        build_phasor_diagram(x0=0.0, xf=1.0, T=1.0, m=1.0, hbar=0.1, potential="harmonic", n_slices=10, n_paths=5, sigma=0.1, seed=0)
+    with pytest.raises(ValueError):
+        build_phasor_diagram(x0=0.0, xf=1.0, T=1.0, m=1.0, hbar=0.1, potential="quartic", n_slices=10, n_paths=5, sigma=0.1, seed=0)

physicskit/semiclassical/tests/test_propagators.py

@@ -0,0 +1,82 @@
+"""Physics-correctness tests for physicskit.semiclassical.core.propagators."""
+
+from __future__ import annotations
+
+import numpy as np
+from numba import njit
+
+from physicskit.semiclassical.core.propagators import (
+    coherent_state_overlap,
+    frozen_gaussian_1d,
+    herman_kluk_propagate_wavepacket,
+    propagate_trajectory_monodromy_action,
+    van_vleck_propagator_1d,
+)
+
+_ZERO = njit(lambda q, params: 0.0, cache=False)
+
+
+def _harmonic_potential(m=1.0, omega=1.0):
+    params = np.array([m * omega**2])
+    dVdx = njit(lambda q, params: params[0] * q, cache=False)
+    d2Vdx2 = njit(lambda q, params: params[0], cache=False)
+    V = njit(lambda q, params: 0.5 * params[0] * q**2, cache=False)
+    return dVdx, d2Vdx2, V, params
+
+
+def test_monodromy_matches_harmonic_oscillator_and_preserves_liouville():
+    dVdx, d2Vdx2, V, params = _harmonic_potential()
+    t = 2.1
+    _q, _p, M, _S, _hist = propagate_trajectory_monodromy_action(1.0, 0.4, dVdx, d2Vdx2, V, m=1.0, dt=t / 3000, steps=3000, params=params)
+    M_exact = np.array([[np.cos(t), np.sin(t)], [-np.sin(t), np.cos(t)]])
+    assert np.max(np.abs(M - M_exact)) < 1e-5
+    assert abs(np.linalg.det(M) - 1.0) < 1e-8
+
+
+def test_van_vleck_propagator_matches_free_particle_exact():
+    q_t, K = van_vleck_propagator_1d(q0=0.0, p0=1.0, dVdx=_ZERO, d2Vdx2=_ZERO, V=_ZERO, m=1.0, dt=1.0 / 2000, steps=2000)
+    K_exact = np.sqrt(1.0 / (2j * np.pi)) * np.exp(1j * q_t**2 / 2.0)
+    assert abs(K - K_exact) < 1e-6
+
+
+def test_van_vleck_propagator_matches_harmonic_oscillator_exact():
+    dVdx, d2Vdx2, V, params = _harmonic_potential()
+    q0, p0, t = 1.0, 0.3, 1.3
+    q_t, K = van_vleck_propagator_1d(q0, p0, dVdx, d2Vdx2, V, m=1.0, dt=t / 4000, steps=4000, hbar=1.0, params=params)
+    K_exact = np.sqrt(1.0 / (2j * np.pi * np.sin(t))) * np.exp(1j / (2 * np.sin(t)) * ((q_t**2 + q0**2) * np.cos(t) - 2 * q_t * q0))
+    assert abs(K - K_exact) < 1e-5
+
+
+def test_coherent_state_overlap_matches_numerical_integration():
+    x = np.linspace(-40, 40, 20000)
+    gamma = 1.0
+    psi1 = frozen_gaussian_1d(x, qc=0.3, pc=0.7, gamma=gamma)
+    psi2 = frozen_gaussian_1d(x, qc=-0.2, pc=1.1, gamma=gamma)
+    numeric = np.trapezoid(np.conj(psi1) * psi2, x)
+    closed_form = coherent_state_overlap(0.3, 0.7, -0.2, 1.1, gamma)
+    assert abs(numeric - closed_form) < 1e-6
+
+
+def test_herman_kluk_reconstructs_initial_state_at_short_time():
+    dVdx, d2Vdx2, V, params = _harmonic_potential()
+    x = np.linspace(-6, 6, 400)
+    psi = herman_kluk_propagate_wavepacket(
+        1.0, 0.0, gamma=1.0, dVdx=dVdx, d2Vdx2=d2Vdx2, V=V, m=1.0, dt=1e-6, steps=1, x_eval=x, n_grid=61, n_sigma=7.0, params=params
+    )
+    psi0 = frozen_gaussian_1d(x, qc=1.0, pc=0.0, gamma=1.0)
+    fidelity = abs(np.trapezoid(np.conj(psi0) * psi, x)) ** 2
+    assert fidelity > 0.999
+
+
+def test_herman_kluk_defaults_params_to_an_empty_array_when_unused():
+    # dVdx/d2Vdx2/V here never index into params, so params=None (->
+    # np.empty(0) internally) is fine and should behave like the
+    # explicit-params harmonic oscillator above (m=omega=1).
+    dVdx = njit(lambda q, params: q, cache=False)
+    d2Vdx2 = njit(lambda q, params: 1.0, cache=False)
+    V = njit(lambda q, params: 0.5 * q**2, cache=False)
+    x = np.linspace(-6, 6, 400)
+    psi = herman_kluk_propagate_wavepacket(1.0, 0.0, gamma=1.0, dVdx=dVdx, d2Vdx2=d2Vdx2, V=V, m=1.0, dt=1e-6, steps=1, x_eval=x, n_grid=61, n_sigma=7.0)
+    psi0 = frozen_gaussian_1d(x, qc=1.0, pc=0.0, gamma=1.0)
+    fidelity = abs(np.trapezoid(np.conj(psi0) * psi, x)) ** 2
+    assert fidelity > 0.999

physicskit/semiclassical/tests/test_scarring.py

@@ -0,0 +1,43 @@
+"""Physics-correctness tests for physicskit.semiclassical.systems.scarring."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from physicskit.semiclassical.systems.scarring import bouncing_ball_energies, husimi_projection_1d, scar_enhancement
+
+
+def test_bouncing_ball_energies_match_infinite_well_of_width_2R():
+    R = 0.7
+    energies = bouncing_ball_energies(R, n_max=4)
+    exact = (np.arange(1, 5) * np.pi) ** 2 / (2 * (2 * R) ** 2)
+    assert np.allclose(energies, exact)
+
+
+def test_scar_enhancement_of_uniform_density_is_one():
+    x = np.linspace(-2, 2, 200)
+    y = np.linspace(-1, 1, 100)
+    X, Y = np.meshgrid(x, y, indexing="ij")
+    mask = np.ones_like(X, dtype=bool)
+    uniform = np.ones_like(X)
+    eta = scar_enhancement(uniform, X, Y, mask, x0=0.3, half_width=0.2)
+    assert abs(eta - 1.0) < 1e-10
+
+
+def test_scar_enhancement_detects_concentrated_ridge():
+    x = np.linspace(-2, 2, 400)
+    y = np.linspace(-1, 1, 200)
+    X, Y = np.meshgrid(x, y, indexing="ij")
+    mask = np.ones_like(X, dtype=bool)
+    density = np.exp(-((X - 0.3) ** 2) / (2 * 0.05**2))
+    eta = scar_enhancement(density, X, Y, mask, x0=0.3, half_width=0.15)
+    assert eta > 5.0
+
+
+def test_husimi_projection_1d_default_sigma_and_ranges_peaks_near_source():
+    s = np.linspace(-10, 10, 500)
+    s0_true, p0_true, w = 2.0, 3.0, 1.0
+    psi = np.exp(-((s - s0_true) ** 2) / (2 * w**2)) * np.exp(1j * p0_true * s / 1.0)
+    S0, P0, H = husimi_projection_1d(psi, s, hbar=1.0, resolution=30)
+    i, j = np.unravel_index(np.argmax(H), H.shape)
+    assert abs(S0[i, j] - s0_true) < 1.0

physicskit/semiclassical/tests/test_visualizers.py

@@ -0,0 +1,89 @@
+"""Coverage for physicskit.semiclassical.visualizers's ax=<given> branch
+(previously untested for every function here): plot_density_of_states,
+plot_wkb_wavefunction, plot_classical_trajectory_on_wigner, plot_scar_map,
+plot_husimi_1d. All pure matplotlib, small inputs -- fast."""
+
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from physicskit.semiclassical.core.gutzwiller import gutzwiller_density_of_states
+from physicskit.semiclassical.core.path_integral import build_phasor_diagram
+from physicskit.semiclassical.core.propagators import frozen_gaussian_1d
+from physicskit.semiclassical.core.wkb import bohr_sommerfeld_energies, wkb_wavefunction
+from physicskit.semiclassical.systems.scarring import bouncing_ball_orbit_points, husimi_projection_1d
+from physicskit.semiclassical.visualizers.gutzwiller import plot_density_of_states
+from physicskit.semiclassical.visualizers.path_integral import animate_feynman_phasor_spiral, plot_phasor_convergence
+from physicskit.semiclassical.visualizers.propagators import plot_classical_trajectory_on_wigner
+from physicskit.semiclassical.visualizers.scarring import plot_husimi_1d, plot_scar_map
+from physicskit.semiclassical.visualizers.wkb import plot_wkb_wavefunction
+
+
+def test_plot_density_of_states_given_ax():
+    V = lambda x: 0.5 * x**2
+    E_grid = np.linspace(0.2, 3.5, 50)
+    dos = gutzwiller_density_of_states(E_grid, V, m=1.0, x_min=-20, x_max=20)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_density_of_states(E_grid, dos, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_wkb_wavefunction_given_ax():
+    V = lambda x: 0.5 * x**2
+    E2 = bohr_sommerfeld_energies(V, m=1.0, x_min=-20, x_max=20, n_max=3)[2]
+    x = np.linspace(-6, 6, 200)
+    psi = wkb_wavefunction(x, E2, V)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_wkb_wavefunction(x, psi, V=V, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_classical_trajectory_on_wigner_given_ax():
+    x = np.linspace(-8, 8, 100)
+    psi = frozen_gaussian_1d(x, qc=1.0, pc=0.5, gamma=1.0)
+    theta = np.linspace(0, 2 * np.pi, 30)
+    q_hist, p_hist = np.cos(theta), np.sin(theta)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_classical_trajectory_on_wigner(x, psi, q_hist, p_hist, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_scar_map_given_ax():
+    x = np.linspace(-2, 2, 30)
+    y = np.linspace(-1, 1, 20)
+    X, Y = np.meshgrid(x, y, indexing="ij")
+    density = np.exp(-((X - 0.3) ** 2) / (2 * 0.2**2))
+    orbit_x, orbit_y = bouncing_ball_orbit_points(x0=0.3, R=1.0)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_scar_map(X, Y, density, orbit_x=orbit_x, orbit_y=orbit_y, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_husimi_1d_given_ax():
+    s = np.linspace(-10, 10, 200)
+    psi = np.exp(-((s - 2.0) ** 2) / 2.0) * np.exp(1j * 3.0 * s)
+    S0, P0, H = husimi_projection_1d(psi, s, sigma=1.0, resolution=20, s0_range=(-2, 6), p0_range=(-2, 8))
+    fig, ax = plt.subplots()
+    _, ax_given = plot_husimi_1d(S0, P0, H, ax=ax)
+    assert ax_given is ax
+
+
+def test_animate_feynman_phasor_spiral_builds_and_draws_a_frame():
+    paths, actions, x_cl, S_cl, _ = build_phasor_diagram(x0=0.0, xf=1.0, T=1.0, m=1.0, hbar=0.1, potential="free", n_slices=20, n_paths=10, sigma=0.3, seed=0)
+    t = np.linspace(0.0, 1.0, x_cl.shape[0])
+    anim = animate_feynman_phasor_spiral(paths, actions, x_cl, S_cl, hbar=0.1, t_array=t)
+    anim._init_draw()
+    anim._draw_frame(0)
+    anim._draw_frame(1)  # a second frame exercises removing the previous frame's path lines
+    assert "2/" in anim._fig.axes[1].get_title()
+
+
+def test_plot_phasor_convergence_given_ax():
+    sums, hbars = [], [0.5, 0.05]
+    for hbar in hbars:
+        _, _, _, _, partial = build_phasor_diagram(x0=0.0, xf=1.0, T=1.0, m=1.0, hbar=hbar, potential="free", n_slices=10, n_paths=20, sigma=0.3, seed=0)
+        sums.append(partial)
+    fig, ax = plt.subplots()
+    _, ax_given = plot_phasor_convergence(sums, hbars, ax=ax)
+    assert ax_given is ax

physicskit/semiclassical/tests/test_wkb.py

@@ -0,0 +1,30 @@
+"""Physics-correctness tests for physicskit.semiclassical.core.wkb."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from physicskit.semiclassical.core.wkb import bohr_sommerfeld_energies, turning_points, wkb_wavefunction
+
+
+def test_bohr_sommerfeld_matches_harmonic_oscillator_exact_spectrum():
+    V = lambda x: 0.5 * x**2
+    energies = bohr_sommerfeld_energies(V, m=1.0, x_min=-20, x_max=20, n_max=6)
+    assert np.allclose(energies, np.arange(6) + 0.5, atol=1e-8)
+
+
+def test_turning_points_finds_root_landing_exactly_on_a_search_grid_point():
+    # n_search=3 puts a grid point exactly at x=0, where E - V(x) = 0 - 0 = 0.
+    roots = turning_points(E=0.0, V=lambda x: x, x_min=-1.0, x_max=1.0, n_search=3)
+    assert roots == [0.0]
+
+
+def test_wkb_wavefunction_has_n_nodes():
+    V = lambda x: 0.5 * x**2
+    for n in range(4):
+        E_n = bohr_sommerfeld_energies(V, m=1.0, x_min=-20, x_max=20, n_max=n + 1)[n]
+        x = np.linspace(-10, 10, 4000)
+        psi = wkb_wavefunction(x, E_n, V, m=1.0)
+        nonzero = psi[psi != 0]
+        n_nodes = np.sum(np.diff(np.sign(nonzero)) != 0)
+        assert n_nodes == n

physicskit/semiclassical/visualizers/__init__.py

@@ -0,0 +1,17 @@
+from physicskit.semiclassical.visualizers.gutzwiller import plot_density_of_states
+from physicskit.semiclassical.visualizers.propagators import plot_classical_trajectory_on_wigner
+from physicskit.semiclassical.visualizers.scarring import (
+    plot_husimi_1d,
+    plot_scar_map,
+    plot_scar_map_interactive,
+)
+from physicskit.semiclassical.visualizers.wkb import plot_wkb_wavefunction
+
+__all__ = [
+    "plot_wkb_wavefunction",
+    "plot_classical_trajectory_on_wigner",
+    "plot_density_of_states",
+    "plot_scar_map",
+    "plot_husimi_1d",
+    "plot_scar_map_interactive",
+]

physicskit/semiclassical/visualizers/gutzwiller.py

@@ -0,0 +1,55 @@
+r"""Plotting helpers for the Gutzwiller trace-formula density of states."""
+
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+__all__ = ["plot_density_of_states"]
+
+
+def plot_density_of_states(E_grid: np.ndarray, dos: np.ndarray, exact_energies: np.ndarray | None = None, ax=None):
+    r"""Plot the Gutzwiller trace-formula density of states, with exact levels marked.
+
+    Parameters
+    ----------
+    E_grid : ndarray
+        Energy grid.
+    dos : ndarray
+        Density of states, e.g. from
+        :func:`physicskit.semiclassical.core.gutzwiller.gutzwiller_density_of_states`.
+    exact_energies : ndarray, optional
+        Exact (or Bohr-Sommerfeld) energy levels to mark with vertical
+        dashed lines.
+    ax : matplotlib.axes.Axes, optional
+        Axes to draw into; a new figure is created if omitted.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    ax : matplotlib.axes.Axes
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from physicskit.semiclassical.core.gutzwiller import gutzwiller_density_of_states
+    >>> from physicskit.semiclassical.core.wkb import bohr_sommerfeld_energies
+    >>> V = lambda x: 0.5 * x ** 2
+    >>> E_grid = np.linspace(0.2, 3.5, 200)
+    >>> dos = gutzwiller_density_of_states(E_grid, V, m=1.0, x_min=-20, x_max=20)
+    >>> energies = bohr_sommerfeld_energies(V, m=1.0, x_min=-20, x_max=20, n_max=3)
+    >>> fig, ax = plot_density_of_states(E_grid, dos, exact_energies=energies)
+    >>> isinstance(fig, plt.Figure)
+    True
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.figure
+    ax.plot(E_grid, dos, color="C0")
+    if exact_energies is not None:
+        for E_n in exact_energies:
+            ax.axvline(E_n, color="C3", linestyle="--", linewidth=0.8)
+    ax.set_xlabel("E")
+    ax.set_ylabel("g(E)")
+    return fig, ax