wannier_calc.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_calc_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 global_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 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 wannier_opts_oct_m
  use wfs_elec_oct_m, only: wfs_elec_t
  use ylm_wannier_oct_m
 
  implicit none
  private
 
  public :: &
    wannier_calc_load_restart, &
    wannier_calc_proj_t, &
    wannier_calc_create_amn, &
    wannier_calc_create_mmn, &
    wannier_calc_read_centers, &
    wannier_calc_read_u_matrices, &
    wannier_calc_write_u_matrices, &
    wannier_calc_write_centers
 
  type wannier_calc_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 wannier_calc_proj_t
 
contains
 
  subroutine wannier_calc_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)
    if (.not. (restart_states_ == restart_gs .or. restart_states_ == restart_td)) then
      message(1) = 'Invalid restart_states value; expected RESTART_GS or RESTART_TD.'
      call messages_fatal(1)
    end if
 
    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) then
      write(message(1),'(a)') 'Restart structure not commensurate, since spin dimensions differ.'
      call messages_fatal(1)
    else if (nik /= kpoints%reduced%npoints) then
      write(message(1),'(a)') 'Restart structure not commensurate, since number of k-points differ.'
      call messages_fatal(1)
    else if (nst /= st%nst) then
      write(message(1),'(a)') 'Restart structure not commensurate, since number of states differ.'
      call messages_fatal(1)
    end if
 
    call states_elec_load(restart, namespace, space, st, gr, kpoints, fixed_occ=st%restart_fixed_occ , &
      ierr=ierr, label = ": wannier90")
    call restart%end()
 
  end subroutine wannier_calc_load_restart
 
  ! TODO: Removed SCDM - Should reintroduce once the code is working
  subroutine wannier_calc_create_amn(wan_opts, exclude_list, &
    band_index, space, mesh, latt, st, kpoints, &
    spin_channel, num_proj, proj_input, projection)
    type(wannier_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(wannier_calc_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_calc_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, 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, 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_calc_create_amn)
 
  end subroutine wannier_calc_create_amn
 
  subroutine wannier_calc_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_calc_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_calc_create_mmn)
 
  end subroutine wannier_calc_create_mmn
 
  subroutine wannier_calc_read_centers(wan_opts, centers)
    type(wannier_opts_t), intent(in)    :: wan_opts
    real(real64),         intent(inout) :: centers(:, :)
 
    integer :: w90_xyz, iw
    character(len=2) :: dum
    logical :: exist
 
    push_sub(wannier_calc_read_centers)
 
    assert(size(centers, 2) == wan_opts%num_wann)
    assert(size(centers, 1) == 3)
 
    inquire(file=trim(adjustl(wan_opts%prefix))//'_centres.xyz',exist=exist)
    if (.not. exist) then
      message(1) = 'Cannot find the Wannier90 file seedname_centres.xyz.'
      call messages_fatal(1)
    end if
 
    w90_xyz = io_open(trim(adjustl(wan_opts%prefix))//'_centres.xyz', global_namespace, action='read')
 
    !Skip two lines
    read(w90_xyz, *)
    read(w90_xyz, *)
    do iw = 1, wan_opts%num_wann
      read(w90_xyz, *) dum, centers(1:3, iw)
      ! Wannier90 outputs the coordinates in angstrom
      centers(1:3, iw) = units_to_atomic(unit_angstrom, centers(1:3, iw))
    end do
    call io_close(w90_xyz)
 
    pop_sub(wannier_calc_read_centers)
  end subroutine wannier_calc_read_centers
 
  subroutine wannier_calc_read_u_matrices(wan_opts, kpoints, u_matrix, u_dis_matrix)
    type(wannier_opts_t), intent(in)    :: wan_opts
    type(kpoints_t),      intent(in)    :: kpoints
    complex(real64),      intent(inout) :: u_matrix(:, :, :)
    complex(real64),      intent(inout) :: u_dis_matrix(:, :, :)
 
    integer :: w90_u_mat, w90_u_dis, ik, ib, iw, iw2, nik, nwann, nib
    integer :: num_kpts
    logical :: exist, have_disentangled
 
    num_kpts = product(kpoints%nik_axis(1:3))
 
    push_sub(wannier_calc_read_u_matrices)
 
    inquire(file=trim(adjustl(wan_opts%prefix))//'_u.mat',exist=exist)
    if (.not. exist) then
      message(1) = 'Cannot find the Wannier90 seedname_u.mat file.'
      call messages_fatal(1)
    end if
    w90_u_mat = io_open(trim(adjustl(wan_opts%prefix))//'_u.mat', global_namespace, action='read')
 
    !Skip one line
    read(w90_u_mat, *)
    !Read num_kpts, num_wann, num_wann for consistency check
    read(w90_u_mat, *) nik, nwann, nwann
    if (nik /= num_kpts .or. nwann /= wan_opts%num_wann) then
      message(1) = "The U matrix has inconsistent shape."
      call messages_fatal(1)
    end if
 
    do ik = 1, num_kpts
      !Skip one line
      read(w90_u_mat, *)
      !Skip one line (k-point coordinate)
      read(w90_u_mat, *)
      read(w90_u_mat, '(f15.10,sp,f15.10)') ((u_matrix(iw, iw2, ik), iw=1, wan_opts%num_wann), iw2=1, wan_opts%num_wann)
    end do
 
    call io_close(w90_u_mat)
 
    inquire(file=trim(adjustl(wan_opts%prefix))//'_u_dis.mat',exist=have_disentangled)
    if (have_disentangled) then
      ! Read u_dis.mat: dimensions are (w90_num_bands x w90_num_wann) per k-point.
      ! The band rows follow the compressed, non-excluded band ordering used by
      ! the rest of the Wannier90 interface (i.e. after applying exclude_bands).
      w90_u_dis = io_open(trim(adjustl(wan_opts%prefix))//'_u_dis.mat', global_namespace, action='read')
      !Skip one line
      read(w90_u_dis, *)
      !Read num_kpts, num_wann, num_bands for consistency check
      read(w90_u_dis, *) nik, nwann, nib
      if (nik /= num_kpts .or. nwann /= wan_opts%num_wann .or. nib /= wan_opts%num_bands) then
        message(1) = 'The u_dis matrix has inconsistent shape.'
        call messages_fatal(1)
      end if
 
      do ik = 1, num_kpts
        !Skip one line
        read(w90_u_dis, *)
        !Skip one line (k-point coordinate)
        read(w90_u_dis, *)
        read(w90_u_dis, '(f15.10,sp,f15.10)') ((u_dis_matrix(ib, iw, ik), ib=1, wan_opts%num_bands), iw=1, wan_opts%num_wann)
      end do
 
      call io_close(w90_u_dis)
 
      ! In case of disentanglement, u_dis matrix needs to be reordered to the correct layout.
      ! See issue #1465.
    end if
 
    pop_sub(wannier_calc_read_u_matrices)
  end subroutine wannier_calc_read_u_matrices
 
  subroutine wannier_calc_write_centers(dir, prefix, centers, ions)
    character(len=*), intent(in) :: dir
    character(len=*), intent(in) :: prefix
    real(real64),     intent(in) :: centers(:, :)
    class(ions_t),    intent(in) :: ions
 
    integer :: w90_xyz, iw, ia, ind
    character(len=512) :: header
    integer :: date_time(8)
 
    push_sub(wannier_calc_write_centers)
 
    w90_xyz = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_centres.xyz', global_namespace, action='write', die=.false.)
    if (w90_xyz == -1) then
      message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_centres.xyz for writing.'
      call messages_fatal(1)
    end if
 
    ! Get current date and time
    call date_and_time(values=date_time)
 
    write(w90_xyz, '(i6)') size(centers, 2) + ions%natoms
    write(header, '(a,i4,a1,i2.2,a1,i2.2,a,i2.2,a1,i2.2,a1,i2.2)') &
      'Wannier centers, written by octopus on ', date_time(1), '/', date_time(2), '/', date_time(3), &
      ' at ', date_time(5), ':', date_time(6), ':', date_time(7)
    write (w90_xyz, *) trim(header)
 
    do iw = 1, size(centers, 2)
      write(w90_xyz, '("X",6x,3(f14.8,3x))') &
        (units_from_atomic(unit_angstrom, centers(ind, iw)), ind=1, 3)
    end do
    do ia = 1, ions%natoms
      write(w90_xyz, '(a2,5x,3(f14.8,3x))') trim(ions%atom(ia)%label), &
        (units_from_atomic(unit_angstrom, ions%pos(ind, ia)), ind=1, 3)
    end do
    call io_close(w90_xyz)
 
    pop_sub(wannier_calc_write_centers)
 
  end subroutine wannier_calc_write_centers
 
  subroutine wannier_calc_write_u_matrices(dir, prefix, kpoints, u_matrix, u_dis_matrix)
    character(len=*), intent(in) :: dir
    character(len=*), intent(in) :: prefix
    type(kpoints_t),  intent(in) :: kpoints
    complex(real64),  intent(in) :: u_matrix(:, :, :)
    complex(real64),  intent(in) :: u_dis_matrix(:, :, :)
 
    integer :: w90_u_mat, w90_u_dis
    logical :: have_disentangled
    integer :: i, j, nkp, num_wann, num_bands, num_kpts
    character(len=512) :: header
    integer :: date_time(8)
 
    num_kpts = size(u_matrix, 3)
    num_wann = size(u_matrix, 1)
    num_bands = size(u_dis_matrix, 1)
    assert(num_kpts == size(kpoints%reduced%red_point, 2))
 
    push_sub(wannier_calc_write_u_matrices)
 
    ! Get current date and time
    call date_and_time(values=date_time)
 
    w90_u_mat = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_u.mat', global_namespace, action='write', die=.false.)
    if (w90_u_mat == -1) then
      message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_u.mat for writing.'
      call messages_fatal(1)
    end if
 
    write(header, '(a,i4,a1,i2.2,a1,i2.2,a,i2.2,a1,i2.2,a1,i2.2)') &
      'written on ', date_time(1), '/', date_time(2), '/', date_time(3), &
      ' at ', date_time(5), ':', date_time(6), ':', date_time(7)
    write (w90_u_mat, *) trim(header)
    write (w90_u_mat, *) num_kpts, num_wann, num_wann
 
    do nkp = 1, num_kpts
      write (w90_u_mat, *)
      write (w90_u_mat, '(f15.10,sp,f15.10,sp,f15.10)') -kpoints%reduced%red_point(1:3, nkp)
      write (w90_u_mat, '(f15.10,sp,f15.10)') ((u_matrix(i, j, nkp), i=1, num_wann), j=1, num_wann)
    end do
 
    call io_close(w90_u_mat)
 
    inquire(file=trim(trim(adjustl(prefix))//'_u_dis.mat'),exist=have_disentangled)
    if (have_disentangled) then
 
      w90_u_dis = io_open(trim(adjustl(dir))//'/'//trim(adjustl(prefix))//'_u_dis.mat', &
        global_namespace, action='write', die=.false.)
      if (w90_u_dis == -1) then
        message(1) = 'Unable to open output file '//trim(adjustl(prefix))//'_u_dis.mat for writing.'
        call messages_fatal(1)
      end if
 
      write (w90_u_dis, *) trim(header)
      write (w90_u_dis, *) num_kpts, num_wann, num_bands
      do nkp = 1, num_kpts
        write (w90_u_dis, *)
        write (w90_u_dis, '(f15.10,sp,f15.10,sp,f15.10)') -kpoints%reduced%red_point(1:3, nkp)
        write (w90_u_dis, '(f15.10,sp,f15.10)') ((u_dis_matrix(i, j, nkp), i=1, num_bands), j=1, num_wann)
      end do
 
      call io_close(w90_u_dis)
    end if
 
    pop_sub(wannier_calc_write_u_matrices)
 
  end subroutine wannier_calc_write_u_matrices
 
end module wannier_calc_oct_m