Skip to content

Commit

Permalink
COSX Grid Handling Improvements (#3032)
Browse files Browse the repository at this point in the history 
* Small rework of grid classification in COSX

* Clean up comments and grid debug code

* Change CompositeJK timer to separate SCF iteration times per grid

* General changes and refactors

* Whitespace cleanup

* Have grid info print on debug specifically

* Further adjustments to support cleaner COSX model

* Small compiler error fix

* Some more small cleanups

* Post-rebase conde modifications

* Post-rebase cleanup

* Implement some of Susi's suggestions

* Implement Lori's suggestions for gridopt and clarity comments

* Remove spare early_screening_ dependency in CompositeJK.cc

* Remove early_screening_ C++-side, now fully in Python
  • Loading branch information
@davpoolechem
davpoolechem committed Nov 13, 2023
1 parent 325402f commit 289d095
Show file tree
Hide file tree
Showing 7 changed files with 132 additions and 127 deletions.
18 changes: 13 additions & 5 deletions psi4/driver/procrouting/scf_proc/scf_iterator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,9 +267,14 @@ def scf_iterate(self, e_conv=None, d_conv=None):
soscf_enabled = _validate_soscf()
frac_enabled = _validate_frac()
efp_enabled = hasattr(self.molecule(), 'EFP')
cosx_enabled = "COSX" in core.get_option('SCF', 'SCF_TYPE')

# does the JK algorithm use severe screening approximations for early SCF iterations?
early_screening = self.jk().get_early_screening()
early_screening = False
if cosx_enabled:
early_screening = True
self.jk().set_COSX_grid("Initial")

# maximum number of scf iterations to run after early screening is disabled
scf_maxiter_post_screening = core.get_option('SCF', 'COSX_MAXITER_FINAL')

Expand Down Expand Up @@ -457,7 +462,7 @@ def scf_iterate(self, e_conv=None, d_conv=None):

if incfock_performed:
status.append("INCFOCK")

# Reset occupations if necessary
if (self.iteration_ == 0) and self.reset_occ_:
self.reset_occupation()
Expand Down Expand Up @@ -513,13 +518,16 @@ def scf_iterate(self, e_conv=None, d_conv=None):

if early_screening:

# we've reached convergence with early screning enabled; disable it on the JK object
# we've reached convergence with early screning enabled; disable it
early_screening = False
self.jk().set_early_screening(early_screening)

# make note of the change to early screening; next SCF iteration will be the last
# make note of the change to early screening; next SCF iteration(s) will be the last
early_screening_disabled = True

# cosx uses the largest grid for its final SCF iteration(s)
if cosx_enabled:
self.jk().set_COSX_grid("Final")

# clear any cached matrices associated with incremental fock construction
# the change in the screening spoils the linearity in the density matrix
if hasattr(self.jk(), 'clear_D_prev'):
Expand Down
6 changes: 3 additions & 3 deletions psi4/src/export_fock.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,6 @@ void export_fock(py::module &m) {
.def("get_omega_alpha", &JK::get_omega_alpha, "Weight for HF exchange term in range-separated DFT")
.def("set_omega_beta", &JK::set_omega_beta, "Weight for dampened exchange term in range-separated DFT", "beta"_a)
.def("get_omega_beta", &JK::get_omega_beta, "Weight for dampened exchange term in range-separated DFT")
.def("set_early_screening", &JK::set_early_screening, "Use severe screening techniques? Useful in early SCF iterations.", "early_screening"_a)
.def("get_early_screening", &JK::get_early_screening, "Use severe screening techniques? Useful in early SCF iterations.")
.def("compute", &JK::compute)
.def("finalize", &JK::finalize)
.def("C_clear",
Expand Down Expand Up @@ -200,7 +198,9 @@ void export_fock(py::module &m) {

py::class_<CompositeJK, std::shared_ptr<CompositeJK>, JK>(m, "CompositeJK", "docstring")
.def("do_incfock_iter", &CompositeJK::do_incfock_iter, "Was the last Fock build incremental?")
.def("clear_D_prev", &CompositeJK::clear_D_prev, "Clear previous D matrices.");
.def("clear_D_prev", &CompositeJK::clear_D_prev, "Clear previous D matrices.")
.def("set_COSX_grid", &CompositeJK::set_COSX_grid, "Set grid to use for COSX for this SCF iteration.")
.def("get_COSX_grid", &CompositeJK::get_COSX_grid, "Return grid used for COSX for this SCF iteration.");

py::class_<scf::SADGuess, std::shared_ptr<scf::SADGuess>>(m, "SADGuess", "docstring")
.def_static("build_SAD",
Expand Down
153 changes: 65 additions & 88 deletions psi4/src/psi4/libfock/COSK.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,6 @@ COSK::COSK(std::shared_ptr<BasisSet> primary, Options& options) : SplitJK(primar
#endif

// set options
early_screening_ = false;
lr_symmetric_ = true;

kscreen_ = options.get_double("COSX_INTS_TOLERANCE");
Expand All @@ -185,85 +184,63 @@ COSK::COSK(std::shared_ptr<BasisSet> primary, Options& options) : SplitJK(primar

timer_on("COSK: COSX Grid Construction");

// TODO: specify bool "DFT_REMOVE_DISTANT_POINTS" in the DFTGrid constructors

// Create a small DFTGrid for the initial SCF iterations
std::map<std::string, std::string> grid_init_str_options = {
{"DFT_PRUNING_SCHEME", options.get_str("COSX_PRUNING_SCHEME")},
{"DFT_RADIAL_SCHEME", "TREUTLER"},
{"DFT_NUCLEAR_SCHEME", "TREUTLER"},
{"DFT_GRID_NAME", ""},
{"DFT_BLOCK_SCHEME", "OCTREE"},
};
std::map<std::string, int> grid_init_int_options = {
{"DFT_SPHERICAL_POINTS", options.get_int("COSX_SPHERICAL_POINTS_INITIAL")},
{"DFT_RADIAL_POINTS", options.get_int("COSX_RADIAL_POINTS_INITIAL")},
{"DFT_BLOCK_MIN_POINTS", 100},
{"DFT_BLOCK_MAX_POINTS", 256},
};
std::map<std::string, double> grid_init_float_options = {
{"DFT_BASIS_TOLERANCE", basis_tol_},
{"DFT_BS_RADIUS_ALPHA", 1.0},
{"DFT_PRUNING_ALPHA", 1.0},
{"DFT_BLOCK_MAX_RADIUS", 3.0},
{"DFT_WEIGHTS_TOLERANCE", 1e-15},
};
grid_init_ = std::make_shared<DFTGrid>(primary_->molecule(), primary_, grid_init_int_options, grid_init_str_options, grid_init_float_options, options);

// Create a large DFTGrid for the final SCF iteration
std::map<std::string, std::string> grid_final_str_options = {
{"DFT_PRUNING_SCHEME", options.get_str("COSX_PRUNING_SCHEME")},
{"DFT_RADIAL_SCHEME", "TREUTLER"},
{"DFT_NUCLEAR_SCHEME", "TREUTLER"},
{"DFT_GRID_NAME", ""},
{"DFT_BLOCK_SCHEME", "OCTREE"},
};
std::map<std::string, int> grid_final_int_options = {
{"DFT_SPHERICAL_POINTS", options.get_int("COSX_SPHERICAL_POINTS_FINAL")},
{"DFT_RADIAL_POINTS", options.get_int("COSX_RADIAL_POINTS_FINAL")},
{"DFT_BLOCK_MIN_POINTS", 100},
{"DFT_BLOCK_MAX_POINTS", 256},
};
std::map<std::string, double> grid_final_float_options = {
{"DFT_BASIS_TOLERANCE", basis_tol_},
{"DFT_BS_RADIUS_ALPHA", 1.0},
{"DFT_PRUNING_ALPHA", 1.0},
{"DFT_BLOCK_MAX_RADIUS", 3.0},
{"DFT_WEIGHTS_TOLERANCE", 1e-15},
};
grid_final_ = std::make_shared<DFTGrid>(primary_->molecule(), primary_, grid_final_int_options, grid_final_str_options, grid_final_float_options, options);

// Print out warning if grid with negative grid weights is used
// Original Nesse COSX formulation does not support negative grid weights
// which can happen with certain grid configurations
// the Psi4 COSX implementation is slightly modified to work with negative grid weights
// See https://github.com/psi4/psi4/issues/2890 for discussion
auto warning_printed_init = false;
for (const auto &init_block : grid_init_->blocks()) {
const auto w = init_block->w();
for (int ipoint = 0; ipoint < init_block->npoints(); ++ipoint) {
if (w[ipoint] < 0.0) {
outfile->Printf(" INFO: The definition of the current initial grid includes negative weights, which the original COSX formulation does not support!\n If this is of concern, please choose another initial grid through adjusting either COSX_PRUNING_SCHEME or COSX_SPHERICAL_POINTS_INITIAL.\n\n");
warning_printed_init = true;
break;
}
}
if (warning_printed_init) break;
}
// for now, we use two COSX grids:
// - a small DFTGrid for the pre-converged SCF iterations
// - a large DFTGrid for the final SCF iteration
grids_["Initial"] = nullptr;
grids_["Final"] = nullptr;

for (auto& [ gridname, grid ] : grids_) {
std::string gridname_uppercase = gridname;
std::transform(gridname.begin(), gridname.end(), gridname_uppercase.begin(), ::toupper);

std::string gridname_lowercase = gridname;
std::transform(gridname.begin(), gridname.end(), gridname_lowercase.begin(), ::tolower);

// initialize grid
// TODO: specify bool "DFT_REMOVE_DISTANT_POINTS" in the DFTGrid constructors
std::map<std::string, std::string> grid_str_options = {
{"DFT_PRUNING_SCHEME", options_.get_str("COSX_PRUNING_SCHEME")},
{"DFT_RADIAL_SCHEME", "TREUTLER"},
{"DFT_NUCLEAR_SCHEME", "TREUTLER"},
{"DFT_GRID_NAME", ""},
{"DFT_BLOCK_SCHEME", "OCTREE"},
};

std::map<std::string, int> grid_int_options = {
{"DFT_SPHERICAL_POINTS", options_.get_int("COSX_SPHERICAL_POINTS_" + gridname_uppercase)},
{"DFT_RADIAL_POINTS", options_.get_int("COSX_RADIAL_POINTS_" + gridname_uppercase)},
{"DFT_BLOCK_MIN_POINTS", 100},
{"DFT_BLOCK_MAX_POINTS", 256},
};

auto warning_printed_final = false;
for (const auto &final_block : grid_final_->blocks()) {
const auto w = final_block->w();
for (int ipoint = 0; ipoint < final_block->npoints(); ++ipoint) {
if (w[ipoint] < 0.0) {
outfile->Printf(" INFO: The definition of the current final grid includes negative weights, which the original COSX formulation does not support!\n If this is of concern, please choose another final grid through adjusting either COSX_PRUNING_SCHEME or COSX_SPHERICAL_POINTS_FINAL.\n\n");
warning_printed_final = true;
break;
}
std::map<std::string, double> grid_float_options = {
{"DFT_BASIS_TOLERANCE", options_.get_double("COSX_BASIS_TOLERANCE")},
{"DFT_BS_RADIUS_ALPHA", 1.0},
{"DFT_PRUNING_ALPHA", 1.0},
{"DFT_BLOCK_MAX_RADIUS", 3.0},
{"DFT_WEIGHTS_TOLERANCE", 1e-15},
};
grid = std::make_shared<DFTGrid>(primary_->molecule(), primary_, grid_int_options, grid_str_options, grid_float_options, options_);

// Print out specific grid info upon request
if (options_.get_int("DEBUG")) {
outfile->Printf(" ==> COSX: ");
outfile->Printf(gridname);
outfile->Printf(" Grid Details <==\n\n");

auto npoints = grid->npoints();
auto nblocks = grid->blocks().size();
auto natoms = primary_->molecule()->natom();
double npoints_per_batch = static_cast<double>(npoints) / static_cast<double>(nblocks);
double npoints_per_atom = static_cast<double>(npoints) / static_cast<double>(natoms);

outfile->Printf(" Total number of grid points: %d \n", npoints);
outfile->Printf(" Total number of batches: %d \n", nblocks);
outfile->Printf(" Average number of points per batch: %f \n", npoints_per_batch);
outfile->Printf(" Average number of grid points per atom: %f \n\n", npoints_per_atom);
}
if (warning_printed_final) break;
}

timer_off("COSK: COSX Grid Construction");

// => Overlap Fitting Metric <= //
Expand All @@ -278,8 +255,10 @@ COSK::COSK(std::shared_ptr<BasisSet> primary, Options& options) : SplitJK(primar
timer_on("COSK: COSX Numeric Overlap");

// compute the numeric overlap matrix for each grid
auto S_num_init = compute_numeric_overlap(*grid_init_, primary_);
auto S_num_final = compute_numeric_overlap(*grid_final_, primary_ );
std::unordered_map<std::string, Matrix> S_num;
for (auto& [ gridname, grid ] : grids_) {
S_num[gridname] = compute_numeric_overlap(*grid, primary_);
}

timer_off("COSK: COSX Numeric Overlap");

Expand All @@ -298,13 +277,11 @@ COSK::COSK(std::shared_ptr<BasisSet> primary, Options& options) : SplitJK(primar
int nbf = primary_->nbf();
std::vector<int> ipiv(nbf);

// solve: Q_init_ = S_an @ S_num_init_^{-1}
Q_init_ = S_an->clone();
C_DGESV(nbf, nbf, S_num_init.pointer()[0], nbf, ipiv.data(), Q_init_->pointer()[0], nbf);

// solve: Q_final_ = S_an @ S_num_final_^{-1}
Q_final_ = S_an->clone();
C_DGESV(nbf, nbf, S_num_final.pointer()[0], nbf, ipiv.data(), Q_final_->pointer()[0], nbf);
// solve: Q_mat_ = S_an @ S_num_^{-1} for each grid
for (auto& [ gridname, grid ] : grids_) {
Q_mat_[gridname] = S_an->clone();
C_DGESV(nbf, nbf, S_num[gridname].pointer()[0], nbf, ipiv.data(), Q_mat_[gridname]->pointer()[0], nbf);
}

timer_off("COSK: COSX Overlap Metric Solve");

Expand Down Expand Up @@ -343,8 +320,8 @@ void COSK::build_G_component(std::vector<std::shared_ptr<Matrix>>& D, std::vecto

// use a small DFTGrid grid (and overlap metric) for early SCF iterations
// otherwise use a large DFTGrid
auto grid = early_screening_ ? grid_init_ : grid_final_;
auto Q = early_screening_ ? Q_init_ : Q_final_;
auto grid = grids_[current_grid_];
auto Q = Q_mat_[current_grid_];

// => Initialization <= //

Expand Down
26 changes: 17 additions & 9 deletions psi4/src/psi4/libfock/CompositeJK.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
#include <vector>
#include <map>
#include <algorithm>
#include <cctype>
#ifdef _OPENMP
#include <omp.h>
#endif
Expand Down Expand Up @@ -93,11 +94,10 @@ void CompositeJK::common_init() {
// other options
density_screening_ = options_.get_str("SCREENING") == "DENSITY";
set_cutoff(options_.get_double("INTS_TOLERANCE"));
early_screening_ = k_type == "COSX" ? true : false;

// pre-construct per-thread TwoBodyAOInt objects for computing 3- and 4-index ERIs
timer_on("CompositeJK: ERI Computers");

auto zero = BasisSet::zero_ao_basis_set();

// initialize 4-Center ERIs
Expand Down Expand Up @@ -249,11 +249,9 @@ void CompositeJK::compute_JK() {
if (do_wK_) throw PSIEXCEPTION("CompositeJK algorithms do not support wK integrals yet!");

// set compute()-specific parameters
j_algo_->set_early_screening(early_screening_);
j_algo_->set_lr_symmetric(lr_symmetric_);

if (do_K_) {
k_algo_->set_early_screening(early_screening_);
k_algo_->set_lr_symmetric(lr_symmetric_);
}

Expand Down Expand Up @@ -288,28 +286,38 @@ void CompositeJK::compute_JK() {

// Coulomb Matrix
if (do_J_) {
timer_on("CompositeJK: J");
timer_on("CompositeJK: " + j_algo_->name());

j_algo_->build_G_component(D_ref_, J_ao_, eri_computers_["3-Center"]);

if (get_bench()) {
computed_shells_per_iter_["Triplets"].push_back(j_algo_->num_computed_shells());
}

timer_off("CompositeJK: J");
timer_off("CompositeJK: " + j_algo_->name());
}

// Exchange Matrix
if (do_K_) {
timer_on("CompositeJK: K");
timer_on("CompositeJK: " + k_algo_->name());

if (k_algo_->name() == "COSX") {
std::string gridname = k_algo_->get_COSX_grid();
timer_on("COSX " + gridname + " Grid");
}

k_algo_->build_G_component(D_ref_, K_ao_, eri_computers_["4-Center"]);

if (get_bench()) {
computed_shells_per_iter_["Quartets"].push_back(k_algo_->num_computed_shells());
}

timer_off("CompositeJK: K");
if (k_algo_->name() == "COSX") {
std::string gridname = k_algo_->get_COSX_grid();
timer_off("COSX " + gridname + " Grid");
}

timer_off("CompositeJK: " + k_algo_->name());
}

// => Finalize Incremental Fock if required <= //
Expand Down
Loading

0 comments on commit 289d095

Please sign in to comment.