def test_heom(fname=None):
ph_dims = list(np.repeat(model.ph_dims, 2))
n_dims = ph_dims if model.bath_dims is None else ph_dims + model.bath_dims
print(n_dims)
root = tensor_tree_template(rho_0, n_dims, rank=rank_heom)
leaves = root.leaves()
h_list = model.heom_h_list(leaves[0], leaves[1], leaves[2:], beta=beta)
solver = MultiLayer(root, h_list)
solver.ode_method = 'RK45'
solver.cmf_steps = solver.max_ode_steps # use constant mean-field
solver.ps_method = 'split'
#solver.svd_err = 1.0e-10
# 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,
)):
# renormalized by the trace of rho
norm = np.trace(np.reshape(np.reshape(r.array, (4, -1))[:, 0], (2, 2)))
r.set_array(r.array / norm)
if n % callback_interval == 0:
t = Quantity(time).convert_to(unit='fs').value
rho = np.reshape(r.array, (4, -1))[:, 0]
logger.info("{} {} {} {} {}".format(t, rho[0], rho[1], rho[2],
rho[3]))
en = np.trace(np.reshape(rho, (2, 2)) @ model.h)
logger2.info('{} {}'.format(t, en))
return
def __partial_product(array1, i, array2, j):
r"""Times a matrix to a tensor.
|
-- 1 -i--j- 2 --
|
Parameters
----------
array1 : ndarray
i : int
array2 : 2-d ndarray
j : int
Returns
-------
ans : ndarray
"""
shape_1, shape_2 = map(list, (array1.shape, array2.shape))
n, m = shape_1.pop(i), shape_2.pop(j)
new_shape = shape_1[:i] + shape_2 + shape_1[i:]
array1 = np.moveaxis(array1, i, -1)
array1 = np.reshape(array1, (-1, n))
array2 = np.moveaxis(array2, j, 0)
array2 = np.reshape(array2, (m, -1))
ans = np.dot(array1, array2)
ans = np.reshape(ans, shape_1 + [-1])
ans = np.moveaxis(ans, -1, i)
ans = np.reshape(ans, new_shape)
return ans
def _partial_product(i, a, b):
r"""
Parameters
----------
i : {int, None}
a : (..., n, ...) ndarray
b : (..., m, ...) ndarray
Returns
-------
mat : (n, m) ndarray
Or return a float if i is None.
"""
if i is None:
a = np.reshape(a, -1)
b = np.reshape(b, -1)
else:
n = a.shape[i]
m = b.shape[i]
a = np.moveaxis(a, i, 0)
a = np.reshape(a, (n, -1))
b = np.moveaxis(b, i, -1)
b = np.reshape(b, (-1, m))
mat = np.dot(a, b)
return mat
def opt_one_site(env_tensor, A):
r"""Solve eigenvalue problem at one site A.
Parameters
----------
env_tensor : LinearOperator
A linear operator, which is the environment tensor of A s. t.
(t, j, l) ndarray -> (s, i, k) ndarray.
A : (t, j, l) ndarray
A MPS matrix.
Returns
-------
E : float
Rayleigh quotient
V : (t, j, l) ndarray
optimized one-site MPS matrix
"""
A = np.reshape(A, env_tensor.size)
E, V = eigsh(env_tensor, 1, v0=A, which='SA')
E, V = E[0], np.reshape(V[:, 0], env_tensor.io_shape)
# A = np.reshape(V, env_tensor.size)
# E_squared = np.dot(A, env_tensor.matvec(env_tensor.matvec(A)))
# E_rms = math.sqrt(E_squared)
# logging.debug("INFO: sqrt(<E^2>): {}".format(E_rms))
return E, V
def partial_trace(array1, i, array2, j):
r"""Partial trace of 2 tensors, return a matrix.
+----+
| |
-i- 1 -- 2 -j-
Parameters
----------
array1 : ndarray
i : {int, None}
if i is None then j must be None.
array2 : ndarray
j : {int, None}
if j is None then i must be None.
Returns
-------
ans : ndarray
Of shape `(n, m)`
"""
if i is not None and j is not None:
shape_1, shape_2 = map(list, (array1.shape, array2.shape))
n, m = shape_1[i], shape_2[j]
array1 = np.moveaxis(array1, i, 0)
array1 = np.reshape(array1, (n, -1))
array2 = np.moveaxis(array2, j, -1)
array2 = np.reshape(array2, (-1, m))
elif i is None and j is None:
array1 = np.reshape(array1, -1)
array2 = np.reshape(array2, -1)
else:
raise TypeError('Invalid parameters i={} and j={}!'.format(i, j))
ans = np.dot(array1, array2)
return ans
def normalize(self, forced=False):
"""Normalize the array of self. Only work when self.normalized.
Set `forced` to `True` to normalize any way.
"""
array = self.array
if array is None or (not self.normalized and not forced):
return
axis = self.axis
if axis is None:
norm = np.array(self.local_norm())
self.set_array(array / norm)
ans = norm
else:
norm = linalg.norm
shape = self.shape
dim = shape.pop(axis)
array = np.reshape(np.moveaxis(array, axis, 0), (dim, -1))
vecs = []
norm_list = []
for vec_i in array:
for vec_j in vecs:
vec_i -= vec_j * np.dot(np.conj(vec_j), vec_i)
norm_ = norm(vec_i)
vecs.append(vec_i / norm_)
norm_list.append(norm_)
array = np.array(vecs)
array = np.moveaxis(np.reshape(array, [-1] + shape), 0, axis)
self.set_array(array)
ans = norm_list
return ans
def diff(t, y):
"""This function will not change the arrays in tensor network.
"""
tensor.set_array(np.reshape(y, tensor.shape))
ans = np.zeros_like(y)
for n in self.term_visitor():
ans += np.reshape(self._single_diff(tensor, n), -1)
return ans
def energy_expection(self, vec=None):
if vec is None:
vec = self.vec
a_tensor = self.get_sub_vec(-1, vec)
mod_terms = self.update_mod_terms(vec=vec, write=False)
ans = 0.
for r, term in enumerate(mod_terms):
h_a = self._coeff_op(a_tensor, term)
h_a = np.reshape(h_a, -1)
a_h = np.conj(np.reshape(a_tensor, -1))
ans += np.dot(a_h, h_a)
return ans
def _move_right(mat, mat_next, _trunc):
shape = np.shape(mat)
# SVD on mat[si, j]
mat = np.reshape(mat, (shape[0] * shape[1], shape[2]))
U, S, V = np.linalg.svd(mat, full_matrices=0)
# truncated to compress.
U, S, V, compress_error = compress_svd(U, S, V, _trunc)
mat = np.reshape(U, (shape[0], shape[1], -1)) # U[si, m]
SV = np.matmul(np.diag(S), V)
mat_next = np.einsum('mj,tjk->tmk', SV, mat_next)
return mat, mat_next
def diff(t, y):
"""This function will not change the arrays in tensor network.
"""
origin = tensor.array
tensor.set_array(np.reshape(y, tensor.shape))
ans = np.zeros_like(y)
self.time = t if not imaginary else None
for n in self.term_visitor(use_cache=True):
ans += np.reshape(self._single_eom(tensor, n, cache=cache), -1)
if self.init_energy is not None:
ans -= self.init_energy * y
ans /= self.coefficient()
tensor.set_array(origin)
return ans
def _matvec(self, vec):
v = np.reshape(vec, self.io_sizes)
ans = np.zeros_like(v)
for i, h_i in enumerate(self.h_list):
v_i = np.swapaxes(v, -1, i)
size_i = (self.io_sizes[:i] + self.io_sizes[i + 1:] + [self.io_sizes[i]])
v_i = np.reshape(v_i, (-1, self.io_sizes[i]))
tmp = np.array(list(map(h_i.dot, v_i)))
tmp = np.reshape(tmp, size_i)
ans += np.swapaxes(tmp, -1, i)
ans = np.reshape(ans, -1)
if self.v_rst is not None:
ans = ans + self.v_rst * vec
return ans
def _move_left(mat, mat_prev, _trunc):
shape = np.shape(mat)
mat = np.reshape(
np.transpose(mat, (1, 0, 2)), # mat[i, s, j]
(shape[1], shape[0] * shape[2])) # SVD on mat[i, sj]
U, S, V = np.linalg.svd(mat, full_matrices=0)
# truncated to compress.
U, S, V, compress_error = compress_svd(U, S, V, _trunc)
mat = np.reshape(V, (-1, shape[0], shape[2])) # V[m, sj]
mat = np.transpose(mat, (1, 0, 2)) # mat[s, m, j]
US = np.matmul(U, np.diag(S))
mat_prev = np.einsum('rhi,im->rhm', mat_prev, US)
return mat, mat_prev
def plot_wf(self, vec, msg=None):
"""Plot 2D wavefunction.
Parameters
----------
vec : (N,) ndarray
msg : string
Notes
-----
This is a generator. Use .next() to plot, and .send(vec, msg) to
plot next wavefunction with the same figure parameters.
"""
x_lim, y_lim = self.bound_list[:2]
x, y = self.grid_points_list[:2]
shape = self.n_list[:2]
x, y = np.meshgrid(x, y)
z_lim = int(np.max(np.abs(vec)) * 15) / 10
bound = np.linspace(-z_lim, z_lim, 100)
while vec is not None:
vec = np.reshape(vec, shape)
if msg is None:
msg = int(time.time())
with figure() as fig:
plt.contourf(x, y, vec, bound, cmap='seismic')
plt.colorbar(boundaries=bound)
plt.savefig('functions-{}.pdf'.format(msg))
received = yield
if received is None:
raise StopIteration
else:
vec, msg = received
def init_state(self):
r"""Form the initial vector according to shape list::
n_0| | |n_p-1
C_0 ... C_p-1
\ | /
m_0 \|/ m_p-1
A
Returns
-------
init : (self.size,) ndarray
Formally, init = np.concatenate([C_0, ..., C_p-1, A], axis=None),
where C_i is a (n_i * m_i,) ndarray, i <- {0, ..., p-1},
A is a (M,) ndarray, and M = m_0 * ... * m_p-1, m_i < n_i.
"""
dvr_list = self.dvr_list
c_list = []
for i, (_, m_i) in enumerate(self.shape_list[:-1]):
_, v_i = dvr_list[i].solve(n_state=m_i)
v_i = np.transpose(v_i)
c_list.append(np.reshape(v_i, -1))
vec_a = np.zeros(self.size_list[-1])
vec_a[0] = 1.0
vec_list = c_list + [vec_a]
init = np.concatenate(vec_list, axis=None)
self.vec = init
self.update_mod_terms()
return init
def _split_prop(self, tensor, tau=0.01, imaginary=False, cache=False):
def diff(t, y):
"""This function will not change the arrays in tensor network.
"""
origin = tensor.array
tensor.set_array(np.reshape(y, tensor.shape))
ans = np.zeros_like(y)
self.time = t if not imaginary else None
for n in self.term_visitor(use_cache=True):
ans += np.reshape(self._single_eom(tensor, n, cache=cache), -1)
if self.init_energy is not None:
ans -= self.init_energy * y
ans /= self.coefficient()
tensor.set_array(origin)
return ans
def reformer():
self._form_env(update=True)
return
def updater(y):
tensor.set_array(np.reshape(y, tensor.shape))
tensor.normalize(forced=(not imaginary))
if tensor.axis is None:
self.root = tensor
else:
raise RuntimeError("Cannot propagate on Tensor {}" "which is not a root node!".format(tensor))
logging.debug(__("* Propagating at {} ({}) for {}", tensor, tensor.shape, tau))
y0 = np.reshape(tensor.array, -1)
_re = None if self.f_list is None else reformer
with self.log_inner_product(level=logging.DEBUG):
self._solve_ode(diff, y0, tau, _re, updater)
return tensor
def tensorize(self, vec, use_aux=False, shape_dict=None):
"""Read a vector and set the arrays in the tensors in the network of
self.
Parameters
----------
vec : (n,) ndarray
A vector which should have the same shape with self.vectorize()
use_aux : bool
`True` to write the arrays in .aux, else to write arrays in
.array.
shape_dict: dict
Where to load the shape of each tensor.
"""
start = 0
for t in self.visitor(leaf=False):
shape = t.shape if shape_dict is None else shape_dict[t]
end = start + np.prod(shape)
array = np.reshape(vec[start:end], shape)
if use_aux:
t.aux = array
else:
t.set_array(array)
start = end
return
def vectorize(self, use_aux=False, shape_dict=None):
"""Return a vector from the arrays in the tensors in the network of
self.
Parameters
----------
use_aux : bool
`True` to vectorize the arrays in .aux, else to use arrays in
.array.
shape_dict: dict
Where to save the shape of each tensor.
Return
------
vec : (n,) ndarray
"""
vec_list = []
for t in self.visitor(leaf=False):
array = t.aux if use_aux else t.array
vec = np.reshape(array, -1)
vec_list.append(vec)
if shape_dict is not None:
shape_dict[t] = t.shape
ans = np.concatenate(vec_list, axis=None)
return ans
def init_state(self):
v = 1.
for i in range(self.rank):
_, v_i = self.dvr_list[i].solve(n_state=1)
v = np.tensordot(v, v_i[0], axes=0)
v = np.reshape(v, -1)
return v
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 coarse_grain_mpo(W, X):
"""Coarse-graining of two site MPO into one site.::
|su |s |u
-a- R -c- = -a- W -b- X -c-
|tv |t |v
Parameters
----------
W : (a, b, s, t) ndarray
A MPS matrix.
X : (b, c, u, v) ndarray
A MPS matrix.
Returns
-------
R : (a, c, su, tv) ndarray
A MPS matrix.
"""
R = np.einsum("abst,bcuv->acsutv", W, X)
sh = [
W.shape[0], # a
X.shape[1], # c
W.shape[2] * X.shape[2], # su
W.shape[3] * X.shape[3]
] # tv
return np.reshape(R, sh)
def test_simple(fname=None):
# Type settings
corr.print()
n_dims = [max_tier] * max_terms
heom = Hierachy(n_dims, H, V, corr)
# Adopt MCTDH
root = simple_heom(rho_0, n_dims)
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 = ProjectorSplitting(root, all_terms)
solver = MultiLayer(root, all_terms)
solver.ode_method = 'RK45'
solver.snd_order = False
# Define the obersevable of interest
logger = Logger(filename=fname, level='info').logger
for n, (time, r) in enumerate(solver.propagator(
steps=count,
ode_inter=dt_unit,
)):
try:
if n % callback_interval == 0:
rho = np.reshape(r.array, (-1, 4))[0]
logger.info("{} {} {} {} {}".format(time, rho[0], rho[1], rho[2], rho[3]))
except:
break
return
def fine_grain_mps(C, dims, direction, _trunc=False):
"""Fine-graining of one-site MPS into three site by SVD.::
|st |s |t
-i- C -k- = -i- A -m- B -k-
Parameters
----------
C : (st, i, k) ndarray
A MPS matrix.
dims : (int, int)
[s, t].
direction : {'>', '<'}
'>': move to right; '<' move to left
_trunc : int, optional
Set m in compress_svd to _trunc. Not compressed if not _trunc.
Returns
-------
A : (s, i, m) ndarray
A MPS matrix.
B : (t, m, k) ndarray
A MPS matrix.
Notes
-----
If direction == '>', A is (left-)canonical;
if direction == '<', B is (right-)canonical.
"""
sh = dims + [C.shape[1], C.shape[2]] # [s, t, i, k]
mat = np.reshape(C, sh)
mat = np.transpose(mat, (0, 2, 1, 3)) # mat[s, i, t, k]
mat = np.reshape(mat, (sh[0] * sh[2], sh[1] * sh[3])) # mat[si, tk]
U, S, V = np.linalg.svd(mat, full_matrices=0)
if _trunc:
U, S, V, compress_error = compress_svd(U, S, V, _trunc)
if direction == '>':
A = U
B = np.matmul(np.diag(S), V) # [m, tk]
elif direction == '<':
A = np.matmul(U, np.diag(S))
B = V
A = np.reshape(A, (sh[0], sh[2], -1)) # [s, i, m]
B = np.reshape(B, (-1, sh[1], sh[3])) # [m, t, k]
B = np.transpose(B, (1, 0, 2)) # [t, m, k]
return A, B
def _partial_transform(i, tensor, mat):
r"""
Parameters
----------
i : int
tensor : (..., m_i, ...) ndarray
mat : (n_i, m_i) ndarray
Returns
-------
tensor : (..., n_i, ...) ndarray
"""
tensor_shape = list(tensor.shape)
shape = tensor_shape[:i] + tensor_shape[i + 1:] + [mat.shape[0]]
v_i = np.swapaxes(tensor, -1, i)
v_i = np.reshape(v_i, (-1, mat.shape[1]))
v_i = np.array(list(map(mat.dot, v_i)))
v_i = np.reshape(v_i, shape)
v_i = np.swapaxes(v_i, -1, i)
return v_i
def fbr2dvr_vec(self, vec):
r"""Transform a vector from FBR to DVR.
Parameters
----------
mat : (n,) ndarray
Returns
-------
(n,) ndarray
"""
vec = np.reshape(vec, -1)
return np.dot(np.transpose(self._u_mat), vec)
def _normalizer(vec):
ans = []
for i in range(self.rank + 1):
vec_i = self.get_sub_vec(i, vec)
vec_i = np.reshape(vec_i, -1)
if i == self.rank:
std_norm = sqrt(1.0)
else:
std_norm = sqrt(self.shape_list[i][1])
norm = (linalg.norm(vec_i) / std_norm)
logging.debug(__('norm {}: {}', i, norm))
ans.append(vec_i / norm)
ans = np.concatenate(ans, axis=None)
return ans
def projector(self, comp=False):
"""[Deprecated] Return the projector corresponding to self.
Returns
-------
ans : ndarray
"""
axis = self.axis
if axis is not None:
array = self.array
shape = self.shape
dim = shape.pop(self.axis)
comp_dim = np.prod(shape)
array = np.moveaxis(array, axis, -1)
array = np.reshape(array, (-1, dim))
array_h = np.conj(np.transpose(array))
ans = np.dot(array, array_h)
if comp:
identity = np.identity(comp_dim)
ans = identity - ans
ans = np.reshape(ans, shape * 2)
return ans
else:
raise RuntimeError('Need to specific the normalization axis!')
def _single_prop(self, tensor, tau=0.01):
if tensor.axis is None:
self.root = tensor
else:
raise RuntimeError("Cannot propagate on Tensor {}:" "Not a root node!".format(tensor))
def diff(t, y):
"""This function will not change the arrays in tensor network.
"""
tensor.set_array(np.reshape(y, tensor.shape))
ans = np.zeros_like(y)
for n in self.term_visitor():
ans += np.reshape(self._single_diff(tensor, n), -1)
return ans
y0 = np.reshape(tensor.array, -1)
solver = solve_ivp(diff, (self.time, self.time + tau),
y0,
method=self.ode_method,
atol=self.atol,
rtol=self.rtol)
tensor.set_array(np.reshape(solver.y[:, -1], tensor.shape))
return tensor
def gen_extended_rho(self, rho):
"""Get rho_n from rho with the conversion:
rho[n_0, ..., n_(k-1), i, j]
Parameters
----------
rho : np.ndarray
"""
shape = list(rho.shape)
assert len(shape) == 2 and shape[0] == shape[1]
# Let: rho_n[0, i, j] = rho and rho_n[n, i, j] = 0
ext = np.zeros((np.prod(self.n_dims), ))
ext[0] = 1
rho_n = np.reshape(np.tensordot(ext, rho, axes=0),
list(self.n_dims) + shape)
return np.array(rho_n, dtype=DTYPE)
def get_sub_vec(self, i, vec=None):
"""Get C_i mat or A tensor from MCTDH vec.
Parameters
----------
i : int
If 0 <= i < len(self.shape_list) - 1, return C_i mat, else return
A tensor.
"""
if vec is None:
vec = self.vec
size_list = self.size_list
i = i % len(size_list)
start = sum(size_list[:i])
end = sum(size_list[:i + 1])
sub = vec[start:end]
return np.reshape(sub, self.shape_list[i])
def coarse_grain_mps(A, B):
"""Coarse-graining of two-site MPS into one site.::
|st |s |t
-i- C -k- = -i- A -j- B -k-
Parameters
----------
A : (s, i, j) ndarray
A MPS matrix.
B : (t, j, k) ndarray
A MPS matrix.
Returns
-------
C : (st, i, k) ndarray
A MPS matrix.
"""
shape_C = [A.shape[0] * B.shape[0], A.shape[1], B.shape[2]]
C = np.einsum("sij,tjk->stik", A, B)
return np.reshape(C, shape_C)
