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)