update doc

Author: fishjojo

doc/make.py

@@ -202,7 +202,7 @@ def main():
         "command", nargs="?", default="html", help=f"command to run: {joined}"
     )
     argparser.add_argument(
-        "--num-jobs", default="auto", help="number of jobs used by sphinx-build"
+        "--num-jobs", default="1", help="number of jobs used by sphinx-build"
     )
     argparser.add_argument(
         "--python-path", type=str, default=os.path.dirname(DOC_PATH), help="path"

doc/requirements.txt

@@ -3,10 +3,11 @@ jinja2
 sphinx
 numpydoc
 sphinx_copybutton
+sphinx-autodoc-typehints
 matplotlib
 sphinx_design
 nbsphinx
-sphinxemoji
+sphinxemoji>=0.3.2
 pydata-sphinx-theme
 jupytext
 myst-nb

doc/source/api_reference/index.rst

@@ -26,3 +26,4 @@ All classes and functions exposed in pyscfad.* namespace are public.
    tdscf
    tools
    util
+   ml

doc/source/api_reference/ml.rst

@@ -0,0 +1,13 @@
+.. _ml:
+
+================
+Machine learning
+================
+
+The :mod:`pyscfad.ml` subpackage implements methods used in machine learning tasks.
+
+.. autosummary::
+   :toctree: api/
+   :recursive:
+
+   pyscfad.ml

doc/source/conf.py

@@ -10,8 +10,18 @@
 import warnings
 import inspect
 
+# Workaround to avoid expanding type aliases. Copied from jax
+from typing import ForwardRef
+def _do_not_evaluate(
+    self, globalns, *args, _evaluate=ForwardRef._evaluate, **kwargs,
+):
+  if globalns.get('__name__', '').startswith('pyscfad'):
+    return self
+  return _evaluate(self, globalns, *args, **kwargs)
+ForwardRef._evaluate = _do_not_evaluate
+
 project = "pyscfad"
-copyright = "2021-2025, Xing Zhang"
+copyright = "2021-2026, Xing Zhang"
 author = "Xing Zhang"
 
 import pyscfad
@@ -27,11 +37,13 @@
 #    "IPython.sphinxext.ipython_directive",
 #    "IPython.sphinxext.ipython_console_highlighting",
     "matplotlib.sphinxext.plot_directive",
-    "numpydoc",
+#    "numpydoc",
     "sphinx_copybutton",
     "sphinx_design",
     "sphinx.ext.autodoc",
     "sphinx.ext.autosummary",
+    "sphinx.ext.napoleon", # if using Google/NumPy docstrings
+    #"sphinx_autodoc_typehints",
     "sphinx.ext.coverage",
     "sphinx.ext.doctest",
     "sphinx.ext.extlinks",
@@ -56,11 +68,22 @@
 ]
 
 autosummary_generate = True
-autodoc_typehints = "none"
 
-numpydoc_show_class_members = False
-numpydoc_show_inherited_class_members = False
-numpydoc_attributes_as_param_list = False
+napolean_use_rtype = False
+autodoc_typehints_format = "short"
+
+autodoc_typehints = "description"
+autodoc_typehints_description_target = "all"
+autodoc_type_aliases = {
+    'ArrayLike': 'pyscfad.typing.ArrayLike',
+}
+
+#always_document_param_types = True
+#set_type_checking_flag = True
+
+#numpydoc_show_class_members = False
+#numpydoc_show_inherited_class_members = False
+#numpydoc_attributes_as_param_list = False
 
 pygments_style = "sphinx"
 

pyproject.toml

@@ -15,7 +15,7 @@ authors = [
 ]
 
 dependencies = [
-  "jax>=0.7.1,<0.9",
+  "jax>=0.7.1,<0.10",
   "pyscfadlib>=0.2.0",
   "pyscf>=2.3",
   "pyscf-properties",
@@ -33,7 +33,7 @@ classifiers = [
 
 [project.optional-dependencies]
 cuda12 = [
-  "jax[cuda12]>=0.7.1,<0.9",
+  "jax[cuda12]>=0.7.1,<0.10",
   "pyscfad-cuda12-plugin[with_cuda]>=0.2.0",
 ]
 

pyscfad/gto/mole_lite.py

@@ -11,9 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""
+Lightweight ``mole`` module.
+"""
 
 from __future__ import annotations
-from typing import Any
 from collections.abc import Callable
 import contextlib
 
@@ -46,13 +48,12 @@
     is_au,
 )
 
+from pyscfad.typing import ArrayLike, Array
 from pyscfad import numpy as np
 from pyscfad import ops
 from pyscfad.gto.mole import energy_nuc
 from pyscfad.gto import moleintor_lite
 
-Array = Any
-
 def _format_basis(basis, uniq_symbols):
     if isinstance(basis, dict):
         for k, v in basis.items():
@@ -85,44 +86,32 @@ def _format_symbols(symbols):
         symbols = [symbols,]
     return tuple(_atom_symbol(symb) for symb in symbols)
 
-class Mole(MoleBase):
-    """Molecular information.
-
-    Parameters
-    ----------
-    symbols : tuple of str
-        Atomic symbols (mutually exclusive with ``numbers``).
-    coords : array
-        Atomic coordinates (in Bohr).
-    basis : dict or str
-        A string indicating the Gaussian basis set to use, or
-        a dictionary including the basis set parameters
-        (exponents and contraction coefficients).
-    numbers : tuple of ints
-        Atomic numbers (mutually exclusive with ``symbols``).
-    charge : int
-        Total charge.
-    spin : int
-        2S (number of alpha electrons minus number of beta electrons).
-    cart : bool
-        Whether to use Cartesian Gaussian basis.
-    verbose : int
-        Printing level.
-    trace_coords : bool
-        Whether to trace atomic coordinates for gradient calculations.
-    trace_basis : bool
-        Whether to trace basis set parameters for gradient calculations.
-
-    Notes
-    -----
-    The molecular composition (i.e., ``symbols`` or ``numbers``)
-    must be static as input. For dynamic molecular composition,
-    refer to :class:`pyscfad.ml.gto.MolePad`.
+class MoleLite(MoleBase):
+    """Molecular information (lightweight version).
+
+    Parameters:
+        symbols: Atomic symbols (mutually exclusive with ``numbers``).
+        coords: Atomic coordinates (in Bohr).
+        basis: A string indicating the Gaussian basis set to use, or
+            a dictionary including the basis set parameters
+            (exponents and contraction coefficients).
+        numbers: Atomic numbers (mutually exclusive with ``symbols``).
+        charge: Total charge.
+        spin: 2S (number of alpha electrons minus number of beta electrons).
+        cart: Whether to use Cartesian Gaussian basis.
+        verbose: Printing level.
+        trace_coords: Whether to trace atomic coordinates for gradient calculations.
+        trace_basis: Whether to trace basis set parameters for gradient calculations.
+
+    Notes:
+        The molecular composition (i.e., ``symbols`` or ``numbers``)
+        must be static as input. For dynamic molecular composition,
+        refer to :class:`~pyscfad.ml.gto.mole_pad.MolePad`.
     """
     def __init__(
         self,
         symbols: tuple[str, ...] | None = None,
-        coords: Array | None = None,
+        coords: ArrayLike | None = None,
         basis: dict | str | None = None,
         numbers: tuple[int, ...] | None = None,
         charge: int = 0,
@@ -171,7 +160,12 @@ def atom_pure_symbol(
     def atom_coords(
         self,
         unit: str = "Bohr",
-    ):
+    ) -> Array:
+        """Atom coordinates.
+
+        Parameters:
+            unit: Units of returned coordinates, can be ``Bohr`` or ``Angstrom``.
+        """
         if not is_au(unit):
             return self.coords * BOHR
         else:
@@ -180,7 +174,7 @@ def atom_coords(
     def copy(
         self,
         deep: bool = True,
-    ) -> Mole:
+    ) -> MoleLite:
         import copy
         newmol = self.view(self.__class__)
         if not deep:
@@ -193,7 +187,10 @@ def copy(
         newmol._env = np.copy(self._env)
         return newmol
 
-    def build(self, *args, **kwargs):
+    def build(self, *args, **kwargs) -> MoleLite:
+        """Placeholder for post-init processing.
+        Can be used for caching one-time computed quantities.
+        """
         return self
 
     def intor(
@@ -202,10 +199,32 @@ def intor(
         comp: int | None = None,
         hermi: int = 0,
         aosym: str = "s1",
-        out: Array | None = None,
+        out: ArrayLike | None = None,
         shls_slice: tuple[int, ...] | None = None,
-        grids: Array | None = None,
+        grids: ArrayLike | None = None,
     ) -> Array:
+        """Integral generator.
+
+        Parameters:
+            intor_name: Name of the integral.
+            comp: Components of the integrals, e.g. ``int1e_ovlp_dr10`` has 3 components.
+            hermi: Hermitian symmetry of 1e integrals, can be
+                `0` (no symmetry), `1` (hermitian), and `2` (anti-hermitian).
+            aosym: Permutation symmetry of 2e integrals, can be
+                `s1`, `s2`, `s4`, and `s8`.
+            out: Not used.
+            shls_slice: Label the start-stop shells for each index in the integral.
+                For example, the 8-element tuple for the 2e integral
+                tensor ``(ij|kl) = intor('int2e')`` are specified as
+                ``(i0, i1, j0, j1, k0, k1, l0, l1)``.
+            grids: Not used.
+
+        Returns:
+            Computed integral as an array.
+
+        See also:
+            :func:`pyscf.gto.Mole.intor`
+        """
         del out, grids
 
         intor_name = self._add_suffix(intor_name)
@@ -230,8 +249,8 @@ def intor(
 
     def set_common_origin(
         self,
-        coord: Array,
-    ) -> Mole:
+        coord: ArrayLike,
+    ) -> MoleLite:
         if self._env is None:
             raise RuntimeError("{self}._env is not initialized, "
                                "possibly because basis is not set.")
@@ -245,15 +264,15 @@ def set_common_origin(
 
     def with_common_origin(
         self,
-        coord: Array,
+        coord: ArrayLike,
     ):
         coord0 = np.copy(self._env[PTR_COMMON_ORIG:PTR_COMMON_ORIG+3])
         return self._TemporaryMoleContext(self.set_common_origin, (coord,), (coord0,))
 
     def set_rinv_origin(
         self,
-        coord: Array,
-    ) -> Mole:
+        coord: ArrayLike,
+    ) -> MoleLite:
         if self._env is None:
             raise RuntimeError("{self}._env is not initialized, "
                                "possibly because basis is not set.")
@@ -267,7 +286,7 @@ def set_rinv_origin(
 
     def with_rinv_origin(
         self,
-        coord: Array,
+        coord: ArrayLike,
     ):
         coord0 = np.copy(self._env[PTR_RINV_ORIG:PTR_RINV_ORIG+3])
         return self._TemporaryMoleContext(self.set_rinv_origin, (coord,), (coord0,))
@@ -291,7 +310,9 @@ def from_pyscf(
         mol: MoleBase,
         trace_coords: bool = False,
         trace_basis: bool = False,
-    ) -> Mole:
+    ) -> MoleLite:
+        """Initialize from :class:`pyscf.gto.Mole` object.
+        """
         if not mol._built:
             raise KeyError(f"{mol} not built")
         if mol.ecp or mol.pseudo:
@@ -327,6 +348,8 @@ def to_pyscf(
         output: str | None = None,
         max_memory: int | None = None,
     ) -> MoleBase:
+        """Convert to :class:`pyscf.gto.Mole` object.
+        """
         coords = ops.to_numpy(self.coords)
         atom = [[a, tuple(x.tolist())] for a, x in zip(self.symbols, coords)]
 
@@ -353,14 +376,13 @@ def to_pyscf(
 
     energy_nuc = energy_nuc
 
-MoleLite = Mole
 
 def gaussian_int(
     n: int | numpy.ndarray,
-    alpha: Array,
+    alpha: ArrayLike,
 ) -> Array:
     r"""Gaussian integral.
-    Computes :math:`\int_0^\infty x^n exp(-alpha x^2) dx`.
+    Computes :math:`\int_0^\infty x^n exp(-\alpha x^2) dx`.
     """
     from pyscfad.scipy.special import gamma
     n1 = (n + 1) * .5
@@ -369,7 +391,7 @@ def gaussian_int(
 @with_doc(pyscf.gto.mole.gto_norm.__doc__)
 def gto_norm(
     l: int | numpy.ndarray,
-    expnt: Array,
+    expnt: ArrayLike,
 ) -> Array:
     assert numpy.all(l >= 0)
     return 1. / np.sqrt(gaussian_int(l*2+2, 2*expnt))
@@ -432,7 +454,7 @@ def make_bas_env(
     return _bas, _env
 
 def make_env(
-    mol: Mole,
+    mol: MoleLite,
 ) -> tuple[numpy.ndarray, numpy.ndarray, Array]:
     """Make ``_atm``, ``_bas``, and ``_env`` for
     interfacing with libcint.

pyscfad/ml/gto/__init__.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from .basis_array import make_basis_array, BasisArray
-from .mole_pad import MolePad
+"""
+Molecular sturcture and Gaussian type orbital integrals
+"""
+from pyscfad.ml.gto.basis_array import make_basis_array, BasisArray
+from pyscfad.ml.gto.mole_pad import MolePad