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))