!> Routines for calculating and writing various physical diagnostics
module stella_diagnostics

  implicit none

  public :: init_stella_diagnostics, finish_stella_diagnostics
  public :: diagnose_stella
  public :: nsave

  private

  integer :: ntg_out
  integer :: nwrite, nsave, navg
  integer :: stdout_unit, fluxes_unit, omega_unit
  logical :: save_for_restart
  logical :: write_omega
  logical :: write_moments
  logical :: write_phi_vs_time
  logical :: write_gvmus
  logical :: write_gzvs
  logical :: write_kspectra
  logical :: write_radial_fluxes
  logical :: write_fluxes_kxkyz
  logical :: flux_norm

  !> Arrays needed for averaging in x,y,z
  real, dimension (:), allocatable :: pflux_avg, vflux_avg, qflux_avg, heat_avg
  real, dimension (:,:,:), allocatable :: pflux, vflux, qflux, exchange

  !> Needed for calculating growth rates and frequencies
  complex, dimension (:,:,:), allocatable :: omega_vs_time

  !> Current maximum index of the time dimension in the netCDF file
  integer :: nout = 1
  !> Has this module been initialised?
  logical :: diagnostics_initialized = .false.

  !> Debugging
  logical :: debug = .false.

contains

  !> Initialise the [[stella_diagnostics]] module
  !>
  !> Broadcast the parameters from the namelist "stella_diagnostics_knobs"
  !> and open/append the netcdf file and the ascii files.
  subroutine init_stella_diagnostics (restart, tstart)

    use zgrid, only: init_zgrid
    use kt_grids, only: init_kt_grids
    use physics_parameters, only: init_physics_parameters
    use run_parameters, only: init_run_parameters
    use species, only: init_species
    use dist_fn, only: init_dist_fn
    use init_g, only: init_init_g
    use stella_io, only: init_stella_io, get_nout
    use mp, only: broadcast, proc0

    implicit none

    !> Has this simulation been restarted?
    logical, intent (in) :: restart
    !> Current simulation time
    real, intent (in) :: tstart

    ! Only initialize the diagnostics once
    if (diagnostics_initialized) return
    diagnostics_initialized = .true.

    ! Only debug on the first processor
    debug = debug .and. proc0

    ! Make sure the other routines are intialized
    call init_zgrid
    call init_physics_parameters
    call init_kt_grids
    call init_run_parameters
    call init_species
    call init_init_g
    call init_dist_fn

    ! Read the namelist "stella_diagnostics_knobs" in the input file
    call read_parameters
    call allocate_arrays

    ! Broadcast the variables to all processors
    call broadcast (nwrite)
    call broadcast (navg)
    call broadcast (nsave)
    call broadcast (save_for_restart)
    call broadcast (write_omega)
    call broadcast (write_kspectra)
    call broadcast (write_moments)
    call broadcast (write_phi_vs_time)
    call broadcast (write_gvmus)
    call broadcast (write_gzvs)
    call broadcast (write_radial_fluxes)
    call broadcast (write_fluxes_kxkyz)
    call broadcast (flux_norm)

    ! Initiate the netcdf file with extension '.out.nc'
    call init_stella_io (restart, write_phi_vs_time, write_kspectra, &
         write_gvmus, write_gzvs, write_moments, write_radial_fluxes, &
         write_fluxes_kxkyz)

    ! Open the '.out', '.fluxes' and '.omega' file
    if (proc0) call open_loop_ascii_files (restart)

    ! Get the final position [[nout]] of the time axis in the netcdf file
    if (proc0) call get_nout (tstart,nout)
    call broadcast (nout)

  end subroutine init_stella_diagnostics

  !> Read the diagnostic input parameters from the input file
  !>
  !> Namelist: `stella_diagnostics_knobs`
  subroutine read_parameters

    use mp, only: proc0
    use file_utils, only: input_unit_exist
    use zgrid, only: nperiod, nzed
    use physics_flags, only: radial_variation

    implicit none

    logical :: exist
    integer :: in_file

    namelist /stella_diagnostics_knobs/ nwrite, navg, nsave, &
         save_for_restart, write_phi_vs_time, write_gvmus, write_gzvs, &
         write_omega, write_kspectra, write_moments, write_radial_fluxes, &
         write_fluxes_kxkyz, flux_norm

    if (proc0) then
       nwrite = 50
       navg = 50
       nsave = -1
       save_for_restart = .false.
       write_omega = .false.
       write_phi_vs_time = .false.
       write_gvmus = .false.
       write_gzvs = .false.
       write_kspectra = .false.
       write_moments = .false.
       write_radial_fluxes = radial_variation
       write_fluxes_kxkyz = .false.
       flux_norm = .true.

       in_file = input_unit_exist ("stella_diagnostics_knobs", exist)
       if (exist) read (unit=in_file, nml=stella_diagnostics_knobs)

       if (.not. save_for_restart) nsave = -1
    end if
    ntg_out = nzed/2 + (nperiod-1)*nzed

  end subroutine read_parameters

  !> Allocate the module-level arrays
  subroutine allocate_arrays

    use species, only: nspec
    use kt_grids, only: nakx, naky

    implicit none

    if (.not.allocated(pflux)) allocate(pflux (nakx,naky,nspec)) ; pflux = 0.
    if (.not.allocated(qflux)) allocate(qflux (nakx,naky,nspec)) ; qflux = 0.
    if (.not.allocated(vflux)) allocate(vflux (nakx,naky,nspec)) ; vflux = 0.
    if (.not.allocated(exchange)) allocate(exchange (nakx,naky,nspec)) ; exchange = 0.
    if (.not.allocated(pflux_avg)) allocate(pflux_avg(nspec)) ; pflux_avg = 0.
    if (.not.allocated(qflux_avg)) allocate(qflux_avg(nspec)) ; qflux_avg = 0.
    if (.not.allocated(vflux_avg)) allocate(vflux_avg(nspec)) ; vflux_avg = 0.
    if (.not.allocated(heat_avg)) allocate(heat_avg(nspec)) ; heat_avg = 0.
    if (.not.allocated(omega_vs_time)) then
       if (write_omega) then
          allocate (omega_vs_time(navg,naky,nakx))
          omega_vs_time = 0.
       else
          allocate (omega_vs_time(1,1,1))
          navg=1
       end if
    end if

  end subroutine allocate_arrays

  !> Open the '.out' and the '.fluxes' file.
  !>
  !> When running a new simulation, create a new file or replace an old file.
  !> When restarting a simulation, append the old files.
  subroutine open_loop_ascii_files(restart)

    use file_utils, only: open_output_file
    use species, only: nspec

    implicit none

    logical, intent (in) :: restart
    character (3) :: nspec_str
    character (100) :: str
    logical :: overwrite

    ! Do not overwrite, but append files, when we restart the simulation.
    overwrite = .not.restart

    ! Open the '.out' and the '.fluxes' files.
    call open_output_file (stdout_unit,'.out',overwrite)
    call open_output_file (fluxes_unit,'.fluxes',overwrite)

    ! Create the header for the .fluxes file.
    ! Every column is made up of 12 spaces, so make sure the headers
    ! are placed correctly since we have the following columns for nspec=3:
    ! #time pflx1 pflx2 pflx3 vflx1 vflx2 vflx32 qflx1 qflx2 qflx3
    if (.not.restart) then
      write (nspec_str,'(i3)') nspec*12
      str = trim('(2a12,2a'//trim(nspec_str)//')')
      write (fluxes_unit,str) '#time', 'pflx', 'vflx', 'qflx'
    end if

    ! Open the '.omega' file and create its header.
    if (write_omega) then
       call open_output_file (omega_unit,'.omega',overwrite)
       if (.not.restart) then
         write (omega_unit,'(7a12)') '#time', 'ky', 'kx', &
                'Re[om]', 'Im[om]', 'Re[omavg]', 'Im[omavg]'
       end if
    end if

  end subroutine open_loop_ascii_files

  !> Close the text files opened by [[open_loop_ascii_files]]
  subroutine close_loop_ascii_files

    use file_utils, only: close_output_file

    implicit none

    call close_output_file (stdout_unit)
    call close_output_file (fluxes_unit)
    if (write_omega) call close_output_file (omega_unit)

  end subroutine close_loop_ascii_files

  !> Calculate and write diagnostics
  subroutine diagnose_stella (istep)

    use mp, only: proc0
    use constants, only: zi
    use redistribute, only: scatter
    use fields_arrays, only: phi, apar
    use fields_arrays, only: phi_old, phi_corr_QN
    use fields, only: fields_updated, advance_fields
    use dist_fn_arrays, only: gvmu, gnew
    use g_tofrom_h, only: g_to_h
    use stella_io, only: write_time_nc
    use stella_io, only: write_phi2_nc
    use stella_io, only: write_phi_nc
    use stella_io, only: write_gvmus_nc
    use stella_io, only: write_gzvs_nc
    use stella_io, only: write_kspectra_nc
    use stella_io, only: write_moments_nc
    use stella_io, only: write_radial_fluxes_nc
    use stella_io, only: write_fluxes_kxkyz_nc
    use stella_io, only: sync_nc
!    use stella_io, only: write_symmetry_nc
    use stella_time, only: code_time, code_dt
    use zgrid, only: nztot, nzgrid, ntubes
    use vpamu_grids, only: nmu, nvpa
    use species, only: nspec
    use kt_grids, only: naky, nakx
    use dist_redistribute, only: kxkyz2vmu
    use physics_flags, only: radial_variation
    use volume_averages, only: volume_average, fieldline_average

    implicit none

    !> The current timestep
    integer, intent (in) :: istep

    real :: phi2, apar2
    real :: zero
    real, dimension (:,:,:), allocatable :: gvmus
    real, dimension (:,:,:,:), allocatable :: gzvs
!    real, dimension (:,:,:), allocatable :: pflx_zvpa, vflx_zvpa, qflx_zvpa
    real, dimension (:), allocatable :: part_flux, mom_flux, heat_flux
    real, dimension (:,:), allocatable :: part_flux_x, mom_flux_x, heat_flux_x
    real, dimension (:,:), allocatable :: phi2_vs_kxky
    real, dimension (:,:,:,:,:), allocatable :: pflx_kxkyz, vflx_kxkyz, qflx_kxkyz
    complex, dimension (:,:,:,:,:), allocatable :: density, upar, temperature

    complex, dimension (:,:), allocatable :: omega_avg
    complex, dimension (:,:), allocatable :: phiavg, phioldavg
    complex, dimension (:,:,:,:), allocatable :: phi_out

    ! calculation of omega requires computation of omega more
    ! frequently than every nwrite time steps
    if (write_omega .and. proc0) then
       zero = 100.*epsilon(0.)
       if (istep > 0) then
          allocate (phiavg(naky,nakx))
          allocate (phioldavg(naky,nakx))
          call fieldline_average (phi, phiavg)
          call fieldline_average (phi_old, phioldavg)
          where (abs(phiavg) < zero .or. abs(phioldavg) < zero)
             omega_vs_time(mod(istep,navg)+1,:,:) = 0.0
          elsewhere
             omega_vs_time(mod(istep,navg)+1,:,:) = log(phiavg/phioldavg)*zi/code_dt
          end where
          deallocate (phiavg, phioldavg)
       end if
    end if

    ! only write data to file every nwrite time steps
    if (mod(istep,nwrite) /= 0) return

    if(radial_variation) fields_updated = .false.
    call advance_fields(gnew, phi, apar, dist='gbar')

    allocate(phi_out(naky,nakx,-nzgrid:nzgrid,ntubes))
    phi_out = phi
    if(radial_variation) then
      phi_out = phi_out + phi_corr_QN
    endif

    allocate (part_flux(nspec))
    allocate (mom_flux(nspec))
    allocate (heat_flux(nspec))

    allocate (pflx_kxkyz(naky,nakx,nztot,ntubes,nspec))
    allocate (vflx_kxkyz(naky,nakx,nztot,ntubes,nspec))
    allocate (qflx_kxkyz(naky,nakx,nztot,ntubes,nspec))

    if(write_radial_fluxes) then
      allocate (part_flux_x(nakx,nspec))
      allocate (mom_flux_x(nakx,nspec))
      allocate (heat_flux_x(nakx,nspec))
    endif

    ! obtain turbulent fluxes
    if(radial_variation.or.write_radial_fluxes) then
!     handle g_to_h in get_fluxes_vmulo to eliminate x^2 terms
!     call g_to_h (gnew, phi, fphi, phi_corr_QN)
      call get_fluxes_vmulo (gnew, phi_out, part_flux, mom_flux, heat_flux, &
                                            part_flux_x,mom_flux_x, heat_flux_x)
!     call g_to_h (gnew, phi, -fphi, phi_corr_QN)
    else
      call scatter (kxkyz2vmu, gnew, gvmu)
!     call g_to_h (gvmu, phi, fphi)
      call get_fluxes (gvmu, part_flux, mom_flux, heat_flux, &
           pflx_kxkyz, vflx_kxkyz, qflx_kxkyz)
!     call g_to_h (gvmu, phi, -fphi)
    endif

    if (proc0) then
       if(write_omega) then
         allocate (omega_avg(naky,nakx))
         omega_avg = sum(omega_vs_time,dim=1)/real(navg)
       else
         allocate (omega_avg(1,1))
       endif
       call volume_average (phi_out, phi2)
       call volume_average (apar, apar2)
       ! Print information to stella.out, the header is printed in stella.f90
       write (*,'(A2,I7,A2,ES12.4,A2,ES12.4,A2,ES12.4)') " ",istep," ",code_time," ",code_dt," ", phi2
       call write_loop_ascii_files (istep, phi2, apar2, part_flux, mom_flux, heat_flux, &
            omega_vs_time(mod(istep,navg)+1,:,:), omega_avg)

       ! do not need omega_avg again this time step
       deallocate (omega_avg)
    end if

    if (proc0) then
       if (debug) write (*,*) 'stella_diagnostics::write_time_nc'
       call write_time_nc (nout, code_time)
       call write_phi2_nc (nout, phi2)
       if (write_phi_vs_time) then
          if (debug) write (*,*) 'stella_diagnostics::diagnose_stella::write_phi_nc'
          call write_phi_nc (nout, phi_out)
       end if
       if (write_kspectra) then
          if (debug) write (*,*) 'stella_diagnostics::diagnose_stella::write_kspectra'
          allocate (phi2_vs_kxky(naky,nakx))
          call fieldline_average (real(phi_out*conjg(phi_out)),phi2_vs_kxky)
          call write_kspectra_nc (nout, phi2_vs_kxky)
          deallocate (phi2_vs_kxky)
       end if
       if (write_radial_fluxes) then
          call write_radial_fluxes_nc(nout,part_flux_x,mom_flux_x,heat_flux_x)
       endif
    end if
    if (write_moments) then
       if (debug) write (*,*) 'stella_diagnostics::diagnose_stella::write_moments'
       allocate (density(naky,nakx,nztot,ntubes,nspec))
       allocate (upar(naky,nakx,nztot,ntubes,nspec))
       allocate (temperature(naky,nakx,nztot,ntubes,nspec))
       call get_moments (gnew, density, upar, temperature)
       if (proc0) call write_moments_nc (nout, density, upar, temperature)
       deallocate (density, upar, temperature)
    end if
    IF (write_fluxes_kxkyz) then
       IF (debug) write (*,*) 'stella_diagnostics::diagnose_stella::write_fluxes_kxkyz'
       IF (proc0) call write_fluxes_kxkyz_nc (nout, pflx_kxkyz, vflx_kxkyz, qflx_kxkyz)
    ENDIF
    if (write_gvmus) then
       allocate (gvmus(nvpa,nmu,nspec))
       if (debug) write (*,*) 'stella_diagnostics::diagnose_stella::get_gvmus'
       ! note that gvmus is h at this point
       call get_gvmus (gvmu, gvmus)
       if (debug) write (*,*) 'stella_diagnostics::diagnose_stella::write_gvmus_nc'
       if (proc0) call write_gvmus_nc (nout, gvmus)
       deallocate (gvmus)
    end if
    if (write_gzvs) then
       allocate (gzvs(ntubes,nztot,nvpa,nspec))
       if (debug) write (*,*) 'stella_diagnostics::diagnose_stella::get_gzvs'
       call get_gzvs (gnew, gzvs)
       if (debug) write (*,*) 'stella_diagnostics::diagnose_stella::write_gzvs_nc'
       if (proc0) call write_gzvs_nc (nout, gzvs)
       deallocate (gzvs)
    end if
!     if (write_symmetry) then
!        allocate (pflx_zvpa(nztot,nvpa,nspec))
!        allocate (vflx_zvpa(nztot,nvpa,nspec))
!        allocate (qflx_zvpa(nztot,nvpa,nspec))
!        call get_fluxes_vs_zvpa (gnew, pflx_zvpa, vflx_zvpa, qflx_zvpa)
!        deallocate (pflx_zvpa, vflx_zvpa, qflx_zvpa)
!     end if

    if (proc0) call sync_nc

    deallocate (part_flux, mom_flux, heat_flux)
    if(allocated(part_flux_x)) deallocate (part_flux_x)
    if(allocated(mom_flux_x))  deallocate (mom_flux_x)
    if(allocated(heat_flux_x)) deallocate (heat_flux_x)
    deallocate(phi_out)

    nout = nout + 1

  end subroutine diagnose_stella

  !> Calculate fluxes
  !>
  !> Assumes that the non-Boltzmann part of df is passed in (aka h)
  subroutine get_fluxes (g, pflx, vflx, qflx,&
       pflx_vs_kxkyz, vflx_vs_kxkyz, qflx_vs_kxkyz)

    use mp, only: sum_reduce
    use constants, only: zi
    use fields_arrays, only: phi, apar
    use stella_layouts, only: kxkyz_lo
    use stella_layouts, only: iky_idx, ikx_idx, iz_idx, it_idx, is_idx
    use species, only: spec, nspec
    use stella_geometry, only: jacob, grho, bmag, btor
    use stella_geometry, only: gds21, gds22
    use stella_geometry, only: geo_surf
    use zgrid, only: delzed, nzgrid, ntubes
    use vpamu_grids, only: nvpa, nmu
    use vpamu_grids, only: vperp2, vpa
    use run_parameters, only: fphi, fapar
    use kt_grids, ONLY: nakx, naky
    use kt_grids, only: aky, theta0
    use gyro_averages, only: gyro_average, gyro_average_j1

    implicit none

    complex, dimension (:,:,kxkyz_lo%llim_proc:), intent (in) :: g
    real, dimension (:), intent (out) :: pflx, vflx, qflx
    real, dimension (:,:,-nzgrid:,:,:), intent (out) :: pflx_vs_kxkyz, vflx_vs_kxkyz, qflx_vs_kxkyz

    integer :: ikxkyz, iky, ikx, iz, it, is, ia
    real, dimension (:), allocatable :: flx_norm
    real :: flx_norm_partial
    complex, dimension (:,:), allocatable :: gtmp1, gtmp2, gtmp3

    allocate (flx_norm(-nzgrid:nzgrid))
    allocate (gtmp1(nvpa,nmu), gtmp2(nvpa,nmu), gtmp3(nvpa,nmu))

    pflx = 0. ; vflx = 0. ; qflx = 0.
    pflx_vs_kxkyz = 0. ; vflx_vs_kxkyz = 0. ; qflx_vs_kxkyz = 0.

    flx_norm = jacob(1,:)*delzed
    flx_norm(-nzgrid) = 0.5*flx_norm(-nzgrid)
    flx_norm( nzgrid) = 0.5*flx_norm( nzgrid)

    if (flux_norm) then
       ! Flux definition with an extra factor 1/<\nabla\rho> in front.
       flx_norm_partial = sum(flx_norm)/sum(flx_norm*grho(1,:))
       flx_norm = flx_norm/sum(flx_norm*grho(1,:))
    else
       ! Flux definition witou the extra factor.
       flx_norm_partial = 1.0
       flx_norm = flx_norm/sum(flx_norm)
    endif

    ia = 1
    ! get electrostatic contributions to fluxes
    if (fphi > epsilon(0.0)) then
       do ikxkyz = kxkyz_lo%llim_proc, kxkyz_lo%ulim_proc
          iky = iky_idx(kxkyz_lo,ikxkyz)
          ikx = ikx_idx(kxkyz_lo,ikxkyz)
          iz = iz_idx(kxkyz_lo,ikxkyz)
          it = it_idx(kxkyz_lo,ikxkyz)
          is = is_idx(kxkyz_lo,ikxkyz)

          ! get particle flux
          call gyro_average (g(:,:,ikxkyz), ikxkyz, gtmp1)
          call get_one_flux (iky, iz, flx_norm(iz), gtmp1, phi(iky,ikx,iz,it), pflx(is))
          call get_one_flux (iky, iz, flx_norm_partial, gtmp1, phi(iky,ikx,iz,it), pflx_vs_kxkyz(iky,ikx,iz,it,is))

          ! get heat flux
          ! NEEDS TO BE MODIFIED TO TREAT ENERGY = ENERGY(ALPHA)
          gtmp1 = gtmp1*(spread(vpa**2,2,nmu)+spread(vperp2(1,iz,:),1,nvpa))
          call get_one_flux (iky, iz, flx_norm(iz), gtmp1, phi(iky,ikx,iz,it), qflx(is))
          call get_one_flux (iky, iz, flx_norm_partial, gtmp1, phi(iky,ikx,iz,it), qflx_vs_kxkyz(iky,ikx,iz,it,is))

          ! get momentum flux
          ! parallel component
          gtmp1 = g(:,:,ikxkyz)*spread(vpa,2,nmu)*geo_surf%rmaj*btor(iz)/bmag(ia,iz)
          call gyro_average (gtmp1, ikxkyz, gtmp2)
          gtmp1 = -g(:,:,ikxkyz)*zi*aky(iky)*spread(vperp2(ia,iz,:),1,nvpa)*geo_surf%rhoc &
               * (gds21(ia,iz)+theta0(iky,ikx)*gds22(ia,iz))*spec(is)%smz &
               / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2)
          call gyro_average_j1 (gtmp1, ikxkyz, gtmp3)
          gtmp1 = gtmp2 + gtmp3

          call get_one_flux (iky, iz, flx_norm(iz), gtmp1, phi(iky,ikx,iz,it), vflx(is))
          call get_one_flux (iky, iz, flx_norm_partial, gtmp1, phi(iky,ikx,iz,it), vflx_vs_kxkyz(iky,ikx,iz,it,is))
       end do
    end if

    if (fapar > epsilon(0.0)) then
       ! particle flux
       do ikxkyz = kxkyz_lo%llim_proc, kxkyz_lo%ulim_proc
          iky = iky_idx(kxkyz_lo,ikxkyz)
          ikx = ikx_idx(kxkyz_lo,ikxkyz)
          iz = iz_idx(kxkyz_lo,ikxkyz)
          it = it_idx(kxkyz_lo,ikxkyz)
          is = is_idx(kxkyz_lo,ikxkyz)

          ! Apar contribution to particle flux
          gtmp1 = -g(:,:,ikxkyz)*spec(is)%stm*spread(vpa,2,nmu)
          call gyro_average (gtmp1, ikxkyz, gtmp2)
          call get_one_flux (iky, iz, flx_norm(iz), gtmp2, apar(iky,ikx,iz,it), pflx(is))

          ! Apar contribution to heat flux
          gtmp2 = gtmp2*(spread(vpa**2,2,nmu)+spread(vperp2(ia,iz,:),1,nvpa))
          call get_one_flux (iky, iz, flx_norm(iz), gtmp2, apar(iky,ikx,iz,it), qflx(is))

          ! Apar contribution to momentum flux
          ! parallel component
          gtmp1 = -spread(vpa**2,2,nmu)*spec(is)%stm*g(:,:,ikxkyz) &
               * geo_surf%rmaj*btor(iz)/bmag(1,iz)
          call gyro_average (gtmp1, ikxkyz, gtmp2)
          ! perp component
          gtmp1 = spread(vpa,2,nmu)*spec(is)%stm*g(:,:,ikxkyz) &
               * zi*aky(iky)*spread(vperp2(ia,iz,:),1,nvpa)*geo_surf%rhoc &
               * (gds21(ia,iz)+theta0(iky,ikx)*gds22(ia,iz))*spec(is)%smz &
               / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2)
          call gyro_average_j1 (gtmp1, ikxkyz, gtmp3)
          ! FLAG -- NEED TO ADD IN CONTRIBUTION FROM BOLTZMANN PIECE !!

          gtmp1 = gtmp2 + gtmp3

          call get_one_flux (iky, iz, flx_norm(iz), gtmp1, apar(iky,ikx,iz,it), vflx(is))
       end do
    end if

    call sum_reduce (pflx, 0) ; pflx = pflx*spec%dens_psi0
    call sum_reduce (qflx, 0) ; qflx = qflx*spec%dens_psi0*spec%temp_psi0
    call sum_reduce (vflx, 0) ; vflx = vflx*spec%dens_psi0*sqrt(spec%mass*spec%temp_psi0)

    ! normalise to account for contributions from multiple flux tubes
    ! in flux tube train
    pflx = pflx/real(ntubes)
    qflx = qflx/real(ntubes)
    vflx = vflx/real(ntubes)

    call sum_reduce (pflx_vs_kxkyz, 0)
    call sum_reduce (qflx_vs_kxkyz, 0)
    call sum_reduce (vflx_vs_kxkyz, 0)

    do is = 1, nspec
      pflx_vs_kxkyz(:,:,:,:,is) = pflx_vs_kxkyz(:,:,:,:,is)*spec(is)%dens_psi0
      qflx_vs_kxkyz(:,:,:,:,is) = qflx_vs_kxkyz(:,:,:,:,is)*spec(is)%dens_psi0*spec(is)%temp_psi0
      vflx_vs_kxkyz(:,:,:,:,is) = vflx_vs_kxkyz(:,:,:,:,is)*spec(is)%dens_psi0*sqrt(spec(is)%mass*spec(is)%temp_psi0)
    enddo

    deallocate (gtmp1, gtmp2, gtmp3)
    deallocate (flx_norm)

  end subroutine get_fluxes

  !==============================================
  !============ GET FLUXES VMULO ================
  !==============================================
  subroutine get_fluxes_vmulo (g, phi, pflx, vflx, qflx, pflx_x, vflx_x, qflx_x)

    use mp, only: sum_reduce
    use constants, only: zi
    use dist_fn_arrays, only: g1, g2, kperp2, dkperp2dr
    use stella_layouts, only: vmu_lo
    use stella_layouts, only: iv_idx, imu_idx, is_idx
    use species, only: spec
    use stella_geometry, only: jacob, grho, bmag, btor
    use stella_geometry, only: drhodpsi
    use stella_geometry, only: gds21, gds22
    use stella_geometry, only: dgds21dr, dgds22dr
    use stella_geometry, only: geo_surf
    use stella_geometry, only: dBdrho, dIdrho
    use stella_geometry, only: dl_over_b, d_dl_over_b_drho
    use zgrid, only: delzed, nzgrid, ntubes
    use vpamu_grids, only: vperp2, vpa, mu
    use vpamu_grids, only: maxwell_vpa, maxwell_mu, maxwell_fac
    use run_parameters, only: fphi
    use kt_grids, only: aky, theta0, naky, nakx, multiply_by_rho
    use physics_flags, only: radial_variation
    use gyro_averages, only: gyro_average, gyro_average_j1, aj0x, aj1x

    implicit none

    complex, dimension (:,:,-nzgrid:,:,vmu_lo%llim_proc:), intent (in) :: g
    complex, dimension (:,:,-nzgrid:,:), intent (in) :: phi
    real, dimension (:), intent (out) :: pflx, vflx, qflx
    real, dimension (:,:), intent (out) :: pflx_x, vflx_x, qflx_x

    integer :: ivmu, imu, iv, iz, it, is, ia
    real, dimension (:), allocatable :: flx_norm, dflx_norm
    complex, dimension (:,:), allocatable :: g0k,g1k

    allocate (flx_norm(-nzgrid:nzgrid))
    allocate (dflx_norm(-nzgrid:nzgrid))

    pflx = 0. ; vflx = 0. ; qflx = 0.
    pflx_x = 0. ; vflx_x = 0. ; qflx_x = 0.

    ia = 1

    flx_norm = jacob(1,:)*delzed
    flx_norm(-nzgrid) = 0.5*flx_norm(-nzgrid)
    flx_norm( nzgrid) = 0.5*flx_norm( nzgrid)

    flx_norm = flx_norm/sum(flx_norm*grho(1,:))

    allocate (g0k(naky,nakx))
    allocate (g1k(naky,nakx))
    if(radial_variation) then
      where (dl_over_b(ia,:) .gt. epsilon(0.0))
        dflx_norm = d_dl_over_b_drho(ia,:)/dl_over_b(ia,:)
      elsewhere
        dflx_norm = 0.
      endwhere
    endif

    ! FLAG - electrostatic for now
    ! get electrostatic contributions to fluxes

    if (fphi > epsilon(0.0)) then
       ia = 1

       !get particle flux
       do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
          iv = iv_idx(vmu_lo,ivmu)
          imu = imu_idx(vmu_lo,ivmu)
          is = is_idx(vmu_lo,ivmu)

          call gyro_average (g(:,:,:,:,ivmu), ivmu, g1(:,:,:,:,ivmu))

          do it = 1, ntubes
            do iz= -nzgrid, nzgrid
              if(radial_variation) then
                g0k = g1(:,:,iz,it,ivmu) &
                  * (-0.5*aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                  * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                  * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                  + dBdrho(iz)/bmag(ia,iz) + dflx_norm(iz))

                call multiply_by_rho(g0k)
                g1(:,:,iz,it,ivmu) = g1(:,:,iz,it,ivmu) + g0k
              endif

              !subtract adiabatic contribution part of g
              g0k = spec(is)%zt*fphi*phi(:,:,iz,it)*aj0x(:,:,iz,ivmu)**2 &
                    *maxwell_vpa(iv,is)*maxwell_mu(ia,iz,imu,is)*maxwell_fac(is)
              if(radial_variation) then
                g1k = g0k*( -spec(is)%tprim*(vpa(iv)**2+vperp2(ia,iz,imu) - 2.5) &
                            -spec(is)%fprim - 2.0*dBdrho(iz)*mu(imu)  &
                            -aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                                 * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                                 * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                            + dBdrho(iz)/bmag(ia,iz)+dflx_norm(iz))

                call multiply_by_rho(g1k)

                g0k = g0k + g1k
              endif
              g1(:,:,iz,it,ivmu) = g1(:,:,iz,it,ivmu) + g0k

            enddo
          enddo
       enddo
       call get_one_flux_vmulo (flx_norm, spec%dens_psi0, g1, phi, pflx)

       if(write_radial_fluxes) then
         call get_one_flux_radial (flx_norm, spec%dens_psi0, g1, phi, pflx_x)
       endif

       !get heat flux
       do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
          iv = iv_idx(vmu_lo,ivmu)
          imu = imu_idx(vmu_lo,ivmu)
          is = is_idx(vmu_lo,ivmu)

          call gyro_average (g(:,:,:,:,ivmu), ivmu, g1(:,:,:,:,ivmu))
          do it = 1, ntubes
            do iz= -nzgrid, nzgrid

              g1(:,:,iz,it,ivmu) = g1(:,:,iz,it,ivmu)*(vpa(iv)**2+vperp2(ia,iz,imu))

              if(radial_variation) then
                g0k = g1(:,:,iz,it,ivmu) &
                  * (-0.5*aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                     * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                     * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                     + dBdrho(iz)/bmag(ia,iz) &
                     + 2.0*mu(imu)*dBdrho(iz)/(vpa(iv)**2+vperp2(ia,iz,imu)) &
                     + dflx_norm(iz))

                call multiply_by_rho(g0k)

                g1(:,:,iz,it,ivmu) = g1(:,:,iz,it,ivmu) + g0k

              endif

              !subtract adiabatic contribution part of g
              g0k = spec(is)%zt*fphi*phi(:,:,iz,it)*aj0x(:,:,iz,ivmu)**2 &
                    *(vpa(iv)**2+vperp2(ia,iz,imu)) &
                    *maxwell_vpa(iv,is)*maxwell_mu(ia,iz,imu,is)*maxwell_fac(is)
              if(radial_variation) then
                g1k = g0k*( -spec(is)%tprim*(vpa(iv)**2+vperp2(ia,iz,imu) - 2.5) &
                            -spec(is)%fprim - 2.0*dBdrho(iz)*mu(imu)  &
                            -aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                                 * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                                 * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                            + dBdrho(iz)/bmag(ia,iz)+dflx_norm(iz) &
                            + 2.0*mu(imu)*dBdrho(iz)/(vpa(iv)**2+vperp2(ia,iz,imu)))

                call multiply_by_rho(g1k)

                g0k = g0k + g1k
              endif
              g1(:,:,iz,it,ivmu) = g1(:,:,iz,it,ivmu) + g0k
            enddo
          enddo
       enddo
       call get_one_flux_vmulo (flx_norm,spec%dens_psi0*spec%temp_psi0, g1, phi, qflx)

       if(write_radial_fluxes) then
         call get_one_flux_radial (flx_norm,spec%dens_psi0*spec%temp_psi0, g1, phi, qflx_x)
       endif

       ! get momentum flux
       do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
          iv = iv_idx(vmu_lo,ivmu)
          imu = imu_idx(vmu_lo,ivmu)
          is = is_idx(vmu_lo,ivmu)
          do it = 1, ntubes
            do iz= -nzgrid, nzgrid
            ! parallel component
              g0k = g(:,:,iz,it,ivmu)*vpa(iv)*geo_surf%rmaj*btor(iz)/bmag(ia,iz)
              call gyro_average (g0k, iz, ivmu, g1(:,:,iz,it,ivmu))

              if(radial_variation) then
                g0k = g1(:,:,iz,it,ivmu) &
                  * (-0.5*aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                  * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                  * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                  + dIdrho/(geo_surf%rmaj*btor(iz)) &
                  + dflx_norm(iz))

                call multiply_by_rho(g0k)

                g1(:,:,iz,it,ivmu) = g1(:,:,iz,it,ivmu) + g0k

              endif
              !subtract adiabatic contribution part of g
              g0k = spec(is)%zt*fphi*phi(:,:,iz,it)*aj0x(:,:,iz,ivmu)**2 &
                    *vpa(iv)*geo_surf%rmaj*btor(iz)/bmag(ia,iz) &
                    *maxwell_vpa(iv,is)*maxwell_mu(ia,iz,imu,is)*maxwell_fac(is)
              if(radial_variation) then
                g1k = g0k*( -spec(is)%tprim*(vpa(iv)**2+vperp2(ia,iz,imu) - 2.5) &
                            -spec(is)%fprim - 2.0*dBdrho(iz)*mu(imu)  &
                            -aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                                 * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                                 * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                            + dflx_norm(iz) + dIdrho/(geo_surf%rmaj*btor(iz)))

                call multiply_by_rho(g1k)

                g0k = g0k + g1k
              endif
              g1(:,:,iz,it,ivmu) = g1(:,:,iz,it,ivmu) + g0k

              ! perpendicular component
              g0k = -g(:,:,iz,it,ivmu)*zi*spread(aky,2,nakx)*vperp2(ia,iz,imu)*geo_surf%rhoc &
                * (gds21(ia,iz)+theta0*gds22(ia,iz))*spec(is)%smz &
                / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2)

              call gyro_average_j1 (g0k, iz, ivmu, g2(:,:,iz,it,ivmu))
              if(radial_variation) then
                g0k =     - g(:,:,iz,it,ivmu)*zi*spread(aky,2,nakx)*vperp2(ia,iz,imu)*geo_surf%rhoc &
                     * (dgds21dr(ia,iz)+theta0*dgds22dr(ia,iz))*aj1x(:,:,iz,ivmu)*spec(is)%smz &
                     / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2)

                g0k = g0k - g(:,:,iz,it,ivmu)*zi*spread(aky,2,nakx)*vperp2(ia,iz,imu)*geo_surf%rhoc &
                            * (gds21(ia,iz)+theta0*gds22(ia,iz))*spec(is)%smz &
                            / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2) &
                            * (0.5*aj0x(:,:,iz,ivmu) - aj1x(:,:,iz,ivmu))  &
                            * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz))

                g0k = g0k + g2(:,:,iz,it,ivmu) &
                    * (- geo_surf%d2qdr2*geo_surf%rhoc/(geo_surf%shat*geo_surf%qinp) &
                       - geo_surf%d2psidr2*drhodpsi + dflx_norm(iz))

                call multiply_by_rho(g0k)

                g2(:,:,iz,it,ivmu) = g2(:,:,iz,it,ivmu) + g0k
              endif

              !subtract adiabatic contribution part of g
              g0k = -spec(is)%zt*fphi*phi(:,:,iz,it)*aj0x(:,:,iz,ivmu)*aj1x(:,:,iz,ivmu) &
                    *maxwell_vpa(iv,is)*maxwell_mu(ia,iz,imu,is)*maxwell_fac(is) &
                    *zi*spread(aky,2,nakx)*vperp2(ia,iz,imu)*geo_surf%rhoc &
                    * (gds21(ia,iz)+theta0*gds22(ia,iz))*spec(is)%smz &
                      / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2)

              if(radial_variation) then
                g1k = -spec(is)%zt*fphi*phi(:,:,iz,it)*aj0x(:,:,iz,ivmu)*aj1x(:,:,iz,ivmu) &
                      *maxwell_vpa(iv,is)*maxwell_mu(ia,iz,imu,is)*maxwell_fac(is) &
                      *zi*spread(aky,2,nakx)*vperp2(ia,iz,imu)*geo_surf%rhoc &
                      * (dgds21dr(ia,iz)+theta0*dgds22dr(ia,iz))*spec(is)%smz &
                      / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2)

                g1k = g1k -spec(is)%zt*fphi*phi(:,:,iz,it)*aj0x(:,:,iz,ivmu) &
                       *maxwell_vpa(iv,is)*maxwell_mu(ia,iz,imu,is)*maxwell_fac(is) &
                       *zi*spread(aky,2,nakx)*vperp2(ia,iz,imu)*geo_surf%rhoc &
                       * (gds21(ia,iz)+theta0*gds22(ia,iz))*spec(is)%smz &
                       / (geo_surf%qinp*geo_surf%shat*bmag(ia,iz)**2) &
                            * (0.5*aj0x(:,:,iz,ivmu) - aj1x(:,:,iz,ivmu))  &
                            * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz))

                g1k = g1k + &
                      g0k*(-spec(is)%tprim*(vpa(iv)**2+vperp2(ia,iz,imu) - 2.5) &
                           -spec(is)%fprim - 2.0*dBdrho(iz)*mu(imu)  &
                           -0.5*aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                              * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                              * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                           - geo_surf%d2qdr2*geo_surf%rhoc/(geo_surf%shat*geo_surf%qinp) &
                           - geo_surf%d2psidr2*drhodpsi + dflx_norm(iz))

                call multiply_by_rho(g1k)

                g0k = g0k + g1k
              endif
              g2(:,:,iz,it,ivmu) = g2(:,:,iz,it,ivmu) + g0k
          enddo
        enddo
      enddo

      g1 = g1 + g2
      call get_one_flux_vmulo (flx_norm,spec%dens_psi0*sqrt(spec%mass*spec%temp_psi0), g1, phi, vflx)

      if(write_radial_fluxes) then
        call get_one_flux_radial (flx_norm, spec%dens_psi0*sqrt(spec%mass*spec%temp_psi0), g1, phi, vflx_x)
      endif

    end if

    ! normalise to account for contributions from multiple flux tubes
    ! in flux tube train
    pflx = pflx/real(ntubes)
    qflx = qflx/real(ntubes)
    vflx = vflx/real(ntubes)

    deallocate (flx_norm)
    if(allocated(dflx_norm)) deallocate(dflx_norm)
    if(allocated(g0k)) deallocate(g0k)
    if(allocated(g1k)) deallocate(g1k)

  end subroutine get_fluxes_vmulo

  !==============================================
  !=============== GET ONE FLUX =================
  !==============================================
  subroutine get_one_flux (iky, iz, norm, gin, fld, flxout)

    use vpamu_grids, only: integrate_vmu
    use volume_averages, only: mode_fac
    use kt_grids, only: aky

    implicit none

    integer, intent (in) :: iky, iz
    real, intent (in) :: norm
    complex, dimension (:,:), intent (in) :: gin
    complex, intent (in) :: fld
    real, intent (in out) :: flxout

    complex :: flx

    call integrate_vmu (gin,iz,flx)
    flxout = flxout &
         + 0.5*mode_fac(iky)*aky(iky)*aimag(flx*conjg(fld))*norm

  end subroutine get_one_flux

  !==============================================
  !============ GET ONE FLUX VMULO ==============
  !==============================================
  subroutine get_one_flux_vmulo (norm, weights, gin, fld, flxout)

    use vpamu_grids, only: integrate_vmu
    use stella_layouts, only: vmu_lo
    use kt_grids, only: aky, nakx, naky
    use zgrid, only: nzgrid, ntubes
    use species, only: nspec
    use volume_averages, only: mode_fac

    implicit none

    real, dimension (-nzgrid:), intent (in) :: norm
    real, dimension (:), intent (in) :: weights
    complex, dimension (:,:,-nzgrid:,:,vmu_lo%llim_proc:), intent (in) :: gin
    complex, dimension (:,:,-nzgrid:,:), intent (in) :: fld
    real, dimension (:), intent (in out) :: flxout

    complex, dimension (:,:,:,:,:), allocatable :: totals

    integer :: is, it, iz, ikx

    allocate (totals(naky,nakx,-nzgrid:nzgrid,ntubes,nspec))

    call integrate_vmu (gin,weights,totals)
    do is = 1, nspec
      do it = 1, ntubes
        do iz= -nzgrid, nzgrid
          do ikx = 1, nakx
            flxout(is) = flxout(is) &
              + sum(0.5*mode_fac*aky*aimag(totals(:,ikx,iz,it,is)*conjg(fld(:,ikx,iz,it)))*norm(iz))
          enddo
        enddo
      enddo
    enddo

    deallocate (totals)

  end subroutine get_one_flux_vmulo

  !==============================================
  !=========== GET ONE FLUX RADIAL ==============
  !==============================================
  subroutine get_one_flux_radial (norm, weights, gin, fld, flxout)

    use vpamu_grids, only: integrate_vmu
    use stella_layouts, only: vmu_lo
    use kt_grids, only: aky, nakx, naky
    use zgrid, only: nzgrid, ntubes
    use species, only: nspec
    use volume_averages, only: mode_fac

    use stella_transforms, only: transform_kx2x_unpadded

    implicit none

    real, dimension (-nzgrid:), intent (in) :: norm
    real, dimension (:), intent (in) :: weights
    complex, dimension (:,:,-nzgrid:,:,vmu_lo%llim_proc:), intent (in) :: gin
    complex, dimension (:,:,-nzgrid:,:), intent (in) :: fld
    real, dimension (:,:), intent (in out) :: flxout

    complex, dimension (:,:,:,:,:), allocatable :: totals

    complex, dimension (:,:), allocatable :: g0x, g1x

    integer :: is, it, iz, ikx

    allocate (g0x(naky,nakx))
    allocate (g1x(naky,nakx))

    allocate (totals(naky,nakx,-nzgrid:nzgrid,ntubes,nspec))

    call integrate_vmu (gin,weights,totals)
    do is = 1, nspec
      do it = 1, ntubes
        do iz= -nzgrid, nzgrid
          call transform_kx2x_unpadded(totals(:,:,iz,it,is),g0x)
          call transform_kx2x_unpadded(fld(:,:,iz,it)      ,g1x)
          do ikx = 1, nakx
            flxout(ikx,is) = flxout(ikx,is) &
              + sum(0.5*mode_fac*aky*aimag(g0x(:,ikx)*conjg(g1x(:,ikx)))*norm(iz))
          enddo
        enddo
      enddo
    enddo

    deallocate (totals)

  end subroutine get_one_flux_radial

  !==============================================
  !=============== GET MOMENTS ==================
  !==============================================
  subroutine get_moments (g, dens, upar, temp)

    use zgrid, only: nzgrid, ntubes
    use species, only: spec
    use vpamu_grids, only: integrate_vmu
    use vpamu_grids, only: vpa, vperp2, mu
    use vpamu_grids, only: maxwell_mu, ztmax, maxwell_fac
    use kt_grids, only: naky, nakx, multiply_by_rho
    use stella_layouts, only: vmu_lo
    use stella_layouts, only: iv_idx, imu_idx, is_idx
    use dist_fn_arrays, only: g1, g2, kperp2, dkperp2dr
    use stella_geometry, only: bmag, dBdrho
    use gyro_averages, only: aj0x, aj1x, gyro_average
    use fields_arrays, only: phi, phi_corr_QN
    use run_parameters, only: fphi
    use physics_flags, only: radial_variation

    implicit none

    complex, dimension (:,:,-nzgrid:,:,vmu_lo%llim_proc:), intent (in) :: g
    complex, dimension (:,:,:,:,:), intent (out) :: dens, upar, temp

    complex, dimension (:,:), allocatable :: g0k
    real :: zero

    integer :: ivmu, iv, imu, is, ia
    integer :: iz, it

    if(radial_variation) then
      allocate (g0k(naky,nakx))
    endif

    ! Hack below. Works since J0^2 - 1 and its derivative are zero at the origin
    zero = 100.*epsilon(0.)

    ! h is gyrophase independent, but is in gyrocenter coordinates,
    ! so requires a J_0 to get to particle coordinates
    ! <f>_r = h J_0 - Ze*phi/T * F0
    ! g     = h     - Ze*<phi>_R/T * F0
    ! <f>_r = g J_0 + Ze*(J_0<phi>_R-phi)/T * F0

    ia = 1
    do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
       iv = iv_idx(vmu_lo,ivmu)
       imu = imu_idx(vmu_lo,ivmu)
       is = is_idx(vmu_lo,ivmu)
       call gyro_average (g(:,:,:,:,ivmu), ivmu, g1(:,:,:,:,ivmu))
       ! FLAG -- AJ0X NEEDS DEALING WITH BELOW
       g2(:,:,:,:,ivmu) = g1(:,:,:,:,ivmu) + ztmax(iv,is) &
            * spread(spread(spread(maxwell_mu(ia,:,imu,is),1,naky),2,nakx) &
            * maxwell_fac(is)*(aj0x(:,:,:,ivmu)**2-1.0),4,ntubes)*fphi*phi

       if(radial_variation) then
         do it = 1, ntubes
           do iz= -nzgrid, nzgrid
             !phi
             g0k = ztmax(iv,is)*maxwell_mu(ia,iz,imu,is) &
                * maxwell_fac(is)*(aj0x(:,:,iz,ivmu)**2-1.0)*fphi*phi(:,:,iz,it) &
                *(-spec(is)%tprim*(vpa(iv)**2+vperp2(ia,iz,imu)-2.5) &
                  -spec(is)%fprim+(dBdrho(iz)/bmag(ia,iz))*(1.0 - 2.0*mu(imu)*bmag(ia,iz)) &
                  -aj1x(:,:,iz,ivmu)*aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                  * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                  * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                     /(aj0x(:,:,iz,ivmu)**2 - 1.0 + zero))

             !g
             g0k = g0k + g1(:,:,iz,it,ivmu) &
               * (-0.5*aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
               * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
               * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
               + dBdrho(iz)/bmag(ia,iz))

             call multiply_by_rho(g0k)

            ! g0k(1,1) = 0.0

             !phi QN
             g0k = g0k + ztmax(iv,is)*maxwell_mu(ia,iz,imu,is)*fphi &
                 * maxwell_fac(is)*(aj0x(:,:,iz,ivmu)**2-1.0)*phi_corr_QN(:,:,iz,it)

             g2(:,:,iz,it,ivmu) = g2(:,:,iz,it,ivmu) + g0k
           enddo
         enddo
       endif
    end do
    call integrate_vmu (g2, spec%dens_psi0, dens)

    do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
       iv = iv_idx(vmu_lo,ivmu)
       imu = imu_idx(vmu_lo,ivmu)
       is = is_idx(vmu_lo,ivmu)
       g2(:,:,:,:,ivmu) = (g1(:,:,:,:,ivmu) + ztmax(iv,is) &
            * spread(spread(spread(maxwell_mu(ia,:,imu,is),1,naky),2,nakx) &
            * maxwell_fac(is)*(aj0x(:,:,:,ivmu)**2-1.0),4,ntubes)*phi*fphi) &
            *(vpa(iv)**2+spread(spread(spread(vperp2(1,:,imu),1,naky),2,nakx),4,ntubes))/1.5
       if(radial_variation) then
         do it = 1, ntubes
           do iz= -nzgrid, nzgrid
             !phi
             g0k = ztmax(iv,is)*maxwell_mu(ia,iz,imu,is)*maxwell_fac(is) &
               *(vpa(iv)**2 + vperp2(ia,iz,imu))/1.5 &
               *(aj0x(:,:,iz,ivmu)**2-1.0)*phi(:,:,iz,it)*fphi &
               *(-spec(is)%tprim*(vpa(iv)**2+vperp2(ia,iz,imu)-2.5) &
                 -spec(is)%fprim+(dBdrho(iz)/bmag(ia,iz))*(1.0 - 2.0*mu(imu)*bmag(ia,iz)) &
                 -aj1x(:,:,iz,ivmu)*aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
                 * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
                 * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                    /(aj0x(:,:,iz,ivmu)**2 - 1.0 + zero) &
                 + 2.0*mu(imu)*dBdrho(iz)/(vpa(iv)**2+vperp2(ia,iz,imu)))

             !g
             g0k = g0k + g1(:,:,iz,it,ivmu)*(vpa(iv)**2+vperp2(ia,iz,imu))/1.5 &
               * (-0.5*aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
               * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
               * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
                 + dBdrho(iz)/bmag(ia,iz) &
                 + 2.0*mu(imu)*dBdrho(iz)/(vpa(iv)**2+vperp2(ia,iz,imu)))

             call multiply_by_rho(g0k)

             !phi QN
             g0k = g0k + fphi*ztmax(iv,is)*maxwell_mu(ia,iz,imu,is) &
               * (vpa(iv)**2 + vperp2(ia,iz,imu))/1.5 &
               * maxwell_fac(is)*(aj0x(:,:,iz,ivmu)**2-1.0)*phi_corr_QN(:,:,iz,it)

             g2(:,:,iz,it,ivmu) = g2(:,:,iz,it,ivmu) + g0k
           enddo
         enddo
       endif
    end do
    ! integrate to get dTs/Tr
!    call integrate_vmu (g2, spec%temp, temp)
    call integrate_vmu (g2, spec%temp_psi0*spec%dens_psi0, temp)

    do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
       iv = iv_idx(vmu_lo,ivmu)
       imu = imu_idx(vmu_lo,ivmu)
       is = is_idx(vmu_lo,ivmu)
       g2(:,:,:,:,ivmu) = vpa(iv)*g1(:,:,:,:,ivmu)
       if(radial_variation) then
         do it = 1, ntubes
           do iz= -nzgrid, nzgrid
             !g
             g0k = vpa(iv)*g1(:,:,iz,it,ivmu) &
               * (-0.5*aj1x(:,:,iz,ivmu)/aj0x(:,:,iz,ivmu)*(spec(is)%smz)**2 &
               * (kperp2(:,:,ia,iz)*vperp2(ia,iz,imu)/bmag(ia,iz)**2) &
               * (dkperp2dr(:,:,ia,iz) - dBdrho(iz)/bmag(ia,iz)) &
               + dBdrho(iz)/bmag(ia,iz))

             call multiply_by_rho(g0k)

             g2(:,:,iz,it,ivmu) = g2(:,:,iz,it,ivmu) + g0k
           enddo
         enddo
       endif
    end do
    call integrate_vmu (g2, spec%stm_psi0, upar)

    if(allocated(g0k)) deallocate(g0k)

  end subroutine get_moments

  !==============================================
  !================ GET GVMUS ===================
  !==============================================
  ! get_gvmus takes g(kx,ky,z) and returns average over z of int dxdy g(x,y,z)^2
  ! SHOULD MODIFY TO TAKE ADVANTAGE OF FACT THAT G(KY,KX,Z) LOCAL IS AVAILABLE
  subroutine get_gvmus (g, gv)

    use mp, only: nproc, sum_reduce
    use stella_layouts, only: kxkyz_lo
    use stella_layouts, only: is_idx, iky_idx, iz_idx
    use zgrid, only: ntubes
    use vpamu_grids, only: nvpa, nmu
    use stella_geometry, only: dl_over_b
    use volume_averages, only: mode_fac

    implicit none

    complex, dimension (:,:,kxkyz_lo%llim_proc:), intent (in) :: g
    real, dimension (:,:,:), intent (out) :: gv

    integer :: ikxkyz, iv, is, imu, iz, iky, ia!, ivp

    ia = 1

    ! when doing volume averages, note the following:
    ! int dxdy g(x,y)^2 = sum_ky |g(ky=0,kx)|^2 + 2 * sum_{kx,ky} |g(ky>0,kx)|^2
    ! factor of 2 accounted for in fac

    gv = 0.
    do ikxkyz = kxkyz_lo%llim_proc, kxkyz_lo%ulim_proc
       is = is_idx(kxkyz_lo,ikxkyz)
       iky = iky_idx(kxkyz_lo,ikxkyz)
       iz = iz_idx(kxkyz_lo,ikxkyz)
       do imu = 1, nmu
          do iv = 1, nvpa
             gv(iv,imu,is) = gv(iv,imu,is) + real(g(iv,imu,ikxkyz)*conjg(g(iv,imu,ikxkyz)))*mode_fac(iky)*dl_over_b(ia,iz)
          end do
       end do
    end do
    gv = gv/real(ntubes)

    if (nproc > 1) call sum_reduce (gv,0)

  end subroutine get_gvmus

  !==============================================
  !================= GET GVZS ===================
  !==============================================
  ! get_gzvs takes g(kx,ky,z,vpa,mu,s) and returns int dmudxdy g(x,y,z,vpa,mu,s)^2
  subroutine get_gzvs (g, gz)

    use stella_layouts, only: vmu_lo
    use zgrid, only: nzgrid, ntubes
    use vpamu_grids, only: integrate_mu
    use kt_grids, only: nakx, naky
    use volume_averages, only: mode_fac

    implicit none

    complex, dimension (:,:,-nzgrid:,:,vmu_lo%llim_proc:), intent (in) :: g
    real, dimension (:,:,:,:), intent (out) :: gz

    integer :: ivmu, iz, it, ikx, iky, izp

    real, dimension (:,:,:), allocatable :: gtmp

    allocate (gtmp(-nzgrid:nzgrid,ntubes,vmu_lo%llim_proc:vmu_lo%ulim_alloc))

    ! when doing volume averages, note the following:
    ! int dxdy g(x,y)^2 = sum_kx |g(kx,ky=0)|^2 + 2 * sum_{kx,ky} |g(kx,ky>0)|^2
    ! factor of 2 accounted for in mode_fac

    gtmp = 0.
    do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
       do ikx = 1, nakx
          do iky = 1, naky
             gtmp(:,:,ivmu) = gtmp(:,:,ivmu) + real(g(iky,ikx,:,:,ivmu)*conjg(g(iky,ikx,:,:,ivmu)))*mode_fac(iky)
          end do
       end do
    end do

    do it = 1, ntubes
       do iz = -nzgrid, nzgrid
          izp = iz+nzgrid+1
          call integrate_mu (iz, gtmp(iz,it,:), gz(it,izp,:,:))
       end do
    end do

    deallocate (gtmp)

  end subroutine get_gzvs

  !==============================================
  !======== FINISH STELLA DIAGNOSTIC ============
  !==============================================
  subroutine finish_stella_diagnostics(istep)

    use mp, only: proc0
    use redistribute, only: scatter
    use stella_io, only: finish_stella_io
    use stella_time, only: code_dt, code_time
    use stella_save, only: stella_save_for_restart
    use dist_redistribute, only: kxkyz2vmu
    use dist_fn_arrays, only: gnew, gvmu

    implicit none

    integer :: istatus
    integer, intent (in) :: istep

    if (proc0) then
       call write_final_ascii_files
       call close_loop_ascii_files
    end if
    if (save_for_restart) then
        call scatter (kxkyz2vmu, gnew, gvmu)
        call stella_save_for_restart (gvmu, istep, code_time, code_dt, istatus, .true.)
    end if
    call finish_stella_io
    call deallocate_arrays

    nout = 1
    diagnostics_initialized = .false.

  end subroutine finish_stella_diagnostics

  !==============================================
  !========= WRITE LOOP ASCII FILES =============
  !==============================================
  subroutine write_loop_ascii_files (istep, phi2, apar2, pflx, vflx, qflx, om, om_avg)

    use stella_time, only: code_time
    use species, only: nspec
    use kt_grids, only: naky, nakx
    use kt_grids, only: aky, akx

    implicit none

    integer, intent (in) :: istep
    real, intent (in) :: phi2, apar2
    real, dimension (:), intent (in) :: pflx, vflx, qflx
    complex, dimension (:,:), intent (in) :: om, om_avg

    character (3) :: nspec_str
    character (100) :: str
    integer :: ikx, iky

    write (stdout_unit,'(a7,i7,a6,e12.4,a10,e12.4,a11,e12.4)') 'istep=', istep, &
         'time=', code_time, '|phi|^2=', phi2, '|apar|^2= ', apar2

    call flush(stdout_unit)

    write (nspec_str,'(i3)') 3*nspec+1
    str = trim('('//trim(nspec_str)//'es15.4e3)')
    write (fluxes_unit,str) code_time, pflx, vflx, qflx

    call flush(stdout_unit)
    call flush(fluxes_unit)

    if (write_omega .and. istep > 0) then
       do iky = 1, naky
          do ikx = 1, nakx
             write (omega_unit,'(7e16.8)') code_time, aky(iky), akx(ikx),&
                  real(om(iky,ikx)), aimag(om(iky,ikx)), &
                  real(om_avg(iky,ikx)), aimag(om_avg(iky,ikx))
          end do
          if (nakx > 1) write (omega_unit,*)
       end do
       if (naky > 1) write (omega_unit,*)
       call flush(omega_unit)
    end if

  end subroutine write_loop_ascii_files

  !==============================================
  !========= WRITE FINAL ASCII FILES ============
  !==============================================
  subroutine write_final_ascii_files

    use file_utils, only: open_output_file, close_output_file
    use fields_arrays, only: phi, apar
    use zgrid, only: nzgrid, ntubes
    use zgrid, only: zed
    use kt_grids, only: naky, nakx
    use kt_grids, only: aky, akx, zed0
    use stella_geometry, only: zed_eqarc
    USE dist_fn_arrays, ONLY: kperp2

    implicit none

    integer :: tmpunit
    integer :: iky, ikx, iz, it

    call open_output_file (tmpunit,'.final_fields')
    write (tmpunit,'(10a14)') '# z', 'z-zed0', 'aky', 'akx', &
         'real(phi)', 'imag(phi)', 'real(apar)', 'imag(apar)', &
         'z_eqarc-zed0', 'kperp2'
    do iky = 1, naky
       do ikx = 1, nakx
          do it = 1, ntubes
             do iz = -nzgrid, nzgrid
                write (tmpunit,'(10es15.4e3,i3)') zed(iz), zed(iz)-zed0(iky,ikx), aky(iky), akx(ikx), &
                  real(phi(iky,ikx,iz,it)), aimag(phi(iky,ikx,iz,it)), &
                  real(apar(iky,ikx,iz,it)), aimag(apar(iky,ikx,iz,it)), zed_eqarc(iz)-zed0(iky,ikx), &
                  kperp2(iky,ikx,it,iz), it
             end do
             write (tmpunit,*)
          end do
       end do
    end do
    call close_output_file (tmpunit)

  end subroutine write_final_ascii_files

  !==============================================
  !============ DEALLCOATE ARRAYS ===============
  !==============================================
  subroutine deallocate_arrays

    implicit none

    if (allocated(pflux)) deallocate (pflux)
    if (allocated(qflux)) deallocate (qflux)
    if (allocated(vflux)) deallocate (vflux)
    if (allocated(exchange)) deallocate (exchange)
    if (allocated(pflux_avg)) deallocate (pflux_avg)
    if (allocated(qflux_avg)) deallocate (qflux_avg)
    if (allocated(vflux_avg)) deallocate (vflux_avg)
    if (allocated(heat_avg)) deallocate (heat_avg)
    if (allocated(omega_vs_time)) deallocate (omega_vs_time)

  end subroutine deallocate_arrays

end module stella_diagnostics


text:
#include "rundeck_opts.h"

      SUBROUTINE CONDSE
!@sum   CONDSE driver for moist convection AND large-scale condensation
!@auth  M.S.Yao/A. Del Genio (modularisation by Gavin Schmidt)
!@ver   1.0 (taken from CB265)
!@calls CLOUDS:MSTCNV,CLOUDS:LSCOND

      USE CONSTANT, only : bygrav,lhm,rgas,grav,tf,lhe,lhs,sha,deltx
     *     ,teeny
      USE MODEL_COM, only : im,jm,lm,p,u,v,t,q,wm,JHOUR,fearth
     *     ,ls1,psf,ptop,dsig,bydsig,jeq,sig,DTsrc,ftype
     *     ,ntype,itime,fim,airx,lmc,focean,fland,flice
      USE QUSDEF, only : nmom
      USE SOMTQ_COM, only : t3mom=>tmom,q3mom=>qmom
      USE GEOM, only : bydxyp,dxyp,imaxj,kmaxj,ravj,idij,idjj,dyp
      USE RANDOM
      USE CLOUDS_COM, only : ttold,qtold,svlhx,svlat,rhsav,cldsav
     *     ,pbltop,tauss,taumc,cldss,cldmc,csizmc,csizss
      USE DAGCOM, only : aj,areg,aij,ajl,ail,adiurn,jreg,ij_pscld,
     *     ij_pdcld,ij_scnvfrq,ij_dcnvfrq,ij_wmsum,ij_snwf,ij_prec,
     *     ij_neth,ij_f0oc,j_eprcp,j_prcpmc,j_prcpss,il_mceq,j5s,j5n,
     *     ijdd,idd_pr,idd_ecnd,idd_mcp,idd_dmc,idd_smc,idd_ssp,
     &     jl_mcmflx,jl_sshr,jl_mchr,jl_dammc,jl_rhe,jl_mchphas
     *     ,jl_mcdtotw,jl_mcldht,jl_mcheat,jl_mcdry,ij_ctpi,ij_taui
     *     ,ij_lcldi,ij_mcldi,ij_hcldi,ij_tcldi,ij_sstabx,isccp_diags
     *     ,ndiupt,jl_cldmc,jl_cldss,jl_csizmc,jl_csizss,
     *     ntau,npres,aisccp,isccp_reg
#ifdef TRACERS_ON
      USE TRACER_COM, only: itime_tr0,TRM,TRMOM,NTM
#ifdef TRACERS_WATER
     *     ,trwm,trw0
#endif
      USE TRACER_DIAG_COM,only: tajln,jlnt_mc,jlnt_lscond,itcon_mc
     *     ,itcon_ss
#ifdef TRACERS_WATER
     *     ,jls_source,taijn,tajls,tij_prec
#endif
      USE CLOUDS, only : tm,tmom ! local  (i,j)
     *     ,ntx,ntix              ! global (same for all i,j)
#ifdef TRACERS_WATER
     *     ,trwml,trsvwml,trprmc,trprss
#endif
#ifdef TRACERS_SPECIAL_Shindell
      USE LIGHTNING, only : i_lgt,j_lgt,RNOx_lgt
#endif
#endif
      USE CLOUDS, only : BYDTsrc,mstcnv,lscond ! glb var & subroutines
     *     ,airm,byam,etal,sm,smom,qm,qmom,isc,dxypj,lp50
     *     ,tl,ris,ri1,ri2,mcflx,sshr,dgdsm,dphase,dtotw,dqcond,dctei
     *     ,wml,sdl,u_0,v_0,um,vm,qs,us,vs,dcl,airxl,prcpss,hcndss
     *     ,prcpmc,pearth,ts,taumcl,cldmcl,svwmxl,svlatl,svlhxl,dgdqm
     *     ,cldslwij,clddepij,csizel,precnvl,vsubl,lmcmax,lmcmin,wmsum
     *     ,aq,dpdt,th,ql,wmx,ttoldl,rh,taussl,cldssl,cldsavl,rh1
     *     ,kmax,ra,pl,ple,plk,rndssl,lhp,debug
      USE PBLCOM, only : tsavg,qsavg,usavg,vsavg,tgvavg,qgavg,dclev
      USE DYNAMICS, only : pk,pek,pmid,pedn,sd_clouds,gz,ptold,pdsig
     *     ,ltropo,dke
      USE SEAICE_COM, only : rsi
      USE GHYCOM, only : snoage
      USE LAKES_COM, only : flake
      USE FLUXES, only : prec,eprec,precss,gtemp
#ifdef TRACERS_WATER
     *     ,trprec
#endif
      IMPLICIT NONE

#ifdef TRACERS_ON
!@var tmsave holds tracer value (for diagnostics)
      REAL*8 tmsave(lm,ntm),dtr_mc(jm,ntm),dtr_ss(jm,ntm)
      INTEGER NX
#endif

!@var UC,VC,UZM,VZM velocity work arrays
      REAL*8, DIMENSION(IM,JM,LM) :: UC,VC
      REAL*8, DIMENSION(2,LM) :: UZM,VZM

!@param ENTCON fractional rate of entrainment (km**-1)
      REAL*8,  PARAMETER :: ENTCON = .2d0

      INTEGER I,J,K,L,N  !@var I,J,K,L,N loop variables
      INTEGER JR,KR,ITYPE,IT,IH,LP850,LP600
!@var JR = JREG(I,J)
!@var KR index for regional diagnostics
!@var ITYPE index for snow age
!@var IT index for surface types
!@var LP850 layer near 850 mb
!@var LP600 layer near 600 mb
!@var LERR,IERR error reporting
      INTEGER :: LERR, IERR
      INTEGER, DIMENSION(IM) :: IDI,IDJ    !@var ID

      REAL*8 :: HCNDMC,PRCP,TPRCP,EPRCP,ENRGP,WMERR,ALPHA1,ALPHA2,ALPHAS
      REAL*8 :: DTDZ,DTDZS,DUDZ,DVDZ,DUDZS,DVDZS,THSV,THV1,THV2,QG,TGV
      REAL*8 :: DH1S,BYDH1S,DH12,BYDH12,DTDZG,DUDZG,DVDZG,SSTAB,DIFT,CSC
     *     ,E,E1,ep
!@var HCNDMC heating due to moist convection
!@var PRCP precipitation
!@var TPRCP temperature of mc. precip  (deg. C)
!@var EPRCP sensible heat of precip
!@var ENRGP total energy of precip
!@var WMERR DH12,BYDH12,DH1S,BYDH1S,SSTAB dummy variable
!@var THSV,THV1,THV2 vertual potential temperatures
!@var QG,TGV ground humidity,virt.temperature from pbl
!@var ALPHA1,ALPHA2,ALPHAS,DIFT,CSC dummy variables
!@var DTDZ,DTDZS,DTDZG vertical potential temperature gradients
!@var DUDZ,DVDZ,DUDZS,DVDZS,DUDZG,DVDZG vertical wind gradients

C**** parameters and variables for isccp diags
      real*8, parameter :: bywc = 1./2.56d0 , byic= 1./2.13d0
      real*8 skt,conv(lm),qv(lm)
      real*8 pfull(lm),at(lm),cc(lm),dtau_s(lm),dtau_c(lm)
      real*8 dem_s(lm),dem_c(lm),phalf(lm+1)
      real*8 fq_isccp(ntau,npres),ctp,tauopt
      integer itau,itrop,nbox
C****

C
Cred*                       Reduced Arrays 1                 *********
C        not clear yet whether they still speed things up
      REAL*8  GZIL(IM,LM), SD_CLDIL(IM,LM), WMIL(IM,LM)
      REAL*8  TMOMIL(NMOM,IM,LM),  QMOMIL(NMOM,IM,LM)
Cred*                   end Reduced Arrays 1
      INTEGER ICKERR, JCKERR, JERR, seed, NR
      REAL*8  RNDSS(3,LM,IM,JM),xx
      REAL*8  AJEQIL(J5N-J5S+1,IM,JM), AREGIJ(IM,JM,3)
      REAL*8  UKP1(IM,LM), VKP1(IM,LM), UKPJM(IM,LM),VKPJM(IM,LM)
      REAL*8  UKM(4,IM,2:JM-1,LM), VKM(4,IM,2:JM-1,LM)

C
C     OBTAIN RANDOM NUMBERS FOR PARALLEL REGION
C
      DO J=1,JM
      DO I=1,IMAXJ(J)
        DO L=LP50,1,-1
          DO NR=1,3
            RNDSS(NR,L,I,J) = RANDU(xx)
          END DO
        END DO
C     Do not bother to save random numbers for isccp_clouds
      END DO
      END DO
C     But save the current seed in case isccp_routine is activated
      if (isccp_diags.eq.1) CALL RFINAL(seed)
C
C**** SAVE UC AND VC, AND ZERO OUT CLDSS AND CLDMC
      UC=U
      VC=V
C**** COMPUTE ZONAL MEAN U AND V AT POLES
      DO L=1,LM
        UZM(1,L)=0.
        UZM(2,L)=0.
        VZM(1,L)=0.
        VZM(2,L)=0.
      ENDDO
      DO L=1,LM
        DO I=1,IM
          UZM(1,L)=UZM(1,L)+UC(I,2,L)
          UZM(2,L)=UZM(2,L)+UC(I,JM,L)
          VZM(1,L)=VZM(1,L)+VC(I,2,L)
          VZM(2,L)=VZM(2,L)+VC(I,JM,L)
        ENDDO
        UZM(1,L)=UZM(1,L)/FIM
        UZM(2,L)=UZM(2,L)/FIM
        VZM(1,L)=VZM(1,L)/FIM
        VZM(2,L)=VZM(2,L)/FIM
      ENDDO
      IH=JHOUR+1
#ifdef TRACERS_ON
C**** Find the ntx active tracers ntix(1->ntx)
      nx = 0
      do n=1,ntm
        if (itime.lt.itime_tr0(n)) cycle
        nx = nx+1
        ntix(nx) = n
      end do
      ntx = nx
#endif
C****
C**** MAIN J LOOP
C****
       ICKERR=0
       JCKERR=0
!$OMP  PARALLEL DO PRIVATE (
#ifdef TRACERS_ON
!$OMP*  NX,tmsave,
#endif
!$OMP*  ALPHAS,ALPHA1,ALPHA2,AT,BYDH1S,BYDH12, CC,CONV,CTP,
!$OMP*  DH1S,DH12,DTDZ,DTDZG,DTDZS,DUDZ,DUDZG,DUDZS,DVDZ,DVDZG,DVDZS,
!$OMP*  DTAU_S,DTAU_C,DEM_S,DEM_C, FQ_ISCCP, ENRGP,EPRCP,
!$OMP*  HCNDMC, I,ITYPE,IT,ITAU, IDI,IDJ,
!$OMP*  ITROP,IERR, J,JERR, K,KR, L,LERR, N,NBOX, PRCP,PFULL,PHALF,
!$OMP*  GZIL, SD_CLDIL, WMIL, TMOMIL, QMOMIL,        ! reduced arrays
!$OMP*  QG,QV, SKT,SSTAB, TGV,TPRCP,THSV,THV1,THV2,TAUOPT, WMERR,
!$OMP*  LP600,LP850,CSC,DIFT, E,E1,ep)
!$OMP*    SCHEDULE(DYNAMIC,2)
!$OMP*    REDUCTION(+:ICKERR,JCKERR)
C
      DO J=1,JM
C
Cred* Reduced Arrays 2
C
        DXYPJ=DXYP(J)
      DO L=1,LM
         GZIL(:,L) = GZ(:,J,L)
      END DO
      DO L=1,LM
         SD_CLDIL(:,L) = SD_CLOUDS(:,J,L)
      END DO
      DO L=1,LM
         WMIL(:,L) = WM(:,J,L)
      END DO
      DO L=1,LM
         TMOMIL(:,:,L) = T3MOM(:,:,J,L)
      END DO
      DO L=1,LM
         QMOMIL(:,:,L) = Q3MOM(:,:,J,L)
      END DO
Cred* end Reduced Arrays 2
#ifdef TRACERS_ON
      dtr_mc(j,:)=0. ; dtr_ss(j,:)=0.
#endif
      kmax = kmaxj(j)
C****
C**** MAIN I LOOP
C****
      DO I=1,IMAXJ(J)
cc       JR=JREG(I,J)  ! summing done outside parallel region
C****
C**** SET UP VERTICAL ARRAYS, OMITTING THE J AND I SUBSCRIPTS
C****
      DEBUG = .FALSE.   ! use for individual box diags in clouds
      PEARTH=FEARTH(I,J)
      TS=TSAVG(I,J)
      QS=QSAVG(I,J)
      US=USAVG(I,J)
      VS=VSAVG(I,J)
      TGV=TGVAVG(I,J)
      QG=QGAVG(I,J)
!!!   DCL=NINT(DCLEV(I,J))   ! prevented by openMP bug
      DCL=INT(DCLEV(I,J)+.5)

      DO K=1,KMAX
        RA(K)=RAVJ(K,J)
        IDI(K)=IDIJ(K,I,J)
        IDJ(K)=IDJJ(K,J)
      END DO
C**** PRESSURES, AND PRESSURE TO THE KAPA
      PL(:) =PMID(:,I,J)
      PLE(:)=PEDN(:,I,J)
      PLK(:)=PK(:,I,J)
      AIRM(:)=PDSIG(:,I,J)
      BYAM(:)=1./AIRM(:)
C**** other fields where L is the leading index
      SVLHXL(:)=SVLHX(:,I,J)
      TTOLDL(:)=TTOLD(:,I,J)
      CLDSAVL(:)=CLDSAV(:,I,J)
      RH(:)=RHSAV(:,I,J)
      DPDT(1:LS1-1)=SIG(1:LS1-1)*(P(I,J)-PTOLD(I,J))*BYDTsrc
      DPDT(LS1:LM)=0.
      DO L=1,LM
C**** TEMPERATURES
        SM(L)  =T(I,J,L)*AIRM(L)
Cred    SMOM(:,L) =T3MOM(:,I,J,L)*AIRM(L)
        SMOM(:,L) =TMOMIL(:,I,L)*AIRM(L)
        TL(L)=T(I,J,L)*PLK(L)
C**** MOISTURE (SPECIFIC HUMIDITY)
        QM(L)  =Q(I,J,L)*AIRM(L)
Cred    QMOM(:,L) =Q3MOM(:,I,J,L)*AIRM(L)
        QMOM(:,L) =QMOMIL(:,I,L)*AIRM(L)
Cred    WML(L)=WM(I,J,L)
        WML(L)=WMIL(I,L)
C**** others
Cred    SDL(L)=SD_CLOUDS(I,J,L)*BYDXYP(J)
        SDL(L)=SD_CLDIL(I,L)*BYDXYP(J)
        IF(L.LE.LM-2)
Cred *    ETAL(L+1)=.5*ENTCON*(GZ(I,J,L+2)-GZ(I,J,L))*1.d-3*BYGRAV
     *    ETAL(L+1)=.5*ENTCON*(GZIL(I,L+2)-GZIL(I,L))*1.d-3*BYGRAV
      END DO
      ETAL(LM)=ETAL(LM-1)
      ETAL(1)=0.     ! not used
#ifdef TRACERS_ON
C**** TRACERS: Use only the active ones
      do nx=1,ntx
      do l=1,lm
        tm(l,nx) = trm(i,j,l,ntix(nx))
        tmom(:,l,nx) = trmom(:,i,j,l,ntix(nx))
      end do
      end do
#endif
C**** SURROUNDING WINDS
      DO L=1,LM
        DO K=1,KMAX
          U_0(K,L) = UC(IDI(K),IDJ(K),L)
          V_0(K,L) = VC(IDI(K),IDJ(K),L)
          UM(K,L) = U_0(K,L)*AIRM(L)
          VM(K,L) = V_0(K,L)*AIRM(L)
        END DO
      END DO

C**** INITIALISE PRECIPITATION AND LATENT HEAT
      PRCP=0.
      ENRGP=0.
C**** temperature of precip is based on pre-mstcnv profile
      TPRCP=T(I,J,1)*PK(1,I,J)-TF
#ifdef TRACERS_WATER
      TRPREC(:,I,J) = 0.
#endif

C**** SET DEFAULT FOR AIR MASS FLUX (STRAT MODEL)
      AIRX(I,J)=0.

#ifdef TRACERS_SPECIAL_Shindell
C**** Save current i,j for lightning calculation in MSTCNV:
      i_lgt = i
      j_lgt = j
      RNOx_lgt(i,j) = 0.
#endif

C****
C**** Energy conservation note: For future reference the energy function
C**** for these column calculations (assuming energy reference level
C**** of 0 K for air, and 0 C liquid for water) is:
C****  E = SH + LH_vapour + LH_clw + ENRGP
C****    =  (sum(TL(:)*AIRM(:))*SHA + sum(QM(:))*LHE +
C****        sum(WML(:)*(LHE-SVLHXL(:))*AIRM(:)))*100.*BYGRAV
C**** The LH_clw term is slightly different after MSTCNV:
C****   LH_clw = sum((WML(:)*(LHE-SVLHXL(:))+SVWMXL(:)*(LHE-SVLATL(:)))
C****                *AIRM(:))*100.*BYGRAV
C**** and again after LSCOND:
C****          = sum(WMX(:)*(LHE-SVLHXL(:))*AIRM(:))*100.*BYGRAV
C****

C**** MOIST CONVECTION
      CALL MSTCNV(IERR,LERR)

C**** Error reports
      if (ierr.gt.0) then
        write(6,*) "Error in moist conv: i,j,l=",i,j,lerr
ccc     if (ierr.eq.2) call stop_model("Subsid error: abs(c) > 1",255)
        if (ierr.eq.2) ickerr = ickerr + 1
      end if

C**** ACCUMULATE MOIST CONVECTION DIAGNOSTICS
      IF (LMCMIN.GT.0) THEN
        AIJ(I,J,IJ_PSCLD)=AIJ(I,J,IJ_PSCLD)+CLDSLWIJ
        AIJ(I,J,IJ_PDCLD)=AIJ(I,J,IJ_PDCLD)+CLDDEPIJ
        IF(CLDSLWIJ.GT.1e-6) AIJ(I,J,IJ_SCNVFRQ)=AIJ(I,J,IJ_SCNVFRQ)+1.
        IF(CLDDEPIJ.GT.1e-6) AIJ(I,J,IJ_DCNVFRQ)=AIJ(I,J,IJ_DCNVFRQ)+1.
        AIJ(I,J,IJ_WMSUM)=AIJ(I,J,IJ_WMSUM)+WMSUM
        HCNDMC=0.
        DO L=1,LMCMAX
          HCNDMC=HCNDMC+DGDSM(L)+DPHASE(L)
          AJL(J,L,JL_MCHR)=AJL(J,L,JL_MCHR)+DGDSM(L)*BYDSIG(L)
          AJL(J,L,JL_MCHPHAS)=AJL(J,L,JL_MCHPHAS)+DPHASE(L)*BYDSIG(L)
          AJL(J,L,JL_MCDTOTW)=AJL(J,L,JL_MCDTOTW)+DTOTW(L)*BYDSIG(L)
CCC       IF(J.GE.J5S.AND.J.LE.J5N) AIL(I,L,IL_MCEQ)=AIL(I,L,IL_MCEQ)+
CCC  *         (DGDSM(L)+DPHASE(L))*(DXYP(J)*BYDSIG(L))
          IF(J.GE.J5S.AND.J.LE.J5N)     ! add in after parallel region
     *      AJEQIL(J-J5S+1,I,L) = (DGDSM(L)+DPHASE(L))*
     *                            (DXYP(J)*BYDSIG(L))
          AJL(J,L,JL_MCHEAT)=AJL(J,L,JL_MCHEAT)+
     &         (DPHASE(L)+DGDSM(L))*BYDSIG(L)
          AJL(J,L,JL_MCDRY)=AJL(J,L,JL_MCDRY)+
     &         (DQCOND(L)-DGDQM(L))*BYDSIG(L)
          AJL(J,L,JL_MCMFLX)=AJL(J,L,JL_MCMFLX)+MCFLX(L)
          AJL(J,L,JL_CLDMC) =AJL(J,L,JL_CLDMC) +CLDMCL(L)
          AJL(J,L,JL_CSIZMC)=AJL(J,L,JL_CSIZMC)+CSIZEL(L)*CLDMCL(L)
        END DO
        DO IT=1,NTYPE
          AJ(J,J_PRCPMC,IT)=AJ(J,J_PRCPMC,IT)+PRCPMC*FTYPE(IT,I,J)
        END DO
CCC     AREG(JR,J_PRCPMC)=AREG(JR,J_PRCPMC)+PRCPMC*DXYP(J)
        AREGIJ(I,J,1)=PRCPMC*DXYP(J)  ! add in after parallel region
        DO KR=1,NDIUPT
          IF(I.EQ.IJDD(1,KR).AND.J.EQ.IJDD(2,KR)) THEN
            ADIURN(IH,IDD_PR  ,KR)=ADIURN(IH,IDD_PR  ,KR)+PRCPMC
            ADIURN(IH,IDD_ECND,KR)=ADIURN(IH,IDD_ECND,KR)+HCNDMC
            ADIURN(IH,IDD_MCP ,KR)=ADIURN(IH,IDD_MCP ,KR)+PRCPMC
            ADIURN(IH,IDD_DMC ,KR)=ADIURN(IH,IDD_DMC ,KR)+CLDDEPIJ
            ADIURN(IH,IDD_SMC ,KR)=ADIURN(IH,IDD_SMC ,KR)+CLDSLWIJ
          END IF
        END DO

C**** ACCUMULATE PRECIP
        PRCP=PRCPMC*100.*BYGRAV
C**** CALCULATE PRECIPITATION HEAT FLUX (FALLS AT 0 DEGREES CENTIGRADE)
C**** NEED TO TAKE ACCOUNT OF LATENT HEAT THOUGH
        IF (TPRCP.gt.0) THEN
C         EPRCP=PRCP*TPRCP*SHW
          EPRCP=0.
          ENRGP=ENRGP+EPRCP
        ELSE
C         EPRCP=PRCP*TPRCP*SHI
          EPRCP=0.
          ENRGP=ENRGP+EPRCP-PRCP*LHM
          AIJ(I,J,IJ_SNWF)=AIJ(I,J,IJ_SNWF)+PRCP
        END IF

        DO L=1,LMCMAX
          T(I,J,L)=  SM(L)*BYAM(L)
          Q(I,J,L)=  QM(L)*BYAM(L)
        END DO
        CSIZMC(1:LMCMAX,I,J)=CSIZEL(1:LMCMAX)
        AIRX(I,J) = AIRXL*DXYP(J)
      END IF                    ! should this be after tracers....????
#ifdef TRACERS_ON
C**** TRACERS: Use only the active ones
      do nx=1,ntx
        n = ntix(nx)
        do l=1,lm
          dtr_mc(j,nx)=dtr_mc(j,nx)+(tm(l,nx)-trm(i,j,l,n))
#ifdef TRACERS_WATER
     *         + trsvwml(nx,l)
#endif
          tajln(j,l,jlnt_mc,n) = tajln(j,l,jlnt_mc,n) +
     &          (tm(l,nx)-trm(i,j,l,n))
#ifdef TRACERS_WATER
     *         + trsvwml(nx,l)
          trwml(nx,l) = trwm(i,j,l,n)+trsvwml(nx,l)
#endif
          tmsave(l,nx) = tm(l,nx) ! save for tajln(large-scale condense)
        end do
#ifdef TRACERS_WATER
        trprec(n,i,j) = trprmc(nx)
#endif
      end do
#endif
      LMC(1,I,J) = LMCMIN
      LMC(2,I,J) = LMCMAX+1
C****
C**** SET UP VERTICAL ARRAYS, OMITTING THE J AND I SUBSCRIPTS
C****
      DO L=1,LM
        TL(L)=T(I,J,L)*PLK(L)
        TH(L)=T(I,J,L)
        QL(L)=Q(I,J,L)
      END DO
      WMX(:)=WML(:)+SVWMXL(:)
      AQ(:)=(QL(:)-QTOLD(:,I,J))*BYDTsrc
      RNDSSL(:,1:LP50)=RNDSS(:,1:LP50,I,J)
C****
C**** COMPUTE STRATOCUMULUS CLOUDS USING PHILANDER'S FORMULA
C****
      IF (ISC.EQ.1.AND.FOCEAN(I,J).GT..5) THEN
        CSC=0.D0
        LP600=LM
        LP850=LM
        DO L=2,LM
          IF(L.GT.LP600) EXIT
          IF(PL(L).LT.600.) THEN
            LP600=L
            IF(600.-PL(L).GT.PL(L-1)-600.) LP600=L-1
          ENDIF
        ENDDO
        DO L=2,LM
          IF(L.GT.LP850) EXIT
          IF(PL(L).LT.850.) THEN
            LP850=L
            IF(850.-PL(L).GT.PL(L-1)-850.) LP850=L-1
          ENDIF
        ENDDO
        IF(SDL(LP600)+SDL(LP600+1).GT.0.) THEN
          DIFT=TL(LP850)-TGV/(1.+DELTX*QG)
          CSC=.031D0*DIFT+.623D0
          IF(CSC.LT.0.) CSC=0.
        ENDIF
        CLDMCL(1)=CLDMCL(1)+CSC
        IF(CSC.GT.0.) TAUMCL(1)=AIRM(1)*.08D0
        IF(CLDMCL(1).GT.1.) CLDMCL(1)=1.
C     IF(CSC.GT.0.) WRITE (6,*) I,J,DCL,TL(LP850),TGV/(1.+DELTX*QG),CSC
      ENDIF

C**** COMPUTE RICHARDSON NUMBER FROM SURFACE CONDITIONS WHEN DEPTH OF
C**** BOUNDARY LAYER IS AT OR BELOW FIRST LAYER (E.G. AT NIGHT)
      IF(DCL.LE.1) THEN
        THSV=TS*(1.+DELTX*QS)/PEK(1,I,J)
        THV1=TH(1)*(1.+DELTX*QL(1))
        THV2=TH(2)*(1.+DELTX*QL(2))
        ALPHAS=2./((TGV/(1.+DELTX*QG)+TS)/PEK(1,I,J))
        ALPHA1=2./(TH(1)+TS/PEK(1,I,J))
        ALPHA2=2./(TH(1)+TH(2))
        DH1S=(PLE(1)-PL(1))*TL(1)*RGAS/(GRAV*PL(1))
        BYDH1S=1./DH1S
        DH12=(GZIL(I,2)-GZIL(I,1))*BYGRAV
        BYDH12=1./DH12
        DTDZS=(THV1-THSV)*BYDH1S
        DTDZ=(THV2-THV1)*BYDH12
        IF (J.EQ.1) THEN
          DUDZ=(UZM(1,2)-UZM(1,1))*BYDH12
          DVDZ=(VZM(1,2)-VZM(1,1))*BYDH12
          DUDZS=(UZM(1,1)-US)*BYDH1S
          DVDZS=(VZM(1,1)-VS)*BYDH1S
        ENDIF
        IF (J.EQ.JM) THEN
          DUDZ=(UZM(2,2)-UZM(2,1))*BYDH12
          DVDZ=(VZM(2,2)-VZM(2,1))*BYDH12
          DUDZS=(UZM(2,1)-US)*BYDH1S
          DVDZS=(VZM(2,1)-VS)*BYDH1S
        ENDIF
        IF(J.GT.1.AND.J.LT.JM) THEN
          DUDZ=(UC(IDI(1),IDJ(1),2)+UC(IDI(2),IDJ(2),2)+
     *         UC(IDI(3),IDJ(3),2)+UC(IDI(4),IDJ(4),2)-
     *         UC(IDI(1),IDJ(1),1)-UC(IDI(2),IDJ(2),1)-
     *         UC(IDI(3),IDJ(3),1)-UC(IDI(4),IDJ(4),1))*.25*BYDH12
          DVDZ=(VC(IDI(1),IDJ(1),2)+VC(IDI(2),IDJ(2),2)+
     *         VC(IDI(3),IDJ(3),2)+VC(IDI(4),IDJ(4),2)-
     *         VC(IDI(1),IDJ(1),1)-VC(IDI(2),IDJ(2),1)-
     *         VC(IDI(3),IDJ(3),1)-VC(IDI(4),IDJ(4),1))*.25*BYDH12
          DUDZS=(UC(IDI(1),IDJ(1),1)+UC(IDI(2),IDJ(2),1)+
     *         UC(IDI(3),IDJ(3),1)+UC(IDI(4),IDJ(4),1)-
     *         4.*US)*.25*BYDH1S
          DVDZS=(VC(IDI(1),IDJ(1),1)+VC(IDI(2),IDJ(2),1)+
     *         VC(IDI(3),IDJ(3),1)+VC(IDI(4),IDJ(4),1)-
     *         4.*VS)*.25*BYDH1S
        ENDIF
        DUDZG=.1d0*US
        DVDZG=.1d0*VS
        DTDZG=.1d0*(THSV-TGV/PEK(1,I,J))
        RIS=(GRAV*ALPHAS*DTDZG)/(DUDZG*DUDZG+DVDZG*DVDZG+teeny)
        RI1=(GRAV*ALPHA1*DTDZS)/(DUDZS*DUDZS+DVDZS*DVDZS)
        RI2=(GRAV*ALPHA2*DTDZ)/(DUDZ*DUDZ+DVDZ*DVDZ)
C       WRITE (6,*)'I,J,QG,TGV,THSV,RIS,RI1=',I,J,QG,TGV,THSV,RIS,RI1
      ENDIF

c**** uncomment lines marked ECON to check energy conservation
cECON  E = ( sum(TL(1:LP50)*AIRM(1:LP50))*SHA + sum(QL(1:LP50)
cECON *     *AIRM(1:LP50))*LHE +sum( (WML(1:LP50)*(LHE-SVLHXL(1:LP50))
cECON *     +SVWMXL(1:LP50)*(LHE-SVLATL(1:LP50)))*AIRM(1:LP50))  )*100.
cECON *     *BYGRAV

C**** LARGE-SCALE CLOUDS AND PRECIPITATION
      CALL LSCOND(IERR,WMERR,LERR)

cECON  E1 = ( sum(TL(1:LP50)*AIRM(1:LP50))*SHA + sum(QL(1:LP50)
cECON *     *AIRM(1:LP50))*LHE +sum(WMX(1:LP50)*(LHE-SVLHXL(1:LP50))
cECON *     *AIRM(1:LP50)) )*100.*BYGRAV

C**** Error reports
      IF (IERR.ne.0) WRITE(99,'(I10,3I4,A,D14.5,A)')
     *       Itime,I,J,LERR,' CONDSE:H2O<0',WMERR,' ->0'

C**** Accumulate diagnostics of LSCOND
         AIJ(I,J,IJ_WMSUM)=AIJ(I,J,IJ_WMSUM)+WMSUM
         DO IT=1,NTYPE
           AJ(J,J_PRCPSS,IT)=AJ(J,J_PRCPSS,IT)+PRCPSS*FTYPE(IT,I,J)
         END DO
CCC      AREG(JR,J_PRCPSS)=AREG(JR,J_PRCPSS)+PRCPSS*DXYP(J)
         AREGIJ(I,J,2)=PRCPSS*DXYP(J)  ! add in after parallel region
         DO KR=1,NDIUPT
           IF(I.EQ.IJDD(1,KR).AND.J.EQ.IJDD(2,KR)) THEN
             ADIURN(IH,IDD_PR  ,KR)=ADIURN(IH,IDD_PR  ,KR)+PRCPSS
             ADIURN(IH,IDD_ECND,KR)=ADIURN(IH,IDD_ECND,KR)+HCNDSS
             ADIURN(IH,IDD_SSP ,KR)=ADIURN(IH,IDD_SSP ,KR)+PRCPSS
           END IF
         END DO

C**** TOTAL PRECIPITATION AND AGE OF SNOW
      PRCP=PRCP+PRCPSS*100.*BYGRAV
C**** CALCULATE PRECIPITATION HEAT FLUX (FALLS AT 0 DEGREES CENTIGRADE)
C**** NEED TO TAKE ACCOUNT OF LATENT HEAT THOUGH
      IF (LHP(1).ne.LHS) THEN
C       EPRCP=PRCPSS*100.*BYGRAV*TPRCP*SHW
        EPRCP=0.
        ENRGP=ENRGP+EPRCP
cECON    ep=0.
      ELSE
C       EPRCP=PRCPSS*100.*BYGRAV*TPRCP*SHI
        EPRCP=0.
        ENRGP=ENRGP+EPRCP-PRCPSS*100.*BYGRAV*LHM
cECON    ep=-PRCPSS*100.*BYGRAV*LHM
        AIJ(I,J,IJ_SNWF)=AIJ(I,J,IJ_SNWF)+PRCPSS*100.*BYGRAV
      END IF

cECON  if (abs(E-E1-ep).gt.0.01) print*,"energy err",i,j,E-E1-ep,
cECON *     E,E1,ep,prcpss,lhp(1)

C**** PRECIPITATION DIAGNOSTICS
        DO IT=1,NTYPE
          AJ(J,J_EPRCP,IT)=AJ(J,J_EPRCP,IT)+ENRGP*FTYPE(IT,I,J)
        END DO
CCC     AREG(JR,J_EPRCP)=AREG(JR,J_EPRCP)+ENRGP*DXYP(J)
        AREGIJ(I,J,3)=ENRGP*DXYP(J)  ! add in after parallel region
        AIJ(I,J,IJ_PREC)=AIJ(I,J,IJ_PREC)+PRCP
        AIJ(I,J,IJ_NETH)=AIJ(I,J,IJ_NETH)+ENRGP
        IF (FOCEAN(I,J).gt.0) AIJ(I,J,IJ_F0OC)=AIJ(I,J,IJ_F0OC)+
     *       ENRGP*FOCEAN(I,J)*(1.-RSI(I,J))

      IF(ENRGP.LT.0.) THEN ! MODIFY SNOW AGES AFTER SNOW FALL
        DO ITYPE=1,3
          SNOAGE(ITYPE,I,J)=SNOAGE(ITYPE,I,J)*EXP(-PRCP)
        END DO
      END IF

C**** Calculate ISCCP cloud diagnostics if required
      if (isccp_diags.eq.1) then
        do l=1,lm
          cc(l)=cldmcl(LM+1-L)+cldssl(LM+1-L)
          if(cc(l) .gt. 1.) then
            cc(l)=1.
          endif
          conv(l)=cldmcl(LM+1-L)
          if(conv(l) .gt. 1.) then
            conv(l)=1.
          endif

          dtau_s(l)=taussl(LM+1-L)
          dtau_c(l)=taumcl(LM+1-L)
          pfull(l)=pl(LM+1-L)*100.
          phalf(l)=ple(LM+2-L)*100.
          at(l)=tl(LM+1-L)  ! in situ temperature

C**** set tg1 from GTEMP array (or save in SURFACE?)
c          skt=tf+tg1(i,j)
          skt=tf + (focean(i,j)+flake(i,j))*(1.-rsi(i,j))*gtemp(1,1,i,j)
     *         + (focean(i,j)+flake(i,j))*rsi(i,j)*gtemp(1,2,i,j)
     *         + flice(i,j)*gtemp(1,3,i,j)+fearth(i,j)*gtemp(1,4,i,j)
          dem_s(l)=0.
          dem_c(l)=0.
          if(svlhxl(LM+1-L) .eq. lhe )   ! large-scale water cloud
     *      dem_s(l)=1.-exp(-taussl(LM+1-L)*bywc)
          if(svlatl(LM+1-L) .eq. lhe )   ! convective water cloud
     *      dem_c(l)=1.-exp(-taumcl(LM+1-L)*bywc)
          if(svlhxl(LM+1-L) .eq. lhs )   ! large-scale ice cloud
     *      dem_s(l)=1.-exp(-taussl(LM+1-L)*byic)
          if(svlatl(LM+1-L) .eq. lhs )   ! convective ice cloud
     *      dem_c(l)=1.-exp(-taumcl(LM+1-L)*byic)

          qv(l)=ql(LM+1-L)
        end do
        phalf(lm+1)=ple(1)*100.
        itrop = LM+1-LTROPO(I,J)

        call ISCCP_CLOUD_TYPES(pfull,phalf,qv,
     &       cc,conv,dtau_s,dtau_c,skt,
     &       at,dem_s,dem_c,itrop,fq_isccp,ctp,tauopt,nbox,jerr)
        if(jerr.ne.0) jckerr = jckerr + 1
C**** set ISCCP diagnostics
        if (nbox.gt.0) then
          AIJ(I,J,IJ_CTPI) = AIJ(I,J,IJ_CTPI) + ctp
          AIJ(I,J,IJ_TAUI) = AIJ(I,J,IJ_TAUI) + tauopt
          AIJ(I,J,IJ_TCLDI)= AIJ(I,J,IJ_TCLDI)+ 1.
        end if
C**** note LOW CLOUDS:       ipres=6,7, MID-LEVEL CLOUDS: ipres=4,5 ,
C****      HIGH CLOUDS:      ipres=1,2,3
C**** Sum over itau=2,ntau (itau=1 is no cloud)
        do itau=2,ntau
          AIJ(I,J,IJ_LCLDI)= AIJ(I,J,IJ_LCLDI) + fq_isccp(itau,6)
     *         + fq_isccp(itau,7)
          AIJ(I,J,IJ_MCLDI)= AIJ(I,J,IJ_MCLDI) + fq_isccp(itau,4)
     *         + fq_isccp(itau,5)
          AIJ(I,J,IJ_HCLDI)= AIJ(I,J,IJ_HCLDI) + fq_isccp(itau,1)
     *         + fq_isccp(itau,2) + fq_isccp(itau,3)
        end do
C**** Save area weighted isccp histograms
        n=isccp_reg(j)
        if (n.gt.0) then
          AISCCP(:,:,n) = AISCCP(:,:,n) + fq_isccp(:,:)*DYP(j)
        end if
      end if

C**** Peak static stability diagnostic
      SSTAB=-1.d30
      DO L=1,DCL
Cred    IF(SSTAB.lt.(TH(L+1)-TH(L))/(GZ(I,J,L+1)-GZ(I,J,L)))
Cred *     SSTAB =  (TH(L+1)-TH(L))/(GZ(I,J,L+1)-GZ(I,J,L))
        IF(SSTAB.lt.(TH(L+1)-TH(L))/(GZIL(I,L+1)-GZIL(I,L)))
     *     SSTAB =  (TH(L+1)-TH(L))/(GZIL(I,L+1)-GZIL(I,L))
      END DO
      AIJ(I,J,ij_sstabx) = AIJ(I,J,ij_sstabx) + SSTAB

C**** WRITE TO GLOBAL ARRAYS
      TAUMC(:,I,J)=TAUMCL(:)
      CLDMC(:,I,J)=CLDMCL(:)
      SVLAT(:,I,J)=SVLATL(:)

      TAUSS(:,I,J)=TAUSSL(:)
      CLDSS(:,I,J)=CLDSSL(:)
      CLDSAV(:,I,J)=CLDSAVL(:)
      SVLHX(:,I,J)=SVLHXL(:)
      CSIZSS(:,I,J)=CSIZEL(:)

      RHSAV(:,I,J)=RH(:)
      TTOLD(:,I,J)=TH(:)
      QTOLD(:,I,J)=QL(:)

      PREC(I,J)=PRCP            ! total precip mass (kg/m^2)
      EPREC(I,J)=ENRGP          ! energy of precipitation (J/m^2)
C**** The PRECSS array is only used if a distinction is being made
C**** between kinds of rain in the ground hydrology.
      PRECSS(I,J)=PRCPSS*100.*BYGRAV  ! large scale precip (kg/m^2)

      DO L=1,LM
        AJL(J,L,JL_SSHR)=AJL(J,L,JL_SSHR)+SSHR(L)
        AJL(J,L,JL_MCLDHT)=AJL(J,L,JL_MCLDHT)+DCTEI(L)
        AJL(J,L,JL_RHE)=AJL(J,L,JL_RHE)+RH1(L)
        AJL(J,L,JL_CLDSS) =AJL(J,L,JL_CLDSS) +CLDSSL(L)
        AJL(J,L,JL_CSIZSS)=AJL(J,L,JL_CSIZSS)+CSIZEL(L)*CLDSSL(L)

        T(I,J,L)=TH(L)
        Q(I,J,L)=QL(L)
C**** update moment changes
Cred    T3MOM(:,I,J,L)=SMOM(:,L)*BYAM(L)
Cred    Q3MOM(:,I,J,L)=QMOM(:,L)*BYAM(L)
Cred    WM(I,J,L)=WMX(L)
        TMOMIL(:,I,L)=SMOM(:,L)*BYAM(L)
        QMOMIL(:,I,L)=QMOM(:,L)*BYAM(L)
        WMIL(I,L)=WMX(L)

C**** UPDATE MODEL WINDS
CCC     DO K=1,KMAX       !  add in after parallel region (order)
CCC       U(IDI(K),IDJ(K),L)=U(IDI(K),IDJ(K),L)
CCC  &         +(UM(K,L)*BYAM(L)-UC(IDI(K),IDJ(K),L))
CCC       V(IDI(K),IDJ(K),L)=V(IDI(K),IDJ(K),L)
CCC  &         +(VM(K,L)*BYAM(L)-VC(IDI(K),IDJ(K),L))
CCC     ENDDO
         IF(J.EQ.1)  THEN
            DO K=1,IM ! KMAX
               UKP1(K,L)=(UM(K,L)*BYAM(L)-UC(IDI(K),IDJ(K),L))
               VKP1(K,L)=(VM(K,L)*BYAM(L)-VC(IDI(K),IDJ(K),L))
            END DO
         ELSE IF(J.EQ.JM)  THEN
            DO K=1,IM ! KMAX
               UKPJM(K,L)=(UM(K,L)*BYAM(L)-UC(IDI(K),IDJ(K),L))
               VKPJM(K,L)=(VM(K,L)*BYAM(L)-VC(IDI(K),IDJ(K),L))
            END DO
         ELSE
            DO K=1,4 ! KMAX
               UKM(K,I,J,L)=(UM(K,L)*BYAM(L)-UC(IDI(K),IDJ(K),L))
               VKM(K,I,J,L)=(VM(K,L)*BYAM(L)-VC(IDI(K),IDJ(K),L))
            END DO
         END IF
      ENDDO

#ifdef TRACERS_ON
C**** TRACERS: Use only the active ones
      do nx=1,ntx
        n = ntix(nx)
        do l=1,lp50
          dtr_ss(j,nx)=dtr_ss(j,nx)+(tm(l,nx)-tmsave(l,nx))
#ifdef TRACERS_WATER
     &         + (trwml(nx,l)-trwm(i,j,l,n)-trsvwml(nx,l))
#endif
          trm(i,j,l,n) = tm(l,nx)
          trmom(:,i,j,l,n) = tmom(:,l,nx)
          tajln(j,l,jlnt_lscond,n) = tajln(j,l,jlnt_lscond,n) +
     &          (tm(l,nx)-tmsave(l,nx))
#ifdef TRACERS_WATER
     &         + (trwml(nx,l)-trwm(i,j,l,n)-trsvwml(nx,l))
          trwm(i,j,l,n) = trwml(nx,l)
#endif
        end do
#ifdef TRACERS_WATER
        trprec(n,i,j) = trprec(n,i,j)+trprss(nx)
C**** diagnostics
        tajls(j,1,jls_source(3,n))=tajls(j,1,jls_source(3,n))
     *       +trprec(n,i,j)*bydxyp(j)
        tajls(j,1,jls_source(4,n))=tajls(j,1,jls_source(4,n))
     *       +trprec(n,i,j)*focean(i,j)*bydxyp(j)
        taijn(i,j,tij_prec,n) =taijn(i,j,tij_prec,n)+trprec(n,i,j)
     *       *bydxyp(j)
#endif
      end do
#endif

      END DO
C**** END OF MAIN LOOP FOR INDEX I

C****
Cred*           Reduced Arrays 3
C****
         DO L=1,LM
            WM(:,J,L) = WMIL(:,L)
         END DO
         DO L=1,LM
            T3MOM(:,:,J,L) = TMOMIL(:,:,L)
         END DO
         DO L=1,LM
            Q3MOM(:,:,J,L) = QMOMIL(:,:,L)
         END DO
Cred*       end Reduced Arrays 3
      END DO
C**** END OF MAIN LOOP FOR INDEX J
!$OMP  END PARALLEL DO
C****
C
C     WAS THERE AN ERROR IN SUBSID ??
C
      IF(ICKERR.NE.0)  THEN
         WRITE(6,*)  'SUBSID ERROR: ABS(C) > 1'
         call stop_model('SUBSID ERROR: ABS(C) > 1',255)
      END IF
C
C     WAS THERE AN ERROR IN ISCCP CLOUD TYPING ??
C
      IF(JCKERR.NE.0)  THEN
         WRITE(6,*)  'ISCCP CLOUD TYPING ERROR'
         call stop_model('ISCCP CLOUD TYPING ERROR',255)
      END IF

#ifdef TRACERS_ON
C**** Save the conservation quantities for tracers
      do nx=1,ntx
        n=ntix(nx)
        call diagtcb(dtr_mc(1,nx),itcon_mc(n),n)
        call diagtcb(dtr_ss(1,nx),itcon_ss(n),n)
      end do
#endif

C**** Delayed summations (to control order of summands)
      DO J=J5S,J5N
      DO I=1,IM
        IF(LMC(1,I,J).GT.0) THEN
          DO L=1,LMC(2,I,J)-1
            AIL(I,L,IL_MCEQ)=AIL(I,L,IL_MCEQ)+AJEQIL(J-J5S+1,I,L)
          END DO
        END IF
      END DO
      END DO
C
      DO J=1,JM
      DO I=1,IMAXJ(J)
         JR=JREG(I,J)
         IF(LMC(1,I,J).GT.0)
     *     AREG(JR,J_PRCPMC)=AREG(JR,J_PRCPMC)+AREGIJ(I,J,1)
         AREG(JR,J_PRCPSS)=AREG(JR,J_PRCPSS)+AREGIJ(I,J,2)
         AREG(JR,J_EPRCP) =AREG(JR,J_EPRCP) +AREGIJ(I,J,3)
      END DO
      END DO
C
C     NOW REALLY UPDATE THE MODEL WINDS
C
      J=1
      DO K=1,IM ! KMAXJ(J)
         IDI(K)=IDIJ(K,1,J)
         IDJ(K)=IDJJ(K,J)
      END DO
      DO L=1,LM
      DO K=1,IM ! KMAXJ(J)
         U(IDI(K),IDJ(K),L)=U(IDI(K),IDJ(K),L)+UKP1(K,L)
         V(IDI(K),IDJ(K),L)=V(IDI(K),IDJ(K),L)+VKP1(K,L)
      END DO
      END DO
C
!$OMP  PARALLEL DO PRIVATE(I,J,K,L,IDI,IDJ)
      DO L=1,LM
      DO J=2,JM-1
         DO K=1,4  !  KMAXJ(J)
            IDJ(K)=IDJJ(K,J)
         END DO
         DO I=1,IM
         DO K=1,4 ! KMAXJ(J)
            IDI(K)=IDIJ(K,I,J)
            U(IDI(K),IDJ(K),L)=U(IDI(K),IDJ(K),L)+UKM(K,I,J,L)
            V(IDI(K),IDJ(K),L)=V(IDI(K),IDJ(K),L)+VKM(K,I,J,L)
         END DO
         END DO
      END DO
      END DO
!$OMP  END PARALLEL DO
C
      J=JM
      DO K=1,IM  !  KMAXJ(J)
         IDI(K)=IDIJ(K,1,J)
         IDJ(K)=IDJJ(K,J)
      END DO
      DO L=1,LM
      DO K=1,IM  !  KMAXJ(J)
         U(IDI(K),IDJ(K),L)=U(IDI(K),IDJ(K),L)+UKPJM(K,L)
         V(IDI(K),IDJ(K),L)=V(IDI(K),IDJ(K),L)+VKPJM(K,L)
      END DO
      END DO
C
C**** ADD IN CHANGE OF MOMENTUM BY MOIST CONVECTION AND CTEI
C**** and save changes in KE for addition as heat later
!$OMP  PARALLEL DO PRIVATE(I,J,L)
      DO L=1,LM
      DO J=2,JM
      DO I=1,IM
         AJL(J,L,JL_DAMMC)=AJL(J,L,JL_DAMMC)+
     &         (U(I,J,L)-UC(I,J,L))*PDSIG(L,I,J)
         DKE(I,J,L)=0.5*(U(I,J,L)*U(I,J,L)+V(I,J,L)*V(I,J,L)
     *       -UC(I,J,L)*UC(I,J,L)-VC(I,J,L)*VC(I,J,L))
      END DO
      END DO
      END DO
!$OMP  END PARALLEL DO

      if (isccp_diags.eq.1) CALL RINIT(seed) ! reset random number sequ.

      RETURN
      END SUBROUTINE CONDSE

      SUBROUTINE init_CLD
!@sum  init_CLD initialises parameters for MSTCNV and LSCOND
!@auth M.S.Yao/A. Del Genio (modularisation by Gavin Schmidt)
!@ver  1.0 (taken from CB265)
      USE CONSTANT, only : grav,by3
      USE MODEL_COM, only : dtsrc,ls1,sige,lm,psfmpt,ptop,plbot,jm
      USE GEOM, only : lat_dg
      USE CLOUDS, only : lmcm,bydtsrc,xmass,brcld,bybr,U00wtrX,U00ice
     *  ,HRMAX,ISC,lp50,RICldX,RWCldOX,xRIcld
      USE CLOUDS_COM, only : llow,lmid,lhi
      USE DAGCOM, only : nisccp,isccp_reg
      USE PARAM

      IMPLICIT NONE
      REAL*8 PLE
      INTEGER L,J

      call sync_param( 'U00wtrX', U00wtrX )
      call sync_param( 'U00ice', U00ice )
      call sync_param( "LMCM", LMCM )
      call sync_param( "HRMAX", HRMAX )
      call sync_param( "RICldX", RICldX )
      xRIcld = .001d0*(RICldX-1.d0)
      call sync_param( "RWCldOX", RWCldOX )
      call sync_param( "ISC", ISC)

      IF(LMCM.LT.0) LMCM = LS1-1
      call set_param( "LMCM", LMCM, 'o' )

      BYDTsrc=1./DTsrc
      XMASS=0.1d0*DTsrc*GRAV

      BYBR=((1.-BRCLD)*(1.-2.*BRCLD))**BY3

C**** SEARCH FOR THE 50 MB LEVEL
      LP50=LM
      DO L=LM-1,1,-1
        PLE=.25*(SIGE(L)+2.*SIGE(L+1)+SIGE(L+2))*PSFMPT+PTOP
        IF (PLE.LT.50.) LP50=L
      END DO
      write(6,*) "Maximum level for LSCOND calculations (50mb): ",LP50

C**** CLOUD LAYER INDICES USED FOR DIAGNOSTICS
      DO L=1,LM
        LLOW=L
        IF (.5*(PLbot(L+1)+PLbot(L+2)).LT.750.) EXIT ! was 786. 4/16/97
      END DO
      DO L=LLOW+1,LM
        LMID=L
        IF (.5*(PLbot(L+1)+PLbot(L+2)).LT.430.) EXIT
      END DO
      LHI=LM
      IF (LMID+1.GT.LHI) LHI=LMID+1
      WRITE (6,47) LLOW,LLOW+1,LMID,LMID+1,LHI
 47   FORMAT (' LOW CLOUDS IN LAYERS 1-',I2,'   MID LEVEL CLOUDS IN',
     *     ' LAYERS',I3,'-',I2,'   HIGH CLOUDS IN LAYERS',I3,'-',I2)

C**** Define regions for ISCCP diagnostics
      do j=1,jm
        isccp_reg(j)=0.
        if (lat_dg(j,1).ge.-60. .and. lat_dg(j,1).lt.-45.)
     *       isccp_reg(j)=1
        if (lat_dg(j,1).ge.-45. .and. lat_dg(j,1).lt.-30.)
     *       isccp_reg(j)=2
        if (lat_dg(j,1).ge.-30. .and. lat_dg(j,1).lt.-15.)
     *       isccp_reg(j)=3
        if (lat_dg(j,1).ge.-15. .and. lat_dg(j,1).lt.15.)
     *       isccp_reg(j)=4
        if (lat_dg(j,1).ge.15. .and. lat_dg(j,1).lt.30.)
     *       isccp_reg(j)=5
        if (lat_dg(j,1).ge.30. .and. lat_dg(j,1).lt.60.)
     *       isccp_reg(j)=6
      end do

      END SUBROUTINE init_CLD

ext: f


text:
FUNCTION factln(n)
!
IMPLICIT NONE
!
! Function arguments
!
INTEGER :: n
REAL :: factln
!
! Local variables
!
REAL,DIMENSION(100) :: a
REAL :: gammln
!
!     + + + purpose + + +
!     computes the ln(n!) for n > 0
!
!     based on:
!     'numerical recipes - the art of scientific computing',
!     w.h. press, b.p. flannery, s.a. teukolsky, w.t. vetterling
!     cambridge university press, 1986
!     pg 159
!
!     + + + keywords + + +
!     factorial
!
!     + + + argument declarations + + +
!
!
!     + + + argument definitions + + +
!     n      - real value for values >= 0
!
!     + + + local variables + + +
!
!
DATA a/100* - 1/
!
!     + + + end specifications + + +
!
IF (n.LT.0) WRITE (*,*) 'in function factln(n): negative factorial!!'
IF (n.LE.99) THEN
  IF (a(n+1).LT.0.0) a(n+1) = gammln(n+1.0)
  factln = a(n+1)
ELSE
  factln = gammln(n+1.0)
END IF
!
END FUNCTION factln

max_stars_repo_path: Project Documents/Source Code Original/Factln.f90

max_stars_repo_name: USDA-ARS-WMSRU/upgm-standalone

max_issues_repo_name: USDA-ARS-WMSRU/upgm-standalone

max_forks_repo_name: USDA-ARS-WMSRU/upgm-standalone


text:
+incdir+./
+incdir+../../

top.v

max_stars_repo_name: testdrive-profiling-master/profiles

max_issues_repo_name: testdrive-profiling-master/profiles

max_forks_repo_name: testdrive-profiling-master/profiles


text:
program findqmesh

implicit none
  integer, parameter :: DP = kind(1.0d0)
  integer, parameter :: DPC = kind((1.0d0,1.0d0))
integer :: n1,n2,i,m1,m2,mm,mden,p,q,nx1,nx2,mx1,mx2,my1,my2
integer :: components(3),qkv(3)
real(DP) :: xm1,xm2,scqx,scqy,ucqx,ucqy
real(DP) :: qpt(3)
real(DP), parameter :: TOL_Small = 1.0d-6

  open(14,file='qlist',status='replace')
  open(15,file='qlistshift',status='replace')
read(*,*) p, q, n1,n2,m1,m2,scqx,scqy
  mm = n1*m2 - n2*m1
if(mm .lt. 0) then
  mm = -mm
endif
mden = mm*(p)-1
ucqx = (m2*scqx-n2*scqy)/mm
ucqy = (-m1*scqx+n1*scqy)/mm
write(*,*)ucqx,ucqy
i=0
mx1=MAX(i,n1 ,n2 ,n1+n2)
mx2=MAX(i,m1 ,m2 ,m1+m2)
my1=MIN(i,n1 ,n2 ,n1+n2)
my2=MIN(i,m1 ,m2 ,m1+m2)
qpt(3) = 0.0d0
do nx1 = my1*p, mx1*p-1
 do nx2 = my2*q,mx2*q-1
 xm1 = (m2*(1.0d0*nx1/p) - n2*(1.0d0*nx2/q))/mm
 xm2 =(-m1*(1.0d0*nx1/p) + n1*(1.0d0*nx2/q))/mm
  if(xm1 .gt. -0.000000001d0 .and. xm2 .gt. -0.000000001d0)then
   if(xm1 .le. 0.999999999d0 .and. xm2 .le. 0.999999999d0)then
    qpt(1) = xm1
    qpt(2) = xm2
    write(14,'(3(f13.9))') qpt(:)
    qpt(1) = xm1+ucqx
    qpt(2) = xm2+ucqy
    write(15,'(3(f13.9))') qpt(:)
   endif
  endif
 enddo
enddo
  close(14)
  close(15)
end program findqmesh

max_stars_repo_name: quek-group/XAF-GW

max_issues_repo_name: quek-group/XAF-GW

max_forks_repo_name: quek-group/XAF-GW


! Name: Anantha Rao
! Reg no: 20181044
! This program provides the Runge-Kutta4 Method to solve the differential equation y' = f(x,y)

program rungekutta
    implicit none
    real:: x,y,xwant,h,m1,m2,m3,m4,f
    integer:: niter,i

    x=0.0; y=0.0; xwant=1.57; h=0.01;
    niter=int((xwant-x)/h)
    open(unit=21, file="q1_rk4.dat")

    do i=1,niter
        m1=h*f(x,y)
        m2=h*f(x+0.5*h,y+0.5*m1)
        m3=h*f(x+0.5*h,y+0.5*m2)
        m4=h*f(x+h,y+m3)
        x=x+h
        y=y+(m1+2.0*m2+2.0*m3+m4)*(1/6.0)
        write(21,*) dfloat(i)*h,y
    end do
end program

!define your function here
real function f(x,y)
real:: x,y
f=1+y**2

end function


text:
! Copyright 2018 Xanadu Quantum Technologies Inc.

! Licensed under the Apache License, Version 2.0 (the "License");
! you may not use this file except in compliance with the License.
! You may obtain a copy of the License at

!     http://www.apache.org/licenses/LICENSE-2.0

! Unless required by applicable law or agreed to in writing, software
! distributed under the License is distributed on an "AS IS" BASIS,
! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
! See the License for the specific language governing permissions and
! limitations under the License.
module vars
    use kinds
    implicit none

    integer(ip)         :: n, ell
    real(wp), parameter :: eps = 1.0d-6
    real(wp), parameter :: pi = acos(-1.0_wp)
end module vars

max_stars_repo_name: XanaduAI/torontonian-sampling

max_issues_repo_name: XanaduAI/torontonian-sampling

max_forks_repo_name: XanaduAI/torontonian-sampling


text:
program plot_potential
use precisions, only : dp
use delta_shell, only : nn_model
use exp_data, only : nn_experiment
use optimization, only : setup_optimization, covariance_matrix
use chi_square, only : total_chi_square
use read_write, only : plot_potential_components
implicit none
type(nn_model) :: model
real(dp), allocatable, dimension(:) :: parameters
logical, allocatable, dimension(:) :: mask
type(nn_experiment), allocatable, dimension(:) :: database
logical :: save_results
character(len=1024) :: output_file
integer :: n_points
real(dp) :: chi2
real(dp), allocatable, dimension(:, :) :: alpha, covariance
real(dp), allocatable, dimension(:) :: beta
real(dp), parameter :: r_min = 0.0_dp, r_max = 2.0_dp, r_step = 0.0078125_dp

call setup_optimization(model, parameters, mask, database, save_results, output_file)
call total_chi_square(database, parameters, mask, model, n_points, chi2, alpha, beta)
covariance = covariance_matrix(alpha, mask)
call plot_potential_components(model, parameters, covariance, r_min, r_max, r_step, output_file)

end program plot_potential

max_stars_repo_name: rnavarroperez/nn-scattering-fit

max_issues_repo_name: rnavarroperez/nn-scattering-fit

max_forks_repo_name: rnavarroperez/nn-scattering-fit


module rsdft_allreduce_module

  use communicator_module, only: set_communicator

  private
  public :: rsdft_allreduce

  interface rsdft_allreduce
    module procedure d_rsdft_allreduce_0 &
                    ,d_rsdft_allreduce_1, z_rsdft_allreduce_1 &
                    ,d_rsdft_allreduce_2, z_rsdft_allreduce_2
  end interface rsdft_allreduce

  integer :: n_opt, n_opt_h

contains

  subroutine d_rsdft_allreduce_0( a, comm )
    implicit none
    real(8),intent(inout) :: a
    character(*),optional,intent(in) :: comm
    integer :: communicator,ierr
    real(8) :: b
    include 'mpif.h'
    communicator = MPI_COMM_WORLD
    if ( present(comm) ) communicator = set_communicator( comm )
    b=a
    call MPI_Allreduce( b, a, 1, MPI_REAL8, MPI_SUM, communicator, ierr )
  end subroutine d_rsdft_allreduce_0

  subroutine d_rsdft_allreduce_1( a, comm )
    implicit none
    real(8),intent(inout) :: a(:)
    character(*),optional,intent(in) :: comm
    integer :: communicator,n,ierr
    real(8),allocatable :: b(:)
    include 'mpif.h'
    communicator = MPI_COMM_WORLD
    if ( present(comm) ) communicator = set_communicator( comm )
    n=size(a)
#ifdef _NO_MPI_INPLACE_
    allocate( b(n) ); b=a
    call MPI_Allreduce( b, a, n, MPI_REAL8, MPI_SUM, communicator, ierr )
#else
    call MPI_Allreduce( MPI_IN_PLACE,a,n,MPI_REAL8,MPI_SUM,communicator,ierr )
#endif
  end subroutine d_rsdft_allreduce_1

  subroutine d_rsdft_allreduce_2( a, comm )
    implicit none
    real(8),intent(inout) :: a(:,:)
    character(*),optional,intent(in) :: comm
    integer :: communicator,n,ierr
    real(8),allocatable :: b(:,:)
    include 'mpif.h'
    communicator = MPI_COMM_WORLD
    if ( present(comm) ) communicator = set_communicator( comm )
    n=size(a)
#ifdef _NO_MPI_INPLACE_
    allocate( b(size(a,1),size(a,2)) ); b=a
    call MPI_Allreduce( b, a, n, MPI_REAL8, MPI_SUM, communicator, ierr )
#else
    call MPI_Allreduce( MPI_IN_PLACE,a,n,MPI_REAL8,MPI_SUM,communicator,ierr )
#endif
  end subroutine d_rsdft_allreduce_2

  subroutine z_rsdft_allreduce_1( a, comm )
    implicit none
    complex(8),intent(inout) :: a(:)
    character(*),optional,intent(in) :: comm
    integer :: communicator,n,ierr
    complex(8),allocatable :: b(:)
    include 'mpif.h'
    communicator = MPI_COMM_WORLD
    if ( present(comm) ) communicator = set_communicator( comm )
    n=size(a)
#ifdef _NO_MPI_INPLACE_
    allocate( b(n) ); b=a
    call MPI_Allreduce( b,a,n,MPI_COMPLEX16,MPI_SUM,communicator,ierr )
#else
    call MPI_Allreduce( MPI_IN_PLACE,a,n,MPI_COMPLEX16,MPI_SUM,communicator,ierr )
#endif
  end subroutine z_rsdft_allreduce_1

  subroutine z_rsdft_allreduce_2( a, comm )
    implicit none
    complex(8),intent(inout) :: a(:,:)
    character(*),optional,intent(in) :: comm
    integer :: communicator,n,ierr
    complex(8),allocatable :: b(:,:)
    include 'mpif.h'
    communicator = MPI_COMM_WORLD
    if ( present(comm) ) communicator = set_communicator( comm )
    n=size(a)
#ifdef _NO_MPI_INPLACE_
    allocate( b(size(a,1),size(a,2)) ); b=a
    call MPI_Allreduce( b,a,n,MPI_COMPLEX16,MPI_SUM,communicator,ierr )
#else
    call MPI_Allreduce( MPI_IN_PLACE,a,n,MPI_COMPLEX16,MPI_SUM,communicator,ierr )
#endif
  end subroutine z_rsdft_allreduce_2

  subroutine test_allreduce
    implicit none
    integer :: i,j,k,m,mt,n,i0,ierr,myrank
    real(8),allocatable :: a(:), b(:)
    real(8) :: ct,ct0,ct1,ctmin,et,et0,et1,etmin
    include 'mpif.h'

    !return

    call write_border( 0, " test_allreduce(start)" )

    call MPI_Comm_rank(MPI_COMM_WORLD,myrank,ierr)
    n=2**16
    allocate( a(n) ); a(:)=1.0d0
    allocate( b(n) ); b(:)=0.0d0
    ctmin=1.d100
    etmin=1.d100
    do j=16,0,-1
      m = 2**j
      mt=0
      do i=1,2**(16-j)
        mt=mt+m
      end do
      b=0.0d0
      call MPI_Barrier( MPI_COMM_WORLD, ierr )
      call cpu_time(ct0) ; et0=mpi_wtime()
      do k=1,16
        do i=1,2**(16-j)
          i0=(i-1)*m+1
          call MPI_Allreduce(a(i0),b(i0),m,MPI_REAL8,MPI_SUM,MPI_COMM_WORLD,ierr)
        end do
      end do ! k
      call MPI_Barrier( MPI_COMM_WORLD, ierr )
      call cpu_time(ct1) ; et1=mpi_wtime()
      ct=ct1-ct0
      et=et1-et0
      ctmin=min(ct,ctmin)
      if ( et < etmin ) then
        n_opt = m
        etmin = et
      end if
      if ( myrank == 0 ) then
        write(*,'(1x,3i12,4f10.5,2x,g15.8)') m,i-1,mt,ct,ctmin,et,etmin,sum(b(1:mt))
      end if
    end do ! j
    deallocate( b )
    deallocate( a )
    call MPI_Bcast(n_opt,1,MPI_INTEGER,0,MPI_COMM_WORLD,ierr)
    n_opt_h=n_opt/2
    if ( myrank == 0 ) write(*,*) "n_opt, n_opt_h=",n_opt,n_opt_h
    call write_border( 0, " test_allreduce(end)" )
  end subroutine test_allreduce

end module rsdft_allreduce_module


text:
      SUBROUTINE MXCIDS(*,Z,MSET,MSZE,NWDS,USET,NSTRT,SNAM)
C-----
C   THIS SUBROUTINE CREATES A LIST OF SUBSTRUCTURE NAMES AT Z(1)
C AND AN ARRAY AT Z(NSTRT) OF LENGTH MSZE*NWDS WITH THE FIRST TWO WORDS
C OF EACH ENTRY AS,
C     1 - EXTERNAL ID * 10 + COMPONENT
C     2 - POINTER TO FIRST WORD OF SUBSTRUCTURE NAME...
C
C INPUT
C     Z     = OPEN CORE - Z(1) TO Z(NSTRT-1)
C     MSET  = ONE WORD BCD IDENTIFING SET
C     MSZE  = NUMBER OF COLUMNS IN MATRIX
C     NWDS  = NUMBER WORDS/ENTRY DESIRED (2 MINIMUM)
C     USET  = USET GINO FILE NAME
C     NSTRT = LOCATION OF FIRST OF FOUR (4) BUFFERS.
C     SNAM  = SUBSTRUCTURE NAME BEING SOLVED (2 WORD BCD).
C OUTPUT
C     *     = TASK NOT COMPLETED
C     Z     = SUBSTRUCTURE NAMES + COLUMN IDENTIFIERS (SEE ABOVE)
C     NSTRT = FIRST WORD OF COLUMN IDENTIFIERS
C-----
C
C       N O T   C O D E D   Y E T
C
        IF (MSET.EQ.1) RETURN
      RETURN 1
        END


text:
! { dg-do compile }
! PR 38425 I/O: POS= compile-time diagnostics
!----------------------------------------------------------
character(len=30) :: str
open(3,access='stream')

! C919 (R913) If io-unit is not a file-unit-number, the
! io-control-spec-list shall not contain a REC= specifier
! or a POS= specifier.
write(str,*, pos=4) 5 ! { dg-error "incompatible with internal" }

! C927 (R913) If a POS= specifier appears, the
! io-control-spec-list shall not contain a REC= specifier.
write(3,pos=5,rec=4) 5 ! { dg-error "POS= is not allowed with REC=" }
write(3,rec=4,pos=5) 5 ! { dg-error "POS= is not allowed with REC=" }

!Fortran runtime error: REC=specifier not allowed with STREAM access
write(3,rec=4) 5
!Fortran runtime error: REC=specifier must be positive
write(3,rec=-3) 44
!Fortran runtime error: POS=specifier must be positive
write(3,pos=-4) 44
end

max_issues_repo_name: sujankh/rose-matlab

max_forks_repo_name: sujankh/rose-matlab