Skip to content

Commit

Permalink
Merge branch 'master' of https://github.com/block-hczhai/block2-preview
Browse files Browse the repository at this point in the history
  • Loading branch information
@hczhai
hczhai committed Jun 15, 2023
2 parents 59f9db3 + 6293ceb commit f155bbe
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 28 deletions.
4 changes: 4 additions & 0 deletions pyblock2/driver/core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1689,6 +1689,8 @@ def td_dmrg(
hermitian=False,
iprint=0,
cutoff=1e-20,
krylov_conv_thrd=5E-6,
krylov_subspace_size=20,
):
bw = self.bw
import numpy as np
Expand Down Expand Up @@ -1736,6 +1738,8 @@ def td_dmrg(
te.n_sub_sweeps = n_sub_sweeps
te.normalize_mps = normalize_mps
te.cutoff = cutoff
te.krylov_conv_thrd = krylov_conv_thrd
te.krylov_subspace_size = krylov_subspace_size

te_times = []
te_energies = []
Expand Down
12 changes: 8 additions & 4 deletions src/dmrg/effective_functions.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -405,7 +405,8 @@ struct EffectiveFunctions<
static tuple<FL, FP, int, size_t, double> expo_apply(
const shared_ptr<EffectiveHamiltonian<S, FL, MultiMPS<S, FL>>> &h_eff,
FC beta, typename const_fl_type<FL>::FL const_e, bool iprint = false,
const shared_ptr<ParallelRule<S>> &para_rule = nullptr) {
const shared_ptr<ParallelRule<S>> &para_rule = nullptr,
FP conv_thrd = 5E-6, int deflation_max_size = 20) {
assert(h_eff->compute_diag);
assert(h_eff->ket.size() == 2);
FP anorm = GMatrixFunctions<FL>::norm(GMatrix<FL>(
Expand All @@ -423,10 +424,12 @@ struct EffectiveFunctions<
(h_eff->tf->opf->seq->mode & SeqTypes::Tasked))
? GMatrixFunctions<FC>::expo_apply(
*h_eff->tf, beta, anorm, vr, vi, (FL)const_e, iprint,
para_rule == nullptr ? nullptr : para_rule->comm)
para_rule == nullptr ? nullptr : para_rule->comm,
conv_thrd, deflation_max_size)
: GMatrixFunctions<FC>::expo_apply(
*h_eff, beta, anorm, vr, vi, (FL)const_e, iprint,
para_rule == nullptr ? nullptr : para_rule->comm);
para_rule == nullptr ? nullptr : para_rule->comm,
conv_thrd, deflation_max_size);
FP norm_re = GMatrixFunctions<FL>::norm(vr);
FP norm_im = GMatrixFunctions<FL>::norm(vi);
FP norm = sqrt(norm_re * norm_re + norm_im * norm_im);
Expand Down Expand Up @@ -683,7 +686,8 @@ struct EffectiveFunctions<S, FL,
static tuple<FL, FP, int, size_t, double> expo_apply(
const shared_ptr<EffectiveHamiltonian<S, FL, MultiMPS<S, FL>>> &h_eff,
FC beta, typename const_fl_type<FL>::FL const_e, bool iprint = false,
const shared_ptr<ParallelRule<S>> &para_rule = nullptr) {
const shared_ptr<ParallelRule<S>> &para_rule = nullptr,
FP conv_thrd = 5E-6, int deflation_max_size = 20) {
assert(false);
return make_tuple(0.0, 0.0, 0, (size_t)0, 0.0);
}
Expand Down
9 changes: 6 additions & 3 deletions src/dmrg/effective_hamiltonian.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -891,7 +891,8 @@ struct EffectiveHamiltonian<S, FL, MPS<S, FL>> {
tuple<FL, FP, int, size_t, double>
expo_apply(FL beta, typename const_fl_type<FL>::FL const_e, bool symmetric,
bool iprint = false,
const shared_ptr<ParallelRule<S>> &para_rule = nullptr) {
const shared_ptr<ParallelRule<S>> &para_rule = nullptr,
FP conv_thrd = 5E-6, int deflation_max_size = 20) {
assert(compute_diag);
FP anorm = GMatrixFunctions<FL>::norm(
GMatrix<FL>(diag->data, (MKL_INT)diag->total_memory, 1));
Expand All @@ -905,10 +906,12 @@ struct EffectiveHamiltonian<S, FL, MPS<S, FL>> {
(tf->opf->seq->mode & SeqTypes::Tasked))
? IterativeMatrixFunctions<FL>::expo_apply(
*tf, beta, anorm, v, (FL)const_e, symmetric, iprint,
para_rule == nullptr ? nullptr : para_rule->comm)
para_rule == nullptr ? nullptr : para_rule->comm,
conv_thrd, deflation_max_size)
: IterativeMatrixFunctions<FL>::expo_apply(
*this, beta, anorm, v, (FL)const_e, symmetric, iprint,
para_rule == nullptr ? nullptr : para_rule->comm);
para_rule == nullptr ? nullptr : para_rule->comm,
conv_thrd, deflation_max_size);
FP norm = GMatrixFunctions<FL>::norm(v);
GMatrix<FL> tmp(nullptr, (MKL_INT)ket->total_memory, 1);
tmp.allocate();
Expand Down
58 changes: 38 additions & 20 deletions src/dmrg/sweep_algorithm_td.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -818,6 +818,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
bool store_wfn_spectra = false;
vector<vector<FPS>> sweep_wfn_spectra;
vector<FPS> wfn_spectra;
FPS krylov_conv_thrd = 5E-6;
int krylov_subspace_size = 20;
TimeEvolution(const shared_ptr<MovingEnvironment<S, FL, FLS>> &me,
const vector<ubond_t> &bond_dims,
TETypes mode = TETypes::TangentSpace, int n_sub_sweeps = 1)
Expand Down Expand Up @@ -898,7 +900,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
memcpy(tmp.data, h_eff->ket->data,
h_eff->ket->total_memory * sizeof(FLS));
pdi = h_eff->expo_apply(-beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
memcpy(h_eff->ket->data, tmp.data,
h_eff->ket->total_memory * sizeof(FLS));
tmp.deallocate();
Expand All @@ -907,7 +910,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
pdpf = pdp.first;
} else if (effective_mode == TETypes::TangentSpace)
pdi = h_eff->expo_apply(-beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
else if (effective_mode == TETypes::RK4) {
auto pdp =
h_eff->rk4_apply(-beta, me->mpo->const_e, false, me->para_rule);
Expand Down Expand Up @@ -1046,8 +1050,9 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
me->move_to(i + 1, true);
shared_ptr<EffectiveHamiltonian<S, FL>> k_eff = me->eff_ham(
FuseTypes::NoFuseL, forward, true, right, right);
auto pdk = k_eff->expo_apply(beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
auto pdk = k_eff->expo_apply(
beta, me->mpo->const_e, hermitian, iprint >= 3,
me->para_rule, krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
if (me->para_rule == nullptr || me->para_rule->is_root()) {
if (normalize_mps)
Expand All @@ -1064,8 +1069,9 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
me->move_to(i - 1, true);
shared_ptr<EffectiveHamiltonian<S, FL>> k_eff =
me->eff_ham(FuseTypes::NoFuseR, forward, true, left, left);
auto pdk = k_eff->expo_apply(beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
auto pdk = k_eff->expo_apply(
beta, me->mpo->const_e, hermitian, iprint >= 3,
me->para_rule, krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
if (me->para_rule == nullptr || me->para_rule->is_root()) {
if (normalize_mps)
Expand Down Expand Up @@ -1178,7 +1184,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
memcpy(tmp.data, h_eff->ket->data,
h_eff->ket->total_memory * sizeof(FLS));
pdi = h_eff->expo_apply(-beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
memcpy(h_eff->ket->data, tmp.data,
h_eff->ket->total_memory * sizeof(FLS));
tmp.deallocate();
Expand All @@ -1187,7 +1194,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
pdpf = pdp.first;
} else if (effective_mode == TETypes::TangentSpace)
pdi = h_eff->expo_apply(-beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
else if (effective_mode == TETypes::RK4) {
auto pdp =
h_eff->rk4_apply(-beta, me->mpo->const_e, false, me->para_rule);
Expand Down Expand Up @@ -1288,8 +1296,9 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
shared_ptr<EffectiveHamiltonian<S, FL>> k_eff =
me->eff_ham(FuseTypes::FuseR, forward, true,
me->bra->tensors[i + 1], me->ket->tensors[i + 1]);
auto pdk = k_eff->expo_apply(beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
auto pdk = k_eff->expo_apply(
beta, me->mpo->const_e, hermitian, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
if (me->para_rule == nullptr || me->para_rule->is_root()) {
if (normalize_mps)
Expand All @@ -1305,8 +1314,9 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
shared_ptr<EffectiveHamiltonian<S, FL>> k_eff =
me->eff_ham(FuseTypes::FuseL, forward, true,
me->bra->tensors[i], me->ket->tensors[i]);
auto pdk = k_eff->expo_apply(beta, me->mpo->const_e, hermitian,
iprint >= 3, me->para_rule);
auto pdk = k_eff->expo_apply(
beta, me->mpo->const_e, hermitian, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
if (me->para_rule == nullptr || me->para_rule->is_root()) {
if (normalize_mps)
Expand Down Expand Up @@ -1392,7 +1402,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
memcpy(tmp_im.data, h_eff->ket[1]->data,
h_eff->ket[1]->total_memory * sizeof(FLS));
pdi = EffectiveFunctions<S, FL>::expo_apply(
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule);
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
memcpy(h_eff->ket[0]->data, tmp_re.data,
h_eff->ket[0]->total_memory * sizeof(FLS));
memcpy(h_eff->ket[1]->data, tmp_im.data,
Expand All @@ -1404,7 +1415,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
pdpf = pdp.first;
} else if (effective_mode == TETypes::TangentSpace)
pdi = EffectiveFunctions<S, FL>::expo_apply(
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule);
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
else if (effective_mode == TETypes::RK4) {
auto pdp =
h_eff->rk4_apply(-beta, me->mpo->const_e, false, me->para_rule);
Expand Down Expand Up @@ -1576,7 +1588,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
shared_ptr<EffectiveHamiltonian<S, FL, MultiMPS<S, FL>>> k_eff =
me->multi_eff_ham(FuseTypes::NoFuseL, forward, true);
auto pdk = EffectiveFunctions<S, FL>::expo_apply(
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule);
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
mket->wfns = kwfns;
if (me->para_rule == nullptr || me->para_rule->is_root()) {
Expand All @@ -1597,7 +1610,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
shared_ptr<EffectiveHamiltonian<S, FL, MultiMPS<S, FL>>> k_eff =
me->multi_eff_ham(FuseTypes::NoFuseR, forward, true);
auto pdk = EffectiveFunctions<S, FL>::expo_apply(
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule);
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
mket->wfns = kwfns;
if (me->para_rule == nullptr || me->para_rule->is_root()) {
Expand Down Expand Up @@ -1749,7 +1763,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
memcpy(tmp_im.data, h_eff->ket[1]->data,
h_eff->ket[1]->total_memory * sizeof(FLS));
pdi = EffectiveFunctions<S, FL>::expo_apply(
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule);
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
memcpy(h_eff->ket[0]->data, tmp_re.data,
h_eff->ket[0]->total_memory * sizeof(FLS));
memcpy(h_eff->ket[1]->data, tmp_im.data,
Expand All @@ -1761,7 +1776,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
pdpf = pdp.first;
} else if (effective_mode == TETypes::TangentSpace)
pdi = EffectiveFunctions<S, FL>::expo_apply(
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule);
h_eff, -beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
else if (effective_mode == TETypes::RK4) {
auto pdp =
h_eff->rk4_apply(-beta, me->mpo->const_e, false, me->para_rule);
Expand Down Expand Up @@ -1875,7 +1891,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
shared_ptr<EffectiveHamiltonian<S, FL, MultiMPS<S, FL>>> k_eff =
me->multi_eff_ham(FuseTypes::FuseR, forward, true);
auto pdk = EffectiveFunctions<S, FL>::expo_apply(
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule);
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
if (me->para_rule == nullptr || me->para_rule->is_root()) {
if (normalize_mps)
Expand All @@ -1891,7 +1908,8 @@ template <typename S, typename FL, typename FLS> struct TimeEvolution {
shared_ptr<EffectiveHamiltonian<S, FL, MultiMPS<S, FL>>> k_eff =
me->multi_eff_ham(FuseTypes::FuseL, forward, true);
auto pdk = EffectiveFunctions<S, FL>::expo_apply(
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule);
k_eff, beta, me->mpo->const_e, iprint >= 3, me->para_rule,
krylov_conv_thrd, krylov_subspace_size);
k_eff->deallocate();
if (me->para_rule == nullptr || me->para_rule->is_root()) {
if (normalize_mps)
Expand Down
8 changes: 7 additions & 1 deletion src/pybind/pybind_dmrg.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -600,7 +600,9 @@ template <typename S, typename FL> void bind_fl_partition(py::module &m) {
py::arg("eval_energy") = false, py::arg("para_rule") = nullptr)
.def("expo_apply", &EffectiveHamiltonian<S, FL>::expo_apply,
py::arg("beta"), py::arg("const_e"), py::arg("symmetric"),
py::arg("iprint") = false, py::arg("para_rule") = nullptr)
py::arg("iprint") = false, py::arg("para_rule") = nullptr,
py::arg("conv_thrd") = (typename GMatrix<FL>::FP)5E-6,
py::arg("deflation_max_size") = 20)
.def("deallocate", &EffectiveHamiltonian<S, FL>::deallocate);

py::class_<EffectiveFunctions<S, FL>,
Expand Down Expand Up @@ -1349,6 +1351,10 @@ void bind_fl_td_dmrg(py::module &m) {
.def_readwrite("wfn_spectra", &TimeEvolution<S, FL, FLS>::wfn_spectra)
.def_readwrite("sweep_wfn_spectra",
&TimeEvolution<S, FL, FLS>::sweep_wfn_spectra)
.def_readwrite("krylov_conv_thrd",
&TimeEvolution<S, FL, FLS>::krylov_conv_thrd)
.def_readwrite("krylov_subspace_size",
&TimeEvolution<S, FL, FLS>::krylov_subspace_size)
.def("update_one_dot", &TimeEvolution<S, FL, FLS>::update_one_dot)
.def("update_two_dot", &TimeEvolution<S, FL, FLS>::update_two_dot)
.def("update_multi_one_dot",
Expand Down

0 comments on commit f155bbe

Please sign in to comment.