Skip to content

Commit

Permalink
fix disjoint svd test
Browse files Browse the repository at this point in the history
  • Loading branch information
@hczhai
hczhai committed Oct 10, 2022
1 parent 3388c85 commit f069353
Show file tree
Hide file tree
Showing 7 changed files with 74 additions and 35 deletions.
4 changes: 2 additions & 2 deletions .github/workflows/build.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
8 changes: 8 additions & 0 deletions pyblock2/driver/core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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:
Expand All @@ -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)
Expand Down
10 changes: 5 additions & 5 deletions pyblock2/unit_test/dmrg_mpo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,10 @@ def system_def(request):
"Bipartite",
"SVD",
"DisjointSVD",
"BlockedSumDisjointSVD",
"FastBlockedSumDisjointSVD",
"BlockedRescaledSumDisjointSVD",
"FastBlockedRescaledSumDisjointSVD",
# "BlockedSumDisjointSVD",
# "FastBlockedSumDisjointSVD",
# "BlockedRescaledSumDisjointSVD",
# "FastBlockedRescaledSumDisjointSVD",
"BlockedDisjointSVD",
"FastBlockedDisjointSVD",
"BlockedRescaledDisjointSVD",
Expand Down Expand Up @@ -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
Expand Down
28 changes: 16 additions & 12 deletions src/core/iterative_matrix_functions.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
}
};

Expand Down
48 changes: 33 additions & 15 deletions src/dmrg/general_mpo.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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!");
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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,
Expand All @@ -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),
Expand All @@ -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),
Expand Down Expand Up @@ -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++)
Expand Down Expand Up @@ -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 +
Expand Down Expand Up @@ -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++)
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 4 additions & 0 deletions src/pybind/pybind_dmrg.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>(),
Expand Down
7 changes: 6 additions & 1 deletion unit_test/debug_test_dmrg_auto.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit f069353

Please sign in to comment.