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

238 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-01-30 12:35 +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) 

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 default=False, 

104 ) 

105 useAmatrix = pexConfig.Field( 

106 dtype=bool, 

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

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

109 default=False, 

110 ) 

111 

112 useCovModelSample = pexConfig.Field( 

113 dtype=bool, 

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

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

116 default=False, 

117 ) 

118 

119 covModelFluxSample = pexConfig.DictField( 

120 keytype=str, 

121 itemtype=float, 

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

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

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

125 default={'ALL_AMPS': 25000.0}, 

126 ) 

127 maxIterSuccessiveOverRelaxation = pexConfig.Field( 

128 dtype=int, 

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

130 default=10000 

131 ) 

132 eLevelSuccessiveOverRelaxation = pexConfig.Field( 

133 dtype=float, 

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

135 default=5.0e-14 

136 ) 

137 correlationQuadraticFit = pexConfig.Field( 

138 dtype=bool, 

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

140 default=False, 

141 ) 

142 correlationModelRadius = pexConfig.Field( 

143 dtype=int, 

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

145 default=100, 

146 ) 

147 correlationModelSlope = pexConfig.Field( 

148 dtype=float, 

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

150 default=-1.35, 

151 ) 

152 

153 

154class BrighterFatterKernelSolveTask(pipeBase.PipelineTask): 

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

156 """ 

157 

158 ConfigClass = BrighterFatterKernelSolveConfig 

159 _DefaultName = 'cpBfkMeasure' 

160 

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

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

163 

164 Parameters 

165 ---------- 

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

167 Butler to operate on. 

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

169 Input data refs to load. 

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

171 Output data refs to persist. 

172 """ 

173 inputs = butlerQC.get(inputRefs) 

174 

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

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

177 

178 outputs = self.run(**inputs) 

179 butlerQC.put(outputs, outputRefs) 

180 

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

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

183 kernels. 

184 

185 Parameters 

186 ---------- 

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

188 PTC data containing per-amplifier covariance measurements. 

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

190 The exposure used to select the appropriate PTC dataset. 

191 In almost all circumstances, one of the input exposures 

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

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

194 Camera to use for camera geometry information. 

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

196 DataIds to use to populate the output calibration. 

197 

198 Returns 

199 ------- 

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

201 The resulst struct containing: 

202 

203 ``outputBfk`` 

204 Resulting Brighter-Fatter Kernel 

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

206 """ 

207 if len(dummy) == 0: 

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

209 

210 detector = camera[inputDims['detector']] 

211 detName = detector.getName() 

212 

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

214 detectorCorrList = list() 

215 detectorFluxes = list() 

216 

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

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

219 

220 # Get flux sample dictionary 

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

222 for ampName in inputPtc.ampNames: 

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

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

225 elif ampName in self.config.covModelFluxSample: 

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

227 

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

229 bfk.rawMeans = inputPtc.rawMeans # ADU 

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

231 bfk.expIdMask = inputPtc.expIdMask 

232 

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

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

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

236 # ordering, as is the aMatrix. 

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

238 bfk.badAmps = inputPtc.badAmps 

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

240 bfk.gain = inputPtc.gain 

241 bfk.noise = inputPtc.noise 

242 bfk.meanXcorrs = dict() 

243 bfk.valid = dict() 

244 bfk.updateMetadataFromExposures([inputPtc]) 

245 

246 for amp in detector: 

247 ampName = amp.getName() 

248 gain = bfk.gain[ampName] 

249 noiseMatrix = inputPtc.noiseMatrix[ampName] 

250 mask = inputPtc.expIdMask[ampName] 

251 if gain <= 0: 

252 # We've received very bad data. 

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

254 ampName, gain) 

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

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

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

258 bfk.valid[ampName] = False 

259 continue 

260 

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

262 # covariances that were not masked after PTC. The 

263 # covariances may now have the mask already applied. 

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

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

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

267 

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

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

270 # Only apply the mask if needed. 

271 xCorrList = xCorrList[mask] 

272 

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

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

275 

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

277 # (arxiv:1711.06273) 

278 scaledCorrList = list() 

279 corrList = list() 

280 truncatedFluxes = list() 

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

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

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

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

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

286 

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

288 # component attributable to Poisson noise. This 

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

290 # Coulton et al. 2017 equation 29 

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

292 

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

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

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

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

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

298 # drop the flux entry as well. 

299 continue 

300 

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

302 # Coulton et al. 2017 equation 29. 

303 # The quadratic fit option needs the correlations unscaled 

304 q /= -2.0 

305 unscaled = self._tileArray(q) 

306 q /= flux**2 

307 scaled = self._tileArray(q) 

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

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

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

311 ampName, xcorrNum, xcorrCheck) 

312 continue 

313 

314 scaledCorrList.append(scaled) 

315 corrList.append(unscaled) 

316 truncatedFluxes.append(flux) 

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

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

319 

320 fluxes = np.array(truncatedFluxes) 

321 

322 if len(scaledCorrList) == 0: 

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

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

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

326 bfk.valid[ampName] = False 

327 continue 

328 

329 if self.config.useAmatrix: 

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

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

332 elif self.config.correlationQuadraticFit: 

333 # Use a quadratic fit to the correlations as a 

334 # function of flux. 

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

336 elif self.config.useCovModelSample: 

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

338 # Use the non-truncated fluxes for this 

339 mu = bfk.rawMeans[ampName] 

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

341 covModelList, fluxSampleDict[ampName], 

342 f"Amp: {ampName}") 

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

344 else: 

345 # Use a simple average of the measured correlations. 

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

347 

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

349 

350 if self.config.forceZeroSum: 

351 totalSum = np.sum(preKernel) 

352 

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

354 # Assume a correlation model of 

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

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

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

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

359 

360 preKernel[center, center] -= totalSum 

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

362 

363 finalSum = np.sum(preKernel) 

364 bfk.meanXcorrs[ampName] = preKernel 

365 

366 postKernel = self.successiveOverRelax(preKernel) 

367 bfk.ampKernels[ampName] = postKernel 

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

369 detectorCorrList.extend(scaledCorrList) 

370 detectorFluxes.extend(fluxes) 

371 bfk.valid[ampName] = True 

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

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

374 

375 # Assemble a detector kernel? 

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

377 if self.config.correlationQuadraticFit: 

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

379 else: 

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

381 finalSum = np.sum(preKernel) 

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

383 

384 postKernel = self.successiveOverRelax(preKernel) 

385 bfk.detKernels[detName] = postKernel 

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

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

388 

389 return pipeBase.Struct( 

390 outputBFK=bfk, 

391 ) 

392 

393 def averageCorrelations(self, xCorrList, name): 

394 """Average input correlations. 

395 

396 Parameters 

397 ---------- 

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

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

400 square arrays. 

401 name : `str` 

402 Name for log messages. 

403 

404 Returns 

405 ------- 

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

407 The averaged cross-correlation. 

408 """ 

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

410 xCorrList = np.array(xCorrList) 

411 

412 sctrl = afwMath.StatisticsControl() 

413 sctrl.setNumSigmaClip(self.config.nSigmaClip) 

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

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

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

417 afwMath.MEANCLIP, sctrl).getValue() 

418 

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

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

421 

422 return meanXcorr 

423 

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

425 """Measure a quadratic correlation model. 

426 

427 Parameters 

428 ---------- 

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

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

431 square arrays. 

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

433 Associated list of fluxes. 

434 name : `str` 

435 Name for log messages. 

436 

437 Returns 

438 ------- 

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

440 The averaged cross-correlation. 

441 """ 

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

443 fluxList = np.square(fluxList) 

444 xCorrList = np.array(xCorrList) 

445 

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

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

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

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

450 # inversion. 

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

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

453 scaleResidual=False) 

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

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

456 

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

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

459 

460 return meanXcorr 

461 

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

463 """Sample the correlation model and measure 

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

465 

466 Parameters 

467 ---------- 

468 fluxes : `list` [`float`] 

469 List of fluxes (in ADU) 

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

471 Noise matrix 

472 gain : `float` 

473 Amplifier gain 

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

475 List of covariance model matrices. These are 

476 expected to be square arrays. 

477 flux : `float` 

478 Flux in electrons at which to sample the 

479 covariance model. 

480 name : `str` 

481 Name for log messages. 

482 

483 Returns 

484 ------- 

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

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

487 """ 

488 

489 # Get the index of the flux sample 

490 # (this must be done in electron units) 

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

492 assert len(fluxes) == len(covModelList) 

493 

494 # Find the nearest measured flux level 

495 # and the full covariance model at that point 

496 nearestFlux = fluxes[ix] 

497 covModelSample = covModelList[ix] 

498 

499 # Calculate flux sample 

500 # covTilde returned in ADU units 

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

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

503 

504 return covTilde 

505 

506 @staticmethod 

507 def _tileArray(in_array): 

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

509 

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

511 

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

513 [4, 5, 6], 

514 [7, 8, 9]]) 

515 

516 return an array of size 2n-1 as 

517 

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

519 [ 6, 5, 4, 5, 6], 

520 [ 3, 2, 1, 2, 3], 

521 [ 6, 5, 4, 5, 6], 

522 [ 9, 8, 7, 8, 9]]) 

523 

524 Parameters 

525 ---------- 

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

527 The square input quarter-array 

528 

529 Returns 

530 ------- 

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

532 The full, tiled array 

533 """ 

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

535 length = in_array.shape[0] - 1 

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

537 

538 for i in range(length + 1): 

539 for j in range(length + 1): 

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

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

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

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

544 return output 

545 

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

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

548 

549 A numerical method for solving a system of linear equations 

550 with faster convergence than the Gauss-Seidel method. 

551 

552 Parameters 

553 ---------- 

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

555 The input array. 

556 maxIter : `int`, optional 

557 Maximum number of iterations to attempt before aborting. 

558 eLevel : `float`, optional 

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

560 occurred. 

561 

562 Returns 

563 ------- 

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

565 The solution. 

566 """ 

567 if not maxIter: 

568 maxIter = self.config.maxIterSuccessiveOverRelaxation 

569 if not eLevel: 

570 eLevel = self.config.eLevelSuccessiveOverRelaxation 

571 

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

573 # initialize, and set boundary conditions 

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

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

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

577 

578 # Calculate the initial error 

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

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

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

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

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

584 

585 # Iterate until convergence 

586 # We perform two sweeps per cycle, 

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

588 nIter = 0 

589 omega = 1.0 

590 dx = 1.0 

591 while nIter < maxIter*2: 

592 outError = 0 

593 if nIter%2 == 0: 

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

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

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

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

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

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

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

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

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

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

604 else: 

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

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

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

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

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

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

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

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

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

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

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

616 if outError < inError*eLevel: 

617 break 

618 if nIter == 0: 

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

620 else: 

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

622 nIter += 1 

623 

624 if nIter >= maxIter*2: 

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

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

627 else: 

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

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

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