wannier.F90

Go to the documentation of this file.

!! Copyright (C) A Buccheri, J Reimann. 2026
!!
!! This Source Code Form is subject to the terms of the Mozilla Public
!! License, v. 2.0. If a copy of the MPL was not distributed with this
!! file, You can obtain one at https://mozilla.org/MPL/2.0/.
!!
 
#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 mesh_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 restart_oct_m, only: restart_gs, restart_td
  use space_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_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(wan_opts_t)            :: wan_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
 
  integer                     :: idir, ii, itemp, ik, ist
 
  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, int(option__calculationmode__dummy, int32))
  call sys%init_parallelization(mpi_world)
 
  assert(sys%space%dim == 3)
 
  ! 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 wan_opts_parse(namespace, wan_opts)
 
  ! Initialize wannier90 lib object
  call wannier90lib_init_w90main(namespace, sys%space, sys%gr, sys%ions, sys%kpoints, sys%st, &
    wan_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_load_restart(global_namespace, sys%mc, sys%space, sys%st, sys%gr, sys%kpoints, wan_opts%restart_from)
 
  ! Convert the exclude bands to exclude list
  safe_allocate(band_index(1:sys%st%nst))
  safe_allocate(exclude_list(1:sys%st%nst))
  exclude_list(1:sys%st%nst) = .false.
 
  if (allocated(w90main%exclude_bands)) then
    do itemp = 1, size(w90main%exclude_bands)
      exclude_list(w90main%exclude_bands(itemp)) = .true.
    end do
  end if
 
  itemp = 0
  do ii = 1, sys%st%nst
    if (exclude_list(ii)) cycle
    itemp = itemp + 1
    band_index(ii) = itemp
  end do
 
  ! Initial setup (previously in nnkp file)
  call w90_get_nn(w90main, nn, stdout, stderr, ierr)
  safe_allocate(nnkp(w90main%num_kpts, nn)) 
  safe_allocate(gkpb(3, w90main%num_kpts, 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
  safe_allocate(m_matrix(wan_opts%num_bands, wan_opts%num_bands, nn, w90main%num_kpts))
  safe_allocate(u_matrix_opt(wan_opts%num_bands, wan_opts%num_wann, w90main%num_kpts))
  safe_allocate(eigenvals(wan_opts%num_bands, w90main%num_kpts))
  safe_allocate(u_matrix(wan_opts%num_wann, wan_opts%num_wann, 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 (wan_opts%calc_overlaps) then
    call wannier90lib_create_wannier90_mmn(w90main, exclude_list, band_index, sys%gr, sys%st, sys%ions, m_matrix)
    call wannier90lib_create_wannier90_amn(w90main, wan_opts, exclude_list, 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 (wan_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