xas_tdp_integrals.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2026 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
! **************************************************************************************************
!> \brief 3-center integrals machinery for the XAS_TDP method
!> \author A. Bussy (03.2020)
! **************************************************************************************************
MODULE xas_tdp_integrals
   USE r12_fundamentals, ONLY: r12_fundamentals_init
   USE coulomb_integral_interface, ONLY: coulomb_integral_context_type, coulomb_integral_init, &
                                         coulomb_integral_cleanup, compute_coulomb_2c, compute_coulomb_3c, &
                                         compute_coulomb_2c_derivs, compute_coulomb_3c_derivs
   USE integral_library_types, ONLY: library_libint
   USE omp_lib, ONLY: omp_get_max_threads, &
                      omp_get_num_threads, &
                      omp_get_thread_num
   USE ai_contraction_sphi, ONLY: ab_contract, &
                                  abc_contract_xsmm
   USE atomic_kind_types, ONLY: atomic_kind_type
   USE basis_set_types, ONLY: get_gto_basis_set, &
                              gto_basis_set_p_type, &
                              gto_basis_set_type
   USE cell_types, ONLY: cell_type
   USE constants_operator, ONLY: operator_coulomb
   USE cp_array_utils, ONLY: cp_1d_i_p_type, &
                             cp_2d_r_p_type
   USE cp_blacs_env, ONLY: cp_blacs_env_type
   USE cp_dbcsr_api, ONLY: dbcsr_distribution_get, &
                           dbcsr_distribution_release, &
                           dbcsr_distribution_type
   USE cp_dbcsr_operations, ONLY: cp_dbcsr_dist2d_to_dist
   USE cp_eri_mme_interface, ONLY: cp_eri_mme_param, &
                                   cp_eri_mme_set_params
   USE cp_files, ONLY: close_file, &
                       open_file
   USE dbt_api, ONLY: &
      dbt_create, dbt_distribution_destroy, dbt_distribution_new, dbt_distribution_type, &
      dbt_finalize, dbt_pgrid_create, dbt_pgrid_destroy, dbt_pgrid_type, dbt_put_block, &
      dbt_reserve_blocks, dbt_type
   USE distribution_1d_types, ONLY: distribution_1d_type
   USE distribution_2d_types, ONLY: distribution_2d_create, &
                                    distribution_2d_type
   USE eri_mme_integrate, ONLY: eri_mme_2c_integrate
   USE eri_mme_types, ONLY: eri_mme_init, &
                            eri_mme_release
   USE generic_os_integrals, ONLY: int_operators_r12_ab_os
   USE input_constants, ONLY: do_potential_coulomb, &
                              do_potential_id, &
                              do_potential_short, &
                              do_potential_truncated
   USE input_section_types, ONLY: section_vals_val_get
   USE kinds, ONLY: dp
   USE integral_library_types, ONLY: cutoff_screen_factor, &
                                     coulomb_operator_type
   USE mathlib, ONLY: invmat_symm
   USE message_passing, ONLY: mp_comm_type, &
                              mp_para_env_type
   USE molecule_types, ONLY: molecule_type
   USE orbital_pointers, ONLY: ncoset
   USE particle_methods, ONLY: get_particle_set
   USE particle_types, ONLY: particle_type
   USE qs_environment_types, ONLY: get_qs_env, &
                                   qs_environment_type
   USE qs_integral_utils, ONLY: basis_set_list_setup
   USE qs_kind_types, ONLY: get_qs_kind, &
                            qs_kind_type
   USE qs_neighbor_list_types, ONLY: &
      get_iterator_info, get_neighbor_list_set_p, neighbor_list_iterate, &
      neighbor_list_iterator_create, neighbor_list_iterator_p_type, &
      neighbor_list_iterator_release, neighbor_list_set_p_type, nl_set_sub_iterator, &
      nl_sub_iterate, release_neighbor_list_sets
   USE qs_neighbor_lists, ONLY: atom2d_build, &
                                atom2d_cleanup, &
                                build_neighbor_lists, &
                                local_atoms_type, &
                                pair_radius_setup
   USE qs_o3c_types, ONLY: get_o3c_iterator_info, &
                           init_o3c_container, &
                           o3c_container_type, &
                           o3c_iterate, &
                           o3c_iterator_create, &
                           o3c_iterator_release, &
                           o3c_iterator_type, &
                           release_o3c_container
   USE xas_tdp_types, ONLY: xas_tdp_control_type, &
                            xas_tdp_env_type
#if defined(__LIBXS)
   USE iso_c_binding, ONLY: c_loc
   USE libxs, ONLY: libxs_matcopy, &
                    libxs_otrans
#endif
#include "./base/base_uses.f90"
 
   IMPLICIT NONE
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'xas_tdp_integrals'
 
   PUBLIC :: create_pqx_tensor, fill_pqx_tensor, compute_ri_3c_coulomb, compute_ri_3c_exchange, &
             build_xas_tdp_3c_nl, build_xas_tdp_ovlp_nl, get_opt_3c_dist2d, &
             compute_ri_coulomb2_int, compute_ri_exchange2_int
 
CONTAINS
 
! **************************************************************************************************
!> \brief Prepares a tensor to hold 3-center integrals (pq|X), where p,q are distributed according
!>        to the given 2d dbcsr distribution of the given . The third dimension of the tensor is
!>        iteslf not distributed (i.e. the t_pgrid's third dimension has size 1). The blocks are
!>        reserved according to the neighbor lists
!> \param pq_X the tensor to store the integrals
!> \param ab_nl the 1st and 2nd center neighbor list
!> \param ac_nl the 1st and 3rd center neighbor list
!> \param matrix_dist ...
!> \param blk_size_1 the block size in the first dimension
!> \param blk_size_2 the block size in the second dimension
!> \param blk_size_3 the block size in the third dimension
!> \param only_bc_same_center only keep block if, for the corresponding integral (ab|c), b and c
!>        share the same center, i.e. r_bc = 0.0
! **************************************************************************************************
   SUBROUTINE create_pqx_tensor(pq_X, ab_nl, ac_nl, matrix_dist, blk_size_1, blk_size_2, &
                                blk_size_3, only_bc_same_center)
 
      TYPE(dbt_type), INTENT(OUT)                        :: pq_x
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: ab_nl, ac_nl
      TYPE(dbcsr_distribution_type), INTENT(IN)          :: matrix_dist
      INTEGER, DIMENSION(:), INTENT(IN)                  :: blk_size_1, blk_size_2, blk_size_3
      LOGICAL, INTENT(IN), OPTIONAL                      :: only_bc_same_center
 
      CHARACTER(len=*), PARAMETER                        :: routinen = 'create_pqX_tensor'
 
      INTEGER                                            :: a, b, group_handle, handle, i, iatom, &
                                                            ikind, jatom, katom, kkind, nblk, &
                                                            nblk_3, nblk_per_thread, nkind
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: idx1, idx2, idx3
      INTEGER, DIMENSION(3)                              :: pdims
      INTEGER, DIMENSION(:), POINTER                     :: col_dist, row_dist
      INTEGER, DIMENSION(:, :), POINTER                  :: mat_pgrid
      LOGICAL                                            :: my_sort_bc, symmetric
      REAL(dp), DIMENSION(3)                             :: rab, rac, rbc
      TYPE(dbt_distribution_type)                        :: t_dist
      TYPE(dbt_pgrid_type)                               :: t_pgrid
      TYPE(mp_comm_type)                                 :: group
      TYPE(neighbor_list_iterator_p_type), &
         DIMENSION(:), POINTER                           :: ab_iter, ac_iter
 
      NULLIFY (ab_iter, ac_iter, col_dist, row_dist, mat_pgrid)
 
      CALL timeset(routinen, handle)
 
      my_sort_bc = .false.
      IF (PRESENT(only_bc_same_center)) my_sort_bc = only_bc_same_center
 
      CALL get_neighbor_list_set_p(ab_nl, symmetric=symmetric)
      cpassert(symmetric)
 
      !get 2D distribution info from matrix_dist
      CALL dbcsr_distribution_get(matrix_dist, pgrid=mat_pgrid, group=group_handle, &
                                  row_dist=row_dist, col_dist=col_dist)
      CALL group%set_handle(group_handle)
 
      !create the corresponding tensor proc grid and dist
      pdims(1) = SIZE(mat_pgrid, 1); pdims(2) = SIZE(mat_pgrid, 2); pdims(3) = 1
      CALL dbt_pgrid_create(group, pdims, t_pgrid)
 
      nblk_3 = SIZE(blk_size_3)
      CALL dbt_distribution_new(t_dist, t_pgrid, nd_dist_1=row_dist, nd_dist_2=col_dist, &
                                nd_dist_3=[(0, i=1, nblk_3)])
 
      !create the tensor itself.
      CALL dbt_create(pq_x, name="(pq|X)", dist=t_dist, map1_2d=[1, 2], map2_2d=[3], &
                      blk_size_1=blk_size_1, blk_size_2=blk_size_2, blk_size_3=blk_size_3)
 
      !count the blocks to reserve !note: dbcsr takes care of only keeping unique indices
      CALL neighbor_list_iterator_create(ab_iter, ab_nl)
      CALL neighbor_list_iterator_create(ac_iter, ac_nl, search=.true.)
      nblk = 0
      DO WHILE (neighbor_list_iterate(ab_iter) == 0)
         CALL get_iterator_info(ab_iter, ikind=ikind, iatom=iatom, nkind=nkind, r=rab)
 
         DO kkind = 1, nkind
            CALL nl_set_sub_iterator(ac_iter, ikind, kkind, iatom)
 
            DO WHILE (nl_sub_iterate(ac_iter) == 0)
 
               IF (my_sort_bc) THEN
                  !we check for rbc or rac because of symmetry in ab_nl
                  CALL get_iterator_info(ac_iter, r=rac)
                  rbc(:) = rac(:) - rab(:)
                  IF (.NOT. (all(abs(rbc) <= 1.0e-8_dp) .OR. all(abs(rac) <= 1.0e-8_dp))) cycle
 
               END IF
 
               nblk = nblk + 1
            END DO !ac_iter
         END DO !kkind
      END DO !ab_iter
      CALL neighbor_list_iterator_release(ab_iter)
      CALL neighbor_list_iterator_release(ac_iter)
 
      ALLOCATE (idx1(nblk), idx2(nblk), idx3(nblk))
 
      !actually reserve the blocks
      CALL neighbor_list_iterator_create(ab_iter, ab_nl)
      CALL neighbor_list_iterator_create(ac_iter, ac_nl, search=.true.)
      nblk = 0
      DO WHILE (neighbor_list_iterate(ab_iter) == 0)
         CALL get_iterator_info(ab_iter, ikind=ikind, iatom=iatom, jatom=jatom, nkind=nkind, r=rab)
 
         DO kkind = 1, nkind
            CALL nl_set_sub_iterator(ac_iter, ikind, kkind, iatom)
 
            DO WHILE (nl_sub_iterate(ac_iter) == 0)
               CALL get_iterator_info(ac_iter, jatom=katom, r=rac)
 
               IF (my_sort_bc) THEN
                  !we check for rbc or rac because of symmetry in ab_nl
                  CALL get_iterator_info(ac_iter, r=rac)
                  rbc(:) = rac(:) - rab(:)
                  IF (.NOT. (all(abs(rbc) <= 1.0e-8_dp) .OR. all(abs(rac) <= 1.0e-8_dp))) cycle
 
               END IF
 
               nblk = nblk + 1
 
               idx1(nblk) = iatom
               idx2(nblk) = jatom
               idx3(nblk) = katom
 
            END DO !ac_iter
         END DO !kkind
      END DO !ab_iter
      CALL neighbor_list_iterator_release(ab_iter)
      CALL neighbor_list_iterator_release(ac_iter)
 
!TODO: Parallelize creation of block list.
!$OMP PARALLEL DEFAULT(NONE) SHARED(pq_X, nblk, idx1, idx2, idx3) PRIVATE(nblk_per_thread,A,b)
      nblk_per_thread = nblk/omp_get_num_threads() + 1
      a = omp_get_thread_num()*nblk_per_thread + 1
      b = min(a + nblk_per_thread, nblk)
      CALL dbt_reserve_blocks(pq_x, idx1(a:b), idx2(a:b), idx3(a:b))
!$OMP END PARALLEL
      CALL dbt_finalize(pq_x)
 
      !clean-up
      CALL dbt_distribution_destroy(t_dist)
      CALL dbt_pgrid_destroy(t_pgrid)
 
      CALL timestop(handle)
 
   END SUBROUTINE create_pqx_tensor
 
! **************************************************************************************************
!> \brief Fills the given 3 dimensional (pq|X) tensor with integrals
!> \param pq_X the tensor to fill
!> \param ab_nl the neighbor list for the first 2 centers
!> \param ac_nl the neighbor list for the first and third centers
!> \param basis_set_list_a basis sets for first center
!> \param basis_set_list_b basis sets for second center
!> \param basis_set_list_c basis sets for third center
!> \param potential_parameter the operator for the integrals
!> \param qs_env ...
!> \param only_bc_same_center same as in create_pqX_tensor
!> \param eps_screen threshold for possible screening
!> \note The following indices are happily mixed within this routine: First center i,a,p
!>       Second center: j,b,q       Third center: k,c,X
! **************************************************************************************************
   SUBROUTINE fill_pqx_tensor(pq_X, ab_nl, ac_nl, basis_set_list_a, basis_set_list_b, &
                              basis_set_list_c, potential_parameter, qs_env, &
                              only_bc_same_center, eps_screen)
 
      TYPE(dbt_type)                                     :: pq_x
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: ab_nl, ac_nl
      TYPE(gto_basis_set_p_type), DIMENSION(:), POINTER  :: basis_set_list_a, basis_set_list_b, &
                                                            basis_set_list_c
      TYPE(coulomb_operator_type)                        :: potential_parameter
      TYPE(qs_environment_type), POINTER                 :: qs_env
      LOGICAL, INTENT(IN), OPTIONAL                      :: only_bc_same_center
      REAL(dp), INTENT(IN), OPTIONAL                     :: eps_screen
 
      CHARACTER(len=*), PARAMETER                        :: routinen = 'fill_pqX_tensor'
 
      INTEGER :: egfa, egfb, egfc, handle, i, iatom, ibasis, ikind, ilist, imax, iset, jatom, &
         jkind, jset, katom, kkind, kset, m_max, max_nset, maxli, maxlj, maxlk, mepos, nbasis, &
         ncoa, ncob, ncoc, ni, nj, nk, nseta, nsetb, nsetc, nthread, sgfa, sgfb, sgfc
      INTEGER, DIMENSION(:), POINTER                     :: la_max, la_min, lb_max, lb_min, lc_max, &
                                                            lc_min, npgfa, npgfb, npgfc, nsgfa, &
                                                            nsgfb, nsgfc
      INTEGER, DIMENSION(:, :), POINTER                  :: first_sgfa, first_sgfb, first_sgfc
      LOGICAL                                            :: do_screen, my_sort_bc
      REAL(dp)                                           :: dij, dik, djk, my_eps_screen, ri(3), &
                                                            rij(3), rik(3), rj(3), rjk(3), rk(3), &
                                                            sabc_ext, screen_radius
      REAL(dp), ALLOCATABLE, DIMENSION(:)                :: ccp_buffer, cpp_buffer
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: max_contr, max_contra, max_contrb, &
                                                            max_contrc
      REAL(dp), ALLOCATABLE, DIMENSION(:, :, :)          :: iabc, sabc, work
      REAL(dp), DIMENSION(:), POINTER                    :: set_radius_a, set_radius_b, set_radius_c
      REAL(dp), DIMENSION(:, :), POINTER                 :: rpgf_a, rpgf_b, rpgf_c, zeta, zetb, zetc
      TYPE(coulomb_integral_context_type)                :: integral_context
      TYPE(cp_2d_r_p_type), DIMENSION(:, :), POINTER     :: spb, spc, tspa
      TYPE(gto_basis_set_p_type), DIMENSION(:), POINTER  :: basis_set_list
      TYPE(gto_basis_set_type), POINTER                  :: basis_set, basis_set_a, basis_set_b, &
                                                            basis_set_c
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(o3c_container_type), POINTER                  :: o3c
      TYPE(o3c_iterator_type)                            :: o3c_iterator
 
      NULLIFY (basis_set, basis_set_list, para_env, la_max, la_min)
      NULLIFY (lb_max, lb_min, lc_max, lc_min, npgfa, npgfb, npgfc, nsgfa, nsgfb, nsgfc)
      NULLIFY (first_sgfa, first_sgfb, first_sgfc, set_radius_a, set_radius_b, set_radius_c)
      NULLIFY (rpgf_a, rpgf_b, rpgf_c, zeta, zetb, zetc)
      NULLIFY (basis_set_a, basis_set_b, basis_set_c, tspa, spb, spc)
 
      CALL timeset(routinen, handle)
 
      !Need the max l for each basis for libint (and overall max #of sets for screening)
      nbasis = SIZE(basis_set_list_a)
      max_nset = 0
      maxli = 0
      DO ibasis = 1, nbasis
         CALL get_gto_basis_set(gto_basis_set=basis_set_list_a(ibasis)%gto_basis_set, &
                                maxl=imax, nset=iset, nsgf_set=nsgfa)
         maxli = max(maxli, imax)
         max_nset = max(max_nset, iset)
      END DO
      maxlj = 0
      DO ibasis = 1, nbasis
         CALL get_gto_basis_set(gto_basis_set=basis_set_list_b(ibasis)%gto_basis_set, &
                                maxl=imax, nset=iset, nsgf_set=nsgfb, npgf=npgfb)
         maxlj = max(maxlj, imax)
         max_nset = max(max_nset, iset)
      END DO
      maxlk = 0
      DO ibasis = 1, nbasis
         CALL get_gto_basis_set(gto_basis_set=basis_set_list_c(ibasis)%gto_basis_set, &
                                maxl=imax, nset=iset, npgf=npgfc)
         maxlk = max(maxlk, imax)
         max_nset = max(max_nset, iset)
      END DO
      m_max = maxli + maxlj + maxlk
 
      !Screening
      do_screen = .false.
      IF (PRESENT(eps_screen)) THEN
         do_screen = .true.
         my_eps_screen = eps_screen
      END IF
      screen_radius = 0.0_dp
      IF (potential_parameter%potential_type == do_potential_truncated .OR. &
          potential_parameter%potential_type == do_potential_short) THEN
 
         screen_radius = potential_parameter%cutoff_radius*cutoff_screen_factor
      ELSE IF (potential_parameter%potential_type == do_potential_coulomb) THEN
 
         screen_radius = 1000000.0_dp
      END IF
 
      !get maximum contraction values for abc_contract screening
      IF (do_screen) THEN
 
         !Allocate max_contraction arrays such that we have a specific value for each set/kind
         ALLOCATE (max_contr(max_nset, nbasis), max_contra(max_nset, nbasis), &
                   max_contrb(max_nset, nbasis), max_contrc(max_nset, nbasis))
 
         !Not the most elegent, but better than copying 3 times the same
         DO ilist = 1, 3
 
            IF (ilist == 1) basis_set_list => basis_set_list_a
            IF (ilist == 2) basis_set_list => basis_set_list_b
            IF (ilist == 3) basis_set_list => basis_set_list_c
 
            max_contr = 0.0_dp
 
            DO ibasis = 1, nbasis
               basis_set => basis_set_list(ibasis)%gto_basis_set
 
               DO iset = 1, basis_set%nset
 
                  ncoa = basis_set%npgf(iset)*ncoset(basis_set%lmax(iset))
                  sgfa = basis_set%first_sgf(1, iset)
                  egfa = sgfa + basis_set%nsgf_set(iset) - 1
 
                  max_contr(iset, ibasis) = &
                     maxval([(sum(abs(basis_set%sphi(1:ncoa, i))), i=sgfa, egfa)])
 
               END DO !iset
            END DO !ibasis
 
            IF (ilist == 1) max_contra(:, :) = max_contr(:, :)
            IF (ilist == 2) max_contrb(:, :) = max_contr(:, :)
            IF (ilist == 3) max_contrc(:, :) = max_contr(:, :)
         END DO !ilist
         DEALLOCATE (max_contr)
      END IF !do_screen
 
      !To minimize memory ops in contraction, we need to pre-allocate buffers, pre-tranpose sphi_a
      !and also trim sphi in general to have contiguous arrays
      ALLOCATE (tspa(max_nset, nbasis), spb(max_nset, nbasis), spc(max_nset, nbasis))
      DO ibasis = 1, nbasis
         DO iset = 1, max_nset
            NULLIFY (tspa(iset, ibasis)%array)
            NULLIFY (spb(iset, ibasis)%array)
            NULLIFY (spc(iset, ibasis)%array)
         END DO
      END DO
 
      DO ilist = 1, 3
 
         DO ibasis = 1, nbasis
            IF (ilist == 1) basis_set => basis_set_list_a(ibasis)%gto_basis_set
            IF (ilist == 2) basis_set => basis_set_list_b(ibasis)%gto_basis_set
            IF (ilist == 3) basis_set => basis_set_list_c(ibasis)%gto_basis_set
 
            DO iset = 1, basis_set%nset
 
               ncoa = basis_set%npgf(iset)*ncoset(basis_set%lmax(iset))
               sgfa = basis_set%first_sgf(1, iset)
               egfa = sgfa + basis_set%nsgf_set(iset) - 1
 
               IF (ilist == 1) THEN
                  ALLOCATE (tspa(iset, ibasis)%array(basis_set%nsgf_set(iset), ncoa))
#if defined(__LIBXS)
                  CALL libxs_otrans(c_loc(tspa(iset, ibasis)%array), &
                                    c_loc(basis_set%sphi(1, sgfa)), &
                                    8, ncoa, egfa - sgfa + 1, &
                                    SIZE(basis_set%sphi, 1), &
                                    basis_set%nsgf_set(iset))
#else
                  tspa(iset, ibasis)%array(:, :) = transpose(basis_set%sphi(1:ncoa, sgfa:egfa))
#endif
               ELSE IF (ilist == 2) THEN
                  ALLOCATE (spb(iset, ibasis)%array(ncoa, basis_set%nsgf_set(iset)))
#if defined(__LIBXS)
                  CALL libxs_matcopy(c_loc(spb(iset, ibasis)%array), &
                                     c_loc(basis_set%sphi(1, sgfa)), &
                                     8, ncoa, egfa - sgfa + 1, &
                                     SIZE(basis_set%sphi, 1), ncoa)
#else
                  spb(iset, ibasis)%array(:, :) = basis_set%sphi(1:ncoa, sgfa:egfa)
#endif
               ELSE
                  ALLOCATE (spc(iset, ibasis)%array(ncoa, basis_set%nsgf_set(iset)))
#if defined(__LIBXS)
                  CALL libxs_matcopy(c_loc(spc(iset, ibasis)%array), &
                                     c_loc(basis_set%sphi(1, sgfa)), &
                                     8, ncoa, egfa - sgfa + 1, &
                                     SIZE(basis_set%sphi, 1), ncoa)
#else
                  spc(iset, ibasis)%array(:, :) = basis_set%sphi(1:ncoa, sgfa:egfa)
#endif
               END IF
 
            END DO !iset
         END DO !ibasis
      END DO !ilist
 
      my_sort_bc = .false.
      IF (PRESENT(only_bc_same_center)) my_sort_bc = only_bc_same_center
 
      CALL get_qs_env(qs_env, para_env=para_env)
      CALL r12_fundamentals_init(m_max, potential_parameter, para_env)
 
      !Strategy: we use the o3c iterator because it is OMP parallelized and also takes the
      !          only_bc_same_center argument. Only the dbcsr_put_block is critical
 
      nthread = 1
!$    nthread = omp_get_max_threads()
 
      ALLOCATE (o3c)
      CALL init_o3c_container(o3c, 1, basis_set_list_a, basis_set_list_b, basis_set_list_c, &
                              ab_nl, ac_nl, only_bc_same_center=my_sort_bc)
      CALL o3c_iterator_create(o3c, o3c_iterator, nthread=nthread)
 
!$OMP PARALLEL DEFAULT(NONE) &
!$OMP SHARED (pq_X,do_screen,max_nset,basis_set_list_a,max_contra,max_contrb,max_contrc,&
!$OMP         basis_set_list_b, basis_set_list_c,ncoset,screen_radius,potential_parameter,&
!$OMP         my_eps_screen,maxli,maxlj,maxlk,my_sort_bc,nthread,o3c,o3c_iterator,tspa,spb,spc) &
!$OMP PRIVATE (integral_context,i,mepos,work,iset,ncoa,sgfa,egfa,nseta,iatom,ikind,jatom,jkind,katom,kkind,&
!$OMP          rij,rik,rjk,basis_set_a,nsetb,la_max,la_min,lb_max,lb_min,lc_max,lc_min,npgfa,&
!$OMP          npgfb,npgfc,nsgfa,nsgfb,nsgfc,ri,rk,first_sgfa,first_sgfb,first_sgfc,nsetc,&
!$OMP          set_radius_a,set_radius_b,set_radius_c,rj,rpgf_a,rpgf_b,rpgf_c,zeta,zetb,&
!$OMP          zetc,basis_set_b,basis_set_c,dij,dik,djk,ni,nj,nk,iabc,sabc,jset,kset,&
!$OMP          ncob,ncoc,sgfb,sgfc,egfb,egfc,sabc_ext,cpp_buffer,ccp_buffer)
 
      mepos = 0
!$    mepos = omp_get_thread_num()
 
      !each thread need its own libint object (internals may change at different rates)
      integral_context = coulomb_integral_context_type()
      CALL coulomb_integral_init(integral_context, max_am_3c=max(maxli, maxlj, maxlk), &
                                 derivative_order=0, library=library_libint)
 
      DO WHILE (o3c_iterate(o3c_iterator, mepos=mepos) == 0)
         CALL get_o3c_iterator_info(o3c_iterator, mepos=mepos, ikind=ikind, jkind=jkind, kkind=kkind, &
                                    iatom=iatom, jatom=jatom, katom=katom, rij=rij, rik=rik)
 
         !get first center basis info
         basis_set_a => basis_set_list_a(ikind)%gto_basis_set
         first_sgfa => basis_set_a%first_sgf
         la_max => basis_set_a%lmax
         la_min => basis_set_a%lmin
         npgfa => basis_set_a%npgf
         nseta = basis_set_a%nset
         nsgfa => basis_set_a%nsgf_set
         zeta => basis_set_a%zet
         rpgf_a => basis_set_a%pgf_radius
         set_radius_a => basis_set_a%set_radius
         ni = sum(nsgfa)
         !second center basis info
         basis_set_b => basis_set_list_b(jkind)%gto_basis_set
         first_sgfb => basis_set_b%first_sgf
         lb_max => basis_set_b%lmax
         lb_min => basis_set_b%lmin
         npgfb => basis_set_b%npgf
         nsetb = basis_set_b%nset
         nsgfb => basis_set_b%nsgf_set
         zetb => basis_set_b%zet
         rpgf_b => basis_set_b%pgf_radius
         set_radius_b => basis_set_b%set_radius
         nj = sum(nsgfb)
         !third center basis info
         basis_set_c => basis_set_list_c(kkind)%gto_basis_set
         first_sgfc => basis_set_c%first_sgf
         lc_max => basis_set_c%lmax
         lc_min => basis_set_c%lmin
         npgfc => basis_set_c%npgf
         nsetc = basis_set_c%nset
         nsgfc => basis_set_c%nsgf_set
         zetc => basis_set_c%zet
         rpgf_c => basis_set_c%pgf_radius
         set_radius_c => basis_set_c%set_radius
         nk = sum(nsgfc)
 
         !position and distances, only relative pos matter for libint
         rjk = rik - rij
         ri = 0.0_dp
         rj = rij ! ri + rij
         rk = rik ! ri + rik
 
         djk = norm2(rjk)
         dij = norm2(rij)
         dik = norm2(rik)
 
         !sgf integrals
         ALLOCATE (iabc(ni, nj, nk))
         iabc(:, :, :) = 0.0_dp
 
         DO iset = 1, nseta
            ncoa = npgfa(iset)*ncoset(la_max(iset))
            sgfa = first_sgfa(1, iset)
            egfa = sgfa + nsgfa(iset) - 1
 
            DO jset = 1, nsetb
               ncob = npgfb(jset)*ncoset(lb_max(jset))
               sgfb = first_sgfb(1, jset)
               egfb = sgfb + nsgfb(jset) - 1
 
               !screening (overlap)
               IF (set_radius_a(iset) + set_radius_b(jset) < dij) cycle
 
               DO kset = 1, nsetc
                  ncoc = npgfc(kset)*ncoset(lc_max(kset))
                  sgfc = first_sgfc(1, kset)
                  egfc = sgfc + nsgfc(kset) - 1
 
                  !screening (potential)
                  IF (set_radius_a(iset) + set_radius_c(kset) + screen_radius < dik) cycle
                  IF (set_radius_b(jset) + set_radius_c(kset) + screen_radius < djk) cycle
 
                  !pgf integrals
                  ALLOCATE (sabc(ncoa, ncob, ncoc))
                  sabc(:, :, :) = 0.0_dp
 
                  IF (do_screen) THEN
                     CALL compute_coulomb_3c(integral_context, la_min(iset), la_max(iset), lb_min(jset), &
                                             lb_max(jset), lc_min(kset), lc_max(kset), npgfa(iset), npgfb(jset), npgfc(kset), &
                                             zeta(:, iset), zetb(:, jset), zetc(:, kset), rpgf_a(:, iset), rpgf_b(:, jset), &
                                             rpgf_c(:, kset), ri, rj, rk, sabc, potential_parameter, dab=dij, dac=dik, dbc=djk, &
                                             int_abc_ext=sabc_ext)
                     IF (my_eps_screen > sabc_ext*(max_contra(iset, ikind)* &
                                                   max_contrb(jset, jkind)* &
                                                   max_contrc(kset, kkind))) THEN
                        DEALLOCATE (sabc)
                        cycle
                     END IF
                  ELSE
                     CALL compute_coulomb_3c(integral_context, la_min(iset), la_max(iset), lb_min(jset), &
                                             lb_max(jset), lc_min(kset), lc_max(kset), npgfa(iset), npgfb(jset), npgfc(kset), &
                                             zeta(:, iset), zetb(:, jset), zetc(:, kset), rpgf_a(:, iset), rpgf_b(:, jset), &
                                             rpgf_c(:, kset), ri, rj, rk, sabc, potential_parameter, dab=dij, dac=dik, dbc=djk)
                  END IF
 
                  ALLOCATE (work(nsgfa(iset), nsgfb(jset), nsgfc(kset)))
 
                  CALL abc_contract_xsmm(work, sabc, tspa(iset, ikind)%array, spb(jset, jkind)%array, &
                                         spc(kset, kkind)%array, ncoa, ncob, ncoc, nsgfa(iset), &
                                         nsgfb(jset), nsgfc(kset), cpp_buffer, ccp_buffer)
 
                  iabc(sgfa:egfa, sgfb:egfb, sgfc:egfc) = work(:, :, :)
                  DEALLOCATE (sabc, work)
 
               END DO !kset
            END DO !jset
         END DO !iset
 
         !Add the integral to the proper tensor block
!$OMP CRITICAL(omp_fill_pqX_tensor)
         CALL dbt_put_block(pq_x, [iatom, jatom, katom], shape(iabc), iabc, summation=.true.)
!$OMP END CRITICAL(omp_fill_pqX_tensor)
 
         DEALLOCATE (iabc)
      END DO !o3c_iterator
 
      IF (ALLOCATED(ccp_buffer)) DEALLOCATE (ccp_buffer)
      IF (ALLOCATED(cpp_buffer)) DEALLOCATE (cpp_buffer)
 
      CALL coulomb_integral_cleanup(integral_context)
 
!$OMP END PARALLEL
      CALL o3c_iterator_release(o3c_iterator)
      CALL release_o3c_container(o3c)
      DEALLOCATE (o3c)
 
      DO iset = 1, max_nset
         DO ibasis = 1, nbasis
            IF (ASSOCIATED(tspa(iset, ibasis)%array)) DEALLOCATE (tspa(iset, ibasis)%array)
            IF (ASSOCIATED(spb(iset, ibasis)%array)) DEALLOCATE (spb(iset, ibasis)%array)
            IF (ASSOCIATED(spc(iset, ibasis)%array)) DEALLOCATE (spc(iset, ibasis)%array)
         END DO
      END DO
      DEALLOCATE (tspa, spb, spc)
 
      CALL timestop(handle)
 
   END SUBROUTINE fill_pqx_tensor
 
! **************************************************************************************************
!> \brief Builds a neighbor lists set for overlaping 2-center S_ab, where b is restricted on a
!>        a given list of atoms. Used for Coulomb RI where (aI|P) = sum_b C_bI (ab|P), where
!>        contraction coeff C_bI is assumed to be non-zero only on excited atoms
!> \param ab_list the neighbor list
!> \param basis_a basis set list for atom a
!> \param basis_b basis set list for atom b
!> \param qs_env ...
!> \param excited_atoms the indices of the excited atoms on which b is centered
!> \param ext_dist2d use an external distribution2d
! **************************************************************************************************
   SUBROUTINE build_xas_tdp_ovlp_nl(ab_list, basis_a, basis_b, qs_env, excited_atoms, ext_dist2d)
 
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: ab_list
      TYPE(gto_basis_set_p_type), DIMENSION(:), POINTER  :: basis_a, basis_b
      TYPE(qs_environment_type), POINTER                 :: qs_env
      INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL        :: excited_atoms
      TYPE(distribution_2d_type), OPTIONAL, POINTER      :: ext_dist2d
 
      INTEGER                                            :: ikind, nkind
      LOGICAL                                            :: my_restrictb
      LOGICAL, ALLOCATABLE, DIMENSION(:)                 :: a_present, b_present
      REAL(dp)                                           :: subcells
      REAL(dp), ALLOCATABLE, DIMENSION(:)                :: a_radius, b_radius
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: pair_radius
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(cell_type), POINTER                           :: cell
      TYPE(distribution_1d_type), POINTER                :: distribution_1d
      TYPE(distribution_2d_type), POINTER                :: distribution_2d
      TYPE(local_atoms_type), ALLOCATABLE, DIMENSION(:)  :: atom2d
      TYPE(molecule_type), DIMENSION(:), POINTER         :: molecule_set
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
 
      NULLIFY (atomic_kind_set, distribution_1d, distribution_2d, molecule_set, particle_set, cell)
 
!  Initialization
      CALL get_qs_env(qs_env, nkind=nkind)
      CALL section_vals_val_get(qs_env%input, "DFT%SUBCELLS", r_val=subcells)
 
      my_restrictb = .false.
      IF (PRESENT(excited_atoms)) THEN
         my_restrictb = .true.
      END IF
 
      ALLOCATE (a_present(nkind), b_present(nkind))
      a_present = .false.
      b_present = .false.
      ALLOCATE (a_radius(nkind), b_radius(nkind))
      a_radius = 0.0_dp
      b_radius = 0.0_dp
 
!  Set up the radii
      DO ikind = 1, nkind
         IF (ASSOCIATED(basis_a(ikind)%gto_basis_set)) THEN
            a_present(ikind) = .true.
            CALL get_gto_basis_set(basis_a(ikind)%gto_basis_set, kind_radius=a_radius(ikind))
         END IF
 
         IF (ASSOCIATED(basis_b(ikind)%gto_basis_set)) THEN
            b_present(ikind) = .true.
            CALL get_gto_basis_set(basis_b(ikind)%gto_basis_set, kind_radius=b_radius(ikind))
         END IF
      END DO !ikind
 
      ALLOCATE (pair_radius(nkind, nkind))
      pair_radius = 0.0_dp
      CALL pair_radius_setup(a_present, b_present, a_radius, b_radius, pair_radius)
 
!  Set up the nl
      CALL get_qs_env(qs_env, atomic_kind_set=atomic_kind_set, cell=cell, &
                      distribution_2d=distribution_2d, local_particles=distribution_1d, &
                      particle_set=particle_set, molecule_set=molecule_set)
 
      !use an external distribution_2d if required
      IF (PRESENT(ext_dist2d)) distribution_2d => ext_dist2d
 
      ALLOCATE (atom2d(nkind))
      CALL atom2d_build(atom2d, distribution_1d, distribution_2d, atomic_kind_set, &
                        molecule_set, .false., particle_set)
 
      IF (my_restrictb) THEN
 
         CALL build_neighbor_lists(ab_list, particle_set, atom2d, cell, pair_radius, subcells, &
                                   atomb_to_keep=excited_atoms, nlname="XAS_TDP_ovlp_nl")
 
      ELSE
 
         CALL build_neighbor_lists(ab_list, particle_set, atom2d, cell, pair_radius, subcells, &
                                   nlname="XAS_TDP_ovlp_nl")
 
      END IF
!  Clean-up
      CALL atom2d_cleanup(atom2d)
 
   END SUBROUTINE build_xas_tdp_ovlp_nl
 
! **************************************************************************************************
!> \brief Builds a neighbor lists set taylored for 3-center integral within XAS TDP, such that only
!>        excited atoms are taken into account for the list_c
!> \param ac_list the neighbor list ready for 3-center integrals
!> \param basis_a basis set list for atom a
!> \param basis_c basis set list for atom c
!> \param op_type to indicate whther the list should be built with overlap, Coulomb or else in mind
!> \param qs_env ...
!> \param excited_atoms the indices of the excited atoms to consider (if not given, all atoms are taken)
!> \param x_range in case some truncated/screened operator is used, gives its range
!> \param ext_dist2d external distribution_2d to be used
!> \note Based on setup_neighbor_list with added features
! **************************************************************************************************
   SUBROUTINE build_xas_tdp_3c_nl(ac_list, basis_a, basis_c, op_type, qs_env, excited_atoms, &
                                  x_range, ext_dist2d)
 
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: ac_list
      TYPE(gto_basis_set_p_type), DIMENSION(:), POINTER  :: basis_a, basis_c
      INTEGER, INTENT(IN)                                :: op_type
      TYPE(qs_environment_type), POINTER                 :: qs_env
      INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL        :: excited_atoms
      REAL(dp), INTENT(IN), OPTIONAL                     :: x_range
      TYPE(distribution_2d_type), OPTIONAL, POINTER      :: ext_dist2d
 
      INTEGER                                            :: ikind, nkind
      LOGICAL                                            :: sort_atoms
      LOGICAL, ALLOCATABLE, DIMENSION(:)                 :: a_present, c_present
      REAL(dp)                                           :: subcells
      REAL(dp), ALLOCATABLE, DIMENSION(:)                :: a_radius, c_radius
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: pair_radius
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(cell_type), POINTER                           :: cell
      TYPE(distribution_1d_type), POINTER                :: distribution_1d
      TYPE(distribution_2d_type), POINTER                :: distribution_2d
      TYPE(local_atoms_type), ALLOCATABLE, DIMENSION(:)  :: atom2d
      TYPE(molecule_type), DIMENSION(:), POINTER         :: molecule_set
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
 
      NULLIFY (atomic_kind_set, distribution_1d, distribution_2d, molecule_set, particle_set, cell)
 
!  Initialization
      CALL get_qs_env(qs_env, nkind=nkind)
      CALL section_vals_val_get(qs_env%input, "DFT%SUBCELLS", r_val=subcells)
      sort_atoms = .false.
      IF ((PRESENT(excited_atoms))) sort_atoms = .true.
 
      ALLOCATE (a_present(nkind), c_present(nkind))
      a_present = .false.
      c_present = .false.
      ALLOCATE (a_radius(nkind), c_radius(nkind))
      a_radius = 0.0_dp
      c_radius = 0.0_dp
 
!  Set up the radii, depending on the operator type
      IF (op_type == do_potential_id) THEN
 
         !overlap => use the kind radius for both a and c
         DO ikind = 1, nkind
            !orbital basis set
            IF (ASSOCIATED(basis_a(ikind)%gto_basis_set)) THEN
               a_present(ikind) = .true.
               CALL get_gto_basis_set(basis_a(ikind)%gto_basis_set, kind_radius=a_radius(ikind))
            END IF
            !RI_XAS basis set
            IF (ASSOCIATED(basis_c(ikind)%gto_basis_set)) THEN
               c_present(ikind) = .true.
               CALL get_gto_basis_set(basis_c(ikind)%gto_basis_set, kind_radius=c_radius(ikind))
            END IF
         END DO !ikind
 
      ELSE IF (op_type == do_potential_coulomb) THEN
 
         !Coulomb operator, virtually infinite range => set c_radius to arbitrarily large number
         DO ikind = 1, nkind
            IF (ASSOCIATED(basis_c(ikind)%gto_basis_set)) THEN
               c_present(ikind) = .true.
               c_radius(ikind) = 1000000.0_dp
            END IF
            IF (ASSOCIATED(basis_a(ikind)%gto_basis_set)) a_present(ikind) = .true.
         END DO !ikind
 
      ELSE IF (op_type == do_potential_truncated .OR. op_type == do_potential_short) THEN
 
         !Truncated coulomb/short range: set c_radius to x_range + the kind_radii
         DO ikind = 1, nkind
            IF (ASSOCIATED(basis_a(ikind)%gto_basis_set)) THEN
               a_present(ikind) = .true.
               CALL get_gto_basis_set(basis_a(ikind)%gto_basis_set, kind_radius=a_radius(ikind))
            END IF
            IF (ASSOCIATED(basis_c(ikind)%gto_basis_set)) THEN
               c_present(ikind) = .true.
               CALL get_gto_basis_set(basis_c(ikind)%gto_basis_set, kind_radius=c_radius(ikind))
               c_radius(ikind) = c_radius(ikind) + x_range
            END IF
         END DO !ikind
 
      ELSE
         cpabort("Operator not known")
      END IF
 
      ALLOCATE (pair_radius(nkind, nkind))
      pair_radius = 0.0_dp
      CALL pair_radius_setup(a_present, c_present, a_radius, c_radius, pair_radius)
 
!  Actually setup the list
      CALL get_qs_env(qs_env, atomic_kind_set=atomic_kind_set, cell=cell, &
                      distribution_2d=distribution_2d, local_particles=distribution_1d, &
                      particle_set=particle_set, molecule_set=molecule_set)
 
      !use an external distribution_2d if required
      IF (PRESENT(ext_dist2d)) distribution_2d => ext_dist2d
 
      ALLOCATE (atom2d(nkind))
      CALL atom2d_build(atom2d, distribution_1d, distribution_2d, atomic_kind_set, &
                        molecule_set, .false., particle_set)
 
      IF (sort_atoms) THEN
         CALL build_neighbor_lists(ac_list, particle_set, atom2d, cell, pair_radius, subcells, &
                                   operator_type="ABC", atomb_to_keep=excited_atoms, &
                                   nlname="XAS_TDP_3c_nl")
      ELSE
         CALL build_neighbor_lists(ac_list, particle_set, atom2d, cell, pair_radius, subcells, &
                                   operator_type="ABC", nlname="XAS_TDP_3c_nl")
      END IF
 
!  Clean-up
      CALL atom2d_cleanup(atom2d)
 
   END SUBROUTINE build_xas_tdp_3c_nl
 
! **************************************************************************************************
!> \brief Returns an optimized distribution_2d for the given neighbor lists based on an evaluation
!>        of the cost of the corresponding 3-center integrals
!> \param opt_3c_dist2d the optimized distribution_2d
!> \param ab_list ...
!> \param ac_list ...
!> \param basis_set_a ...
!> \param basis_set_b ...
!> \param basis_set_c ...
!> \param qs_env ...
!> \param only_bc_same_center ...
! **************************************************************************************************
   SUBROUTINE get_opt_3c_dist2d(opt_3c_dist2d, ab_list, ac_list, basis_set_a, basis_set_b, &
                                basis_set_c, qs_env, only_bc_same_center)
 
      TYPE(distribution_2d_type), POINTER                :: opt_3c_dist2d
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: ab_list, ac_list
      TYPE(gto_basis_set_p_type), DIMENSION(:), POINTER  :: basis_set_a, basis_set_b, basis_set_c
      TYPE(qs_environment_type), POINTER                 :: qs_env
      LOGICAL, INTENT(IN), OPTIONAL                      :: only_bc_same_center
 
      CHARACTER(len=*), PARAMETER                        :: routinen = 'get_opt_3c_dist2d'
 
      INTEGER                                            :: handle, i, iatom, ikind, ip, jatom, &
                                                            jkind, kkind, mypcol, myprow, n, &
                                                            natom, nkind, npcol, nprow, nsgfa, &
                                                            nsgfb, nsgfc
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: nparticle_local_col, nparticle_local_row
      INTEGER, DIMENSION(:, :), POINTER                  :: col_dist, row_dist
      LOGICAL                                            :: my_sort_bc
      REAL(dp)                                           :: cost, rab(3), rac(3), rbc(3)
      REAL(dp), ALLOCATABLE, DIMENSION(:)                :: col_cost, col_proc_cost, row_cost, &
                                                            row_proc_cost
      TYPE(cp_1d_i_p_type), DIMENSION(:), POINTER        :: local_particle_col, local_particle_row
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(neighbor_list_iterator_p_type), &
         DIMENSION(:), POINTER                           :: ab_iter, ac_iter
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
 
      NULLIFY (para_env, col_dist, row_dist, blacs_env, qs_kind_set, particle_set)
      NULLIFY (local_particle_col, local_particle_row, ab_iter, ac_iter)
 
      CALL timeset(routinen, handle)
 
      !Idea: create a,b and a,c nl_iterators in the original dist, then loop over them and compute the
      !      cost of each ab pairs and project on the col/row. Then distribute the atom col/row to
      !      the proc col/row in order to spread out the cost as much as possible
 
      my_sort_bc = .false.
      IF (PRESENT(only_bc_same_center)) my_sort_bc = only_bc_same_center
 
      CALL get_qs_env(qs_env, natom=natom, para_env=para_env, blacs_env=blacs_env, &
                      qs_kind_set=qs_kind_set, particle_set=particle_set, nkind=nkind)
 
      myprow = blacs_env%mepos(1) + 1
      mypcol = blacs_env%mepos(2) + 1
      nprow = blacs_env%num_pe(1)
      npcol = blacs_env%num_pe(2)
 
      ALLOCATE (col_cost(natom), row_cost(natom))
      col_cost = 0.0_dp; row_cost = 0.0_dp
 
      ALLOCATE (row_dist(natom, 2), col_dist(natom, 2))
      row_dist = 1; col_dist = 1
 
      CALL neighbor_list_iterator_create(ab_iter, ab_list)
      CALL neighbor_list_iterator_create(ac_iter, ac_list, search=.true.)
      DO WHILE (neighbor_list_iterate(ab_iter) == 0)
         CALL get_iterator_info(ab_iter, ikind=ikind, jkind=jkind, iatom=iatom, jatom=jatom, r=rab)
 
         DO kkind = 1, nkind
            CALL nl_set_sub_iterator(ac_iter, ikind, kkind, iatom)
 
            DO WHILE (nl_sub_iterate(ac_iter) == 0)
 
               IF (my_sort_bc) THEN
                  !only take a,b,c if b,c (or a,c because of symmetry) share the same center
                  CALL get_iterator_info(ac_iter, r=rac)
                  rbc(:) = rac(:) - rab(:)
                  IF (.NOT. (all(abs(rbc) <= 1.0e-8_dp) .OR. all(abs(rac) <= 1.0e-8_dp))) cycle
 
               END IF
 
               !Use the size of integral as measure as contraciton cost seems to dominate
               nsgfa = basis_set_a(ikind)%gto_basis_set%nsgf
               nsgfb = basis_set_b(jkind)%gto_basis_set%nsgf
               nsgfc = basis_set_c(kkind)%gto_basis_set%nsgf
 
               cost = real(nsgfa*nsgfb*nsgfc, dp)
 
               row_cost(iatom) = row_cost(iatom) + cost
               col_cost(jatom) = col_cost(jatom) + cost
 
            END DO !ac_iter
         END DO !kkind
      END DO !ab_iter
      CALL neighbor_list_iterator_release(ab_iter)
      CALL neighbor_list_iterator_release(ac_iter)
 
      CALL para_env%sum(row_cost)
      CALL para_env%sum(col_cost)
 
      !Distribute the cost as evenly as possible
      ALLOCATE (col_proc_cost(npcol), row_proc_cost(nprow))
      col_proc_cost = 0.0_dp; row_proc_cost = 0.0_dp
      DO i = 1, natom
         iatom = maxloc(row_cost, 1)
         ip = minloc(row_proc_cost, 1)
         row_proc_cost(ip) = row_proc_cost(ip) + row_cost(iatom)
         row_dist(iatom, 1) = ip
         row_cost(iatom) = 0.0_dp
 
         iatom = maxloc(col_cost, 1)
         ip = minloc(col_proc_cost, 1)
         col_proc_cost(ip) = col_proc_cost(ip) + col_cost(iatom)
         col_dist(iatom, 1) = ip
         col_cost(iatom) = 0.0_dp
      END DO
 
      !the usual stuff
      ALLOCATE (local_particle_col(nkind), local_particle_row(nkind))
      ALLOCATE (nparticle_local_row(nkind), nparticle_local_col(nkind))
      nparticle_local_row = 0; nparticle_local_col = 0
 
      DO iatom = 1, natom
         ikind = particle_set(iatom)%atomic_kind%kind_number
 
         IF (row_dist(iatom, 1) == myprow) nparticle_local_row(ikind) = nparticle_local_row(ikind) + 1
         IF (col_dist(iatom, 1) == mypcol) nparticle_local_col(ikind) = nparticle_local_col(ikind) + 1
      END DO
 
      DO ikind = 1, nkind
         n = nparticle_local_row(ikind)
         ALLOCATE (local_particle_row(ikind)%array(n))
 
         n = nparticle_local_col(ikind)
         ALLOCATE (local_particle_col(ikind)%array(n))
      END DO
 
      nparticle_local_row = 0; nparticle_local_col = 0
      DO iatom = 1, natom
         ikind = particle_set(iatom)%atomic_kind%kind_number
 
         IF (row_dist(iatom, 1) == myprow) THEN
            nparticle_local_row(ikind) = nparticle_local_row(ikind) + 1
            local_particle_row(ikind)%array(nparticle_local_row(ikind)) = iatom
         END IF
         IF (col_dist(iatom, 1) == mypcol) THEN
            nparticle_local_col(ikind) = nparticle_local_col(ikind) + 1
            local_particle_col(ikind)%array(nparticle_local_col(ikind)) = iatom
         END IF
      END DO
 
      !Finally create the dist_2d
      row_dist(:, 1) = row_dist(:, 1) - 1
      col_dist(:, 1) = col_dist(:, 1) - 1
      CALL distribution_2d_create(opt_3c_dist2d, row_distribution_ptr=row_dist, &
                                  col_distribution_ptr=col_dist, local_rows_ptr=local_particle_row, &
                                  local_cols_ptr=local_particle_col, blacs_env=blacs_env)
 
      CALL timestop(handle)
 
   END SUBROUTINE get_opt_3c_dist2d
 
! **************************************************************************************************
!> \brief Computes the RI exchange 3-center integrals (ab|c), where c is from the RI_XAS basis and
!>        centered on excited atoms and kind. The operator used is that of the RI metric
!> \param ex_atoms excited atoms on which the third center is located
!> \param xas_tdp_env ...
!> \param xas_tdp_control ...
!> \param qs_env ...
!> \note  This routine is called once for each excited atom. Because there are many different a,b
!>        pairs involved, load balance is ok. This allows memory saving
! **************************************************************************************************
   SUBROUTINE compute_ri_3c_exchange(ex_atoms, xas_tdp_env, xas_tdp_control, qs_env)
 
      INTEGER, DIMENSION(:), INTENT(IN)                  :: ex_atoms
      TYPE(xas_tdp_env_type), POINTER                    :: xas_tdp_env
      TYPE(xas_tdp_control_type), POINTER                :: xas_tdp_control
      TYPE(qs_environment_type), POINTER                 :: qs_env
 
      CHARACTER(len=*), PARAMETER :: routinen = 'compute_ri_3c_exchange'
 
      INTEGER                                            :: handle, natom, nkind
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: blk_size_orb, blk_size_ri
      TYPE(dbcsr_distribution_type)                      :: opt_dbcsr_dist
      TYPE(gto_basis_set_p_type), DIMENSION(:), POINTER  :: basis_set_orb, basis_set_ri
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: ab_list, ac_list
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
 
      NULLIFY (basis_set_ri, basis_set_orb, ac_list, ab_list, qs_kind_set, para_env, particle_set)
 
      CALL timeset(routinen, handle)
 
!  Take what we need from the qs_env
      CALL get_qs_env(qs_env, nkind=nkind, qs_kind_set=qs_kind_set, para_env=para_env, &
                      natom=natom, particle_set=particle_set)
 
!  Build the basis set lists
      ALLOCATE (basis_set_ri(nkind))
      ALLOCATE (basis_set_orb(nkind))
      CALL basis_set_list_setup(basis_set_ri, "RI_XAS", qs_kind_set)
      CALL basis_set_list_setup(basis_set_orb, "ORB", qs_kind_set)
 
!  Get the optimized distribution 2d for theses integrals (and store it in xas_tdp_env)
      CALL build_xas_tdp_ovlp_nl(ab_list, basis_set_orb, basis_set_orb, qs_env)
      CALL build_xas_tdp_3c_nl(ac_list, basis_set_orb, basis_set_ri, &
                               xas_tdp_control%ri_m_potential%potential_type, qs_env, &
                               excited_atoms=ex_atoms, x_range=xas_tdp_control%ri_m_potential%cutoff_radius)
 
      CALL get_opt_3c_dist2d(xas_tdp_env%opt_dist2d_ex, ab_list, ac_list, basis_set_orb, &
                             basis_set_orb, basis_set_ri, qs_env)
      CALL release_neighbor_list_sets(ab_list)
      CALL release_neighbor_list_sets(ac_list)
 
!  Build the ab and ac centers neighbor lists based on the optimized distribution
      CALL build_xas_tdp_ovlp_nl(ab_list, basis_set_orb, basis_set_orb, qs_env, &
                                 ext_dist2d=xas_tdp_env%opt_dist2d_ex)
      CALL build_xas_tdp_3c_nl(ac_list, basis_set_orb, basis_set_ri, &
                               xas_tdp_control%ri_m_potential%potential_type, qs_env, &
                               excited_atoms=ex_atoms, x_range=xas_tdp_control%ri_m_potential%cutoff_radius, &
                               ext_dist2d=xas_tdp_env%opt_dist2d_ex)
 
!  Allocate, init and compute the integrals.
      ALLOCATE (blk_size_orb(natom), blk_size_ri(natom))
      CALL cp_dbcsr_dist2d_to_dist(xas_tdp_env%opt_dist2d_ex, opt_dbcsr_dist)
      CALL get_particle_set(particle_set, qs_kind_set, nsgf=blk_size_orb, basis=basis_set_orb)
      CALL get_particle_set(particle_set, qs_kind_set, nsgf=blk_size_ri, basis=basis_set_ri)
 
      ALLOCATE (xas_tdp_env%ri_3c_ex)
      CALL create_pqx_tensor(xas_tdp_env%ri_3c_ex, ab_list, ac_list, opt_dbcsr_dist, blk_size_orb, &
                             blk_size_orb, blk_size_ri)
      CALL fill_pqx_tensor(xas_tdp_env%ri_3c_ex, ab_list, ac_list, basis_set_orb, basis_set_orb, &
                           basis_set_ri, xas_tdp_control%ri_m_potential, qs_env, &
                           eps_screen=xas_tdp_control%eps_screen)
 
! Clean-up
      CALL release_neighbor_list_sets(ab_list)
      CALL release_neighbor_list_sets(ac_list)
      CALL dbcsr_distribution_release(opt_dbcsr_dist)
      DEALLOCATE (basis_set_ri, basis_set_orb)
 
      !not strictly necessary but avoid having any load unbalance here being reported in the
      !timings for other routines
      CALL para_env%sync()
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_ri_3c_exchange
 
! **************************************************************************************************
!> \brief Computes the RI Coulomb 3-center integrals (ab|c), where c is from the RI_XAS basis and
!>        centered on the excited atoms of xas_tdp_env
!> \param xas_tdp_env ...
!> \param qs_env ...
!> \note  The ri_3c_coul tensor of xas_tdp_env is defined and allocated here. Only computed once
!>        for the whole system (for optimized load balance). Ok because not too much memory needed
! **************************************************************************************************
   SUBROUTINE compute_ri_3c_coulomb(xas_tdp_env, qs_env)
 
      TYPE(xas_tdp_env_type), POINTER                    :: xas_tdp_env
      TYPE(qs_environment_type), POINTER                 :: qs_env
 
      CHARACTER(len=*), PARAMETER :: routinen = 'compute_ri_3c_coulomb'
 
      INTEGER                                            :: handle, natom, nkind
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: blk_size_orb, blk_size_ri
      TYPE(coulomb_operator_type)                        :: pot
      TYPE(dbcsr_distribution_type)                      :: opt_dbcsr_dist
      TYPE(gto_basis_set_p_type), DIMENSION(:), POINTER  :: basis_set_orb, basis_set_ri
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: ab_list, ac_list
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
 
      NULLIFY (basis_set_ri, basis_set_orb, ac_list, ab_list, qs_kind_set, para_env, particle_set)
 
      CALL timeset(routinen, handle)
 
!  Take what we need from the qs_env
      CALL get_qs_env(qs_env, nkind=nkind, qs_kind_set=qs_kind_set, para_env=para_env, &
                      natom=natom, particle_set=particle_set)
 
!  Build the basis set lists
      ALLOCATE (basis_set_ri(nkind))
      ALLOCATE (basis_set_orb(nkind))
      CALL basis_set_list_setup(basis_set_ri, "RI_XAS", qs_kind_set)
      CALL basis_set_list_setup(basis_set_orb, "ORB", qs_kind_set)
 
!  Get the optimized distribution 2d for these integrals (and store it in xas_tdp_env)
      CALL build_xas_tdp_ovlp_nl(ab_list, basis_set_orb, basis_set_orb, qs_env, &
                                 excited_atoms=xas_tdp_env%ex_atom_indices)
      CALL build_xas_tdp_3c_nl(ac_list, basis_set_orb, basis_set_ri, do_potential_id, &
                               qs_env, excited_atoms=xas_tdp_env%ex_atom_indices)
      CALL get_opt_3c_dist2d(xas_tdp_env%opt_dist2d_coul, ab_list, ac_list, basis_set_orb, &
                             basis_set_orb, basis_set_ri, qs_env, only_bc_same_center=.true.)
      CALL release_neighbor_list_sets(ab_list)
      CALL release_neighbor_list_sets(ac_list)
 
!  Build a neighbor list for the ab centers. Assume (aI|c) = sum_b c_bI (ab|c), with c_bI only
!  non-zero for b centered on the same atom as c => build overlap nl, but only keeping b if centered
!  on an excited atom
      CALL build_xas_tdp_ovlp_nl(ab_list, basis_set_orb, basis_set_orb, qs_env, &
                                 excited_atoms=xas_tdp_env%ex_atom_indices, &
                                 ext_dist2d=xas_tdp_env%opt_dist2d_coul)
 
!  Build a neighbor list for the ac centers. Since we later contract as (aI|c) and we assume I is
!  very localized on the same atom as c, we take a,c as neighbors if they overlap
      CALL build_xas_tdp_3c_nl(ac_list, basis_set_orb, basis_set_ri, do_potential_id, &
                               qs_env, excited_atoms=xas_tdp_env%ex_atom_indices, &
                               ext_dist2d=xas_tdp_env%opt_dist2d_coul)
 
!  Allocate, init and compute the integrals
      ALLOCATE (blk_size_orb(natom), blk_size_ri(natom))
      CALL cp_dbcsr_dist2d_to_dist(xas_tdp_env%opt_dist2d_coul, opt_dbcsr_dist)
      CALL get_particle_set(particle_set, qs_kind_set, nsgf=blk_size_orb, basis=basis_set_orb)
      CALL get_particle_set(particle_set, qs_kind_set, nsgf=blk_size_ri, basis=basis_set_ri)
      pot%potential_type = do_potential_coulomb
 
      ALLOCATE (xas_tdp_env%ri_3c_coul)
      CALL create_pqx_tensor(xas_tdp_env%ri_3c_coul, ab_list, ac_list, opt_dbcsr_dist, blk_size_orb, &
                             blk_size_orb, blk_size_ri, only_bc_same_center=.true.)
      CALL fill_pqx_tensor(xas_tdp_env%ri_3c_coul, ab_list, ac_list, basis_set_orb, basis_set_orb, &
                           basis_set_ri, pot, qs_env, only_bc_same_center=.true.)
 
! Clean-up
      CALL release_neighbor_list_sets(ab_list)
      CALL release_neighbor_list_sets(ac_list)
      CALL dbcsr_distribution_release(opt_dbcsr_dist)
      DEALLOCATE (basis_set_ri, basis_set_orb)
 
      !not strictly necessary but avoid having any load unbalance here being reported in the
      !timings for other routines
      CALL para_env%sync()
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_ri_3c_coulomb
 
! **************************************************************************************************
!> \brief Computes the two-center Exchange integral needed for the RI in kernel calculation. Stores
!>        the integrals in the xas_tdp_env as global (small) arrays. Does that for a given excited
!>        kind. The quantity stored is M^-1 (P|Q) M^-1, where M is the RI metric. If the metric is
!>        the same as the exchange potential, then we end up with the V-approximation (P|Q)^-1
!>        By default (if no metric), the ri_m_potential is a copy of the x_potential
!> \param ex_kind ...
!> \param xas_tdp_env ...
!> \param xas_tdp_control ...
!> \param qs_env ...
!> \note Computes all these integrals in non-PBCs as we assume that the range is short enough that
!>       atoms do not exchange with their periodic images
! **************************************************************************************************
   SUBROUTINE compute_ri_exchange2_int(ex_kind, xas_tdp_env, xas_tdp_control, qs_env)
 
      INTEGER, INTENT(IN)                                :: ex_kind
      TYPE(xas_tdp_env_type), POINTER                    :: xas_tdp_env
      TYPE(xas_tdp_control_type), POINTER                :: xas_tdp_control
      TYPE(qs_environment_type), POINTER                 :: qs_env
 
      INTEGER                                            :: egfp, egfq, maxl, ncop, ncoq, nset, &
                                                            nsgf, pset, qset, sgfp, sgfq
      INTEGER, DIMENSION(:), POINTER                     :: lmax, lmin, npgf_set, nsgf_set
      INTEGER, DIMENSION(:, :), POINTER                  :: first_sgf
      REAL(dp)                                           :: r(3)
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: metric, pq, work
      REAL(dp), DIMENSION(:, :), POINTER                 :: rpgf, sphi, zet
      TYPE(coulomb_integral_context_type)                :: integral_context
      TYPE(gto_basis_set_type), POINTER                  :: ri_basis
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
 
      NULLIFY (ri_basis, qs_kind_set, para_env, lmin, lmax, npgf_set, zet, rpgf, first_sgf)
      NULLIFY (sphi, nsgf_set)
 
!  Initialization
      CALL get_qs_env(qs_env, qs_kind_set=qs_kind_set, para_env=para_env)
      IF (ASSOCIATED(xas_tdp_env%ri_inv_ex)) THEN
         DEALLOCATE (xas_tdp_env%ri_inv_ex)
      END IF
 
!  Get the RI basis of interest and its quantum numbers
      CALL get_qs_kind(qs_kind_set(ex_kind), basis_set=ri_basis, basis_type="RI_XAS")
      CALL get_gto_basis_set(ri_basis, nsgf=nsgf, maxl=maxl, npgf=npgf_set, lmin=lmin, &
                             lmax=lmax, zet=zet, pgf_radius=rpgf, first_sgf=first_sgf, &
                             nsgf_set=nsgf_set, sphi=sphi, nset=nset)
      ALLOCATE (metric(nsgf, nsgf))
      metric = 0.0_dp
 
      r = 0.0_dp
 
      !init the libint 2-center object
      CALL coulomb_integral_init(integral_context, max_am_2c=maxl, &
                                 derivative_order=0, library=library_libint)
 
      CALL r12_fundamentals_init(2*maxl + 1, xas_tdp_control%ri_m_potential, para_env)
 
!  Compute the RI metric
      DO pset = 1, nset
         ncop = npgf_set(pset)*ncoset(lmax(pset))
         sgfp = first_sgf(1, pset)
         egfp = sgfp + nsgf_set(pset) - 1
 
         DO qset = 1, nset
            ncoq = npgf_set(qset)*ncoset(lmax(qset))
            sgfq = first_sgf(1, qset)
            egfq = sgfq + nsgf_set(qset) - 1
 
            ALLOCATE (work(ncop, ncoq))
            work = 0.0_dp
 
            CALL compute_coulomb_2c(integral_context, lmin(pset), lmax(pset), lmin(qset), lmax(qset), &
                                   npgf_set(pset), npgf_set(qset), zet(:, pset), zet(:, qset), rpgf(:, pset), rpgf(:, qset), r, r, &
                                    work, xas_tdp_control%ri_m_potential, dab=0.0_dp)
 
            CALL ab_contract(metric(sgfp:egfp, sgfq:egfq), work, sphi(:, sgfp:), sphi(:, sgfq:), &
                             ncop, ncoq, nsgf_set(pset), nsgf_set(qset))
 
            DEALLOCATE (work)
         END DO !qset
      END DO !pset
 
!  Inverting (to M^-1)
      CALL invmat_symm(metric)
 
      IF (.NOT. xas_tdp_control%do_ri_metric) THEN
 
         !If no metric, then x_pot = ri_m_pot and (P|Q)^-1 = M^-1 (V-approximation)
         ALLOCATE (xas_tdp_env%ri_inv_ex(nsgf, nsgf))
         xas_tdp_env%ri_inv_ex(:, :) = metric(:, :)
         CALL coulomb_integral_cleanup(integral_context)
         RETURN
 
      END IF
 
      CALL r12_fundamentals_init(2*maxl + 1, xas_tdp_control%x_potential, para_env)
 
!  Compute the proper exchange 2-center
      ALLOCATE (pq(nsgf, nsgf))
      pq = 0.0_dp
 
      DO pset = 1, nset
         ncop = npgf_set(pset)*ncoset(lmax(pset))
         sgfp = first_sgf(1, pset)
         egfp = sgfp + nsgf_set(pset) - 1
 
         DO qset = 1, nset
            ncoq = npgf_set(qset)*ncoset(lmax(qset))
            sgfq = first_sgf(1, qset)
            egfq = sgfq + nsgf_set(qset) - 1
 
            ALLOCATE (work(ncop, ncoq))
            work = 0.0_dp
 
            CALL compute_coulomb_2c(integral_context, lmin(pset), lmax(pset), lmin(qset), lmax(qset), &
                                   npgf_set(pset), npgf_set(qset), zet(:, pset), zet(:, qset), rpgf(:, pset), rpgf(:, qset), r, r, &
                                    work, xas_tdp_control%x_potential, dab=0.0_dp)
 
            CALL ab_contract(pq(sgfp:egfp, sgfq:egfq), work, sphi(:, sgfp:), sphi(:, sgfq:), &
                             ncop, ncoq, nsgf_set(pset), nsgf_set(qset))
 
            DEALLOCATE (work)
         END DO !qset
      END DO !pset
 
!  Compute and store M^-1 (P|Q) M^-1
      ALLOCATE (xas_tdp_env%ri_inv_ex(nsgf, nsgf))
      xas_tdp_env%ri_inv_ex = 0.0_dp
 
      CALL dgemm('N', 'N', nsgf, nsgf, nsgf, 1.0_dp, metric, nsgf, pq, nsgf, &
                 0.0_dp, xas_tdp_env%ri_inv_ex, nsgf)
      CALL dgemm('N', 'N', nsgf, nsgf, nsgf, 1.0_dp, xas_tdp_env%ri_inv_ex, nsgf, metric, nsgf, &
                 0.0_dp, pq, nsgf)
      xas_tdp_env%ri_inv_ex(:, :) = pq(:, :)
 
      CALL coulomb_integral_cleanup(integral_context)
 
   END SUBROUTINE compute_ri_exchange2_int
 
! **************************************************************************************************
!> \brief Computes the two-center Coulomb integral needed for the RI in kernel calculation. Stores
!>        the integrals (P|Q)^-1 in the xas_tdp_env as global (small) arrays. Does that for a given
!>        excited kind
!> \param ex_kind ...
!> \param xas_tdp_env ...
!> \param xas_tdp_control ...
!> \param qs_env ...
! **************************************************************************************************
   SUBROUTINE compute_ri_coulomb2_int(ex_kind, xas_tdp_env, xas_tdp_control, qs_env)
 
      INTEGER, INTENT(IN)                                :: ex_kind
      TYPE(xas_tdp_env_type), POINTER                    :: xas_tdp_env
      TYPE(xas_tdp_control_type), POINTER                :: xas_tdp_control
      TYPE(qs_environment_type), POINTER                 :: qs_env
 
      INTEGER                                            :: nsgf
      TYPE(gto_basis_set_type), POINTER                  :: ri_basis
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
 
      NULLIFY (ri_basis, qs_kind_set)
 
!  Initialization
      CALL get_qs_env(qs_env, qs_kind_set=qs_kind_set)
      IF (ASSOCIATED(xas_tdp_env%ri_inv_coul)) THEN
         DEALLOCATE (xas_tdp_env%ri_inv_coul)
      END IF
 
!  Get the RI basis of interest and its quantum numbers
      CALL get_qs_kind(qs_kind_set(ex_kind), basis_set=ri_basis, basis_type="RI_XAS")
      CALL get_gto_basis_set(ri_basis, nsgf=nsgf)
      ALLOCATE (xas_tdp_env%ri_inv_coul(nsgf, nsgf))
      xas_tdp_env%ri_inv_coul = 0.0_dp
 
      IF (.NOT. xas_tdp_control%is_periodic) THEN
         CALL int_operators_r12_ab_os(r12_operator=operator_coulomb, vab=xas_tdp_env%ri_inv_coul, &
                                      rab=[0.0_dp, 0.0_dp, 0.0_dp], fba=ri_basis, fbb=ri_basis, &
                                      calculate_forces=.false.)
         cpassert(ASSOCIATED(xas_tdp_control))
      ELSE
         CALL periodic_ri_coulomb2(xas_tdp_env%ri_inv_coul, ri_basis, qs_env)
      END IF
 
!  Inverting
      CALL invmat_symm(xas_tdp_env%ri_inv_coul)
 
   END SUBROUTINE compute_ri_coulomb2_int
 
! **************************************************************************************************
!> \brief Computes the two-center inverse coulomb integral in the case of PBCs
!> \param ri_coul2 ...
!> \param ri_basis ...
!> \param qs_env ...
! **************************************************************************************************
   SUBROUTINE periodic_ri_coulomb2(ri_coul2, ri_basis, qs_env)
 
      REAL(dp), DIMENSION(:, :), INTENT(INOUT)           :: ri_coul2
      TYPE(gto_basis_set_type), POINTER                  :: ri_basis
      TYPE(qs_environment_type), POINTER                 :: qs_env
 
      INTEGER                                            :: maxco, ncop, ncoq, nset, op, oq, ppgf, &
                                                            pset, qpgf, qset, sgfp, sgfq
      INTEGER, DIMENSION(:), POINTER                     :: lmax, lmin, npgf, nsgf
      INTEGER, DIMENSION(:, :), POINTER                  :: first_sgf
      REAL(dp)                                           :: r(3)
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: hpq
      REAL(dp), DIMENSION(:, :), POINTER                 :: sphi, zet
      TYPE(cell_type), POINTER                           :: cell
      TYPE(cp_eri_mme_param)                             :: mme_param
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
 
      NULLIFY (cell, qs_kind_set, lmin, lmax, nsgf, npgf, zet, sphi, first_sgf)
 
      ! Use eri_mme for this. Don't want to annoy the user with a full input section just for this
      ! tiny bit => initialize our own eri_mme section with the defaults
 
      CALL get_qs_env(qs_env, cell=cell, qs_kind_set=qs_kind_set, para_env=para_env)
 
      CALL eri_mme_init(mme_param%par, n_minimax=20, cutoff=300.0_dp, do_calib_cutoff=.true., &
                        cutoff_min=10.0_dp, cutoff_max=10000.0_dp, cutoff_eps=0.01_dp, &
                        cutoff_delta=0.9_dp, sum_precision=1.0e-12_dp, debug=.false., &
                        debug_delta=1.0e-6_dp, debug_nsum=1000000, unit_nr=0, print_calib=.false., &
                        do_error_est=.false.)
      mme_param%do_calib = .true.
 
      CALL cp_eri_mme_set_params(mme_param, cell, qs_kind_set, basis_type_1="RI_XAS", para_env=para_env)
 
      CALL get_gto_basis_set(ri_basis, lmax=lmax, npgf=npgf, zet=zet, lmin=lmin, nset=nset, &
                             nsgf_set=nsgf, sphi=sphi, first_sgf=first_sgf, maxco=maxco)
 
      r = 0.0_dp
      ALLOCATE (hpq(nset*maxco, nset*maxco))
      hpq = 0.0_dp
 
      DO pset = 1, nset
         ncop = npgf(pset)*ncoset(lmax(pset))
         sgfp = first_sgf(1, pset)
 
         DO qset = 1, nset
            ncoq = npgf(qset)*ncoset(lmax(qset))
            sgfq = first_sgf(1, qset)
 
            DO ppgf = 1, npgf(pset)
               op = (pset - 1)*maxco + (ppgf - 1)*ncoset(lmax(pset))
               DO qpgf = 1, npgf(qset)
                  oq = (qset - 1)*maxco + (qpgf - 1)*ncoset(lmax(qset))
 
                  CALL eri_mme_2c_integrate(mme_param%par, lmin(pset), lmax(pset), lmin(qset), &
                                            lmax(qset), zet(ppgf, pset), zet(qpgf, qset), r, hpq, &
                                            op, oq)
 
               END DO !qpgf
            END DO ! ppgf
 
            !contraction into sgfs
            op = (pset - 1)*maxco + 1
            oq = (qset - 1)*maxco + 1
 
            CALL ab_contract(ri_coul2(sgfp:sgfp + nsgf(pset) - 1, sgfq:sgfq + nsgf(qset) - 1), &
                             hpq(op:op + ncop - 1, oq:oq + ncoq - 1), sphi(:, sgfp:), sphi(:, sgfq:), &
                             ncop, ncoq, nsgf(pset), nsgf(qset))
 
         END DO !qset
      END DO !pset
 
      !celan-up
      CALL eri_mme_release(mme_param%par)
 
   END SUBROUTINE periodic_ri_coulomb2
 
END MODULE xas_tdp_integrals