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"
 
module wannier_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 fft_oct_m
  use global_oct_m
  use grid_oct_m
  use io_oct_m
  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 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 submesh_oct_m
  use types_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use wfs_elec_oct_m, only: wfs_elec_t
  use ylm_wannier_oct_m
 
  implicit none
  private
 
  public :: wannier_load_restart, &
    wan_opts_parse, &
    wan_opts_t, &
    wan_proj_t, &
    wannier_create_amn, &
    wannier_create_mmn
 
  type :: wan_opts_t
    integer           :: num_bands
    integer           :: num_wann
    logical           :: use_scdm
    real(real64)      :: scdm_mu
    real(real64)      :: scdm_sigma
    real(real64)      :: threshold
    character(len=80) :: prefix
    integer           :: restart_from
    logical           :: plot, wannierize, calc_overlaps, dump_inputs
    logical           :: spinors = .false. 
    integer,      allocatable  :: spin_proj_component(:)
  end type wan_opts_t
 
  type wan_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 wan_proj_t
 
contains
 
  subroutine wannier_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
    ! Defaults to ground state
 
    integer         :: restart_states_, ierr, nik, dim, nst
    type(restart_t) :: restart
 
    restart_states_ = optional_default(restart_states, restart_gs)
 
    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 .or. nik /= kpoints%reduced%npoints .or. nst /= st%nst) then
      write(message(1),'(a)') 'Restart structure not commensurate.'
      call messages_fatal(1)
    end if
 
    call states_elec_load(restart, namespace, space, st, gr, kpoints, ierr, label = ": wannier90")
    call restart%end()
 
  end subroutine wannier_load_restart
 
 
  subroutine wan_opts_parse(namespace, wannier_opts)
    type(namespace_t),  intent(in)    :: namespace
    type(wan_opts_t), intent(inout)   :: wannier_opts
 
    logical :: read_from_td
 
    !%Variable Wannier90NumWann
    !%Type integer
    !%Default 1
    !%Section Utilities::oct-wannier90
    !%Description
    !% Number of wannier function
    !%End
    call parse_variable(namespace, 'Wannier90NumWann', 1, wannier_opts%num_wann)
 
    !%Variable Wannier90NumBands
    !%Type integer
    !%Default 1
    !%Section Utilities::oct-wannier90
    !%Description
    !% Number of bands considered for wannierization
    !%End
    call parse_variable(namespace, 'Wannier90NumBands', 1, wannier_opts%num_bands)
 
    !%Variable Wannier90Prefix
    !%Type string
    !%Default w90
    !%Section Utilities::oct-wannier90
    !%Description
    !% Prefix for wannier90 files
    !%End
    call parse_variable(namespace, 'Wannier90Prefix', 'w90', wannier_opts%prefix)
    if (wannier_opts%prefix == 'w90') then
      message(1) = "oct-wannier90: the prefix is set by default to w90"
      call messages_info(1)
    end if
 
    !%Variable Wannier90UseTD
    !%Type logical
    !%Default no
    !%Section Utilities::oct-wannier90
    !%Description
    !% By default oct-wannier90 uses the ground-state states to compute the necessary information.
    !% By setting this variable to yes, oct-wannier90 will use the TD states instead.
    !%End
    call parse_variable(global_namespace, 'Wannier90UseTD', .false., read_from_td)
    if (read_from_td) then
      wannier_opts%restart_from = restart_td
    else
      wannier_opts%restart_from = restart_gs
    end if
 
    !%Variable Wannier90DumpFiles
    !%Type logical
    !%Default no
    !%Section Utilities::oct-wannier90
    !%Description
    !% Whether or not to dump intermediate files to disk.
    !% Files are: .amn, .mmn, .eig and .win with inputs.
    !%End
    call parse_variable(global_namespace, 'Wannier90DumpFiles', .true., wannier_opts%dump_inputs)
 
    !%Variable Wannier90Wannierize
    !%Type logical
    !%Default yes
    !%Section Utilities::oct-wannier90
    !%Description
    !% Whether or not to perform the wannierization step (including disentanglement).
    !%End
    call parse_variable(global_namespace, 'Wannier90Wannierize', .true., wannier_opts%wannierize)
 
    !%Variable Wannier90Plot
    !%Type logical
    !%Default no
    !%Section Utilities::oct-wannier90
    !%Description
    !% If set to yes, oct-wannier90 will generate the wannier functions of a real space
    !% grid and write them to file.
    !%End
    call parse_variable(global_namespace, 'Wannier90Plot', .true., wannier_opts%plot)
 
    !%Variable Wannier90CalcOverlaps
    !%Type logical
    !%Default yes
    !%Section Utilities::oct-wannier90
    !%Description
    !% Whether or not to calculate the amn and mmn matrices from octopus states.
    !% By setting this variable to no, oct-wannier90 will attempt to read them from files
    !%End
    call parse_variable(global_namespace, 'Wannier90CalcOverlaps', .true., wannier_opts%calc_overlaps)
 
    !%Variable Wannier90UseSCDM
    !%Type logical
    !%Default no
    !%Section Utilities::oct-wannier90
    !%Description
    !% By default oct-wannier90 uses the projection method to generate the .amn file.
    !% By setting this variable to yes, oct-wannier90 will use SCDM method instead.
    !%End
    call parse_variable(namespace, 'Wannier90UseSCDM', .false., wannier_opts%use_scdm)
    if (wannier_opts%use_scdm) then
      !%Variable SCDMsigma
      !%Type float
      !%Default 0.2
      !%Section Utilities::oct-wannier90
      !%Description
      !% Broadening of SCDM smearing function.
      !%End
      call parse_variable(namespace, 'SCDMsigma', 0.2_real64, wannier_opts%scdm_sigma)
 
      !%Variable SCDMmu
      !%Type float
      !%Section Utilities::oct-wannier90
      !%Description
      !% Energy range up to which states are considered for SCDM.
      !%End
      call parse_variable(namespace, 'SCDMmu', m_huge, wannier_opts%scdm_mu)
    end if
 
    ! We use the variable AOThreshold to deterine the threshold on the radii of the atomic orbitals
    ! TODO: Have the threshold for wannier functions different than for AO
    call parse_variable(namespace, 'AOThreshold', 0.001_real64, wannier_opts%threshold)
 
  end subroutine wan_opts_parse
 
  ! TODO: Removed SCDM - Should reintroduce once the code is working
  
  subroutine wannier_create_amn(wan_opts, exclude_list, &
    band_index, space, mesh, latt, st, kpoints, &
    spin_channel, num_proj, proj_input, projection)
    type(wan_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
    logical,                 intent(in)    :: exclude_list(:)
    integer,                 intent(in)    :: band_index(:)
    integer,                 intent(in)    :: spin_channel
    integer,                 intent(in)    :: num_proj
    type(wan_proj_t),        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_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 = product(kpoints%nik_axis(1:3))
 
    ! 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))
      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, 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
 
      !For spin-polarized calculations, we select the right k-point
      if (st%d%ispin == spin_polarized) then
        ik_real = (ik-1)*2 + spin_channel
      else
        ik_real = ik
      end if
 
      ! TODO(Alex) ist and iw loops need interchanging to optimise array access
      do ist = 1, st%nst
        if (exclude_list(ist)) cycle
        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, ist, 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(iband, :, ik))
          call profiling_out("W90_AMN_REDUCE")
        end if
 
        if(st%d%kpt%parallel) then
          call comm_allreduce(st%d%kpt%mpi_grp, projection(iband, :, 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_create_amn)
 
  end subroutine wannier_create_amn
 
  subroutine wannier_create_mmn(exclude_list, band_index, mesh, st, ions, &
    spin_channel, nnlist, nncell, overlap)
    logical,                     intent(in) :: exclude_list(:)
    integer,                     intent(in) :: band_index(:)
    class(mesh_t),               intent(in) :: mesh
    type(states_elec_t), target, intent(in) :: st
    type(ions_t),                intent(in) :: ions
    integer,                     intent(in) :: spin_channel
    integer,                     intent(in) :: nnlist(:,:)
    integer,                     intent(in) :: nncell(:,:,:)
 
    complex(real64), contiguous, intent(out) :: overlap(:, :, :, :)
 
    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_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 wannerisation"
    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))
 
    if (st%d%kpt%parallel) ik_loc = 0
 
    ! loop over the pairs specified in the nnk tuple
    do ik = 1, num_kpts
      do inn = 1, nntot
 
        iknn = nnlist(ik, inn)
        g(1:3) = nncell(:, ik, inn)
 
        ! For spin-polarized calculations, we select the right k-point
        if (st%d%ispin == spin_polarized) then
          ik_oct = (ik-1)*2 + spin_channel
          iknn_oct = (iknn-1)*2 + spin_channel
        else
          ik_oct = ik
          iknn_oct = iknn
        end if
 
        ! 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(ions%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, st%nst
          if (exclude_list(jst)) cycle
          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, jst, iknn_oct)) then
              call states_elec_get_state(st, mesh, jst, 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, jst, 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 ist = 1, st%nst
              if (exclude_list(ist)) cycle
              iband = band_index(ist)
 
              batch => st%group%psib(st%group%iblock(ist), ik_oct)
 
              select case (batch%status())
              case (batch_not_packed)
                overlap(iband, jband, inn, ik) = m_z0
                do idim = 1, st%d%dim
                  ibind = batch%inv_index((/ist, idim/))
                  overlap(iband, jband, inn, ik) = overlap(iband, jband, inn, ik) + &
                    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, jband, inn, ik) = zmf_dotp(mesh, st%d%dim, psim, psin, reduce = .false.)
              end select
            end do !ist
          end if
 
          call profiling_in("W90_MMN_REDUCE")
          call mesh%allreduce(overlap(:, jband, inn, ik))
          call profiling_out("W90_MMN_REDUCE")
 
          if(st%d%kpt%parallel) then
            call comm_allreduce(st%d%kpt%mpi_grp, overlap(:, jband, inn, ik))
          end if
        end do !jst
      end do !inn
 
      ! This contains a necessary remap: wannier90 only stores the local m_matrix, so it stores
      ! only the first nk_rank kpoints in the matrix. The following block remaps from the global
      ! kpoint to the local kpoint
      if (st%d%kpt%parallel) then
        if (st%d%ispin == spin_polarized) then
          ik_oct = (ik-1)*2 + spin_channel
        else
          ik_oct = ik
        end if
        if (ik_oct >= st%d%kpt%start .and. ik_oct <= st%d%kpt%end) then
          ik_loc = ik_loc + 1
          if (ik_loc /= ik) then
            overlap(:, :, :, ik_loc) = overlap(:, :, :, ik)
          end if
        end if
      end if
    end do !ik
 
    safe_deallocate_a(psim)
    safe_deallocate_a(psin)
    safe_deallocate_a(phase)
 
    call profiling_out("W90_MMN")
 
    pop_sub(wannier_create_mmn)
 
  end subroutine wannier_create_mmn
 
end module wannier_oct_m