def tensor_train_template(init_rho, pb_index, rank=1):
"""Get rho_n from rho in a Tensor Train representation.
Parameters
----------
rho : np.ndarray
"""
n_vec = np.zeros((rank, ), dtype=DTYPE)
n_vec[0] = 1.0
root_array = np.tensordot(init_rho, n_vec, axes=0)
root = Tensor(name='root', array=root_array, axis=None)
max_terms = len(pb_index)
# +2: i and j
root[0] = (Leaf(name=max_terms), 0)
root[1] = (Leaf(name=max_terms + 1), 0)
for i in pb_index:
assert rank <= i
train = [root]
for k in range(max_terms):
if k < max_terms - 1:
array = np.eye(rank, pb_index[k] * rank)
array = np.reshape(array, (rank, -1, rank))
else:
array = np.eye(rank, pb_index[k])
spf = Tensor(name=k, array=array, axis=0)
l = Leaf(name=k)
spf[0] = (train[-1], 2)
spf[1] = (l, 0)
train.append(spf)
return root
def test_train(fname=None):
# Type settings
corr = Correlation(k_max=max_terms)
corr.symm_coeff = np.diag(corr_dict['s'].toarray())
corr.asymm_coeff = np.diag(corr_dict['a'].toarray())
corr.exp_coeff = np.diag(corr_dict['gamma'].toarray())
corr.delta_coeff = 0.0 # delta_coeff
corr.print()
n_dims = [max_tier] * max_terms
heom = Hierachy(n_dims, H, V, corr)
# Adopt TT
tensor_train = tensor_train_template(rho_0, n_dims)
root = tensor_train[0]
leaves_dict = {leaf.name: leaf for leaf in root.leaves()}
all_terms = []
for term in heom.diff():
all_terms.append([(leaves_dict[str(fst)], snd) for fst, snd in term])
solver = MultiLayer(root, all_terms)
#solver = ProjectorSplitting(root, all_terms)
solver.ode_method = 'RK45'
solver.snd_order = False
solver.atol = 1.e-7
solver.rtol = 1.e-7
solver.ps_method = 'split-unite'
projector = np.zeros((max_tier, 1))
projector[0] = 1.0
logger = Logger(filename=fname, level='info').logger
for n, (time, _) in enumerate(
solver.propagator(steps=count, ode_inter=dt_unit, split=False)):
if n % callback_interval == 0:
head = root.array
for t in tensor_train[1:]:
spf = Tensor.partial_product(t.array, 1, projector, 0)
head = Tensor.partial_product(head, head.ndim - 1, spf, 0)
rho = np.reshape(head, (4, -1))[:, 0]
logger.info("{} {} {} {} {}".format(time, rho[0], rho[1], rho[2],
rho[3]))
return
def tensor_tree_template(init_rho, pb_index, rank=1, nbranch=2):
"""Get rho_n from rho in a Tensor Tree representation.
Parameters
----------
rho : np.ndarray
"""
n_state = get_n_state(init_rho)
n_vec = np.zeros((rank, ), dtype=DTYPE)
n_vec[0] = 1.0
root_array = np.tensordot(init_rho, n_vec, axes=0)
max_terms = len(pb_index)
for i in pb_index:
assert rank <= i
# generate leaves
leaves = list(range(max_terms))
class new_spf(object):
counter = 0
prefix = 'SPF'
def __new__(cls):
name = cls.prefix + str(cls.counter)
cls.counter += 1
return name
importance = list(reversed(range(len(pb_index))))
graph, spf_root = huffman_tree(
leaves,
importances=importance,
obj_new=new_spf,
n_branch=nbranch,
)
root = 'root'
graph[root] = [str(max_terms), str(max_terms + 1), spf_root]
print(graph, root)
root = Tensor.generate(graph, root)
root.set_array(root_array)
bond_dict = {}
# Leaves
l_range = list(pb_index) + [n_state] * 2
for s, i, t, j in root.linkage_visitor():
if isinstance(t, Leaf):
bond_dict[(s, i, t, j)] = l_range[int(t.name)]
else:
bond_dict[(s, i, t, j)] = rank
autocomplete(root, bond_dict)
return root
def test_train(fname=None):
# HEOM metas
corr.print()
n_dims = [max_tier] * max_terms
heom = Hierachy(n_dims, H, V, corr)
# 2-site TT
tensor_train = tensor_train_template(rho_0, n_dims, rank=1)
root = tensor_train[0]
leaves_dict = {leaf.name: leaf for leaf in root.leaves()}
all_terms = []
for term in heom.diff():
all_terms.append([(leaves_dict[str(fst)], snd) for fst, snd in term])
solver = MultiLayer(root, all_terms)
solver.ode_method = 'RK45'
solver.snd_order = False
solver.svd_err = 1.e-8
solver.svd_rank = max_tier
solver.ps_method = 'unite'
projector = np.zeros((max_tier, 1))
projector[0] = 1.0
logger = Logger(filename=fname, level='info').logger
logger2 = Logger(filename=fname + '_norm', level='info').logger
for n, (time, _) in enumerate(solver.propagator(steps=count, ode_inter=dt_unit, split=True)):
#print('n = {}: '.format(n))
#for t in tensor_train:
# print('{}: {}'.format(t, t.shape))
if n % callback_interval == 0:
head = root.array
for t in tensor_train[1:]:
spf = Tensor.partial_product(t.array, 1, projector, 0)
head = Tensor.partial_product(head, head.ndim - 1, spf, 0)
rho = np.reshape(head, (4, -1))[:, 0]
logger2.warning("{} {}".format(time, rho[0] + rho[3]))
logger.info("{} {} {} {} {}".format(time, rho[0], rho[1], rho[2], rho[3]))
return
def simple_heom(init_rho, n_indices):
"""Get rho_n from rho with the conversion:
rho[i, j, n_0, ..., n_(k-1)]
Parameters
----------
rho : np.ndarray
"""
n_state = get_n_state(init_rho)
# Let: rho_n[0, :, :] = rho and rho_n[n, :, :] = 0
ext = np.zeros((np.prod(n_indices), ))
ext[0] = 1.0
new_shape = [n_state, n_state] + list(n_indices)
rho_n = np.reshape(np.tensordot(init_rho, ext, axes=0), new_shape)
root = Tensor(name='root', array=rho_n, axis=None)
d = len(n_indices)
root[0] = (Leaf(name=d), 0)
root[1] = (Leaf(name=d + 1), 0)
for k in range(d): # +2: i and j
root[k + 2] = (Leaf(name=k), 0)
return root
def test_drude_train():
eta = 0.05 # reorganization energy (dimensionless)
gamma_c = 0.05 # vibrational frequency (dimensionless)
max_tier = 10
max_terms = 3
J = pyheom.Drudian(eta, gamma_c)
corr_dict = pyheom.noise_decomposition(
J,
T=1, # temperature (dimensionless)
type_LTC='PSD',
n_PSD=max_terms - 1,
type_PSD='N-1/N')
s = corr_dict['s'].toarray()
a = corr_dict['a'].toarray()
gamma = corr_dict['gamma'].toarray()
delta = 0
omega_1 = 0.05
omega_2 = 0.02
H = np.array([[omega_1, omega_2], [omega_2, 0]])
V = np.array([[0, 0], [0, 1]])
corr = Correlation(k_max=max_terms, beta=1)
corr.symm_coeff = np.diag(s)
corr.asymm_coeff = np.diag(a)
corr.exp_coeff = np.diag(gamma)
corr.delta_coeff = delta
corr.print()
heom = Hierachy([max_tier] * max_terms, H, V, corr)
rho_0 = np.zeros((2, 2))
rho_0[0, 0] = 1
# TT HEOM
tensor_train = tensor_train_template(rho_0, [max_tier] * max_terms,
rank=max_tier)
root = tensor_train[0]
leaves_dict = {leaf.name: leaf for leaf in root.leaves()}
all_terms = []
for term in heom.diff():
all_terms.append([(leaves_dict[str(fst)], snd) for fst, snd in term])
solver = MultiLayer(root, all_terms)
solver.ode_method = 'RK45'
solver.snd_order = False
solver.max_ode_steps = 100000
# Define the obersevable of interest
projector = np.zeros((max_tier, 1))
projector[0] = 1.0
dat = []
for n, (time,
r) in enumerate(solver.propagator(
steps=20000,
ode_inter=0.01,
)):
head = root.array
for t in tensor_train[1:]:
spf = Tensor.partial_product(t.array, 1, projector, 0)
head = Tensor.partial_product(head, head.ndim - 1, spf, 0)
rho = np.reshape(head, (4, -1))[:, 0]
flat_data = [time] + list(rho)
dat.append(flat_data)
print("Time {} | Pop_1 {} | Total {}".format(time, rho[0],
rho[0] + rho[-1]))
return np.array(dat)
def test_mctdh(fname=None):
sys_leaf = Leaf(name='sys0')
ph_leaves = []
for n, (omega, g) in enumerate(ph_parameters, 1):
ph_leaf = Leaf(name='ph{}'.format(n))
ph_leaves.append(ph_leaf)
def ph_spf():
t = Tensor(axis=0)
t.name = 'spf' + str(hex(id(t)))[-4:]
return t
graph, root = huffman_tree(ph_leaves, obj_new=ph_spf, n_branch=2)
try:
graph[root].insert(0, sys_leaf)
except KeyError:
ph_leaf = root
root = Tensor()
graph[root] = [sys_leaf, ph_leaf]
finally:
root.name = 'wfn'
root.axis = None
stack = [root]
while stack:
parent = stack.pop()
for child in graph[parent]:
parent.link_to(parent.order, child, 0)
if child in graph:
stack.append(child)
# Define the detailed parameters for the ML-MCTDH tree
h_list = model.wfn_h_list(sys_leaf, ph_leaves)
solver = MultiLayer(root, h_list)
bond_dict = {}
# Leaves
for s, i, t, j in root.linkage_visitor():
if t.name.startswith('sys'):
bond_dict[(s, i, t, j)] = 2
else:
if isinstance(t, Leaf):
bond_dict[(s, i, t, j)] = max_tier
else:
bond_dict[(s, i, t, j)] = rank_wfn
solver.autocomplete(bond_dict)
# set initial root array
init_proj = np.array([[A, 0.0], [B, 0.0]]) / np.sqrt(A**2 + B**2)
root_array = Tensor.partial_product(root.array, 0, init_proj, 1)
root.set_array(root_array)
solver = MultiLayer(root, h_list)
solver.ode_method = 'RK45'
solver.cmf_steps = solver.max_ode_steps # constant mean-field
solver.ps_method = 'split'
solver.svd_err = 1.0e-14
# Define the obersevable of interest
logger = Logger(filename=prefix + fname, level='info').logger
logger2 = Logger(filename=prefix + 'en_' + fname, level='info').logger
for n, (time, r) in enumerate(
solver.propagator(
steps=count,
ode_inter=dt_unit,
split=True,
)):
if n % callback_interval == 0:
t = Quantity(time).convert_to(unit='fs').value
rho = r.partial_env(0, proper=False)
logger.info("{} {} {} {} {}".format(t, rho[0, 0], rho[0, 1],
rho[1, 0], rho[1, 1]))
en = np.trace(rho @ model.h)
logger2.info('{} {}'.format(t, en))
def ph_spf():
t = Tensor(axis=0)
t.name = 'spf' + str(hex(id(t)))[-4:]
return t
def ph_spf():
t = Tensor(axis=0, normalized=True)
t.name = str(hex(id(t)))[-4:]
return t
def ml(fname,
e,
v,
primitive_dim,
spf_dim,
ph_parameters,
steps=2000,
ode_inter=0.1):
logger = Logger(filename=fname).logger
# define parameters
sys_hamiltonian = np.array([[e, v], [v, -e]], dtype=DTYPE)
projector = np.array([[1.0, 0.0], [0.0, -1.0]],
dtype=DTYPE) # S in H_SB = S x B
primitive_dim = primitive_dim
spf_dim = spf_dim
# Define all Leaf tensors and hamiltonian we need
h_list = []
sys_leaf = Leaf(name='sys0')
h_list.append([(sys_leaf, -1.0j * sys_hamiltonian)])
ph_parameters = ph_parameters
leaves = []
for n, (omega, g) in enumerate(ph_parameters, 1):
ph = Phonon(primitive_dim, omega)
ph_leaf = Leaf(name='ph{}'.format(n))
leaves.append(ph_leaf)
# hamiltonian ph part
h_list.append([(ph_leaf, -1.0j * ph.hamiltonian)])
# e-ph part
op = ph.annihilation_operator + ph.creation_operator
h_list.append([(ph_leaf, g * op), (sys_leaf, -1.0j * projector)])
def ph_spf():
t = Tensor(axis=0, normalized=True)
t.name = str(hex(id(t)))[-4:]
return t
graph, root = huffman_tree(leaves, obj_new=ph_spf, n_branch=2)
try:
graph[root].insert(0, sys_leaf)
except KeyError:
ph_leaf = root
root = Tensor()
graph[root] = [sys_leaf, ph_leaf]
finally:
root.name = 'wfn'
root.axis = None
root.normalized = True
stack = [root]
while stack:
parent = stack.pop()
for child in graph[parent]:
parent.link_to(parent.order, child, 0)
if child in graph:
stack.append(child)
logger.info(f"graph:{graph}")
# Define the detailed parameters for the ML-MCTDH tree
solver = MultiLayer(root, h_list)
bond_dict = {}
# Leaves
for s, i, t, j in root.linkage_visitor():
if t.name and t.name.startswith('sys'):
bond_dict[(s, i, t, j)] = 2
else:
if isinstance(t, Leaf):
bond_dict[(s, i, t, j)] = primitive_dim
else:
bond_dict[(s, i, t, j)] = spf_dim
solver.autocomplete(bond_dict)
logger.info(f"bond_dict:{bond_dict}")
# set initial root array
a, b = 1.0, 0
init_proj = np.array([[a, 0.0], [b, 0.0]]) / np.sqrt(a**2 + b**2)
root_array = Tensor.partial_product(root.array, 0, init_proj, 1)
root.set_array(root_array)
# Define the computation details
solver.ode_method = 'RK45'
solver.snd_order = False
solver.cmf_steps = 100
# Define the obersevable of interest
logger.info('''# time rho00 rho01 rho10 rho11''')
for time, _ in solver.propagator(
steps=steps,
ode_inter=ode_inter,
split=True,
):
t = time
for tensor in root.visitor(axis=None):
tensor.reset()
tensor.normalize(forced=True)
rho = root.partial_env(0, proper=False)
for tensor in root.visitor(axis=None):
tensor.reset()
flat_data = [t] + list(np.reshape(rho, -1))
logger.info('{} {} {} {} {}'.format(*flat_data))
def sbm_ft(including_bath=False, snd=False):
# Define parameters of the model.
sbm = SpinBosonModel(
including_bath=including_bath,
e1=0.,
e2=Quantity(6500, 'cm-1').value_in_au,
v=Quantity(500, 'cm-1').value_in_au,
omega_list=[Quantity(2100, 'cm-1').value_in_au],
lambda_list=([Quantity(750, 'cm-1').value_in_au]),
dim_list=[10],
stop=Quantity(10000, 'cm-1').value_in_au,
n=32,
dim=30,
lambda_g=Quantity(2250, 'cm-1').value_in_au,
omega_g=Quantity(500, 'cm-1').value_in_au,
lambda_d=Quantity(1250, 'cm-1').value_in_au,
omega_d=Quantity(50, 'cm-1').value_in_au,
mu=Quantity(250, 'cm-1').value_in_au,
tau=Quantity(3, 'fs').value_in_au,
t_d=Quantity(6, 'fs').value_in_au,
omega=Quantity(13000, 'cm-1').value_in_au,
)
# Define the topological structure of the ML-MCTDH tree
graph, root = {
'ROOT': ['ELECs', 'I0s'],
'ELECs': ['ELEC', "ELEC'"],
'I0s': ['I0', "I0'"],
}, 'ROOT'
root = Tensor.generate(graph, root)
# Define the detailed parameters for the MC-MCTDH tree
solver = MultiLayer(root, sbm.h_list, f_list=sbm.f_list, use_str_name=True)
bond_dict = {}
# Leaves
for s, i, t, j in root.linkage_visitor():
if isinstance(t, Leaf):
try:
dim = sbm.dimensions[t.name]
except KeyError:
dim = sbm.dimensions[t.name[:-1]]
bond_dict[(s, i, t, j)] = dim
s_ax = s.axis
p, p_ax = s[s_ax]
bond_dict[(p, p_ax, s, s_ax)] = dim**2 if dim < 9 else 50
solver.autocomplete(bond_dict, max_entangled=True)
# Define the computation details
solver.settings(cmf_steps=10,
ode_method='RK45',
ps_method='split-unite',
snd_order=snd)
logging.info("Size of a wfn: {} complexes".format(len(root.vectorize())))
# Do the imaginary time propogation
inv_tem = 1 / 1000
steps = 100
for time, _ in solver.propagator(
steps=steps,
ode_inter=Quantity(inv_tem / steps / 2, unit='K-1').value_in_au,
split=True,
imaginary=True):
t = 2 * Quantity(time).convert_to(unit='K-1').value
z = solver.relative_partition_function
kelvin = 'inf' if abs(t) < 1.e-14 else 1.0 / t
logging.warning('Temperatue: {} K; ln(Z/Z_0): {}'.format(
kelvin, np.log(z)))
# Define the obersevable of interest
projector = np.array([[0., 0.], [0., 1.]])
op = [[[root[0][0][1][0], projector]]]
# Do the real time propogation
tp_list = []
steps = 100
root.is_normalized = True
for time, _ in solver.propagator(steps=steps,
ode_inter=Quantity(10 / steps,
'fs').value_in_au,
split=True,
imaginary=False):
t = Quantity(time).convert_to(unit='fs').value
p = solver.expection(op=op)
logging.warning('Time: {:.2f} fs; P2: {}'.format(t, p))
tp_list.append((t, p))
return np.array(tp_list)
def sbm_zt(including_bath=False, split=False, snd=False):
sbm_para_dict = {
'including_bath':
including_bath,
'e1':
0.,
'e2':
Quantity(6500, 'cm-1').value_in_au,
'v':
Quantity(500, 'cm-1').value_in_au,
"omega_list": [
Quantity(2100, 'cm-1').value_in_au,
Quantity(650, 'cm-1').value_in_au,
Quantity(400, 'cm-1').value_in_au,
Quantity(150, 'cm-1').value_in_au
],
'lambda_list': ([Quantity(750, 'cm-1').value_in_au] * 4),
'dim_list': [10, 14, 20, 30],
'stop':
Quantity(3 * 2250, 'cm-1').value_in_au,
'n':
32,
'dim':
30,
'lambda_g':
Quantity(2250, 'cm-1').value_in_au,
'omega_g':
Quantity(500, 'cm-1').value_in_au,
'lambda_d':
Quantity(1250, 'cm-1').value_in_au,
'omega_d':
Quantity(50, 'cm-1').value_in_au,
'mu':
Quantity(250, 'cm-1').value_in_au,
'tau':
Quantity(30, 'fs').value_in_au,
't_d':
Quantity(60, 'fs').value_in_au,
'omega':
Quantity(13000, 'cm-1').value_in_au,
}
sbm = SpinBosonModel(**sbm_para_dict)
# Define the topological structure of the ML-MCTDH tree
graph, root = sbm.autograph(n_branch=2)
root = Tensor.generate(graph, root)
# Define the detailed parameters for the MC-MCTDH tree
solver = MultiLayer(root, sbm.h_list, f_list=sbm.f_list, use_str_name=True)
bond_dict = {}
# Leaves
for s, i, t, j in root.linkage_visitor():
if isinstance(t, Leaf):
bond_dict[(s, i, t, j)] = sbm.dimensions[t.name]
# ELEC part
elec_r = root[0][0]
for s, i, t, j in elec_r.linkage_visitor(leaf=False):
raise NotImplementedError()
# INNER part
inner_r = root[1][0] if including_bath else root
if including_bath:
bond_dict[(root, 1, inner_r, 0)] = 60
for s, i, t, j in inner_r.linkage_visitor(leaf=False):
bond_dict[(s, i, t, j)] = 50
# OUTER part
if including_bath:
outer_r = root[2][0]
bond_dict[(root, 2, outer_r, 0)] = 20
for s, i, t, j in root[2][0].linkage_visitor(leaf=False):
bond_dict[(s, i, t, j)] = 10
solver.autocomplete(bond_dict, max_entangled=False)
# Define the computation details
solver.settings(
max_ode_steps=100,
cmf_steps=(1 if split else 10),
ode_method='RK45',
ps_method='s',
snd_order=snd,
)
root.is_normalized = True
# Define the obersevable of interest
projector = np.array([[0., 0.], [0., 1.]])
op = [[[root[0][0], projector]]]
t_p = []
for time, _ in solver.propagator(
steps=20,
ode_inter=Quantity(0.2, 'fs').value_in_au,
split=split,
move_energy=True,
):
t, p = (Quantity(time).convert_to(unit='fs').value,
solver.expection(op=op))
t_p.append((t, p))
logging.warning('Time: {:.2f} fs, P2: {}'.format(t, p))
if np.abs(p) > 0.5:
break
# Save the results
msg = 'split' if split else 'origin'
msg2 = 'snd' if snd else 'fst'
np.savetxt('sbm-zt-{}-{}.dat'.format(msg, msg2), t_p)
def test_drude_tree():
eta = 0.05 # reorganization energy (dimensionless)
gamma_c = 0.05 # vibrational frequency (dimensionless)
max_tier = 10
max_terms = 3
J = pyheom.Drudian(eta, gamma_c)
corr_dict = pyheom.noise_decomposition(
J,
T=1, # temperature (dimensionless)
type_LTC='PSD',
n_PSD=max_terms - 1,
type_PSD='N-1/N')
s = corr_dict['s'].toarray()
a = corr_dict['a'].toarray()
gamma = corr_dict['gamma'].toarray()
delta = 0
omega_1 = 0.05
omega_2 = 0.02
H = np.array([[omega_1, omega_2], [omega_2, 0]])
V = np.array([[0, 0], [0, 1]])
corr = Correlation(k_max=max_terms, beta=1)
corr.symm_coeff = np.diag(s)
corr.asymm_coeff = np.diag(a)
corr.exp_coeff = np.diag(gamma)
corr.delta_coeff = delta
corr.print()
heom = Hierachy([max_tier] * max_terms, H, V, corr)
rho_0 = np.zeros((2, 2))
rho_0[0, 0] = 1
root = tensor_tree_template(rho_0, [max_tier] * max_terms,
rank=max_tier // 2)
solver = MultiLayer(root, heom.diff(), use_str_name=True)
solver.ode_method = 'RK45'
solver.snd_order = False
solver.max_ode_steps = 100000
dat = []
for n, (time,
r) in enumerate(solver.propagator(
steps=20000,
ode_inter=0.01,
)):
if n % 100 == 0:
head = root.array
print(head.shape)
rho = Tensor.partial_product(r.array, 0, r[0][0].array, 0)
rho = np.reshape(rho, (-1, 4))
flat_data = [time] + list(rho[0])
dat.append(flat_data)
print("Time: {} | Pop 0: {} | Total: {}".format(
flat_data[0], flat_data[1], flat_data[1] + flat_data[-1]))
return np.array(dat)
def sbm_zt(including_bath=False, split=False, snd=False):
omega = 13000
sbm = SpinBosonModel(
including_bath=including_bath,
e1=0.,
e2=Quantity(6500, 'cm-1').value_in_au,
v=Quantity(500, 'cm-1').value_in_au,
omega_list=[
Quantity(2100, 'cm-1').value_in_au,
Quantity(650, 'cm-1').value_in_au,
Quantity(400, 'cm-1').value_in_au,
Quantity(150, 'cm-1').value_in_au
],
lambda_list=([Quantity(750, 'cm-1').value_in_au] * 4),
dim_list=[10, 14, 20, 30],
stop=Quantity(3 * 2250, 'cm-1').value_in_au,
n=32,
dim=30,
lambda_g=Quantity(2250, 'cm-1').value_in_au,
omega_g=Quantity(500, 'cm-1').value_in_au,
lambda_d=Quantity(1250, 'cm-1').value_in_au,
omega_d=Quantity(50, 'cm-1').value_in_au,
mu=Quantity(250, 'cm-1').value_in_au,
tau=Quantity(30, 'fs').value_in_au,
t_d=Quantity(60, 'fs').value_in_au,
omega=Quantity(omega, 'cm-1').value_in_au,
)
# Define the topological structure of the ML-MCTDH tree
graph, root = sbm.autograph(n_branch=2)
root = Tensor.generate(graph, root)
# Define the detailed parameters for the MC-MCTDH tree
solver = MultiLayer(root, sbm.h_list, f_list=sbm.f_list, use_str_name=True)
bond_dict = {}
# Leaves
for s, i, t, j in root.linkage_visitor():
if isinstance(t, Leaf):
bond_dict[(s, i, t, j)] = sbm.dimensions[t.name]
# ELEC part
elec_r = root[0][0]
for s, i, t, j in elec_r.linkage_visitor(leaf=False):
raise NotImplementedError()
# INNER part
inner_r = root[1][0] if including_bath else root
if including_bath:
bond_dict[(root, 1, inner_r, 0)] = 60
for s, i, t, j in inner_r.linkage_visitor(leaf=False):
bond_dict[(s, i, t, j)] = 50
# OUTER part
if including_bath:
outer_r = root[2][0]
bond_dict[(root, 2, outer_r, 0)] = 20
for s, i, t, j in root[2][0].linkage_visitor(leaf=False):
bond_dict[(s, i, t, j)] = 10
solver.autocomplete(bond_dict, max_entangled=False)
# Define the computation details
solver.settings(
max_ode_steps=100,
cmf_steps=(1 if split else 10),
ode_method='RK45',
ps_method='s',
snd_order=snd,
)
print("Size of a wfn: {} complexes".format(len(root.vectorize())))
root.is_normalized = True
# Define the obersevable of interest
projector = np.array([[0., 0.], [0., 1.]])
op = [[[root[0][0], projector]]]
t_p = []
for time, _ in solver.propagator(
steps=20,
ode_inter=Quantity(0.2, 'fs').value_in_au,
split=split,
move_energy=True,
):
t, p = (Quantity(time).convert_to(unit='fs').value,
solver.expection(op=op))
t_p.append((t, p))
logging.warning('Time: {:.2f} fs, P2: {}'.format(t, p))
if np.abs(p) > 0.5:
break
return np.array(t_p)
def ml(dof, e, v, eta, cutoff, scale=5, loc=None, steps=2000, ode_inter=0.1):
f_ = 'dof{}-eta{}.log'.format(dof, eta)
logger = Logger(filename=f_).logger
# define parameters
e = Quantity(e, 'cm-1').value_in_au
v = Quantity(v, 'cm-1').value_in_au
eta = Quantity(eta, 'cm-1').value_in_au
omega0 = Quantity(cutoff, 'cm-1').value_in_au
sys_hamiltonian = np.array([[-e / 2.0, v], [v, e / 2.0]], dtype=DTYPE)
projector = np.array([[0.0, 0.0], [0.0, 1.0]],
dtype=DTYPE) # S in H_SB = S x B
primitive_dim = 100
spf_dim = 20
# Spectrum function
def spec_func(omega):
if 0 < omega < omega0:
return eta
else:
return 0.0
# Define all Leaf tensors and hamiltonian we need
h_list = []
sys_leaf = Leaf(name='sys0')
h_list.append([(sys_leaf, -1.0j * sys_hamiltonian)])
ph_parameters = linear_discretization(spec_func, omega0, dof)
if loc is not None:
adj_pair = (ph_parameters[loc][0], ph_parameters[loc][1] * scale)
ph_parameters[loc] = adj_pair
leaves = []
for n, (omega, g) in enumerate(ph_parameters, 1):
ph = Phonon(primitive_dim, omega)
ph_leaf = Leaf(name='ph{}'.format(n))
leaves.append(ph_leaf)
# hamiltonian ph part
h_list.append([(ph_leaf, -1.0j * ph.hamiltonian)])
# e-ph part
op = ph.annihilation_operator + ph.creation_operator
h_list.append([(ph_leaf, g * op), (sys_leaf, -1.0j * projector)])
def ph_spf(n=0):
n += 1
return Tensor(name='spf{}'.format(n), axis=0)
graph, root = huffman_tree(leaves, obj_new=ph_spf, n_branch=2)
try:
graph[root].insert(0, sys_leaf)
except KeyError:
ph_leaf = root
root = Tensor()
graph[root] = [sys_leaf, ph_leaf]
finally:
root.name = 'wfn'
root.axis = None
print(graph)
stack = [root]
while stack:
parent = stack.pop()
for child in graph[parent]:
parent.link_to(parent.order, child, 0)
if child in graph:
stack.append(child)
# Define the detailed parameters for the ML-MCTDH tree
solver = MultiLayer(root, h_list)
bond_dict = {}
# Leaves
for s, i, t, j in root.linkage_visitor():
if t.name.startswith('sys'):
bond_dict[(s, i, t, j)] = 2
else:
if isinstance(t, Leaf):
bond_dict[(s, i, t, j)] = primitive_dim
else:
bond_dict[(s, i, t, j)] = spf_dim
solver.autocomplete(bond_dict)
# set initial root array
a, b = 1.0, 1.0
init_proj = np.array([[a, 0.0], [b, 0.0]]) / np.sqrt(a**2 + b**2)
root_array = Tensor.partial_product(root.array, 0, init_proj, 1)
root.set_array(root_array)
# Define the computation details
solver.ode_method = 'RK45'
solver.snd_order = True
solver.cmf_steps = 1
root.is_normalized = True
# Define the obersevable of interest
logger.info('''# time/fs rho00 rho01 rho10 rho11''')
for time, _ in solver.propagator(
steps=steps,
ode_inter=Quantity(ode_inter, 'fs').value_in_au,
split=True,
):
t = Quantity(time).convert_to(unit='fs').value
for tensor in root.visitor(axis=None):
tensor.reset()
rho = root.partial_env(0, proper=False)
for tensor in root.visitor(axis=None):
tensor.reset()
flat_data = [t] + list(np.reshape(rho, -1))
logger.info('{} {} {} {} {}'.format(*flat_data))
def ph_spf(n=0):
n += 1
return Tensor(name='spf{}'.format(n), axis=0)
