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"
 
program wannier
  use, intrinsic :: iso_fortran_env
 
  use calc_mode_par_oct_m
  use debug_oct_m
  use electrons_oct_m, only: electrons_t
  use electron_space_oct_m
  use fft_oct_m
  use global_oct_m
  use grid_oct_m
  use io_oct_m
  use kpoints_oct_m
  use ions_oct_m
  use kpoints_oct_m
  use mpi_oct_m
  use messages_oct_m
  use multicomm_oct_m
  use namespace_oct_m
  use parser_oct_m
  use profiling_oct_m
  use states_elec_oct_m
  use states_elec_io_oct_m
  use states_elec_restart_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use wannier_calc_oct_m
  use wannier_opts_oct_m
  use wannier90lib_interface_oct_m
 
  ! External Wannier 90 library
  use w90_library
  use w90_library_extra, only:overlaps
 
  implicit none
 
  type(namespace_t)           :: namespace
  type(electrons_t), pointer  :: sys
  type(wannier_opts_t)        :: wannier_opts
  integer,        allocatable :: band_index(:)
 
  type(lib_common_type)             :: w90main
 
  complex(real64), allocatable :: m_matrix(:, :, :, :)
  complex(real64), allocatable :: u_matrix(:, :, :)
  complex(real64), allocatable :: u_matrix_opt(:, :, :)
  real(real64), allocatable    :: eigenvals(:,:)
 
  integer, allocatable        :: nnkp(:,:), gkpb(:,:,:)
  logical, allocatable        :: exclude_list(:)
  integer                     :: nn, nk_local
 
  integer                     :: idir, ik, ist, iband
 
  integer :: ierr
  real(real64)                 :: factor(3)
 
  ! Initialise all Octopus singletons
  call global_init()
  call parser_init()
  call messages_init()
  call io_init()
  namespace = namespace_t("Wannier90Lib", parent = global_namespace)
  call profiling_init(namespace)
  call fft_all_init(namespace)
  call unit_system_init(namespace)
 
  ! Initialise a system object
  call calc_mode_par%set_parallelization(p_strategy_states, default = .false.)
  sys => electrons_t(global_namespace, mpi_world, int(option__calculationmode__dummy, int32))
 
  if (sys%space%dim /= 3) then ! see issue #1466
    call messages_not_implemented("oct-wannier90lib: only 3D is supported")
  end if
 
  ! Sanity checks
  if (sys%kpoints%use_symmetries) then
    message(1) = 'oct-wannier90: k-points symmetries are not allowed'
    call messages_fatal(1)
  end if
  if (sys%kpoints%use_time_reversal) then
    message(1) = 'oct-wannier90: time-reversal symmetry is not allowed'
    call messages_fatal(1)
  end if
  if (sys%kpoints%reduced%nshifts > 1) then
    message(1) = 'oct-wannier90: Wannier90 does not allow for multiple shifts of the k-point grid'
    call messages_fatal(1)
  end if
  if (sys%st%d%nspin == 4) then
    call messages_not_implemented('oct-wannier90: Spinors currently not supported', namespace)
  end if
 
  ! "Correct" rlattice/klattice for Wannier90 (3D periodicity assumed here).
  ! Note: this is how the lattice is scaled in src/utils/wannier90_interface.F90
  factor = m_one
  do idir = sys%space%periodic_dim+1, sys%space%dim
    factor(idir) = m_two * sys%gr%box%bounding_box_l(idir)
  end do
  call sys%ions%latt%scale(factor)
 
  ! Parsing all input related to wannier90
  call wannier_opts%parse_oct(namespace)
  call wannier_opts%parse_win()
 
  ! Initialize wannier90 lib object
  call wannier90lib_init_w90main(namespace, sys%ions, sys%ions%latt, sys%kpoints, sys%st, &
    wannier_opts, w90main, stdout, stderr, ierr)
  if (ierr /= 0) then
    write(message(1),'(a,i0)') 'Error initializing wannier90 library: ', ierr
    call messages_fatal(1)
  end if
 
  ! Check for consistency of spin polarization
  if (sys%st%d%ispin /= unpolarized) then
    call messages_experimental("oct-wannier90 with SpinComponents /= unpolarized")
  elseif(sys%st%d%ispin == unpolarized .and. w90main%wvfn_read%spin_channel == 2) then
    w90main%wvfn_read%spin_channel = 1
    message(1) = "octopus was run with SpinComponents == unpolarized. Ignoring wannier90 spin input variable"
    call messages_warning(1)
  end if
 
  ! Load octopus states
  call wannier_calc_load_restart(global_namespace, sys%mc, sys%space, sys%st, sys%gr, sys%kpoints, wannier_opts%restart_from)
 
  ! Update exclude_list and band_index based on wannier90 input from .win file (if any)
  safe_allocate(exclude_list(1:sys%st%nst))
  exclude_list = .false.
  if (allocated(w90main%exclude_bands)) then
    do iband = 1, size(w90main%exclude_bands)
      ist = w90main%exclude_bands(iband)
      exclude_list(ist) = .true.
    end do
  end if
  ! Currently TD Wannier is not supported for exclude_band case
  if (any(exclude_list)) then
    call messages_not_implemented('TD Wannier: exclude_bands not supported.')
  end if
  safe_allocate(band_index(1:count(.not. exclude_list)))
  band_index = pack([(ist, ist=1,sys%st%nst)], exclude_list .eqv. .false.)
 
  ! Initial setup (previously in nnkp file)
  call w90_get_nn(w90main, nn, stdout, stderr, ierr)
  safe_allocate(nnkp(1:w90main%num_kpts, 1:nn)) 
  safe_allocate(gkpb(1:3, 1:w90main%num_kpts, 1:nn))
  call w90_get_nnkp(w90main, nnkp, stdout, stderr, ierr)  
  call w90_get_gkpb(w90main, gkpb, stdout, stderr, ierr)
 
  ! Prepare computation of quantities required by Wannier90
  ! m_matrix is distributed in k, matching wannier90''s m_matrix_local
  nk_local = sys%st%d%kpt%nlocal()
  if (sys%st%d%ispin == spin_polarized) nk_local = nk_local / 2
  safe_allocate(m_matrix(1:wannier_opts%num_bands, 1:wannier_opts%num_bands, 1:nn, 1:nk_local))
  safe_allocate(u_matrix_opt(1:wannier_opts%num_bands, 1:wannier_opts%num_wann, 1:w90main%num_kpts))
  safe_allocate(eigenvals(1:wannier_opts%num_bands, 1:w90main%num_kpts))
  safe_allocate(u_matrix(1:wannier_opts%num_wann, 1:wannier_opts%num_wann, 1:w90main%num_kpts))
  call w90_set_m_local(w90main, m_matrix)
  call w90_set_u_opt(w90main, u_matrix_opt)
  call w90_set_u_matrix(w90main, u_matrix)
 
  ! Get eigenvalues from octopus, shape is (num_bands, nk)
  do ik = 1, w90main%num_kpts
    do ist = 1, sys%st%nst
      if (exclude_list(ist)) cycle
      if (sys%st%d%ispin /= spin_polarized) then
        eigenvals(band_index(ist), ik) = units_from_atomic(unit_ev, sys%st%eigenval(ist, ik))
      else
        eigenvals(band_index(ist), ik) = units_from_atomic(unit_ev, sys%st%eigenval(ist, (ik-1)*2+w90main%wvfn_read%spin_channel))
      end if
    end do
  end do
  call w90_set_eigval(w90main, eigenvals)
 
  ! Compute or read quantities required by wannier90
  if (wannier_opts%calc_overlaps) then
    call wannier90lib_create_wannier90_mmn(w90main, band_index, sys%gr, sys%st, sys%ions%latt, m_matrix)
    call wannier90lib_create_wannier90_amn(w90main, wannier_opts, band_index, sys%space, &
      sys%gr, sys%ions%latt, sys%st, sys%kpoints, u_matrix_opt)
    ! TODO: Revisit spin-polarised functionality once spin-unpolarised works
    ! call create_wannier90_spn(sys%gr, sys%st)
    ! TODO: Add write functions here and remove dump_inputs variable in wannier90
  else
    call overlaps(w90main, stdout, stderr, ierr)
    if (ierr /= 0) then
      write(message(1), '(a,i0)') 'Error while reading overlaps with wannier90: ', ierr
      call messages_fatal(1)
    end if
  end if
 
  ! Perform wannierization
  if (wannier_opts%wannierize) then
    call w90_disentangle(w90main, stdout, stderr, ierr)
    if (ierr /= 0) then
      write(message(1), '(a,i0)') 'Error in w90_disentangle: ', ierr
      call messages_fatal(1)
    end if
    call w90_project_overlap(w90main, stdout, stderr, ierr)
    
    if (ierr /= 0) then
      write(message(1), '(a,i0)') 'Error in w90_project_overlap: ', ierr
      call messages_fatal(1)
    end if
    call w90_wannierise(w90main, stdout, stderr, ierr)
    if (ierr /= 0) then
      write(message(1), '(a,i0)') 'Error in w90_wannierise: ', ierr
      call messages_fatal(1)
    end if
    ! Avoid plotting from Wannier90
    w90main%w90_calculation%wannier_plot = .false.
    ! Output U matrices, hr.dat and more output files for plotting with octopus
    call w90_plot(w90main, stdout, stderr, ierr)
    if (ierr /= 0) then
      write(message(1), '(a,i0)') 'Error in w90_plot: ', ierr
      call messages_fatal(1)
    end if
  end if
  ! if (wan_opts%plot) then
  ! Here we generate wannier functions
  ! Currently this is only meaningful after performing the wannierization
  ! TODO: Check functionality to read U matrix from file
  ! generate_wannier_states(space, mesh, ions, st, kpoints)
  ! TODO: Port function to generate wannier states from old utility
  ! end if
 
  ! Free memory of variables in this scope
  safe_deallocate_a(m_matrix)
  safe_deallocate_a(u_matrix)
  safe_deallocate_a(u_matrix_opt)
  safe_deallocate_a(eigenvals)
  safe_deallocate_a(exclude_list)
  safe_deallocate_a(band_index)
  safe_deallocate_a(nnkp)
  safe_deallocate_a(gkpb)
  safe_deallocate_p(sys)
 
  ! End all singletons
  call fft_all_end()
  call io_end()
  call profiling_end(namespace)
  call messages_end()
  call parser_end()
  call global_end()
 
end program wannier