ompi_kronlininv.f08

Go to the documentation of this file.

!------------------------------------------------------------------------
!
!    Copyright 2018, Andrea Zunino 
!
!    This file is part of KronLinInv.
!
!    KronLinInv is free software: you can redistribute it and/or modify
!    it under the terms of the GNU General Public License as published by
!    the Free Software Foundation, either version 3 of the License, or
!    (at your option) any later version.
!
!    KronLinInv is distributed in the hope that it will be useful,
!    but WITHOUT ANY WARRANTY; without even the implied warranty of
!    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!    GNU General Public License for more details.
!
!    You should have received a copy of the GNU General Public License
!    along with KronLinInv.  If not, see <http://www.gnu.org/licenses/>.
!
!------------------------------------------------------------------------



  
!-------------------------------------------------------------
!
!
!-------------------------------------------------------------
module realprec

  ! http://fortranwiki.org/fortran/show/Real+precision
  integer,parameter  :: pdigits=15  !15 !! used also by MPI
  integer,parameter  :: rexprange=307  !307 !! used also by MPI
    
  ! RESULT = SELECTED_REAL_KIND([P, R, RADIX]) 
  integer,parameter  :: dp = selected_real_kind(pdigits, rexprange)

  ! integer, parameter :: sp = selected_real_kind(6, 37)
  ! integer, parameter :: dp = selected_real_kind(15, 307)
  ! integer, parameter :: qp = selected_real_kind(33, 4931)

  public  :: pdigits,rexprange
  public  :: dp

end module realprec


!################################################################
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!################################################################

!-------------------------------------------------------------------------
!
!
!-------------------------------------------------------------------------
module parautil

  use realprec
  use mpi  
  !!use mpi_f08
  implicit none
  !!include 'mpif.h'
    
  
  integer,parameter :: masterrank=0
  integer,protected :: myrank,numcpus,dp_real_type
  
  public :: masterrank,myrank,numcpus

contains

  !!====================================================
  
  subroutine para_init() 

    use, intrinsic :: iso_c_binding
     
    integer :: mpierr

    !! start MPI section
    call mpi_init(mpierr)
    if (mpierr /= 0) then
       write(*,*) "call MPI_INIT(mpierr): mpierr /= 0"
       stop
    end if

    call mpi_comm_rank(mpi_comm_world, myrank, mpierr)
    call mpi_comm_size(mpi_comm_world, numcpus, mpierr)
    !! create mpi real and complex types
    call mpi_type_create_f90_real(pdigits,rexprange,dp_real_type,mpierr)
    !!call MPI_Type_create_f90_complex(pdigits,rexprange,complextype,error )

    if (myrank==masterrank) print*,"###>>>> Number of tasks:",numcpus

    
  end subroutine para_init

  !!====================================================

  subroutine para_finish() 
    ! finalize mpi
    integer :: ierror
    call mpi_finalize(ierror)
  end subroutine para_finish

  !!=====================================================

  subroutine para_barrier()
    integer :: mpierr
    call mpi_barrier(mpi_comm_world,mpierr)
  end subroutine para_barrier

  !!=====================================================
  
  ! subroutine spreadwork1(nit,nunits,scheduling,looping)
  !   integer :: nit,nunits
  !   integer,allocatable :: scheduling(:),looping(:,:)
  !   integer :: i,nitcpu,resto
  !   if (nit<=nunits) then
  !      write(*,*) myrank,": spredwork(): nit<=nunits"
  !      write(*,*) nit,nunits
  !      stop
  !   end if
  !   if (allocated(scheduling) ) deallocate(scheduling)
  !   if (allocated(looping)) deallocate(looping)
  !   allocate(scheduling(nunits),looping(nunits,2))

  !   !! compute distribution of workload for Nb
  !   nitcpu = nit/nunits
  !   scheduling(:) = nitcpu
  !   resto = mod(nit,nunits)

  !   !! spread the remaining workload
  !   scheduling(1:resto) =  scheduling(1:resto) + 1
  !   looping(1,:) = [1,scheduling(1)] 
  !   do i=2,nunits
  !      looping(i,1) = sum(scheduling(1:i-1)) + 1
  !      looping(i,2) = sum(scheduling(1:i)) 
  !   end do

  !   !!print*,myrank,": loop 1",looping(:,1),&
  !   !! "| loop 2",looping(:,2)," | sched ",scheduling

  ! end subroutine spreadwork1

  !!=====================================================
  
  subroutine spreadwork(nit,nunits,scheduling,looping,startpoint)
    integer,intent(in) :: nit,nunits,startpoint
    integer,allocatable,intent(out) :: scheduling(:),looping(:,:)
    integer :: i,nitcpu,resto
    if (nit<=nunits) then
       write(*,*) myrank,": spredwork(): nit<=nunits"
       write(*,*) nit,nunits
       stop
    end if
    if (allocated(scheduling) ) deallocate(scheduling)
    if (allocated(looping)) deallocate(looping)
    allocate(scheduling(nunits),looping(nunits,2))

    !! compute distribution of workload for Nb
    nitcpu = nit/nunits
    scheduling(:) = nitcpu
    resto = mod(nit,nunits)

    !! spread the remaining workload
    scheduling(1:resto) = scheduling(1:resto) + 1
    looping(1,:) = [startpoint,startpoint+scheduling(1)-1] 
    do i=2,nunits
       looping(i,1) = sum(scheduling(1:i-1)) + startpoint
       looping(i,2) = sum(scheduling(1:i)) + startpoint - 1
    end do

    !!print*,myrank,": loop 1",looping(:,1),&
    !! "| loop 2",looping(:,2)," | sched ",scheduling

  end subroutine spreadwork

  !!=====================================================

  subroutine para_srint(source,dest,tag,sint,rint)
    integer,intent(in) :: source,dest,tag
    integer,intent(in) :: sint
    integer,intent(inout) :: rint
    integer :: nelem,mpierr
    integer :: mpistat(mpi_status_size)
    
    if ((myrank==source) .and. (myrank==dest))  then
       rint = sint
    else
       if (myrank==source) then
          call mpi_send(sint, nelem, mpi_integer, dest, tag, &
               mpi_comm_world, mpierr)
       end if
       if (myrank==dest) then
          call mpi_recv(rint, nelem, mpi_integer, source, tag,&
               mpi_comm_world, mpistat, mpierr)
       end if
    end if
    
  end subroutine para_srint
  
  !!=====================================================

  subroutine para_srarr1ddp(source,dest,tag,sendrow,recvrow)
   
    integer,intent(in) :: source,dest,tag
    real(dp),intent(in) :: sendrow(:)
    real(dp),intent(inout) :: recvrow(:)

    integer :: mpierr,nelem
    integer :: mpistat(mpi_status_size)

    nelem=size(sendrow)  ! with no dim gives the total array size

    ! to avoid deadlocks
    if ((myrank==source) .and. (myrank==dest))  then
       !print*,"same send recv in ",myrank
       recvrow = sendrow
    else
       if (myrank==source) then
          !print*,"sending from ",myrank," to ",dest," tag=",tag
          call mpi_send(sendrow, nelem, dp_real_type, dest, tag, &
               mpi_comm_world, mpierr)
       end if
       if (myrank==dest) then
          !print*,"receiving in ",myrank," from ",source," tag=",tag
          call mpi_recv(recvrow, nelem, dp_real_type, source, tag,&
               mpi_comm_world, mpistat, mpierr)
       end if
    end if
    
  end subroutine para_srarr1ddp

  !!=====================================================
  
  subroutine para_bc2ddp(arr,source)
    real(dp) :: arr(:,:)
    integer,intent(in) :: source
    integer :: mpierr
    call mpi_bcast(arr,size(arr), dp_real_type, source, mpi_comm_world, mpierr)    
  end subroutine para_bc2ddp

  !!=====================================================

  subroutine para_bc1ddp(arr,source)
    real(dp) :: arr(:)
    integer,intent(in) :: source
    integer :: mpierr
    call mpi_bcast(arr,size(arr), dp_real_type, source, mpi_comm_world, mpierr)    
  end subroutine para_bc1ddp

 !!=====================================================

  subroutine para_allgathv1ddp(  scheduling, displs, sendbuf, recvbuf )
    real(dp),intent(out) :: recvbuf(:)
    integer,intent(in) :: scheduling(:),displs(:)
    real(dp),intent(in) :: sendbuf(:)
    integer :: mpierr
    !! SENDBUF, SENDCOUNT, SENDTYPE, RECVBUF, RECVCOUNTS,
    !!     DISPLS, RECVTYPE,  COMM, IERROR
    !print*,myrank," : size(sendbuf)",size(sendbuf),size(recvbuf)," loop ",displs," | sched ",scheduling
    ! if (displs(1)/=0) then
    !    write(*,*) myrank,": para_allgathv1ddp(): displs(1)/=0  potential problem (should be 0)"
    !    stop
    ! end if

    call mpi_allgatherv(sendbuf,size(sendbuf), dp_real_type, recvbuf, scheduling, &
         displs, dp_real_type, mpi_comm_world, mpierr)
    !print*,myrank," mpierr",mpierr
  end subroutine para_allgathv1ddp

  !!=====================================================

  subroutine para_gathv1ddp( recvcounts, displs, receiver, sendbuf, recvbuf )
    real(dp),intent(out) :: recvbuf(:)
    integer,intent(in) :: recvcounts(:),receiver,displs(:)
    real(dp),intent(in) :: sendbuf(:)
    integer :: mpierr

    ! print*,myrank," : size(sendbuf)",size(sendbuf)," size(recvbuf)",size(recvbuf)," displs ",displs
    
    ! if (displs(1)/=0) then
    !    write(*,*) myrank,": para_allgathv1ddp(): displs(1)/=0  potential problem (should be 0)"
    !    stop
    ! end if

    !! SENDBUF, SENDCOUNT, SENDTYPE, RECVBUF, RECVCOUNTS,
    !!     DISPLS, RECVTYPE, ROOT, COMM, IERROR

    call mpi_gatherv(sendbuf,size(sendbuf), dp_real_type, recvbuf, recvcounts, &
         displs, dp_real_type, receiver, mpi_comm_world, mpierr)
    
  end subroutine para_gathv1ddp
  
  !!=====================================================
  
  subroutine timeinfo(totit,curit,startt,loopinfo)
    use, intrinsic :: iso_fortran_env, only: output_unit
    character(len=30) :: loopinfo
    integer :: totit,curit
    real(dp) :: frac,eta,startt
    
    frac = real(curit,dp)/(real(totit,dp))
    eta = ( (mpi_wtime()-startt) / real(curit,dp) ) * &
         (real(totit-curit,dp))
    !! char(27) -> escape
    !! ESC[#A -> moves cursor up # lines (ANSI)
    !! achar(13) -> carriage return (ASCII)
    write(output_unit,fmt='(a5,a34,f7.3,a7,f12.2,1x,a3)') &
         char(27)//'[1A'//achar(13),&
         adjustl(loopinfo)//':  ',&
         frac*100_dp,"%,  ETA:",eta/60.0,"min"
    flush(output_unit) !! to make sure it prints immediately 
  end subroutine timeinfo

  !!=====================================================

end module parautil


!################################################################
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!################################################################

!-------------------------------------------------------------------------
!
!
!-------------------------------------------------------------------------
module ompi_kronlininv
  
  use realprec
  use parautil
  use, intrinsic :: iso_fortran_env, only: output_unit

  implicit none
  
  private :: solvels,symgeneigvv,symsolvels
  public :: calcfactors,blockpostcov,posteriormean,bandpostcov
  
contains

  !!!=====================================================
  !--------------------------------------------------
  ! DESCRIPTION: 
  !  REVISION HISTORY:
  !
  !--------------------------------------------------
  subroutine calcfactors(G1,G2,G3,Cm1,Cm2,Cm3,Cd1,Cd2,Cd3,&
       U1,U2,U3,diaginvlambda,iUCm1,iUCm2,iUCm3,&
       iUCmGtiCd1,iUCmGtiCd2,iUCmGtiCd3 ) 
 
    real(dp),intent(in) :: G1(:,:),G2(:,:),G3(:,:),&
         Cm1(:,:),Cm2(:,:),Cm3(:,:),&
         Cd1(:,:),Cd2(:,:),Cd3(:,:)
    
    real(dp),intent(out) :: U1(:,:),U2(:,:),U3(:,:),&
         iUCm1(:,:),iUCm2(:,:),iUCm3(:,:),&
         iUCmGtiCd1(:,:),iUCmGtiCd2(:,:),iUCmGtiCd3(:,:),&
         diaginvlambda(:)
    
    real(dp),allocatable :: lambda1(:),lambda2(:),lambda3(:)
    real(dp),allocatable :: iCdG1(:,:),iCdG2(:,:),iCdG3(:,:),&
         GtiCdG1(:,:),GtiCdG2(:,:),GtiCdG3(:,:),&
         ide1(:,:),ide2(:,:),ide3(:,:),&
         iCd1(:,:),iCd2(:,:),iCd3(:,:)
    integer :: i,j,k,p,nm,nm1,nm2,nm3,nd1,nd2,nd3
    character :: uplo
   
    !! Lapack routines
    external dsygv

    nm1 = size(cm1,1)
    nm2 = size(cm2,1)
    nm3 = size(cm3,1)
    nd1 = size(cd1,1)
    nd2 = size(cd2,1)
    nd3 = size(cd3,1)
    nm = nm1*nm2*nm3

    !! compute preliminary stuff
    print*,'calcfactors(): compute preliminary stuff'
    allocate(icdg1(nd1,nm1),icdg2(nd2,nm2),icdg3(nd3,nm3))
    call symsolvels( cd1, g1, icdg1 )
    call symsolvels( cd2, g2, icdg2 )
    call symsolvels( cd3, g3, icdg3 )

    allocate(gticdg1(nm1,nm1),gticdg2(nm2,nm2),gticdg3(nm3,nm3))
    gticdg1  = matmul( transpose(g1), icdg1 )
    gticdg2  = matmul( transpose(g2), icdg2 ) 
    gticdg3  = matmul( transpose(g3), icdg3 )

    !!--------------------------
    !!          fa
    !!--------------------------
    !! compute eigendecomposition
    print*,'calcfactors(): compute fa'
    allocate(lambda1(nm1),lambda2(nm2),lambda3(nm3))
    uplo = 'L'
    call symgeneigvv(gticdg1,uplo,cm1,lambda1,u1)
    call symgeneigvv(gticdg2,uplo,cm2,lambda2,u2)
    call symgeneigvv(gticdg3,uplo,cm3,lambda3,u3)

    !!--------------------------
    !!          fb
    !!--------------------------
    !! calculates diagonal central factor
    print*,'calcfactors(): compute fb'
    p=1
    do i=1,nm1
       do j=1,nm2
          do k=1,nm3
             diaginvlambda(p) = 1.0_dp/(1.0_dp+lambda1(i)*lambda2(j)*lambda3(k))
             p = p+1
          end do 
       end do
    end do

    !!--------------------------
    !!          fc
    !!--------------------------
    !! computes U1^-1 Cmx, U2^-1 Cmy, U3^-1 Cmz
    print*,'calcfactors(): compute fc'
    call solvels( u1, cm1, iucm1 )
    call solvels( u2, cm2, iucm2 )
    call solvels( u3, cm3, iucm3 )

    
    !!--------------------------
    !!          fd
    !!--------------------------
    print*,'calcfactors(): compute fd'
    allocate(ide1(nd1,nd1),ide2(nd2,nd2),ide3(nd3,nd3))
    ide1 = 0.0_dp  ! Initialize the array.
    ide2 = 0.0_dp
    ide3 = 0.0_dp
    forall(i = 1:nd1) ide1(i,i) = 1.0_dp  ! Set the diagonal.
    forall(i = 1:nd2) ide2(i,i) = 1.0_dp 
    forall(i = 1:nd3) ide3(i,i) = 1.0_dp
    allocate(icd1(nd1,nd1),icd2(nd2,nd2),icd3(nd3,nd3))
    call symsolvels( cd1, ide1, icd1 )
    call symsolvels( cd2, ide2, icd2 )
    call symsolvels( cd3, ide3, icd3 )

    iucmgticd1 = matmul( iucm1, matmul( transpose(g1), icd1 ))     
    iucmgticd2 = matmul( iucm2, matmul( transpose(g2), icd2 )) 
    iucmgticd3 = matmul( iucm3, matmul( transpose(g3), icd3 ))

    print*,"calcfactors(): end "
  end subroutine calcfactors

  !!!==================================================
  !-------------------------------------------------
  !
  !
  !
  ! REVISION HISTORY:
  !
  !-------------------------------------------------  
  subroutine blockpostcov(U1,U2,U3, diaginvlambda, &
       iUCm1,iUCm2,iUCm3, astart,aend,bstart,bend, postC) 
    !!
    !! Calculates a block of the posterior covariance
    !! 
    real(dp),intent(in) :: U1(:,:),U2(:,:),U3(:,:),iUCm1(:,:), &
         iUCm2(:,:),iUCm3(:,:)
    !! diaginvlambda = (I + lam1 x lam2 x lam3 )^-1
    real(dp),intent(in) :: diaginvlambda(:) !! diagonal/vector central factor
    real(dp),intent(out) :: postC(:,:)
    integer,intent(in) :: astart,aend,bstart,bend
    
    integer :: a,b,irow,i,j,sender
    real(dp),allocatable :: row2(:),col1(:)
    integer :: Nr12,Nc1
    integer :: Ni,Nj,Nk,Nl,Nm,Nn,Na,Nb,ntota,ntotb   

    integer,allocatable :: av(:)!,bv(:)
    integer,allocatable :: iv(:),jv(:),kv(:) !,lv(:),mv(:),nv(:) 
    integer :: p,everynit,impicount,mytotit,rirow
    real(dp) :: startt,firststartt,endt
    integer,allocatable :: scheduling(:),looping(:,:)
    real(dp),allocatable :: postcmpi(:,:),recvrow(:),sendrow(:)
    character(len=30) :: loopinfo

    if (myrank==masterrank) then
       write(output_unit,*) "posteriormean(): calculating posterior mean... [using OpenMPI]"
    end if
    
    ni = size(u1,1)
    nl = size(u1,2)
    nj = size(u2,1)
    nm = size(u2,2)
    nk = size(u3,1)
    nn = size(u3,2)
    na = ni*nj*nk
    nb = nl*nm*nn
    
    !! check shape of output array
    if (na /= nb) then
       write(*,*) '(Na /= Nb)', na,nb
       stop
    end if
    ntota = aend-astart+1
    ntotb = bend-bstart+1
    if ( (size(postc,1)/=ntota) .or. (size(postc,2)/=ntotb)) then
       write(*,*) "Wrong size of the posterior covariance block array."
       stop
    end if
    !! check limits of requested block
    if ( (astart<1) .or. (aend>na) .or. (astart>aend) .or. &
         (bstart<1) .or. (bend>na) .or. (bstart>bend) ) then
       write(*,*) "Wrong size of the requested block array."
       stop
    end if
        
    !! vectorize row and col calculations for Kron prod AxBxC
    allocate(av(na),iv(na),jv(na),kv(na))
    !allocate(bv(Nb),lv(Nb),mv(Nb),nv(Nb))
    forall(p = 1:na) av(p) = p
    !forall(p = 1:Nb) bv(p) = p

    !! vectors containing all possible indices for 
    !!    row calculations of Kron prod AxBxC
    iv =  (av-1)/(nk*nj)+1 
    jv =  (av-1-(iv-1)*nk*nj)/nk + 1 
    kv =  av-(jv-1)*nk-(iv-1)*nk*nj 
     
    !! allocate stuff
    nr12 = size(u1,2)* size(u2,2)* size(u3,2)
    nc1  = size(iucm1,1)* size(iucm2,1)* size(iucm3,1)
    allocate(row2(nr12),col1(nc1))

    !!------------------------
    if (ntota<1000) then
       everynit = 20
    else 
       everynit = ntota/1000
    end if
        
    !!---------------------------------
    firststartt = mpi_wtime()
    loopinfo = "blockpostcov()  "
    if (myrank==masterrank) write(output_unit,*)
    !! calculate scheduling for a
    call spreadwork(ntota,numcpus,scheduling,looping,1)
    allocate(postcmpi(scheduling(myrank+1),ntotb),recvrow(ntotb),sendrow(ntotb))
    mytotit = scheduling(myrank+1)

    !!print*,myrank,":",looping(myrank+1,1),"|",looping(myrank+1,2),"|",scheduling
    
    !!-----------------------
    startt = mpi_wtime()
    impicount = 0
    do irow=looping(myrank+1,1),looping(myrank+1,2)
       !! set a
       a=irow+astart-1 
       !! info
       if ( (myrank==masterrank) .and. (mod(irow,everynit)==0) ) then
          call timeinfo(mytotit,irow,startt,loopinfo)
       end if

       !! row first two factors
       !! a row x diag matrix 
       !!row2 =  U1(iv(a),lv) * U2(jv(a),mv) * U3(kv(a),nv) * diaginvlambda
       row2 =  u1(iv(a),iv) * u2(jv(a),jv) * u3(kv(a),kv) * diaginvlambda

       j=0
       do b=bstart,bend
          j=j+1
          !! calculate one row of first TWO factors
          !!call columnAxBxC(Ni,Nj,Nk,Nl,Nm,Nn, iUCm1,iUCm2,iUCm3,b,col1)
          col1 = iucm1(iv,iv(b)) * iucm2(jv,jv(b)) * iucm3(kv,kv(b))

          !! calculate one element of the posterior covariance
          i = irow-looping(myrank+1,1)+1
          postcmpi(i,j) = sum(row2*col1)

       end do     
     
    end do

    !! collect all results in masterrank
    call para_barrier()
    do sender=0,numcpus-1   
       if ( (myrank==sender) .or. (myrank==masterrank) ) then
             do i=looping(sender+1,1),looping(sender+1,2)
             !!----------------------
             !! send/rec each row
             !print*,"==>",myrank,": sender=",sender,"i=",i
             if (myrank==sender) sendrow = postcmpi(i,:)
             call para_srarr1ddp(sender,masterrank,i,sendrow,recvrow)
             !call para_srint(myrank,masterrank,impicount,irow,rirow)
             if (myrank==masterrank)  postc(i,:) = recvrow
          end do
       end if
    end do
    
    call para_barrier()
    if (myrank==masterrank) then
       write(output_unit,*)
       endt = MPI_Wtime()
       print*,"blockpostcov():",endt-firststartt," MPI wall clock time"
    end if    

  end subroutine blockpostcov

  !!!==================================================
  !-------------------------------------------------
  !
  !
  ! REVISION HISTORY:
  !
  !
  !-------------------------------------------------  
  subroutine bandpostcov(U1,U2,U3, diaginvlambda, &
       iUCm1,iUCm2,iUCm3, lowdiag, updiag, bandpostC) 
    !!
    !! Calculate a band of the posterior covariance
    !! 
    real(dp),intent(in) :: U1(:,:),U2(:,:),U3(:,:),iUCm1(:,:), &
         iUCm2(:,:),iUCm3(:,:)
    !! diaginvlambda = (I + lam1 x lam2 x lam3 )^-1
    real(dp),intent(in) :: diaginvlambda(:) !! diagonal/vector central factor
    real(dp),intent(inout) :: bandpostC(:,:)
    integer,intent(in) :: lowdiag, updiag
    
    integer :: a,b
    real(dp),allocatable :: row2(:),col1(:),recvrow(:),bandpostCmpi(:)
    integer :: Nr12,Nc1
    integer :: Ni,Nj,Nk,Nl,Nm,Nn,Na,Nb    

    real(dp) :: firststartt,startt
    integer,allocatable :: av(:)!,bv(:)
    integer,allocatable :: iv(:),jv(:),kv(:) !,lv(:),mv(:),nv(:) 
    integer :: p,aband,aend,bband,astart,d,bband1,bband2,ntota,everynit,mytotit
    integer,allocatable :: scheduling(:),looping(:,:)
    character(len=30) :: loopinfo
    character(len=12) :: dnum
    
    ni = size(u1,1)
    nl = size(u1,2)
    nj = size(u2,1)
    nm = size(u2,2)
    nk = size(u3,1)
    nn = size(u3,2)
    na = ni*nj*nk 
    nb = nl*nm*nn    

    if (na /= nb) then
       write(*,*) '(Na /= Nb)', na,nb
       stop
    end if
    if ( (updiag>=na) .or. (lowdiag>=na) .or. (lowdiag<0) .or. (updiag<0) ) then
       write(*,*) .or."(updiag<Na)  (lowdiag<Na)"
       write(*,*) "updiag",updiag,"Na",na,"lowdiag",lowdiag,"Na",na
       stop
    end if

    !! vectorize row and col calculations for Kron prod AxBxC
    allocate(av(na),iv(na),jv(na),kv(na))
    !allocate(bv(Nb),lv(Nb),mv(Nb),nv(Nb))
    forall(p = 1:na) av(p) = p
    !forall(p = 1:Nb) bv(p) = p

    !! vectors containing all possible indices for 
    !!    row calculations of Kron prod AxBxC
    iv =  (av-1)/(nk*nj)+1 
    jv =  (av-1-(iv-1)*nk*nj)/nk+1 
    kv =  av-(jv-1)*nk-(iv-1)*nk*nj 
    !! vectors containing all possible indices for
    !!    column calculations of Kron prod AxBxC
    ! lv =  (bv-1)/(Nn*Nm) + 1 
    ! mv =  (bv-1-(lv-1)*Nn*Nm)/Nn + 1 
    ! nv =  bv-(mv-1)*Nn-(lv-1)*Nn*Nm
    
    !! allocate stuff
    nr12 = size(u1,2)* size(u2,2)* size(u3,2)
    nc1  = size(iucm1,1)* size(iucm2,1)* size(iucm3,1)
    allocate(row2(nr12),col1(nc1))

    ! Lapack: http://www.netlib.org/lapack/lug/node124.html
    ! aij is stored in AB(ku+1+i-j,j) for max(1,j-ku) <= i <= \min(m,j+kl).
    !------------
    ! Diagonals of a matrix
    ! i + d = j
    ! main diag d = 0
    ! upper d > 0
    ! lower d < 0
    ! diagonals of a matrix and indices of related band matrix
    ! ONLY for square matrix

    !!---------------------------------
    firststartt = mpi_wtime()
    ! !! calculate scheduling for a
    ! call spreadwork(ntota,numcpus,scheduling,looping,1)
    ! allocate(postcmpi(scheduling(myrank+1),ntotb),recvrow(ntotb),sendrow(ntotb))
    ! mytotit = looping(myrank+1,2)-looping(myrank+1,1)+1

    allocate(bandpostcmpi(size(bandpostc,2)), recvrow(size(bandpostc,2)))
    ! initialize postC
    
    bandpostc = 0.0_dp       

    do d=-lowdiag,updiag
       call para_barrier()
       
       !!--------------------------
       if (d<0) then
          astart = abs(d)+1
       else
          astart = 1
       end if
       if (d>0) then
          aend = na-abs(d)
       else
          aend = na
       end if
       !print*,'diagonal',d

       !! calculate scheduling for a
       ntota = aend-astart+1
       call spreadwork(ntota,numcpus,scheduling,looping,astart)
       if (ntota<1000) then
          everynit = ntota/100
       else
          everynit = ntota/1000
       end if
       
       write(dnum,"(i9)") d
       loopinfo = "bandpostcov()  diag "//trim(dnum)//" "
       startt = mpi_wtime()
       if (myrank==masterrank) write(output_unit,*)
       !! indices of normal matrix
       bandpostcmpi(:) = 0.0_dp
       aband = updiag+1-d !! aband is constant within this loop
       do a=looping(myrank+1,1),looping(myrank+1,2) !!astart,aend
          !! info
          if ( (myrank==masterrank) .and. (mod(a,everynit)==0) ) then
             call timeinfo(scheduling(myrank+1),a-astart+1,startt,loopinfo)
          end if

          !if ( (myrank==masterrank) .and. (mod(a,250)==0) ) print*,"d",d,"a",a
          !!--------------------------------------------------------
          b = a+d
          !! indices of the band matrix
          !aband = updiag+1+a-b
          !bband = b
          !!--------------------------------------------------------
                         
          !!----------------------------------------------------------------
          ! row first two factors
          row2 = diaginvlambda * u1(iv(a),iv) * u2(jv(a),jv) * u3(kv(a),kv)

          !! calculate one row of first TWO factors
          !!call columnAxBxC(Ni,Nj,Nk,Nl,Nm,Nn, iUCm1,iUCm2,iUCm3,b,col1)
          col1 = iucm1(iv,iv(b)) * iucm2(jv,jv(b)) * iucm3(kv,kv(b))

          !! calculate one element of the posterior covariance
          !! store it in the band storage format
          !!bandpostC(aband,bband) = sum(row2*col1)
          bandpostcmpi(b) = sum(row2*col1)
          !print*, a,b,aband,bband,bandpostC(aband,bband)
          !!----------------------------------------------------------------
          
       end do ! a=astart,aend
       
       !! send/recv rows                                                      
       !! collect all results in masterrank
       

       bband1 = looping(myrank+1,1)+d !!b = a+d
       bband2 = looping(myrank+1,2)+d
       
       call para_gathv1ddp(scheduling, looping(:,1)-1, masterrank, &
            bandpostcmpi(bband1:bband2), recvrow )
       if (myrank==masterrank) bandpostc(aband,:) = recvrow
      
       
    end do

  end subroutine bandpostcov

  !!!===============================================================
  !-------------------------------------------------
  !
  !
  !
  ! REVISION HISTORY:
  ! 
  !-------------------------------------------------
  subroutine posteriormean(U1,U2,U3, diaginvlambda, Z1,Z2,Z3,&
       G1,G2,G3, mprior, dobs, postm)

    !!
    use parautil
    
    !!
    !! Calculate the posterior mean model
    !! 
    real(dp),intent(in)  :: U1(:,:),U2(:,:),U3(:,:), G1(:,:),G2(:,:),G3(:,:)
    real(dp),intent(in)  :: Z1(:,:),Z2(:,:),Z3(:,:)
    real(dp),intent(in)  :: mprior(:),dobs(:),diaginvlambda(:)
    real(dp),intent(out) :: postm(:)

    real(dp),allocatable :: ddiff(:),ddiffmpi(:),postmmpi(:)
    integer :: a,b,Na,Nb
    integer :: Ni,Nj,Nk,Nl,Nm,Nn

    integer,allocatable :: av(:),bv(:)
    integer,allocatable :: iv(:),jv(:),kv(:),lv(:),mv(:),nv(:)
        !! ivo(:),jvo(:),kvo(:)
    integer :: p,everynit

    real(dp),allocatable :: Zh(:),elUDZh(:),Zhmpi(:),elUDZhmpi(:)

    integer :: ia,ib
    real(dp) :: startt,firststartt,endt
    character(len=30) :: loopinfo
    
    integer ::curit,totit
    integer,allocatable :: scheduling(:),looping(:,:)

    !! sizes
    ni = size(z1,1)
    nl = size(z1,2)
    nj = size(z2,1)
    nm = size(z2,2)
    nk = size(z3,1)
    nn = size(z3,2)

    !! sizes
    ! Nr12 = size(U1,2)*size(U2,2)*size(U3,2)
    ! Nc1  = size(Z1,1)*size(Z2,1)*size(Z3,1)
    na = size(mprior,1) 
    nb = size(dobs,1)
    
    !! allocate stuff
    allocate(ddiff(nb))
    allocate(zh(na),eludzh(na))
 
    !! vectorize row and col calculations for Kron prod AxBxC
    !!allocate(ivo(Nb),jvo(Nb),kvo(Nb))
    allocate(av(na),iv(na),jv(na),kv(na))
    allocate(bv(nb),lv(nb),mv(nb),nv(nb))
    forall(p = 1:na) av(p) = p
    forall(p = 1:nb) bv(p) = p

    !! vectors containing all possible indices for 
    !!    row calculations of Kron prod AxBxC
    iv =  (av-1)/(nk*nj)+1 
    jv =  (av-1-(iv-1)*nk*nj)/nk+1 
    kv =  av-(jv-1)*nk-(iv-1)*nk*nj 
    !! vectors containing all possible indices for
    !!    column calculations of Kron prod AxBxC
    lv =  (bv-1)/(nn*nm) + 1 
    mv =  (bv-1-(lv-1)*nn*nm)/nn + 1 
    nv =  bv-(mv-1)*nn-(lv-1)*nn*nm 
    !!  Gs have different shape than Us !!

    !!#######################
    if (na<1000)then
       everynit = na/50
    else if (na<100000) then
       everynit = na/100
    else
       everynit = na/1000
    end if

    !!---------------------------------
    firststartt = mpi_wtime()

    !!#######################
    !!#    dobs - dcalc     #
    !!#######################
    startt = mpi_wtime()    
    loopinfo = "posteriormean(): loop 1/3"
    if (myrank==masterrank) write(output_unit,*)
    !! calculate scheduling for Nb
    call spreadwork(nb,numcpus,scheduling,looping,1)
    allocate(ddiffmpi(scheduling(myrank+1)))
    !! difference obs-calc data
    totit = looping(myrank+1,2)-looping(myrank+1,1)
    do b=looping(myrank+1,1),looping(myrank+1,2) !! rank starts from 0
       curit = (b-looping(myrank+1,1)+1)
       if ( (myrank==masterrank) .and. (mod( curit,everynit)==0) ) then          
          call timeinfo(totit,curit,startt,loopinfo)
       end if
       
       !!--------------------------------------------------------------------------------
       ib = b - looping(myrank+1,1) + 1 
       ddiffmpi(ib) = dobs(b) - sum(mprior * g1(lv(b),iv) * g2(mv(b),jv) * g3(nv(b),kv))
       !!--------------------------------------------------------------------------------

    end do

    startt = mpi_wtime()    
    !! assemble full ddfiff and broadcast it
    !! displacements  must start from 0, not 1, so looping-1
    call para_allgathv1ddp( scheduling,looping(:,1)-1, ddiffmpi, ddiff )
    if (myrank==masterrank) then
       print*,"time for first allgatherv:",mpi_wtime()-startt,"s"
    end if

    
    !!#######################
    !!#   === U d Z h ===   #
    !!#######################
    startt = mpi_wtime()    
    !! calculate scheduling for Na
    loopinfo = "posteriormean(): loop 2/3"
    if (myrank==masterrank) write(output_unit,*)
    call spreadwork(na,numcpus,scheduling,looping,1)
    allocate(zhmpi(scheduling(myrank+1)))
    totit = looping(myrank+1,2)-looping(myrank+1,1)
    do a=looping(myrank+1,1),looping(myrank+1,2) !! rank starts from 0
       curit =(a-looping(myrank+1,1)+1)
       if ( ( myrank==masterrank ) .and. (mod( curit,everynit)==0) ) then       
          call timeinfo(totit,curit,startt,loopinfo)
       end if

       !!---------------------------------------------------------------------------
       ia = a - looping(myrank+1,1) + 1
       zhmpi(ia) = sum( z1(iv(a),lv) * z2(jv(a),mv) * z3(kv(a),nv) * ddiff )
       !!---------------------------------------------------------------------------
       
    end do

    startt = mpi_wtime()
    !! assemble full Zh and broadcast it
    call para_allgathv1ddp( scheduling,looping(:,1)-1, zhmpi, zh )
    if (myrank==masterrank) then
       print*,"time for second allgatherv:",mpi_wtime()-startt,"s"
    end if
    
    !!-----------------------------------------------
    !! Second loop
    !!### need to re-loop because full Zh is needed
    !!#######################
    !!#       post(a)       #
    !!#######################
    startt = mpi_wtime()
    !! calculate scheduling for Na
    loopinfo = "posteriormean(): loop 3/3"
    if (myrank==masterrank) write(output_unit,*)
    call spreadwork(na,numcpus,scheduling,looping,1)
    allocate(eludzhmpi(scheduling(myrank+1)),postmmpi(scheduling(myrank+1)))
    totit = looping(myrank+1,2)-looping(myrank+1,1)
    do a=looping(myrank+1,1),looping(myrank+1,2) !! rank starts from 0
       curit = (a-looping(myrank+1,1)+1)
       if ( ( myrank==masterrank ) .and. (mod(curit,everynit)==0) ) then
          call timeinfo(totit,curit,startt,loopinfo)
       end if

       !!--------------------------------------------------------------------------------------
       ia = a - looping(myrank+1,1) + 1
       !! 
       eludzhmpi(ia) = sum( u1(iv(a),iv) * u2(jv(a),jv) * u3(kv(a),kv) * diaginvlambda * zh )
       !! element of the posterior mean
       postmmpi(ia) = mprior(a) + eludzhmpi(ia)
       !!--------------------------------------------------------------------------------------
       
    end do

    startt = mpi_wtime()    
    !! assemble full postm only for masterrank
    call para_gathv1ddp( scheduling,looping(:,1)-1, masterrank, postmmpi, postm )
    if (myrank==masterrank) then
       print*,"time for third allgatherv:",mpi_wtime()-startt,"s"
    end if


    if (myrank==masterrank) then
       write(output_unit,*)
       write(output_unit,*)
       endt = MPI_Wtime()
       print*,"posteriormean():",endt-firststartt," MPI wall clock time"
    end if

    !print*,myrank,":",shape(ddiffmpi),shape(Zhmpi),shape(elUDZhmpi),shape(postmmpi),&
     !    shape(ddiff),shape(Zh),shape(elUDZh),shape(postm)

    return
  end subroutine posteriormean
  
  !!!==================================================
  !-------------------------------------------------
  !
  !
  ! REVISION HISTORY:
  ! 
  !-------------------------------------------------  
  subroutine symgeneigvv(A,uplo,Bpd,lambda,U)
    real(dp),intent(in) :: A(:,:),Bpd(:,:)
    character,intent(in) :: uplo
    real(dp),intent(out) :: lambda(:),U(:,:)
    integer :: n,lwork,itype,info
    character :: jobz
    real(dp),allocatable :: work(:),tmpB(:,:)
    ! DSYGV computes all the eigenvalues, and optionally, the eigenvectors
    ! of a real generalized symmetric-definite eigenproblem, of the form
    ! A*x=(lambda)*B*x,  A*Bx=(lambda)*x,  or B*A*x=(lambda)*x.
    ! Here A and B are assumed to be symmetric and B is also
    ! positive definite.
    itype = 3
    jobz = 'V'
    n=size(a,1)
    if ( (size(a,2)/=n) .or. (size(bpd,1)/=n) .or. (size(bpd,2)/=n) ) then
       write(*,*) "symgeneigvv(): Matrices have wrong sizes"
       stop 
    end if
    allocate(tmpb(n,n))
    u = a
    tmpb = bpd
    lwork=3*n-1
    ! lwork=-1 
    ! allocate(work(1))
    ! call DSYGV(itype,jobz,uplo,nm1,U1,nm1,choCm1,nm1,lambda1,work,lwork,info)
    ! deallocate(work)
    allocate(work(lwork))
    call dsygv(itype,jobz,uplo,n,u,n,tmpb,n,lambda,work,lwork,info)
    if (info /= 0) stop 'symgeneigvv(): Matrix eigendecomposition failed!'
  end subroutine symgeneigvv

  !!!==================================================
  !-------------------------------------------------
  !
  !
  !
  ! REVISION HISTORY:
  ! 
  !-------------------------------------------------
  subroutine solvels( A, B, sol )
    real(dp),intent(in) :: A(:,:),B(:,:)
    real(dp),intent(out) :: sol(:,:)
    integer,allocatable :: ipiv(:)
    integer :: n,nrhs,lda,ldb
    integer :: info
    real(dp),allocatable :: A2(:,:)
    ! External procedures defined in LAPACK
    external dgesv
    n = size(a,1)
    nrhs = size(b,2)
    allocate(a2(n,n))
    a2 = a
    lda = size(a,1)
    allocate(ipiv(n))
    ldb = size(b,1)
    sol = b
    call dgesv(n,nrhs,a2,lda,ipiv,sol,ldb,info)
    if (info/=0) then
       write(*,*) "linear system solver failed...'"
       print*,'info: ',info
       stop
    end if
  end subroutine solvels

  !!!============================================================
  !-------------------------------------------------
  !
  !
  ! REVISION HISTORY:
  ! 
  !-------------------------------------------------
  subroutine symsolvels( A, B, sol )
    real(dp),intent(in) :: A(:,:),B(:,:)
    real(dp),intent(out) :: sol(:,:)
    integer,parameter :: wmax=3000 ! max size allowed for work
    integer,allocatable :: ipiv(:)
    integer :: n,nrhs,lda,ldb,lwork
    character :: uplo
    real(dp),allocatable :: work(:),A2(:,:)
    integer :: info
    ! External procedures defined in LAPACK
    external dsysv
    uplo = 'L'
    n = size(a,1)
    nrhs = size(b,2)
    allocate(a2(n,n))
    a2 = a
    lda = size(a,1)
    allocate(ipiv(n))
    ldb = size(b,1)
    ! copy B to sol to avoid it to be overwritten
    sol = b
    ! apparently work must be allocated for the query to dsysv to work...
    allocate(work(1)) 
    lwork = -1
    call dsysv(uplo,n,nrhs,a2,lda,ipiv,sol,ldb,work,lwork,info)
    lwork = min(wmax,int(work(1)))
    !! now lwork is assigned, so actually calculate eigvec/val
    deallocate(work)
    allocate(work(lwork))
    call dsysv(uplo,n,nrhs,a2,lda,ipiv,sol,ldb,work,lwork,info)
    if (info/=0) then
       write(*,*) "linear system solver failed...'"
       print*,'info: ',info
       stop
    end if
  end subroutine symsolvels

  !!====================================================
 
end module ompi_kronlininv

!!====================================================
!!====================================================