physicskit 0.1.0__py3-none-any.whl

physicskit/optics/tests/test_quantum_optics.py

@@ -0,0 +1,147 @@
+import numpy as np
+import pytest
+
+from physicskit.optics.quantum_optics import (
+    JaynesCummingsModel,
+    coherent_state,
+    compute_wigner_function,
+    fock_state,
+    squeezed_state,
+    wigner_negativity,
+)
+
+
+def test_fock_state_is_normalized_and_correctly_placed():
+    psi = fock_state(3, 8)
+    assert np.linalg.norm(psi) == pytest.approx(1.0)
+    assert psi[3] == pytest.approx(1.0)
+    assert np.sum(np.abs(psi) ** 2) - np.abs(psi[3]) ** 2 == pytest.approx(0.0)
+
+
+def test_fock_state_raises_when_n_exceeds_cutoff():
+    with pytest.raises(ValueError):
+        fock_state(5, 5)
+
+
+def test_coherent_state_is_normalized():
+    psi = coherent_state(1.5 + 0.5j, cutoff=40)
+    assert np.linalg.norm(psi) == pytest.approx(1.0, rel=1e-6)
+
+
+def test_coherent_state_of_zero_amplitude_is_the_vacuum_fock_state():
+    psi = coherent_state(0.0, cutoff=8)
+    np.testing.assert_allclose(psi, fock_state(0, 8))
+
+
+def test_coherent_state_mean_photon_number_matches_alpha_squared():
+    alpha = 1.2
+    cutoff = 40
+    psi = coherent_state(alpha, cutoff)
+    n = np.arange(cutoff)
+    mean_n = np.sum(n * np.abs(psi) ** 2)
+    assert mean_n == pytest.approx(alpha**2, rel=1e-2)
+
+
+def test_squeezed_state_is_normalized():
+    psi = squeezed_state(xi=0.5 + 0.2j, alpha=0.3, cutoff=30)
+    assert np.linalg.norm(psi) == pytest.approx(1.0, rel=1e-6)
+
+
+def test_squeezed_state_reduces_to_vacuum_when_xi_and_alpha_zero():
+    psi = squeezed_state(xi=0.0, alpha=0.0, cutoff=10)
+    expected = fock_state(0, 10)
+    assert np.allclose(psi, expected)
+
+
+def test_squeezed_state_reduces_to_coherent_state_when_xi_zero():
+    alpha = 0.8 - 0.3j
+    cutoff = 30
+    psi = squeezed_state(xi=0.0, alpha=alpha, cutoff=cutoff)
+    expected = coherent_state(alpha, cutoff)
+    assert np.allclose(psi, expected, atol=1e-6)
+
+
+def test_wigner_function_integrates_to_one_for_coherent_state():
+    cutoff = 25
+    psi = coherent_state(1.0 + 0.5j, cutoff)
+    xg = np.linspace(-6, 6, 121)
+    pg = np.linspace(-6, 6, 121)
+    W = compute_wigner_function(psi, xg, pg)
+    integral = np.trapezoid(np.trapezoid(W, pg, axis=1), xg)
+    assert integral == pytest.approx(1.0, rel=0.03)
+
+
+def test_wigner_function_integrates_to_one_for_fock_state_one():
+    cutoff = 15
+    psi = fock_state(1, cutoff)
+    xg = np.linspace(-6, 6, 121)
+    pg = np.linspace(-6, 6, 121)
+    W = compute_wigner_function(psi, xg, pg)
+    integral = np.trapezoid(np.trapezoid(W, pg, axis=1), xg)
+    assert integral == pytest.approx(1.0, rel=0.03)
+    assert W.shape == (len(xg), len(pg))
+
+
+def test_wigner_negativity_zero_for_coherent_state_positive_for_fock_one():
+    xg = np.linspace(-6, 6, 121)
+    pg = np.linspace(-6, 6, 121)
+
+    psi_coh = coherent_state(1.0, 25)
+    W_coh = compute_wigner_function(psi_coh, xg, pg)
+    neg_coh = wigner_negativity(W_coh, xg, pg)
+    assert neg_coh == pytest.approx(0.0, abs=1e-9)
+
+    psi_fock1 = fock_state(1, 15)
+    W_fock1 = compute_wigner_function(psi_fock1, xg, pg)
+    neg_fock1 = wigner_negativity(W_fock1, xg, pg)
+    assert neg_fock1 > 0.0
+
+
+def test_fock_state_one_wigner_function_is_negative_at_origin():
+    cutoff = 10
+    psi = fock_state(1, cutoff)
+    xg = np.linspace(-4, 4, 81)
+    pg = np.linspace(-4, 4, 81)
+    W = compute_wigner_function(psi, xg, pg)
+    origin = (np.abs(xg).argmin(), np.abs(pg).argmin())
+    assert W[origin] == pytest.approx(-1.0 / np.pi, rel=1e-3)
+
+
+def test_jaynes_cummings_hamiltonian_is_hermitian():
+    jc = JaynesCummingsModel(omega_c=1.0, omega_a=1.2, g=0.3, cutoff=6)
+    H = jc.hamiltonian()
+    assert H.shape == (12, 12)
+    assert np.allclose(H, H.conj().T)
+
+
+def test_resonant_vacuum_rabi_oscillation_matches_cos_squared():
+    omega0 = 2.0
+    g = 0.4
+    cutoff = 12
+    jc = JaynesCummingsModel(omega_c=omega0, omega_a=omega0, g=g, cutoff=cutoff)
+
+    t_array = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0])
+    Pe = jc.excited_state_population(t_array, n_photons=0)
+    expected = np.cos(g * t_array) ** 2
+
+    assert np.allclose(Pe, expected, atol=1e-3)
+
+
+def test_excited_state_population_conserves_probability():
+    jc = JaynesCummingsModel(omega_c=1.0, omega_a=1.5, g=0.2, cutoff=8)
+    psi0 = fock_state(0, 2 * jc.cutoff)  # |e, 0>
+    t_array = np.linspace(0.0, 20.0, 15)
+    psi_t = jc.evolve(psi0, t_array)
+    norms = np.sum(np.abs(psi_t) ** 2, axis=1)
+    assert np.allclose(norms, 1.0, atol=1e-8)
+
+
+def test_off_resonant_jaynes_cummings_reduces_rabi_amplitude():
+    # Far off resonance, the atom barely exchanges excitation with the
+    # cavity: the excited-state population should stay close to 1.
+    g = 0.1
+    cutoff = 10
+    jc = JaynesCummingsModel(omega_c=1.0, omega_a=1.0 + 50.0 * g, g=g, cutoff=cutoff)
+    t_array = np.linspace(0.0, 20.0, 50)
+    Pe = jc.excited_state_population(t_array, n_photons=0)
+    assert Pe.min() > 0.9

physicskit/optics/tests/test_ray.py

@@ -0,0 +1,177 @@
+import numpy as np
+import pytest
+
+from physicskit.optics.ray import (
+    OpticalElement,
+    OpticalSystem,
+    cavity_round_trip_matrix,
+    cavity_stability,
+    curved_interface,
+    flat_interface,
+    free_space,
+    grin_medium,
+    spherical_mirror,
+    thick_lens,
+    thin_lens,
+)
+
+
+def test_free_space_matrix():
+    M = free_space(3.0)
+    assert np.allclose(M, [[1.0, 3.0], [0.0, 1.0]])
+
+
+def test_thin_lens_matrix():
+    M = thin_lens(0.25)
+    assert np.allclose(M, [[1.0, 0.0], [-4.0, 1.0]])
+
+
+def test_flat_interface_matrix():
+    M = flat_interface(1.0, 1.5)
+    assert np.allclose(M, [[1.0, 0.0], [0.0, 1.0 / 1.5]])
+
+
+def test_curved_interface_reduces_to_flat_interface_for_large_radius():
+    M_curved = curved_interface(R=1.0e10, n1=1.0, n2=1.5)
+    M_flat = flat_interface(1.0, 1.5)
+    assert np.allclose(M_curved, M_flat, atol=1e-8)
+
+
+def test_spherical_mirror_matrix():
+    M = spherical_mirror(4.0)
+    assert np.allclose(M, [[1.0, 0.0], [-0.5, 1.0]])
+
+
+def test_thick_lens_composition_order_and_unit_determinant():
+    R1, R2, t, n = 0.1, -0.1, 0.01, 1.5
+    M = thick_lens(R1, R2, t, n)
+    expected = curved_interface(R2, n, 1.0) @ free_space(t) @ curved_interface(R1, 1.0, n)
+    assert np.allclose(M, expected)
+    # a lossless system's ABCD matrix always has unit determinant
+    assert np.linalg.det(M) == pytest.approx(1.0, abs=1e-10)
+
+
+def test_thin_lens_focuses_parallel_ray_bundle_to_a_point():
+    # A bundle of rays parallel to the axis (theta0=0), of any height y0,
+    # converges to the axis after propagating one focal length past the lens.
+    f = 0.2
+    M = free_space(f) @ thin_lens(f)
+    for y0 in (-0.05, -0.01, 0.0, 0.01, 0.05, 1.0):
+        y_out, theta_out = M @ np.array([y0, 0.0])
+        assert y_out == pytest.approx(0.0, abs=1e-12)
+        assert theta_out == pytest.approx(-y0 / f)
+
+
+def test_grin_medium_reduces_to_free_space_as_n2_coeff_vanishes():
+    d = 2.5
+    M_grin = grin_medium(n0=1.0, n2_coeff=1e-9, d=d)
+    M_free = free_space(d)
+    assert np.allclose(M_grin, M_free, atol=1e-6)
+
+
+def test_grin_medium_zero_n2_coeff_is_exactly_homogeneous_propagation():
+    n0, d = 1.3, 2.5
+    assert np.array_equal(grin_medium(n0=n0, n2_coeff=0.0, d=d), np.array([[1.0, d / n0], [0.0, 1.0]]))
+
+
+def test_grin_medium_matches_closed_form_for_finite_gradient():
+    n0, n2_coeff, d = 1.5, 4.0, 0.3
+    M = grin_medium(n0, n2_coeff, d)
+    sqrt_n2 = np.sqrt(n2_coeff)
+    arg = sqrt_n2 * d
+    expected = np.array(
+        [
+            [np.cos(arg), np.sin(arg) / (n0 * sqrt_n2)],
+            [-n0 * sqrt_n2 * np.sin(arg), np.cos(arg)],
+        ]
+    )
+    assert np.allclose(M, expected)
+
+
+def test_optical_element_stores_attributes():
+    M = thin_lens(0.5)
+    elem = OpticalElement(M, name="my lens", length=0.0)
+    assert elem.name == "my lens"
+    assert elem.length == 0.0
+    assert np.allclose(elem.matrix, M)
+
+
+def test_optical_system_matrix_multiplication_order():
+    # elements[0] is hit first, so it must be the RIGHTMOST factor:
+    # M = M2 @ M1, not M1 @ M2 (these differ since matrices don't commute).
+    m1 = thin_lens(0.3)
+    m2 = free_space(1.0)
+    system = OpticalSystem([OpticalElement(m1, name="lens"), OpticalElement(m2, name="gap")])
+    expected = m2 @ m1
+    assert np.allclose(system.system_matrix(), expected)
+    assert not np.allclose(system.system_matrix(), m1 @ m2)
+
+
+def test_optical_system_trace_ray_is_cumulative():
+    f = 0.5
+    system = OpticalSystem(
+        [
+            OpticalElement(thin_lens(f), name="lens"),
+            OpticalElement(free_space(f), name="gap"),
+        ]
+    )
+    y0, theta0 = 0.02, 0.0
+    trajectory = system.trace_ray(y0, theta0)
+    assert trajectory.shape == (3, 2)
+    assert np.allclose(trajectory[0], [y0, theta0])
+    assert np.allclose(trajectory[1], thin_lens(f) @ [y0, theta0])
+    assert np.allclose(trajectory[2], free_space(f) @ thin_lens(f) @ [y0, theta0])
+    # the whole bundle refocuses to the axis after one focal length
+    assert trajectory[-1, 0] == pytest.approx(0.0, abs=1e-12)
+
+
+def test_stability_agrees_with_cavity_stability_for_stable_cavity():
+    # Two R=2.0 mirrors separated by d=1.0: g = 1 - d/R = 0.5, g1*g2 = 0.25,
+    # which lies in [0, 1] -- a stable resonator by the standard g-parameter
+    # criterion (Kogelnik & Li 1966).
+    R, d = 2.0, 1.0
+    elements = [
+        OpticalElement(spherical_mirror(R), name="M1"),
+        OpticalElement(free_space(d), name="gap1"),
+        OpticalElement(spherical_mirror(R), name="M2"),
+        OpticalElement(free_space(d), name="gap2"),
+    ]
+    system = OpticalSystem(elements)
+    M = system.system_matrix()
+    assert M is not None
+    assert system.stability_parameter == pytest.approx((M[0, 0] + M[1, 1]) / 2.0)
+    assert system.is_stable() is True
+    assert cavity_stability(cavity_round_trip_matrix(elements)) is True
+    assert cavity_stability(system.system_matrix()) == system.is_stable()
+
+
+def test_stability_agrees_with_cavity_stability_for_unstable_cavity():
+    # Two R=0.5 mirrors separated by d=2.0: g = 1 - d/R = -3, g1*g2 = 9,
+    # well outside [0, 1] -- an unstable resonator.
+    R, d = 0.5, 2.0
+    elements = [
+        OpticalElement(spherical_mirror(R), name="M1"),
+        OpticalElement(free_space(d), name="gap1"),
+        OpticalElement(spherical_mirror(R), name="M2"),
+        OpticalElement(free_space(d), name="gap2"),
+    ]
+    system = OpticalSystem(elements)
+    assert system.is_stable() is False
+    assert cavity_stability(cavity_round_trip_matrix(elements)) is False
+    assert cavity_stability(system.system_matrix()) == system.is_stable()
+
+
+def test_cavity_round_trip_matrix_matches_optical_system():
+    elements = [
+        OpticalElement(spherical_mirror(1.0), name="M1"),
+        OpticalElement(free_space(0.4), name="gap1"),
+        OpticalElement(spherical_mirror(1.5), name="M2"),
+        OpticalElement(free_space(0.4), name="gap2"),
+    ]
+    assert np.allclose(cavity_round_trip_matrix(elements), OpticalSystem(elements).system_matrix())
+
+
+def test_cavity_stability_boundary_cases():
+    assert cavity_stability(np.array([[1.0, 0.0], [0.0, 1.0]])) is True  # trace=2
+    assert cavity_stability(np.array([[-1.0, 0.0], [0.0, -1.0]])) is True  # trace=-2
+    assert cavity_stability(np.array([[2.0, 0.0], [0.0, 1.01]])) is False  # trace=3.01

physicskit/optics/tests/test_visualizers.py

@@ -0,0 +1,58 @@
+"""Fast coverage of physicskit.optics.visualizers: the ax=None/ax=<given>
+branch of every function, plus the log_scale=True branches of
+plot_diffraction_pattern and animate_diffraction_propagation. See
+test_diffraction_animation.py for the real anim.save()-to-disk smoke test
+(marked slow)."""
+
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from physicskit.optics.gaussian import GaussianBeam
+from physicskit.optics.ray import OpticalElement, OpticalSystem, free_space, thin_lens
+from physicskit.optics.visualizers import animate_diffraction_propagation, plot_beam_envelope, plot_diffraction_pattern, plot_ray_trace
+from physicskit.optics.wave import double_slit_aperture
+
+
+def test_plot_ray_trace_given_ax():
+    system = OpticalSystem(
+        [
+            OpticalElement(free_space(1.0), name="d1", length=1.0),
+            OpticalElement(thin_lens(1.0), name="lens"),
+        ]
+    )
+    fig, ax = plot_ray_trace(system, y0=0.5, theta0=0.0)
+    assert isinstance(fig, plt.Figure)
+    _, ax_given = plot_ray_trace(system, y0=0.5, theta0=0.0, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_beam_envelope_given_ax():
+    beam = GaussianBeam(wavelength=0.5e-3, w0=0.1, z0=0.0)
+    fig, ax = plot_beam_envelope(beam, z_range=(-5, 5), n_points=20)
+    assert isinstance(fig, plt.Figure)
+    _, ax_given = plot_beam_envelope(beam, z_range=(-5, 5), n_points=20, ax=ax)
+    assert ax_given is ax
+
+
+def test_plot_diffraction_pattern_log_scale_and_given_ax():
+    dx = 1e-3
+    U = double_slit_aperture((16, 16), dx=dx, width=2e-3, separation=1e-2).astype(complex)
+    fig, ax = plot_diffraction_pattern(U, dx=dx, log_scale=True)
+    assert isinstance(fig, plt.Figure)
+    assert ax.get_ylabel() == "y"
+    _, ax_given = plot_diffraction_pattern(U, dx=dx, ax=ax)
+    assert ax_given is ax
+
+
+def test_animate_diffraction_propagation_log_scale_and_given_ax_builds_and_draws_a_frame():
+    dx = 1e-3
+    ap = double_slit_aperture((16, 16), dx=dx, width=2e-3, separation=1e-2).astype(complex)
+    z_values = np.linspace(0.01, 1.0, 3)
+
+    fig, ax = plt.subplots()
+    anim = animate_diffraction_propagation(ap, wavelength=0.5e-3, z_values=z_values, dx=dx, log_scale=True, ax=ax)
+    assert anim._fig is fig
+    anim._draw_frame(1)
+    assert ax.get_title() == f"z = {z_values[1]:.4g}"

physicskit/optics/tests/test_wave.py

@@ -0,0 +1,146 @@
+import numpy as np
+import pytest
+
+from physicskit.optics.wave import (
+    angular_spectrum_propagate,
+    circular_aperture,
+    double_slit_aperture,
+    field_grid,
+    fraunhofer_diffraction,
+    fresnel_diffraction,
+    intensity,
+    single_slit_aperture,
+)
+
+
+def test_field_grid_is_centered():
+    X, Y = field_grid((4, 6), dx=0.5)
+    assert X.shape == (4, 6)
+    assert Y.shape == (4, 6)
+    assert X[0, 3] == pytest.approx(0.0)
+    assert Y[2, 0] == pytest.approx(0.0)
+
+
+def test_circular_aperture_is_open_at_center_and_closed_at_corner():
+    ap = circular_aperture((65, 65), dx=1.0, radius=10.0)
+    assert ap[32, 32] == pytest.approx(1.0)
+    assert ap[0, 0] == pytest.approx(0.0)
+    assert set(np.unique(ap)) <= {0.0, 1.0}
+
+
+def test_single_slit_aperture_width_and_full_height():
+    shape = (10, 101)
+    dx = 1.0
+    width = 20.0
+    ap = single_slit_aperture(shape, dx, width)
+    # every row should be identical (open along full height)
+    assert np.all(ap == ap[0])
+    n_open = int(ap[0].sum())
+    assert n_open == pytest.approx(width / dx, abs=1)
+
+
+def test_double_slit_aperture_has_two_separated_openings():
+    shape = (5, 201)
+    dx = 1.0
+    width, separation = 8.0, 60.0
+    ap = double_slit_aperture(shape, dx, width, separation)
+    row = ap[0]
+    X, _ = field_grid(shape, dx)
+    x = X[0]
+    # sample the intensity mask around each expected slit center and the gap between them
+    left_open = row[np.argmin(np.abs(x - (-separation / 2)))]
+    right_open = row[np.argmin(np.abs(x - (separation / 2)))]
+    middle_closed = row[np.argmin(np.abs(x - 0.0))]
+    assert left_open == pytest.approx(1.0)
+    assert right_open == pytest.approx(1.0)
+    assert middle_closed == pytest.approx(0.0)
+
+
+def test_intensity_is_squared_modulus():
+    U = np.array([3.0 + 4.0j, 1.0 + 0.0j, 0.0 + 2.0j])
+    assert np.allclose(intensity(U), [25.0, 1.0, 4.0])
+
+
+def test_angular_spectrum_propagate_conserves_power_over_short_distance():
+    # A circular aperture propagated a short distance, with fine enough
+    # sampling (dx well below wavelength/2, so no aperture spatial
+    # frequency content is evanescent) should conserve total power to
+    # within a small grid/aliasing tolerance.
+    wavelength = 0.5e-3  # e.g. millimeters, so this is 500 nm
+    dx = 0.2e-3
+    N = 256
+    radius = 20 * dx
+    z = 2.0
+
+    aperture = circular_aperture((N, N), dx, radius)
+    U_out = angular_spectrum_propagate(aperture, wavelength, z, dx)
+
+    power_in = intensity(aperture).sum()
+    power_out = intensity(U_out).sum()
+    assert power_out == pytest.approx(power_in, rel=0.02)
+
+
+def test_angular_spectrum_propagate_matches_input_grid_shape():
+    ap = circular_aperture((32, 48), dx=1.0e-3, radius=5.0e-3)
+    U = angular_spectrum_propagate(ap, wavelength=0.6e-3, z=1.0, dx=1.0e-3)
+    assert U.shape == ap.shape
+    assert np.iscomplexobj(U)
+
+
+def test_fresnel_diffraction_matches_input_grid_shape():
+    ap = circular_aperture((64, 64), dx=0.01, radius=0.05)
+    U = fresnel_diffraction(ap, wavelength=0.5e-3, z=5.0, dx=0.01)
+    assert U.shape == ap.shape
+    assert np.iscomplexobj(U)
+
+
+def test_fresnel_diffraction_conserves_power_over_short_distance():
+    # The single-FFT Fresnel method resamples onto an output grid with a
+    # different pixel spacing (dx_out = wavelength*z/(N*dx_in)), so total
+    # power must be compared as intensity-sum times each grid's own pixel
+    # area, not a bare intensity-sum comparison (which implicitly assumes
+    # equal pixel areas, valid for angular_spectrum_propagate but not here).
+    wavelength = 0.5e-3
+    dx = 0.2e-3
+    N = 256
+    radius = 20 * dx
+    z = 2.0
+
+    aperture = circular_aperture((N, N), dx, radius)
+    U_out = fresnel_diffraction(aperture, wavelength, z, dx)
+    dx_out = wavelength * z / (N * dx)
+
+    power_in = intensity(aperture).sum() * dx**2
+    power_out = intensity(U_out).sum() * dx_out**2
+    assert power_out == pytest.approx(power_in, rel=0.05)
+
+
+def test_fraunhofer_single_slit_first_minimum_matches_analytic_formula():
+    # Standard single-slit diffraction: the first intensity zero occurs at
+    # sin(theta) = wavelength / width. We locate the first minimum of the
+    # Fraunhofer pattern's central row and convert its transverse position
+    # x' at distance z to an angle via the small-angle relation
+    # sin(theta) ~= x' / z (valid here since x' << z in the far field).
+    wavelength = 0.5e-3  # mm
+    width = 0.05  # mm
+    dx = 0.002  # mm
+    N = 512
+    z = 500.0  # mm; far enough for the Fraunhofer approximation to hold
+
+    aperture = single_slit_aperture((N, N), dx, width)
+    U = fraunhofer_diffraction(aperture, wavelength, z, dx)
+    row = intensity(U)[N // 2]
+
+    fx = np.fft.fftshift(np.fft.fftfreq(N, d=dx))
+    x_prime = wavelength * z * fx
+
+    # walk outward from the central peak until intensity starts rising again
+    center = N // 2
+    i = center
+    while row[i + 1] < row[i]:
+        i += 1
+    x_min = x_prime[i]
+
+    sin_theta_measured = x_min / z
+    sin_theta_predicted = wavelength / width
+    assert sin_theta_measured == pytest.approx(sin_theta_predicted, rel=0.05)