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

1"""Preconditioned Conjugate Gradient Solver Package.

2

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

4"""

5

6 import numba

7

8

9 class 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

168 def 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

297 def 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

394 def 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

476 def 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

509 def 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

534 def 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

631 def 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