diff --git a/python/pyabacus/src/ModuleBase/py_base_math.cpp b/python/pyabacus/src/ModuleBase/py_base_math.cpp index 89791d1ab37..df1cbea86f5 100644 --- a/python/pyabacus/src/ModuleBase/py_base_math.cpp +++ b/python/pyabacus/src/ModuleBase/py_base_math.cpp @@ -1,11 +1,10 @@ -#include -#include - -#include "source_base/math_sphbes.h" +#include "../utils/pybind_utils.h" #include "source_base/math_integral.h" +#include "source_base/math_sphbes.h" #include "source_base/sph_bessel_tf.h" -#include "../utils/pybind_utils.h" +#include +#include namespace py = pybind11; using namespace pybind11::literals; @@ -25,21 +24,13 @@ void bind_base_math(py::module& m) [](const int n, py::array_t r, const double q, const int l, py::array_t jl) { check_1d_array(r, "r"); check_1d_array(jl, "jl"); - ModuleBase::Sphbes::sphbesj(n, - get_array_ptr(r), - q, - l, - get_array_ptr(jl)); + ModuleBase::Sphbes::sphbesj(n, get_array_ptr(r), q, l, get_array_ptr(jl)); }) .def_static("dsphbesj", [](const int n, py::array_t r, const double q, const int l, py::array_t djl) { check_1d_array(r, "r"); check_1d_array(djl, "djl"); - ModuleBase::Sphbes::dsphbesj(n, - get_array_ptr(r), - q, - l, - get_array_ptr(djl)); + ModuleBase::Sphbes::dsphbesj(n, get_array_ptr(r), q, l, get_array_ptr(djl)); }) .def_static("sphbes_zeros", [](const int l, const int n, py::array_t zeros) { check_1d_array(zeros, "zeros"); @@ -49,76 +40,68 @@ void bind_base_math(py::module& m) // python binding for class Integral py::class_(m, "Integral") .def(py::init<>()) - .def_static("Simpson_Integral", [](const int mesh, py::array_t func, py::array_t rab, double asum) { - check_1d_array(func, "func"); - check_1d_array(rab, "rab"); + .def_static("Simpson_Integral", + [](const int mesh, py::array_t func, py::array_t rab, double asum) { + check_1d_array(func, "func"); + check_1d_array(rab, "rab"); - double isum = asum; - ModuleBase::Integral::Simpson_Integral(mesh, - get_array_ptr(func), - get_array_ptr(rab), - isum); - return isum; - }) - .def_static("Simpson_Integral", [](const int mesh, py::array_t func, const double dr, double asum){ - check_1d_array(func, "func"); + double isum = asum; + ModuleBase::Integral::Simpson_Integral(mesh, get_array_ptr(func), get_array_ptr(rab), isum); + return isum; + }) + .def_static("Simpson_Integral", + [](const int mesh, py::array_t func, const double dr, double asum) { + check_1d_array(func, "func"); - double isum = asum; - ModuleBase::Integral::Simpson_Integral(mesh, - get_array_ptr(func), - dr, - isum); - return isum; - }) - .def_static("Simpson_Integral_0toall", [](const int mesh, py::array_t func, py::array_t rab, py::array_t asum){ - check_1d_array(func, "func"); - check_1d_array(rab, "rab"); - check_1d_array(asum, "asum"); - ModuleBase::Integral::Simpson_Integral_0toall(mesh, - get_array_ptr(func), - get_array_ptr(rab), - get_array_ptr(asum)); - }) - .def_static("Simpson_Integral_alltoinf", [](const int mesh, py::array_t func, py::array_t rab, py::array_t asum){ - check_1d_array(func, "func"); - check_1d_array(rab, "rab"); - check_1d_array(asum, "asum"); - ModuleBase::Integral::Simpson_Integral_alltoinf(mesh, - get_array_ptr(func), - get_array_ptr(rab), - get_array_ptr(asum)); - }) - .def_static("simpson", [](const int n, py::array_t f, const double dx){ - check_1d_array(f, "f"); - return ModuleBase::Integral::simpson(n, - get_array_ptr(f), - dx); - }) - .def_static("simpson", [](const int n, py::array_t f, py::array_t h){ - check_1d_array(f, "f"); - check_1d_array(h, "h"); - return ModuleBase::Integral::simpson(n, - get_array_ptr(f), - get_array_ptr(h)); - }) - .def_static("Gauss_Legendre_grid_and_weight", [](const int n, py::array_t x, py::array_t w){ - check_1d_array(x, "x"); - check_1d_array(w, "w"); - ModuleBase::Integral::Gauss_Legendre_grid_and_weight(n, - get_array_ptr(x), - get_array_ptr(w)); - }) - .def_static("Gauss_Legendre_grid_and_weight", [](const double xmin, const double xmax, const int n, py::array_t x, py::array_t w){ - check_1d_array(x, "x"); - check_1d_array(w, "w"); - ModuleBase::Integral::Gauss_Legendre_grid_and_weight(xmin, - xmax, - n, - get_array_ptr(x), - get_array_ptr(w)); - }); - py::class_(m, "SphericalBesselTransformer") - .def(py::init<>()); + double isum = asum; + ModuleBase::Integral::Simpson_Integral(mesh, get_array_ptr(func), dr, isum); + return isum; + }) + .def_static("Simpson_Integral_0toall", + [](const int mesh, py::array_t func, py::array_t rab, py::array_t asum) { + check_1d_array(func, "func"); + check_1d_array(rab, "rab"); + check_1d_array(asum, "asum"); + ModuleBase::Integral::Simpson_Integral_0toall(mesh, + get_array_ptr(func), + get_array_ptr(rab), + get_array_ptr(asum)); + }) + .def_static("Simpson_Integral_alltoinf", + [](const int mesh, py::array_t func, py::array_t rab, py::array_t asum) { + check_1d_array(func, "func"); + check_1d_array(rab, "rab"); + check_1d_array(asum, "asum"); + ModuleBase::Integral::Simpson_Integral_alltoinf(mesh, + get_array_ptr(func), + get_array_ptr(rab), + get_array_ptr(asum)); + }) + .def_static("simpson", + [](const int n, py::array_t f, const double dx) { + check_1d_array(f, "f"); + return ModuleBase::Integral::simpson(n, get_array_ptr(f), dx); + }) + .def_static("simpson", + [](const int n, py::array_t f, py::array_t h) { + check_1d_array(f, "f"); + check_1d_array(h, "h"); + return ModuleBase::Integral::simpson(n, get_array_ptr(f), get_array_ptr(h)); + }) + .def_static("Gauss_Legendre_grid_and_weight", + [](const int n, py::array_t x, py::array_t w) { + check_1d_array(x, "x"); + check_1d_array(w, "w"); + ModuleBase::Integral::Gauss_Legendre_grid_and_weight(n, get_array_ptr(x), get_array_ptr(w)); + }) + .def_static( + "Gauss_Legendre_grid_and_weight", + [](const double xmin, const double xmax, const int n, py::array_t x, py::array_t w) { + check_1d_array(x, "x"); + check_1d_array(w, "w"); + ModuleBase::Integral::Gauss_Legendre_grid_and_weight(xmin, xmax, n, get_array_ptr(x), get_array_ptr(w)); + }); + py::class_(m, "SphericalBesselTransformer").def(py::init<>()); } PYBIND11_MODULE(_base_pack, m) diff --git a/python/pyabacus/src/ModuleDriver/py_driver.cpp b/python/pyabacus/src/ModuleDriver/py_driver.cpp index afd62caf425..26b02bf49f8 100644 --- a/python/pyabacus/src/ModuleDriver/py_driver.cpp +++ b/python/pyabacus/src/ModuleDriver/py_driver.cpp @@ -8,31 +8,31 @@ #include "py_driver.hpp" -#include #include +#include #include // ABACUS headers -#include "source_main/driver.h" -#include "source_cell/unitcell.h" +#include "source_base/global_file.h" +#include "source_base/global_variable.h" +#include "source_base/matrix.h" +#include "source_base/memory.h" +#include "source_base/timer.h" #include "source_cell/check_atomic_stru.h" +#include "source_cell/unitcell.h" #include "source_esolver/esolver.h" -#include "source_io/read_input.h" #include "source_io/input_conv.h" #include "source_io/module_parameter/parameter.h" -#include "source_base/global_variable.h" -#include "source_base/global_file.h" -#include "source_base/timer.h" -#include "source_base/memory.h" -#include "source_base/matrix.h" +#include "source_io/read_input.h" +#include "source_main/driver.h" #include "source_relax/relax_driver.h" +#include +#include #include -#include #include -#include +#include #include -#include namespace py = pybind11; namespace fs = std::filesystem; @@ -48,7 +48,7 @@ namespace py_driver */ class GlobalStateGuard { -public: + public: GlobalStateGuard() { // Save current state @@ -63,7 +63,7 @@ class GlobalStateGuard GlobalV::NPROC = saved_nproc_; } -private: + private: int saved_my_rank_ = 0; int saved_nproc_ = 1; }; @@ -73,9 +73,12 @@ class GlobalStateGuard */ class PyDriver::Impl { -public: + public: Impl() = default; - ~Impl() { cleanup(); } + ~Impl() + { + cleanup(); + } void cleanup() { @@ -120,7 +123,7 @@ PyDriver::~PyDriver() void PyDriver::initialize_context() { // Set up for serial mode (no MPI in Python context) - PARAM.set_pal_param(0, 1, 1); // rank=0, nproc=1, nthread=1 + PARAM.set_pal_param(0, 1, 1); // rank=0, nproc=1, nthread=1 GlobalV::MY_RANK = 0; GlobalV::NPROC = 1; @@ -218,14 +221,13 @@ void PyDriver::setup_output(const std::string& output_dir, int verbosity) } } -void PyDriver::read_input( - const std::string& input_dir, - const std::string& input_file, - const std::string& stru_file, - const std::string& kpt_file, - const std::string& pseudo_dir, - const std::string& orbital_dir, - const std::string& output_dir) +void PyDriver::read_input(const std::string& input_dir, + const std::string& input_file, + const std::string& stru_file, + const std::string& kpt_file, + const std::string& pseudo_dir, + const std::string& orbital_dir, + const std::string& output_dir) { // Save original working directory impl_->original_cwd_ = fs::current_path().string(); @@ -260,7 +262,7 @@ void PyDriver::read_input( // Read INPUT file // Note: ReadInput will set PARAM.globalv.global_in_card internally - ModuleIO::ReadInput reader(0); // rank 0 + ModuleIO::ReadInput reader(0); // rank 0 std::string input_filename = fs::path(input_path).filename().string(); reader.read_parameters(PARAM, input_filename); @@ -341,15 +343,10 @@ CalculationResult PyDriver::collect_results(bool calculate_force, bool calculate if (!result.output_dir.empty() && fs::exists(result.output_dir)) { // Look for running_*.log files - std::vector log_patterns = { - "running_scf.log", - "running_relax.log", - "running_cell-relax.log", - "running_nscf.log", - "running_md.log" - }; - - for (const auto& log_name : log_patterns) + std::vector log_patterns + = {"running_scf.log", "running_relax.log", "running_cell-relax.log", "running_nscf.log", "running_md.log"}; + + for (const auto& log_name: log_patterns) { std::string log_path = result.output_dir + "/" + log_name; if (fs::exists(log_path)) @@ -362,7 +359,7 @@ CalculationResult PyDriver::collect_results(bool calculate_force, bool calculate // Iterate directory to populate output_files map try { - for (const auto& entry : fs::directory_iterator(result.output_dir)) + for (const auto& entry: fs::directory_iterator(result.output_dir)) { if (entry.is_regular_file()) { @@ -381,17 +378,16 @@ CalculationResult PyDriver::collect_results(bool calculate_force, bool calculate return result; } -CalculationResult PyDriver::run( - const std::string& input_dir, - const std::string& input_file, - const std::string& stru_file, - const std::string& kpt_file, - const std::string& pseudo_dir, - const std::string& orbital_dir, - const std::string& output_dir, - bool calculate_force, - bool calculate_stress, - int verbosity) +CalculationResult PyDriver::run(const std::string& input_dir, + const std::string& input_file, + const std::string& stru_file, + const std::string& kpt_file, + const std::string& pseudo_dir, + const std::string& orbital_dir, + const std::string& output_dir, + bool calculate_force, + bool calculate_stress, + int verbosity) { // Use RAII guard for global state GlobalStateGuard state_guard; @@ -411,25 +407,34 @@ CalculationResult PyDriver::run( try { // Read input files - read_input(input_dir, input_file, stru_file, kpt_file, - pseudo_dir, orbital_dir, output_dir); + read_input(input_dir, input_file, stru_file, kpt_file, pseudo_dir, orbital_dir, output_dir); // Create UnitCell impl_->ucell_ = std::make_unique(); - impl_->ucell_->setup( - PARAM.inp.latname, - PARAM.inp.ntype, - PARAM.inp.lmaxmax, - PARAM.inp.init_vel, - PARAM.inp.fixed_axes - ); + impl_->ucell_->setup(PARAM.inp.latname, + PARAM.inp.ntype, + PARAM.inp.lmaxmax, + PARAM.inp.init_vel, + PARAM.inp.fixed_axes); // Read structure - impl_->ucell_->setup_cell(PARAM.globalv.global_in_stru, GlobalV::ofs_running, PARAM.inp.symmetry_prec, PARAM.inp.dfthalf_type, PARAM.inp.pseudo_dir, PARAM.inp.nspin, - PARAM.inp.basis_type, PARAM.inp.orbital_dir, PARAM.inp.init_wfc, - PARAM.inp.onsite_radius, PARAM.globalv.deepks_setorb, PARAM.inp.rpa, - PARAM.inp.fixed_atoms, PARAM.inp.noncolin, PARAM.inp.calculation, PARAM.inp.esolver_type, - std::stoi(PARAM.inp.symmetry)); + impl_->ucell_->setup_cell(PARAM.globalv.global_in_stru, + GlobalV::ofs_running, + PARAM.inp.symmetry_prec, + PARAM.inp.dfthalf_type, + PARAM.inp.pseudo_dir, + PARAM.inp.nspin, + PARAM.inp.basis_type, + PARAM.inp.orbital_dir, + PARAM.inp.init_wfc, + PARAM.inp.onsite_radius, + PARAM.globalv.deepks_setorb, + PARAM.inp.rpa, + PARAM.inp.fixed_atoms, + PARAM.inp.noncolin, + PARAM.inp.calculation, + PARAM.inp.esolver_type, + std::stoi(PARAM.inp.symmetry)); // Check atomic structure unitcell::check_atomic_stru(*impl_->ucell_, PARAM.inp.min_dist_coef); @@ -511,8 +516,9 @@ PYBIND11_MODULE(_driver_pack, m) )pbdoc"; // Bind CalculationResult - py::class_(m, "CalculationResult", - R"pbdoc( + py::class_(m, + "CalculationResult", + R"pbdoc( Container for DFT calculation results. Attributes @@ -531,72 +537,46 @@ PYBIND11_MODULE(_driver_pack, m) Stress tensor (3, 3) in kbar )pbdoc") .def(py::init<>()) - .def_readonly("converged", &py_driver::CalculationResult::converged, - "Whether SCF converged") - .def_readonly("niter", &py_driver::CalculationResult::niter, - "Number of SCF iterations") - .def_readonly("drho", &py_driver::CalculationResult::drho, - "Final charge density difference") - .def_readonly("etot", &py_driver::CalculationResult::etot, - "Total energy (Ry)") - .def_readonly("eband", &py_driver::CalculationResult::eband, - "Band energy (Ry)") - .def_readonly("hartree_energy", &py_driver::CalculationResult::hartree_energy, - "Hartree energy (Ry)") - .def_readonly("etxc", &py_driver::CalculationResult::etxc, - "Exchange-correlation energy (Ry)") - .def_readonly("ewald_energy", &py_driver::CalculationResult::ewald_energy, - "Ewald energy (Ry)") - .def_readonly("demet", &py_driver::CalculationResult::demet, - "-TS term for metals (Ry)") - .def_readonly("exx", &py_driver::CalculationResult::exx, - "Exact exchange energy (Ry)") - .def_readonly("evdw", &py_driver::CalculationResult::evdw, - "van der Waals energy (Ry)") - .def_readonly("forces", &py_driver::CalculationResult::forces, - "Forces on atoms (nat, 3) in Ry/Bohr") - .def_readonly("has_forces", &py_driver::CalculationResult::has_forces, - "Whether forces are available") - .def_readonly("stress", &py_driver::CalculationResult::stress, - "Stress tensor (3, 3) in kbar") - .def_readonly("has_stress", &py_driver::CalculationResult::has_stress, - "Whether stress is available") - .def_readonly("fermi_energy", &py_driver::CalculationResult::fermi_energy, - "Fermi energy (eV)") - .def_readonly("bandgap", &py_driver::CalculationResult::bandgap, - "Band gap (eV)") - .def_readonly("nat", &py_driver::CalculationResult::nat, - "Number of atoms") - .def_readonly("ntype", &py_driver::CalculationResult::ntype, - "Number of atom types") - .def_readonly("nbands", &py_driver::CalculationResult::nbands, - "Number of bands") - .def_readonly("nks", &py_driver::CalculationResult::nks, - "Number of k-points") - .def_readonly("output_dir", &py_driver::CalculationResult::output_dir, - "Path to output directory (OUT.$suffix)") - .def_readonly("log_file", &py_driver::CalculationResult::log_file, - "Path to the main log file") - .def_readonly("output_files", &py_driver::CalculationResult::output_files, - "Dictionary of output files (filename -> full path)") - .def("etot_eV", &py_driver::CalculationResult::etot_eV, - "Get total energy in eV") - .def("get_energies", &py_driver::CalculationResult::get_energies, - "Get all energies as a dictionary") - .def("get_forces_eV_Ang", &py_driver::CalculationResult::get_forces_eV_Ang, - "Get forces in eV/Angstrom") - .def("summary", &py_driver::CalculationResult::summary, - "Get a summary string of the calculation result") + .def_readonly("converged", &py_driver::CalculationResult::converged, "Whether SCF converged") + .def_readonly("niter", &py_driver::CalculationResult::niter, "Number of SCF iterations") + .def_readonly("drho", &py_driver::CalculationResult::drho, "Final charge density difference") + .def_readonly("etot", &py_driver::CalculationResult::etot, "Total energy (Ry)") + .def_readonly("eband", &py_driver::CalculationResult::eband, "Band energy (Ry)") + .def_readonly("hartree_energy", &py_driver::CalculationResult::hartree_energy, "Hartree energy (Ry)") + .def_readonly("etxc", &py_driver::CalculationResult::etxc, "Exchange-correlation energy (Ry)") + .def_readonly("ewald_energy", &py_driver::CalculationResult::ewald_energy, "Ewald energy (Ry)") + .def_readonly("demet", &py_driver::CalculationResult::demet, "-TS term for metals (Ry)") + .def_readonly("exx", &py_driver::CalculationResult::exx, "Exact exchange energy (Ry)") + .def_readonly("evdw", &py_driver::CalculationResult::evdw, "van der Waals energy (Ry)") + .def_readonly("forces", &py_driver::CalculationResult::forces, "Forces on atoms (nat, 3) in Ry/Bohr") + .def_readonly("has_forces", &py_driver::CalculationResult::has_forces, "Whether forces are available") + .def_readonly("stress", &py_driver::CalculationResult::stress, "Stress tensor (3, 3) in kbar") + .def_readonly("has_stress", &py_driver::CalculationResult::has_stress, "Whether stress is available") + .def_readonly("fermi_energy", &py_driver::CalculationResult::fermi_energy, "Fermi energy (eV)") + .def_readonly("bandgap", &py_driver::CalculationResult::bandgap, "Band gap (eV)") + .def_readonly("nat", &py_driver::CalculationResult::nat, "Number of atoms") + .def_readonly("ntype", &py_driver::CalculationResult::ntype, "Number of atom types") + .def_readonly("nbands", &py_driver::CalculationResult::nbands, "Number of bands") + .def_readonly("nks", &py_driver::CalculationResult::nks, "Number of k-points") + .def_readonly("output_dir", &py_driver::CalculationResult::output_dir, "Path to output directory (OUT.$suffix)") + .def_readonly("log_file", &py_driver::CalculationResult::log_file, "Path to the main log file") + .def_readonly("output_files", + &py_driver::CalculationResult::output_files, + "Dictionary of output files (filename -> full path)") + .def("etot_eV", &py_driver::CalculationResult::etot_eV, "Get total energy in eV") + .def("get_energies", &py_driver::CalculationResult::get_energies, "Get all energies as a dictionary") + .def("get_forces_eV_Ang", &py_driver::CalculationResult::get_forces_eV_Ang, "Get forces in eV/Angstrom") + .def("summary", &py_driver::CalculationResult::summary, "Get a summary string of the calculation result") .def("__repr__", [](const py_driver::CalculationResult& r) { std::ostringstream ss; - ss << ""; + ss << ""; return ss.str(); }); // Bind PyDriver - py::class_(m, "PyDriver", - R"pbdoc( + py::class_(m, + "PyDriver", + R"pbdoc( Python wrapper for ABACUS Driver. This class provides a Python interface for running complete ABACUS @@ -614,8 +594,9 @@ PYBIND11_MODULE(_driver_pack, m) >>> print(f"Converged: {result.converged}") )pbdoc") .def(py::init<>()) - .def("run", &py_driver::PyDriver::run, - R"pbdoc( + .def("run", + &py_driver::PyDriver::run, + R"pbdoc( Run a complete DFT calculation. Parameters @@ -646,19 +627,19 @@ PYBIND11_MODULE(_driver_pack, m) CalculationResult Container with all calculation results )pbdoc", - py::arg("input_dir") = ".", - py::arg("input_file") = "", - py::arg("stru_file") = "", - py::arg("kpt_file") = "", - py::arg("pseudo_dir") = "", - py::arg("orbital_dir") = "", - py::arg("output_dir") = "", - py::arg("calculate_force") = true, - py::arg("calculate_stress") = false, - py::arg("verbosity") = 1) - .def("is_ready", &py_driver::PyDriver::is_ready, - "Check if the driver is ready for calculation") - .def("get_last_result", &py_driver::PyDriver::get_last_result, - py::return_value_policy::reference_internal, - "Get the last calculation result"); + py::arg("input_dir") = ".", + py::arg("input_file") = "", + py::arg("stru_file") = "", + py::arg("kpt_file") = "", + py::arg("pseudo_dir") = "", + py::arg("orbital_dir") = "", + py::arg("output_dir") = "", + py::arg("calculate_force") = true, + py::arg("calculate_stress") = false, + py::arg("verbosity") = 1) + .def("is_ready", &py_driver::PyDriver::is_ready, "Check if the driver is ready for calculation") + .def("get_last_result", + &py_driver::PyDriver::get_last_result, + py::return_value_policy::reference_internal, + "Get the last calculation result"); } diff --git a/python/pyabacus/src/ModuleDriver/py_driver.hpp b/python/pyabacus/src/ModuleDriver/py_driver.hpp index db426cee0b6..4df414d88cf 100644 --- a/python/pyabacus/src/ModuleDriver/py_driver.hpp +++ b/python/pyabacus/src/ModuleDriver/py_driver.hpp @@ -9,16 +9,15 @@ #ifndef PY_DRIVER_HPP #define PY_DRIVER_HPP -#include +#include +#include +#include #include +#include #include - +#include #include -#include -#include #include -#include -#include namespace py = pybind11; @@ -57,24 +56,27 @@ struct CalculationResult bool has_stress = false; // Electronic structure info - double fermi_energy = 0.0; // in eV - double bandgap = 0.0; // in eV + double fermi_energy = 0.0; // in eV + double bandgap = 0.0; // in eV int nat = 0; int ntype = 0; int nbands = 0; int nks = 0; // Output file tracking - std::string output_dir = ""; // Path to OUT.$suffix folder - std::string log_file = ""; // Path to the main log file - std::map output_files; // filename -> full path + std::string output_dir = ""; // Path to OUT.$suffix folder + std::string log_file = ""; // Path to the main log file + std::map output_files; // filename -> full path // Unit conversion constants static constexpr double Ry_to_eV = 13.605693122994; static constexpr double Bohr_to_Ang = 0.529177249; // Convenience methods - double etot_eV() const { return etot * Ry_to_eV; } + double etot_eV() const + { + return etot * Ry_to_eV; + } py::dict get_energies() const { @@ -177,7 +179,7 @@ struct CalculationResult */ class PyDriver { -public: + public: PyDriver(); ~PyDriver(); @@ -200,30 +202,34 @@ class PyDriver * @param verbosity Output verbosity level (0=silent, 1=normal, 2=verbose) * @return CalculationResult containing all results */ - CalculationResult run( - const std::string& input_dir = ".", - const std::string& input_file = "", - const std::string& stru_file = "", - const std::string& kpt_file = "", - const std::string& pseudo_dir = "", - const std::string& orbital_dir = "", - const std::string& output_dir = "", - bool calculate_force = true, - bool calculate_stress = false, - int verbosity = 1 - ); + CalculationResult run(const std::string& input_dir = ".", + const std::string& input_file = "", + const std::string& stru_file = "", + const std::string& kpt_file = "", + const std::string& pseudo_dir = "", + const std::string& orbital_dir = "", + const std::string& output_dir = "", + bool calculate_force = true, + bool calculate_stress = false, + int verbosity = 1); /** * @brief Check if the driver is ready for calculation */ - bool is_ready() const { return initialized_; } + bool is_ready() const + { + return initialized_; + } /** * @brief Get the last calculation result */ - const CalculationResult& get_last_result() const { return last_result_; } + const CalculationResult& get_last_result() const + { + return last_result_; + } -private: + private: class Impl; std::unique_ptr impl_; diff --git a/python/pyabacus/src/ModuleESolver/components/charge_mixer_wrapper.hpp b/python/pyabacus/src/ModuleESolver/components/charge_mixer_wrapper.hpp index 0fc299ddf6c..c64c45c86d2 100644 --- a/python/pyabacus/src/ModuleESolver/components/charge_mixer_wrapper.hpp +++ b/python/pyabacus/src/ModuleESolver/components/charge_mixer_wrapper.hpp @@ -8,14 +8,16 @@ #ifndef PYABACUS_ESOLVER_CHARGE_MIXER_WRAPPER_HPP #define PYABACUS_ESOLVER_CHARGE_MIXER_WRAPPER_HPP -#include "../interfaces/i_charge_mixer.hpp" #include "../../utils/pybind_utils.h" +#include "../interfaces/i_charge_mixer.hpp" #include #include -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief Wrapper for ABACUS Charge_Mixing class @@ -25,11 +27,10 @@ namespace esolver { */ class ChargeMixerWrapper : public IChargeMixer { -public: + public: ChargeMixerWrapper() = default; - ChargeMixerWrapper(int nspin, int nrxx) - : nspin_(nspin), nrxx_(nrxx) + ChargeMixerWrapper(int nspin, int nrxx) : nspin_(nspin), nrxx_(nrxx) { // Initialize history buffers for Pulay mixing rho_history_.reserve(config_.ndim); @@ -40,8 +41,7 @@ class ChargeMixerWrapper : public IChargeMixer // ==================== Core Mixing Operations ==================== - py::array_t mix(const py::array_t& rho_in, - const py::array_t& rho_out) override + py::array_t mix(const py::array_t& rho_in, const py::array_t& rho_out) override { using namespace pyabacus::utils; @@ -53,8 +53,7 @@ class ChargeMixerWrapper : public IChargeMixer const double* out_ptr = get_array_ptr(rho_out); // Create output array - py::array_t rho_mixed({static_cast(nspin_), - static_cast(nrxx_)}); + py::array_t rho_mixed({static_cast(nspin_), static_cast(nrxx_)}); double* mixed_ptr = get_array_ptr(rho_mixed); // Calculate drho @@ -69,20 +68,20 @@ class ChargeMixerWrapper : public IChargeMixer // Perform mixing based on method switch (config_.method) { - case MixingMethod::Plain: - mix_plain(in_ptr, out_ptr, mixed_ptr); - break; - case MixingMethod::Pulay: - mix_pulay(in_ptr, out_ptr, mixed_ptr); - break; - case MixingMethod::Broyden: - mix_broyden(in_ptr, out_ptr, mixed_ptr); - break; - case MixingMethod::Anderson: - mix_anderson(in_ptr, out_ptr, mixed_ptr); - break; - default: - mix_plain(in_ptr, out_ptr, mixed_ptr); + case MixingMethod::Plain: + mix_plain(in_ptr, out_ptr, mixed_ptr); + break; + case MixingMethod::Pulay: + mix_pulay(in_ptr, out_ptr, mixed_ptr); + break; + case MixingMethod::Broyden: + mix_broyden(in_ptr, out_ptr, mixed_ptr); + break; + case MixingMethod::Anderson: + mix_anderson(in_ptr, out_ptr, mixed_ptr); + break; + default: + mix_plain(in_ptr, out_ptr, mixed_ptr); } iteration_++; @@ -99,9 +98,15 @@ class ChargeMixerWrapper : public IChargeMixer // ==================== State Queries ==================== - double get_drho() const override { return drho_; } + double get_drho() const override + { + return drho_; + } - int get_iteration() const override { return iteration_; } + int get_iteration() const override + { + return iteration_; + } // ==================== Configuration ==================== @@ -111,7 +116,10 @@ class ChargeMixerWrapper : public IChargeMixer reset(); // Reset history when config changes } - MixingConfig get_config() const override { return config_; } + MixingConfig get_config() const override + { + return config_; + } void set_mixing_beta(double beta) override { @@ -122,7 +130,10 @@ class ChargeMixerWrapper : public IChargeMixer config_.beta = beta; } - double get_mixing_beta() const override { return config_.beta; } + double get_mixing_beta() const override + { + return config_.beta; + } void set_mixing_method(MixingMethod method) override { @@ -130,7 +141,10 @@ class ChargeMixerWrapper : public IChargeMixer reset(); } - MixingMethod get_mixing_method() const override { return config_.method; } + MixingMethod get_mixing_method() const override + { + return config_.method; + } // ==================== Dimension Setters ==================== @@ -141,7 +155,7 @@ class ChargeMixerWrapper : public IChargeMixer reset(); } -private: + private: // Plain linear mixing: rho_new = (1-beta)*rho_in + beta*rho_out void mix_plain(const double* rho_in, const double* rho_out, double* rho_mixed) { @@ -215,8 +229,7 @@ class ChargeMixerWrapper : public IChargeMixer { for (size_t k = 0; k < size; ++k) { - rho_mixed[k] += coeff[i] * (rho_history_[i][k] + - config_.beta * residual_history_[i][k]); + rho_mixed[k] += coeff[i] * (rho_history_[i][k] + config_.beta * residual_history_[i][k]); } } } @@ -236,9 +249,7 @@ class ChargeMixerWrapper : public IChargeMixer } // Simple linear system solver (Gaussian elimination with partial pivoting) - std::vector solve_linear_system(std::vector& A, - std::vector& b, - int n) + std::vector solve_linear_system(std::vector& A, std::vector& b, int n) { std::vector x(n, 0.0); @@ -270,7 +281,8 @@ class ChargeMixerWrapper : public IChargeMixer // Eliminate for (int i = k + 1; i < n; ++i) { - if (std::abs(A[k * n + k]) < 1e-12) continue; + if (std::abs(A[k * n + k]) < 1e-12) + continue; double factor = A[i * n + k] / A[k * n + k]; for (int j = k; j < n; ++j) { diff --git a/python/pyabacus/src/ModuleESolver/components/diagonalizer_wrapper.hpp b/python/pyabacus/src/ModuleESolver/components/diagonalizer_wrapper.hpp index 4b22fcc53c4..d28a75b7203 100644 --- a/python/pyabacus/src/ModuleESolver/components/diagonalizer_wrapper.hpp +++ b/python/pyabacus/src/ModuleESolver/components/diagonalizer_wrapper.hpp @@ -8,14 +8,16 @@ #ifndef PYABACUS_ESOLVER_DIAGONALIZER_WRAPPER_HPP #define PYABACUS_ESOLVER_DIAGONALIZER_WRAPPER_HPP -#include "../interfaces/i_diagonalizer.hpp" -#include "../../utils/pybind_utils.h" #include "../../hsolver/diago_adapter.hpp" +#include "../../utils/pybind_utils.h" +#include "../interfaces/i_diagonalizer.hpp" #include -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief Wrapper for diagonalization solvers @@ -28,11 +30,10 @@ namespace esolver { template class DiagonalizerWrapper : public IDiagonalizer { -public: + public: DiagonalizerWrapper() = default; - DiagonalizerWrapper(int nbasis, int nbands) - : nbasis_(nbasis), nbands_(nbands) + DiagonalizerWrapper(int nbasis, int nbands) : nbasis_(nbasis), nbands_(nbands) { } @@ -52,13 +53,11 @@ class DiagonalizerWrapper : public IDiagonalizer // the appropriate ABACUS solver // Create hpsi and spsi functions from matrices - auto hpsi_func = [&Hk, this](const py::array_t& psi) -> py::array_t { - return matrix_vector_multiply(Hk, psi); - }; + auto hpsi_func + = [&Hk, this](const py::array_t& psi) -> py::array_t { return matrix_vector_multiply(Hk, psi); }; - auto spsi_func = [&Sk, this](const py::array_t& psi) -> py::array_t { - return matrix_vector_multiply(Sk, psi); - }; + auto spsi_func + = [&Sk, this](const py::array_t& psi) -> py::array_t { return matrix_vector_multiply(Sk, psi); }; // Use iterative method as fallback py::array_t precond = compute_preconditioner(Hk); @@ -67,28 +66,27 @@ class DiagonalizerWrapper : public IDiagonalizer // ==================== Iterative Diagonalization ==================== - DiagResult diagonalize_iterative( - int ik, - std::function(const py::array_t&)> hpsi_func, - std::function(const py::array_t&)> spsi_func, - const py::array_t& psi_init, - const py::array_t& precond) override + DiagResult diagonalize_iterative(int ik, + std::function(const py::array_t&)> hpsi_func, + std::function(const py::array_t&)> spsi_func, + const py::array_t& psi_init, + const py::array_t& precond) override { DiagResult result; switch (config_.method) { - case DiagMethod::Davidson: - result = diagonalize_davidson(hpsi_func, psi_init, precond); - break; - case DiagMethod::DavSubspace: - result = diagonalize_dav_subspace(hpsi_func, psi_init, precond); - break; - case DiagMethod::CG: - result = diagonalize_cg(hpsi_func, psi_init, precond); - break; - default: - result = diagonalize_davidson(hpsi_func, psi_init, precond); + case DiagMethod::Davidson: + result = diagonalize_davidson(hpsi_func, psi_init, precond); + break; + case DiagMethod::DavSubspace: + result = diagonalize_dav_subspace(hpsi_func, psi_init, precond); + break; + case DiagMethod::CG: + result = diagonalize_cg(hpsi_func, psi_init, precond); + break; + default: + result = diagonalize_davidson(hpsi_func, psi_init, precond); } return result; @@ -101,7 +99,10 @@ class DiagonalizerWrapper : public IDiagonalizer config_ = config; } - DiagConfig get_config() const override { return config_; } + DiagConfig get_config() const override + { + return config_; + } void set_tolerance(double tol) override { @@ -115,20 +116,31 @@ class DiagonalizerWrapper : public IDiagonalizer // ==================== Dimension Queries ==================== - int get_nbasis() const override { return nbasis_; } + int get_nbasis() const override + { + return nbasis_; + } - int get_nbands() const override { return nbands_; } + int get_nbands() const override + { + return nbands_; + } - void set_nbands(int nbands) override { nbands_ = nbands; } + void set_nbands(int nbands) override + { + nbands_ = nbands; + } - void set_nbasis(int nbasis) { nbasis_ = nbasis; } + void set_nbasis(int nbasis) + { + nbasis_ = nbasis; + } -private: + private: // Davidson diagonalization - DiagResult diagonalize_davidson( - std::function(const py::array_t&)> hpsi_func, - const py::array_t& psi_init, - const py::array_t& precond) + DiagResult diagonalize_davidson(std::function(const py::array_t&)> hpsi_func, + const py::array_t& psi_init, + const py::array_t& precond) { DiagResult result; @@ -149,15 +161,13 @@ class DiagonalizerWrapper : public IDiagonalizer ::hsolver::diag_comm_info comm_info(0, 1); // Run diagonalization - int niter = david.diag( - hpsi_func, - precond_vec, - config_.dav_ndim, - config_.tolerance, - diag_ethr, - config_.max_iterations, - comm_info - ); + int niter = david.diag(hpsi_func, + precond_vec, + config_.dav_ndim, + config_.tolerance, + diag_ethr, + config_.max_iterations, + comm_info); // Get results result.psi = david.get_psi(); @@ -169,10 +179,9 @@ class DiagonalizerWrapper : public IDiagonalizer } // Davidson-Subspace diagonalization - DiagResult diagonalize_dav_subspace( - std::function(const py::array_t&)> hpsi_func, - const py::array_t& psi_init, - const py::array_t& precond) + DiagResult diagonalize_dav_subspace(std::function(const py::array_t&)> hpsi_func, + const py::array_t& psi_init, + const py::array_t& precond) { DiagResult result; @@ -193,18 +202,17 @@ class DiagonalizerWrapper : public IDiagonalizer ::hsolver::diag_comm_info comm_info(0, 1); // Run diagonalization - int niter = dav_sub.diag( - hpsi_func, - precond_vec, - config_.dav_ndim, - config_.tolerance, - config_.max_iterations, - false, // need_subspace - diag_ethr, - true, // scf_type - comm_info, - 0, // diag_subspace (LAPACK) - 1 // nb2d + int niter = dav_sub.diag(hpsi_func, + precond_vec, + config_.dav_ndim, + config_.tolerance, + config_.max_iterations, + false, // need_subspace + diag_ethr, + true, // scf_type + comm_info, + 0, // diag_subspace (LAPACK) + 1 // nb2d ); // Get results @@ -217,10 +225,9 @@ class DiagonalizerWrapper : public IDiagonalizer } // CG diagonalization - DiagResult diagonalize_cg( - std::function(const py::array_t&)> hpsi_func, - const py::array_t& psi_init, - const py::array_t& precond) + DiagResult diagonalize_cg(std::function(const py::array_t&)> hpsi_func, + const py::array_t& psi_init, + const py::array_t& precond) { DiagResult result; @@ -237,15 +244,13 @@ class DiagonalizerWrapper : public IDiagonalizer std::vector diag_ethr(nbands_, config_.tolerance); // Run diagonalization - cg.diag( - hpsi_func, - config_.dav_ndim, - config_.tolerance, - diag_ethr, - false, // need_subspace - true, // scf_type - config_.nproc_in_pool - ); + cg.diag(hpsi_func, + config_.dav_ndim, + config_.tolerance, + diag_ethr, + false, // need_subspace + true, // scf_type + config_.nproc_in_pool); // Get results result.psi = cg.get_psi(); @@ -260,8 +265,7 @@ class DiagonalizerWrapper : public IDiagonalizer } // Helper: Matrix-vector multiplication - py::array_t matrix_vector_multiply(const py::array_t& matrix, - const py::array_t& vec) + py::array_t matrix_vector_multiply(const py::array_t& matrix, const py::array_t& vec) { using namespace pyabacus::utils; diff --git a/python/pyabacus/src/ModuleESolver/components/hamiltonian_builder_lcao.hpp b/python/pyabacus/src/ModuleESolver/components/hamiltonian_builder_lcao.hpp index 97965674609..fffd3b2c382 100644 --- a/python/pyabacus/src/ModuleESolver/components/hamiltonian_builder_lcao.hpp +++ b/python/pyabacus/src/ModuleESolver/components/hamiltonian_builder_lcao.hpp @@ -8,16 +8,18 @@ #ifndef PYABACUS_ESOLVER_HAMILTONIAN_BUILDER_LCAO_HPP #define PYABACUS_ESOLVER_HAMILTONIAN_BUILDER_LCAO_HPP -#include "../interfaces/i_hamiltonian_builder.hpp" #include "../../utils/pybind_utils.h" +#include "../interfaces/i_hamiltonian_builder.hpp" -#include -#include #include +#include #include +#include -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief LCAO Hamiltonian builder implementation @@ -31,7 +33,7 @@ namespace esolver { template class HamiltonianBuilderLCAO : public IHamiltonianBuilder { -public: + public: HamiltonianBuilderLCAO() = default; HamiltonianBuilderLCAO(int nbasis, int nks, int nrow, int ncol) @@ -72,8 +74,7 @@ class HamiltonianBuilderLCAO : public IHamiltonianBuilder nrxx_ = static_cast(buf.shape[1]); // Store rho - rho_data_.assign(static_cast(buf.ptr), - static_cast(buf.ptr) + nspin_ * nrxx_); + rho_data_.assign(static_cast(buf.ptr), static_cast(buf.ptr) + nspin_ * nrxx_); // Mark H(k) as needing update hk_valid_.assign(nks_, false); @@ -106,22 +107,18 @@ class HamiltonianBuilderLCAO : public IHamiltonianBuilder { validate_kpoint(ik); - return utils::numpy_from_ptr_copy_2d( - Hk_matrices_[ik].data(), - static_cast(nrow_), - static_cast(ncol_) - ); + return utils::numpy_from_ptr_copy_2d(Hk_matrices_[ik].data(), + static_cast(nrow_), + static_cast(ncol_)); } py::array_t get_Sk(int ik) const override { validate_kpoint(ik); - return utils::numpy_from_ptr_copy_2d( - Sk_matrices_[ik].data(), - static_cast(nrow_), - static_cast(ncol_) - ); + return utils::numpy_from_ptr_copy_2d(Sk_matrices_[ik].data(), + static_cast(nrow_), + static_cast(ncol_)); } // ==================== Real-space Matrix Access ==================== @@ -161,16 +158,25 @@ class HamiltonianBuilderLCAO : public IHamiltonianBuilder // ==================== Dimension Queries ==================== - int get_nbasis() const override { return nbasis_; } + int get_nbasis() const override + { + return nbasis_; + } - int get_nks() const override { return nks_; } + int get_nks() const override + { + return nks_; + } std::pair get_local_dims() const override { return {nrow_, ncol_}; } - bool is_valid() const override { return valid_; } + bool is_valid() const override + { + return valid_; + } // ==================== Data Setters (for testing/compatibility) ==================== @@ -227,7 +233,7 @@ class HamiltonianBuilderLCAO : public IHamiltonianBuilder } } -private: + private: void validate_kpoint(int ik) const { if (!valid_) @@ -240,14 +246,12 @@ class HamiltonianBuilderLCAO : public IHamiltonianBuilder } } - py::array_t matrix_multiply(const std::vector& matrix, - const py::array_t& vec) const + py::array_t matrix_multiply(const std::vector& matrix, const py::array_t& vec) const { auto vec_buf = vec.request(); const TK* vec_ptr = static_cast(vec_buf.ptr); - const ssize_t nvec = (vec_buf.ndim == 1) ? 1 : - (vec_buf.ndim == 2) ? vec_buf.shape[1] : 1; + const ssize_t nvec = (vec_buf.ndim == 1) ? 1 : (vec_buf.ndim == 2) ? vec_buf.shape[1] : 1; const ssize_t vec_rows = (vec_buf.ndim == 1) ? vec_buf.shape[0] : vec_buf.shape[0]; if (vec_rows != ncol_) diff --git a/python/pyabacus/src/ModuleESolver/components/scf_controller_lcao.hpp b/python/pyabacus/src/ModuleESolver/components/scf_controller_lcao.hpp index c5dac087fb3..52c549f845d 100644 --- a/python/pyabacus/src/ModuleESolver/components/scf_controller_lcao.hpp +++ b/python/pyabacus/src/ModuleESolver/components/scf_controller_lcao.hpp @@ -8,19 +8,21 @@ #ifndef PYABACUS_ESOLVER_SCF_CONTROLLER_LCAO_HPP #define PYABACUS_ESOLVER_SCF_CONTROLLER_LCAO_HPP -#include "../interfaces/i_scf_controller.hpp" -#include "../interfaces/i_hamiltonian_builder.hpp" #include "../interfaces/i_charge_mixer.hpp" #include "../interfaces/i_diagonalizer.hpp" -#include "hamiltonian_builder_lcao.hpp" +#include "../interfaces/i_hamiltonian_builder.hpp" +#include "../interfaces/i_scf_controller.hpp" #include "charge_mixer_wrapper.hpp" #include "diagonalizer_wrapper.hpp" +#include "hamiltonian_builder_lcao.hpp" -#include #include +#include -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief LCAO SCF controller implementation @@ -34,19 +36,14 @@ namespace esolver { template class SCFControllerLCAO : public ISCFController { -public: + public: SCFControllerLCAO() = default; SCFControllerLCAO(int nbasis, int nks, int nbands, int nspin, int nrxx) - : nbasis_(nbasis) - , nks_(nks) - , nbands_(nbands) - , nspin_(nspin) - , nrxx_(nrxx) + : nbasis_(nbasis), nks_(nks), nbands_(nbands), nspin_(nspin), nrxx_(nrxx) { // Initialize components - hamilt_builder_ = std::make_unique>( - nbasis, nks, nbasis, nbasis); + hamilt_builder_ = std::make_unique>(nbasis, nks, nbasis, nbasis); charge_mixer_ = std::make_unique(nspin, nrxx); diagonalizer_ = std::make_unique>(nbasis, nbands); } @@ -109,8 +106,7 @@ class SCFControllerLCAO : public ISCFController return status_; } - SCFStatus run_scf(const SCFConvergenceCriteria& criteria, - SCFIterationCallback callback) override + SCFStatus run_scf(const SCFConvergenceCriteria& criteria, SCFIterationCallback callback) override { initialize(istep_); status_ = SCFStatus::Running; @@ -153,15 +149,27 @@ class SCFControllerLCAO : public ISCFController return status_ == SCFStatus::Converged; } - SCFStatus get_status() const override { return status_; } + SCFStatus get_status() const override + { + return status_; + } // ==================== State Queries ==================== - int get_iteration() const override { return iteration_; } + int get_iteration() const override + { + return iteration_; + } - double get_drho() const override { return drho_; } + double get_drho() const override + { + return drho_; + } - double get_energy() const override { return energy_; } + double get_energy() const override + { + return energy_; + } // ==================== Component Access ==================== @@ -220,7 +228,7 @@ class SCFControllerLCAO : public ISCFController } } -private: + private: // Dimensions int nbasis_ = 0; int nks_ = 0; diff --git a/python/pyabacus/src/ModuleESolver/interfaces/i_charge_mixer.hpp b/python/pyabacus/src/ModuleESolver/interfaces/i_charge_mixer.hpp index 43ac40fc58e..f139aa2375b 100644 --- a/python/pyabacus/src/ModuleESolver/interfaces/i_charge_mixer.hpp +++ b/python/pyabacus/src/ModuleESolver/interfaces/i_charge_mixer.hpp @@ -8,25 +8,26 @@ #ifndef PYABACUS_ESOLVER_I_CHARGE_MIXER_HPP #define PYABACUS_ESOLVER_I_CHARGE_MIXER_HPP -#include #include - +#include #include namespace py = pybind11; -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief Mixing method types */ enum class MixingMethod { - Plain, ///< Simple linear mixing - Pulay, ///< Pulay mixing (DIIS) - Broyden, ///< Broyden mixing - Anderson ///< Anderson mixing + Plain, ///< Simple linear mixing + Pulay, ///< Pulay mixing (DIIS) + Broyden, ///< Broyden mixing + Anderson ///< Anderson mixing }; /** @@ -35,11 +36,11 @@ enum class MixingMethod struct MixingConfig { MixingMethod method = MixingMethod::Pulay; - double beta = 0.7; ///< Mixing parameter - int ndim = 8; ///< Number of history steps for Pulay/Broyden - double gg0 = 1.0; ///< Kerker mixing parameter - bool mix_gg0 = false; ///< Whether to use Kerker mixing - bool mix_rho = true; ///< Mix charge density (vs potential) + double beta = 0.7; ///< Mixing parameter + int ndim = 8; ///< Number of history steps for Pulay/Broyden + double gg0 = 1.0; ///< Kerker mixing parameter + bool mix_gg0 = false; ///< Whether to use Kerker mixing + bool mix_rho = true; ///< Mix charge density (vs potential) }; /** @@ -50,7 +51,7 @@ struct MixingConfig */ class IChargeMixer { -public: + public: virtual ~IChargeMixer() = default; // ==================== Core Mixing Operations ==================== @@ -61,8 +62,7 @@ class IChargeMixer * @param rho_out Output charge density (from current iteration) * @return Mixed charge density for next iteration */ - virtual py::array_t mix(const py::array_t& rho_in, - const py::array_t& rho_out) = 0; + virtual py::array_t mix(const py::array_t& rho_in, const py::array_t& rho_out) = 0; /** * @brief Reset mixer state (clear history) @@ -129,11 +129,16 @@ inline std::string mixing_method_to_string(MixingMethod method) { switch (method) { - case MixingMethod::Plain: return "plain"; - case MixingMethod::Pulay: return "pulay"; - case MixingMethod::Broyden: return "broyden"; - case MixingMethod::Anderson: return "anderson"; - default: return "unknown"; + case MixingMethod::Plain: + return "plain"; + case MixingMethod::Pulay: + return "pulay"; + case MixingMethod::Broyden: + return "broyden"; + case MixingMethod::Anderson: + return "anderson"; + default: + return "unknown"; } } @@ -142,10 +147,14 @@ inline std::string mixing_method_to_string(MixingMethod method) */ inline MixingMethod string_to_mixing_method(const std::string& str) { - if (str == "plain") return MixingMethod::Plain; - if (str == "pulay") return MixingMethod::Pulay; - if (str == "broyden") return MixingMethod::Broyden; - if (str == "anderson") return MixingMethod::Anderson; + if (str == "plain") + return MixingMethod::Plain; + if (str == "pulay") + return MixingMethod::Pulay; + if (str == "broyden") + return MixingMethod::Broyden; + if (str == "anderson") + return MixingMethod::Anderson; return MixingMethod::Pulay; // default } diff --git a/python/pyabacus/src/ModuleESolver/interfaces/i_diagonalizer.hpp b/python/pyabacus/src/ModuleESolver/interfaces/i_diagonalizer.hpp index 3f626b06b31..be8364dd379 100644 --- a/python/pyabacus/src/ModuleESolver/interfaces/i_diagonalizer.hpp +++ b/python/pyabacus/src/ModuleESolver/interfaces/i_diagonalizer.hpp @@ -9,31 +9,32 @@ #ifndef PYABACUS_ESOLVER_I_DIAGONALIZER_HPP #define PYABACUS_ESOLVER_I_DIAGONALIZER_HPP -#include -#include -#include - #include #include +#include +#include +#include #include #include namespace py = pybind11; -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief Diagonalization method types */ enum class DiagMethod { - Davidson, ///< Davidson iterative method - DavSubspace, ///< Davidson with subspace rotation - CG, ///< Conjugate gradient - LAPACK, ///< Direct LAPACK diagonalization - ScaLAPACK, ///< Parallel ScaLAPACK - ELPA ///< ELPA eigensolver + Davidson, ///< Davidson iterative method + DavSubspace, ///< Davidson with subspace rotation + CG, ///< Conjugate gradient + LAPACK, ///< Direct LAPACK diagonalization + ScaLAPACK, ///< Parallel ScaLAPACK + ELPA ///< ELPA eigensolver }; /** @@ -42,11 +43,11 @@ enum class DiagMethod template struct DiagResult { - py::array_t psi; ///< Eigenvectors + py::array_t psi; ///< Eigenvectors py::array_t eigenvalues; ///< Eigenvalues - int iterations = 0; ///< Number of iterations (for iterative methods) - bool converged = false; ///< Whether converged - double residual = 0.0; ///< Final residual + int iterations = 0; ///< Number of iterations (for iterative methods) + bool converged = false; ///< Whether converged + double residual = 0.0; ///< Final residual }; /** @@ -55,10 +56,10 @@ struct DiagResult struct DiagConfig { DiagMethod method = DiagMethod::Davidson; - double tolerance = 1e-6; ///< Convergence tolerance - int max_iterations = 100; ///< Maximum iterations - int dav_ndim = 4; ///< Davidson subspace dimension multiplier - int nproc_in_pool = 1; ///< Number of processes in pool + double tolerance = 1e-6; ///< Convergence tolerance + int max_iterations = 100; ///< Maximum iterations + int dav_ndim = 4; ///< Davidson subspace dimension multiplier + int nproc_in_pool = 1; ///< Number of processes in pool }; /** @@ -72,7 +73,7 @@ struct DiagConfig template class IDiagonalizer { -public: + public: virtual ~IDiagonalizer() = default; // ==================== Direct Diagonalization ==================== @@ -104,12 +105,11 @@ class IDiagonalizer * @param precond Preconditioner (diagonal approximation to H) * @return Diagonalization result */ - virtual DiagResult diagonalize_iterative( - int ik, - std::function(const py::array_t&)> hpsi_func, - std::function(const py::array_t&)> spsi_func, - const py::array_t& psi_init, - const py::array_t& precond) = 0; + virtual DiagResult diagonalize_iterative(int ik, + std::function(const py::array_t&)> hpsi_func, + std::function(const py::array_t&)> spsi_func, + const py::array_t& psi_init, + const py::array_t& precond) = 0; // ==================== Configuration ==================== @@ -169,13 +169,20 @@ inline std::string diag_method_to_string(DiagMethod method) { switch (method) { - case DiagMethod::Davidson: return "davidson"; - case DiagMethod::DavSubspace: return "dav_subspace"; - case DiagMethod::CG: return "cg"; - case DiagMethod::LAPACK: return "lapack"; - case DiagMethod::ScaLAPACK: return "scalapack"; - case DiagMethod::ELPA: return "elpa"; - default: return "unknown"; + case DiagMethod::Davidson: + return "davidson"; + case DiagMethod::DavSubspace: + return "dav_subspace"; + case DiagMethod::CG: + return "cg"; + case DiagMethod::LAPACK: + return "lapack"; + case DiagMethod::ScaLAPACK: + return "scalapack"; + case DiagMethod::ELPA: + return "elpa"; + default: + return "unknown"; } } @@ -184,12 +191,18 @@ inline std::string diag_method_to_string(DiagMethod method) */ inline DiagMethod string_to_diag_method(const std::string& str) { - if (str == "davidson") return DiagMethod::Davidson; - if (str == "dav_subspace") return DiagMethod::DavSubspace; - if (str == "cg") return DiagMethod::CG; - if (str == "lapack") return DiagMethod::LAPACK; - if (str == "scalapack") return DiagMethod::ScaLAPACK; - if (str == "elpa") return DiagMethod::ELPA; + if (str == "davidson") + return DiagMethod::Davidson; + if (str == "dav_subspace") + return DiagMethod::DavSubspace; + if (str == "cg") + return DiagMethod::CG; + if (str == "lapack") + return DiagMethod::LAPACK; + if (str == "scalapack") + return DiagMethod::ScaLAPACK; + if (str == "elpa") + return DiagMethod::ELPA; return DiagMethod::Davidson; // default } diff --git a/python/pyabacus/src/ModuleESolver/interfaces/i_hamiltonian_builder.hpp b/python/pyabacus/src/ModuleESolver/interfaces/i_hamiltonian_builder.hpp index 794b8c119ee..a192412c990 100644 --- a/python/pyabacus/src/ModuleESolver/interfaces/i_hamiltonian_builder.hpp +++ b/python/pyabacus/src/ModuleESolver/interfaces/i_hamiltonian_builder.hpp @@ -9,16 +9,17 @@ #ifndef PYABACUS_ESOLVER_I_HAMILTONIAN_BUILDER_HPP #define PYABACUS_ESOLVER_I_HAMILTONIAN_BUILDER_HPP -#include -#include - #include #include +#include +#include namespace py = pybind11; -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief Abstract interface for Hamiltonian builder @@ -32,7 +33,7 @@ namespace esolver { template class IHamiltonianBuilder { -public: + public: virtual ~IHamiltonianBuilder() = default; // ==================== Build/Update ==================== diff --git a/python/pyabacus/src/ModuleESolver/interfaces/i_scf_controller.hpp b/python/pyabacus/src/ModuleESolver/interfaces/i_scf_controller.hpp index 401a7e6dac1..c91c84bbe3a 100644 --- a/python/pyabacus/src/ModuleESolver/interfaces/i_scf_controller.hpp +++ b/python/pyabacus/src/ModuleESolver/interfaces/i_scf_controller.hpp @@ -12,8 +12,10 @@ #include #include -namespace pyabacus { -namespace esolver { +namespace pyabacus +{ +namespace esolver +{ /** * @brief SCF convergence status @@ -52,9 +54,11 @@ struct SCFConvergenceCriteria using SCFIterationCallback = std::function; // Forward declarations for component interfaces -template class IHamiltonianBuilder; +template +class IHamiltonianBuilder; class IChargeMixer; -template class IDiagonalizer; +template +class IDiagonalizer; /** * @brief Abstract interface for SCF controller @@ -65,7 +69,7 @@ template class IDiagonalizer; */ class ISCFController { -public: + public: virtual ~ISCFController() = default; // ==================== Lifecycle ==================== @@ -97,8 +101,7 @@ class ISCFController * @param callback Optional callback for each iteration * @return Final SCF status */ - virtual SCFStatus run_scf(const SCFConvergenceCriteria& criteria, - SCFIterationCallback callback = nullptr) = 0; + virtual SCFStatus run_scf(const SCFConvergenceCriteria& criteria, SCFIterationCallback callback = nullptr) = 0; /** * @brief Check if SCF is converged diff --git a/python/pyabacus/src/ModuleESolver/py_esolver_lcao.cpp b/python/pyabacus/src/ModuleESolver/py_esolver_lcao.cpp index cc899b55c20..38dab915317 100644 --- a/python/pyabacus/src/ModuleESolver/py_esolver_lcao.cpp +++ b/python/pyabacus/src/ModuleESolver/py_esolver_lcao.cpp @@ -6,26 +6,26 @@ * enabling Python-controlled SCF workflows with breakpoint support. */ -#include -#include -#include +#include "py_esolver_lcao.hpp" + #include #include - -#include "py_esolver_lcao.hpp" +#include +#include +#include // ABACUS headers for actual implementation -#include "source_estate/module_charge/charge.h" -#include "source_estate/fp_energy.h" +#include "source_basis/module_ao/parallel_orbitals.h" #include "source_estate/elecstate.h" +#include "source_estate/fp_energy.h" +#include "source_estate/module_charge/charge.h" #include "source_estate/module_dm/density_matrix.h" -#include "source_lcao/hamilt_lcao.h" #include "source_hamilt/module_hcontainer/hcontainer.h" -#include "source_basis/module_ao/parallel_orbitals.h" +#include "source_lcao/hamilt_lcao.h" #include -#include #include +#include namespace py = pybind11; using namespace pybind11::literals; @@ -50,7 +50,7 @@ void PyChargeAccessor::set_from_charge(const Charge* chr) } chr_ptr_ = chr; - rho_ptr_ = nullptr; // Use chr_ptr_ instead + rho_ptr_ = nullptr; // Use chr_ptr_ instead nspin_ = chr->nspin; nrxx_ = chr->nrxx; ngmc_ = chr->ngmc; @@ -58,7 +58,7 @@ void PyChargeAccessor::set_from_charge(const Charge* chr) void PyChargeAccessor::set_data(const double* rho_ptr, int nspin, int nrxx) { - chr_ptr_ = nullptr; // Not using Charge object + chr_ptr_ = nullptr; // Not using Charge object rho_ptr_ = rho_ptr; nspin_ = nspin; nrxx_ = nrxx; @@ -178,9 +178,14 @@ void PyEnergyAccessor::set_from_fenergy(const elecstate::fenergy* f_en) evdw_ = f_en->evdw; } -void PyEnergyAccessor::set_energies(double etot, double eband, double hartree, - double etxc, double ewald, double demet, - double exx, double evdw) +void PyEnergyAccessor::set_energies(double etot, + double eband, + double hartree, + double etxc, + double ewald, + double demet, + double exx, + double evdw) { etot_ = etot; eband_ = eband; @@ -211,7 +216,9 @@ py::dict PyEnergyAccessor::get_all_energies() const // ============================================================================ template -void PyHamiltonianAccessor::set_from_hamilt(hamilt::HamiltLCAO* hamilt_lcao, int nks, const Parallel_Orbitals* pv) +void PyHamiltonianAccessor::set_from_hamilt(hamilt::HamiltLCAO* hamilt_lcao, + int nks, + const Parallel_Orbitals* pv) { hamilt_ptr_ = hamilt_lcao; pv_ = pv; @@ -649,15 +656,17 @@ template class PyESolverLCAO, double>; void bind_charge_accessor(py::module& m) { - py::class_(m, "ChargeAccessor", - R"pbdoc( + py::class_(m, + "ChargeAccessor", + R"pbdoc( Accessor for charge density data. Provides access to real-space charge density (rho) and related quantities. )pbdoc") .def(py::init<>()) - .def("get_rho", &py_esolver::PyChargeAccessor::get_rho, - R"pbdoc( + .def("get_rho", + &py_esolver::PyChargeAccessor::get_rho, + R"pbdoc( Get real-space charge density as numpy array. Returns @@ -665,8 +674,9 @@ void bind_charge_accessor(py::module& m) numpy.ndarray Charge density with shape (nspin, nrxx) )pbdoc") - .def("get_rhog", &py_esolver::PyChargeAccessor::get_rhog, - R"pbdoc( + .def("get_rhog", + &py_esolver::PyChargeAccessor::get_rhog, + R"pbdoc( Get reciprocal-space charge density as numpy array. Returns @@ -674,8 +684,9 @@ void bind_charge_accessor(py::module& m) numpy.ndarray Charge density in G-space with shape (nspin, ngmc) )pbdoc") - .def("get_rho_core", &py_esolver::PyChargeAccessor::get_rho_core, - R"pbdoc( + .def("get_rho_core", + &py_esolver::PyChargeAccessor::get_rho_core, + R"pbdoc( Get core charge density as numpy array. Returns @@ -683,44 +694,36 @@ void bind_charge_accessor(py::module& m) numpy.ndarray Core charge density with shape (nrxx,) )pbdoc") - .def_property_readonly("nspin", &py_esolver::PyChargeAccessor::get_nspin, - "Number of spin channels") - .def_property_readonly("nrxx", &py_esolver::PyChargeAccessor::get_nrxx, - "Number of real-space grid points") - .def_property_readonly("ngmc", &py_esolver::PyChargeAccessor::get_ngmc, - "Number of G-vectors for charge density") - .def("is_valid", &py_esolver::PyChargeAccessor::is_valid, - "Check if charge data is available"); + .def_property_readonly("nspin", &py_esolver::PyChargeAccessor::get_nspin, "Number of spin channels") + .def_property_readonly("nrxx", &py_esolver::PyChargeAccessor::get_nrxx, "Number of real-space grid points") + .def_property_readonly("ngmc", + &py_esolver::PyChargeAccessor::get_ngmc, + "Number of G-vectors for charge density") + .def("is_valid", &py_esolver::PyChargeAccessor::is_valid, "Check if charge data is available"); } void bind_energy_accessor(py::module& m) { - py::class_(m, "EnergyAccessor", - R"pbdoc( + py::class_(m, + "EnergyAccessor", + R"pbdoc( Accessor for energy data. Provides access to various energy components from the calculation. All energies are in Rydberg units. )pbdoc") .def(py::init<>()) - .def_property_readonly("etot", &py_esolver::PyEnergyAccessor::get_etot, - "Total energy (Ry)") - .def_property_readonly("eband", &py_esolver::PyEnergyAccessor::get_eband, - "Band energy (Ry)") - .def_property_readonly("hartree_energy", &py_esolver::PyEnergyAccessor::get_hartree_energy, - "Hartree energy (Ry)") - .def_property_readonly("etxc", &py_esolver::PyEnergyAccessor::get_etxc, - "Exchange-correlation energy (Ry)") - .def_property_readonly("ewald_energy", &py_esolver::PyEnergyAccessor::get_ewald_energy, - "Ewald energy (Ry)") - .def_property_readonly("demet", &py_esolver::PyEnergyAccessor::get_demet, - "-TS term for metals (Ry)") - .def_property_readonly("exx", &py_esolver::PyEnergyAccessor::get_exx, - "Exact exchange energy (Ry)") - .def_property_readonly("evdw", &py_esolver::PyEnergyAccessor::get_evdw, - "van der Waals energy (Ry)") - .def("get_all_energies", &py_esolver::PyEnergyAccessor::get_all_energies, - "Get all energies as a dictionary"); + .def_property_readonly("etot", &py_esolver::PyEnergyAccessor::get_etot, "Total energy (Ry)") + .def_property_readonly("eband", &py_esolver::PyEnergyAccessor::get_eband, "Band energy (Ry)") + .def_property_readonly("hartree_energy", + &py_esolver::PyEnergyAccessor::get_hartree_energy, + "Hartree energy (Ry)") + .def_property_readonly("etxc", &py_esolver::PyEnergyAccessor::get_etxc, "Exchange-correlation energy (Ry)") + .def_property_readonly("ewald_energy", &py_esolver::PyEnergyAccessor::get_ewald_energy, "Ewald energy (Ry)") + .def_property_readonly("demet", &py_esolver::PyEnergyAccessor::get_demet, "-TS term for metals (Ry)") + .def_property_readonly("exx", &py_esolver::PyEnergyAccessor::get_exx, "Exact exchange energy (Ry)") + .def_property_readonly("evdw", &py_esolver::PyEnergyAccessor::get_evdw, "van der Waals energy (Ry)") + .def("get_all_energies", &py_esolver::PyEnergyAccessor::get_all_energies, "Get all energies as a dictionary"); } template @@ -730,19 +733,19 @@ void bind_hamiltonian_accessor(py::module& m, const std::string& suffix) std::string class_name = "HamiltonianAccessor" + suffix; - py::class_(m, class_name.c_str(), - R"pbdoc( + py::class_(m, + class_name.c_str(), + R"pbdoc( Accessor for Hamiltonian matrix data. Provides access to H(k), S(k), H(R), and S(R) matrices. )pbdoc") .def(py::init<>()) - .def_property_readonly("nbasis", &HamiltAccessor::get_nbasis, - "Number of basis functions") - .def_property_readonly("nks", &HamiltAccessor::get_nks, - "Number of k-points") - .def("get_Hk", &HamiltAccessor::get_Hk, - R"pbdoc( + .def_property_readonly("nbasis", &HamiltAccessor::get_nbasis, "Number of basis functions") + .def_property_readonly("nks", &HamiltAccessor::get_nks, "Number of k-points") + .def("get_Hk", + &HamiltAccessor::get_Hk, + R"pbdoc( Get H(k) matrix for specific k-point. Parameters @@ -754,9 +757,11 @@ void bind_hamiltonian_accessor(py::module& m, const std::string& suffix) ------- numpy.ndarray Hamiltonian matrix at k-point ik - )pbdoc", "ik"_a) - .def("get_Sk", &HamiltAccessor::get_Sk, - R"pbdoc( + )pbdoc", + "ik"_a) + .def("get_Sk", + &HamiltAccessor::get_Sk, + R"pbdoc( Get S(k) overlap matrix for specific k-point. Parameters @@ -768,13 +773,11 @@ void bind_hamiltonian_accessor(py::module& m, const std::string& suffix) ------- numpy.ndarray Overlap matrix at k-point ik - )pbdoc", "ik"_a) - .def("get_HR", &HamiltAccessor::get_HR, - "Get H(R) in sparse format") - .def("get_SR", &HamiltAccessor::get_SR, - "Get S(R) in sparse format") - .def("is_valid", &HamiltAccessor::is_valid, - "Check if Hamiltonian data is available"); + )pbdoc", + "ik"_a) + .def("get_HR", &HamiltAccessor::get_HR, "Get H(R) in sparse format") + .def("get_SR", &HamiltAccessor::get_SR, "Get S(R) in sparse format") + .def("is_valid", &HamiltAccessor::is_valid, "Check if Hamiltonian data is available"); } template @@ -784,21 +787,20 @@ void bind_density_matrix_accessor(py::module& m, const std::string& suffix) std::string class_name = "DensityMatrixAccessor" + suffix; - py::class_(m, class_name.c_str(), - R"pbdoc( + py::class_(m, + class_name.c_str(), + R"pbdoc( Accessor for density matrix data. Provides access to DM(k) and DM(R) matrices. )pbdoc") .def(py::init<>()) - .def_property_readonly("nks", &DMAccessor::get_nks, - "Number of k-points") - .def_property_readonly("nrow", &DMAccessor::get_nrow, - "Number of rows in density matrix") - .def_property_readonly("ncol", &DMAccessor::get_ncol, - "Number of columns in density matrix") - .def("get_dmk", &DMAccessor::get_dmk, - R"pbdoc( + .def_property_readonly("nks", &DMAccessor::get_nks, "Number of k-points") + .def_property_readonly("nrow", &DMAccessor::get_nrow, "Number of rows in density matrix") + .def_property_readonly("ncol", &DMAccessor::get_ncol, "Number of columns in density matrix") + .def("get_dmk", + &DMAccessor::get_dmk, + R"pbdoc( Get DM(k) for specific k-point. Parameters @@ -810,13 +812,11 @@ void bind_density_matrix_accessor(py::module& m, const std::string& suffix) ------- numpy.ndarray Density matrix at k-point ik - )pbdoc", "ik"_a) - .def("get_DMK_all", &DMAccessor::get_DMK_all, - "Get all DM(k) matrices as a list") - .def("get_DMR", &DMAccessor::get_DMR, - "Get DM(R) in sparse format") - .def("is_valid", &DMAccessor::is_valid, - "Check if density matrix data is available"); + )pbdoc", + "ik"_a) + .def("get_DMK_all", &DMAccessor::get_DMK_all, "Get all DM(k) matrices as a list") + .def("get_DMR", &DMAccessor::get_DMR, "Get DM(R) in sparse format") + .def("is_valid", &DMAccessor::is_valid, "Check if density matrix data is available"); } template @@ -826,8 +826,9 @@ void bind_esolver_lcao(py::module& m, const std::string& suffix) std::string class_name = "ESolverLCAO" + suffix; - py::class_(m, class_name.c_str(), - R"pbdoc( + py::class_(m, + class_name.c_str(), + R"pbdoc( Python wrapper for ESolver_KS_LCAO. This class provides a Python interface for LCAO calculations @@ -853,101 +854,92 @@ void bind_esolver_lcao(py::module& m, const std::string& suffix) .def(py::init<>()) // Initialization - .def("initialize", &ESolver::initialize, - R"pbdoc( + .def("initialize", + &ESolver::initialize, + R"pbdoc( Initialize ESolver from INPUT file. Parameters ---------- input_dir : str Directory containing INPUT, STRU, and other input files - )pbdoc", "input_dir"_a) - .def("before_all_runners", &ESolver::before_all_runners, - "Initialize calculation environment") + )pbdoc", + "input_dir"_a) + .def("before_all_runners", &ESolver::before_all_runners, "Initialize calculation environment") // SCF Control - .def("before_scf", &ESolver::before_scf, - R"pbdoc( + .def("before_scf", + &ESolver::before_scf, + R"pbdoc( Prepare for SCF calculation. Parameters ---------- istep : int, optional Ion step index (default: 0) - )pbdoc", "istep"_a = 0) - .def("run_scf_iteration", &ESolver::run_scf_iteration, - R"pbdoc( + )pbdoc", + "istep"_a = 0) + .def("run_scf_iteration", + &ESolver::run_scf_iteration, + R"pbdoc( Run a single SCF iteration. Parameters ---------- iter : int Iteration number (1-based) - )pbdoc", "iter"_a) - .def("run_scf", &ESolver::run_scf, - R"pbdoc( + )pbdoc", + "iter"_a) + .def("run_scf", + &ESolver::run_scf, + R"pbdoc( Run complete SCF loop. Parameters ---------- max_iter : int, optional Maximum number of iterations (default: 100) - )pbdoc", "max_iter"_a = 100) - .def("after_scf", &ESolver::after_scf, - R"pbdoc( + )pbdoc", + "max_iter"_a = 100) + .def("after_scf", + &ESolver::after_scf, + R"pbdoc( Finalize SCF calculation. Parameters ---------- istep : int, optional Ion step index (default: 0) - )pbdoc", "istep"_a = 0) + )pbdoc", + "istep"_a = 0) // Status - .def("is_converged", &ESolver::is_converged, - "Check if SCF is converged") - .def_property_readonly("niter", &ESolver::get_niter, - "Current iteration number") - .def_property_readonly("drho", &ESolver::get_drho, - "Charge density difference") - .def_property_readonly("istep", &ESolver::get_istep, - "Current ion step") + .def("is_converged", &ESolver::is_converged, "Check if SCF is converged") + .def_property_readonly("niter", &ESolver::get_niter, "Current iteration number") + .def_property_readonly("drho", &ESolver::get_drho, "Charge density difference") + .def_property_readonly("istep", &ESolver::get_istep, "Current ion step") // Data Accessors - .def("get_charge", &ESolver::get_charge, - "Get charge density accessor") - .def("get_energy", &ESolver::get_energy, - "Get energy accessor") - .def("get_hamiltonian", &ESolver::get_hamiltonian, - "Get Hamiltonian accessor") - .def("get_density_matrix", &ESolver::get_density_matrix, - "Get density matrix accessor") + .def("get_charge", &ESolver::get_charge, "Get charge density accessor") + .def("get_energy", &ESolver::get_energy, "Get energy accessor") + .def("get_hamiltonian", &ESolver::get_hamiltonian, "Get Hamiltonian accessor") + .def("get_density_matrix", &ESolver::get_density_matrix, "Get density matrix accessor") // Wave functions - .def("get_psi", &ESolver::get_psi, - "Get wave function coefficients for k-point ik", "ik"_a) - .def("get_eigenvalues", &ESolver::get_eigenvalues, - "Get eigenvalues for k-point ik", "ik"_a) - .def("get_occupations", &ESolver::get_occupations, - "Get occupation numbers for k-point ik", "ik"_a) + .def("get_psi", &ESolver::get_psi, "Get wave function coefficients for k-point ik", "ik"_a) + .def("get_eigenvalues", &ESolver::get_eigenvalues, "Get eigenvalues for k-point ik", "ik"_a) + .def("get_occupations", &ESolver::get_occupations, "Get occupation numbers for k-point ik", "ik"_a) // K-points - .def_property_readonly("nks", &ESolver::get_nks, - "Number of k-points") - .def("get_kvec_d", &ESolver::get_kvec_d, - "Get k-vector in direct coordinates", "ik"_a) - .def("get_wk", &ESolver::get_wk, - "Get k-point weights") + .def_property_readonly("nks", &ESolver::get_nks, "Number of k-points") + .def("get_kvec_d", &ESolver::get_kvec_d, "Get k-vector in direct coordinates", "ik"_a) + .def("get_wk", &ESolver::get_wk, "Get k-point weights") // System info - .def_property_readonly("nbasis", &ESolver::get_nbasis, - "Number of basis functions") - .def_property_readonly("nbands", &ESolver::get_nbands, - "Number of bands") - .def_property_readonly("nspin", &ESolver::get_nspin, - "Number of spin channels") - .def_property_readonly("nat", &ESolver::get_nat, - "Number of atoms"); + .def_property_readonly("nbasis", &ESolver::get_nbasis, "Number of basis functions") + .def_property_readonly("nbands", &ESolver::get_nbands, "Number of bands") + .def_property_readonly("nspin", &ESolver::get_nspin, "Number of spin channels") + .def_property_readonly("nat", &ESolver::get_nat, "Number of atoms"); } PYBIND11_MODULE(_esolver_pack, m) diff --git a/python/pyabacus/src/ModuleESolver/py_esolver_lcao.hpp b/python/pyabacus/src/ModuleESolver/py_esolver_lcao.hpp index 1f6d5ce78b7..c3662ca87a1 100644 --- a/python/pyabacus/src/ModuleESolver/py_esolver_lcao.hpp +++ b/python/pyabacus/src/ModuleESolver/py_esolver_lcao.hpp @@ -1,44 +1,53 @@ #ifndef PY_ESOLVER_LCAO_HPP #define PY_ESOLVER_LCAO_HPP -#include -#include -#include -#include +#include "../utils/pybind_utils.h" +#include "components/scf_controller_lcao.hpp" +#include "interfaces/i_charge_mixer.hpp" +#include "interfaces/i_diagonalizer.hpp" +#include "interfaces/i_hamiltonian_builder.hpp" +#include "interfaces/i_scf_controller.hpp" #include +#include #include +#include +#include +#include +#include #include -#include -#include #include - -#include "../utils/pybind_utils.h" -#include "interfaces/i_scf_controller.hpp" -#include "interfaces/i_hamiltonian_builder.hpp" -#include "interfaces/i_charge_mixer.hpp" -#include "interfaces/i_diagonalizer.hpp" -#include "components/scf_controller_lcao.hpp" +#include // Forward declarations for ABACUS types class UnitCell; class Charge; class Parallel_Orbitals; -namespace elecstate { - struct fenergy; - class ElecState; -} -namespace module_dm { - template class DensityMatrix; -} -namespace hamilt { - template class HContainer; - template class AtomPair; - template class BaseMatrix; - template class HamiltLCAO; +namespace elecstate +{ +struct fenergy; +class ElecState; +} // namespace elecstate +namespace module_dm +{ +template +class DensityMatrix; } -namespace ModuleESolver { - template class ESolver_KS_LCAO; +namespace hamilt +{ +template +class HContainer; +template +class AtomPair; +template +class BaseMatrix; +template +class HamiltLCAO; +} // namespace hamilt +namespace ModuleESolver +{ +template +class ESolver_KS_LCAO; } namespace py = pybind11; @@ -53,7 +62,7 @@ namespace py_esolver */ class PyChargeAccessor { -public: + public: PyChargeAccessor() = default; /// Set internal pointers from Charge object @@ -72,20 +81,32 @@ class PyChargeAccessor py::array_t get_rho_core() const; /// Get number of spin channels - int get_nspin() const { return nspin_; } + int get_nspin() const + { + return nspin_; + } /// Get number of real-space grid points - int get_nrxx() const { return nrxx_; } + int get_nrxx() const + { + return nrxx_; + } /// Get number of G-vectors - int get_ngmc() const { return ngmc_; } + int get_ngmc() const + { + return ngmc_; + } /// Check if data is valid - bool is_valid() const { return (chr_ptr_ != nullptr || rho_ptr_ != nullptr) && nspin_ > 0; } + bool is_valid() const + { + return (chr_ptr_ != nullptr || rho_ptr_ != nullptr) && nspin_ > 0; + } -private: + private: const Charge* chr_ptr_ = nullptr; - const double* rho_ptr_ = nullptr; // Direct pointer for compatibility + const double* rho_ptr_ = nullptr; // Direct pointer for compatibility int nspin_ = 0; int nrxx_ = 0; int ngmc_ = 0; @@ -98,45 +119,74 @@ class PyChargeAccessor */ class PyEnergyAccessor { -public: + public: PyEnergyAccessor() = default; /// Set from fenergy structure void set_from_fenergy(const elecstate::fenergy* f_en); /// Set energies directly (for compatibility) - void set_energies(double etot, double eband, double hartree, - double etxc, double ewald, double demet, - double exx, double evdw); + void set_energies(double etot, + double eband, + double hartree, + double etxc, + double ewald, + double demet, + double exx, + double evdw); /// Get total energy (Ry) - double get_etot() const { return etot_; } + double get_etot() const + { + return etot_; + } /// Get band energy (Ry) - double get_eband() const { return eband_; } + double get_eband() const + { + return eband_; + } /// Get Hartree energy (Ry) - double get_hartree_energy() const { return hartree_energy_; } + double get_hartree_energy() const + { + return hartree_energy_; + } /// Get exchange-correlation energy (Ry) - double get_etxc() const { return etxc_; } + double get_etxc() const + { + return etxc_; + } /// Get Ewald energy (Ry) - double get_ewald_energy() const { return ewald_energy_; } + double get_ewald_energy() const + { + return ewald_energy_; + } /// Get -TS term for metals (Ry) - double get_demet() const { return demet_; } + double get_demet() const + { + return demet_; + } /// Get exact exchange energy (Ry) - double get_exx() const { return exx_; } + double get_exx() const + { + return exx_; + } /// Get van der Waals energy (Ry) - double get_evdw() const { return evdw_; } + double get_evdw() const + { + return evdw_; + } /// Get all energies as a dictionary py::dict get_all_energies() const; -private: + private: double etot_ = 0.0; double eband_ = 0.0; double hartree_energy_ = 0.0; @@ -155,7 +205,7 @@ class PyEnergyAccessor template class PyHamiltonianAccessor { -public: + public: PyHamiltonianAccessor() = default; /// Set from HamiltLCAO object @@ -171,13 +221,22 @@ class PyHamiltonianAccessor void set_Sk_data(int ik, const TK* data, int nrow, int ncol); /// Get number of basis functions - int get_nbasis() const { return nbasis_; } + int get_nbasis() const + { + return nbasis_; + } /// Get number of k-points - int get_nks() const { return nks_; } + int get_nks() const + { + return nks_; + } /// Get local matrix size (for 2D distribution) - int get_nloc() const { return nloc_; } + int get_nloc() const + { + return nloc_; + } /// Get H(k) matrix for specific k-point (local part in 2D distribution) py::array_t get_Hk(int ik) const; @@ -198,9 +257,12 @@ class PyHamiltonianAccessor py::dict get_SR() const; /// Check if data is valid - bool is_valid() const { return (hamilt_ptr_ != nullptr || nbasis_ > 0) && nks_ > 0; } + bool is_valid() const + { + return (hamilt_ptr_ != nullptr || nbasis_ > 0) && nks_ > 0; + } -private: + private: hamilt::HamiltLCAO* hamilt_ptr_ = nullptr; const Parallel_Orbitals* pv_ = nullptr; int nbasis_ = 0; @@ -223,7 +285,7 @@ class PyHamiltonianAccessor template class PyDensityMatrixAccessor { -public: + public: PyDensityMatrixAccessor() = default; /// Set from DensityMatrix object @@ -245,18 +307,30 @@ class PyDensityMatrixAccessor py::dict get_DMR() const; /// Get number of k-points - int get_nks() const { return nks_; } + int get_nks() const + { + return nks_; + } /// Get matrix row dimension - int get_nrow() const { return nrow_; } + int get_nrow() const + { + return nrow_; + } /// Get matrix column dimension - int get_ncol() const { return ncol_; } + int get_ncol() const + { + return ncol_; + } /// Check if data is valid - bool is_valid() const { return (dm_ptr_ != nullptr || nks_ > 0); } + bool is_valid() const + { + return (dm_ptr_ != nullptr || nks_ > 0); + } -private: + private: module_dm::DensityMatrix* dm_ptr_ = nullptr; int nks_ = 0; int nrow_ = 0; @@ -279,7 +353,7 @@ class PyDensityMatrixAccessor template class PyESolverLCAO { -public: + public: PyESolverLCAO(); ~PyESolverLCAO(); @@ -308,16 +382,28 @@ class PyESolverLCAO // ==================== Status Queries ==================== /// Check if SCF is converged - bool is_converged() const { return conv_esolver_; } + bool is_converged() const + { + return conv_esolver_; + } /// Get current iteration number - int get_niter() const { return niter_; } + int get_niter() const + { + return niter_; + } /// Get charge density difference (drho) - double get_drho() const { return drho_; } + double get_drho() const + { + return drho_; + } /// Get current SCF step - int get_istep() const { return istep_; } + int get_istep() const + { + return istep_; + } // ==================== Data Accessors ==================== @@ -403,8 +489,7 @@ class PyESolverLCAO { if (scf_controller_) { - return static_cast*>( - scf_controller_->get_diagonalizer()); + return static_cast*>(scf_controller_->get_diagonalizer()); } return nullptr; } @@ -443,7 +528,7 @@ class PyESolverLCAO } } -private: + private: // Internal state bool initialized_ = false; bool scf_started_ = false; diff --git a/python/pyabacus/src/ModuleNAO/py_m_nao.cpp b/python/pyabacus/src/ModuleNAO/py_m_nao.cpp index c6c21e5fdf5..3ae4e6a089c 100644 --- a/python/pyabacus/src/ModuleNAO/py_m_nao.cpp +++ b/python/pyabacus/src/ModuleNAO/py_m_nao.cpp @@ -1,12 +1,11 @@ -#include -#include -#include - +#include "../utils/pybind_utils.h" #include "source_base/vector3.h" #include "source_basis/module_nao/radial_collection.h" #include "source_basis/module_nao/two_center_integrator.h" -#include "../utils/pybind_utils.h" +#include +#include +#include namespace py = pybind11; using namespace pybind11::literals; @@ -447,22 +446,18 @@ void bind_m_nao(py::module& m) .def_property_readonly("pr", &NumericalRadial::pr) .def_property_readonly("pk", &NumericalRadial::pk) .def_property_readonly("sbt", &NumericalRadial::sbt) - .def_property_readonly("rgrid", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.rgrid(), static_cast(self.nr())); - }) - .def_property_readonly("kgrid", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.kgrid(), static_cast(self.nk())); - }) - .def_property_readonly("rvalue", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.rvalue(), static_cast(self.nr())); - }) - .def_property_readonly("kvalue", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.kvalue(), static_cast(self.nk())); - }) + .def_property_readonly( + "rgrid", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.rgrid(), static_cast(self.nr())); }) + .def_property_readonly( + "kgrid", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.kgrid(), static_cast(self.nk())); }) + .def_property_readonly( + "rvalue", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.rvalue(), static_cast(self.nr())); }) + .def_property_readonly( + "kvalue", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.kvalue(), static_cast(self.nk())); }) .def_property_readonly("is_fft_compliant", overload_cast_<>()(&NumericalRadial::is_fft_compliant, py::const_)); } diff --git a/python/pyabacus/src/hsolver/diago_adapter.hpp b/python/pyabacus/src/hsolver/diago_adapter.hpp index 0493e474a69..2d6e652ba46 100644 --- a/python/pyabacus/src/hsolver/diago_adapter.hpp +++ b/python/pyabacus/src/hsolver/diago_adapter.hpp @@ -9,23 +9,24 @@ #ifndef PYABACUS_HSOLVER_DIAGO_ADAPTER_HPP #define PYABACUS_HSOLVER_DIAGO_ADAPTER_HPP +#include "../utils/pybind_utils.h" +#include "diago_traits.hpp" + #include -#include #include - -#include -#include +#include #include #include +#include +#include #include -#include "diago_traits.hpp" -#include "../utils/pybind_utils.h" - namespace py = pybind11; -namespace pyabacus { -namespace hsolver { +namespace pyabacus +{ +namespace hsolver +{ // ============================================================================ // PyDiagoDavid Adapter @@ -36,13 +37,12 @@ namespace hsolver { */ class PyDiagoDavidAdapter { -public: + public: using Traits = DiagoDavidTraits; using T = typename Traits::T; using SolverType = typename Traits::SolverType; - PyDiagoDavidAdapter(int nbasis, int nband) - : nbasis_(nbasis), nband_(nband) + PyDiagoDavidAdapter(int nbasis, int nband) : nbasis_(nbasis), nband_(nband) { storage_.allocate(nbasis, nband); } @@ -51,9 +51,7 @@ class PyDiagoDavidAdapter PyDiagoDavidAdapter& operator=(const PyDiagoDavidAdapter&) = delete; PyDiagoDavidAdapter(PyDiagoDavidAdapter&& other) noexcept - : storage_(std::move(other.storage_)) - , nbasis_(other.nbasis_) - , nband_(other.nband_) + : storage_(std::move(other.storage_)), nbasis_(other.nbasis_), nband_(other.nband_) { } @@ -77,36 +75,22 @@ class PyDiagoDavidAdapter return storage_.get_eigenvalue(); } - int diag( - std::function(py::array_t)> mm_op, - std::vector& precond_vec, - int dav_ndim, - double tol, - std::vector& diag_ethr, - int max_iter, - ::hsolver::diag_comm_info comm_info) + int diag(std::function(py::array_t)> mm_op, + std::vector& precond_vec, + int dav_ndim, + double tol, + std::vector& diag_ethr, + int max_iter, + ::hsolver::diag_comm_info comm_info) { const PyHSOperator op(mm_op); - solver_ = std::make_unique( - precond_vec.data(), - nband_, - nbasis_, - dav_ndim, - comm_info - ); - - return solver_->diag( - op, - nbasis_, - storage_.psi_ptr(), - storage_.eigenvalue_ptr(), - diag_ethr, - max_iter - ); + solver_ = std::make_unique(precond_vec.data(), nband_, nbasis_, dav_ndim, comm_info); + + return solver_->diag(op, nbasis_, storage_.psi_ptr(), storage_.eigenvalue_ptr(), diag_ethr, max_iter); } -private: + private: RawPointerStorage storage_; std::unique_ptr solver_; int nbasis_; @@ -122,13 +106,12 @@ class PyDiagoDavidAdapter */ class PyDiagoDavSubspaceAdapter { -public: + public: using Traits = DiagoDavSubspaceTraits; using T = typename Traits::T; using SolverType = typename Traits::SolverType; - PyDiagoDavSubspaceAdapter(int nbasis, int nband) - : nbasis_(nbasis), nband_(nband) + PyDiagoDavSubspaceAdapter(int nbasis, int nband) : nbasis_(nbasis), nband_(nband) { storage_.allocate(nbasis, nband); } @@ -137,9 +120,7 @@ class PyDiagoDavSubspaceAdapter PyDiagoDavSubspaceAdapter& operator=(const PyDiagoDavSubspaceAdapter&) = delete; PyDiagoDavSubspaceAdapter(PyDiagoDavSubspaceAdapter&& other) noexcept - : storage_(std::move(other.storage_)) - , nbasis_(other.nbasis_) - , nband_(other.nband_) + : storage_(std::move(other.storage_)), nbasis_(other.nbasis_), nband_(other.nband_) { } @@ -163,44 +144,34 @@ class PyDiagoDavSubspaceAdapter return storage_.get_eigenvalue(); } - int diag( - std::function(py::array_t)> mm_op, - std::vector& precond_vec, - int dav_ndim, - double tol, - int max_iter, - bool need_subspace, - std::vector& diag_ethr, - bool scf_type, - ::hsolver::diag_comm_info comm_info, - int diag_subspace, - int nb2d) + int diag(std::function(py::array_t)> mm_op, + std::vector& precond_vec, + int dav_ndim, + double tol, + int max_iter, + bool need_subspace, + std::vector& diag_ethr, + bool scf_type, + ::hsolver::diag_comm_info comm_info, + int diag_subspace, + int nb2d) { const PyHSOperator op(mm_op); - solver_ = std::make_unique( - precond_vec, - nband_, - nbasis_, - dav_ndim, - tol, - max_iter, - comm_info, - diag_subspace, - nb2d - ); - - return solver_->diag( - op, - storage_.psi_ptr(), - nbasis_, - storage_.eigenvalue_ptr(), - diag_ethr, - scf_type - ); + solver_ = std::make_unique(precond_vec, + nband_, + nbasis_, + dav_ndim, + tol, + max_iter, + comm_info, + diag_subspace, + nb2d); + + return solver_->diag(op, storage_.psi_ptr(), nbasis_, storage_.eigenvalue_ptr(), diag_ethr, scf_type); } -private: + private: RawPointerStorage storage_; std::unique_ptr solver_; int nbasis_; @@ -217,13 +188,12 @@ class PyDiagoDavSubspaceAdapter */ class PyDiagoCGAdapter { -public: + public: using Traits = DiagoCGTraits; using T = typename Traits::T; using SolverType = typename Traits::SolverType; - PyDiagoCGAdapter(int dim, int num_eigs) - : dim_(dim), num_eigs_(num_eigs) + PyDiagoCGAdapter(int dim, int num_eigs) : dim_(dim), num_eigs_(num_eigs) { storage_.allocate(dim, num_eigs); } @@ -232,9 +202,7 @@ class PyDiagoCGAdapter PyDiagoCGAdapter& operator=(const PyDiagoCGAdapter&) = delete; PyDiagoCGAdapter(PyDiagoCGAdapter&& other) noexcept - : storage_(std::move(other.storage_)) - , dim_(other.dim_) - , num_eigs_(other.num_eigs_) + : storage_(std::move(other.storage_)), dim_(other.dim_), num_eigs_(other.num_eigs_) { } @@ -263,14 +231,13 @@ class PyDiagoCGAdapter storage_.set_preconditioner(prec_in); } - void diag( - std::function(py::array_t)> mm_op, - int diag_ndim, - double tol, - const std::vector& diag_ethr, - bool need_subspace, - bool scf_type, - int nproc_in_pool = 1) + void diag(std::function(py::array_t)> mm_op, + int diag_ndim, + double tol, + const std::vector& diag_ethr, + bool need_subspace, + bool scf_type, + int nproc_in_pool = 1) { const std::string basis_type = "pw"; const std::string calculation = scf_type ? "scf" : "nscf"; @@ -281,27 +248,23 @@ class PyDiagoCGAdapter // Do nothing - placeholder }; - solver_ = std::make_unique( - basis_type, - calculation, - need_subspace, - subspace_func, - tol, - diag_ndim, - nproc_in_pool - ); - - solver_->diag( - hpsi_func, - spsi_func, - *storage_.psi_tensor(), - *storage_.eig_tensor(), - diag_ethr, - *storage_.prec_tensor() - ); + solver_ = std::make_unique(basis_type, + calculation, + need_subspace, + subspace_func, + tol, + diag_ndim, + nproc_in_pool); + + solver_->diag(hpsi_func, + spsi_func, + *storage_.psi_tensor(), + *storage_.eig_tensor(), + diag_ethr, + *storage_.prec_tensor()); } -private: + private: TensorStorage storage_; std::unique_ptr solver_; int dim_; @@ -313,7 +276,8 @@ class PyDiagoCGAdapter // Backward Compatibility Aliases // ============================================================================ -namespace py_hsolver_compat { +namespace py_hsolver_compat +{ using PyDiagoDavid = PyDiagoDavidAdapter; using PyDiagoDavSubspace = PyDiagoDavSubspaceAdapter; diff --git a/python/pyabacus/src/hsolver/diago_traits.hpp b/python/pyabacus/src/hsolver/diago_traits.hpp index c73b6f3bea9..b183aced8f1 100644 --- a/python/pyabacus/src/hsolver/diago_traits.hpp +++ b/python/pyabacus/src/hsolver/diago_traits.hpp @@ -10,22 +10,21 @@ #ifndef PYABACUS_HSOLVER_DIAGO_TRAITS_HPP #define PYABACUS_HSOLVER_DIAGO_TRAITS_HPP +#include "source_base/module_device/memory_op.h" +#include "source_hsolver/diago_cg.h" +#include "source_hsolver/diago_dav_subspace.h" +#include "source_hsolver/diago_david.h" +#include "source_hsolver/hs_operator.h" + #include -#include -#include #include - -#include -#include +#include #include #include +#include +#include #include - -#include "source_hsolver/diago_david.h" -#include "source_hsolver/hs_operator.h" -#include "source_hsolver/diago_dav_subspace.h" -#include "source_hsolver/diago_cg.h" -#include "source_base/module_device/memory_op.h" +#include #ifdef __ENABLE_ATEN #include @@ -35,8 +34,10 @@ namespace py = pybind11; -namespace pyabacus { -namespace hsolver { +namespace pyabacus +{ +namespace hsolver +{ // ============================================================================ // Storage Policies @@ -50,7 +51,7 @@ namespace hsolver { template class RawPointerStorage { -public: + public: using value_type = T; RawPointerStorage() = default; @@ -70,10 +71,7 @@ class RawPointerStorage // Move semantics RawPointerStorage(RawPointerStorage&& other) noexcept - : psi_(other.psi_) - , eigenvalue_(other.eigenvalue_) - , nbasis_(other.nbasis_) - , nband_(other.nband_) + : psi_(other.psi_), eigenvalue_(other.eigenvalue_), nbasis_(other.nbasis_), nband_(other.nband_) { other.psi_ = nullptr; other.eigenvalue_ = nullptr; @@ -143,12 +141,24 @@ class RawPointerStorage } // Accessors for solver - T* psi_ptr() { return psi_; } - double* eigenvalue_ptr() { return eigenvalue_; } - int nbasis() const { return nbasis_; } - int nband() const { return nband_; } + T* psi_ptr() + { + return psi_; + } + double* eigenvalue_ptr() + { + return eigenvalue_; + } + int nbasis() const + { + return nbasis_; + } + int nband() const + { + return nband_; + } -private: + private: void cleanup() { if (psi_ != nullptr) @@ -178,7 +188,7 @@ class RawPointerStorage template class TensorStorage { -public: + public: using value_type = T; TensorStorage() = default; @@ -197,11 +207,7 @@ class TensorStorage // Move semantics TensorStorage(TensorStorage&& other) noexcept - : psi_(other.psi_) - , eig_(other.eig_) - , prec_(other.prec_) - , dim_(other.dim_) - , num_eigs_(other.num_eigs_) + : psi_(other.psi_), eig_(other.eig_), prec_(other.prec_), dim_(other.dim_), num_eigs_(other.num_eigs_) { other.psi_ = nullptr; other.eig_ = nullptr; @@ -234,12 +240,10 @@ class TensorStorage py::buffer_info buf = psi_in.request(); T* ptr = static_cast(buf.ptr); - psi_ = new ct::TensorMap( - ptr, - ct::DataType::DT_COMPLEX_DOUBLE, - ct::DeviceType::CpuDevice, - ct::TensorShape({num_eigs_, dim_}) - ); + psi_ = new ct::TensorMap(ptr, + ct::DataType::DT_COMPLEX_DOUBLE, + ct::DeviceType::CpuDevice, + ct::TensorShape({num_eigs_, dim_})); } py::array_t get_psi() const @@ -281,22 +285,32 @@ class TensorStorage py::buffer_info buf = prec_in.request(); double* ptr = static_cast(buf.ptr); - prec_ = new ct::TensorMap( - ptr, - ct::DataType::DT_DOUBLE, - ct::DeviceType::CpuDevice, - ct::TensorShape({dim_}) - ); + prec_ = new ct::TensorMap(ptr, ct::DataType::DT_DOUBLE, ct::DeviceType::CpuDevice, ct::TensorShape({dim_})); } // Accessors for solver - ct::Tensor* psi_tensor() { return psi_; } - ct::Tensor* eig_tensor() { return eig_; } - ct::Tensor* prec_tensor() { return prec_; } - int dim() const { return dim_; } - int num_eigs() const { return num_eigs_; } + ct::Tensor* psi_tensor() + { + return psi_; + } + ct::Tensor* eig_tensor() + { + return eig_; + } + ct::Tensor* prec_tensor() + { + return prec_; + } + int dim() const + { + return dim_; + } + int num_eigs() const + { + return num_eigs_; + } -private: + private: void cleanup() { if (psi_ != nullptr) @@ -338,12 +352,11 @@ struct DiagoDavidTraits using StoragePolicy = RawPointerStorage; static constexpr const char* name = "diago_david"; - static constexpr bool uses_f_style = true; // Column-major arrays + static constexpr bool uses_f_style = true; // Column-major arrays static constexpr bool has_preconditioner = true; // Memory synchronization operation - using syncmem_op = base_device::memory::synchronize_memory_op< - T, base_device::DEVICE_CPU, base_device::DEVICE_CPU>; + using syncmem_op = base_device::memory::synchronize_memory_op; }; /** @@ -356,12 +369,11 @@ struct DiagoDavSubspaceTraits using StoragePolicy = RawPointerStorage; static constexpr const char* name = "diago_dav_subspace"; - static constexpr bool uses_f_style = true; // Column-major arrays + static constexpr bool uses_f_style = true; // Column-major arrays static constexpr bool has_preconditioner = true; // Memory synchronization operation - using syncmem_op = base_device::memory::synchronize_memory_op< - T, base_device::DEVICE_CPU, base_device::DEVICE_CPU>; + using syncmem_op = base_device::memory::synchronize_memory_op; }; #ifdef __ENABLE_ATEN @@ -375,12 +387,11 @@ struct DiagoCGTraits using StoragePolicy = TensorStorage; static constexpr const char* name = "diago_cg"; - static constexpr bool uses_f_style = false; // Row-major arrays + static constexpr bool uses_f_style = false; // Row-major arrays static constexpr bool has_preconditioner = true; // Memory synchronization operation for tensor interface - using syncmem_op = base_device::memory::synchronize_memory_op< - T, base_device::DEVICE_CPU, base_device::DEVICE_CPU>; + using syncmem_op = base_device::memory::synchronize_memory_op; }; #endif // __ENABLE_ATEN @@ -397,9 +408,8 @@ struct DiagoCGTraits template class PyHSOperator : public ::hsolver::HSOperator { -public: - explicit PyHSOperator(std::function(py::array_t)> mm_op) - : mm_op_(std::move(mm_op)) + public: + explicit PyHSOperator(std::function(py::array_t)> mm_op) : mm_op_(std::move(mm_op)) { } @@ -429,7 +439,7 @@ class PyHSOperator : public ::hsolver::HSOperator std::copy(x, x + static_cast(nvec) * ld, sx); } -private: + private: std::function(py::array_t)> mm_op_; }; @@ -438,8 +448,7 @@ class PyHSOperator : public ::hsolver::HSOperator * @brief Create hpsi_func lambda for tensor interface */ template -auto make_hpsi_func_tensor( - std::function(py::array_t)> mm_op) +auto make_hpsi_func_tensor(std::function(py::array_t)> mm_op) { return [mm_op](const ct::Tensor& psi_in, ct::Tensor& hpsi_out) { const auto ndim = psi_in.shape().ndim(); @@ -477,11 +486,7 @@ auto make_spsi_func_tensor_identity() REQUIRES_OK(ndim <= 2, "dims of psi_in should be less than or equal to 2"); const int nrow = ndim == 1 ? psi_in.NumElements() : psi_in.shape().dim_size(1); const int nbands = ndim == 1 ? 1 : psi_in.shape().dim_size(0); - syncmem_op()( - spsi_out.data(), - psi_in.data(), - static_cast(nrow * nbands) - ); + syncmem_op()(spsi_out.data(), psi_in.data(), static_cast(nrow * nbands)); }; } #endif // __ENABLE_ATEN diff --git a/python/pyabacus/src/hsolver/py_diago_cg.hpp b/python/pyabacus/src/hsolver/py_diago_cg.hpp index 0ab2cb25dab..487573851cb 100644 --- a/python/pyabacus/src/hsolver/py_diago_cg.hpp +++ b/python/pyabacus/src/hsolver/py_diago_cg.hpp @@ -1,22 +1,20 @@ #ifndef PYTHON_PYABACUS_SRC_PY_DIAGO_CG_HPP #define PYTHON_PYABACUS_SRC_PY_DIAGO_CG_HPP +#include "source_base/module_device/memory_op.h" +#include "source_hsolver/diago_cg.h" + +#include +#include +#include #include #include - -#include #include #include #include +#include #include -#include -#include -#include - -#include "source_hsolver/diago_cg.h" -#include "source_base/module_device/memory_op.h" - namespace py = pybind11; namespace py_hsolver @@ -24,8 +22,10 @@ namespace py_hsolver class PyDiagoCG { -public: - PyDiagoCG(int dim, int num_eigs) : dim{dim}, num_eigs{num_eigs} { } + public: + PyDiagoCG(int dim, int num_eigs) : dim{dim}, num_eigs{num_eigs} + { + } PyDiagoCG(const PyDiagoCG&) = delete; PyDiagoCG& operator=(const PyDiagoCG&) = delete; PyDiagoCG(PyDiagoCG&& other) @@ -37,9 +37,9 @@ class PyDiagoCG other.eig = nullptr; } - ~PyDiagoCG() + ~PyDiagoCG() { - if (psi != nullptr) + if (psi != nullptr) { delete psi; psi = nullptr; @@ -52,19 +52,20 @@ class PyDiagoCG } } - void init_eig() + void init_eig() { eig = new ct::Tensor(ct::DataType::DT_DOUBLE, {num_eigs}); eig->zero(); } - py::array_t get_eig() + py::array_t get_eig() { py::array_t eig_out(eig->NumElements()); py::buffer_info eig_buf = eig_out.request(); double* eig_out_ptr = static_cast(eig_buf.ptr); - if (eig == nullptr) { + if (eig == nullptr) + { throw std::runtime_error("eig is not initialized"); } double* eig_ptr = eig->data(); @@ -78,12 +79,10 @@ class PyDiagoCG py::buffer_info psi_buf = psi_in.request(); std::complex* psi_ptr = static_cast*>(psi_buf.ptr); - psi = new ct::TensorMap( - psi_ptr, - ct::DataType::DT_COMPLEX_DOUBLE, - ct::DeviceType::CpuDevice, - ct::TensorShape({num_eigs, dim}) - ); + psi = new ct::TensorMap(psi_ptr, + ct::DataType::DT_COMPLEX_DOUBLE, + ct::DeviceType::CpuDevice, + ct::TensorShape({num_eigs, dim})); } py::array_t> get_psi() @@ -92,7 +91,8 @@ class PyDiagoCG py::buffer_info psi_buf = psi_out.request(); std::complex* psi_out_ptr = static_cast*>(psi_buf.ptr); - if (psi == nullptr) { + if (psi == nullptr) + { throw std::runtime_error("psi is not initialized"); } std::complex* psi_ptr = psi->data>(); @@ -106,12 +106,7 @@ class PyDiagoCG py::buffer_info prec_buf = prec_in.request(); double* prec_ptr = static_cast(prec_buf.ptr); - prec = new ct::TensorMap( - prec_ptr, - ct::DataType::DT_DOUBLE, - ct::DeviceType::CpuDevice, - ct::TensorShape({dim}) - ); + prec = new ct::TensorMap(prec_ptr, ct::DataType::DT_DOUBLE, ct::DeviceType::CpuDevice, ct::TensorShape({dim})); } void diag(std::function>(py::array_t>)> mm_op, @@ -120,15 +115,15 @@ class PyDiagoCG const std::vector& diag_ethr, bool need_subspace, bool scf_type, - int nproc_in_pool = 1 - ) { + int nproc_in_pool = 1) + { const std::string basis_type = "pw"; const std::string calculation = scf_type ? "scf" : "nscf"; - auto hpsi_func = [mm_op] (const ct::Tensor& psi_in, ct::Tensor& hpsi_out) { + auto hpsi_func = [mm_op](const ct::Tensor& psi_in, ct::Tensor& hpsi_out) { const auto ndim = psi_in.shape().ndim(); REQUIRES_OK(ndim <= 2, "dims of psi_in should be less than or equal to 2"); - const int nvec = ndim == 1 ? 1 : psi_in.shape().dim_size(0); + const int nvec = ndim == 1 ? 1 : psi_in.shape().dim_size(0); const int ld_psi = ndim == 1 ? psi_in.NumElements() : psi_in.shape().dim_size(1); // Note: numpy's py::array_t is row-major, and @@ -136,7 +131,9 @@ class PyDiagoCG py::array_t> psi({ld_psi, nvec}); py::buffer_info psi_buf = psi.request(); std::complex* psi_ptr = static_cast*>(psi_buf.ptr); - std::copy(psi_in.data>(), psi_in.data>() + nvec * ld_psi, psi_ptr); + std::copy(psi_in.data>(), + psi_in.data>() + nvec * ld_psi, + psi_ptr); py::array_t> hpsi = mm_op(psi); @@ -147,32 +144,28 @@ class PyDiagoCG auto subspace_func = [](const ct::Tensor& psi_in, ct::Tensor& psi_out, const bool S_orth) { /*do nothing*/ }; - auto spsi_func = [this] (const ct::Tensor& psi_in, ct::Tensor& spsi_out) { + auto spsi_func = [this](const ct::Tensor& psi_in, ct::Tensor& spsi_out) { const auto ndim = psi_in.shape().ndim(); REQUIRES_OK(ndim <= 2, "dims of psi_in should be less than or equal to 2"); - const int nrow = ndim == 1 ? psi_in.NumElements() : psi_in.shape().dim_size(1); + const int nrow = ndim == 1 ? psi_in.NumElements() : psi_in.shape().dim_size(1); const int nbands = ndim == 1 ? 1 : psi_in.shape().dim_size(0); - syncmem_z2z_h2h_op()( - spsi_out.data>(), - psi_in.data>(), - static_cast(nrow * nbands) - ); + syncmem_z2z_h2h_op()(spsi_out.data>(), + psi_in.data>(), + static_cast(nrow * nbands)); }; - cg = std::make_unique, base_device::DEVICE_CPU>>( - basis_type, - calculation, - need_subspace, - subspace_func, - tol, - diag_ndim, - nproc_in_pool - ); + cg = std::make_unique, base_device::DEVICE_CPU>>(basis_type, + calculation, + need_subspace, + subspace_func, + tol, + diag_ndim, + nproc_in_pool); cg->diag(hpsi_func, spsi_func, *psi, *eig, diag_ethr, *prec); } -private: + private: base_device::DEVICE_CPU* ctx = {}; int dim; diff --git a/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp b/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp index c258de44061..19969c8ef0f 100644 --- a/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp +++ b/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp @@ -1,17 +1,16 @@ #ifndef PYTHON_PYABACUS_SRC_PY_DIAGO_DAV_SUBSPACE_HPP #define PYTHON_PYABACUS_SRC_PY_DIAGO_DAV_SUBSPACE_HPP +#include "source_hsolver/diago_dav_subspace.h" + #include #include - -#include #include #include #include +#include #include -#include "source_hsolver/diago_dav_subspace.h" - namespace py = pybind11; namespace py_hsolver @@ -19,7 +18,7 @@ namespace py_hsolver class PyDiagoDavSubspace { -public: + public: PyDiagoDavSubspace(int nbasis, int nband) : nbasis(nbasis), nband(nband) { psi = new std::complex[nbasis * nband]; @@ -39,15 +38,15 @@ class PyDiagoDavSubspace ~PyDiagoDavSubspace() { - if (psi != nullptr) - { - delete[] psi; - psi = nullptr; + if (psi != nullptr) + { + delete[] psi; + psi = nullptr; } - if (eigenvalue != nullptr) - { - delete[] eigenvalue; - eigenvalue = nullptr; + if (eigenvalue != nullptr) + { + delete[] eigenvalue; + eigenvalue = nullptr; } } @@ -73,7 +72,7 @@ class PyDiagoDavSubspace psi_out_ptr[i] = psi[i]; } - return psi_out; + return psi_out; } void init_eigenvalue() @@ -99,60 +98,53 @@ class PyDiagoDavSubspace return eigenvalue_out; } - int diag( - std::function>(py::array_t>)> mm_op, - std::vector& precond_vec, - int dav_ndim, - double tol, - int max_iter, - bool need_subspace, - std::vector& diag_ethr, - bool scf_type, - hsolver::diag_comm_info comm_info, - int diag_subspace, - int nb2d - ) { - auto hpsi_func = [mm_op] ( - std::complex *psi_in, - std::complex *hpsi_out, - const int ld_psi, - const int nvec - ) { - // Note: numpy's py::array_t is row-major, but - // our raw pointer-array is column-major - py::array_t, py::array::f_style> psi({ld_psi, nvec}); - py::buffer_info psi_buf = psi.request(); - std::complex* psi_ptr = static_cast*>(psi_buf.ptr); - std::copy(psi_in, psi_in + nvec * ld_psi, psi_ptr); - - py::array_t, py::array::f_style> hpsi = mm_op(psi); - - py::buffer_info hpsi_buf = hpsi.request(); - std::complex* hpsi_ptr = static_cast*>(hpsi_buf.ptr); - std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out); - }; + int diag(std::function>(py::array_t>)> mm_op, + std::vector& precond_vec, + int dav_ndim, + double tol, + int max_iter, + bool need_subspace, + std::vector& diag_ethr, + bool scf_type, + hsolver::diag_comm_info comm_info, + int diag_subspace, + int nb2d) + { + auto hpsi_func + = [mm_op](std::complex* psi_in, std::complex* hpsi_out, const int ld_psi, const int nvec) { + // Note: numpy's py::array_t is row-major, but + // our raw pointer-array is column-major + py::array_t, py::array::f_style> psi({ld_psi, nvec}); + py::buffer_info psi_buf = psi.request(); + std::complex* psi_ptr = static_cast*>(psi_buf.ptr); + std::copy(psi_in, psi_in + nvec * ld_psi, psi_ptr); + + py::array_t, py::array::f_style> hpsi = mm_op(psi); + + py::buffer_info hpsi_buf = hpsi.request(); + std::complex* hpsi_ptr = static_cast*>(hpsi_buf.ptr); + std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out); + }; auto spsi_func = [this](const std::complex* psi_in, std::complex* spsi_out, const int ld_psi, const int nvec) { syncmem_op()(spsi_out, psi_in, static_cast(ld_psi * nvec)); }; - obj = std::make_unique, base_device::DEVICE_CPU>>( - precond_vec, - nband, - nbasis, - dav_ndim, - tol, - max_iter, - comm_info, - diag_subspace, - nb2d - ); + obj = std::make_unique, base_device::DEVICE_CPU>>(precond_vec, + nband, + nbasis, + dav_ndim, + tol, + max_iter, + comm_info, + diag_subspace, + nb2d); return obj->diag(hpsi_func, spsi_func, psi, nbasis, eigenvalue, diag_ethr, scf_type); } -private: + private: std::complex* psi = nullptr; double* eigenvalue = nullptr; diff --git a/python/pyabacus/src/hsolver/py_diago_david.hpp b/python/pyabacus/src/hsolver/py_diago_david.hpp index 40261b01b93..e1c5d2641da 100644 --- a/python/pyabacus/src/hsolver/py_diago_david.hpp +++ b/python/pyabacus/src/hsolver/py_diago_david.hpp @@ -1,17 +1,16 @@ #ifndef PYTHON_PYABACUS_SRC_PY_DIAGO_DAVID_HPP #define PYTHON_PYABACUS_SRC_PY_DIAGO_DAVID_HPP +#include "source_hsolver/diago_david.h" + #include #include - -#include #include #include #include +#include #include -#include "source_hsolver/diago_david.h" - namespace py = pybind11; namespace py_hsolver @@ -19,7 +18,7 @@ namespace py_hsolver class PyDiagoDavid { -public: + public: PyDiagoDavid(int nbasis, int nband) : nbasis(nbasis), nband(nband) { psi = new std::complex[nbasis * nband]; @@ -39,15 +38,15 @@ class PyDiagoDavid ~PyDiagoDavid() { - if (psi != nullptr) - { - delete[] psi; - psi = nullptr; + if (psi != nullptr) + { + delete[] psi; + psi = nullptr; } - if (eigenvalue != nullptr) - { - delete[] eigenvalue; - eigenvalue = nullptr; + if (eigenvalue != nullptr) + { + delete[] eigenvalue; + eigenvalue = nullptr; } } @@ -73,7 +72,7 @@ class PyDiagoDavid psi_out_ptr[i] = psi[i]; } - return psi_out; + return psi_out; } void init_eigenvalue() @@ -99,56 +98,46 @@ class PyDiagoDavid return eigenvalue_out; } - int diag( - std::function>(py::array_t>)> mm_op, - std::vector& precond_vec, - int dav_ndim, - double tol, - std::vector& diag_ethr, - int max_iter, - hsolver::diag_comm_info comm_info - ) { - auto hpsi_func = [mm_op] ( - std::complex *psi_in, - std::complex *hpsi_out, - const int ld_psi, - const int nvec - ) { - // Note: numpy's py::array_t is row-major, but - // our raw pointer-array is column-major - py::array_t, py::array::f_style> psi({ld_psi, nvec}); - py::buffer_info psi_buf = psi.request(); - std::complex* psi_ptr = static_cast*>(psi_buf.ptr); - std::copy(psi_in, psi_in + nvec * ld_psi, psi_ptr); - - py::array_t, py::array::f_style> hpsi = mm_op(psi); - - py::buffer_info hpsi_buf = hpsi.request(); - std::complex* hpsi_ptr = static_cast*>(hpsi_buf.ptr); - std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out); - }; - - auto spsi_func = [this] ( - const std::complex *psi_in, - std::complex *spsi_out, - const int nrow, - const int nbands - ) { - syncmem_op()(spsi_out, psi_in, static_cast(nbands * nrow)); - }; - - obj = std::make_unique, base_device::DEVICE_CPU>>( - precond_vec.data(), - nband, - nbasis, - dav_ndim, - comm_info - ); + int diag(std::function>(py::array_t>)> mm_op, + std::vector& precond_vec, + int dav_ndim, + double tol, + std::vector& diag_ethr, + int max_iter, + hsolver::diag_comm_info comm_info) + { + auto hpsi_func + = [mm_op](std::complex* psi_in, std::complex* hpsi_out, const int ld_psi, const int nvec) { + // Note: numpy's py::array_t is row-major, but + // our raw pointer-array is column-major + py::array_t, py::array::f_style> psi({ld_psi, nvec}); + py::buffer_info psi_buf = psi.request(); + std::complex* psi_ptr = static_cast*>(psi_buf.ptr); + std::copy(psi_in, psi_in + nvec * ld_psi, psi_ptr); + + py::array_t, py::array::f_style> hpsi = mm_op(psi); + + py::buffer_info hpsi_buf = hpsi.request(); + std::complex* hpsi_ptr = static_cast*>(hpsi_buf.ptr); + std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out); + }; + + auto spsi_func + = [this](const std::complex* psi_in, + std::complex* spsi_out, + const int nrow, + const int nbands) { syncmem_op()(spsi_out, psi_in, static_cast(nbands * nrow)); }; + + obj = std::make_unique, base_device::DEVICE_CPU>>(precond_vec.data(), + nband, + nbasis, + dav_ndim, + comm_info); return obj->diag(hpsi_func, spsi_func, nbasis, psi, eigenvalue, diag_ethr, max_iter); } -private: + private: std::complex* psi = nullptr; double* eigenvalue = nullptr; @@ -158,7 +147,8 @@ class PyDiagoDavid std::unique_ptr, base_device::DEVICE_CPU>> obj; base_device::DEVICE_CPU* ctx = {}; - using syncmem_op = base_device::memory::synchronize_memory_op, base_device::DEVICE_CPU, base_device::DEVICE_CPU>; + using syncmem_op = base_device::memory:: + synchronize_memory_op, base_device::DEVICE_CPU, base_device::DEVICE_CPU>; }; } // namespace py_hsolver diff --git a/python/pyabacus/src/hsolver/py_hsolver.cpp b/python/pyabacus/src/hsolver/py_hsolver.cpp index 02f91c9441e..71f4c9ce5a2 100644 --- a/python/pyabacus/src/hsolver/py_hsolver.cpp +++ b/python/pyabacus/src/hsolver/py_hsolver.cpp @@ -6,18 +6,17 @@ * using the unified adapter template approach. */ +#include "diago_adapter.hpp" +#include "source_base/kernels/math_kernel_op.h" +#include "source_base/module_device/types.h" +#include "source_hsolver/diago_dav_subspace.h" + #include #include -#include -#include #include #include - -#include "source_hsolver/diago_dav_subspace.h" -#include "source_base/kernels/math_kernel_op.h" -#include "source_base/module_device/types.h" - -#include "diago_adapter.hpp" +#include +#include namespace py = pybind11; using namespace pybind11::literals; @@ -34,7 +33,8 @@ void bind_hsolver(py::module& m) // Bind PyDiagoDavSubspace using adapter py::class_(m, "diago_dav_subspace") - .def(py::init(), R"pbdoc( + .def(py::init(), + R"pbdoc( Constructor of diago_dav_subspace, a class for diagonalizing a linear operator using the Davidson-Subspace Method. @@ -48,8 +48,12 @@ void bind_hsolver(py::module& m) The number of basis functions. nband : int The number of bands to be calculated. - )pbdoc", "nbasis"_a, "nband"_a) - .def("diag", &PyDiagoDavSubspaceAdapter::diag, R"pbdoc( + )pbdoc", + "nbasis"_a, + "nband"_a) + .def("diag", + &PyDiagoDavSubspaceAdapter::diag, + R"pbdoc( Diagonalize the linear operator using the Davidson-Subspace Method. Parameters @@ -85,20 +89,23 @@ void bind_hsolver(py::module& m) nb2d : int The block size in 2d block cyclic distribution if use elpa or scalapack. )pbdoc", - "mm_op"_a, - "precond_vec"_a, - "dav_ndim"_a, - "tol"_a, - "max_iter"_a, - "need_subspace"_a, - "diag_ethr"_a, - "scf_type"_a, - "comm_info"_a, - "diago_subspace"_a, - "nb2d"_a) - .def("set_psi", &PyDiagoDavSubspaceAdapter::set_psi, R"pbdoc( + "mm_op"_a, + "precond_vec"_a, + "dav_ndim"_a, + "tol"_a, + "max_iter"_a, + "need_subspace"_a, + "diag_ethr"_a, + "scf_type"_a, + "comm_info"_a, + "diago_subspace"_a, + "nb2d"_a) + .def("set_psi", + &PyDiagoDavSubspaceAdapter::set_psi, + R"pbdoc( Set the initial guess of the eigenvectors, i.e. the wave functions. - )pbdoc", "psi_in"_a) + )pbdoc", + "psi_in"_a) .def("get_psi", &PyDiagoDavSubspaceAdapter::get_psi, R"pbdoc( Get the eigenvectors. )pbdoc") @@ -111,7 +118,8 @@ void bind_hsolver(py::module& m) // Bind PyDiagoDavid using adapter py::class_(m, "diago_david") - .def(py::init(), R"pbdoc( + .def(py::init(), + R"pbdoc( Constructor of diago_david, a class for diagonalizing a linear operator using the Davidson Method. @@ -125,8 +133,12 @@ void bind_hsolver(py::module& m) The number of basis functions. nband : int The number of bands to be calculated. - )pbdoc", "nbasis"_a, "nband"_a) - .def("diag", &PyDiagoDavidAdapter::diag, R"pbdoc( + )pbdoc", + "nbasis"_a, + "nband"_a) + .def("diag", + &PyDiagoDavidAdapter::diag, + R"pbdoc( Diagonalize the linear operator using the Davidson Method. Parameters @@ -146,16 +158,19 @@ void bind_hsolver(py::module& m) max_iter : int The maximum number of iterations. )pbdoc", - "mm_op"_a, - "precond_vec"_a, - "dav_ndim"_a, - "tol"_a, - "diag_ethr"_a, - "max_iter"_a, - "comm_info"_a) - .def("set_psi", &PyDiagoDavidAdapter::set_psi, R"pbdoc( + "mm_op"_a, + "precond_vec"_a, + "dav_ndim"_a, + "tol"_a, + "diag_ethr"_a, + "max_iter"_a, + "comm_info"_a) + .def("set_psi", + &PyDiagoDavidAdapter::set_psi, + R"pbdoc( Set the initial guess of the eigenvectors, i.e. the wave functions. - )pbdoc", "psi_in"_a) + )pbdoc", + "psi_in"_a) .def("get_psi", &PyDiagoDavidAdapter::get_psi, R"pbdoc( Get the eigenvectors. )pbdoc") @@ -197,33 +212,38 @@ void bind_hsolver(py::module& m) Whether to use the SCF type, which is used to determine the convergence criterion. )pbdoc", - "mm_op"_a, - "max_iter"_a, - "tol"_a, - "diag_ethr"_a, - "need_subspace"_a, - "scf_type"_a, - "nproc_in_pool"_a) + "mm_op"_a, + "max_iter"_a, + "tol"_a, + "diag_ethr"_a, + "need_subspace"_a, + "scf_type"_a, + "nproc_in_pool"_a) .def("init_eig", &PyDiagoCGAdapter::init_eig, R"pbdoc( Initialize the eigenvalues. )pbdoc") .def("get_eig", &PyDiagoCGAdapter::get_eig, R"pbdoc( Get the eigenvalues. )pbdoc") - .def("set_psi", &PyDiagoCGAdapter::set_psi, R"pbdoc( + .def("set_psi", + &PyDiagoCGAdapter::set_psi, + R"pbdoc( Set the eigenvectors. - )pbdoc", "psi_in"_a) + )pbdoc", + "psi_in"_a) .def("get_psi", &PyDiagoCGAdapter::get_psi, R"pbdoc( Get the eigenvectors. )pbdoc") - .def("set_prec", &PyDiagoCGAdapter::set_prec, R"pbdoc( + .def("set_prec", + &PyDiagoCGAdapter::set_prec, + R"pbdoc( Set the preconditioner. - )pbdoc", "prec_in"_a); + )pbdoc", + "prec_in"_a); #else // Provide stub binding when ATen is not available // This allows the module to load but will raise an error if used - m.def("diago_cg_available", []() { return false; }, - "Check if diago_cg is available (requires ATen)"); + m.def("diago_cg_available", []() { return false; }, "Check if diago_cg is available (requires ATen)"); #endif } diff --git a/python/pyabacus/src/py_numerical_radial.cpp b/python/pyabacus/src/py_numerical_radial.cpp index 18500641cf0..c3417acb626 100644 --- a/python/pyabacus/src/py_numerical_radial.cpp +++ b/python/pyabacus/src/py_numerical_radial.cpp @@ -1,10 +1,9 @@ -#include -#include - #include "source_basis/module_nao/numerical_radial.h" - #include "utils/pybind_utils.h" +#include +#include + namespace py = pybind11; using namespace pybind11::literals; using namespace pyabacus::utils; @@ -122,22 +121,18 @@ void bind_numerical_radial(py::module& m) .def_property_readonly("nk", &NumericalRadial::nk) .def_property_readonly("rcut", &NumericalRadial::rcut) .def_property_readonly("kcut", &NumericalRadial::kcut) - .def_property_readonly("rgrid", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.rgrid(), static_cast(self.nr())); - }) - .def_property_readonly("kgrid", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.kgrid(), static_cast(self.nk())); - }) - .def_property_readonly("rvalue", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.rvalue(), static_cast(self.nr())); - }) - .def_property_readonly("kvalue", - [](NumericalRadial& self) { - return numpy_from_ptr_copy(self.kvalue(), static_cast(self.nk())); - }) + .def_property_readonly( + "rgrid", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.rgrid(), static_cast(self.nr())); }) + .def_property_readonly( + "kgrid", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.kgrid(), static_cast(self.nk())); }) + .def_property_readonly( + "rvalue", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.rvalue(), static_cast(self.nr())); }) + .def_property_readonly( + "kvalue", + [](NumericalRadial& self) { return numpy_from_ptr_copy(self.kvalue(), static_cast(self.nk())); }) .def_property_readonly("pr", &NumericalRadial::pr) .def_property_readonly("pk", &NumericalRadial::pk) .def_property_readonly("is_fft_compliant", overload_cast_<>()(&NumericalRadial::is_fft_compliant, py::const_)); diff --git a/python/pyabacus/src/utils/pybind_utils.h b/python/pyabacus/src/utils/pybind_utils.h index 1fc24fbf2be..8785c204738 100644 --- a/python/pyabacus/src/utils/pybind_utils.h +++ b/python/pyabacus/src/utils/pybind_utils.h @@ -13,21 +13,22 @@ #ifndef PYABACUS_UTILS_PYBIND_UTILS_H #define PYABACUS_UTILS_PYBIND_UTILS_H -#include -#include -#include - #include +#include #include +#include +#include +#include #include #include #include -#include namespace py = pybind11; -namespace pyabacus { -namespace utils { +namespace pyabacus +{ +namespace utils +{ // ============================================================================ // Array Validation Functions @@ -45,8 +46,7 @@ void check_1d_array(const py::array_t& arr, const std::string& name = "array" py::buffer_info info = arr.request(); if (info.ndim != 1) { - throw std::runtime_error(name + " must be 1-dimensional, got " + - std::to_string(info.ndim) + " dimensions"); + throw std::runtime_error(name + " must be 1-dimensional, got " + std::to_string(info.ndim) + " dimensions"); } } @@ -62,8 +62,7 @@ void check_2d_array(const py::array_t& arr, const std::string& name = "array" py::buffer_info info = arr.request(); if (info.ndim != 2) { - throw std::runtime_error(name + " must be 2-dimensional, got " + - std::to_string(info.ndim) + " dimensions"); + throw std::runtime_error(name + " must be 2-dimensional, got " + std::to_string(info.ndim) + " dimensions"); } } @@ -75,14 +74,12 @@ void check_2d_array(const py::array_t& arr, const std::string& name = "array" * @throws std::runtime_error if size doesn't match */ template -void check_array_size(const py::array_t& arr, size_t expected_size, - const std::string& name = "array") +void check_array_size(const py::array_t& arr, size_t expected_size, const std::string& name = "array") { if (static_cast(arr.size()) != expected_size) { - throw std::runtime_error(name + " must have size " + - std::to_string(expected_size) + ", got " + - std::to_string(arr.size())); + throw std::runtime_error(name + " must have size " + std::to_string(expected_size) + ", got " + + std::to_string(arr.size())); } } @@ -107,8 +104,7 @@ void check_array_shape(const py::array_t& arr, { if (info.shape[i] != expected_shape[i]) { - throw std::runtime_error(name + " has wrong shape at dimension " + - std::to_string(i)); + throw std::runtime_error(name + " has wrong shape at dimension " + std::to_string(i)); } } } @@ -122,14 +118,14 @@ void check_array_shape(const py::array_t& arr, * @throws std::runtime_error if sizes don't match */ template -void check_same_size(const py::array_t& arr1, const py::array_t& arr2, +void check_same_size(const py::array_t& arr1, + const py::array_t& arr2, const std::string& name1 = "array1", const std::string& name2 = "array2") { if (arr1.size() != arr2.size()) { - throw std::runtime_error(name1 + " and " + name2 + - " must have the same size"); + throw std::runtime_error(name1 + " and " + name2 + " must have the same size"); } } @@ -228,8 +224,7 @@ py::array_t numpy_from_ptr_copy(const T* ptr, size_t size) template py::array_t numpy_from_ptr_copy_2d(const T* ptr, size_t nrow, size_t ncol) { - std::vector shape = {static_cast(nrow), - static_cast(ncol)}; + std::vector shape = {static_cast(nrow), static_cast(ncol)}; py::array_t result(shape); std::memcpy(result.mutable_data(), ptr, nrow * ncol * sizeof(T)); return result; @@ -246,7 +241,7 @@ py::array_t numpy_from_ptr_copy_nd(const T* ptr, const std::vector& { py::array_t result(shape); size_t total_size = 1; - for (auto dim : shape) + for (auto dim: shape) { total_size *= static_cast(dim); } @@ -270,9 +265,7 @@ py::array_t numpy_from_ptr_copy_nd(const T* ptr, const std::vector& * @return Numpy array view */ template -py::array_t numpy_view_with_keepalive(T* ptr, - const std::vector& shape, - std::shared_ptr owner_ptr) +py::array_t numpy_view_with_keepalive(T* ptr, const std::vector& shape, std::shared_ptr owner_ptr) { // Calculate strides for C-contiguous array std::vector strides(shape.size()); @@ -303,10 +296,7 @@ py::array_t numpy_view_with_keepalive(T* ptr, template py::array_t numpy_view_1d(T* ptr, size_t size) { - return py::array_t({static_cast(size)}, - {sizeof(T)}, - ptr, - py::none()); + return py::array_t({static_cast(size)}, {sizeof(T)}, ptr, py::none()); } /** @@ -322,12 +312,10 @@ py::array_t numpy_view_1d(T* ptr, size_t size) template py::array_t numpy_view_2d(T* ptr, size_t nrow, size_t ncol) { - return py::array_t( - {static_cast(nrow), static_cast(ncol)}, - {static_cast(ncol * sizeof(T)), sizeof(T)}, - ptr, - py::none() - ); + return py::array_t({static_cast(nrow), static_cast(ncol)}, + {static_cast(ncol * sizeof(T)), sizeof(T)}, + ptr, + py::none()); } // ============================================================================ @@ -400,9 +388,7 @@ struct BufferHelper * @param arr The numpy array */ explicit BufferHelper(py::array_t& arr) - : info(arr.request()) - , ptr(static_cast(info.ptr)) - , size(static_cast(arr.size())) + : info(arr.request()), ptr(static_cast(info.ptr)), size(static_cast(arr.size())) { } @@ -411,9 +397,7 @@ struct BufferHelper * @param arr The numpy array */ explicit BufferHelper(const py::array_t& arr) - : info(arr.request()) - , ptr(static_cast(info.ptr)) - , size(static_cast(arr.size())) + : info(arr.request()), ptr(static_cast(info.ptr)), size(static_cast(arr.size())) { } @@ -453,9 +437,8 @@ struct BufferHelper { if (size != expected) { - throw std::runtime_error(name + " must have size " + - std::to_string(expected) + ", got " + - std::to_string(size)); + throw std::runtime_error(name + " must have size " + std::to_string(expected) + ", got " + + std::to_string(size)); } } @@ -465,8 +448,7 @@ struct BufferHelper * @param name Name for error messages * @throws std::runtime_error if shape doesn't match */ - void require_shape(const std::vector& expected_shape, - const std::string& name = "array") const + void require_shape(const std::vector& expected_shape, const std::string& name = "array") const { if (info.ndim != static_cast(expected_shape.size())) { @@ -476,8 +458,7 @@ struct BufferHelper { if (info.shape[i] != expected_shape[i]) { - throw std::runtime_error(name + " has wrong shape at dimension " + - std::to_string(i)); + throw std::runtime_error(name + " has wrong shape at dimension " + std::to_string(i)); } } } @@ -486,26 +467,38 @@ struct BufferHelper * @brief Get number of dimensions * @return Number of dimensions */ - ssize_t ndim() const { return info.ndim; } + ssize_t ndim() const + { + return info.ndim; + } /** * @brief Get shape at dimension i * @param i Dimension index * @return Size at dimension i */ - ssize_t shape(size_t i) const { return info.shape[i]; } + ssize_t shape(size_t i) const + { + return info.shape[i]; + } /** * @brief Get number of rows (for 2D arrays) * @return Number of rows */ - ssize_t nrow() const { return info.shape[0]; } + ssize_t nrow() const + { + return info.shape[0]; + } /** * @brief Get number of columns (for 2D arrays) * @return Number of columns */ - ssize_t ncol() const { return info.shape[1]; } + ssize_t ncol() const + { + return info.shape[1]; + } }; // ============================================================================