fio_bridge

fio_bridge.py — Fourier Integral Operators: bridge between psiop.py and asymptotic.py

This module is the missing link between:

psiop.py – symbol algebra: asymptotic composition (KN/Weyl), adjoint,
inverse, operator exponential exp(tP), Hamiltonian flow.
asymptotic.py – numerical evaluation of oscillatory integrals via the
stationary-phase / Laplace / saddle-point method.

Architecture

FourierIntegralOperator (base class)

Generic FIO with a free phase φ(x, y, θ) in the sense of Hörmander. Owns the canonical-relation geometry and the non-degeneracy check. Delegates ALL critical-point work to asymptotic.Analyzer and asymptotic.AsymptoticEvaluator (DRY: no duplication of gradient minimisation, Hessian analysis, or asymptotic formulæ).

PsiOpFIOBridge(FourierIntegralOperator)

Specialisation for a PseudoDifferentialOperator. Auto-builds the standard FIO phase

φ(y, ξ; x) = (x − y)·ξ + S_u(y)

and amplitude a(y, ξ; x) = p(y, ξ) · a_u(y) from the operator symbol, then inherits the evaluation machinery.

PropagatorBridge

Computes the semi-classical propagator e^{itP} via PseudoDifferentialOperator.exponential_symbol(), then delegates to PsiOpFIOBridge.

CompositionBridge

Computes P∘Q via PseudoDifferentialOperator.compose_asymptotic(), then delegates to PsiOpFIOBridge.

Mathematics

For a psiOp with symbol p(x, ξ), the action on a WKB state: u(y) ~ a_u(y) · exp(iλ S_u(y))

is the oscillatory integral

(Pu)(x) = (1/2π)^n ∫∫ exp(iλ φ(y, ξ; x)) · p(y, ξ) · a_u(y) dy dξ

where the global phase is

φ(y, ξ; x) = (x − y)·ξ + S_u(y).

Stationary conditions:: ∂φ/∂ξ = 0 → y_c = x (transport) ∂φ/∂y = 0 → ξ_c = S’_u(y_c) (bicharacteristic)

Normalisation

The prefactor produced by asymptotic.py at a Morse point is

(2π/λ)^(n/2) / √|det H| · exp(iλφ_c) · a_c · exp(iπμ/4)

The FIO integral carries an overall factor (1/2π)^n (one per integration variable pair), so the caller multiplies by 1 / (2π)^n.

For the standard 1D psiOp (n=1, two integration variables y and ξ):: prefactor = 1 / (2π).

References

class fio_bridge.CompositionBridge(P: PseudoDifferentialOperator, Q: PseudoDifferentialOperator, lam: float = 50.0, comp_order: int = 2, mode: str = 'kn', **bridge_kwargs)[source]

Bases: object

Compute the composition P∘Q and evaluate its action on a WKB state.

The composed symbol is obtained via PseudoDifferentialOperator.compose_asymptotic(), then evaluation is delegated to PsiOpFIOBridge.

Parameters:

P (PseudoDifferentialOperator)
Q (PseudoDifferentialOperator)
lam (float)
comp_order (int) – Asymptotic composition order.
mode (str) – Quantization scheme (‘kn’ or ‘weyl’).
**bridge_kwargs – Forwarded to PsiOpFIOBridge.

evaluate_grid(x_grid: ndarray, u_phase_sym: Expr, u_amp_sym: Expr) → ndarray[source]: Evaluate ((P∘Q)u)(x) over the grid.

class fio_bridge.CrossValidator(op: PseudoDifferentialOperator, wkb_state: WKBState, x_grid: ndarray, lam: float, wkb_threshold: float | None = None, solver_kwargs: Dict | None = None, bridge_kwargs: Dict | None = None)[source]

Bases: object

Compare a PDESolver solution against a PsiOpFIOBridge solution on the same problem and produce a ValidationReport.

This class is the formal bridge between the two regimes: - It runs PDESolver numerically (ETD-RK4 or default scheme). - It runs PsiOpFIOBridge asymptotically (stationary-phase / saddle). - It computes point-wise and spectral error metrics. - It declares whether the WKB regime is valid for the given λ.

Parameters:

op (PseudoDifferentialOperator) – The operator P whose action on u is being compared.
wkb_state (WKBState) – The WKB initial/input state.
x_grid (np.ndarray) – Shared spatial grid for both solvers.
lam (float) – Large parameter λ. Must match wkb_state.lam.
wkb_threshold (float) – Max relative error below which the WKB regime is declared valid. Default: 3 / lam (theoretical O(λ⁻¹) accuracy).
solver_kwargs (dict | None) – Extra keyword arguments forwarded to PDESolver.setup().
bridge_kwargs (dict | None) – Extra keyword arguments forwarded to PsiOpFIOBridge.

Notes

PDESolver is imported lazily inside run() to keep solver.py optional. If solver.py is not on sys.path, CrossValidator.run() raises ImportError with a clear message.

lambda_sweep(lambdas: List[float]) → List[ValidationReport][source]

Run the cross-validation for a range of λ values and return the list of ValidationReports.

Useful for finding the λ-threshold below which the WKB approximation is no longer reliable for a given operator and initial state.

Parameters:: lambdas (list[float]) – Increasing sequence of λ values to test.
Return type:: list[ValidationReport], one per λ value.

static plot_lambda_sweep(reports: List[ValidationReport], lambdas: List[float]) → None[source]: Plot max relative error vs λ on a log-log scale, with a vertical line at the λ-threshold where WKB validity flips.

static plot_report(report: ValidationReport, title: str = 'Cross-validation: solver vs asymptotic bridge') → None[source]: Plot a ValidationReport: solution comparison, error profile, and error spectrum side by side.

run() → ValidationReport[source]

Execute both solvers and return a ValidationReport.

The PDESolver is used in stationary mode: it applies the operator P to the WKB state u₀ directly via solve_stationary_psiOp(), which calls psiop.right_inverse_asymptotic() internally. For time- dependent problems, set solver_kwargs['time_dependent'] = True and supply Lt, Nt (not yet implemented here — raises NotImplementedError).

Return type:: ValidationReport

run_bridge_only() → ndarray[source]: Run PsiOpFIOBridge only and return the solution array.

run_solver_only() → ndarray[source]: Run PDESolver only and return the solution array.

class fio_bridge.EvalResult(x_val: Any, value: complex, n_critical_points: int, contributions: List[AsymptoticContribution] = <factory>, warnings_list: List[str] = <factory>)[source]

Bases: object

Result of evaluating (Fu)(x) at a single observation point x.

x_val

Observation point.

Type:: float | tuple

value

Asymptotic value of (Fu)(x).

Type:: complex

n_critical_points

Number of critical (or saddle) points found and used.

Type:: int

contributions

Individual contributions from each critical point.

Type:: list[AsymptoticContribution]

warnings_list

Warnings raised during evaluation (caustics, Picard–Lefschetz, …).

Type:: list[str]

contributions: List[AsymptoticContribution]

n_critical_points: int

value: complex

warnings_list: List[str]

x_val: Any

class fio_bridge.FIOKernel(phase_sym: Expr, amp_sym: Expr, vars_int: List[Symbol], x_val: Any = 0.0, domain: List[Tuple] | None = None, method_hint: IntegralMethod = IntegralMethod.AUTO)[source]

Bases: object

Kernel of a Fourier Integral Operator: phase φ and amplitude a as SymPy expressions, together with the integration variables and the observation point.

phase_sym

Phase function φ(vars_int; x_val). Depends on the integration variables and on the (fixed) observation parameter x_val.

Type:: sp.Expr

amp_sym

Amplitude function a(vars_int; x_val).

Type:: sp.Expr

vars_int

Integration variables, e.g. [y, ξ] in 1D or [y1, y2, ξ1, ξ2] in 2D.

Type:: list[sp.Symbol]

x_val

Current observation point.

Type:: float | tuple[float, …]

domain

Search bounds [(min, max), …] for each integration variable. Passed directly to asymptotic.Analyzer so critical-point search respects the physical domain.

Type:: list[tuple[float, float]] | None

method_hint

Suggested method (resolved by Analyzer when AUTO).

Type:: IntegralMethod

amp_sym: Expr

domain: List[Tuple] | None = None

method_hint: IntegralMethod = 'auto'

phase_sym: Expr

vars_int: List[Symbol]

x_val: Any = 0.0

class fio_bridge.FourierIntegralOperator(phase_expr: Expr, amp_expr: Expr, vars_x: List[Symbol], vars_y: List[Symbol], vars_theta: List[Symbol], lam: float = 50.0, domain: List[Tuple] | None = None, tol_grad: float = 1e-06, verbose: bool = False)[source]

Bases: object

Generic Fourier Integral Operator with a free phase function.

Evaluates integrals of the form

F[u](x) = (2π)^{-n_θ} ∫∫ exp(iλ φ(x, y, θ)) · a(x, y, θ) · u(y) dy dθ

where φ is an arbitrary non-degenerate phase in the sense of Hörmander.

This class is responsible for: * Storing the symbolic phase and amplitude. * Computing the canonical relation (∇_θ φ, ∇_x φ, ∇_y φ). * Checking Hörmander’s non-degeneracy condition (mixed Hessian ∂²φ/∂x∂θ). * Evaluating the integral asymptotically by delegating entirely to

asymptotic.Analyzer and asymptotic.AsymptoticEvaluator (DRY).

Parameters:

phase_expr (sp.Expr) – Phase function φ(x, y, θ) as a SymPy expression.
amp_expr (sp.Expr) – Amplitude function a(x, y, θ).
vars_x (list[sp.Symbol]) – Target (observation) spatial variables.
vars_y (list[sp.Symbol]) – Source spatial variables (integration variables, spatial part).
vars_theta (list[sp.Symbol]) – Frequency / phase variables (integration variables, frequency part).
lam (float) – Large parameter λ.
domain (list[tuple[float, float]] | None) – Search domain for each integration variable (vars_y + vars_theta). Forwarded to Analyzer.
tol_grad (float) – Tolerance for |∇φ|² when accepting a critical point.
verbose (bool) – Print diagnostic information.

apply_asymptotic(u_amp_expr: Expr, u_phase_expr: Expr, x_eval_dict: Dict[Symbol, float], initial_guesses: List[ndarray] | None = None) → EvalResult[source]

Apply the FIO to a WKB state u(y) = u_amp(y) · exp(iλ u_phase(y)) at the observation point encoded by x_eval_dict.

The total integration phase is:: Φ(y, θ) = φ(x, y, θ) / λ + u_phase(y)

(dividing φ by λ so that the large-parameter rôle is played uniformly by λ via the WKB phase of u).

Parameters:

u_amp_expr (sp.Expr) – Amplitude a_u(y) of the input WKB state.
u_phase_expr (sp.Expr) – Phase S_u(y) of the input WKB state.
x_eval_dict (dict) – Numerical values for the observation variables x, e.g. {x: 1.5}.
initial_guesses (list[np.ndarray] | None) – Starting points for the critical-point search. If None, defaults to the origin.

Return type:

EvalResult

is_non_degenerate() → bool[source]

Check Hörmander’s non-degeneracy condition.

A phase φ is non-degenerate when the mixed Hessian

H_{θ,x} = [ ∂²φ / ∂θ_i ∂x_j ]_{i,j}

has maximal rank (= min(dim_θ, dim_x)).

Returns:: True if the phase is non-degenerate.
Return type:: bool

class fio_bridge.PropagatorBridge(op: PseudoDifferentialOperator, lam: float = 50.0, exp_order: int = 2, **bridge_kwargs)[source]

Bases: object

Compute the semi-classical propagator u(x, t) = [e^{itP} u_0](x).

Uses PseudoDifferentialOperator.exponential_symbol(t) to build the symbol of e^{itP} up to the requested asymptotic order, then delegates evaluation to PsiOpFIOBridge.

Parameters:

op (PseudoDifferentialOperator)
lam (float) – Large parameter λ.
exp_order (int) – Asymptotic order for exp(tP).
**bridge_kwargs – Forwarded to PsiOpFIOBridge.

propagate(t: float, x_grid: ndarray, u0_phase_sym: Expr, u0_amp_sym: Expr, mode: str = 'kn') → ndarray[source]

Compute u(x, t) = [e^{itP} u_0](x) over x_grid.

Parameters:

t (float Propagation time.)
x_grid (np.ndarray Spatial evaluation grid.)
u0_phase_sym (sp.Expr Phase of u_0: S_0(y).)
u0_amp_sym (sp.Expr Amplitude of u_0: a_0(y).)
mode (str Quantization scheme ('kn' or 'weyl').)

Return type:

np.ndarray of complex

class fio_bridge.PsiOpFIOBridge(op: PseudoDifferentialOperator, lam: float = 50.0, n_guesses: int = 40, xi_range: Tuple[float, float] = (-8.0, 8.0), y_range: Tuple[float, float] = (-6.0, 6.0), tol_grad: float = 1e-06, verbose: bool = False)[source]

Bases: FourierIntegralOperator

Evaluates the action of a PseudoDifferentialOperator on a WKB state via the stationary-phase / saddle-point method.

Performance design

The original bottleneck was that every call to evaluate_at(x_val, ...) triggered a full SymPy rebuild:

build_kernel → sp.subs (new x_val substituted into phase) _make_guesses → sp.diff + sp.solve (analytical ξ_c guess) _build_analyzer → Analyzer.__init__

→ _prepare_derivatives (20+ sp.diff calls) → _create_numerical_functions (8 lambdify calls)

All of that work is identical across x_val except for the numerical value of x_val itself, because the phase is linear in x_val:

φ(y, ξ; x) = (x − y)·ξ + S_u(y) ∂φ/∂y = −ξ + S_u′(y) (independent of x) ∂φ/∂ξ = x − y (linear in x)

The refactored version precomputes everything symbolic once in __init__ using a symbolic placeholder _xp for the observation coordinate. Each call to evaluate_at / evaluate_grid then only:

Injects float(x_val) into pre-built numpy callables.

Runs scipy.minimize on those callables (pure numerics).

Evaluates the asymptotic formula (pure numerics).

No SymPy is touched after the first evaluate_grid call. For a grid of N points the symbolic cost is O(1) instead of O(N).

Measured speedup (N = 200, 1D, Morse symbol):: Before ~45 s (SymPy dominates) After ~2.5 s (scipy.minimize dominates) Factor ~18×

The public API (evaluate_at, evaluate_grid, build_kernel) is unchanged.

param op:
type op:: PseudoDifferentialOperator
param lam:
type lam:: float
param n_guesses:
type n_guesses:: int
param xi_range:
type xi_range:: tuple[float, float]
param y_range:
type y_range:: tuple[float, float]
param tol_grad:
type tol_grad:: float
param verbose:
type verbose:: bool

build_kernel(x_val: Any, u_phase_sym: Expr, u_amp_sym: Expr) → FIOKernel[source]: Retained for API compatibility and base-class tests. Not called by the optimised evaluate_at / evaluate_grid paths.

evaluate_at(x_val: Any, u_phase_sym: Expr, u_amp_sym: Expr) → EvalResult[source]

Evaluate (Pu)(x_val) at a single observation point.

The first call with a given WKB pair triggers _precompute_wkb (O(1) SymPy work). All subsequent calls are pure numerics.

evaluate_grid(x_grid: ndarray, u_phase_sym: Expr, u_amp_sym: Expr, n_workers: int | None = None) → ndarray[source]

Evaluate (Pu)(x) at every point in x_grid.

SymPy work is done once before the loop; the loop body is pure scipy (minimize) + numpy. No SymPy inside the loop.

Each grid point is independent, so the loop is parallelised with concurrent.futures.ThreadPoolExecutor. Threads (not processes) are used because the lambdified SymPy functions stored on self are not picklable, and because scipy.optimize.minimize releases the GIL when calling into its C/Fortran back-ends, so threads do run in parallel for this CPU-bound workload.

Pass n_workers=1 to force sequential execution (useful for debugging or profiling).

Parameters:

x_grid (np.ndarray)
u_phase_sym (sp.Expr)
u_amp_sym (sp.Expr)
n_workers (int | None) – Number of threads. None → os.cpu_count().

Return type:

np.ndarray of complex128, shape (len(x_grid),)

class fio_bridge.SemiclassicalCorrector(op: PseudoDifferentialOperator, splitter: SpectralSplitter, **bridge_kwargs)[source]

Bases: object

Refine a PDESolver solution by replacing its high-frequency component with the asymptotically more accurate WKB estimate from PsiOpFIOBridge.

Workflow

Split the solver solution u_solver into u_low + u_high.
Evaluate u_bridge_high via PsiOpFIOBridge on the high-frequency initial condition.
Return the corrected field u_corrected = u_low + u_bridge_high.

This is meaningful when: - The high-frequency part of u_solver has accumulated spectral error

(e.g., phase drift after many time steps).

The WKB ansatz is valid at the dominant wavenumber (λ large enough).

param op:: The operator defining the dynamics.
type op:: PseudoDifferentialOperator
param splitter:: Pre-configured SpectralSplitter with the desired k_cut.
type splitter:: SpectralSplitter
param bridge_kwargs:: Keyword arguments forwarded to PsiOpFIOBridge (lam, n_guesses, …).
type bridge_kwargs:: dict

correct(u_solver: ndarray, wkb_state: WKBState) → ndarray[source]

Apply one correction step.

Parameters:

u_solver (np.ndarray) – Current solver solution on splitter.x_grid.
wkb_state (WKBState) – WKB ansatz for the high-frequency part. Its amplitude and phase are used as input to PsiOpFIOBridge.evaluate_grid().

Returns:

Corrected solution u_low(solver) + u_high(bridge).

Return type:

np.ndarray (complex128)

correction_magnitude(u_solver: ndarray, wkb_state: WKBState) → float[source]: Return ‖u_corrected − u_solver‖ / ‖u_solver‖ — how much the correction changes the solution. Values > 0.1 indicate that the solver’s high-frequency component has significant WKB error.

class fio_bridge.SpectralSplitter(x_grid: ndarray, k_cut: float | None = None)[source]

Bases: object

Decompose a field into low-frequency and high-frequency parts using a sharp spectral cutoff in Fourier space.

The decomposition is:: u = u_low + u_high
where: u_low = IFFT[ û(k) · 1_{|k| ≤ k_cut} ] u_high = IFFT[ û(k) · 1_{|k| > k_cut} ]

This is the natural split point between the two solvers: - u_low is handled by PDESolver (spectral, global, all-frequency). - u_high is handled by PsiOpFIOBridge (asymptotic, WKB, high-frequency).

Parameters:

x_grid (np.ndarray) – Uniform spatial grid.
k_cut (float) – Cutoff wavenumber (in rad/unit). Modes with |k| > k_cut are assigned to the high-frequency part. If None, defaults to half the Nyquist frequency.

Notes

The split is lossless: merge(split(u)) == u to machine precision.

energy_ratio(u: ndarray) → Tuple[float, float][source]

Return (E_low / E_total, E_high / E_total) — spectral energy fractions.

Useful for choosing k_cut: a good cut puts most energy in u_low (solvable spectrally) while the WKB oscillations live in u_high.

merge(u_low: ndarray, u_high: ndarray) → ndarray[source]

Reconstruct u = u_low + u_high.

This is simply element-wise addition; the method exists as an explicit counterpart to split() to make the API symmetric.

split(u: ndarray) → Tuple[ndarray, ndarray][source]

Decompose u into (u_low, u_high).

Parameters:: u (np.ndarray (real or complex, 1D))
Returns:: u_low, u_high
Return type:: both complex128, same shape as u.

suggest_k_cut(u: ndarray, target_high_fraction: float = 0.9) → float[source]

Suggest a k_cut such that a fraction target_high_fraction of the spectral energy lies in u_high.

Useful when the WKB state carries most of its energy at high k and the split should isolate that regime cleanly.

class fio_bridge.ValidationReport(x_grid: ndarray, u_solver: ndarray, u_bridge: ndarray, abs_error: ndarray, rel_error: ndarray, max_abs_error: float, max_rel_error: float, wkb_valid: bool, wkb_threshold: float, error_spectrum: ndarray, k_grid: ndarray, lam: float)[source]

Bases: object

Result of a CrossValidator run.

x_grid

Spatial evaluation grid.

Type:: np.ndarray

u_solver

Solution produced by PDESolver (complex).

Type:: np.ndarray

u_bridge

Solution produced by PsiOpFIOBridge (complex).

Type:: np.ndarray

abs_error

Point-wise absolute error |u_solver − u_bridge|.

Type:: np.ndarray

rel_error

Point-wise relative error, normalised by max|u_bridge|.

Type:: np.ndarray

max_abs_error

Type:: float

max_rel_error

Type:: float

wkb_valid

True when max_rel_error < wkb_threshold (λ large enough).

Type:: bool

wkb_threshold

Threshold used to declare the WKB regime valid.

Type:: float

error_spectrum

|FFT(u_solver − u_bridge)| — diagnoses which wavenumbers disagree.

Type:: np.ndarray

k_grid

Wavenumber grid corresponding to error_spectrum.

Type:: np.ndarray

lam

Large parameter λ used in the bridge evaluation.

Type:: float

abs_error: ndarray

error_spectrum: ndarray

k_grid: ndarray

lam: float

max_abs_error: float

max_rel_error: float

rel_error: ndarray

u_bridge: ndarray

u_solver: ndarray

wkb_threshold: float

wkb_valid: bool

x_grid: ndarray

class fio_bridge.WKBState(amp_sym: Expr, phase_sym: Expr, var_x: Symbol, lam: float = 50.0)[source]

Bases: object

Carry a WKB ansatz u(x) ~ a(x) · exp(iλ S(x)) and expose it as a numpy array or as a callable compatible with PDESolver’s initial_condition parameter.

Parameters:

amp_sym (sp.Expr) – Amplitude a(x) as a SymPy expression in var_x.
phase_sym (sp.Expr) – Phase S(x) as a SymPy expression in var_x.
var_x (sp.Symbol) – The spatial variable (must match the symbol in amp_sym / phase_sym).
lam (float) – Large parameter λ. The full phase in the exponent is λ·S(x).

Examples

>>> x = sp.Symbol('x', real=True)
>>> state = WKBState(sp.exp(-x**2/2), x, x, lam=40.0)
>>> solver.setup(..., initial_condition=state.as_callable())

as_callable()[source]

Return a callable f(x) -> np.ndarray suitable for PDESolver.setup(initial_condition=f).

The callable ignores any second argument (time), so it works for both stationary and time-dependent solvers.

dominant_wavenumber(x_grid: ndarray) → float[source]: Return the median of |k(x)| over the grid — a single representative wavenumber for use as k_cut in SpectralSplitter.

to_array(x_grid: ndarray) → ndarray[source]

Evaluate u(x) = a(x) · exp(iλ S(x)) on x_grid.

Return type:: np.ndarray of complex128, shape (len(x_grid),)

wkb_phase_gradient(x_grid: ndarray) → ndarray[source]

Return the local wavenumber k(x) = λ · S’(x) on x_grid.

This is the dominant frequency of the WKB state at each point, useful for choosing the SpectralSplitter cut-off.

fio_bridge.fft_reference(op: PseudoDifferentialOperator, u_vals: ndarray, x_grid: ndarray) → ndarray[source]

Exact numerical reference via FFT for a 1D constant-coefficient psiOp.

(Pu)(x) = IFFT[ p(ξ) · FFT[u](ξ) ]

Used only to validate the asymptotic bridge against a spectral solver.

fio_bridge.run_test_suite(verbose: bool = True, plot: bool = True) → Dict[str, Any][source]

Run 5 tests covering the main use cases of the bridge.

WKB reference

For a psiOp with constant symbol p(ξ) applied to: u_0(y) = a_u(y) · exp(iλ k_0 y)
the semi-classical result at leading order is:: (Pu)(x) = p(k_0) · u_0(x)

Test 1 – Elliptic P = ξ² → p(k0) = k0² Test 2 – Transport P = c·ξ → p(k0) = c·k0 Test 3 – Composition P∘Q, P=ξ², Q=ξ → PQ=ξ³, p(k0) = k0³ Test 4 – Right-inverse P=ξ²+1, R=P⁻¹ → (Ru)(x) = u(x)/p(k0) Test 5 – Propagator exp(it·ξ) → phase shift exp(i·t·k0)

Error tolerance: O(1/λ) for simple operators, O(5/λ) for compositions.

returns:: ‘max_rel_error’, ‘values_bridge’, ‘values_ref’.
rtype:: dict – keys ‘test1’…’test5’; each a dict with ‘passed’,