diff --git a/.circleci/config.yml b/.circleci/config.yml index a0d8ea36..6788ded6 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -41,8 +41,12 @@ jobs: pip install pytest-xdist pytest-cov export RENO_NUM_THREADS=1 pytest -n 4 --durations=0 --cov=renormalizer renormalizer - pip install primme==3.2.* + - run: + name: Run optional tests + command: | + pip install primme==3.2.* pyscf==2.4.0 pytest --durations=0 renormalizer/mps/tests/test_gs.py::test_multistate --cov=renormalizer --cov-append + pytest --durations=0 renormalizer/mps/tests/test_gs.py::test_pyscf_solver --cov=renormalizer --cov-append - run: name: Run examples command: | diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d7e4c202..6bc33ef8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,13 +19,16 @@ jobs: run: | python -m pip install --upgrade pip pip install --no-cache-dir -r requirements.txt - - name: test scripts + - name: run tests run: | pip install pytest-xdist pytest-cov export RENO_NUM_THREADS=1 pytest -n 4 --durations=0 --cov=renormalizer renormalizer - pip install primme==3.2.* + - name: run optional tests + run: | + pip install primme==3.2.* pyscf==2.4.0 pytest --durations=0 renormalizer/mps/tests/test_gs.py::test_multistate --cov=renormalizer --cov-append + pytest --durations=0 renormalizer/mps/tests/test_gs.py::test_pyscf_solver --cov=renormalizer --cov-append - name: run examples run: | cd example; bash run.sh diff --git a/renormalizer/__init__.py b/renormalizer/__init__.py index 768846a2..034fbb97 100644 --- a/renormalizer/__init__.py +++ b/renormalizer/__init__.py @@ -38,5 +38,5 @@ from renormalizer.model import Model, HolsteinModel, SpinBosonModel, TI1DModel, Op, OpSum from renormalizer.model.basis import BasisSHO, BasisHopsBoson, BasisSineDVR, BasisMultiElectron, \ BasisMultiElectronVac, BasisSimpleElectron, BasisHalfSpin, BasisDummy -from renormalizer.mps import Mpo, Mps, optimize_mps +from renormalizer.mps import Mpo, Mps, optimize_mps, DmrgFCISolver from renormalizer.utils import Quantity diff --git a/renormalizer/model/h_qc.py b/renormalizer/model/h_qc.py index b1eaf781..0dc5c1bc 100644 --- a/renormalizer/model/h_qc.py +++ b/renormalizer/model/h_qc.py @@ -1,5 +1,6 @@ # -*- coding: utf-8 -*- +from functools import partial import itertools import logging @@ -67,6 +68,62 @@ def int_to_h(h, eri): return sh, aseri + +def generate_ladder_operator(norbs): + # construct electronic creation/annihilation operators by Jordan-Wigner transformation + a_ops = [] + a_dag_ops = [] + for j in range(norbs): + # qn for each op will be processed in `process_op` + sigma_z_list = [Op("Z", l) for l in range(j)] + a_ops.append(Op.product(sigma_z_list + [Op("+", j)])) + a_dag_ops.append(Op.product(sigma_z_list + [Op("-", j)])) + + return a_ops, a_dag_ops + + +def simplify_op(old_op: Op, norbs: int, conserve_qn:bool=True): + # helper function to process operators. + # Remove "sigma_z sigma_z". Use {sigma_z, sigma_+} = 0 + # and {sigma_z, sigma_-} = 0 to simplify operators, + # and set quantum number + dof_to_siteidx = dict(zip(range(norbs), range(norbs))) + if conserve_qn: + qn_dict0 = {"+": [-1, 0], "-": [1, 0], "Z": [0, 0]} + qn_dict1 = {"+": [0, -1], "-": [0, 1], "Z": [0, 0]} + else: + qn_dict0 = {"+": 0, "-": 0, "Z": 0} + + old_ops, _ = old_op.split_elementary(dof_to_siteidx) + new_ops = [] + for elem_op in old_ops: + # move all sigma_z to the start of the operator + # and cancel as many as possible + n_sigma_z = elem_op.split_symbol.count("Z") + n_non_sigma_z = 0 + n_permute = 0 + for simple_elem_op in elem_op.split_symbol: + if simple_elem_op != "Z": + n_non_sigma_z += 1 + else: + n_permute += n_non_sigma_z + # remove as many "sigma_z" as possible + new_symbol = [s for s in elem_op.split_symbol if s != "Z"] + if n_sigma_z % 2 == 1: + new_symbol.insert(0, "Z") + # this op is identity, discard it + if not new_symbol: + continue + new_dof_name = elem_op.dofs[0] + if conserve_qn and new_dof_name % 2 == 1: + qn_dict = qn_dict1 + else: + qn_dict = qn_dict0 + new_qn = [qn_dict[s] for s in new_symbol] + new_ops.append(Op(" ".join(new_symbol), new_dof_name, (-1) ** n_permute, new_qn)) + return Op.product(new_ops) + + def qc_model(h1e, h2e, stacked=False, conserve_qn=True): """ Ab initio electronic Hamiltonian in spin-orbitals @@ -91,59 +148,11 @@ def qc_model(h1e, h2e, stacked=False, conserve_qn=True): assert np.all(np.array(h1e.shape) == norbs) assert np.all(np.array(h2e.shape) == norbs) - # construct electronic creation/annihilation operators by Jordan-Wigner transformation - a_ops = [] - a_dag_ops = [] - for j in range(norbs): - # qn for each op will be processed in `process_op` - sigma_z_list = [Op("Z", l) for l in range(j)] - a_ops.append( Op.product(sigma_z_list + [Op("+", j)]) ) - a_dag_ops.append( Op.product(sigma_z_list + [Op("-", j)]) ) - ham_terms = [] - - # helper function to process operators. - # Remove "sigma_z sigma_z". Use {sigma_z, sigma_+} = 0 - # and {sigma_z, sigma_-} = 0 to simplify operators, - # and set quantum number - dof_to_siteidx = dict(zip(range(norbs), range(norbs))) - if conserve_qn: - qn_dict0 = {"+": [-1, 0], "-": [1, 0], "Z": [0, 0]} - qn_dict1 = {"+": [0, -1], "-": [0, 1], "Z": [0, 0]} - else: - qn_dict0 = {"+": 0, "-": 0, "Z": 0} - def process_op(old_op: Op): - old_ops, _ = old_op.split_elementary(dof_to_siteidx) - new_ops = [] - for elem_op in old_ops: - # move all sigma_z to the start of the operator - # and cancel as many as possible - n_sigma_z = elem_op.split_symbol.count("Z") - n_non_sigma_z = 0 - n_permute = 0 - for simple_elem_op in elem_op.split_symbol: - if simple_elem_op != "Z": - n_non_sigma_z += 1 - else: - n_permute += n_non_sigma_z - # remove as many "sigma_z" as possible - new_symbol = [s for s in elem_op.split_symbol if s != "Z"] - if n_sigma_z % 2 == 1: - new_symbol.insert(0, "Z") - # this op is identity, discard it - if not new_symbol: - continue - new_dof_name = elem_op.dofs[0] - if conserve_qn and new_dof_name % 2 == 1: - qn_dict = qn_dict1 - else: - qn_dict = qn_dict0 - new_qn = [qn_dict[s] for s in new_symbol] - new_ops.append(Op(" ".join(new_symbol), new_dof_name, (-1) ** n_permute, new_qn)) - return Op.product(new_ops) - + process_op = partial(simplify_op, norbs=norbs, conserve_qn=conserve_qn) pairs1 = np.argwhere(h1e!=0) pairs2 = np.argwhere(h2e!=0) + a_ops, a_dag_ops = generate_ladder_operator(norbs) if stacked is False: # 1-e terms for p, q in pairs1: @@ -185,9 +194,3 @@ def process_op(old_op: Op): basis.append(b) return basis, ham_terms - - - - - - diff --git a/renormalizer/mps/__init__.py b/renormalizer/mps/__init__.py index 277b108b..42748931 100644 --- a/renormalizer/mps/__init__.py +++ b/renormalizer/mps/__init__.py @@ -3,5 +3,5 @@ from renormalizer.mps.mps import Mps, BraKetPair from renormalizer.mps.mpdm import MpDm from renormalizer.mps.thermalprop import ThermalProp, load_thermal_state -from renormalizer.mps.gs import optimize_mps +from renormalizer.mps.gs import optimize_mps, DmrgFCISolver from renormalizer.mps.tda import TDA diff --git a/renormalizer/mps/gs.py b/renormalizer/mps/gs.py index 6251b76d..143a2384 100644 --- a/renormalizer/mps/gs.py +++ b/renormalizer/mps/gs.py @@ -7,6 +7,9 @@ """ from typing import Tuple, List, Union +from functools import partial +from itertools import product +from collections import deque import logging import numpy as np @@ -14,6 +17,8 @@ import opt_einsum as oe from renormalizer.lib import davidson +from renormalizer.model.h_qc import qc_model, int_to_h, generate_ladder_operator, simplify_op +from renormalizer.model import Model, Op from renormalizer.mps.backend import xp, OE_BACKEND, primme, IMPORT_PRIMME_EXCEPTION from renormalizer.mps.matrix import multi_tensor_contract, tensordot, asnumpy, asxp from renormalizer.mps.hop_expr import hop_expr @@ -22,6 +27,7 @@ from renormalizer.mps.lib import Environ, cvec2cmat from renormalizer.utils import Quantity, CompressConfig, CompressCriteria + logger = logging.getLogger(__name__) @@ -567,4 +573,173 @@ def hop(x): else: assert False logger.debug(f"use {algo}, HC hops: {count}") - return e, sign_fix(c, nroots) \ No newline at end of file + return e, sign_fix(c, nroots) + + +class DmrgFCISolver: + """ + DMRG interface for PySCF FCI/CASCI/CASSCF + """ + def __init__(self): + self.mps: Mps = None + self.nsorb: int = None + self.bond_dimension: int = 32 + self.procedure = None + self.rdm1_mpos = [] + self.rdm2_mpos = [] + + def kernel(self, h1, h2, norb, nelec, ci0=None, ecore=0, **kwargs): + if self.nsorb is None: + self.nsorb = norb * 2 + else: + assert norb * 2 == self.nsorb + + import pyscf + h2 = pyscf.ao2mo.restore(1, h2, norb) + # spatial orbital to spin orbital + h1, h2 = int_to_h(h1, h2) + + basis, ham_terms = qc_model(h1, h2) + + + model = Model(basis, ham_terms) + mpo = Mpo(model) + logger.info(f"mpo_bond_dims:{mpo.bond_dims}") + + if isinstance(nelec, (int, np.integer)): + nelec = [nelec - nelec//2, nelec//2] + + M = self.bond_dimension + mps = Mps.random(model, nelec, M, percent=1.0) + + if self.procedure is None: + mps.optimize_config.procedure = [[M, 0.4], [M, 0.2], [M, 0.1], [M, 0], [M, 0], [M, 0], [M, 0]] + else: + mps.optimize_config.procedure = self.procedure + mps.optimize_config.method = "2site" + energies, mps = optimize_mps(mps.copy(), mpo) + gs_e = min(energies) + ecore + + self.mps = mps + return gs_e, mps + + def _make_rdm1_mpos(self, model: Model, norb: int): + assert norb == self.nsorb // 2 + assert not self.rdm1_mpos + a_ops, a_dag_ops = generate_ladder_operator(self.nsorb) + process_op = partial(simplify_op, norbs=self.nsorb, conserve_qn=True) + for i in range(norb): + for j in range(i + 1): + opaa = process_op(a_dag_ops[2*i] * a_ops[2*j]) + opbb = process_op(a_dag_ops[2*i+1] * a_ops[2*j+1]) + mpo = Mpo(model, terms=[opaa, opbb]) + self.rdm1_mpos.append(mpo) + + def make_rdm1(self, params, norb, nelec): + r""" + Evaluate the spin-traced one-body reduced density matrix (1RDM). + + .. math:: + + \textrm{1RDM}[p,q] = \langle p_{\alpha}^\dagger q_{\alpha} \rangle + + \langle p_{\beta}^\dagger q_{\beta} \rangle + + Returns + ------- + rdm1: np.ndarray + The spin-traced one-body RDM. + + See Also + -------- + make_rdm2: Evaluate the spin-traced two-body reduced density matrix (2RDM). + """ + if params is None: + mps = self.mps + else: + mps = params + + if not self.rdm1_mpos: + self._make_rdm1_mpos(self.mps.model, norb) + + expectations = deque(mps.expectations(self.rdm1_mpos)) + rdm1 = np.zeros([norb] * 2) + for i in range(norb): + for j in range(i + 1): + rdm1[i, j] = rdm1[j, i] = expectations.popleft() + return rdm1 + + def _make_rdm2_mpos(self, model: Model, norb: int): + assert norb == self.nsorb // 2 + assert not self.rdm2_mpos + a_ops, a_dag_ops = generate_ladder_operator(self.nsorb) + process_op = partial(simplify_op, norbs=self.nsorb, conserve_qn=True) + calculated_indices = set() + # a^\dagger_p a^\dagger_q a_r a_s + # possible spins: aaaa, abba, baab, bbbb + for p, q, r, s in product(range(norb), repeat=4): + if (p, q, r, s) in calculated_indices: + continue + opaaaa = process_op(a_dag_ops[2 * p] * a_dag_ops[2 * q] * a_ops[2 * r] * a_ops[2 * s]) + opabba = process_op(a_dag_ops[2 * p] * a_dag_ops[2 * q + 1] * a_ops[2 * r + 1] * a_ops[2 * s]) + opbaab = process_op(a_dag_ops[2 * p + 1] * a_dag_ops[2 * q] * a_ops[2 * r] * a_ops[2 * s + 1]) + opbbbb = process_op(a_dag_ops[2 * p + 1] * a_dag_ops[2 * q + 1] * a_ops[2 * r + 1] * a_ops[2 * s + 1]) + mpo = Mpo(model, terms=[opaaaa, opabba, opbaab, opbbbb]) + self.rdm2_mpos.append(mpo) + indices = [(p, q, r, s), (s, r, q, p), (q, p, s, r), (r, s, p, q)] + for idx in indices: + calculated_indices.add(idx) + + def make_rdm2(self, params, norb, nelec): + r""" + Evaluate the spin-traced two-body reduced density matrix (2RDM). + + .. math:: + + \begin{aligned} + \textrm{2RDM}[p,q,r,s] & = \langle p_{\alpha}^\dagger r_{\alpha}^\dagger + s_{\alpha} q_{\alpha} \rangle + + \langle p_{\beta}^\dagger r_{\alpha}^\dagger s_{\alpha} q_{\beta} \rangle \\ + & \quad + \langle p_{\alpha}^\dagger r_{\beta}^\dagger s_{\beta} q_{\alpha} \rangle + + \langle p_{\beta}^\dagger r_{\beta}^\dagger s_{\beta} q_{\beta} \rangle + \end{aligned} + + Returns + ------- + rdm2: np.ndarray + The spin-traced two-body RDM. + + See Also + -------- + make_rdm1: Evaluate the spin-traced one-body reduced density matrix (1RDM). + """ + if params is None: + mps = self.mps + else: + mps = params + + if not self.rdm2_mpos: + self._make_rdm2_mpos(self.mps.model, norb) + + expectations = deque(mps.expectations(self.rdm2_mpos)) + rdm2 = np.zeros([norb] * 4) + calculated_indices = set() + for p, q, r, s in product(range(norb), repeat=4): + if (p, q, r, s) in calculated_indices: + continue + v = expectations.popleft() + indices = [(p, q, r, s), (s, r, q, p), (q, p, s, r), (r, s, p, q)] + for idx in indices: + calculated_indices.add(idx) + rdm2[idx] = v + # transpose to PySCF notation: rdm2[p,q,r,s] = + rdm2 = rdm2.transpose(0, 3, 1, 2) + return rdm2 + + def make_rdm12(self, params, norb, nelec): + rdm1 = self.make_rdm1(params, norb, nelec) + rdm2 = self.make_rdm2(params, norb, nelec) + return rdm1, rdm2 + + def spin_square(self, params, norb, nelec): + # S^2 = (S+ * S- + S- * S+)/2 + Sz * Sz + raise NotImplementedError diff --git a/renormalizer/mps/mps.py b/renormalizer/mps/mps.py index 1dadb782..32cdcc55 100644 --- a/renormalizer/mps/mps.py +++ b/renormalizer/mps/mps.py @@ -1599,6 +1599,9 @@ def calc_edof_rdm(self) -> np.ndarray: r"""Calculate the reduced density matrix of electronic DoF :math:`\rho_{ij} = \langle \Psi | a_i^\dagger a_j | \Psi \rangle` + + .. note:: + This function is designed for single-electron system. Fermionic commutation relation is not considered. """ diff --git a/renormalizer/mps/tests/test_gs.py b/renormalizer/mps/tests/test_gs.py index 11183e32..429224d1 100644 --- a/renormalizer/mps/tests/test_gs.py +++ b/renormalizer/mps/tests/test_gs.py @@ -8,7 +8,7 @@ from renormalizer.model import Model, h_qc from renormalizer.mps.backend import primme -from renormalizer.mps.gs import construct_mps_mpo, optimize_mps +from renormalizer.mps.gs import construct_mps_mpo, optimize_mps, DmrgFCISolver from renormalizer.mps import Mpo, Mps, StackedMpo from renormalizer.tests.parameter import holstein_model from renormalizer.utils.configs import OFS @@ -156,3 +156,40 @@ def test_stackedmpo(): energies1, _ = optimize_mps(mps.copy(), mpo) energies2, _ = optimize_mps(mps.copy(), StackedMpo([mpo, mpo])) assert np.all(np.abs(np.array(energies2) - np.array(energies1) * 2) < 1e-8) + + +def test_pyscf_solver(): + try: + from pyscf import M, mcscf, fci + except ImportError: + pytest.skip("pyscf not installed") + return # make IDE happy + + # the unused code is for debugging + # the test might not pass based on the unused code + if True: + mol = M(atom=[["N", 0, 0, 0], ["N", 0, 0, 1]], basis="sto3g") + active_space = (10, 8) + active_space = (6, 6) + else: + mol = M(atom=[["Li", 0, 0, 0], ["H", 0, 0, 1]], basis="sto3g") + active_space = (2, 5) + hf = mol.HF() + hf.kernel() + + casci = mcscf.CASCI(hf, active_space[1], active_space[0]) + casci.kernel() + e_ref = casci.e_tot + rdm1_ref = casci.make_rdm1() + rdm2_ref = fci.direct_spin1.make_rdm12(casci.ci, casci.ncas, casci.nelecas)[1] + + casci = mcscf.CASCI(hf, active_space[1], active_space[0]) + casci.fcisolver = DmrgFCISolver() + casci.kernel() + rdm1 = casci.make_rdm1() + rdm2 = casci.fcisolver.make_rdm2(casci.fcisolver.mps, casci.ncas, casci.nelecas) + + # reasonable error since the bond dimension is only 32 + assert casci.e_tot == pytest.approx(e_ref, abs=1e-2) + np.testing.assert_allclose(rdm1, rdm1_ref, atol=1e-2) + np.testing.assert_allclose(rdm2, rdm2_ref, atol=1e-2)