Merge pull request #4 from Panadestein/as

Faster gradient and laplacian in 3body Jastrow

Makefile

@@ -1,4 +1,4 @@
-IRPF90 = irpf90 --codelet=factor_een:100000 #-a -d
+IRPF90 = irpf90 --codelet=jastrow_full:1000 #-s nelec:10 -s nnuc:2 -s ncord:5 #-a -d
 FC     = ifort -xHost -g -mkl=sequential
 FCFLAGS= -O2 -ffree-line-length-none -I .
 NINJA  = ninja

README.org

@@ -4,34 +4,35 @@
 
   Original equation:
   
- $$
+  $$
- \sum_{i=2}^{Ne} \sum_{j=1}^i \sum_{pkl} \sum_a^{Nn} c_{apkl}\, r_{ij}^k\, ( R_{ia}^l + R_{ja}^l) ( R_{ia} R_{ja})^m 
+  \sum_{i=2}^{Ne} \sum_{j=1}^i \sum_{pkl} \sum_a^{Nn} c_{apkl}\, r_{ij}^k\, ( R_{ia}^l + R_{ja}^l) ( R_{ia} R_{ja})^m 
- $$
+  $$
 
- Expanding, one obtains:
+  Expanding, one obtains:
  
- $$
+  $$
- \sum_{i=2}^{Ne} \sum_{j=1}^i \sum_{pkl} \sum_a^{Nn} c_{apkl} R_{ia}^{p-k-l}\, r_{ij}^k\, R_{ja}^{p-k+l} + c_{apkl} R_{ia}^{p-k+l}\, r_{ij}^k\, R_{ja}^{p-k-l}
+  \sum_{i=2}^{Ne} \sum_{j=1}^i \sum_{pkl} \sum_a^{Nn} c_{apkl} R_{ia}^{p-k-l}\, r_{ij}^k\, R_{ja}^{p-k+l} + c_{apkl} R_{ia}^{p-k+l}\, r_{ij}^k\, R_{ja}^{p-k-l}
- $$
+  $$
 
- The equation is symmetric if we exchange $i$ and $j$, so we can rewrite
+  The equation is symmetric if we exchange $i$ and $j$, so we can rewrite
 
- $$
+  $$
- \sum_{i=1}^{Ne} \sum_{j=1}^{Ne} \sum_{pkl} \sum_a^{Nn} c_{apkl} R_{ia}^{p-k-l}\, r_{ij}^k\, R_{ja}^{p-k+l} 
+  \sum_{i=1}^{Ne} \sum_{j=1}^{Ne} \sum_{pkl} \sum_a^{Nn} c_{apkl} R_{ia}^{p-k-l}\, r_{ij}^k\, R_{ja}^{p-k+l} 
- $$
+  $$
  
- If we reshape $R_{ja}^p$ as a matrix $R_{j,al}$ of size 
+  If we reshape $R_{ja}^p$ as a matrix $R_{j,al}$ of size 
- $N_e \times N_n(N_c+1)$, 
+  $N_e \times N_n(N_c+1)$, 
- for every $k$, we can pre-compute the matrix product
+  for every $k$, we can pre-compute the matrix product
 
- $$
+  $$
- C_{i,al}^{k} = \sum_j r_{ij}^k\, R_{i,al}
+  C_{i,al}^{k} = \sum_j r_{ij}^k\, R_{i,al}
- $$
+  $$
- which can be computed efficiently with BLAS.
+  which can be computed efficiently with BLAS.
- We can express the total Jastrow as:
+  We can express the total Jastrow as:
 
- $$
+  $$
- \sum_{i=1}^{Ne} \sum_{pkl} \sum_a^{Nn}
+  \sum_{i=1}^{Ne} \sum_{pkl} \sum_a^{Nn}
- c_{apkl} R_{ia}^{p-k-l}\, C_{i,a(p-k+l)}^k
+  c_{apkl} R_{ia}^{p-k-l}\, C_{i,a(p-k+l)}^k
- $$
+  $$
 
+  

el_nuc_el.irp.f

@@ -7,8 +7,6 @@ BEGIN_PROVIDER [ double precision, factor_een ]
   END_DOC
   integer                        :: i, j, a, p, k, l, lmax, m, n
   double precision               :: cn, accu2, accu
-  double precision               :: f(nnuc,0:ncord-2,0:ncord-2)
-  double precision               :: tmp_c(nelec,nnuc,0:ncord,0:ncord-1)
 
 !  factor_een = factor_een_blas
 !  return
@@ -35,10 +33,9 @@ BEGIN_PROVIDER [ double precision, factor_een ]
            do i = 1, nelec
              accu = accu +                                           &
                  rescale_een_e(i,j,k) *                              &
-                 rescale_een_n(i,a,m) *                              &
+                 rescale_een_n(i,a,m)
-                 rescale_een_n(j,a,m+l)
            enddo
-           accu2 = accu2 + accu
+           accu2 = accu2 + accu*rescale_een_n(j,a,m+l)
          enddo
          factor_een = factor_een + accu2 * cn
        enddo
@@ -50,6 +47,66 @@ BEGIN_PROVIDER [ double precision, factor_een ]
 END_PROVIDER
 
 BEGIN_PROVIDER [ double precision, factor_een_deriv_e, (4, nelec) ]
+  implicit none
+  integer                        :: i, j, a, p, k, l, lmax, m, n
+  double precision               :: cn, accu, accu2, daccu(1:4), daccu2(1:4)
+
+ factor_een_deriv_e(1:4,1:nelec) = 0.0d0
+
+ do p = 2, ncord
+   do k = 0, p - 1
+     if (k /= 0) then
+       lmax = p - k
+     else
+       lmax = p - k - 2
+     endif
+     do l = 0, lmax
+       if ( iand(p - k - l, 1) == 1) cycle
+
+       m = (p - k - l) / 2
+
+       do a = 1, nnuc
+         cn = cord_vect_lkp(l, k, p, typenuc_arr(a))
+         do j = 1, nelec
+           accu=0.d0
+           accu2 = 0.d0
+           daccu (1:4) = 0.d0
+           daccu2(1:4) = 0.d0
+           do i = 1, nelec
+             accu = accu +                                           &
+                 rescale_een_e(i,j,k) *                              &
+                 rescale_een_n(i,a,m)
+             accu2 = accu2 +                             &
+                 rescale_een_e(i,j,k) *                              &
+                 rescale_een_n(i,a,m+l)
+             daccu(1:4) = daccu(1:4) +                               &
+                 rescale_een_e_deriv_e(1:4,j,i,k) *                  &
+                 rescale_een_n(i,a,m)
+             daccu2(1:4) = daccu2(1:4) +                             &
+                 rescale_een_e_deriv_e(1:4,j,i,k) *                  &
+                 rescale_een_n(i,a,m+l)
+
+           enddo
+           factor_een_deriv_e(1:4,j) = factor_een_deriv_e(1:4,j) +   &
+               (accu * rescale_een_n_deriv_e(1:4,j,a,m+l) + daccu(1:4) * rescale_een_n(j,a,m+l) + &
+                daccu2(1:4)* rescale_een_n(j,a,m) + accu2*rescale_een_n_deriv_e(1:4,j,a,m)) * cn
+
+           factor_een_deriv_e(4,j) = factor_een_deriv_e(4,j) + 2.d0*( &
+             daccu (1) * rescale_een_n_deriv_e(1,j,a,m+l) + &
+             daccu (2) * rescale_een_n_deriv_e(2,j,a,m+l) + &
+             daccu (3) * rescale_een_n_deriv_e(3,j,a,m+l) + &
+             daccu2(1) * rescale_een_n_deriv_e(1,j,a,m  ) + &
+             daccu2(2) * rescale_een_n_deriv_e(2,j,a,m  ) + &
+             daccu2(3) * rescale_een_n_deriv_e(3,j,a,m  ) )*cn
+         enddo
+       enddo
+     enddo
+   enddo
+ enddo
+
+END_PROVIDER
+
+BEGIN_PROVIDER [ double precision, factor_een_deriv_e_ref, (4, nelec) ]
  implicit none
  BEGIN_DOC
  ! Dimensions 1-3 : dx, dy, dz
@@ -60,7 +117,7 @@ BEGIN_PROVIDER [ double precision, factor_een_deriv_e, (4, nelec) ]
  double precision, dimension(4) :: driam, drjam_cn, drial, drjal, drijk
  double precision :: cn, v1, v2, d1, d2, lap1, lap2
 
- factor_een_deriv_e = 0.0d0
+ factor_een_deriv_e_ref = 0.0d0
 
  do p = 2, ncord
     do k = 0 , p - 1
@@ -105,14 +162,14 @@ BEGIN_PROVIDER [ double precision, factor_een_deriv_e, (4, nelec) ]
                       d2 = drjam_cn(ii) * riam
                       lap1 = lap1 + d1 * d2
                       lap2 = lap2 + drijk(ii) * drjal(ii)
-                      factor_een_deriv_e(ii, j) = factor_een_deriv_e(ii, j) + v1 * d2 + d1 * v2
+                      factor_een_deriv_e_ref(ii, j) = factor_een_deriv_e_ref(ii, j) + v1 * d2 + d1 * v2
                    enddo
 
                    ! v(x) u''(x) + 2 * u'(x) v'(x) + u(x) v''(x)
                    ii = 4
                    d1 = drijk(ii) * rial + rijk * drjal(ii) + lap2 + lap2
                    d2 = drjam_cn(ii) * riam
-                   factor_een_deriv_e(ii, j) = factor_een_deriv_e(ii, j) + v1 * d2 + d1 * v2 + lap1 + lap1
+                   factor_een_deriv_e_ref(ii, j) = factor_een_deriv_e_ref(ii, j) + v1 * d2 + d1 * v2 + lap1 + lap1
 
                 enddo
              enddo

el_nuc_el_blas.irp.f

@@ -57,70 +57,57 @@ BEGIN_PROVIDER [ double precision, factor_een_deriv_e_blas, (4, nelec) ]
  ! Dimensions 1-3 : dx, dy, dz
  ! Dimension 4 : d2x + d2y + d2z
  END_DOC
- integer :: i, ii, j, a, p, k, l, lmax, m
+
- double precision :: riam, rjam_cn, rial, rjal, rijk
+ integer                        :: i, j, a, p, k, l, lmax, m, n
- double precision, dimension(4) :: driam, drjam_cn, drial, drjal, drijk
+ double precision               :: cn(nnuc), accu(4)
- double precision :: cn, v1, v2, d1, d2, lap1, lap2
+ double precision               :: f(nnuc,0:ncord-2,0:ncord-2)
+ double precision               :: tmp_c(nelec,nnuc,0:ncord,0:ncord-1)
+ double precision               :: dtmp_c(4,nelec,nnuc,0:ncord,0:ncord-1)
 
  factor_een_deriv_e_blas = 0.0d0
 
- do p = 2, ncord
+ ! r_{ij}^k . R_{ja}^l -> tmp_c_{ia}^{kl}
-    do k = 0 , p - 1
+ do k=0,ncord-1
-       if (k /= 0) then
+   call dgemm('N','N', nelec, nnuc*(ncord+1), nelec, 1.d0,           &
-         lmax = p - k
+       rescale_een_e(1,1,k), size(rescale_een_e,1),                  &
-       else
+       rescale_een_n(1,1,0), size(rescale_een_n,1), 0.d0,            &
-         lmax = p - k - 2
+       tmp_c(1,1,0,k), size(tmp_c,1))
-       endif
-
-       do l = 0, lmax
-          if ( iand(p - k - l, 1) == 1) cycle
-          m = (p - k - l) / 2
-
-          do a = 1, nnuc
-             cn = cord_vect_lkp(l, k, p, typenuc_arr(a))
-
-             do j = 1, nelec
-                rjal = rescale_een_n(j, a, l)
-                rjam_cn = rescale_een_n(j, a, m) * cn
-
-                do ii = 1, 4
-                   drjal(ii) = rescale_een_n_deriv_e(ii, j, a, l)
-                   drjam_cn(ii) = rescale_een_n_deriv_e(ii, j, a, m) * cn
-                enddo
-
-                do i = 1, nelec
-                   rial = rescale_een_n(i, a, l) + rjal
-                   riam = rescale_een_n(i, a, m)
-                   rijk = rescale_een_e(i, j, k)
-
-                   do ii = 1, 4
-                      drijk(ii) = rescale_een_e_deriv_e(ii, j, i, k)
-                   enddo
-
-                   v1 = rijk * rial    ! v(x)
-                   v2 = rjam_cn * riam ! u(x)
-
-                   lap1 = 0.0d0
-                   lap2 = 0.0d0
-                   do ii = 1, 3
-                      d1 = drijk(ii) * rial + rijk * drjal(ii)
-                      d2 = drjam_cn(ii) * riam
-                      lap1 = lap1 + d1 * d2
-                      lap2 = lap2 + drijk(ii) * drjal(ii)
-                      factor_een_deriv_e_blas(ii, j) = factor_een_deriv_e_blas(ii, j) + v1 * d2 + d1 * v2
-                   enddo
-
-                   ! v(x) u''(x) + 2 * u'(x) v'(x) + u(x) v''(x)
-                   ii = 4
-                   d1 = drijk(ii) * rial + rijk * drjal(ii) + lap2 + lap2
-                   d2 = drjam_cn(ii) * riam
-                   factor_een_deriv_e_blas(ii, j) = factor_een_deriv_e_blas(ii, j) + v1 * d2 + d1 * v2 + lap1 + lap1
-
-                enddo
-             enddo
-          enddo
-       enddo
-    enddo
  enddo
 
+ ! dr_{ij}^k . R_{ja}^l -> dtmp_c_{ia}^{kl}
+ do k=0,ncord-1
+   call dgemm('N','N', 4*nelec, nnuc*(ncord+1), nelec, 1.d0,         &
+       rescale_een_e_deriv_e(1,1,1,k), 4*size(rescale_een_e_deriv_e,2),&
+       rescale_een_n(1,1,0), size(rescale_een_n,1), 0.d0,            &
+       dtmp_c(1,1,1,0,k), 4*size(dtmp_c,2))
+ enddo
+
+ do p = 2, ncord
+   do k = 0, p - 1
+     m = p-k
+     if (k > 0) then
+       lmax = m
+     else
+       lmax = m - 2
+     endif
+
+     n = shiftr(m,1)
+     do l = iand(m, 1), lmax, 2
+
+       do a = 1, nnuc
+         cn(a) = cord_vect_lkp(l, k, p, typenuc_arr(a))
+       enddo
+
+       do a = 1, nnuc
+         do i=1,nelec
+           accu(1:4) = rescale_een_n(i,a,n) * dtmp_c(1:4,i,a,n+l,k) &
+                       + rescale_een_n_deriv_e(1:4,i,a,n) * tmp_c(i,a,n+l,k)
+           factor_een_deriv_e_blas(1:4,i) = factor_een_deriv_e_blas(1:4,i) + accu(1:4) * cn(a)
+         enddo
+       enddo
+       n = n-1
+
+     enddo
+   enddo
+ enddo
 END_PROVIDER

jastrow.irp.f

@@ -2,6 +2,12 @@ program jastrow
   implicit none
   print *, 'The total Jastrow factor'
   print *, jastrow_full
+  print *, 'REF'
+  print *, factor_een_deriv_e_ref
+  print *, 'X'
+  print *, factor_een_deriv_e
+  print *, 'BLAS'
+  print *, factor_een_deriv_e_blas
   !PROVIDE jastrow_full
 
 end program

rescale.irp.f

@@ -279,6 +279,7 @@ BEGIN_PROVIDER [double precision, rescale_een_e_deriv_e, (4, nelec, nelec, 0:nco
  enddo
 END_PROVIDER
 
+
 BEGIN_PROVIDER [double precision, elec_dist_deriv_e, (4, nelec, nelec)]
  BEGIN_DOC
  ! Derivative of R_{ij} wrt x