wannier Program

program wannier
  !! The main Wannier90 program

#ifdef W90_MPI08
  use mpi_f08
#endif
#ifdef W90_MPI90
  use mpi
#endif

  use w90_library
  use w90_library_extra ! for input_reader_special, overlaps, etc

  use w90_comms, only: w90_comm_type, comms_sync_error
  use w90_io, only: io_commandline, io_date, io_time, prterr
  use w90_sitesym, only: sitesym_read
  use w90_error, only: w90_error_type, set_error_input

  implicit none

#ifdef W90_MPIH
  include 'mpif.h'
#endif

! nvfortran fails to correctly handle the dp private attribute in module w90_library
!       nvfortran 25.7-0 64-bit target on x86-64 Linux -tp znve
! if needed in practice, use "only" or rename dp in use directive
#ifndef __NVCOMPILER
  integer, parameter :: dp = kind(0.d0)
#endif

  character(len=:), allocatable :: seedname, progname
  character(len=20) :: cpstatus ! checkpoint file status
  character(len=:), pointer :: restart
  complex(kind=dp), allocatable :: m_matrix_loc(:, :, :, :)
  complex(kind=dp), allocatable :: u_matrix(:, :, :)
  complex(kind=dp), allocatable :: u_matrix_opt(:, :, :)
  real(kind=dp), allocatable :: eigval(:, :)
  integer, allocatable :: dist_k(:)
  integer :: mpisize, rank, ierr, nkl
  integer, pointer :: nb, nk, nw, nn
  integer :: stdout, stderr
  logical, pointer :: pp
  logical :: ld, lovlp, ldsnt, lwann, lplot, ltran, need_eigvals
  type(lib_common_type), target :: common_data
  type(w90_error_type), allocatable :: error
  character(len=9) :: cdate, ctime

  pp => common_data%w90_calculation%postproc_setup
  restart => common_data%w90_calculation%restart
  nw => common_data%num_wann
  nb => common_data%num_bands
  nk => common_data%num_kpts
  nn => common_data%kmesh_info%nntot

  progname = 'wannier90' ! https://gcc.gnu.org/bugzilla/show_bug.cgi?id=91442
  call io_commandline(progname, ld, pp, seedname)

  call w90_get_fortran_stderr(stderr) ! for early error cases; maybe overwritten later
#ifdef W90_MPI
  call mpi_init(ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: mpi_init() returned an error!'
    stop 1
  end if
  call mpi_comm_rank(mpi_comm_world, rank, ierr) ! the type of comm_world depends on interface used
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: mpi_comm_rank() returned an error!'
    stop 1
  end if
  call mpi_comm_size(mpi_comm_world, mpisize, ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: mpi_comm_size() returned an error!'
    stop 1
  end if
  call w90_set_comm(common_data, mpi_comm_world)
#else
  rank = 0
  mpisize = 1
#endif

  ! open main output file
  if (rank == 0) open (newunit=stdout, file=seedname//'.wout', status="replace")

  ! open main error file
  if (rank == 0) open (newunit=stderr, file=seedname//'.werr', status="replace")

  call io_date(cdate, ctime)
  if (rank == 0) write (stderr, *) 'Wannier90: Execution started on ', cdate, ' at ', ctime

  ! read key parameters from .win file
  call input_reader_special(common_data, seedname, stdout, stderr, ierr)
  if (ierr /= 0) stop 1

  ! read all remaining parameters from .win file
  call w90_input_reader(common_data, stdout, stderr, ierr)
  if (ierr /= 0) stop 1

  ! write useful info (includes jazzy header info)
  call w90_print_info(common_data, stdout, stderr, ierr)
  if (ierr /= 0) stop 1

  ! special branch for writing nnkp file
  ! exit immediately after writing the nnkp file
  if (pp) then
    call write_kmesh(common_data, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
    if (rank == 0) close (unit=stderr, status='delete')
    if (rank == 0) write (stdout, '(1x,a25,f11.3,a)') 'Time to write kmesh      ', io_time(), ' (sec)'
    if (rank == 0) write (stdout, '(/a)') ' Exiting... '//trim(seedname)//'.nnkp written.'
#ifdef W90_MPI
    call mpi_finalize(ierr)
#endif
    stop
  end if

  ! test mpi error handling using "unlucky" input token
  if (rank == -common_data%print_output%timing_level) then
    call set_error_input(error, 'received unlucky_rank', common_data%comm)
  else
    ! this is necessary since non-root may never enter an mpi collective if root has exited here
    call comms_sync_error(common_data%comm, error, 0)
  end if
  if (allocated(error)) then ! applies (is t) for all ranks now
    call prterr(error, ierr, stdout, stderr, common_data%comm)
#ifdef W90_MPI
    call mpi_finalize(ierr) ! let's be nice
#endif
    stop 1
  end if
  ! end unlucky code

  ! setup kpoint distribution
  allocate (dist_k(nk), stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to allocate dist_k array!'
    stop 1
  end if
  ! get a basic k-point/rank distribution
  call w90_distribute_kpts(common_data, nk, mpisize, dist_k, stdout, stderr, ierr)
  if (ierr /= 0) stop 1

  ! copy distribution to library
  call set_kpoint_distribution(common_data, dist_k, stdout, stderr, ierr)
  if (ierr /= 0) stop 1

  ! setup SAWF data
  if (common_data%lsitesymmetry) then
    call sitesym_read(common_data%sitesym, nb, nk, nw, seedname, error, common_data%comm) ! (not a library call)
    if (allocated(error)) then
      write (stderr, *) 'Wannier90: failed to setup symmetry!'
      deallocate (error)
      stop 1
    end if
  end if

  call w90_get_nn(common_data, nn, stdout, stderr, ierr)
  nkl = count(dist_k == rank) ! number of kpoints this rank
  !write (*, *) 'rank, nw, nb, nk, nn, nk(rank): ', rank, nw, nb, nk, nn, nkl

  allocate (m_matrix_loc(nb, nb, nn, nkl), stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to allocate m_matrix_loc!'
    stop 1
  end if
  call w90_set_m_local(common_data, m_matrix_loc)  ! we don't need global m

  allocate (u_matrix(nw, nw, nk), stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to allocate u_matrix!'
    stop 1
  end if
  call w90_set_u_matrix(common_data, u_matrix)

  allocate (u_matrix_opt(nb, nw, nk), stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to allocate u_matrix_opt!'
    stop 1
  end if
  call w90_set_u_opt(common_data, u_matrix_opt)

! restart system
  lovlp = .true.
  ldsnt = .true.
  lwann = .true.
  lplot = .true.
  ltran = .false.

  if (restart == '') then
    if (rank == 0) write (stdout, '(1x,a/)') 'Starting a new Wannier90 calculation ...'
  else
    cpstatus = ''
    call read_chkpt(common_data, cpstatus, stdout, stderr, ierr)
    if (ierr /= 0) stop 1

    if (restart == 'wannierise' .or. (restart == 'default' .and. cpstatus == 'postdis')) then
      if (rank == 0) write (stdout, '(1x,a/)') 'Restarting Wannier90 from wannierisation ...'
      lovlp = .false.
      ldsnt = .false.
      lwann = .true.
      lplot = .true.
      ltran = .false.
    elseif (restart == 'plot' .or. (restart == 'default' .and. cpstatus == 'postwann')) then
      if (rank == 0) write (stdout, '(1x,a/)') 'Restarting Wannier90 from plotting routines ...'
      lovlp = .false.
      ldsnt = .false.
      lwann = .false.
      lplot = .true.
      ltran = .false.
    elseif (restart == 'transport') then
      if (rank == 0) write (stdout, '(1x,a/)') 'Restarting Wannier90 from transport routines ...'
      lovlp = .false.
      ldsnt = .false.
      lwann = .false.
      lplot = .false.
      ltran = .true.
      !else
      ! illegitimate restart choice, should declaim the acceptable choices
    end if
  end if
  ltran = (ltran .or. common_data%w90_calculation%transport)
  ldsnt = (ldsnt .and. (nw < nb)) ! disentanglement only needed if space reduced

  ! circumstances where eigenvalues are needed are a little overcomplicated
  need_eigvals = .false.
  need_eigvals = common_data%w90_calculation%bands_plot
  need_eigvals = (need_eigvals .or. common_data%w90_calculation%fermi_surface_plot)
  need_eigvals = (need_eigvals .or. common_data%output_file%write_hr)
  need_eigvals = (need_eigvals .or. common_data%output_file%write_tb)
  need_eigvals = (need_eigvals .or. ldsnt) ! disentanglement anyway requires evals

  if (need_eigvals) then
    allocate (eigval(nb, nk), stat=ierr)
    if (ierr /= 0) then
      write (stderr, *) 'Wannier90: failed to allocate eigval array!'
      stop 1
    end if
    call read_eigvals(common_data, eigval, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
    call w90_set_eigval(common_data, eigval)
  end if

  ! ends setup

  if (lovlp) then
    call overlaps(common_data, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
  end if

  if (ldsnt) then
    call w90_disentangle(common_data, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
    call write_chkpt(common_data, 'postdis', stdout, stderr, ierr)
    if (ierr /= 0) stop 1
  end if

  if (lwann) then
    call w90_project_overlap(common_data, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
    call w90_wannierise(common_data, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
    call write_chkpt(common_data, 'postwann', stdout, stderr, ierr)
    if (ierr /= 0) stop 1
  end if

  if (lplot) then
    call w90_plot(common_data, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
  end if

  if (ltran) then
    call w90_transport(common_data, stdout, stderr, ierr)
    if (ierr /= 0) stop 1
  end if

  ! cleanup

  if (need_eigvals) then
    deallocate (eigval, stat=ierr)
    if (ierr /= 0) then
      write (stderr, *) 'Wannier90: failed to deallocate eigval array!'
      stop 1
    end if
  end if
  deallocate (dist_k, stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to deallocate dist_k array!'
    stop 1
  end if
  deallocate (m_matrix_loc, stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to deallocate m_matrix_loc!'
    stop 1
  end if
  deallocate (u_matrix, stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to deallocate u_matrix!'
    stop 1
  end if
  deallocate (u_matrix_opt, stat=ierr)
  if (ierr /= 0) then
    write (stderr, *) 'Wannier90: failed to deallocate u_matrix_opt!'
    stop 1
  end if

  call print_times(common_data, stdout)

  if (rank == 0) then
    close (unit=stderr, status='delete')
    write (stdout, '(1x,a)') 'All done: wannier90 exiting'
    close (unit=stdout)
  end if

#ifdef W90_MPI
  call mpi_finalize(ierr)
#endif
end program wannier