dis_extract Subroutine

  subroutine dis_extract(dis_control, kmesh_info, sitesym, print_output, dis_manifold, &
                         m_matrix_orig_local, u_matrix_opt, eigval_opt, omega_invariant, &
                         indxnfroz, ndimfroz, my_node_id, num_bands, num_kpts, num_wann, &
                         lsitesymmetry, timer, ranknk, global_k, error, stdout, comm)
    !================================================!
    !
    !! Extracts an num_wann-dimensional subspace at each k by
    !! minimizing Omega_I
    !
    !================================================!

    ! MODIFIED:
    !           u_matrix_opt (At input it contains the initial guess for the optima
    ! subspace (expressed in terms of the original states inside the window)
    ! output it contains the  states that diagonalize the hamiltonian inside
    ! optimal subspace (again expressed in terms of the original window stat
    ! Giving out states that diagonalize the hamiltonian inside the optimal
    ! subspace (instead of the eigenstates of the Z matrix) is useful for
    ! performing the Wannier interpolation of the energy bands as described
    ! Sec. III.F of SMV)
    !
    !           eigval (At input: original energy eigenvalues.
    ! At output: eigenvalues of the hamiltonian inside optimal subspace)

    ! ----------------------------------------------------------------------
    ! TO DO: The complement subspace is computed but is not saved anywhere!
    ! (Check what was done with it in original code space.f)
    ! Diagonalize Z matrix only at those k points where ndimwin>num_wann?
    ! ----------------------------------------------------------------------

    ! *******************
    ! SHELLS OF K-VECTORS
    ! *******************
    ! nshells           number of shells of k-points to be used in the
    !                   finite-difference formulas for the k-derivatives
    ! aam: wb is now wb(1:nntot) 09/04/2006
    ! wb(nkp,nnx)       weight of the nnx-th b-vector (ordered along shells
    !                   of increasing length) associated with the nkp-th k-p
    ! wbtot             sum of the weights of all b-vectors associated with
    !                   given k-point (k-point 1 is used in calculation)
    ! nnlist(nkp,nnx)   vkpt(1:3,nnlist(nkp,nnx)) is the nnx-th neighboring
    !                   k-point of the nkp-th k-point vkpt(1:3,nkp) (or its
    !                   periodic image in the "home Brillouin zone")
    ! cm(n,m,nkp,nnx)   Overlap matrix <u_nk|u_{m,k+b}>
    use w90_comms, only: comms_bcast, comms_allreduce, w90_comm_type, mpirank
    use w90_constants, only: dp, cmplx_0, cmplx_1
    use w90_error
    use w90_io, only: io_stopwatch_start, io_stopwatch_stop, io_wallclocktime
    use w90_sitesym, only: sitesym_symmetrize_u_matrix, sitesym_symmetrize_zmatrix, &
                           sitesym_dis_extract_symmetry
    use w90_types, only: dis_manifold_type, kmesh_info_type, print_output_type, timer_list_type
    use w90_wannier90_types, only: dis_control_type, dis_spheres_type, sitesym_type

    implicit none

    ! arguments
    type(dis_control_type), intent(in) :: dis_control
    type(dis_manifold_type), intent(in) :: dis_manifold
    type(kmesh_info_type), intent(in) :: kmesh_info
    type(print_output_type), intent(in) :: print_output
    type(sitesym_type), intent(in) :: sitesym
    type(timer_list_type), intent(inout) :: timer
    type(w90_comm_type), intent(in) :: comm
    type(w90_error_type), allocatable, intent(out) :: error

    integer, intent(in) :: my_node_id
    integer, intent(in) :: stdout
    integer, intent(in) :: num_bands, num_kpts, num_wann
    integer, intent(in) :: ndimfroz(:) ! (num_kpts)
    integer, intent(in) :: indxnfroz(:, :) ! (num_bands,num_kpts)
    integer, intent(in) :: ranknk, global_k(:)

    real(kind=dp), intent(inout) :: eigval_opt(:, :)
    real(kind=dp), intent(out) :: omega_invariant

    complex(kind=dp), intent(inout) :: m_matrix_orig_local(:, :, :, :)
    complex(kind=dp), intent(inout) :: u_matrix_opt(:, :, :)

    logical, intent(in) :: lsitesymmetry

    ! Internal variables
    integer :: i, j, l, m, n, nn, nkp, nkp2, info, ierr, ndimk, p
    integer :: icompflag, iter, ndiff
    real(kind=dp) :: womegai, wkomegai, womegai1, rsum, delta_womegai
    real(kind=dp), allocatable :: wkomegai1(:)
    real(kind=dp), allocatable :: history(:)
    real(kind=dp), allocatable :: rwork(:)
    real(kind=dp), allocatable :: w(:)
    real(kind=dp), allocatable :: wkomegai1_loc(:)

    complex(kind=dp), allocatable :: camp_loc(:, :, :)
    complex(kind=dp), allocatable :: u_matrix_opt_loc(:, :, :)
    complex(kind=dp), allocatable :: ceamp(:, :, :) ! (alloc on root rank only)
    complex(kind=dp), allocatable :: camp(:, :, :)
    complex(kind=dp), allocatable :: czmat_in(:, :, :)
    complex(kind=dp), allocatable :: czmat_out(:, :, :)
    ! the z-matrices are now stored in local arrays
    complex(kind=dp), allocatable :: czmat_in_loc(:, :, :)
    complex(kind=dp), allocatable :: czmat_out_loc(:, :, :)
    complex(kind=dp), allocatable :: cham(:, :, :) ! (alloc on root rank only)

    complex(kind=dp), allocatable :: cap(:)
    complex(kind=dp), allocatable :: cwb(:, :), cww(:, :), cbw(:, :)
    complex(kind=dp), allocatable :: cwork(:)
    complex(kind=dp), allocatable :: cz(:, :)
    complex(kind=dp) :: lambda(num_wann, num_wann) !RS:

    integer, allocatable :: ifail(:)
    integer, allocatable :: iwork(:)
    integer :: nkp_loc

    logical :: dis_converged
    logical :: on_root = .false.

    on_root = (my_node_id == 0)

    if (print_output%timing_level > 1) call io_stopwatch_start('dis: extract', timer)

    if (print_output%iprint > 0) write (stdout, '(/1x,a)') &
      '                  Extraction of optimally-connected subspace                  '
    if (print_output%iprint > 0) write (stdout, '(1x,a)') &
      '                  ------------------------------------------                  '

    allocate (cwb(num_wann, num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating cwb in dis_extract', comm)
      return
    end if
    allocate (cww(num_wann, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating cww in dis_extract', comm)
      return
    end if
    allocate (cbw(num_bands, num_wann), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating cbw in dis_extract', comm)
      return
    end if
    allocate (iwork(5*num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating iwork in dis_extract', comm)
      return
    end if
    allocate (ifail(num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating ifail in dis_extract', comm)
      return
    end if
    allocate (w(num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating w in dis_extract', comm)
      return
    end if
    allocate (rwork(7*num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating rwork in dis_extract', comm)
      return
    end if
    allocate (cap((num_bands*(num_bands + 1))/2), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating cap in dis_extract', comm)
      return
    end if
    allocate (cwork(2*num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating cwork in dis_extract', comm)
      return
    end if
    allocate (cz(num_bands, num_bands), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating cz in dis_extract', comm)
      return
    end if
    allocate (u_matrix_opt_loc(num_bands, num_wann, ranknk), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating u_matrix_opt_loc in dis_extract', comm)
      return
    end if
    allocate (wkomegai1_loc(ranknk), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating wkomegai1_loc in dis_extract', comm)
      return
    end if
    allocate (czmat_in_loc(num_bands, num_bands, ranknk), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating czmat_in_loc in dis_extract', comm)
      return
    end if
    allocate (czmat_out_loc(num_bands, num_bands, ranknk), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating czmat_out_loc in dis_extract', comm)
      return
    end if
    allocate (wkomegai1(num_kpts), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating wkomegai1 in dis_extract', comm)
      return
    end if
    if (lsitesymmetry) then !we only need these large arrays for the symmetry code
      allocate (czmat_in(num_bands, num_bands, num_kpts), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error allocating czmat_in in dis_extract', comm)
        return
      end if
      allocate (czmat_out(num_bands, num_bands, num_kpts), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error allocating czmat_out in dis_extract', comm)
        return
      end if
    end if
    allocate (history(dis_control%conv_window), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error allocating history in dis_extract', comm)
      return
    end if

    cwb = cmplx_0; cww = cmplx_0; cbw = cmplx_0

    ! Copy matrix elements from global U matrix to local U matrix
    do nkp_loc = 1, ranknk
      nkp = global_k(nkp_loc)
      u_matrix_opt_loc(:, :, nkp_loc) = u_matrix_opt(:, :, nkp)
    end do

    ! ********************************************
    ! ENERGY WINDOWS AND SUBSPACES AT EACH K-POINT
    ! ********************************************
    ! num_wann             dimensionality of the subspace at each k-point
    !                   (number of Wannier functions per unit cell that we w
    ! NDIMWIN(NKP)      number of bands at the nkp-th k-point that fall
    !                   within the outer energy window
    ! NDIMFROZ(NKP)     number of frozen bands at the nkp-th k-point
    ! INDXNFROZ(I,NKP)  INDEX (BETWEEN 1 AND NDIMWIN(NKP)) OF THE I-TH NON-F
    !                   ORIGINAL BAND STATE AT THE NKP-TH K-POINT
    ! U_MATRIX_OPT(J,L,NKP)    AMPLITUDE OF THE J-TH ENERGY EIGENVECTOR INSIDE THE
    !                   ENERGY WINDOW AT THE NKP-TH K-POINT IN THE EXPANSION
    !                   THE L-TH LEADING RLAMBDA EIGENVECTOR AT THE SAME K-P
    !                   If there are M_k frozen states, they occupy the lowe
    !                   entries of the second index of u_matrix_opt, and the leadin
    !                   nabnds-M_k eigenvectors of the Z matrix occupy the
    !                   remaining slots
    ! CAMP(J,L,NKP)     SAME AS U_MATRIX_OPT, BUT FOR THE COMPLEMENT SUBSPACE INSID
    !                   ENERGY WINDOW (I.E., THE NON-LEADING RLAMBDA EIGENVE
    ! CEAMP(J,L,NKPTS)  SAME AS U_MATRIX_OPT, BUT INSTEAD OF RLAMBDA EIGENVECTOR, I
    !                   FOR THE ENERGY EIGENVECTOR OBTAINED BY DIAGONALIZING
    !                   HAMILTONIAN IN THE OPTIMIZED SUBSPACE
    ! CZMAT_IN(M,N,NKP) Z-MATRIX [Eq. (21) SMV]
    ! CZMAT_OUT(M,N,NKP) OUTPUT Z-MATRIX FROM THE PRESENT ITERATION
    ! RLAMBDA           An eigenvalue of the Z matrix
    ! womegai           Gauge-invariant Wannier spread, computed usinf all s
    !                   from current iteration
    ! wkomegai1(NKP)    Eq. (18) of SMV
    ! womegai1          Eq.(11) of SMV (like wowmegai, but neighboring state
    !                   for computing overlaps are from previous iteration.
    !                   become equal at self-consistency)
    ! alphafixe         mixing parameter for the iterative procedure
    ! nitere            total number of iterations

    ! DEBUG
    if (print_output%iprint > 2) then
      if (on_root) then
        write (stdout, '(a,/)') '  Original eigenvalues inside outer window:'
        do nkp = 1, num_kpts
          write (stdout, '(a,i3,3x,20(f9.5,1x))') '  K-point ', nkp, &
            (eigval_opt(i, nkp), i=1, dis_manifold%ndimwin(nkp))
        end do
      end if
    end if
    ! ENDDEBUG

    ! TO DO: Check if this is the best place to initialize icompflag
    icompflag = 0

    if (print_output%iprint > 0) write (stdout, '(1x,a)') &
      '+---------------------------------------------------------------------+<-- DIS'
    if (print_output%iprint > 0) write (stdout, '(1x,a)') &
      '|  Iter     Omega_I(i-1)      Omega_I(i)      Delta (frac.)    Time   |<-- DIS'
    if (print_output%iprint > 0) write (stdout, '(1x,a)') &
      '+---------------------------------------------------------------------+<-- DIS'

    dis_converged = .false.

    ! ------------------
    ! BIG ITERATION LOOP
    ! ------------------
    do iter = 1, dis_control%num_iter

      if (print_output%timing_level > 1) call io_stopwatch_start('dis: extract_1', timer)

      if (iter .eq. 1) then
        ! Initialize Z matrix at k points w/ non-frozen states
        do nkp_loc = 1, ranknk
          nkp = global_k(nkp_loc)
          if (num_wann .gt. ndimfroz(nkp)) then
            call internal_zmatrix(cbw, czmat_in_loc(:, :, nkp_loc), m_matrix_orig_local, &
                                  u_matrix_opt, kmesh_info%wb, indxnfroz, ndimfroz, &
                                  dis_manifold%ndimwin, kmesh_info%nnlist, nkp, nkp_loc, &
                                  kmesh_info%nntot, num_bands, num_wann, print_output%timing_level, &
                                  timer)
          end if
        end do

        if (lsitesymmetry) then
          czmat_in = cmplx_0
          do nkp_loc = 1, ranknk
            nkp = global_k(nkp_loc)
            czmat_in(:, :, nkp) = czmat_in_loc(:, :, nkp_loc)
          end do
          call comms_allreduce(czmat_in(1, 1, 1), num_bands*num_bands*num_kpts, 'SUM', error, comm)
          if (allocated(error)) return
          call sitesym_symmetrize_zmatrix(sitesym, czmat_in, num_bands, num_kpts, &
                                          dis_manifold%lwindow)
          do nkp_loc = 1, ranknk
            nkp = global_k(nkp_loc)
            czmat_in_loc(:, :, nkp_loc) = czmat_in(:, :, nkp)
          end do
        end if

      else
        ! [iter.ne.1]
        ! Update Z matrix at k points with non-frozen states, using a mixing sch
        do nkp_loc = 1, ranknk
          nkp = global_k(nkp_loc)
          if (lsitesymmetry) then
            if (sitesym%ir2ik(sitesym%ik2ir(nkp)) .ne. nkp) cycle
          end if
          if (num_wann .gt. ndimfroz(nkp)) then
            ndimk = dis_manifold%ndimwin(nkp) - ndimfroz(nkp)
            do i = 1, ndimk
              do j = 1, i
                czmat_in_loc(j, i, nkp_loc) = &
                  cmplx(dis_control%mix_ratio, 0.0_dp, dp)*czmat_out_loc(j, i, nkp_loc) &
                  + cmplx(1.0_dp - dis_control%mix_ratio, 0.0_dp, dp)*czmat_in_loc(j, i, nkp_loc)
                ! hermiticity
                czmat_in_loc(i, j, nkp_loc) = conjg(czmat_in_loc(j, i, nkp_loc))
              end do
            end do
          end if
        end do
      end if

      if (print_output%timing_level > 1) call io_stopwatch_stop('dis: extract_1', timer)
      if (print_output%timing_level > 1) call io_stopwatch_start('dis: extract_2', timer)

      womegai1 = 0.0_dp
      ! wkomegai1 is defined by Eq. (18) of SMV.
      ! Contribution to wkomegai1 from frozen states should be calculated now
      ! every k (before updating any k), so that for iter>1 overlaps are with
      ! non-frozen neighboring states from the previous iteration
      wkomegai1 = real(num_wann, dp)*kmesh_info%wbtot

      if (lsitesymmetry) then
        do nkp = 1, sitesym%nkptirr
          wkomegai1(sitesym%ir2ik(nkp)) = wkomegai1(sitesym%ir2ik(nkp))* &
                                          sitesym%nsymmetry/count(sitesym%kptsym(:, nkp) .eq. sitesym%ir2ik(nkp))
        end do
      end if
      do nkp_loc = 1, ranknk
        nkp = global_k(nkp_loc)
        wkomegai1_loc(nkp_loc) = wkomegai1(nkp)
      end do

      do nkp_loc = 1, ranknk
        nkp = global_k(nkp_loc)
        if (ndimfroz(nkp) .gt. 0) then
          if (lsitesymmetry) then
            call set_error_fatal(error, 'not implemented in symmetry-adapted mode', comm)
            return
          end if
          do nn = 1, kmesh_info%nntot
            nkp2 = kmesh_info%nnlist(nkp, nn)
            call zgemm('C', 'N', ndimfroz(nkp), dis_manifold%ndimwin(nkp2), dis_manifold%ndimwin(nkp), &
                       cmplx_1, u_matrix_opt(:, :, nkp), num_bands, &
                       m_matrix_orig_local(:, :, nn, nkp_loc), num_bands, cmplx_0, cwb, num_wann)
            call zgemm('N', 'N', ndimfroz(nkp), num_wann, dis_manifold%ndimwin(nkp2), cmplx_1, cwb, &
                       num_wann, u_matrix_opt(:, :, nkp2), num_bands, cmplx_0, cww, num_wann)
            rsum = 0.0_dp
            do n = 1, num_wann
              do m = 1, ndimfroz(nkp)
                rsum = rsum + real(cww(m, n), dp)**2 + aimag(cww(m, n))**2
              end do
            end do
            wkomegai1_loc(nkp_loc) = wkomegai1_loc(nkp_loc) - kmesh_info%wb(nn)*rsum
          end do
        end if
      end do

      if (print_output%timing_level > 1) call io_stopwatch_stop('dis: extract_2', timer)
      if (print_output%timing_level > 1) call io_stopwatch_start('dis: extract_3', timer)

      !! ! send chunks of wkomegai1 to root node
      !! call comms_gatherv(wkomegai1_loc, counts(my_node_id), wkomegai1, counts, displs)
      !! ! send back the whole wkomegai1 array to other nodes
      !! call comms_bcast(wkomegai1(1), num_kpts)

      ! Refine optimal subspace at k points w/ non-frozen states
      do nkp_loc = 1, ranknk
        nkp = global_k(nkp_loc)
        if (lsitesymmetry) then
          if (sitesym%ir2ik(sitesym%ik2ir(nkp)) .ne. nkp) cycle

          call sitesym_dis_extract_symmetry(sitesym, lambda, u_matrix_opt_loc(:, :, nkp_loc), &
                                            czmat_in_loc(:, :, nkp_loc), nkp, &
                                            dis_manifold%ndimwin(nkp), num_bands, num_wann, &
                                            stdout, error, comm)
          if (allocated(error)) return
          do j = 1, num_wann
            wkomegai1_loc(nkp_loc) = wkomegai1_loc(nkp_loc) - real(lambda(j, j), kind=dp)
          end do
        else
          if (num_wann .gt. ndimfroz(nkp)) then
            ! Diagonalize Z matrix
            do j = 1, dis_manifold%ndimwin(nkp) - ndimfroz(nkp)
              do i = 1, j
                cap(i + ((j - 1)*j)/2) = czmat_in_loc(i, j, nkp_loc)
              end do
            end do
            ndiff = dis_manifold%ndimwin(nkp) - ndimfroz(nkp)
            call ZHPEVX('V', 'A', 'U', ndiff, cap, 0.0_dp, 0.0_dp, 0, 0, &
                        -1.0_dp, m, w, cz, num_bands, cwork, rwork, iwork, ifail, info)
            if (info .lt. 0) then
              if (on_root) then
                write (stdout, *) ' *** ERROR *** ZHPEVX WHILE DIAGONALIZING Z MATRIX'
                write (stdout, *) ' THE ', -info, ' ARGUMENT OF ZHPEVX HAD AN ILLEGAL VALUE'
              end if
              call set_error_fatal(error, ' dis_extract: error', comm)
              return
            end if
            if (info .gt. 0) then
              if (on_root) write (stdout, *) ' *** ERROR *** ZHPEVX WHILE DIAGONALIZING Z MATRIX'
              if (on_root) write (stdout, *) info, ' EIGENVECTORS FAILED TO CONVERGE'
              call set_error_fatal(error, ' dis_extract: error', comm)
              return
            end if

            ! Update the optimal subspace by incorporating the num_wann-ndimfroz(nkp) l
            ! eigenvectors of the Z matrix into u_matrix_opt. Also, add contribution from
            ! non-frozen states to wkomegai1(nkp) (minus the corresponding eigenvalu
            m = ndimfroz(nkp)
            do j = dis_manifold%ndimwin(nkp) - num_wann + 1, dis_manifold%ndimwin(nkp) - ndimfroz(nkp)
              m = m + 1
              wkomegai1_loc(nkp_loc) = wkomegai1_loc(nkp_loc) - w(j)
              u_matrix_opt_loc(1:dis_manifold%ndimwin(nkp), m, nkp_loc) = cmplx_0
              ndimk = dis_manifold%ndimwin(nkp) - ndimfroz(nkp)
              do i = 1, ndimk
                p = indxnfroz(i, nkp)
                u_matrix_opt_loc(p, m, nkp_loc) = cz(i, j)
              end do
            end do
          end if
          ! [if num_wann>ndimfroz(nkp)]
        end if !RS:

        ! Now that we have contribs. from both frozen and non-frozen states to
        ! wkomegai1(nkp), add it to womegai1
        womegai1 = womegai1 + wkomegai1_loc(nkp_loc)

        !if (index(print_output%devel_flag, 'compspace') > 0) then

        ! AT THE LAST ITERATION FIND A BASIS FOR THE (NDIMWIN(NKP)-num_wann)-DIMENS
        ! COMPLEMENT SPACE

        !  if (iter .eq. dis_control%num_iter) then
        !    allocate (camp(num_bands, num_bands, num_kpts), stat=ierr)
        !    if (ierr /= 0) then
        !      call set_error_alloc(error, 'Error allocating camp in dis_extract', comm)
        !      return
        !    endif
        !    allocate (camp_loc(num_bands, num_bands, max(1, counts(my_node_id))), stat=ierr)
        !    if (ierr /= 0) then
        !      call set_error_alloc(error, 'Error allocating ucamp_loc in dis_extract', comm)
        !      return
        !    endif
        !    if (dis_manifold%ndimwin(nkp) .gt. num_wann) then
        !      do j = 1, dis_manifold%ndimwin(nkp) - num_wann
        !        if (num_wann .gt. ndimfroz(nkp)) then
        !          ! USE THE NON-LEADING EIGENVECTORS OF THE Z-MATRIX
        !          camp_loc(1:dis_manifold%ndimwin(nkp), j, nkp_loc) = cz(1:dis_manifold%ndimwin(nkp), j)
        !        else
        ! Then num_wann=NDIMFROZ(NKP)
        ! USE THE ORIGINAL NON-FROZEN BLOCH EIGENSTATES
        !          do i = 1, dis_manifold%ndimwin(nkp)
        !            camp_loc(i, j, nkp_loc) = cmplx_0
        !            if (i .eq. indxnfroz(j, nkp)) camp_loc(i, j, nkp_loc) = cmplx_1
        !          enddo
        !        endif
        !      enddo
        !    else
        !      icompflag = 1
        !    endif
        !  endif

        !end if ! index(print_output%devel_flag,'compspace')>0

      end do
      ! [Loop over k points (nkp)]

      !! ! send chunks of wkomegai1 to root node
      !! call comms_gatherv(wkomegai1_loc, counts(my_node_id), wkomegai1, counts, displs)
      !! ! send back the whole wkomegai1 array to other nodes
      !! call comms_bcast(wkomegai1(1), num_kpts)

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

      u_matrix_opt(:, :, :) = cmplx_0
      do nkp_loc = 1, ranknk
        nkp = global_k(nkp_loc)
        u_matrix_opt(:, :, nkp) = u_matrix_opt_loc(:, :, nkp_loc)
      end do
      call comms_allreduce(u_matrix_opt(1, 1, 1), num_bands*num_wann*num_kpts, 'SUM', error, comm)
      if (allocated(error)) return

      if (lsitesymmetry) then
        call sitesym_symmetrize_u_matrix(sitesym, u_matrix_opt, num_bands, num_bands, num_kpts, &
                                         num_wann, stdout, error, comm, dis_manifold%lwindow)
        if (allocated(error)) return
        do nkp_loc = 1, ranknk
          nkp = global_k(nkp_loc)
          u_matrix_opt_loc(:, :, nkp_loc) = u_matrix_opt(:, :, nkp)
        end do
      end if

      if (print_output%timing_level > 1) call io_stopwatch_stop('dis: extract_3', timer)

      womegai1 = womegai1/real(num_kpts, dp)

      ! DEBUG
      ! Orthonormality check
      !         do nkp=1,nkpts
      !           write(*,*) ' '
      !           write(*,'(a8,i4)') 'k-point ',nkp
      !           do l=1,num_wann
      !           do m=1,l
      !             ctmp=czero
      !             do j=1,ndimwin(nkp)
      !               ctmp=ctmp+conjg(u_matrix_opt(j,m,nkp))*u_matrix_opt(j,l,nkp)
      !             enddo
      !             write(*,'(i2,2x,i2,f16.12,1x,f16.12)') l,m,ctmp
      !             if(l.eq.m) then
      !               if(abs(ctmp-cmplx(1.0d0,0.0d0)).gt.1.0e-8) then
      !                 write(*,'(a49,i4)')
      !     1           '*** ERROR *** with iterative subspace at k-point ',
      !     2           nkp
      !                 write(*,*) 'vectors in u_matrix_opt not orthonormal'
      !                 stop
      !               endif
      !             else
      !               if(abs(ctmp).gt.1.0e-8) then
      !                 write(*,'(a49,i4)')
      !     1           '*** ERROR *** with iterative subspace at k-point ',
      !     2           nkp
      !                 write(*,*) 'vectors in u_matrix_opt not orthonormal'
      !                 stop
      !               endif
      !             endif
      !           enddo
      !           enddo
      !         enddo
      ! ENDDEBUG

      ! Compute womegai  using the updated subspaces at all k, i.e.,
      ! replacing (i-1) by (i) in Eq. (12) SMV
      if (print_output%timing_level > 1) call io_stopwatch_start('dis: extract_4', timer)

      womegai = 0.0_dp
      do nkp_loc = 1, ranknk
        nkp = global_k(nkp_loc)
        wkomegai = 0.0_dp
        do nn = 1, kmesh_info%nntot
          nkp2 = kmesh_info%nnlist(nkp, nn) ! nkp2 here may not be local to this processor hence operate on global u
          call zgemm('C', 'N', num_wann, dis_manifold%ndimwin(nkp2), dis_manifold%ndimwin(nkp), cmplx_1, &
                     u_matrix_opt(:, :, nkp), num_bands, m_matrix_orig_local(:, :, nn, nkp_loc), &
                     num_bands, cmplx_0, cwb, num_wann)
          call zgemm('N', 'N', num_wann, num_wann, dis_manifold%ndimwin(nkp2), cmplx_1, cwb, num_wann, &
                     u_matrix_opt(:, :, nkp2), num_bands, cmplx_0, cww, num_wann)
          rsum = 0.0_dp
          do n = 1, num_wann
            do m = 1, num_wann
              rsum = rsum + real(cww(m, n), dp)**2 + aimag(cww(m, n))**2
            end do
          end do
          wkomegai = wkomegai + kmesh_info%wb(nn)*rsum
        end do
        wkomegai = real(num_wann, dp)*kmesh_info%wbtot - wkomegai
        womegai = womegai + wkomegai
      end do

      call comms_allreduce(womegai, 1, 'SUM', error, comm)
      if (allocated(error)) return
      womegai = womegai/real(num_kpts, dp)
      delta_womegai = womegai1/womegai - 1.0_dp

      if (print_output%timing_level > 1) call io_stopwatch_stop('dis: extract_4', timer)

      if (print_output%iprint > 0) then
        write (stdout, 124) iter, womegai1*print_output%lenconfac**2, &
          womegai*print_output%lenconfac**2, delta_womegai, io_wallclocktime()
      end if
124   format(2x, i6, 3x, f14.8, 3x, f14.8, 6x, es10.3, 2x, f8.2, 4x, '<-- DIS')

      ! Construct the updated Z matrix, CZMAT_OUT, at k points w/ non-frozen s
      do nkp_loc = 1, ranknk
        nkp = global_k(nkp_loc)
        if (num_wann .gt. ndimfroz(nkp)) then
          call internal_zmatrix(cbw, czmat_out_loc(:, :, nkp_loc), m_matrix_orig_local, &
                                u_matrix_opt, kmesh_info%wb, indxnfroz, ndimfroz, &
                                dis_manifold%ndimwin, kmesh_info%nnlist, nkp, nkp_loc, &
                                kmesh_info%nntot, num_bands, num_wann, print_output%timing_level, &
                                timer)
        end if
      end do

      if (lsitesymmetry) then
        czmat_out = cmplx_0
        do nkp_loc = 1, ranknk
          nkp = global_k(nkp_loc)
          czmat_out(:, :, nkp) = czmat_out_loc(:, :, nkp_loc)
        end do
        call comms_allreduce(czmat_out(1, 1, 1), num_bands*num_bands*num_kpts, 'SUM', error, comm)
        if (allocated(error)) return
        call sitesym_symmetrize_zmatrix(sitesym, czmat_out, num_bands, num_kpts, dis_manifold%lwindow)
        do nkp_loc = 1, ranknk
          nkp = global_k(nkp_loc)
          czmat_out_loc(:, :, nkp_loc) = czmat_out(:, :, nkp)
        end do
      end if

      call internal_test_convergence(history, delta_womegai, dis_control%conv_tol, iter, &
                                     dis_control%conv_window, dis_converged, error, comm)
      if (allocated(error)) return

      if (dis_converged) then
        if (print_output%iprint > 0) then
          write (stdout, '(/13x,a,es10.3,a,i2,a)') '<<<      Delta <', dis_control%conv_tol, &
            '  over ', dis_control%conv_window, ' iterations     >>>'
          write (stdout, '(13x,a)') '<<< Disentanglement convergence criteria satisfied >>>'
        end if
        exit
      end if

    end do
    ! [BIG ITERATION LOOP (iter)]

    if (lsitesymmetry) then
      deallocate (czmat_out, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating czmat_out in dis_extract', comm)
        return
      end if
      deallocate (czmat_in, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating czmat_in in dis_extract', comm)
        return
      end if
    end if
    deallocate (czmat_out_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating czmat_out_loc in dis_extract', comm)
      return
    end if
    deallocate (czmat_in_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating czmat_in_loc in dis_extract', comm)
      return
    end if

    if (on_root) then
      allocate (ceamp(num_bands, num_bands, num_kpts), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error allocating ceamp in dis_extract', comm)
        return
      end if
      allocate (cham(num_bands, num_bands, num_kpts), stat=ierr)
      if (ierr /= 0) then
        call set_error_alloc(error, 'Error allocating cham in dis_extract', comm)
        return
      end if
    end if

    if (.not. dis_converged) then
      if (on_root) then
        write (stdout, '(/5x,a)') &
          '<<< Warning: Maximum number of disentanglement iterations reached >>>'
        write (stdout, '(10x,a)') '<<< Disentanglement convergence criteria not satisfied >>>'
      end if
    end if

    !if (index(print_output%devel_flag, 'compspace') > 0) then

    !  if (icompflag .eq. 1) then
    !    if (print_output%iprint > 2) then
    !      if (on_root) write (stdout, ('(/4x,a)')) &
    !        'WARNING: Complement subspace has zero dimensions at the following k-points:'
    !      i = 0
    !      if (on_root) write (stdout, '(4x)', advance='no')
    !      do nkp = 1, num_kpts
    !        if (dis_manifold%ndimwin(nkp) .eq. num_wann) then
    !          i = i + 1
    !          if (i .le. 12) then
    !            if (on_root) write (stdout, '(i6)', advance='no') nkp
    !          else
    !            i = 1
    !            if (on_root) write (stdout, '(/4x)', advance='no')
    !            if (on_root) write (stdout, '(i6)', advance='no') nkp
    !          endif
    !        endif
    !      enddo
    !    endif
    !  endif

    !endif

    ! Write the final womegai. This should remain unchanged during the
    ! subsequent minimization of Omega_tilde in wannierise.f90
    ! We store it in the checkpoint file as a sanity check
    if (print_output%iprint > 0) write (stdout, '(/8x,a,f14.8,a/)') 'Final Omega_I ', &
      womegai*print_output%lenconfac**2, ' ('//trim(print_output%length_unit)//'^2)'

    ! Set public variable omega_invariant
    omega_invariant = womegai

    ! Currently, this part is not parallelized; thus, we perform the task only on root and then broadcast the result.
    if (on_root) then
      ! DIAGONALIZE THE HAMILTONIAN WITHIN THE OPTIMIZED SUBSPACES
      do nkp = 1, num_kpts

        do j = 1, num_wann
          do i = 1, num_wann
            cham(i, j, nkp) = cmplx_0
            do l = 1, dis_manifold%ndimwin(nkp)
              cham(i, j, nkp) = cham(i, j, nkp) + conjg(u_matrix_opt(l, i, nkp)) &
                                *u_matrix_opt(l, j, nkp)*eigval_opt(l, nkp)
            end do
          end do
        end do

        do j = 1, num_wann
          do i = 1, j
            cap(i + ((j - 1)*j)/2) = cham(i, j, nkp)
          end do
        end do

        call ZHPEVX('V', 'A', 'U', num_wann, cap, 0.0_dp, 0.0_dp, 0, 0, -1.0_dp, m, w, cz, &
                    num_bands, cwork, rwork, iwork, ifail, info)

        if (info .lt. 0) then
          if (on_root) write (stdout, *) ' *** ERROR *** ZHPEVX WHILE DIAGONALIZING HAMILTONIAN'
          if (on_root) write (stdout, *) ' THE ', -info, ' ARGUMENT OF ZHPEVX HAD AN ILLEGAL VALUE'
          call set_error_fatal(error, ' dis_extract: error', comm)
          return
        end if
        if (info .gt. 0) then
          if (on_root) write (stdout, *) ' *** ERROR *** ZHPEVX WHILE DIAGONALIZING HAMILTONIAN'
          if (on_root) write (stdout, *) info, 'EIGENVECTORS FAILED TO CONVERGE'
          call set_error_fatal(error, ' dis_extract: error', comm)
          return
        end if

        ! Store the energy eigenvalues of the optimal subspace (used in wann_ban
        eigval_opt(1:num_wann, nkp) = w(1:num_wann)

        ! CALCULATE AMPLITUDES OF THE CORRESPONDING ENERGY EIGENVECTORS IN TERMS
        ! THE ORIGINAL ("WINDOW SPACE") ENERGY EIGENVECTORS
        do j = 1, num_wann
          do i = 1, dis_manifold%ndimwin(nkp)
            ceamp(i, j, nkp) = cmplx_0
            do l = 1, num_wann
              ceamp(i, j, nkp) = ceamp(i, j, nkp) + cz(l, j)*u_matrix_opt(i, l, nkp)
            end do
          end do
        end do
        ! NKP
      end do

      if (print_output%iprint > 2) then
        if (on_root) write (stdout, '(/,a,/)') '  Eigenvalues inside optimal subspace:'
        do nkp = 1, num_kpts
          if (on_root) write (stdout, '(a,i3,2x,20(f9.5,1x))') '  K-point ', &
            nkp, (eigval_opt(i, nkp), i=1, num_wann)
        end do
      end if

      ! Replace u_matrix_opt by ceamp. Both span the
      ! same space, but the latter is more convenient for the purpose of obtai
      ! an optimal Fourier-interpolated band structure: see Sec. III.E of SMV.
      if (.not. lsitesymmetry) then                                                                         !YN:
        do nkp = 1, num_kpts
          do j = 1, num_wann
            u_matrix_opt(1:dis_manifold%ndimwin(nkp), j, nkp) = ceamp(1:dis_manifold%ndimwin(nkp), j, nkp)
          end do
        end do
        !else                                                                                                        !YN:
        ! Above is skipped as we require Uopt(Rk) to be related to Uopt(k)                                        !YN: RS:
        !write(stdout,"(a)")  &                                                                                   !YN: RS:
        !  'Note(symmetry-adapted mode): u_matrix_opt are no longer the eigenstates of the subspace Hamiltonian.' !RS:
      end if                                                                                                        !YN:
    end if
    call comms_bcast(eigval_opt(1, 1), num_bands*num_kpts, error, comm)
    if (allocated(error)) return
    call comms_bcast(u_matrix_opt(1, 1, 1), num_bands*num_wann*num_kpts, error, comm)
    if (allocated(error)) return

    !if (index(print_output%devel_flag, 'compspace') > 0) then

    !  if (icompflag .eq. 1) then
    !    if (print_output%iprint > 2) then
    !      if (on_root) then
    !        write (stdout, *) 'AT SOME K-POINT(S) COMPLEMENT SUBSPACE HAS ZERO DIMENSIONALITY'
    !        write (stdout, *) '=> DID NOT CREATE FILE COMPSPACE.DAT'
    !      endif
    !    endif
    !  else
    ! DIAGONALIZE THE HAMILTONIAN IN THE COMPLEMENT SUBSPACE, WRITE THE
    ! CORRESPONDING EIGENFUNCTIONS AND ENERGY EIGENVALUES
    !    do nkp = 1, num_kpts
    !      do j = 1, dis_manifold%ndimwin(nkp) - num_wann
    !        do i = 1, dis_manifold%ndimwin(nkp) - num_wann
    !          cham(i, j, nkp) = cmplx_0
    !          do l = 1, dis_manifold%ndimwin(nkp)
    !            cham(i, j, nkp) = cham(i, j, nkp) + conjg(camp(l, i, nkp)) &
    !                              *camp(l, j, nkp)*eigval_opt(l, nkp)
    !          enddo
    !        enddo
    !      enddo
    !      do j = 1, dis_manifold%ndimwin(nkp) - num_wann
    !        do i = 1, j
    !          cap(i + ((j - 1)*j)/2) = cham(i, j, nkp)
    !        enddo
    !      enddo
    !      ndiff = dis_manifold%ndimwin(nkp) - num_wann
    !      call ZHPEVX('V', 'A', 'U', ndiff, cap, 0.0_dp, 0.0_dp, 0, 0, &
    !                  -1.0_dp, m, w, cz, num_bands, cwork, rwork, iwork, ifail, info)
    !      if (info .lt. 0) then
    !        if (on_root) write (stdout, *) '*** ERROR *** ZHPEVX WHILE DIAGONALIZING HAMILTONIAN'
    !        if (on_root) write (stdout, *) 'THE ', -info, ' ARGUMENT OF ZHPEVX HAD AN ILLEGAL VALUE'
    !        call io_error(' dis_extract: error', stdout, seedname)
    !      endif
    !      if (info .gt. 0) then
    !        if (on_root) write (stdout, *) '*** ERROR *** ZHPEVX WHILE DIAGONALIZING HAMILTONIAN'
    !        if (on_root) write (stdout, *) info, 'EIGENVECTORS FAILED TO CONVERGE'
    !        call io_error(' dis_extract: error', stdout, seedname)
    !      endif
    ! CALCULATE AMPLITUDES OF THE ENERGY EIGENVECTORS IN THE COMPLEMENT SUBS
    ! TERMS OF THE ORIGINAL ENERGY EIGENVECTORS
    !      do j = 1, dis_manifold%ndimwin(nkp) - num_wann
    !        do i = 1, dis_manifold%ndimwin(nkp)
    !          camp(i, j, nkp) = cmplx_0
    !          do l = 1, dis_manifold%ndimwin(nkp) - num_wann
!write(stdout,*) 'i=',i,'   j=',j,'   l=',l
!write(stdout,*) '           camp(i,j,nkp)=',camp(i,j,nkp)
!write(stdout,*) '           cz(l,j)=',cz(l,j)
!write(stdout,*) '           u_matrix_opt(i,l,nkp)=',u_matrix_opt(i,l,nkp)

! aam: 20/10/2006 -- the second dimension of u_matrix_opt is out of bounds (allocated as num_wann)!
! commenting this line out.
!                     camp(i,j,nkp) = camp(i,j,nkp) + cz(l,j) * u_matrix_opt(i,l,nkp)
    !         enddo
    !       enddo
    !     enddo
    !   enddo   ! [loop over k points (nkp)]

    ! endif   ! [if icompflag=1]

    !endif     ![if(index(devel_flag,'compspace')>0)]

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

    if (on_root) then
      deallocate (cham, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating cham in dis_extract', comm)
        return
      end if
    end if
    if (allocated(camp)) then
      deallocate (camp, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating camp in dis_extract', comm)
        return
      end if
    end if
    if (allocated(camp_loc)) then
      deallocate (camp_loc, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating camp_loc in dis_extract', comm)
        return
      end if
    end if
    if (on_root) then
      deallocate (ceamp, stat=ierr)
      if (ierr /= 0) then
        call set_error_dealloc(error, 'Error deallocating ceamp in dis_extract', comm)
        return
      end if
    end if
    deallocate (u_matrix_opt_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating u_matrix_opt_loc in dis_extract', comm)
      return
    end if
    deallocate (wkomegai1_loc, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating wkomegai1_loc in dis_extract', comm)
      return
    end if
    deallocate (wkomegai1, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating wkomegai1 in dis_extract', comm)
      return
    end if

    deallocate (cz, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating cz in dis_extract', comm)
      return
    end if
    deallocate (cwork, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating cwork in dis_extract', comm)
      return
    end if
    deallocate (cap, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating cap in dis_extract', comm)
      return
    end if
    deallocate (rwork, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating rwork in dis_extract', comm)
      return
    end if
    deallocate (w, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating w in dis_extract', comm)
      return
    end if
    deallocate (ifail, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating ifail in dis_extract', comm)
      return
    end if
    deallocate (iwork, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating iwork in dis_extract', comm)
      return
    end if

    deallocate (cbw, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating cbw in dis_extract', comm)
      return
    end if
    deallocate (cww, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating cww in dis_extract', comm)
      return
    end if
    deallocate (cwb, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating cwb in dis_extract', comm)
      return
    end if

    if (print_output%iprint > 0) write (stdout, '(1x,a/)') &
      '+----------------------------------------------------------------------------+'

    if (print_output%timing_level > 1) call io_stopwatch_stop('dis: extract', timer)

    return
    !================================================!
  end subroutine dis_extract