Refactor: implementation of two-center integral table in module_nao

source/module_base/math_sphbes.cpp

@@ -775,15 +775,34 @@ double Sphbes::sphbesj(const int l, const double x)
     }
 }
 
+double Sphbes::dsphbesj(const int l, const double x)
+{
+    assert( l >= 0 );
+    assert( x >= 0 );
+    return l == 0 ? -sphbesj(1, x) : ( l * sphbesj(l - 1, x) - (l + 1) * sphbesj(l + 1, x) ) / (2 * l + 1);
+}
+
 void Sphbes::sphbesj(const int n,
+                     const double* const r,
+                     const double q,
+                     const int l,
+                     double* const jl)
+{
+    for (int i = 0; i != n; ++i)
+    {
+        jl[i] = Sphbes::sphbesj(l, q * r[i]);
+    }
+}
+
+void Sphbes::dsphbesj(const int n,
              const double* const r,
              const double q,
              const int l,
-             double* const jl)
+             double* const djl)
 {
     for (int i = 0; i != n; ++i)
     {
-        jl[i] = Sphbes::sphbesj(l, q * r[i]);
+        djl[i] = Sphbes::dsphbesj(l, q * r[i]);
     }
 }
 

source/module_base/math_sphbes.h

@@ -97,11 +97,25 @@ class Sphbes
                           const double x //!< [in] argument
     );
 
-    static void sphbesj(const int n,
-                        const double* const r,
-                        const double q,
-                        const int l,
-                        double* const jl
+    //! derivative of spherical Bessel function
+    static double dsphbesj(const int l,   //!< [in] order
+                          const double x  //!< [in] argument
+    );
+
+    //! computes the values of l-th order spherical Bessel function at q*r[ir]
+    static void sphbesj(const int n,           //!< [in] number of r grid points
+                        const double* const r, //!< [in] r grid
+                        const double q,        //!< [in] wave vector
+                        const int l,           //!< [in] order of the spherical Bessel function
+                        double* const jl       //!< [out] results
+    );
+
+    //! computes the derivative of l-th order spherical Bessel function at q*r[ir]
+    static void dsphbesj(const int n,          //!< [in] number of r grid points
+                        const double* const r, //!< [in] r grid
+                        const double q,        //!< [in] wave vector
+                        const int l,           //!< [in] order of the spherical Bessel function
+                        double* const djl      //!< [out] results
     );
 
 private:

source/module_base/spherical_bessel_transformer.cpp

@@ -19,6 +19,9 @@ SphericalBesselTransformer::~SphericalBesselTransformer()
 {
     fftw_destroy_plan(rfft_plan_);
     fftw_free(f_);
+    delete[] grid_in_;
+    delete[] grid_out_;
+    delete[] jl_;
 }
 
 long long int SphericalBesselTransformer::spherical_bessel_sincos_polycoef(const bool get_sine,
@@ -73,49 +76,22 @@ long long int SphericalBesselTransformer::spherical_bessel_sincos_polycoef(const
            - (n >= 2 ? spherical_bessel_sincos_polycoef(get_sine, l - 2, n - 2) : 0);
 }
 
-void SphericalBesselTransformer::radrfft(const int l,
-                                         const int ngrid,
-                                         const double cutoff,
-                                         const double* const in,
-                                         double* const out,
-                                         const int p)
+void SphericalBesselTransformer::_radrfft_base(const int l,
+                                               const int ngrid,
+                                               const double cutoff,
+                                               const double* const in,
+                                               double* const out,
+                                               const int p)
 {
-    /*
-     * An l-th order spherical Bessel transform F(x) -> G(y) can be expressed in terms of Fourier transforms:
-     *
-     *           l
-     *          ---    1        / +inf            -iyx
-     * G(y) =   \   -------  Re |      dr f(n,x) e
-     *          /     l+1-n     / -inf
-     *          ---  y
-     *          n=0
-     *
-     * where
-     *
-     *              1                              F(x)
-     * f(n,x) = ---------- [ c(l,n) + i*s(l,n) ] --------
-     *          sqrt(2*pi)                         l-n-1   ,
-     *                                            x
-     *
-     * c(l,n) / s(l,n) are sin / cos coefficients from spherical_bessel_sincos_polycoef, and
-     * the domain of F(x) is extended to negative values by defining F(-x) = pow(-1,l)*F(x).
-     *
-     * With an appropriate grid, the Fourier transform can be approximated by a discrete Fourier transform.
-     * Note that given l & n, c(l,n) and s(l,n) cannot be both non-zero. Therefore, each FFT input
-     * array is either purely real or purely imaginary, which suggests the use of real-input FFT.
-     *                                                                                                  */
-    assert(l >= 0);
-    assert(ngrid > 1);
-    assert(p <= 2);
+    // this function does not support in-place transform (but radrfft does)
+    assert(in != out);
 
     int n = ngrid - 1;
     double dx = cutoff / n;
     double dy = PI / cutoff;
     double pref = dx / std::sqrt(2. * PI);
 
-    // use a temporary output array in case the transform is in-place
-    double* out_tmp = new double[ngrid];
-    std::fill(out_tmp, out_tmp + ngrid, 0.0);
+    std::fill(out, out + ngrid, 0.0);
 
     // rfft buffer (f_) allocation and plan creation
     rfft_prepare(2 * n);
@@ -125,12 +101,11 @@ void SphericalBesselTransformer::radrfft(const int l,
 
     for (int m = 0; m <= l; ++m)
     {
-
         // m even --> sin; f[2*n-i] = -f[i]; out += -imag(rfft(f)) / y^(l+1-m)
         // m odd  --> cos; f[2*n-i] = +f[i]; out += +real(rfft(f)) / y^(l+1-m)
         bool flag = (m % 2 == 0);
 
-        int sincos_coef = spherical_bessel_sincos_polycoef(flag, l, m);
+        long long int sincos_coef = spherical_bessel_sincos_polycoef(flag, l, m);
 
         f[0] = f[n] = 0.0;
         for (int i = 1; i != n; ++i)
@@ -145,7 +120,7 @@ void SphericalBesselTransformer::radrfft(const int l,
         // out[0] is handled independently
         for (int j = 1; j <= n; ++j)
         {
-            out_tmp[j] = (out_tmp[j] + (flag ? -f_[j][1] : f_[j][0])) / (j * dy);
+            out[j] = (out[j] + (flag ? -f_[j][1] : f_[j][0])) / (j * dy);
         }
     }
 
@@ -155,20 +130,90 @@ void SphericalBesselTransformer::radrfft(const int l,
     {
         for (int i = 0; i <= n; ++i)
         {
-            out_tmp[0] += 2. * pref * in[i] * std::pow(i * dx, 2 - p); // p <= 2 is required!
+            out[0] += 2. * pref * in[i] * std::pow(i * dx, 2 - p); // p <= 2 is required!
+        }
+    }
+}
+
+void SphericalBesselTransformer::radrfft(const int l,
+                                         const int ngrid,
+                                         const double cutoff,
+                                         const double* const in,
+                                         double* const out,
+                                         const int p,
+                                         const bool deriv)
+{
+    /*
+     * An l-th order spherical Bessel transform F(x) -> G(y) can be expressed in terms of Fourier transforms:
+     *
+     *           l
+     *          ---    1        / +inf            -iyx
+     * G(y) =   \   -------  Re |      dr f(n,x) e
+     *          /     l+1-n     / -inf
+     *          ---  y
+     *          n=0
+     *
+     * where
+     *
+     *              1                              F(x)
+     * f(n,x) = ---------- [ c(l,n) + i*s(l,n) ] --------
+     *          sqrt(2*pi)                         l-n-1   ,
+     *                                            x
+     *
+     * c(l,n) / s(l,n) are sin / cos coefficients from spherical_bessel_sincos_polycoef, and
+     * the domain of F(x) is extended to negative values by defining F(-x) = pow(-1,l)*F(x).
+     *
+     * With an appropriate grid, the Fourier transform can be approximated by a discrete Fourier transform.
+     * Note that given l & n, c(l,n) and s(l,n) cannot be both non-zero. Therefore, each FFT input
+     * array is either purely real or purely imaginary, which suggests the use of real-input FFT.
+     *
+     * If deriv == true, dG(y)/dy is calculated instead of G(y). This is done by using the recurrence
+     * relation of j_l(x) to convert the original transform to two Spherical Bessel transforms of order
+     * l-1 and l+1 respectively.
+     *                                                                                                  */
+    assert(l >= 0);
+    assert(ngrid > 1);
+    assert(p <= 2);
+
+    double* out_tmp = new double[deriv && l > 0 ? 2 * ngrid : ngrid];
+    int n_direct = 0;
+
+    if (deriv)
+    {
+        if (l == 0)
+        {
+            _radrfft_base(1, ngrid, cutoff, in, out_tmp, p - 1);
+            std::for_each(out_tmp, out_tmp + ngrid, [](double& x) { x = -x; });
+        }
+        else
+        {
+            _radrfft_base(l - 1, ngrid, cutoff, in, out_tmp, p - 1);
+            _radrfft_base(l + 1, ngrid, cutoff, in, &out_tmp[ngrid], p - 1);
+            std::transform(out_tmp, out_tmp + ngrid, &out_tmp[ngrid], out_tmp,
+                    [l](double x, double y) { return (l * x - (l + 1) * y) / (2 * l + 1); });
         }
+
+        n_direct = static_cast<int>(ngrid * std::pow(1e-8, 1. / (l + 1)));
+    }
+    else
+    {
+        _radrfft_base(l, ngrid, cutoff, in, out_tmp, p);
+        n_direct = (l == 0) ? 0 : static_cast<int>(ngrid * std::pow(1e-8, 1. / l));
     }
 
-    // FFT-based method is not accurate for small k at high l
+    // FFT-based method does not yield accurate results for small y at high l
     // use numerical integration in this case
-    int n_direct = (l == 0) ? 0 : static_cast<int>(ngrid * std::pow(1e-8, 1. / l));
     if (n_direct > 0)
     {
+        double dx = cutoff / (ngrid - 1);
+        double dy = PI / cutoff;
         double* grid_in = new double[ngrid];
         double* grid_out = new double[n_direct];
         std::for_each(grid_in, grid_in + ngrid, [&](double& x) { x = (&x - grid_in) * dx; });
         std::for_each(grid_out, grid_out + n_direct, [&](double& y) { y = ((&y - grid_out) + 1) * dy; });
-        direct(l, ngrid, grid_in, in, n_direct, grid_out, &out_tmp[1], p);
+
+        direct(l, ngrid, grid_in, in, n_direct, grid_out, &out_tmp[1], p, deriv);
+
         delete[] grid_in;
         delete[] grid_out;
     }
@@ -184,49 +229,43 @@ void SphericalBesselTransformer::direct(const int l,
                                         const int ngrid_out,
                                         const double* const grid_out,
                                         double* const out,
-                                        const int p)
+                                        const int p,
+                                        const bool deriv)
 {
-    // TODO:
-    // 1. avoid new/delete rab/r2f/integrand/jl every time; use some buffer class member instead
-    // 2. cache the spherical Bessel function values for given (l, grid_in, grid_out). This should be done
-    //    alongside the change of for-loop sequence in building TwoCenterTable.
-
     assert(p <= 2);
     assert(grid_in[0] >= 0.0 && std::is_sorted(grid_in, grid_in + ngrid_in, std::less_equal<double>()));
     assert(grid_out[0] >= 0.0 && std::is_sorted(grid_out, grid_out + ngrid_out, std::less_equal<double>()));
 
-    double pref = std::sqrt(2.0 / ModuleBase::PI);
+    double* buffer = new double[3 * ngrid_in];
 
     // NOTE: std::adjacent_difference returns an array of the same size as the input array
     // with rab[0] = grid_in[0] and rab[i] = grid_in[i] - grid[i-1] for i > 0
-    double* rab = new double[ngrid_in];
+    double* rab = buffer;
     std::adjacent_difference(grid_in, grid_in + ngrid_in, rab);
 
-    // r^2*f(r) (integrand without j_l)
-    double* r2f = new double[ngrid_in];
-    std::memcpy(r2f, in, ngrid_in * sizeof(double));
-    std::for_each(r2f, r2f + ngrid_in, [&](double& x) { x *= std::pow(grid_in[&x - r2f], 2 - p); });
+    // integrand without the jl (or jl') part
+    // note that (d/dy)jl(y*x) = x * (jl')(y*x), hence the extra power of x
+    double* itgpart = buffer + ngrid_in;
+    std::memcpy(itgpart, in, ngrid_in * sizeof(double));
+    std::for_each(itgpart, itgpart + ngrid_in, [&](double& x) { x *= std::pow(grid_in[&x - itgpart], 2 - p + deriv); });
 
     // there's no need to use a temporary output array even if the transform is in-place
-    // because "in" is not referenced after the construction of r2f.
+    // because "in" is not referenced after the construction of itgpart.
 
-    double* integrand = new double[ngrid_in];
+    // cache the jl or jl' values for potential subsequent transforms on the same grid
+    cache(l, ngrid_in, grid_in, ngrid_out, grid_out, deriv);
 
-    // TODO cache the calculated jl
-    double* jl = new double[ngrid_in];
+    double* integrand = buffer + 2 * ngrid_in;
     for (int i_out = 0; i_out != ngrid_out; ++i_out)
     {
-        ModuleBase::Sphbes::sphbesj(ngrid_in, grid_in, grid_out[i_out], l, jl);
-        std::transform(r2f, r2f + ngrid_in, jl, integrand, std::multiplies<double>());
+        std::transform(itgpart, itgpart + ngrid_in, &jl_[i_out * ngrid_in_], integrand, std::multiplies<double>());
         out[i_out] = ModuleBase::Integral::simpson(ngrid_in, integrand, &rab[1]);
     }
-    delete[] jl;
 
+    double pref = std::sqrt(2.0 / ModuleBase::PI);
     std::for_each(out, out + ngrid_out, [pref](double& x) { x *= pref; });
 
-    delete[] r2f;
-    delete[] integrand;
-    delete[] rab;
+    delete[] buffer;
 }
 
 void SphericalBesselTransformer::rfft_in_place()
@@ -264,4 +303,45 @@ void SphericalBesselTransformer::set_fftw_plan_flag(const unsigned new_flag)
     }
 }
 
+void SphericalBesselTransformer::cache(const int l,
+                                       const int ngrid_in,
+                                       const double* const grid_in,
+                                       const int ngrid_out,
+                                       const double* const grid_out,
+                                       const bool deriv)
+{
+    // TODO: if the given input grid matchs the cached output grid and vice versa,
+    // we can just swap the cached grids and transpose jl_.
+    bool is_cached = deriv == is_deriv_ && l == l_ && ngrid_in <= ngrid_in_ && ngrid_out <= ngrid_out_
+                     && std::equal(grid_in, grid_in + ngrid_in, grid_in_)
+                     && std::equal(grid_out, grid_out + ngrid_out, grid_out_);
+
+    if (!is_cached)
+    {
+        delete[] jl_;
+        delete[] grid_in_;
+        delete[] grid_out_;
+
+        is_deriv_ = deriv;
+        l_ = l;
+        ngrid_in_ = ngrid_in;
+        ngrid_out_ = ngrid_out;
+
+        grid_in_ = new double[ngrid_in];
+        grid_out_ = new double[ngrid_out];
+        std::memcpy(grid_in_, grid_in, ngrid_in * sizeof(double));
+        std::memcpy(grid_out_, grid_out, ngrid_out * sizeof(double));
+
+        typedef void(*FuncType)(int, const double*, double, int, double*);
+        FuncType f = is_deriv_ ? static_cast<FuncType>(ModuleBase::Sphbes::dsphbesj) :
+                                 static_cast<FuncType>(ModuleBase::Sphbes::sphbesj);
+
+        jl_ = new double[ngrid_in * ngrid_out];
+        for (int i_out = 0; i_out != ngrid_out; ++i_out)
+        {
+            f(ngrid_in, grid_in, grid_out[i_out], l, &jl_[i_out * ngrid_in]);
+        }
+    }
+}
+
 } // namespace ModuleBase