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 Optimize automatically initialized atom-centred RI-RS grids.
10 : !> \par History
11 : !> 09.2026 created [Jan Wilhelm]
12 : ! **************************************************************************************************
13 : MODULE ri_rs_grid_optimization
14 : USE basis_set_types, ONLY: gto_basis_set_type
15 : USE cell_types, ONLY: cell_type
16 : USE cp_lbfgs, ONLY: setulb
17 : USE gw_utils_compute_integrals, ONLY: build_3c_integral_block_ctx,&
18 : gw_3c_ctx_create,&
19 : gw_3c_ctx_release,&
20 : gw_3c_ctx_type,&
21 : gw_3c_ws_create,&
22 : gw_3c_ws_release,&
23 : gw_3c_ws_type
24 : USE kinds, ONLY: dp
25 : USE machine, ONLY: m_flush,&
26 : m_walltime
27 : USE message_passing, ONLY: mp_mem_avail_per_rank_GB,&
28 : mp_para_env_type
29 : USE particle_types, ONLY: particle_type
30 : USE post_scf_bandstructure_types, ONLY: post_scf_bandstructure_type
31 : USE ri_rs_fitting, ONLY: evaluate_ri_rs_fit
32 : USE ri_rs_grid_initialization, ONLY: initialize_ri_rs_grid
33 : USE ri_rs_types, ONLY: rirs_grid_type
34 : USE ri_rs_utils, ONLY: evaluate_ao_basis_on_points,&
35 : filter_grid_to_voronoi,&
36 : get_rirs_cluster_atoms
37 : #include "./base/base_uses.f90"
38 :
39 : IMPLICIT NONE
40 : PRIVATE
41 :
42 : PUBLIC :: optimize_ri_rs_grid
43 :
44 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'ri_rs_grid_optimization'
45 :
46 : TYPE, PRIVATE :: local_cluster_3c_integrals_type
47 : INTEGER, ALLOCATABLE :: atom_indices(:)
48 : REAL(KIND=dp), ALLOCATABLE :: Int_3c(:, :, :)
49 : END TYPE local_cluster_3c_integrals_type
50 :
51 : CONTAINS
52 :
53 : ! **************************************************************************************************
54 : !> \brief Initialize Lebedev grids and subsequently optimize their grid-point coordinates.
55 : !> \param bs_env ...
56 : ! **************************************************************************************************
57 10 : SUBROUTINE optimize_ri_rs_grid(bs_env)
58 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
59 :
60 : CHARACTER(len=*), PARAMETER :: routineN = 'optimize_ri_rs_grid'
61 :
62 : INTEGER :: handle
63 10 : TYPE(local_cluster_3c_integrals_type), ALLOCATABLE :: cluster_3c_int(:)
64 :
65 10 : CALL timeset(routineN, handle)
66 :
67 : ! Validate all state and user input before the initialization consumes them.
68 10 : CALL validate_ri_rs_grid_optimization_input(bs_env)
69 :
70 : ! setup Lebedev initial grid inside Voronoi volume, select points with Cholesky decomp.
71 10 : CALL initialize_ri_rs_grid(bs_env)
72 :
73 : ! Validate the initialized RI-RS state and report the start of the optimization.
74 10 : CALL prepare_ri_rs_grid_optimization(bs_env)
75 :
76 : ! C_A = {B: |R_A - R_B| < R_cut}; store (μν|P) for μ, ν, and P centered in C_A.
77 10 : CALL build_cluster_3c_int(bs_env, cluster_3c_int)
78 :
79 : ! E_A = Σ_{μνP∈C_A} [(μν|P) - Σ_{l∈G_A} ϕ_μ(r_l) ϕ_ν(r_l) Z_lP^(A)]².
80 : ! Minimize E_loc(norm) = N_atom^(-1) Σ_A [E_A / Σ_{μνP∈C_A} |(μν|P)|²].
81 10 : CALL optimize_grid_coordinates(bs_env, cluster_3c_int)
82 :
83 10 : CALL timestop(handle)
84 :
85 25 : END SUBROUTINE optimize_ri_rs_grid
86 :
87 : ! **************************************************************************************************
88 : !> \brief Validate state and user input needed to initialize and optimize RI-RS grids.
89 : !> \param bs_env ...
90 : ! **************************************************************************************************
91 10 : SUBROUTINE validate_ri_rs_grid_optimization_input(bs_env)
92 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
93 :
94 : INTEGER :: iatom, ikind
95 :
96 10 : CPASSERT(ASSOCIATED(bs_env%ri_rs%cell))
97 10 : CPASSERT(ASSOCIATED(bs_env%ri_rs%particle_set))
98 10 : CPASSERT(ASSOCIATED(bs_env%para_env))
99 10 : CPASSERT(ALLOCATED(bs_env%basis_set_AO))
100 10 : CPASSERT(ALLOCATED(bs_env%basis_set_RI))
101 10 : CPASSERT(ALLOCATED(bs_env%sizes_AO))
102 10 : CPASSERT(SIZE(bs_env%ri_rs%particle_set) == bs_env%n_atom)
103 10 : CPASSERT(SIZE(bs_env%sizes_AO) == bs_env%n_atom)
104 10 : CPASSERT(SIZE(bs_env%basis_set_AO) == SIZE(bs_env%basis_set_RI))
105 :
106 10 : IF (bs_env%ri_rs%grid_opt%rs_ao_ratio <= 0.0_dp) THEN
107 0 : CPABORT("GRID_OPTIMIZATION%RS_AO_RATIO must be positive")
108 : END IF
109 10 : IF (bs_env%ri_rs%grid_opt%cutoff_atomic_cluster <= 0.0_dp) THEN
110 0 : CPABORT("GRID_OPTIMIZATION%CUTOFF_ATOMIC_CLUSTER must be positive")
111 : END IF
112 10 : IF (bs_env%ri_rs%grid_opt%max_iter < 1) THEN
113 0 : CPABORT("GRID_OPTIMIZATION%MAX_ITER must be positive")
114 : END IF
115 10 : IF (bs_env%ri_rs%tikhonov < 0.0_dp) THEN
116 0 : CPABORT("RI_RS%TIKHONOV must not be negative")
117 : END IF
118 :
119 40 : DO iatom = 1, bs_env%n_atom
120 30 : ikind = bs_env%ri_rs%particle_set(iatom)%atomic_kind%kind_number
121 30 : CPASSERT(ikind >= 1 .AND. ikind <= SIZE(bs_env%basis_set_AO))
122 30 : CPASSERT(ASSOCIATED(bs_env%basis_set_AO(ikind)%gto_basis_set))
123 30 : CPASSERT(ASSOCIATED(bs_env%basis_set_RI(ikind)%gto_basis_set))
124 30 : IF (bs_env%sizes_AO(iatom) < 1 .OR. &
125 10 : bs_env%basis_set_RI(ikind)%gto_basis_set%nsgf < 1) THEN
126 0 : CPABORT("Every atom in GRID_OPTIMIZATION needs ORB and RI_AUX functions")
127 : END IF
128 : END DO
129 10 : END SUBROUTINE validate_ri_rs_grid_optimization_input
130 :
131 : ! **************************************************************************************************
132 : !> \brief Validate initialized RI-RS state and announce the grid optimization.
133 : !> \param bs_env ...
134 : ! **************************************************************************************************
135 10 : SUBROUTINE prepare_ri_rs_grid_optimization(bs_env)
136 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
137 :
138 : INTEGER :: iatom, unit_nr
139 :
140 10 : CPASSERT(ALLOCATED(bs_env%ri_rs%atomic_grids))
141 10 : CPASSERT(SIZE(bs_env%ri_rs%atomic_grids) == bs_env%n_atom)
142 :
143 40 : DO iatom = 1, bs_env%n_atom
144 30 : CPASSERT(ALLOCATED(bs_env%ri_rs%atomic_grids(iatom)%raw_points))
145 30 : CPASSERT(SIZE(bs_env%ri_rs%atomic_grids(iatom)%raw_points, 1) == 3)
146 40 : CPASSERT(SIZE(bs_env%ri_rs%atomic_grids(iatom)%raw_points, 2) > 0)
147 : END DO
148 :
149 10 : unit_nr = bs_env%unit_nr
150 10 : IF (unit_nr > 0) THEN
151 5 : WRITE (unit_nr, '(T2,A)') 'Started RI-RS grid optimization'
152 5 : CALL m_flush(unit_nr)
153 : END IF
154 10 : END SUBROUTINE prepare_ri_rs_grid_optimization
155 :
156 : ! **************************************************************************************************
157 : !> \brief Evaluate the normalized local objective E_loc and its gradient for one L-BFGS trial vector.
158 : !> \param grid_coordinates Flattened atom-relative grid coordinates.
159 : !> \param atom_coordinate_offsets Starting coordinate offset for each atom.
160 : !> \param bs_env ...
161 : !> \param cluster_3c_int Rank-local cluster three-centre integrals.
162 : !> \param fit_memory_budget Per-rank bytes available for cluster workspaces; zero means unknown.
163 : !> \param normalized_error Mean normalized squared three-centre-integral error.
164 : !> \param coordinate_gradient Optional derivative; omitting it requests final error diagnostics only.
165 : !> \param max_abs_error Largest absolute integral error in diagnostic mode; zero in gradient mode.
166 : !> \param fit_successful Whether every local-cluster evaluation succeeded.
167 : ! **************************************************************************************************
168 256 : SUBROUTINE grid_objective(grid_coordinates, atom_coordinate_offsets, bs_env, cluster_3c_int, &
169 256 : fit_memory_budget, normalized_error, coordinate_gradient, &
170 : max_abs_error, fit_successful)
171 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: grid_coordinates
172 : INTEGER, DIMENSION(:), INTENT(IN) :: atom_coordinate_offsets
173 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
174 : TYPE(local_cluster_3c_integrals_type), &
175 : DIMENSION(:), INTENT(IN) :: cluster_3c_int
176 : REAL(KIND=dp), INTENT(IN) :: fit_memory_budget
177 : REAL(KIND=dp), INTENT(OUT) :: normalized_error
178 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT), OPTIONAL :: coordinate_gradient
179 : REAL(KIND=dp), INTENT(OUT) :: max_abs_error
180 : LOGICAL, INTENT(OUT) :: fit_successful
181 :
182 : INTEGER :: evaluated_cluster_count, icluster, &
183 : successful_ranks
184 : LOGICAL :: cluster_fit_successful
185 : REAL(KIND=dp) :: cluster_max_abs_error, &
186 : cluster_normalized_error
187 256 : REAL(KIND=dp), ALLOCATABLE :: cluster_coordinate_gradient(:, :)
188 : TYPE(mp_para_env_type), POINTER :: para_env
189 :
190 256 : para_env => bs_env%para_env
191 256 : IF (PRESENT(coordinate_gradient)) THEN
192 246 : CPASSERT(SIZE(coordinate_gradient) == SIZE(grid_coordinates))
193 18174 : coordinate_gradient(:) = 0.0_dp
194 : END IF
195 : CALL unpack_lbfgs_grid_coordinates(grid_coordinates, atom_coordinate_offsets, &
196 256 : bs_env%ri_rs%atomic_grids)
197 256 : normalized_error = 0.0_dp
198 256 : evaluated_cluster_count = 0
199 256 : max_abs_error = 0.0_dp
200 256 : fit_successful = .TRUE.
201 640 : DO icluster = 1, SIZE(cluster_3c_int)
202 384 : IF (PRESENT(coordinate_gradient)) THEN
203 : CALL evaluate_local_cluster(cluster_3c_int(icluster), bs_env, fit_memory_budget, cluster_normalized_error, &
204 369 : cluster_max_abs_error, cluster_fit_successful, cluster_coordinate_gradient)
205 : ELSE
206 : CALL evaluate_local_cluster(cluster_3c_int(icluster), bs_env, fit_memory_budget, cluster_normalized_error, &
207 15 : cluster_max_abs_error, cluster_fit_successful)
208 : END IF
209 384 : IF (.NOT. cluster_fit_successful) THEN
210 0 : fit_successful = .FALSE.
211 0 : EXIT
212 : END IF
213 384 : normalized_error = normalized_error + cluster_normalized_error
214 384 : IF (PRESENT(coordinate_gradient)) THEN
215 : CALL accumulate_atom_gradient(cluster_3c_int(icluster)%atom_indices, &
216 : bs_env%ri_rs%atomic_grids, &
217 : atom_coordinate_offsets, cluster_coordinate_gradient, &
218 369 : coordinate_gradient)
219 369 : DEALLOCATE (cluster_coordinate_gradient)
220 : END IF
221 384 : evaluated_cluster_count = evaluated_cluster_count + 1
222 : max_abs_error = &
223 640 : MAX(max_abs_error, cluster_max_abs_error)
224 : END DO
225 256 : successful_ranks = MERGE(1, 0, fit_successful)
226 256 : CALL para_env%sum(normalized_error)
227 36112 : IF (PRESENT(coordinate_gradient)) CALL para_env%sum(coordinate_gradient)
228 256 : CALL para_env%sum(evaluated_cluster_count)
229 256 : CALL para_env%sum(successful_ranks)
230 256 : CALL para_env%max(max_abs_error)
231 : fit_successful = evaluated_cluster_count == bs_env%n_atom .AND. &
232 256 : successful_ranks == para_env%num_pe
233 256 : IF (.NOT. fit_successful) THEN
234 0 : normalized_error = HUGE(normalized_error)
235 0 : IF (PRESENT(coordinate_gradient)) coordinate_gradient(:) = 0.0_dp
236 0 : RETURN
237 : END IF
238 : ! E_loc(norm) = (1/N_atom) Σ_A [E_A / Σ_{μνP ∈ C_A} |(μν|P)|²].
239 256 : normalized_error = normalized_error/REAL(bs_env%n_atom, dp)
240 18184 : IF (PRESENT(coordinate_gradient)) coordinate_gradient(:) = coordinate_gradient/REAL(bs_env%n_atom, dp)
241 256 : END SUBROUTINE grid_objective
242 :
243 : ! **************************************************************************************************
244 : !> \brief Drive L-BFGS, repeatedly evaluating E_loc and its gradient, and return the best trial.
245 : !> \param bs_env ...
246 : !> \param cluster_3c_int Rank-local cluster three-centre integrals.
247 : ! **************************************************************************************************
248 10 : SUBROUTINE optimize_grid_coordinates(bs_env, cluster_3c_int)
249 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
250 : TYPE(local_cluster_3c_integrals_type), &
251 : DIMENSION(:), INTENT(IN) :: cluster_3c_int
252 :
253 : INTEGER, PARAMETER :: lbfgs_history = 7
254 : REAL(KIND=dp), PARAMETER :: lbfgs_factr = 0.0_dp, &
255 : lbfgs_pgtol = 1.0E-9_dp, &
256 : mem_safety = 0.8_dp
257 :
258 : CHARACTER(LEN=60) :: line_search_state, optimizer_task
259 : INTEGER :: evaluation_count, handle, iatom, &
260 : molecular_outside, npoints, unit_nr
261 10 : INTEGER, ALLOCATABLE :: atom_coordinate_offsets(:), &
262 10 : bound_types(:), integer_workspace(:)
263 : INTEGER, DIMENSION(44) :: integer_state
264 : LOGICAL :: evaluation_successful, fit_successful, &
265 : have_best
266 10 : LOGICAL, ALLOCATABLE :: inside(:)
267 : LOGICAL, DIMENSION(4) :: logical_state
268 : REAL(KIND=dp) :: best_normalized_error, &
269 : fit_memory_budget, max_abs_error, &
270 : mem_avail_GB, normalized_error, &
271 : start_time
272 10 : REAL(KIND=dp), ALLOCATABLE :: best_grid_coordinates(:), coordinate_gradient(:), &
273 10 : grid_coordinates(:), lower_bounds(:), points(:, :), upper_bounds(:), workspace(:)
274 : REAL(KIND=dp), DIMENSION(29) :: real_state
275 :
276 10 : CALL timeset("rirs_grid_LBFGS", handle)
277 10 : start_time = m_walltime()
278 :
279 : CALL pack_lbfgs_grid_coordinates(bs_env%ri_rs%atomic_grids, grid_coordinates, &
280 10 : atom_coordinate_offsets)
281 10 : CPASSERT(SIZE(grid_coordinates) > 0)
282 :
283 : ALLOCATE (best_grid_coordinates(SIZE(grid_coordinates)), &
284 : coordinate_gradient(SIZE(grid_coordinates)), &
285 : lower_bounds(SIZE(grid_coordinates)), &
286 : upper_bounds(SIZE(grid_coordinates)), &
287 : bound_types(SIZE(grid_coordinates)), &
288 : integer_workspace(3*SIZE(grid_coordinates)), &
289 : workspace(2*lbfgs_history*SIZE(grid_coordinates) + &
290 120 : 5*SIZE(grid_coordinates) + 11*lbfgs_history**2 + 8*lbfgs_history))
291 : ! L-BFGS-B ignores lower_bounds and upper_bounds when bound_types is zero.
292 10 : lower_bounds = 0.0_dp
293 10 : upper_bounds = 0.0_dp
294 10 : bound_types = 0
295 10 : optimizer_task = 'START'
296 10 : line_search_state = ''
297 10 : normalized_error = HUGE(normalized_error)
298 10 : coordinate_gradient = 0.0_dp
299 10 : workspace = 0.0_dp
300 10 : integer_workspace = 0
301 10 : logical_state = .FALSE.
302 10 : integer_state = 0
303 10 : real_state = 0.0_dp
304 10 : evaluation_count = 0
305 10 : have_best = .FALSE.
306 10 : best_normalized_error = HUGE(best_normalized_error)
307 2296 : best_grid_coordinates(:) = grid_coordinates
308 :
309 : ! This query is collective. Take one snapshot after resident integrals and L-BFGS
310 : ! workspaces exist, before ranks enter their unequal numbers of local cluster evaluations.
311 : ! Use the same safety fraction as the production Z_lP fitting classification.
312 10 : CALL mp_mem_avail_per_rank_GB(bs_env%para_env, mem_avail_GB)
313 : ! This budgets all incremental cluster workspaces, including the caller's AO arrays.
314 : ! The product cache receives only the remainder after those workspaces are reserved.
315 10 : fit_memory_budget = mem_safety*mem_avail_GB*1.0E9_dp
316 10 : IF (mem_avail_GB <= 0.0_dp .AND. bs_env%unit_nr > 0) THEN
317 0 : CPWARN("RI-RS grid optimization: no meminfo; streaming one RI column")
318 : END IF
319 :
320 : DO
321 : CALL setulb(SIZE(grid_coordinates), lbfgs_history, grid_coordinates, &
322 : lower_bounds, upper_bounds, bound_types, &
323 : normalized_error, coordinate_gradient, lbfgs_factr, lbfgs_pgtol, &
324 : workspace, integer_workspace, optimizer_task, -1, line_search_state, &
325 450 : logical_state, integer_state, real_state, -1.0_dp)
326 460 : IF (optimizer_task(1:2) == 'FG') THEN
327 252 : IF (evaluation_count >= bs_env%ri_rs%grid_opt%max_iter) EXIT
328 : CALL grid_objective(grid_coordinates, atom_coordinate_offsets, bs_env, cluster_3c_int, &
329 : fit_memory_budget, normalized_error, coordinate_gradient, &
330 246 : max_abs_error, evaluation_successful)
331 246 : evaluation_count = evaluation_count + 1
332 246 : IF (.NOT. evaluation_successful) EXIT
333 246 : IF (evaluation_count == 1 .AND. bs_env%unit_nr > 0) THEN
334 : WRITE (bs_env%unit_nr, '(T2,A,T61,ES20.12)') &
335 5 : 'RI-RS grid initial normalized error', normalized_error
336 5 : CALL m_flush(bs_env%unit_nr)
337 : END IF
338 246 : IF (.NOT. have_best .OR. normalized_error < best_normalized_error) THEN
339 216 : have_best = .TRUE.
340 216 : best_normalized_error = normalized_error
341 16308 : best_grid_coordinates(:) = grid_coordinates
342 : END IF
343 198 : ELSE IF (optimizer_task(1:5) == 'NEW_X') THEN
344 : CYCLE
345 : ELSE
346 194 : EXIT
347 : END IF
348 : END DO
349 10 : IF (.NOT. have_best) CPABORT("RI-RS grid optimization produced no regularized fit")
350 2296 : grid_coordinates(:) = best_grid_coordinates
351 : ! grid_objective restores the best coordinates to atomic_grids and evaluates only the
352 : ! final errors. No best-fit coefficients are retained across L-BFGS trials.
353 : CALL grid_objective(grid_coordinates, atom_coordinate_offsets, bs_env, cluster_3c_int, &
354 10 : fit_memory_budget, normalized_error, max_abs_error=max_abs_error, fit_successful=fit_successful)
355 :
356 10 : IF (.NOT. fit_successful) CPABORT("RI-RS grid optimization produced no regularized fit")
357 :
358 40 : DO iatom = 1, SIZE(bs_env%ri_rs%atomic_grids)
359 : bs_env%ri_rs%atomic_grids(iatom)%npts = &
360 40 : SIZE(bs_env%ri_rs%atomic_grids(iatom)%raw_points, 2)
361 : END DO
362 10 : bs_env%ri_rs%Z_lP_exists = .FALSE.
363 :
364 10 : unit_nr = bs_env%unit_nr
365 10 : IF (unit_nr > 0) THEN
366 5 : npoints = 0
367 5 : molecular_outside = 0
368 20 : DO iatom = 1, SIZE(bs_env%ri_rs%atomic_grids)
369 0 : ALLOCATE (points(3, bs_env%ri_rs%atomic_grids(iatom)%npts), &
370 75 : inside(bs_env%ri_rs%atomic_grids(iatom)%npts))
371 1539 : points(:, :) = bs_env%ri_rs%atomic_grids(iatom)%raw_points
372 15 : CALL filter_grid_to_voronoi(points, iatom, bs_env%ri_rs%particle_set, mask=inside)
373 396 : molecular_outside = molecular_outside + COUNT(.NOT. inside)
374 15 : npoints = npoints + SIZE(inside)
375 20 : DEALLOCATE (points, inside)
376 : END DO
377 : WRITE (unit_nr, '(T2,A,T69,F10.1,A)') &
378 5 : 'RI-RS grid optimization completed, execution time:', m_walltime() - start_time, ' s'
379 5 : WRITE (unit_nr, '(T2,A,T72,ES9.1)') 'Normalized 3C error:', normalized_error
380 5 : WRITE (unit_nr, '(T2,A,T72,ES9.1)') 'Maximum absolute 3C error:', max_abs_error
381 : WRITE (unit_nr, '(T2,A,T60,I9,A,I9)') &
382 5 : 'Optimized points outside molecular Voronoi cells:', molecular_outside, ' / ', npoints
383 5 : FLUSH (unit_nr)
384 : END IF
385 10 : CALL timestop(handle)
386 20 : END SUBROUTINE optimize_grid_coordinates
387 :
388 : ! **************************************************************************************************
389 : !> \brief Build each C_A and store the upper AO triangle of (μν|P); this routine owns the integral context.
390 : !> \param bs_env ...
391 : !> \param cluster_3c_int Rank-local cluster three-centre integrals.
392 : ! **************************************************************************************************
393 10 : SUBROUTINE build_cluster_3c_int(bs_env, cluster_3c_int)
394 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
395 : TYPE(local_cluster_3c_integrals_type), &
396 : ALLOCATABLE, INTENT(OUT) :: cluster_3c_int(:)
397 :
398 : INTEGER :: handle, i_cluster_atom_j, i_cluster_atom_k, i_cluster_atom_p, iatom, iatom_j, &
399 : iatom_k, iatom_p, icenter_atom, icluster_atom, ikind, ilocal_cluster, ip, &
400 : n_local_clusters, nAO_cluster, nRI_cluster, nRI_nonzero
401 20 : INTEGER, ALLOCATABLE :: AO_offset(:), RI_offset(:), &
402 10 : sizes_ref_RI(:)
403 : LOGICAL :: screened
404 : REAL(KIND=dp) :: cluster_radius
405 10 : REAL(KIND=dp), ALLOCATABLE :: Int_3c_nonzero(:, :, :)
406 : TYPE(cell_type), POINTER :: cell
407 130 : TYPE(gw_3c_ctx_type) :: integral_context
408 10 : TYPE(gw_3c_ws_type) :: workspace
409 : TYPE(mp_para_env_type), POINTER :: para_env
410 10 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
411 :
412 10 : CALL timeset("rirs_grid_build_clusters", handle)
413 :
414 10 : cell => bs_env%ri_rs%cell
415 10 : para_env => bs_env%para_env
416 10 : particle_set => bs_env%ri_rs%particle_set
417 10 : cluster_radius = bs_env%ri_rs%grid_opt%cutoff_atomic_cluster
418 30 : ALLOCATE (sizes_ref_RI(bs_env%n_atom))
419 40 : DO iatom = 1, bs_env%n_atom
420 30 : ikind = particle_set(iatom)%atomic_kind%kind_number
421 40 : sizes_ref_RI(iatom) = bs_env%basis_set_RI(ikind)%gto_basis_set%nsgf
422 : END DO
423 : n_local_clusters = COUNT([(MOD(icenter_atom - 1, para_env%num_pe) == para_env%mepos, &
424 80 : icenter_atom=1, bs_env%n_atom)])
425 0 : ALLOCATE (cluster_3c_int(n_local_clusters), &
426 75 : AO_offset(bs_env%n_atom), RI_offset(bs_env%n_atom))
427 : CALL gw_3c_ctx_create(integral_context, bs_env, bs_env%ri_metric, &
428 : basis_j=bs_env%basis_set_AO, basis_k=bs_env%basis_set_AO, &
429 10 : basis_i=bs_env%basis_set_RI)
430 10 : CALL gw_3c_ws_create(workspace, integral_context)
431 10 : ilocal_cluster = 0
432 40 : DO icenter_atom = 1, bs_env%n_atom
433 30 : IF (MOD(icenter_atom - 1, para_env%num_pe) /= para_env%mepos) CYCLE
434 15 : ilocal_cluster = ilocal_cluster + 1
435 : CALL get_rirs_cluster_atoms(particle_set, cell, icenter_atom, cluster_radius, &
436 15 : cluster_3c_int(ilocal_cluster)%atom_indices)
437 15 : AO_offset = 0
438 15 : RI_offset = 0
439 15 : nAO_cluster = 0
440 15 : nRI_cluster = 0
441 60 : DO icluster_atom = 1, SIZE(cluster_3c_int(ilocal_cluster)%atom_indices)
442 45 : iatom = cluster_3c_int(ilocal_cluster)%atom_indices(icluster_atom)
443 45 : AO_offset(iatom) = nAO_cluster
444 45 : RI_offset(iatom) = nRI_cluster
445 45 : nAO_cluster = nAO_cluster + bs_env%sizes_AO(iatom)
446 60 : nRI_cluster = nRI_cluster + sizes_ref_RI(iatom)
447 : END DO
448 0 : ALLOCATE (cluster_3c_int(ilocal_cluster)%Int_3c(nAO_cluster, nAO_cluster, nRI_cluster), &
449 92244 : source=0.0_dp)
450 60 : DO i_cluster_atom_p = 1, SIZE(cluster_3c_int(ilocal_cluster)%atom_indices)
451 45 : iatom_p = cluster_3c_int(ilocal_cluster)%atom_indices(i_cluster_atom_p)
452 195 : DO i_cluster_atom_k = 1, SIZE(cluster_3c_int(ilocal_cluster)%atom_indices)
453 135 : iatom_k = cluster_3c_int(ilocal_cluster)%atom_indices(i_cluster_atom_k)
454 : ! The evaluator reads only the upper AO triangle (DSYMM with UPLO='U').
455 : ! Global atom ordering also orders the cluster AO blocks; lower blocks stay zero.
456 450 : DO i_cluster_atom_j = 1, i_cluster_atom_k
457 270 : iatom_j = cluster_3c_int(ilocal_cluster)%atom_indices(i_cluster_atom_j)
458 : CALL build_3c_integral_block_ctx(cluster_3c_int(ilocal_cluster)%Int_3c, &
459 : integral_context, workspace, &
460 : atom_j=iatom_j, atom_k=iatom_k, atom_i=iatom_p, &
461 : j_offset=AO_offset(iatom_j), &
462 : k_offset=AO_offset(iatom_k), &
463 405 : i_offset=RI_offset(iatom_p), screened=screened)
464 : END DO
465 : END DO
466 : END DO
467 : ! Zero auxiliary columns contribute neither to the objective nor its derivatives.
468 : nRI_nonzero = COUNT([(ANY(cluster_3c_int(ilocal_cluster)%Int_3c(:, :, ip) /= 0.0_dp), &
469 3348 : ip=1, nRI_cluster)])
470 25 : IF (nRI_nonzero < nRI_cluster) THEN
471 0 : ALLOCATE (Int_3c_nonzero(nAO_cluster, nAO_cluster, nRI_nonzero))
472 0 : nRI_nonzero = 0
473 0 : DO ip = 1, nRI_cluster
474 0 : IF (.NOT. ANY(cluster_3c_int(ilocal_cluster)%Int_3c(:, :, ip) /= 0.0_dp)) CYCLE
475 0 : nRI_nonzero = nRI_nonzero + 1
476 0 : Int_3c_nonzero(:, :, nRI_nonzero) = cluster_3c_int(ilocal_cluster)%Int_3c(:, :, ip)
477 : END DO
478 0 : CALL MOVE_ALLOC(Int_3c_nonzero, cluster_3c_int(ilocal_cluster)%Int_3c)
479 : END IF
480 : END DO
481 10 : CALL gw_3c_ws_release(workspace)
482 10 : CALL gw_3c_ctx_release(integral_context)
483 10 : CALL timestop(handle)
484 20 : END SUBROUTINE build_cluster_3c_int
485 :
486 : ! **************************************************************************************************
487 : !> \brief Evaluate ϕ_μ(r_l) on G_A and compute the normalized local error E_A and its gradient.
488 : !> \param cluster_3c_int Local atom indices and exact three-centre integrals.
489 : !> \param bs_env ...
490 : !> \param fit_memory_budget Per-rank incremental workspace budget before allocating this cluster.
491 : !> \param normalized_error Normalized squared three-centre-integral error.
492 : !> \param max_abs_error Largest absolute integral error in diagnostic mode; zero in gradient mode.
493 : !> \param fit_successful Whether the regularized fitting equations were solved.
494 : !> \param coordinate_gradient Optional derivative with respect to the cluster's physical grid coordinates.
495 : ! **************************************************************************************************
496 384 : SUBROUTINE evaluate_local_cluster(cluster_3c_int, bs_env, fit_memory_budget, normalized_error, max_abs_error, &
497 : fit_successful, coordinate_gradient)
498 : TYPE(local_cluster_3c_integrals_type), INTENT(IN) :: cluster_3c_int
499 : TYPE(post_scf_bandstructure_type), POINTER :: bs_env
500 : REAL(KIND=dp), INTENT(IN) :: fit_memory_budget
501 : REAL(KIND=dp), INTENT(OUT) :: normalized_error, max_abs_error
502 : LOGICAL, INTENT(OUT) :: fit_successful
503 : REAL(KIND=dp), ALLOCATABLE, INTENT(OUT), OPTIONAL :: coordinate_gradient(:, :)
504 :
505 : INTEGER :: AO_offset, grid_point_offset, iatom, &
506 : icartesian, icluster_atom, igrid, &
507 : ikind, n_grid_points, nAO, nAO_atom
508 : REAL(KIND=dp) :: evaluation_budget
509 384 : REAL(KIND=dp), ALLOCATABLE :: ao_gradient(:, :), &
510 384 : dPhi_alpha_l_mu(:, :, :), &
511 384 : grid_points(:, :), Phi_l_mu(:, :)
512 : TYPE(cell_type), POINTER :: cell
513 : TYPE(gto_basis_set_type), POINTER :: AO_basis
514 384 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
515 :
516 384 : cell => bs_env%ri_rs%cell
517 384 : particle_set => bs_env%ri_rs%particle_set
518 384 : nAO = SIZE(cluster_3c_int%Int_3c, 1)
519 384 : n_grid_points = 0
520 1536 : DO icluster_atom = 1, SIZE(cluster_3c_int%atom_indices)
521 1152 : iatom = cluster_3c_int%atom_indices(icluster_atom)
522 1536 : n_grid_points = n_grid_points + SIZE(bs_env%ri_rs%atomic_grids(iatom)%raw_points, 2)
523 : END DO
524 384 : evaluation_budget = 0.0_dp
525 384 : IF (PRESENT(coordinate_gradient) .AND. fit_memory_budget > 0.0_dp) THEN
526 : ! phi, dphi and its adjoint use 5*g*a doubles; points and coordinate gradients use 6*g.
527 : evaluation_budget = fit_memory_budget - &
528 369 : 8.0_dp*REAL(n_grid_points, dp)*(5.0_dp*REAL(nAO, dp) + 6.0_dp)
529 369 : IF (evaluation_budget <= 0.0_dp) THEN
530 0 : CPABORT("RI-RS cluster AO workspace exceeds memory budget")
531 : END IF
532 : END IF
533 2304 : ALLOCATE (grid_points(3, n_grid_points), Phi_l_mu(n_grid_points, nAO))
534 384 : IF (PRESENT(coordinate_gradient)) THEN
535 : ALLOCATE (dPhi_alpha_l_mu(3, n_grid_points, nAO), ao_gradient(n_grid_points, nAO), &
536 2583 : coordinate_gradient(3, n_grid_points))
537 369 : dPhi_alpha_l_mu = 0.0_dp
538 369 : coordinate_gradient(:, :) = 0.0_dp
539 : END IF
540 384 : Phi_l_mu = 0.0_dp
541 384 : grid_point_offset = 0
542 1536 : DO icluster_atom = 1, SIZE(cluster_3c_int%atom_indices)
543 1152 : iatom = cluster_3c_int%atom_indices(icluster_atom)
544 : grid_points(:, grid_point_offset + 1:grid_point_offset + &
545 : SIZE(bs_env%ri_rs%atomic_grids(iatom)%raw_points, 2)) = &
546 : SPREAD(particle_set(iatom)%r, 2, &
547 : SIZE(bs_env%ri_rs%atomic_grids(iatom)%raw_points, 2)) + &
548 41580 : bs_env%ri_rs%atomic_grids(iatom)%raw_points
549 1536 : grid_point_offset = grid_point_offset + SIZE(bs_env%ri_rs%atomic_grids(iatom)%raw_points, 2)
550 : END DO
551 : AO_offset = 0
552 1536 : DO icluster_atom = 1, SIZE(cluster_3c_int%atom_indices)
553 1152 : iatom = cluster_3c_int%atom_indices(icluster_atom)
554 1152 : ikind = particle_set(iatom)%atomic_kind%kind_number
555 1152 : AO_basis => bs_env%basis_set_AO(ikind)%gto_basis_set
556 1152 : nAO_atom = AO_basis%nsgf
557 1152 : IF (PRESENT(coordinate_gradient)) THEN
558 : CALL evaluate_ao_basis_on_points(Phi_l_mu(:, AO_offset + 1:AO_offset + nAO_atom), &
559 : grid_points, AO_basis, particle_set(iatom)%r, cell, &
560 1107 : dphi=dPhi_alpha_l_mu(:, :, AO_offset + 1:AO_offset + nAO_atom))
561 : ELSE
562 : CALL evaluate_ao_basis_on_points(Phi_l_mu(:, AO_offset + 1:AO_offset + nAO_atom), &
563 45 : grid_points, AO_basis, particle_set(iatom)%r, cell)
564 : END IF
565 1536 : AO_offset = AO_offset + nAO_atom
566 : END DO
567 384 : max_abs_error = 0.0_dp
568 384 : IF (PRESENT(coordinate_gradient)) THEN
569 : CALL evaluate_ri_rs_fit(Phi_l_mu, cluster_3c_int%Int_3c, bs_env%ri_rs%tikhonov, &
570 : normalized_error, fit_successful, ao_gradient=ao_gradient, &
571 369 : workspace_budget_bytes=evaluation_budget)
572 369 : IF (.NOT. fit_successful) RETURN
573 9333 : DO igrid = 1, n_grid_points
574 36225 : DO icartesian = 1, 3
575 : coordinate_gradient(icartesian, igrid) = &
576 224100 : DOT_PRODUCT(ao_gradient(igrid, :), dPhi_alpha_l_mu(icartesian, igrid, :))
577 : END DO
578 : END DO
579 : ELSE
580 : CALL evaluate_ri_rs_fit(Phi_l_mu, cluster_3c_int%Int_3c, bs_env%ri_rs%tikhonov, &
581 15 : normalized_error, fit_successful, max_abs_error=max_abs_error)
582 : END IF
583 768 : END SUBROUTINE evaluate_local_cluster
584 :
585 : ! **************************************************************************************************
586 : !> \brief Sum gradient contributions from overlapping C_A into each atom-centred grid coordinate.
587 : !> \param cluster_atoms Atom indices in the local cluster.
588 : !> \param grids Atom-centred RI-RS grids.
589 : !> \param atom_coordinate_offsets Starting coordinate offset for each atom.
590 : !> \param cluster_coordinate_gradient Gradient for the cluster's physical grid points.
591 : !> \param coordinate_gradient Global flattened atom-relative gradient to update.
592 : ! **************************************************************************************************
593 369 : SUBROUTINE accumulate_atom_gradient(cluster_atoms, grids, atom_coordinate_offsets, &
594 369 : cluster_coordinate_gradient, coordinate_gradient)
595 : INTEGER, DIMENSION(:), INTENT(IN) :: cluster_atoms
596 : TYPE(rirs_grid_type), DIMENSION(:), INTENT(IN) :: grids
597 : INTEGER, DIMENSION(:), INTENT(IN) :: atom_coordinate_offsets
598 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: cluster_coordinate_gradient
599 : REAL(KIND=dp), DIMENSION(:), INTENT(INOUT) :: coordinate_gradient
600 :
601 : INTEGER :: grid_point_offset, iatom, icartesian, &
602 : icluster_atom, igrid_point
603 :
604 369 : grid_point_offset = 0
605 1476 : DO icluster_atom = 1, SIZE(cluster_atoms)
606 1107 : iatom = cluster_atoms(icluster_atom)
607 10071 : DO igrid_point = 1, SIZE(grids(iatom)%raw_points, 2)
608 36963 : DO icartesian = 1, 3
609 : coordinate_gradient(atom_coordinate_offsets(iatom) + 3*(igrid_point - 1) + icartesian) = &
610 : coordinate_gradient(atom_coordinate_offsets(iatom) + 3*(igrid_point - 1) + icartesian) + &
611 35856 : cluster_coordinate_gradient(icartesian, grid_point_offset + igrid_point)
612 : END DO
613 : END DO
614 1476 : grid_point_offset = grid_point_offset + SIZE(grids(iatom)%raw_points, 2)
615 : END DO
616 369 : END SUBROUTINE accumulate_atom_gradient
617 :
618 : ! **************************************************************************************************
619 : !> \brief Pack all atom-centred grids into the single Cartesian vector required by L-BFGS.
620 : !> \param grids Atom-centred RI-RS grids.
621 : !> \param grid_coordinates Flattened atom-relative grid coordinates.
622 : !> \param atom_coordinate_offsets Starting coordinate offset for each atom.
623 : !>
624 : !> SIZE(grid_coordinates) = 3 Σ_A N_grid,A. Coordinates are ordered Cartesian component first,
625 : !> then grid point, then atom, matching the column-major layout of raw_points(3,N_grid,A).
626 : ! **************************************************************************************************
627 10 : SUBROUTINE pack_lbfgs_grid_coordinates(grids, grid_coordinates, atom_coordinate_offsets)
628 : TYPE(rirs_grid_type), DIMENSION(:), INTENT(IN) :: grids
629 : REAL(KIND=dp), ALLOCATABLE, INTENT(OUT) :: grid_coordinates(:)
630 : INTEGER, ALLOCATABLE, INTENT(OUT) :: atom_coordinate_offsets(:)
631 :
632 : INTEGER :: atom_coordinate_count, iatom, &
633 : n_grid_coordinates
634 :
635 30 : ALLOCATE (atom_coordinate_offsets(SIZE(grids)))
636 10 : n_grid_coordinates = 0
637 40 : DO iatom = 1, SIZE(grids)
638 30 : atom_coordinate_offsets(iatom) = n_grid_coordinates
639 100 : n_grid_coordinates = n_grid_coordinates + SIZE(grids(iatom)%raw_points)
640 : END DO
641 30 : ALLOCATE (grid_coordinates(n_grid_coordinates))
642 :
643 40 : DO iatom = 1, SIZE(grids)
644 90 : atom_coordinate_count = SIZE(grids(iatom)%raw_points)
645 : grid_coordinates(atom_coordinate_offsets(iatom) + 1: &
646 : atom_coordinate_offsets(iatom) + atom_coordinate_count) = &
647 70 : RESHAPE(grids(iatom)%raw_points, [atom_coordinate_count])
648 : END DO
649 10 : END SUBROUTINE pack_lbfgs_grid_coordinates
650 :
651 : ! **************************************************************************************************
652 : !> \brief Restore the optimized Cartesian vector to the persistent atom-centred RI-RS grids.
653 : !> \param grid_coordinates Flattened atom-relative grid coordinates.
654 : !> \param atom_coordinate_offsets Starting coordinate offset for each atom.
655 : !> \param grids Atom-centred RI-RS grids to update.
656 : ! **************************************************************************************************
657 256 : SUBROUTINE unpack_lbfgs_grid_coordinates(grid_coordinates, atom_coordinate_offsets, grids)
658 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: grid_coordinates
659 : INTEGER, DIMENSION(:), INTENT(IN) :: atom_coordinate_offsets
660 : TYPE(rirs_grid_type), DIMENSION(:), INTENT(INOUT) :: grids
661 :
662 : INTEGER :: atom_coordinate_count, iatom
663 :
664 256 : CPASSERT(SIZE(atom_coordinate_offsets) == SIZE(grids))
665 1024 : DO iatom = 1, SIZE(grids)
666 2304 : atom_coordinate_count = SIZE(grids(iatom)%raw_points)
667 : grids(iatom)%raw_points(:, :) = RESHAPE( &
668 : grid_coordinates(atom_coordinate_offsets(iatom) + 1: &
669 : atom_coordinate_offsets(iatom) + &
670 : atom_coordinate_count), &
671 29512 : SHAPE(grids(iatom)%raw_points))
672 : END DO
673 256 : END SUBROUTINE unpack_lbfgs_grid_coordinates
674 :
675 : END MODULE ri_rs_grid_optimization
|
