add print for conventional mpo

block-hczhai · Oct 12, 2022 · b65c9e0 · b65c9e0
1 parent 2247339
commit b65c9e0
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/pyblock2/driver/core.py b/pyblock2/driver/core.py
@@ -478,10 +478,12 @@ def su2_to_sgf(self):
                 [self.orb_sym[i // 2] for i in range(self.n_sites)]
             )
 
-    def get_conventional_qc_mpo(self, fcidump, algo_type=None):
+    def get_conventional_qc_mpo(self, fcidump, algo_type=None, iprint=1):
         """This method cannot take care of parallelization!"""
         bw = self.bw
         hamil = bw.bs.HamiltonianQC(self.vacuum, self.n_sites, self.orb_sym, fcidump)
+        import time
+        tt = time.perf_counter()
         if algo_type is not None and MPOAlgorithmTypes.NC in algo_type:
             mpo = bw.bs.MPOQC(hamil, bw.b.QCTypes.NC, "HQC")
         elif algo_type is not None and MPOAlgorithmTypes.CN in algo_type:
@@ -497,6 +499,12 @@ def get_conventional_qc_mpo(self, fcidump, algo_type=None):
         use_r_intermediates = (
             algo_type is None or MPOAlgorithmTypes.NoRIntermed not in algo_type
         )
+
+        if iprint:
+            nnz, sz, bdim = mpo.get_summary()
+            print('Ttotal = %10.3f MPO method = %s bond dimension = %7d NNZ = %12d SIZE = %12d SPT = %6.4f'
+                % (time.perf_counter() - tt, algo_type.name, bdim, nnz, sz, (1.0 * sz - nnz) / sz))
+
         mpo = bw.bs.SimplifiedMPO(
             mpo,
             rule,
@@ -694,7 +702,7 @@ def get_qc_mpo(
 
         if algo_type is not None and MPOAlgorithmTypes.Conventional in algo_type:
             fd = self.write_fcidump(h1e, g2e, ecore=ecore)
-            return self.get_conventional_qc_mpo(fd, algo_type=algo_type)
+            return self.get_conventional_qc_mpo(fd, algo_type=algo_type, iprint=iprint)
 
         # build Hamiltonian expression
         b = self.expr_builder()

diff --git a/src/dmrg/mpo.hpp b/src/dmrg/mpo.hpp
@@ -282,6 +282,18 @@ template <typename S, typename FL> struct MPO {
             if (tensors[m] != nullptr)
                 tensors[m]->deallocate();
     }
+    // nnz, size, bond dimension
+    tuple<size_t, size_t, int> get_summary() const {
+        size_t lnnz = 0, lsz = 0, rnnz = 0, rsz = 0;
+        int bdim = 0;
+        for (int ii = 0; ii < n_sites; ii++) {
+            shared_ptr<OperatorTensor<S, FL>> opt = tensors[ii];
+            lnnz += opt->lmat->nnz(), rnnz += opt->rmat->nnz();
+            lsz += opt->lmat->size(), rsz += opt->rmat->size();
+            bdim = max(max(bdim, (int)opt->lmat->n), (int)opt->lmat->m);
+        }
+        return make_tuple(max(lnnz, rnnz), max(lsz, rsz), bdim);
+    }
     string get_filename(int i, int ixtag, const string &dir = "") const {
         const static string xtag[] = {"TENSOR", "LEFT.OP", "RIGHT.OP",
                                       "MIDDLE.OP"};

diff --git a/src/pybind/pybind_dmrg.hpp b/src/pybind/pybind_dmrg.hpp
@@ -1576,6 +1576,7 @@ template <typename S, typename FL> void bind_fl_mpo(py::module &m) {
         .def_readwrite("tag", &MPO<S, FL>::tag)
         .def_readwrite("tread", &MPO<S, FL>::tread)
         .def_readwrite("twrite", &MPO<S, FL>::twrite)
+        .def("get_summary", &MPO<S, FL>::get_summary)
         .def("get_filename", &MPO<S, FL>::get_filename)
         .def("load_left_operators", &MPO<S, FL>::load_left_operators)
         .def("save_left_operators", &MPO<S, FL>::save_left_operators)