Refactored function by adding numpy.newaxis

lsst · Feb 12, 2020 · cc75af8 · cc75af8
1 parent 3f0cff5
commit cc75af8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 17 deletions.
diff --git a/python/lsst/ip/isr/isrFunctions.py b/python/lsst/ip/isr/isrFunctions.py
@@ -539,27 +539,24 @@ def overscanCorrection(ampMaskedImage, overscanImage, fitType='MEDIAN', order=1,
             for row in biasArray:
                 rowMedian = afwMath.makeStatistics(row, fitTypeStats, statControl).getValue()
                 collapsed.append(rowMedian)
-            if shortInd == 0:
-                collapsed = numpy.transpose(collapsed)
-            shortDimAmpImage = int(ampImage.array.shape[0]*ampImage.array.shape[1]/len(collapsed))
-            collapsedReshaped = numpy.repeat(collapsed, shortDimAmpImage).reshape(ampImage.array.shape)
-            offImage = afwImage.ImageF(ampImage.getDimensions())
-            overscanFit = afwImage.ImageF(overscanImageI.getDimensions())
-            if shortInd == 0:
-                offArray = numpy.transpose(offImage.getArray())
-                ovFitArray = numpy.transpose(overscanFit.getArray())
+            collapsed = numpy.array(collapsed)
+            offImage = ampImage.Factory(ampImage.getDimensions())
+            offArray = offImage.getArray()
+            overscanFit = afwImage.ImageF(overscanImage.getDimensions())
+            overscanArray = overscanFit.getArray()
+
+            if shortInd == 1:
+                offArray[:, :] = collapsed[:, numpy.newaxis]
+                overscanArray[:, :] = collapsed[:, numpy.newaxis]
             else:
-                offArray = offImage.getArray()
-                ovFitArray = overscanFit.getArray()
-            offArray[:, :] = collapsedReshaped
-            ovFitArray[:, :] = numpy.repeat(collapsed, biasArray.shape[1]).reshape(biasArray.shape)
-            del collapsedReshaped, collapsed, biasArray
+                offArray[:, :] = collapsed[numpy.newaxis, :]
+                overscanArray[:, :] = collapsed[numpy.newaxis, :]
+
+            del collapsed, biasArray
 
         if overscanIsInt:
             del overscanImageI
-    elif fitType in ('POLY', 'CHEB', 'LEG', 'NATURAL_SPLINE', 'CUBIC_SPLINE',
-                     'AKIMA_SPLINE'):
-
+    elif fitType in ('POLY', 'CHEB', 'LEG', 'NATURAL_SPLINE', 'CUBIC_SPLINE', 'AKIMA_SPLINE'):
         if hasattr(overscanImage, "getImage"):
             biasArray = overscanImage.getImage().getArray()
             biasArray = numpy.ma.masked_where(overscanImage.getMask().getArray() & statControl.getAndMask(),

diff --git a/tests/test_overscanCorrection.py b/tests/test_overscanCorrection.py
@@ -82,6 +82,7 @@ def checkOverscanCorrectionX(self, **kwargs):
         biassec = '[11:13,1:10]'
         overscan = afwImage.MaskedImageF(maskedImage, bbox)
         overscan.set(2, 0x0, 1)
+
         exposure = afwImage.ExposureF(maskedImage, None)
         metadata = exposure.getMetadata()
         metadata.setString(self.overscanKeyword, biassec)