Add DMRG pyscf solver interface

.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: |

.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

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

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
 
-
-
-
-
-
-

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

renormalizer/mps/gs.py

@@ -7,13 +7,18 @@
 """
 
 from typing import Tuple, List, Union
+from functools import partial
+from itertools import product
+from collections import deque
 import logging
 
 import numpy as np
 import scipy
 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)
+    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] = <p^+ r^+ s q>
+        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