wannier_calc.F90

Go to the documentation of this file.

!! Copyright (C) 2017-2019 H. Huebener, N. Tancogne-Dejean
!! Copyright (C) 2026 A. Buccheri, J. Reimann
!!
!! This program is free software; you can redistribute it and/or modify
!! it under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 2, or (at your option)
!! any later version.
!!
!! This program is distributed in the hope that it will be useful,
!! but WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!! GNU General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with this program; if not, write to the Free Software
!! Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
!! 02110-1301, USA.
!!
 
#include "global.h"
 
module wannier_calc_oct_m
  use, intrinsic :: iso_fortran_env
 
  use batch_oct_m, only: batch_not_packed, batch_packed, batch_device_packed
  use comm_oct_m
  use debug_oct_m
  use electron_space_oct_m, only: spin_polarized
  use global_oct_m
  use hamiltonian_elec_oct_m, only: hamiltonian_elec_t
  use io_oct_m
  use io_function_oct_m, only: write_standard_xyz_file
  use ions_oct_m
  use lattice_vectors_oct_m, only: lattice_vectors_t
  use kpoints_oct_m
  use mesh_oct_m
  use mesh_function_oct_m, only: zmf_dotp
  use messages_oct_m
  use mpi_oct_m
  use multicomm_oct_m
  use namespace_oct_m
  use orbitalset_oct_m
  use profiling_oct_m
  use restart_oct_m
  use space_oct_m
  use states_elec_oct_m
  use states_elec_restart_oct_m
  use subspace_oct_m, only: zsubspace_diag_hamiltonian
  use submesh_oct_m
  use types_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use wannier_opts_oct_m
  use wfs_elec_oct_m, only: wfs_elec_t
  use ylm_wannier_oct_m
 
  implicit none
  private
 
  public :: &
    wannier_calc_load_restart, &
    wannier_calc_proj_t, &
    wannier_calc_create_amn, &
    wannier_calc_create_mmn, &
    wannier_calc_update_phase, &
    wannier_calc_read_centers, &
    wannier_calc_read_u_matrices, &
    wannier_calc_write_u_matrices, &
    wannier_calc_write_centers, &
    wannier_calc_write_spreads, &
    wannier_calc_write_hr_file, &
    wannier_calc_wigner_seitz_count, &
    wannier_calc_wigner_seitz_fill, &
    wannier_calc_compute_ham_r
 
  type wannier_calc_proj_t
    integer :: l, m, s
    integer :: radial = 1
    real(real64) :: site(3)
    real(real64) :: s_qaxis(3) = [0.0_real64, 0.0_real64, 1.0_real64]
    real(real64) :: z(3) = [0.0_real64, 0.0_real64, 1.0_real64]
    real(real64) :: x(3) = [1.0_real64, 0.0_real64, 0.0_real64]
    real(real64) :: zona = 1.0_real64
  end type wannier_calc_proj_t
 
contains
 
  subroutine wannier_calc_load_restart(namespace, mc, space, st, gr, kpoints, restart_states)
    type(namespace_t),   intent(in)    :: namespace
    type(multicomm_t),   intent(in)    :: mc
    class(space_t),      intent(inout) :: space
    type(states_elec_t), intent(inout) :: st
    class(mesh_t),       intent(inout) :: gr
    type(kpoints_t),     intent(inout) :: kpoints
    integer,   optional, intent(in)    :: restart_states ! RESTART_GS or RESTART_TD
    ! Defaults to ground state
 
    integer         :: restart_states_, ierr, nik, dim, nst
    type(restart_t) :: restart
 
    restart_states_ = optional_default(restart_states, restart_gs)
    if (.not. (restart_states_ == restart_gs .or. restart_states_ == restart_td)) then
      message(1) = 'Invalid restart_states value; expected RESTART_GS or RESTART_TD.'
      call messages_fatal(1)
    end if
 
    call states_elec_allocate_wfns(st, gr, wfs_type = type_cmplx)
    call restart%init(namespace, restart_states_, restart_type_load, mc, ierr, gr)
 
    if (ierr /= 0) then
      write(message(1),'(a)') 'Error opening restart file'
      call messages_fatal(1)
    endif
 
    call states_elec_look(restart, nik, dim, nst, ierr)
    if (st%d%ispin == spin_polarized) nik = nik / 2
    if (dim /= st%d%dim) then
      write(message(1),'(a)') 'Restart structure not commensurate, since spin dimensions differ.'
      call messages_fatal(1)
    else if (nik /= kpoints%reduced%npoints) then
      write(message(1),'(a)') 'Restart structure not commensurate, since number of k-points differ.'
      call messages_fatal(1)
    else if (nst /= st%nst) then
      write(message(1),'(a)') 'Restart structure not commensurate, since number of states differ.'
      call messages_fatal(1)
    end if
 
    call states_elec_load(restart, namespace, space, st, gr, kpoints, fixed_occ=st%restart_fixed_occ , &
      ierr=ierr, label = ": wannier90")
    call restart%end()
 
  end subroutine wannier_calc_load_restart
 
  ! TODO: Removed SCDM - Should reintroduce once the code is working
  subroutine wannier_calc_create_amn(wan_opts, &
    band_index, space, mesh, latt, st, kpoints, &
    spin_channel, num_proj, proj_input, projection)
    type(wannier_opts_t),                  intent(in)    :: wan_opts
    class(space_t),                        intent(in)    :: space
    class(mesh_t),                         intent(in)    :: mesh
    type(lattice_vectors_t),               intent(in)    :: latt
    type(states_elec_t),                   intent(in)    :: st
    type(kpoints_t),                       intent(in)    :: kpoints
    integer, contiguous,                   intent(in)    :: band_index(:)
    integer,                               intent(in)    :: spin_channel
    integer,                               intent(in)    :: num_proj
    type(wannier_calc_proj_t), contiguous, intent(in)    :: proj_input(:)
 
    complex(real64), contiguous, intent(out) :: projection(:, :, :)
 
    integer        ::  ist, ik, idim, iw, ip, ik_real, iband, num_kpts
    real(real64)   ::  center(3),  kpoint(3)
 
    complex(real64),    allocatable :: psi(:,:), phase(:)
    real(real64),       allocatable :: ylm(:)
    type(orbitalset_t), allocatable :: orbitals(:)
 
    push_sub(wannier_calc_create_amn)
    call profiling_in("W90_AMN")
 
    if (st%parallel_in_states) then
      call messages_not_implemented("w90_amn output with states parallelization")
    end if
 
    message(1) = "Info: Computing the projection matrix"
    call messages_info(1)
 
    assert(num_proj == size(proj_input))
 
    num_kpts = kpoints%full%npoints
 
    ! Precompute orbitals
    safe_allocate(orbitals(1:num_proj))
    do iw=1, num_proj
      call orbitalset_init(orbitals(iw))
 
      orbitals(iw)%norbs = 1
      orbitals(iw)%ndim = 1
      orbitals(iw)%radius = -log(wan_opts%threshold)
      orbitals(iw)%use_submesh = .false.
 
      ! cartesian coordinate of orbital center
      center(1:3) = latt%red_to_cart(proj_input(iw)%site(1:3))
      call submesh_init(orbitals(iw)%sphere, space, mesh, latt, center, orbitals(iw)%radius)
 
      ! get dorb as submesh points
      safe_allocate(ylm(1:orbitals(iw)%sphere%np))
      ! (this is a routine from pwscf)
      call ylm_wannier(ylm, proj_input(iw)%l, proj_input(iw)%m, &
        orbitals(iw)%sphere%r, orbitals(iw)%sphere%rel_x, orbitals(iw)%sphere%np)
 
      if (proj_input(iw)%radial /= 1) then
        call messages_not_implemented("oct-wannier90: r/=1 for the radial part")
      end if
 
      ! apply radial function
      do ip = 1, orbitals(iw)%sphere%np
        ylm(ip) = ylm(ip) * m_two * exp(-orbitals(iw)%sphere%r(ip))
      end do
 
      ! NOTE(Alex) One can probably fill directly into zorb, and avoid allocating ylm
      safe_allocate(orbitals(iw)%zorb(1:orbitals(iw)%sphere%np, 1:1, 1:1))
      orbitals(iw)%zorb(1:orbitals(iw)%sphere%np, 1, 1) = ylm(1:orbitals(iw)%sphere%np)
      safe_deallocate_a(ylm)
 
      safe_allocate(orbitals(iw)%phase(1:orbitals(iw)%sphere%np, st%d%kpt%start:st%d%kpt%end))
      orbitals(iw)%phase = m_z0
      safe_allocate(orbitals(iw)%eorb_mesh(1:mesh%np, 1:1, 1:1, st%d%kpt%start:st%d%kpt%end))
      orbitals(iw)%eorb_mesh = m_z0
 
      call orbitalset_update_phase(orbitals(iw), space%dim, st%d%kpt, kpoints, st%d%ispin == spin_polarized, &
        kpt_max = num_kpts)
    end do
 
    ! Compute the projections
    safe_allocate(psi(1:mesh%np, 1:st%d%dim))
    safe_allocate(phase(1:mesh%np))
    projection = m_zero
 
    do ik = 1, num_kpts
      kpoint(1:space%dim) = kpoints%get_point(ik)
      do ip = 1, mesh%np
        phase(ip) = exp(-m_zi * sum(mesh%x(1:space%dim, ip) * kpoint(1:space%dim)))
      end do
 
      ik_real = st%d%get_k_spin_index(ik, spin_channel)
 
      ! TODO(Alex) ist and iw loops need interchanging to optimise array access
      do ist = 1, size(band_index)
        iband = band_index(ist)
 
        if(ik_real >= st%d%kpt%start .and. ik_real <= st%d%kpt%end) then
          call states_elec_get_state(st, mesh, iband, ik_real, psi)
 
          do idim = 1, st%d%dim
            do ip = 1, mesh%np
              psi(ip, idim) = psi(ip, idim)*phase(ip)
            end do
          end do
 
          do iw = 1, num_proj
            idim = 1
            if (wan_opts%spinors) idim = wan_opts%spin_proj_component(iw)
            !At the moment the orbitals do not depend on idim
            !The idim index for eorb_mesh would be for a spin-resolved orbital like j=1/2
            projection(iband, iw, ik) = zmf_dotp(mesh, psi(1:mesh%np, idim), &
              orbitals(iw)%eorb_mesh(1:mesh%np, 1, 1, ik_real), reduce = .false.)
          end do
 
          call profiling_in("W90_AMN_REDUCE")
          call mesh%allreduce(projection(ist, :, ik))
          call profiling_out("W90_AMN_REDUCE")
        end if
 
        if(st%d%kpt%parallel) then
          call comm_allreduce(st%d%kpt%mpi_grp, projection(ist, :, ik))
        end if
 
      end do !ist
    end do !ik
 
    safe_deallocate_a(psi)
    safe_deallocate_a(phase)
 
    do iw = 1, num_proj
      call orbitalset_end(orbitals(iw))
    end do
    safe_deallocate_a(orbitals)
 
    call profiling_out("W90_AMN")
 
    pop_sub(wannier_calc_create_amn)
 
  end subroutine wannier_calc_create_amn
 
  subroutine wannier_calc_create_mmn(band_index, mesh, st, latt, &
    spin_channel, nnlist, nncell, overlap)
    integer, contiguous,         intent(in) :: band_index(:)
    class(mesh_t),               intent(in) :: mesh
    type(states_elec_t), target, intent(in) :: st
    type(lattice_vectors_t),     intent(in) :: latt   ! can differ from ions%latt for non-periodic systems
    integer,                     intent(in) :: spin_channel
    integer, contiguous,         intent(in) :: nnlist(:,:)
    integer, contiguous,         intent(in) :: nncell(:,:,:)
 
    complex(real64), contiguous, intent(out) :: overlap(:, :, :, :) ! distributed in k-points
 
    integer       ::  ist, jst, ik, ip, iknn, idim, ibind
    integer       ::  ik_oct, iknn_oct, inn
    integer       ::  ik_loc
    integer       ::  iband, jband, inode, node_fr, node_to
    real(real64)  ::  gcart(3)
    integer       ::  g(3)
 
    integer       ::  nntot, num_kpts
 
    complex(real64), allocatable :: psim(:,:), psin(:,:), phase(:)
    type(wfs_elec_t),    pointer :: batch
    type(mpi_request)            :: send_req
 
    push_sub(wannier_calc_create_mmn)
 
    call profiling_in("W90_MMN")
 
    
    nntot = size(nnlist, 2)
    num_kpts = size(nnlist, 1)
    assert(size(nnlist, 1) == num_kpts)
    assert(size(nncell, 1) == 3)
    assert(size(nncell, 2) == num_kpts)
    assert(size(nncell, 3) == nntot)
 
    if (st%parallel_in_states) then
      call messages_not_implemented("w90_mmn output with states parallelization")
    end if
 
    message(1) = "Info: Computing the overlap matrix for wannierisation"
    call messages_info(1)
 
    overlap = m_zero
    safe_allocate(psim(1:mesh%np, 1:st%d%dim))
    safe_allocate(psin(1:mesh%np, 1:st%d%dim))
    safe_allocate(phase(1:mesh%np))
 
    ! overlap is distributed in k like wannier90''s m_matrix_local
    ik_loc = 0
 
    ! loop over the pairs specified in the nnk tuple
    do ik = 1, num_kpts
      ik_oct = st%d%get_k_spin_index(ik, spin_channel)
      if (ik_oct >= st%d%kpt%start .and. ik_oct <= st%d%kpt%end) ik_loc = ik_loc + 1
 
      do inn = 1, nntot
 
        iknn = nnlist(ik, inn)
        g(1:3) = nncell(:, ik, inn)
        iknn_oct = st%d%get_k_spin_index(iknn, spin_channel)
 
        ! If the k-point is local, compute the phase e^{-i G.r}
        if(ik_oct >= st%d%kpt%start .and. ik_oct <= st%d%kpt%end) then
          ! Wannier90 treats everything fully periodic
          ! Conversion is done with the 3D "correct" klattice
          call kpoints_to_absolute(latt, real(g, real64) , gcart)
 
          ! Phase that gives u_{n,k+G}(r) from u_{nk}(r)
          ! Here the minus sign of Octopus cancels with minus sign of the input G
          ! (ik and iknn correspond in the list to minus the k-points in Octopus)
          if (any(g /= 0)) then
            do ip = 1, mesh%np
              phase(ip) = exp(-m_zi*dot_product(mesh%x(1:3, ip), gcart(1:3)))
            end do
          end if
        end if
 
        ! Loop over distributed bands
        do jst = 1, size(band_index)
          jband = band_index(jst)
 
          if (st%d%kpt%parallel) then
            node_fr = -1
            node_to = -1
            do inode = 0, st%d%kpt%mpi_grp%size-1
              if(iknn_oct >= st%st_kpt_task(inode,3) .and. iknn_oct <= st%st_kpt_task(inode,4)) then
                node_fr = inode
              end if
              if(ik_oct >= st%st_kpt_task(inode,3) .and. ik_oct <= st%st_kpt_task(inode,4)) then
                node_to = inode
              end if
            end do
            assert(node_fr > -1)
            assert(node_to > -1)
            send_req = mpi_request_null
            ! We have locally the k-point
            if (state_kpt_is_local(st, jband, iknn_oct)) then
              call states_elec_get_state(st, mesh, jband, iknn_oct, psin)
              ! We send it only if we don`t want to use it locally
              if(node_to /= st%d%kpt%mpi_grp%rank) then
                call st%d%kpt%mpi_grp%isend(psin, mesh%np*st%d%dim, mpi_double_complex, node_to, send_req)
              end if
            end if
            ! We receive the desired state, only if it is not a local one
            if(node_to == st%d%kpt%mpi_grp%rank .and. node_to /= node_fr) then
              call st%d%kpt%mpi_grp%recv(psin, mesh%np*st%d%dim, mpi_double_complex, node_fr)
            end if
            if (send_req /= mpi_request_null) then
              call st%d%kpt%mpi_grp%wait(send_req)
            end if
          else
            call states_elec_get_state(st, mesh, jband, iknn_oct, psin)
          end if
 
          if(ik_oct >= st%d%kpt%start .and. ik_oct <= st%d%kpt%end) then
 
            ! Do not apply the phase if the phase factor is 1
            if (any(g /= 0)) then
              do idim = 1, st%d%dim
                do ip = 1, mesh%np
                  psin(ip, idim) = psin(ip, idim) * phase(ip)
                end do
              end do
            end if
 
            ! See Eq. (25) in PRB 56, 12847 (1997)
            ! Loop over local k-points
            do iband = 1, size(band_index)
              ist = band_index(iband)
 
              batch => st%group%psib(st%group%iblock(ist), ik_oct)
 
              select case (batch%status())
              case (batch_not_packed)
                overlap(iband, jst, inn, ik_loc) = m_z0
                do idim = 1, st%d%dim
                  ibind = batch%inv_index((/ist, idim/))
                  overlap(iband, jst, inn, ik_loc) = overlap(iband, jst, inn, ik_loc) + &
                    zmf_dotp(mesh, batch%zff_linear(:, ibind), psin(:,idim), reduce = .false.)
                end do
                !Not properly done at the moment
              case (batch_packed, batch_device_packed)
                call states_elec_get_state(st, mesh, ist, ik_oct, psim)
                overlap(iband, jst, inn, ik_loc) = zmf_dotp(mesh, st%d%dim, psim, psin, reduce = .false.)
              end select
            end do !iband
 
            call profiling_in("W90_MMN_REDUCE")
            call mesh%allreduce(overlap(:, jst, inn, ik_loc))
            call profiling_out("W90_MMN_REDUCE")
          end if
        end do !jst
      end do !inn
    end do !ik
 
    assert(ik_loc == size(overlap, 4))
 
    safe_deallocate_a(psim)
    safe_deallocate_a(psin)
    safe_deallocate_a(phase)
 
    call profiling_out("W90_MMN")
 
    pop_sub(wannier_calc_create_mmn)
 
  end subroutine wannier_calc_create_mmn
 
  subroutine wannier_calc_read_centers(wan_opts, centers)
    type(wannier_opts_t),      intent(in)    :: wan_opts
    real(real64), contiguous,  intent(inout) :: centers(:, :)
 
    integer :: w90_xyz, iw
    character(len=2) :: dum
    logical :: exist
 
    push_sub(wannier_calc_read_centers)
 
    assert(size(centers, 2) == wan_opts%num_wann)
    assert(size(centers, 1) == 3)
 
    inquire(file=trim(adjustl(wan_opts%prefix))//'_centres.xyz',exist=exist)
    if (.not. exist) then
      message(1) = 'Cannot find the Wannier90 file seedname_centres.xyz.'
      call messages_fatal(1)
    end if
 
    w90_xyz = io_open(trim(adjustl(wan_opts%prefix))//'_centres.xyz', global_namespace, action='read')
 
    !Skip two lines
    read(w90_xyz, *)
    read(w90_xyz, *)
    do iw = 1, wan_opts%num_wann
      read(w90_xyz, *) dum, centers(1:3, iw)
      ! Wannier90 outputs the coordinates in angstrom
      centers(1:3, iw) = units_to_atomic(unit_angstrom, centers(1:3, iw))
    end do
    call io_close(w90_xyz)
 
    pop_sub(wannier_calc_read_centers)
  end subroutine wannier_calc_read_centers
 
  subroutine wannier_calc_read_u_matrices(wan_opts, kpoints, u_matrix, u_dis_matrix)
    type(wannier_opts_t), intent(in)    :: wan_opts
    type(kpoints_t),      intent(in)    :: kpoints
    complex(real64), contiguous, intent(inout) :: u_matrix(:, :, :)
    complex(real64), contiguous, intent(inout) :: u_dis_matrix(:, :, :)
 
    integer :: w90_u_mat, w90_u_dis, ik, ib, iw, iw2, nik, nwann, nib
    integer :: num_kpts
    logical :: exist, have_disentangled
 
    num_kpts = product(kpoints%nik_axis(1:3))
 
    push_sub(wannier_calc_read_u_matrices)
 
    inquire(file=trim(adjustl(wan_opts%prefix))//'_u.mat',exist=exist)
    if (.not. exist) then
      message(1) = 'Cannot find the Wannier90 seedname_u.mat file.'
      call messages_fatal(1)
    end if
    w90_u_mat = io_open(trim(adjustl(wan_opts%prefix))//'_u.mat', global_namespace, action='read')
 
    !Skip one line
    read(w90_u_mat, *)
    !Read num_kpts, num_wann, num_wann for consistency check
    read(w90_u_mat, *) nik, nwann, nwann
    if (nik /= num_kpts .or. nwann /= wan_opts%num_wann) then
      message(1) = "The U matrix has inconsistent shape."
      call messages_fatal(1)
    end if
 
    do ik = 1, num_kpts
      !Skip one line
      read(w90_u_mat, *)
      !Skip one line (k-point coordinate)
      read(w90_u_mat, *)
      read(w90_u_mat, '(f15.10,sp,f15.10)') ((u_matrix(iw, iw2, ik), iw=1, wan_opts%num_wann), iw2=1, wan_opts%num_wann)
    end do
 
    call io_close(w90_u_mat)
 
    inquire(file=trim(adjustl(wan_opts%prefix))//'_u_dis.mat',exist=have_disentangled)
    if (have_disentangled) then
      ! Read u_dis.mat: dimensions are (w90_num_bands x w90_num_wann) per k-point.
      ! The band rows follow the compressed, non-excluded band ordering used by
      ! the rest of the Wannier90 interface (i.e. after applying exclude_bands).
      w90_u_dis = io_open(trim(adjustl(wan_opts%prefix))//'_u_dis.mat', global_namespace, action='read')
      !Skip one line
      read(w90_u_dis, *)
      !Read num_kpts, num_wann, num_bands for consistency check
      read(w90_u_dis, *) nik, nwann, nib
      if (nik /= num_kpts .or. nwann /= wan_opts%num_wann .or. nib /= wan_opts%num_bands) then
        message(1) = 'The u_dis matrix has inconsistent shape.'
        call messages_fatal(1)
      end if
 
      do ik = 1, num_kpts
        !Skip one line
        read(w90_u_dis, *)
        !Skip one line (k-point coordinate)
        read(w90_u_dis, *)
        read(w90_u_dis, '(f15.10,sp,f15.10)') ((u_dis_matrix(ib, iw, ik), ib=1, wan_opts%num_bands), iw=1, wan_opts%num_wann)
      end do
 
      call io_close(w90_u_dis)
 
      ! In case of disentanglement, u_dis matrix needs to be reordered to the correct layout.
      ! See issue #1465.
    end if
 
    pop_sub(wannier_calc_read_u_matrices)
  end subroutine wannier_calc_read_u_matrices
 
  subroutine wannier_calc_get_header(header_start, header)
    character(len=*), intent(in)  :: header_start
    character(len=*), intent(out) :: header
 
    integer :: date_time(8)
 
    push_sub(wannier_calc_get_header)
 
    call date_and_time(values=date_time)
    write(header, '(a,i4,a1,i2.2,a1,i2.2,a,i2.2,a1,i2.2,a1,i2.2)') &
      header_start, date_time(1), '/', date_time(2), '/', date_time(3), &
      ' at ', date_time(5), ':', date_time(6), ':', date_time(7)
 
    pop_sub(wannier_calc_get_header)
  end subroutine wannier_calc_get_header
 
  subroutine wannier_calc_write_centers(dir, prefix, centers, ions)
    character(len=*), intent(in) :: dir
    character(len=*), intent(in) :: prefix
    real(real64),     intent(in) :: centers(:, :)
    class(ions_t),    intent(in) :: ions
 
    integer :: w90_xyz, iw, ia, ind
    character(len=512) :: header
 
    push_sub(wannier_calc_write_centers)
 
    w90_xyz = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_centres.xyz', global_namespace, action='write', die=.false.)
    if (w90_xyz == -1) then
      message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_centres.xyz for writing.'
      call messages_fatal(1)
    end if
 
    write(w90_xyz, '(i6)') size(centers, 2) + ions%natoms
    call wannier_calc_get_header('Wannier centers, written by octopus on ', header)
    write (w90_xyz, *) trim(header)
 
    do iw = 1, size(centers, 2)
      write(w90_xyz, '("X",6x,3(f14.8,3x))') &
        (units_from_atomic(unit_angstrom, centers(ind, iw)), ind=1, 3)
    end do
    do ia = 1, ions%natoms
      write(w90_xyz, '(a2,5x,3(f14.8,3x))') trim(ions%atom(ia)%label), &
        (units_from_atomic(unit_angstrom, ions%pos(ind, ia)), ind=1, 3)
    end do
    call io_close(w90_xyz)
 
    pop_sub(wannier_calc_write_centers)
 
  end subroutine wannier_calc_write_centers
 
  subroutine wannier_calc_write_u_matrices(dir, prefix, kpoints, u_matrix, u_dis_matrix)
    character(len=*), intent(in) :: dir
    character(len=*), intent(in) :: prefix
    type(kpoints_t),  intent(in) :: kpoints
    complex(real64), contiguous,  intent(in) :: u_matrix(:, :, :)
    complex(real64), contiguous,  intent(in) :: u_dis_matrix(:, :, :)
 
    integer :: w90_u_mat, w90_u_dis
    logical :: have_disentangled
    integer :: i, j, nkp, num_wann, num_bands, num_kpts
    character(len=512) :: header
 
    num_kpts = size(u_matrix, 3)
    num_wann = size(u_matrix, 1)
    num_bands = size(u_dis_matrix, 1)
    assert(num_kpts == size(kpoints%reduced%red_point, 2))
 
    push_sub(wannier_calc_write_u_matrices)
 
    call wannier_calc_get_header('written on ', header)
 
    w90_u_mat = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_u.mat', global_namespace, action='write', die=.false.)
    if (w90_u_mat == -1) then
      message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_u.mat for writing.'
      call messages_fatal(1)
    end if
 
    write (w90_u_mat, *) trim(header)
    write (w90_u_mat, *) num_kpts, num_wann, num_wann
 
    do nkp = 1, num_kpts
      write (w90_u_mat, *)
      write (w90_u_mat, '(f15.10,sp,f15.10,sp,f15.10)') -kpoints%reduced%red_point(1:3, nkp)
      write (w90_u_mat, '(f15.10,sp,f15.10)') ((u_matrix(i, j, nkp), i=1, num_wann), j=1, num_wann)
    end do
 
    call io_close(w90_u_mat)
 
    inquire(file=trim(trim(adjustl(prefix))//'_u_dis.mat'),exist=have_disentangled)
    if (have_disentangled) then
 
      w90_u_dis = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_u_dis.mat', &
        global_namespace, action='write', die=.false.)
      if (w90_u_dis == -1) then
        message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_u_dis.mat for writing.'
        call messages_fatal(1)
      end if
 
      write (w90_u_dis, *) trim(header)
      write (w90_u_dis, *) num_kpts, num_wann, num_bands
      do nkp = 1, num_kpts
        write (w90_u_dis, *)
        write (w90_u_dis, '(f15.10,sp,f15.10,sp,f15.10)') -kpoints%reduced%red_point(1:3, nkp)
        write (w90_u_dis, '(f15.10,sp,f15.10)') ((u_dis_matrix(i, j, nkp), i=1, num_bands), j=1, num_wann)
      end do
 
      call io_close(w90_u_dis)
    end if
 
    pop_sub(wannier_calc_write_u_matrices)
 
  end subroutine wannier_calc_write_u_matrices
 
  subroutine wannier_calc_write_spreads(dir, prefix, spreads)
    character(len=*), intent(in) :: dir
    character(len=*), intent(in) :: prefix
    real(real64),     intent(in) :: spreads(:)
 
    integer :: w90_xyz, iw
    character(len=512) :: header
 
    push_sub(wannier_calc_write_spreads)
 
    w90_xyz = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_spreads.xyz', global_namespace, action='write', die=.false.)
    if (w90_xyz == -1) then
      message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_spreads.xyz for writing.'
      call messages_fatal(1)
    end if
 
    call wannier_calc_get_header('Wannier spreads, written by octopus on ', header)
    write (w90_xyz, *) trim(header)
 
    do iw = 1, size(spreads, 1)
      write(w90_xyz, '("X",6x,f14.8)') spreads(iw) * units_from_atomic(unit_angstrom, 1.0_real64)**2
    end do
    call io_close(w90_xyz)
 
    pop_sub(wannier_calc_write_spreads)
 
  end subroutine wannier_calc_write_spreads
 
  subroutine wannier_calc_write_hr_file(dir, prefix, irvec, ndegen, ham_r)
    character(len=*),            intent(in) :: dir
    character(len=*),            intent(in) :: prefix
    integer, contiguous,         intent(in) :: irvec(:, :)
    integer, contiguous,         intent(in) :: ndegen(:)
    complex(real64), contiguous, intent(in) :: ham_r(:, :, :)
 
    integer :: w90_hr_dat
    integer :: i, j, irpt
    character(len=512) :: header
    integer :: nrpts, n_wannier
 
    push_sub(wannier_calc_write_hr_file)
 
    nrpts = size(ham_r, 3)
    n_wannier = size(ham_r, 1)
    assert(size(irvec, 2) == nrpts)
    assert(size(ndegen, 1) == nrpts)
 
    w90_hr_dat = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_hr.dat', global_namespace, action='write', die=.false.)
    if (w90_hr_dat == -1) then
      message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_hr.dat for writing.'
      call messages_fatal(1)
    end if
 
    call wannier_calc_get_header('written on ', header)
    write (w90_hr_dat, *) trim(header)
 
    write (w90_hr_dat, *) size(ham_r, 1)
    write (w90_hr_dat, *) nrpts
    write (w90_hr_dat, '(15I5)') (ndegen(i), i=1, nrpts)
    do irpt = 1, nrpts
      do i = 1, n_wannier
        do j = 1, n_wannier
          write (w90_hr_dat, '(5I5,2F12.6)') irvec(:, irpt), j, i, &
            units_from_atomic(unit_ev, ham_r(j, i, irpt))
        end do
      end do
    end do
 
    call io_close(w90_hr_dat)
 
    pop_sub(wannier_calc_write_hr_file)
  end subroutine wannier_calc_write_hr_file
 
  logical function wannier_calc_r_in_wigner_seitz_cell(rvec, real_metric, mp_grid, ws_search_size, &
    ws_distance_tol, ndegen) result(in_cell)
    integer,           intent(in)  :: rvec(3)
    real(real64),      intent(in)  :: real_metric(3, 3)
    integer,           intent(in)  :: mp_grid(3)
    integer,           intent(in)  :: ws_search_size(3)
    real(real64),      intent(in)  :: ws_distance_tol
    integer, optional, intent(out) :: ndegen
 
    integer      :: i1, i2, i3, icnt
    integer      :: ndiff(3)
    real(real64) :: dist_min, dist_to_origin
    real(real64) :: dist(product(2*ws_search_size + 3))
 
    ! no push_sub, called too frequently
 
    ! Squared distances from rvec to all superlattice points within the search range
    icnt = 0
    do i1 = -ws_search_size(1) - 1, ws_search_size(1) + 1
      do i2 = -ws_search_size(2) - 1, ws_search_size(2) + 1
        do i3 = -ws_search_size(3) - 1, ws_search_size(3) + 1
          icnt = icnt + 1
          ndiff = rvec - [i1, i2, i3] * mp_grid
          dist(icnt) = dot_product(real(ndiff, real64), matmul(real_metric, real(ndiff, real64)))
        end do
      end do
    end do
 
    dist_to_origin = dot_product(real(rvec, real64), matmul(real_metric, real(rvec, real64)))
    dist_min = minval(dist)
 
    in_cell = abs(dist_to_origin - dist_min) < ws_distance_tol**2
    if (present(ndegen)) ndegen = count(abs(dist - dist_min) < ws_distance_tol**2)
 
  end function wannier_calc_r_in_wigner_seitz_cell
 
  subroutine wannier_calc_wigner_seitz_count(real_lattice, mp_grid, ws_search_size, ws_distance_tol, nrpts)
    real(real64), intent(in)  :: real_lattice(3,3)
    integer,      intent(in)  :: mp_grid(3)
    integer,      intent(in)  :: ws_search_size(3)
    real(real64), intent(in)  :: ws_distance_tol
    integer,      intent(out) :: nrpts
 
    integer      :: n1, n2, n3
    real(real64) :: real_metric(3, 3)
 
    push_sub(wannier_calc_wigner_seitz_count)
 
    real_metric = matmul(transpose(real_lattice), real_lattice)
 
    nrpts = 0
    do n1 = -ws_search_size(1)*mp_grid(1), ws_search_size(1)*mp_grid(1)
      do n2 = -ws_search_size(2)*mp_grid(2), ws_search_size(2)*mp_grid(2)
        do n3 = -ws_search_size(3)*mp_grid(3), ws_search_size(3)*mp_grid(3)
          if (wannier_calc_r_in_wigner_seitz_cell([n1, n2, n3], real_metric, mp_grid, &
            ws_search_size, ws_distance_tol)) then
            nrpts = nrpts + 1
          end if
        end do
      end do
    end do
 
    pop_sub(wannier_calc_wigner_seitz_count)
  end subroutine wannier_calc_wigner_seitz_count
 
  subroutine wannier_calc_wigner_seitz_fill(real_lattice, mp_grid, ws_search_size, ws_distance_tol, &
    nrpts, irvec, ndegen)
    real(real64), intent(in)  :: real_lattice(3,3)
    integer,      intent(in)  :: mp_grid(3)
    integer,      intent(in)  :: ws_search_size(3)
    real(real64), intent(in)  :: ws_distance_tol
    integer,      intent(in)  :: nrpts
    integer,      intent(out) :: irvec(3, nrpts)
    integer,      intent(out) :: ndegen(nrpts)
 
    integer      :: n1, n2, n3, irpt, ndeg
    real(real64) :: real_metric(3, 3)
 
    push_sub(wannier_calc_wigner_seitz_fill)
 
    real_metric = matmul(transpose(real_lattice), real_lattice)
 
    irpt = 0
    do n1 = -ws_search_size(1)*mp_grid(1), ws_search_size(1)*mp_grid(1)
      do n2 = -ws_search_size(2)*mp_grid(2), ws_search_size(2)*mp_grid(2)
        do n3 = -ws_search_size(3)*mp_grid(3), ws_search_size(3)*mp_grid(3)
          if (wannier_calc_r_in_wigner_seitz_cell([n1, n2, n3], real_metric, mp_grid, &
            ws_search_size, ws_distance_tol, ndeg)) then
            irpt = irpt + 1
            irvec(:, irpt) = [n1, n2, n3]
            ndegen(irpt) = ndeg
          end if
        end do
      end do
    end do
    assert(irpt == nrpts)
 
    ! Sum rule from wannier90: the R-points weighted by their inverse degeneracy
    ! must add up to the number of k-points
    if (abs(sum(m_one/ndegen) - product(mp_grid)) > 1e-8_real64) then
      message(1) = "Error computing the Wigner-Seitz points."
      call messages_fatal(1)
    end if
 
    pop_sub(wannier_calc_wigner_seitz_fill)
  end subroutine wannier_calc_wigner_seitz_fill
 
  subroutine wannier_calc_update_phase(band_index, st, dt, phase)
    integer, contiguous,         intent(in)    :: band_index(:)
    type(states_elec_t),         intent(in)    :: st
    real(real64),                intent(in)    :: dt
    complex(real64), contiguous, intent(inout) :: phase(:,:)
 
    integer :: ik, ist
 
    push_sub(wannier_calc_update_phase)
 
    do ik = 1, size(phase, 2)
      do ist = 1, size(band_index)
        phase(ist, ik) = phase(ist, ik) * exp(m_zi * st%eigenval(band_index(ist), ik) * dt)
      end do
    end do
 
    pop_sub(wannier_calc_update_phase)
  end subroutine wannier_calc_update_phase
 
  subroutine wannier_calc_compute_ham_r(namespace, mesh, st, hm, kpoints, wannier_opts, &
    u_matrix, irvec, ham_r)
    type(namespace_t),           intent(in)    :: namespace
    class(mesh_t),               intent(in)    :: mesh
    type(states_elec_t), target, intent(inout) :: st
    type(hamiltonian_elec_t),    intent(in)    :: hm
    type(kpoints_t),             intent(in)    :: kpoints
    type(wannier_opts_t),        intent(in)    :: wannier_opts
    complex(real64), contiguous, intent(in)    :: u_matrix(:, :, :)
    integer,         contiguous, intent(in)    :: irvec(:, :)
    complex(real64), contiguous, intent(out)   :: ham_r(:, :, :)
 
    integer :: nk, nrpts, ik, irpt, iq, ikl
    real(real64) :: rdotk
    complex(real64), allocatable :: hmss(:, :), ham_k_local(:, :, :)
 
    push_sub(wannier_calc_compute_ham_r)
 
    nk = kpoints%full%npoints       
    nrpts = size(ham_r, 3)
 
    safe_allocate(hmss(1:wannier_opts%num_bands, 1:wannier_opts%num_bands))
    safe_allocate(ham_k_local(1:wannier_opts%num_wann, 1:wannier_opts%num_wann, 1:st%d%kpt%nlocal()))
 
    assert(wannier_opts%num_bands == st%nst)
    if (wannier_opts%num_bands /= wannier_opts%num_wann) then
      call messages_not_implemented('TD Wannier: num_bands unequal num_wann', namespace)
    end if
 
    ! Hamiltonian matrix elements in the subspace of the (time-dependent) KS states,
    ! rotated into the Wannier basis: H_W(k) = U(k)^dagger eps(k) U(k)
    ! Computed once per local k-point and cached
    do iq = st%d%kpt%start, st%d%kpt%end
      ik = st%d%get_kpoint_index(iq)
      ikl = iq - st%d%kpt%start + 1
      call zsubspace_diag_hamiltonian(namespace, mesh, st, hm, iq, hmss)
      ham_k_local(:, :, ikl) = matmul(transpose(conjg(u_matrix(:, :, ik))), matmul(hmss, u_matrix(:, :, ik)))
    end do
 
    ! Adapted from Wannier90: Fourier transform to real space: H(R) = (1/N_k) sum_k exp(-i k.R) H_W(k)
    ! Each rank sums only its own local k-points
    ham_r = m_z0
    do irpt = 1, nrpts
      do iq = st%d%kpt%start, st%d%kpt%end
        ik = st%d%get_kpoint_index(iq)
        ikl = iq - st%d%kpt%start + 1
        rdotk = m_two*m_pi*dot_product(kpoints%reduced%red_point(1:3, ik), real(irvec(1:3, irpt), real64))
        ham_r(:, :, irpt) = ham_r(:, :, irpt) + (exp(m_zi*rdotk)/real(nk, real64)) * ham_k_local(:, :, ikl)
        ! sign mismatch here because of mismatch in k-point convention between Octopus and Wannier90
      end do
    end do
 
    if (st%d%kpt%parallel) then
      call comm_allreduce(st%d%kpt%mpi_grp, ham_r)
    end if
 
    safe_deallocate_a(hmss)
    safe_deallocate_a(ham_k_local)
 
    pop_sub(wannier_calc_compute_ham_r)
  end subroutine wannier_calc_compute_ham_r
 
end module wannier_calc_oct_m