Analytical 1D Doppler Solver¶
This notebook demonstrates how to use rydiqule’s prototype implementation of the analytical 1D doppler solver described in Omar Nagib and Thad G. Walker, Exact steady state of perturbed open quantum systems, arXiv 2501.06134 (2025) http://arxiv.org/abs/2501.06134 It currently contains development work, including detailed profiling performance, as the features are being worked on.
We will look at a simple case of Autler-Townes splitting in Rubidium-87 with two optical fields (blue and red). Applying an RF tone at different field strengths causes the central peak to split and shift as we scan the probe detuning. We show excellent agreement between numerically integrating the effect of Doppler broadening and the new exact method.
%load_ext autoreload
%autoreload 2
%load_ext line_profiler
import numpy as np
import matplotlib.pyplot as plt
import rydiqule as rq
# parameters for Cell
atom = 'Rb87'
states = [
rq.ground_state(atom),
rq.D2_excited(atom),
rq.A_QState(41,2,5/2),
rq.A_QState(40,3,7/2)
]
cell = rq.Cell(atom, states)
# laser parameters
detunings = 2*np.pi*np.linspace(-30,30,201)
Omega_r = 2*np.pi*2
Omega_b = 2*np.pi*5
Omega_rf = 2*np.pi*np.array([0,5,40])
kunit1 = np.array([1,0,0])
kunit2 = np.array([-1,0,0])
red = {'states': (states[0],states[1]), 'detuning': detunings, 'rabi_frequency': Omega_r, 'kunit': kunit1}
blue = {'states': (states[1],states[2]), 'detuning': 0, 'rabi_frequency': Omega_b, 'kunit': kunit2}
rf = {'states': (states[2],states[3]), 'detuning': 0, 'rabi_frequency': Omega_rf}
cell.add_couplings(red, blue, rf)
def compare_accuracy(sol1: np.ndarray, sol2: np.ndarray,
rtol: float = 1e-5, atol: float = 1e-7):
"""Helper function for summarizing relative and absolute differences between
density matrix solutions.
Note that sol1 is considered the 'correct' solution in the comparison.
Tolerances are passed to numpy.isclose for defining how close something is."""
assert sol1.shape == sol2.shape, 'solutions must have same shape to be compared'
abs_diff = np.abs(sol2 - sol1)
null_elem = np.isclose(sol1, 0.0)
sol_ref = sol1.copy()
sol_ref[null_elem] = 1.0
rel_diff = abs_diff/np.abs(sol_ref)
rel_diff[null_elem] = 0.0
print(f'Abs(diff) max {abs_diff.max():.3e}, mean {abs_diff.mean():.3g}, std {abs_diff.std():.3g}')
print(f'Rel(diff) max {rel_diff.max():.3e}, mean {rel_diff.mean():.3g}, std {rel_diff.std():.3g}')
close = np.isclose(sol2, sol1, rtol=rtol, atol=atol) # element-wise close
close_sys = close.all(axis=-1) # density-matrix wise close
if not close_sys.all():
print(f'Not close matrix elements {(~close).sum():d} out of {close.size:d} total')
print(f'Not close solutions {(~close_sys).sum():d} out of {close_sys.size:d} total')
not_close_inds = (~close).nonzero()
vals, counts = np.unique(not_close_inds[2], return_counts=True)
for l, v, c in zip(cell.dm_basis(), vals, counts):
print(f'\tdm element {l:s}-[{v:d}] has {c:d} misses: ' +
f'Abs-diff (max, mean, diff) {abs_diff[...,v].max():.3e}, {abs_diff[...,v].mean():.3e}, {abs_diff[...,v].std():.3e}')
return not_close_inds
%%time
sol_riemann = rq.solve_steady_state(cell, doppler=True)
CPU times: total: 1.92 s
Wall time: 1.02 s
dop_mesh_method = {'method': 'split',
'width_doppler': 2.0,
'n_doppler': 201,
'width_coherent': 0.28,
'n_coherent': 1001}
%%time
sol_riemann_finer = rq.solve_steady_state(cell, doppler=True, doppler_mesh_method=dop_mesh_method)
CPU times: total: 3.05 s
Wall time: 2.03 s
%%time
sol_exact = rq.doppler_1d_exact(cell)
CPU times: total: 1 s
Wall time: 266 ms
bad_inds = compare_accuracy(sol_exact.rho, sol_riemann.rho, rtol=1e-5, atol=7e-5)
Abs(diff) max 4.852e-04, mean 3.08e-05, std 6.43e-05
Rel(diff) max 1.961e+03, mean 1.19, std 30.3
Not close matrix elements 1313 out of 9045 total
Not close solutions 507 out of 603 total
dm element 01_real-[1] has 409 misses: Abs-diff (max, mean, diff) 4.852e-04, 1.435e-04, 1.140e-04
dm element 02_real-[2] has 238 misses: Abs-diff (max, mean, diff) 4.385e-04, 7.979e-05, 9.618e-05
dm element 03_real-[9] has 327 misses: Abs-diff (max, mean, diff) 3.313e-04, 8.965e-05, 7.331e-05
dm element 01_imag-[10] has 118 misses: Abs-diff (max, mean, diff) 2.891e-04, 3.663e-05, 4.974e-05
dm element 11_real-[14] has 221 misses: Abs-diff (max, mean, diff) 3.348e-04, 6.056e-05, 7.086e-05
bad_inds_finer = compare_accuracy(sol_exact.rho, sol_riemann_finer.rho, rtol=1e-5, atol=7e-5)
Abs(diff) max 6.090e-05, mean 3.31e-06, std 7.83e-06
Rel(diff) max 1.279e+02, mean 0.14, std 2.1
bad_inds_riemann = compare_accuracy(sol_riemann_finer.rho, sol_riemann.rho, rtol=1e-5, atol=7e-5)
Abs(diff) max 4.882e-04, mean 3.1e-05, std 6.47e-05
Rel(diff) max 2.134e+02, mean 0.53, std 4.1
Not close matrix elements 1345 out of 9045 total
Not close solutions 507 out of 603 total
dm element 01_real-[1] has 407 misses: Abs-diff (max, mean, diff) 4.882e-04, 1.432e-04, 1.138e-04
dm element 02_real-[2] has 256 misses: Abs-diff (max, mean, diff) 4.490e-04, 8.210e-05, 9.933e-05
dm element 03_real-[9] has 325 misses: Abs-diff (max, mean, diff) 3.406e-04, 8.887e-05, 7.339e-05
dm element 01_imag-[10] has 134 misses: Abs-diff (max, mean, diff) 3.044e-04, 3.857e-05, 5.220e-05
dm element 11_real-[14] has 223 misses: Abs-diff (max, mean, diff) 3.407e-04, 5.994e-05, 7.061e-05
%timeit rq.doppler_1d_exact(cell)
262 ms ± 12.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%lprun -f rq.doppler_1d_exact rq.doppler_1d_exact(cell)
Timer unit: 1e-07 s
Total time: 0.229973 s
File: C:\Users\naqsL\src\Rydiqule\src\rydiqule\doppler_exact.py
Function: doppler_1d_exact at line 69
Line # Hits Time Per Hit % Time Line Contents
==============================================================
69 def doppler_1d_exact(sensor: Sensor, rtol: float = 1e-5, atol: float = 1e-9) -> Solution:
70 """
71 Analytically solves a sensor in steady-state in the presence of 1 dimensional
72 Doppler broadening.
73
74 Uses the method outlined in Ref [1]_.
75 In particular, it uses Eq. 12 to analytically evaluate the Doppler average in 1D.
76
77 This solver is considered more accurate than :func:`~.solve_steady_state`
78 since it replaces direct sampling and solving of the velocity classes
79 with a few tensor inversions and calculation of the numerical prefactor.
80 This also leads to faster solves,
81 approximately dictated by the ratio of samples along the doppler axis
82 relative to the other parameter dimensions.
83
84 Parameters
85 ----------
86 sensor : :class:`~.Sensor`
87 The sensor for which the solution will be calculated.
88 It must define 1 and only 1 dimension of doppler shifts
89 (ie one or more couplings with `kvec` with non-zero values on the same dimension).
90 rtol: float, optional
91 Relative tolerance parameter for checking 0-eigenvalues when calculating the doppler prefactor.
92 Passed to :external+numpy:func:`~numpy.isclose`.
93 Defaults to 1e-5.
94 atol: float, optional
95 Absolute tolerance parameter for checking 0-eigenvalues when calculating the doppler prefactor.
96 Passed to :external+numpy:func:`~numpy.isclose`.
97 Defaults to 1e-9.
98
99 Returns
100 -------
101 :class:`~.Solution`
102 An object containing the solution and related information.
103
104 Raises
105 ------
106 RydiquleError
107 If the `sensor` does not have exactly 1 dimension of doppler shifts to average over.
108 AssertionError
109 If the initial rho0 calculation results in an unphysical result.
110
111 Warns
112 -----
113 RydiquleWarning
114 If the averaged result is not real within tolerances.
115 While rydiqule's computational basis is real,
116 the method employed here involves complex number floating point calculations.
117 If all is well, the complex parts should cancel to return a real result,
118 but imprecision in floating point operations can occur.
119 PopulationNotConservedWarning
120 Before removing the ground state in the final solution,
121 population conservation is confirmed.
122 If the resulting density matrices do not preserve trace, this warning is raised
123 indicating an issue in the calculation.
124
125 References
126 ----------
127 .. [1] Omar Nagib and Thad G. Walker,
128 Exact steady state of perturbed open quantum systems,
129 arXiv 2501.06134 (2025)
130 http://arxiv.org/abs/2501.06134
131 """
132
133 1 4933.0 4933.0 0.2 if sensor.spatial_dim() != 1:
134 raise RydiquleError(f'Sensor must have 1 spatial dimension of Doppler shifts, found {sensor.spatial_dim():d}')
135
136 1 4007.0 4007.0 0.2 stack_shape = sensor._stack_shape()
137 1 27.0 27.0 0.0 n = sensor.basis_size
138 # Liouvillian superoperator for the non-doppler-broadened components
139 2 265358.0 132679.0 11.5 L0, dummy_const = generate_eom(sensor.get_hamiltonian(), sensor.decoherence_matrix(),
140 1 3.0 3.0 0.0 remove_ground_state=False,
141 1 2.0 2.0 0.0 real_eom=True)
142
143 ### construct L0^minus
144 1 192831.0 192831.0 8.4 el0, ev0 = np.linalg.eig(L0)
145 #pseudo invert the eigenvalues
146 1 2993.0 2993.0 0.1 l_mask = np.isclose(el0, 0.0)
147 1 170.0 170.0 0.0 el0[l_mask] = 1 # prevent divide by zero
148 1 532.0 532.0 0.0 li = 1/el0
149 1 88.0 88.0 0.0 li[l_mask] = 0 # set masked elements to zero
150 # calculate Eq B2
151 1 432368.0 432368.0 18.8 L0m = (ev0*li[...,None, :])@np.linalg.pinv(ev0)
152
153 ### calculate rho0 (doppler-free solution) from nullspace eigenvectors in L0m calculation
154 1 1550.0 1550.0 0.1 assert np.all(np.count_nonzero(l_mask, axis=-1) == 1), 'rho0 solution not unique'
155 # select right eigenvector corresponding to 0-eigenvalue (marked by l_mask) for each equation set
156 # using reshape restores stack axes after binary array indexing flatten
157 1 3047.0 3047.0 0.1 rho0 = np.real_if_close(np.swapaxes(ev0, axis1=-1, axis2=-2)[l_mask].reshape(*stack_shape, -1))
158 1 81.0 81.0 0.0 assert not np.iscomplexobj(rho0), 'rho0 solution is not real; it is unphysical'
159 1 334.0 334.0 0.0 rho0 *= np.sign(rho0[...,0])[...,None] # remove arbitrary sign from null-vector so all pops are positive
160 1 332.0 332.0 0.0 pops = np.sum(rho0[...,::n+1], axis=-1) # calculate trace of each vector
161 1 252.0 252.0 0.0 rho0 /= pops[..., None] # normalize vectors by trace
162
163 ### Liouvillian superoperator for doppler only
164 # these are already multiplied by sqrt(2)*sigma_v by rydiqule
165 # as such, lambdas are redefined as sqrt(2)*sigma_v*lambdas of Eq12 in the paper
166 1 22532.0 22532.0 1.0 dopp = sensor.get_doppler_shifts().squeeze()
167 1 614.0 614.0 0.0 Lv_complex = _hamiltonian_term(dopp)
168 1 425.0 425.0 0.0 Lv, _ = make_real(Lv_complex, dummy_const, ground_removed=False)
169
170 ### Calculate doppler averaged steady-state density matrix from propagator
171 1 423962.0 423962.0 18.4 el, R = np.linalg.eig(L0m@Lv)
172 1 468139.0 468139.0 20.4 L = np.linalg.pinv(np.swapaxes(R,-1,-2))
173
174 # calculate Eq 12
175 1 38861.0 38861.0 1.7 prefix = _doppler_eigvec_array(el, rtol, atol)
176 1 288326.0 288326.0 12.5 rho_dopp_complex = np.einsum('...j,...ij,...kj,...k->...i', prefix, R, L, rho0)
177 # confirm that result is approximately real
178 1 18.0 18.0 0.0 imag_tol = 10000
179 1 1854.0 1854.0 0.1 rho_dopp = np.real_if_close(rho_dopp_complex, tol=imag_tol) # chop complex parts if all are smaller than 10000*f64_eps
180 1 105.0 105.0 0.0 if np.iscomplexobj(rho_dopp):
181 rho_dopp_imag = np.abs(rho_dopp_complex.imag)
182 count = np.count_nonzero(rho_dopp_imag > np.finfo(float).eps*imag_tol)
183 warnings.warn('Doppler-averaged solution has complex parts outside of tolerance, solution is suspect. ' +
184 f'{count:d} of {rho_dopp.size:d} elments larger than cutoff {np.finfo(float).eps*imag_tol:.3g}. ' +
185 f'Max Abs(Imag): {rho_dopp_imag.max():.3g}, Std Abs(Imag): {np.std(rho_dopp_imag):.3g}',
186 RydiquleWarning)
187 # ensure trace is preserved before dropping ground state, which will implicitely enforce it
188 1 675.0 675.0 0.0 pops_dopp = np.sum(rho_dopp[...,::n+1], axis=-1)
189 1 2635.0 2635.0 0.1 if not np.isclose(pops_dopp, 1.0).all():
190 warnings.warn('Doppler-averaged solution has populations not conserved, solution is suspect. ' +
191 f'{np.count_nonzero(~np.isclose(pops_dopp, 1.0)):d} of {pops_dopp.size:d} have non-unit trace. ' +
192 f'Min trace is {pops_dopp.min():f}',
193 PopulationNotConservedWarning)
194
195 ### make into a Solution object
196 1 23.0 23.0 0.0 solution = Solution()
197 # match rydiqule convention for return (ie no ground population)
198 1 26.0 26.0 0.0 solution.rho = rho_dopp[...,1:]
199
200 # specific to observable calculations
201 1 30604.0 30604.0 1.3 solution._eta = sensor.eta
202 1 24795.0 24795.0 1.1 solution._kappa = sensor.kappa
203 1 31.0 31.0 0.0 solution._cell_length = sensor.cell_length
204 1 16.0 16.0 0.0 solution._beam_area = sensor.beam_area
205 1 1372.0 1372.0 0.1 solution._probe_freq = sensor.probe_freq
206 1 23.0 23.0 0.0 solution._probe_tuple = sensor.probe_tuple
207
208 # store the graph with fully_expanded dimensionality
209 1 4417.0 4417.0 0.2 sensor._expand_dims()
210 1 28246.0 28246.0 1.2 solution.couplings = deepcopy(sensor.couplings)
211 1 2831.0 2831.0 0.1 _squeeze_dims(sensor.couplings)
212
213 1 3800.0 3800.0 0.2 solution.axis_labels = (sensor.axis_labels() + ["density_matrix"])
214 2 3026.0 1513.0 0.1 solution.axis_values = ([val for _,_,val,_ in sensor.variable_parameters()]
215 1 980.0 980.0 0.0 + [sensor.dm_basis()])
216 1 776.0 776.0 0.0 solution.dm_basis = sensor.dm_basis()
217 1 41708.0 41708.0 1.8 solution.rq_version = version("rydiqule")
218
219 1 5.0 5.0 0.0 return solution
colors=['slateblue', 'lime', 'coral']
plt.figure(figsize=(8,6))
for i in range(len(Omega_rf)):
plt.plot(detunings/(2*np.pi), rq.get_rho_ij(sol_riemann,1,0)[:,i].imag,
label=f'RF={Omega_rf[i]/(2*np.pi):.1f} MHz (Riemann)', c=colors[i], linestyle='dashed')
plt.plot(detunings/(2*np.pi), rq.get_rho_ij(sol_exact,1,0)[:,i].imag,
label=f'RF={Omega_rf[i]/(2*np.pi):.1f} MHz (Exact)', c=colors[i], alpha=0.5)
plt.xlabel('Detuning [MHz]')
plt.ylabel('Absorption [a.u.]')
plt.legend()
plt.title('Autler-Townes Splitting with Doppler Broadening')
plt.show()
rq.about()
Rydiqule
================
Rydiqule Version: 2.1.0.dev23+gc957192.d20250226
Installation Path: ~\src\Rydiqule\src\rydiqule
Dependencies
================
NumPy Version: 1.26.4
SciPy Version: 1.12.0
Matplotlib Version: 3.8.0
ARC Version: 3.6.0
Python Version: 3.11.8
Python Install Path: ~\Miniconda3\envs\rq
Platform Info: Windows (AMD64)
CPU Count and Freq: 12 @ 3.60 GHz
Total System Memory: 128 GB