fix pbcintor

Author: fishjojo

pyscfad/gto/mole_lite.py

@@ -47,10 +47,9 @@
 )
 
 from pyscfad import numpy as np
-#from pyscfad import pytree
 from pyscfad import ops
-#from pyscfad.ops import jit
 from pyscfad.gto.mole import energy_nuc
+from pyscfad.gto import moleintor_lite
 
 Array = Any
 
@@ -92,12 +91,13 @@ class Mole(MoleBase):
     Parameters
     ----------
     symbols : tuple of str
-        Atomic symbols.
+        Atomic symbols (mutually exclusive with ``numbers``).
     coords : array
         Atomic coordinates (in Bohr).
     basis : dict or str
-        Atom-centered contracted Gaussian basis set parameters
-        (including exponents and contraction coefficients).
+        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
@@ -106,10 +106,18 @@ class Mole(MoleBase):
         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`.
     """
     def __init__(
         self,
@@ -127,7 +135,6 @@ def __init__(
         if numbers is not None:
             if symbols is not None:
                 raise KeyError("Only one of 'symbols' and 'numbers' can be specified.")
-            #numbers = numpy.asarray(numbers, dtype=int)
             self.symbols = tuple(_symbol(i) for i in numbers)
         else:
             self.symbols = _format_symbols(symbols)
@@ -187,7 +194,7 @@ def copy(
         return newmol
 
     def build(self, *args, **kwargs):
-        pass
+        return self
 
     def intor(
         self,
@@ -199,7 +206,8 @@ def intor(
         shls_slice: tuple[int, ...] | None = None,
         grids: Array | None = None,
     ) -> Array:
-        from pyscfad.gto import moleintor_lite
+        del out, grids
+
         intor_name = self._add_suffix(intor_name)
         if "ECP" in intor_name:
             raise NotImplementedError
@@ -215,7 +223,6 @@ def intor(
             comp=comp,
             hermi=hermi,
             aosym=aosym,
-            out=out,
             trace_coords=self.trace_coords,
             trace_basis=self.trace_basis,
         )
@@ -309,14 +316,14 @@ def from_pyscf(
             charge=mol.charge,
             spin=mol.spin,
             cart=mol.cart,
+            verbose=mol.verbose,
             trace_coords=trace_coords,
             trace_basis=trace_basis,
         )
         return dmol
 
     def to_pyscf(
         self,
-        verbose: int | None = None,
         output: str | None = None,
         max_memory: int | None = None,
     ) -> MoleBase:
@@ -336,7 +343,7 @@ def to_pyscf(
             spin=self.spin,
             cart=self.cart,
             unit="AU",
-            verbose=verbose,
+            verbose=self.verbose,
             output=output,
             max_memory=max_memory,
             dump_input=False,

pyscfad/gto/moleintor_lite.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 from typing import Any
-import warnings
 from functools import partial
 import numpy
 
+from jax.custom_derivatives import SymbolicZero
 from pyscf.gto.mole import (
     ATOM_OF,
     PTR_COORD,
@@ -161,8 +161,6 @@ def _get_shape(
         "hermi",
         "aosym",
         "ao_loc",
-        "cintopt",
-        "out",
         "trace_coords",
         "trace_basis",
         "aoslices",
@@ -178,15 +176,11 @@ def getints(
     hermi: int = 0,
     aosym: str = "s1",
     ao_loc: Array | None = None,
-    cintopt: Any | None = None,
-    out: Any | None = None,
     trace_coords: bool = False,
     trace_basis: bool = False,
     aoslices: Array | None = None, # for padding
 ):
     from pyscfad.gto._pyscf_moleintor import getints as callback
-    if out is not None:
-        warnings.warn("Pre-allocated 'out' is not used")
 
     shape = _get_shape(
         intor_name,
@@ -211,8 +205,6 @@ def getints(
         comp=comp,
         hermi=hermi,
         ao_loc=ao_loc,
-        cintopt=cintopt,
-        out=None,
     )
     return out
 
@@ -226,8 +218,6 @@ def getints_jvp(
     hermi,
     aosym,
     ao_loc,
-    cintopt,
-    out,
     trace_coords,
     trace_basis,
     aoslices,
@@ -250,8 +240,6 @@ def getints_jvp(
         hermi=hermi,
         aosym=aosym,
         ao_loc=ao_loc,
-        cintopt=cintopt,
-        out=out,
         trace_coords=trace_coords,
         trace_basis=trace_basis,
         aoslices=aoslices,
@@ -262,34 +250,35 @@ def getints_jvp(
     intor_ip_bra = intor_ip_ket = None
     intor_ip_bra, intor_ip_ket = int1e_dr1_name(intor_name)
 
-    if trace_coords and (intor_ip_bra or intor_ip_ket):
-        if intor_name.startswith("int1e_rinv"):
-            rc_deriv = PTR_RINV_ORIG
-        elif intor_name.startswith("int1e_r"):
-            rc_deriv = PTR_COMMON_ORIG
-        else:
-            rc_deriv = None
-
-        tangent_out += _gen_int1e_jvp_r0(
-            intor_ip_bra,
-            intor_ip_ket,
-            atm,
-            bas,
-            env,
-            env_dot,
-            shls_slice,
-            comp,
-            hermi,
-            aosym,
-            ao_loc,
-            trace_coords,
-            trace_basis,
-            aoslices,
-            rc_deriv,
-        ).reshape(tangent_out.shape)
-
-    if trace_basis:
-        raise NotImplementedError("basis set parameter derivative not supported")
+    if not isinstance(env_dot, SymbolicZero):
+        if trace_coords and (intor_ip_bra or intor_ip_ket):
+            if intor_name.startswith("int1e_rinv"):
+                rc_deriv = PTR_RINV_ORIG
+            elif intor_name.startswith("int1e_r"):
+                rc_deriv = PTR_COMMON_ORIG
+            else:
+                rc_deriv = None
+
+            tangent_out += _gen_int1e_jvp_r0(
+                intor_ip_bra,
+                intor_ip_ket,
+                atm,
+                bas,
+                env,
+                env_dot,
+                shls_slice,
+                comp,
+                hermi,
+                aosym,
+                ao_loc,
+                trace_coords,
+                trace_basis,
+                aoslices,
+                rc_deriv,
+            ).reshape(tangent_out.shape)
+
+        if trace_basis:
+            raise NotImplementedError("basis set parameter derivative not supported")
     return primal_out, tangent_out
 
 
@@ -323,8 +312,6 @@ def _gen_int1e_jvp_r0(
         hermi=0,
         aosym=aosym,
         ao_loc=ao_loc,
-        cintopt=None,
-        out=None,
         trace_coords=trace_coords,
         trace_basis=trace_basis,
         aoslices=aoslices,
@@ -364,8 +351,6 @@ def _gen_int1e_jvp_r0(
             hermi=0,
             aosym=aosym,
             ao_loc=ao_loc,
-            cintopt=None,
-            out=None,
             trace_coords=trace_coords,
             trace_basis=trace_basis,
             aoslices=aoslices,

pyscfad/ml/gto/mole_pad.py

@@ -32,6 +32,7 @@
 
 from pyscfad import numpy as np
 from pyscfad import ops
+from pyscfad.gto import moleintor_lite
 from pyscfad.gto.mole_lite import MoleLite
 from pyscfad.ml.gto.basis_array import BasisArray
 
@@ -160,7 +161,8 @@ def intor(
         shls_slice: tuple[int, ...] | None = None,
         grids: Array | None = None,
     ) -> Array:
-        from pyscfad.gto import moleintor_lite
+        del out, grids
+
         intor_name = self._add_suffix(intor_name)
         if "ECP" in intor_name:
             raise NotImplementedError
@@ -180,7 +182,6 @@ def intor(
             hermi=hermi,
             aosym=aosym,
             ao_loc=ao_loc,
-            out=out,
             trace_coords=self.trace_coords,
             trace_basis=self.trace_basis,
             aoslices=aoslices,

pyscfad/pbc/gto/_pbcintor.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2025 Xing Zhang
+# Copyright 2021-2025 The PySCFAD Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -24,7 +24,79 @@
 @partial(custom_jvp, nondiff_argnums=tuple(range(1,7)))
 def _pbc_intor(mol, intor, comp=None, hermi=0, kpts=None, kpt=None,
                shls_slice=None):
-    return Cell.pbc_intor(mol.view(Cell), intor, comp, hermi, kpts, kpt, shls_slice)
+    import ctypes
+    from pyscf.gto import moleintor
+    from pyscf import lib
+    from pyscf.gto.mole import conc_env
+    from pyscf.pbc.gto import _pbcintor
+    libpbc = _pbcintor.libpbc
+
+    cell1 = cell2 = mol
+    intor, comp = moleintor._get_intor_and_comp(cell1._add_suffix(intor), comp)
+
+    if kpts is None:
+        if kpt is not None:
+            kpts_lst = numpy.reshape(kpt, (1,3))
+        else:
+            kpts_lst = numpy.zeros((1,3))
+    else:
+        kpts_lst = numpy.reshape(kpts, (-1,3))
+    nkpts = len(kpts_lst)
+
+    pcell = cell1.copy(deep=False)
+    pcell.precision = min(cell1.precision, cell2.precision)
+    pcell._atm, pcell._bas, pcell._env = \
+            atm, bas, env = conc_env(cell1._atm, cell1._bas, cell1._env,
+                                     cell2._atm, cell2._bas, cell2._env)
+
+    if shls_slice is None:
+        shls_slice = (0, cell1.nbas, 0, cell2.nbas)
+    i0, i1, j0, j1 = shls_slice[:4]
+    j0 += cell1.nbas
+    j1 += cell1.nbas
+    ao_loc = moleintor.make_loc(bas, intor)
+    ni = ao_loc[i1] - ao_loc[i0]
+    nj = ao_loc[j1] - ao_loc[j0]
+    out = numpy.empty((nkpts,comp,ni,nj), dtype=numpy.complex128)
+
+    if hermi == 0:
+        aosym = 's1'
+    else:
+        aosym = 's2'
+    fill = getattr(libpbc, 'PBCnr2c_fill_k'+aosym)
+    fintor = getattr(moleintor.libcgto, intor)
+    cintopt = lib.c_null_ptr()
+
+    rcut = max(cell1.rcut, cell2.rcut)
+    Ls = numpy.asarray(cell1.get_lattice_Ls(rcut=rcut), order='C')
+    expkL = numpy.asarray(numpy.exp(1j*numpy.dot(kpts_lst, Ls.T)), order='C')
+    drv = libpbc.PBCnr2c_drv
+
+    drv(fintor, fill, out.ctypes.data_as(ctypes.c_void_p),
+        ctypes.c_int(nkpts), ctypes.c_int(comp), ctypes.c_int(len(Ls)),
+        Ls.ctypes.data_as(ctypes.c_void_p),
+        expkL.ctypes.data_as(ctypes.c_void_p),
+        (ctypes.c_int*4)(i0, i1, j0, j1),
+        ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt,
+        atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(pcell.natm),
+        bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(pcell.nbas),
+        env.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(env.size))
+
+    mat = []
+    for k, kpt in enumerate(kpts_lst):
+        v = out[k]
+        if hermi != 0:
+            for ic in range(comp):
+                lib.hermi_triu(v[ic], hermi=hermi, inplace=True)
+        if comp == 1:
+            v = v[0]
+        if abs(kpt).sum() < 1e-9:  # gamma_point
+            v = v.real
+        mat.append(v)
+
+    if kpts is None or numpy.shape(kpts) == (3,):  # A single k-point
+        mat = mat[0]
+    return mat
 
 @_pbc_intor.defjvp
 def _pbc_intor_jvp(intor, comp, hermi, kpts, kpt, shls_slice,
@@ -74,6 +146,7 @@ def _int1e_jvp_r0(mol, mol_t, intor, hermi, kpts, kpt, shls_slice):
         tangent_out = tangent_out[0]
     return tangent_out
 
+# FIXME use pyscfad's Ls
 @partial(custom_vjp, nondiff_argnums=tuple(range(1,7)))
 def _pbc_intor_rev(mol, intor, comp=None, hermi=0, kpts=None, kpt=None,
                    shls_slice=None):