wkb
wkb.py — Multidimensional WKB approximation with caustic corrections
Overview
The wkb module provides a comprehensive implementation of the Wentzel–Kramers–Brillouin (WKB) method for constructing asymptotic solutions to linear partial differential equations of the form
P(x, –iε∇) u(x) = 0, ε → 0,
where P(x,ξ) is the (pseudo‑differential) symbol of the operator. The solution is sought as a sum over rays (bicharacteristics):
u(x) ≈ ∑_j A_j(x) e^{i S_j(x)/ε},
with the phase S satisfying the eikonal equation P(x,∇S)=0 and the amplitudes A_j determined by transport equations along the rays.
Key features:
Automatic dimension detection – works for both 1D and 2D problems without changing the calling interface.
Ray tracing – integrates Hamilton’s equations for the bicharacteristics, including the variational equations for the stability matrix J (used to detect caustics).
Multi‑order amplitude transport – computes amplitudes up to arbitrary order (0,1,2,3,…) by solving ODEs along each ray.
- Caustic detection and correction – monitors det(J) to locate caustics (folds, cusps) and applies:
Maslov phase shifts (π/2 per caustic),
Uniform Airy (fold) or Pearcey (cusp) corrections near caustics.
Interpolation onto regular grids – uses scipy.interpolate (linear, griddata) to map the ray‑based solution to a uniform spatial grid.
Rich visualisation suite – phase portraits, amplitude decomposition, caustic analysis, ray plots, comparison of WKB orders.
Utilities for generating initial data – line segments, circles, point sources.
Mathematical background
The WKB ansatz u = e^{iS/ε} (a₀ + ε a₁ + ε² a₂ + …) is inserted into the equation P(x,–iε∇)u = 0. Expanding in powers of ε yields:
Order ε⁰ (eikonal equation): P(x, ∇S) = 0. This is a Hamilton–Jacobi equation solved by the method of characteristics (rays). The rays satisfy
dx/dt = ∂P/∂ξ, dξ/dt = –∂P/∂x, dS/dt = ξ·∂P/∂ξ – P.
Order ε¹ (transport equation for a₀): (∂P/∂ξ)·∇a₀ + ½ (∇_ξ·∇_x P) a₀ = 0. Along a ray this becomes an ODE for the amplitude.
Higher orders: Similar ODEs for a₁, a₂, … involving derivatives of P up to order three.
The stability matrix J = ∂x(t)/∂q (where q parametrises the initial data) measures the focusing of nearby rays. Its determinant vanishes at caustics. Near a fold caustic the standard WKB amplitude blows up; the correct uniform approximation involves the Airy function. Near a cusp, the Pearcey integral is required. The module implements both corrections using the companion caustics module.
References
- wkb.compare_orders(symbol, initial_phase, max_order=3, **kwargs)[source]
Compare WKB approximations at different orders. Works for both 1D and 2D automatically.
- wkb.create_initial_data_circle(radius=1.0, n_points=30, outward=True)[source]
Create initial data for WKB on a circle.
- Parameters:
radius (float) – Radius of the circle.
n_points (int) – Number of points on the circle.
outward (bool) – If True, rays point outward; if False, inward.
- Returns:
Initial data for wkb_approximation.
- Return type:
dict
Examples
>>> # Circle with outward rays >>> ic = create_initial_data_circle(radius=1.0, n_points=30, outward=True)
- wkb.create_initial_data_line(x_range, n_points=20, direction=(1, 0), y_intercept=0.0)[source]
Create initial data for WKB on a line segment.
- Parameters:
x_range (tuple) – Range (x_min, x_max) for the line segment.
n_points (int) – Number of points on the line.
direction (tuple) – Direction of rays (ξ₀, η₀).
y_intercept (float) – y-coordinate of the line.
- Returns:
Initial data for wkb_approximation.
- Return type:
dict
Examples
>>> # Horizontal line with rays going upward >>> ic = create_initial_data_line((-1, 1), n_points=20, ... direction=(0, 1), y_intercept=0)
- wkb.create_initial_data_point_source(x0=0.0, y0=0.0, n_rays=20)[source]
Create initial data for WKB from a point source.
- Parameters:
x0 (float) – Source location.
y0 (float) – Source location.
n_rays (int) – Number of rays emanating from source.
- Returns:
Initial data for wkb_approximation.
- Return type:
dict
Examples
>>> # Point source at origin >>> ic = create_initial_data_point_source(0, 0, n_rays=24)
- wkb.plot_amplitude_decomposition(solution)[source]
Plot individual amplitude orders aₖ and their contributions.
- wkb.plot_phase_space(solution, time_slice=None)[source]
Plot phase space (position-momentum) trajectories.
- Parameters:
solution (dict) – Output from wkb_approximation
time_slice (float or None) – Time at which to sample (None = final time)
- wkb.plot_with_caustics(solution, component='abs', highlight_caustics=True)[source]
Plot WKB solution with caustics highlighted.
- Parameters:
solution (dict) – Output of wkb_approximation()
component ({'abs','real','imag','phase'}) – Which component of u to visualize.
highlight_caustics (bool) – Whether to mark caustic locations.
- wkb.visualize_wkb_rays(wkb_result, plot_type='phase', n_rays_plot=None)[source]
Visualize WKB solution with rays.
- Parameters:
wkb_result (dict) – Output from wkb_approximation.
plot_type (str) – What to visualize: ‘phase’, ‘amplitude’, ‘real’, ‘rays’.
n_rays_plot (int, optional) – Number of rays to plot (if None, plot all).
Examples
>>> wkb = wkb_approximation(...) >>> visualize_wkb_rays(wkb, plot_type='phase')
- wkb.wkb_approximation(symbol, initial_phase, order=1, domain=None, resolution=50, epsilon=0.1, dimension=None, caustic_correction='auto', caustic_threshold=0.001)[source]
Compute multidimensional WKB approximation (1D or 2D).
u(x) ≈ exp(iS/ε) · [a₀ + ε·a₁ + ε²·a₂ + …]
Automatically detects dimension from initial_phase or uses dimension parameter.
- Parameters:
symbol (sympy expression) – Principal symbol p(x, ξ) for 1D or p(x, y, ξ, η) for 2D.
initial_phase (dict) –
Initial data on a curve/point:
1D: Keys ‘x’, ‘S’, ‘p_x’, optionally ‘a’ (dict or array) 2D: Keys ‘x’, ‘y’, ‘S’, ‘p_x’, ‘p_y’, optionally ‘a’ (dict or array)
order (int) – WKB order (0, 1, 2, or 3).
domain (tuple or None) – 1D: (x_min, x_max) 2D: ((x_min, x_max), (y_min, y_max)) If None, inferred from initial data.
resolution (int or tuple) – Grid resolution (single int or (nx, ny) for 2D).
epsilon (float) – Small parameter for asymptotic expansion.
dimension (int or None) – Force dimension (1 or 2). If None, auto-detect.
- Returns:
WKB solution with keys adapted to dimension.
- Return type:
dict
Examples
>>> from sympy import symbols, sqrt >>> >>> # 1D harmonic oscillator >>> x, xi = symbols('x xi', real=True) >>> symbol = xi**2 + x**2 # p(x,ξ) = ξ² + x² >>> >>> # Initial conditions: Gaussian wave packet >>> n_rays = 20 >>> x0 = np.linspace(-2, 2, n_rays) >>> initial = { ... 'x': x0, ... 'p_x': np.ones(n_rays), # momentum = 1 ... 'S': 0.5 * x0**2, # phase ... 'a': np.exp(-x0**2) # Gaussian amplitude ... } >>> >>> result = wkb_approximation( ... symbol, initial, order=2, epsilon=0.05, resolution=100 ... ) >>> >>> # Extract solution >>> x_grid = result['x'] >>> u = result['u'] >>> plt.plot(x_grid, np.abs(u))
>>> # 2D wave equation >>> x, y, xi, eta = symbols('x y xi eta', real=True) >>> c = 1.0 # wave speed >>> symbol = xi**2 + eta**2 - c**2 >>> >>> # Point source initial conditions >>> theta = np.linspace(0, 2*np.pi, 30, endpoint=False) >>> r0 = 0.5 >>> initial = { ... 'x': r0 * np.cos(theta), ... 'y': r0 * np.sin(theta), ... 'p_x': np.cos(theta), ... 'p_y': np.sin(theta), ... 'S': np.zeros(30), ... 'a': np.ones(30) ... } >>> >>> result = wkb_approximation( ... symbol, initial, order=1, epsilon=0.1, resolution=(80, 80) ... ) >>> >>> # Visualize 2D solution >>> X, Y = result['x'], result['y'] >>> U = result['u'] >>> plt.pcolormesh(X, Y, np.abs(U))