Gather strategy seems to work

william-dawson · Apr 18, 2024 · 5053a6a · 5053a6a
1 parent d02dc11
commit 5053a6a
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 13 deletions.
diff --git a/Source/Fortran/FermiOperatorModule.F90 b/Source/Fortran/FermiOperatorModule.F90
@@ -7,17 +7,19 @@ MODULE FermiOperatorModule
   USE LoggingModule, ONLY : WriteElement, WriteHeader, &
        & EnterSubLog, ExitSubLog, WriteListElement, WriteComment
   USE PSMatrixAlgebraModule, ONLY : MatrixMultiply, SimilarityTransform, &
-       & IncrementMatrix, MatrixNorm, ScaleMatrix, DotMatrix, MatrixTrace
+       & IncrementMatrix, MatrixNorm, ScaleMatrix, DotMatrix, MatrixTrace, &
+       & MatrixDiagonalScale
   USE PSMatrixModule, ONLY : Matrix_ps, ConstructEmptyMatrix, &
        & FillMatrixFromTripletList, GetMatrixTripletList, &
-       & TransposeMatrix, ConjugateMatrix, DestructMatrix, &
-       & FillMatrixIdentity, PrintMatrixInformation, CopyMatrix, GetMatrixSize
+       & TransposeMatrix, ConjugateMatrix, DestructMatrix, FilterMatrix, &
+       & FillMatrixIdentity, PrintMatrixInformation, CopyMatrix, GetMatrixSize, printmatrix
   USE PMatrixMemoryPoolModule, ONLY : MatrixMemoryPool_p, &
        & DestructMatrixMemoryPool
   USE SolverParametersModule, ONLY : SolverParameters_t, &
        & PrintParameters, DestructSolverParameters, CopySolverParameters, &
        & ConstructSolverParameters
-  USE TripletListModule, ONLY : TripletList_r, DestructTripletList
+  USE TripletListModule, ONLY : TripletList_r, DestructTripletList, &
+       & AllGatherTripletList
   USE NTMPIModule
   IMPLICIT NONE
   PRIVATE
@@ -54,7 +56,7 @@ SUBROUTINE ComputeDenseFOE(H, ISQ, trace, K, inv_temp_in, &
     TYPE(Matrix_ps) :: WD
     TYPE(Matrix_ps) :: vecs, vecsT, vals, Temp
     TYPE(MatrixMemoryPool_p) :: pool
-    TYPE(TripletList_r) :: tlist
+    TYPE(TripletList_r) :: tlist, gathered_list
     REAL(NTREAL) :: chemical_potential, energy_value
     REAL(NTREAL), DIMENSION(:), ALLOCATABLE :: eigs, occ
     REAL(NTREAL) :: sval, sv, occ_temp
@@ -148,7 +150,7 @@ SUBROUTINE ComputeDenseFOE(H, ISQ, trace, K, inv_temp_in, &
        CALL ExitSubLog
     END IF
 
-    !! Map
+    !! Map - note that we store the square root of the occupation numbers
     energy_value = 0.0_NTREAL
     DO II = 1, tlist%CurrentSize
        IF (.NOT. do_smearing) THEN
@@ -159,7 +161,7 @@ SUBROUTINE ComputeDenseFOE(H, ISQ, trace, K, inv_temp_in, &
              occ_temp = trace - FLOOR(trace)
              energy_value = energy_value + &
                   & occ_temp * tlist%DATA(II)%point_value
-             tlist%DATA(II)%point_value = occ_temp
+             tlist%DATA(II)%point_value = SQRT(occ_temp)
           ELSE
              tlist%DATA(II)%point_value = 0.0_NTREAL
           ENDIF
@@ -168,21 +170,26 @@ SUBROUTINE ComputeDenseFOE(H, ISQ, trace, K, inv_temp_in, &
           ! occ_temp = 0.5_NTREAL * (1.0_NTREAL - ERF(inv_temp * sval))
           occ_temp = 1.0_NTREAL / (1.0_NTREAL + EXP(inv_temp * sval))
           energy_value = energy_value + occ_temp * tlist%DATA(II)%point_value
-          tlist%DATA(II)%point_value = occ_temp
+          IF (occ_temp .LT. 0) THEN  ! for safety
+             tlist%DATA(II)%point_value = 0
+          ELSE
+             tlist%DATA(II)%point_value = SQRT(occ_temp)
+          END IF
        END IF
     END DO
     CALL MPI_ALLREDUCE(MPI_IN_PLACE, energy_value, 1, MPINTREAL, MPI_SUM, &
          & H%process_grid%within_slice_comm, ierr)
+
+    !! Gather the triplet lists
+    CALL AllGatherTripletList(tlist, H%process_grid%column_comm, gathered_list)
 
-    !! Fill
-    CALL ConstructEmptyMatrix(vals, H)
-    CALL FillMatrixFromTripletList(vals, tlist, preduplicated_in = .TRUE.)
+    !! Scale the eigenvectors
+    CALL MatrixDiagonalScale(vecs, gathered_list)
 
     !! Multiply Back Together
-    CALL MatrixMultiply(vecs, vals, temp, threshold_in = params%threshold)
     CALL TransposeMatrix(vecs, vecsT)
     CALL ConjugateMatrix(vecsT)
-    CALL MatrixMultiply(temp, vecsT, WD, &
+    CALL MatrixMultiply(vecs, vecsT, WD, &
          & threshold_in = params%threshold)
 
     !! Compute the density matrix in the non-orthogonalized basis
@@ -208,6 +215,7 @@ SUBROUTINE ComputeDenseFOE(H, ISQ, trace, K, inv_temp_in, &
     CALL DestructMatrix(vals)
     CALL DestructMatrix(temp)
     CALL DestructTripletList(tlist)
+    CALL DestructTripletList(gathered_list)
     CALL DestructMatrixMemoryPool(pool)
     IF (ALLOCATED(occ)) THEN
        DEALLOCATE(occ)

diff --git a/Source/Fortran/TripletListModule.F90 b/Source/Fortran/TripletListModule.F90
@@ -39,6 +39,7 @@ MODULE TripletListModule
   PUBLIC :: SymmetrizeTripletList
   PUBLIC :: GetTripletListSize
   PUBLIC :: RedistributeTripletLists
+  PUBLIC :: AllGatherTripletList
   PUBLIC :: ShiftTripletList
   PUBLIC :: ConvertTripletListType
   PUBLIC :: MergeTripletLists
@@ -91,6 +92,10 @@ MODULE TripletListModule
      MODULE PROCEDURE RedistributeTripletLists_r
      MODULE PROCEDURE RedistributeTripletLists_c
   END INTERFACE RedistributeTripletLists
+  INTERFACE AllGatherTripletList
+     MODULE PROCEDURE AllGatherTripletList_r
+     MODULE PROCEDURE AllGatherTripletList_c
+  END INTERFACE AllGatherTripletList
   INTERFACE ShiftTripletList
      MODULE PROCEDURE ShiftTripletList_r
      MODULE PROCEDURE ShiftTripletList_c
@@ -365,6 +370,44 @@ SUBROUTINE RedistributeTripletLists_c(triplet_lists, comm, local_data_out)
 #undef MPIDATATYPE
 
   END SUBROUTINE RedistributeTripletLists_c
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+  !> Gather triplet lists from a set of processors.
+  SUBROUTINE AllGatherTripletList_r(triplet_in, comm, gathered_out)
+    !> Locally held triplet list
+    TYPE(TripletList_r), INTENT(IN) :: triplet_in
+    !> The mpi communicator to gather along.
+    INTEGER, INTENT(IN) :: comm
+    !> The resulting gathered triplet list.
+    TYPE(TripletList_r), INTENT(INOUT) :: gathered_out
+    !! Local data (type specific)
+    REAL(NTREAL), DIMENSION(:), ALLOCATABLE :: send_buffer_val
+    REAL(NTREAL), DIMENSION(:), ALLOCATABLE :: recv_buffer_val
+    TYPE(Triplet_r) :: temp_triplet
+
+#define MPIDATATYPE MPINTREAL
+#include "triplet_includes/GatherTripletList.f90"
+#undef MPIDATATYPE
+
+  END SUBROUTINE AllGatherTripletList_r
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+  !> Gather triplet lists from a set of processors.
+  SUBROUTINE AllGatherTripletList_c(triplet_in, comm, gathered_out)
+    !> Locally held triplet list
+    TYPE(TripletList_c), INTENT(IN) :: triplet_in
+    !> The mpi communicator to gather along.
+    INTEGER, INTENT(IN) :: comm
+    !> The resulting gathered triplet list.
+    TYPE(TripletList_c), INTENT(INOUT) :: gathered_out
+    !! Local data (type specific)
+    COMPLEX(NTCOMPLEX), DIMENSION(:), ALLOCATABLE :: send_buffer_val
+    COMPLEX(NTCOMPLEX), DIMENSION(:), ALLOCATABLE :: recv_buffer_val
+    TYPE(Triplet_c) :: temp_triplet
+
+#define MPIDATATYPE MPINTCOMPLEX
+#include "triplet_includes/GatherTripletList.f90"
+#undef MPIDATATYPE
+
+  END SUBROUTINE AllGatherTripletList_c
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
   !> Shift the rows and columns of a triplet list by set values.
   !> Frequently, we have a triplet list that comes from the global matrix which