revamp libcint to return c_contig view with zero copy

Author: brianz98

forte2/helpers/ndarray.h

@@ -158,3 +158,18 @@ auto make_zeros(const std::array<size_t, N>& shape) {
 
     return array;
 }
+
+template <typename Type, typename T, int N>
+auto fortran_to_c(nb::ndarray<Type, T, nb::ndim<N>, nb::f_contig>& arr) {
+    std::array<std::size_t, N> C_shape;
+    std::array<std::int64_t, N> C_strides;
+    for (std::size_t i = 0; i < static_cast<std::size_t>(N); ++i) {
+        C_shape[i] = arr.shape(N - 1 - i);
+        C_strides[i] = arr.stride(N - 1 - i);
+    }
+    T* data = arr.data();
+
+    nb::object owner = arr.cast();
+
+    return nb::ndarray<Type, T, nb::ndim<N>, nb::c_contig>(data, static_cast<std::size_t>(N), C_shape.data(), owner, C_strides.data());
+}

forte2/integrals/integrals.py

@@ -715,13 +715,13 @@ def _parse_basis_args_cint_3c2e(system, basis1, basis2, basis3, origin=None):
             system, system.basis, common_origin=origin
         )
         nsh_bas = system.basis.nshells
-        shell_slice = [0, nsh_bas, 0, nsh_bas, nsh_bas, nsh_bas + nsh_aux]
+        shell_slice = [nsh_bas, nsh_bas + nsh_aux, 0, nsh_bas, 0, nsh_bas]
     elif basis2 is not None and basis3 is None:
         bas_atm, bas_bas, bas_env = basis_to_cint_envs(
             system, basis2, common_origin=origin
         )
         nsh_bas = basis2.nshells
-        shell_slice = [0, nsh_bas, 0, nsh_bas, nsh_bas, nsh_bas + nsh_aux]
+        shell_slice = [nsh_bas, nsh_bas + nsh_aux, 0, nsh_bas, 0, nsh_bas]
     elif basis2 is not None and basis3 is not None and basis2 != basis3:
         raise ValueError(
             "libcint doesn't support (P|QR) with Q and R being different basis sets."
@@ -734,10 +734,10 @@ def _parse_basis_args_cint_3c2e(system, basis1, basis2, basis3, origin=None):
 
 
 def _f2c(arr):
-    if arr.shape[-1] == 1:
-        return np.ascontiguousarray(arr[..., 0])
+    if arr.shape[0] == 1:
+        return arr[0]
     else:
-        return np.ascontiguousarray(np.rollaxis(arr, -1, 0))
+        return arr
 
 
 def cint_overlap(system, basis1=None, basis2=None):
@@ -783,6 +783,7 @@ def cint_overlap_spinor(system, basis1=None, basis2=None):
     _require_libcint()
     atm, bas, env, shell_slice = _parse_basis_args_cint_1e(system, basis1, basis2)
     res = ints.cint_int1e_ovlp_spinor(shell_slice, atm, bas, env)
+    np.conjugate(res, out=res)
     return _f2c(res)
 
 
@@ -857,7 +858,10 @@ def cint_opVop(system, basis1=None, basis2=None):
     _require_libcint()
     atm, bas, env, shell_slice = _parse_basis_args_cint_1e(system, basis1, basis2)
     res = ints.cint_int1e_spnucsp_sph(shell_slice, atm, bas, env)
+    # the x/y/z components are antisymmetric due to the cross produce (swapping the two basis functions changes the sign)
+    # libcint returns the transposed version, hence the sign flip
     # C-layout, first index is the integral component (slowest changing)
+    res[:-1] *= -1
     return _f2c(res)
 
 
@@ -885,6 +889,7 @@ def cint_opVop_spinor(system, basis1=None, basis2=None):
     _require_libcint()
     atm, bas, env, shell_slice = _parse_basis_args_cint_1e(system, basis1, basis2)
     res = ints.cint_int1e_spnucsp_spinor(shell_slice, atm, bas, env)
+    np.conj(res, out=res)
     return _f2c(res)
 
 
@@ -958,4 +963,4 @@ def cint_coulomb_3c(system, basis1=None, basis2=None, basis3=None):
         system, basis1, basis2, basis3
     )
     res = ints.cint_int3c2e_sph(shell_slice, atm, bas, env)
-    return res.transpose(2, 0, 1).copy()  # copy makes sure it's C-contiguous
+    return res

forte2/integrals/libcint_compute.h

@@ -29,13 +29,13 @@ namespace forte2 {
 // templated function to compute two-center integrals with M components
 // (i.e., 1 for scalar integrals, 3 for dipoles integrals, etc.)
 template <std::size_t M>
-np_tensor3_f cint_int2c(CIntorFunc intor, const std::vector<int>& shell_slice, np_matrix_int atm,
+np_tensor3_c cint_int2c(CIntorFunc intor, const std::vector<int>& shell_slice, np_matrix_int atm,
                         np_matrix_int bas, np_vector env) {
 
-    const int ish_0 = static_cast<int>(shell_slice[0]);
-    const int ish_1 = static_cast<int>(shell_slice[1]);
-    const int jsh_0 = static_cast<int>(shell_slice[2]);
-    const int jsh_1 = static_cast<int>(shell_slice[3]);
+    const int jsh_0 = static_cast<int>(shell_slice[0]);
+    const int jsh_1 = static_cast<int>(shell_slice[1]);
+    const int ish_0 = static_cast<int>(shell_slice[2]);
+    const int ish_1 = static_cast<int>(shell_slice[3]);
     const int nish = ish_1 - ish_0;
     const int njsh = jsh_1 - jsh_0;
 
@@ -76,17 +76,17 @@ np_tensor3_f cint_int2c(CIntorFunc intor, const std::vector<int>& shell_slice, n
         }
     }
 
-    return ints;
+    return fortran_to_c<nb::numpy, double, 3>(ints);
 }
 
 template <std::size_t M>
-np_tensor3_complex_f cint_int2c_spinor(CIntorFuncSpinor intor, const std::vector<int>& shell_slice,
+np_tensor3_complex_c cint_int2c_spinor(CIntorFuncSpinor intor, const std::vector<int>& shell_slice,
                                        np_matrix_int atm, np_matrix_int bas, np_vector env) {
 
-    const int ish_0 = static_cast<int>(shell_slice[0]);
-    const int ish_1 = static_cast<int>(shell_slice[1]);
-    const int jsh_0 = static_cast<int>(shell_slice[2]);
-    const int jsh_1 = static_cast<int>(shell_slice[3]);
+    const int jsh_0 = static_cast<int>(shell_slice[0]);
+    const int jsh_1 = static_cast<int>(shell_slice[1]);
+    const int ish_0 = static_cast<int>(shell_slice[2]);
+    const int ish_1 = static_cast<int>(shell_slice[3]);
     const int nish = ish_1 - ish_0;
     const int njsh = jsh_1 - jsh_0;
 
@@ -127,7 +127,7 @@ np_tensor3_complex_f cint_int2c_spinor(CIntorFuncSpinor intor, const std::vector
         }
     }
 
-    return ints;
+    return fortran_to_c<nb::numpy, std::complex<double>, 3>(ints);
 }
 
 void fill_3c_sym(np_tensor3_f& ints) {
@@ -166,14 +166,14 @@ void fill_3c_sym(np_tensor3_f& ints) {
 }
 
 // function to compute three-center integrals
-np_tensor3_f cint_int3c(CIntorFunc intor, const std::vector<int>& shell_slice, np_matrix_int atm,
+np_tensor3_c cint_int3c(CIntorFunc intor, const std::vector<int>& shell_slice, np_matrix_int atm,
                         np_matrix_int bas, np_vector env) {
-    const int ish_0 = static_cast<int>(shell_slice[0]);
-    const int ish_1 = static_cast<int>(shell_slice[1]);
+    const int ksh_0 = static_cast<int>(shell_slice[0]);
+    const int ksh_1 = static_cast<int>(shell_slice[1]);
     const int jsh_0 = static_cast<int>(shell_slice[2]);
     const int jsh_1 = static_cast<int>(shell_slice[3]);
-    const int ksh_0 = static_cast<int>(shell_slice[4]);
-    const int ksh_1 = static_cast<int>(shell_slice[5]);
+    const int ish_0 = static_cast<int>(shell_slice[4]);
+    const int ish_1 = static_cast<int>(shell_slice[5]);
 
     // Whether i and j shells are identical, if they are we can exploit symmetry
     bool ij_sym = (ish_0 == jsh_0) && (ish_1 == jsh_1);
@@ -264,7 +264,7 @@ np_tensor3_f cint_int3c(CIntorFunc intor, const std::vector<int>& shell_slice, n
         fill_3c_sym(ints);
     }
 
-    return ints;
+    return fortran_to_c<nb::numpy, double, 3>(ints);
 }
 
 } // namespace forte2

forte2/integrals/libcint_three_center.h

@@ -42,7 +42,7 @@ namespace forte2 {
 //     return cint_int2c<1>(int1e_ovlp_sph, shell_slice, atm, bas, env);
 // }
 #define DECL_CINT_FORTE2_FUNC_SPH(name)                                                            \
-    np_tensor3_f cint_##name##_sph(const std::vector<int>& shell_slice, np_matrix_int atm,         \
+    np_tensor3_c cint_##name##_sph(const std::vector<int>& shell_slice, np_matrix_int atm,         \
                                    np_matrix_int bas, np_vector env) {                             \
         return cint_int3c(name##_sph, shell_slice, atm, bas, env);                               \
     }

forte2/integrals/libcint_two_center.h

@@ -55,7 +55,7 @@ namespace forte2 {
 //     return cint_int2c<1>(int1e_ovlp_sph, shell_slice, atm, bas, env);
 // }
 #define DECL_CINT_FORTE2_FUNC_SPH(name, comp) \
-np_tensor3_f cint_##name##_sph(const std::vector<int>& shell_slice, np_matrix_int atm, np_matrix_int bas, \
+np_tensor3_c cint_##name##_sph(const std::vector<int>& shell_slice, np_matrix_int atm, np_matrix_int bas, \
                         np_vector env) { \
     return cint_int2c<comp>(name##_sph, shell_slice, atm, bas, env); \
 }
@@ -69,7 +69,7 @@ DECL_CINT_FORTE2_FUNC_SPH(int2c2e, 1)
 #undef DECL_CINT_FORTE2_FUNC_SPH
 
 #define DECL_CINT_FORTE2_FUNC_SPINOR(name, comp) \
-np_tensor3_complex_f cint_##name##_spinor(const std::vector<int>& shell_slice, np_matrix_int atm, \
+np_tensor3_complex_c cint_##name##_spinor(const std::vector<int>& shell_slice, np_matrix_int atm, \
                                    np_matrix_int bas, np_vector env) { \
     return cint_int2c_spinor<comp>(name##_spinor, shell_slice, atm, bas, env); \
 }

tests/integrals/test_libcint.py

@@ -70,10 +70,10 @@ def test_libcint_spnucsp_sph():
     system = System(xyz, basis_set="sto-3g", minao_basis_set=None)
     s_cint = integrals.cint_opVop(system)
     c_int2 = integrals.opVop(system)
-    assert np.linalg.norm(s_cint[3] - c_int2[0]) < 1e-6  # I2
-    assert np.linalg.norm(s_cint[0] - c_int2[1]) < 1e-6  # sigma_x
-    assert np.linalg.norm(s_cint[1] - c_int2[2]) < 1e-6  # sigma_y
-    assert np.linalg.norm(s_cint[2] - c_int2[3]) < 1e-6  # sigma_z
+    assert np.linalg.norm(s_cint[3] - c_int2[0]) < 1e-6  # sigma_x
+    assert np.linalg.norm(s_cint[0] - c_int2[1]) < 1e-6  # sigma_y
+    assert np.linalg.norm(s_cint[1] - c_int2[2]) < 1e-6  # sigma_z
+    assert np.linalg.norm(s_cint[2] - c_int2[3]) < 1e-6  # I2
     assert np.linalg.norm(s_cint) == pytest.approx(5982385.234519612, rel=1e-6)