Skip to content

Commit

Permalink
core: 5pdm
Browse files Browse the repository at this point in the history
  • Loading branch information
@hczhai
hczhai committed Oct 19, 2024
1 parent a7f7da9 commit d8c67e8
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 51 deletions.
83 changes: 57 additions & 26 deletions pyblock2/driver/core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5544,12 +5544,11 @@ def get_npdm(
op_str = "(C+D)0"
for _ in range(pdm_type - 1):
op_str = su2_coupling % op_str
if mask is None:
perm = bw.b.SpinPermScheme.initialize_su2(pdm_type * 2, op_str, True)
else:
perm = bw.b.SpinPermScheme.initialize_su2(
pdm_type * 2, op_str, True, mask=bw.b.VectorUInt16(mask)
)
perm = bw.b.SpinPermScheme.initialize_su2(
pdm_type * 2, op_str, True,
mask=bw.b.VectorUInt16() if mask is None else bw.b.VectorUInt16(mask),
max_n_sites=ket.n_sites,
)
perms = bw.b.VectorSpinPermScheme([perm])
elif SymmetryTypes.SZ in bw.symm_type:
if npdm_expr is not None and isinstance(npdm_expr, str):
Expand All @@ -5563,10 +5562,10 @@ def get_npdm(
if mask is None:
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sz(pdm_type * 2, cd, True)
if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(pdm_type * 2, cd, fermionic_ops)
for cd in op_str
bw.b.SpinPermScheme.initialize_sz(pdm_type * 2, cd, True,
mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(pdm_type * 2, cd, fermionic_ops,
mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) for cd in op_str
]
)
elif len(mask) != 0 and not isinstance(mask[0], int):
Expand All @@ -5577,10 +5576,10 @@ def get_npdm(
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sz(
pt * 2, cd, True, mask=bw.b.VectorUInt16(xm)
pt * 2, cd, True, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
) if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(xm)
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
)
for cd, xm, pt in zip(op_str, mask, pts)
]
Expand All @@ -5592,10 +5591,10 @@ def get_npdm(
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sz(
pt * 2, cd, True, mask=bw.b.VectorUInt16(mask)
pt * 2, cd, True, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
) if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(mask)
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
)
for cd, pt in zip(op_str, pts)
]
Expand All @@ -5610,10 +5609,10 @@ def get_npdm(
if mask is None:
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sz(pdm_type * 2, cd, True)
if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(pdm_type * 2, cd, fermionic_ops)
for cd in op_str
bw.b.SpinPermScheme.initialize_sz(pdm_type * 2, cd, True,
mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(pdm_type * 2, cd, fermionic_ops,
mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) for cd in op_str
]
)
elif len(mask) != 0 and not isinstance(mask[0], int):
Expand All @@ -5624,10 +5623,10 @@ def get_npdm(
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sz(
pt * 2, cd, True, mask=bw.b.VectorUInt16(xm)
pt * 2, cd, True, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
) if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(xm)
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
)
for cd, xm, pt in zip(op_str, mask, pts)
]
Expand All @@ -5639,10 +5638,10 @@ def get_npdm(
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sz(
pt * 2, cd, True, mask=bw.b.VectorUInt16(mask)
pt * 2, cd, True, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
) if fermionic_ops is None else
bw.b.SpinPermScheme.initialize_sany(
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(mask)
pt * 2, cd, fermionic_ops, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
)
for cd, pt in zip(op_str, pts)
]
Expand All @@ -5654,16 +5653,16 @@ def get_npdm(
if mask is None:
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sany(len(cd), cd, fermionic_ops)
for cd in op_str
bw.b.SpinPermScheme.initialize_sany(len(cd), cd, fermionic_ops,
mask=bw.b.VectorUInt16(), max_n_sites=ket.n_sites) for cd in op_str
]
)
elif len(mask) != 0 and not isinstance(mask[0], int):
assert len(mask) == len(op_str)
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sany(
len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(xm)
len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(xm), max_n_sites=ket.n_sites
)
for cd, xm in zip(op_str, mask)
]
Expand All @@ -5672,7 +5671,7 @@ def get_npdm(
perms = bw.b.VectorSpinPermScheme(
[
bw.b.SpinPermScheme.initialize_sany(
len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(mask)
len(cd), cd, fermionic_ops, mask=bw.b.VectorUInt16(mask), max_n_sites=ket.n_sites
)
for cd in op_str
]
Expand Down Expand Up @@ -5914,6 +5913,20 @@ def get_4pdm(self, ket, *args, **kwargs):
"""
return self.get_npdm(ket, pdm_type=4, *args, **kwargs)

def get_5pdm(self, ket, *args, **kwargs):
"""
Compute the 5-Particle Density Matrix (5PDM) for the given MPS.
See ``DMRGDriver.get_npdm``.
"""
return self.get_npdm(ket, pdm_type=5, *args, **kwargs)

def get_6pdm(self, ket, *args, **kwargs):
"""
Compute the 6-Particle Density Matrix (6PDM) for the given MPS.
See ``DMRGDriver.get_npdm``.
"""
return self.get_npdm(ket, pdm_type=6, *args, **kwargs)

def get_trans_1pdm(self, bra, ket, *args, **kwargs):
"""
Compute the Transition 1-Particle Density Matrix (T-1PDM)
Expand Down Expand Up @@ -5950,6 +5963,24 @@ def get_trans_4pdm(self, bra, ket, *args, **kwargs):
"""
return self.get_npdm(ket, pdm_type=4, bra=bra, *args, **kwargs)

def get_trans_5pdm(self, bra, ket, *args, **kwargs):
"""
Compute the Transition 5-Particle Density Matrix (T-5PDM)
between the given bra and ket MPSs.
Note that there can be an overall phase uncertainty for transition NPDMs.
See ``DMRGDriver.get_npdm``.
"""
return self.get_npdm(ket, pdm_type=5, bra=bra, *args, **kwargs)

def get_trans_6pdm(self, bra, ket, *args, **kwargs):
"""
Compute the Transition 6-Particle Density Matrix (T-6PDM)
between the given bra and ket MPSs.
Note that there can be an overall phase uncertainty for transition NPDMs.
See ``DMRGDriver.get_npdm``.
"""
return self.get_npdm(ket, pdm_type=6, bra=bra, *args, **kwargs)

def get_csf_coefficients(
self, ket, cutoff=0.1, given_dets=None, max_print=200, fci_conv=False, iprint=1
):
Expand Down
38 changes: 24 additions & 14 deletions src/core/spin_permutation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -938,8 +938,9 @@ struct SpinPermPattern {
vector<uint16_t> mask;
vector<uint16_t> data;
SpinPermPattern(uint16_t n,
const vector<uint16_t> &mask = vector<uint16_t>())
: n(n), mask(mask), data(initialize(n, mask)) {}
const vector<uint16_t> &mask = vector<uint16_t>(),
int max_n_sites = 0)
: n(n), mask(mask), data(initialize(n, mask, max_n_sites)) {}
static vector<uint16_t>
all_reordering(const vector<uint16_t> &x,
const vector<uint16_t> &mask = vector<uint16_t>()) {
Expand Down Expand Up @@ -1007,13 +1008,16 @@ struct SpinPermPattern {
return r;
}
static vector<uint16_t>
initialize(uint16_t n, const vector<uint16_t> &mask = vector<uint16_t>()) {
initialize(uint16_t n, const vector<uint16_t> &mask = vector<uint16_t>(),
int max_n_sites = 0) {
int mi = n;
if (mask.size() != 0) {
mi = 1;
for (uint16_t j = 1; j < n; j++)
mi += mask[j] != mask[j - 1];
}
if (max_n_sites != 0)
mi = mi < max_n_sites ? mi : max_n_sites;
map<pair<uint16_t, uint16_t>, vector<vector<uint16_t>>> mp;
for (uint16_t i = 0; i <= n; i++)
mp[make_pair(0, i)] = vector<vector<uint16_t>>();
Expand Down Expand Up @@ -1190,12 +1194,14 @@ struct SpinPermScheme {
SpinPermScheme() {}
SpinPermScheme(string spin_str, bool su2 = true, bool is_fermion = true,
bool is_npdm = false, bool is_drt = false,
const vector<uint16_t> &mask = vector<uint16_t>()) {
const vector<uint16_t> &mask = vector<uint16_t>(),
int max_n_sites = 0) {
int nn = SpinPermRecoupling::count_cds(spin_str);
SpinPermScheme r =
su2 ? SpinPermScheme::initialize_su2(nn, spin_str, is_npdm, is_drt,
mask)
: SpinPermScheme::initialize_sz(nn, spin_str, is_fermion, mask);
mask, max_n_sites)
: SpinPermScheme::initialize_sz(nn, spin_str, is_fermion, mask,
max_n_sites);
index_patterns = r.index_patterns;
this->mask = r.mask;
data = r.data;
Expand All @@ -1204,10 +1210,11 @@ struct SpinPermScheme {
}
static SpinPermScheme
initialize_sz(int nn, const string &spin_str, bool is_fermion = true,
const vector<uint16_t> &mask = vector<uint16_t>()) {
const vector<uint16_t> &mask = vector<uint16_t>(),
int max_n_sites = 0) {
using T = SpinPermTensor;
using R = SpinPermRecoupling;
SpinPermPattern spat(nn, mask);
SpinPermPattern spat(nn, mask, max_n_sites);
vector<double> mptr;
SpinPermScheme r;
r.index_patterns.resize(spat.count());
Expand Down Expand Up @@ -1241,10 +1248,11 @@ struct SpinPermScheme {
static SpinPermScheme
initialize_sany(int nn, const string &spin_str,
const string &fermionic_ops = "cdCD",
const vector<uint16_t> &mask = vector<uint16_t>()) {
const vector<uint16_t> &mask = vector<uint16_t>(),
int max_n_sites = 0) {
using T = SpinPermTensor;
using R = SpinPermRecoupling;
SpinPermPattern spat(nn, mask);
SpinPermPattern spat(nn, mask, max_n_sites);
vector<double> mptr;
SpinPermScheme r;
r.index_patterns.resize(spat.count());
Expand Down Expand Up @@ -1463,7 +1471,8 @@ struct SpinPermScheme {
static SpinPermScheme
initialize_su2(int nn, const string &spin_str, bool is_npdm = false,
bool is_drt = false,
const vector<uint16_t> &mask = vector<uint16_t>()) {
const vector<uint16_t> &mask = vector<uint16_t>(),
int max_n_sites = 0) {
using T = SpinPermTensor;
using R = SpinPermRecoupling;
SU2CG cg;
Expand All @@ -1473,7 +1482,7 @@ struct SpinPermScheme {
string spin_pat_str = R::split_cds(spin_str, cds);
bool heis = cds.size() != 0 && cds[0] == 2;
int target_twos = R::get_target_twos(spin_pat_str);
SpinPermPattern spat(nn, mask);
SpinPermPattern spat(nn, mask, max_n_sites);
SpinPermScheme r;
r.index_patterns.resize(spat.count());
r.data.resize(spat.count());
Expand Down Expand Up @@ -1523,7 +1532,8 @@ struct SpinPermScheme {
static SpinPermScheme
initialize_su2_old2(int nn, string spin_str, bool is_npdm = false,
bool is_drt = false,
const vector<uint16_t> &mask = vector<uint16_t>()) {
const vector<uint16_t> &mask = vector<uint16_t>(),
int max_n_sites = 0) {
using T = SpinPermTensor;
using R = SpinPermRecoupling;
SU2CG cg;
Expand All @@ -1533,7 +1543,7 @@ struct SpinPermScheme {
spin_str = R::split_cds(spin_str, cds);
bool heis = cds.size() != 0 && cds[0] == 2;
int target_twos = R::get_target_twos(spin_str);
SpinPermPattern spat(nn, mask);
SpinPermPattern spat(nn, mask, max_n_sites);
SpinPermScheme r;
r.index_patterns.resize(spat.count());
r.data.resize(spat.count());
Expand Down
28 changes: 17 additions & 11 deletions src/pybind/pybind_core.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2323,10 +2323,12 @@ template <typename S = void> void bind_io(py::module &m) {
.def_readwrite("data", &SpinPermPattern::data)
.def(py::init<uint16_t>())
.def(py::init<uint16_t, const vector<uint16_t> &>())
.def(py::init<uint16_t, const vector<uint16_t> &, int>())
.def_static("all_reordering", &SpinPermPattern::all_reordering,
py::arg("x"), py::arg("mask") = vector<uint16_t>())
.def_static("initialize", &SpinPermPattern::initialize, py::arg("n"),
py::arg("mask") = vector<uint16_t>())
py::arg("mask") = vector<uint16_t>(),
py::arg("max_n_sites") = 0)
.def_static("get_unique", &SpinPermPattern::get_unique)
.def_static("make_matrix", &SpinPermPattern::make_matrix)
.def("count", &SpinPermPattern::count)
Expand Down Expand Up @@ -2359,30 +2361,34 @@ template <typename S = void> void bind_io(py::module &m) {
.def(py::init<string, bool, bool, bool, bool>())
.def(py::init<string, bool, bool, bool, bool,
const vector<uint16_t> &>())
.def(py::init<string, bool, bool, bool, bool, const vector<uint16_t> &,
int>())
.def_readwrite("index_patterns", &SpinPermScheme::index_patterns)
.def_readwrite("data", &SpinPermScheme::data)
.def_readwrite("is_su2", &SpinPermScheme::is_su2)
.def_readwrite("left_vacuum", &SpinPermScheme::left_vacuum)
.def_readwrite("mask", &SpinPermScheme::mask)
.def_static("initialize_sz", &SpinPermScheme::initialize_sz,
py::arg("nn"), py::arg("spin_str"),
py::arg("is_fermion") = true,
py::arg("mask") = vector<uint16_t>())
.def_static("initialize_sany", &SpinPermScheme::initialize_sany,
py::arg("nn"), py::arg("spin_str"),
py::arg("fermionic_ops") = string("cdCD"),
py::arg("mask") = vector<uint16_t>())
.def_static(
"initialize_sz", &SpinPermScheme::initialize_sz, py::arg("nn"),
py::arg("spin_str"), py::arg("is_fermion") = true,
py::arg("mask") = vector<uint16_t>(), py::arg("max_n_sites") = 0)
.def_static(
"initialize_sany", &SpinPermScheme::initialize_sany, py::arg("nn"),
py::arg("spin_str"), py::arg("fermionic_ops") = string("cdCD"),
py::arg("mask") = vector<uint16_t>(), py::arg("max_n_sites") = 0)
.def_static("initialize_su2_old", &SpinPermScheme::initialize_su2_old,
py::arg("nn"), py::arg("spin_str"),
py::arg("is_npdm") = false)
.def_static("initialize_su2", &SpinPermScheme::initialize_su2,
py::arg("nn"), py::arg("spin_str"),
py::arg("is_npdm") = false, py::arg("is_drt") = false,
py::arg("mask") = vector<uint16_t>())
py::arg("mask") = vector<uint16_t>(),
py::arg("max_n_sites") = 0)
.def_static("initialize_su2_old2", &SpinPermScheme::initialize_su2_old2,
py::arg("nn"), py::arg("spin_str"),
py::arg("is_npdm") = false, py::arg("is_drt") = false,
py::arg("mask") = vector<uint16_t>())
py::arg("mask") = vector<uint16_t>(),
py::arg("max_n_sites") = 0)
.def("to_str", &SpinPermScheme::to_str);

py::bind_vector<vector<shared_ptr<SpinPermScheme>>>(m,
Expand Down

0 comments on commit d8c67e8

Please sign in to comment.