Integration and utilities reference

This section documents the lower-level modules that power numerical simulations: integration of the Bloch Hamiltonian, helper functions for setting up simulation initial conditions, and utilities for precomputing Bragg pulse lookup tables.

Integrate module

The mwave.integrate module contains functions for integrating the Bloch Hamiltonian to simulate Bragg diffraction pulses. Most functions use numba for performance.

class mwave.integrate.PropagateResult(phi_final, kvec, omega, omega_args, phase, phase_args, dt=None, error=None, scipy_sol=None)

Result returned by propagate().

Attributes always available:

phi_final: Final wavefunction, (N,) for single-atom or (natoms, N) for batch mode.
kvec: Momentum-state grid used in the simulation.

Attributes set only for the 'numba' backend:

dt: Step size used by the integrator.
error: Estimated integration error.

Attribute set only for the 'scipy' backend:

scipy_sol: The full scipy.integrate.OdeResult. When dense=True was requested the interpolant is accessible via scipy_sol.sol.

plot()

Three-panel plot of populations, Rabi frequency, and phase vs time.

Only available when the 'scipy' backend was used, because the full time trajectory is required. Raises ValueError otherwise.

Returns:: The matplotlib Figure.

population(k): Return \(|\langle k | \phi \rangle|^2\) at the momentum state nearest to k. For batch results returns an (natoms,) array.

populations(): Return \(|\phi|^2\) for every momentum state.

mwave.integrate.make_kvec(n0, nf, npad=10)

Generates a vector of \(k\)-states. Note that neighboring \(k\)-states are spaced by 2 photon recoils.

Parameters:

n0 – The initial momentum state to include in the space.
nf – The final momentums tate to include in the space.
npad – The padding to include on each side of the initial and final momentum states.

Returns:

A tuple containing a vector of momentum states spaced by \(2\hbar k\), the index of n0 in the vector, and the index of nf in the vector.

Example

>>> from mwave.integrate import make_kvec
>>> n0, nf = 0, 5
>>> make_kvec(n0,nf)
(array([-20., -18., -16., -14., -12., -10.,  -8.,  -6.,  -4.,  -2.,   0.,
         2.,   4.,   6.,   8.,  10.,  12.,  14.,  16.,  18.,  20.,  22.,
        24.,  26.,  28.,  30.]), 10, 15)

mwave.integrate.make_phi(kvec, n0)

Creates a vector the same length as kvec filled with all zeros aside from in the index where kvec==2*n0, which is set to 1.

Parameters:

kvec – The momentum state values at which phi is defined.
n0 – The state to initialize all amplitude in.

mwave.integrate.multi_omega_fnc(t, args)

Function defining a multifrequency Gaussian pulse profile in time, i.e.

\[\Omega(t)=2\Omega\cos(\omega_\text{mod}t)\exp\left(-\frac{(t-t_0)^2}{2\sigma^2}\right)\]

where \(\Omega\), \(\sigma\), \(t_0\), and \(\omega_\text{mod}\) are given by args[0], args[1], args[2], and args[3], respectively.

Parameters:

t – The time at which to evalute the Gaussian.
args – A tuple of four parameters defining \(\Omega\), \(\sigma\), \(t_0\), and omega_text{mod}.

Returns:

The function value at the provided time.

mwave.integrate.omega_fnc_gaussian(t, args)

Function defining a Gaussian pulse profile in time, i.e.

\[\Omega(t)=\Omega\exp\left(-\frac{(t-t_0)^2}{2\sigma^2}\right)\]

where \(\Omega\), \(\sigma\), and \(t_0\) are given by args[0], args[1], and args[2], respectively.

Parameters:

t – The time at which to evalute the Gaussian.
args – A tuple of three parameters defining \(\Omega\), \(\sigma\), and \(t_0\).

Returns:

The function value at the provided time.

mwave.integrate.phase_fnc_constant(t, args)

Function defining a constant phase as a function of time.

Parameters:

t – The time at which to evalute the phase. Since the phase is constant this parameter has no effect.
args – A tuple of one parameters defining the value of the constant phase.

Returns:

The phase value at the provided time.

mwave.integrate.propagate(kvec, phi0, tfinal, delta, omega, omega_args, phase, phase_args, omegas=None, t0=0.0, backend=None, dense=False, method='DOP853', atol=1e-10, rtol=1e-10, max_step=0.1, transformed=False, Gamma_sps=None, tol=1e-10, cache=None)

Evolves the provided wavefunction using the equations of motion described in Integration backends. The user can provide a single atom wavefunction (in which case the scipy backend is used), or a batch of wavefunctions to be integrated in parallel (in which case the numba backend is used). The wavefunction batching provided by the function is significantly more efficient than looping over a single atom wavefunction call multiple times.

Backends

'scipy' — adaptive solve_ivp (default for single-atom). Supports dense output, the transformed frame, and spontaneous emission via Gamma_sps. Unavailable for batch mode.
'numba' — Numba RK45 with prange (default for batch).

Parameters:

kvec – The vector of momentum states to simulate (N,) float64.
phi0 – The initial value of phi. (N,) for single-atom or (natoms, N) for batch mode.
tfinal – The final time to integrate to.
delta – The two-photon detuning. Scalar for single-atom, (natoms,) for batch mode.
omega – Callable omega(t, omega_args) -> float.
omega_args – Extra arguments forwarded to omega.
phase – Callable phase(t, phase_args) -> float.
phase_args – Extra arguments forwarded to phase.
omegas – Per-atom Rabi frequency scale. Scalar or (natoms,). Defaults to 1.0 for single-atom, np.ones(natoms) for batch.
t0 – Integration start time (default 0.0).
backend – 'scipy', 'numba', or None. If None, scipy will be selected for a single wavefunction and numba will be selected for a batch of wavefunctions.
dense – If true dense output is returned (i.e. the integration result can be queried for any intermediate time). Requires backend='scipy'.
method – ODE method for backend='scipy' (default 'DOP853'). Ignored for other backends.
atol – Absolute tolerance for backend='scipy' (default 1e-10). Ignored for other backends.
rtol – Relative tolerance for backend='scipy' (default 1e-10). Ignored for other backends.
max_step – Maximum step size for backend='scipy' (default 0.1). Ignored for other backends.
transformed – Use the transformed frame for backend='scipy' (default False). Ignored for other backends.
Gamma_sps – Single-photon scattering rate for density-matrix evolution for backend='scipy' (default None). Ignored for other backends.
tol – Error tolerance for the numba RK45 backend (default 1e-10).
cache – Optional dict for memoising results from the numba backend. The cache key incorporates every input that affects the output (phi0, delta, omegas, kvec, omega_args, phase_args, t0, tfinal, backend, tol) along with the omega and phase callables themselves. The callables are hashed by object identity, so to get cache hits across calls you must reuse the same function object. Wrapping the same underlying function in a fresh lambda on each call will produce a distinct object and miss the cache. Define the wrapper once and reuse it.

Returns:

A PropagateResult.

mwave.integrate.score_backends(n0=0, nf=5, natoms=10000, tol=1e-10, repeat=3)

Benchmark propagate() in single-atom and batch modes.

Two timed runs are performed on a fixed Gaussian-pulse Bragg scenario:

Single-atom run — both scipy and numba are timed, and the numba result is compared against the scipy reference.
Batch run — only numba is timed; scipy does not support batch input so there is no reference to compare against.

Parameters:

n0 – Lower momentum-state order (default 0).
nf – Upper momentum-state order (default 5).
natoms – Number of atoms in the batch run (default 1000).
tol – Integration tolerance (default 1e-10).
repeat – Number of timed repetitions per backend; the minimum is used (default 3).

Returns:

A dict with timings (and the single-atom error) for each run.

Simulation utils module

The mwave.simulation_utils module provides helpers for initializing atom clouds and computing interferometer observables.

mwave.simulation_utils.cloud_init(natoms, sigma_cloud, sigma_transverse_v, sigma_vertical_v, x_offset=0, y_offset=0, z_offset=0, vx_offset=0, vy_offset=0, vz_offset=0, seed=None)

Basic cloud initialization. For more complicated atom clouds (i.e. clouds where positions/velocities are correlated with each other), the user should write their own initialization function.

Parameters:

natoms – The number of atoms to initialize.
sigma_cloud – The standard deviation of atom position within the cloud. The user can either provide a single number (in which case the standard deviation is the same in all directions), or three numbers (in which case the standard deviation is defined as sigma_cloud=(sigma_x, sigma_y, sigma_z)).
sigma_transverse_v – The standard deviation of atom velocity within the cloud. The user can either provide a single number (in which case the standard deviation is same in both transverse directions), or two numbers (in which case the standard deviation is defined as sigma_transverse_v=(sigma_x, sigma_y)).
sigma_vertical_v – The standard deviation of atom velocity with the vertical direction of the cloud.
x_offset – The offset of the cloud from \(x=0\).
y_offset – The offset of the cloud from \(y=0\).
z_offset – The offset of the cloud from \(z=0\).
vx_offset – The offset of the cloud velocity from \(v_x=0\).
vy_offset – The offset of the cloud velocity from \(v_y=0\).
vz_offset – The offset of the cloud velocity from \(v_z=0\).
seed – The seed to use when drawing the positions and velocities. If this function is called within an optimization loop the seed should remain the same across function calls!

Returns:

A tuple of (x0, y0, z0, vz, vx, vy), where each element of the tuple is a numpy array of length natoms.