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 Shared numerical operations for computing the RI-RS matrix Z_lP.
10 : !> \par History
11 : !> 09.2026 created Jan Wilhelm
12 : ! **************************************************************************************************
13 : MODULE ri_rs_fitting
14 : USE cp_blacs_env, ONLY: cp_blacs_env_type
15 : USE cp_dbcsr_api, ONLY: dbcsr_put_block,&
16 : dbcsr_type
17 : USE cp_fm_cholesky, ONLY: cp_fm_cholesky_decompose,&
18 : cp_fm_cholesky_solve
19 : USE cp_fm_struct, ONLY: cp_fm_struct_create,&
20 : cp_fm_struct_release,&
21 : cp_fm_struct_type
22 : USE cp_fm_types, ONLY: cp_fm_create,&
23 : cp_fm_get_info,&
24 : cp_fm_get_submatrix,&
25 : cp_fm_release,&
26 : cp_fm_set_submatrix,&
27 : cp_fm_type
28 : USE kinds, ONLY: dp
29 : USE message_passing, ONLY: mp_para_env_type
30 : #include "./base/base_uses.f90"
31 :
32 : IMPLICIT NONE
33 : PRIVATE
34 :
35 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'ri_rs_fitting'
36 : REAL(KIND=dp), PARAMETER, PRIVATE :: jacobi_floor = 1.0E-16_dp
37 :
38 : PUBLIC :: evaluate_ri_rs_fit, &
39 : solve_ri_rs_fit, &
40 : store_Z_lP_columns
41 :
42 : CONTAINS
43 :
44 : ! **************************************************************************************************
45 : !> \brief Computes the regularized RI-RS fit from an AO collocation matrix and a replicated RHS.
46 : !>
47 : !> D_ll' = [Σ_μ ϕ_μ(r_l) ϕ_μ(r_l')]²,
48 : !> d_l = 1/sqrt(max(D_ll, jacobi_floor)),
49 : !> D'_ll' = d_l D_ll' d_l' + λ δ_ll',
50 : !> Σ_l' D'_ll' Z'_l'P = d_l d_lP,
51 : !> Z_lP = d_l Z'_lP.
52 : !>
53 : !> The RI-RS least-squares construction and balancing are described by Duchemin and
54 : !> Blase, arXiv:1902.00066, Sec. II.3. Here jacobi_floor bounds D_ll from below in
55 : !> both the LAPACK and ScaLAPACK paths. The regularization λ is added after Jacobi
56 : !> scaling; in the unscaled system its diagonal contribution is λ/d_l².
57 : !> \param phi_local AO values on the fitting grid, replicated within the solve subgroup
58 : !> \param d_lp replicated RHS d_lP on entry and fitted coefficients Z_lP on return
59 : !> \param tikhonov regularization parameter λ in the Jacobi-scaled space
60 : !> \param para_env_sub optional subgroup communicator for the distributed solve
61 : !> \param blacs_env_sub optional BLACS context, supplied together with para_env_sub
62 : ! **************************************************************************************************
63 54 : SUBROUTINE solve_ri_rs_fit(phi_local, d_lp, tikhonov, para_env_sub, blacs_env_sub)
64 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi_local
65 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: d_lp
66 : REAL(KIND=dp), INTENT(IN) :: tikhonov
67 : TYPE(mp_para_env_type), OPTIONAL, POINTER :: para_env_sub
68 : TYPE(cp_blacs_env_type), OPTIONAL, POINTER :: blacs_env_sub
69 :
70 : CHARACTER(LEN=*), PARAMETER :: routineN = 'solve_ri_rs_fit'
71 :
72 : INTEGER :: handle, handle2, i, info, n_grid, n_rhs
73 54 : REAL(KIND=dp), ALLOCATABLE :: gram(:, :), jacobi_scale(:)
74 :
75 54 : CALL timeset(routineN, handle)
76 :
77 54 : n_grid = SIZE(phi_local, 1)
78 54 : n_rhs = SIZE(d_lp, 2)
79 54 : CPASSERT(SIZE(d_lp, 1) == n_grid)
80 54 : CPASSERT(PRESENT(para_env_sub) .EQV. PRESENT(blacs_env_sub))
81 54 : IF (PRESENT(para_env_sub)) THEN
82 : ! Only the row norms are replicated; the distributed solve builds its own Gram tiles.
83 0 : ALLOCATE (jacobi_scale(n_grid))
84 : !$OMP PARALLEL DO DEFAULT(NONE) &
85 : !$OMP SHARED(n_grid, phi_local, jacobi_scale) &
86 0 : !$OMP PRIVATE(i) SCHEDULE(STATIC)
87 : DO i = 1, n_grid
88 : jacobi_scale(i) = 1.0_dp/MAX(DOT_PRODUCT(phi_local(i, :), phi_local(i, :)), SQRT(jacobi_floor))
89 : END DO
90 : !$OMP END PARALLEL DO
91 : ELSE
92 54 : CALL factor_ri_rs_fit(phi_local, tikhonov, gram, jacobi_scale, info)
93 54 : IF (info /= 0) CPABORT('RI-RS Cholesky factorization failed')
94 : END IF
95 :
96 : ! Pre-scale the right-hand side: d_lP <- d_l d_lP.
97 54 : CALL scale_rows_by_diag(d_lp, jacobi_scale)
98 54 : IF (PRESENT(para_env_sub)) THEN
99 : CALL solve_D_lp_distributed(phi_local, jacobi_scale, d_lp, tikhonov, &
100 0 : para_env_sub, blacs_env_sub)
101 : ELSE
102 : ! Solve Σ_l' D'_ll' Z'_l'P = d_l d_lP, overwriting d_lp with Z'_lP.
103 54 : CALL timeset(routineN//'_dpotrs', handle2)
104 54 : CALL dpotrs('L', n_grid, n_rhs, gram, n_grid, d_lp, n_grid, info)
105 54 : CALL timestop(handle2)
106 54 : IF (info /= 0) CPABORT('RI-RS Cholesky solve failed')
107 : END IF
108 : ! Undo the conditioning on the solution: Z_lP = d_l Z'_lP.
109 54 : CALL scale_rows_by_diag(d_lp, jacobi_scale)
110 :
111 54 : CALL timestop(handle)
112 :
113 54 : END SUBROUTINE solve_ri_rs_fit
114 :
115 : ! **************************************************************************************************
116 : !> \brief Build and factor the Jacobi-scaled squared AO Gram matrix, K = d D d + λ I.
117 : !> The lower triangle is used throughout; the factor is retained for subsequent solves.
118 : !> See solve_ri_rs_fit for the RI-RS balancing convention and reference.
119 : !> \param phi AO values indexed (grid point, AO)
120 : !> \param tikhonov regularization in the Jacobi-scaled space
121 : !> \param factor lower Cholesky factor of K on successful return
122 : !> \param jacobi_scale pair-feature scales d_l
123 : !> \param info LAPACK factorization status; the caller chooses how to handle failure
124 : !> \param overlap optional lower triangle of the scaled AO overlap V V^T, V_l = sqrt(d_l) phi_l
125 : ! **************************************************************************************************
126 438 : SUBROUTINE factor_ri_rs_fit(phi, tikhonov, factor, jacobi_scale, info, overlap)
127 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi
128 : REAL(KIND=dp), INTENT(IN) :: tikhonov
129 : REAL(KIND=dp), ALLOCATABLE, INTENT(OUT) :: factor(:, :), jacobi_scale(:)
130 : INTEGER, INTENT(OUT) :: info
131 : REAL(KIND=dp), ALLOCATABLE, INTENT(OUT), OPTIONAL :: overlap(:, :)
132 :
133 : CHARACTER(LEN=*), PARAMETER :: routineN = 'factor_ri_rs_fit'
134 :
135 : INTEGER :: handle, handle2, i, j, n_grid
136 :
137 438 : CALL timeset(routineN, handle)
138 438 : n_grid = SIZE(phi, 1)
139 2628 : ALLOCATE (factor(n_grid, n_grid), jacobi_scale(n_grid))
140 : ! First form S_ll' = Σ_μ phi_lμ phi_l'μ; D_ll' = S_ll'^2.
141 438 : CALL timeset(routineN//'_dsyrk', handle2)
142 438 : CALL dsyrk('L', 'N', n_grid, SIZE(phi, 2), 1.0_dp, phi, n_grid, 0.0_dp, factor, n_grid)
143 438 : CALL timestop(handle2)
144 29152 : DO i = 1, n_grid
145 29152 : jacobi_scale(i) = 1.0_dp/MAX(factor(i, i), SQRT(jacobi_floor))
146 : END DO
147 438 : IF (PRESENT(overlap)) THEN
148 1476 : ALLOCATE (overlap(n_grid, n_grid))
149 9333 : DO j = 1, n_grid
150 228693 : DO i = j, n_grid
151 228324 : overlap(i, j) = factor(i, j)*SQRT(jacobi_scale(i))*SQRT(jacobi_scale(j))
152 : END DO
153 : END DO
154 : END IF
155 : ! Add λ after Jacobi scaling; LAPACK needs only the lower triangle.
156 : !$OMP PARALLEL DO DEFAULT(NONE) &
157 : !$OMP SHARED(n_grid, factor, jacobi_scale, tikhonov) &
158 438 : !$OMP PRIVATE(j, i) SCHEDULE(DYNAMIC)
159 : DO j = 1, n_grid
160 : DO i = j, n_grid
161 : factor(i, j) = factor(i, j)**2*jacobi_scale(i)*jacobi_scale(j)
162 : END DO
163 : factor(j, j) = factor(j, j) + tikhonov
164 : END DO
165 : !$OMP END PARALLEL DO
166 438 : CALL timeset(routineN//'_dpotrf', handle2)
167 438 : CALL dpotrf('L', n_grid, factor, n_grid, info)
168 438 : CALL timestop(handle2)
169 438 : CALL timestop(handle)
170 1314 : END SUBROUTINE factor_ri_rs_fit
171 :
172 : ! **************************************************************************************************
173 : !> \brief Evaluate the normalized cluster residual and its exact derivative in AO space.
174 : !>
175 : !> V_l = sqrt(d_l) phi_l, S = V V^T, K = S∘S + λI, and K Z' = b, with
176 : !> b_lP = V_l^T B_P V_l. B_P is the symmetric matrix defined by the upper AO triangle.
177 : !> The objective is ||B - A Z'||²/||B||², A_:l = svec(V_l V_l^T), without a penalty term.
178 : !> For λ>0 the derivative includes the response of the regularized coefficients.
179 : !> Each RI-column panel reuses V B_P for the RHS and gradient before releasing the panel.
180 : !> The balancing convention and RI-RS reference are given in solve_ri_rs_fit.
181 : !> \param phi unscaled AO values indexed (grid point, AO)
182 : !> \param int_3c three-center integrals indexed (AO, AO, RI); only the upper AO triangle is used
183 : !> \param tikhonov regularization in the Jacobi-scaled space
184 : !> \param normalized_error normalized squared residual
185 : !> \param fit_successful false if the target is empty or the regularized fit cannot be solved
186 : !> \param ao_gradient optional derivative with respect to the unscaled AO values
187 : !> \param max_abs_error optional largest absolute unweighted integral error
188 : !> \param workspace_budget_bytes optional total budget for evaluator-owned gradient workspaces;
189 : !> the caller deducts its own arrays first. Reserve all fixed workspaces before sizing
190 : !> the product cache. An absent or nonpositive budget streams one RI column.
191 : ! **************************************************************************************************
192 384 : SUBROUTINE evaluate_ri_rs_fit(phi, int_3c, tikhonov, normalized_error, fit_successful, &
193 384 : ao_gradient, max_abs_error, workspace_budget_bytes)
194 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi
195 : REAL(KIND=dp), DIMENSION(:, :, :), INTENT(IN) :: int_3c
196 : REAL(KIND=dp), INTENT(IN) :: tikhonov
197 : REAL(KIND=dp), INTENT(OUT) :: normalized_error
198 : LOGICAL, INTENT(OUT) :: fit_successful
199 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT), &
200 : OPTIONAL :: ao_gradient
201 : REAL(KIND=dp), INTENT(OUT), OPTIONAL :: max_abs_error
202 : REAL(KIND=dp), INTENT(IN), OPTIONAL :: workspace_budget_bytes
203 :
204 : CHARACTER(LEN=*), PARAMETER :: routineN = 'evaluate_ri_rs_fit'
205 : INTEGER, PARAMETER :: residual_block_size = 64
206 :
207 : INTEGER :: first, handle, handle2, i, icache, imu, &
208 : info, inu, ip, j, last, n_ao, &
209 : n_cached, n_grid, n_panel, n_ri, &
210 : panel_first, panel_last
211 : REAL(KIND=dp) :: coefficient_norm2, error_scale, fixed_bytes, integral_norm2, na, ng, &
212 : projection, residual_norm2, roundoff_estimate, row_norm2, value, weight
213 384 : REAL(KIND=dp), ALLOCATABLE :: coefficient_work(:, :), factor(:, :), jacobi_scale(:), &
214 384 : overlap(:, :), products(:, :, :), residual(:, :), v(:, :), z_prime(:, :)
215 :
216 384 : CALL timeset(routineN, handle)
217 384 : n_grid = SIZE(phi, 1)
218 384 : n_ao = SIZE(phi, 2)
219 384 : n_ri = SIZE(int_3c, 3)
220 384 : CPASSERT(SIZE(int_3c, 1) == n_ao .AND. SIZE(int_3c, 2) == n_ao)
221 384 : CPASSERT(tikhonov >= 0.0_dp)
222 384 : IF (PRESENT(ao_gradient)) THEN
223 1107 : CPASSERT(ALL(SHAPE(ao_gradient) == SHAPE(phi)))
224 65700 : ao_gradient(:, :) = 0.0_dp
225 : END IF
226 384 : normalized_error = HUGE(normalized_error)
227 384 : IF (PRESENT(max_abs_error)) max_abs_error = HUGE(max_abs_error)
228 384 : fit_successful = .FALSE.
229 384 : IF (n_grid < 1 .OR. n_ri < 1) THEN
230 0 : CALL timestop(handle)
231 0 : RETURN
232 : END IF
233 : integral_norm2 = 0.0_dp
234 81360 : DO ip = 1, n_ri
235 648192 : DO inu = 1, n_ao
236 : integral_norm2 = integral_norm2 + int_3c(inu, inu, ip)**2 + &
237 2348304 : 2.0_dp*DOT_PRODUCT(int_3c(1:inu - 1, inu, ip), int_3c(1:inu - 1, inu, ip))
238 : END DO
239 : END DO
240 384 : IF (integral_norm2 <= TINY(1.0_dp)) THEN
241 0 : CALL timestop(handle)
242 0 : RETURN
243 : END IF
244 :
245 384 : n_panel = n_ri
246 384 : n_cached = 1
247 384 : IF (PRESENT(ao_gradient)) THEN
248 369 : n_panel = 1
249 369 : IF (PRESENT(workspace_budget_bytes)) THEN
250 369 : IF (workspace_budget_bytes > 0.0_dp) THEN
251 369 : ng = REAL(n_grid, dp)
252 369 : na = REAL(n_ao, dp)
253 : ! Reserve factor, overlap, V, two coefficient arrays, Jacobi scales and the
254 : ! largest residual panel before assigning all remaining budget to V B_P.
255 : fixed_bytes = 8.0_dp*(2.0_dp*ng**2 + ng*na + 2.0_dp*ng*REAL(n_ri, dp) + ng + &
256 369 : na*MIN(na, REAL(residual_block_size, dp)))
257 : n_panel = INT(MIN(REAL(n_ri, dp), MAX(0.0_dp, &
258 369 : (workspace_budget_bytes - fixed_bytes)/(8.0_dp*ng*na))))
259 369 : IF (n_panel < 1) THEN
260 0 : CPABORT("RI-RS minimum fit workspace exceeds memory budget")
261 : END IF
262 : END IF
263 : END IF
264 369 : n_cached = n_panel
265 369 : CALL factor_ri_rs_fit(phi, tikhonov, factor, jacobi_scale, info, overlap)
266 : ELSE
267 15 : CALL factor_ri_rs_fit(phi, tikhonov, factor, jacobi_scale, info)
268 : END IF
269 384 : IF (info /= 0) THEN
270 0 : CALL timestop(handle)
271 0 : RETURN
272 : END IF
273 : ALLOCATE (v(n_grid, n_ao), products(n_grid, n_ao, n_cached), &
274 4992 : coefficient_work(n_grid, n_ri), z_prime(n_grid, n_ri))
275 3072 : DO imu = 1, n_ao
276 73821 : v(:, imu) = SQRT(jacobi_scale)*phi(:, imu)
277 : END DO
278 : ! K is independent of P. Solve a panel of RHSs and consume its cached products immediately.
279 : ! Keep Z' and the coefficient response for the final grid-space symmetric product.
280 384 : value = 0.0_dp
281 384 : coefficient_norm2 = 0.0_dp
282 384 : CALL timeset(routineN//'_ri_panels', handle2)
283 768 : DO panel_first = 1, n_ri, n_panel
284 384 : panel_last = MIN(panel_first + n_panel - 1, n_ri)
285 81360 : DO ip = panel_first, panel_last
286 80976 : icache = 1
287 80976 : IF (PRESENT(ao_gradient)) icache = ip - panel_first + 1
288 : CALL dsymm('R', 'U', n_grid, n_ao, 1.0_dp, int_3c(:, :, ip), n_ao, v, n_grid, &
289 80976 : 0.0_dp, products(:, :, icache), n_grid)
290 1975140 : DO i = 1, n_grid
291 15231216 : coefficient_work(i, ip) = DOT_PRODUCT(v(i, :), products(i, :, icache))
292 : END DO
293 : END DO
294 1975140 : z_prime(:, panel_first:panel_last) = coefficient_work(:, panel_first:panel_last)
295 : CALL dpotrs('L', n_grid, panel_last - panel_first + 1, factor, n_grid, &
296 384 : z_prime(:, panel_first:panel_last), n_grid, info)
297 384 : IF (info /= 0) THEN
298 0 : CALL timestop(handle2)
299 0 : CALL timestop(handle)
300 0 : RETURN
301 : END IF
302 81360 : DO ip = panel_first, panel_last
303 1974756 : value = value + DOT_PRODUCT(z_prime(:, ip), coefficient_work(:, ip))
304 1975140 : coefficient_norm2 = coefficient_norm2 + DOT_PRODUCT(z_prime(:, ip), z_prime(:, ip))
305 : END DO
306 :
307 768 : IF (PRESENT(ao_gradient)) THEN
308 : ! Reuse the RHS for Y = K^-1 Z' and then C = Z' + λY; λ=0 needs no response solve.
309 1908492 : coefficient_work(:, panel_first:panel_last) = z_prime(:, panel_first:panel_last)
310 369 : IF (tikhonov > 0.0_dp) THEN
311 : CALL dpotrs('L', n_grid, panel_last - panel_first + 1, factor, n_grid, &
312 369 : coefficient_work(:, panel_first:panel_last), n_grid, info)
313 369 : IF (info /= 0) THEN
314 0 : CALL timestop(handle2)
315 0 : CALL timestop(handle)
316 0 : RETURN
317 : END IF
318 : coefficient_work(:, panel_first:panel_last) = z_prime(:, panel_first:panel_last) + &
319 1908492 : tikhonov*coefficient_work(:, panel_first:panel_last)
320 : END IF
321 : ! Accumulate -Σ_P C_lP (V B_P)_l while the panel products are still available.
322 79728 : DO ip = panel_first, panel_last
323 79359 : icache = ip - panel_first + 1
324 635241 : DO imu = 1, n_ao
325 13436220 : ao_gradient(:, imu) = ao_gradient(:, imu) - coefficient_work(:, ip)*products(:, imu, icache)
326 : END DO
327 : END DO
328 369 : IF (tikhonov > 0.0_dp) THEN
329 : ! Retain T = C - Z'/2 = Z'/2 + λY for M = Z'T^T + TZ'^T below.
330 : coefficient_work(:, panel_first:panel_last) = coefficient_work(:, panel_first:panel_last) - &
331 1908492 : 0.5_dp*z_prime(:, panel_first:panel_last)
332 : END IF
333 : END IF
334 : END DO
335 384 : CALL timestop(handle2)
336 : ! K Z' = b implies ||B-AZ'||² = ||B||² - tr(Z'^T b) - λ ||Z'||².
337 384 : weight = tikhonov*coefficient_norm2
338 384 : residual_norm2 = integral_norm2 - value - weight
339 384 : error_scale = integral_norm2 + ABS(value) + weight
340 : ! Large coefficients can also amplify the factorization's backward error. The scaled
341 : ! Gram has diagonal <= 1, hence ||K||_2 <= n_grid + λ; use a conservative roundoff estimate.
342 : roundoff_estimate = 8.0_dp*EPSILON(1.0_dp)*REAL(n_grid, dp)* &
343 384 : (REAL(n_grid, dp) + tikhonov)*coefficient_norm2
344 : ! Near an exact fit, use bounded AO-column panels to avoid cancellation. This path also
345 : ! supplies the final maximum-error diagnostic, without storing a pair-by-RI residual.
346 384 : IF (PRESENT(max_abs_error) .OR. &
347 : residual_norm2 <= MAX(SQRT(EPSILON(1.0_dp))*error_scale, roundoff_estimate)) THEN
348 45 : CALL timeset(routineN//'_residual', handle2)
349 180 : ALLOCATE (residual(n_ao, MIN(residual_block_size, n_ao)))
350 45 : residual_norm2 = 0.0_dp
351 45 : IF (PRESENT(max_abs_error)) max_abs_error = 0.0_dp
352 8151 : DO ip = 1, n_ri
353 16257 : DO first = 1, n_ao, residual_block_size
354 8106 : last = MIN(first + residual_block_size - 1, n_ao)
355 64848 : DO inu = first, last
356 2532684 : products(:, inu - first + 1, 1) = z_prime(:, ip)*v(:, inu)
357 : END DO
358 : CALL dgemm('T', 'N', n_ao, last - first + 1, n_grid, -1.0_dp, v, n_grid, products(:, :, 1), n_grid, &
359 8106 : 0.0_dp, residual, n_ao)
360 72954 : DO inu = first, last
361 291816 : DO imu = 1, inu
362 226968 : value = int_3c(imu, inu, ip) + residual(imu, inu - first + 1)
363 226968 : weight = REAL(2 - MERGE(1, 0, imu == inu), dp)
364 226968 : residual_norm2 = residual_norm2 + weight*value**2
365 283710 : IF (PRESENT(max_abs_error)) max_abs_error = MAX(max_abs_error, ABS(value))
366 : END DO
367 : END DO
368 : END DO
369 : END DO
370 90 : CALL timestop(handle2)
371 : END IF
372 384 : DEALLOCATE (products)
373 384 : normalized_error = residual_norm2/integral_norm2
374 :
375 384 : IF (PRESENT(ao_gradient)) THEN
376 369 : CALL timeset(routineN//'_gradient', handle2)
377 : ! All panel solves are complete. Reuse the factor for M = Z'Z'^T + λ(Z'Y^T + YZ'^T).
378 369 : IF (tikhonov > 0.0_dp) THEN
379 : CALL dsyr2k('L', 'N', n_grid, n_ri, 1.0_dp, z_prime, n_grid, coefficient_work, n_grid, &
380 369 : 0.0_dp, factor, n_grid)
381 : ELSE
382 0 : CALL dsyrk('L', 'N', n_grid, n_ri, 1.0_dp, z_prime, n_grid, 0.0_dp, factor, n_grid)
383 : END IF
384 9333 : DO j = 1, n_grid
385 228693 : DO i = j, n_grid
386 228324 : factor(i, j) = factor(i, j)*overlap(i, j)
387 : END DO
388 : END DO
389 : ! Add (M∘S)V to the accumulated negative term, then normalize the AO gradient.
390 : CALL dsymm('L', 'L', n_grid, n_ao, 1.0_dp, factor, n_grid, v, n_grid, &
391 369 : 1.0_dp, ao_gradient, n_grid)
392 65700 : ao_gradient(:, :) = (4.0_dp/integral_norm2)*ao_gradient
393 : ! V_l = sqrt(d_l) phi_l. Differentiate the normalization only above the floor;
394 : ! on the clamped branch d_l is constant, including at the branch boundary.
395 9333 : DO i = 1, n_grid
396 71712 : row_norm2 = DOT_PRODUCT(phi(i, :), phi(i, :))
397 8964 : IF (row_norm2 > SQRT(jacobi_floor)) THEN
398 69000 : projection = DOT_PRODUCT(phi(i, :), ao_gradient(i, :))/row_norm2
399 69000 : ao_gradient(i, :) = ao_gradient(i, :) - projection*phi(i, :)
400 : END IF
401 72081 : ao_gradient(i, :) = SQRT(jacobi_scale(i))*ao_gradient(i, :)
402 : END DO
403 369 : CALL timestop(handle2)
404 : END IF
405 384 : fit_successful = .TRUE.
406 384 : CALL timestop(handle)
407 1152 : END SUBROUTINE evaluate_ri_rs_fit
408 :
409 : ! **************************************************************************************************
410 : !> \brief Multiplies every matrix row by the corresponding diagonal entry:
411 : !> A(l, :) <- d_l A(l, :).
412 : !> \param matrix ...
413 : !> \param diagonal ...
414 : ! **************************************************************************************************
415 108 : SUBROUTINE scale_rows_by_diag(matrix, diagonal)
416 :
417 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: matrix
418 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: diagonal
419 :
420 : CHARACTER(LEN=*), PARAMETER :: routineN = 'scale_rows_by_diag'
421 :
422 : INTEGER :: handle, i, j
423 :
424 108 : CALL timeset(routineN, handle)
425 108 : CPASSERT(SIZE(diagonal) == SIZE(matrix, 1))
426 :
427 : !$OMP PARALLEL DO DEFAULT(NONE) &
428 : !$OMP SHARED(matrix, diagonal) &
429 : !$OMP PRIVATE(j, i) &
430 108 : !$OMP SCHEDULE(STATIC)
431 : DO j = 1, SIZE(matrix, 2)
432 : DO i = 1, SIZE(matrix, 1)
433 : matrix(i, j) = matrix(i, j)*diagonal(i)
434 : END DO
435 : END DO
436 : !$OMP END PARALLEL DO
437 :
438 108 : CALL timestop(handle)
439 :
440 108 : END SUBROUTINE scale_rows_by_diag
441 :
442 : ! **************************************************************************************************
443 : !> \brief Stores dense Z_lP columns in the distributed block-sparse matrix.
444 : !> \param mat_Z_lP ...
445 : !> \param Z_local ...
446 : !> \param local_grid_idx ...
447 : !> \param n_local_grid ...
448 : !> \param n_loc_ri ...
449 : !> \param atom_P ...
450 : !> \param r_blk_sizes ...
451 : !> \param row_offset ...
452 : !> \param eps_filter ...
453 : ! **************************************************************************************************
454 50 : SUBROUTINE store_Z_lP_columns(mat_Z_lP, Z_local, local_grid_idx, n_local_grid, n_loc_ri, &
455 50 : atom_P, r_blk_sizes, row_offset, eps_filter)
456 :
457 : TYPE(dbcsr_type), INTENT(INOUT) :: mat_Z_lP
458 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: Z_local
459 : INTEGER, DIMENSION(:), INTENT(IN) :: local_grid_idx
460 : INTEGER, INTENT(IN) :: n_local_grid, n_loc_ri, atom_P
461 : INTEGER, DIMENSION(:), INTENT(IN) :: r_blk_sizes, row_offset
462 : REAL(KIND=dp), INTENT(IN) :: eps_filter
463 :
464 : CHARACTER(LEN=*), PARAMETER :: routineN = 'store_Z_lP_columns'
465 :
466 : INTEGER :: current_chunk_size, g_pt, handle, i_blk, &
467 : loc_ptr, r_end, r_start
468 50 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: Z_blk
469 :
470 50 : CALL timeset(routineN, handle)
471 :
472 882 : ALLOCATE (Z_blk(MAXVAL(r_blk_sizes), n_loc_ri))
473 : ! local_grid_idx is ascending, so one walking pointer maps sphere rows to global grid blocks.
474 50 : loc_ptr = 1
475 :
476 732 : DO i_blk = 1, SIZE(r_blk_sizes)
477 682 : r_start = row_offset(i_blk) + 1
478 682 : r_end = row_offset(i_blk) + r_blk_sizes(i_blk)
479 682 : current_chunk_size = r_blk_sizes(i_blk)
480 682 : Z_blk = 0.0_dp
481 :
482 19163 : DO WHILE (loc_ptr <= n_local_grid)
483 19110 : g_pt = local_grid_idx(loc_ptr)
484 19110 : IF (g_pt > r_end) EXIT
485 258312 : Z_blk(g_pt - r_start + 1, 1:n_loc_ri) = Z_local(loc_ptr, 1:n_loc_ri)
486 19110 : loc_ptr = loc_ptr + 1
487 : END DO
488 :
489 : ! DBCSR routes each retained block to its global owner when the caller finalizes mat_Z_lP.
490 252928 : IF (MAXVAL(ABS(Z_blk(1:current_chunk_size, 1:n_loc_ri))) > eps_filter) THEN
491 : CALL dbcsr_put_block(mat_Z_lP, row=i_blk, col=atom_P, &
492 673 : block=Z_blk(1:current_chunk_size, 1:n_loc_ri))
493 : END IF
494 : END DO
495 :
496 50 : DEALLOCATE (Z_blk)
497 :
498 50 : CALL timestop(handle)
499 :
500 50 : END SUBROUTINE store_Z_lP_columns
501 :
502 : ! **************************************************************************************************
503 : !> \brief Solves the Jacobi-scaled, regularized RI-RS system with ScaLAPACK.
504 : !>
505 : !> Each rank builds its block-cyclic part of
506 : !>
507 : !> D'_ll' = d_l [Σ_μ ϕ_μ(r_l) ϕ_μ(r_l')]² d_l' + λ δ_ll'.
508 : !>
509 : !> The Cholesky solution Z'_lP is gathered into the replicated right-hand side.
510 : !> The caller then restores Z_lP = d_l Z'_lP.
511 : !> \param phi_local replicated AO values on the fitting grid
512 : !> \param d_vec Jacobi scaling vector d_l, with the same floor as in the dense solve
513 : !> \param d_lp replicated scaled RHS d_l d_lP on entry, intermediate solution Z'_lP on return
514 : !> \param tikhonov ...
515 : !> \param para_env_sub ...
516 : !> \param blacs_env_sub ...
517 : ! **************************************************************************************************
518 0 : SUBROUTINE solve_D_lp_distributed(phi_local, d_vec, d_lp, tikhonov, para_env_sub, blacs_env_sub)
519 :
520 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: phi_local
521 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: d_vec
522 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: d_lp
523 : REAL(KIND=dp), INTENT(IN) :: tikhonov
524 : TYPE(mp_para_env_type), POINTER :: para_env_sub
525 : TYPE(cp_blacs_env_type), POINTER :: blacs_env_sub
526 :
527 : CHARACTER(LEN=*), PARAMETER :: routineN = 'solve_D_lp_distributed'
528 :
529 : INTEGER :: handle, i_loc, ig, info, j_loc, jg, &
530 : n_ao, n_loc, n_rhs, ncol_local, &
531 : nrow_local
532 0 : INTEGER, DIMENSION(:), POINTER :: col_indices, row_indices
533 : REAL(KIND=dp), CONTIGUOUS, DIMENSION(:, :), &
534 0 : POINTER :: local_data
535 : TYPE(cp_fm_struct_type), POINTER :: fm_struct_b, fm_struct_D
536 : TYPE(cp_fm_type) :: fm_b, fm_D
537 :
538 0 : CALL timeset(routineN, handle)
539 0 : n_loc = SIZE(phi_local, 1)
540 0 : n_ao = SIZE(phi_local, 2)
541 0 : n_rhs = SIZE(d_lp, 2)
542 :
543 0 : NULLIFY (fm_struct_D, fm_struct_b)
544 : CALL cp_fm_struct_create(fm_struct_D, para_env=para_env_sub, &
545 : context=blacs_env_sub, &
546 0 : nrow_global=n_loc, ncol_global=n_loc)
547 : CALL cp_fm_struct_create(fm_struct_b, para_env=para_env_sub, &
548 : context=blacs_env_sub, &
549 0 : nrow_global=n_loc, ncol_global=n_rhs)
550 0 : CALL cp_fm_create(fm_D, fm_struct_D)
551 0 : CALL cp_fm_create(fm_b, fm_struct_b)
552 :
553 : CALL cp_fm_get_info(fm_D, nrow_local=nrow_local, ncol_local=ncol_local, &
554 : row_indices=row_indices, col_indices=col_indices, &
555 0 : local_data=local_data)
556 :
557 : ! Build only the rank-owned Gram tiles, avoiding a replicated n_grid × n_grid matrix.
558 0 : IF (nrow_local > 0 .AND. ncol_local > 0) THEN
559 0 : BLOCK
560 : INTEGER, PARAMETER :: ntile = 1024
561 : INTEGER :: ib, ie, jb, je, mb, kb, ti, tj, handle_dgemm
562 0 : REAL(KIND=dp), ALLOCATABLE :: gram_t(:, :), phi_cols_t(:, :), phi_rows_t(:, :)
563 0 : ALLOCATE (phi_rows_t(ntile, n_ao), phi_cols_t(n_ao, ntile), gram_t(ntile, ntile))
564 : ! Pack AO values for the grid rows/columns selected by the block-cyclic FM layout.
565 0 : DO ib = 1, nrow_local, ntile
566 0 : ie = MIN(ib + ntile - 1, nrow_local)
567 0 : mb = ie - ib + 1
568 : !$OMP PARALLEL DO DEFAULT(NONE) &
569 : !$OMP SHARED(mb, n_ao, phi_rows_t, phi_local, row_indices, ib) &
570 0 : !$OMP PRIVATE(ti, j_loc) SCHEDULE(STATIC)
571 : DO j_loc = 1, n_ao
572 : DO ti = 1, mb
573 : phi_rows_t(ti, j_loc) = phi_local(row_indices(ib + ti - 1), j_loc)
574 : END DO
575 : END DO
576 : !$OMP END PARALLEL DO
577 0 : DO jb = 1, ncol_local, ntile
578 0 : je = MIN(jb + ntile - 1, ncol_local)
579 0 : kb = je - jb + 1
580 : !$OMP PARALLEL DO DEFAULT(NONE) &
581 : !$OMP SHARED(kb, n_ao, phi_cols_t, phi_local, col_indices, jb) &
582 0 : !$OMP PRIVATE(tj, i_loc) SCHEDULE(STATIC)
583 : DO tj = 1, kb
584 : DO i_loc = 1, n_ao
585 : phi_cols_t(i_loc, tj) = phi_local(col_indices(jb + tj - 1), i_loc)
586 : END DO
587 : END DO
588 : !$OMP END PARALLEL DO
589 0 : CALL timeset(routineN//'_dgemm', handle_dgemm)
590 : CALL dgemm('N', 'N', mb, kb, n_ao, &
591 : 1.0_dp, phi_rows_t, ntile, phi_cols_t, n_ao, &
592 0 : 0.0_dp, gram_t, ntile)
593 0 : CALL timestop(handle_dgemm)
594 : ! Square the AO Gram entries, apply d_l d_l', and add λ on the global diagonal.
595 : !$OMP PARALLEL DO DEFAULT(NONE) &
596 : !$OMP SHARED(mb, kb, gram_t, d_vec, row_indices, col_indices, ib, jb) &
597 : !$OMP SHARED(local_data, tikhonov) &
598 0 : !$OMP PRIVATE(ti, tj, ig, jg) SCHEDULE(STATIC)
599 : DO tj = 1, kb
600 : jg = col_indices(jb + tj - 1)
601 : DO ti = 1, mb
602 : ig = row_indices(ib + ti - 1)
603 : local_data(ib + ti - 1, jb + tj - 1) = &
604 : gram_t(ti, tj)*gram_t(ti, tj)*d_vec(ig)*d_vec(jg)
605 : IF (ig == jg) THEN
606 : local_data(ib + ti - 1, jb + tj - 1) = &
607 : local_data(ib + ti - 1, jb + tj - 1) + tikhonov
608 : END IF
609 : END DO
610 : END DO
611 : !$OMP END PARALLEL DO
612 : END DO
613 : END DO
614 0 : DEALLOCATE (phi_rows_t, phi_cols_t, gram_t)
615 : END BLOCK
616 : END IF
617 :
618 : ! Distribute the pre-scaled RHS and solve for Z'_lP within the BLACS subgroup.
619 0 : CALL cp_fm_set_submatrix(fm_b, d_lp)
620 :
621 0 : CALL cp_fm_cholesky_decompose(fm_D, n=n_loc, info_out=info)
622 0 : IF (info /= 0) CPABORT("pdpotrf failed in solve_D_lp_distributed")
623 :
624 0 : CALL cp_fm_cholesky_solve(fm_D, fm_b, n=n_loc, info_out=info)
625 0 : IF (info /= 0) CPABORT("pdpotrs failed in solve_D_lp_distributed")
626 :
627 : ! Replicate Z'_lP within the subgroup; the caller applies the final Jacobi rescaling.
628 0 : CALL cp_fm_get_submatrix(fm_b, d_lp)
629 :
630 0 : CALL cp_fm_release(fm_D)
631 0 : CALL cp_fm_release(fm_b)
632 0 : CALL cp_fm_struct_release(fm_struct_D)
633 0 : CALL cp_fm_struct_release(fm_struct_b)
634 :
635 0 : CALL timestop(handle)
636 :
637 0 : END SUBROUTINE solve_D_lp_distributed
638 :
639 : END MODULE ri_rs_fitting
|
