def test_to_from_json(self):
import os
filepath = "test_ob_solve_02.json"
ob_solve_02 = ob_solve.OBSolve().from_json_str(JSON_STR_02)
ob_solve_02.to_json(filepath)
ob_solve_03 = ob_solve.OBSolve().from_json(filepath)
os.remove(filepath)
self.assertEqual(ob_solve_02.to_json_str(), ob_solve_03.to_json_str())
def test_to_from_json_str_00(self):
ob_solve_00 = ob_solve.OBSolve()
ob_solve_01 = ob_solve.OBSolve.from_json_str(ob_solve_00.to_json_str())
self.assertEqual(ob_solve_00.to_json_str.__repr__(),
ob_solve_01.to_json_str.__repr__())
def test_set_field_rabi_t_func_1(self):
""" Test that a custom double pulse Rabi freq time functions can be
set.
"""
ob_solve_02 = ob_solve.OBSolve().from_json_str(JSON_STR_02)
two_pulse_t_func = lambda t, args: (t_funcs.gaussian(0)
(t, args) + t_funcs.gaussian(1)
(t, args))
two_pulse_t_args = {
"ampl_0": 1.0,
"centre_0": 0.0,
"fwhm_0": 0.1,
"ampl_1": 2.0,
"centre_1": 0.5,
"fwhm_1": 0.1,
}
ob_solve_02.set_field_rabi_freq_t_func(0, two_pulse_t_func)
ob_solve_02.set_field_rabi_freq_t_args(0, two_pulse_t_args)
field_0 = ob_solve_02.atom.fields[0]
self.assertAlmostEqual(
field_0.rabi_freq_t_func(0.0, field_0.rabi_freq_t_args), 1.0)
self.assertAlmostEqual(
field_0.rabi_freq_t_func(0.5, field_0.rabi_freq_t_args), 2.0)
self.assertAlmostEqual(
field_0.rabi_freq_t_func(1.0, field_0.rabi_freq_t_args), 0.0)
def test_two_level_rabi_oscillations(self):
""" Solve the optical Bloch equations for the two-level atom.
Notes:
See https://en.wikipedia.org/wiki/Rabi_cycle
"""
RABI_FREQ = 5.0
atom_dict = {
"fields": [{
"coupled_levels": [[0, 1]],
"rabi_freq": RABI_FREQ
}],
"num_states": 2
}
obs = ob_solve.OBSolve(atom=atom_dict,
t_min=0.0,
t_max=1.0,
t_steps=100)
obs.solve()
# Get the populations
pop_0 = np.absolute(obs.states_t()[:, 0, 0])
pop_1 = np.absolute(obs.states_t()[:, 1, 1])
# The solution is known, we should have Rabi cycling at the frequency.
known_0 = np.cos(2.0 * np.pi * RABI_FREQ * obs.tlist / 2.0)**2
known_1 = np.sin(2.0 * np.pi * RABI_FREQ * obs.tlist / 2.0)**2
self.assertTrue(np.allclose(pop_0, known_0, rtol=1.e-5, atol=1.e-5))
self.assertTrue(np.allclose(pop_1, known_1, rtol=1.e-5, atol=1.e-5))
def test_two_level_with_inital_state(self):
""" Same as test_two_level_rabi_oscillations() but with the initial
state set so that the population starts in the upper level.
"""
RABI_FREQ = 5.0
atom_dict = {
"fields": [{
"coupled_levels": [[0, 1]],
"rabi_freq": RABI_FREQ
}],
"num_states": 2,
"initial_state": [0., 1.]
}
obs = ob_solve.OBSolve(atom=atom_dict,
t_min=0.0,
t_max=1.0,
t_steps=100)
obs.solve()
# Get the populations
pop_0 = np.absolute(obs.states_t()[:, 0, 0])
pop_1 = np.absolute(obs.states_t()[:, 1, 1])
# The solution is as test_two_level_rabi_oscillations() but swapped
known_0 = np.sin(2.0 * np.pi * RABI_FREQ * obs.tlist / 2.0)**2
known_1 = np.cos(2.0 * np.pi * RABI_FREQ * obs.tlist / 2.0)**2
self.assertTrue(np.allclose(pop_0, known_0, rtol=1.e-5, atol=1.e-5))
self.assertTrue(np.allclose(pop_1, known_1, rtol=1.e-5, atol=1.e-5))
def test_two_two_fields(self):
""" Test a two-level atom addressed by multiple fields.
Notes:
- Test for bug in #159, where multiple fields coupling the same
levels isn't working.
- The first square pi-pulse drives all the population to the excited
state, the second square pi-pulse drives all pop back to ground.
- Before fix, only the second field is driving the atoms.
"""
json_path = os.path.join(JSON_DIR, "obs-two-two-fields.json")
obs = ob_solve.OBSolve().from_json(json_path)
obs.solve()
# Get the populations
pop_0 = np.absolute(obs.states_t()[:, 0, 0])
pop_1 = np.absolute(obs.states_t()[:, 1, 1])
# All population should start in the ground state
self.assertAlmostEqual(pop_0[0], 1.0, places=5)
self.assertAlmostEqual(pop_1[0], 0.0, places=5)
# The first pi-pulse between t = 0.2 and t = 0.3 should drive all the
# population to the exited state
self.assertAlmostEqual(pop_0[len(pop_0) // 2], 0.0, places=5)
self.assertAlmostEqual(pop_1[len(pop_0) // 2], 1.0, places=5)
# The second pi-pulse between t = 0.6 and t = 0.7 should drive all the
# population back to the ground state
self.assertAlmostEqual(pop_0[-1], 1.0, places=5)
self.assertAlmostEqual(pop_1[-1], 0.0, places=5)
def test_from_json_str(self):
ob_solve_02 = ob_solve.OBSolve().from_json_str(JSON_STR_02)
self.assertEqual(ob_solve_02.t_min, 0.0)
self.assertEqual(ob_solve_02.t_max, 1.0)
self.assertEqual(ob_solve_02.t_steps, 100)
self.assertEqual(ob_solve_02.method, "mesolve")
def test_to_from_json_str_03(self):
json_path = os.path.join(JSON_DIR, "ob_solve_03.json")
obs = ob_solve.OBSolve().from_json(json_path)
obs_test = ob_solve.OBSolve.from_json_str(obs.to_json_str())
self.assertEqual(obs.to_json_str(), obs_test.to_json_str())
def test_save_load_01(self):
""" Solve a basic OBSolve problem. Save the results to file. Set the
results in the OBSolve object to null. Load the results from file
and check that they match the original values.
"""
json_path = os.path.join(JSON_DIR, "ob_solve_02.json")
ob_solve_02 = ob_solve.OBSolve().from_json(json_path)
states_t = ob_solve_02.solve()
states_t_loaded = ob_solve_02.solve(recalc=False)
self.assertTrue((states_t == states_t_loaded).all())
def test_vee_cw_weak_sech_2pi(self):
""" Test a three-level vee config atom where the probe transition is
addressed by a weak cw and the drive is a sech pulse.
Notes:
- Test for bug in #222, where the `t_args can contain ampl or n_pi,
not both` exception is raised even though those args are in
different fields.
- Bug is due to field_idxs not being set correctly, see also #159.
"""
json_path = os.path.join(JSON_DIR, "obs-vee-cw-weak-sech-2pi.json")
obs = ob_solve.OBSolve().from_json(json_path)
# Test that solve does not throw any exceptions.
obs.solve()
def test_two_level_with_opts(self):
""" Same as test_two_level_rabi_oscillations() but with opts set such
that the tolerances are lower. The results will be less
accurate.
"""
RABI_FREQ = 5.0
atom_dict = {
"fields": [{
"coupled_levels": [[0, 1]],
"rabi_freq": RABI_FREQ
}],
"num_states": 2,
"initial_state": [1., 0.]
}
obs = ob_solve.OBSolve(atom=atom_dict,
t_min=0.0,
t_max=1.0,
t_steps=100,
opts={
'atol': 1e-6,
'rtol': 1e-4
})
obs.solve()
# Get the populations
pop_0 = np.absolute(obs.states_t()[:, 0, 0])
pop_1 = np.absolute(obs.states_t()[:, 1, 1])
# The solution is known, we should have Rabi cycling at the frequency.
known_0 = np.cos(2.0 * np.pi * RABI_FREQ * obs.tlist / 2.0)**2
known_1 = np.sin(2.0 * np.pi * RABI_FREQ * obs.tlist / 2.0)**2
# Compared with test_two_level_rabi_oscillations() we can only assert
# a lower tolerance to the known solution.
self.assertTrue(np.allclose(pop_0, known_0, rtol=1.e-3, atol=1.e-3))
self.assertTrue(np.allclose(pop_1, known_1, rtol=1.e-3, atol=1.e-3))