add docs for pyabacus

python/pyabacus/CMakeLists.txt

@@ -77,7 +77,6 @@ list(APPEND _naos
     # dependency
     ${ABACUS_SOURCE_DIR}/module_base/kernels/math_op.cpp
     ${ABACUS_SOURCE_DIR}/module_psi/kernels/psi_memory_op.cpp
-    ${ABACUS_SOURCE_DIR}/module_io/output_radial.cpp
     )
 add_library(naopack SHARED 
     ${_naos}

python/pyabacus/src/py_abacus.cpp

@@ -8,6 +8,7 @@ void bind_m_nao(py::module& m);
 
 PYBIND11_MODULE(_core, m)
 {
+    m.doc() = "Python extension for ABACUS built with pybind11 and scikit-build.";
     bind_base_math(m);
     bind_m_nao(m);
 }

python/pyabacus/src/py_m_nao.cpp

@@ -14,10 +14,17 @@ void bind_m_nao(py::module& m)
 {
     // Create the submodule for Module NAO
     py::module m_nao = m.def_submodule("ModuleNAO");
+    m_nao.doc() = "Module for Numerical Atomic Orbitals (NAO) in ABACUS";
 
     // Bind the RadialCollection class
     py::class_<RadialCollection>(m_nao, "RadialCollection")
-        .def(py::init<>())
+        .def(py::init<>(), R"pbdoc(
+            A class that holds all numerical radial functions of the same kind. 
+            
+            An instance of this class could be the collection of all radial functions 
+            of numerical atomic orbitals, or all Kleinman-Bylander beta functions from 
+            all elements involved in a calculation.
+            )pbdoc")
         .def("build", [](RadialCollection& self, int nfile, const py::list &file_list, char ftype){
              std::vector<std::string> files;
              files.reserve(nfile);
@@ -26,17 +33,17 @@ void bind_m_nao(py::module& m)
                 files.push_back(file.cast<std::string>());
             }
             self.build(nfile, files.data(), ftype);
-        }, "nfile"_a, "file_list"_a, "ftype"_a = '\0')
-        .def("set_transformer", &RadialCollection::set_transformer, "sbt"_a, "update"_a = 0)
-        .def("set_uniform_grid", &RadialCollection::set_uniform_grid, "for_r_space"_a, "ngrid"_a, "cutoff"_a, "mode"_a = 'i', "enable_fft"_a = false)
+        }, "Builds the collection from (orbital) files", "nfile"_a, "file_list"_a, "ftype"_a = '\0')
+        .def("set_transformer", &RadialCollection::set_transformer, "Sets a spherical Bessel transformers for all RadialSet objects.","sbt"_a, "update"_a = 0)
+        .def("set_uniform_grid", &RadialCollection::set_uniform_grid, "Sets a common uniform grid for all RadialSet objects","for_r_space"_a, "ngrid"_a, "cutoff"_a, "mode"_a = 'i', "enable_fft"_a = false)
         .def("set_grid", [](RadialCollection& self, const bool for_r_space, const int ngrid, py::array_t<double> grid, const char mode = 'i'){
             py::buffer_info grid_info = grid.request();
             if (grid_info.ndim != 1)
             {
                 throw std::runtime_error("Input array must be 1-dimensional");
             }
             self.set_grid(for_r_space, ngrid, static_cast<double*>(grid_info.ptr), mode);
-        }, "for_r_space"_a, "ngrid"_a, "grid"_a, "mode"_a = 'i')
+        }, "Sets a common grid for all RadialSet objects", "for_r_space"_a, "ngrid"_a, "grid"_a, "mode"_a = 'i')
         // Getters
         .def("symbol", &RadialCollection::symbol, "itype"_a)
         .def_property_readonly("ntype", &RadialCollection::ntype)
@@ -51,8 +58,36 @@ void bind_m_nao(py::module& m)
         .def_property_readonly("nchi", overload_cast_<>()(&RadialCollection::nchi, py::const_));
     //Bind the TwoCenterIntegrator class
     py::class_<TwoCenterIntegrator>(m_nao, "TwoCenterIntegrator")
-        .def(py::init<>())
-        .def("tabulate", &TwoCenterIntegrator::tabulate, "bra"_a, "ket"_a, "op"_a, "nr"_a, "cutoff"_a)
+        .def(py::init<>(),  R"pbdoc(
+    A class to compute two-center integrals.
+
+    This class computes two-center integrals of the form:
+
+                        /    
+                 I(R) = | dr phi1(r) (op) phi2(r - R)
+                        /               
+
+    as well as their gradients, where op is 1 (overlap) or minus Laplacian (kinetic), and phi1, phi2 are "atomic-orbital-like" functions of the form:
+
+                 phi(r) = chi(|r|) * Ylm(r/|r|)
+
+    where chi is some numerical radial function and Ylm is some real spherical harmonics.
+
+    This class is designed to efficiently compute the two-center integrals between two "collections" of the above functions with various R, e.g., the overlap integrals between all numerical atomic orbitals and all Kleinman-Bylander nonlocal projectors, the overlap & kinetic integrals between all numerical atomic orbitals, etc.
+    This is done by tabulating the radial part of the integrals on an r-space grid and the real Gaunt coefficients in advance.
+    )pbdoc")
+        .def("tabulate", &TwoCenterIntegrator::tabulate, 
+        R"pbdoc(
+        Tabulates the radial part of a two-center integral.
+
+        Parameters:
+        bra (RadialFunctions): The radial functions of the first collection.
+        ket (RadialFunctions): The radial functions of the second collection.
+        op (char): Operator, could be 'S' (overlap) or 'T' (kinetic).
+        nr (int): Number of r-space grid points.
+        cutoff (float): r-space cutoff radius.
+        )pbdoc",
+        "bra"_a, "ket"_a, "op"_a, "nr"_a, "cutoff"_a)
         .def("calculate", [](TwoCenterIntegrator& self, const int itype1, const int l1, const int izeta1, const int m1, const int itype2, const int l2, const int izeta2,  const int m2, py::array_t<double> pvR, bool cal_grad = false)
         {
             py::buffer_info pvR_info = pvR.request();
@@ -81,7 +116,49 @@ void bind_m_nao(py::module& m)
                 return py::make_tuple(out_array, grad_out_array);
             }
 
-        }, "itype1"_a, "l1"_a, "izeta1"_a, "m1"_a, "itype2"_a, "l2"_a, "izeta2"_a, "m2"_a, "pvR"_a, "cal_grad"_a = false)
+        }, 
+        R"pbdoc(
+    Compute the two-center integrals.
+
+    This function calculates the two-center integral
+
+                        /    
+                 I(R) = | dr phi1(r) (op_) phi2(r - R)
+                        /               
+
+    or its gradient by using the tabulated radial part and real Gaunt coefficients.
+
+    Parameters
+    ----------
+    itype1 : int
+        Element index of orbital 1.
+    l1 : int
+        Angular momentum of orbital 1.
+    izeta1 : int
+        Zeta number of orbital 1.
+    m1 : int
+        Magnetic quantum number of orbital 1.
+    itype2 : int
+        Element index of orbital 2.
+    l2 : int
+        Angular momentum of orbital 2.
+    izeta2 : int
+        Zeta number of orbital 2.
+    m2 : int
+        Magnetic quantum number of orbital 2.
+    pvR : array_like
+        R2 - R1, the displacement vector between the two centers.
+    cal_grad : bool, optional
+        The gradient will not be computed if cal_grad is false.
+    
+    Returns
+    -------
+    out_array : array_like
+        The two-center integral.
+    grad_out_array : array_like
+        Gradient of the two-center integral.
+        )pbdoc",
+         "itype1"_a, "l1"_a, "izeta1"_a, "m1"_a, "itype2"_a, "l2"_a, "izeta2"_a, "m2"_a, "pvR"_a, "cal_grad"_a = false)
         .def("snap", [](TwoCenterIntegrator& self, const int itype1, const int l1, const int izeta1, const int m1, const int itype2, py::array_t<double> pvR, const bool deriv){
             py::buffer_info pvR_info = pvR.request();
             if (pvR_info.size != 3)
@@ -94,6 +171,13 @@ void bind_m_nao(py::module& m)
             std::vector<std::vector<double>> out;
             self.snap(itype1, l1, izeta1, m1, itype2, vR, deriv, out);
             return out;
-        }, "itype1"_a, "l1"_a, "izeta1"_a, "m1"_a, "itype2"_a, "pvR"_a, "deriv"_a = false);
+        },
+        R"pbdoc(
+    Compute a batch of two-center integrals.
+
+    This function calculates the two-center integrals (and optionally their gradients)
+    between one orbital and all orbitals of a certain type from the other collection.
+        )pbdoc", 
+        "itype1"_a, "l1"_a, "izeta1"_a, "m1"_a, "itype2"_a, "pvR"_a, "deriv"_a = false);
 
 }

python/pyabacus/src/pyabacus/__init__.py

@@ -1,4 +1,3 @@
 from __future__ import annotations
-# from ._core import __doc__, __version__, NumericalRadial, ModuleBase
 from ._core import ModuleBase, ModuleNAO
 __all__ = ["ModuleBase", "ModuleNAO"]