From dd14a4f50d21eae44c2fd20341cae2d3f0837439 Mon Sep 17 00:00:00 2001
From: Anthony Scemama <scemama@irsamc.ups-tlse.fr>
Date: Sun, 3 Feb 2019 17:30:28 +0100
Subject: [PATCH] Better handling of pt2 tasks

---
 src/cipsi/pt2_stoch_routines.irp.f | 30 +++++++++++++++++++++---------
 src/cipsi/run_pt2_slave.irp.f      |  4 ++--
 src/davidson/diagonalize_ci.irp.f  |  4 +++-
 3 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/cipsi/pt2_stoch_routines.irp.f b/src/cipsi/pt2_stoch_routines.irp.f
index 83924bb6..ea339bee 100644
--- a/src/cipsi/pt2_stoch_routines.irp.f
+++ b/src/cipsi/pt2_stoch_routines.irp.f
@@ -10,13 +10,20 @@ END_PROVIDER
 &BEGIN_PROVIDER [ integer, pt2_n_tasks_max ]
   implicit none
   logical, external :: testTeethBuilding
-  integer :: i
-  integer :: e
-  e = elec_num - n_core_orb * 2
-  pt2_n_tasks_max = 1+min((e*(e-1)), int(dsqrt(dble(N_det_selectors)))/10)
-  do i=1,N_det_generators
-    pt2_F(i) = 1 + int(dble(pt2_n_tasks_max)*dabs(maxval(psi_coef_sorted_gen(i,:)))**(0.75d0))
+  integer :: i,j
+  pt2_n_tasks_max = elec_beta_num*elec_beta_num + elec_alpha_num*elec_beta_num  - n_core_orb*2 
+  pt2_n_tasks_max = min(pt2_n_tasks_max,1+N_det_generators/10000)
+  call write_int(6,pt2_n_tasks_max,'pt2_n_tasks_max')
+
+  pt2_F(:) = int(sqrt(float(pt2_n_tasks_max)))
+  do i=1,pt2_n_0(pt2_N_teeth/4)
+    pt2_F(i) = pt2_n_tasks_max
   enddo
+  do i=pt2_n_0(pt2_N_teeth-pt2_N_teeth/4), N_det_generators
+    pt2_F(i) = 1
+  enddo
+
+
 END_PROVIDER
 
  BEGIN_PROVIDER [ integer, pt2_N_teeth ]
@@ -29,7 +36,7 @@ END_PROVIDER
     pt2_N_teeth = 1
   else
     pt2_minDetInFirstTeeth = min(5, N_det_generators)
-    do pt2_N_teeth=100,2,-1
+    do pt2_N_teeth=400,2,-1
       if(testTeethBuilding(pt2_minDetInFirstTeeth, pt2_N_teeth)) exit
     end do
   end if
@@ -74,10 +81,14 @@ logical function testTeethBuilding(minF, N)
 
   n0 = 0
   testTeethBuilding = .false.
+  double precision :: f
+  integer :: minFN
+  minFN = N_det_generators - minF * N
+  f = 1.d0/dble(N)
   do
     u0 = tilde_cW(n0)
     r = tilde_cW(n0 + minF)
-    Wt = (1d0 - u0) / dble(N)
+    Wt = (1d0 - u0) * f 
     if (dabs(Wt) <= 1.d-3) then
       return
     endif
@@ -86,7 +97,8 @@ logical function testTeethBuilding(minF, N)
        return
     end if
     n0 += 1
-    if(N_det_generators - n0 < minF * N) then
+!    if(N_det_generators - n0 < minF * N) then
+    if(n0 > minFN) then
       return
     end if
   end do
diff --git a/src/cipsi/run_pt2_slave.irp.f b/src/cipsi/run_pt2_slave.irp.f
index b4df2def..be3184c9 100644
--- a/src/cipsi/run_pt2_slave.irp.f
+++ b/src/cipsi/run_pt2_slave.irp.f
@@ -54,7 +54,6 @@ subroutine run_pt2_slave(thread,iproc,energy)
 
   sending = .False.
   done = .False.
-  n_tasks = 1
   do while (.not.done)
 
     n_tasks = max(1,n_tasks)
@@ -108,7 +107,8 @@ subroutine run_pt2_slave(thread,iproc,energy)
     b%cur=0
 
     ! Try to adjust n_tasks around nproc/2 seconds per job
-    n_tasks = min(2*n_tasks,int( dble(n_tasks * nproc/2) / (time1 - time0 + 1.d0)))
+!    n_tasks = min(2*n_tasks,int( dble(n_tasks * nproc/2) / (time1 - time0 + 1.d0)))
+    n_tasks = 1
   end do
   call push_pt2_results_async_recv(zmq_socket_push,b%mini,sending)
 
diff --git a/src/davidson/diagonalize_ci.irp.f b/src/davidson/diagonalize_ci.irp.f
index ab551084..d4b3b48c 100644
--- a/src/davidson/diagonalize_ci.irp.f
+++ b/src/davidson/diagonalize_ci.irp.f
@@ -202,6 +202,8 @@ subroutine diagonalize_CI
       psi_coef(i,j) = CI_eigenvectors(i,j)
     enddo
   enddo
+  psi_energy(1:N_states) = CI_electronic_energy(1:N_states)
+  psi_s2(1:N_states) = CI_s2(1:N_states)
 
-  SOFT_TOUCH psi_coef CI_electronic_energy CI_energy CI_eigenvectors CI_s2
+  SOFT_TOUCH psi_coef CI_electronic_energy CI_energy CI_eigenvectors CI_s2 psi_energy psi_s2
 end