wannier.F90

Go to the documentation of this file.

!! 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_oct_m
  use, intrinsic :: iso_fortran_env
 
  use debug_oct_m
  use global_oct_m
  use electron_space_oct_m
  use hamiltonian_elec_oct_m, only: hamiltonian_elec_t
  use io_oct_m
  use ions_oct_m
  use kpoints_oct_m
  use lattice_vectors_oct_m, only: lattice_vectors_t
  use mesh_oct_m
  use messages_oct_m
  use mpi_oct_m
  use multicomm_oct_m
  use namespace_oct_m
  use output_low_oct_m
  use parser_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 types_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use wannier_calc_oct_m
  use wannier_opts_oct_m
  use ylm_wannier_oct_m
 
#ifdef HAVE_WANNIER90
  use w90_library
  use wannier90lib_interface_oct_m, &
    only: &
    wannier90lib_init_w90main, &
    wannier90lib_initialize_matrix_data, &
    wannier90lib_wannierization
#endif
 
  implicit none
  private
 
  public :: wannier_t
  type :: wannier_t
    type(wannier_opts_t)         :: options
    type(lattice_vectors_t)      :: latt
    complex(real64), allocatable :: u_matrix(:,:,:)
    complex(real64), allocatable :: u_dis_matrix(:,:,:)
    real(real64),    allocatable :: centers(:,:)
    real(real64),    allocatable :: spreads(:)
    complex(real64), allocatable :: m_matrix(:,:,:,:)
    complex(real64), allocatable :: phase(:,:)
    complex(real64), allocatable :: ham_r(:, : ,:)
    integer,         allocatable :: irvec(:, :)
    integer,         allocatable :: ndegen(:)
    integer                      :: last_iter = 0
#ifdef HAVE_WANNIER90
    type(lib_common_type) :: w90main
#endif
 
  contains
    procedure :: init_from_input => wannier_init_from_input
    procedure :: init_parallelization => wannier_init_parallelization
    procedure :: init_data_from_file => wannier_init_data_from_file
    procedure :: restart_write_data => wannier_restart_write_data
    procedure :: restart_read_data => wannier_restart_read_data
    procedure :: write_iter => wannier_write_iter
    procedure :: update_matrices => wannier_update_matrices
    final :: wannier_finalize
  end type wannier_t
 
contains
 
  subroutine wannier_init_from_input(this, namespace, kpoints, ions, space, st, gr)
    class(wannier_t), intent(inout)     :: this
    type(namespace_t), intent(in)       :: namespace
    type(kpoints_t), intent(in)         :: kpoints
    type(ions_t), intent(in)            :: ions
    class(space_t), intent(in)          :: space
    type(states_elec_t), intent(in)     :: st
    class(mesh_t), intent(in)           :: gr
 
    integer :: nrpts
 
    push_sub(wannier_init_from_input)
 
    ! Parse all Wannier related options
    call this%options%parse_oct(namespace)
    if (this%options%td_method == td_wannier_method_none) then
      ! No TD wannier functionality requested
      pop_sub(wannier_init_from_input)
      return
    end if
 
    call this%options%parse_win()
    ! Currently TD Wannier is not supported for the disentanglement case
    if (this%options%num_bands /= this%options%num_wann) then
      call messages_not_implemented('TD Wannier: num_bands not equal num_wann')
    end if
    ! also not for spin polarized or spinors
    if (st%d%ispin /= unpolarized) then
      call messages_not_implemented('TD Wannier: spin polarized and spinors')
    end if
    ! and only for 3D systems
    if (space%dim /= 3) then     ! see issue #1466
      call messages_not_implemented("TD wannier: dimensionality other than 3D")
    end if
    ! wannierization is not supported for k-point symmetries
    if (kpoints%use_symmetries .eqv. .true.) then
      call messages_not_implemented("TD wannier: k-point symmetries")
    end if
    ! nst /= num_bands can occur without excluded bands
    if (this%options%num_bands /= st%nst) then
      call messages_not_implemented("TD wannier: Setting nst /= num_bands to exclude bands")
    end if
 
    ! Build the effective lattice used for all TD Wannier geometry, and return in this%latt
    call wannier_build_effective_lattice(ions, space, gr, this%latt)
 
    safe_allocate(this%u_matrix(1:this%options%num_wann, 1:this%options%num_wann, 1:kpoints%full%npoints))
    safe_allocate(this%u_dis_matrix(1:this%options%num_bands, 1:this%options%num_wann, 1:kpoints%full%npoints))
    safe_allocate(this%centers(1:3, 1:this%options%num_wann))
    safe_allocate(this%spreads(1:this%options%num_wann))
    safe_allocate(this%phase(1:this%options%num_bands, 1:kpoints%full%npoints))
    this%u_matrix = m_zero
    this%u_dis_matrix = m_zero
    this%centers = m_zero
    this%spreads = m_zero
    this%phase = m_z1
 
    !Initialize translation vectors needed for calculation of H(R) in Wannier basis
    nrpts = 1
    call wannier_calc_wigner_seitz_count(this%latt%rlattice, kpoints%nik_axis(1:3), &
      this%options%ws_search_size, this%options%ws_distance_tol, nrpts)
    safe_allocate(this%irvec(1:3, 1:nrpts))
    safe_allocate(this%ndegen(1:nrpts))
    call wannier_calc_wigner_seitz_fill(this%latt%rlattice, kpoints%nik_axis(1:3), &
      this%options%ws_search_size, this%options%ws_distance_tol, nrpts, this%irvec, this%ndegen)
    safe_allocate(this%ham_r(1:this%options%num_wann, 1:this%options%num_wann, 1:nrpts))
    this%ham_r = m_zero
 
 
    pop_sub(wannier_init_from_input)
  end subroutine wannier_init_from_input
 
  subroutine wannier_build_effective_lattice(ions, space, gr, latt)
    type(ions_t),             intent(in)  :: ions
    class(space_t),           intent(in)  :: space
    class(mesh_t),            intent(in)  :: gr
    type(lattice_vectors_t),  intent(out) :: latt
 
    real(real64) :: factor(space%dim)
    integer :: idir
 
    push_sub(wannier_build_effective_lattice)
 
    call latt%copy(ions%latt)
 
    factor = m_one
    do idir = space%periodic_dim + 1, space%dim
      factor(idir) = m_two * gr%box%bounding_box_l(idir)
    end do
    call latt%scale(factor)
 
    pop_sub(wannier_build_effective_lattice)
  end subroutine wannier_build_effective_lattice
 
  subroutine wannier_init_parallelization(this, namespace, st, ions, kpoints)
    class(wannier_t), intent(inout)     :: this
    type(namespace_t), intent(in)       :: namespace
    type(states_elec_t), intent(in)     :: st
    type(ions_t), intent(in)            :: ions
    type(kpoints_t), intent(in)         :: kpoints
 
#ifdef HAVE_WANNIER90
    integer :: ierr, nn, iband, ist, nk_local
#endif
 
    push_sub(wannier_init_parallelization)
 
    if (this%options%td_method == td_wannier_method_none) then
      pop_sub(wannier_init_parallelization)
      return
    end if
 
#ifdef HAVE_WANNIER90
    !Initialize wannier90 library object, including kpoint parallelization
    call wannier90lib_init_w90main(namespace, ions, this%latt, kpoints, st, &
      this%options, this%w90main, stdout, stderr, ierr)
    if (ierr /= 0) then
      write(message(1),'(a,i0)') 'Error initializing wannier90 library object: ', ierr
      call messages_fatal(1)
    end if
 
    ! Suppress all disk IO from wannier90 side. The needed output is controlled within octopus
    this%w90main%output_file = output_file_type()   ! resets all file output to false
 
    ! allocation of m_matrix is done here, since nn HAS to be known (which comes from wannier90 after initialization)
    ! m_matrix is distributed in k, matching wannier90''s m_matrix_local
    call w90_get_nn(this%w90main, nn, stdout, stderr, ierr)
    nk_local = st%d%kpt%nlocal()
    if (st%d%ispin == spin_polarized) nk_local = nk_local / 2
    safe_allocate(this%m_matrix(1:this%options%num_bands, 1:this%options%num_bands, 1:nn, 1:nk_local))
 
    ! Update exclude_list and band_index based on wannier90 input from .win file (if any)
    safe_allocate(this%options%exclude_list(1:st%nst))
    this%options%exclude_list = .false.
    if (allocated(this%w90main%exclude_bands)) then
      do iband = 1, size(this%w90main%exclude_bands)
        ist = this%w90main%exclude_bands(iband)
        this%options%exclude_list(ist) = .true.
      end do
    end if
    ! Currently TD Wannier is not supported for exclude_band case
    if (any(this%options%exclude_list)) then
      call messages_not_implemented('TD Wannier: exclude_bands')
    end if
    this%options%band_index = pack([(ist, ist=1,st%nst)], this%options%exclude_list .eqv. .false.)
#endif
 
    pop_sub(wannier_init_parallelization)
 
  end subroutine wannier_init_parallelization
 
  subroutine wannier_init_data_from_file(this, kpoints)
    class(wannier_t),       intent(inout) :: this
    type(kpoints_t),        intent(in)    :: kpoints
 
    push_sub(wannier_init_data_from_file)
 
    if (this%options%td_method == td_wannier_method_none) then
      pop_sub(wannier_init_data_from_file)
      return
    end if
 
    call wannier_calc_read_centers(this%options, this%centers)
    call wannier_calc_read_u_matrices(this%options, kpoints, this%u_matrix, this%u_dis_matrix)
#ifdef HAVE_WANNIER90
    call wannier90lib_initialize_matrix_data(this%w90main, this%u_matrix, &
      this%u_dis_matrix, this%m_matrix, this%centers)
#endif
 
    pop_sub(wannier_init_data_from_file)
  end subroutine wannier_init_data_from_file
 
  subroutine wannier_restart_write_data(this, namespace, mc, gr)
    class(wannier_t), intent(inout) :: this
    type(namespace_t), intent(in)   :: namespace
    type(multicomm_t), intent(in)   :: mc
    class(mesh_t), intent(in)       :: gr
 
    type(restart_t) :: restart
    integer :: ierr
 
    push_sub(wannier_restart_write_data)
 
    if (this%options%td_method == td_wannier_method_none) then
      pop_sub(wannier_restart_write_data)
      return
    end if
 
    call restart%init(namespace, restart_td, restart_type_dump, mc, ierr, mesh=gr)
    if (ierr /= 0) then
      message(1) = 'Unable to initialize TD restart for Wannier write.'
      call messages_fatal(1)
    end if
 
    call restart%write_binary('wannier_u_matrix', size(this%u_matrix), this%u_matrix, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to write Wannier restart data (u_matrix).'
      call messages_fatal(1)
    end if
 
    call restart%write_binary('wannier_u_dis_matrix', size(this%u_dis_matrix), this%u_dis_matrix, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to write Wannier restart data (u_dis_matrix).'
      call messages_fatal(1)
    end if
 
    call restart%write_binary('wannier_centers', size(this%centers), this%centers, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to write Wannier restart data (centers).'
      call messages_fatal(1)
    end if
 
    ! Restart dumps do not generally coincide with a wannierization, so the accumulated
    ! phase is needed to keep it consistent with the U matrix written above
    call restart%write_binary('wannier_phase', size(this%phase), this%phase, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to write Wannier restart data (phase).'
      call messages_fatal(1)
    end if
 
    call restart%end()
 
    pop_sub(wannier_restart_write_data)
  end subroutine wannier_restart_write_data
 
  subroutine wannier_restart_read_data(this, namespace, mc, gr)
    class(wannier_t), intent(inout) :: this
    type(namespace_t), intent(in)   :: namespace
    type(multicomm_t), intent(in)   :: mc
    class(mesh_t), intent(in)       :: gr
 
    type(restart_t) :: restart
    integer :: ierr
 
    push_sub(wannier_restart_read_data)
 
    if (this%options%td_method == td_wannier_method_none) then
      pop_sub(wannier_restart_read_data)
      return
    end if
 
    call restart%init(namespace, restart_td, restart_type_load, mc, ierr, mesh=gr)
    if (ierr /= 0) then
      message(1) = 'Unable to initialize TD restart for Wannier read.'
      call messages_fatal(1)
    end if
 
    call restart%read_binary('wannier_u_matrix', size(this%u_matrix), this%u_matrix, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to read Wannier restart data (u_matrix).'
      call messages_fatal(1)
    end if
 
    call restart%read_binary('wannier_u_dis_matrix', size(this%u_dis_matrix), this%u_dis_matrix, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to read Wannier restart data (u_dis_matrix).'
      call messages_fatal(1)
    end if
 
    call restart%read_binary('wannier_centers', size(this%centers), this%centers, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to read Wannier restart data (centers).'
      call messages_fatal(1)
    end if
 
    call restart%read_binary('wannier_phase', size(this%phase), this%phase, ierr)
    if (ierr /= 0) then
      message(1) = 'Unable to read Wannier restart data (phase).'
      call messages_fatal(1)
    end if
 
#ifdef HAVE_WANNIER90
    call wannier90lib_initialize_matrix_data(this%w90main, this%u_matrix, &
      this%u_dis_matrix, this%m_matrix, this%centers)
#endif
 
    call restart%end()
 
    pop_sub(wannier_restart_read_data)
  end subroutine wannier_restart_read_data
 
  subroutine wannier_write_iter(this, namespace, outp, ions, kpoints)
    class(wannier_t), intent(in) :: this
    type(namespace_t), intent(in)   :: namespace
    type(output_t), intent(inout)   :: outp
    class(ions_t), intent(in)       :: ions
    type(kpoints_t), intent(in)     :: kpoints
 
    character(len=MAX_PATH_LEN) :: dir, raw_dir
 
    push_sub(wannier_write_iter)
 
    ! Only write wannier90 output every n iterations, as specified in input
    ! This must always coincide with the update, since there is only one variable for both update and output interval.
    if (mod(this%last_iter, this%options%td_update_interval) == 0) then
      write(dir, '(a,a,i7.7)') trim(outp%iter_dir),"td.", this%last_iter
      raw_dir = io_workpath(trim(dir), namespace)
 
      call io_mkdir(trim(dir), namespace, parents=.true.)
 
      select case (this%options%td_method)
      case (td_wannier_method_wannier90)
        call wannier_calc_write_u_matrices(trim(raw_dir), trim(this%options%prefix), kpoints, this%u_matrix, this%u_dis_matrix)
        call wannier_calc_write_centers(trim(raw_dir), trim(this%options%prefix), this%centers, ions)
        call wannier_calc_write_spreads(trim(raw_dir), trim(this%options%prefix), this%spreads)
        if (this%options%td_write_hr) then
          call wannier_calc_write_hr_file(trim(raw_dir), trim(this%options%prefix), this%irvec, this%ndegen, this%ham_r)
        end if
      case default
        ! Do nothing
      end select
    end if
 
    pop_sub(wannier_write_iter)
  end subroutine wannier_write_iter
 
  subroutine wannier_update_matrices(this, iter, namespace, gr, space, st, kpoints, hm, dt)
    class(wannier_t),            intent(inout) :: this
    integer,                     intent(in)    :: iter
    type(namespace_t),           intent(in)    :: namespace
    class(mesh_t),               intent(in)    :: gr
    class(space_t),              intent(in)    :: space
    type(states_elec_t),         intent(inout) :: st            ! inout because calculation of H(R) requires it
    type(kpoints_t),             intent(in)    :: kpoints
    type(hamiltonian_elec_t),    intent(in)    :: hm
    real(real64),                intent(in)    :: dt
 
#ifdef HAVE_WANNIER90
    integer :: iw, ik
    complex(real64), allocatable :: u_start(:,:,:)
#endif
 
    push_sub(wannier_update_matrices)
 
    if (this%options%td_method == td_wannier_method_none) then
      pop_sub(wannier_update_matrices)
      return
    end if
 
    ! Keep track of iteration for which the update was performed
    this%last_iter = iter
 
    ! This is needed at every step, so that the accumulated phase covers the whole
    ! interval between two wannierizations
    call wannier_calc_update_phase(this%options%band_index, st, dt, this%phase)
 
    if (mod(iter, this%options%td_update_interval) /= 0) then
      pop_sub(wannier_update_matrices)
      return
    end if
 
    select case (this%options%td_method)
    case (td_wannier_method_wannier90)
      ! Brute-force wannierization of TD KS states by calling libwannier90
#ifdef HAVE_WANNIER90
      ! Warm start from V^dagger U, i.e. the previous U with the phase a stationary state
      ! would have picked up since then
      safe_allocate(u_start(1:this%options%num_wann, 1:this%options%num_wann, 1:kpoints%full%npoints))
      do ik = 1, kpoints%full%npoints
        do iw = 1, this%options%num_wann
          u_start(:, iw, ik) = this%phase(:, ik) * this%u_matrix(:, iw, ik)
        end do
      end do
 
      call wannier90lib_wannierization(this%w90main, gr, space, st, this%latt, kpoints, &
        this%options, this%options%band_index, &
        this%m_matrix, this%u_dis_matrix, u_start)
 
      safe_deallocate_a(u_start)
 
      ! The new U refers to the states at the current time, so the phase starts over
      this%phase = m_z1
 
      ! Convert units and fold the centers into the unit cell
      do iw = 1, size(this%centers, 2)
        this%centers(1:3, iw) = units_to_atomic(unit_angstrom, this%w90main%wannier_data%centres(1:3, iw))
        this%centers(1:3, iw) = this%latt%fold_into_cell(this%centers(1:3, iw)) !using minimum-image convention
      end do
 
      !Update the spreads
      this%spreads = this%w90main%wannier_data%spreads * units_to_atomic(unit_angstrom, 1.0_real64)**2
 
      !Update lattice model
      if (this%options%td_write_hr) then
        call wannier_calc_compute_ham_r(namespace, gr, st, hm, kpoints, this%options, &
          this%u_matrix, this%irvec, this%ham_r)
      end if
#endif
    case default
      ! Do nothing
    end select
 
    pop_sub(wannier_update_matrices)
  end subroutine wannier_update_matrices
 
  subroutine wannier_finalize(wan)
    type(wannier_t), intent(inout) :: wan
 
    push_sub(wannier_finalize)
 
    if (wan%options%td_method == td_wannier_method_none) then
      pop_sub(wannier_finalize)
      return
    end if
 
    safe_deallocate_a(wan%m_matrix)
    safe_deallocate_a(wan%u_matrix)
    safe_deallocate_a(wan%u_dis_matrix)
    safe_deallocate_a(wan%centers)
    safe_deallocate_a(wan%spreads)
    safe_deallocate_a(wan%phase)
    safe_deallocate_a(wan%ham_r)
    safe_deallocate_a(wan%irvec)
    safe_deallocate_a(wan%ndegen)
 
#ifdef HAVE_WANNIER90
    ! Currently parsing exclude_bands from .win is only supported
    ! when using libwannier90
    safe_deallocate_a(wan%options%exclude_list)
    safe_deallocate_a(wan%options%band_index)
#endif
 
    ! Attributes of `lib_common_type` are just pointers to memory we own, so no additional finalisation required
 
    pop_sub(wannier_finalize)
  end subroutine wannier_finalize
 
end module wannier_oct_m