Coverage for python/lsst/cp/pipe/makeBrighterFatterKernel.py: 12%
253 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-15 07:06 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-15 07:06 -0700
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, extractCalibDate)
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 "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 )
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 )
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 )
155class BrighterFatterKernelSolveTask(pipeBase.PipelineTask):
156 """Measure appropriate Brighter-Fatter Kernel from the PTC dataset.
157 """
159 ConfigClass = BrighterFatterKernelSolveConfig
160 _DefaultName = 'cpBfkMeasure'
162 def runQuantum(self, butlerQC, inputRefs, outputRefs):
163 """Ensure that the input and output dimensions are passed along.
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)
176 # Use the dimensions to set calib/provenance information.
177 inputs['inputDims'] = dict(inputRefs.inputPtc.dataId.required)
179 # Add calibration provenance info to header.
180 kwargs = dict()
181 reference = getattr(inputRefs, "inputPtc", None)
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)
192 kwargs[runKey] = runValue
193 kwargs[idKey] = idValue
194 kwargs[dateKey] = dateValue
196 self.log.info("Using " + str(reference.run))
198 outputs = self.run(**inputs)
199 outputs.outputBFK.updateMetadata(setDate=False, **kwargs)
201 butlerQC.put(outputs, outputRefs)
203 def run(self, inputPtc, dummy, camera, inputDims):
204 """Combine covariance information from PTC into brighter-fatter
205 kernels.
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.
220 Returns
221 -------
222 results : `lsst.pipe.base.Struct`
223 The resulst struct containing:
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.")
232 detector = camera[inputDims['detector']]
233 detName = detector.getName()
235 if self.config.level == 'DETECTOR':
236 detectorCorrList = list()
237 detectorFluxes = list()
239 if not inputPtc.ptcFitType == "FULLCOVARIANCE" and self.config.useCovModelSample:
240 raise ValueError("ptcFitType must be FULLCOVARIANCE if useCovModelSample=True.")
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]
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
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])
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
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])
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]
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-
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])
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)
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
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
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)
342 fluxes = np.array(truncatedFluxes)
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
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}")
370 center = int((bfk.shape[0] - 1) / 2)
372 if self.config.forceZeroSum:
373 totalSum = np.sum(preKernel)
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)
382 preKernel[center, center] -= totalSum
383 self.log.info("%s Zero-Sum Scale: %g", ampName, totalSum)
385 finalSum = np.sum(preKernel)
386 bfk.meanXcorrs[ampName] = preKernel
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])
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)
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])
411 return pipeBase.Struct(
412 outputBFK=bfk,
413 )
415 def averageCorrelations(self, xCorrList, name):
416 """Average input correlations.
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.
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)
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()
441 # To match previous definitions, pad by one element.
442 meanXcorr = np.pad(meanXcorr, ((1, 1)))
444 return meanXcorr
446 def quadraticCorrelations(self, xCorrList, fluxList, name):
447 """Measure a quadratic correlation model.
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.
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)
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])
479 # To match previous definitions, pad by one element.
480 meanXcorr = np.pad(meanXcorr, ((1, 1)))
482 return meanXcorr
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)
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.
505 Returns
506 -------
507 covTilde : `numpy.array`, (N, N)
508 The calculated C-tilde from Broughton et al. 2023 (eq. 4).
509 """
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)
516 # Find the nearest measured flux level
517 # and the full covariance model at that point
518 nearestFlux = fluxes[ix]
519 covModelSample = covModelList[ix]
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)
526 return covTilde
528 @staticmethod
529 def _tileArray(in_array):
530 """Given an input quarter-image, tile/mirror it and return full image.
532 Given a square input of side-length n, of the form
534 input = array([[1, 2, 3],
535 [4, 5, 6],
536 [7, 8, 9]])
538 return an array of size 2n-1 as
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]])
546 Parameters
547 ----------
548 input : `np.array`, (N, N)
549 The square input quarter-array
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))
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
568 def successiveOverRelax(self, source, maxIter=None, eLevel=None):
569 """An implementation of the successive over relaxation (SOR) method.
571 A numerical method for solving a system of linear equations
572 with faster convergence than the Gauss-Seidel method.
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.
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
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
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))
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
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]