isdf_serial.F90

Go to the documentation of this file.

!! Copyright (C) 2024 - 2025. Alexander Buccheri
!!
!! 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, U
 
#include "global.h"
 
module isdf_serial_oct_m
  use, intrinsic :: iso_fortran_env, only: real64, int64
  use batch_oct_m
  use batch_ops_oct_m
  use blas_oct_m
  use debug_oct_m
  use electron_space_oct_m
  use exchange_operator_oct_m, only: exchange_operator_t
  use face_splitting_oct_m
  use global_oct_m
  use grid_oct_m
  use ions_oct_m
  use io_function_oct_m
  use isdf_options_oct_m,   only: isdf_options_t
  use isdf_utils_oct_m,     only: output_matrix, isdf_potential
  use kmeans_clustering_oct_m
  use kpoints_oct_m
  use lalg_basic_oct_m
  use lalg_adv_oct_m
  use math_oct_m
  use mesh_oct_m
  use mesh_batch_oct_m
  use messages_oct_m
  use mpi_oct_m,            only: mpi_world
  use namespace_oct_m
  use parser_oct_m,         only: parse_variable
  use profiling_oct_m
  use poisson_oct_m
  use states_abst_oct_m,    only: states_are_real
  use states_elec_oct_m
  use states_elec_parallel_oct_m
  use space_oct_m
  use unit_system_oct_m,    only: unit_one, unit_angstrom
  use wfs_elec_oct_m
  use xc_cam_oct_m,         only: xc_cam_t
 
  implicit none
  private
 
  public :: &
    isdf_serial_ace_compute_potentials,  &
    isdf_serial_interpolation_vectors,   &
    quantify_error_and_visualise
 
  ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised systems
  integer, parameter   :: ik = 1
 
contains
 
  subroutine isdf_serial_interpolation_vectors(isdf, namespace, mesh, st, indices, phi_mu, P_r_mu, isdf_vectors)
    type(isdf_options_t),       intent(in ) :: isdf
    type(namespace_t),          intent(in ) :: namespace
    class(mesh_t),              intent(in ) :: mesh
    type(states_elec_t),        intent(in ) :: st
    integer(int64), contiguous, intent(in ) :: indices(:)
 
    real(real64), allocatable,  intent(out) :: phi_mu(:, :)
    ! defined at interpolation points: \f$ \varphi_i(\mathbf{r}_\mu) \f$
    real(real64), allocatable,  intent(out) :: P_r_mu(:, :)
    ! \f$P_{\mathbf{r},\mu}\f$, with size (np, n_int)
    real(real64), allocatable,  intent(out) :: isdf_vectors(:, :)
 
    real(real64), allocatable :: phi(:, :), cct(:, :)
    integer                   :: n_states, n_int, rank
    logical                   :: data_is_packed
 
    push_sub_with_profile(isdf_serial_interpolation_vectors)
    call messages_write(1)
 
    ! Reference serial implementation not parallel in states or domain
    if (st%parallel_in_states .or. mesh%parallel_in_domains) then
      message(1) = "Serial ISDF called when running state or domain-parallel"
      call messages_fatal(1)
    endif
 
    ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised systems
    if (st%d%nspin > 1) then
      call messages_not_implemented("ISDF Serial for SPIN_POLARIZED and SPINOR calculations", namespace)
    endif
 
    ! TODO(Alex) Issue #1196 Template ISDF handle both real and complex states
    if (.not. states_are_real(st)) then
      call messages_not_implemented("ISDF Serial handling of complex states", namespace)
    endif
 
    ! TODO(Alex) Implement algorithms for unpacked data structure
    data_is_packed = st%group%psib(st%group%block_start, 1)%status() == batch_packed
 
    if (.not. data_is_packed) then
      message(1) = "Serial ISDF only implemented for BATCH_PACKED"
      call messages_fatal(1)
    endif
 
    n_states = isdf%n_ks_states
    n_int = size(indices)
 
    call collate_batches_get_state(mesh, st, n_states, phi)
    if (debug%info) call output_matrix(namespace, "phi_r_serial.txt", phi)
 
    safe_allocate(phi_mu(1:n_states, 1:n_int))
    call sample_phi_at_centroids(phi, indices, phi_mu)
    if (debug%info) call output_matrix(namespace, "phi_mu_serial.txt", phi_mu)
 
    safe_allocate(p_r_mu(1:mesh%np, 1:n_int))
    call construct_density_matrix_packed(phi, phi_mu, p_r_mu)
    if (debug%info) call output_matrix(namespace, "p_r_mu_serial.txt", p_r_mu)
 
    ! Mutate P_r_mu to [ZC^T] = P_r_mu o P_r_mu
    call construct_zct(p_r_mu)
    if (debug%info) call output_matrix(namespace, "zct_serial.txt", p_r_mu)
 
    ! [CC^T] = ZC^T[indices, :]
    safe_allocate(cct(1:n_int, 1:n_int))
    call construct_cct(indices, p_r_mu, cct)
    if (debug%info) call output_matrix(namespace, "cct_serial.txt", cct)
    assert(is_symmetric(cct))
 
    ! Note, in principle one could use lalg_pseudo_inverse and just add an optional
    ! arg for returning the rank, but annoyingly the criterion in that is > rather than >=.
    ! Quantify the optimal number of interpolation points to use.
    ! If the rank of CC^T is <= ISDFNpoints, this does not give any indication how many
    ! additional points are required. It is only indicative when oversampling.
    if (isdf%check_n_interp) then
      rank = lalg_matrix_rank_svd(cct, preserve_mat=.true.)
      write(message(1),'(a, I4)') "ISDF Serial: Rank of CC^T is ", rank
      if (rank <  n_int) then
        write(message(2),'(a)') " - This rank is the optimal ISDFNpoints to run the calculation with"
      else
        write(message(2),'(a)') " - This suggests that ISDFNpoints is either optimal, or could be larger"
      endif
      call messages_info(2, namespace=namespace)
    endif
 
    ! [CC^T]^{-1}, mutating cct in-place
    ! NOTE, if the number of interpolation points exceeds the rank of CC^T, CC^T is by definition
    ! ill-conditioned, and requires inverting with the pseudo-inverse (SVD).
    ! Tests show that this is a much better solution than regularisation of either the diagonals
    ! or off-diagonals of CC^T. If one limits the number of interpolation points, there is no problem
    ! with inversion, but this limits the total accuracy achievable with the method.
 
    ! As CC^T and its inverse should be symmetric, ultimately want to:
    ! * Only compute the inverse of the upper (or lower) triangle
    ! * Modify the GEMM operation below to only use the upper (or lower) triangle of [CC^T]^{-1}
    write(message(1),'(a)') "ISDF Serial: Inverting [CC^T]"
    call messages_info(1, namespace=namespace, debug_only=.true.)
 
    ! Invert [CC^T] and symmetrise
    call lalg_svd_inverse(n_int, n_int, cct)
    call symmetrize_matrix(n_int, cct)
 
    ! Compute interpolation vectors, [ZC^T] [CC^T]^{-1}
    safe_allocate(isdf_vectors(1:mesh%np, 1:n_int))
    ! CC^T is by definition symmetric, implying [CC^T]^{-1} also is
    call lalg_gemm(mesh%np, n_int, n_int, 1.0_real64, p_r_mu, cct, 0.0_real64, isdf_vectors)
    if (debug%info) call output_matrix(namespace, "isdf_serial.txt", isdf_vectors)
    safe_deallocate_a(cct)
 
    ! Rebuild P_r_mu, with occupation numbers absorbed into it. Used in construction of W_ace
    call construct_density_matrix_with_occ_packed(st, phi, phi_mu, p_r_mu)
    safe_deallocate_a(phi)
    if (debug%info) call output_matrix(namespace, "OccP_r_mu_serial.txt", p_r_mu)
 
    pop_sub_with_profile(isdf_serial_interpolation_vectors)
 
  end subroutine isdf_serial_interpolation_vectors
 
 
  subroutine collate_batches_get_state(mesh, st, max_state, psi)
    class(mesh_t),              intent(in ) :: mesh
    type(states_elec_t),        intent(in ) :: st
    integer,                    intent(in ) :: max_state
    real(real64),  allocatable, intent(out) :: psi(:, :)
 
    integer :: istate, ib, ist, minst, maxst, block_end
 
    push_sub_with_profile(collate_batches_get_state)
 
    assert(max_state > 0 .and. max_state <= st%nst)
 
    safe_allocate(psi(1:max_state, 1:mesh%np))
    block_end = st%group%iblock(max_state)
 
    istate = 0
    do ib = 1, block_end
      ! Normalisation did not affect condition number of CC^T matrix
      !call dmesh_batch_normalize(mesh, st%group%psib(ib, ik))
      minst = states_elec_block_min(st, ib)
      maxst = min(states_elec_block_max(st, ib), max_state)
      do ist = minst, maxst
        istate = istate + 1
        call states_elec_get_state(st, mesh, st%d%dim, ist, ik, psi(istate, :))
      enddo
    enddo
 
    pop_sub_with_profile(collate_batches_get_state)
 
  end subroutine collate_batches_get_state
 
 
  subroutine sample_phi_at_centroids(phi_r, indices, phi_mu)
    real(real64),   contiguous, intent(in ) :: phi_r(:, :)
    integer(int64), contiguous, intent(in ) :: indices(:)
    real(real64),   contiguous, intent(out) :: phi_mu(:, :)
 
    integer        :: ic, is, nst, n_int
    integer(int64) :: ipg
 
    push_sub_with_profile(sample_phi_at_centroids)
 
    write(message(1),'(a)') "ISDF Serial: Sampling phi(r) at mu"
    call messages_info(1, debug_only=.true.)
 
    nst = size(phi_r, 1)
    assert(size(phi_mu, 1) == nst)
 
    n_int = size(indices)
    assert(size(phi_mu, 2) == n_int)
 
    do ic = 1, n_int
      ipg = indices(ic)
      do is = 1, nst
        phi_mu(is, ic) = phi_r(is, ipg)
      enddo
    enddo
 
    pop_sub_with_profile(sample_phi_at_centroids)
 
  end subroutine sample_phi_at_centroids
 
 
  subroutine construct_zct(zct)
    real(real64), contiguous, intent(inout)  :: zct(:, :)
    ! Out: Contraction of Z and C^T matrices == element-wise square of quasi-density matrix
 
    integer :: i, j
 
    push_sub_with_profile(construct_zct)
 
    write(message(1),'(a)') "ISDF Serial: Constructing ZC^T"
    call messages_info(1, debug_only=.true.)
 
    !$omp parallel do collapse(2)
    do j = 1, size(zct, 2)
      do i = 1, size(zct, 1)
        zct(i, j) = zct(i, j)**2
      end do
    enddo
    !$omp end parallel do
 
    pop_sub_with_profile(construct_zct)
 
  end subroutine construct_zct
 
 
  subroutine construct_cct(indices, zct, cct)
    integer(int64), contiguous, intent(in ) :: indices(:)
    real(real64),   contiguous, intent(in ) :: zct(:, :)
 
    real(real64),   contiguous, intent(out) :: cct(:, :)
 
    integer(int64) :: ipg
    integer        :: i_mu, i_nu, n_int
 
    push_sub_with_profile(construct_cct)
 
    write(message(1),'(a)') "ISDF Serial: Constructing CC^T by sampling ZC^T"
    call messages_info(1, debug_only=.true.)
 
    n_int = size(indices)
    assert(all(shape(cct) == [n_int, n_int]))
    assert(size(zct, 1) > n_int)
    assert(size(zct, 2) == n_int)
 
    ! Mask ZC^T to obtain CC^T
    do i_nu = 1, n_int
      do i_mu = 1, n_int
        ipg = indices(i_mu)
        cct(i_mu, i_nu) = zct(ipg, i_nu)
      enddo
    enddo
 
    pop_sub_with_profile(construct_cct)
 
  end subroutine construct_cct
 
 
  subroutine construct_density_matrix_packed(phi, phi_mu, P_r_mu)
    real(real64), contiguous, intent(in ) :: phi(:, :)
    !                                                         of shape (m_states, np)
    real(real64), contiguous, intent(in ) :: phi_mu(:, :)
 
    real(real64), contiguous, intent(out) :: P_r_mu(:, :)
 
    integer :: np
    integer :: n_int
    integer :: m_states
 
    push_sub_with_profile(construct_density_matrix_packed)
 
    write(message(1),'(a)') "ISDF Serial: Constructing P_r_mu"
    call messages_info(1, debug_only=.true.)
 
    m_states = size(phi, 1)
    np = size(phi, 2)
    n_int = size(phi_mu, 2)
 
    assert(size(phi_mu, 1) == m_states)
    assert(size(p_r_mu, 1) == np)
    assert(size(p_r_mu, 2) == n_int)
 
    ! Contract over the state index, P = phi^T @ phi_mu, with shape (np, m_states) (m_states, n_int)
    call lalg_gemm(phi, phi_mu, p_r_mu, transa='T')
 
    pop_sub_with_profile(construct_density_matrix_packed)
 
  end subroutine construct_density_matrix_packed
 
 
  subroutine construct_density_matrix_with_occ_packed(st, phi, phi_mu, P_r_mu)
    type(states_elec_t), intent(in ) :: st
    real(real64),        intent(in ) :: phi(:, :)
    !                                                  of shape (m_states, np)
    real(real64),        intent(in ) :: phi_mu(:, :)
 
    real(real64),        intent(out) :: P_r_mu(:, :)
 
    integer                          :: np, n_int, m_states, imu, ist
    real(real64), allocatable        :: focc_phi_mu(:, :)
 
    push_sub_with_profile(construct_p_with_occ_packed)
 
    write(message(1),'(a)') "ISDF Serial: Constructing P_r_mu with occupations"
    call messages_info(1, debug_only=.true.)
 
    m_states = size(phi_mu, 1)
    np = size(phi, 2)
    n_int = size(phi_mu, 2)
 
    assert(size(phi, 1) == m_states)
    assert(size(p_r_mu, 1) == np)
    assert(size(p_r_mu, 2) == n_int)
 
    ! Element-wise multiply the occupations with the smaller of the two arrays
    safe_allocate(focc_phi_mu(1:m_states, 1:n_int))
    do imu = 1, n_int
      do ist = 1, m_states
        focc_phi_mu(ist, imu) = st%kweights(ik) * st%occ(ist, ik) * phi_mu(ist, imu)
      enddo
    enddo
 
    ! Contract over the state index, P = phi^T @ focc_phi_mu, with shape (np, m_states) (m_states, n_int)
    call lalg_gemm(phi, focc_phi_mu, p_r_mu, transa='T')
    safe_deallocate_a(focc_phi_mu)
 
    pop_sub_with_profile(construct_p_with_occ_packed)
 
  end subroutine construct_density_matrix_with_occ_packed
 
 
  subroutine quantify_error_and_visualise(isdf, namespace, st, space, mesh, ions, indices, isdf_vectors, output_cubes)
    type(isdf_options_t),          intent(in)    :: isdf
    type(namespace_t),             intent(in)    :: namespace
    type(states_elec_t),           intent(in)    :: st
    class(space_t),                intent(in)    :: space
    class(mesh_t),                 intent(in)    :: mesh
    class(ions_t),    pointer,     intent(in)    :: ions
    integer(int64),   contiguous,  intent(in)    :: indices(:)
    real(real64),     allocatable, intent(inout) :: isdf_vectors(:, :)
    logical,                       intent(in)    :: output_cubes
 
    real(real64), allocatable :: product_basis(:, :), approx_product_basis(:, :)
    real(real64), allocatable :: phi(:, :), phi_mu(:, :)
    real(real64), allocatable :: product_error(:)
    integer                   :: n_occ, n_products, n_int, i, j, ij, unit
    real(real64)              :: mean_error
 
    push_sub_with_profile(quantify_error_and_visualise)
 
    write(message(1),'(a)') "ISDF Serial: Computing exact pair products"
    call messages_info(1, debug_only=.true.)
 
    ! Rebuild phi matrix
    n_occ = isdf%n_ks_states
    call collate_batches_get_state(mesh, st, n_occ, phi)
    assert(size(phi, 2) == mesh%np)
 
    n_products = n_occ * n_occ
    safe_allocate(product_basis(1:n_products, 1:mesh%np))
    call column_wise_khatri_rao_product(phi, phi, product_basis)
 
    if (output_cubes) then
      call generate_product_state_cubes(namespace, space, mesh, ions, "exact_product_", &
        product_basis)
    endif
 
    write(message(1),'(a)') "ISDF Serial Test: Computing approximate pair products"
    call messages_info(1, namespace=namespace, debug_only=.true.)
 
    ! Rebuild phi_mu, again only for occupied states
    n_int = size(indices)
    safe_allocate(phi_mu(1:n_occ, 1:n_int))
    call sample_phi_at_centroids(phi, indices, phi_mu)
    safe_deallocate_a(phi)
 
    safe_allocate(approx_product_basis(1:n_products, 1:mesh%np))
    call approximate_pair_products(phi_mu, isdf_vectors, approx_product_basis)
    ! if (debug%info) call output_matrix(namespace, "approx_product_blas.txt", approx_product_basis)
 
    safe_deallocate_a(phi_mu)
    safe_deallocate_a(isdf_vectors)
 
    if (output_cubes) then
      call generate_product_state_cubes(namespace, space, mesh, ions, "approx_product_", &
        approx_product_basis)
    endif
 
    ! Quantify the error
    safe_allocate(product_error(1:n_products))
    call error_in_product_basis(mesh, product_basis, approx_product_basis, product_error, mean_error)
    safe_deallocate_a(product_basis)
    safe_deallocate_a(approx_product_basis)
 
    if (mpi_world%is_root()) then
      open(newunit=unit, file="isdf_error_serial.txt")
      write(unit, *) 'Mean error', mean_error
      ij = 0
      do i = 1, n_occ
        do j = 1, n_occ
          ij = ij + 1
          write(unit, *) i, j, product_error(ij)
        enddo
      enddo
      close(unit)
    endif
 
    safe_deallocate_a(product_error)
 
    pop_sub_with_profile(quantify_error_and_visualise)
 
  end subroutine quantify_error_and_visualise
 
 
  subroutine approximate_pair_products(psi_mu, zeta, product_basis)
    real(real64), contiguous,  intent(in ) :: psi_mu(:, :)
    real(real64), contiguous,  intent(in ) :: zeta(:, :)
    real(real64), contiguous,  intent(out) :: product_basis(:, :)
 
    real(real64), allocatable              :: psi_ij_mu(:, :)
    integer                                :: mn_states, n_int, np
 
    push_sub_with_profile(approximate_pair_products)
 
    mn_states = size(psi_mu, 1)**2
    np = size(zeta, 1)
    n_int = size(zeta, 2)
 
    assert(size(product_basis, 1) == mn_states)
    assert(size(product_basis, 2) == np)
 
    safe_allocate(psi_ij_mu(1:mn_states, 1:n_int))
    call column_wise_khatri_rao_product(psi_mu, psi_mu, psi_ij_mu)
 
    ! Contract product_basis = [psi_ij_mu] [zeta]^T over interpolation vector index
    call lalg_gemm(psi_ij_mu, zeta, product_basis, transb='T')
 
    safe_deallocate_a(psi_ij_mu)
 
    pop_sub_with_profile(approximate_pair_products)
 
  end subroutine approximate_pair_products
 
 
  subroutine error_in_product_basis(mesh, product_basis, approx_product_basis, error, mean_error)
    class(mesh_t),             intent(in ) :: mesh
    real(real64),  contiguous, intent(in ) :: product_basis(:, :)
    real(real64),  contiguous, intent(in ) :: approx_product_basis(:, :)
 
    real(real64),  contiguous, intent(out) :: error(:)
    real(real64),              intent(out) :: mean_error
 
    integer      :: mn_states, np, ij, ip
 
    push_sub_with_profile(error_in_product_basis)
 
    mn_states = size(product_basis, 1)
    np = size(product_basis, 2)
 
    ! product_basis shape is not as expected
    assert(mesh%np == np)
 
    ! Two arrays should be the same dimensions
    assert(all(shape(product_basis) == shape(approx_product_basis)))
 
    ! error should be allocated, and with the correct size
    assert(size(error) == mn_states)
 
    ! Initialise error with first point from the grid
    do ij = 1, mn_states
      error(ij) = (product_basis(ij, 1) - approx_product_basis(ij, 1))**2
    enddo
 
    do ip = 2, np
      do ij = 1, mn_states
        error(ij) = error(ij) + (product_basis(ij, ip) - approx_product_basis(ij, ip))**2
      enddo
    enddo
 
    mean_error = 0.0_real64
    do ij = 1, mn_states
      error(ij) = sqrt(mesh%volume_element * error(ij))
      mean_error = mean_error + error(ij)
    enddo
 
    mean_error = mean_error / real(mn_states, real64)
 
    pop_sub_with_profile(error_in_product_basis)
 
  end subroutine error_in_product_basis
 
 
  subroutine generate_product_state_cubes(namespace, space, mesh, ions, file_prefix, data, limits)
    type(namespace_t),        intent(in) :: namespace
    class(space_t),           intent(in) :: space
    class(mesh_t),            intent(in) :: mesh
    class(ions_t), pointer,   intent(in) :: ions
    character(len=*),         intent(in) :: file_prefix
    real(real64), contiguous, intent(in) :: data(:, :)
    integer, optional,        intent(in) :: limits(2)
 
    integer            :: m_states, limit_j, limit_i, i, j, ij, ierr
    real(real64)       :: size_data
    character(len=4)   :: i_char, j_char
    character(len=120) :: file_name
 
    ! product basis size is currently defined as (m_states * m_states)
    size_data = real(size(data, 1), real64)
    m_states = int(sqrt(size_data))
 
    if (present(limits)) then
      limit_j = limits(1)
      limit_i = limits(2)
    else
      limit_j = m_states
      limit_i = m_states
    endif
 
    do i = 1, limit_i
      do j = 1, limit_j
        ij = j + (i - 1) * m_states
        write(i_char, '(I4)') i
        write(j_char, '(I4)') j
        file_name = trim(adjustl(file_prefix)) // trim(adjustl(i_char)) // '_' // trim(adjustl(j_char))
        call dio_function_output(option__outputformat__cube, "./cubes", trim(adjustl(file_name)), namespace, space, mesh, &
          data(ij,:) , unit_one, ierr, pos=ions%pos, atoms=ions%atom)
      enddo
    enddo
 
  end subroutine generate_product_state_cubes
 
 
  subroutine isdf_serial_ace_compute_potentials(exxop, namespace, space, mesh, st, Vx_on_st, kpoints)
    type(exchange_operator_t), intent(in   ) :: exxop
    !                                                           ISDF interpolation points, and cam parameters.
    ! An ISDF instance is not passed directly so this API is consistent with the other "compute_potential" routines.
    type(namespace_t),         intent(in   ) :: namespace
    class(space_t),            intent(in   ) :: space
    ! with the existing routines
    class(mesh_t),             intent(in   ) :: mesh
    type(states_elec_t),       intent(in   ) :: st
    type(kpoints_t),           intent(in   ) :: kpoints
    ! with the existing routines
 
    type(states_elec_t),       intent(out) :: Vx_on_st
 
    real(real64),  allocatable :: psi_mu(:, :), P_r_mu(:, :), W_ace(:, :), isdf_vectors(:, :)
    integer(int64), allocatable :: indices(:)
 
    push_sub(isdf_serial_ace_compute_potentials)
 
    ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised and periodic systems
    assert(kpoints%gamma_only())
    assert(.not. space%is_periodic())
    assert(st%d%nspin == 1)
 
    indices = exxop%isdf%centroids%global_mesh_indices()
 
    call isdf_serial_interpolation_vectors(exxop%isdf, namespace, mesh, st, &
      indices, psi_mu, p_r_mu, isdf_vectors)
 
    safe_allocate(w_ace(1:mesh%np, exxop%isdf%n_ks_states))
    call isdf_serial_ace_w_unpacked(namespace, p_r_mu, isdf_vectors, psi_mu, exxop%psolver, exxop%cam, st, w_ace)
    safe_deallocate_a(psi_mu)
    safe_deallocate_a(p_r_mu)
    safe_deallocate_a(isdf_vectors)
 
    call isdf_serial_ace_batch_w(exxop%isdf, st, w_ace, vx_on_st)
    safe_deallocate_a(w_ace)
 
    pop_sub(isdf_serial_ace_compute_potentials)
 
  end subroutine isdf_serial_ace_compute_potentials
 
 
  subroutine isdf_serial_ace_w_unpacked(namespace, P_r_mu, isdf_vectors, psi_mu, poisson_solver, cam, st, W_ace)
    type(namespace_t), intent(in )             :: namespace
    real(real64),      intent(in ), contiguous :: P_r_mu(:, :)
    real(real64),      intent(in ), contiguous :: isdf_vectors(:, :)
    real(real64),      intent(in ), contiguous :: psi_mu(:, :)
    type(poisson_t),   intent(in )             :: poisson_solver
    type(xc_cam_t),    intent(in)              :: cam
    type(states_elec_t), intent(in   )         :: st
 
    real(real64),      intent(out), contiguous :: W_ace(:, :)
 
    integer                   :: ip, i_mu, ist, np, n_int, nst
    real(real64)              :: psi_ist_mu
    real(real64), allocatable :: V_r_nu(:, :)
    logical                   :: use_external_kernel
    real(real64)              :: exx_weight
    real(real64)              :: weight
 
    push_sub_with_profile(isdf_serial_ace_w_unpacked)
 
    ! Number of states defines the number used in ISDF, which is typically ~ N occupied states
    nst = size(psi_mu, 1)
    np = size(p_r_mu, 1)
    n_int = size(p_r_mu, 2)
 
    assert(all(shape(p_r_mu) == shape(isdf_vectors)))
    assert(size(psi_mu, 2) == n_int)
    assert(size(w_ace, 1) == np)
    ! Implies a size issue with either W_ace or psi_mu
    assert(size(w_ace, 2) == nst)
 
    use_external_kernel = (st%nik > st%d%spin_channels .or. cam%omega > m_epsilon)
    if (use_external_kernel) then
      message(1) = "External kernel not supported in ISDF"
      call messages_fatal(1)
    endif
    exx_weight = cam%alpha
    weight = exx_weight / st%smear%el_per_state
 
    safe_allocate(v_r_nu(1:np, 1:n_int))
    call isdf_potential(namespace, poisson_solver, isdf_vectors, v_r_nu)
 
    write(message(1),'(a)') "ISDF: Writing V from isdf_ace_w_unpacked"
    call messages_info(1, namespace=namespace, debug_only=.true.)
 
    ! Initialise elements of W_ace with data from the first interpolation point
    i_mu = 1
    do ist = 1, nst
      psi_ist_mu = weight * psi_mu(ist, i_mu)
      do ip = 1, np
        w_ace(ip, ist) = - (p_r_mu(ip, i_mu) * v_r_nu(ip, i_mu) * psi_ist_mu)
      enddo
    enddo
 
    ! Construct W_ace
    do i_mu = 2, n_int
      do ist = 1, nst
        psi_ist_mu = weight * psi_mu(ist, i_mu)
        do ip = 1, np
          w_ace(ip, ist) = w_ace(ip, ist) - (p_r_mu(ip, i_mu) * v_r_nu(ip, i_mu) * psi_ist_mu)
        enddo
      enddo
    enddo
 
    safe_deallocate_a(v_r_nu)
 
    pop_sub_with_profile(isdf_serial_ace_w_unpacked)
 
  end subroutine isdf_serial_ace_w_unpacked
 
 
  subroutine isdf_serial_ace_batch_w(isdf, st, W_ace, Vx_on_st)
    type(isdf_options_t), intent(in   )              :: isdf
    type(states_elec_t),  intent(in   )              :: st
    real(real64),         intent(in   ), contiguous  :: w_ace(:, :)
 
    type(states_elec_t),  intent(out)              :: vx_on_st
 
    integer :: ist, ib, ist_b, np, max_state, minst, maxst, block_end, block_size
 
    push_sub_with_profile(isdf_serial_ace_batch_w)
 
    assert(size(w_ace, 2) == isdf%n_ks_states)
    assert(st%d%dim == 1)
 
    ! Copy memory layout, without data
    call states_elec_copy(vx_on_st, st)
    call states_elec_set_zero(vx_on_st)
    np = size(w_ace, 1)
 
    ! Ensure we do not go beyond the total number of occupied states
    max_state = min(isdf%n_ks_states, st%st_end)
    block_end = min(st%group%block_end, st%group%iblock(max_state))
 
    do ib = st%group%block_start, block_end
      minst = states_elec_block_min(st, ib)
      maxst = min(states_elec_block_max(st, ib), max_state)
      block_size = maxst - minst + 1
      ! States in a block
      do ist_b = 1, block_size
        ! Global state index
        ist = minst - 1 + ist_b
        call batch_set_state(vx_on_st%group%psib(ib, 1), ist_b, np, w_ace(:, ist))
      enddo
    enddo
 
    pop_sub_with_profile(isdf_serial_ace_batch_w)
 
  end subroutine isdf_serial_ace_batch_w
 
end module isdf_serial_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: