gw_auto_ri_generate_RI_basis.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 Automatic RI basis set optimization for molecular GW.
!> \par History
!>      09.2026 created [Jan Wilhelm]
! **************************************************************************************************
MODULE gw_auto_ri_generate_ri_basis
   USE ai_contraction_sphi,             ONLY: ab_contract,&
                                              abc_contract
   USE atomic_kind_types,               ONLY: get_atomic_kind_set
   USE basis_set_types,                 ONLY: get_gto_basis_set,&
                                              gto_basis_set_type
   USE constants_operator,              ONLY: operator_coulomb,&
                                              operator_truncated
   USE coulomb_integral_interface,      ONLY: compute_coulomb_2c,&
                                              compute_coulomb_3c,&
                                              coulomb_integral_cleanup,&
                                              coulomb_integral_context_type,&
                                              coulomb_integral_init,&
                                              coulomb_operator_type
   USE cp_blacs_env,                    ONLY: cp_blacs_env_type
   USE cp_dbcsr_api,                    ONLY: &
        dbcsr_create, dbcsr_distribution_get, dbcsr_distribution_new, dbcsr_distribution_release, &
        dbcsr_distribution_type, dbcsr_finalize, dbcsr_put_block, dbcsr_release, dbcsr_type, &
        dbcsr_type_symmetric
   USE cp_dbcsr_operations,             ONLY: copy_dbcsr_to_fm
   USE cp_files,                        ONLY: close_file,&
                                              discover_file,&
                                              open_file
   USE cp_fm_basic_linalg,              ONLY: cp_fm_column_scale,&
                                              cp_fm_syrk,&
                                              cp_fm_uplo_to_full
   USE cp_fm_diag,                      ONLY: cp_fm_geeig
   USE cp_fm_struct,                    ONLY: cp_fm_struct_create,&
                                              cp_fm_struct_release,&
                                              cp_fm_struct_type
   USE cp_fm_types,                     ONLY: cp_fm_create,&
                                              cp_fm_get_info,&
                                              cp_fm_get_submatrix,&
                                              cp_fm_release,&
                                              cp_fm_type
   USE generic_os_integrals,            ONLY: int_operators_r12_ab_os
   USE gw_auto_ri_types,                ONLY: auto_ri_type
   USE gw_integrals,                    ONLY: build_3c_integral_block_ctx,&
                                              gw_3c_ctx_create,&
                                              gw_3c_ctx_release,&
                                              gw_3c_ctx_type,&
                                              gw_3c_ws_create,&
                                              gw_3c_ws_release,&
                                              gw_3c_ws_type
   USE ieee_arithmetic,                 ONLY: ieee_is_finite
   USE input_constants,                 ONLY: do_potential_coulomb,&
                                              do_potential_truncated,&
                                              small_cell_full_kp
   USE integral_library_types,          ONLY: active_integral_library,&
                                              library_libint
   USE kinds,                           ONLY: default_path_length,&
                                              dp
   USE libint_2c_3c,                    ONLY: libint_potential_type
   USE local_gemm_api,                  ONLY: local_gemm_pu_host,&
                                              local_gemm_ctxt_type
   USE machine,                         ONLY: m_walltime
   USE mathlib,                         ONLY: diamat_all
   USE message_passing,                 ONLY: mp_para_env_type
   USE orbital_pointers,                ONLY: init_orbital_pointers,&
                                              ncoset
   USE orbital_transformation_matrices, ONLY: init_spherical_harmonics
   USE post_scf_bandstructure_types,    ONLY: post_scf_bandstructure_type
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type
   USE t_c_g0,                          ONLY: get_lmax_init,&
                                              init_t_c_g0 => init
   USE util,                            ONLY: sort
#include "base/base_uses.f90"
 
   IMPLICIT NONE
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'gw_auto_ri_generate_RI_basis'
 
   ! Numerical cutoffs for the reference M_PQ eigenspace and final M_pq pseudoinverse.
   REAL(KIND=dp), PARAMETER :: reference_m_eigenval_threshold = 1.0e-10_dp
   REAL(KIND=dp), PARAMETER :: optimized_m_eigenval_threshold = 1.0e-6_dp
   ! Discard null target directions and reject non-positive-semidefinite M_PQ or B_PQ spectra.
   REAL(KIND=dp), PARAMETER :: relative_target_eigenval_threshold = 1.0e-12_dp
   REAL(KIND=dp), PARAMETER :: ab_absolute_target_eigenval_threshold = 1.0e-12_dp
   REAL(KIND=dp), PARAMETER :: negative_eigenval_tolerance = 1.0e-10_dp
 
   PUBLIC :: generate_auto_ri_basis
 
   TYPE ab_block_list_type
      INTEGER, ALLOCATABLE                               :: block(:)
      INTEGER, ALLOCATABLE                               :: first_p_AB(:)
      INTEGER, ALLOCATABLE                               :: first_p_atom(:)
      INTEGER, ALLOCATABLE                               :: size_opt_RI(:)
   END TYPE ab_block_list_type
 
   TYPE optimized_atom_u_type
      INTEGER, ALLOCATABLE                               :: ref_atom(:)
      INTEGER, ALLOCATABLE                               :: first_P(:)
      REAL(KIND=dp), ALLOCATABLE                         :: u_pp(:, :)
   END TYPE optimized_atom_u_type
 
CONTAINS
 
!**************************************************************************************************!
! The full derivation of the AUTO_RI algorithm will be published separately.                       !
!                                                                                                  !
! The equations below define the notation used throughout this module.                             !
!                                                                                                  !
! The RI scalar product uses the truncated Coulomb operator:                                       !
!                                                                                                  !
! (f|g) = ∫dr∫dr' f(r) V_rc(r-r') g(r').                                                   (1)     !
!                                                                                                  !
! For the RI basis-set optimization, the orbital-product density formed by                         !
! AOs on atoms A and B is                                                                          !
!                                                                                                  !
! ρ_mn^AB(r) = Σ_{μ∈A} Σ_{ν∈B} C_μm C_νn Φ_μ(r) Φ_ν(r).                                    (2)     !
!                                                                                                  !
! The atom-local correlation target B_PQ^{A,C} mimics contributions to the density response χ,     !
! and the exchange target B_PQ^{A,X} mimics contributions to the exchange self-energy:             !
!                                                                                                  !
! B_PQ^{A,C} = (1/N_C^A) Σ_ia^oe w_ia (φ_P^A|ρ_ia^AA)(ρ_ia^AA|φ_Q^A).                      (3)     !
!                                                                                                  !
! B_PQ^{A,X} = (1/N_X^A) Σ_{n∈W_occ} Σ_i^occ (φ_P^A|ρ_ni^AA)(ρ_ni^AA|φ_Q^A).                 (4)     !
!                                                                                                  !
! W_occ = [ε_HOMO-E_win, ε_LUMO+E_win], where E_win is OCC_EMPTY_FRONTIER_ORBITAL_WINDOW.          !
!                                                                                                  !
! Here w_ia gives larger weight to frontier-orbital products close to the HOMO and LUMO:           !
!                                                                                                  !
! w_ia = 1/(|ε_i-E_F| |ε_a-E_F|),   E_F = (ε_HOMO+ε_LUMO)/2.                               (5)     !
!                                                                                                  !
! Diagonalizing B^A = B^{A,C}+B^{A,X} identifies the most important atom-local RI functions        !
! that we use in the optimized RI basis set. The contraction coefficients U_Pp^A [see Eq. (7)]     !
! used for optimized RI basis functions are obtained from the eigenvalue problem                   !
!                                                                                                  !
! B^A U^A = M^A U^A λ^A,   Σ_Q B_PQ^A U_Qp^A = λ_p^A Σ_Q M_PQ^A U_Qp^A.                    (6)     !
!                                                                                                  !
! The eigenvectors with the largest eigenvalues are the contracted RI functions φ_p^A(r)           !
! used for the optimized RI basis set:                                                             !
!                                                                                                  !
! φ_p^A(r) = Σ_P U_Pp^A φ_P^A(r).                                                          (7)     !
!                                                                                                  !
! We also use superpositions of Gaussians on neighboring atoms in the RI basis. Before optimizing  !
! the expansion coefficients, the AA from Eq. (7) used in AB contractions are projected out:       !
!                                                                                                  !
! span{φ_P^{AB,⊥}} = span{φ_P^A,φ_Q^B} \ span{φ_p^A,φ_q^B}.                                (8)     !
!                                                                                                  !
! For obtaining the AB contraction coefficients, we use the symmetric off-diagonal density:        !
!                                                                                                  !
! ρ̅_mn^AB(r) = ρ_mn^AB(r)+ρ_mn^BA(r).                                                      (9)     !
!                                                                                                  !
! To identify the most important functions for a neighboring pair A-B, define                      !
! correlation and exchange targets in this projected basis:                                        !
!                                                                                                  !
! B_PQ^{AB,C} = (1/N_C^{AB}) Σ_ia^oe w_ia (φ_P^{AB,⊥}|ρ̅_ia^AB)(ρ̅_ia^AB|φ_Q^{AB,⊥}).       (10)     !
!                                                                                                  !
! B_PQ^{AB,X} = (1/N_X^{AB}) Σ_{n∈W_occ} Σ_i^occ (φ_P^{AB,⊥}|ρ̅_ni^AB)(ρ̅_ni^AB|φ_Q^{AB,⊥}).  (11)     !
!                                                                                                  !
! Diagonalizing their sum gives the AB-neighbor-pair contraction coefficients:                     !
!                                                                                                  !
! B^AB U^AB = U^AB λ^AB,   B^AB = B^{AB,C}+B^{AB,X}.                                      (12)     !
!                                                                                                  !
! The eigenvectors with the largest eigenvalues are used as AB contractions in the optimized       !
! RI basis set:                                                                                    !
!                                                                                                  !
! φ_p^AB(r) = Σ_P U_Pp^AB φ_P^{AB,⊥}(r).                                                  (13)     !
!                                                                                                  !
! Eq. (13) means that the RI basis is not an atom-centered basis any more, but can be delocalized  !
! over two neighboring atoms AB.                                                                   !
!                                                                                                  !
! AUTO_RI uses the same number of AA and AB contractions in the final optimized RI basis:          !
!                                                                                                  !
! Σ_A N_A = Σ_{A>B} N_AB = N_RI/2.                                                        (14)     !
!                                                                                                  !
!**************************************************************************************************!
 
! **************************************************************************************************
!> \brief Executes the AUTO_RI algorithm defined by Eqs. (1)-(14).
!> \param qs_env ...
!> \param bs_env ...
! **************************************************************************************************
   SUBROUTINE generate_auto_ri_basis(qs_env, bs_env)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'generate_auto_ri_basis'
 
      INTEGER                                            :: handle
      REAL(kind=dp)                                      :: time_start
      REAL(kind=dp), ALLOCATABLE                         :: u_pp_a(:, :, :)
 
      CALL timeset(routinen, handle)
 
      time_start = m_walltime()
 
      ! Compute AA target matrices B^{AA,C}, B^{AA,X} and contraction U_Pp^A in Eqs. (3)-(7).
      CALL compute_aa_optimized_ri_basis(bs_env, u_pp_a)
 
      ! Compute AB target matrices B^{AB,C}, B^{AB,X} and contraction U_Pp^AB in Eqs. (8)-(13).
      CALL compute_ab_optimized_ri_basis(qs_env, bs_env, u_pp_a)
      DEALLOCATE (u_pp_a)
 
      ! Compute and invert M_pq = (φ_p|φ_q) of the optimized RI basis in Eqs. (7) and (13).
      CALL compute_m_pq_inv_and_v_pq(qs_env, bs_env)
 
      ! Print the optimized RI basis dimension and total optimization time.
      CALL print_auto_ri_summary(bs_env, time_start)
 
      CALL timestop(handle)
 
   END SUBROUTINE generate_auto_ri_basis
 
! **************************************************************************************************
!> \brief Computes the atom-local contraction coefficients U^A from Eqs. (3)-(7):
!>
!>          B_PQ^{A,C} = (1/N_C^A) Σ_ia^oe w_ia (φ_P^A|ρ_ia^AA)(ρ_ia^AA|φ_Q^A),         (3)
!>
!>          B_PQ^{A,X} = (1/N_X^A) Σ_{n∈W_occ} Σ_i^occ (φ_P^A|ρ_ni^AA)(ρ_ni^AA|φ_Q^A),    (4)
!>
!>          w_ia = 1/(|ε_i-E_F| |ε_a-E_F|),  E_F = (ε_HOMO+ε_LUMO)/2,                   (5)
!>
!>          B^A U^A = M^A U^A λ^A,                                                      (6)
!>
!>          Contracted RI basis functions for optimized RI basis:
!>
!>          φ_p^A(r) = Σ_P U_Pp^A φ_P^A(r).                                             (7)
!>
!> \param bs_env ...
!> \param U_Pp_A temporary atom-local contraction coefficients used to construct the AB space
! **************************************************************************************************
   SUBROUTINE compute_aa_optimized_ri_basis(bs_env, U_Pp_A)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      REAL(kind=dp), ALLOCATABLE, INTENT(OUT)            :: u_pp_a(:, :, :)
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_AA_optimized_RI_basis'
 
      INTEGER :: handle, iatom, ikind, max_nao_kind, max_nri_kind, nao_kind, natom, nkind, &
         nri_kind, ri_size_opt_aa_requested, root_rank
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: kind_of_atom, sizes_opt_ri_aa
      REAL(kind=dp)                                      :: target_norm
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:)           :: eigenval
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: b_pq, eigenval_aa, m_pq, u_pp_a_atom
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: m_pq_by_kind
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :, :)  :: int_3c
      TYPE(mp_para_env_type), POINTER                    :: para_env
 
      CALL timeset(routinen, handle)
 
      ! Determine reference dimensions and the number of AA RI functions (i.e. number of
      ! eigenvectors U_Pp^A used for RI-basis contractions).
      CALL initialize_aa_optimization(bs_env, para_env, kind_of_atom, natom, nkind, &
                                      max_nri_kind, ri_size_opt_aa_requested)
      ALLOCATE (eigenval_aa(max_nri_kind, natom), source=0.0_dp)
      ALLOCATE (sizes_opt_ri_aa(natom))
 
      CALL initialize_auto_ri_integrals(bs_env)
      max_nao_kind = maxval(bs_env%sizes_AO)
      ALLOCATE (u_pp_a(max_nri_kind, max_nri_kind, natom), &
                m_pq_by_kind(max_nri_kind, max_nri_kind, nkind), source=0.0_dp)
      ALLOCATE (int_3c(max_nao_kind, max_nao_kind, max_nri_kind, nkind), source=0.0_dp)
 
      ! The few atom-kind integral blocks are inexpensive; compute them on every rank.
      DO ikind = 1, nkind
         nao_kind = bs_env%basis_set_AO(ikind)%gto_basis_set%nsgf
         nri_kind = bs_env%basis_set_RI(ikind)%gto_basis_set%nsgf
         CALL compute_int_3c( &
            ao_basis=bs_env%basis_set_AO(ikind)%gto_basis_set, &
            ri_basis=bs_env%basis_set_RI(ikind)%gto_basis_set, &
            potential=bs_env%ri_metric, m_pq=m_pq, &
            int_3c=int_3c(1:nao_kind, 1:nao_kind, 1:nri_kind, ikind))
         CALL add_ri_metric_regularization(bs_env%basis_set_RI(ikind)%gto_basis_set, &
                                           bs_env%regularization_RI, m_pq)
         m_pq_by_kind(1:nri_kind, 1:nri_kind, ikind) = m_pq
         DEALLOCATE (m_pq)
      END DO
 
      bs_env%auto_ri%sizes_opt_RI = 0
      DO iatom = 1, natom
         root_rank = modulo(iatom - 1, para_env%num_pe)
         ikind = kind_of_atom(iatom)
         nri_kind = bs_env%basis_set_RI(ikind)%gto_basis_set%nsgf
         nao_kind = bs_env%sizes_AO(iatom)
         ALLOCATE (m_pq(nri_kind, nri_kind))
         m_pq(:, :) = m_pq_by_kind(1:nri_kind, 1:nri_kind, ikind)
         ! Eqs. (3)-(5): form B^{A,C}, B^{A,X}, and their weighted sum B^A.
         CALL compute_b_pq_aa(bs_env, iatom, root_rank, &
                              int_3c(1:nao_kind, 1:nao_kind, 1:nri_kind, ikind), &
                              m_pq, b_pq, target_norm)
         DEALLOCATE (m_pq)
         IF (root_rank /= para_env%mepos) THEN
            DEALLOCATE (b_pq)
            cycle
         END IF
         ! Eq. (6): solve B^A U^A = M^A U^A λ^A.
         CALL solve_aa_eigenproblem(m_pq_by_kind(1:nri_kind, 1:nri_kind, ikind), &
                                    b_pq, target_norm, u_pp_a_atom, eigenval)
         eigenval_aa(1:SIZE(u_pp_a_atom, 2), iatom) = eigenval(1:SIZE(u_pp_a_atom, 2))
         bs_env%auto_ri%sizes_opt_RI(iatom) = SIZE(u_pp_a_atom, 2)
         u_pp_a(1:SIZE(u_pp_a_atom, 1), 1:SIZE(u_pp_a_atom, 2), iatom) = u_pp_a_atom
         DEALLOCATE (u_pp_a_atom, eigenval, b_pq)
      END DO
      CALL para_env%sum(bs_env%auto_ri%sizes_opt_RI)
      CALL para_env%sum(u_pp_a)
      CALL para_env%sum(eigenval_aa)
      CALL select_ri_sizes(eigenval_aa, ri_size_opt_aa_requested, 1, sizes_opt_ri_aa)
      bs_env%auto_ri%sizes_opt_RI(:) = sizes_opt_ri_aa
      DO iatom = 1, natom
         u_pp_a(:, sizes_opt_ri_aa(iatom) + 1:, iatom) = 0.0_dp
      END DO
      DEALLOCATE (int_3c, m_pq_by_kind)
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_aa_optimized_ri_basis
 
! **************************************************************************************************
!> \brief Initializes dimensions and atom-kind mappings for the AA optimization.
!>
!>        This routine does not evaluate an equation of the objective. It maps every atom
!>        to its reference RI basis, allocates the persistent per-atom arrays, and determines
!>        the AA part of the requested basis size from
!>
!>          Σ_A N_A = N_RI/2.                                             (14)
!>
!>        At least one AA contraction is assigned to every atom.
!> \param bs_env ...
!> \param para_env ...
!> \param kind_of_atom ...
!> \param natom ...
!> \param nkind ...
!> \param max_nRI_kind ...
!> \param RI_size_opt_AA_requested ...
! **************************************************************************************************
   SUBROUTINE initialize_aa_optimization(bs_env, para_env, kind_of_atom, &
                                         natom, nkind, max_nRI_kind, &
                                         RI_size_opt_AA_requested)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(mp_para_env_type), POINTER                    :: para_env
      INTEGER, ALLOCATABLE, INTENT(OUT)                  :: kind_of_atom(:)
      INTEGER, INTENT(OUT)                               :: natom, nkind, max_nri_kind, &
                                                            ri_size_opt_aa_requested
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'initialize_AA_optimization'
 
      INTEGER                                            :: handle, ikind, ri_size_requested
      INTEGER, ALLOCATABLE                               :: sizes_ref_ri_by_kind(:)
 
      CALL timeset(routinen, handle)
 
      para_env => bs_env%para_env
      natom = bs_env%n_atom
      nkind = SIZE(bs_env%basis_set_RI)
      cpassert(ASSOCIATED(bs_env%ri_rs%atomic_kind_set))
      cpassert(ASSOCIATED(bs_env%ri_rs%particle_set))
      cpassert(ASSOCIATED(para_env) .AND. natom > 0 .AND. nkind > 0)
      cpassert(ALLOCATED(bs_env%basis_set_AO))
      cpassert(ALLOCATED(bs_env%basis_set_RI))
      cpassert(SIZE(bs_env%basis_set_AO) == nkind)
      cpassert(SIZE(bs_env%basis_set_RI) == nkind)
      cpassert(bs_env%do_gw_ri_rs)
      cpassert(.NOT. bs_env%do_periodic)
      cpassert(bs_env%small_cell_full_kp_or_large_cell_Gamma /= small_cell_full_kp)
 
      CALL get_atomic_kind_set(atomic_kind_set=bs_env%ri_rs%atomic_kind_set, &
                               kind_of=kind_of_atom)
      ALLOCATE (bs_env%auto_ri%sizes_opt_RI(natom))
      ALLOCATE (sizes_ref_ri_by_kind(nkind))
      DO ikind = 1, nkind
         cpassert(ASSOCIATED(bs_env%basis_set_AO(ikind)%gto_basis_set))
         cpassert(ASSOCIATED(bs_env%basis_set_RI(ikind)%gto_basis_set))
         sizes_ref_ri_by_kind(ikind) = bs_env%basis_set_RI(ikind)%gto_basis_set%nsgf
      END DO
      max_nri_kind = maxval(sizes_ref_ri_by_kind)
      ri_size_requested = nint( &
                          min(bs_env%auto_ri%ri_ao_ratio, &
                              REAL(sum(sizes_ref_ri_by_kind(kind_of_atom)), dp)/ &
                              REAL(SUM(bs_env%sizes_AO), dp))*REAL(SUM(bs_env%sizes_AO), dp))
      ri_size_opt_aa_requested = max(natom, nint(0.5_dp*real(ri_size_requested, dp)))
      DEALLOCATE (sizes_ref_ri_by_kind)
 
      CALL timestop(handle)
 
   END SUBROUTINE initialize_aa_optimization
 
! **************************************************************************************************
!> \brief Sets the RI-RS block sizes to the optimized RI basis dimensions.
!> \param bs_env ...
! **************************************************************************************************
   SUBROUTINE set_optimized_ri_basis_sizes(bs_env)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'set_optimized_RI_basis_sizes'
 
      INTEGER                                            :: handle
 
      CALL timeset(routinen, handle)
 
      bs_env%sizes_RI(:) = bs_env%auto_ri%sizes_opt_RI
      bs_env%auto_ri%ready = .true.
 
      CALL timestop(handle)
 
   END SUBROUTINE set_optimized_ri_basis_sizes
 
! **************************************************************************************************
!> \brief Prints the optimized RI basis size and total optimization time.
!> \param bs_env ...
!> \param time_start ...
! **************************************************************************************************
   SUBROUTINE print_auto_ri_summary(bs_env, time_start)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      REAL(kind=dp), INTENT(IN)                          :: time_start
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'print_auto_ri_summary'
 
      INTEGER                                            :: handle, ri_size_optimized, unit_nr
      REAL(kind=dp)                                      :: elapsed
 
      CALL timeset(routinen, handle)
 
      ri_size_optimized = sum(bs_env%auto_ri%sizes_opt_RI)
      elapsed = m_walltime() - time_start
      CALL bs_env%para_env%max(elapsed)
      unit_nr = bs_env%unit_nr
      IF (unit_nr > 0) THEN
         WRITE (unit_nr, '(A)') ' '
         WRITE (unit_nr, '(T2,A,T74,I9)') &
            'AUTO_RI| Number of automatic RI functions for χ, ε, W:', ri_size_optimized
         WRITE (unit_nr, '(T2,A,T71,F8.1,A)') &
            'AUTO_RI| Total RI basis optimization time:', elapsed, ' s'
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE print_auto_ri_summary
 
! **************************************************************************************************
!> \brief Adds REGULARIZATION_RI to a two-center RI metric matrix (P|Q).
!> \param basis ...
!> \param regularization ...
!> \param M_PQ reference RI metric M_PQ
!> \param offset ...
! **************************************************************************************************
   SUBROUTINE add_ri_metric_regularization(basis, regularization, M_PQ, offset)
      USE basis_set_types, ONLY: gto_basis_set_type
      TYPE(gto_basis_set_type), INTENT(IN)               :: basis
      REAL(kind=dp), INTENT(IN)                          :: regularization
      REAL(kind=dp), DIMENSION(:, :), INTENT(INOUT)      :: m_pq
      INTEGER, INTENT(IN), OPTIONAL                      :: offset
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'add_RI_metric_regularization'
 
      INTEGER                                            :: first, handle, i, iset, nset, start
      INTEGER, DIMENSION(:), POINTER                     :: npgf, nsgf_set
      INTEGER, DIMENSION(:, :), POINTER                  :: first_sgf
      REAL(kind=dp)                                      :: shift
      REAL(kind=dp), DIMENSION(:, :), POINTER            :: zet
 
      CALL timeset(routinen, handle)
 
      start = 0
      IF (PRESENT(offset)) start = offset
      CALL get_gto_basis_set(basis, first_sgf=first_sgf, npgf=npgf, nset=nset, &
                             nsgf_set=nsgf_set, zet=zet)
      cpassert(start >= 0 .AND. start + basis%nsgf <= SIZE(m_pq, 1))
      cpassert(start + basis%nsgf <= SIZE(m_pq, 2))
      DO iset = 1, nset
         shift = regularization*max(1.0_dp, 1.0_dp/minval(zet(1:npgf(iset), iset)))
         first = start + first_sgf(1, iset)
         DO i = first, first + nsgf_set(iset) - 1
            m_pq(i, i) = m_pq(i, i) + shift
         END DO
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE add_ri_metric_regularization
 
! **************************************************************************************************
!> \brief Copies atom-local diagonal blocks M^AA from a distributed metric M.
!> \param M_pq_local locally stored entries of the optimized RI metric M_pq
!> \param M_pq_AA_local locally stored entries of the atom-block-diagonal metric
!> \param row_indices global indices of the locally owned rows
!> \param col_indices global indices of the locally owned columns
!> \param sizes_RI optimized RI functions in each atomic matrix block
! **************************************************************************************************
   SUBROUTINE copy_m_pq_aa_blocks(M_pq_local, M_pq_AA_local, row_indices, &
                                  col_indices, sizes_RI)
      REAL(kind=dp), DIMENSION(:, :), INTENT(IN)         :: m_pq_local
      REAL(kind=dp), DIMENSION(:, :), INTENT(OUT)        :: m_pq_aa_local
      INTEGER, DIMENSION(:), INTENT(IN)                  :: row_indices, col_indices, sizes_ri
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'copy_M_pq_AA_blocks'
 
      INTEGER                                            :: handle, i, j, last_index
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: block_ends, col_atom, row_atom
 
      CALL timeset(routinen, handle)
 
      cpassert(SIZE(m_pq_local, 1) == SIZE(row_indices))
      cpassert(SIZE(m_pq_local, 2) == SIZE(col_indices))
      cpassert(all(shape(m_pq_aa_local) == shape(m_pq_local)))
      cpassert(SIZE(sizes_ri) > 0 .AND. all(sizes_ri > 0))
 
      ALLOCATE (block_ends(SIZE(sizes_ri)))
      block_ends(1) = sizes_ri(1)
      DO i = 2, SIZE(sizes_ri)
         block_ends(i) = block_ends(i - 1) + sizes_ri(i)
      END DO
      last_index = block_ends(SIZE(block_ends))
      cpassert(all(row_indices >= 1 .AND. row_indices <= last_index))
      cpassert(all(col_indices >= 1 .AND. col_indices <= last_index))
 
      ALLOCATE (row_atom(SIZE(row_indices)), col_atom(SIZE(col_indices)))
      DO i = 1, SIZE(row_indices)
         row_atom(i) = m_pq_index_atom(row_indices(i), block_ends)
      END DO
      DO j = 1, SIZE(col_indices)
         col_atom(j) = m_pq_index_atom(col_indices(j), block_ends)
      END DO
 
      m_pq_aa_local = 0.0_dp
      DO j = 1, SIZE(col_indices)
         DO i = 1, SIZE(row_indices)
            IF (row_atom(i) == col_atom(j)) THEN
               m_pq_aa_local(i, j) = m_pq_local(i, j)
            END IF
         END DO
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE copy_m_pq_aa_blocks
 
! **************************************************************************************************
!> \brief Forms f_i=1/sqrt(λ_i) for eigenvalues of the RI metric matrix M_PQ.
!> \param eigenval generalized metric eigenvalues
!> \param factors inverse-square-root weights, or zero for discarded modes
! **************************************************************************************************
   SUBROUTINE compute_optimized_m_pq_inv_factors(eigenval, factors)
      REAL(kind=dp), DIMENSION(:), INTENT(IN)            :: eigenval
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:), &
         INTENT(OUT)                                     :: factors
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_optimized_M_pq_inv_factors'
 
      INTEGER                                            :: handle, i, n_dependent
      REAL(kind=dp)                                      :: largest, threshold
 
      CALL timeset(routinen, handle)
 
      cpassert(SIZE(eigenval) > 0)
 
      largest = maxval(eigenval)
      cpassert(largest > 0.0_dp)
      threshold = optimized_m_eigenval_threshold*largest
      ALLOCATE (factors(SIZE(eigenval)))
      factors = 0.0_dp
      n_dependent = 0
      DO i = 1, SIZE(eigenval)
         IF (eigenval(i) > 0.0_dp .AND. eigenval(i) >= threshold) THEN
            factors(i) = 1.0_dp/sqrt(eigenval(i))
         ELSE
            n_dependent = n_dependent + 1
         END IF
      END DO
      cpassert(n_dependent < SIZE(eigenval))
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_optimized_m_pq_inv_factors
 
! **************************************************************************************************
!> \brief Finds the atom owning a global RI metric index.
!> \param global_index global metric index
!> \param block_ends inclusive end index of each consecutive atom block
!> \return atom index
! **************************************************************************************************
   PURE FUNCTION m_pq_index_atom(global_index, block_ends) RESULT(atom_index)
      INTEGER, INTENT(IN)                                :: global_index
      INTEGER, DIMENSION(:), INTENT(IN)                  :: block_ends
      INTEGER                                            :: atom_index
 
      INTEGER                                            :: high, low, middle
 
      low = 1
      high = SIZE(block_ends)
      DO WHILE (low < high)
         middle = (low + high)/2
         IF (global_index <= block_ends(middle)) THEN
            high = middle
         ELSE
            low = middle + 1
         END IF
      END DO
      atom_index = low
 
   END FUNCTION m_pq_index_atom
 
! **************************************************************************************************
!> \brief Computes X=U diag(1/sqrt(λ)) with X^T M X=I.
!>
!>        The columns of X span the non-null eigenspace of the symmetric
!>        two-center RI metric matrix M=U diag(λ) U^T.
!> \param M_PQ ...
!> \param X_Pi ...
!> \param rank ...
! **************************************************************************************************
   SUBROUTINE compute_reference_m_pq_inv_factor(M_PQ, X_Pi, rank)
      REAL(kind=dp), DIMENSION(:, :), INTENT(IN)         :: m_pq
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
         INTENT(OUT)                                     :: x_pi
      INTEGER, INTENT(OUT)                               :: rank
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_reference_M_PQ_inv_factor'
 
      INTEGER                                            :: column, handle, i, n
      REAL(kind=dp)                                      :: largest
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:)           :: eigenval
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: eigenvec
 
      CALL timeset(routinen, handle)
 
      n = SIZE(m_pq, 1)
      cpassert(n > 0 .AND. SIZE(m_pq, 2) == n)
      ALLOCATE (eigenvec, source=0.5_dp*(m_pq + transpose(m_pq)))
      ALLOCATE (eigenval(n))
      CALL diamat_all(eigenvec, eigenval, dac=.true.)
      largest = maxval(eigenval)
      IF (largest <= 0.0_dp) cpabort("AUTO_RI metric has no positive non-null range")
      IF (minval(eigenval) < -negative_eigenval_tolerance*largest) THEN
         cpabort("AUTO_RI metric has no positive non-null range")
      END IF
      rank = count(eigenval > reference_m_eigenval_threshold*largest)
      IF (rank == 0) cpabort("AUTO_RI metric has no positive non-null range")
      ALLOCATE (x_pi(n, rank))
      column = 0
      DO i = 1, n
         IF (eigenval(i) <= reference_m_eigenval_threshold*largest) cycle
         column = column + 1
         x_pi(:, column) = eigenvec(:, i)/sqrt(eigenval(i))
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_reference_m_pq_inv_factor
 
! **************************************************************************************************
!> \brief Solves eigenvalue problem B^A U^A=M^A U^A λ^A in Eq. (6).
!>
!>        Columns are ordered from the largest to the smallest λ^A.
!>        Columns are retained when λ_i/λ_max >= eigenval_threshold.
!> \param M_PQ metric of the large reference RI basis
!> \param B_PQ normalized atom-local target matrix B^A
!> \param target_normalization ...
!> \param U_Pp_A selected AA contractions from the large RI basis P to optimized functions p
!> \param eigenval normalized target eigenvalues in descending order
! **************************************************************************************************
   SUBROUTINE solve_aa_eigenproblem(M_PQ, B_PQ, target_normalization, U_Pp_A, eigenval)
      REAL(kind=dp), DIMENSION(:, :), INTENT(IN)         :: m_pq, b_pq
      REAL(kind=dp), INTENT(IN)                          :: target_normalization
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
         INTENT(OUT)                                     :: u_pp_a
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:), &
         INTENT(OUT)                                     :: eigenval
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'solve_AA_eigenproblem'
 
      INTEGER                                            :: handle, i, j, m_pq_rank, n, nret
      REAL(kind=dp)                                      :: largest_score, relative_fitting_defect
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:)           :: eigenval_b, eigenval_m
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: b_ij, m_pp_check, u_ij, u_pp_a_rotated, &
                                                            x_pi
 
      CALL timeset(routinen, handle)
 
      n = SIZE(m_pq, 1)
      cpassert(n > 0)
      cpassert(SIZE(m_pq, 2) == n)
      cpassert(all(shape(b_pq) == [n, n]))
      cpassert(target_normalization > 0.0_dp)
      CALL compute_reference_m_pq_inv_factor(m_pq, x_pi, m_pq_rank)
 
      ALLOCATE (b_ij(m_pq_rank, m_pq_rank))
      b_ij(:, :) = matmul(transpose(x_pi), &
                          matmul(0.5_dp*(b_pq + transpose(b_pq)), x_pi))/target_normalization
      ALLOCATE (u_ij, source=b_ij)
      ALLOCATE (eigenval_b(m_pq_rank))
      CALL diamat_all(u_ij, eigenval_b, dac=.true.)
      largest_score = maxval(eigenval_b)
      IF (largest_score <= 0.0_dp) THEN
         cpabort("AUTO_RI conditioned atom-local eigensolver failed")
      END IF
      IF (minval(eigenval_b) < -negative_eigenval_tolerance*largest_score) THEN
         cpabort("AUTO_RI conditioned atom-local eigensolver failed")
      END IF
      nret = count(eigenval_b/largest_score >= relative_target_eigenval_threshold)
      ALLOCATE (u_pp_a(n, nret), eigenval(m_pq_rank))
      DO i = 1, m_pq_rank
         eigenval(i) = eigenval_b(m_pq_rank - i + 1)
      END DO
      DO i = 1, nret
         j = m_pq_rank - i + 1
         u_pp_a(:, i) = matmul(x_pi, u_ij(:, j))
      END DO
 
      relative_fitting_defect = 1.0_dp
      DO i = 1, nret
         j = m_pq_rank - i + 1
         relative_fitting_defect = relative_fitting_defect - &
                                   dot_product(u_ij(:, j), matmul(b_ij, u_ij(:, j)))
      END DO
      IF (abs(relative_fitting_defect) < 100.0_dp*epsilon(1.0_dp)) THEN
         relative_fitting_defect = 0.0_dp
      END IF
      ! A projection cannot capture more than the exact product norm.
      IF (relative_fitting_defect < -1.0e-8_dp .OR. &
          relative_fitting_defect > 1.0_dp + 1.0e-8_dp) THEN
         cpabort("AUTO_RI conditioned atom-local eigensolver failed")
      END IF
      IF (nret > 0) THEN
         ALLOCATE (m_pp_check(nret, nret), eigenval_m(nret), u_pp_a_rotated(n, nret))
         m_pp_check(:, :) = matmul(transpose(u_pp_a), &
                                   matmul(0.5_dp*(m_pq + transpose(m_pq)), u_pp_a))
         CALL diamat_all(m_pp_check, eigenval_m, dac=.true.)
         IF (minval(eigenval_m) <= reference_m_eigenval_threshold) THEN
            cpabort("AUTO_RI conditioned atom-local eigensolver failed")
         END IF
         u_pp_a_rotated(:, :) = matmul(u_pp_a, m_pp_check)
         DO i = 1, nret
            u_pp_a_rotated(:, i) = u_pp_a_rotated(:, i)/sqrt(eigenval_m(i))
         END DO
         u_pp_a(:, :) = matmul(u_pp_a_rotated, transpose(m_pp_check))
         m_pp_check(:, :) = matmul(transpose(u_pp_a), &
                                   matmul(0.5_dp*(m_pq + transpose(m_pq)), u_pp_a))
         DO i = 1, nret
            m_pp_check(i, i) = m_pp_check(i, i) - 1.0_dp
         END DO
         IF (maxval(abs(m_pp_check)) >= 1.0e-7_dp) THEN
            cpabort("AUTO_RI conditioned atom-local eigensolver failed")
         END IF
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE solve_aa_eigenproblem
 
! **************************************************************************************************
!> \brief Solves eigenvalue problem B^AB U^AB=U^AB λ^AB of Eq. (8).
!>
!>        Columns are ordered from the largest to the smallest λ^AB.
!> \param M_PQ combined AB reference-basis metric
!> \param B_PQ normalized two-center target matrix B^AB
!> \param U_Pp_AA retained AA columns embedded in the combined AB reference space
!> \param U_Pp_AB selected AB contractions from the large RI basis P to optimized functions p
!> \param eigenval normalized scores in descending order
! **************************************************************************************************
   SUBROUTINE solve_ab_eigenproblem(M_PQ, B_PQ, U_Pp_AA, U_Pp_AB, eigenval)
      REAL(kind=dp), INTENT(IN)                          :: m_pq(:, :), b_pq(:, :), u_pp_aa(:, :)
      REAL(kind=dp), ALLOCATABLE, INTENT(OUT)            :: u_pp_ab(:, :), eigenval(:)
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'solve_AB_eigenproblem'
 
      INTEGER                                            :: handle, i, k, n, nkeep, rank
      REAL(kind=dp)                                      :: threshold
      REAL(kind=dp), ALLOCATABLE :: b_ij(:, :), b_ij_residual(:, :), eigenval_residual(:), &
         m_aa(:, :), p_ij(:, :), q_ia(:, :), u_aa(:, :), u_ia(:, :), x_pi(:, :)
 
      CALL timeset(routinen, handle)
 
      n = SIZE(m_pq, 1)
      CALL compute_reference_m_pq_inv_factor(m_pq, x_pi, rank)
 
      ! Xᵀ M X=I; Q spans the retained atomic columns in these orthonormal coordinates.
      ALLOCATE (q_ia(rank, SIZE(u_pp_aa, 2)))
      q_ia(:, :) = matmul(transpose(x_pi), matmul(m_pq, u_pp_aa))
      ALLOCATE (p_ij(rank, rank), source=0.0_dp)
      DO i = 1, rank
         p_ij(i, i) = 1.0_dp
      END DO
      IF (SIZE(u_pp_aa, 2) > 0) THEN
         ALLOCATE (m_aa(SIZE(u_pp_aa, 2), SIZE(u_pp_aa, 2)))
         m_aa(:, :) = matmul(transpose(q_ia), q_ia)
         CALL compute_reference_m_pq_inv_factor(m_aa, u_aa, k)
         ALLOCATE (u_ia(rank, SIZE(u_aa, 2)))
         u_ia(:, :) = matmul(q_ia, u_aa)
         p_ij(:, :) = p_ij - matmul(u_ia, transpose(u_ia))
         DEALLOCATE (u_ia)
      END IF
      ALLOCATE (b_ij(rank, rank))
      b_ij(:, :) = matmul(transpose(x_pi), matmul(b_pq, x_pi))
      ALLOCATE (b_ij_residual(rank, rank))
      b_ij_residual(:, :) = matmul(p_ij, matmul(b_ij, p_ij))
      ALLOCATE (eigenval_residual(rank))
      CALL diamat_all(b_ij_residual, eigenval_residual, dac=.true.)
      threshold = max(ab_absolute_target_eigenval_threshold, &
                      relative_target_eigenval_threshold*maxval(eigenval_residual), &
                      100.0_dp*epsilon(1.0_dp)*maxval(abs(b_ij)))
      nkeep = count(eigenval_residual > threshold)
      ALLOCATE (u_pp_ab(n, nkeep), eigenval(nkeep))
      k = 0
      DO i = rank, 1, -1
         IF (eigenval_residual(i) <= threshold) cycle
         k = k + 1
         u_pp_ab(:, k) = matmul(x_pi, matmul(p_ij, b_ij_residual(:, i)))
         eigenval(k) = eigenval_residual(i)
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE solve_ab_eigenproblem
 
! **************************************************************************************************
!> \brief Selects the requested RI sizes from the largest positive eigenvalues.
!> \param eigenval_by_block descending eigenvalues; zero pads unavailable modes
!> \param RI_size_requested requested total number of functions
!> \param RI_size_minimum minimum retained size in every nonempty block
!> \param sizes_RI_selected number retained from each block
! **************************************************************************************************
   SUBROUTINE select_ri_sizes(eigenval_by_block, RI_size_requested, RI_size_minimum, &
                              sizes_RI_selected)
      REAL(kind=dp), INTENT(IN)                          :: eigenval_by_block(:, :)
      INTEGER, INTENT(IN)                                :: ri_size_requested, ri_size_minimum
      INTEGER, INTENT(OUT)                               :: sizes_ri_selected(:)
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'select_RI_sizes'
 
      INTEGER                                            :: block, handle, i, n, nmode, remaining
      INTEGER, ALLOCATABLE                               :: indices(:)
      REAL(kind=dp), ALLOCATABLE                         :: values(:)
 
      CALL timeset(routinen, handle)
 
      nmode = SIZE(eigenval_by_block, 1)
      n = SIZE(eigenval_by_block)
      cpassert(SIZE(sizes_ri_selected) == SIZE(eigenval_by_block, 2))
      cpassert(ri_size_minimum >= 0)
      sizes_ri_selected(:) = 0
      IF (n > 0) THEN
         ALLOCATE (values(n), indices(n))
         values(:) = -reshape(eigenval_by_block, [n])
         DO block = 1, SIZE(sizes_ri_selected)
            sizes_ri_selected(block) = &
               min(ri_size_minimum, count(eigenval_by_block(:, block) > 0.0_dp))
            values((block - 1)*nmode + 1: &
                   (block - 1)*nmode + sizes_ri_selected(block)) = 0.0_dp
         END DO
         CALL sort(values, n, indices)
         remaining = min(max(0, ri_size_requested - sum(sizes_ri_selected)), count(values < 0.0_dp))
         DO i = 1, remaining
            block = (indices(i) - 1)/nmode + 1
            sizes_ri_selected(block) = sizes_ri_selected(block) + 1
         END DO
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE select_ri_sizes
 
! **************************************************************************************************
!> \brief Computes orbital factors for the occupied-empty weight in Eq. (5).
!>
!>        weights(n) = 1/|ε_n-E_F|, so weights(i)*weights(a) = w_ia.
!> \param eigenval_MO molecular-orbital eigenvalues, occupied first
!> \param homo Last occupied state in this spin channel.
!> \param weights factors whose occupied-empty products give w_ia
! **************************************************************************************************
   SUBROUTINE compute_orbital_weights(eigenval_MO, homo, weights)
      REAL(kind=dp), INTENT(IN)                          :: eigenval_mo(:)
      INTEGER, INTENT(IN)                                :: homo
      REAL(kind=dp), INTENT(OUT)                         :: weights(:)
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_orbital_weights'
 
      INTEGER                                            :: handle, n
      REAL(kind=dp)                                      :: gap, middle
 
      CALL timeset(routinen, handle)
 
      cpassert(SIZE(weights) == SIZE(eigenval_mo))
      IF (homo < 1 .OR. homo >= SIZE(eigenval_mo)) THEN
         cpabort("AUTO_RI orbital weights require occupied and empty states")
      END IF
      IF (.NOT. all(ieee_is_finite(eigenval_mo))) THEN
         cpabort("AUTO_RI orbital weights require finite eigenvalues")
      END IF
      gap = eigenval_mo(homo + 1) - eigenval_mo(homo)
      IF (.NOT. ieee_is_finite(gap) .OR. gap <= 0.0_dp) THEN
         cpabort("AUTO_RI orbital weights require a finite positive gap")
      END IF
      middle = 0.5_dp*(eigenval_mo(homo) + eigenval_mo(homo + 1))
      DO n = 1, SIZE(eigenval_mo)
         weights(n) = 1.0_dp/abs(eigenval_mo(n) - middle)
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_orbital_weights
 
! **************************************************************************************************
!> \brief Computes the occupied-empty contribution B_PQ^C in Eqs. (3) or (10).
!>
!>        B_PQ^C = Σ_ia w_ia (P|ρ_ia)(ρ_ia|Q).
!>
!> \param Int_3c atom-local or atom-pair three-center integrals (μν|P)
!> \param C_mu_n molecular-orbital coefficients restricted to the relevant AO rows
!> \param weights Per-orbital weights for this spin channel.
!> \param homo Last occupied state.
!> \param B_PQ unnormalized transition matrix for this spin channel
! **************************************************************************************************
   SUBROUTINE compute_b_pq_transition(Int_3c, C_mu_n, weights, homo, B_PQ)
      REAL(kind=dp), CONTIGUOUS, INTENT(IN)              :: int_3c(:, :, :), c_mu_n(:, :), weights(:)
      INTEGER, INTENT(IN)                                :: homo
      REAL(kind=dp), INTENT(OUT)                         :: b_pq(:, :)
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_B_PQ_transition'
 
      INTEGER                                            :: handle, n, nao, nri, nstates
      REAL(kind=dp), ALLOCATABLE                         :: empty(:, :), occupied(:, :), scaled(:, :)
      TYPE(local_gemm_ctxt_type)                         :: gemm_context
 
      CALL timeset(routinen, handle)
 
      CALL gemm_context%create(local_gemm_pu_host)
      nao = SIZE(c_mu_n, 1)
      nstates = SIZE(c_mu_n, 2)
      nri = SIZE(int_3c, 3)
      cpassert(SIZE(int_3c, 1) == nao .AND. SIZE(int_3c, 2) == nao)
      cpassert(SIZE(b_pq, 1) == nri .AND. SIZE(b_pq, 2) == nri)
      cpassert(SIZE(weights) == nstates .AND. homo > 0 .AND. homo < nstates)
      ALLOCATE (scaled(nao, nstates), occupied(nao, nao), empty(nao, nao))
      DO n = 1, nstates
         scaled(:, n) = c_mu_n(:, n)*sqrt(weights(n))
      END DO
      CALL gemm_context%gemm( &
         'N', 'T', nao, nao, homo, 1.0_dp, scaled, nao, scaled, nao, &
         0.0_dp, occupied, nao)
      CALL gemm_context%gemm('N', 'T', nao, nao, nstates - homo, 1.0_dp, &
                             scaled(:, homo + 1:), nao, scaled(:, homo + 1:), nao, &
                             0.0_dp, empty, nao)
      CALL contract_b_pq(int_3c, occupied, empty, b_pq)
      CALL gemm_context%destroy()
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_b_pq_transition
 
! **************************************************************************************************
!> \brief Computes the frontier-exchange contribution B_PQ^X in Eqs. (4) or (11).
!>
!>        B_PQ^X = Σ_{n∈W_occ} Σ_i^occ (P|ρ_ni)(ρ_ni|Q).
!>
!> \param Int_3c atom-local or atom-pair three-center integrals (μν|P)
!> \param C_mu_n molecular-orbital coefficients restricted to the relevant AO rows
!> \param eigenval_MO molecular-orbital eigenvalues, occupied first, in Hartree
!> \param homo Last occupied state in this spin channel.
!> \param occ_energy_window frontier-window extension in Hartree
!> \param B_PQ unnormalized exchange matrix
! **************************************************************************************************
   SUBROUTINE compute_b_pq_exchange(Int_3c, C_mu_n, eigenval_MO, homo, occ_energy_window, B_PQ)
      REAL(kind=dp), CONTIGUOUS, INTENT(IN)              :: int_3c(:, :, :), c_mu_n(:, :), &
                                                            eigenval_mo(:)
      INTEGER, INTENT(IN)                                :: homo
      REAL(kind=dp), INTENT(IN)                          :: occ_energy_window
      REAL(kind=dp), INTENT(OUT)                         :: b_pq(:, :)
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_B_PQ_exchange'
 
      INTEGER                                            :: handle, n, nao, nstates
      REAL(kind=dp), ALLOCATABLE                         :: frontier(:, :), occupied(:, :), &
                                                            selected(:, :)
      TYPE(local_gemm_ctxt_type)                         :: gemm_context
 
      CALL timeset(routinen, handle)
 
      b_pq(:, :) = 0.0_dp
      nstates = SIZE(eigenval_mo)
      IF (homo < 1 .OR. homo >= nstates) THEN
         cpabort("AUTO_RI exchange target requires occupied and empty states")
      END IF
      IF (.NOT. ieee_is_finite(occ_energy_window) .OR. &
          .NOT. all(ieee_is_finite(eigenval_mo))) THEN
         cpabort("AUTO_RI exchange target requires finite input")
      END IF
      IF (occ_energy_window < 0.0_dp .OR. eigenval_mo(homo + 1) <= eigenval_mo(homo)) THEN
         cpabort("AUTO_RI exchange target requires a valid frontier window")
      END IF
      nao = SIZE(c_mu_n, 1)
      cpassert(SIZE(c_mu_n, 2) == nstates)
      ALLOCATE (occupied(nao, nao), frontier(nao, nao), selected(nao, nstates), source=0.0_dp)
      DO n = 1, nstates
         IF (eigenval_mo(n) < eigenval_mo(homo) - occ_energy_window .OR. &
             eigenval_mo(n) > eigenval_mo(homo + 1) + occ_energy_window) cycle
         selected(:, n) = c_mu_n(:, n)
      END DO
      CALL gemm_context%create(local_gemm_pu_host)
      ! D_μν = Σ_i^occ C_μi C_νi.
      ! F_μν = Σ_{n: ε_n ∈ [ε_H−E_win,ε_L+E_win]} C_μn C_νn.
      CALL gemm_context%gemm('N', 'T', nao, nao, homo, 1.0_dp, c_mu_n, nao, c_mu_n, nao, &
                             0.0_dp, occupied, nao)
      CALL gemm_context%gemm('N', 'T', nao, nao, nstates, 1.0_dp, selected, nao, selected, nao, &
                             0.0_dp, frontier, nao)
      CALL contract_b_pq(int_3c, occupied, frontier, b_pq)
      CALL gemm_context%destroy()
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_b_pq_exchange
 
! **************************************************************************************************
!> \brief Contracts AO density matrices and (μν|P) to form B_PQ.
!>
!>        B_PQ = Σ_μν (μν|P) [left (..|Q) right]_μν.
!>
!> \param Int_3c atom-local or atom-pair three-center integrals (μν|P)
!> \param left First density block, for occupied partners.
!> \param right Second density block, for empty or hard-window target states.
!> \param B_PQ symmetric RI-basis product matrix
! **************************************************************************************************
   SUBROUTINE contract_b_pq(Int_3c, left, right, B_PQ)
      REAL(kind=dp), CONTIGUOUS, INTENT(IN)              :: int_3c(:, :, :), left(:, :), right(:, :)
      REAL(kind=dp), INTENT(OUT)                         :: b_pq(:, :)
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'contract_B_PQ'
 
      INTEGER                                            :: b, handle, mu, naux, nu, p, q
      REAL(kind=dp), ALLOCATABLE                         :: weighted(:, :), work(:, :)
      TYPE(local_gemm_ctxt_type)                         :: gemm_context
 
      CALL timeset(routinen, handle)
 
      b = SIZE(int_3c, 1)
      naux = SIZE(int_3c, 3)
      cpassert(SIZE(int_3c, 2) == b)
      cpassert(SIZE(left, 1) == b .AND. SIZE(left, 2) == b)
      cpassert(SIZE(right, 1) == b .AND. SIZE(right, 2) == b)
      cpassert(SIZE(b_pq, 1) == naux .AND. SIZE(b_pq, 2) == naux)
      ALLOCATE (work(b, b), weighted(b, b))
      CALL gemm_context%create(local_gemm_pu_host)
      DO q = 1, naux
         ! B_PQ = Σ_μν (μν|P) [left (..|Q) right]_μν.
         CALL gemm_context%gemm('N', 'N', b, b, b, 1.0_dp, left, b, int_3c(:, :, q), b, &
                                0.0_dp, work, b)
         CALL gemm_context%gemm('N', 'N', b, b, b, 1.0_dp, work, b, right, b, 0.0_dp, weighted, b)
         DO p = 1, q
            b_pq(p, q) = 0.0_dp
            DO nu = 1, b
               DO mu = 1, b
                  b_pq(p, q) = b_pq(p, q) + int_3c(mu, nu, p)*weighted(mu, nu)
               END DO
            END DO
            b_pq(q, p) = b_pq(p, q)
         END DO
      END DO
      CALL gemm_context%destroy()
 
      CALL timestop(handle)
 
   END SUBROUTINE contract_b_pq
 
! **************************************************************************************************
!> \brief Adds separately normalized transition and exchange matrices.
!> \param X_Pi inverse metric factor, X_Pi X_Qi = M_PQ^+
!> \param B_PQ_transition transition matrix, replaced by the combined target
!> \param B_PQ_exchange unnormalized hard-window exchange covariance
!> \param target_norm number of nonzero normalized terms
! **************************************************************************************************
   SUBROUTINE combine_transition_and_exchange_targets( &
      X_Pi, B_PQ_transition, B_PQ_exchange, target_norm)
      REAL(kind=dp), CONTIGUOUS, INTENT(IN)              :: x_pi(:, :)
      REAL(kind=dp), INTENT(INOUT)                       :: b_pq_transition(:, :)
      REAL(kind=dp), CONTIGUOUS, INTENT(IN)              :: b_pq_exchange(:, :)
      REAL(kind=dp), INTENT(OUT)                         :: target_norm
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'combine_transition_and_exchange_targets'
 
      INTEGER                                            :: handle, k
      REAL(kind=dp)                                      :: exchange_norm, transition_norm
      REAL(kind=dp), ALLOCATABLE                         :: work(:, :)
 
      CALL timeset(routinen, handle)
 
      ALLOCATE (work(SIZE(x_pi, 1), SIZE(x_pi, 2)))
      work(:, :) = matmul(b_pq_transition, x_pi)
      transition_norm = 0.0_dp
      DO k = 1, SIZE(x_pi, 2)
         transition_norm = transition_norm + dot_product(x_pi(:, k), work(:, k))
      END DO
      work(:, :) = matmul(b_pq_exchange, x_pi)
      exchange_norm = 0.0_dp
      DO k = 1, SIZE(x_pi, 2)
         exchange_norm = exchange_norm + dot_product(x_pi(:, k), work(:, k))
      END DO
      target_norm = 0.0_dp
      IF (.NOT. ieee_is_finite(transition_norm) .OR. &
          .NOT. ieee_is_finite(exchange_norm) .OR. &
          transition_norm < 0.0_dp .OR. exchange_norm < 0.0_dp) THEN
         cpabort("AUTO_RI target normalization is invalid")
      END IF
      ! B = B_ia / Tr(M^+ B_ia) + B_ni / Tr(M^+ B_ni), with zero-norm terms omitted.
      IF (transition_norm > 0.0_dp) THEN
         b_pq_transition(:, :) = b_pq_transition/transition_norm
         target_norm = target_norm + 1.0_dp
      ELSE
         b_pq_transition(:, :) = 0.0_dp
      END IF
      IF (exchange_norm > 0.0_dp) THEN
         b_pq_transition(:, :) = b_pq_transition + b_pq_exchange/exchange_norm
         target_norm = target_norm + 1.0_dp
      END IF
      IF (target_norm <= 0.0_dp .OR. .NOT. all(ieee_is_finite(b_pq_transition))) THEN
         cpabort("AUTO_RI target normalization is invalid")
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE combine_transition_and_exchange_targets
 
! **************************************************************************************************
!> \brief Computes the atom-local matrix B_PQ^A in Eqs. (3)-(6).
!>
!>        B_PQ^A = B_PQ^{A,C}/N_C^A + B_PQ^{A,X}/N_X^A.
!>
!>        The molecular-orbital coefficient rows belonging to one atom are
!>        collected on one MPI rank. No global AO-pair matrix is formed.
!> \param bs_env molecular spectrum and distributed MO coefficients
!> \param iatom Atom whose AO rows are needed.
!> \param root_rank MPI rank performing the atomic contractions.
!> \param Int_3c atom-local three-center integrals (μν|P)
!> \param M_PQ atom-local metric of the large reference RI basis
!> \param B_PQ normalized atom-local matrix B_PQ^A
!> \param target_norm number of nonzero normalized contributions
! **************************************************************************************************
   SUBROUTINE compute_b_pq_aa(bs_env, iatom, root_rank, Int_3c, M_PQ, &
                              B_PQ, target_norm)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      INTEGER, INTENT(IN)                                :: iatom, root_rank
      REAL(kind=dp), CONTIGUOUS, INTENT(IN)              :: int_3c(:, :, :), m_pq(:, :)
      REAL(kind=dp), ALLOCATABLE, INTENT(OUT)            :: b_pq(:, :)
      REAL(kind=dp), INTENT(OUT)                         :: target_norm
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'compute_B_PQ_AA'
 
      INTEGER                                            :: column, first, handle, i, j, nao, nri, &
                                                            nstates, rank, row, spin
      INTEGER, POINTER                                   :: cols(:), rows(:)
      REAL(kind=dp), ALLOCATABLE                         :: b_pq_exchange(:, :), b_pq_spin(:, :), &
                                                            c_mu_n(:, :), weights(:), x_pi(:, :)
 
      CALL timeset(routinen, handle)
      NULLIFY (rows, cols)
      nao = SIZE(int_3c, 1)
      nri = SIZE(int_3c, 3)
      nstates = bs_env%n_mo_retained
      cpassert(ALLOCATED(bs_env%eigenval_scf_Gamma))
      first = 1 + sum(bs_env%sizes_AO(:iatom - 1))
      ALLOCATE (b_pq(nri, nri), source=0.0_dp)
      ALLOCATE (c_mu_n(nao, nstates), weights(nstates), b_pq_spin(nri, nri))
      target_norm = 0.0_dp
      ALLOCATE (b_pq_exchange(nri, nri), source=0.0_dp)
      DO spin = 1, bs_env%n_spin
         CALL compute_orbital_weights(bs_env%eigenval_scf_Gamma(1:nstates, spin), &
                                      bs_env%n_occ(spin), weights)
         CALL cp_fm_get_info(bs_env%fm_mo_coeff_Gamma(spin), row_indices=rows, col_indices=cols)
         c_mu_n(:, :) = 0.0_dp
         DO i = 1, SIZE(rows)
            row = rows(i) - first + 1
            IF (row < 1 .OR. row > nao) cycle
            DO j = 1, SIZE(cols)
               column = cols(j)
               IF (column > nstates) cycle
               c_mu_n(row, column) = bs_env%fm_mo_coeff_Gamma(spin)%local_data(i, j)
            END DO
         END DO
         CALL bs_env%para_env%sum(c_mu_n, root_rank)
         IF (bs_env%para_env%mepos /= root_rank) cycle
         CALL compute_b_pq_transition(int_3c, c_mu_n, weights, bs_env%n_occ(spin), b_pq_spin)
         b_pq(:, :) = b_pq + b_pq_spin
         CALL compute_b_pq_exchange(int_3c, c_mu_n, &
                                    bs_env%eigenval_scf_Gamma(1:nstates, spin), &
                                    bs_env%n_occ(spin), bs_env%auto_ri%occ_energy_window, b_pq_spin)
         b_pq_exchange(:, :) = b_pq_exchange + b_pq_spin
      END DO
      IF (bs_env%para_env%mepos /= root_rank) THEN
         CALL timestop(handle)
         RETURN
      END IF
      CALL compute_reference_m_pq_inv_factor(m_pq, x_pi, rank)
      CALL combine_transition_and_exchange_targets(x_pi, b_pq, b_pq_exchange, target_norm)
      CALL timestop(handle)
 
   END SUBROUTINE compute_b_pq_aa
 
! **************************************************************************************************
!> \brief Normalizes each contraction coefficient vector U_Pp.
!>
!>        U_Pp <- U_Pp/sqrt(Σ_P |U_Pp|^2).
!> \param U_Pp contractions from the large RI basis P to optimized functions p
! **************************************************************************************************
   SUBROUTINE normalize_ri_coefficients(U_Pp)
      REAL(kind=dp), DIMENSION(:, :), INTENT(INOUT)      :: u_pp
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'normalize_RI_coefficients'
 
      INTEGER                                            :: handle, j
      REAL(kind=dp)                                      :: column_norm
 
      CALL timeset(routinen, handle)
 
      DO j = 1, SIZE(u_pp, 2)
         column_norm = norm2(u_pp(:, j))
         IF (column_norm <= tiny(1.0_dp)) THEN
            cpabort("AUTO_RI produced a numerically zero localized column")
         END IF
         u_pp(:, j) = u_pp(:, j)/column_norm
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE normalize_ri_coefficients
 
! **************************************************************************************************
!> \brief Creates a matrix in the existing molecular BLACS context.
!> \param qs_env ...
!> \param matrix ...
!> \param nr ...
!> \param nc ...
! **************************************************************************************************
   SUBROUTINE create_distributed_matrix(qs_env, matrix, nr, nc)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(cp_fm_type), INTENT(OUT)                      :: matrix
      INTEGER, INTENT(IN)                                :: nr, nc
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'create_distributed_matrix'
 
      INTEGER                                            :: handle
      TYPE(cp_blacs_env_type), POINTER                   :: blacs
      TYPE(cp_fm_struct_type), POINTER                   :: layout
      TYPE(mp_para_env_type), POINTER                    :: para
 
      CALL timeset(routinen, handle)
 
      NULLIFY (blacs, para, layout)
      CALL get_qs_env(qs_env, blacs_env=blacs, para_env=para)
      CALL cp_fm_struct_create(layout, context=blacs, nrow_global=nr, ncol_global=nc, para_env=para)
      CALL cp_fm_create(matrix, layout, set_zero=.true.)
      CALL cp_fm_struct_release(layout)
 
      CALL timestop(handle)
 
   END SUBROUTINE create_distributed_matrix
 
! **************************************************************************************************
!> \brief Computes and stores the two-center contractions U^AB in Eqs. (8)-(13).
!>
!>        The selected AA space is projected out according to Eq. (8). The symmetric
!>        pair density in Eq. (9) defines B^{AB,C} and B^{AB,X} in Eqs. (10) and (11).
!>        Diagonalizing their sum gives U^AB in Eq. (12), which expands the contracted
!>        pair functions in Eq. (13).
!> \param qs_env ...
!> \param bs_env ...
!> \param U_Pp_A temporary atom-local contractions used to project the AA space
! **************************************************************************************************
   SUBROUTINE compute_ab_optimized_ri_basis(qs_env, bs_env, U_Pp_A)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      REAL(kind=dp), DIMENSION(:, :, :), INTENT(IN)      :: u_pp_a
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_AB_optimized_RI_basis'
 
      INTEGER :: ab_pair_count, center_index, current_atom, expected_ri_size, first_ab_pair, &
         handle, handle_pair_batches, iatom, jatom, k, last_ab_pair, nao_kind, nri_a, nri_ab_ref, &
         nri_b, ri_size_opt_aa, ri_size_opt_ab, ri_size_optimized, ri_size_ref, ri_size_requested, &
         root_rank, row, size_opt_ri_ab, size_u_pp_ab, size_u_pp_ab_storage, spin
      INTEGER, ALLOCATABLE :: ab_atom_a(:), ab_atom_b(:), ab_size_opt_ri_to_a(:), block_next(:), &
         sizes_opt_ri_aa(:), sizes_opt_ri_ab(:), sizes_ref_ri(:), u_pp_ab_offset(:)
      REAL(kind=dp), ALLOCATABLE :: b_pq(:, :), c_mu_n_ab(:, :, :), c_mu_n_atom(:, :), &
         eigenval(:), eigenval_ab(:, :), m_pq(:, :), u_pp_aa(:, :), u_pp_ab(:, :), &
         u_pp_ab_packed(:)
      TYPE(gw_3c_ctx_type)                               :: ctx
      TYPE(gw_3c_ws_type)                                :: ws
      TYPE(mp_para_env_type), POINTER                    :: para
 
      CALL timeset(routinen, handle)
 
      NULLIFY (para)
      CALL get_qs_env(qs_env, para_env=para)
      ALLOCATE (sizes_ref_ri(bs_env%n_atom), sizes_opt_ri_aa(bs_env%n_atom))
      DO iatom = 1, bs_env%n_atom
         sizes_ref_ri(iatom) = get_ref_ri_size(bs_env, iatom)
      END DO
      sizes_opt_ri_aa(:) = bs_env%auto_ri%sizes_opt_RI
      ri_size_ref = sum(sizes_ref_ri)
      ri_size_opt_aa = sum(sizes_opt_ri_aa)
 
      ab_pair_count = 0
      DO iatom = 1, bs_env%n_atom - 1
         DO jatom = iatom + 1, bs_env%n_atom
            IF (norm2(bs_env%ri_rs%particle_set(iatom)%r - bs_env%ri_rs%particle_set(jatom)%r) <= &
                bs_env%auto_ri%neighbor_radius) THEN
               ab_pair_count = ab_pair_count + 1
            END IF
         END DO
      END DO
      ALLOCATE (ab_atom_a(ab_pair_count), ab_atom_b(ab_pair_count), &
                ab_size_opt_ri_to_a(ab_pair_count), sizes_opt_ri_ab(ab_pair_count), &
                u_pp_ab_offset(ab_pair_count))
      ab_size_opt_ri_to_a = 0
      u_pp_ab_offset = 0
      ALLOCATE (eigenval_ab(2*maxval(sizes_ref_ri), ab_pair_count), source=0.0_dp)
 
      CALL gw_3c_ctx_create(ctx, qs_env, bs_env%ri_metric, bs_env%basis_set_AO, &
                            bs_env%basis_set_AO, bs_env%basis_set_RI)
      CALL gw_3c_ws_create(ws, ctx)
      k = 0
      DO iatom = 1, bs_env%n_atom - 1
         DO jatom = iatom + 1, bs_env%n_atom
            IF (norm2(bs_env%ri_rs%particle_set(iatom)%r - bs_env%ri_rs%particle_set(jatom)%r) > &
                bs_env%auto_ri%neighbor_radius) cycle
            k = k + 1
            ab_atom_a(k) = iatom
            ab_atom_b(k) = jatom
         END DO
      END DO
      size_u_pp_ab_storage = 0
      DO k = 1, ab_pair_count
         root_rank = modulo(k - 1, para%num_pe)
         IF (para%mepos /= root_rank) cycle
         nri_ab_ref = sizes_ref_ri(ab_atom_a(k)) + sizes_ref_ri(ab_atom_b(k))
         u_pp_ab_offset(k) = size_u_pp_ab_storage + 1
         size_u_pp_ab_storage = size_u_pp_ab_storage + nri_ab_ref*nri_ab_ref
      END DO
      ALLOCATE (u_pp_ab_packed(size_u_pp_ab_storage))
      CALL timeset(routinen//'_pair_batches', handle_pair_batches)
      DO first_ab_pair = 1, ab_pair_count, para%num_pe
         last_ab_pair = min(ab_pair_count, first_ab_pair + para%num_pe - 1)
         ! Gather only one pair's MO rows per rank, then solve all batch pairs concurrently.
         DO k = first_ab_pair, last_ab_pair
            root_rank = k - first_ab_pair
            iatom = ab_atom_a(k)
            jatom = ab_atom_b(k)
            IF (para%mepos == root_rank) THEN
               ALLOCATE (c_mu_n_ab( &
                         bs_env%sizes_AO(iatom) + bs_env%sizes_AO(jatom), &
                         bs_env%n_mo_retained, bs_env%n_spin))
            END IF
            DO spin = 1, bs_env%n_spin
               row = 0
               DO center_index = 1, 2
                  current_atom = iatom
                  IF (center_index == 2) current_atom = jatom
                  nao_kind = bs_env%sizes_AO(current_atom)
                  ALLOCATE (c_mu_n_atom(nao_kind, bs_env%n_mo_retained))
                  CALL cp_fm_get_submatrix(bs_env%fm_mo_coeff_Gamma(spin), c_mu_n_atom, &
                                           start_row=1 + sum(bs_env%sizes_AO(:current_atom - 1)))
                  IF (para%mepos == root_rank) THEN
                     c_mu_n_ab(row + 1:row + nao_kind, :, spin) = c_mu_n_atom
                  END IF
                  DEALLOCATE (c_mu_n_atom)
                  row = row + nao_kind
               END DO
            END DO
         END DO
         k = first_ab_pair + para%mepos
         IF (k <= last_ab_pair) THEN
            iatom = ab_atom_a(k)
            jatom = ab_atom_b(k)
            nri_a = sizes_ref_ri(iatom)
            nri_b = sizes_ref_ri(jatom)
            nri_ab_ref = nri_a + nri_b
            ! Eqs. (9)-(11): form the symmetric pair density and B^{AB,C}+B^{AB,X}.
            CALL compute_b_pq_ab(bs_env, ctx, ws, iatom, jatom, c_mu_n_ab, m_pq, b_pq)
            ALLOCATE (u_pp_aa(nri_ab_ref, sizes_opt_ri_aa(iatom) + sizes_opt_ri_aa(jatom)), &
                      source=0.0_dp)
            u_pp_aa(1:nri_a, 1:sizes_opt_ri_aa(iatom)) = &
               u_pp_a(1:nri_a, 1:sizes_opt_ri_aa(iatom), iatom)
            u_pp_aa(nri_a + 1:, sizes_opt_ri_aa(iatom) + 1:) = &
               u_pp_a(1:nri_b, 1:sizes_opt_ri_aa(jatom), jatom)
            ! Eqs. (8) and (12): project out the AA space and diagonalize the remainder.
            CALL solve_ab_eigenproblem(m_pq, b_pq, u_pp_aa, u_pp_ab, eigenval)
            eigenval_ab(1:SIZE(eigenval), k) = eigenval
            DEALLOCATE (eigenval)
            u_pp_ab_packed(u_pp_ab_offset(k):u_pp_ab_offset(k) + SIZE(u_pp_ab) - 1) = &
               reshape(u_pp_ab, [SIZE(u_pp_ab)])
            DEALLOCATE (u_pp_ab)
            DEALLOCATE (u_pp_aa, b_pq, m_pq, c_mu_n_ab)
         END IF
      END DO
      CALL timestop(handle_pair_batches)
      k = ab_pair_count
      CALL gw_3c_ws_release(ws)
      CALL gw_3c_ctx_release(ctx)
      ri_size_requested = nint(min(bs_env%auto_ri%ri_ao_ratio, &
                                   REAL(ri_size_ref, dp)/REAL(SUM(bs_env%sizes_AO), dp))* &
                               REAL(SUM(bs_env%sizes_AO), dp))
      ri_size_requested = max(ri_size_requested, bs_env%n_atom)
 
      CALL para%sum(eigenval_ab)
      CALL select_ri_sizes(eigenval_ab, ri_size_requested - sum(sizes_opt_ri_aa), 0, &
                           sizes_opt_ri_ab)
      ri_size_opt_ab = sum(sizes_opt_ri_ab)
      bs_env%auto_ri%sizes_opt_RI(:) = sizes_opt_ri_aa
      DO k = 1, ab_pair_count
         iatom = ab_atom_a(k)
         jatom = ab_atom_b(k)
         size_opt_ri_ab = sizes_opt_ri_ab(k)
         ab_size_opt_ri_to_a(k) = balance_ab_columns(bs_env%auto_ri%sizes_opt_RI(iatom), &
                                                     bs_env%auto_ri%sizes_opt_RI(jatom), &
                                                     size_opt_ri_ab)
         bs_env%auto_ri%sizes_opt_RI(iatom) = &
            bs_env%auto_ri%sizes_opt_RI(iatom) + ab_size_opt_ri_to_a(k)
         bs_env%auto_ri%sizes_opt_RI(jatom) = &
            bs_env%auto_ri%sizes_opt_RI(jatom) + size_opt_ri_ab - ab_size_opt_ri_to_a(k)
      END DO
 
      ri_size_optimized = sum(bs_env%auto_ri%sizes_opt_RI)
      expected_ri_size = ri_size_opt_aa + ri_size_opt_ab
      cpassert(ri_size_optimized == expected_ri_size)
      cpassert(all(bs_env%auto_ri%sizes_opt_RI > 0))
      IF (ri_size_optimized > ri_size_ref) THEN
         cpabort("AUTO_RI request exceeds independent reference RI directions")
      END IF
      size_u_pp_ab = 0
      DO iatom = 1, bs_env%n_atom
         size_u_pp_ab = size_u_pp_ab + sizes_ref_ri(iatom)*sizes_opt_ri_aa(iatom)
      END DO
      DO k = 1, ab_pair_count
         size_u_pp_ab = size_u_pp_ab + &
                        (sizes_ref_ri(ab_atom_a(k)) + &
                         sizes_ref_ri(ab_atom_b(k)))*sizes_opt_ri_ab(k)
      END DO
      ALLOCATE (bs_env%auto_ri%AB_atom_A(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%AB_atom_B(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%AB_first_p_A(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%AB_first_p_B(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%AB_size_opt_RI_to_A(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%AB_size_ref_RI(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%AB_size_opt_RI(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%U_Pp_AB_offset(bs_env%n_atom + ab_pair_count), source=0)
      ALLOCATE (bs_env%auto_ri%U_Pp_AB(size_u_pp_ab), source=0.0_dp)
      bs_env%auto_ri%AB_block_count = 0
 
      ALLOCATE (block_next(bs_env%n_atom))
      block_next(1) = 1
      DO iatom = 2, bs_env%n_atom
         block_next(iatom) = block_next(iatom - 1) + bs_env%auto_ri%sizes_opt_RI(iatom - 1)
      END DO
      DO iatom = 1, bs_env%n_atom
         IF (sizes_opt_ri_aa(iatom) > 0) bs_env%auto_ri%AB_block_count = &
            bs_env%auto_ri%AB_block_count + 1
         CALL store_aa_columns(bs_env, iatom, &
                               u_pp_a(1:sizes_ref_ri(iatom), 1:sizes_opt_ri_aa(iatom), iatom), &
                               block_next)
      END DO
 
      DO k = 1, ab_pair_count
         iatom = ab_atom_a(k)
         jatom = ab_atom_b(k)
         size_opt_ri_ab = sizes_opt_ri_ab(k)
         IF (size_opt_ri_ab == 0) cycle
         root_rank = modulo(k - 1, para%num_pe)
         nri_ab_ref = sizes_ref_ri(iatom) + sizes_ref_ri(jatom)
         ALLOCATE (u_pp_ab(nri_ab_ref, size_opt_ri_ab))
         IF (para%mepos == root_rank) THEN
            u_pp_ab(:, :) = reshape( &
                            u_pp_ab_packed(u_pp_ab_offset(k): &
                                           u_pp_ab_offset(k) + &
                                           nri_ab_ref*size_opt_ri_ab - 1), &
                            [nri_ab_ref, size_opt_ri_ab])
         END IF
         CALL para%bcast(u_pp_ab, root_rank)
         bs_env%auto_ri%AB_block_count = bs_env%auto_ri%AB_block_count + 1
         ! Eq. (13): store the selected neighbor-pair contractions in the global transform.
         CALL store_ab_columns(bs_env, u_pp_ab, iatom, jatom, ab_size_opt_ri_to_a(k), block_next)
         DEALLOCATE (u_pp_ab)
      END DO
 
      DO iatom = 1, bs_env%n_atom
         expected_ri_size = 1 + sum(bs_env%auto_ri%sizes_opt_RI(:iatom))
         cpassert(block_next(iatom) == expected_ri_size)
      END DO
      DEALLOCATE (block_next, ab_atom_a, ab_atom_b, ab_size_opt_ri_to_a, &
                  u_pp_ab_offset, u_pp_ab_packed)
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_ab_optimized_ri_basis
 
! **************************************************************************************************
!> \brief Divides AB columns between atoms A and B for load balance.
!> \param load_a current number of columns owned by A
!> \param load_b current number of columns owned by B
!> \param RI_size_opt_AB number of new pair RI functions
!> \return ...
! **************************************************************************************************
   PURE INTEGER FUNCTION balance_ab_columns(load_a, load_b, RI_size_opt_AB) RESULT(RI_size_to_a)
      INTEGER, INTENT(IN)                                :: load_a, load_b, ri_size_opt_ab
 
      ri_size_to_a = (load_b + ri_size_opt_ab - load_a + 1)/2
      ri_size_to_a = max(0, min(ri_size_opt_ab, ri_size_to_a))
 
   END FUNCTION balance_ab_columns
 
! **************************************************************************************************
!> \brief Stores the AA contraction columns for one atom.
!> \param bs_env ...
!> \param atom ...
!> \param U_Pp_A atom-local contraction coefficients
!> \param block_next ...
! **************************************************************************************************
   SUBROUTINE store_aa_columns(bs_env, atom, U_Pp_A, block_next)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      INTEGER, INTENT(IN)                                :: atom
      REAL(kind=dp), DIMENSION(:, :), INTENT(IN)         :: u_pp_a
      INTEGER, DIMENSION(:), INTENT(INOUT)               :: block_next
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'store_AA_columns'
 
      INTEGER                                            :: ab_block, first, handle, ri_size_opt_aa
      REAL(kind=dp), ALLOCATABLE                         :: u_pp_a_normalized(:, :)
 
      CALL timeset(routinen, handle)
 
      first = block_next(atom)
      ab_block = bs_env%auto_ri%AB_block_count
      ri_size_opt_aa = SIZE(u_pp_a, 2)
      ALLOCATE (u_pp_a_normalized(SIZE(u_pp_a, 1), ri_size_opt_aa))
      IF (ri_size_opt_aa > 0) THEN
         u_pp_a_normalized(:, :) = u_pp_a
         CALL normalize_ri_coefficients(u_pp_a_normalized)
         bs_env%auto_ri%AB_atom_A(ab_block) = atom
         bs_env%auto_ri%AB_atom_B(ab_block) = atom
         bs_env%auto_ri%AB_first_p_A(ab_block) = &
            first - sum(bs_env%auto_ri%sizes_opt_RI(:atom - 1))
         bs_env%auto_ri%AB_size_opt_RI_to_A(ab_block) = ri_size_opt_aa
         CALL store_u_pp_ab(bs_env%auto_ri, ab_block, u_pp_a_normalized)
      END IF
      DEALLOCATE (u_pp_a_normalized)
      block_next(atom) = first + ri_size_opt_aa
 
      CALL timestop(handle)
 
   END SUBROUTINE store_aa_columns
 
! **************************************************************************************************
!> \brief Stores an AB contraction in the blocks assigned to atoms A and B.
!> \param bs_env ...
!> \param U_Pp_AB neighbor-pair contractions
!> \param atom_A ...
!> \param atom_B ...
!> \param RI_size_to_a number of optimized pair RI functions assigned to atom A
!> \param block_next ...
! **************************************************************************************************
   SUBROUTINE store_ab_columns(bs_env, U_Pp_AB, atom_A, atom_B, RI_size_to_a, block_next)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
         INTENT(INOUT)                                   :: u_pp_ab
      INTEGER, INTENT(IN)                                :: atom_a, atom_b, ri_size_to_a
      INTEGER, DIMENSION(:), INTENT(INOUT)               :: block_next
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'store_AB_columns'
 
      INTEGER                                            :: ab_block, first_a, first_b, handle, &
                                                            ri_size_opt_ab
 
      CALL timeset(routinen, handle)
 
      ri_size_opt_ab = SIZE(u_pp_ab, 2)
      ab_block = bs_env%auto_ri%AB_block_count
      first_a = block_next(atom_a)
      first_b = block_next(atom_b)
      CALL normalize_ri_coefficients(u_pp_ab)
 
      bs_env%auto_ri%AB_atom_A(ab_block) = atom_a
      bs_env%auto_ri%AB_atom_B(ab_block) = atom_b
      bs_env%auto_ri%AB_first_p_A(ab_block) = first_a - sum(bs_env%auto_ri%sizes_opt_RI(:atom_a - 1))
      bs_env%auto_ri%AB_first_p_B(ab_block) = first_b - sum(bs_env%auto_ri%sizes_opt_RI(:atom_b - 1))
      bs_env%auto_ri%AB_size_opt_RI_to_A(ab_block) = ri_size_to_a
      CALL store_u_pp_ab(bs_env%auto_ri, ab_block, u_pp_ab)
      block_next(atom_a) = first_a + ri_size_to_a
      block_next(atom_b) = first_b + ri_size_opt_ab - ri_size_to_a
 
      CALL timestop(handle)
 
   END SUBROUTINE store_ab_columns
 
! **************************************************************************************************
!> \brief Packs one AB U_Pp matrix into the persistent AUTO_RI data; B=A denotes an AA block.
!> \param auto_ri ...
!> \param AB_block ...
!> \param U_Pp ...
! **************************************************************************************************
   SUBROUTINE store_u_pp_ab(auto_ri, AB_block, U_Pp)
      TYPE(auto_ri_type), INTENT(INOUT)                  :: auto_ri
      INTEGER, INTENT(IN)                                :: ab_block
      REAL(kind=dp), DIMENSION(:, :), INTENT(IN)         :: u_pp
 
      INTEGER                                            :: first, last
 
      first = 1
      IF (ab_block > 1) THEN
         first = auto_ri%U_Pp_AB_offset(ab_block - 1) + &
                 auto_ri%AB_size_ref_RI(ab_block - 1)* &
                 auto_ri%AB_size_opt_RI(ab_block - 1)
      END IF
      last = first + SIZE(u_pp) - 1
      cpassert(last <= SIZE(auto_ri%U_Pp_AB))
      auto_ri%U_Pp_AB_offset(ab_block) = first
      auto_ri%AB_size_ref_RI(ab_block) = SIZE(u_pp, 1)
      auto_ri%AB_size_opt_RI(ab_block) = SIZE(u_pp, 2)
      auto_ri%U_Pp_AB(first:last) = reshape(u_pp, [SIZE(u_pp)])
 
   END SUBROUTINE store_u_pp_ab
 
! **************************************************************************************************
!> \brief Computes the two-center matrix B_PQ^AB in Eqs. (9)-(12).
!>
!>        Its numerator uses the symmetric off-diagonal density
!>        ρ̅_mn^AB=ρ_mn^AB+ρ_mn^BA after projection from the retained AA
!>        space. Its normalization uses the complete two-atom density.
!> \param bs_env ...
!> \param ctx ...
!> \param ws ...
!> \param iatom ...
!> \param jatom ...
!> \param C_mu_n_AB molecular-orbital coefficients on the AO rows of atoms A and B
!> \param M_PQ combined metric of the large reference RI bases on A and B
!> \param B_PQ normalized two-center matrix B_PQ^AB
! **************************************************************************************************
   SUBROUTINE compute_b_pq_ab(bs_env, ctx, ws, iatom, jatom, C_mu_n_AB, M_PQ, B_PQ)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(gw_3c_ctx_type), INTENT(IN)                   :: ctx
      TYPE(gw_3c_ws_type), INTENT(INOUT)                 :: ws
      INTEGER, INTENT(IN)                                :: iatom, jatom
      REAL(kind=dp), INTENT(IN)                          :: c_mu_n_ab(:, :, :)
      REAL(kind=dp), ALLOCATABLE, INTENT(OUT)            :: m_pq(:, :), b_pq(:, :)
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'compute_B_PQ_AB'
 
      INTEGER :: atoms(2), handle, homo, i, j, k, kind_a, kind_b, nao_ab, nri_ab_ref, &
         offset_ao(2), offset_ref_ri(2), rank, sizes_ao(2), sizes_ref_ri(2), spin
      LOGICAL                                            :: screened
      REAL(kind=dp)                                      :: norm_exchange, norm_transition
      REAL(kind=dp), ALLOCATABLE :: b_pi(:, :), b_pq_exchange_ab(:, :), b_pq_exchange_full(:, :), &
         b_pq_spin(:, :), b_pq_transition_ab(:, :), b_pq_transition_full(:, :), &
         int_3c_ab(:, :, :), int_3c_all(:, :, :), m_pq_block(:, :), weights(:), x_pi(:, :)
 
      CALL timeset(routinen, handle)
 
      atoms(1) = iatom
      atoms(2) = jatom
      sizes_ao = bs_env%sizes_AO(atoms)
      DO i = 1, 2
         sizes_ref_ri(i) = get_ref_ri_size(bs_env, atoms(i))
      END DO
      offset_ao(1) = 0
      offset_ao(2) = sizes_ao(1)
      offset_ref_ri(1) = 0
      offset_ref_ri(2) = sizes_ref_ri(1)
      nao_ab = sum(sizes_ao)
      nri_ab_ref = sum(sizes_ref_ri)
      ALLOCATE (b_pq(nri_ab_ref, nri_ab_ref), &
                m_pq(nri_ab_ref, nri_ab_ref), source=0.0_dp)
      ALLOCATE (int_3c_all(nao_ab, nao_ab, nri_ab_ref), source=0.0_dp)
      DO i = 1, 2
         DO j = 1, 2
            CALL compute_m_pq_ab( &
               bs_env, atoms(i), atoms(j), m_pq_block, bs_env%ri_metric)
            m_pq(offset_ref_ri(i) + 1:offset_ref_ri(i) + sizes_ref_ri(i), &
                 offset_ref_ri(j) + 1:offset_ref_ri(j) + sizes_ref_ri(j)) = m_pq_block
            DEALLOCATE (m_pq_block)
            DO k = 1, 2
               CALL build_3c_integral_block_ctx( &
                  int_3c_all, ctx, ws, atom_j=atoms(j), atom_k=atoms(k), &
                  atom_i=atoms(i), j_offset=offset_ao(j), k_offset=offset_ao(k), &
                  i_offset=offset_ref_ri(i), screened=screened)
            END DO
         END DO
      END DO
      kind_a = bs_env%ri_rs%particle_set(iatom)%atomic_kind%kind_number
      kind_b = bs_env%ri_rs%particle_set(jatom)%atomic_kind%kind_number
      CALL add_ri_metric_regularization(bs_env%basis_set_RI(kind_a)%gto_basis_set, &
                                        bs_env%regularization_RI, m_pq)
      CALL add_ri_metric_regularization(bs_env%basis_set_RI(kind_b)%gto_basis_set, &
                                        bs_env%regularization_RI, m_pq, &
                                        offset=sizes_ref_ri(1))
      ALLOCATE (int_3c_ab, source=int_3c_all)
      int_3c_ab(1:sizes_ao(1), 1:sizes_ao(1), :) = 0.0_dp
      int_3c_ab(sizes_ao(1) + 1:, sizes_ao(1) + 1:, :) = 0.0_dp
      ALLOCATE (weights(bs_env%n_mo_retained), b_pq_spin(nri_ab_ref, nri_ab_ref))
      ALLOCATE (b_pq_transition_ab(nri_ab_ref, nri_ab_ref), &
                b_pq_exchange_ab(nri_ab_ref, nri_ab_ref), &
                b_pq_transition_full(nri_ab_ref, nri_ab_ref), &
                b_pq_exchange_full(nri_ab_ref, nri_ab_ref), &
                source=0.0_dp)
      DO spin = 1, bs_env%n_spin
         homo = bs_env%n_occ(spin)
         CALL compute_orbital_weights( &
            eigenval_mo=bs_env%eigenval_scf_Gamma(1:bs_env%n_mo_retained, spin), &
            homo=homo, weights=weights)
         CALL compute_b_pq_transition(int_3c_all, c_mu_n_ab(:, :, spin), weights, homo, b_pq_spin)
         b_pq_transition_full(:, :) = b_pq_transition_full + b_pq_spin
         CALL compute_b_pq_transition(int_3c_ab, c_mu_n_ab(:, :, spin), weights, homo, b_pq_spin)
         b_pq_transition_ab(:, :) = b_pq_transition_ab + b_pq_spin
         CALL compute_b_pq_exchange(int_3c_all, c_mu_n_ab(:, :, spin), &
                                    bs_env%eigenval_scf_Gamma(1:bs_env%n_mo_retained, spin), &
                                    homo, bs_env%auto_ri%occ_energy_window, b_pq_spin)
         b_pq_exchange_full(:, :) = b_pq_exchange_full + b_pq_spin
         CALL compute_b_pq_exchange(int_3c_ab, c_mu_n_ab(:, :, spin), &
                                    bs_env%eigenval_scf_Gamma(1:bs_env%n_mo_retained, spin), &
                                    homo, bs_env%auto_ri%occ_energy_window, b_pq_spin)
         b_pq_exchange_ab(:, :) = b_pq_exchange_ab + b_pq_spin
      END DO
      CALL compute_reference_m_pq_inv_factor(m_pq, x_pi, rank)
      DEALLOCATE (b_pq_spin)
      ALLOCATE (b_pi(nri_ab_ref, rank))
      b_pi(:, :) = matmul(b_pq_transition_full, x_pi)
      norm_transition = sum(x_pi*b_pi)
      b_pi(:, :) = matmul(b_pq_exchange_full, x_pi)
      norm_exchange = sum(x_pi*b_pi)
      ! Eqs. (10)-(12): B_PQ^AB = (B_PQ^{AB,C}/N_C^{AB}+B_PQ^{AB,X}/N_X^{AB})/2.
      IF (norm_transition > 0.0_dp) THEN
         b_pq(:, :) = b_pq + 0.5_dp*b_pq_transition_ab/norm_transition
      END IF
      IF (norm_exchange > 0.0_dp) THEN
         b_pq(:, :) = b_pq + 0.5_dp*b_pq_exchange_ab/norm_exchange
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_b_pq_ab
 
! **************************************************************************************************
!> \brief Builds lists of AB blocks whose columns belong to each atom; B=A denotes an AA block.
!> \param bs_env ...
!> \param lists ...
! **************************************************************************************************
   SUBROUTINE build_ab_block_lists(bs_env, lists)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(ab_block_list_type), ALLOCATABLE, INTENT(OUT) :: lists(:)
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'build_AB_block_lists'
 
      INTEGER                                            :: ab_block, atom_a, atom_b, handle, &
                                                            position, ri_size_to_a, size_opt_ri_ab
      INTEGER, ALLOCATABLE                               :: n_ab_blocks_by_atom(:)
 
      CALL timeset(routinen, handle)
 
      ALLOCATE (lists(bs_env%n_atom))
      ALLOCATE (n_ab_blocks_by_atom(bs_env%n_atom), source=0)
      DO ab_block = 1, bs_env%auto_ri%AB_block_count
         atom_a = bs_env%auto_ri%AB_atom_A(ab_block)
         atom_b = bs_env%auto_ri%AB_atom_B(ab_block)
         size_opt_ri_ab = bs_env%auto_ri%AB_size_opt_RI(ab_block)
         ri_size_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(ab_block)
         IF (ri_size_to_a > 0) n_ab_blocks_by_atom(atom_a) = n_ab_blocks_by_atom(atom_a) + 1
         IF (atom_b /= atom_a .AND. ri_size_to_a < size_opt_ri_ab) THEN
            n_ab_blocks_by_atom(atom_b) = n_ab_blocks_by_atom(atom_b) + 1
         END IF
      END DO
 
      DO atom_a = 1, bs_env%n_atom
         ALLOCATE (lists(atom_a)%block(n_ab_blocks_by_atom(atom_a)))
         ALLOCATE (lists(atom_a)%first_p_AB(n_ab_blocks_by_atom(atom_a)))
         ALLOCATE (lists(atom_a)%first_p_atom(n_ab_blocks_by_atom(atom_a)))
         ALLOCATE (lists(atom_a)%size_opt_RI(n_ab_blocks_by_atom(atom_a)))
      END DO
      n_ab_blocks_by_atom = 0
      DO ab_block = 1, bs_env%auto_ri%AB_block_count
         atom_a = bs_env%auto_ri%AB_atom_A(ab_block)
         atom_b = bs_env%auto_ri%AB_atom_B(ab_block)
         size_opt_ri_ab = bs_env%auto_ri%AB_size_opt_RI(ab_block)
         ri_size_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(ab_block)
         IF (ri_size_to_a > 0) THEN
            position = n_ab_blocks_by_atom(atom_a) + 1
            n_ab_blocks_by_atom(atom_a) = position
            lists(atom_a)%block(position) = ab_block
            lists(atom_a)%first_p_AB(position) = 1
            lists(atom_a)%first_p_atom(position) = bs_env%auto_ri%AB_first_p_A(ab_block)
            lists(atom_a)%size_opt_RI(position) = ri_size_to_a
         END IF
         IF (atom_b /= atom_a .AND. ri_size_to_a < size_opt_ri_ab) THEN
            position = n_ab_blocks_by_atom(atom_b) + 1
            n_ab_blocks_by_atom(atom_b) = position
            lists(atom_b)%block(position) = ab_block
            lists(atom_b)%first_p_AB(position) = ri_size_to_a + 1
            lists(atom_b)%first_p_atom(position) = &
               bs_env%auto_ri%AB_first_p_B(ab_block)
            lists(atom_b)%size_opt_RI(position) = size_opt_ri_ab - ri_size_to_a
         END IF
      END DO
      DO atom_a = 1, bs_env%n_atom
         IF (sum(lists(atom_a)%size_opt_RI) /= bs_env%auto_ri%sizes_opt_RI(atom_a)) THEN
            cpabort("AUTO_RI AB-block columns do not fill the optimized atom block")
         END IF
      END DO
      DEALLOCATE (n_ab_blocks_by_atom)
 
      CALL timestop(handle)
 
   END SUBROUTINE build_ab_block_lists
 
! **************************************************************************************************
!> \brief Returns the location of atom in atoms(:count), or zero if it is absent.
!> \param atom ...
!> \param atoms ...
!> \param count ...
!> \return ...
! **************************************************************************************************
   PURE INTEGER FUNCTION find_atom_position(atom, atoms, count) RESULT(position)
      INTEGER, INTENT(IN)                                :: atom
      INTEGER, DIMENSION(:), INTENT(IN)                  :: atoms
      INTEGER, INTENT(IN)                                :: count
 
      INTEGER                                            :: i
 
      position = 0
      DO i = 1, count
         IF (atoms(i) == atom) THEN
            position = i
            RETURN
         END IF
      END DO
 
   END FUNCTION find_atom_position
 
! **************************************************************************************************
!> \brief Collects U_Pp from the large reference basis for one optimized atom block.
!> \param bs_env ...
!> \param atom ...
!> \param list ...
!> \param atom_U ...
! **************************************************************************************************
   SUBROUTINE build_optimized_atom_u(bs_env, atom, list, atom_U)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      INTEGER, INTENT(IN)                                :: atom
      TYPE(ab_block_list_type), INTENT(IN)               :: list
      TYPE(optimized_atom_u_type), INTENT(OUT)           :: atom_u
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'build_optimized_atom_U'
 
      INTEGER :: ab_block, atom_a, atom_b, first, first_p_ab, first_p_atom, first_row, handle, i, &
         last, nri_ref_a, nri_ref_block, nsupport, position, size_opt_ri_ab
      INTEGER, ALLOCATABLE                               :: support_atoms(:)
      REAL(kind=dp), ALLOCATABLE                         :: u_pp_ab(:, :)
 
      CALL timeset(routinen, handle)
 
      ALLOCATE (support_atoms(2*SIZE(list%block)))
      nsupport = 0
      DO i = 1, SIZE(list%block)
         ab_block = list%block(i)
         atom_a = bs_env%auto_ri%AB_atom_A(ab_block)
         atom_b = bs_env%auto_ri%AB_atom_B(ab_block)
         position = find_atom_position(atom_a, support_atoms, nsupport)
         IF (position == 0) THEN
            nsupport = nsupport + 1
            support_atoms(nsupport) = atom_a
         END IF
         IF (atom_b /= atom_a) THEN
            position = find_atom_position(atom_b, support_atoms, nsupport)
            IF (position == 0) THEN
               nsupport = nsupport + 1
               support_atoms(nsupport) = atom_b
            END IF
         END IF
      END DO
      cpassert(nsupport > 0)
      ALLOCATE (atom_u%ref_atom(nsupport), atom_u%first_P(nsupport))
      atom_u%ref_atom(:) = support_atoms(:nsupport)
      first_row = 1
      DO i = 1, nsupport
         atom_u%first_P(i) = first_row
         first_row = first_row + get_ref_ri_size(bs_env, atom_u%ref_atom(i))
      END DO
      ALLOCATE (atom_u%U_Pp(first_row - 1, bs_env%auto_ri%sizes_opt_RI(atom)), &
                source=0.0_dp)
 
      DO i = 1, SIZE(list%block)
         ab_block = list%block(i)
         atom_a = bs_env%auto_ri%AB_atom_A(ab_block)
         atom_b = bs_env%auto_ri%AB_atom_B(ab_block)
         nri_ref_block = bs_env%auto_ri%AB_size_ref_RI(ab_block)
         size_opt_ri_ab = bs_env%auto_ri%AB_size_opt_RI(ab_block)
         first = bs_env%auto_ri%U_Pp_AB_offset(ab_block)
         last = first + nri_ref_block*size_opt_ri_ab - 1
         ALLOCATE (u_pp_ab(nri_ref_block, size_opt_ri_ab))
         u_pp_ab(:, :) = reshape(bs_env%auto_ri%U_Pp_AB(first:last), &
                                 [nri_ref_block, size_opt_ri_ab])
         first_p_ab = list%first_p_AB(i)
         first_p_atom = list%first_p_atom(i)
         size_opt_ri_ab = list%size_opt_RI(i)
         IF (first_p_ab < 1 .OR. &
             first_p_ab + size_opt_ri_ab - 1 > SIZE(u_pp_ab, 2)) THEN
            cpabort("AUTO_RI AB-block column range is invalid")
         END IF
         IF (first_p_atom < 1 .OR. &
             first_p_atom + size_opt_ri_ab - 1 > bs_env%auto_ri%sizes_opt_RI(atom)) THEN
            cpabort("AUTO_RI optimized atom-block column range is invalid")
         END IF
 
         nri_ref_a = get_ref_ri_size(bs_env, atom_a)
         position = find_atom_position(atom_a, atom_u%ref_atom, nsupport)
         cpassert(position > 0)
         first_row = atom_u%first_P(position)
         atom_u%U_Pp(first_row:first_row + nri_ref_a - 1, &
                     first_p_atom:first_p_atom + size_opt_ri_ab - 1) = &
            u_pp_ab(1:nri_ref_a, first_p_ab:first_p_ab + size_opt_ri_ab - 1)
         IF (atom_b /= atom_a) THEN
            position = find_atom_position(atom_b, atom_u%ref_atom, nsupport)
            cpassert(position > 0)
            first_row = atom_u%first_P(position)
            atom_u%U_Pp(first_row:first_row + get_ref_ri_size(bs_env, atom_b) - 1, &
                        first_p_atom:first_p_atom + size_opt_ri_ab - 1) = &
               u_pp_ab(nri_ref_a + 1:, first_p_ab:first_p_ab + size_opt_ri_ab - 1)
         END IF
         DEALLOCATE (u_pp_ab)
      END DO
      DEALLOCATE (support_atoms)
 
      CALL timestop(handle)
 
   END SUBROUTINE build_optimized_atom_u
 
! **************************************************************************************************
!> \brief Adds REGULARIZATION_RI to one atom-local reference metric block.
!> \param bs_env ...
!> \param atom ...
!> \param M_PQ atom-local metric block in the large reference RI basis
! **************************************************************************************************
   SUBROUTINE regularize_m_pq_block(bs_env, atom, M_PQ)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      INTEGER, INTENT(IN)                                :: atom
      REAL(kind=dp), DIMENSION(:, :), INTENT(INOUT)      :: m_pq
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'regularize_M_PQ_block'
 
      INTEGER                                            :: handle, i, iset, kind, nset
      INTEGER, DIMENSION(:), POINTER                     :: npgf, nsgf_set
      INTEGER, DIMENSION(:, :), POINTER                  :: first_sgf
      REAL(kind=dp), DIMENSION(:, :), POINTER            :: zet
 
      CALL timeset(routinen, handle)
 
      cpassert(SIZE(m_pq, 1) == SIZE(m_pq, 2))
      kind = bs_env%ri_rs%particle_set(atom)%atomic_kind%kind_number
      CALL get_gto_basis_set(bs_env%basis_set_RI(kind)%gto_basis_set, &
                             first_sgf=first_sgf, npgf=npgf, nset=nset, &
                             nsgf_set=nsgf_set, zet=zet)
      DO iset = 1, nset
         DO i = first_sgf(1, iset), first_sgf(1, iset) + nsgf_set(iset) - 1
            m_pq(i, i) = m_pq(i, i) + &
                         bs_env%regularization_RI* &
                         max(1.0_dp, 1.0_dp/minval(zet(1:npgf(iset), iset)))
         END DO
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE regularize_m_pq_block
 
! **************************************************************************************************
!> \brief Directly contracts one optimized atom-pair metric block.
!>
!>        M_pq^{AB} = Σ_IJ (U_Ip^A)^T M_IJ U_Jq^B. Each reference
!>        atom-pair block M_IJ is discarded immediately after this contraction.
!> \param bs_env ...
!> \param atom_U_a contractions U_Ip for optimized block A
!> \param atom_U_b contractions U_Jq for optimized block B
!> \param coulomb true for Coulomb; false for the GW fitting operator
!> \param symmetric true when both transforms describe the same optimized block
!> \param matrix_optimized contracted M_pq or V_pq block in the optimized RI basis
! **************************************************************************************************
   SUBROUTINE compute_optimized_ri_matrix_block( &
      bs_env, atom_U_a, atom_U_b, coulomb, symmetric, matrix_optimized)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(optimized_atom_u_type), INTENT(IN)            :: atom_u_a, atom_u_b
      LOGICAL, INTENT(IN)                                :: coulomb, symmetric
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
         INTENT(OUT)                                     :: matrix_optimized
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_optimized_RI_matrix_block'
 
      INTEGER                                            :: atom_a, atom_b, first_a, first_b, &
                                                            handle, i, j, last_a, last_b, &
                                                            nri_small_a, nri_small_b
      REAL(kind=dp), ALLOCATABLE :: matrix_pq_contribution(:, :), matrix_reference(:, :), &
         matrix_reference_times_u_b(:, :)
 
      CALL timeset(routinen, handle)
 
      nri_small_a = SIZE(atom_u_a%U_Pp, 2)
      nri_small_b = SIZE(atom_u_b%U_Pp, 2)
      IF (symmetric) THEN
         cpassert(nri_small_a == nri_small_b)
      END IF
      ALLOCATE (matrix_optimized(nri_small_a, nri_small_b), source=0.0_dp)
      DO i = 1, SIZE(atom_u_a%ref_atom)
         atom_a = atom_u_a%ref_atom(i)
         first_a = atom_u_a%first_P(i)
         last_a = first_a + get_ref_ri_size(bs_env, atom_a) - 1
         DO j = 1, SIZE(atom_u_b%ref_atom)
            IF (symmetric .AND. j > i) cycle
            atom_b = atom_u_b%ref_atom(j)
            first_b = atom_u_b%first_P(j)
            last_b = first_b + get_ref_ri_size(bs_env, atom_b) - 1
            IF (coulomb) THEN
               CALL compute_m_pq_ab(bs_env, atom_a, atom_b, matrix_reference)
            ELSE
               CALL compute_m_pq_ab( &
                  bs_env, atom_a, atom_b, matrix_reference, bs_env%ri_metric)
               IF (atom_a == atom_b .AND. bs_env%regularization_RI > 0.0_dp) THEN
                  CALL regularize_m_pq_block(bs_env, atom_a, matrix_reference)
               END IF
            END IF
            ALLOCATE (matrix_reference_times_u_b(SIZE(matrix_reference, 1), nri_small_b))
            matrix_reference_times_u_b(:, :) = &
               matmul(matrix_reference, atom_u_b%U_Pp(first_b:last_b, :))
            ALLOCATE (matrix_pq_contribution(nri_small_a, nri_small_b))
            matrix_pq_contribution(:, :) = &
               matmul(transpose(atom_u_a%U_Pp(first_a:last_a, :)), &
                      matrix_reference_times_u_b)
            matrix_optimized(:, :) = matrix_optimized + matrix_pq_contribution
            IF (symmetric .AND. i /= j) THEN
               matrix_optimized(:, :) = matrix_optimized + transpose(matrix_pq_contribution)
            END IF
            DEALLOCATE (matrix_pq_contribution, matrix_reference, matrix_reference_times_u_b)
         END DO
      END DO
      IF (.NOT. all(ieee_is_finite(matrix_optimized))) THEN
         cpabort("AUTO_RI optimized metric block contains invalid values")
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_optimized_ri_matrix_block
 
! **************************************************************************************************
!> \brief Builds M_opt or V_opt directly from contracted atom-pair blocks.
!> \param qs_env ...
!> \param bs_env ...
!> \param lists ...
!> \param matrix_pq contracted M_pq or V_pq in the optimized RI basis
!> \param coulomb true for Coulomb; false for the GW fitting operator
! **************************************************************************************************
   SUBROUTINE compute_contracted_ri_matrix(qs_env, bs_env, lists, matrix_pq, coulomb)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(ab_block_list_type), DIMENSION(:), INTENT(IN) :: lists
      TYPE(cp_fm_type), INTENT(OUT)                      :: matrix_pq
      LOGICAL, INTENT(IN)                                :: coulomb
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_contracted_RI_matrix'
 
      INTEGER                                            :: handle, iatom, jatom, npcol, nprow, &
                                                            pair_index, ri_size_opt
      INTEGER, POINTER                                   :: col_dist(:), row_dist(:), sizes_opt_ri(:)
      LOGICAL                                            :: atom_u_a_ready
      REAL(kind=dp), ALLOCATABLE                         :: matrix_pq_block(:, :)
      TYPE(dbcsr_distribution_type)                      :: dist
      TYPE(dbcsr_distribution_type), POINTER             :: template
      TYPE(dbcsr_type)                                   :: blocks
      TYPE(optimized_atom_u_type)                        :: atom_u_a, atom_u_b
 
      CALL timeset(routinen, handle)
 
      ALLOCATE (sizes_opt_ri(bs_env%n_atom), row_dist(bs_env%n_atom), col_dist(bs_env%n_atom))
      sizes_opt_ri(:) = bs_env%auto_ri%sizes_opt_RI
      CALL get_qs_env(qs_env, dbcsr_dist=template)
      CALL dbcsr_distribution_get(template, nprows=nprow, npcols=npcol)
      DO iatom = 1, SIZE(sizes_opt_ri)
         row_dist(iatom) = mod(iatom - 1, nprow)
         col_dist(iatom) = mod(iatom - 1, npcol)
      END DO
      ri_size_opt = sum(sizes_opt_ri)
      CALL create_distributed_matrix(qs_env, matrix_pq, ri_size_opt, ri_size_opt)
      CALL dbcsr_distribution_new(dist, template=template, row_dist=row_dist, col_dist=col_dist)
      CALL dbcsr_create(blocks, name='AUTO_RI contracted metric', dist=dist, &
                        matrix_type=dbcsr_type_symmetric, &
                        row_blk_size=sizes_opt_ri, col_blk_size=sizes_opt_ri)
 
      pair_index = 0
      DO iatom = 1, SIZE(sizes_opt_ri)
         atom_u_a_ready = .false.
         DO jatom = 1, iatom
            pair_index = pair_index + 1
            IF (mod(pair_index - 1, bs_env%para_env%num_pe) /= bs_env%para_env%mepos) cycle
            IF (.NOT. atom_u_a_ready) THEN
               CALL build_optimized_atom_u(bs_env, iatom, lists(iatom), atom_u_a)
               atom_u_a_ready = .true.
            END IF
            IF (iatom == jatom) THEN
               CALL compute_optimized_ri_matrix_block( &
                  bs_env, atom_u_a, atom_u_a, coulomb, .true., &
                  matrix_pq_block)
               matrix_pq_block(:, :) = 0.5_dp*(matrix_pq_block + transpose(matrix_pq_block))
            ELSE
               CALL build_optimized_atom_u(bs_env, jatom, lists(jatom), atom_u_b)
               CALL compute_optimized_ri_matrix_block( &
                  bs_env, atom_u_a, atom_u_b, coulomb, .false., &
                  matrix_pq_block)
            END IF
            CALL dbcsr_put_block(blocks, iatom, jatom, matrix_pq_block)
            DEALLOCATE (matrix_pq_block)
         END DO
      END DO
      CALL dbcsr_finalize(blocks)
      CALL copy_dbcsr_to_fm(blocks, matrix_pq)
      CALL dbcsr_release(blocks)
      CALL dbcsr_distribution_release(dist)
      DEALLOCATE (sizes_opt_ri, row_dist, col_dist)
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_contracted_ri_matrix
 
! **************************************************************************************************
!> \brief Computes the pseudoinverse of the optimized fitting metric M_opt.
!>
!>        The generalized problem M_opt X=M_AA X λ separates linearly
!>        dependent directions. The result is M_opt^+=X diag(1/λ) X^T.
!> \param qs_env ...
!> \param bs_env ...
!> \param M_pq optimized RI fitting metric
! **************************************************************************************************
   SUBROUTINE compute_m_pq_inv(qs_env, bs_env, M_pq)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(cp_fm_type), INTENT(INOUT)                    :: m_pq
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'compute_M_pq_inv'
 
      INTEGER                                            :: handle, ncol_local, nrow_local, &
                                                            ri_size_opt
      INTEGER, DIMENSION(:), POINTER                     :: col_indices, row_indices
      REAL(kind=dp), ALLOCATABLE                         :: eigenval(:), factors(:)
      REAL(kind=dp), CONTIGUOUS, DIMENSION(:, :), &
         POINTER                                         :: m_pq_aa_local, m_pq_local
      TYPE(cp_fm_type)                                   :: m_pq_aa, work, x_pq
 
      CALL timeset(routinen, handle)
 
      ri_size_opt = sum(bs_env%auto_ri%sizes_opt_RI)
      CALL create_distributed_matrix(qs_env, bs_env%auto_ri%M_pq_inv, &
                                     ri_size_opt, ri_size_opt)
      CALL create_distributed_matrix(qs_env, x_pq, ri_size_opt, ri_size_opt)
      CALL create_distributed_matrix(qs_env, m_pq_aa, ri_size_opt, ri_size_opt)
      CALL create_distributed_matrix(qs_env, work, ri_size_opt, ri_size_opt)
      CALL cp_fm_get_info(m_pq, nrow_local=nrow_local, ncol_local=ncol_local, &
                          row_indices=row_indices, col_indices=col_indices, &
                          local_data=m_pq_local)
      CALL cp_fm_get_info(m_pq_aa, local_data=m_pq_aa_local)
      m_pq_aa_local = 0.0_dp
      CALL copy_m_pq_aa_blocks(m_pq_local(:nrow_local, :ncol_local), &
                               m_pq_aa_local(:nrow_local, :ncol_local), &
                               row_indices(:nrow_local), col_indices(:ncol_local), &
                               bs_env%auto_ri%sizes_opt_RI)
 
      ALLOCATE (eigenval(ri_size_opt))
      CALL cp_fm_geeig(m_pq, m_pq_aa, x_pq, eigenval, work)
      CALL compute_optimized_m_pq_inv_factors(eigenval, factors)
      CALL cp_fm_column_scale(x_pq, factors)
      CALL cp_fm_syrk("U", "N", ri_size_opt, 1.0_dp, x_pq, 1, 1, 0.0_dp, &
                      bs_env%auto_ri%M_pq_inv)
      CALL cp_fm_uplo_to_full(bs_env%auto_ri%M_pq_inv, work)
 
      DEALLOCATE (eigenval, factors)
      CALL cp_fm_release(x_pq)
      CALL cp_fm_release(m_pq_aa)
      CALL cp_fm_release(work)
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_m_pq_inv
 
! **************************************************************************************************
!> \brief Computes M_opt and V_opt without constructing M_ref or V_ref.
!> \param qs_env ...
!> \param bs_env ...
! **************************************************************************************************
   SUBROUTINE compute_m_pq_inv_and_v_pq(qs_env, bs_env)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'compute_M_pq_inv_and_V_pq'
 
      INTEGER                                            :: handle
      TYPE(ab_block_list_type), ALLOCATABLE              :: lists(:)
      TYPE(cp_fm_type)                                   :: m_pq
 
      CALL timeset(routinen, handle)
 
      CALL build_ab_block_lists(bs_env, lists)
      CALL initialize_auto_ri_integrals(bs_env)
      CALL compute_contracted_ri_matrix(qs_env, bs_env, lists, m_pq, .false.)
      CALL compute_m_pq_inv(qs_env, bs_env, m_pq)
      CALL cp_fm_release(m_pq)
      CALL compute_contracted_ri_matrix(qs_env, bs_env, lists, bs_env%auto_ri%V_pq, .true.)
 
      ! Activate the optimized per-atom RI block sizes for RI-RS.
      CALL set_optimized_ri_basis_sizes(bs_env)
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_m_pq_inv_and_v_pq
 
! **************************************************************************************************
!> \brief Computes the atom-local two-center RI metric (P|Q) and the
!>        three-center integrals (μν|P) used in Eqs. (3) and (4).
!> \param ao_basis ...
!> \param ri_basis ...
!> \param potential fitting operator used for both integral types
!> \param M_PQ two-center metric of the large reference RI basis
!> \param Int_3c three-center integrals (μν|P)
! **************************************************************************************************
   SUBROUTINE compute_int_3c(ao_basis, ri_basis, potential, M_PQ, Int_3c)
      TYPE(gto_basis_set_type), INTENT(IN)               :: ao_basis, ri_basis
      TYPE(coulomb_operator_type), INTENT(IN)            :: potential
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
         INTENT(OUT)                                     :: m_pq
      REAL(kind=dp), INTENT(OUT)                         :: int_3c(:, :, :)
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'compute_int_3c'
 
      INTEGER                                            :: handle, iset, jset, kset, max_ao_l, &
                                                            max_l, nao, ncoa, ncob, ncoc, nri, &
                                                            nsgfa, nsgfb, nsgfc, sgfa, sgfb, sgfc
      INTEGER, DIMENSION(:), POINTER                     :: lmax_a, lmax_c, lmin_a, lmin_c, npgf_a, &
                                                            npgf_c, nsgf_set_a, nsgf_set_c
      INTEGER, DIMENSION(:, :), POINTER                  :: first_sgf_a, first_sgf_c
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: int_2c_cart, int_2c_spherical
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: int_3c_cart, int_3c_spherical_all, &
                                                            int_3c_spherical_block
      REAL(kind=dp), DIMENSION(3)                        :: center
      REAL(kind=dp), DIMENSION(:), POINTER               :: radius_a, radius_b, radius_c
      REAL(kind=dp), DIMENSION(:, :), POINTER            :: pgf_radius_a, pgf_radius_c, sphi_a, &
                                                            sphi_c, zeta_a, zeta_c
      TYPE(coulomb_integral_context_type)                :: context
 
      CALL timeset(routinen, handle)
 
      CALL get_gto_basis_set(ao_basis, nsgf=nao, lmax=lmax_a, lmin=lmin_a, &
                             npgf=npgf_a, nsgf_set=nsgf_set_a, first_sgf=first_sgf_a, &
                             pgf_radius=pgf_radius_a, sphi=sphi_a, zet=zeta_a)
      CALL get_gto_basis_set(ri_basis, nsgf=nri, lmax=lmax_c, lmin=lmin_c, &
                             npgf=npgf_c, nsgf_set=nsgf_set_c, first_sgf=first_sgf_c, &
                             pgf_radius=pgf_radius_c, sphi=sphi_c, zet=zeta_c)
      cpassert(nao > 0 .AND. nri > 0)
      max_ao_l = maxval(lmax_a)
      max_l = max(max_ao_l, maxval(lmax_c))
      CALL init_orbital_pointers(max(max_l, 2*max_ao_l))
      CALL init_spherical_harmonics(max_l, -1)
      IF (potential%potential_type == do_potential_truncated) THEN
         IF (active_integral_library%coulomb2_library /= library_libint .OR. &
             active_integral_library%coulomb3_library /= library_libint) THEN
            CALL cp_abort(__location__, &
                          "AUTO_RI TRUNCATED fitting requires LIBINT for "// &
                          "COULOMB_2C and COULOMB_3C")
         END IF
         CALL coulomb_integral_init(context, max_am_2c=max(max_l, 2*max_ao_l), max_am_3c=max_l)
      ELSE
         CALL coulomb_integral_init(context, max_am_2c=max_l, max_am_3c=max_l)
      END IF
      center = 0.0_dp
 
      ALLOCATE (m_pq(nri, nri), int_3c_spherical_all(nao, nao, nri))
      m_pq = 0.0_dp
      int_3c_spherical_all = 0.0_dp
 
      DO iset = 1, ri_basis%nset
         ncoa = npgf_c(iset)*ncoset(lmax_c(iset))
         sgfa = first_sgf_c(1, iset)
         nsgfa = nsgf_set_c(iset)
         radius_a => pgf_radius_c(1:npgf_c(iset), iset)
         DO jset = 1, ri_basis%nset
            ncob = npgf_c(jset)*ncoset(lmax_c(jset))
            sgfb = first_sgf_c(1, jset)
            nsgfb = nsgf_set_c(jset)
            radius_b => pgf_radius_c(1:npgf_c(jset), jset)
            ALLOCATE (int_2c_cart(ncoa, ncob), int_2c_spherical(nsgfa, nsgfb))
            int_2c_cart = 0.0_dp
            CALL compute_coulomb_2c(context, lmin_c(iset), lmax_c(iset), &
                                    lmin_c(jset), lmax_c(jset), &
                                    npgf_c(iset), npgf_c(jset), zeta_c(1:npgf_c(iset), iset), &
                                    zeta_c(1:npgf_c(jset), jset), &
                                    radius_a, radius_b, center, center, int_2c_cart, potential)
            CALL ab_contract(int_2c_spherical, int_2c_cart, sphi_c(:, sgfa:), sphi_c(:, sgfb:), &
                             ncoa, ncob, nsgfa, nsgfb)
            m_pq(sgfa:sgfa + nsgfa - 1, sgfb:sgfb + nsgfb - 1) = int_2c_spherical
            DEALLOCATE (int_2c_spherical, int_2c_cart)
         END DO
      END DO
 
      DO iset = 1, ao_basis%nset
         ncoa = npgf_a(iset)*ncoset(lmax_a(iset))
         sgfa = first_sgf_a(1, iset)
         nsgfa = nsgf_set_a(iset)
         radius_a => pgf_radius_a(1:npgf_a(iset), iset)
         DO jset = 1, ao_basis%nset
            ncob = npgf_a(jset)*ncoset(lmax_a(jset))
            sgfb = first_sgf_a(1, jset)
            nsgfb = nsgf_set_a(jset)
            radius_b => pgf_radius_a(1:npgf_a(jset), jset)
            DO kset = 1, ri_basis%nset
               ncoc = npgf_c(kset)*ncoset(lmax_c(kset))
               sgfc = first_sgf_c(1, kset)
               nsgfc = nsgf_set_c(kset)
               radius_c => pgf_radius_c(1:npgf_c(kset), kset)
               ALLOCATE (int_3c_cart(ncoa, ncob, ncoc), &
                         int_3c_spherical_block(nsgfa, nsgfb, nsgfc))
               int_3c_cart = 0.0_dp
               CALL compute_coulomb_3c(context, lmin_a(iset), lmax_a(iset), &
                                       lmin_a(jset), lmax_a(jset), &
                                       lmin_c(kset), lmax_c(kset), npgf_a(iset), npgf_a(jset), &
                                       npgf_c(kset), zeta_a(1:npgf_a(iset), iset), &
                                       zeta_a(1:npgf_a(jset), jset), zeta_c(1:npgf_c(kset), kset), &
                                       radius_a, radius_b, radius_c, center, center, center, &
                                       int_3c_cart, potential)
               CALL abc_contract(int_3c_spherical_block, int_3c_cart, &
                                 sphi_a(:, sgfa:), sphi_a(:, sgfb:), &
                                 sphi_c(:, sgfc:), ncoa, ncob, ncoc, nsgfa, nsgfb, nsgfc)
               int_3c_spherical_all(sgfa:sgfa + nsgfa - 1, sgfb:sgfb + nsgfb - 1, &
                                    sgfc:sgfc + nsgfc - 1) = int_3c_spherical_block
               DEALLOCATE (int_3c_spherical_block, int_3c_cart)
            END DO
         END DO
      END DO
 
      int_3c(:, :, :) = int_3c_spherical_all
      CALL coulomb_integral_cleanup(context)
      IF (.NOT. all(ieee_is_finite(m_pq))) THEN
         cpabort("AUTO_RI metric contains invalid values")
      END IF
      IF (.NOT. all(ieee_is_finite(int_3c))) THEN
         cpabort("AUTO_RI three-center integrals contain invalid values")
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_int_3c
 
! **************************************************************************************************
!> \brief Initializes the integral kernels used by the AUTO_RI optimization.
!> \param bs_env ...
! **************************************************************************************************
   SUBROUTINE initialize_auto_ri_integrals(bs_env)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'initialize_auto_ri_integrals'
 
      CHARACTER(LEN=default_path_length)                 :: table_file
      INTEGER                                            :: handle, ikind, max_l, table_lmax, unit_id
 
      CALL timeset(routinen, handle)
 
      max_l = 0
      DO ikind = 1, SIZE(bs_env%basis_set_RI)
         cpassert(ASSOCIATED(bs_env%basis_set_RI(ikind)%gto_basis_set))
         max_l = max(max_l, maxval(bs_env%basis_set_RI(ikind)%gto_basis_set%lmax))
      END DO
      ! The OS two-center kernel accesses l+1 even without force evaluation.
      CALL init_orbital_pointers(max_l + 1)
      CALL init_spherical_harmonics(max_l, -1)
 
      IF (bs_env%ri_metric%potential_type == do_potential_truncated) THEN
         ! The truncated-Coulomb kernel needs auxiliary functions through l_A+l_B+1.
         table_lmax = 2*max_l + 1
         IF (get_lmax_init() < table_lmax) THEN
            unit_id = -1
            IF (bs_env%para_env%is_source()) THEN
               table_file = discover_file(bs_env%ri_metric%filename)
               CALL open_file(unit_number=unit_id, file_name=trim(table_file))
            END IF
            CALL init_t_c_g0(table_lmax, unit_id, bs_env%para_env%mepos, bs_env%para_env)
            IF (bs_env%para_env%is_source()) CALL close_file(unit_id)
         END IF
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE initialize_auto_ri_integrals
 
! **************************************************************************************************
!> \brief Computes the atom-pair two-center RI metric matrix (P_A|Q_B).
!> \param bs_env ...
!> \param atom_A ...
!> \param atom_B ...
!> \param M_PQ metric block in the large reference RI basis
!> \param ri_potential ...
! **************************************************************************************************
   SUBROUTINE compute_m_pq_ab(bs_env, atom_A, atom_B, M_PQ, ri_potential)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      INTEGER, INTENT(IN)                                :: atom_a, atom_b
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
         INTENT(OUT)                                     :: m_pq
      TYPE(libint_potential_type), INTENT(IN), OPTIONAL  :: ri_potential
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'compute_M_PQ_AB'
 
      INTEGER                                            :: atom_1, atom_2, handle, kind_1, kind_2
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: m_pq_canonical
      TYPE(libint_potential_type)                        :: potential
 
      CALL timeset(routinen, handle)
 
      cpassert(ASSOCIATED(bs_env%ri_rs%particle_set))
      cpassert(atom_a >= 1 .AND. atom_a <= bs_env%n_atom)
      cpassert(atom_b >= 1 .AND. atom_b <= bs_env%n_atom)
      IF (atom_a >= atom_b) THEN
         atom_1 = atom_a
         atom_2 = atom_b
      ELSE
         atom_1 = atom_b
         atom_2 = atom_a
      END IF
      kind_1 = bs_env%ri_rs%particle_set(atom_1)%atomic_kind%kind_number
      kind_2 = bs_env%ri_rs%particle_set(atom_2)%atomic_kind%kind_number
 
      IF (PRESENT(ri_potential)) THEN
         cpassert(ri_potential%potential_type == do_potential_truncated)
         potential = ri_potential
      ELSE
         potential%potential_type = do_potential_coulomb
      END IF
      CALL compute_m_pq( &
         basis_a=bs_env%basis_set_RI(kind_1)%gto_basis_set, &
         basis_b=bs_env%basis_set_RI(kind_2)%gto_basis_set, &
         center_a=bs_env%ri_rs%particle_set(atom_1)%r, &
         center_b=bs_env%ri_rs%particle_set(atom_2)%r, &
         potential=potential, m_pq=m_pq_canonical)
 
      IF (atom_a >= atom_b) THEN
         CALL move_alloc(m_pq_canonical, m_pq)
      ELSE
         ALLOCATE (m_pq(SIZE(m_pq_canonical, 2), SIZE(m_pq_canonical, 1)))
         m_pq(:, :) = transpose(m_pq_canonical)
         DEALLOCATE (m_pq_canonical)
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_m_pq_ab
 
! **************************************************************************************************
!> \brief Computes one contracted two-center RI metric block (P_A|Q_B).
!> \param basis_a ...
!> \param basis_b ...
!> \param center_a ...
!> \param center_b ...
!> \param potential ...
!> \param M_PQ metric block in the large reference RI basis
! **************************************************************************************************
   SUBROUTINE compute_m_pq(basis_a, basis_b, center_a, center_b, potential, M_PQ)
      TYPE(gto_basis_set_type), POINTER                  :: basis_a, basis_b
      REAL(kind=dp), DIMENSION(3), INTENT(IN)            :: center_a, center_b
      TYPE(libint_potential_type), INTENT(IN)            :: potential
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
         INTENT(OUT)                                     :: m_pq
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'compute_M_PQ'
 
      INTEGER                                            :: handle
 
      CALL timeset(routinen, handle)
 
      ALLOCATE (m_pq(basis_a%nsgf, basis_b%nsgf), source=0.0_dp)
      SELECT CASE (potential%potential_type)
      CASE (do_potential_truncated)
         CALL int_operators_r12_ab_os(operator_truncated, m_pq, rab=center_b - center_a, &
                                      fba=basis_a, fbb=basis_b, &
                                      r_cutoff=potential%cutoff_radius, calculate_forces=.false.)
      CASE (do_potential_coulomb)
         CALL int_operators_r12_ab_os(operator_coulomb, m_pq, rab=center_b - center_a, &
                                      fba=basis_a, fbb=basis_b, calculate_forces=.false.)
      CASE DEFAULT
         cpabort("AUTO_RI two-center metric requires COULOMB or TRUNCATED")
      END SELECT
 
      CALL timestop(handle)
 
   END SUBROUTINE compute_m_pq
 
! **************************************************************************************************
!> \brief Returns the reference RI basis size of one atom.
!> \param bs_env ...
!> \param iatom ...
!> \return ...
! **************************************************************************************************
   INTEGER FUNCTION get_ref_ri_size(bs_env, iatom) RESULT(nRI_ref)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      INTEGER, INTENT(IN)                                :: iatom
 
      INTEGER                                            :: ikind
 
      ikind = bs_env%ri_rs%particle_set(iatom)%atomic_kind%kind_number
      nri_ref = bs_env%basis_set_RI(ikind)%gto_basis_set%nsgf
 
   END FUNCTION get_ref_ri_size
 
END MODULE gw_auto_ri_generate_ri_basis