Coverage for C:\src\imod-python\imod\prepare\pcg.py: 93%

343 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1"""Preconditioned Conjugate Gradient Solver Package. 

2 

3This is a translation of the fortran77 modflow2000 pcg2.f module. 

4""" 

5 

6import numba 

7 

8 

9class PreconditionedConjugateGradientSolver(object): 

10 __slots__ = ("hclose", "rclose", "mxiter", "iter1", "relax") 

11 

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 

32 

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. 

37 

38 Mutates hnew, converged, cr, cc, cv, ibound, and the work arrays. 

39 

40 Parameters 

41 ---------- 

42 hnew : ndarray (3D) of floats 

43 Initial guess, will be updated in place. 

44 

45 The linear operator exists only virtually, and is formed by cr, cc, cv, 

46 ibound. 

47 

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 

70 

71 Returns 

72 ------- 

73 converged : boolean 

74 convergence flag 

75 """ 

76 

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 

85 

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() 

94 

95 # Clear work arrays 

96 res[:] = 0.0 

97 cd[:] = 0.0 

98 v[:] = 0.0 

99 ss[:] = 0.0 

100 p[:] = 0.0 

101 

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 ) 

107 

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 

115 

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 ) 

132 

133 ss = back_substition(ibound, nlay, nrow, ncol, nrc, cr, cc, cv, v, cd, ss) 

134 

135 p, srnew = calculate_cg_p(ibound, nodes, ss, res, srnew, iiter, p) 

136 

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 ) 

153 

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 ) 

157 

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 

163 

164 return converged 

165 

166 

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. 

173 

174 For a dense form: 

175 Residual = np.dot(A, x) - b 

176 

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] 

191 

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]) 

272 

273 if i == nrow - 1: 

274 cc[n] = 0.0 

275 if j == ncol - 1: 

276 cr[n] = 0.0 

277 

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 

286 

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 

294 

295 

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 

323 

324 Returns 

325 ------- 

326 delta 

327 cd1 

328 cd 

329 """ 

330 

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] 

343 

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) 

350 

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]) 

360 

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] 

379 

380 hhcof = hcof[n] - z - b - d - f - h - s 

381 cd[n] = (1.0 + delta) * hhcof - cdcr - cdcc - cdcv - relax * (fcr + fcc + fcv) 

382 

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 

391 

392 

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. 

402 

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 

418 

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] 

427 

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] 

435 

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] 

443 

444 v[n] = res[n] - vcr - vcc - vcv 

445 

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 ) 

471 

472 return cd, v 

473 

474 

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. 

480 

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 

506 

507 

508@numba.njit 

509def calculate_cg_p(ibound, nodes, ss, res, srnew, iiter, p): 

510 """ 

511 Calculate p of the conjugate gradient algorithm 

512 

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 

531 

532 

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. 

541 

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 

557 

558 nrn = n + ncol 

559 nrl = n - ncol 

560 ncn = n + 1 

561 ncl = n - 1 

562 nln = n + nrc 

563 nll = n - nrc 

564 

565 ncf = n 

566 ncd = ncl 

567 nrb = nrl 

568 nrh = n 

569 nls = n 

570 nlz = nll 

571 

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] 

590 

591 pn = p[n] 

592 

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) 

611 

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 

618 

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 

628 

629 

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 

651 

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