def test_mb_solve(self):
""" Basic test of mb_solve method. """
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_00 = mb_solve.MBSolve().from_json(json_path)
mb_solve_00.mbsolve()
def test_00(self):
mb_solve_00 = mb_solve.MBSolve()
zlist = np.array([0., .1, .2, .3, .4, .5, .6, .7, .8, .9, 1.])
self.assertTrue(np.allclose(mb_solve_00.zlist, zlist, rtol=1.0e-6))
def test_one_field(self):
""" For the case of a single field """
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_00 = mb_solve.MBSolve().from_json(json_path)
self.assertEqual(mb_solve_00.get_Omegas_intp_t_funcs(), ['intp'])
def test_two_spectral(self):
json_path = os.path.join(JSON_DIR, "mbs-two-spectral.json")
mbs = mb_solve.MBSolve().from_json(json_path)
mbs.mbsolve()
# We're not asserting anything about these calls, apart from that they
# run without exception.
freq_list = spectral.freq_list(mbs)
rabi_freq_fft = spectral.rabi_freq(mbs, 0)
dis = spectral.dispersion(mbs, 0, -1)
abs = spectral.absorption(mbs, 0, -1)
hm, r1, r2 = utility.half_max_roots(freq_list, abs)
# The absorption profile should have a peak at 1.0 and have a FWHM of
# 1.0 centred at 0.0.
self.assertAlmostEqual(hm, 0.5, places=1)
self.assertAlmostEqual(r1, -0.5, places=3)
self.assertAlmostEqual(r2, 0.5, places=3)
abs_linear_known = spectral.absorption_two_linear_known(
freq_list, mbs.interaction_strengths[0],
mbs.atom.decays[0]['rate'])
# Assert that the max of the abs residuals between the absorption
# profile and the known absorption profile for linear two-level systems
# is below a tolerance
self.assertTrue(np.max(np.abs(abs - abs_linear_known)) < 0.05)
def test_twolevel_shape(self):
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mbs = mb_solve.MBSolve().from_json(json_path)
coh = mbs.coherences_field(field_idx=0)
np.testing.assert_allclose(
coh, np.zeros((mbs.z_steps + 1, mbs.t_steps + 1)))
def test_two_fields(self):
""" For the case of two fields """
json_path = os.path.join(JSON_DIR, "mb_solve_lamda.json")
mb_solve_lamda = mb_solve.MBSolve().from_json(json_path)
self.assertEqual(mb_solve_lamda.get_Omegas_intp_t_funcs(),
['intp', 'intp'])
def test_no_decays(self):
""" Empty decay list. """
json_path = os.path.join(JSON_DIR, "mb_solve_no_decays.json")
mb_solve_nd = mb_solve.MBSolve().from_json(json_path)
mb_solve_nd.mbsolve()
def test_00(self):
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_00 = mb_solve.MBSolve().from_json(json_path)
result_Delta = mb_solve_00.solve_over_thermal_detunings()
self.assertEqual(len(result_Delta),
len(mb_solve_00.thermal_delta_list))
def test_no_vel_classes(self):
""" Empty velocity class dict. """
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mbs = mb_solve.MBSolve().from_json(json_path)
vc = {}
mbs.build_velocity_classes(vc)
mbs.mbsolve()
def test_two_sech_2pi(self):
""" Test of a 2pi soliton propagating through a two-level system.
"""
json_path = os.path.join(JSON_DIR, "mbs_two_sech_2pi.json")
mbs = mb_solve.MBSolve().from_json(json_path)
mbs.mbsolve()
# Input pulse is 2pi
self.assertAlmostEqual(mbs.fields_area()[0][0]/(np.pi), 2.0, places=1)
# Output pulse is 2pi
self.assertAlmostEqual(mbs.fields_area()[0][-1]/(np.pi), 2.0, places=1)
def test_no_rabi_freq_t_func(self):
""" Empty decay list. TODO: No mbsolve, should be in init"""
json_path = os.path.join(JSON_DIR, "mb_solve_no_rabi_freq_t_func.json")
mbs = mb_solve.MBSolve().from_json(json_path)
# self.assertEqual(mbs.ob_atom.fields[0].rabi_freq_t_func,
# t_funcs.square_1)
self.assertDictEqual(mbs.atom.fields[0].rabi_freq_t_args,
{"ampl_0": 1.0, "on_0": 0.0, "off_0": 1.0})
def test_twolevel_shape(self):
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mbs = mb_solve.MBSolve().from_json(json_path)
pop_upper = mbs.populations_field(field_idx=0, upper=True)
pop_lower = mbs.populations_field(field_idx=0, upper=False)
np.testing.assert_allclose(
pop_lower, np.zeros((mbs.z_steps + 1, mbs.t_steps + 1)))
np.testing.assert_allclose(
pop_upper, np.zeros((mbs.z_steps + 1, mbs.t_steps + 1)))
def test_tlist_max(self):
""" When the speed of light is 1.5, for t_max = 1.0 and z_max = 1.0,
the t_max_fixed should be 1 + 1/1.5 = 5/3. """
mb_solve_00 = mb_solve.MBSolve()
speed_of_light = 1.5 # speed of light
t_max_fixed_expected = mb_solve_00.t_max + \
1 / (speed_of_light * mb_solve_00.z_max)
t_max_fixed = fixed.t_list(mb_solve_00, speed_of_light)[-1]
self.assertAlmostEqual(t_max_fixed, t_max_fixed_expected)
def test_no_atoms_ab(self):
""" Setting the number density to 0.0, i.e. no atoms, with AB step. """
json_path = os.path.join(JSON_DIR, "mb_solve_no_atoms.json")
mbs = mb_solve.MBSolve().from_json(json_path)
mbs.mbsolve(step='ab')
# Check that the field at the end of the medium matches the field
# at the start of the medium.
self.assertTrue(np.allclose(mbs.Omegas_zt[:, 0, :],
mbs.Omegas_zt[:, -1, :], rtol=1.0e-6))
def test_one_field(self):
""" For the case of a single field """
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_00 = mb_solve.MBSolve().from_json(json_path)
Omegas_z = mb_solve_00.Omegas_zt[:, 0, :]
t_args = mb_solve_00.get_Omegas_intp_t_args(Omegas_z)
self.assertEqual(len(t_args), 1)
self.assertTrue(np.all(t_args[0]['tlist'] == mb_solve_00.tlist))
self.assertTrue(np.all(t_args[0]['ylist'] == Omegas_z / (2.0 * np.pi)))
def test_spectral_twolevel(self):
""" Check the spectral methods for exceptions. """
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_00 = mb_solve.MBSolve().from_json(json_path)
mb_solve_00.mbsolve()
freq_list = spectral.freq_list(mb_solve_00)
rabi_freq_fft = spectral.rabi_freq(mb_solve_00, 0)
abs = spectral.absorption(mb_solve_00, 0, -1)
dis = spectral.dispersion(mb_solve_00, 0, -1)
def test_vel_classes(self):
"""Tests that for a linear two-level system with velocity classes, the
absorption matches the known Voigt profile.
"""
json_path = os.path.join(JSON_DIR, "velocity-classes.json")
mbs = mb_solve.MBSolve().from_json(json_path)
mbs.mbsolve()
freq_list = spectral.freq_list(mbs)
abs = spectral.absorption(mbs, 0, -1)
voigt = spectral.voigt_two_linear_known(freq_list, 1.0, 0.05).imag
# Assert that the max of the abs residuals between the absorption
# profile and the known broadened Voigt absorption profile for linear
# two-level systems is below a tolerance
self.assertTrue(np.max(np.abs(abs - voigt)) < 0.05)
def test_no_vel_classes_inner(self):
""" No inner delta values in dict. TODO: No mbsolve, should be init"""
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mbs = mb_solve.MBSolve().from_json(json_path)
vc = {
"thermal_delta_min": -1.0,
"thermal_delta_max": 1.0,
"thermal_delta_steps": 2,
"thermal_width": 1.0
}
mbs.build_velocity_classes(vc)
mbs.mbsolve()
def test_zero_thermal_width(self):
"""TODO: No mbsolve, should be in init"""
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mbs = mb_solve.MBSolve().from_json(json_path)
vc = {
"thermal_delta_min": -1.0,
"thermal_delta_max": 1.0,
"thermal_delta_steps": 2,
"thermal_delta_inner_min": 0.0,
"thermal_delta_inner_max": 0.0,
"thermal_delta_inner_steps": 0,
"thermal_width": 0.0
}
self.assertRaises(ValueError, mbs.build_velocity_classes, vc)
def test_no_atoms(self):
""" Setting the number density ampl to 0.0, i.e. no atoms. The end
pulse should be the same as the start. """
json_path = os.path.join(JSON_DIR, "mb_solve_no_atoms.json")
mbs = mb_solve.MBSolve().from_json(json_path)
mbs.mbsolve(step='euler')
self.assertEqual(mbs.Omegas_zt.shape, (1, 5, 101))
# Check that the field at the end of the medium matches the field
# at the start of the medium.
self.assertTrue(np.allclose(mbs.Omegas_zt[:, 0, :],
mbs.Omegas_zt[:, -1, :], rtol=1.0e-6))
def test_rabi_freq_abs(self):
json_path = os.path.join(JSON_DIR, "mb_solve_no_atoms.json")
mb_solve_01 = mb_solve.MBSolve().from_json(json_path)
mb_solve_01.mbsolve(recalc=False)
speed_of_light = 0.5
t_list_fixed = fixed.t_list(mb_solve_01, speed_of_light)
rabi_freq_abs_fixed = fixed.rabi_freq_abs(mb_solve_01,
0,
speed_of_light,
interp_kind='linear')
# The max of the the field in both frames should be the same
self.assertAlmostEqual(np.max(rabi_freq_abs_fixed),
np.max(mb_solve_01.Omegas_zt),
places=0)
def test_save_load_no_recalc(self):
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_01 = mb_solve.MBSolve().from_json(json_path)
Omegas_zt, states_zt = mb_solve_01.mbsolve()
mb_solve_01.save_results()
mb_solve_01.Omegas_zt = None
mb_solve_01.states_zt = None
Omegas_zt, states_zt = mb_solve_01.mbsolve(recalc=False)
Omegas_zt_loaded = mb_solve_01.Omegas_zt
states_zt_loaded = mb_solve_01.states_zt
self.assertTrue((Omegas_zt == Omegas_zt_loaded).all())
self.assertTrue((states_zt == states_zt_loaded).all())
def test_save_load(self):
json_path = os.path.join(JSON_DIR, 'mb_solve_01.json')
mb_solve_00 = mb_solve.MBSolve().from_json(json_path)
# Don't solve, just check initial state
fileio.save_csv_rabi_freq(mb_solve_00,
field_idx=0,
filename='rabi_freq.csv')
fileio.save_csv_rabi_freq_abs(mb_solve_00,
field_idx=0,
filename='rabi_freq_abs.csv')
test_rf = qufileio.file_data_read(filename='rabi_freq.csv')
test_rf_abs = qufileio.file_data_read(filename='rabi_freq_abs.csv')
self.assertTrue(np.allclose(test_rf, mb_solve_00.Omegas_zt[0]))
self.assertTrue(np.allclose(test_rf_abs, mb_solve_00.Omegas_zt[0]))
def test_save_load_01(self):
""" Solve a basic MBSolve problem. Save the results to file. Set the
results in the MBSolve object to null. Load the results from
file and check that they equal the original values.
"""
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_01 = mb_solve.MBSolve().from_json(json_path)
Omegas_zt, states_zt = mb_solve_01.mbsolve()
mb_solve_01.save_results()
mb_solve_01.Omegas_zt = None
mb_solve_01.states_zt = None
mb_solve_01.load_results()
Omegas_zt_loaded = mb_solve_01.Omegas_zt
states_zt_loaded = mb_solve_01.states_zt
self.assertTrue((Omegas_zt == Omegas_zt_loaded).all())
self.assertTrue((states_zt == states_zt_loaded).all())
def test_init_default(self):
""" Test Default Initialise """
mb_solve_00 = mb_solve.MBSolve()
self.assertEqual(mb_solve_00.atom.num_states, 1)
# TODO: CHECK FILE EXISTS
# TODO: CHECK SAVE PATH EXISTS
speed_of_light = float(opts.speed_of_light)
y_min = float(opts.y_min)
y_max = float(opts.y_max)
y_max_1 = float(opts.y_max_1)
z = int(opts.zoom)
fps = float(opts.fps)
atoms_alpha = float(opts.atoms_alpha)
show_c_line = opts.c_line
show_peak_line = opts.peak_line
mb_solve_00 = mb_solve.MBSolve().\
from_json(json_file)
mb_solve_00.mbsolve(recalc=False, pbar_chunk_size=2)
### PLOT FIRST
fig = plt.figure(1) #, figsize=(16, 12))
# Probe
ax = fig.add_subplot(211)
cmap_range = np.linspace(0.0, 1.0e-3, 11)
cf = ax.contourf(mb_solve_00.tlist,
mb_solve_00.zlist,
np.abs(mb_solve_00.Omegas_zt[0] / (2 * np.pi)),
cmap_range,
cmap=plt.cm.Blues)
def test_init_00(self):
json_path = os.path.join(JSON_DIR, "mb_solve_01.json")
mb_solve_01 = mb_solve.MBSolve().from_json(json_path)
