Coverage for C:\src\imod-python\imod\prepare\pcg.py: 93%
343 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-16 11:41 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-16 11:41 +0200
1"""Preconditioned Conjugate Gradient Solver Package.
3This is a translation of the fortran77 modflow2000 pcg2.f module.
4"""
6import numba
9class PreconditionedConjugateGradientSolver(object):
10 __slots__ = ("hclose", "rclose", "mxiter", "iter1", "relax")
12 def __init__(self, hclose, rclose, mxiter, iter1, relax=1.0):
13 """
14 hclose : float
15 Head tolerance for stopping condition
16 rclose : float
17 Residual tolerance for stopping condition
18 iter1 : integer
19 Maximum number of inner iterations
20 mxiter : integer
21 Maximum number of outer iterations
22 relax : float
23 Relaxation parameter for incomplete Cholesky preconditioning
24 damp : float
25 Dampening parameter for head change of outer iterations
26 """
27 self.hclose = hclose
28 self.rclose = rclose
29 self.mxiter = mxiter
30 self.iter1 = iter1
31 self.relax = relax
33 def solve(self, hnew, cc, cr, cv, ibound, rhs, hcof, res, cd, v, ss, p):
34 """
35 Solves the groundwater equations via the method of conjugated gradients,
36 with incomplete Cholesky preconditioning.
38 Mutates hnew, converged, cr, cc, cv, ibound, and the work arrays.
40 Parameters
41 ----------
42 hnew : ndarray (3D) of floats
43 Initial guess, will be updated in place.
45 The linear operator exists only virtually, and is formed by cr, cc, cv,
46 ibound.
48 cr : ndarray (3D) of floats
49 Row conductance.
50 cc : ndarray (3D) of floats
51 Column conductance
52 cv : ndarray (3D) of floats
53 Vertical conductance (layer)
54 ibound : ndarray (3D) of integers
55 Active (1), inactive (0), fixed head (-1)
56 rhs : ndarray (3D) of floats
57 Right-hand side
58 hcof : ndarray (3D)
59 Work array, coefficient of head
60 res : ndarray (1D)
61 Work array, residual
62 cd : ndarray (1D)
63 Work array, Incomplete Cholesky diagonal
64 v : ndarray (1D)
65 Work array, intermediate solution
66 ss : ndarray (1D)
67 Work array
68 p : ndarray (1D)
69 Work array
71 Returns
72 -------
73 converged : boolean
74 convergence flag
75 """
77 # Initialize
78 nlay, nrow, ncol = ibound.shape
79 nrc = nrow * ncol
80 nodes = nrc * nlay
81 srnew = 0.0
82 iiter = 0
83 converged = False
84 pcg_converged = False
86 # Ravel 3D arrays
87 hnew = hnew.ravel()
88 cc = cc.ravel()
89 cr = cr.ravel()
90 cv = cv.ravel()
91 ibound = ibound.ravel()
92 rhs = rhs.ravel()
93 hcof = hcof.ravel()
95 # Clear work arrays
96 res[:] = 0.0
97 cd[:] = 0.0
98 v[:] = 0.0
99 ss[:] = 0.0
100 p[:] = 0.0
102 # Calculate residual
103 # Will also get rid of dry cells in ibound, cc, cr, cv
104 ibound, cc, cr, cv, hcof, res = calculate_residual(
105 ibound, nlay, nrow, ncol, nrc, nodes, cr, cc, cv, hnew, rhs, hcof, res
106 )
108 while not pcg_converged and iiter < self.iter1:
109 # Start internal iterations
110 iiter = iiter + 1
111 # Initialize variables that track maximum head change
112 # and residual value during each iteration
113 bigh = 0.0
114 bigr = 0.0
116 cd, v = precondition_incomplete_cholesky(
117 ibound,
118 nlay,
119 nrow,
120 ncol,
121 nrc,
122 cc,
123 cr,
124 cv,
125 iiter,
126 hcof,
127 self.relax,
128 res,
129 cd,
130 v,
131 )
133 ss = back_substition(ibound, nlay, nrow, ncol, nrc, cr, cc, cv, v, cd, ss)
135 p, srnew = calculate_cg_p(ibound, nodes, ss, res, srnew, iiter, p)
137 v, alpha = calculate_cg_alpha(
138 ibound,
139 nlay,
140 nrow,
141 ncol,
142 nrc,
143 nodes,
144 v,
145 cc,
146 cr,
147 cv,
148 p,
149 hcof,
150 srnew,
151 self.mxiter,
152 )
154 hnew, res, bigh, bigr, indices_bigh, indices_bigr = calculate_heads(
155 ibound, nlay, nrow, ncol, nrc, alpha, p, bigh, bigr, hnew, res, v
156 )
158 # check the convergence criterion
159 if abs(bigh) <= self.hclose and abs(bigr) <= self.rclose:
160 pcg_converged = True
161 if iiter == 1 and pcg_converged:
162 converged = True
164 return converged
167@numba.njit
168def calculate_residual(
169 ibound, nlay, nrow, ncol, nrc, nodes, cr, cc, cv, hnew, rhs, hcof, res
170):
171 """
172 Calculate the residual.
174 For a dense form:
175 Residual = np.dot(A, x) - b
177 Parameters
178 ----------
179 ibound
180 nlay
181 nrow
182 ncol
183 nrc
184 nodes
185 cr
186 cc
187 cv
188 hnew
189 rhs
190 hcof : array[nodes]
192 Returns
193 -------
194 ibound
195 ibound with cells remove that are surrounded by inactive cells
196 hcof
197 coefficient of head
198 res
199 residuals
200 """
201 for k in range(nlay):
202 for i in range(nrow):
203 for j in range(ncol):
204 # Calculate 1 dimensional subscript of current cell
205 # and skip calculcations if cell is inactive
206 n = j + i * ncol + k * nrc
207 if ibound[n] == 0:
208 cc[n] = 0.0
209 cr[n] = 0.0
210 if n <= (nodes - nrc - 1):
211 cv[n] = 0.0
212 if n >= 1:
213 cr[n - 1] = 0.0
214 if n >= ncol - 1:
215 cc[n - ncol] = 0.0
216 if n <= (nodes - nrc - 1) and n >= nrc:
217 cv[n - nrc] = 0.0
218 continue
219 # Calculate 1 dimensional subscript for locating
220 # the 6 surrounding cells
221 nrn = n + ncol
222 nrl = n - ncol
223 ncn = n + 1
224 ncl = n - 1
225 nln = n + nrc
226 nll = n - nrc
227 # Calculate 1 dimensional subscripts for conductance to
228 # the 6 surrounding cells
229 ncf = n
230 ncd = n - 1
231 nrb = n - ncol
232 nrh = n
233 nls = n
234 nlz = n - nrc
235 # Get conductances to neighboring cells
236 # neighbor is 1 row back
237 b = 0.0
238 bhnew = 0.0
239 if i != 0:
240 b = cc[nrb]
241 bhnew = b * (hnew[nrl] - hnew[n])
242 # neighbour is 1 row ahead
243 h = 0.0
244 hhnew = 0.0
245 if i != nrow - 1:
246 h = cc[nrh]
247 hhnew = h * (hnew[nrn] - hnew[n])
248 # neighbour is 1 column back
249 d = 0.0
250 dhnew = 0.0
251 if j != 0:
252 d = cr[ncd]
253 dhnew = d * (hnew[ncl] - hnew[n])
254 # neighbour is 1 column ahead
255 f = 0.0
256 fhnew = 0.0
257 if j != ncol - 1:
258 f = cr[ncf]
259 fhnew = f * (hnew[ncn] - hnew[n])
260 # neighbour is 1 layer behind
261 z = 0.0
262 zhnew = 0.0
263 if k != 0:
264 z = cv[nlz]
265 zhnew = z * (hnew[nll] - hnew[n])
266 # neighbour is 1 layer ahead
267 s = 0.0
268 shnew = 0.0
269 if k != nlay - 1:
270 s = cv[nls]
271 shnew = s * (hnew[nln] - hnew[n])
273 if i == nrow - 1:
274 cc[n] = 0.0
275 if j == ncol - 1:
276 cr[n] = 0.0
278 # Skip calculations and make cell inactive if all
279 # surrounding cells are inactive
280 if ibound[n] == 1:
281 if (b + h + d + f + z + s) == 0.0:
282 ibound[n] = 0.0
283 hcof[n] = 0.0
284 rhs[n] = 0.0
285 continue
287 # Calculate the residual and store it in res.
288 rrhs = rhs[n]
289 hhcof = hnew[n] * hcof[n]
290 res[n] = rrhs - zhnew - bhnew - dhnew - hhcof - fhnew - hhnew - shnew
291 if ibound[n] < 0:
292 res[n] = 0.0
293 return ibound, cc, cr, cv, hcof, res
296@numba.njit
297def cholesky_diagonal(
298 ir,
299 ic,
300 il,
301 i,
302 j,
303 k,
304 n,
305 f,
306 h,
307 s,
308 ncol,
309 nrow,
310 nlay,
311 cc,
312 cr,
313 cv,
314 hcof,
315 relax,
316 delta,
317 cd1,
318 cd,
319):
320 """
321 Calculate one value of cd
322 first interal iteration only
324 Returns
325 -------
326 delta
327 cd1
328 cd
329 """
331 cdcr = 0.0
332 cdcc = 0.0
333 cdcv = 0.0
334 fcc = 0.0
335 fcr = 0.0
336 fcv = 0.0
337 if ir >= 0 and cd[ir] != 0.0:
338 cdcr = (f**2.0) / cd[ir]
339 if ic >= 0 and cd[ic] != 0.0:
340 cdcc = (h**2.0) / cd[ic]
341 if il >= 0 and cd[il] != 0.0:
342 cdcv = (s**2.0) / cd[il]
344 if ir >= 0:
345 fv = cv[ir]
346 if k == nlay - 1:
347 fv = 0.0
348 if cd[ir] != 0.0:
349 fcr = (f / cd[ir]) * (cc[ir] + fv)
351 if ic >= 0:
352 fv = cv[ic]
353 if k == nlay - 1 and i > 0:
354 fv = 0.0
355 if cd[ic] != 0.0:
356 fcc = (h / cd[ic]) * (cr[ic] + fv)
357 if il >= 0:
358 if cd[il] != 0.0:
359 fcv = (s / cd[il]) * (cr[il] + cc[il])
361 b = 0.0
362 h = 0.0
363 d = 0.0
364 f = 0.0
365 z = 0.0
366 s = 0.0
367 if i != 0:
368 b = cc[ic]
369 if i != nrow:
370 h = cc[n]
371 if j != 0:
372 d = cr[ir]
373 if j != ncol:
374 f = cr[n]
375 if k != 0:
376 z = cv[il]
377 if k != nlay:
378 s = cv[n]
380 hhcof = hcof[n] - z - b - d - f - h - s
381 cd[n] = (1.0 + delta) * hhcof - cdcr - cdcc - cdcv - relax * (fcr + fcc + fcv)
383 if cd1 == 0.0 and cd[n] != 0.0:
384 cd1 = cd[n]
385 if (cd[n] * cd1) < 0:
386 delta = 1.5 * delta + 0.001
387 raise RuntimeError(
388 "Matrix is severely non-diagonally dominant. Check input. Stopping execution."
389 )
390 return cd, cd1, delta
393@numba.njit
394def precondition_incomplete_cholesky(
395 ibound, nlay, nrow, ncol, nrc, cc, cr, cv, iiter, hcof, relax, res, cd, v
396):
397 """
398 Incomplete Cholesky preconditioning
399 Step through cells to calculate the diagonal of the cholesky
400 matrix (first internal iteration only) and the intermediate
401 solution. Store them in cd and v, respectively.
403 Returns
404 -------
405 cd
406 Cholesky diagonal
407 v
408 intermediate solution
409 """
410 delta = 0.0
411 cd1 = 0.0
412 for k in range(nlay):
413 for i in range(nrow):
414 for j in range(ncol):
415 n = j + i * ncol + k * nrc
416 if ibound[n] < 1:
417 continue
419 # Calculate v
420 h = 0.0
421 vcc = 0.0
422 ic = n - ncol # won't be used if i == 0
423 if i != 0:
424 h = cc[ic]
425 if cd[ic] != 0.0:
426 vcc = h * v[ic] / cd[ic]
428 f = 0.0
429 vcr = 0.0
430 ir = n - 1 # wont'be used if j == 0
431 if j != 0:
432 f = cr[ir]
433 if cd[ir] != 0.0:
434 vcr = f * v[ir] / cd[ir]
436 s = 0.0
437 vcv = 0.0
438 il = n - nrc # won't be used if k == 0
439 if k != 0:
440 s = cv[il]
441 if cd[il] != 0.0:
442 vcv = s * v[il] / cd[il]
444 v[n] = res[n] - vcr - vcc - vcv
446 # Calculate Cholesky diagonal for the first internal iteration.
447 if iiter == 1:
448 cd, cd1, delta = cholesky_diagonal(
449 ir,
450 ic,
451 il,
452 i,
453 j,
454 k,
455 n,
456 f,
457 h,
458 s,
459 ncol,
460 nrow,
461 nlay,
462 cc,
463 cr,
464 cv,
465 hcof,
466 relax,
467 delta,
468 cd1,
469 cd,
470 )
472 return cd, v
475@numba.njit
476def back_substition(ibound, nlay, nrow, ncol, nrc, cr, cc, cv, v, cd, ss):
477 """
478 Step through each cell and solve for s of the conjugate
479 gradient algorithm by back substition. Store the result in ss.
481 Returns
482 -------
483 ss
484 """
485 for kk in range(nlay - 1, -1, -1):
486 for ii in range(nrow - 1, -1, -1):
487 for jj in range(ncol - 1, -1, -1):
488 n = jj + ii * ncol + kk * nrc
489 if ibound[n] <= 0:
490 continue
491 nc = n + 1
492 nr = n + ncol
493 nl = n + nrc
494 sscr = 0.0
495 sscc = 0.0
496 sscv = 0.0
497 if jj != ncol - 1:
498 sscr = cr[n] * ss[nc] / cd[n]
499 if ii != nrow - 1:
500 sscc = cc[n] * ss[nr] / cd[n]
501 if kk != nlay - 1:
502 sscv = cv[n] * ss[nl] / cd[n]
503 vn = v[n] / cd[n]
504 ss[n] = vn - sscr - sscc - sscv
505 return ss
508@numba.njit
509def calculate_cg_p(ibound, nodes, ss, res, srnew, iiter, p):
510 """
511 Calculate p of the conjugate gradient algorithm
513 Returns
514 -------
515 p
516 srnew
517 """
518 srold = srnew
519 srnew = 0.0
520 for n in range(nodes):
521 if ibound[n] <= 0:
522 continue
523 srnew = srnew + ss[n] * res[n]
524 if iiter == 1:
525 for n in range(nodes):
526 p[n] = ss[n]
527 else:
528 for n in range(nodes):
529 p[n] = ss[n] + (srnew / srold) * p[n]
530 return p, srnew
533@numba.njit
534def calculate_cg_alpha(
535 ibound, nlay, nrow, ncol, nrc, nodes, v, cc, cr, cv, p, hcof, srnew, mxiter
536):
537 """
538 Calculate alpha of the conjugate routine.
539 For the denominator of alpha, multiply the matrix a by the
540 vector p, and store in v; then multiply p by v. Store in pap.
542 Returns
543 -------
544 v
545 intermediate solution
546 alpha
547 conjugate gradient alpha
548 """
549 pap = 0.0
550 for k in range(nlay):
551 for i in range(nrow):
552 for j in range(ncol):
553 n = j + i * ncol + k * nrc
554 v[n] = 0.0
555 if ibound[n] < 1:
556 continue
558 nrn = n + ncol
559 nrl = n - ncol
560 ncn = n + 1
561 ncl = n - 1
562 nln = n + nrc
563 nll = n - nrc
565 ncf = n
566 ncd = ncl
567 nrb = nrl
568 nrh = n
569 nls = n
570 nlz = nll
572 b = 0.0
573 if i != 0:
574 b = cc[nrb]
575 h = 0.0
576 if i != nrow - 1:
577 h = cc[nrh]
578 d = 0.0
579 if j != 0:
580 d = cr[ncd]
581 f = 0.0
582 if j != ncol - 1:
583 f = cr[ncf]
584 z = 0.0
585 if k != 0:
586 z = cv[nlz]
587 s = 0.0
588 if k != nlay - 1:
589 s = cv[nls]
591 pn = p[n]
593 bhnew = 0.0
594 hhnew = 0.0
595 dhnew = 0.0
596 fhnew = 0.0
597 zhnew = 0.0
598 shnew = 0.0
599 if nrl >= 0:
600 bhnew = b * (p[nrl] - pn)
601 if nrn <= nodes - 1:
602 hhnew = h * (p[nrn] - pn)
603 if ncl >= 0:
604 dhnew = d * (p[ncl] - pn)
605 if ncn <= nodes - 1:
606 fhnew = f * (p[ncn] - pn)
607 if nll >= 0:
608 zhnew = z * (p[nll] - pn)
609 if nln <= nodes - 1:
610 shnew = s * (p[nln] - pn)
612 # Calculate the product of a matrix and vector p and store
613 # result in v.
614 pn = hcof[n] * p[n]
615 vn = zhnew + bhnew + dhnew + pn + fhnew + hhnew + shnew
616 v[n] = vn
617 pap = pap + p[n] * vn
619 # Calculate alpha
620 alpha = 1.0
621 if pap == 0.0 and mxiter == 1:
622 raise RuntimeError(
623 "Conjugate gradient method failed. Set mxiter greater than one and try again. Stopping execution."
624 )
625 if pap != 0.0:
626 alpha = srnew / pap
627 return v, alpha
630@numba.njit
631def calculate_heads(ibound, nlay, nrow, ncol, nrc, alpha, p, bigh, bigr, hnew, res, v):
632 """
633 Calculate new heads and residuals, and save the largest
634 change in head and the largest value of the residual.
635 """
636 ih = jh = kh = -1
637 ir = jr = kr = -1
638 for k in range(nlay):
639 for i in range(nrow):
640 for j in range(ncol):
641 n = j + i * ncol + k * nrc
642 if ibound[n] <= 0:
643 continue
644 hchgn = alpha * p[n]
645 if abs(hchgn) > abs(bigh):
646 bigh = hchgn
647 ih = i
648 jh = j
649 kh = k
650 hnew[n] = hnew[n] + hchgn
652 # residual (v is the product of matrix a and vector p)
653 rchgn = -alpha * v[n]
654 res[n] = res[n] + rchgn
655 if abs(res[n]) > abs(bigr):
656 bigr = res[n]
657 ir = i
658 jr = j
659 kr = k
660 indices_bigh = (kh, ih, jh)
661 indices_bigr = (kr, ir, jr)
662 return hnew, res, bigh, bigr, indices_bigh, indices_bigr