fix disjoint svd test

.github/workflows/build.yml

@@ -113,14 +113,14 @@ jobs:
  run: |
  python -m pip install pytest 'pyscf==2.1.0'
  export PYTHONPATH=$(pwd)/build:$(pwd):${PYTHONPATH}
- py.test pyblock2/unit_test/*.py
+ py.test -s pyblock2/unit_test/*.py
 
  - name: run test (serial-pytest, macos)
  if: matrix.parallel == 'serial-pytest' && matrix.os == 'macos-latest'
  run: |
  python -m pip install pytest 'pyscf==2.0.1'
  export PYTHONPATH=$(pwd)/build:$(pwd):${PYTHONPATH}
- py.test pyblock2/unit_test/*.py
+ py.test -s pyblock2/unit_test/*.py
 
  - name: build test (mpi)
  if: matrix.parallel == 'mpi'

pyblock2/driver/core.py

@@ -530,6 +530,8 @@ def get_qc_mpo(
  csvd_eps=1e-10,
  csvd_max_iter=1000,
  disjoint_levels=None,
+ disjoint_all_blocks=False,
+ disjoint_multiplier=1.0,
  iprint=1,
  ):
  import numpy as np
@@ -779,6 +781,8 @@ def get_qc_mpo(
  csvd_eps=csvd_eps,
  csvd_max_iter=csvd_max_iter,
  disjoint_levels=disjoint_levels,
+ disjoint_all_blocks=disjoint_all_blocks,
+ disjoint_multiplier=disjoint_multiplier,
  )
 
  def get_mpo(
@@ -794,6 +798,8 @@ def get_mpo(
  csvd_eps=1e-10,
  csvd_max_iter=1000,
  disjoint_levels=None,
+ disjoint_all_blocks=False,
+ disjoint_multiplier=1.0,
  ):
  bw = self.bw
  if left_vacuum is None:
@@ -810,6 +816,8 @@ def get_mpo(
  mpo.csvd_max_iter = csvd_max_iter
  if disjoint_levels is not None:
  mpo.disjoint_levels = bw.VectorFP(disjoint_levels)
+ mpo.disjoint_all_blocks = disjoint_all_blocks
+ mpo.disjoint_multiplier = disjoint_multiplier
  mpo.build()
  mpo = bw.bs.SimplifiedMPO(mpo, bw.bs.Rule(), False, False)
  mpo = bw.bs.IdentityAddedMPO(mpo)

pyblock2/unit_test/dmrg_mpo.py

@@ -24,10 +24,10 @@ def system_def(request):
  "Bipartite",
  "SVD",
  "DisjointSVD",
- "BlockedSumDisjointSVD",
- "FastBlockedSumDisjointSVD",
- "BlockedRescaledSumDisjointSVD",
- "FastBlockedRescaledSumDisjointSVD",
+ # "BlockedSumDisjointSVD",
+ # "FastBlockedSumDisjointSVD",
+ # "BlockedRescaledSumDisjointSVD",
+ # "FastBlockedRescaledSumDisjointSVD",
  "BlockedDisjointSVD",
  "FastBlockedDisjointSVD",
  "BlockedRescaledDisjointSVD",
@@ -101,7 +101,7 @@ def test_rhf(self, tmp_path, system_def, mpo_algo_type):
  reorder='irrep',
  cutoff=1e-12,
  algo_type=mpo_algo_type,
- iprint=1,
+ iprint=2,
  )
  ket = driver.get_random_mps(tag="GS", bond_dim=250, nroots=1)
  bond_dims = [250] * 8

src/core/iterative_matrix_functions.hpp

@@ -3208,7 +3208,9 @@ template <typename FL> struct IterativeMatrixFunctions : GMatrixFunctions<FL> {
  // the block-diagonal can have arbitrarily index permutation
  static void disjoint_svd(GMatrix<FL> x, GMatrix<FL> l, GMatrix<FP> s,
  GMatrix<FL> r, vector<FP> levels = vector<FP>(),
- bool ensure_ortho = true) {
+ bool ensure_ortho = true, bool iprint = false) {
+ if (x.m == 0 || x.n == 0)
+ return;
  // cout << "x = " << x << endl;
  sort(levels.begin(), levels.end(), greater<FP>());
  vector<DSU> dsus(levels.size() + 1, DSU(x.m + x.n));
@@ -3264,11 +3266,13 @@ template <typename FL> struct IterativeMatrixFunctions : GMatrixFunctions<FL> {
  }
  // number of singular values may exceed the maximal number
  // when needed, remove some levels to avoid this
+ MKL_INT ssk = s.n;
+ assert(ssk >= min(x.m, x.n));
  for (;;) {
  MKL_INT k = 0;
  for (auto &r : sub_k)
  k += r;
- if (k <= min(x.m, x.n))
+ if (k <= ssk)
  break;
  assert(sub_k.size() > 1);
  DSU &dsua = dsus[dsus.size() - 2], &dsub = dsus.back();
@@ -3298,17 +3302,17 @@ template <typename FL> struct IterativeMatrixFunctions : GMatrixFunctions<FL> {
  make_shared<VectorAllocator<FL>>();
  shared_ptr<VectorAllocator<FP>> d_alloc =
  make_shared<VectorAllocator<FP>>();
- size_t lwork = (size_t)max(x.m * x.m, x.n * x.n) +
-  (size_t)(x.m + x.n) * min(x.m, x.n) * 2;
+ size_t lwork =
+ (size_t)max(x.m * x.m, x.n * x.n) + (size_t)(x.m + x.n) * ssk * 2;
  FL *xwork = c_alloc->allocate(lwork);
  FL *xlwork = xwork + (size_t)max(x.m * x.m, x.n * x.n);
- FL *xrwork = xlwork + (size_t)x.m * min(x.m, x.n);
- FL *gwork = xrwork + (size_t)x.n * min(x.m, x.n);
- FP *swork = d_alloc->allocate(max(x.m, x.n) + min(x.m, x.n));
+ FL *xrwork = xlwork + (size_t)x.m * ssk;
+ FL *gwork = xrwork + (size_t)x.n * ssk;
+ FP *swork = d_alloc->allocate(max(x.m, x.n) + ssk);
  size_t iacc = 0;
- GMatrix<FL> glmat(gwork, x.m, min(x.m, x.n));
- GMatrix<FL> grmat(gwork + glmat.size(), min(x.m, x.n), x.n);
- GMatrix<FP> gsmat(swork + max(x.m, x.n), 1, min(x.m, x.n));
+ GMatrix<FL> glmat(gwork, x.m, ssk);
+ GMatrix<FL> grmat(gwork + glmat.size(), ssk, x.n);
+ GMatrix<FP> gsmat(swork + max(x.m, x.n), 1, ssk);
  glmat.clear();
  grmat.clear();
  gsmat.clear();
@@ -3391,7 +3395,7 @@ template <typename FL> struct IterativeMatrixFunctions : GMatrixFunctions<FL> {
  // cout << "gl = " << glmat << endl;
  // cout << "gs = " << gsmat << endl;
  // cout << "gr = " << grmat << endl;
- assert(gxk <= min(x.m, x.n));
+ assert(gxk <= ssk);
  l.clear(), r.clear(), s.clear();
  vector<MKL_INT> iwork(max(x.m, x.n));
  if (ensure_ortho) {
@@ -3477,7 +3481,7 @@ template <typename FL> struct IterativeMatrixFunctions : GMatrixFunctions<FL> {
  // cout << s << endl;
  // cout << r << endl;
  c_alloc->deallocate(xwork, lwork);
- d_alloc->deallocate(swork, max(x.m, x.n) + min(x.m, x.n));
+ d_alloc->deallocate(swork, max(x.m, x.n) + ssk);
  }
 };
 

src/dmrg/general_mpo.hpp

@@ -1415,6 +1415,8 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  FP csvd_eps = (FP)1E-10;
  int csvd_max_iter = 1000;
  vector<FP> disjoint_levels;
+ bool disjoint_all_blocks = false;
+ FP disjoint_multiplier = (FP)1.0;
  static inline size_t expr_index_hash(const string &expr,
  const uint16_t *terms, int n,
  const uint16_t init = 0) noexcept {
@@ -1439,6 +1441,8 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  bool sum_mpo = algo_type & MPOAlgorithmTypes::Sum;
  bool constrain = algo_type & MPOAlgorithmTypes::Constrained;
  bool disjoint = algo_type & MPOAlgorithmTypes::Disjoint;
+ if (!disjoint)
+ disjoint_multiplier = (FP)1.0;
  if (!(algo_type & MPOAlgorithmTypes::SVD) && max_bond_dim != -1)
  throw runtime_error(
  "Max bond dimension can only be used together with SVD!");
@@ -1632,6 +1636,8 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  for (int k = sr[j]; k < sr[j + 1]; k++)
  ip_sparse[k] = isr;
  }
+ FP eff_disjoint_multiplier =
+ ii == n_sites - 1 ? (FP)1.0 : disjoint_multiplier;
  // Part 1: iter over all mpos
  for (int ip = 0; ip < (int)cur_values.size(); ip++) {
  LL cn = (LL)cur_terms[ip].size(), cnr = cn;
@@ -1854,10 +1860,12 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  else if (algo_type & MPOAlgorithmTypes::CN)
  szm = ii == 0 ? szl : szr;
  else // bipartitie / SVD
- szm = min(szl, szr);
+ szm = (int)(min(szl, szr) * eff_disjoint_multiplier);
  if (!(algo_type & MPOAlgorithmTypes::Bipartite)) {
  if (delayed_term != -1 && iq == 0)
- szm = min(szl - 1, szr) + 1;
+ szm =
+ (int)(min(szl - 1, szr) * eff_disjoint_multiplier) +
+ 1;
  svds[iq].first[0].resize((size_t)szm * szl);
  svds[iq].second.resize(szm);
  svds[iq].first[1].resize((size_t)szm * szr);
@@ -1933,7 +1941,7 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  svds[iq].second[0] = 1;
  svds[iq].first[0][0] = 1;
  threading->activate_global_mkl();
- if (pqx[iq] < 2 && disjoint)
+ if ((pqx[iq] >= 2 || disjoint_all_blocks) && disjoint)
  IterativeMatrixFunctions<FL>::disjoint_svd(
  GMatrix<FL>(mat.data(), szl, szr),
  GMatrix<FL>(svds[iq].first[0].data() + 1 + szm,
@@ -1942,7 +1950,7 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  szm - 1),
  GMatrix<FL>(svds[iq].first[1].data() + szr,
  szm - 1, szr),
- disjoint_levels, false);
+ disjoint_levels, false, iprint >= 2);
  else
  GMatrixFunctions<FL>::svd(
  GMatrix<FL>(mat.data(), szl, szr),
@@ -1961,13 +1969,13 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  // cout << "mat = " << GMatrix<FL>(mat.data(), szl, szr)
  // << endl;
  threading->activate_global_mkl();
- if (pqx[iq] < 2 && disjoint)
+ if ((pqx[iq] >= 2 || disjoint_all_blocks) && disjoint)
  IterativeMatrixFunctions<FL>::disjoint_svd(
  GMatrix<FL>(mat.data(), szl, szr),
  GMatrix<FL>(svds[iq].first[0].data(), szl, szm),
  GMatrix<FP>(svds[iq].second.data(), 1, szm),
  GMatrix<FL>(svds[iq].first[1].data(), szm, szr),
- disjoint_levels, false);
+ disjoint_levels, false, iprint >= 2);
  else
  GMatrixFunctions<FL>::svd(
  GMatrix<FL>(mat.data(), szl, szr),
@@ -2025,9 +2033,11 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  int iq = mq.second;
  auto &nm = nms[iq];
  int szl = nm.first, szr = nm.second;
- int szm = min(szl, szr);
+ int szm = (int)(min(szl, szr) * eff_disjoint_multiplier);
  if (delayed_term != -1 && iq == 0)
- szm = min(szl - 1, szr) + 1;
+ szm =
+ (int)(min(szl - 1, szr) * eff_disjoint_multiplier) +
+ 1;
  for (int i = 0; i < szm; i++)
  svds[iq].second[i] /= res_factor;
  for (int i = 0; i < szm; i++)
@@ -2081,13 +2091,15 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  if (pqx[iq] < 2)
  continue;
  int szl = nm.first, szr = nm.second;
- int szm = min(szl, szr);
+ int szm = (int)(min(szl, szr) * eff_disjoint_multiplier);
  int rank = svds[iq].second.size();
  if ((szl == 1 && szr == 1) || rank == 0)
  continue;
  vector<FL> mat((size_t)szl * szr, 0);
  if (delayed_term != -1 && iq == 0) {
- szm = min(szl - 1, szr) + 1;
+ szm =
+ (int)(min(szl - 1, szr) * eff_disjoint_multiplier) +
+ 1;
  for (auto &lrv : matvs)
  if (lrv.first.first != 0)
  mat[(lrv.first.first - 1) * szr +
@@ -2132,12 +2144,16 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  int iq = mq.second;
  auto &nm = nms[iq];
  auto &matvs = mats[iq];
- int szl = nm.first, szr = nm.second, szm = min(szl, szr);
+ int szl = nm.first, szr = nm.second,
+ szm = (int)(min(szl, szr) * eff_disjoint_multiplier);
  if (delayed_term != -1 && iq == 0)
- szm = min(szl - 1, szr) + 1;
+ szm =
+ (int)(min(szl - 1, szr) * eff_disjoint_multiplier) +
+ 1;
  int s_kept = svds[iq].second.size();
  vector<FL> smat, stmp;
- smat.reserve((size_t)szl * szr);
+ smat.reserve(
+ max((size_t)szl * szr, (size_t)s_kept * s_kept));
  stmp.reserve((size_t)s_kept * szr);
  memset(smat.data(), 0, sizeof(FL) * s_kept * s_kept);
  for (int i = 0; i < s_kept; i++)
@@ -2256,9 +2272,11 @@ template <typename S, typename FL> struct GeneralMPO : MPO<S, FL> {
  else if (algo_type & MPOAlgorithmTypes::Bipartite)
  szm = (int)mvcs[iq][0].size() + (int)mvcs[iq][1].size();
  else { // SVD
- szm = min(szl, szr);
+ szm = (int)(min(szl, szr) * eff_disjoint_multiplier);
  if (delayed_term != -1 && iq == 0)
- szm = min(szl - 1, szr) + 1;
+ szm =
+ (int)(min(szl - 1, szr) * eff_disjoint_multiplier) +
+ 1;
  }
  vector<shared_ptr<OpElement<S, FL>>> site_mp(szl);
  for (auto &vls : mpl)

src/pybind/pybind_dmrg.hpp

@@ -1849,6 +1849,10 @@ template <typename S, typename FL> void bind_fl_general(py::module &m) {
  .def_readwrite("csvd_eps", &GeneralMPO<S, FL>::csvd_eps)
  .def_readwrite("csvd_max_iter", &GeneralMPO<S, FL>::csvd_max_iter)
  .def_readwrite("disjoint_levels", &GeneralMPO<S, FL>::disjoint_levels)
+ .def_readwrite("disjoint_all_blocks",
+ &GeneralMPO<S, FL>::disjoint_all_blocks)
+ .def_readwrite("disjoint_multiplier",
+ &GeneralMPO<S, FL>::disjoint_multiplier)
  .def(py::init<const shared_ptr<GeneralHamiltonian<S, FL>> &,
  const shared_ptr<GeneralFCIDUMP<FL>> &,
  MPOAlgorithmTypes>(),

unit_test/debug_test_dmrg_auto.cpp

@@ -162,7 +162,12 @@ TYPED_TEST(TestDMRG, Test) {
  // MPO construction
  cout << "MPO start" << endl;
  shared_ptr<MPO<S, FL>> mpo = make_shared<GeneralMPO<S, FL>>(
- gham, gfd, MPOAlgorithmTypes::FastBipartite, 1E-7, -1);
+ gham, gfd, MPOAlgorithmTypes::FastBlockedSumDisjointSVD, 1E-7, -1);
+ // dynamic_pointer_cast<GeneralMPO<S, FL>>(mpo)->disjoint_all_blocks = true;
+ // dynamic_pointer_cast<GeneralMPO<S, FL>>(mpo)->disjoint_multiplier = 4.0;
+ // dynamic_pointer_cast<GeneralMPO<S, FL>>(mpo)->disjoint_levels.push_back(1E-3);
+ // dynamic_pointer_cast<GeneralMPO<S, FL>>(mpo)->disjoint_levels.push_back(1E-4);
+ // dynamic_pointer_cast<GeneralMPO<S, FL>>(mpo)->disjoint_levels.push_back(1E-5);
  mpo->build();
  // shared_ptr<MPO<S, FL>> mpo =
  // make_shared<MPOQC<S, FL>>(hamil, QCTypes::Conventional);