diff --git a/renormalizer/tn/tests/test_tn.py b/renormalizer/tn/tests/test_tn.py index 7af77335..bffc6b63 100644 --- a/renormalizer/tn/tests/test_tn.py +++ b/renormalizer/tn/tests/test_tn.py @@ -233,7 +233,8 @@ def test_rdm_entropy_holstein(): mps_idx1, mps_idx2 = 1, 3 dof1 = model.basis[mps_idx1].dof dof2 = model.basis[mps_idx2].dof - ttns_mutual_info = ttns.calc_2dof_mutual_info(dof1, dof2) + ttns_mutual_infos = ttns.calc_2dof_mutual_info((dof1, dof2)) + ttns_mutual_info = ttns_mutual_infos[(dof1, dof2)] np.testing.assert_allclose(ttns_mutual_info, mps_mutual_info[mps_idx1, mps_idx2], atol=1e-4) @@ -258,7 +259,7 @@ def test_2dof_rdm(basis_tree, dofs): # test 2 dof rdm dof1, dof2 = dofs - rdm1 = ttns.calc_2dof_rdm(dof1, dof2).reshape(4, 4) + rdm1 = ttns.calc_2dof_rdm((dof1, dof2))[(dof1, dof2)].reshape(4, 4) rdm2 = mps.calc_2site_rdm()[(dof1, dof2)].reshape(4, 4) #np.testing.assert_allclose(rdm1, rdm2, atol=1e-10) diff --git a/renormalizer/tn/tree.py b/renormalizer/tn/tree.py index 6a002332..61ed310f 100644 --- a/renormalizer/tn/tree.py +++ b/renormalizer/tn/tree.py @@ -1020,13 +1020,14 @@ def calc_1dof_entropy(self, dof: Union[Any, List[Any]]=None) -> Dict[Any, float] rdm = self.calc_1dof_rdm(dof) return {key: calc_vn_entropy_dm(dm) for key, dm in rdm.items()} - def calc_2site_rdm(self, idx1, idx2) -> np.ndarray: + def calc_2site_rdm(self, idxs: Union[Tuple[int, int], List[Tuple[int, int]]]=None) -> Dict[Tuple[int, int], np.ndarray]: r""" Calculate 2-site reduced density matrix :math:`\rho_{ij} = \textrm{Tr}_{k \neq i, k \neq j} | \Psi \rangle \langle \Psi |`. Parameters ---------- + idxs: list(tuple), optional idx1: int site index (in terms of ``self.node_list``) of the first site. idx2: int @@ -1039,106 +1040,157 @@ def calc_2site_rdm(self, idx1, idx2) -> np.ndarray: """ ttno_dummy = TTNO.dummy(self.basis) ttne = TTNEnviron(self, ttno_dummy) - path = self.find_path(self.node_list[idx1], self.node_list[idx2]) - assert path[0] is self.node_list[idx1] - assert path[-1] is self.node_list[idx2] - args = [] - # put the nodes for RDM in the arguments - for snode in [path[0], path[-1]]: - args.append(snode.tensor.conj()) - args.append(self.get_node_indices(snode, conj=True)) - args.append(snode.tensor) - args.append(self.get_node_indices(snode)) + if isinstance(idxs, tuple): + idxs = [idxs] + else: + assert isinstance(idxs, list) - # put the nodes in the path in the arguments - for snode in path[1:-1]: - args.append(snode.tensor.conj()) - args.append(self.get_node_indices(snode, conj=True)) + rdm = {} + for idx_pair in idxs: + idx1 = idx_pair[0] + idx2 = idx_pair[1] - args.append(snode.tensor) - # set ttno to ttno_dummy so that the physical indices are contracted directly - args.append(self.get_node_indices(snode, ttno=ttno_dummy)) + path = self.find_path(self.node_list[idx1], self.node_list[idx2]) + assert path[0] is self.node_list[idx1] + assert path[-1] is self.node_list[idx2] + args = [] + # put the nodes for RDM in the arguments + for snode in [path[0], path[-1]]: + args.append(snode.tensor.conj()) + args.append(self.get_node_indices(snode, conj=True)) + + args.append(snode.tensor) + args.append(self.get_node_indices(snode)) + + # put the nodes in the path in the arguments + for snode in path[1:-1]: + args.append(snode.tensor.conj()) + args.append(self.get_node_indices(snode, conj=True)) + + args.append(snode.tensor) + # set ttno to ttno_dummy so that the physical indices are contracted directly + args.append(self.get_node_indices(snode, ttno=ttno_dummy)) + + # put all environment tensors in the arguments + for i, node in enumerate(path): + # skip some of the environments because they are included in the path + + if i == 0: + neighbour_nodes = [path[i+1]] + elif i == len(path)-1: + neighbour_nodes = [path[i-1]] + else: + neighbour_nodes = [path[i-1], path[i+1]] + + skip_child_idx_list: List[int] = [] + skip_parent: bool = False + for neighbour_node in neighbour_nodes: + if neighbour_node.parent is node: + skip_child_idx_list.append(neighbour_node.idx_as_child) + elif node.parent is neighbour_node: + skip_parent = True + + enode = ttne.node_list[self.node_idx[node]] + # put all children environments in the arguments + for j, child_tensor in enumerate(enode.environ_children): + if j in skip_child_idx_list: + continue + indices = ttne.get_child_indices(enode, j, self, ttno_dummy) + args.extend([child_tensor, indices]) + # put the parent environment in the arguments + if not skip_parent: + args.append(enode.environ_parent) + args.append(ttne.get_parent_indices(enode, self, ttno_dummy)) - # put all environment tensors in the arguments - for i, node in enumerate(path): - # skip some of the environments because they are included in the path + # the indices for the output tensor + indices_ket = [] + indices_bra = [] + for snode in [path[0], path[-1]]: + for dofs in self.tn2dofs[snode]: + indices_ket.append(("down", str(dofs))) + indices_bra.append(("up", str(dofs))) + args.append(indices_ket + indices_bra) + res = oe.contract(*asxp_oe_args(args)) + rdm[idx_pair] = res + # perform the contraction + return rdm - if i == 0: - neighbour_nodes = [path[i+1]] - elif i == len(path)-1: - neighbour_nodes = [path[i-1]] - else: - neighbour_nodes = [path[i-1], path[i+1]] + def calc_2site_entropy(self, idxs: Union[Tuple[int, int], List[Tuple[int, int]]]) -> Dict[tuple, float]: - skip_child_idx_list: List[int] = [] - skip_parent: bool = False - for neighbour_node in neighbour_nodes: - if neighbour_node.parent is node: - skip_child_idx_list.append(neighbour_node.idx_as_child) - elif node.parent is neighbour_node: - skip_parent = True + if isinstance(idxs, tuple): + idxs = [idxs] + else: + assert isinstance(idxs, list) - enode = ttne.node_list[self.node_idx[node]] - # put all children environments in the arguments - for j, child_tensor in enumerate(enode.environ_children): - if j in skip_child_idx_list: - continue - indices = ttne.get_child_indices(enode, j, self, ttno_dummy) - args.extend([child_tensor, indices]) - # put the parent environment in the arguments - if not skip_parent: - args.append(enode.environ_parent) - args.append(ttne.get_parent_indices(enode, self, ttno_dummy)) - - # the indices for the output tensor - indices_ket = [] - indices_bra = [] - for snode in [path[0], path[-1]]: - for dofs in self.tn2dofs[snode]: - indices_ket.append(("down", str(dofs))) - indices_bra.append(("up", str(dofs))) - args.append(indices_ket + indices_bra) + rdm = self.calc_2site_rdm(idxs) + entropy = {key: calc_vn_entropy_dm(dm) for key, dm in rdm.items()} + return entropy - # perform the contraction - return oe.contract(*asxp_oe_args(args)) + def calc_2dof_rdm(self, dofs: Union[Tuple[Any, Any], List[Tuple[Any, Any]]]) -> Dict[Tuple[Any, Any], np.ndarray]: - def calc_2site_entropy(self, idx1, idx2) -> float: - rdm = self.calc_2site_rdm(idx1, idx2) - return calc_vn_entropy_dm(rdm) - - def calc_2dof_rdm(self, dof1, dof2) -> np.ndarray: - site_idx1 = self.basis.dof2idx[dof1] - site_idx2 = self.basis.dof2idx[dof2] - if site_idx1 == site_idx2: - # two dofs on the same site - rdm = self.calc_1site_rdm(site_idx1)[site_idx1] - basis_node: TreeNodeBasis = self.basis.node_list[site_idx1] - n_sets = basis_node.n_sets - basis_idx1 = basis_node.basis_sets.index(self.basis.dof2basis[dof1]) - basis_idx2 = basis_node.basis_sets.index(self.basis.dof2basis[dof2]) - assert basis_idx1 != basis_idx2 + if isinstance(dofs, tuple): + dofs = [dofs] else: - # two dofs on different sites - rdm = self.calc_2site_rdm(site_idx1, site_idx2) - basis_node1: TreeNodeBasis = self.basis.node_list[site_idx1] - basis_node2: TreeNodeBasis = self.basis.node_list[site_idx2] - n_sets = basis_node1.n_sets + basis_node2.n_sets - basis_idx1 = basis_node1.basis_sets.index(self.basis.dof2basis[dof1]) - basis_idx2 = basis_node1.n_sets + basis_node2.basis_sets.index(self.basis.dof2basis[dof2]) - - indices = [(0, i) for i in range(n_sets)] * 2 - indices[basis_idx1] = (1, 0) - indices[basis_idx2] = (1, 1) - indices[n_sets + basis_idx1] = (1, 2) - indices[n_sets + basis_idx2] = (1, 3) - return oe.contract(rdm, indices, [(1, i) for i in range(4)]) - - def calc_2dof_entropy(self, dof1, dof2) -> float: - rdm = self.calc_2dof_rdm(dof1, dof2) - return calc_vn_entropy_dm(rdm) - - def calc_2dof_mutual_info(self, dof1, dof2) -> float: + assert isinstance(dofs, list) + + rdm_ = {} + rdm_1site_idx_lst = [] + rdm_2site_idx_lst = [] + for dof_pair in dofs: + dof1 = dof_pair[0] + dof2 = dof_pair[1] + + site_idx1 = self.basis.dof2idx[dof1] + site_idx2 = self.basis.dof2idx[dof2] + if site_idx1 == site_idx2: + rdm_1site_idx_lst.append(site_idx1) + rdm_1site_idx_lst.append(site_idx2) + else: + rdm_2site_idx_lst.append((site_idx1, site_idx2)) + if len(rdm_1site_idx_lst) > 0: + rdm_1sites = self.calc_1site_rdm(rdm_1site_idx_lst) + if len(rdm_2site_idx_lst) > 0: + rdm_2sites = self.calc_2site_rdm(rdm_2site_idx_lst) + + for dof_pair in dofs: + dof1 = dof_pair[0] + dof2 = dof_pair[1] + + site_idx1 = self.basis.dof2idx[dof1] + site_idx2 = self.basis.dof2idx[dof2] + if site_idx1 == site_idx2: + # two dofs on the same site + rdm = rdm_1sites[site_idx1] + basis_node: TreeNodeBasis = self.basis.node_list[site_idx1] + n_sets = basis_node.n_sets + basis_idx1 = basis_node.basis_sets.index(self.basis.dof2basis[dof1]) + basis_idx2 = basis_node.basis_sets.index(self.basis.dof2basis[dof2]) + assert basis_idx1 != basis_idx2 + else: + # two dofs on different sites + rdm = rdm_2sites[(site_idx1, site_idx2)] + basis_node1: TreeNodeBasis = self.basis.node_list[site_idx1] + basis_node2: TreeNodeBasis = self.basis.node_list[site_idx2] + n_sets = basis_node1.n_sets + basis_node2.n_sets + basis_idx1 = basis_node1.basis_sets.index(self.basis.dof2basis[dof1]) + basis_idx2 = basis_node1.n_sets + basis_node2.basis_sets.index(self.basis.dof2basis[dof2]) + + indices = [(0, i) for i in range(n_sets)] * 2 + indices[basis_idx1] = (1, 0) + indices[basis_idx2] = (1, 1) + indices[n_sets + basis_idx1] = (1, 2) + indices[n_sets + basis_idx2] = (1, 3) + res = oe.contract(rdm, indices, [(1, i) for i in range(4)]) + rdm_[dof_pair] = res + return rdm_ + + def calc_2dof_entropy(self, dofs: Union[Tuple[Any, Any], List[Tuple[Any, Any]]]) -> Dict[Tuple[Any, Any], float]: + rdm = self.calc_2dof_rdm(dofs) + entropy = {key: calc_vn_entropy_dm(dm) for key, dm in rdm.items()} + return entropy + + def calc_2dof_mutual_info(self, dofs: Union[Tuple[Any, Any], List[Tuple[Any, Any]]]) -> Dict[Tuple[Any, Any], float]: r""" Calculate mutual information between two DOFs. @@ -1150,10 +1202,26 @@ def calc_2dof_mutual_info(self, dof1, dof2) -> float: ------- mutual_info : float mutual information between the two DOFs + mutual_infos : Dict[Any, float] """ - entropy_1site = self.calc_1dof_entropy([dof1, dof2]) - entropy_2site = self.calc_2dof_entropy(dof1, dof2) - return (entropy_1site[dof1] + entropy_1site[dof2] - entropy_2site) / 2 + if isinstance(dofs, tuple): + dofs = [dofs] + else: + assert isinstance(dofs, list) + + mutual_infos = {} + dofs_lst = [] + for dof_pair in dofs: + dofs_lst.append(dof_pair[0]) + dofs_lst.append(dof_pair[1]) + entropy_1site = self.calc_1dof_entropy(dofs_lst) + entropy_2site = self.calc_2dof_entropy(dofs) + for dof_pair in dofs: + dof1 = dof_pair[0] + dof2 = dof_pair[1] + mutual_info = (entropy_1site[dof1] + entropy_1site[dof2] - entropy_2site[dof_pair]) / 2 + mutual_infos[dof_pair] = mutual_info + return mutual_infos def calc_bond_entropy(self) -> np.ndarray: r""" diff --git a/setup.py b/setup.py index 4293ec50..53556cc9 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,8 @@ "scipy", "h5py", "opt_einsum", - "sympy" + "sympy", + "print-tree2" ] setuptools.setup(