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
  use batch_oct_m
  use batch_ops_oct_m
  use blas_oct_m
  use debug_oct_m
  use electrons_oct_m,       only: electrons_t
  use electron_space_oct_m
  use face_splitting_oct_m
  use global_oct_m
  use grid_oct_m
  use ions_oct_m
  use io_function_oct_m
  use isdf_utils_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 smear_oct_m,          only: smear_semiconductor
  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
 
  implicit none
  private
 
  ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised systems
  integer,      parameter   :: ik = 1
 
  type isdf_options_t
    integer :: n_interp
  contains
    procedure :: init => isdf_options_init
  end type isdf_options_t
 
  public ::                                                 &
    isdf_options_t,                                         &
    serial_interpolative_separable_density_fitting_vectors, &
    quantify_error_and_visualise
 
contains
 
  subroutine isdf_options_init(this, namespace, st)
    class(isdf_options_t),     intent(out) :: this
    type(namespace_t),         intent(in ) :: namespace
    type(states_elec_t),       intent(in ) :: st
 
    integer :: default_n_interp
 
    push_sub(isdf_options_init)
 
    default_n_interp = 10 * highest_occupied_index(st, ik)
 
    !%Variable NCentroidPoints
    !%Type integer
    !%Default 10*Nocc
    !%Section ISDF
    !%Description
    !% The total number of interpolation points used in ISDF density fitting.
    !% This number should be considerably smaller than the number of grid points.
    !% By default, Octopus sets the number of interpolation points to 10x the number
    !% of occupied states.
    !%End
    call parse_variable(namespace, 'NCentroidPoints', default_n_interp, this%n_interp)
 
    pop_sub(isdf_options_init)
 
  end subroutine isdf_options_init
 
 
  subroutine serial_interpolative_separable_density_fitting_vectors(namespace, mesh, st, int_indices, phi, isdf_vectors)
    type(namespace_t),          intent(in ) :: namespace
    class(mesh_t),              intent(in ) :: mesh
    type(states_elec_t),        intent(in ) :: st
    integer(int64), contiguous, intent(in ) :: int_indices(:)
    real(real64),  allocatable, intent(out) :: phi(:, :)
    !                                                             of shape either:
    !                                                             Packed (nst, npg)
    !                                                             Unpacked (npg, nst)
    real(real64),  allocatable, intent(out) :: isdf_vectors(:, :)
 
    integer :: nocc
 
    push_sub_with_profile(serial_interpolative_separable_density_fitting_vectors)
 
    message(1) = "Serial ISDF"
    call messages_write(1)
 
    ! Reference 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
 
    nocc = highest_occupied_index(st, ik)
 
    ! Alternate means of constructing phi from a batch - retained for reference.
    ! Dress in a fake code path so the compiler does not warn that the routine is unused.
    if (.false.) then
      call collate_batches(mesh%np, st, nocc, phi)
      if (debug%info) call output_matrix(namespace, "phi_r_serial.txt", phi)
      safe_deallocate_a(phi)
    endif
 
    ! This gives identical results to the above
    call collate_batches_get_state(mesh, st, nocc, phi)
    if (debug%info) call output_matrix(namespace, "phi_r_serial3.txt", phi)
 
    select case (st%group%psib(st%group%block_start, 1)%status())
    case (batch_device_packed)
      message(1) = "Serial ISDF not implemented for BATCH_DEVICE_PACKED"
      call messages_fatal(1)
 
    case (batch_packed)
      call isdf_construct_interpolation_vectors_packed(namespace, mesh, phi, int_indices, isdf_vectors)
 
    case (batch_not_packed)
      message(1) = "Serial ISDF not implemented for BATCH_NOT_PACKED"
      call messages_fatal(1)
      assert(all(shape(phi) == [mesh%np, st%nst]))
 
    end select
 
    pop_sub_with_profile(serial_interpolative_separable_density_fitting_vectors)
 
  end subroutine serial_interpolative_separable_density_fitting_vectors
 
 
  subroutine isdf_construct_interpolation_vectors_packed(namespace, mesh, phi, indices, isdf_vectors)
    type(namespace_t),           intent(in ) :: namespace
    class(mesh_t),               intent(in ) :: mesh
    real(real64),   contiguous,  intent(in ) :: phi(:, :)
    integer(int64), contiguous,  intent(in ) :: indices(:)
    real(real64),   allocatable, intent(out) :: isdf_vectors(:, :)
 
    real(real64),  allocatable :: phi_mu(:, :)
    real(real64),  allocatable :: P_mu_nu(:, :)
    real(real64),  allocatable :: zct(:, :)
    real(real64),  allocatable :: cct(:, :)
 
    integer            :: n_int, i, j, n_states
    logical, parameter :: construct_P_mu_nu = .false.              
 
    push_sub_with_profile(isdf_construct_interpolation_vectors_packed)
 
    assert(size(phi, 2) == mesh%np)
 
    n_states = size(phi, 1)
    n_int = size(indices)
 
    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)
 
    ! Note, this actually constructs P_r_mu, but we want to mutate the memory in the subsequent call
    safe_allocate(zct(1:mesh%np, 1:n_int))
    call construct_density_matrix_packed(phi, phi_mu, zct)
    if (debug%info) call output_matrix(namespace, "p_r_mu_serial.txt", zct)
 
    ! [ZC^T] = P_r_mu o P_r_mu
    call construct_zct(zct)
    if (debug%info) call output_matrix(namespace, "zct_serial.txt", zct)
 
    ! [CC^T] = ZC^T[indices, :]
    safe_allocate(cct(1:n_int, 1:n_int))
    call construct_cct(indices, zct, cct)
    if (debug%info) call output_matrix(namespace, "cct_serial.txt", cct)
    assert(is_symmetric(cct))
 
    ! For testing, alternately, construct [CC^T] = P_mu_nu o P_mu_nu
    if (construct_p_mu_nu) then
      safe_allocate(p_mu_nu(1:n_int, 1:n_int))
      call construct_density_matrix_all_centroids_packed(phi_mu, p_mu_nu)
      assert(is_symmetric(p_mu_nu))
 
      safe_deallocate_a(phi_mu)
      if (debug%info) call output_matrix(namespace, "p_mu_nu_serial.txt", p_mu_nu)
 
      ! Mutate P_mu_nu -> cct
      write(message(1),'(a)') "ISDF Serial: Constructing [CC^T] = P_mu_nu o P_mu_nu"
      call messages_info(1, namespace=namespace, debug_only=.true.)
 
      do j = 1, n_int
        do i = 1, n_int
          p_mu_nu(i, j) = p_mu_nu(i, j) * p_mu_nu(i, j)
        enddo
      enddo
 
      if (debug%info) call output_matrix(namespace, "cct_alt_serial.txt", p_mu_nu)
      safe_deallocate_a(p_mu_nu)
    endif
 
    ! [CC^T]^{-1}, mutating cct in-place
    ! NOTE, CC^T is extremely ill-conditioned, and it is not clear why this is the case.
    ! * Using the pseudo-inverse (SVD) circumvents this problem, without addressing the root cause
    ! 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, zct, cct, 0.0_real64, isdf_vectors)
 
    if (debug%info) call output_matrix(namespace, "isdf_serial.txt", isdf_vectors)
    safe_deallocate_a(zct)
    safe_deallocate_a(cct)
 
    pop_sub_with_profile(isdf_construct_interpolation_vectors_packed)
 
  end subroutine isdf_construct_interpolation_vectors_packed
 
 
  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(collate_batches_get_state)
 
    assert(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
        ! In case there are any states < max_state that are not occupied
        if (abs(st%occ(ist, ik) * st%kweights(ik)) < m_min_occ) then
          psi(istate, :) = 0.0_real64
        else
          call states_elec_get_state(st, mesh, st%d%dim, ist, ik, psi(istate, :))
        endif
      enddo
    enddo
 
    pop_sub(collate_batches_get_state)
 
  end subroutine collate_batches_get_state
 
 
  subroutine collate_batches(np, st, max_state, psi)
    integer,                    intent(in ) :: np
    type(states_elec_t),        intent(in ) :: st
    integer,                    intent(in) ::  max_state
    real(real64),  allocatable, intent(out) :: psi(:, :)
 
    integer                   :: ip, ib, minst, maxst, block_end, ist, ist_local
 
    push_sub(collate_batches)
 
    select case (st%group%psib(1,1)%status())
    case (batch_device_packed)
      ! No GPU implementation
      assert(.false.)
 
    case (batch_packed)
 
      ! Not state distributed, so expect ist to run [1: max_state], contiguously
      safe_allocate(psi(1:max_state, 1:np))
      block_end = st%group%iblock(max_state)
 
      do ip = 1, np
        do ib = 1, block_end
          minst = states_elec_block_min(st, ib)
          ! Truncate maxst to max_state in last block
          maxst = min(states_elec_block_max(st, ib), max_state)
          ! Assign states from current batch
          do ist = minst, maxst
            ist_local = ist - minst + 1
            ! In case there are any states < max_state that are not occupied
            if (abs(st%occ(ist, ik) * st%kweights(ik)) < m_min_occ) then
              psi(ist, ip) = 0.0_real64
            else
              psi(ist, ip) = st%group%psib(ib, ik)%dff_pack(ist_local, ip)
            endif
          enddo
        enddo
      enddo
 
    case (batch_not_packed)
      assert(.false.)
 
    end select
 
    pop_sub(collate_batches)
 
  end subroutine collate_batches
 
 
  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(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(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_all_centroids_packed(phi_mu, P_mu_nu)
    real(real64), contiguous, intent(in ) :: phi_mu(:, :)
    !                                                         of shape (m_states, n_int)
    real(real64), contiguous, intent(out) :: P_mu_nu(:, :)
 
    integer :: n_int
 
    push_sub_with_profile(construct_density_matrix_all_centroids_packed)
 
    write(message(1),'(a)') "ISDF Serial: Constructing P_mu_nu"
    call messages_info(1, debug_only=.true.)
 
    n_int = size(phi_mu, 2)
    assert(size(p_mu_nu, 1) == n_int)
    assert(size(p_mu_nu, 2) == n_int)
 
    ! Contract over the state index
    call lalg_gemm(phi_mu, phi_mu, p_mu_nu, transa='T')
 
    pop_sub_with_profile(construct_density_matrix_all_centroids_packed)
 
  end subroutine construct_density_matrix_all_centroids_packed
 
 
  subroutine quantify_error_and_visualise(namespace, st, space, mesh, ions, phi, indices, isdf_vectors, output_cubes)
    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
    real(real64),     allocatable, intent(inout) :: phi(:, :)
    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_mu(:, :), phi_occ(:, :)
    real(real64), allocatable :: product_error(:)
    integer                   :: n_occ, n_products, n_int, i, j, ij, is, ip, unit
    real(real64)              :: mean_error
 
    push_sub(quantify_error_and_visualise)
 
    write(message(1),'(a)') "ISDF Serial: Computing exact pair products"
    call messages_info(1, debug_only=.true.)
 
    assert(size(phi, 2) == mesh%np)
 
    ! Limit product basis comparison to occupied states
    n_occ = highest_occupied_index(st, ik)
    safe_allocate(phi_occ(1:n_occ, 1:mesh%np))
    do ip = 1, mesh%np
      do is = 1, n_occ
        phi_occ(is, ip) = phi(is, ip)
      enddo
    enddo
    safe_deallocate_a(phi)
 
    n_products = n_occ * n_occ
    safe_allocate(product_basis(1:n_products, 1:mesh%np))
    call column_wise_khatri_rao_product(phi_occ, phi_occ, product_basis)
    ! if (debug%info) call output_matrix(namespace, "exact_product_blas.txt", 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_occ, indices, phi_mu)
    safe_deallocate_a(phi_occ)
 
    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(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(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(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(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(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
 
end module isdf_serial_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: