Coverage for python/lsst/cp/pipe/makeBrighterFatterKernel.py: 12%
238 statements
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-03 12:41 +0000
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-03 12:41 +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."""
24__all__ = ['BrighterFatterKernelSolveTask',
25 'BrighterFatterKernelSolveConfig']
27import numpy as np
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
34from lsst.ip.isr import (BrighterFatterKernel)
35from .utils import (funcPolynomial, irlsFit)
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 )
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 )
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 )
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 )
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 )
154class BrighterFatterKernelSolveTask(pipeBase.PipelineTask):
155 """Measure appropriate Brighter-Fatter Kernel from the PTC dataset.
156 """
158 ConfigClass = BrighterFatterKernelSolveConfig
159 _DefaultName = 'cpBfkMeasure'
161 def runQuantum(self, butlerQC, inputRefs, outputRefs):
162 """Ensure that the input and output dimensions are passed along.
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)
175 # Use the dimensions to set calib/provenance information.
176 inputs['inputDims'] = dict(inputRefs.inputPtc.dataId.required)
178 outputs = self.run(**inputs)
179 butlerQC.put(outputs, outputRefs)
181 def run(self, inputPtc, dummy, camera, inputDims):
182 """Combine covariance information from PTC into brighter-fatter
183 kernels.
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.
198 Returns
199 -------
200 results : `lsst.pipe.base.Struct`
201 The resulst struct containing:
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.")
210 detector = camera[inputDims['detector']]
211 detName = detector.getName()
213 if self.config.level == 'DETECTOR':
214 detectorCorrList = list()
215 detectorFluxes = list()
217 if not inputPtc.ptcFitType == "FULLCOVARIANCE" and self.config.useCovModelSample:
218 raise ValueError("ptcFitType must be FULLCOVARIANCE if useCovModelSample=True.")
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]
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
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])
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
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])
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]
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-
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])
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)
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
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
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)
320 fluxes = np.array(truncatedFluxes)
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
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}")
348 center = int((bfk.shape[0] - 1) / 2)
350 if self.config.forceZeroSum:
351 totalSum = np.sum(preKernel)
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)
360 preKernel[center, center] -= totalSum
361 self.log.info("%s Zero-Sum Scale: %g", ampName, totalSum)
363 finalSum = np.sum(preKernel)
364 bfk.meanXcorrs[ampName] = preKernel
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])
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)
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])
389 return pipeBase.Struct(
390 outputBFK=bfk,
391 )
393 def averageCorrelations(self, xCorrList, name):
394 """Average input correlations.
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.
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)
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()
419 # To match previous definitions, pad by one element.
420 meanXcorr = np.pad(meanXcorr, ((1, 1)))
422 return meanXcorr
424 def quadraticCorrelations(self, xCorrList, fluxList, name):
425 """Measure a quadratic correlation model.
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.
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)
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])
457 # To match previous definitions, pad by one element.
458 meanXcorr = np.pad(meanXcorr, ((1, 1)))
460 return meanXcorr
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)
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.
483 Returns
484 -------
485 covTilde : `numpy.array`, (N, N)
486 The calculated C-tilde from Broughton et al. 2023 (eq. 4).
487 """
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)
494 # Find the nearest measured flux level
495 # and the full covariance model at that point
496 nearestFlux = fluxes[ix]
497 covModelSample = covModelList[ix]
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)
504 return covTilde
506 @staticmethod
507 def _tileArray(in_array):
508 """Given an input quarter-image, tile/mirror it and return full image.
510 Given a square input of side-length n, of the form
512 input = array([[1, 2, 3],
513 [4, 5, 6],
514 [7, 8, 9]])
516 return an array of size 2n-1 as
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]])
524 Parameters
525 ----------
526 input : `np.array`, (N, N)
527 The square input quarter-array
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))
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
546 def successiveOverRelax(self, source, maxIter=None, eLevel=None):
547 """An implementation of the successive over relaxation (SOR) method.
549 A numerical method for solving a system of linear equations
550 with faster convergence than the Gauss-Seidel method.
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.
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
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
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))
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
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]