! $Id: alfven_subbox.f90,v 1.3 2003/12/10 21:23:51 nilshau Exp $
!
!***********************************************************************
program alfven_subbox
  !
  ! The program reads in the data from ALL the different processors, 
  ! and stores it in a global array. If the runs are to large the use of
  ! shared memory computers are necessary. In IRIX machines OpenMP is
  ! used as default, on othere operating systems one can invoke 
  ! OpenMP by adding the line:
  ! FLAGS_OMP=-mp
  ! in Makefile.local
  ! You run the program with the script run_cal.sh.
  !
  !-----------------------------------------------------------------------
  !   22-july-03/nils: coded
  !
  use Cdata
  use Cparam
  use Sub
  use Deriv
  use General
  use Mpicomm
  !
  implicit none
  !
  !
  !  define parameters
  !
  character (len=130) ::  file,result_file,file_ext,input_file
  character (len=4) :: ch,csnap
  integer :: cpu,mom,exp1,exp2,lb_ll,counter
  integer, parameter :: nbox_x=8,nbox_y=8,nbox_z=8
  integer, parameter :: nbins=100
  real :: dummy,sf1
  real :: x1p,y1p,z1p
  real :: x1frac,x2frac,y1frac,y2frac,z1frac,z2frac
  real :: b2,u2,udotb
  real, dimension (mx,my,mz,mvar+maux) :: f
  real, dimension (mx,my,mz,3) :: bb_loc ,dir_loc
  real, dimension (:,:,:,:), allocatable :: bb_glob,dir_glob,vv_glob
  real, dimension (nx,3) :: bb,BnB,b_parallel,normal,binormal
  real, dimension (nx) :: xxx,xxy,xxz,yyx,yyy,yyz,zzx,zzy,zzz
  real, dimension (nx) :: bxx,bxy,bxz,byx,byy,byz,bzx,bzy,bzz,a1,a2
  real, dimension (3) :: B0,B_vec,dir0,B_ext,tmpB,V_vec,V0,bf,uf
  real, dimension (nbins,nbox_x,nbox_y,nbox_z) :: PDF,PDF_local
  real, dimension (3,nbox_x,nbox_y,nbox_z) :: meanB_box, meanU_box
  real, dimension (3,nbox_x,nbox_y,nbox_z) :: meanB_box_local, meanU_box_local
  integer :: snap,snap_start,snap_stop
  integer :: l1_glob,l2_glob,m1_glob,m2_glob,n1_glob,n2_glob
  integer :: xxp1,xxp2,yyp1,yyp2,zzp1,zzp2
  integer :: ll,mmm,nnn,i,dir_nr,dir_count,sizeofbox
  integer :: box_x,box_y,box_z,bin
  logical :: lext,lpri
  !
  !  Stuff for OpenMP
  !
  integer :: my_number=0,total_number=1,ipy_loc,ipz_loc
  integer :: ny_OMP,nz_OMP,ipy_OMP,ipz_OMP,yprocs,zprocs,m_OMP,n_OMP
  !
  ! Allocating allocatable arrays. The large global arrays must be
  ! defined with allocate and NOT in the normal way, in order to 
  ! make sure they go into the heap and not the stack. This is 
  ! important on shared memory computers where there normally is
  ! an upper limit on the stack but not on the heap.
  !
  allocate(bb_glob(nxgrid,nygrid,nzgrid,3))
  allocate(vv_glob(nxgrid,nygrid,nzgrid,3))
  !
  ! Setting some constants. (Should have been set by invoking the different
  ! initilize routines.)
  !
  iux=1 
  iuy=2
  iuz=3
  ilnrho=4
  iax=5
  iay=6
  iaz=7
  !
  ip=10
  !
  ! Reading alfven_subbox.in
  !
  call safe_character_assign(input_file,'alfven_subbox.in')
  if (ip<8) print*,'Reading '//trim(input_file)
  open(1,FILE=input_file,FORM='formatted')
  read(1,*) sizeofbox
  read(1,*) B_ext
  read(1,*) snap_start,snap_stop
  close(1)
  !
  ! Doing all directions (parallel, normal and binormal) if direction_='all',
  ! otherwise doing only direction=direction_
  !
  !  Loop over snapshots
  !
  do snap=snap_start,snap_stop
     !
     !  Loop over number of CPU's to get global array
     !
     !  Beginning parallel session
     !
     !$OMP PARALLEL DEFAULT(none)  &
     !$OMP firstprivate(datadir,snap_stop,ip,lshear) &
     !$OMP shared(iaz,iay,iax,iux,iuy,iuz) &
     !$OMP private(my_number,total_number,cpu,ipy,ipz,ch,snap) &
     !$OMP private(csnap,file,f,x,y,z,deltay,ipy_loc,ipz_loc) &
     !$OMP private(f,bb,b2,bb_loc) &
     !$OMP private(l1_glob,l2_glob,m1_glob,m2_glob) &
     !$OMP private(n1_glob,n2_glob,lpri,lext) &
     !$OMP private(V0,B0,mmm,nnn,ll) &
     !$OMP private(m_OMP,n_OMP) &
     !$OMP private(ny_OMP,nz_OMP,ipy_OMP,ipz_OMP,yprocs,zprocs) &
     !$OMP private(tmpB,PDF_local,bin,u2,b2,uf,bf,box_x,box_y,box_z) &
     !$OMP private(udotb,sizeofbox,meanU_box_local,meanB_box_local) &
     !$OMP shared(meanU_box,meanB_box,B_ext) &
     !$OMP shared(bb_glob,t,dx,dy,dz) &
     !$OMP shared(vv_glob,PDF)
     !
     !  Stopping if ncpus is not divideable by total_number 
     !
     !$ total_number=OMP_GET_NUM_THREADS()
     if (mod(ncpus,total_number) .ne. 0) then
        call stop_it('mod(ncpus,total_number) .ne. 0')
     endif
     !
     ! Who am I?
     !
     !$ my_number=OMP_GET_THREAD_NUM()
     !
     yprocs=1
     if (total_number .ge. 4) yprocs=2
     if (total_number .ge. 16) yprocs=4
     if (total_number .ge. 64) yprocs=8
     zprocs=total_number/yprocs
     if (yprocs*zprocs .ne. total_number) then
        call stop_it('(yprocs*zprocs) .ne. total_number)')
     endif
     ny_OMP=nygrid/yprocs
     nz_OMP=nzgrid/zprocs
     ipy_OMP = modulo(my_number, yprocs)
     ipz_OMP = my_number/(yprocs)
     !
     ! Read data from all processors
     !
     do ipz_loc=0,nprocz/zprocs-1
        do ipy_loc=0,nprocy/yprocs-1
           ipy=ipy_OMP*nprocy/yprocs+ipy_loc
           ipz=ipz_OMP*nprocz/zprocs+ipz_loc
           cpu=ipy+ipz*nprocy
           !
           ! Read var.dat 
           !
           call chn(cpu,ch)
           call chn(snap,csnap)
           if (snap_stop .eq. 0) &
                call safe_character_assign(file,trim(datadir) &
                //'/proc'//trim(ch)//'/var.dat')
           if (snap_stop .ne. 0) &
                call safe_character_assign(file,trim(datadir) &
                //'/proc'//trim(ch)//'/VAR'//trim(csnap))
           if (ip<8) print*,'Reading '//trim(file)
           print*,'Reading ',trim(file)
           !  Use 'my_number' instead of e.g. 1 as file identifier.
           open(my_number,file=file,form='unformatted')
           read(my_number) f(1:mx,1:my,1:mz,1:mvar)
           if (lshear) then
              read(my_number) t,x,y,z,dx,dy,dz,deltay
           else
              read(my_number) t,x,y,z,dx,dy,dz
           endif
           close(my_number)
           !
           ! Finding direction for calculation of the function of interest
           !
           lext=((B_ext(1).ne.0).or.(B_ext(2).ne.0).or.(B_ext(3).ne.0))
           do m_OMP=m1,m2
              do n_OMP=n1,n2                    
                 lpri=((my_number.eq.0).and.(m_OMP.eq.m1).and.(n_OMP.eq.n1))
                 call curl_OMP(f,iax,bb,m_OMP,n_OMP)
                 !
                 ! Adding imposed field
                 !
                 if (lext) then
                    if (lpri) print*,'Adding imposed field'
                    bb(:,1)=B_ext(1)+bb(:,1)
                    bb(:,2)=B_ext(2)+bb(:,2)
                    bb(:,3)=B_ext(3)+bb(:,3)
                 endif
                 bb_loc(l1:l2,m_OMP,n_OMP,:)=bb
              enddo
           enddo
           !
           ! Put local variables in global array  
           !
           l1_glob=1
           l2_glob=nxgrid
           m1_glob=1+ipy*ny
           m2_glob=(1+ipy)*ny
           n1_glob=1+ipz*nz
           n2_glob=(1+ipz)*nz
           !
           bb_glob(l1_glob:l2_glob,m1_glob:m2_glob,n1_glob:n2_glob,:)= &
                bb_loc(l1:l2,m1:m2,n1:n2,:)
           vv_glob(l1_glob:l2_glob,m1_glob:m2_glob,n1_glob:n2_glob,:)= &
                f(l1:l2,m1:m2,n1:n2,iux:iuz)
           !
           !  Finnishing loop over processors
           !
        enddo
     enddo
     !
     ! From now on we only consider the 'fluctuations' of B,
     ! we therefore subtract the imposed field.
     !
     if (lext .and. (my_number .eq. 0)) then
        print*,'Removing imposed field'
        bb_glob(:,:,:,1)=bb_glob(:,:,:,1)-B_ext(1)
        bb_glob(:,:,:,2)=bb_glob(:,:,:,2)-B_ext(2)
        bb_glob(:,:,:,3)=bb_glob(:,:,:,3)-B_ext(3)
     endif
     !
     !  ---------------------------------------------------
     !  Start calculation of the functions of interest
     !  ---------------------------------------------------
     !
     PDF=0
     PDF_local=0
     !
     ! Define which part of the global array to work on (In case of OpenMP)
     !
     l1_glob=1
     l2_glob=nxgrid
     m1_glob=1+ipy_OMP*ny_OMP
     m2_glob=(1+ipy_OMP)*ny_OMP
     n1_glob=1+ipz_OMP*nz_OMP
     n2_glob=(1+ipz_OMP)*nz_OMP
     !
     ! Need a barrier such that the whole global array is produced
     ! before we start the calculations
     !
     !$omp barrier
     !
     if (my_number .eq. 0)  then
        print*,'------------------------------------'
        print*,'Start calculation of '
        print*,'PDF for u dot b'
        print*,'------------------------------------'
     endif
     !
     !  Loop over all grid points in *my* part of the global array. 
     !
     meanB_box=0
     meanU_box=0
     meanB_box_local=0
     meanU_box_local=0
     !
     !  Find average of B and U in subbox
     !
     do ll=l1_glob,l2_glob
        if (my_number .eq. 0) print*,'Doing ',ll,' of ',nxgrid
        do mmm=m1_glob,m2_glob
           do nnn=n1_glob,n2_glob
              B0=bb_glob(ll,mmm,nnn,:)
              V0=vv_glob(ll,mmm,nnn,:)
              box_x=( (ll-1)/sizeofbox)+1
              box_y=((mmm-1)/sizeofbox)+1
              box_z=((nnn-1)/sizeofbox)+1
              meanB_box_local(:,box_x,box_y,box_z)= &
                   meanB_box_local(:,box_x,box_y,box_z)+B0
              meanU_box_local(:,box_x,box_y,box_z)= &
                   meanU_box_local(:,box_x,box_y,box_z)+V0
           end do
        end do
     end do
     !
     ! Summing contribution from all threads
     !
     !$OMP critical
     meanB_box=meanB_box+meanB_box_local
     meanU_box=meanU_box+meanU_box_local
     !$OMP end critical
     !$OMP barrier
     !
     !  Dividing by number of grid points in subbox 
     !
     meanB_box=meanB_box/sizeofbox**3
     meanU_box=meanU_box/sizeofbox**3
     !
     !  Find PDF of (u dot b)/sqrt(u^2*b^2)
     !
     do ll=l1_glob,l2_glob
        if (my_number .eq. 0) print*,'Doing ',ll,' of ',nxgrid
        do mmm=m1_glob,m2_glob
           do nnn=n1_glob,n2_glob
              B0=bb_glob(ll,mmm,nnn,:)
              V0=vv_glob(ll,mmm,nnn,:)
              box_x=(( ll-1)/sizeofbox)+1
              box_y=((mmm-1)/sizeofbox)+1
              box_z=((nnn-1)/sizeofbox)+1
              bf=B0-meanB_box(:,box_x,box_y,box_z)
              uf=V0-meanU_box(:,box_x,box_y,box_z)
              udotb=bf(1)*uf(1)+bf(2)*uf(2)+bf(3)*uf(3)
              u2=uf(1)*uf(1)+uf(2)*uf(2)+uf(3)*uf(3)
              b2=bf(1)*bf(1)+bf(2)*bf(2)+bf(3)*bf(3)
              bin=nint((udotb/sqrt(u2*b2)+1)*(nbins-1)/2)+1
              PDF_local(bin,box_x,box_y,box_z)= &
                   PDF_local(bin,box_x,box_y,box_z)+1
           end do
        end do
     end do


     !
     ! Summing contribution from all threads
     !
     !$OMP critical
     PDF=PDF+PDF_local
     !$OMP end critical
     !
     !  Finalizinig parallel session
     !
     !$OMP END PARALLEL 
     !
     ! Normalizing
     !
     PDF=PDF/sizeofbox**3
     !
     !  Saving the results
     !
     if (snap_stop .eq. 0) call safe_character_assign(result_file, &
          trim(datadir)//'/alfven_subbox'//'.dat')
     if (snap_stop .ne. 0) call safe_character_assign(result_file, &
          trim(datadir)//'/alfven_subbox_'//'_VAR'// &
          trim(csnap)//'.dat')
     close(1)
     if (ip<8) print*,'Writing '//trim(result_file)
     open(1,file=result_file)
     write(1,'(3e12.4)') t
     write(1,'(3e12.4)') PDF
     close(1)
     !
     !  Finnishing loop over snapshots (if snap_stop .ne. 0)
     !
  enddo
  !
  ! Deallocating the allocateable arrays
  !
  deallocate(vv_glob)
  deallocate(bb_glob)
  !
end program alfven_subbox