def test_psi_traj_setter_if():
"""
test to see if the current psi_traj is properly being stored when the input is of
type list
"""
TS = TrajStorage()
TS.psi_traj = [[0, 0], [1, 2]]
TS.psi_traj = [[1, 0], [1, 2]]
traj = TS.psi_traj
known_traj = [[[0, 0], [1, 2]], [[1, 0], [1, 2]]]
assert np.array_equal(traj, known_traj)
def test_psi_traj_setter_else():
"""
test to see if the current psi_traj is properly being stored when the input is not
a list
"""
TS = TrajStorage()
TS.psi_traj = [1, 2, 3]
TS.psi_traj = [4, 5, 6]
traj = TS.psi_traj
known = [[1, 2, 3], [4, 5, 6]]
assert np.array_equal(traj, known)
def test_t_axis_setter_else():
"""
test to see if the current t_axis is properly being stored when the input is not a
list
"""
TS = TrajStorage()
TS.t_axis = 2
TS.t_axis = 4
t_axis = TS.t_axis
known_t_axis = [2, 4]
assert np.array_equal(t_axis, known_t_axis)
def test_t_setter():
"""
test to see if the current time is properly being stored
"""
TS = TrajStorage()
TS.t = 1
t = TS.t
known_t = 1
assert t == known_t
TS.t = 2
t2 = TS.t
known_t2 = 2
assert t2 == known_t2
def test_phi_setter():
"""
test to see if the current phi is properly being stored
"""
TS = TrajStorage()
TS.phi = np.array([0, 0, 1, 0])
phi = TS.phi
known_phi = np.array([0, 0, 1, 0])
assert np.array_equal(phi, known_phi)
TS.phi = np.array([0, 1, 0, 0])
phi = TS.phi
known_phi = np.array([0, 1, 0, 0])
assert np.array_equal(phi, known_phi)
def test_z_mem_setter():
"""
test to see if the current z_mem is properly being stored
"""
TS = TrajStorage()
TS.z_mem = np.array([1])
z_mem = TS.z_mem
known_z_mem = np.array([1])
assert z_mem == known_z_mem
TS.z_mem = np.array([2])
z_mem = TS.z_mem
known_z_mem = np.array([2])
assert z_mem == known_z_mem
def test_n_dim_setter():
"""
test to see if the current N_dim is properly being stored
"""
TS = TrajStorage()
TS.n_dim = 1
N_dim = TS.n_dim
known_n_dim = 1
assert N_dim == known_n_dim
TS.n_dim = 2
n_dim = TS.n_dim
known_n_dim = 2
assert n_dim == known_n_dim
def test_aux_list_setter():
"""
Test to see if aux list if properly being stored. If store_aux = False nothing
gets stored. If store_aux = True list_aux is stored
"""
TS = TrajStorage()
TS.store_aux = False
TS.aux_list = [1]
aux_list = TS.aux_list
known_aux_list = []
assert np.array_equal(aux_list, known_aux_list)
TS.store_aux = True
TS.aux_list = [2]
aux_list = TS.aux_list
known_aux_list = [[2]]
assert np.array_equal(aux_list, known_aux_list)
def test_phi_traj_setter_if():
"""
test to see if the current phi_traj is properly being stored when the input is a
list. If store_aux is false nothing should be getting stored but if store_aux is
True we should store phi_traj
"""
TS = TrajStorage()
TS.store_aux = False
TS.phi_traj = [[0, 0], [1, 2]]
traj = TS.phi_traj
known = []
assert np.array_equal(traj, known)
TS.store_aux = True
TS.phi_traj = [[1, 0], [1, 2]]
traj = TS.phi_traj
known = [[[1, 0], [1, 2]]]
assert np.array_equal(traj, known)
def initialize(self, psi_0, store_aux=False):
"""
This function initializes the trajectory module by ensuring that
each sub-component is prepared to begin propagating a trajectory.
PARAMETERS
----------
1. psi_0 : np.array
Wave function at initial time
RETURNS
-------
None
"""
psi_0 = np.array(psi_0, dtype=np.complex128)
if not self.__initialized__:
# Initialize Noise
# ----------------
if not self.noise1.__locked__:
self.noise1.prepare_noise()
if not self.noise2.__locked__:
self.noise2.prepare_noise()
# Prepare the derivative
# ----------------------
self.dsystem_dt = self.basis.initialize(psi_0)
# Prepare Storage
# ---------------
if self.basis.adaptive:
self.storage = AdaptiveTrajectoryStorage()
else:
self.storage = TrajectoryStorage()
self.storage.store_aux = store_aux
# Initialize System State
# -----------------------
self.storage.n_dim = self.basis.system.param["NSTATES"]
phi_tmp = np.zeros(self.n_hier * self.n_state, dtype=np.complex128)
phi_tmp[:self.n_state] = np.array(psi_0)[self.state_list]
z_mem = np.zeros(len(self.basis.system.param["L_NOISE1"]))
if self.basis.adaptive:
# Update Basis
z_step = self.noise1.get_noise([0])[:, 0]
(state_update,
aux_update) = self.basis.define_basis(phi_tmp, 1, z_step)
(phi_tmp,
dsystem_dt) = self.basis.update_basis(phi_tmp, state_update,
aux_update)
self.dsystem_dt = dsystem_dt
else:
aux_update = [[], [], []]
# Store System State
# ------------------
self._store_step(self.storage, phi_tmp, z_mem, aux_update,
self.state_list, 0)
# Lock Trajectory
# ---------------
self.__initialized__ = True
else:
raise LockedException("HopsTrajectory.initialize()")