wann_omega Subroutine

  subroutine wann_omega(csheet, sheet, rave, r2ave, rave2, wann_spread, num_wann, kmesh_info, &
                        num_kpts, print_output, use_ss_functional, wann_slwf, omega_invariant, ln_tmp_loc, &
                        m_matrix_loc, lambda_loc, first_pass, timer, nkrank, global_k, error, comm)
    !================================================!
    !
    !!   Calculate the Wannier Function spread
    !
    ! Modified by Valerio Vitale for the SLWF+C method (PRB 90, 165125)
    ! Jun 2018, based on previous work by Charles T. Johnson and
    ! Radu Miron at Implerial College London
    !================================================

    use w90_io, only: io_stopwatch_start, io_stopwatch_stop
    use w90_comms, only: comms_allreduce, w90_comm_type, mpirank
    use w90_types, only: kmesh_info_type, print_output_type, timer_list_type
    use w90_wannier90_types, only: wann_slwf_type

    implicit none

    ! arguments
    type(kmesh_info_type), intent(in) :: kmesh_info
    type(localisation_vars_type), intent(out)  :: wann_spread
    type(print_output_type), intent(in) :: print_output
    type(w90_comm_type), intent(in) :: comm
    type(wann_slwf_type), intent(in) :: wann_slwf
    type(timer_list_type), intent(inout) :: timer
    type(w90_error_type), allocatable, intent(out) :: error

    integer, intent(in) :: nkrank, global_k(:)
    integer, intent(in) :: num_kpts
    integer, intent(in) :: num_wann

    complex(kind=dp), intent(in)  :: csheet(:, :, :)
    complex(kind=dp), intent(in) :: m_matrix_loc(:, :, :, :)

    real(kind=dp), intent(in) :: lambda_loc
    real(kind=dp), intent(in) :: omega_invariant
    logical, intent(in) :: use_ss_functional
    real(kind=dp), intent(inout) :: ln_tmp_loc(:, :, :)
    real(kind=dp), intent(in)  :: sheet(:, :, :)
    real(kind=dp), intent(out) :: r2ave(:)
    real(kind=dp), intent(out) :: rave(:, :)
    real(kind=dp), intent(out) :: rave2(:)

    logical, intent(inout) :: first_pass

    ! local variables
    real(kind=dp) :: mnn2
    complex(kind=dp) :: summ
    complex(kind=dp), allocatable :: sum_mnn(:, :)
    real(kind=dp) :: brn
    integer :: ind, nkp, nn, m, n, iw, nkp_loc, cnn
    integer :: my_node_id
    integer :: ierr

    my_node_id = mpirank(comm)

    if (print_output%timing_level > 1 .and. print_output%iprint > 0) call io_stopwatch_start('wann: omega', timer)

    if (use_ss_functional) then

      allocate (sum_mnn(num_wann, kmesh_info%nntot), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error in allocating sum_mnn in wann_omega', comm)
        return
      end if

      sum_mnn = 0.0_dp
      ! JJ, maybe reorder loops?
      do nn = 1, kmesh_info%nntot
        do n = 1, num_wann
          do nkp_loc = 1, nkrank
            nkp = global_k(nkp_loc)
            cnn = kmesh_info%nnord(nn, nkp) ! enforce uniform order of bk vectors
            sum_mnn(n, nn) = sum_mnn(n, nn) + csheet(n, nn, 1)*m_matrix_loc(n, n, cnn, nkp_loc)
          end do
        end do
      end do

      call comms_allreduce(sum_mnn(1, 1), num_wann*kmesh_info%nntot, 'SUM', error, comm)
      if (allocated(error)) return

      sum_mnn = sum_mnn/real(num_kpts, dp)

      ! k-index is always 1 in SS method (k summation alread accomplished)
      ln_tmp_loc(:, :, 1) = aimag(log(sum_mnn(:, :))) - sheet(:, :, 1)

      rave = 0.0_dp
      do iw = 1, num_wann
        do ind = 1, 3
          do nn = 1, kmesh_info%nntot
            rave(ind, iw) = rave(ind, iw) + kmesh_info%wb(nn)*kmesh_info%bk(ind, nn, 1)* &
                            ln_tmp_loc(iw, nn, 1)
          end do
        end do
      end do
      rave = -rave

      rave2 = 0.0_dp
      do iw = 1, num_wann
        rave2(iw) = sum(rave(:, iw)*rave(:, iw))
      end do

      r2ave = 0.0_dp
      do nn = 1, kmesh_info%nntot
        r2ave(:) = r2ave(:) + kmesh_info%wb(nn)*(1.0_dp - abs(sum_mnn(:, nn))**2)
      end do

      r2ave = r2ave + rave2

      deallocate (sum_mnn, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error in deallocating sum_mnn in wann_omega', comm)
        return
      end if

    else ! not Stengel-Spaldin
      do nkp_loc = 1, nkrank
        nkp = global_k(nkp_loc)
        do nn = 1, kmesh_info%nntot
          do n = 1, num_wann
            ! Note that this ln_tmp is defined differently wrt the one in wann_domega
            ln_tmp_loc(n, nn, nkp_loc) = (aimag(log(csheet(n, nn, nkp) &
                                                    *m_matrix_loc(n, n, nn, nkp_loc))) - sheet(n, nn, nkp))
          end do
        end do
      end do

      rave = 0.0_dp
      do iw = 1, num_wann
        do ind = 1, 3
          do nkp_loc = 1, nkrank
            nkp = global_k(nkp_loc)
            do nn = 1, kmesh_info%nntot
              rave(ind, iw) = rave(ind, iw) + kmesh_info%wb(nn)*kmesh_info%bk(ind, nn, nkp) &
                              *ln_tmp_loc(iw, nn, nkp_loc)
            end do
          end do
        end do
      end do

      call comms_allreduce(rave(1, 1), num_wann*3, 'SUM', error, comm)
      if (allocated(error)) return

      rave = -rave/real(num_kpts, dp)

      rave2 = 0.0_dp
      do iw = 1, num_wann
        rave2(iw) = sum(rave(:, iw)*rave(:, iw))
      end do

      r2ave = 0.0_dp
      do iw = 1, num_wann
        do nkp_loc = 1, nkrank
          do nn = 1, kmesh_info%nntot
            mnn2 = real(m_matrix_loc(iw, iw, nn, nkp_loc)* &
                        conjg(m_matrix_loc(iw, iw, nn, nkp_loc)), kind=dp)
            r2ave(iw) = r2ave(iw) + kmesh_info%wb(nn)* &
                        (1.0_dp - mnn2 + ln_tmp_loc(iw, nn, nkp_loc)**2)
          end do
        end do
      end do

      call comms_allreduce(r2ave(1), num_wann, 'SUM', error, comm)
      if (allocated(error)) return

      r2ave = r2ave/real(num_kpts, dp)
    end if ! not Stengel-Spaldin

!~    wann_spread%om_1 = 0.0_dp
!~    do nkp = 1, num_kpts
!~       do nn = 1, nntot
!~          do loop_wann = 1, num_wann
!~             wann_spread%om_1 = wann_spread%om_1 + wb(nn) * &
!~                  ( 1.0_dp - m_matrix(loop_wann,loop_wann,nn,nkp) * &
!~                  conjg(m_matrix(loop_wann,loop_wann,nn,nkp)) )
!~          enddo
!~       enddo
!~    enddo
!~    wann_spread%om_1 = wann_spread%om_1 / real(num_kpts,dp)
!~
!~    wann_spread%om_2 = 0.0_dp
!~    do loop_wann = 1, num_wann
!~       sqim = 0.0_dp
!~       do nkp = 1, num_kpts
!~          do nn = 1, nntot
!~             sqim = sqim + wb(nn) * &
!~                  ( (aimag(log(csheet(loop_wann,nn,nkp) * &
!~                  m_matrix(loop_wann,loop_wann,nn,nkp))) - &
!~                  sheet(loop_wann,nn,nkp))**2 )
!~          enddo
!~       enddo
!~       sqim = sqim / real(num_kpts,dp)
!~       wann_spread%om_2 = wann_spread%om_2 + sqim
!~    enddo
!~
!~    wann_spread%om_3 = 0.0_dp
!~    do loop_wann = 1, num_wann
!~       bim = 0.0_dp
!~       do ind = 1, 3
!~          do nkp = 1, num_kpts
!~             do nn = 1, nntot
!~                bim(ind) = bim(ind) &
!~                     + wb(nn) * bk(ind,nn,nkp) &
!~                     * ( aimag(log(csheet(loop_wann,nn,nkp) &
!~                     * m_matrix(loop_wann,loop_wann,nn,nkp))) &
!~                     - sheet(loop_wann,nn,nkp) )
!~             enddo
!~          enddo
!~       enddo
!~       bim = bim/real(num_kpts,dp)
!~       bim2 = 0.0_dp
!~       do ind = 1, 3
!~          bim2 = bim2 + bim (ind) * bim (ind)
!~       enddo
!~       wann_spread%om_3 = wann_spread%om_3 - bim2
!~    enddo

    !jry: Either the above (om1,2,3) or the following is redundant
    !     keep it in the code base for testing

    if (wann_slwf%selective_loc) then

      if (use_ss_functional) then
        call set_error_alloc(error, 'finish ss_functional and selective_loc combination', comm)
        return
      end if

      wann_spread%om_iod = 0.0_dp
      do nkp_loc = 1, nkrank
        do nn = 1, kmesh_info%nntot
          summ = 0.0_dp
          do n = 1, wann_slwf%slwf_num
            summ = summ + real(m_matrix_loc(n, n, nn, nkp_loc)* &
                               conjg(m_matrix_loc(n, n, nn, nkp_loc)), kind=dp)
            if (wann_slwf%constrain) then
              !! Centre constraint contribution. Zero if slwf_constrain=false
              summ = summ - lambda_loc*ln_tmp_loc(n, nn, nkp_loc)**2
            end if
          end do
          wann_spread%om_iod = wann_spread%om_iod + kmesh_info%wb(nn)* &
                               (real(wann_slwf%slwf_num, dp) - summ)
        end do
      end do

      call comms_allreduce(wann_spread%om_iod, 1, 'SUM', error, comm)
      if (allocated(error)) return

      wann_spread%om_iod = wann_spread%om_iod/real(num_kpts, dp)

      wann_spread%om_d = 0.0_dp
      do nkp_loc = 1, nkrank
        nkp = global_k(nkp_loc)
        do nn = 1, kmesh_info%nntot
          do n = 1, wann_slwf%slwf_num
            brn = sum(kmesh_info%bk(:, nn, nkp)*rave(:, n))
            ! if SS, only nkp_loc = 1 nonzero
            wann_spread%om_d = wann_spread%om_d + (1.0_dp - lambda_loc)*kmesh_info%wb(nn) &
                               *(ln_tmp_loc(n, nn, nkp_loc) + brn)**2
          end do
        end do
      end do

      call comms_allreduce(wann_spread%om_d, 1, 'SUM', error, comm)
      if (allocated(error)) return

      if (.not. use_ss_functional) then !JJ
        wann_spread%om_d = wann_spread%om_d/real(num_kpts, dp)
      end if

      wann_spread%om_nu = 0.0_dp
      !! Contribution from constrains on centres
      if (wann_slwf%constrain) then
        do nkp_loc = 1, nkrank
          nkp = global_k(nkp_loc)
          do nn = 1, kmesh_info%nntot
            do n = 1, wann_slwf%slwf_num
              wann_spread%om_nu = wann_spread%om_nu + 2.0_dp*kmesh_info%wb(nn)* &
                                  ln_tmp_loc(n, nn, nkp_loc)*lambda_loc* &
                                  sum(kmesh_info%bk(:, nn, nkp)*wann_slwf%centres(n, :))
            end do
          end do
        end do

        call comms_allreduce(wann_spread%om_nu, 1, 'SUM', error, comm)
        if (allocated(error)) return

        wann_spread%om_nu = wann_spread%om_nu/real(num_kpts, dp)

        do n = 1, wann_slwf%slwf_num
          wann_spread%om_nu = wann_spread%om_nu &
                              + lambda_loc*sum(wann_slwf%centres(n, :)**2)
        end do

      end if

      wann_spread%om_tot = wann_spread%om_iod + wann_spread%om_d + wann_spread%om_nu
      !! wann_spread%om_c = wann_spread%om_iod + wann_spread%om_d + wann_spread%om_nu
    else ! not selective localisation
      if (first_pass) then
        wann_spread%om_i = 0.0_dp
        !nkp = nkp_loc + displs(my_node_id)
        do nkp_loc = 1, nkrank
          do nn = 1, kmesh_info%nntot
            summ = 0.0_dp
            do m = 1, num_wann
              do n = 1, num_wann
                summ = summ &
                       + real(m_matrix_loc(n, m, nn, nkp_loc) &
                              *conjg(m_matrix_loc(n, m, nn, nkp_loc)), kind=dp)
              end do
            end do
            wann_spread%om_i = wann_spread%om_i &
                               + kmesh_info%wb(nn)*(real(num_wann, dp) - summ)
          end do
        end do

        call comms_allreduce(wann_spread%om_i, 1, 'SUM', error, comm)
        if (allocated(error)) return

        wann_spread%om_i = wann_spread%om_i/real(num_kpts, dp)
        first_pass = .false.
      else
        wann_spread%om_i = omega_invariant
      end if

      wann_spread%om_od = 0.0_dp
      do nkp_loc = 1, nkrank
        do nn = 1, kmesh_info%nntot
          do m = 1, num_wann
            do n = 1, num_wann
              if (m .ne. n) wann_spread%om_od = wann_spread%om_od &
                                                + kmesh_info%wb(nn)*real(m_matrix_loc(n, m, nn, nkp_loc) &
                                                                         *conjg(m_matrix_loc(n, m, nn, nkp_loc)), kind=dp)
            end do
          end do
        end do
      end do

      call comms_allreduce(wann_spread%om_od, 1, 'SUM', error, comm)
      if (allocated(error)) return

      wann_spread%om_od = wann_spread%om_od/real(num_kpts, dp)

      if (use_ss_functional) then

        wann_spread%om_d = 0.0_dp

        do nn = 1, kmesh_info%nntot
          do n = 1, num_wann
            summ = 0.0_dp
            do nkp_loc = 1, nkrank
              nkp = global_k(nkp_loc)
              cnn = kmesh_info%nnord(nn, nkp) ! enforce uniform order of bk vectors
              summ = summ + m_matrix_loc(n, n, cnn, nkp_loc)
            end do

            call comms_allreduce(summ, 1, 'SUM', error, comm)
            if (allocated(error)) return
            summ = summ/real(num_kpts, dp)

            wann_spread%om_d = wann_spread%om_d - kmesh_info%wb(nn)*abs(summ)**2

            summ = 0.0_dp
            do nkp_loc = 1, nkrank
              nkp = global_k(nkp_loc)
              cnn = kmesh_info%nnord(nn, nkp) ! enforce uniform order of bk vectors
              summ = summ + abs(m_matrix_loc(n, n, cnn, nkp_loc))**2
            end do

            call comms_allreduce(summ, 1, 'SUM', error, comm)
            if (allocated(error)) return

            summ = summ/real(num_kpts, dp)

            wann_spread%om_d = wann_spread%om_d + kmesh_info%wb(nn)*summ
          end do
        end do
      else ! not Stengel-Spaldin
        wann_spread%om_d = 0.0_dp
        do nkp_loc = 1, nkrank
          nkp = global_k(nkp_loc)
          do nn = 1, kmesh_info%nntot
            do n = 1, num_wann
              brn = sum(kmesh_info%bk(:, nn, nkp)*rave(:, n))
              wann_spread%om_d = wann_spread%om_d + kmesh_info%wb(nn) &
                                 *(ln_tmp_loc(n, nn, nkp_loc) + brn)**2
            end do
          end do
        end do

        call comms_allreduce(wann_spread%om_d, 1, 'SUM', error, comm)
        if (allocated(error)) return

        wann_spread%om_d = wann_spread%om_d/real(num_kpts, dp)
      end if

      wann_spread%om_tot = wann_spread%om_i + wann_spread%om_d + wann_spread%om_od
    end if

    if (print_output%timing_level > 1 .and. print_output%iprint > 0) call io_stopwatch_stop('wann: omega', timer)
    return

  end subroutine wann_omega