def test_extract_controls():
"""Check that we can extract a list of controls from a list of
objectives"""
# dummy objects
X = qutip.Qobj()
Y = qutip.Qobj()
f = lambda t: 0
g = lambda t: 0
h = lambda t: 0
d = lambda t: 0
# all of these dummy objects should be distinguishable
assert f is not g
assert g is not h
assert h is not d
assert X is not Y
H1 = [X, [X, f], [X, g]]
H2 = [X, [X, f], [X, h]]
H3 = [X, [X, d], X]
# check same Hamiltonian occuring in multiple objectives
objectives = [
krotov.Objective(initial_state=X, target=Y, H=H1),
krotov.Objective(initial_state=Y, target=X, H=H1),
]
controls = extract_controls(objectives)
maps = extract_controls_mapping(objectives, controls)
assert len(controls) == 2
assert f in controls
assert g in controls
assert len(maps) == len(objectives)
assert maps == [[[[1], [2]]], [[[1], [2]]]]
# check same control occuring in multiple Hamiltonians
objectives = [
krotov.Objective(initial_state=X, target=Y, H=H1),
krotov.Objective(initial_state=Y, target=X, H=H2),
krotov.Objective(initial_state=Y, target=X, H=H3),
]
controls = extract_controls(objectives)
assert len(controls) == 4
for c in (f, g, h, d):
assert c in controls
maps = extract_controls_mapping(objectives, controls)
assert maps[0] == [[[1], [2], [], []]]
assert maps[1] == [[[1], [], [2], []]]
assert maps[2] == [[[], [], [], [1]]]
def test_extract_controls_with_arrays():
"""Test extract_controls for controls that are numpy arrays"""
X, Y, Z = qutip.Qobj(), qutip.Qobj(), qutip.Qobj() # dummy Hamiltonians
u1, u2 = np.array([]), np.array([]) # dummy controls
psi0, psi_tgt = qutip.Qobj(), qutip.Qobj() # dummy states
assert X is not Y
assert Y is not Z
assert u1 is not u2
assert psi0 is not psi_tgt
H1 = [X, [Y, u1], [Z, u2]] # ham for first objective
H2 = [X, [Y, u2]] # ham for second objective
objectives = [
krotov.Objective(H1, psi0, psi_tgt),
krotov.Objective(H2, psi0, psi_tgt)]
controls = extract_controls(objectives)
control_map = extract_controls_mapping(objectives, controls)
assert controls == [u1, u2]
assert control_map[0] == [[[1], [2]]]
assert control_map[1] == [[[], [1]]]
