Coverage for python/lsst/cp/pipe/makeBrighterFatterKernel.py: 12%

253 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-22 11:10 +0000

1# This file is part of cp_pipe. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21# 

22"""Calculation of brighter-fatter effect correlations and kernels.""" 

23 

24__all__ = ['BrighterFatterKernelSolveTask', 

25 'BrighterFatterKernelSolveConfig'] 

26 

27import numpy as np 

28 

29import lsst.afw.math as afwMath 

30import lsst.pex.config as pexConfig 

31import lsst.pipe.base as pipeBase 

32import lsst.pipe.base.connectionTypes as cT 

33 

34from lsst.ip.isr import (BrighterFatterKernel) 

35from .utils import (funcPolynomial, irlsFit, extractCalibDate) 

36 

37 

38class BrighterFatterKernelSolveConnections(pipeBase.PipelineTaskConnections, 

39 dimensions=("instrument", "exposure", "detector")): 

40 dummy = cT.Input( 

41 name="raw", 

42 doc="Dummy exposure.", 

43 storageClass='Exposure', 

44 dimensions=("instrument", "exposure", "detector"), 

45 multiple=True, 

46 deferLoad=True, 

47 ) 

48 camera = cT.PrerequisiteInput( 

49 name="camera", 

50 doc="Camera associated with this data.", 

51 storageClass="Camera", 

52 dimensions=("instrument", ), 

53 isCalibration=True, 

54 ) 

55 inputPtc = cT.PrerequisiteInput( 

56 name="ptc", 

57 doc="Photon transfer curve dataset.", 

58 storageClass="PhotonTransferCurveDataset", 

59 dimensions=("instrument", "detector"), 

60 isCalibration=True, 

61 ) 

62 

63 outputBFK = cT.Output( 

64 name="brighterFatterKernel", 

65 doc="Output measured brighter-fatter kernel.", 

66 storageClass="BrighterFatterKernel", 

67 dimensions=("instrument", "detector"), 

68 isCalibration=True, 

69 ) 

70 

71 

72class BrighterFatterKernelSolveConfig(pipeBase.PipelineTaskConfig, 

73 pipelineConnections=BrighterFatterKernelSolveConnections): 

74 level = pexConfig.ChoiceField( 

75 doc="The level at which to calculate the brighter-fatter kernels", 

76 dtype=str, 

77 default="AMP", 

78 allowed={ 

79 "AMP": "Every amplifier treated separately", 

80 "DETECTOR": "One kernel per detector", 

81 } 

82 ) 

83 ignoreAmpsForAveraging = pexConfig.ListField( 

84 dtype=str, 

85 doc="List of amp names to ignore when averaging the amplifier kernels into the detector" 

86 " kernel. Only relevant for level = DETECTOR", 

87 default=[] 

88 ) 

89 xcorrCheckRejectLevel = pexConfig.Field( 

90 dtype=float, 

91 doc="Rejection level for the sum of the input cross-correlations. Arrays which " 

92 "sum to greater than this are discarded before the clipped mean is calculated.", 

93 default=2.0 

94 ) 

95 nSigmaClip = pexConfig.Field( 

96 dtype=float, 

97 doc="Number of sigma to clip when calculating means for the cross-correlation", 

98 default=5 

99 ) 

100 forceZeroSum = pexConfig.Field( 

101 dtype=bool, 

102 doc="Force the correlation matrix to have zero sum by adjusting the (0,0) value?" 

103 "Defaults to true bsed on recommendation of Broughton et al. 2024.", 

104 default=True, 

105 ) 

106 useAmatrix = pexConfig.Field( 

107 dtype=bool, 

108 doc="Use the PTC 'a' matrix (Astier et al. 2019 equation 20) " 

109 "instead of the average of measured covariances?", 

110 default=False, 

111 ) 

112 

113 useCovModelSample = pexConfig.Field( 

114 dtype=bool, 

115 doc="Use the covariance matrix sampled from the full covariance model " 

116 "(Astier et al. 2019 equation 20) instead of the average measured covariances?", 

117 default=False, 

118 ) 

119 

120 covModelFluxSample = pexConfig.DictField( 

121 keytype=str, 

122 itemtype=float, 

123 doc="Flux level in electrons at which to sample the full covariance" 

124 "model if useCovModelSample=True. The same level is applied to all" 

125 "amps if this parameter [`dict`] is passed as {'ALL_AMPS': value}", 

126 default={'ALL_AMPS': 25000.0}, 

127 ) 

128 maxIterSuccessiveOverRelaxation = pexConfig.Field( 

129 dtype=int, 

130 doc="The maximum number of iterations allowed for the successive over-relaxation method", 

131 default=10000 

132 ) 

133 eLevelSuccessiveOverRelaxation = pexConfig.Field( 

134 dtype=float, 

135 doc="The target residual error for the successive over-relaxation method", 

136 default=5.0e-14 

137 ) 

138 correlationQuadraticFit = pexConfig.Field( 

139 dtype=bool, 

140 doc="Use a quadratic fit to find the correlations instead of simple averaging?", 

141 default=False, 

142 ) 

143 correlationModelRadius = pexConfig.Field( 

144 dtype=int, 

145 doc="Build a model of the correlation coefficients for radii larger than this value in pixels?", 

146 default=100, 

147 ) 

148 correlationModelSlope = pexConfig.Field( 

149 dtype=float, 

150 doc="Slope of the correlation model for radii larger than correlationModelRadius", 

151 default=-1.35, 

152 ) 

153 

154 

155class BrighterFatterKernelSolveTask(pipeBase.PipelineTask): 

156 """Measure appropriate Brighter-Fatter Kernel from the PTC dataset. 

157 """ 

158 

159 ConfigClass = BrighterFatterKernelSolveConfig 

160 _DefaultName = 'cpBfkMeasure' 

161 

162 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

163 """Ensure that the input and output dimensions are passed along. 

164 

165 Parameters 

166 ---------- 

167 butlerQC : `lsst.daf.butler.QuantumContext` 

168 Butler to operate on. 

169 inputRefs : `lsst.pipe.base.InputQuantizedConnection` 

170 Input data refs to load. 

171 ouptutRefs : `lsst.pipe.base.OutputQuantizedConnection` 

172 Output data refs to persist. 

173 """ 

174 inputs = butlerQC.get(inputRefs) 

175 

176 # Use the dimensions to set calib/provenance information. 

177 inputs['inputDims'] = dict(inputRefs.inputPtc.dataId.required) 

178 

179 # Add calibration provenance info to header. 

180 kwargs = dict() 

181 reference = getattr(inputRefs, "inputPtc", None) 

182 

183 if reference is not None and hasattr(reference, "run"): 

184 runKey = "PTC_RUN" 

185 runValue = reference.run 

186 idKey = "PTC_UUID" 

187 idValue = str(reference.id) 

188 dateKey = "PTC_DATE" 

189 calib = inputs.get("inputPtc", None) 

190 dateValue = extractCalibDate(calib) 

191 

192 kwargs[runKey] = runValue 

193 kwargs[idKey] = idValue 

194 kwargs[dateKey] = dateValue 

195 

196 self.log.info("Using " + str(reference.run)) 

197 

198 outputs = self.run(**inputs) 

199 outputs.outputBFK.updateMetadata(setDate=False, **kwargs) 

200 

201 butlerQC.put(outputs, outputRefs) 

202 

203 def run(self, inputPtc, dummy, camera, inputDims): 

204 """Combine covariance information from PTC into brighter-fatter 

205 kernels. 

206 

207 Parameters 

208 ---------- 

209 inputPtc : `lsst.ip.isr.PhotonTransferCurveDataset` 

210 PTC data containing per-amplifier covariance measurements. 

211 dummy : `lsst.afw.image.Exposure` 

212 The exposure used to select the appropriate PTC dataset. 

213 In almost all circumstances, one of the input exposures 

214 used to generate the PTC dataset is the best option. 

215 camera : `lsst.afw.cameraGeom.Camera` 

216 Camera to use for camera geometry information. 

217 inputDims : `lsst.daf.butler.DataCoordinate` or `dict` 

218 DataIds to use to populate the output calibration. 

219 

220 Returns 

221 ------- 

222 results : `lsst.pipe.base.Struct` 

223 The resulst struct containing: 

224 

225 ``outputBfk`` 

226 Resulting Brighter-Fatter Kernel 

227 (`lsst.ip.isr.BrighterFatterKernel`). 

228 """ 

229 if len(dummy) == 0: 

230 self.log.warning("No dummy exposure found.") 

231 

232 detector = camera[inputDims['detector']] 

233 detName = detector.getName() 

234 

235 if self.config.level == 'DETECTOR': 

236 detectorCorrList = list() 

237 detectorFluxes = list() 

238 

239 if not inputPtc.ptcFitType == "FULLCOVARIANCE" and self.config.useCovModelSample: 

240 raise ValueError("ptcFitType must be FULLCOVARIANCE if useCovModelSample=True.") 

241 

242 # Get flux sample dictionary 

243 fluxSampleDict = {ampName: 0.0 for ampName in inputPtc.ampNames} 

244 for ampName in inputPtc.ampNames: 

245 if 'ALL_AMPS' in self.config.covModelFluxSample: 

246 fluxSampleDict[ampName] = self.config.covModelFluxSample['ALL_AMPS'] 

247 elif ampName in self.config.covModelFluxSample: 

248 fluxSampleDict[ampName] = self.config.covModelFluxSample[ampName] 

249 

250 bfk = BrighterFatterKernel(camera=camera, detectorId=detector.getId(), level=self.config.level) 

251 bfk.rawMeans = inputPtc.rawMeans # ADU 

252 bfk.rawVariances = inputPtc.rawVars # ADU^2 

253 bfk.expIdMask = inputPtc.expIdMask 

254 

255 # Use the PTC covariances as the cross-correlations. These 

256 # are scaled before the kernel is generated, which performs 

257 # the conversion. The input covariances are in (x, y) index 

258 # ordering, as is the aMatrix. 

259 bfk.rawXcorrs = inputPtc.covariances # ADU^2 

260 bfk.badAmps = inputPtc.badAmps 

261 bfk.shape = (inputPtc.covMatrixSide*2 + 1, inputPtc.covMatrixSide*2 + 1) 

262 bfk.gain = inputPtc.gain 

263 bfk.noise = inputPtc.noise 

264 bfk.meanXcorrs = dict() 

265 bfk.valid = dict() 

266 bfk.updateMetadataFromExposures([inputPtc]) 

267 

268 for amp in detector: 

269 ampName = amp.getName() 

270 gain = bfk.gain[ampName] 

271 noiseMatrix = inputPtc.noiseMatrix[ampName] 

272 mask = inputPtc.expIdMask[ampName] 

273 if gain <= 0: 

274 # We've received very bad data. 

275 self.log.warning("Impossible gain recieved from PTC for %s: %f. Skipping bad amplifier.", 

276 ampName, gain) 

277 bfk.meanXcorrs[ampName] = np.zeros(bfk.shape) 

278 bfk.ampKernels[ampName] = np.zeros(bfk.shape) 

279 bfk.rawXcorrs[ampName] = np.zeros((len(mask), inputPtc.covMatrixSide, inputPtc.covMatrixSide)) 

280 bfk.valid[ampName] = False 

281 continue 

282 

283 # Use inputPtc.expIdMask to get the means, variances, and 

284 # covariances that were not masked after PTC. The 

285 # covariances may now have the mask already applied. 

286 fluxes = np.array(bfk.rawMeans[ampName])[mask] 

287 variances = np.array(bfk.rawVariances[ampName])[mask] 

288 covModelList = np.array(inputPtc.covariancesModel[ampName]) 

289 

290 xCorrList = np.array([np.array(xcorr) for xcorr in bfk.rawXcorrs[ampName]]) 

291 if np.sum(mask) < len(xCorrList): 

292 # Only apply the mask if needed. 

293 xCorrList = xCorrList[mask] 

294 

295 fluxes = np.array([flux*gain for flux in fluxes]) # Now in e^- 

296 variances = np.array([variance*gain*gain for variance in variances]) # Now in e^2- 

297 

298 # This should duplicate Coulton et al. 2017 Equation 22-29 

299 # (arxiv:1711.06273) 

300 scaledCorrList = list() 

301 corrList = list() 

302 truncatedFluxes = list() 

303 for xcorrNum, (xcorr, flux, var) in enumerate(zip(xCorrList, fluxes, variances), 1): 

304 q = np.array(xcorr) * gain * gain # xcorr now in e^- 

305 q *= 2.0 # Remove factor of 1/2 applied in PTC. 

306 self.log.info("Amp: %s %d/%d Flux: %f Var: %f Q(0,0): %g Q(1,0): %g Q(0,1): %g", 

307 ampName, xcorrNum, len(xCorrList), flux, var, q[0][0], q[1][0], q[0][1]) 

308 

309 # Normalize by the flux, which removes the (0,0) 

310 # component attributable to Poisson noise. This 

311 # contains the two "t I delta(x - x')" terms in 

312 # Coulton et al. 2017 equation 29 

313 q[0][0] -= 2.0*(flux) 

314 

315 if q[0][0] > 0.0: 

316 self.log.warning("Amp: %s %d skipped due to value of (variance-mean)=%f", 

317 ampName, xcorrNum, q[0][0]) 

318 # If we drop an element of ``scaledCorrList`` 

319 # (which is what this does), we need to ensure we 

320 # drop the flux entry as well. 

321 continue 

322 

323 # This removes the "t (I_a^2 + I_b^2)" factor in 

324 # Coulton et al. 2017 equation 29. 

325 # The quadratic fit option needs the correlations unscaled 

326 q /= -2.0 

327 unscaled = self._tileArray(q) 

328 q /= flux**2 

329 scaled = self._tileArray(q) 

330 xcorrCheck = np.abs(np.sum(scaled))/np.sum(np.abs(scaled)) 

331 if (xcorrCheck > self.config.xcorrCheckRejectLevel) or not (np.isfinite(xcorrCheck)): 

332 self.log.warning("Amp: %s %d skipped due to value of triangle-inequality sum %f", 

333 ampName, xcorrNum, xcorrCheck) 

334 continue 

335 

336 scaledCorrList.append(scaled) 

337 corrList.append(unscaled) 

338 truncatedFluxes.append(flux) 

339 self.log.info("Amp: %s %d/%d Final: %g XcorrCheck: %f", 

340 ampName, xcorrNum, len(xCorrList), q[0][0], xcorrCheck) 

341 

342 fluxes = np.array(truncatedFluxes) 

343 

344 if len(scaledCorrList) == 0: 

345 self.log.warning("Amp: %s All inputs rejected for amp!", ampName) 

346 bfk.meanXcorrs[ampName] = np.zeros(bfk.shape) 

347 bfk.ampKernels[ampName] = np.zeros(bfk.shape) 

348 bfk.valid[ampName] = False 

349 continue 

350 

351 if self.config.useAmatrix: 

352 # Use the aMatrix, ignoring the meanXcorr generated above. 

353 preKernel = np.pad(self._tileArray(-1.0 * np.array(inputPtc.aMatrix[ampName])), ((1, 1))) 

354 elif self.config.correlationQuadraticFit: 

355 # Use a quadratic fit to the correlations as a 

356 # function of flux. 

357 preKernel = self.quadraticCorrelations(corrList, fluxes, f"Amp: {ampName}") 

358 elif self.config.useCovModelSample: 

359 # Sample the full covariance model at a given flux. 

360 # Use the non-truncated fluxes for this 

361 mu = bfk.rawMeans[ampName] 

362 covTilde = self.sampleCovModel(mu, noiseMatrix, gain, 

363 covModelList, fluxSampleDict[ampName], 

364 f"Amp: {ampName}") 

365 preKernel = np.pad(self._tileArray(-1.0 * covTilde), ((1, 1))) 

366 else: 

367 # Use a simple average of the measured correlations. 

368 preKernel = self.averageCorrelations(scaledCorrList, f"Amp: {ampName}") 

369 

370 center = int((bfk.shape[0] - 1) / 2) 

371 

372 if self.config.forceZeroSum: 

373 totalSum = np.sum(preKernel) 

374 

375 if self.config.correlationModelRadius < (preKernel.shape[0] - 1) / 2: 

376 # Assume a correlation model of 

377 # Corr(r) = -preFactor * r^(2 * slope) 

378 preFactor = np.sqrt(preKernel[center, center + 1] * preKernel[center + 1, center]) 

379 slopeFactor = 2.0 * np.abs(self.config.correlationModelSlope) 

380 totalSum += 2.0*np.pi*(preFactor / (slopeFactor*(center + 0.5))**slopeFactor) 

381 

382 preKernel[center, center] -= totalSum 

383 self.log.info("%s Zero-Sum Scale: %g", ampName, totalSum) 

384 

385 finalSum = np.sum(preKernel) 

386 bfk.meanXcorrs[ampName] = preKernel 

387 

388 postKernel = self.successiveOverRelax(preKernel) 

389 bfk.ampKernels[ampName] = postKernel 

390 if self.config.level == 'DETECTOR' and ampName not in self.config.ignoreAmpsForAveraging: 

391 detectorCorrList.extend(scaledCorrList) 

392 detectorFluxes.extend(fluxes) 

393 bfk.valid[ampName] = True 

394 self.log.info("Amp: %s Sum: %g Center Info Pre: %g Post: %g", 

395 ampName, finalSum, preKernel[center, center], postKernel[center, center]) 

396 

397 # Assemble a detector kernel? 

398 if self.config.level == 'DETECTOR': 

399 if self.config.correlationQuadraticFit: 

400 preKernel = self.quadraticCorrelations(detectorCorrList, detectorFluxes, f"Amp: {ampName}") 

401 else: 

402 preKernel = self.averageCorrelations(detectorCorrList, f"Det: {detName}") 

403 finalSum = np.sum(preKernel) 

404 center = int((bfk.shape[0] - 1) / 2) 

405 

406 postKernel = self.successiveOverRelax(preKernel) 

407 bfk.detKernels[detName] = postKernel 

408 self.log.info("Det: %s Sum: %g Center Info Pre: %g Post: %g", 

409 detName, finalSum, preKernel[center, center], postKernel[center, center]) 

410 

411 return pipeBase.Struct( 

412 outputBFK=bfk, 

413 ) 

414 

415 def averageCorrelations(self, xCorrList, name): 

416 """Average input correlations. 

417 

418 Parameters 

419 ---------- 

420 xCorrList : `list` [`numpy.array`] 

421 List of cross-correlations. These are expected to be 

422 square arrays. 

423 name : `str` 

424 Name for log messages. 

425 

426 Returns 

427 ------- 

428 meanXcorr : `numpy.array`, (N, N) 

429 The averaged cross-correlation. 

430 """ 

431 meanXcorr = np.zeros_like(xCorrList[0]) 

432 xCorrList = np.array(xCorrList) 

433 

434 sctrl = afwMath.StatisticsControl() 

435 sctrl.setNumSigmaClip(self.config.nSigmaClip) 

436 for i in range(np.shape(meanXcorr)[0]): 

437 for j in range(np.shape(meanXcorr)[1]): 

438 meanXcorr[i, j] = afwMath.makeStatistics(xCorrList[:, i, j], 

439 afwMath.MEANCLIP, sctrl).getValue() 

440 

441 # To match previous definitions, pad by one element. 

442 meanXcorr = np.pad(meanXcorr, ((1, 1))) 

443 

444 return meanXcorr 

445 

446 def quadraticCorrelations(self, xCorrList, fluxList, name): 

447 """Measure a quadratic correlation model. 

448 

449 Parameters 

450 ---------- 

451 xCorrList : `list` [`numpy.array`] 

452 List of cross-correlations. These are expected to be 

453 square arrays. 

454 fluxList : `numpy.array`, (Nflux,) 

455 Associated list of fluxes. 

456 name : `str` 

457 Name for log messages. 

458 

459 Returns 

460 ------- 

461 meanXcorr : `numpy.array`, (N, N) 

462 The averaged cross-correlation. 

463 """ 

464 meanXcorr = np.zeros_like(xCorrList[0]) 

465 fluxList = np.square(fluxList) 

466 xCorrList = np.array(xCorrList) 

467 

468 for i in range(np.shape(meanXcorr)[0]): 

469 for j in range(np.shape(meanXcorr)[1]): 

470 # Fit corrlation_i(x, y) = a0 + a1 * (flux_i)^2 We do 

471 # not want to transpose, so use (i, j) without 

472 # inversion. 

473 linearFit, linearFitErr, chiSq, weights = irlsFit([0.0, 1e-4], fluxList, 

474 xCorrList[:, i, j], funcPolynomial, 

475 scaleResidual=False) 

476 meanXcorr[i, j] = linearFit[1] # Discard the intercept. 

477 self.log.info("Quad fit meanXcorr[%d,%d] = %g", i, j, linearFit[1]) 

478 

479 # To match previous definitions, pad by one element. 

480 meanXcorr = np.pad(meanXcorr, ((1, 1))) 

481 

482 return meanXcorr 

483 

484 def sampleCovModel(self, fluxes, noiseMatrix, gain, covModelList, flux, name): 

485 """Sample the correlation model and measure 

486 widetile{C}_{ij} from Broughton et al. 2023 (eq. 4) 

487 

488 Parameters 

489 ---------- 

490 fluxes : `list` [`float`] 

491 List of fluxes (in ADU) 

492 noiseMatrix : `numpy.array`, (N, N) 

493 Noise matrix 

494 gain : `float` 

495 Amplifier gain 

496 covModelList : `numpy.array`, (N, N) 

497 List of covariance model matrices. These are 

498 expected to be square arrays. 

499 flux : `float` 

500 Flux in electrons at which to sample the 

501 covariance model. 

502 name : `str` 

503 Name for log messages. 

504 

505 Returns 

506 ------- 

507 covTilde : `numpy.array`, (N, N) 

508 The calculated C-tilde from Broughton et al. 2023 (eq. 4). 

509 """ 

510 

511 # Get the index of the flux sample 

512 # (this must be done in electron units) 

513 ix = np.argmin((fluxes*gain - flux)**2) 

514 assert len(fluxes) == len(covModelList) 

515 

516 # Find the nearest measured flux level 

517 # and the full covariance model at that point 

518 nearestFlux = fluxes[ix] 

519 covModelSample = covModelList[ix] 

520 

521 # Calculate flux sample 

522 # covTilde returned in ADU units 

523 covTilde = (covModelSample - noiseMatrix/gain**2)/(nearestFlux**2) 

524 covTilde[0][0] -= (nearestFlux/gain)/(nearestFlux**2) 

525 

526 return covTilde 

527 

528 @staticmethod 

529 def _tileArray(in_array): 

530 """Given an input quarter-image, tile/mirror it and return full image. 

531 

532 Given a square input of side-length n, of the form 

533 

534 input = array([[1, 2, 3], 

535 [4, 5, 6], 

536 [7, 8, 9]]) 

537 

538 return an array of size 2n-1 as 

539 

540 output = array([[ 9, 8, 7, 8, 9], 

541 [ 6, 5, 4, 5, 6], 

542 [ 3, 2, 1, 2, 3], 

543 [ 6, 5, 4, 5, 6], 

544 [ 9, 8, 7, 8, 9]]) 

545 

546 Parameters 

547 ---------- 

548 input : `np.array`, (N, N) 

549 The square input quarter-array 

550 

551 Returns 

552 ------- 

553 output : `np.array`, (2*N + 1, 2*N + 1) 

554 The full, tiled array 

555 """ 

556 assert in_array.shape[0] == in_array.shape[1] 

557 length = in_array.shape[0] - 1 

558 output = np.zeros((2*length + 1, 2*length + 1)) 

559 

560 for i in range(length + 1): 

561 for j in range(length + 1): 

562 output[i + length, j + length] = in_array[i, j] 

563 output[-i + length, j + length] = in_array[i, j] 

564 output[i + length, -j + length] = in_array[i, j] 

565 output[-i + length, -j + length] = in_array[i, j] 

566 return output 

567 

568 def successiveOverRelax(self, source, maxIter=None, eLevel=None): 

569 """An implementation of the successive over relaxation (SOR) method. 

570 

571 A numerical method for solving a system of linear equations 

572 with faster convergence than the Gauss-Seidel method. 

573 

574 Parameters 

575 ---------- 

576 source : `numpy.ndarray`, (N, N) 

577 The input array. 

578 maxIter : `int`, optional 

579 Maximum number of iterations to attempt before aborting. 

580 eLevel : `float`, optional 

581 The target error level at which we deem convergence to have 

582 occurred. 

583 

584 Returns 

585 ------- 

586 output : `numpy.ndarray`, (N, N) 

587 The solution. 

588 """ 

589 if not maxIter: 

590 maxIter = self.config.maxIterSuccessiveOverRelaxation 

591 if not eLevel: 

592 eLevel = self.config.eLevelSuccessiveOverRelaxation 

593 

594 assert source.shape[0] == source.shape[1], "Input array must be square" 

595 # initialize, and set boundary conditions 

596 func = np.zeros([source.shape[0] + 2, source.shape[1] + 2]) 

597 resid = np.zeros([source.shape[0] + 2, source.shape[1] + 2]) 

598 rhoSpe = np.cos(np.pi/source.shape[0]) # Here a square grid is assumed 

599 

600 # Calculate the initial error 

601 for i in range(1, func.shape[0] - 1): 

602 for j in range(1, func.shape[1] - 1): 

603 resid[i, j] = (func[i, j - 1] + func[i, j + 1] + func[i - 1, j] 

604 + func[i + 1, j] - 4*func[i, j] - source[i - 1, j - 1]) 

605 inError = np.sum(np.abs(resid)) 

606 

607 # Iterate until convergence 

608 # We perform two sweeps per cycle, 

609 # updating 'odd' and 'even' points separately 

610 nIter = 0 

611 omega = 1.0 

612 dx = 1.0 

613 while nIter < maxIter*2: 

614 outError = 0 

615 if nIter%2 == 0: 

616 for i in range(1, func.shape[0] - 1, 2): 

617 for j in range(1, func.shape[1] - 1, 2): 

618 resid[i, j] = float(func[i, j-1] + func[i, j + 1] + func[i - 1, j] 

619 + func[i + 1, j] - 4.0*func[i, j] - dx*dx*source[i - 1, j - 1]) 

620 func[i, j] += omega*resid[i, j]*.25 

621 for i in range(2, func.shape[0] - 1, 2): 

622 for j in range(2, func.shape[1] - 1, 2): 

623 resid[i, j] = float(func[i, j - 1] + func[i, j + 1] + func[i - 1, j] 

624 + func[i + 1, j] - 4.0*func[i, j] - dx*dx*source[i - 1, j - 1]) 

625 func[i, j] += omega*resid[i, j]*.25 

626 else: 

627 for i in range(1, func.shape[0] - 1, 2): 

628 for j in range(2, func.shape[1] - 1, 2): 

629 resid[i, j] = float(func[i, j - 1] + func[i, j + 1] + func[i - 1, j] 

630 + func[i + 1, j] - 4.0*func[i, j] - dx*dx*source[i - 1, j - 1]) 

631 func[i, j] += omega*resid[i, j]*.25 

632 for i in range(2, func.shape[0] - 1, 2): 

633 for j in range(1, func.shape[1] - 1, 2): 

634 resid[i, j] = float(func[i, j - 1] + func[i, j + 1] + func[i - 1, j] 

635 + func[i + 1, j] - 4.0*func[i, j] - dx*dx*source[i - 1, j - 1]) 

636 func[i, j] += omega*resid[i, j]*.25 

637 outError = np.sum(np.abs(resid)) 

638 if outError < inError*eLevel: 

639 break 

640 if nIter == 0: 

641 omega = 1.0/(1 - rhoSpe*rhoSpe/2.0) 

642 else: 

643 omega = 1.0/(1 - rhoSpe*rhoSpe*omega/4.0) 

644 nIter += 1 

645 

646 if nIter >= maxIter*2: 

647 self.log.warning("Failure: SuccessiveOverRelaxation did not converge in %s iterations." 

648 "\noutError: %s, inError: %s,", nIter//2, outError, inError*eLevel) 

649 else: 

650 self.log.info("Success: SuccessiveOverRelaxation converged in %s iterations." 

651 "\noutError: %s, inError: %s", nIter//2, outError, inError*eLevel) 

652 return func[1: -1, 1: -1]