Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Builds the shared RI-RS grid, AO collocation and fitting representation.
10 : !> \par History
11 : !> 09.2026 created Jan Wilhelm
12 : !> 09.2026 moved code by Ritaj Tyagi from gw_ri_rs_non_periodic.F
13 : !> 09.2026 added routine to compute Z_lP for automatic RI optimization
14 : ! **************************************************************************************************
15 : MODULE ri_rs_setup
16 : USE ai_contraction_sphi, ONLY: abc_contract_xsmm
17 : USE atomic_kind_types, ONLY: atomic_kind_type
18 : USE basis_set_types, ONLY: get_gto_basis_set
19 : USE cell_types, ONLY: cell_type,&
20 : pbc
21 : USE cp_blacs_env, ONLY: cp_blacs_env_create,&
22 : cp_blacs_env_release,&
23 : cp_blacs_env_type
24 : USE cp_dbcsr_api, ONLY: &
25 : dbcsr_binary_read, dbcsr_binary_write, dbcsr_create, dbcsr_distribution_get, &
26 : dbcsr_distribution_new, dbcsr_distribution_release, dbcsr_distribution_type, dbcsr_filter, &
27 : dbcsr_finalize, dbcsr_get_block_p, dbcsr_get_info, dbcsr_multiply, dbcsr_put_block, &
28 : dbcsr_release, dbcsr_set, dbcsr_type, dbcsr_type_no_symmetry
29 : USE cp_dbcsr_contrib, ONLY: dbcsr_reserve_all_blocks
30 : USE cp_log_handling, ONLY: cp_get_default_logger,&
31 : cp_logger_type
32 : USE cp_output_handling, ONLY: cp_p_file,&
33 : cp_print_key_should_output
34 : USE gw_auto_ri_types, ONLY: auto_ri_type
35 : USE gw_utils_compute_integrals, ONLY: build_3c_integral_block_ctx,&
36 : gw_3c_ctx_create,&
37 : gw_3c_ctx_release,&
38 : gw_3c_ctx_type,&
39 : gw_3c_ws_create,&
40 : gw_3c_ws_release,&
41 : gw_3c_ws_type
42 : USE input_section_types, ONLY: section_vals_type
43 : USE kinds, ONLY: dp,&
44 : max_line_length
45 : USE libint_2c_3c, ONLY: eri_3center
46 : USE machine, ONLY: m_flush,&
47 : m_walltime
48 : USE message_passing, ONLY: mp_mem_avail_per_rank_GB,&
49 : mp_para_env_type
50 : USE orbital_pointers, ONLY: ncoset
51 : USE particle_types, ONLY: particle_type
52 : USE physcon, ONLY: angstrom
53 : USE post_scf_bandstructure_types, ONLY: post_scf_bandstructure_type
54 : USE qs_environment_types, ONLY: get_qs_env,&
55 : qs_environment_type
56 : USE qs_kind_types, ONLY: qs_kind_type
57 : USE ri_rs_fitting, ONLY: solve_ri_rs_fit,&
58 : store_Z_lP_columns
59 : USE ri_rs_grid_setup, ONLY: setup_ri_rs_grid
60 : USE ri_rs_matrices, ONLY: atomic_basis_at_grid_point,&
61 : print_matrix_occupation
62 : USE ri_rs_types, ONLY: ri_rs_env
63 : USE ri_rs_utils, ONLY: evaluate_ao_basis_on_points,&
64 : evaluate_ao_on_points,&
65 : precompute_ri_rs_radii
66 : USE util, ONLY: locate,&
67 : sort
68 : #include "./base/base_uses.f90"
69 :
70 : IMPLICIT NONE
71 : PRIVATE
72 :
73 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'ri_rs_setup'
74 :
75 : PUBLIC :: compute_Z_lP, ensure_ri_rs_representation
76 :
77 : CONTAINS
78 :
79 : ! **************************************************************************************************
80 : !> \brief Populate the RI-RS grid, AO collocation and auxiliary fitting coefficients once.
81 : !> Periodic systems use the same neighbouring-image convention for phi and Z.
82 : !> \param qs_env electronic-structure environment
83 : !> \param bs_env geometry, RI-RS settings and owned output matrices
84 : ! **************************************************************************************************
85 84 : SUBROUTINE ensure_ri_rs_representation(qs_env, bs_env)
86 : TYPE(qs_environment_type), POINTER :: qs_env
87 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
88 :
89 : CHARACTER(LEN=*), PARAMETER :: routineN = 'ensure_ri_rs_representation'
90 :
91 : INTEGER :: handle
92 : REAL(KIND=dp) :: build_start
93 :
94 84 : CALL timeset(routineN, handle)
95 84 : IF (.NOT. bs_env%ri_rs%grid_built) THEN
96 48 : build_start = m_walltime()
97 : CALL get_qs_env(qs_env, atomic_kind_set=bs_env%ri_rs%atomic_kind_set, &
98 48 : cell=bs_env%ri_rs%cell, particle_set=bs_env%ri_rs%particle_set)
99 : ! AO/RI support radii determine which atoms and grid points enter each local fit.
100 48 : IF (.NOT. ALLOCATED(bs_env%ri_rs%radius_ao_per_atom)) THEN
101 : CALL precompute_ri_rs_radii(bs_env%ri_rs, bs_env%basis_set_AO, bs_env%basis_set_RI, &
102 48 : bs_env%eps_filter, bs_env%unit_nr)
103 : END IF
104 : ! Place atom-centered RI-RS grids at r_l = R_A + r_l^(A), then flatten them into one grid.
105 48 : CALL setup_ri_rs_grid(bs_env, bs_env%ri_rs%grid_points)
106 : ! Collocation matrix ϕ_μl = ϕ_μ(r_l), including neighbouring AO images for Gamma.
107 : CALL atomic_basis_at_grid_point(bs_env%ri_rs, bs_env%basis_set_AO, &
108 : bs_env%mat_ao_ao%matrix, bs_env%para_env, bs_env%eps_filter, &
109 : bs_env%dbcsr_msg_elem_limit, bs_env%unit_nr, &
110 48 : periodic_images=bs_env%do_periodic)
111 : CALL print_matrix_occupation(bs_env%ri_rs%mat_phi_mu_l, 'ϕ_μ(r_l)', &
112 48 : bs_env%para_env, bs_env%unit_nr, bs_env%dbcsr_msg_elem_limit)
113 : ! Fit Z_lP using the integration domains and RI basis described in compute_Z_lP below.
114 : CALL compute_Z_lP(qs_env, bs_env, bs_env%ri_rs%grid_points, &
115 : bs_env%ri_rs%mat_phi_mu_l, bs_env%ri_rs%mat_Z_lP, &
116 48 : periodic_images=bs_env%do_periodic)
117 48 : bs_env%ri_rs%grid_built = .TRUE.
118 48 : IF (bs_env%unit_nr > 0) THEN
119 : WRITE (bs_env%unit_nr, '(T2,A,T58,A,F7.1,A)') &
120 24 : 'Built RI-RS representation,', ' Execution time', m_walltime() - build_start, ' s'
121 24 : WRITE (bs_env%unit_nr, '(A)') ' '
122 : END IF
123 : END IF
124 84 : CALL timestop(handle)
125 84 : END SUBROUTINE ensure_ri_rs_representation
126 :
127 : ! **************************************************************************************************
128 : !> \brief Computes Z_lP for (1) a tabulated RI basis set or (2) an on-the-fly generated RI basis set.
129 : !>
130 : !> In both cases, the Z_lP coefficients are computed by solving the linear system
131 : !>
132 : !> Σ_l' D_ll' Z_l'P = d_lP,
133 : !>
134 : !> with D_ll' and d_lP given by
135 : !>
136 : !> D_ll' = [Σ_μ ϕ_μ(r_l) ϕ_μ(r_l')]²,
137 : !>
138 : !> d_lP = Σ_μν ϕ_μ(r_l) ϕ_ν(r_l) (μν|P).
139 : !>
140 : !> For an on-the-fly generated RI basis set, an RI function φ_P may be a contraction of
141 : !> Gaussians on neighboring atoms, which requires special computation of d_lP.
142 : !>
143 : !> \param qs_env ...
144 : !> \param bs_env Band-structure environment containing GW parameters.
145 : !> \param ri_rs_grid_points ...
146 : !> \param mat_phi_mu_l ...
147 : !> \param mat_Z_lP ...
148 : !> \param periodic_images use Gamma image sums and minimum-image integration spheres
149 : ! **************************************************************************************************
150 48 : SUBROUTINE compute_Z_lP(qs_env, bs_env, ri_rs_grid_points, mat_phi_mu_l, mat_Z_lP, periodic_images)
151 : TYPE(qs_environment_type), POINTER :: qs_env
152 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
153 : REAL(KIND=dp), INTENT(IN) :: ri_rs_grid_points(:, :)
154 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_phi_mu_l
155 : TYPE(dbcsr_type), INTENT(OUT) :: mat_Z_lP
156 : LOGICAL, INTENT(IN), OPTIONAL :: periodic_images
157 :
158 : CHARACTER(LEN=*), PARAMETER :: routineN = 'compute_Z_lP'
159 :
160 : INTEGER :: handle
161 : LOGICAL :: sum_images
162 : REAL(KIND=dp) :: fitting_stats(2)
163 :
164 48 : CALL timeset(routineN, handle)
165 48 : sum_images = .FALSE.
166 48 : IF (PRESENT(periodic_images)) sum_images = periodic_images
167 :
168 48 : bs_env%ri_rs%fit_max_grid = 0
169 48 : bs_env%ri_rs%fit_peak_bytes = 0.0_dp
170 48 : IF (bs_env%auto_ri%enabled) THEN
171 4 : CALL compute_Z_lP_auto_ri(qs_env, bs_env, ri_rs_grid_points, mat_phi_mu_l, mat_Z_lP)
172 : ! Localized AUTO-RI fits can be owned by different process columns.
173 12 : fitting_stats = [REAL(bs_env%ri_rs%fit_max_grid, dp), bs_env%ri_rs%fit_peak_bytes]
174 4 : CALL bs_env%para_env%max(fitting_stats)
175 4 : bs_env%ri_rs%fit_max_grid = NINT(fitting_stats(1))
176 4 : bs_env%ri_rs%fit_peak_bytes = fitting_stats(2)
177 : ELSE
178 44 : CALL compute_Z_lP_standard(qs_env, bs_env, ri_rs_grid_points, mat_phi_mu_l, mat_Z_lP, sum_images)
179 : END IF
180 :
181 48 : CALL timestop(handle)
182 :
183 48 : END SUBROUTINE compute_Z_lP
184 :
185 : ! **************************************************************************************************
186 : !> \brief Computes the RI-RS fitting coefficients Z_lP by solving, independently for every RI
187 : !> atom P, a Jacobi-conditioned, Tikhonov-regularized linear system restricted to the
188 : !> grid points r_l inside P's integration sphere |r_l - R_P| <= cutoff_ri(P):
189 : !>
190 : !> 1. D_ll' = [ Σ_μ ϕ_μ(r_l) ϕ_μ(r_l') ]²
191 : !> 2. d_l = 1 / sqrt(max(D_ll, jacobi_floor)) (Jacobi conditioning vector)
192 : !> 3. D'_ll' = d_l D_ll' d_l' + λ δ_ll' (λ = TIKHONOV regularization)
193 : !> 4. d_lP = Σ_μν ϕ_μ(r_l) ϕ_ν(r_l) (μν|P)
194 : !> 5. Σ_l' D'_ll' Z'_l'P = d_l d_lP (solve linear system for Z'_lP)
195 : !> 6. Z_lP = d_l Z'_lP (undo the conditioning)
196 : !>
197 : !> See Duchemin and Blase, arXiv:1902.00066, Sec. II.3 for the least-squares construction
198 : !> and balancing, and solve_ri_rs_fit for the common floor and regularization convention.
199 : !>
200 : !> The Gamma path uses image-summed AO values and three-center integrals, with RI
201 : !> atom P fixed in the reference cell. Molecular and Gamma sphere cutoffs are retained.
202 : !>
203 : !> Work is distributed over atoms in two phases (planned by classify_z_lp_atoms and
204 : !> lpt_assign_atoms): Phase A solves "small" atoms with single-rank LAPACK
205 : !> (dpotrf/dpotrs); Phase B solves "big" atoms, whose dense matrix D'_ll' would exceed one
206 : !> rank's memory, with ScaLAPACK (pdpotrf/pdpotrs) over rank subgroups of size G.
207 : !> The solved Z columns are scattered into the sparse global mat_Z_lP.
208 : !> If a Z_lP restart file exists, it is read instead and the solve is skipped entirely.
209 : !> \param qs_env ...
210 : !> \param bs_env ...
211 : !> \param ri_rs_grid_points ...
212 : !> \param mat_phi_mu_l ...
213 : !> \param mat_Z_lP ...
214 : !> \param periodic_images use Gamma image sums and minimum-image integration spheres
215 : ! **************************************************************************************************
216 44 : SUBROUTINE compute_Z_lP_standard(qs_env, bs_env, ri_rs_grid_points, mat_phi_mu_l, mat_Z_lP, periodic_images)
217 :
218 : TYPE(qs_environment_type), POINTER :: qs_env
219 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
220 : REAL(KIND=dp), INTENT(IN) :: ri_rs_grid_points(:, :)
221 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_phi_mu_l
222 : TYPE(dbcsr_type), INTENT(OUT) :: mat_Z_lP
223 : LOGICAL, INTENT(IN) :: periodic_images
224 :
225 : CHARACTER(LEN=*), PARAMETER :: key = 'PROPERTIES%BANDSTRUCTURE%GW%PRINT%RESTART', &
226 : routineN = 'compute_Z_lP_standard'
227 :
228 : INTEGER :: atom_j_mepos, atom_j_stride, atom_P, G, handle, i_blk, iatom, idx, iphase, j, &
229 : max_ao_size, my_group, n_big, n_done, n_groups, n_loc_ri, n_local_grid, n_my_atoms, &
230 : n_small, natom, npcol_phi, num_grid_chunks, phase_hi
231 88 : INTEGER, ALLOCATABLE, DIMENSION(:) :: ao_atoms, ao_col_start, big_list, local_grid_idx, &
232 44 : my_atoms_A, my_atoms_B, n_local_grid_atom, row_offset, small_list
233 88 : INTEGER, DIMENSION(:), POINTER :: col_dist_ri, r_blk_sizes, ri_blk_sizes, &
234 44 : row_dist_grid
235 : LOGICAL :: do_scatter, use_dist
236 44 : LOGICAL, ALLOCATABLE :: ao_nonzero(:, :)
237 : REAL(KIND=dp) :: cutoff_ri, item_start_time, r_c, t1
238 44 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: cutoff_ri_per_atom
239 44 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: d_lp_local, phi_local
240 : TYPE(cell_type), POINTER :: cell
241 : TYPE(cp_blacs_env_type), POINTER :: blacs_env_sub
242 : TYPE(cp_logger_type), POINTER :: logger
243 : TYPE(dbcsr_distribution_type) :: dist_phi, dist_Z
244 572 : TYPE(gw_3c_ctx_type) :: ctx_3c
245 : TYPE(mp_para_env_type), POINTER :: para_env, para_env_sub
246 44 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
247 44 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
248 : TYPE(section_vals_type), POINTER :: input
249 :
250 44 : CALL timeset(routineN, handle)
251 :
252 44 : t1 = m_walltime()
253 :
254 : CALL get_qs_env(qs_env, para_env=para_env, particle_set=particle_set, input=input, &
255 44 : qs_kind_set=qs_kind_set, cell=cell)
256 :
257 44 : NULLIFY (para_env_sub, blacs_env_sub)
258 :
259 44 : natom = bs_env%n_atom
260 :
261 : ! =========================================================================
262 : ! 1. SETUP DBCSR TOPOLOGY & EXACT OFFSETS
263 : ! mat_Z_lP inherits the grid row blocking (and row distribution) of
264 : ! mat_phi_mu_l; its columns are one block per RI atom.
265 : ! =========================================================================
266 44 : CALL dbcsr_get_info(mat_phi_mu_l, row_blk_size=r_blk_sizes, distribution=dist_phi)
267 44 : CALL dbcsr_distribution_get(dist_phi, row_dist=row_dist_grid, npcols=npcol_phi)
268 :
269 44 : num_grid_chunks = SIZE(r_blk_sizes)
270 :
271 132 : ALLOCATE (row_offset(num_grid_chunks))
272 44 : row_offset(1) = 0
273 602 : DO i_blk = 2, num_grid_chunks
274 602 : row_offset(i_blk) = row_offset(i_blk - 1) + r_blk_sizes(i_blk - 1)
275 : END DO
276 :
277 176 : ALLOCATE (ri_blk_sizes(natom), col_dist_ri(natom))
278 150 : DO iatom = 1, natom
279 : ri_blk_sizes(iatom) = &
280 106 : bs_env%i_RI_end_from_atom(iatom) - bs_env%i_RI_start_from_atom(iatom) + 1
281 150 : col_dist_ri(iatom) = MOD(iatom - 1, npcol_phi)
282 : END DO
283 :
284 : CALL dbcsr_distribution_new(dist_Z, template=dist_phi, &
285 44 : row_dist=row_dist_grid, col_dist=col_dist_ri)
286 :
287 44 : IF (bs_env%ri_rs%Z_lP_exists) THEN
288 : CALL dbcsr_binary_read(filepath=TRIM(bs_env%prefix)//"Z_lP.matrix", &
289 : distribution=dist_Z, &
290 2 : matrix_new=mat_Z_lP)
291 2 : IF (bs_env%unit_nr > 0) THEN
292 : WRITE (bs_env%unit_nr, '(T2,A,T57,A,F7.1,A)') &
293 1 : 'Read Z_lP from file ', ' Execution time', m_walltime() - t1, ' s'
294 : ! The grid rows use Morton order (spatial_atom_order). A Z_lP.matrix
295 : ! written with another grid ordering would be silently read into
296 : ! the current row order. Delete stale Z_lP.matrix files and recompute if in doubt.
297 : WRITE (bs_env%unit_nr, '(T2,A)') &
298 1 : '*** NOTE: Z_lP restart must match the current (spatial) grid row ordering ***'
299 1 : WRITE (bs_env%unit_nr, '(A)') ' '
300 : END IF
301 : ELSE
302 :
303 42 : IF (bs_env%unit_nr > 0) THEN
304 21 : WRITE (bs_env%unit_nr, '(A)') ' '
305 21 : WRITE (bs_env%unit_nr, '(T2,A)') 'Started computing Z_lP'
306 21 : CALL m_flush(bs_env%unit_nr)
307 : END IF
308 : CALL dbcsr_create(mat_Z_lP, name="mat_Z_lP", dist=dist_Z, &
309 : matrix_type=dbcsr_type_no_symmetry, &
310 42 : row_blk_size=r_blk_sizes, col_blk_size=ri_blk_sizes)
311 :
312 : ! Largest per-atom AO block, needed to size the 3c-integral work buffers.
313 42 : max_ao_size = 0
314 142 : DO j = 1, bs_env%n_atom
315 : max_ao_size = MAX(max_ao_size, &
316 : bs_env%i_ao_end_from_atom(j) - &
317 142 : bs_env%i_ao_start_from_atom(j) + 1)
318 : END DO
319 :
320 : ! Keep each path's integration domain: r_c + r_RI(P) for molecules,
321 : ! r_c + r_AO(P) for the Gamma image-summed approximation.
322 126 : ALLOCATE (cutoff_ri_per_atom(natom))
323 :
324 42 : r_c = bs_env%ri_metric%cutoff_radius
325 142 : DO iatom = 1, natom
326 142 : IF (periodic_images) THEN
327 0 : cutoff_ri_per_atom(iatom) = r_c + bs_env%ri_rs%radius_ao_per_atom(iatom)
328 : ELSE
329 100 : cutoff_ri_per_atom(iatom) = r_c + bs_env%ri_rs%radius_ri_per_atom(iatom)
330 : END IF
331 : END DO
332 :
333 42 : CALL print_sphere_cutoff_table(bs_env, cutoff_ri_per_atom)
334 :
335 : ! =========================================================================
336 : ! 2. PER-ATOM SOLVER CLASSIFICATION
337 : ! Split the atoms into "small" (single-rank LAPACK, Phase A) and "big"
338 : ! (distributed ScaLAPACK over subgroups of G ranks, Phase B) by comparing
339 : ! each atom's estimated solve peak memory against the measured budget.
340 : ! =========================================================================
341 : CALL classify_z_lp_atoms(bs_env, ri_rs_grid_points, cutoff_ri_per_atom, ri_blk_sizes, &
342 42 : n_local_grid_atom, small_list, n_small, big_list, n_big, G, periodic_images)
343 :
344 : ! LPT scheduling: sort the atoms of each phase by estimated solve cost
345 : ! (n_local_grid^3, Cholesky-dominated) and greedily assign to the least-loaded
346 : ! rank (Phase A) / subgroup (Phase B).
347 : CALL lpt_assign_atoms(small_list, n_small, n_local_grid_atom, para_env%num_pe, &
348 42 : para_env%mepos, my_atoms_A)
349 42 : IF (n_big > 0) THEN
350 0 : n_groups = para_env%num_pe/G
351 0 : my_group = MIN(para_env%mepos/G, n_groups - 1)
352 : CALL lpt_assign_atoms(big_list, n_big, n_local_grid_atom, n_groups, my_group, &
353 0 : my_atoms_B)
354 : ELSE
355 42 : ALLOCATE (my_atoms_B(0))
356 : END IF
357 :
358 : ! Atoms this rank will process across both phases for rank-0 progress
359 42 : n_my_atoms = SIZE(my_atoms_A) + SIZE(my_atoms_B)
360 42 : n_done = 0
361 :
362 : ! Shared context for the three-center integrals (μν|P) of the RHS build
363 : CALL gw_3c_ctx_create(ctx_3c, bs_env, bs_env%ri_metric, &
364 : basis_j=bs_env%basis_set_AO, basis_k=bs_env%basis_set_AO, &
365 42 : basis_i=bs_env%basis_set_RI)
366 :
367 : ! =========================================================================
368 : ! 3. TWO-PHASE LOOP OVER ATOMS
369 : ! Phase A processes the "small" atoms with the single-rank BLAS path
370 : ! Phase B processes the "big" atoms with the distributed ScaLAPACK path
371 : ! over rank subgroups of size G. phi_local for each atom's cutoff sphere
372 : ! is built on the fly to avoid replicating a global grid x AO matrix.
373 : ! =========================================================================
374 126 : DO iphase = 1, 2
375 84 : IF (iphase == 1) THEN
376 42 : use_dist = .FALSE.
377 42 : atom_j_mepos = 0
378 42 : atom_j_stride = 1
379 42 : phase_hi = SIZE(my_atoms_A)
380 : ELSE
381 42 : IF (n_big == 0) CYCLE
382 0 : use_dist = .TRUE.
383 0 : n_groups = para_env%num_pe/G
384 0 : my_group = MIN(para_env%mepos/G, n_groups - 1)
385 0 : ALLOCATE (para_env_sub)
386 0 : CALL para_env_sub%from_split(para_env, my_group)
387 0 : CALL cp_blacs_env_create(blacs_env=blacs_env_sub, para_env=para_env_sub)
388 0 : atom_j_mepos = para_env_sub%mepos
389 0 : atom_j_stride = para_env_sub%num_pe
390 : ! All ranks of a subgroup share my_group, hence the identical my_atoms_B list
391 : ! (the per-atom ScaLAPACK solve is collective over the subgroup).
392 0 : phase_hi = SIZE(my_atoms_B)
393 : END IF
394 :
395 92 : DO idx = 1, phase_hi
396 50 : item_start_time = m_walltime()
397 50 : IF (iphase == 1) THEN
398 50 : atom_P = my_atoms_A(idx)
399 : ELSE
400 0 : atom_P = my_atoms_B(idx)
401 : END IF
402 :
403 50 : n_loc_ri = ri_blk_sizes(atom_P)
404 50 : cutoff_ri = cutoff_ri_per_atom(atom_P)
405 :
406 : ! ---------------------------------------------------------------------
407 : ! A. Sphere-local AO matrix ϕ_μ(r_l): select the grid points with
408 : ! |r_l - R_P| <= cutoff_ri(P), using minimum-image distances for Gamma.
409 : ! The molecular path drops rows below EPS_FILTER; Gamma keeps all sphere rows.
410 : ! ---------------------------------------------------------------------
411 : CALL build_phi_on_sphere(bs_env, qs_kind_set, &
412 : ri_rs_grid_points, atom_P, cutoff_ri, &
413 : local_grid_idx, n_local_grid, phi_local, &
414 50 : ao_atoms, ao_col_start, periodic_images=periodic_images, ao_nonzero=ao_nonzero)
415 :
416 : ! ---------------------------------------------------------------------
417 : ! B. Right-hand side d_lP = Σ_μν ϕ_μ(r_l) ϕ_ν(r_l) (μν|P)
418 : ! ---------------------------------------------------------------------
419 200 : ALLOCATE (d_lp_local(n_local_grid, n_loc_ri))
420 50 : d_lp_local = 0.0_dp
421 :
422 : CALL compute_d_lp(bs_env, ctx_3c, phi_local, ao_atoms, ao_col_start, ao_nonzero, d_lp_local, n_local_grid, &
423 50 : n_loc_ri, atom_P, max_ao_size, atom_j_mepos, atom_j_stride, periodic_images)
424 :
425 50 : DEALLOCATE (ao_atoms, ao_col_start, ao_nonzero)
426 :
427 : ! Reduce per-subgroup-rank partials into the replicated d_lp_local.
428 : ! Skipped for BLAS path: each rank has the full sum locally.
429 50 : IF (use_dist) THEN
430 0 : CALL para_env_sub%sum(d_lp_local)
431 : END IF
432 :
433 : ! ---------------------------------------------------------------------
434 : ! C. Solve the Jacobi-scaled, regularized fit with LAPACK or ScaLAPACK.
435 : ! ---------------------------------------------------------------------
436 : IF (use_dist) THEN
437 : CALL solve_ri_rs_fit(phi_local, d_lp_local, bs_env%ri_rs%tikhonov, &
438 0 : para_env_sub, blacs_env_sub)
439 : ELSE
440 50 : CALL solve_ri_rs_fit(phi_local, d_lp_local, bs_env%ri_rs%tikhonov)
441 : END IF
442 :
443 : ! ---------------------------------------------------------------------
444 : ! D. Scatter the solved Z columns back into the global sparse mat_Z_lP.
445 : ! ---------------------------------------------------------------------
446 50 : do_scatter = .TRUE.
447 50 : IF (use_dist) do_scatter = (para_env_sub%mepos == 0)
448 0 : IF (do_scatter) THEN
449 : CALL store_Z_lP_columns(mat_Z_lP, d_lp_local, local_grid_idx, n_local_grid, &
450 : n_loc_ri, atom_P, r_blk_sizes, row_offset, &
451 50 : bs_env%eps_filter)
452 : END IF
453 :
454 50 : DEALLOCATE (d_lp_local)
455 50 : DEALLOCATE (local_grid_idx, phi_local)
456 :
457 : ! Report each atom completed by the printing rank. No inter-rank communication is
458 : ! needed; the final message below is printed only after the global synchronization.
459 50 : n_done = n_done + 1
460 : CALL print_Z_lP_progress(bs_env, n_done, n_my_atoms, &
461 92 : m_walltime() - item_start_time)
462 : END DO ! idx: atoms of this phase owned by this rank / subgroup
463 :
464 : ! Tear down the Phase-B subgroup (all ranks created it collectively).
465 84 : IF (iphase == 2) THEN
466 0 : CALL cp_blacs_env_release(blacs_env_sub)
467 0 : CALL para_env_sub%free()
468 0 : DEALLOCATE (para_env_sub)
469 : END IF
470 : END DO ! iphase
471 :
472 42 : DEALLOCATE (cutoff_ri_per_atom)
473 42 : DEALLOCATE (small_list, big_list)
474 :
475 42 : CALL gw_3c_ctx_release(ctx_3c)
476 :
477 42 : CALL dbcsr_finalize(mat_Z_lP)
478 :
479 42 : CALL para_env%sync()
480 : CALL print_Z_lP_progress(bs_env, natom, natom, m_walltime() - t1, &
481 42 : all_mpi_ranks=.TRUE.)
482 :
483 42 : logger => cp_get_default_logger()
484 :
485 42 : IF (BTEST(cp_print_key_should_output(logger%iter_info, input, key), cp_p_file)) THEN
486 4 : CALL dbcsr_binary_write(matrix=mat_Z_lP, filepath=TRIM(bs_env%prefix)//"Z_lP.matrix")
487 : END IF
488 :
489 : END IF
490 :
491 44 : DEALLOCATE (row_offset, ri_blk_sizes, col_dist_ri)
492 44 : CALL dbcsr_distribution_release(dist_Z)
493 :
494 44 : CALL timestop(handle)
495 :
496 220 : END SUBROUTINE compute_Z_lP_standard
497 :
498 : ! **************************************************************************************************
499 : !> \brief Prints the per-kind maximum RI-RS integration-sphere cutoff table.
500 : !> \param bs_env ...
501 : !> \param cutoff_ri_per_atom ...
502 : ! **************************************************************************************************
503 42 : SUBROUTINE print_sphere_cutoff_table(bs_env, cutoff_ri_per_atom)
504 :
505 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
506 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: cutoff_ri_per_atom
507 :
508 : CHARACTER(LEN=*), PARAMETER :: routineN = 'print_sphere_cutoff_table'
509 :
510 : INTEGER :: handle, iatom, ikind, nkind
511 42 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: cutoff_ri_per_kind
512 42 : TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
513 42 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
514 :
515 42 : CALL timeset(routineN, handle)
516 :
517 42 : atomic_kind_set => bs_env%ri_rs%atomic_kind_set
518 42 : particle_set => bs_env%ri_rs%particle_set
519 :
520 42 : IF (bs_env%unit_nr <= 0) THEN
521 21 : CALL timestop(handle)
522 : RETURN
523 : END IF
524 :
525 21 : nkind = SIZE(atomic_kind_set)
526 63 : ALLOCATE (cutoff_ri_per_kind(nkind))
527 21 : cutoff_ri_per_kind(:) = 0.0_dp
528 :
529 71 : DO iatom = 1, bs_env%n_atom
530 50 : ikind = particle_set(iatom)%atomic_kind%kind_number
531 71 : cutoff_ri_per_kind(ikind) = MAX(cutoff_ri_per_kind(ikind), cutoff_ri_per_atom(iatom))
532 : END DO
533 :
534 21 : WRITE (bs_env%unit_nr, '(T2,A)') 'Per-kind maximum RI-RS sphere cutoff (Å):'
535 21 : WRITE (bs_env%unit_nr, '(T4,A4,A14)') 'Kind', 'cutoff (Å)'
536 50 : DO ikind = 1, nkind
537 : WRITE (bs_env%unit_nr, '(T4,A4,F14.4)') &
538 29 : atomic_kind_set(ikind)%element_symbol, &
539 79 : cutoff_ri_per_kind(ikind)*angstrom
540 : END DO
541 21 : WRITE (bs_env%unit_nr, '(A)') ' '
542 :
543 21 : DEALLOCATE (cutoff_ri_per_kind)
544 :
545 21 : CALL timestop(handle)
546 :
547 42 : END SUBROUTINE print_sphere_cutoff_table
548 :
549 : ! **************************************************************************************************
550 : !> \brief LPT (longest-processing-time) assignment of the Z_lP atoms to workers (MPI ranks in
551 : !> Phase A, rank subgroups in Phase B): sort by estimated solve cost n_local_grid^3
552 : !> (the per-atom Cholesky dominates; the n^2 assembly terms order the atoms the same
553 : !> way) and greedily give each atom to the least-loaded worker.
554 : !> \param atom_list ...
555 : !> \param n_atoms ...
556 : !> \param n_local_grid_atom ...
557 : !> \param n_workers ...
558 : !> \param my_worker ...
559 : !> \param my_atoms ...
560 : ! **************************************************************************************************
561 42 : SUBROUTINE lpt_assign_atoms(atom_list, n_atoms, n_local_grid_atom, n_workers, my_worker, &
562 : my_atoms)
563 :
564 : INTEGER, DIMENSION(:), INTENT(IN) :: atom_list
565 : INTEGER, INTENT(IN) :: n_atoms
566 : INTEGER, DIMENSION(:), INTENT(IN) :: n_local_grid_atom
567 : INTEGER, INTENT(IN) :: n_workers, my_worker
568 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: my_atoms
569 :
570 : CHARACTER(LEN=*), PARAMETER :: routineN = 'lpt_assign_atoms'
571 :
572 : INTEGER :: handle, i, iw, n_mine, w_min
573 42 : INTEGER, ALLOCATABLE, DIMENSION(:) :: mine_tmp, perm
574 42 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: cost, load
575 :
576 42 : CALL timeset(routineN, handle)
577 :
578 42 : IF (n_atoms <= 0) THEN
579 0 : ALLOCATE (my_atoms(0))
580 0 : CALL timestop(handle)
581 0 : RETURN
582 : END IF
583 :
584 336 : ALLOCATE (cost(n_atoms), perm(n_atoms), mine_tmp(n_atoms), load(n_workers))
585 142 : DO i = 1, n_atoms
586 142 : cost(i) = REAL(n_local_grid_atom(atom_list(i)), dp)**3
587 : END DO
588 42 : CALL sort(cost, n_atoms, perm) ! ascending; walk backwards for largest-first
589 :
590 42 : load(:) = 0.0_dp
591 42 : n_mine = 0
592 142 : DO i = n_atoms, 1, -1
593 : w_min = 1
594 200 : DO iw = 2, n_workers
595 200 : IF (load(iw) < load(w_min)) w_min = iw
596 : END DO
597 100 : load(w_min) = load(w_min) + cost(i)
598 142 : IF (w_min - 1 == my_worker) THEN
599 50 : n_mine = n_mine + 1
600 50 : mine_tmp(n_mine) = atom_list(perm(i))
601 : END IF
602 : END DO
603 :
604 126 : ALLOCATE (my_atoms(n_mine))
605 92 : my_atoms(:) = mine_tmp(1:n_mine)
606 :
607 42 : CALL timestop(handle)
608 :
609 84 : END SUBROUTINE lpt_assign_atoms
610 :
611 : ! **************************************************************************************************
612 : !> \brief Computes the dense localized right-hand side for one RI atom P,
613 : !>
614 : !> d_lP = Σ_μν ϕ_μ(r_l) ϕ_ν(r_l) (μν|P).
615 : !>
616 : !> OMP-threaded over unordered AO-atom pairs: build (mu nu|P) in the reference cell,
617 : !> or sum (mu_R nu_S|P_0) over nearest images for Gamma. Contract grid-chunked pair
618 : !> densities into a thread-private d_lp partial, then reduce the partials into d_lp.
619 : !> The shared integral context caches Libint/spherical-transform tables; each thread
620 : !> owns its integral workspace. The kind-radius triangle screen in
621 : !> build_3c_integral_block_ctx skips negligible image triples via the screened output.
622 : !> \param bs_env ...
623 : !> \param ctx shared three-center integral context
624 : !> \param phi_val AO values on local-sphere grid rows and compact active-AO columns
625 : !> \param ao_atoms global indices of the active AO atoms
626 : !> \param ao_col_start local phi column starts, including the terminal column offset
627 : !> \param ao_nonzero exact nonzero AO support on the retained grid rows
628 : !> \param d_lp RHS (n_grid_total × n_loc_ri), zeroed by the caller and accumulated here
629 : !> \param n_grid_total number of retained local-sphere grid rows
630 : !> \param n_loc_ri number of RI functions on atom_P
631 : !> \param atom_P RI atom fixed in the reference cell (0,0,0)
632 : !> \param max_ao_size ...
633 : !> \param atom_j_mepos ...
634 : !> \param atom_j_stride ...
635 : !> \param periodic_images sum nearest AO images with RI atom P fixed in the reference cell
636 : ! **************************************************************************************************
637 50 : SUBROUTINE compute_d_lp(bs_env, ctx, phi_val, ao_atoms, ao_col_start, ao_nonzero, d_lp, n_grid_total, n_loc_ri, atom_P, &
638 : max_ao_size, atom_j_mepos, atom_j_stride, periodic_images)
639 :
640 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
641 : TYPE(gw_3c_ctx_type), INTENT(IN) :: ctx
642 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi_val
643 : INTEGER, DIMENSION(:), INTENT(IN) :: ao_atoms, ao_col_start
644 : LOGICAL, DIMENSION(:, :), INTENT(IN) :: ao_nonzero
645 : INTEGER, INTENT(IN) :: n_grid_total, n_loc_ri
646 : REAL(KIND=dp), INTENT(INOUT) :: d_lp(n_grid_total, n_loc_ri)
647 : INTEGER, INTENT(IN) :: atom_P, max_ao_size, atom_j_mepos, &
648 : atom_j_stride
649 : LOGICAL, INTENT(IN) :: periodic_images
650 :
651 : CHARACTER(LEN=*), PARAMETER :: routineN = 'compute_d_lp'
652 : INTEGER, PARAMETER :: grid_chunk = 1024
653 :
654 : INTEGER :: atom_j, atom_k, c, cell_R_vec(3), cell_S_vec(3), handle, handle2, &
655 : image_extent(3), ix_R, ix_S, iy_R, iy_S, iz_R, iz_S, j, jj, jk_idx, jsize, jstart, k, kk, &
656 : ksize, kstart, l, l0, n_grid_pair, point, ri
657 50 : INTEGER, ALLOCATABLE :: grid_index(:)
658 : LOGICAL :: any_kept, screened
659 : REAL(KIND=dp) :: pair_factor
660 50 : REAL(KIND=dp), ALLOCATABLE :: grid_result(:, :), int_3c_prv(:, :, :)
661 50 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: d_lp_prv, int_2d_prv, rho_chunk
662 50 : TYPE(gw_3c_ws_type) :: ws
663 :
664 50 : CALL timeset(routineN, handle)
665 50 : image_extent = 0
666 50 : IF (periodic_images) image_extent = bs_env%ri_rs%cell%perd
667 :
668 : !$OMP PARALLEL DEFAULT(NONE) &
669 : !$OMP SHARED(bs_env, ctx, phi_val, ao_atoms, ao_col_start, ao_nonzero, d_lp, &
670 : !$OMP n_grid_total, n_loc_ri, atom_P, max_ao_size, &
671 : !$OMP atom_j_mepos, atom_j_stride, image_extent) &
672 : !$OMP PRIVATE(grid_index, grid_result, n_grid_pair, point, pair_factor, &
673 : !$OMP jj, kk, atom_j, atom_k, c, handle2, j, jk_idx, jsize, jstart, k, ksize, kstart, &
674 : !$OMP l, l0, ri, screened, d_lp_prv, int_2d_prv, rho_chunk, int_3c_prv, ws, &
675 50 : !$OMP any_kept, ix_R, iy_R, iz_R, ix_S, iy_S, iz_S, cell_R_vec, cell_S_vec)
676 :
677 : ! The integral context is shared; all scratch arrays and the accumulating RHS are thread-private.
678 : CALL gw_3c_ws_create(ws, ctx)
679 : ALLOCATE (int_3c_prv(max_ao_size, max_ao_size, n_loc_ri))
680 : ALLOCATE (int_2d_prv(max_ao_size*max_ao_size, n_loc_ri))
681 : ALLOCATE (rho_chunk(grid_chunk, max_ao_size*max_ao_size))
682 : ALLOCATE (d_lp_prv(n_grid_total, n_loc_ri))
683 : ALLOCATE (grid_index(n_grid_total), grid_result(grid_chunk, n_loc_ri))
684 : d_lp_prv(:, :) = 0.0_dp
685 :
686 : ! MPI assigns each unordered pair through its first atom; OpenMP divides those atoms.
687 : ! Skip only exact-zero pair-grid support, without a new screening threshold.
688 : !$OMP DO SCHEDULE(DYNAMIC)
689 : DO jj = 1, SIZE(ao_atoms)
690 : atom_j = ao_atoms(jj)
691 : IF (MOD(atom_j - 1, atom_j_stride) /= atom_j_mepos) CYCLE
692 : DO kk = jj, SIZE(ao_atoms)
693 : atom_k = ao_atoms(kk)
694 : jstart = ao_col_start(jj)
695 : kstart = ao_col_start(kk)
696 : jsize = ao_col_start(jj + 1) - jstart
697 : ksize = ao_col_start(kk + 1) - kstart
698 :
699 : n_grid_pair = 0
700 : DO point = 1, n_grid_total
701 : IF (.NOT. (ao_nonzero(point, jj) .AND. ao_nonzero(point, kk))) CYCLE
702 : n_grid_pair = n_grid_pair + 1
703 : grid_index(n_grid_pair) = point
704 : END DO
705 : IF (n_grid_pair == 0) CYCLE
706 : ! (μν|P) = (νμ|P): distinct atom pairs contribute twice.
707 : pair_factor = 1.0_dp
708 : IF (atom_j /= atom_k) pair_factor = 2.0_dp
709 :
710 : ! For Gamma, sum (mu_R nu_S|P_0); the molecular path has only R=S=0.
711 : ! AO exchange symmetry applies to the complete image sum, so the same
712 : ! unordered atom pairs and pair_factor are used in both paths.
713 : ! The integral builder accumulates in place, including across image triples.
714 : ! Its kind-radius triangle screen rejects image triples beyond the integral cutoff.
715 : int_3c_prv(1:jsize, 1:ksize, 1:n_loc_ri) = 0.0_dp
716 : any_kept = .FALSE.
717 : DO ix_R = -image_extent(1), image_extent(1)
718 : DO iy_R = -image_extent(2), image_extent(2)
719 : DO iz_R = -image_extent(3), image_extent(3)
720 : cell_R_vec = [ix_R, iy_R, iz_R]
721 : DO ix_S = -image_extent(1), image_extent(1)
722 : DO iy_S = -image_extent(2), image_extent(2)
723 : DO iz_S = -image_extent(3), image_extent(3)
724 : cell_S_vec = [ix_S, iy_S, iz_S]
725 : CALL build_3c_integral_block_ctx(int_3c_prv(1:jsize, 1:ksize, 1:n_loc_ri), &
726 : ctx, ws, atom_j=atom_j, atom_k=atom_k, atom_i=atom_P, &
727 : cell_j=cell_R_vec, cell_k=cell_S_vec, screened=screened)
728 : IF (screened) CYCLE
729 : any_kept = .TRUE.
730 : END DO
731 : END DO
732 : END DO
733 : END DO
734 : END DO
735 : END DO
736 : IF (.NOT. any_kept) CYCLE
737 :
738 : ! Flatten (μν|P) to a matrix with combined AO index (ν-1)*jsize + μ for BLAS.
739 : DO ri = 1, n_loc_ri
740 : DO k = 1, ksize
741 : DO j = 1, jsize
742 : jk_idx = (k - 1)*jsize + j
743 : int_2d_prv(jk_idx, ri) = int_3c_prv(j, k, ri)
744 : END DO
745 : END DO
746 : END DO
747 :
748 : ! Pair density ρ(l, μν) = ϕ_μ(r_l) ϕ_ν(r_l) in grid chunks, contracted on the fly:
749 : ! d_lP += Σ_μν ρ(l, μν) (μν|P), with the atom-pair multiplicity included in pair_factor.
750 : DO l0 = 1, n_grid_pair, grid_chunk
751 : c = MIN(grid_chunk, n_grid_pair - l0 + 1)
752 : DO k = 1, ksize
753 : DO j = 1, jsize
754 : jk_idx = (k - 1)*jsize + j
755 : DO l = 1, c
756 : point = grid_index(l0 + l - 1)
757 : rho_chunk(l, jk_idx) = phi_val(point, jstart + j - 1)* &
758 : phi_val(point, kstart + k - 1)
759 : END DO
760 : END DO
761 : END DO
762 : CALL timeset(routineN//"_dgemm", handle2)
763 : CALL dgemm("N", "N", c, n_loc_ri, jsize*ksize, &
764 : pair_factor, rho_chunk, grid_chunk, &
765 : int_2d_prv, max_ao_size*max_ao_size, &
766 : 0.0_dp, grid_result, grid_chunk)
767 : DO ri = 1, n_loc_ri
768 : DO l = 1, c
769 : point = grid_index(l0 + l - 1)
770 : d_lp_prv(point, ri) = d_lp_prv(point, ri) + grid_result(l, ri)
771 : END DO
772 : END DO
773 : CALL timestop(handle2)
774 : END DO
775 : END DO
776 : END DO
777 : !$OMP END DO
778 :
779 : ! Reduce once per thread; distributed fits subsequently sum the rank partials in the subgroup.
780 : !$OMP CRITICAL (compute_d_lp_reduce)
781 : d_lp(1:n_grid_total, 1:n_loc_ri) = d_lp(1:n_grid_total, 1:n_loc_ri) + &
782 : d_lp_prv(1:n_grid_total, 1:n_loc_ri)
783 : !$OMP END CRITICAL (compute_d_lp_reduce)
784 :
785 : DEALLOCATE (int_3c_prv, int_2d_prv, rho_chunk, d_lp_prv, grid_index, grid_result)
786 : CALL gw_3c_ws_release(ws)
787 :
788 : !$OMP END PARALLEL
789 :
790 50 : CALL timestop(handle)
791 :
792 50 : END SUBROUTINE compute_d_lp
793 :
794 : ! **************************************************************************************************
795 : !> \brief Computes the RI-RS matrix Z_lP for an automatically optimized RI basis.
796 : !>
797 : !> The fitting coefficients Z_lP satisfy
798 : !>
799 : !> Σ_l' D_ll' Z_l'P = d_lP,
800 : !> D_ll' = [Σ_μ ϕ_μ(r_l) ϕ_μ(r_l')]²,
801 : !> d_lP = Σ_μν ϕ_μ(r_l) ϕ_ν(r_l) (μν|P).
802 : !>
803 : !> For the standard RI-RS fit, P is an atom-centered RI function. Here an optimized
804 : !> function may contain reference functions on neighboring atoms,
805 : !>
806 : !> φ_p(r) = Σ_A Σ_{P∈A} U_Pp^A φ_P^A(r),
807 : !>
808 : !> so all atomic contributions to d_lp must be accumulated before solving for Z_lp.
809 : !> If every optimized column uses the complete RI-RS grid, all columns are solved in one
810 : !> system. Otherwise, each atom-blocked column set is fitted on its own integration sphere;
811 : !> grid points outside that sphere are excluded from its fitting equations.
812 : !> \param qs_env ...
813 : !> \param bs_env ...
814 : !> \param ri_rs_grid_points ...
815 : !> \param mat_phi_mu_l ...
816 : !> \param mat_Z_lP ...
817 : ! **************************************************************************************************
818 4 : SUBROUTINE compute_Z_lP_auto_ri(qs_env, bs_env, ri_rs_grid_points, mat_phi_mu_l, mat_Z_lP)
819 : TYPE(qs_environment_type), POINTER :: qs_env
820 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
821 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: ri_rs_grid_points
822 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_phi_mu_l
823 : TYPE(dbcsr_type), INTENT(OUT) :: mat_Z_lP
824 :
825 : CHARACTER(LEN=*), PARAMETER :: key = 'PROPERTIES%BANDSTRUCTURE%GW%PRINT%RESTART', &
826 : routineN = 'compute_Z_lP_auto_ri'
827 :
828 : INTEGER :: AB_block, atom_A, atom_B, block_size_b, column_first, first_p_AB, fit_atom, &
829 : handle, last_atom, max_ao_size, max_nRI_ref, mypcol, myprow, n_active, n_candidate, &
830 : n_done, n_my_atoms, n_to_a, natom, ncol, ngrid, nri, nRI_ref_A, nRI_ref_B, output_offset, &
831 : ri_atom, slot
832 4 : INTEGER, ALLOCATABLE, DIMENSION(:) :: ao_atoms, ao_col_start, candidate_order, &
833 8 : candidates, local_grid_idx, ri_atoms, &
834 4 : row_offset
835 8 : INTEGER, DIMENSION(:), POINTER :: col_dist_ri, ri_blk_sizes, &
836 4 : row_dist_grid, row_size_grid
837 : LOGICAL :: AB_block_local, common_grid_available, &
838 : have_fitted_columns, &
839 : reuse_atomic_integrals
840 4 : LOGICAL, ALLOCATABLE :: ao_nonzero(:, :)
841 : REAL(KIND=dp) :: cutoff_ri, item_start_time, &
842 : mem_avail_GB, previous_fit_peak, r_c, &
843 : t1
844 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: d_lp_all, d_lp_local, phi_local, U_Pp_AB
845 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: U_active
846 : REAL(KIND=dp), DIMENSION(3) :: center
847 : TYPE(cell_type), POINTER :: cell
848 : TYPE(cp_logger_type), POINTER :: logger
849 : TYPE(dbcsr_distribution_type) :: dist_Z
850 : TYPE(dbcsr_type) :: mat_rhs
851 52 : TYPE(gw_3c_ctx_type) :: ctx_3c
852 : TYPE(mp_para_env_type), POINTER :: para_env, para_env_col
853 4 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
854 4 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
855 : TYPE(section_vals_type), POINTER :: input
856 :
857 4 : CALL timeset(routineN, handle)
858 4 : t1 = m_walltime()
859 4 : IF (bs_env%unit_nr > 0) THEN
860 2 : WRITE (bs_env%unit_nr, '(A)') ' '
861 2 : WRITE (bs_env%unit_nr, '(T2,A)') 'Started computing Z_lP'
862 2 : CALL m_flush(bs_env%unit_nr)
863 : END IF
864 : CALL get_qs_env(qs_env, para_env=para_env, particle_set=particle_set, input=input, &
865 4 : qs_kind_set=qs_kind_set, cell=cell)
866 4 : NULLIFY (para_env_col)
867 :
868 4 : natom = bs_env%n_atom
869 4 : CPASSERT(bs_env%auto_ri%AB_block_count > 0)
870 4 : CPASSERT(SIZE(particle_set) == natom)
871 :
872 : CALL prepare_Z_lP_auto_ri(bs_env, mat_phi_mu_l, mat_Z_lP, dist_Z, row_size_grid, &
873 : row_dist_grid, col_dist_ri, ri_blk_sizes, row_offset, &
874 4 : myprow, mypcol, max_ao_size)
875 :
876 : CALL gw_3c_ctx_create(ctx_3c, bs_env, bs_env%ri_metric, &
877 : basis_j=bs_env%basis_set_AO, basis_k=bs_env%basis_set_AO, &
878 4 : basis_i=bs_env%basis_set_RI)
879 :
880 : ! Check whether one grid can represent d_lp for every optimized atomic column block.
881 : CALL common_Z_lP_grid_available(bs_env, ri_rs_grid_points, &
882 4 : common_grid_available, have_fitted_columns)
883 :
884 4 : IF (common_grid_available .AND. have_fitted_columns) THEN
885 4 : CALL mp_mem_avail_per_rank_GB(para_env, mem_avail_GB)
886 : ! Evaluate ϕ_μ(r_l) once on the common grid.
887 : CALL build_phi_on_complete_grid(bs_env, qs_kind_set, &
888 : ri_rs_grid_points, local_grid_idx, ngrid, &
889 4 : phi_local, ao_atoms, ao_col_start, ao_nonzero)
890 16 : nri = SUM(ri_blk_sizes)
891 4 : CALL record_fit_memory(bs_env%ri_rs, ngrid, SIZE(phi_local, 2), nri)
892 4 : CALL print_fit_memory_estimate(bs_env, mem_avail_GB)
893 16 : ALLOCATE (d_lp_all(ngrid, nri), source=0.0_dp)
894 : ! d_lp = Σ_A Σ_{P∈A} d_lP U_Pp for every optimized RI function q.
895 : CALL compute_d_lp_auto_ri_batch(bs_env, ctx_3c, phi_local, ao_atoms, ao_col_start, ao_nonzero, d_lp_all, ngrid, &
896 4 : max_ao_size, para_env%mepos, para_env%num_pe)
897 4 : DEALLOCATE (ao_atoms, ao_col_start, ao_nonzero)
898 4 : CALL para_env%sum(d_lp_all)
899 4 : CALL solve_ri_rs_fit(phi_local, d_lp_all, bs_env%ri_rs%tikhonov)
900 :
901 22 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
902 18 : atom_A = bs_env%auto_ri%AB_atom_A(AB_block)
903 18 : atom_B = bs_env%auto_ri%AB_atom_B(AB_block)
904 18 : AB_block_local = col_dist_ri(atom_A) == mypcol
905 18 : IF (atom_B /= atom_A) AB_block_local = AB_block_local .OR. col_dist_ri(atom_B) == mypcol
906 12 : IF (.NOT. AB_block_local) CYCLE
907 18 : ncol = bs_env%auto_ri%AB_size_opt_RI(AB_block)
908 18 : n_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)
909 18 : IF (n_to_a > 0 .AND. col_dist_ri(atom_A) == mypcol) THEN
910 : output_offset = SUM(ri_blk_sizes(:atom_A - 1)) + &
911 28 : bs_env%auto_ri%AB_first_p_A(AB_block) - 1
912 : CALL add_Z_lP_columns(mat_Z_lP, &
913 : d_lp_all(:, output_offset + 1:output_offset + n_to_a), &
914 : local_grid_idx, ngrid, atom_A, &
915 : bs_env%auto_ri%AB_first_p_A(AB_block), &
916 : ri_blk_sizes(atom_A), row_size_grid, row_offset, &
917 : row_dist_grid, &
918 14 : myprow, bs_env%eps_filter)
919 : END IF
920 22 : IF (n_to_a < ncol) THEN
921 4 : CPASSERT(atom_B /= atom_A)
922 4 : IF (col_dist_ri(atom_B) == mypcol) THEN
923 4 : block_size_b = ri_blk_sizes(atom_B)
924 : output_offset = SUM(ri_blk_sizes(:atom_B - 1)) + &
925 10 : bs_env%auto_ri%AB_first_p_B(AB_block) - 1
926 : CALL add_Z_lP_columns(mat_Z_lP, &
927 : d_lp_all(:, output_offset + 1: &
928 : output_offset + ncol - n_to_a), &
929 : local_grid_idx, ngrid, atom_B, &
930 : bs_env%auto_ri%AB_first_p_B(AB_block), block_size_b, &
931 : row_size_grid, row_offset, row_dist_grid, myprow, &
932 4 : bs_env%eps_filter)
933 : END IF
934 : END IF
935 : END DO
936 4 : DEALLOCATE (d_lp_all, local_grid_idx, phi_local)
937 : ELSE
938 : reuse_atomic_integrals = &
939 : bs_env%ri_rs%cutoff_radius_ri_ao <= 0.0_dp .OR. &
940 : bs_env%ri_rs%cutoff_radius_ri_ao <= &
941 0 : MINVAL(bs_env%ri_rs%radius_ao_per_atom)
942 : IF (reuse_atomic_integrals) THEN
943 : ! Form d_lp = Σ_A Σ_{P∈A} d_lP U_Pp before fitting each atomic column block.
944 : CALL compute_auto_ri_d_lp(qs_env, bs_env, ctx_3c, &
945 : ri_rs_grid_points, mat_phi_mu_l, mat_rhs, &
946 0 : max_ao_size)
947 0 : CALL dbcsr_release(mat_Z_lP)
948 : CALL dbcsr_create(mat_Z_lP, name='mat_Z_lP localized AA/AB', dist=dist_Z, &
949 : matrix_type=dbcsr_type_no_symmetry, row_blk_size=row_size_grid, &
950 0 : col_blk_size=ri_blk_sizes)
951 : ! Solve Σ_l' D_ll' Z_l'q = d_lp on each atomic fitting grid.
952 0 : CALL fit_auto_ri_z_lp(qs_env, bs_env, ri_rs_grid_points, mat_rhs, mat_Z_lP)
953 0 : CALL dbcsr_release(mat_rhs)
954 : ELSE
955 : ! All ranks query memory before process columns enter independent fitting loops.
956 0 : CALL mp_mem_avail_per_rank_GB(para_env, mem_avail_GB)
957 0 : ALLOCATE (para_env_col)
958 0 : CALL para_env_col%from_split(para_env, mypcol)
959 0 : n_my_atoms = 0
960 0 : DO fit_atom = 1, natom
961 0 : IF (col_dist_ri(fit_atom) /= mypcol) CYCLE
962 0 : IF (ri_blk_sizes(fit_atom) == 0) CYCLE
963 0 : n_my_atoms = n_my_atoms + 1
964 : END DO
965 0 : n_done = 0
966 : ALLOCATE (candidates(2*bs_env%auto_ri%AB_block_count), &
967 0 : candidate_order(2*bs_env%auto_ri%AB_block_count))
968 0 : DO fit_atom = 1, natom
969 0 : IF (col_dist_ri(fit_atom) /= mypcol) CYCLE
970 0 : ncol = ri_blk_sizes(fit_atom)
971 0 : IF (ncol == 0) CYCLE
972 0 : item_start_time = m_walltime()
973 : ! Only reference RI atoms contributing to this fitted column block are stored.
974 0 : n_candidate = 0
975 0 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
976 0 : atom_A = bs_env%auto_ri%AB_atom_A(AB_block)
977 0 : atom_B = bs_env%auto_ri%AB_atom_B(AB_block)
978 0 : n_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)
979 0 : IF (.NOT. ((fit_atom == atom_A .AND. n_to_a > 0) .OR. &
980 : (fit_atom == atom_B .AND. &
981 : n_to_a < bs_env%auto_ri%AB_size_opt_RI(AB_block)))) CYCLE
982 0 : n_candidate = n_candidate + 1
983 0 : candidates(n_candidate) = atom_A
984 0 : IF (atom_B == atom_A) CYCLE
985 0 : n_candidate = n_candidate + 1
986 0 : candidates(n_candidate) = atom_B
987 : END DO
988 0 : CPASSERT(n_candidate > 0)
989 0 : CALL sort(candidates, n_candidate, candidate_order)
990 0 : n_active = 0
991 0 : last_atom = 0
992 0 : DO slot = 1, n_candidate
993 0 : ri_atom = candidates(slot)
994 0 : IF (ri_atom == last_atom) CYCLE
995 0 : n_active = n_active + 1
996 0 : candidates(n_active) = ri_atom
997 0 : last_atom = ri_atom
998 : END DO
999 0 : ALLOCATE (ri_atoms(n_active), source=candidates(:n_active))
1000 : max_nRI_ref = 0
1001 0 : DO slot = 1, n_active
1002 0 : max_nRI_ref = MAX(max_nRI_ref, get_ref_RI_size(bs_env, ri_atoms(slot)))
1003 : END DO
1004 0 : ALLOCATE (U_active(max_nRI_ref, ncol, n_active), source=0.0_dp)
1005 :
1006 0 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
1007 0 : atom_A = bs_env%auto_ri%AB_atom_A(AB_block)
1008 0 : atom_B = bs_env%auto_ri%AB_atom_B(AB_block)
1009 0 : n_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)
1010 0 : IF (fit_atom == atom_A .AND. n_to_a > 0) THEN
1011 0 : column_first = bs_env%auto_ri%AB_first_p_A(AB_block)
1012 0 : first_p_AB = 1
1013 0 : ncol = n_to_a
1014 0 : ELSE IF (fit_atom == atom_B .AND. &
1015 : n_to_a < bs_env%auto_ri%AB_size_opt_RI(AB_block)) THEN
1016 0 : column_first = bs_env%auto_ri%AB_first_p_B(AB_block)
1017 0 : first_p_AB = n_to_a + 1
1018 0 : ncol = bs_env%auto_ri%AB_size_opt_RI(AB_block) - n_to_a
1019 : ELSE
1020 : CYCLE
1021 : END IF
1022 0 : CALL get_U_Pp_AB(bs_env%auto_ri, AB_block, U_Pp_AB)
1023 :
1024 0 : slot = locate(ri_atoms, atom_A)
1025 0 : nRI_ref_A = get_ref_RI_size(bs_env, atom_A)
1026 : U_active(1:nRI_ref_A, &
1027 : column_first:column_first + ncol - 1, slot) = &
1028 : U_Pp_AB( &
1029 0 : 1:nRI_ref_A, first_p_AB:first_p_AB + ncol - 1)
1030 0 : IF (atom_B /= atom_A) THEN
1031 0 : slot = locate(ri_atoms, atom_B)
1032 0 : nRI_ref_B = get_ref_RI_size(bs_env, atom_B)
1033 : U_active(1:nRI_ref_B, &
1034 : column_first:column_first + ncol - 1, slot) = &
1035 : U_Pp_AB( &
1036 : nRI_ref_A + 1:nRI_ref_A + nRI_ref_B, &
1037 0 : first_p_AB:first_p_AB + ncol - 1)
1038 : END IF
1039 0 : DEALLOCATE (U_Pp_AB)
1040 : END DO
1041 :
1042 0 : center = particle_set(fit_atom)%r
1043 0 : r_c = bs_env%ri_metric%cutoff_radius
1044 0 : cutoff_ri = 0.0_dp
1045 0 : DO slot = 1, n_active
1046 0 : ri_atom = ri_atoms(slot)
1047 : cutoff_ri = MAX(cutoff_ri, &
1048 : r_c + bs_env%ri_rs%radius_ri_per_atom(ri_atom) + &
1049 0 : NORM2(center - particle_set(ri_atom)%r))
1050 : END DO
1051 : CALL build_phi_on_sphere(bs_env, qs_kind_set, &
1052 : ri_rs_grid_points, fit_atom, cutoff_ri, &
1053 : local_grid_idx, ngrid, phi_local, ao_atoms, ao_col_start, &
1054 0 : center=center, ao_nonzero=ao_nonzero)
1055 0 : ncol = ri_blk_sizes(fit_atom)
1056 0 : previous_fit_peak = bs_env%ri_rs%fit_peak_bytes
1057 0 : CALL record_fit_memory(bs_env%ri_rs, ngrid, SIZE(phi_local, 2), ncol)
1058 0 : IF (bs_env%ri_rs%fit_peak_bytes > previous_fit_peak) THEN
1059 0 : CALL print_fit_memory_estimate(bs_env, mem_avail_GB)
1060 : END IF
1061 0 : ALLOCATE (d_lp_local(ngrid, ncol), source=0.0_dp)
1062 : ! d_lp = Σ_A Σ_{P∈A} d_lP U_Pp on the grid of this atomic column block.
1063 : CALL compute_d_lp_auto_ri(bs_env, ctx_3c, phi_local, ao_atoms, ao_col_start, ao_nonzero, &
1064 : d_lp_local, ri_atoms, U_active, max_ao_size, &
1065 0 : para_env_col%mepos, para_env_col%num_pe)
1066 0 : DEALLOCATE (ao_atoms, ao_col_start, ao_nonzero)
1067 0 : CALL para_env_col%sum(d_lp_local)
1068 0 : CALL solve_ri_rs_fit(phi_local, d_lp_local, bs_env%ri_rs%tikhonov)
1069 : CALL add_Z_lP_columns(mat_Z_lP, d_lp_local, local_grid_idx, ngrid, fit_atom, 1, &
1070 : ri_blk_sizes(fit_atom), row_size_grid, row_offset, &
1071 : row_dist_grid, &
1072 0 : myprow, bs_env%eps_filter)
1073 0 : DEALLOCATE (d_lp_local, local_grid_idx, phi_local)
1074 0 : DEALLOCATE (U_active, ri_atoms)
1075 0 : n_done = n_done + 1
1076 : CALL print_Z_lP_progress(bs_env, n_done, n_my_atoms, &
1077 0 : m_walltime() - item_start_time)
1078 : END DO
1079 0 : DEALLOCATE (candidates, candidate_order)
1080 0 : CALL para_env_col%free()
1081 0 : DEALLOCATE (para_env_col)
1082 : END IF
1083 : END IF
1084 4 : CALL gw_3c_ctx_release(ctx_3c)
1085 :
1086 4 : CALL dbcsr_filter(mat_Z_lP, bs_env%eps_filter)
1087 4 : CALL dbcsr_finalize(mat_Z_lP)
1088 4 : CALL para_env%sync()
1089 : CALL print_Z_lP_progress(bs_env, natom, natom, m_walltime() - t1, &
1090 4 : all_mpi_ranks=.TRUE.)
1091 4 : logger => cp_get_default_logger()
1092 4 : IF (BTEST(cp_print_key_should_output(logger%iter_info, input, key), cp_p_file)) THEN
1093 0 : CALL dbcsr_binary_write(matrix=mat_Z_lP, filepath=TRIM(bs_env%prefix)//'Z_lP.matrix')
1094 : END IF
1095 :
1096 4 : DEALLOCATE (row_offset, ri_blk_sizes, col_dist_ri)
1097 4 : CALL dbcsr_distribution_release(dist_Z)
1098 4 : CALL timestop(handle)
1099 :
1100 12 : END SUBROUTINE compute_Z_lP_auto_ri
1101 :
1102 : ! **************************************************************************************************
1103 : !> \brief Prints completed Z_lP atoms and execution time from the printing rank.
1104 : !> \param bs_env ...
1105 : !> \param n_done Number of atoms completed by the printing rank.
1106 : !> \param n_total Total number of atoms assigned to the printing rank.
1107 : !> \param execution_time Wall-clock time used for the reported work.
1108 : !> \param all_mpi_ranks Whether all MPI ranks have completed the calculation.
1109 : ! **************************************************************************************************
1110 102 : SUBROUTINE print_Z_lP_progress(bs_env, n_done, n_total, execution_time, all_mpi_ranks)
1111 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1112 : INTEGER, INTENT(IN) :: n_done, n_total
1113 : REAL(KIND=dp), INTENT(IN) :: execution_time
1114 : LOGICAL, INTENT(IN), OPTIONAL :: all_mpi_ranks
1115 :
1116 : CHARACTER(LEN=*), PARAMETER :: routineN = 'print_Z_lP_progress'
1117 :
1118 : INTEGER :: handle
1119 : LOGICAL :: completed
1120 :
1121 102 : CALL timeset(routineN, handle)
1122 :
1123 102 : IF (bs_env%unit_nr > 0) THEN
1124 56 : completed = .FALSE.
1125 56 : IF (PRESENT(all_mpi_ranks)) completed = all_mpi_ranks
1126 23 : IF (completed) THEN
1127 : WRITE (bs_env%unit_nr, '(T2,A,T58,A,F7.1,A,/)') &
1128 23 : 'Computed Z_lP (all MPI ranks) for all atoms,', &
1129 46 : 'Execution time', execution_time, ' s'
1130 : ELSE
1131 : WRITE (bs_env%unit_nr, '(T2,A,I11,A,I3,A,F7.1,A)') &
1132 33 : 'Computed Z_lP (MPI rank 0) for atom', n_done, ' /', n_total, &
1133 66 : ', Execution time', execution_time, ' s'
1134 : END IF
1135 56 : CALL m_flush(bs_env%unit_nr)
1136 : END IF
1137 :
1138 102 : CALL timestop(handle)
1139 :
1140 102 : END SUBROUTINE print_Z_lP_progress
1141 :
1142 : ! **************************************************************************************************
1143 : !> \brief Prepares the distributed Z_lP matrix and its atom-blocked column layout. Column block A
1144 : !> contains all optimized functions assigned to atom A. Every local block is reserved once
1145 : !> because several AA/AB contraction blocks may contribute to the same block.
1146 : !> \param bs_env ...
1147 : !> \param mat_phi_mu_l ...
1148 : !> \param mat_Z_lP ...
1149 : !> \param dist_Z ...
1150 : !> \param row_size_grid ...
1151 : !> \param row_dist_grid ...
1152 : !> \param col_dist_ri ...
1153 : !> \param ri_blk_sizes ...
1154 : !> \param row_offset ...
1155 : !> \param myprow ...
1156 : !> \param mypcol ...
1157 : !> \param max_ao_size ...
1158 : ! **************************************************************************************************
1159 4 : SUBROUTINE prepare_Z_lP_auto_ri(bs_env, mat_phi_mu_l, mat_Z_lP, dist_Z, row_size_grid, &
1160 : row_dist_grid, col_dist_ri, ri_blk_sizes, row_offset, &
1161 : myprow, mypcol, max_ao_size)
1162 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1163 : TYPE(dbcsr_type), INTENT(IN) :: mat_phi_mu_l
1164 : TYPE(dbcsr_type), INTENT(OUT) :: mat_Z_lP
1165 : TYPE(dbcsr_distribution_type), INTENT(OUT) :: dist_Z
1166 : INTEGER, DIMENSION(:), POINTER :: row_size_grid, row_dist_grid, &
1167 : col_dist_ri, ri_blk_sizes
1168 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: row_offset
1169 : INTEGER, INTENT(OUT) :: myprow, mypcol, max_ao_size
1170 :
1171 : CHARACTER(LEN=*), PARAMETER :: routineN = 'prepare_Z_lP_auto_ri'
1172 :
1173 : INTEGER :: handle, i_blk, iatom, npcol
1174 : TYPE(dbcsr_distribution_type) :: dist_phi
1175 :
1176 4 : CALL timeset(routineN, handle)
1177 :
1178 4 : CALL dbcsr_get_info(mat_phi_mu_l, row_blk_size=row_size_grid, distribution=dist_phi)
1179 : CALL dbcsr_distribution_get(dist_phi, row_dist=row_dist_grid, npcols=npcol, &
1180 4 : myprow=myprow, mypcol=mypcol)
1181 12 : ALLOCATE (row_offset(SIZE(row_size_grid)))
1182 4 : row_offset(1) = 0
1183 12 : DO i_blk = 2, SIZE(row_size_grid)
1184 12 : row_offset(i_blk) = row_offset(i_blk - 1) + row_size_grid(i_blk - 1)
1185 : END DO
1186 :
1187 20 : ALLOCATE (ri_blk_sizes(bs_env%n_atom), col_dist_ri(bs_env%n_atom))
1188 32 : ri_blk_sizes = bs_env%auto_ri%sizes_opt_RI
1189 16 : DO iatom = 1, bs_env%n_atom
1190 16 : col_dist_ri(iatom) = MOD(iatom - 1, npcol)
1191 : END DO
1192 : CALL dbcsr_distribution_new(dist_Z, template=dist_phi, row_dist=row_dist_grid, &
1193 4 : col_dist=col_dist_ri)
1194 : CALL dbcsr_create(mat_Z_lP, name='mat_Z_lP localized AA/AB', dist=dist_Z, &
1195 : matrix_type=dbcsr_type_no_symmetry, row_blk_size=row_size_grid, &
1196 4 : col_blk_size=ri_blk_sizes)
1197 4 : CALL dbcsr_reserve_all_blocks(mat_Z_lP)
1198 4 : CALL dbcsr_set(mat_Z_lP, 0.0_dp)
1199 :
1200 4 : max_ao_size = 0
1201 16 : DO iatom = 1, bs_env%n_atom
1202 : max_ao_size = MAX(max_ao_size, bs_env%i_ao_end_from_atom(iatom) - &
1203 16 : bs_env%i_ao_start_from_atom(iatom) + 1)
1204 : END DO
1205 :
1206 4 : CALL timestop(handle)
1207 :
1208 8 : END SUBROUTINE prepare_Z_lP_auto_ri
1209 :
1210 : ! **************************************************************************************************
1211 : !> \brief Tests whether every fitted atomic block can use the complete RI-RS grid.
1212 : !>
1213 : !> For a block centered on A, every grid point must satisfy
1214 : !>
1215 : !> |r_l - R_A| <= R_A^fit,
1216 : !>
1217 : !> and every atom B whose AOs can contribute must satisfy
1218 : !>
1219 : !> |R_B - R_A| <= R_B^AO + R_A^fit.
1220 : !>
1221 : !> Pair-midpoint grids do not satisfy this atom-centered criterion in general.
1222 : !> \param bs_env ...
1223 : !> \param ri_rs_grid_points ...
1224 : !> \param common_grid_available ...
1225 : !> \param have_fitted_columns ...
1226 : ! **************************************************************************************************
1227 4 : SUBROUTINE common_Z_lP_grid_available(bs_env, ri_rs_grid_points, &
1228 : common_grid_available, have_fitted_columns)
1229 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1230 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: ri_rs_grid_points
1231 : LOGICAL, INTENT(OUT) :: common_grid_available, &
1232 : have_fitted_columns
1233 :
1234 : CHARACTER(LEN=*), PARAMETER :: routineN = 'common_Z_lP_grid_available'
1235 :
1236 : INTEGER :: AB_block, atom_A, atom_B, fit_atom, &
1237 : handle, iatom, igrid, n_to_a, natom
1238 4 : LOGICAL, ALLOCATABLE, DIMENSION(:) :: active_atom
1239 : REAL(KIND=dp) :: cutoff_ri
1240 : REAL(KIND=dp), DIMENSION(3) :: center
1241 4 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
1242 :
1243 4 : CALL timeset(routineN, handle)
1244 :
1245 4 : particle_set => bs_env%ri_rs%particle_set
1246 4 : natom = bs_env%n_atom
1247 4 : common_grid_available = .TRUE.
1248 4 : have_fitted_columns = .FALSE.
1249 12 : ALLOCATE (active_atom(natom))
1250 16 : DO fit_atom = 1, natom
1251 12 : IF (bs_env%auto_ri%sizes_opt_RI(fit_atom) == 0) CYCLE
1252 12 : have_fitted_columns = .TRUE.
1253 12 : active_atom = .FALSE.
1254 66 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
1255 54 : atom_A = bs_env%auto_ri%AB_atom_A(AB_block)
1256 54 : atom_B = bs_env%auto_ri%AB_atom_B(AB_block)
1257 54 : n_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)
1258 54 : IF (.NOT. (fit_atom == atom_A .AND. n_to_a > 0) .AND. &
1259 : .NOT. (fit_atom == atom_B .AND. &
1260 : n_to_a < bs_env%auto_ri%AB_size_opt_RI(AB_block))) CYCLE
1261 18 : active_atom(atom_A) = .TRUE.
1262 30 : IF (atom_B /= atom_A) active_atom(atom_B) = .TRUE.
1263 : END DO
1264 :
1265 48 : center = particle_set(fit_atom)%r
1266 : cutoff_ri = 0.0_dp
1267 48 : DO iatom = 1, natom
1268 36 : IF (.NOT. active_atom(iatom)) CYCLE
1269 : cutoff_ri = MAX(cutoff_ri, bs_env%ri_metric%cutoff_radius + &
1270 : bs_env%ri_rs%radius_ri_per_atom(iatom) + &
1271 102 : NORM2(center - particle_set(iatom)%r))
1272 : END DO
1273 :
1274 390 : DO igrid = 1, bs_env%ri_rs%n_grid_points
1275 1524 : IF (NORM2(ri_rs_grid_points(1:3, igrid) - center) > cutoff_ri) THEN
1276 0 : common_grid_available = .FALSE.
1277 0 : EXIT
1278 : END IF
1279 : END DO
1280 12 : IF (.NOT. common_grid_available) EXIT
1281 :
1282 48 : DO iatom = 1, natom
1283 144 : IF (NORM2(particle_set(iatom)%r - center) > &
1284 12 : bs_env%ri_rs%radius_ao_per_atom(iatom) + cutoff_ri) THEN
1285 0 : common_grid_available = .FALSE.
1286 0 : EXIT
1287 : END IF
1288 : END DO
1289 16 : IF (.NOT. common_grid_available) EXIT
1290 : END DO
1291 4 : DEALLOCATE (active_atom)
1292 :
1293 4 : CALL timestop(handle)
1294 :
1295 4 : END SUBROUTINE common_Z_lP_grid_available
1296 :
1297 : ! **************************************************************************************************
1298 : !> \brief Evaluates the AO collocation matrix on the complete RI-RS grid,
1299 : !>
1300 : !> ϕ_lμ = ϕ_μ(r_l), l = 1, ..., N_grid.
1301 : !>
1302 : !> "Complete common grid" means that the same full set of grid points is valid for every
1303 : !> optimized RI column block. The enclosing sphere is only an implementation device passed
1304 : !> to build_phi_on_sphere; it contains every r_l and every AO center, so its chosen center
1305 : !> does not select an AA or AB contraction.
1306 : !> \param bs_env ...
1307 : !> \param qs_kind_set ...
1308 : !> \param ri_rs_grid_points ...
1309 : !> \param local_grid_idx ...
1310 : !> \param ngrid ...
1311 : !> \param phi_local ...
1312 : !> \param ao_atoms global indices of the active AO atoms
1313 : !> \param ao_col_start local phi column starts, including the terminal column offset
1314 : !> \param ao_nonzero exact nonzero AO support on the retained grid rows
1315 : ! **************************************************************************************************
1316 8 : SUBROUTINE build_phi_on_complete_grid(bs_env, qs_kind_set, &
1317 4 : ri_rs_grid_points, local_grid_idx, ngrid, &
1318 : phi_local, ao_atoms, ao_col_start, ao_nonzero)
1319 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1320 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
1321 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: ri_rs_grid_points
1322 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: local_grid_idx
1323 : INTEGER, INTENT(OUT) :: ngrid
1324 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
1325 : INTENT(OUT) :: phi_local
1326 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: ao_atoms, ao_col_start
1327 : LOGICAL, ALLOCATABLE, INTENT(OUT) :: ao_nonzero(:, :)
1328 :
1329 : CHARACTER(LEN=*), PARAMETER :: routineN = 'build_phi_on_complete_grid'
1330 :
1331 : INTEGER :: handle, iatom, igrid, reference_atom
1332 : REAL(KIND=dp) :: cutoff_ri
1333 : REAL(KIND=dp), DIMENSION(3) :: center
1334 4 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
1335 :
1336 4 : CALL timeset(routineN, handle)
1337 :
1338 4 : particle_set => bs_env%ri_rs%particle_set
1339 4 : reference_atom = 1
1340 16 : center = particle_set(reference_atom)%r
1341 4 : cutoff_ri = 0.0_dp
1342 130 : DO igrid = 1, bs_env%ri_rs%n_grid_points
1343 508 : cutoff_ri = MAX(cutoff_ri, NORM2(ri_rs_grid_points(1:3, igrid) - center))
1344 : END DO
1345 16 : DO iatom = 1, bs_env%n_atom
1346 52 : cutoff_ri = MAX(cutoff_ri, NORM2(particle_set(iatom)%r - center))
1347 : END DO
1348 4 : cutoff_ri = cutoff_ri + 1.0_dp
1349 :
1350 : CALL build_phi_on_sphere(bs_env, qs_kind_set, &
1351 : ri_rs_grid_points, reference_atom, cutoff_ri, &
1352 : local_grid_idx, ngrid, phi_local, ao_atoms, ao_col_start, &
1353 4 : center=center, ao_nonzero=ao_nonzero)
1354 :
1355 4 : CALL timestop(handle)
1356 :
1357 4 : END SUBROUTINE build_phi_on_complete_grid
1358 :
1359 : ! **************************************************************************************************
1360 : !> \brief Writes Z_l,p0+p += z_lp into the process-owned grid-row blocks of one atomic column
1361 : !> block. The routine performs no explicit MPI communication; DBCSR summation combines
1362 : !> contributions when a block receives columns from more than one contraction group.
1363 : !> \param mat_Z_lP ...
1364 : !> \param z_block ...
1365 : !> \param local_grid_idx ...
1366 : !> \param n_local_grid ...
1367 : !> \param atom_index ...
1368 : !> \param first_column ...
1369 : !> \param atom_block_size ...
1370 : !> \param r_blk_sizes ...
1371 : !> \param row_offset ...
1372 : !> \param row_dist ...
1373 : !> \param myprow ...
1374 : !> \param eps_filter ...
1375 : ! **************************************************************************************************
1376 18 : SUBROUTINE add_Z_lP_columns(mat_Z_lP, z_block, local_grid_idx, n_local_grid, atom_index, &
1377 18 : first_column, atom_block_size, r_blk_sizes, row_offset, row_dist, &
1378 : myprow, eps_filter)
1379 :
1380 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_Z_lP
1381 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: z_block
1382 : INTEGER, DIMENSION(:), INTENT(IN) :: local_grid_idx
1383 : INTEGER, INTENT(IN) :: n_local_grid, atom_index, first_column, &
1384 : atom_block_size
1385 : INTEGER, DIMENSION(:), INTENT(IN) :: r_blk_sizes, row_offset, row_dist
1386 : INTEGER, INTENT(IN) :: myprow
1387 : REAL(KIND=dp), INTENT(IN) :: eps_filter
1388 :
1389 : CHARACTER(LEN=*), PARAMETER :: routineN = 'add_Z_lP_columns'
1390 :
1391 : INTEGER :: current_chunk_size, g_pt, handle, i_blk, &
1392 : loc_ptr, ncolumn, r_end, r_start
1393 : LOGICAL :: row_owned
1394 18 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: Z_blk
1395 :
1396 18 : CALL timeset(routineN, handle)
1397 :
1398 18 : ncolumn = SIZE(z_block, 2)
1399 18 : CPASSERT(first_column > 0)
1400 18 : CPASSERT(first_column + ncolumn - 1 <= atom_block_size)
1401 126 : ALLOCATE (Z_blk(MAXVAL(r_blk_sizes), atom_block_size), source=0.0_dp)
1402 18 : loc_ptr = 1
1403 72 : DO i_blk = 1, SIZE(r_blk_sizes)
1404 54 : r_start = row_offset(i_blk) + 1
1405 54 : r_end = row_offset(i_blk) + r_blk_sizes(i_blk)
1406 54 : current_chunk_size = r_blk_sizes(i_blk)
1407 54 : row_owned = row_dist(i_blk) == myprow
1408 54 : Z_blk = 0.0_dp
1409 684 : DO WHILE (loc_ptr <= n_local_grid)
1410 666 : g_pt = local_grid_idx(loc_ptr)
1411 666 : IF (g_pt > r_end) EXIT
1412 630 : IF (row_owned) THEN
1413 : Z_blk(g_pt - r_start + 1, first_column:first_column + ncolumn - 1) = &
1414 1302 : z_block(loc_ptr, 1:ncolumn)
1415 : END IF
1416 666 : loc_ptr = loc_ptr + 1
1417 : END DO
1418 3750 : IF (row_owned .AND. MAXVAL(ABS(Z_blk(1:current_chunk_size, :))) > eps_filter) THEN
1419 : CALL dbcsr_put_block(mat_Z_lP, row=i_blk, col=atom_index, &
1420 27 : block=Z_blk(1:current_chunk_size, :), summation=.TRUE.)
1421 : END IF
1422 : END DO
1423 18 : DEALLOCATE (Z_blk)
1424 :
1425 18 : CALL timestop(handle)
1426 :
1427 18 : END SUBROUTINE add_Z_lP_columns
1428 :
1429 : ! **************************************************************************************************
1430 : !> \brief Computes d_lp = Σ_{μνP} ϕ_μ(r_l)ϕ_ν(r_l)(μν|P)U_Pp for all AA and AB
1431 : !> columns p of one RI-RS matrix block. P can belong to either atom of an AB contraction.
1432 : !> \param bs_env ...
1433 : !> \param ctx ...
1434 : !> \param phi_val ...
1435 : !> \param ao_atoms global indices of the active AO atoms
1436 : !> \param ao_col_start local phi column starts, including the terminal column offset
1437 : !> \param ao_nonzero exact nonzero AO support on the retained grid rows
1438 : !> \param d_lp ...
1439 : !> \param ri_atoms global indices of contributing RI atoms in ascending order
1440 : !> \param ri_coefficients ...
1441 : !> \param max_ao_size ...
1442 : !> \param atom_j_mepos ...
1443 : !> \param atom_j_stride ...
1444 : ! **************************************************************************************************
1445 12 : SUBROUTINE compute_d_lp_auto_ri(bs_env, ctx, phi_val, ao_atoms, ao_col_start, ao_nonzero, d_lp, &
1446 6 : ri_atoms, ri_coefficients, max_ao_size, atom_j_mepos, &
1447 : atom_j_stride)
1448 : !$ USE OMP_LIB, ONLY: omp_get_max_threads, omp_get_thread_num
1449 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1450 : TYPE(gw_3c_ctx_type), INTENT(IN) :: ctx
1451 : REAL(KIND=dp), CONTIGUOUS, DIMENSION(:, :), &
1452 : INTENT(IN) :: phi_val
1453 : INTEGER, DIMENSION(:), INTENT(IN) :: ao_atoms, ao_col_start
1454 : LOGICAL, DIMENSION(:, :), INTENT(IN) :: ao_nonzero
1455 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: d_lp
1456 : REAL(KIND=dp), CONTIGUOUS, DIMENSION(:, :, :), INTENT(IN) :: ri_coefficients
1457 : INTEGER, DIMENSION(:), INTENT(IN) :: ri_atoms
1458 : INTEGER, INTENT(IN) :: max_ao_size, atom_j_mepos, atom_j_stride
1459 :
1460 : CHARACTER(LEN=*), PARAMETER :: routineN = 'compute_d_lp_auto_ri'
1461 : INTEGER, PARAMETER :: grid_chunk = 1024
1462 :
1463 : INTEGER :: jj, kk, active_column, iatom, jatom, katom, c, handle, i_thread, j, jk_idx, jsize, &
1464 : jstart, k, ksize, kstart, l, l0, &
1465 : max_active, nactive, ncol, n_grid, nRI_ref, nthreads, ri, slot, thread_id
1466 6 : INTEGER, ALLOCATABLE, DIMENSION(:) :: atom_ncol
1467 6 : INTEGER, ALLOCATABLE, DIMENSION(:, :) :: atom_column
1468 : LOGICAL :: any_integral, screened, single_atom
1469 : REAL(KIND=dp) :: pair_factor
1470 6 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: int_2d_prv, rho_chunk
1471 6 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: coefficient_compact, d_lp_threads, &
1472 6 : int_3c_prv, int_3c_atom
1473 6 : TYPE(gw_3c_ws_type) :: ws
1474 :
1475 6 : INTEGER, ALLOCATABLE, DIMENSION(:) :: grid_index
1476 : INTEGER :: n_grid_pair, grid_l, point
1477 6 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: grid_result
1478 :
1479 6 : CALL timeset(routineN, handle)
1480 6 : n_grid = SIZE(d_lp, 1)
1481 6 : ncol = SIZE(d_lp, 2)
1482 6 : single_atom = SIZE(ri_atoms) == 1
1483 6 : nthreads = 1
1484 6 : !$ nthreads = omp_get_max_threads()
1485 6 : CPASSERT(SIZE(ri_coefficients, 2) == ncol)
1486 6 : CPASSERT(SIZE(ri_coefficients, 3) == SIZE(ri_atoms))
1487 6 : CPASSERT(SIZE(ri_atoms) > 0)
1488 6 : max_active = 0
1489 6 : IF (.NOT. single_atom) THEN
1490 0 : ALLOCATE (atom_ncol(SIZE(ri_atoms)), atom_column(ncol, SIZE(ri_atoms)))
1491 0 : atom_ncol = 0
1492 0 : atom_column = 0
1493 0 : DO slot = 1, SIZE(ri_atoms)
1494 0 : nRI_ref = get_ref_RI_size(bs_env, ri_atoms(slot))
1495 0 : DO ri = 1, ncol
1496 0 : IF (.NOT. ANY(ri_coefficients(1:nRI_ref, ri, slot) /= 0.0_dp)) CYCLE
1497 0 : atom_ncol(slot) = atom_ncol(slot) + 1
1498 0 : atom_column(atom_ncol(slot), slot) = ri
1499 : END DO
1500 : END DO
1501 0 : max_active = MAXVAL(atom_ncol)
1502 0 : CPASSERT(max_active > 0)
1503 0 : ALLOCATE (coefficient_compact(SIZE(ri_coefficients, 1), max_active, &
1504 0 : SIZE(ri_atoms)), source=0.0_dp)
1505 0 : DO slot = 1, SIZE(ri_atoms)
1506 0 : nRI_ref = get_ref_RI_size(bs_env, ri_atoms(slot))
1507 0 : DO active_column = 1, atom_ncol(slot)
1508 0 : ri = atom_column(active_column, slot)
1509 : coefficient_compact(1:nRI_ref, active_column, slot) = &
1510 0 : ri_coefficients(1:nRI_ref, ri, slot)
1511 : END DO
1512 : END DO
1513 : END IF
1514 30 : ALLOCATE (d_lp_threads(n_grid, ncol, nthreads), source=0.0_dp)
1515 :
1516 : !$OMP PARALLEL DEFAULT(NONE) &
1517 : !$OMP SHARED(bs_env, ctx, phi_val, ao_atoms, ao_col_start, ao_nonzero, d_lp, n_grid, ri_coefficients, &
1518 : !$OMP ri_atoms, single_atom, max_ao_size, atom_j_mepos, atom_j_stride, ncol, &
1519 : !$OMP d_lp_threads, nthreads, atom_ncol, &
1520 : !$OMP atom_column, coefficient_compact, max_active) &
1521 : !$OMP PRIVATE(grid_index, n_grid_pair, grid_l, point, grid_result, &
1522 : !$OMP jj, kk, active_column, iatom, jatom, katom, c, i_thread, j, jk_idx, jsize, jstart, &
1523 : !$OMP k, ksize, kstart, l, l0, &
1524 : !$OMP nactive, nRI_ref, ri, slot, any_integral, screened, pair_factor, &
1525 : !$OMP int_2d_prv, rho_chunk, &
1526 6 : !$OMP int_3c_prv, int_3c_atom, ws, thread_id)
1527 :
1528 : thread_id = 1
1529 : !$ thread_id = omp_get_thread_num() + 1
1530 :
1531 : CALL gw_3c_ws_create(ws, ctx)
1532 : ALLOCATE (int_3c_prv(max_ao_size, max_ao_size, ncol))
1533 : IF (.NOT. single_atom) ALLOCATE (int_3c_atom(max_ao_size, max_ao_size, max_active))
1534 : ALLOCATE (int_2d_prv(max_ao_size*max_ao_size, ncol))
1535 : ALLOCATE (rho_chunk(grid_chunk, max_ao_size*max_ao_size))
1536 : ALLOCATE (grid_index(n_grid), grid_result(grid_chunk, ncol))
1537 :
1538 : !$OMP DO SCHEDULE(DYNAMIC)
1539 : DO jj = 1, SIZE(ao_atoms)
1540 : jatom = ao_atoms(jj)
1541 : IF (MOD(jatom - 1, atom_j_stride) /= atom_j_mepos) CYCLE
1542 : DO kk = jj, SIZE(ao_atoms)
1543 : katom = ao_atoms(kk)
1544 : jstart = ao_col_start(jj)
1545 : kstart = ao_col_start(kk)
1546 : jsize = ao_col_start(jj + 1) - jstart
1547 : ksize = ao_col_start(kk + 1) - kstart
1548 :
1549 : n_grid_pair = 0
1550 : DO grid_l = 1, n_grid
1551 : IF (.NOT. (ao_nonzero(grid_l, jj) .AND. ao_nonzero(grid_l, kk))) CYCLE
1552 : n_grid_pair = n_grid_pair + 1
1553 : grid_index(n_grid_pair) = grid_l
1554 : END DO
1555 : IF (n_grid_pair == 0) CYCLE
1556 : int_3c_prv(1:jsize, 1:ksize, 1:ncol) = 0.0_dp
1557 : IF (single_atom) THEN
1558 : ! A single contributor writes the transformed block without packing or scattering.
1559 : CALL build_3c_integral_block_auto_ri_ctx( &
1560 : int_3c_prv(1:jsize, 1:ksize, 1:ncol), ctx, ws, &
1561 : atom_j=jatom, atom_k=katom, atom_i=ri_atoms(1), &
1562 : transform=ri_coefficients(:, :, 1), transform_row=1, screened=screened)
1563 : IF (screened) CYCLE
1564 : ELSE
1565 : any_integral = .FALSE.
1566 : DO slot = 1, SIZE(ri_atoms)
1567 : iatom = ri_atoms(slot)
1568 : nactive = atom_ncol(slot)
1569 : IF (nactive == 0) CYCLE
1570 : nRI_ref = get_ref_RI_size(bs_env, iatom)
1571 : int_3c_atom(1:jsize, 1:ksize, 1:nactive) = 0.0_dp
1572 : CALL build_3c_integral_block_auto_ri_ctx( &
1573 : int_3c_atom(1:jsize, 1:ksize, 1:nactive), ctx, ws, &
1574 : atom_j=jatom, atom_k=katom, atom_i=iatom, &
1575 : transform=coefficient_compact(1:nRI_ref, 1:nactive, slot), &
1576 : transform_row=1, screened=screened)
1577 : IF (.NOT. screened) THEN
1578 : any_integral = .TRUE.
1579 : DO active_column = 1, nactive
1580 : ri = atom_column(active_column, slot)
1581 : int_3c_prv(1:jsize, 1:ksize, ri) = &
1582 : int_3c_prv(1:jsize, 1:ksize, ri) + &
1583 : int_3c_atom(1:jsize, 1:ksize, active_column)
1584 : END DO
1585 : END IF
1586 : END DO
1587 : IF (.NOT. any_integral) CYCLE
1588 : END IF
1589 :
1590 : DO ri = 1, ncol
1591 : DO k = 1, ksize
1592 : DO j = 1, jsize
1593 : jk_idx = (k - 1)*jsize + j
1594 : int_2d_prv(jk_idx, ri) = int_3c_prv(j, k, ri)
1595 : END DO
1596 : END DO
1597 : END DO
1598 :
1599 : pair_factor = 1.0_dp
1600 : IF (jatom /= katom) pair_factor = 2.0_dp
1601 : DO l0 = 1, n_grid_pair, grid_chunk
1602 : c = MIN(grid_chunk, n_grid_pair - l0 + 1)
1603 : DO k = 1, ksize
1604 : DO j = 1, jsize
1605 : jk_idx = (k - 1)*jsize + j
1606 : DO l = 1, c
1607 : point = grid_index(l0 + l - 1)
1608 : rho_chunk(l, jk_idx) = phi_val(point, jstart + j - 1)* &
1609 : phi_val(point, kstart + k - 1)
1610 : END DO
1611 : END DO
1612 : END DO
1613 : CALL dgemm('N', 'N', c, ncol, jsize*ksize, pair_factor, rho_chunk, grid_chunk, &
1614 : int_2d_prv, max_ao_size*max_ao_size, 0.0_dp, &
1615 : grid_result, grid_chunk)
1616 : DO ri = 1, ncol
1617 : DO l = 1, c
1618 : point = grid_index(l0 + l - 1)
1619 : d_lp_threads(point, ri, thread_id) = &
1620 : d_lp_threads(point, ri, thread_id) + grid_result(l, ri)
1621 : END DO
1622 : END DO
1623 : END DO
1624 : END DO
1625 : END DO
1626 : !$OMP END DO
1627 :
1628 : !$OMP DO COLLAPSE(2) SCHEDULE(STATIC)
1629 : DO ri = 1, ncol
1630 : DO l = 1, n_grid
1631 : DO i_thread = 1, nthreads
1632 : d_lp(l, ri) = d_lp(l, ri) + d_lp_threads(l, ri, i_thread)
1633 : END DO
1634 : END DO
1635 : END DO
1636 : !$OMP END DO
1637 : IF (.NOT. single_atom) DEALLOCATE (int_3c_atom)
1638 : DEALLOCATE (int_3c_prv, int_2d_prv, rho_chunk)
1639 : DEALLOCATE (grid_index, grid_result)
1640 : CALL gw_3c_ws_release(ws)
1641 : !$OMP END PARALLEL
1642 :
1643 6 : IF (.NOT. single_atom) DEALLOCATE (coefficient_compact, atom_column, atom_ncol)
1644 6 : DEALLOCATE (d_lp_threads)
1645 :
1646 6 : CALL timestop(handle)
1647 :
1648 12 : END SUBROUTINE compute_d_lp_auto_ri
1649 :
1650 : ! **************************************************************************************************
1651 : !> \brief Computes d_lp = Σ_A Σ_{P∈A} d_lP U_Pp in rank-sized batches on a common grid.
1652 : !> \param bs_env ...
1653 : !> \param ctx ...
1654 : !> \param phi_val ...
1655 : !> \param ao_atoms global indices of the active AO atoms
1656 : !> \param ao_col_start local phi column starts, including the terminal column offset
1657 : !> \param ao_nonzero exact nonzero AO support on the retained grid rows
1658 : !> \param d_lp ...
1659 : !> \param n_grid ...
1660 : !> \param max_ao_size ...
1661 : !> \param mepos ...
1662 : !> \param num_pe ...
1663 : ! **************************************************************************************************
1664 4 : SUBROUTINE compute_d_lp_auto_ri_batch(bs_env, ctx, phi_val, ao_atoms, ao_col_start, ao_nonzero, d_lp, n_grid, &
1665 : max_ao_size, mepos, num_pe)
1666 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1667 : TYPE(gw_3c_ctx_type), INTENT(IN) :: ctx
1668 : REAL(KIND=dp), CONTIGUOUS, DIMENSION(:, :), &
1669 : INTENT(IN) :: phi_val
1670 : INTEGER, DIMENSION(:), INTENT(IN) :: ao_atoms, ao_col_start
1671 : LOGICAL, DIMENSION(:, :), INTENT(IN) :: ao_nonzero
1672 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: d_lp
1673 : INTEGER, INTENT(IN) :: n_grid, max_ao_size, mepos, num_pe
1674 :
1675 : CHARACTER(LEN=*), PARAMETER :: routineN = 'compute_d_lp_auto_ri_batch'
1676 :
1677 : INTEGER :: column, handle, iatom, n_done, n_total, &
1678 : ncol_batch
1679 4 : INTEGER, ALLOCATABLE, DIMENSION(:) :: global_map
1680 : REAL(KIND=dp) :: item_start_time
1681 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: d_lp_batch
1682 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: U_Pp
1683 :
1684 4 : CALL timeset(routineN, handle)
1685 :
1686 4 : n_total = 0
1687 4 : DO iatom = mepos + 1, bs_env%n_atom, num_pe
1688 6 : n_total = n_total + 1
1689 : END DO
1690 4 : n_done = 0
1691 :
1692 10 : DO iatom = mepos + 1, bs_env%n_atom, num_pe
1693 6 : item_start_time = m_walltime()
1694 6 : CALL collect_auto_ri_columns_for_atom(bs_env, iatom, U_Pp, global_map)
1695 6 : ncol_batch = SIZE(global_map)
1696 6 : IF (ncol_batch == 0) THEN
1697 0 : DEALLOCATE (U_Pp, global_map)
1698 0 : n_done = n_done + 1
1699 : CALL print_Z_lP_progress(bs_env, n_done, n_total, &
1700 0 : m_walltime() - item_start_time)
1701 0 : CYCLE
1702 : END IF
1703 :
1704 24 : ALLOCATE (d_lp_batch(n_grid, ncol_batch), source=0.0_dp)
1705 : CALL compute_d_lp_auto_ri(bs_env, ctx, phi_val, ao_atoms, ao_col_start, ao_nonzero, d_lp_batch, &
1706 12 : [iatom], U_Pp, max_ao_size, 0, 1)
1707 44 : DO column = 1, ncol_batch
1708 1409 : d_lp(:, global_map(column)) = d_lp(:, global_map(column)) + d_lp_batch(:, column)
1709 : END DO
1710 6 : DEALLOCATE (U_Pp, global_map, d_lp_batch)
1711 6 : n_done = n_done + 1
1712 : CALL print_Z_lP_progress(bs_env, n_done, n_total, &
1713 10 : m_walltime() - item_start_time)
1714 : END DO
1715 4 : CALL timestop(handle)
1716 :
1717 4 : END SUBROUTINE compute_d_lp_auto_ri_batch
1718 :
1719 : ! **************************************************************************************************
1720 : !> \brief Collects every optimized column p containing reference functions P on atom A. It returns
1721 : !>
1722 : !> U_Pp^A
1723 : !>
1724 : !> together with the global optimized-column index of q.
1725 : !> \param bs_env ...
1726 : !> \param iatom ...
1727 : !> \param U_Pp ...
1728 : !> \param global_map ...
1729 : ! **************************************************************************************************
1730 6 : SUBROUTINE collect_auto_ri_columns_for_atom(bs_env, iatom, U_Pp, global_map)
1731 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1732 : INTEGER, INTENT(IN) :: iatom
1733 : REAL(KIND=dp), ALLOCATABLE, INTENT(OUT) :: U_Pp(:, :, :)
1734 : INTEGER, ALLOCATABLE, INTENT(OUT) :: global_map(:)
1735 :
1736 : CHARACTER(LEN=*), PARAMETER :: routineN = 'collect_auto_ri_columns_for_atom'
1737 :
1738 : INTEGER :: AB_block, atom_A, atom_B, column, global_column, handle, local_column, ncol, &
1739 : ncol_batch, nRI_ref, nRI_ref_A, output_offset, row_first
1740 6 : REAL(KIND=dp), ALLOCATABLE :: U_Pp_AB(:, :)
1741 :
1742 6 : CALL timeset(routineN, handle)
1743 :
1744 6 : ncol_batch = 0
1745 33 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
1746 27 : atom_A = bs_env%auto_ri%AB_atom_A(AB_block)
1747 27 : atom_B = bs_env%auto_ri%AB_atom_B(AB_block)
1748 33 : IF (iatom == atom_A .OR. iatom == atom_B) THEN
1749 12 : ncol_batch = ncol_batch + bs_env%auto_ri%AB_size_opt_RI(AB_block)
1750 : END IF
1751 : END DO
1752 :
1753 6 : nRI_ref = get_ref_RI_size(bs_env, iatom)
1754 24 : ALLOCATE (U_Pp(nRI_ref, ncol_batch, 1), source=0.0_dp)
1755 18 : ALLOCATE (global_map(ncol_batch))
1756 6 : column = 0
1757 33 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
1758 27 : atom_A = bs_env%auto_ri%AB_atom_A(AB_block)
1759 27 : atom_B = bs_env%auto_ri%AB_atom_B(AB_block)
1760 27 : IF (iatom /= atom_A .AND. iatom /= atom_B) CYCLE
1761 12 : ncol = bs_env%auto_ri%AB_size_opt_RI(AB_block)
1762 12 : CALL get_U_Pp_AB(bs_env%auto_ri, AB_block, U_Pp_AB)
1763 12 : row_first = 1
1764 12 : IF (iatom == atom_B .AND. atom_B /= atom_A) THEN
1765 3 : nRI_ref_A = get_ref_RI_size(bs_env, atom_A)
1766 3 : row_first = 1 + nRI_ref_A
1767 : END IF
1768 : U_Pp(:, column + 1:column + ncol, 1) = &
1769 : U_Pp_AB( &
1770 3892 : row_first:row_first + nRI_ref - 1, 1:ncol)
1771 :
1772 50 : DO local_column = 1, ncol
1773 38 : IF (local_column <= bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)) THEN
1774 : output_offset = SUM(bs_env%auto_ri%sizes_opt_RI(:atom_A - 1)) + &
1775 30 : bs_env%auto_ri%AB_first_p_A(AB_block) - 1
1776 22 : global_column = output_offset + local_column
1777 : ELSE
1778 16 : CPASSERT(atom_B /= atom_A)
1779 : output_offset = SUM(bs_env%auto_ri%sizes_opt_RI(:atom_B - 1)) + &
1780 40 : bs_env%auto_ri%AB_first_p_B(AB_block) - 1
1781 : global_column = output_offset + local_column - &
1782 16 : bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)
1783 : END IF
1784 50 : global_map(column + local_column) = global_column
1785 : END DO
1786 12 : column = column + ncol
1787 33 : DEALLOCATE (U_Pp_AB)
1788 : END DO
1789 6 : CPASSERT(column == ncol_batch)
1790 :
1791 6 : CALL timestop(handle)
1792 :
1793 12 : END SUBROUTINE collect_auto_ri_columns_for_atom
1794 :
1795 : ! **************************************************************************************************
1796 : !> \brief Unpacks one AB contraction matrix U_Pp; B=A denotes an AA block.
1797 : !> \param auto_ri ...
1798 : !> \param AB_block ...
1799 : !> \param U_Pp ...
1800 : ! **************************************************************************************************
1801 12 : SUBROUTINE get_U_Pp_AB(auto_ri, AB_block, U_Pp)
1802 : TYPE(auto_ri_type), INTENT(IN) :: auto_ri
1803 : INTEGER, INTENT(IN) :: AB_block
1804 : REAL(KIND=dp), ALLOCATABLE, INTENT(OUT) :: U_Pp(:, :)
1805 :
1806 : INTEGER :: first, last, ncolumn, nrow
1807 :
1808 12 : nrow = auto_ri%AB_size_ref_RI(AB_block)
1809 12 : ncolumn = auto_ri%AB_size_opt_RI(AB_block)
1810 12 : first = auto_ri%U_Pp_AB_offset(AB_block)
1811 12 : last = first + nrow*ncolumn - 1
1812 48 : ALLOCATE (U_Pp(nrow, ncolumn))
1813 36 : U_Pp(:, :) = RESHAPE(auto_ri%U_Pp_AB(first:last), [nrow, ncolumn])
1814 :
1815 12 : END SUBROUTINE get_U_Pp_AB
1816 :
1817 : ! **************************************************************************************************
1818 : !> \brief Computes the fitting radius for every optimized atomic column block. If q assigned to A
1819 : !> contains reference functions on B, the required radius is
1820 : !>
1821 : !> R_A^fit = max_B [r_c + R_B^RI + |R_A - R_B|].
1822 : !> \param bs_env ...
1823 : !> \param radius ...
1824 : ! **************************************************************************************************
1825 0 : SUBROUTINE compute_auto_ri_grid_radii(bs_env, radius)
1826 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1827 : REAL(KIND=dp), INTENT(OUT) :: radius(:)
1828 :
1829 : CHARACTER(LEN=*), PARAMETER :: routineN = 'compute_auto_ri_grid_radii'
1830 :
1831 : INTEGER :: a, AB_block, b, handle, n_to_a, ncol
1832 : REAL(KIND=dp) :: cutoff, distance
1833 0 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
1834 :
1835 0 : CALL timeset(routineN, handle)
1836 :
1837 0 : particle_set => bs_env%ri_rs%particle_set
1838 0 : radius = 0.0_dp
1839 0 : cutoff = bs_env%ri_metric%cutoff_radius
1840 0 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
1841 0 : a = bs_env%auto_ri%AB_atom_A(AB_block)
1842 0 : b = bs_env%auto_ri%AB_atom_B(AB_block)
1843 0 : n_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)
1844 0 : ncol = bs_env%auto_ri%AB_size_opt_RI(AB_block)
1845 0 : IF (n_to_a > 0) THEN
1846 0 : radius(a) = MAX(radius(a), cutoff + bs_env%ri_rs%radius_ri_per_atom(a))
1847 0 : IF (b /= a) THEN
1848 0 : distance = NORM2(particle_set(a)%r - particle_set(b)%r)
1849 0 : radius(a) = MAX(radius(a), cutoff + bs_env%ri_rs%radius_ri_per_atom(b) + distance)
1850 : END IF
1851 : END IF
1852 0 : IF (n_to_a < ncol) THEN
1853 0 : CPASSERT(b /= a)
1854 0 : distance = NORM2(particle_set(b)%r - particle_set(a)%r)
1855 : radius(b) = MAX(radius(b), cutoff + bs_env%ri_rs%radius_ri_per_atom(b), &
1856 0 : cutoff + bs_env%ri_rs%radius_ri_per_atom(a) + distance)
1857 : END IF
1858 : END DO
1859 :
1860 0 : CALL timestop(handle)
1861 0 : END SUBROUTINE compute_auto_ri_grid_radii
1862 :
1863 : ! **************************************************************************************************
1864 : !> \brief Computes
1865 : !>
1866 : !> d_lp = Σ_A Σ_{P∈A} d_lP U_Pp^A,
1867 : !> d_lP = Σ_μν ϕ_μ(r_l) ϕ_ν(r_l) (μν|P),
1868 : !>
1869 : !> on the union of the atomic RI-RS grids needed by q. Each atom's three-center integrals are
1870 : !> evaluated once and transformed with U_Pp^A.
1871 : !> \param qs_env ...
1872 : !> \param bs_env ...
1873 : !> \param ctx ...
1874 : !> \param grid ...
1875 : !> \param mat_phi ...
1876 : !> \param mat_rhs ...
1877 : !> \param max_ao_size ...
1878 : ! **************************************************************************************************
1879 0 : SUBROUTINE compute_auto_ri_d_lp(qs_env, bs_env, ctx, grid, mat_phi, mat_rhs, max_ao_size)
1880 : TYPE(qs_environment_type), POINTER :: qs_env
1881 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
1882 : TYPE(gw_3c_ctx_type), INTENT(IN) :: ctx
1883 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: grid
1884 : TYPE(dbcsr_type), INTENT(IN) :: mat_phi
1885 : TYPE(dbcsr_type), INTENT(OUT) :: mat_rhs
1886 : INTEGER, INTENT(IN) :: max_ao_size
1887 :
1888 : CHARACTER(LEN=*), PARAMETER :: routineN = 'compute_auto_ri_d_lp'
1889 :
1890 : INTEGER :: AB_block, atom_A, atom_B, first, fit_atom, handle, handle_project, handle_rhs, l, &
1891 : last, n_done, n_first_p_ABs, n_to_a, n_total, n_union, natom, ncol, ngrid, npcol, nprow, &
1892 : ri_atom
1893 0 : INTEGER, ALLOCATABLE, DIMENSION(:) :: ao_atoms, ao_col_start, global_map, &
1894 0 : local_index, row_offset, union_index
1895 0 : INTEGER, DIMENSION(:), POINTER :: AB_row_dist, col_dist, &
1896 0 : first_p_ABs_per_atom, retained_size, &
1897 0 : row_dist, row_size
1898 0 : LOGICAL, ALLOCATABLE :: ao_nonzero(:, :)
1899 0 : LOGICAL, ALLOCATABLE, DIMENSION(:) :: needed_fit_atom, union_mask
1900 : REAL(KIND=dp) :: item_start_time, radius
1901 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: fit_radius
1902 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: column_map, phi, rhs, union_grid
1903 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: U_Pp
1904 : TYPE(cell_type), POINTER :: cell
1905 : TYPE(dbcsr_distribution_type) :: dist_AB, dist_phi, dist_t
1906 : TYPE(dbcsr_type) :: AB_d_lp, transform
1907 : TYPE(mp_para_env_type), POINTER :: para_env
1908 0 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
1909 0 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
1910 :
1911 0 : CALL timeset(routineN, handle)
1912 : CALL get_qs_env(qs_env, para_env=para_env, particle_set=particle_set, &
1913 0 : qs_kind_set=qs_kind_set, cell=cell)
1914 0 : natom = bs_env%n_atom
1915 0 : CALL dbcsr_get_info(mat_phi, row_blk_size=row_size, distribution=dist_phi)
1916 0 : CALL dbcsr_distribution_get(dist_phi, row_dist=row_dist, nprows=nprow, npcols=npcol)
1917 : ALLOCATE (first_p_ABs_per_atom(natom), col_dist(natom), AB_row_dist(natom), &
1918 0 : retained_size(natom), row_offset(SIZE(row_size)))
1919 0 : retained_size(:) = bs_env%auto_ri%sizes_opt_RI
1920 0 : DO ri_atom = 1, natom
1921 0 : first_p_ABs_per_atom(ri_atom) = 0
1922 0 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
1923 0 : IF (ri_atom == bs_env%auto_ri%AB_atom_A(AB_block) .OR. &
1924 0 : ri_atom == bs_env%auto_ri%AB_atom_B(AB_block)) THEN
1925 : first_p_ABs_per_atom(ri_atom) = &
1926 : first_p_ABs_per_atom(ri_atom) + &
1927 0 : bs_env%auto_ri%AB_size_opt_RI(AB_block)
1928 : END IF
1929 : END DO
1930 0 : col_dist(ri_atom) = MOD(ri_atom - 1, npcol)
1931 0 : AB_row_dist(ri_atom) = MOD(ri_atom - 1, nprow)
1932 : END DO
1933 0 : row_offset(1) = 0
1934 0 : DO l = 2, SIZE(row_size)
1935 0 : row_offset(l) = row_offset(l - 1) + row_size(l - 1)
1936 : END DO
1937 : CALL dbcsr_distribution_new(dist_AB, template=dist_phi, &
1938 0 : row_dist=row_dist, col_dist=col_dist)
1939 : CALL dbcsr_create(AB_d_lp, name='AUTO_RI AA/AB d_lp', dist=dist_AB, &
1940 : matrix_type=dbcsr_type_no_symmetry, row_blk_size=row_size, &
1941 0 : col_blk_size=first_p_ABs_per_atom)
1942 : CALL dbcsr_distribution_new(dist_t, template=dist_phi, &
1943 0 : row_dist=AB_row_dist, col_dist=col_dist)
1944 : CALL dbcsr_create(transform, name='AUTO_RI U_Pp', dist=dist_t, &
1945 : matrix_type=dbcsr_type_no_symmetry, row_blk_size=first_p_ABs_per_atom, &
1946 0 : col_blk_size=retained_size)
1947 : CALL dbcsr_create(mat_rhs, name='AUTO_RI optimized d_lp', dist=dist_AB, &
1948 : matrix_type=dbcsr_type_no_symmetry, row_blk_size=row_size, &
1949 0 : col_blk_size=retained_size)
1950 0 : ALLOCATE (fit_radius(natom))
1951 0 : CALL compute_auto_ri_grid_radii(bs_env, fit_radius)
1952 0 : ALLOCATE (needed_fit_atom(natom), union_mask(SIZE(grid, 2)))
1953 0 : n_total = 0
1954 0 : DO ri_atom = para_env%mepos + 1, natom, para_env%num_pe
1955 0 : n_total = n_total + 1
1956 : END DO
1957 0 : n_done = 0
1958 0 : CALL timeset(routineN//'_AB_d_lp', handle_rhs)
1959 0 : DO ri_atom = para_env%mepos + 1, natom, para_env%num_pe
1960 0 : item_start_time = m_walltime()
1961 0 : needed_fit_atom = .FALSE.
1962 0 : DO AB_block = 1, bs_env%auto_ri%AB_block_count
1963 0 : atom_A = bs_env%auto_ri%AB_atom_A(AB_block)
1964 0 : atom_B = bs_env%auto_ri%AB_atom_B(AB_block)
1965 0 : IF (ri_atom /= atom_A .AND. ri_atom /= atom_B) CYCLE
1966 0 : n_to_a = bs_env%auto_ri%AB_size_opt_RI_to_A(AB_block)
1967 0 : IF (n_to_a > 0) needed_fit_atom(atom_A) = .TRUE.
1968 0 : IF (n_to_a < bs_env%auto_ri%AB_size_opt_RI(AB_block)) THEN
1969 0 : needed_fit_atom(atom_B) = .TRUE.
1970 : END IF
1971 : END DO
1972 0 : IF (.NOT. ANY(needed_fit_atom)) THEN
1973 0 : n_done = n_done + 1
1974 : CALL print_Z_lP_progress(bs_env, n_done, n_total, &
1975 0 : m_walltime() - item_start_time)
1976 0 : CYCLE
1977 : END IF
1978 0 : union_mask = .FALSE.
1979 0 : radius = 0.0_dp
1980 0 : DO fit_atom = 1, natom
1981 0 : IF (.NOT. needed_fit_atom(fit_atom)) CYCLE
1982 : radius = MAX(radius, fit_radius(fit_atom) + &
1983 0 : NORM2(particle_set(fit_atom)%r - particle_set(ri_atom)%r))
1984 0 : DO l = 1, SIZE(grid, 2)
1985 0 : IF (NORM2(grid(:, l) - particle_set(fit_atom)%r) <= &
1986 0 : fit_radius(fit_atom)) union_mask(l) = .TRUE.
1987 : END DO
1988 : END DO
1989 0 : n_union = COUNT(union_mask)
1990 0 : ALLOCATE (union_index(n_union), union_grid(3, n_union))
1991 0 : n_union = 0
1992 0 : DO l = 1, SIZE(grid, 2)
1993 0 : IF (.NOT. union_mask(l)) CYCLE
1994 0 : n_union = n_union + 1
1995 0 : union_index(n_union) = l
1996 0 : union_grid(:, n_union) = grid(:, l)
1997 : END DO
1998 : CALL build_phi_on_sphere(bs_env, qs_kind_set, union_grid, ri_atom, &
1999 : radius + 1.0_dp, local_index, ngrid, phi, ao_atoms, ao_col_start, &
2000 0 : ao_nonzero=ao_nonzero)
2001 0 : local_index(1:ngrid) = union_index(local_index(1:ngrid))
2002 0 : n_first_p_ABs = first_p_ABs_per_atom(ri_atom)
2003 0 : CALL collect_auto_ri_columns_for_atom(bs_env, ri_atom, U_Pp, global_map)
2004 0 : CPASSERT(SIZE(global_map) == n_first_p_ABs)
2005 0 : first = 1
2006 0 : DO fit_atom = 1, natom
2007 0 : ncol = retained_size(fit_atom)
2008 0 : last = first + ncol - 1
2009 0 : IF (ANY(global_map >= first .AND. global_map <= last)) THEN
2010 0 : ALLOCATE (column_map(n_first_p_ABs, ncol), source=0.0_dp)
2011 0 : DO l = 1, n_first_p_ABs
2012 0 : IF (global_map(l) < first .OR. global_map(l) > last) CYCLE
2013 0 : column_map(l, global_map(l) - first + 1) = 1.0_dp
2014 : END DO
2015 0 : CALL dbcsr_put_block(transform, ri_atom, fit_atom, column_map)
2016 0 : DEALLOCATE (column_map)
2017 : END IF
2018 0 : first = last + 1
2019 : END DO
2020 0 : ALLOCATE (rhs(ngrid, n_first_p_ABs), source=0.0_dp)
2021 : CALL compute_d_lp_auto_ri(bs_env, ctx, phi, ao_atoms, ao_col_start, ao_nonzero, rhs, &
2022 0 : [ri_atom], U_Pp, max_ao_size, 0, 1)
2023 0 : DEALLOCATE (ao_atoms, ao_col_start, ao_nonzero)
2024 : CALL store_Z_lP_columns(AB_d_lp, rhs, local_index, ngrid, n_first_p_ABs, ri_atom, &
2025 0 : row_size, row_offset, 0.0_dp)
2026 0 : DEALLOCATE (rhs, U_Pp, global_map, phi, local_index, &
2027 0 : union_grid, union_index)
2028 0 : n_done = n_done + 1
2029 : CALL print_Z_lP_progress(bs_env, n_done, n_total, &
2030 0 : m_walltime() - item_start_time)
2031 : END DO
2032 0 : CALL dbcsr_finalize(transform)
2033 0 : CALL dbcsr_finalize(AB_d_lp)
2034 0 : CALL timestop(handle_rhs)
2035 0 : CALL timeset(routineN//'_transform', handle_project)
2036 : CALL dbcsr_multiply('N', 'N', 1.0_dp, AB_d_lp, transform, 0.0_dp, &
2037 0 : mat_rhs, filter_eps=0.0_dp)
2038 0 : CALL timestop(handle_project)
2039 0 : CALL dbcsr_release(AB_d_lp)
2040 0 : CALL dbcsr_release(transform)
2041 0 : CALL dbcsr_distribution_release(dist_AB)
2042 0 : CALL dbcsr_distribution_release(dist_t)
2043 0 : DEALLOCATE (first_p_ABs_per_atom, retained_size, col_dist, AB_row_dist, row_offset, &
2044 0 : fit_radius, needed_fit_atom, union_mask)
2045 0 : CALL timestop(handle)
2046 0 : END SUBROUTINE compute_auto_ri_d_lp
2047 :
2048 : ! **************************************************************************************************
2049 : !> \brief Extracts rhs(i,q) = d_{local_index(i),q} from one optimized atomic DBCSR column block.
2050 : !> \param mat_rhs ...
2051 : !> \param atom_index ...
2052 : !> \param local_index ...
2053 : !> \param row_offset ...
2054 : !> \param rhs ...
2055 : ! **************************************************************************************************
2056 0 : SUBROUTINE extract_atom_d_lp(mat_rhs, atom_index, local_index, row_offset, rhs)
2057 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_rhs
2058 : INTEGER, INTENT(IN) :: atom_index
2059 : INTEGER, DIMENSION(:), INTENT(IN) :: local_index, row_offset
2060 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT) :: rhs
2061 :
2062 : CHARACTER(LEN=*), PARAMETER :: routineN = 'extract_atom_d_lp'
2063 :
2064 : INTEGER :: handle, l, next_row, row
2065 : LOGICAL :: found
2066 0 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: block
2067 :
2068 0 : CALL timeset(routineN, handle)
2069 :
2070 0 : rhs = 0.0_dp
2071 0 : row = 0
2072 0 : next_row = 1
2073 0 : NULLIFY (block)
2074 0 : DO l = 1, SIZE(rhs, 1)
2075 0 : DO WHILE (next_row <= SIZE(row_offset))
2076 0 : IF (row_offset(next_row) >= local_index(l)) EXIT
2077 0 : row = next_row
2078 0 : next_row = next_row + 1
2079 0 : CALL dbcsr_get_block_p(mat_rhs, row, atom_index, block, found)
2080 0 : IF (.NOT. found) NULLIFY (block)
2081 : END DO
2082 0 : IF (ASSOCIATED(block)) rhs(l, :) = block(local_index(l) - row_offset(row), :)
2083 : END DO
2084 :
2085 0 : CALL timestop(handle)
2086 0 : END SUBROUTINE extract_atom_d_lp
2087 :
2088 : ! **************************************************************************************************
2089 : !> \brief Solves Σ_l' D_ll' Z_l'q = d_lp for each optimized atomic column block. Small grids are
2090 : !> gathered within one process column and solved by a local Cholesky factorization.
2091 : !> \param qs_env ...
2092 : !> \param bs_env ...
2093 : !> \param grid ...
2094 : !> \param mat_rhs ...
2095 : !> \param mat_z ...
2096 : ! **************************************************************************************************
2097 0 : SUBROUTINE fit_auto_ri_z_lp(qs_env, bs_env, grid, mat_rhs, mat_z)
2098 : TYPE(qs_environment_type), POINTER :: qs_env
2099 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
2100 : REAL(KIND=dp), INTENT(IN) :: grid(:, :)
2101 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_rhs, mat_z
2102 :
2103 : CHARACTER(LEN=*), PARAMETER :: routineN = 'fit_auto_ri_z_lp'
2104 :
2105 : INTEGER :: atom_index, base, group_size, handle, &
2106 : handle_gather, i, mypcol, n_big, &
2107 : n_small, natom, ncol, ngrid, npcol, &
2108 : slot
2109 0 : INTEGER, ALLOCATABLE :: all_index(:), ao_atoms(:), ao_col_start(:), big_list(:), &
2110 0 : grid_size(:), local_index(:), row_offset(:), small_list(:)
2111 0 : INTEGER, POINTER :: row_size(:)
2112 0 : LOGICAL, ALLOCATABLE :: single_rank(:)
2113 0 : REAL(KIND=dp), ALLOCATABLE :: atom_rhs(:, :), buffer(:, :), &
2114 0 : local_rhs(:, :), phi(:, :), radius(:)
2115 : TYPE(dbcsr_distribution_type) :: distribution
2116 : TYPE(mp_para_env_type), POINTER :: column_env, para_env
2117 0 : TYPE(qs_kind_type), POINTER :: qs_kinds(:)
2118 :
2119 0 : CALL timeset(routineN, handle)
2120 0 : CALL get_qs_env(qs_env, para_env=para_env, qs_kind_set=qs_kinds)
2121 0 : CALL dbcsr_get_info(mat_rhs, distribution=distribution, row_blk_size=row_size)
2122 0 : CALL dbcsr_distribution_get(distribution, mypcol=mypcol, npcols=npcol)
2123 0 : ALLOCATE (column_env)
2124 0 : CALL column_env%from_split(para_env, mypcol)
2125 0 : natom = bs_env%n_atom
2126 0 : ALLOCATE (radius(natom), all_index(SIZE(grid, 2)), row_offset(SIZE(row_size)))
2127 0 : CALL compute_auto_ri_grid_radii(bs_env, radius)
2128 : CALL classify_z_lp_atoms(bs_env, grid, radius, &
2129 : bs_env%auto_ri%sizes_opt_RI, grid_size, &
2130 0 : small_list, n_small, big_list, n_big, group_size)
2131 0 : ALLOCATE (single_rank(natom), source=.FALSE.)
2132 0 : single_rank(small_list(:n_small)) = .TRUE.
2133 0 : DO i = 1, SIZE(all_index)
2134 0 : all_index(i) = i
2135 : END DO
2136 0 : row_offset(1) = 0
2137 0 : DO i = 2, SIZE(row_size)
2138 0 : row_offset(i) = row_offset(i - 1) + row_size(i - 1)
2139 : END DO
2140 0 : DO base = mypcol + 1, natom, npcol*column_env%num_pe
2141 0 : CALL timeset(routineN//'_gather_rhs', handle_gather)
2142 0 : DO slot = 0, column_env%num_pe - 1
2143 0 : atom_index = base + slot*npcol
2144 0 : IF (atom_index > natom) EXIT
2145 0 : IF (.NOT. single_rank(atom_index)) CYCLE
2146 0 : ncol = bs_env%auto_ri%sizes_opt_RI(atom_index)
2147 0 : IF (ncol == 0) CYCLE
2148 0 : ALLOCATE (buffer(SIZE(grid, 2), ncol))
2149 0 : CALL extract_atom_d_lp(mat_rhs, atom_index, all_index, row_offset, buffer)
2150 0 : CALL column_env%sum(buffer, slot)
2151 0 : IF (column_env%mepos == slot) THEN
2152 0 : CALL MOVE_ALLOC(buffer, atom_rhs)
2153 : ELSE
2154 0 : DEALLOCATE (buffer)
2155 : END IF
2156 : END DO
2157 0 : CALL timestop(handle_gather)
2158 0 : atom_index = base + column_env%mepos*npcol
2159 0 : IF (atom_index > natom) CYCLE
2160 0 : IF (.NOT. single_rank(atom_index)) CYCLE
2161 0 : ncol = bs_env%auto_ri%sizes_opt_RI(atom_index)
2162 0 : IF (ncol == 0) CYCLE
2163 : CALL build_phi_on_sphere(bs_env, qs_kinds, grid, atom_index, &
2164 0 : radius(atom_index), local_index, ngrid, phi, ao_atoms, ao_col_start)
2165 0 : DEALLOCATE (ao_atoms, ao_col_start)
2166 0 : ALLOCATE (local_rhs(ngrid, ncol))
2167 0 : DO i = 1, ngrid
2168 0 : local_rhs(i, :) = atom_rhs(local_index(i), :)
2169 : END DO
2170 0 : DEALLOCATE (atom_rhs)
2171 0 : CALL solve_ri_rs_fit(phi, local_rhs, bs_env%ri_rs%tikhonov)
2172 : CALL store_Z_lP_columns(mat_z, local_rhs, local_index, ngrid, ncol, atom_index, &
2173 0 : row_size, row_offset, 0.0_dp)
2174 0 : DEALLOCATE (local_rhs, phi, local_index)
2175 : END DO
2176 0 : IF (n_big > 0) THEN
2177 : CALL fit_distributed_auto_ri_z_lp(qs_env, bs_env, grid, mat_rhs, mat_z, radius, &
2178 : big_list(:n_big), group_size, row_size, row_offset, &
2179 0 : all_index)
2180 : END IF
2181 0 : DEALLOCATE (radius, all_index, row_offset, single_rank, grid_size, small_list, big_list)
2182 0 : CALL column_env%free()
2183 0 : DEALLOCATE (column_env)
2184 0 : CALL dbcsr_finalize(mat_z)
2185 0 : CALL timestop(handle)
2186 0 : END SUBROUTINE fit_auto_ri_z_lp
2187 :
2188 : ! **************************************************************************************************
2189 : !> \brief Solves Σ_l' D_ll' Z_l'q = d_lp for large atomic grids with distributed Cholesky
2190 : !> factorization. Each d_lp block is gathered once and distributed over its assigned rank
2191 : !> group.
2192 : !> \param qs_env ...
2193 : !> \param bs_env ...
2194 : !> \param grid ...
2195 : !> \param mat_rhs ...
2196 : !> \param mat_z ...
2197 : !> \param radius ...
2198 : !> \param atom_list ...
2199 : !> \param group_size ...
2200 : !> \param row_size ...
2201 : !> \param row_offset ...
2202 : !> \param all_index ...
2203 : ! **************************************************************************************************
2204 0 : SUBROUTINE fit_distributed_auto_ri_z_lp(qs_env, bs_env, grid, mat_rhs, mat_z, radius, &
2205 0 : atom_list, group_size, row_size, row_offset, &
2206 0 : all_index)
2207 : TYPE(qs_environment_type), POINTER :: qs_env
2208 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
2209 : REAL(KIND=dp), INTENT(IN) :: grid(:, :)
2210 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_rhs, mat_z
2211 : REAL(KIND=dp), INTENT(IN) :: radius(:)
2212 : INTEGER, INTENT(IN) :: atom_list(:), group_size, row_size(:), &
2213 : row_offset(:), all_index(:)
2214 :
2215 : CHARACTER(LEN=*), PARAMETER :: routineN = 'fit_distributed_auto_ri_z_lp'
2216 :
2217 : INTEGER :: atom_index, base, handle, i, my_group, &
2218 : ncol, ngrid, ngroups, root, slot
2219 0 : INTEGER, ALLOCATABLE :: ao_atoms(:), ao_col_start(:), &
2220 0 : local_index(:)
2221 0 : REAL(KIND=dp), ALLOCATABLE :: atom_rhs(:, :), buffer(:, :), &
2222 0 : local_rhs(:, :), phi(:, :)
2223 : TYPE(cp_blacs_env_type), POINTER :: blacs_env
2224 : TYPE(mp_para_env_type), POINTER :: group_env, para_env
2225 0 : TYPE(qs_kind_type), POINTER :: qs_kinds(:)
2226 :
2227 0 : CALL timeset(routineN, handle)
2228 0 : CALL get_qs_env(qs_env, para_env=para_env, qs_kind_set=qs_kinds)
2229 0 : ngroups = MAX(1, para_env%num_pe/group_size)
2230 0 : my_group = MIN(para_env%mepos/group_size, ngroups - 1)
2231 0 : ALLOCATE (group_env)
2232 0 : CALL group_env%from_split(para_env, my_group)
2233 0 : NULLIFY (blacs_env)
2234 0 : CALL cp_blacs_env_create(blacs_env=blacs_env, para_env=group_env)
2235 0 : DO base = 1, SIZE(atom_list), ngroups
2236 0 : DO slot = 0, ngroups - 1
2237 0 : IF (base + slot > SIZE(atom_list)) EXIT
2238 0 : atom_index = atom_list(base + slot)
2239 0 : ncol = bs_env%auto_ri%sizes_opt_RI(atom_index)
2240 0 : IF (ncol == 0) CYCLE
2241 0 : root = slot*group_size
2242 0 : ALLOCATE (buffer(SIZE(grid, 2), ncol))
2243 0 : CALL extract_atom_d_lp(mat_rhs, atom_index, all_index, row_offset, buffer)
2244 0 : CALL para_env%sum(buffer, root)
2245 0 : IF (para_env%mepos == root) THEN
2246 0 : CALL MOVE_ALLOC(buffer, atom_rhs)
2247 : ELSE
2248 0 : DEALLOCATE (buffer)
2249 : END IF
2250 : END DO
2251 0 : IF (base + my_group > SIZE(atom_list)) CYCLE
2252 0 : atom_index = atom_list(base + my_group)
2253 0 : ncol = bs_env%auto_ri%sizes_opt_RI(atom_index)
2254 0 : IF (ncol == 0) CYCLE
2255 0 : IF (group_env%mepos /= 0) ALLOCATE (atom_rhs(SIZE(grid, 2), ncol))
2256 0 : CALL group_env%bcast(atom_rhs, 0)
2257 : CALL build_phi_on_sphere(bs_env, qs_kinds, grid, atom_index, &
2258 0 : radius(atom_index), local_index, ngrid, phi, ao_atoms, ao_col_start)
2259 0 : DEALLOCATE (ao_atoms, ao_col_start)
2260 0 : ALLOCATE (local_rhs(ngrid, ncol))
2261 0 : DO i = 1, ngrid
2262 0 : local_rhs(i, :) = atom_rhs(local_index(i), :)
2263 : END DO
2264 0 : DEALLOCATE (atom_rhs)
2265 0 : CALL solve_ri_rs_fit(phi, local_rhs, bs_env%ri_rs%tikhonov, group_env, blacs_env)
2266 0 : IF (group_env%mepos == 0) THEN
2267 : CALL store_Z_lP_columns(mat_z, local_rhs, local_index, ngrid, ncol, atom_index, &
2268 0 : row_size, row_offset, 0.0_dp)
2269 : END IF
2270 0 : DEALLOCATE (local_rhs, phi, local_index)
2271 : END DO
2272 0 : CALL cp_blacs_env_release(blacs_env)
2273 0 : CALL group_env%free()
2274 0 : DEALLOCATE (group_env)
2275 0 : CALL timestop(handle)
2276 0 : END SUBROUTINE fit_distributed_auto_ri_z_lp
2277 :
2278 : ! **************************************************************************************************
2279 : !> \brief Splits the atoms of the Z_lP solve into a single-rank list ("small", Phase A: LAPACK
2280 : !> dpotrf/dpotrs on one rank) and a distributed list ("big", Phase B: ScaLAPACK
2281 : !> pdpotrf/pdpotrs over a rank subgroup of size G), and sizes G.
2282 : !> Each atom's single-rank peak memory is estimated as
2283 : !> peak(P) = 8*n_local_grid(P)^2 (dense conditioned Gram matrix)
2284 : !> + 8*n_local_grid(P)*n_ao_used(P) (phi_local)
2285 : !> + 8*n_local_grid(P)*n_RI(P)*(1+n_threads) (d_lp + OMP partials)
2286 : !> and send atoms whose peak exceeds mem_safety * available-memory-per-proc to the
2287 : !> distributed path; G is auto-sized so the biggest atom's distributed Gram matrix slice
2288 : !> fits alongside the replicated phi_local + d_lp. G is floored by the ScaLAPACK 32-bit
2289 : !> index limit (a local block-cyclic slice of ~n_local_grid^2/G elements must stay below
2290 : !> 2^31 or pdpotrf segfaults).
2291 : !> Records and reports the memory estimates before the standard or AUTO_RI solves.
2292 : !> \param bs_env ...
2293 : !> \param ri_rs_grid_points ...
2294 : !> \param cutoff_ri_per_atom ...
2295 : !> \param ri_blk_sizes per-atom ...
2296 : !> \param n_local_grid_atom ...
2297 : !> \param small_list ...
2298 : !> \param n_small ...
2299 : !> \param big_list ...
2300 : !> \param n_big ...
2301 : !> \param G ...
2302 : !> \param periodic_images use minimum-image distances for the Gamma integration domain
2303 : ! **************************************************************************************************
2304 42 : SUBROUTINE classify_z_lp_atoms(bs_env, ri_rs_grid_points, cutoff_ri_per_atom, ri_blk_sizes, &
2305 : n_local_grid_atom, small_list, n_small, big_list, n_big, G, periodic_images)
2306 :
2307 : !$ USE OMP_LIB, ONLY: omp_get_max_threads
2308 :
2309 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
2310 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: ri_rs_grid_points
2311 :
2312 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: cutoff_ri_per_atom
2313 : INTEGER, DIMENSION(:), INTENT(IN) :: ri_blk_sizes
2314 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: n_local_grid_atom, small_list, big_list
2315 : INTEGER, INTENT(OUT) :: n_small, n_big, G
2316 : LOGICAL, INTENT(IN), OPTIONAL :: periodic_images
2317 : CHARACTER(LEN=*), PARAMETER :: routineN = 'classify_z_lp_atoms'
2318 :
2319 : INTEGER :: handle
2320 :
2321 : ! Conservative fraction of measured available memory usable per rank for the Z_lP
2322 : REAL(KIND=dp), PARAMETER :: mem_safety = 0.8_dp
2323 :
2324 : ! ScaLAPACK/BLACS index the per-rank local block-cyclic slice (~n_local_grid^2/G
2325 : ! elements) with 32-bit integers; keep it safely below 2^31 or pdpotrf segfaults.
2326 : REAL(KIND=dp), PARAMETER :: scalapack_loc_limit = 2.0E9_dp
2327 :
2328 : INTEGER :: G_atom, G_int32, l, n_ao_used_atom, &
2329 : n_grid_total, n_local_grid, natom, &
2330 : nthreads_cls, P_loop_atom
2331 : REAL(KIND=dp) :: budget_bytes, cutoff_ri, dlp_bytes, &
2332 : mem_avail_GB, ng, nri, peak_bytes, &
2333 : phi_bytes
2334 : REAL(KIND=dp), DIMENSION(3) :: pos_P, displacement
2335 : LOGICAL :: sum_images
2336 : TYPE(mp_para_env_type), POINTER :: para_env
2337 42 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
2338 :
2339 42 : CALL timeset(routineN, handle)
2340 42 : sum_images = .FALSE.
2341 42 : IF (PRESENT(periodic_images)) sum_images = periodic_images
2342 :
2343 42 : para_env => bs_env%para_env
2344 42 : particle_set => bs_env%ri_rs%particle_set
2345 42 : natom = bs_env%n_atom
2346 42 : n_grid_total = bs_env%ri_rs%n_grid_points
2347 42 : CPASSERT(SIZE(ri_rs_grid_points, 2) == n_grid_total)
2348 :
2349 : ! Per-atom sphere size: n_local_grid(P) = number of l with |r_l - R_P| <= cutoff_ri(P).
2350 : ! It sets both the Gram matrix memory footprint (n_local_grid^2) and the solve
2351 : ! cost (~n_local_grid^3), so it drives classification and the LPT load balancing.
2352 210 : ALLOCATE (n_local_grid_atom(natom), small_list(natom), big_list(natom))
2353 142 : DO P_loop_atom = 1, natom
2354 400 : pos_P(:) = particle_set(P_loop_atom)%r(:)
2355 100 : cutoff_ri = cutoff_ri_per_atom(P_loop_atom)
2356 100 : n_local_grid = 0
2357 38088 : DO l = 1, n_grid_total
2358 151952 : displacement = ri_rs_grid_points(:, l) - pos_P
2359 37988 : IF (sum_images) displacement = pbc(ri_rs_grid_points(:, l), pos_P, bs_env%ri_rs%cell)
2360 152052 : IF (NORM2(displacement) <= cutoff_ri) THEN
2361 37850 : n_local_grid = n_local_grid + 1
2362 : END IF
2363 : END DO
2364 142 : n_local_grid_atom(P_loop_atom) = n_local_grid
2365 : END DO
2366 :
2367 42 : nthreads_cls = 1
2368 42 : !$ nthreads_cls = omp_get_max_threads()
2369 42 : CALL mp_mem_avail_per_rank_GB(bs_env%para_env, mem_avail_GB) ! collective over all ranks
2370 42 : budget_bytes = mem_safety*mem_avail_GB*1.0E9_dp
2371 :
2372 42 : n_small = 0
2373 42 : n_big = 0
2374 42 : G = 1
2375 42 : G_atom = 1 ! max G a big atom needs (memory + ScaLAPACK int32 floor)
2376 42 : IF (mem_avail_GB <= 0.0_dp .AND. bs_env%unit_nr > 0) THEN
2377 0 : CPWARN("RI-RS Z_lP: no meminfo; single-rank solve for all atoms")
2378 : END IF
2379 142 : DO P_loop_atom = 1, natom
2380 100 : ng = REAL(n_local_grid_atom(P_loop_atom), dp)
2381 100 : nri = REAL(ri_blk_sizes(P_loop_atom), dp)
2382 : CALL get_n_ao_in_sphere(bs_env, P_loop_atom, &
2383 100 : cutoff_ri_per_atom(P_loop_atom), n_ao_used_atom, sum_images)
2384 : CALL record_fit_memory(bs_env%ri_rs, n_local_grid_atom(P_loop_atom), &
2385 100 : n_ao_used_atom, ri_blk_sizes(P_loop_atom), peak_bytes)
2386 142 : IF (mem_avail_GB <= 0.0_dp .OR. peak_bytes <= budget_bytes) THEN
2387 100 : n_small = n_small + 1
2388 100 : small_list(n_small) = P_loop_atom
2389 : ELSE
2390 0 : n_big = n_big + 1
2391 0 : big_list(n_big) = P_loop_atom
2392 0 : phi_bytes = 8.0_dp*ng*REAL(n_ao_used_atom, dp)
2393 0 : dlp_bytes = 8.0_dp*ng*nri*REAL(1 + nthreads_cls, dp)
2394 : ! Bound both memory and ScaLAPACK's 32-bit local indexing.
2395 0 : G_int32 = CEILING(ng*ng/scalapack_loc_limit)
2396 : G_atom = MAX(G_atom, G_int32, &
2397 0 : CEILING(8.0_dp*ng*ng/MAX(budget_bytes - phi_bytes - dlp_bytes, 1.0_dp)))
2398 : END IF
2399 : END DO
2400 42 : CALL print_fit_memory_estimate(bs_env, mem_avail_GB)
2401 42 : IF (n_big > 0) THEN
2402 0 : IF (G_atom > para_env%num_pe) THEN
2403 : CALL cp_abort(__LOCATION__, &
2404 : "RI-RS Z_lP: an atom is too large to fit even when "// &
2405 : "distributed over all ranks. Add nodes, use fewer MPI ranks "// &
2406 : "per node, lower CUTOFF_RADIUS_RI, or raise EPS_FILTER "// &
2407 0 : "for more grid screening.")
2408 : END IF
2409 0 : G = MIN(MAX(G_atom, 2), para_env%num_pe)
2410 : END IF
2411 :
2412 42 : CALL timestop(handle)
2413 :
2414 42 : END SUBROUTINE classify_z_lp_atoms
2415 :
2416 : ! **************************************************************************************************
2417 : !> \brief Number of AO basis functions that can be non-zero inside the RI-RS integration sphere
2418 : !> \param bs_env ...
2419 : !> \param atom_P ...
2420 : !> \param cutoff_ri ...
2421 : !> \param n_ao_used ...
2422 : !> \param periodic_images use minimum-image distances to the Gamma integration sphere
2423 : ! **************************************************************************************************
2424 100 : SUBROUTINE get_n_ao_in_sphere(bs_env, atom_P, cutoff_ri, n_ao_used, periodic_images)
2425 :
2426 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
2427 : INTEGER, INTENT(IN) :: atom_P
2428 : REAL(KIND=dp), INTENT(IN) :: cutoff_ri
2429 : INTEGER, INTENT(OUT) :: n_ao_used
2430 : LOGICAL, INTENT(IN) :: periodic_images
2431 :
2432 : CHARACTER(LEN=*), PARAMETER :: routineN = 'get_n_ao_in_sphere'
2433 :
2434 : INTEGER :: handle, ri_atom
2435 : REAL(KIND=dp) :: displacement(3)
2436 100 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
2437 :
2438 100 : CALL timeset(routineN, handle)
2439 :
2440 100 : particle_set => bs_env%ri_rs%particle_set
2441 100 : n_ao_used = 0
2442 348 : DO ri_atom = 1, bs_env%n_atom
2443 992 : displacement = particle_set(ri_atom)%r - particle_set(atom_P)%r
2444 248 : IF (periodic_images) displacement = pbc(particle_set(ri_atom)%r, particle_set(atom_P)%r, bs_env%ri_rs%cell)
2445 992 : IF (NORM2(displacement) > &
2446 : bs_env%ri_rs%radius_ao_per_atom(ri_atom) + cutoff_ri) CYCLE
2447 : n_ao_used = n_ao_used + bs_env%i_ao_end_from_atom(ri_atom) - &
2448 348 : bs_env%i_ao_start_from_atom(ri_atom) + 1
2449 : END DO
2450 :
2451 100 : CALL timestop(handle)
2452 :
2453 100 : END SUBROUTINE get_n_ao_in_sphere
2454 :
2455 : ! **************************************************************************************************
2456 : !> \brief Builds the sphere-local AO matrix phi_local(l, μ) = ϕ_μ(r_l) for one RI atom P
2457 : !> \param bs_env ...
2458 : !> \param qs_kind_set ...
2459 : !> \param ri_rs_grid_points ...
2460 : !> \param atom_P ...
2461 : !> \param cutoff_ri ...
2462 : !> \param local_grid_idx ...
2463 : !> \param n_local_grid ...
2464 : !> \param phi_local ...
2465 : !> \param ao_atoms global indices of the active AO atoms
2466 : !> \param ao_col_start local phi column starts, including the terminal column offset
2467 : !> \param center ...
2468 : !> \param periodic_images select Gamma image-summed AO values and minimum-image sphere distances
2469 : !> \param ao_nonzero exact nonzero AO support on the retained grid rows
2470 : ! **************************************************************************************************
2471 54 : SUBROUTINE build_phi_on_sphere(bs_env, qs_kind_set, ri_rs_grid_points, &
2472 : atom_P, cutoff_ri, local_grid_idx, n_local_grid, &
2473 : phi_local, ao_atoms, ao_col_start, center, periodic_images, ao_nonzero)
2474 :
2475 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
2476 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
2477 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: ri_rs_grid_points
2478 : INTEGER, INTENT(IN) :: atom_P
2479 : REAL(KIND=dp), INTENT(IN) :: cutoff_ri
2480 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: local_grid_idx
2481 : INTEGER, INTENT(OUT) :: n_local_grid
2482 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
2483 : INTENT(OUT) :: phi_local
2484 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: ao_atoms, ao_col_start
2485 : REAL(KIND=dp), DIMENSION(3), INTENT(IN), OPTIONAL :: center
2486 : LOGICAL, INTENT(IN), OPTIONAL :: periodic_images
2487 : LOGICAL, ALLOCATABLE, INTENT(OUT), OPTIONAL :: ao_nonzero(:, :)
2488 :
2489 : CHARACTER(LEN=*), PARAMETER :: routineN = 'build_phi_on_sphere'
2490 :
2491 : INTEGER :: col_end, col_start, handle, iactive, j, &
2492 : k, l, loc_idx, n_active, n_ao_used, &
2493 : n_grid_total, n_keep, ri_atom
2494 54 : INTEGER, ALLOCATABLE :: atom_candidates(:)
2495 : LOGICAL :: sum_images
2496 : REAL(KIND=dp) :: d_sP, dist, r2_threshold
2497 54 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: w_pt
2498 54 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: phi_keep, sphere_grid
2499 : REAL(KIND=dp), DIMENSION(3) :: displacement, pos_P
2500 : TYPE(cell_type), POINTER :: cell
2501 54 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
2502 :
2503 54 : CALL timeset(routineN, handle)
2504 54 : sum_images = .FALSE.
2505 54 : IF (PRESENT(periodic_images)) sum_images = periodic_images
2506 :
2507 54 : cell => bs_env%ri_rs%cell
2508 54 : particle_set => bs_env%ri_rs%particle_set
2509 : ! AUTO_RI may pass only the global grid points needed for the current atom block.
2510 54 : n_grid_total = SIZE(ri_rs_grid_points, 2)
2511 54 : IF (PRESENT(center)) THEN
2512 4 : pos_P(:) = center(:)
2513 : ELSE
2514 200 : pos_P(:) = particle_set(atom_P)%r(:)
2515 : END IF
2516 :
2517 54 : n_local_grid = 0
2518 19174 : DO l = 1, n_grid_total
2519 76480 : displacement = ri_rs_grid_points(:, l) - pos_P
2520 19120 : IF (sum_images) displacement = pbc(ri_rs_grid_points(:, l), pos_P, cell)
2521 76480 : dist = NORM2(displacement)
2522 19174 : IF (dist <= cutoff_ri) n_local_grid = n_local_grid + 1
2523 : END DO
2524 :
2525 162 : ALLOCATE (local_grid_idx(n_local_grid))
2526 :
2527 54 : n_local_grid = 0
2528 19174 : DO l = 1, n_grid_total
2529 76480 : displacement = ri_rs_grid_points(:, l) - pos_P
2530 19120 : IF (sum_images) displacement = pbc(ri_rs_grid_points(:, l), pos_P, cell)
2531 76480 : dist = NORM2(displacement)
2532 19174 : IF (dist <= cutoff_ri) THEN
2533 19051 : n_local_grid = n_local_grid + 1
2534 19051 : local_grid_idx(n_local_grid) = l
2535 : END IF
2536 : END DO
2537 :
2538 162 : ALLOCATE (sphere_grid(3, n_local_grid))
2539 19105 : DO loc_idx = 1, n_local_grid
2540 76258 : sphere_grid(:, loc_idx) = ri_rs_grid_points(:, local_grid_idx(loc_idx))
2541 : END DO
2542 :
2543 : ! Keep atoms in global order; phi columns are contiguous within each active atom.
2544 162 : ALLOCATE (atom_candidates(bs_env%n_atom))
2545 54 : n_active = 0
2546 190 : DO ri_atom = 1, bs_env%n_atom
2547 544 : displacement = particle_set(ri_atom)%r - pos_P
2548 136 : IF (sum_images) displacement = pbc(particle_set(ri_atom)%r, pos_P, cell)
2549 544 : d_sP = NORM2(displacement)
2550 136 : IF (d_sP > bs_env%ri_rs%radius_ao_per_atom(ri_atom) + cutoff_ri) CYCLE
2551 136 : IF (bs_env%i_ao_end_from_atom(ri_atom) < bs_env%i_ao_start_from_atom(ri_atom)) CYCLE
2552 136 : n_active = n_active + 1
2553 190 : atom_candidates(n_active) = ri_atom
2554 : END DO
2555 298 : ALLOCATE (ao_atoms(n_active), SOURCE=atom_candidates(:n_active))
2556 54 : DEALLOCATE (atom_candidates)
2557 162 : ALLOCATE (ao_col_start(n_active + 1))
2558 54 : ao_col_start(1) = 1
2559 190 : DO iactive = 1, n_active
2560 136 : ri_atom = ao_atoms(iactive)
2561 : ao_col_start(iactive + 1) = ao_col_start(iactive) + &
2562 190 : bs_env%i_ao_end_from_atom(ri_atom) - bs_env%i_ao_start_from_atom(ri_atom) + 1
2563 : END DO
2564 54 : n_ao_used = ao_col_start(n_active + 1) - 1
2565 216 : ALLOCATE (phi_local(n_local_grid, n_ao_used), SOURCE=0.0_dp)
2566 :
2567 190 : DO iactive = 1, n_active
2568 136 : ri_atom = ao_atoms(iactive)
2569 136 : col_start = ao_col_start(iactive)
2570 136 : col_end = ao_col_start(iactive + 1) - 1
2571 : ! A positive CUTOFF_RADIUS_RI_AO overrides the per-atom Gaussian radius
2572 : ! with a user-defined hard cutoff.
2573 136 : IF (.NOT. sum_images .AND. bs_env%ri_rs%cutoff_radius_ri_ao > 0.0_dp) THEN
2574 9 : r2_threshold = bs_env%ri_rs%cutoff_radius_ri_ao**2
2575 : ELSE
2576 127 : r2_threshold = bs_env%ri_rs%radius_ao_per_atom(ri_atom)**2
2577 : END IF
2578 :
2579 190 : IF (sum_images) THEN
2580 : CALL evaluate_ao_basis_on_points(phi_local(:, col_start:col_end), sphere_grid, &
2581 : bs_env%basis_set_AO( &
2582 : particle_set(ri_atom)%atomic_kind%kind_number)%gto_basis_set, &
2583 : particle_set(ri_atom)%r, cell, &
2584 0 : cutoff_squared=r2_threshold, periodic_images=.TRUE.)
2585 : ELSE
2586 : CALL evaluate_ao_on_points(phi_local(:, col_start:col_end), sphere_grid, &
2587 : ri_atom, particle_set, qs_kind_set, cell, &
2588 136 : cutoff_squared=r2_threshold)
2589 : END IF
2590 : END DO
2591 :
2592 54 : DEALLOCATE (sphere_grid)
2593 :
2594 54 : IF (.NOT. sum_images .AND. n_local_grid > 0) THEN
2595 162 : ALLOCATE (w_pt(n_local_grid))
2596 : !$OMP PARALLEL DO DEFAULT(NONE) &
2597 : !$OMP SHARED(n_local_grid, n_ao_used, phi_local, w_pt) &
2598 54 : !$OMP PRIVATE(l, j) SCHEDULE(STATIC)
2599 : DO l = 1, n_local_grid
2600 : w_pt(l) = 0.0_dp
2601 : DO j = 1, n_ao_used
2602 : w_pt(l) = MAX(w_pt(l), ABS(phi_local(l, j)))
2603 : END DO
2604 : END DO
2605 : !$OMP END PARALLEL DO
2606 19105 : n_keep = COUNT(w_pt > bs_env%eps_filter)
2607 54 : IF (n_keep < n_local_grid) THEN
2608 24 : ALLOCATE (phi_keep(n_keep, n_ao_used))
2609 6 : k = 0
2610 2601 : DO l = 1, n_local_grid
2611 2601 : IF (w_pt(l) > bs_env%eps_filter) THEN
2612 2151 : k = k + 1
2613 17208 : phi_keep(k, :) = phi_local(l, :)
2614 2151 : local_grid_idx(k) = local_grid_idx(l)
2615 : END IF
2616 : END DO
2617 6 : CALL MOVE_ALLOC(phi_keep, phi_local)
2618 6 : n_local_grid = n_keep
2619 : END IF
2620 54 : DEALLOCATE (w_pt)
2621 : END IF
2622 :
2623 54 : IF (PRESENT(ao_nonzero)) THEN
2624 : ! Build support on the final retained rows, once for all RHS contractions using phi.
2625 216 : ALLOCATE (ao_nonzero(n_local_grid, n_active), SOURCE=.FALSE.)
2626 : !$OMP PARALLEL DO DEFAULT(NONE) SCHEDULE(STATIC) &
2627 : !$OMP SHARED(n_active, n_local_grid, ao_col_start, phi_local, ao_nonzero) &
2628 54 : !$OMP PRIVATE(iactive, j, l)
2629 : DO iactive = 1, n_active
2630 : DO j = ao_col_start(iactive), ao_col_start(iactive + 1) - 1
2631 : DO l = 1, n_local_grid
2632 : ao_nonzero(l, iactive) = ao_nonzero(l, iactive) .OR. phi_local(l, j) /= 0.0_dp
2633 : END DO
2634 : END DO
2635 : END DO
2636 : !$OMP END PARALLEL DO
2637 : END IF
2638 :
2639 54 : CALL timestop(handle)
2640 :
2641 108 : END SUBROUTINE build_phi_on_sphere
2642 :
2643 : ! **************************************************************************************************
2644 : !> \brief Computes a three-center integral block directly in AUTO_RI contraction columns.
2645 : !>
2646 : !> For each angular momentum l, the primitive RI Gaussians on the atoms contributing
2647 : !> to one optimized function are collected before the contraction
2648 : !>
2649 : !> (μν|p) = Σ_P (μν|P) U_Pp .
2650 : !>
2651 : !> This avoids constructing and retaining the complete reference-basis block (μν|P).
2652 : !> \param int_3c ...
2653 : !> \param ctx ...
2654 : !> \param ws ...
2655 : !> \param atom_j ...
2656 : !> \param atom_k ...
2657 : !> \param atom_i ...
2658 : !> \param transform ...
2659 : !> \param transform_row ...
2660 : !> \param screened ...
2661 : ! **************************************************************************************************
2662 108 : SUBROUTINE build_3c_integral_block_auto_ri_ctx(int_3c, ctx, ws, atom_j, atom_k, atom_i, &
2663 36 : transform, transform_row, screened)
2664 : REAL(KIND=dp), DIMENSION(:, :, :), INTENT(INOUT) :: int_3c
2665 : TYPE(gw_3c_ctx_type), INTENT(IN) :: ctx
2666 : TYPE(gw_3c_ws_type), INTENT(INOUT) :: ws
2667 : INTEGER, INTENT(IN) :: atom_j, atom_k, atom_i
2668 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: transform
2669 : INTEGER, INTENT(IN) :: transform_row
2670 : LOGICAL, INTENT(OUT), OPTIONAL :: screened
2671 :
2672 : CHARACTER(LEN=*), PARAMETER :: routineN = 'build_3c_integral_block_auto_ri_ctx'
2673 :
2674 : INTEGER :: handle_contract, handle_eri, ikind, iset, jkind, jset, kkind, kset, l, ncoi, &
2675 : ncoj, ncok, ncol, npgf_group, nseti, nsetj, nsetk, primitive_first, sgfi, sgfj, sgfk
2676 36 : INTEGER, DIMENSION(:), POINTER :: lmax_i, lmax_j, lmax_k, lmin_i, lmin_j, &
2677 36 : lmin_k, npgfi, npgfj, npgfk, nsgfi, &
2678 36 : nsgfj, nsgfk
2679 36 : INTEGER, DIMENSION(:, :), POINTER :: first_sgf_i, first_sgf_j, first_sgf_k
2680 : REAL(KIND=dp) :: dij, dik, djk, group_radius, &
2681 : kind_radius_i, kind_radius_j, &
2682 : kind_radius_k, sijk_ext
2683 36 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: rpgf_group, zet_group
2684 36 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: spi_group
2685 36 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: sijk, sijk_contr
2686 : REAL(KIND=dp), DIMENSION(3) :: ri, rij, rik, rj, rjk, rk
2687 36 : REAL(KIND=dp), DIMENSION(:), POINTER :: set_radius_i, set_radius_j, set_radius_k
2688 36 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: rpgf_i, rpgf_j, rpgf_k, zeti, zetj, zetk
2689 :
2690 36 : IF (PRESENT(screened)) screened = .FALSE.
2691 36 : ri = pbc(ctx%particle_set(atom_i)%r(1:3), ctx%cell)
2692 36 : rj = pbc(ctx%particle_set(atom_j)%r(1:3), ctx%cell)
2693 36 : rk = pbc(ctx%particle_set(atom_k)%r(1:3), ctx%cell)
2694 144 : rjk = rk - rj
2695 144 : rij = rj - ri
2696 144 : rik = rk - ri
2697 144 : djk = NORM2(rjk)
2698 144 : dij = NORM2(rij)
2699 144 : dik = NORM2(rik)
2700 :
2701 36 : ikind = ctx%kind_of(atom_i)
2702 36 : jkind = ctx%kind_of(atom_j)
2703 36 : kkind = ctx%kind_of(atom_k)
2704 : CALL get_gto_basis_set(ctx%basis_i(ikind)%gto_basis_set, first_sgf=first_sgf_i, &
2705 : lmax=lmax_i, lmin=lmin_i, npgf=npgfi, nset=nseti, &
2706 : nsgf_set=nsgfi, pgf_radius=rpgf_i, set_radius=set_radius_i, &
2707 36 : zet=zeti, kind_radius=kind_radius_i)
2708 : CALL get_gto_basis_set(ctx%basis_j(jkind)%gto_basis_set, first_sgf=first_sgf_j, &
2709 : lmax=lmax_j, lmin=lmin_j, npgf=npgfj, nset=nsetj, &
2710 : nsgf_set=nsgfj, pgf_radius=rpgf_j, set_radius=set_radius_j, &
2711 36 : zet=zetj, kind_radius=kind_radius_j)
2712 : CALL get_gto_basis_set(ctx%basis_k(kkind)%gto_basis_set, first_sgf=first_sgf_k, &
2713 : lmax=lmax_k, lmin=lmin_k, npgf=npgfk, nset=nsetk, &
2714 : nsgf_set=nsgfk, pgf_radius=rpgf_k, set_radius=set_radius_k, &
2715 36 : zet=zetk, kind_radius=kind_radius_k)
2716 :
2717 : IF (kind_radius_j + kind_radius_i + ctx%dr_ij < dij .OR. &
2718 36 : kind_radius_j + kind_radius_k + ctx%dr_jk < djk .OR. &
2719 : kind_radius_k + kind_radius_i + ctx%dr_ik < dik) THEN
2720 0 : IF (PRESENT(screened)) screened = .TRUE.
2721 0 : RETURN
2722 : END IF
2723 :
2724 36 : ncol = SIZE(transform, 2)
2725 36 : CPASSERT(SIZE(int_3c, 3) == ncol)
2726 252 : DO l = 0, ctx%maxli
2727 216 : npgf_group = 0
2728 216 : group_radius = 0.0_dp
2729 4392 : DO iset = 1, nseti
2730 4176 : IF (lmin_i(iset) /= l .OR. lmax_i(iset) /= l) CYCLE
2731 696 : npgf_group = npgf_group + npgfi(iset)
2732 4392 : group_radius = MAX(group_radius, set_radius_i(iset))
2733 : END DO
2734 216 : IF (npgf_group == 0) CYCLE
2735 192 : ncoi = npgf_group*ncoset(l)
2736 768 : ALLOCATE (zet_group(npgf_group), rpgf_group(npgf_group))
2737 768 : ALLOCATE (spi_group(ncoi, ncol), source=0.0_dp)
2738 192 : primitive_first = 1
2739 4008 : DO iset = 1, nseti
2740 3816 : IF (lmin_i(iset) /= l .OR. lmax_i(iset) /= l) CYCLE
2741 : zet_group(primitive_first:primitive_first + npgfi(iset) - 1) = &
2742 1392 : zeti(1:npgfi(iset), iset)
2743 : rpgf_group(primitive_first:primitive_first + npgfi(iset) - 1) = &
2744 1392 : rpgf_i(1:npgfi(iset), iset)
2745 696 : sgfi = first_sgf_i(1, iset)
2746 : spi_group((primitive_first - 1)*ncoset(l) + 1: &
2747 : (primitive_first + npgfi(iset) - 1)*ncoset(l), :) = &
2748 696 : MATMUL(ctx%spi(iset, ikind)%array, &
2749 : transform(transform_row + sgfi - 1: &
2750 1437180 : transform_row + sgfi + nsgfi(iset) - 2, :))
2751 4008 : primitive_first = primitive_first + npgfi(iset)
2752 : END DO
2753 :
2754 480 : DO jset = 1, nsetj
2755 288 : IF (set_radius_j(jset) + group_radius + ctx%dr_ij < dij) CYCLE
2756 832 : DO kset = 1, nsetk
2757 352 : IF (set_radius_j(jset) + set_radius_k(kset) + ctx%dr_jk < djk) CYCLE
2758 352 : IF (set_radius_k(kset) + group_radius + ctx%dr_ik < dik) CYCLE
2759 352 : ncoj = npgfj(jset)*ncoset(lmax_j(jset))
2760 352 : ncok = npgfk(kset)*ncoset(lmax_k(kset))
2761 352 : sgfj = first_sgf_j(1, jset)
2762 352 : sgfk = first_sgf_k(1, kset)
2763 352 : IF (ncoj*ncok*ncoi <= 0) CYCLE
2764 1760 : ALLOCATE (sijk(ncoj, ncok, ncoi), source=0.0_dp)
2765 352 : CALL timeset(routineN//'_eri', handle_eri)
2766 : CALL eri_3center(sijk, &
2767 : lmin_j(jset), lmax_j(jset), npgfj(jset), zetj(:, jset), &
2768 : rpgf_j(:, jset), rj, &
2769 : lmin_k(kset), lmax_k(kset), npgfk(kset), zetk(:, kset), &
2770 : rpgf_k(:, kset), rk, l, l, npgf_group, zet_group, &
2771 : rpgf_group, ri, djk, dij, dik, ws%lib, ctx%potential_parameter, &
2772 352 : int_abc_ext=sijk_ext)
2773 352 : CALL timestop(handle_eri)
2774 1760 : ALLOCATE (sijk_contr(nsgfj(jset), nsgfk(kset), ncol))
2775 352 : CALL timeset(routineN//'_contract', handle_contract)
2776 : CALL abc_contract_xsmm(sijk_contr, sijk, ctx%tspj(jset, jkind)%array, &
2777 : ctx%spk(kset, kkind)%array, spi_group, ncoj, ncok, ncoi, &
2778 352 : nsgfj(jset), nsgfk(kset), ncol, ws%cpp_buffer, ws%ccp_buffer)
2779 352 : CALL timestop(handle_contract)
2780 352 : DEALLOCATE (sijk)
2781 : int_3c(sgfj:sgfj + nsgfj(jset) - 1, sgfk:sgfk + nsgfk(kset) - 1, :) = &
2782 : int_3c(sgfj:sgfj + nsgfj(jset) - 1, sgfk:sgfk + nsgfk(kset) - 1, :) + &
2783 14476 : sijk_contr
2784 1344 : DEALLOCATE (sijk_contr)
2785 : END DO
2786 : END DO
2787 252 : DEALLOCATE (zet_group, rpgf_group, spi_group)
2788 : END DO
2789 :
2790 72 : END SUBROUTINE build_3c_integral_block_auto_ri_ctx
2791 :
2792 : ! **************************************************************************************************
2793 : !> \brief Returns the reference RI basis size of one atom.
2794 : !> \param bs_env ...
2795 : !> \param iatom ...
2796 : !> \return ...
2797 : ! **************************************************************************************************
2798 9 : INTEGER FUNCTION get_ref_RI_size(bs_env, iatom) RESULT(n)
2799 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
2800 : INTEGER, INTENT(IN) :: iatom
2801 :
2802 : INTEGER :: ikind
2803 :
2804 9 : ikind = bs_env%ri_rs%particle_set(iatom)%atomic_kind%kind_number
2805 9 : n = bs_env%basis_set_RI(ikind)%gto_basis_set%nsgf
2806 :
2807 9 : END FUNCTION get_ref_RI_size
2808 :
2809 : ! **************************************************************************************************
2810 : !> \brief Record the dense Gram, phi and threaded RHS memory model for an RI-RS fit.
2811 : !> The single-rank estimate is an upper bound for a distributed Gram solve.
2812 : !> \param ri_rs owned fitting statistics, reset with the representation
2813 : !> \param n_grid number of fitting grid rows, before screening when obtained by classification
2814 : !> \param n_ao number of AO columns
2815 : !> \param n_ri number of RHS columns
2816 : !> \param peak_bytes optional single-rank estimate used by solver classification
2817 : ! **************************************************************************************************
2818 104 : SUBROUTINE record_fit_memory(ri_rs, n_grid, n_ao, n_ri, peak_bytes)
2819 : !$ USE OMP_LIB, ONLY: omp_get_max_threads
2820 : TYPE(ri_rs_env), INTENT(INOUT) :: ri_rs
2821 : INTEGER, INTENT(IN) :: n_grid, n_ao, n_ri
2822 : REAL(KIND=dp), INTENT(OUT), OPTIONAL :: peak_bytes
2823 :
2824 : INTEGER :: nthreads
2825 : REAL(KIND=dp) :: ng, peak
2826 :
2827 104 : nthreads = 1
2828 104 : !$ nthreads = omp_get_max_threads()
2829 104 : ng = REAL(n_grid, dp)
2830 : peak = 8.0_dp*ng*ng + 8.0_dp*ng*REAL(n_ao, dp) + &
2831 104 : 8.0_dp*ng*REAL(n_ri, dp)*REAL(1 + nthreads, dp)
2832 104 : ri_rs%fit_max_grid = MAX(ri_rs%fit_max_grid, n_grid)
2833 104 : ri_rs%fit_peak_bytes = MAX(ri_rs%fit_peak_bytes, peak)
2834 104 : IF (PRESENT(peak_bytes)) peak_bytes = peak
2835 104 : END SUBROUTINE record_fit_memory
2836 :
2837 : ! **************************************************************************************************
2838 : !> \brief Prints the fitting memory bound before allocating the local fitting RHS and Gram matrix.
2839 : !> \param bs_env recorded fitting sizes and output unit
2840 : !> \param mem_avail_GB available memory per rank, measured collectively by the caller
2841 : ! **************************************************************************************************
2842 46 : SUBROUTINE print_fit_memory_estimate(bs_env, mem_avail_GB)
2843 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
2844 : REAL(KIND=dp), INTENT(IN) :: mem_avail_GB
2845 :
2846 : CHARACTER(LEN=max_line_length) :: msg
2847 : REAL(KIND=dp) :: mem_gram_GB, mem_Zlp_peak_GB
2848 :
2849 46 : mem_gram_GB = 8.0_dp*REAL(bs_env%ri_rs%fit_max_grid, dp)**2*1.0E-9_dp
2850 46 : mem_Zlp_peak_GB = bs_env%ri_rs%fit_peak_bytes*1.0E-9_dp
2851 46 : IF (bs_env%unit_nr > 0) THEN
2852 23 : WRITE (bs_env%unit_nr, '(T2,A)') 'RI-RS fitting memory estimate per MPI process:'
2853 : WRITE (bs_env%unit_nr, '(T4,A,F37.2,A)') &
2854 23 : 'Available memory per process (system)', mem_avail_GB, ' GB'
2855 : WRITE (bs_env%unit_nr, '(T4,A,F17.2,A)') &
2856 23 : 'Required for Z_lP fit (Gram+phi+RHS, single-rank bound)', mem_Zlp_peak_GB, ' GB'
2857 : WRITE (bs_env%unit_nr, '(T4,A,T69,I12)') &
2858 23 : 'Largest fitting-grid number of grid points:', bs_env%ri_rs%fit_max_grid
2859 : WRITE (bs_env%unit_nr, '(T4,A,T69,F9.2,A)') &
2860 23 : 'Largest dense Gram matrix:', mem_gram_GB, ' GB'
2861 23 : CALL m_flush(bs_env%unit_nr)
2862 : END IF
2863 :
2864 46 : IF (mem_avail_GB > 0.0_dp .AND. mem_Zlp_peak_GB > mem_avail_GB) THEN
2865 : WRITE (msg, '(A,F0.2,A,F0.2,A)') &
2866 0 : "The estimated peak memory of the Z_lP solve, ", mem_Zlp_peak_GB, &
2867 0 : " GB per process, exceeds the available ", mem_avail_GB, &
2868 : " GB. The dense Gram matrix scales as n_local_grid^2; "// &
2869 : "the estimate assumes a single-rank solve. Localized fits may distribute the Gram "// &
2870 : "matrix over a ScaLAPACK subgroup. Use fewer MPI ranks per node if more memory "// &
2871 0 : "per rank is needed."
2872 0 : CPWARN(TRIM(msg))
2873 : END IF
2874 46 : END SUBROUTINE print_fit_memory_estimate
2875 :
2876 : END MODULE ri_rs_setup
|
