From e576340d165f9736245cbf7a11caffe103410cec Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Wed, 3 Jul 2024 16:52:56 +0800 Subject: [PATCH] [tet.py, tetku.py] Add support for trim style hamiltonian. This kind of hamiltonian is useful especially for quantum chemistry models, which is set by attribute["quantum_chemistry_term"]. --- .../sampling_neural_state/observer.py | 168 +++++++++++++++++- .../openfermion_quantum_chemistry/trim.py | 164 +++++++++++++++++ 2 files changed, 329 insertions(+), 3 deletions(-) create mode 100644 tetraku/tetraku/models/openfermion_quantum_chemistry/trim.py diff --git a/tetragono/tetragono/sampling_neural_state/observer.py b/tetragono/tetragono/sampling_neural_state/observer.py index cd497b52..776425b0 100644 --- a/tetragono/tetragono/sampling_neural_state/observer.py +++ b/tetragono/tetragono/sampling_neural_state/observer.py @@ -19,7 +19,7 @@ import numpy as np import torch from ..utility import allreduce_buffer, allreduce_number, show, showln -from .state import Configuration, index_tensor_element +from .state import Configuration, index_tensor_element, torch_grad class Observer(): @@ -284,6 +284,167 @@ def __call__(self, configurations, amplitudes, weights, multiplicities): } for _ in range(batch_size)] whole_result = [{name: 0.0 for name in self._observer} for _ in range(batch_size)] + quantum_chemistry_term = [0 for _ in range(batch_size)] + if "quantum_chemistry_term" in self.owner.attribute: + with torch_grad(False): + # Precalculate hopping + # batch_size * site * site + # we only calculate site a and site b exchange here + # since for quantum chemistry model hopping is just exchange if one has electron and the other is empty + # configurations_cpu : batch * L1 * L2 * 1 + # configurations_precalc : (site * site-1 / 2) * batch * L1 * L2 * 1 + configurations_precalc = configurations.unsqueeze(0).repeat( + [self.owner.site_number * (self.owner.site_number - 1) // 2 + 1, 1, 1, 1, 1]) + for al1, al2 in self.owner.sites(): + for bl1, bl2 in self.owner.sites(): + ai = al1 * self.owner.L2 + al2 + bi = bl1 * self.owner.L2 + bl2 + if ai < bi: + i = (2 * self.owner.site_number - ai - 1) * ai // 2 + (bi - ai - 1) + 1 + a = configurations_precalc[i, :, al1, al2, :].clone() + b = configurations_precalc[i, :, bl1, bl2, :] + configurations_precalc[i, :, al1, al2, :] = b + configurations_precalc[i, :, bl1, bl2, :] = a + amplitudes_precalc = self.owner(configurations_precalc.reshape([-1, self.owner.L1, self.owner.L2, 1]), + enable_grad=False).reshape([ + self.owner.site_number * (self.owner.site_number - 1) // 2 + 1, + batch_size + ]) + amplitudes_precalc_cpu = amplitudes_precalc.cpu() + gradients_precalc_norm = torch.zeros_like(amplitudes_precalc_cpu) + gradients_precalc_conj = torch.zeros_like(amplitudes_precalc_cpu) + + for batch_index in range(batch_size): + configuration_cpu = configurations_cpu[batch_index] + amplitude = amplitudes_cpu[batch_index] + weight = weights_cpu[batch_index] + multiplicity = multiplicities_cpu[batch_index] + reweight = (amplitude.abs()**2 / weight).item() # / p(s) + energy = 0 + for positions, observer, exists in self.owner.attribute["quantum_chemistry_term"][0]: + if not all(configuration_cpu[exist] == 1 for exist in exists): + continue + body = 2 + element_pool = index_tensor_element(observer) + positions_configuration = tuple(configuration_cpu[l1l2o].item() for l1l2o in positions) + if positions_configuration not in element_pool: + continue + sub_energy = 0 + for positions_configuration_s, item in element_pool[positions_configuration].items(): + configuration_cpu_s = configuration_cpu.clone() + for l1l2o, value in zip(positions, positions_configuration_s): + configuration_cpu_s[l1l2o] = value + # we only have hopping item, since position configuration found in element pool, it must be a swap item. + ((al1, al2, _), (bl1, bl2, _)) = positions + ai = al1 * self.owner.L2 + al2 + bi = bl1 * self.owner.L2 + bl2 + ai, bi = sorted([ai, bi]) + i = (2 * self.owner.site_number - ai - 1) * ai // 2 + (bi - ai - 1) + 1 + value = item * self._fermi_sign( + configuration_cpu, configuration_cpu_s, + positions) * amplitudes_precalc_cpu[i, batch_index].conj() / amplitude.conj() + sub_energy = sub_energy + value + if self._enable_gradient: + gradients_precalc_conj[i, batch_index] += multiplicity * reweight * value / 2 + energy = energy + sub_energy + if self._enable_gradient: + gradients_precalc_norm[0, batch_index] += multiplicity * reweight * sub_energy / 2 + for positions_1, observer_1, positions_2, observer_2, exists in self.owner.attribute[ + "quantum_chemistry_term"][1]: + if not all(configuration_cpu[exist] == 1 for exist in exists): + continue + body_1 = body_2 = 2 + element_pool_1 = index_tensor_element(observer_1) + element_pool_2 = index_tensor_element(observer_2) + positions_configuration_1 = tuple(configuration_cpu[l1l2o].item() for l1l2o in positions_1) + positions_configuration_2 = tuple(configuration_cpu[l1l2o].item() for l1l2o in positions_2) + if positions_configuration_1 not in element_pool_1: + continue + if positions_configuration_2 not in element_pool_2: + continue + sub_energy_1 = 0 + for positions_configuration_s_1, item_1 in element_pool_1[positions_configuration_1].items(): + configuration_cpu_s_1 = configuration_cpu.clone() + for l1l2o, value in zip(positions_1, positions_configuration_s_1): + configuration_cpu_s_1[l1l2o] = value + # we only have hopping item, since position configuration found in element pool, it must be a swap item. + ((al1, al2, _), (bl1, bl2, _)) = positions_1 + ai = al1 * self.owner.L2 + al2 + bi = bl1 * self.owner.L2 + bl2 + ai, bi = sorted([ai, bi]) + i = (2 * self.owner.site_number - ai - 1) * ai // 2 + (bi - ai - 1) + 1 + value = item_1 * self._fermi_sign( + configuration_cpu, configuration_cpu_s_1, + positions_1) * amplitudes_precalc_cpu[i, batch_index].conj() / amplitude.conj() + sub_energy_1 = sub_energy_1 + value + sub_energy_2 = 0 + for positions_configuration_s_2, item_2 in element_pool_2[positions_configuration_2].items(): + configuration_cpu_s_2 = configuration_cpu.clone() + for l1l2o, value in zip(positions_2, positions_configuration_s_2): + configuration_cpu_s_2[l1l2o] = value + # we only have hopping item, since position configuration found in element pool, it must be a swap item. + ((al1, al2, _), (bl1, bl2, _)) = positions_2 + ai = al1 * self.owner.L2 + al2 + bi = bl1 * self.owner.L2 + bl2 + ai, bi = sorted([ai, bi]) + i = (2 * self.owner.site_number - ai - 1) * ai // 2 + (bi - ai - 1) + 1 + value = item_2 * self._fermi_sign( + configuration_cpu, configuration_cpu_s_2, + positions_2) * amplitudes_precalc_cpu[i, batch_index].conj() / amplitude.conj() + sub_energy_2 = sub_energy_2 + value + energy = energy + sub_energy_1 * sub_energy_2.conj() + if self._enable_gradient: + for positions_configuration_s_1, item_1 in element_pool_1[positions_configuration_1].items( + ): + configuration_cpu_s_1 = configuration_cpu.clone() + for l1l2o, value in zip(positions_1, positions_configuration_s_1): + configuration_cpu_s_1[l1l2o] = value + # we only have hopping item, since position configuration found in element pool, it must be a swap item. + ((al1, al2, _), (bl1, bl2, _)) = positions_1 + ai = al1 * self.owner.L2 + al2 + bi = bl1 * self.owner.L2 + bl2 + ai, bi = sorted([ai, bi]) + i = (2 * self.owner.site_number - ai - 1) * ai // 2 + (bi - ai - 1) + 1 + value = item_1 * self._fermi_sign( + configuration_cpu, configuration_cpu_s_1, + positions_1) * amplitudes_precalc_cpu[i, batch_index].conj() / amplitude.conj() + gradients_precalc_conj[ + i, batch_index] += multiplicity * reweight * value * sub_energy_2.conj() / 2 + + for positions_configuration_s_2, item_2 in element_pool_2[positions_configuration_2].items( + ): + configuration_cpu_s_2 = configuration_cpu.clone() + for l1l2o, value in zip(positions_2, positions_configuration_s_2): + configuration_cpu_s_2[l1l2o] = value + # we only have hopping item, since position configuration found in element pool, it must be a swap item. + ((al1, al2, _), (bl1, bl2, _)) = positions_2 + ai = al1 * self.owner.L2 + al2 + bi = bl1 * self.owner.L2 + bl2 + ai, bi = sorted([ai, bi]) + i = (2 * self.owner.site_number - ai - 1) * ai // 2 + (bi - ai - 1) + 1 + value = item_2 * self._fermi_sign( + configuration_cpu, configuration_cpu_s_2, + positions_2) * amplitudes_precalc_cpu[i, batch_index].conj() / amplitude.conj() + gradients_precalc_norm[ + i, batch_index] += multiplicity * reweight * sub_energy_1 * value.conj() / 2 + + whole_result[batch_index]["energy"] += complex(energy) + quantum_chemistry_term[batch_index] = complex(energy) + if self._enable_gradient: + edelta = 0 + for i in range(self.owner.site_number * (self.owner.site_number - 1) // 2 + 1): + for j in range(batch_size): + if gradients_precalc_norm[i, j] != 0 or gradients_precalc_conj[i, j] != 0: + configuration = configurations_precalc[i, j] + amplitude = self.owner(configuration.unsqueeze(0), enable_grad=True) + grad = self.owner.holes(amplitude) + edelta = edelta + gradients_precalc_norm[i, j] * grad + edelta = edelta + gradients_precalc_conj[i, j] * grad.conj() + if self._EDelta is None: + self._EDelta = edelta + else: + self._EDelta += edelta + for batch_index in range(batch_size): configuration_cpu = configurations_cpu[batch_index] amplitude = amplitudes_cpu[batch_index] @@ -340,7 +501,8 @@ def __call__(self, configurations, amplitudes, weights, multiplicities): self._total_imaginary_energy_reweight += multiplicity * whole_result[batch_index][ name].imag * reweight if name == "energy" and self._enable_gradient: - Es = whole_result[batch_index][name] + Es = whole_result[batch_index][name] - quantum_chemistry_term[batch_index] + # The EDelta contributed by quantum chemistry term has been collected before. if self.owner.Tensor.is_real: Es = Es.real @@ -571,7 +733,7 @@ def DT(v): def A(v): return DT(D(v)) - b = DT(Energy) + b = ((self._EDelta / self._total_weight) - energy * (self._Delta / self._total_weight)) b_square = torch.dot(torch.conj(b), b).real x = torch.zeros_like(b) diff --git a/tetraku/tetraku/models/openfermion_quantum_chemistry/trim.py b/tetraku/tetraku/models/openfermion_quantum_chemistry/trim.py new file mode 100644 index 00000000..3a2030ea --- /dev/null +++ b/tetraku/tetraku/models/openfermion_quantum_chemistry/trim.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# Copyright (C) 2024 Hao Zhang +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +# + +import numpy as np +import TAT +import tetragono as tet +from tetragono.common_tensor.tensor_toolkit import rename_io, kronecker_product + + +def abstract_state(L1, L2, file_name, T=False, is_complex=False): + """ + Create a quantum chemistry state from Hamiltonian in openfermion format. + Every orbit will be put into a site in lattice L1 * L2. + + Parameters + ---------- + L1, L2 : int + The lattice size. + file_name : str + A file containing Hamiltonian in openfermion format. + T : bool + Whether the total electron number is odd. + """ + if is_complex: + Tensor = TAT.FermiZ2.C.Tensor + else: + Tensor = TAT.FermiZ2.D.Tensor + EF = ([(False, 1), (True, 1)], False) + ET = ([(False, 1), (True, 1)], True) + + CP = Tensor(["O0", "I0", "T"], [EF, ET, ([(True, 1)], False)]).zero() + CP[{"O0": (True, 0), "I0": (False, 0), "T": (True, 0)}] = 1 + CM = Tensor(["O0", "I0", "T"], [EF, ET, ([(True, 1)], True)]).zero() + CM[{"O0": (False, 0), "I0": (True, 0), "T": (True, 0)}] = 1 + I = Tensor(["O0", "I0"], [EF, ET]).identity({("I0", "O0")}) + C0daggerC1 = rename_io(CP, [0]).contract(rename_io(CM, [1]), {("T", "T")}) + C0daggerC1daggerC2C3 = kronecker_product(rename_io(C0daggerC1, [1, 2]), rename_io(C0daggerC1, [0, 3])) + C0daggerC1_conjugated = C0daggerC1.conjugate().edge_rename({"I0": "O0", "I1": "O1", "O0": "I0", "O1": "I1"}) + + state = tet.AbstractState(Tensor, L1, L2) + state.physics_edges[...] = EF + state.total_symmetry = T + + m = lambda x: (x // L2, x % L2, 0) + + data = np.load(file_name, allow_pickle=True).item() + single, double = state.attribute["quantum_chemistry_term"] = [[], []] + for count, [term, coefficient] in enumerate(data.terms.items()): + tet.show(f"reading {count}/{len(data.terms)}: {term}") + match term: + case ((site_0, 1), (site_1, 0)): # c^ c + if site_0 != site_1: + single.append([ + (m(site_0), m(site_1)), + C0daggerC1 * coefficient, + [], + ]) + else: + # single site N + state.hamiltonians[m(site_0), m(site_1)] = C0daggerC1 * coefficient + case ((site_0, 1), (site_1, 1), (site_2, 0), (site_3, 0)): # c^ c^ c c + if len({site_0, site_1, site_2, site_3}) == 4: + # c^0 c^1 c2 c3 + # + c^0 c3 c^1 c2 + # - c^0 c2 c^1 c3 + spins = [site_0 % 2, site_1 % 2, site_2 % 2, site_3 % 2] + if sum(spins) in [2]: + if spins[0] == spins[3]: + double.append([ + (m(site_0), m(site_3)), + (+coefficient) * C0daggerC1, + (m(site_1), m(site_2)), + C0daggerC1_conjugated, + [], + ]) + else: + double.append([ + (m(site_0), m(site_2)), + (-coefficient) * C0daggerC1, + (m(site_1), m(site_3)), + C0daggerC1_conjugated, + [], + ]) + elif sum(spins) in [0, 4]: + double.append([ + (m(site_0), m(site_3)), + (+coefficient / 2) * C0daggerC1, + (m(site_1), m(site_2)), + C0daggerC1_conjugated, + [], + ]) + double.append([ + (m(site_0), m(site_2)), + (-coefficient / 2) * C0daggerC1, + (m(site_1), m(site_3)), + C0daggerC1_conjugated, + [], + ]) + else: + tet.showln(f"unrecognized term: {other}") + elif len({site_0, site_1, site_2, site_3}) == 3: + if site_1 == site_2: + single.append([(m(site_0), m(site_3)), +coefficient * C0daggerC1, [m(site_1)]]) + elif site_0 == site_3: + single.append([(m(site_1), m(site_2)), +coefficient * C0daggerC1, [m(site_0)]]) + elif site_0 == site_2: + single.append([(m(site_1), m(site_3)), -coefficient * C0daggerC1, [m(site_0)]]) + elif site_1 == site_3: + single.append([(m(site_0), m(site_2)), -coefficient * C0daggerC1, [m(site_1)]]) + else: + tet.showln(f"unrecognized term: {other}") + else: + # double site NN + if {site_0, site_1} != {site_2, site_3}: + tet.showln(f"unrecognized term: {other}") + state.hamiltonians[m(site_0), m(site_1), m(site_2), m(site_3)] = C0daggerC1daggerC2C3 * coefficient + case (): + state.hamiltonians[((0, 0, 0),)] = coefficient * I + case other: + tet.showln(f"unrecognized term: {other}") + raise NotImplementedError() + tet.showln(f"read done, total {len(data.terms)}") + state.hamiltonians.trace_repeated().sort_points().check_hermite(1e-15) + tet.showln(f"clean terms, total {len(state.hamiltonians)}") + return state + + +def abstract_lattice(L1, L2, D, file_name, T=False, is_complex=False): + """ + Create a quantum chemistry lattice from Hamiltonian in openfermion format. + Every orbit will be put into a site in lattice L1 * L2. + + Parameters + ---------- + L1, L2 : int + The lattice size. + D : int + The dimension cut for PEPS. + file_name : str + A file containing Hamiltonian in openfermion format. + T : bool + Whether the total electron number is odd. + """ + state = tet.AbstractLattice(abstract_state(L1, L2, file_name, T=T, is_complex=is_complex)) + D1 = D // 2 + D2 = D - D1 + state.virtual_bond["R"] = [(False, D1), (True, D2)] + state.virtual_bond["D"] = [(False, D1), (True, D2)] + return state