Skip to content

Commit

Permalink
Update environment.py docstrings, add tests
Browse files Browse the repository at this point in the history 
Note:  in adding tests, i fixed a couple of bugs in existing tests.
  • Loading branch information
@nwlambert
nwlambert committed Dec 23, 2024
1 parent 912a6ed commit e40fa6e
Show file tree
Hide file tree
Showing 4 changed files with 225 additions and 12 deletions.
8 changes: 7 additions & 1 deletion qutip/core/environment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1530,6 +1530,10 @@ class CFExponent:
"+" and "-" are fermionic exponents.
"Input", "Output_fn_L", "Output_fn_R", "Output_L", "Output_R" are
special types used with InputOutputBaths from bofin_baths, and
are not used with other solvers.
ck : complex
The coefficient of the excitation term.
Expand Down Expand Up @@ -1560,7 +1564,9 @@ class CFExponent:
All of the parameters are also available as attributes.
"""
types = enum.Enum("ExponentType", ["R", "I", "RI", "+", "-", "Input", "Output", "Output_fn_L", "Output_fn_R", "Output_L", "Output_R"])
types = enum.Enum("ExponentType", ["R", "I", "RI", "+", "-", "Input",
"Output_fn_L", "Output_fn_R",
"Output_L", "Output_R"])

def _check_ck2(self, type, ck2):
if type == self.types["RI"]:
Expand Down
13 changes: 7 additions & 6 deletions qutip/solver/heom/bofin_solvers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -801,7 +801,7 @@ def _grad_prev(self, he_n, k):

def _grad_prev_bosonic(self, he_n, k):
if self.ados.exponents[k].type == BathExponent.types.R:
op = _data.mul(
op = _data.mul(
self._s_pre_minus_post_Q[k],
-1j * he_n[k] * self.ados.ck[k],
)
Expand Down Expand Up @@ -952,7 +952,7 @@ def _rhs(self):
[self._sup_shape * self._n_ados], [self._sup_shape * self._n_ados]
]
ops = _GatherHEOMRHS(
self.ados.idx, block=self._sup_shape,
self.ados.idx, block=self._sup_shape,
nhe=self._n_ados, rhs_dims=rhs_dims
)

Expand Down Expand Up @@ -1499,9 +1499,9 @@ def gather_td(self):

self._ops_td.sort(key=lambda x: x[4])
RHStemp = 0
#we group terms by 'exponent'/TD funct given by 'kpos'
#and add them together. Most efficient construction
#we could find for the moment.
# We group terms by 'exponent'/TD funct given by 'kpos'
# and add them together. Most efficient construction
# we could find for the moment.

for k, ops in groupby(self._ops_td, key=lambda x: x[4]):
ops = np.array(list(ops), dtype=[
Expand All @@ -1515,6 +1515,7 @@ def gather_td(self):
RHStemp += QobjEvo([Qobj(_csr._from_csr_blocks(
ops["row"], ops["col"], ops["op"],
self._n_blocks, self._block_size,
)), ops["func"][0]])
), dims=self._rhs_dims),
ops["func"][0]])

return RHStemp
214 changes: 209 additions & 5 deletions qutip/tests/solver/heom/test_bofin_solvers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,14 @@

from qutip import (
basis, destroy, expect, liouvillian, qeye, sigmax, sigmaz, sigmay,
tensor, Qobj, QobjEvo, fidelity, fdestroy
tensor, Qobj, QobjEvo, fidelity, fdestroy, mesolve
)
from qutip.core import data as _data
from qutip.solver.heom.bofin_baths import (
BathExponent,
Bath,
BosonicBath,
InputOutputBath,
DrudeLorentzBath,
DrudeLorentzPadeBath,
UnderDampedBath,
Expand Down Expand Up @@ -348,17 +349,17 @@ def __init__(self, N, Lambda, Omega, gamma_b):

def bath_coefficients(self):
ck_real = [0.5 * self.Lambda**2, 0.5 * self.Lambda**2]
vk_real = [0.5 * 1.0j * self.Lambda**2, -0.5 * 1.0j * self.Lambda**2]
ck_imag = [0.5 * 1.0j * self.Lambda**2, -0.5 * 1.0j * self.Lambda**2]

ck_imag = [
vk_real = [
-1.0j * self.Omega + self.gamma_b / 2,
1.0j * self.Omega + self.gamma_b / 2,
]
vk_imag = [
-1.0j * self.Omega + self.gamma_b / 2,
1.0j * self.Omega + self.gamma_b / 2,
]
return ck_real, ck_imag, vk_real, vk_imag
return ck_real, vk_real, ck_imag, vk_imag


class DiscreteLevelCurrentModel:
Expand Down Expand Up @@ -1193,6 +1194,173 @@ def test_parity(self):
expected = np.diag([0.10623, 0.39376, 0.39376, 0.10623])
np.testing.assert_allclose(rhoss, expected, atol=1e-5)

def test_InputOutput_single_mode_input(self):
# Defining the system Hamiltonian

Del = 2 * np.pi * 1.0
Hsys = 0.5 * Del * sigmaz() # system Hamiltonian
# Initial state of the system.
rho0 = basis(2, 0) * basis(2, 0).dag()
tlist = np.linspace(0, 20/Del, 100)

# System-bath coupling (underdamed spectral density)
Q = sigmax() # coupling operator

# Bath properties:
Gamma = .1 * Del # cut off frequency
lam = .1 * Del # coupling strength
Om = .2 * Del # resonance frequency
NC = 12
bosonic_mode = BosonicMode(
N=4, Omega=Om, Lambda=lam, gamma_b=2*Gamma,
)

bath = BosonicBath(
Q, *bosonic_mode.bath_coefficients(),
tag="SHO"
)

ck_in = [lambda t: lam * np.exp(-1.0j*Om*t - Gamma*t),
lambda t: lam * np.exp(1.0j*Om*t - Gamma*t)]

def input1(t):
return lam * np.exp(-1.0j*Om*t - Gamma*t)

def input2(t):
return lam * np.exp(1.0j*Om*t - Gamma*t)

ck_in = [input1, input2]
bath_input = InputOutputBath(Q, ck_input=ck_in, tag="input")

options = {
"store_ados": True,
}
SHO_model = HEOMSolver(Hsys, [bath, bath_input], NC, options=options)

resultSHO = SHO_model.run(rho0, tlist)
result_input = []

for t in range(len(tlist)):
label = resultSHO.ado_states[t].filter(level=2, tags=["input",
"input"])

state = (resultSHO.ado_states[t].extract(label[0]))

result_input.append(expect(resultSHO.states[t], sigmaz()) -
expect(state, sigmaz()))

Nbos = 12
a = qeye(2) & destroy(Nbos)
H = Hsys & qeye(Nbos)

H = H + lam * tensor(Q,qeye(Nbos)) * (a+a.dag()) + Om * a.dag() * a

# one photon initial condition with mesolve()
resultME1 = mesolve(H, rho0 & (basis(Nbos, 1)*basis(Nbos, 1).dag()),
tlist, [np.sqrt(Gamma*2)*a])
np.testing.assert_allclose(result_input,
expect(resultME1.states,
sigmaz() & qeye(Nbos)),
atol=1e-5)

def test_InputOutput_single_mode_input_output_spectral_decomp(self):
# Defining the system Hamiltonian

Del = 2 * np.pi * 1.0
Hsys = 0.5 * Del * sigmaz() # system Hamiltonian
# Initial state of the system.
rho0 = basis(2, 0) * basis(2, 0).dag()
tlist = np.linspace(0, 20/Del, 100)

# System-bath coupling (underdamed spectral density)
Q = sigmax() # coupling operator

# Bath properties:
Gamma = .1 * Del # cut off frequency
lam = .1 * Del # coupling strength
Om = .2 * Del # resonance frequency
NC = 12
bosonic_mode = BosonicMode(
N=4, Omega=Om, Lambda=lam, gamma_b=2*Gamma,
)

bath = BosonicBath(
Q, *bosonic_mode.bath_coefficients(),
tag="SHO"
)

ck_in_1 = [lambda t: lam * np.exp(-1.0j*Om*t - Gamma*t)]
ck_in_2 = [lambda t: lam * np.exp(1.0j*Om*t - Gamma*t)]

bath_input_1 = InputOutputBath(Q, ck_input=ck_in_1, tag="input1")
bath_input_2 = InputOutputBath(Q, ck_input=ck_in_2, tag="input2")

ck_output_L = [lam]
ck_output_R = [lam]
vk_output_L = [1.0j*Om+Gamma]
vk_output_R = [-1.0j*Om+Gamma]
bath_output_1R = InputOutputBath(Q,
ck_output_R=ck_output_R,
vk_output_R=vk_output_R,
tag="output1")
bath_output_2L = InputOutputBath(Q,
ck_output_L=ck_output_L,
vk_output_L=vk_output_L,
tag="output2")
options = {
"store_ados": True,
}
SHO_model = HEOMSolver(Hsys, [bath, bath_output_1R, bath_output_2L,
bath_input_1, bath_input_2],
NC, options=options)

resultSHO = SHO_model.run(rho0, tlist)

result_output = []

for t in range(len(tlist)):
label = resultSHO.ado_states[t].filter(level=2,
tags=["output2",
"input1"])
s0110 = (np.exp(1.0j*Om*tlist[t] - Gamma*tlist[t]) *
resultSHO.ado_states[t].extract(label[0]).tr())

label = resultSHO.ado_states[t].filter(level=2,
tags=["output1",
"input2"])
s1001 = (np.exp(-1.0j*Om*tlist[t] - Gamma*tlist[t]) *
resultSHO.ado_states[t].extract(label[0]).tr())

label = resultSHO.ado_states[t].filter(level=2,
tags=["output1",
"output2"])
s1100 = (resultSHO.ado_states[t].extract(label[0]).tr())

label = resultSHO.ado_states[t].filter(level=4,
tags=["output1",
"output2",
"input1",
"input2"])
s1111 = (resultSHO.ado_states[t].extract(label[0]).tr())

result_output.append(resultSHO.states[t].tr() *
np.exp(-2.0*Gamma*tlist[t]) -
s0110-s1001-s1100+s1111)

Nbos = 12
a = qeye(2) & destroy(Nbos)
H = Hsys & qeye(Nbos)

H = H + lam * tensor(Q,qeye(Nbos)) * (a+a.dag()) + Om * a.dag() * a

#one photon initial condition:
resultME1 = mesolve(H, rho0 & (basis(Nbos, 1)*basis(Nbos, 1).dag()),
tlist, [np.sqrt(Gamma*2)*a])

np.testing.assert_allclose(result_output,
expect(resultME1.states, a.dag()*a),
atol=1e-5)


class TestHeomsolveFunction:
@pytest.mark.parametrize(['evo'], [
Expand Down Expand Up @@ -1873,7 +2041,7 @@ def test_simple_gather(self):
def f(label):
return int(label.lstrip("o"))

gather_heoms = _GatherHEOMRHS(f, block=2, nhe=3)
gather_heoms = _GatherHEOMRHS(f, block=2, nhe=3, rhs_dims=[2, 2])

for i in range(3):
for j in range(3):
Expand All @@ -1896,3 +2064,39 @@ def f(label):

np.testing.assert_array_equal(op.to_array(), expected_op)
assert isinstance(op, _data.CSR)

def test_simple_gather_td(self):
def f(label):
return int(label.lstrip("o"))

gather_heoms = _GatherHEOMRHS(f, block=2, nhe=3, rhs_dims=[2*3, 2*3])

for i in range(3):
for j in range(3):
base = 10 * (j * 2) + (i * 2)
block_op = _data.to(
_data.CSR,
_data.create(np.array([
[base, base + 10],
[base + 1, base + 11],
]))
)
gather_heoms.add_op(f"o{i}", f"o{j}", block_op,
ck_td_factor=lambda t: np.exp(1.0j * t
- 0.5 * t),
ado_pos=j)

op = gather_heoms.gather_td()

expected_op = QobjEvo([Qobj(_data.to(
_data.CSR,
_data.create(
np.array([
[10 * i + j for i in range(2 * 3)]
for j in range(2 * 3)
], dtype=np.complex128))),
dims=[6, 6]),
lambda t: np.exp(1.0j * t - 0.5 * t)])

np.testing.assert_array_equal(op(0.4), expected_op(0.4))
assert isinstance(op, QobjEvo)
2 changes: 2 additions & 0 deletions qutip/tests/solver/heom/test_heom.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
BathExponent,
Bath,
BosonicBath,
InputOutputBath,
DrudeLorentzBath,
DrudeLorentzPadeBath,
UnderDampedBath,
Expand All @@ -26,6 +27,7 @@ def test_api(self):
assert BathExponent
assert Bath
assert BosonicBath
assert InputOutputBath
assert DrudeLorentzBath
assert DrudeLorentzPadeBath
assert UnderDampedBath
Expand Down

0 comments on commit e40fa6e

Please sign in to comment.