Add Coadd function to PY_SWarp

sfft/CustomizedPacket.py

@@ -1,3 +1,7 @@
+import cupy as cp
+import tracemalloc
+import logging
+import sys
 import time
 import numpy as np
 import os.path as pa
@@ -13,7 +17,7 @@ class Customized_Packet:
     @staticmethod
     def CP(FITS_REF, FITS_SCI, FITS_mREF, FITS_mSCI, ForceConv, GKerHW, \
         FITS_DIFF=None, FITS_Solution=None, KerPolyOrder=2, BGPolyOrder=2, ConstPhotRatio=True, \
-        BACKEND_4SUBTRACT='Cupy', CUDA_DEVICE_4SUBTRACT='0', NUM_CPU_THREADS_4SUBTRACT=8, VERBOSE_LEVEL=2):
+        BACKEND_4SUBTRACT='Cupy', CUDA_DEVICE_4SUBTRACT='0', NUM_CPU_THREADS_4SUBTRACT=8, VERBOSE_LEVEL=2, logger=None):
 
         """
         * Parameters for Customized SFFT
@@ -85,11 +89,20 @@ def CP(FITS_REF, FITS_SCI, FITS_mREF, FITS_mSCI, ForceConv, GKerHW, \
 
         """
 
+        # Configure logger (Rob)
+        _logger = logging.getLogger(f'sfft.CustomizedPacket.CP')
+        if not _logger.hasHandlers():
+            log_out = logging.StreamHandler(sys.stderr)
+            formatter = logging.Formatter(f'[%(asctime)s - %(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
+            log_out.setFormatter(formatter)
+            _logger.addHandler(log_out)
+            _logger.setLevel(logging.DEBUG) # ERROR, WARNING, INFO, or DEBUG (in that order by increasing detail)
+
         # * Read input images
-        PixA_REF = fits.getdata(FITS_REF, ext=0).T
-        PixA_SCI = fits.getdata(FITS_SCI, ext=0).T
-        PixA_mREF = fits.getdata(FITS_mREF, ext=0).T
-        PixA_mSCI = fits.getdata(FITS_mSCI, ext=0).T
+        PixA_REF = fits.getdata(FITS_REF, memmap=False, ext=0).T
+        PixA_SCI = fits.getdata(FITS_SCI, memmap=False, ext=0).T
+        PixA_mREF = fits.getdata(FITS_mREF, memmap=False, ext=0).T
+        PixA_mSCI = fits.getdata(FITS_mSCI, memmap=False, ext=0).T
 
         if not PixA_REF.flags['C_CONTIGUOUS']:
             PixA_REF = np.ascontiguousarray(PixA_REF, np.float64)
@@ -131,15 +144,29 @@ def CP(FITS_REF, FITS_SCI, FITS_mREF, FITS_mSCI, ForceConv, GKerHW, \
             print('MeLOn CheckPoint: TRIGGER Function Compilations of SFFT-SUBTRACTION!')
 
         Tcomp_start = time.time()
+
+        ### I ADDED DEBUG MEM TRACING STATEMENTS
+        # tracemalloc.start()
         SFFTConfig = SingleSFFTConfigure.SSC(NX=PixA_REF.shape[0], NY=PixA_REF.shape[1], KerHW=KerHW, \
             KerPolyOrder=KerPolyOrder, BGPolyOrder=BGPolyOrder, ConstPhotRatio=ConstPhotRatio, \
             BACKEND_4SUBTRACT=BACKEND_4SUBTRACT, NUM_CPU_THREADS_4SUBTRACT=NUM_CPU_THREADS_4SUBTRACT, \
             VERBOSE_LEVEL=VERBOSE_LEVEL)
 
+        # size, peak = tracemalloc.get_traced_memory()
+        # logger.debug(f'MEMORY IN CUSTOMIZEDPACKET FROM SFFTCONFIG size={size}, peak={peak}')
+        # tracemalloc.clear_traces()
+        # tracemalloc.stop()
+
         if VERBOSE_LEVEL in [1, 2]:
             _message = 'Function Compilations of SFFT-SUBTRACTION TAKES [%.3f s]' %(time.time() - Tcomp_start)
             print('\nMeLOn Report: %s' %_message)
 
+        mempool = cp.get_default_memory_pool()
+        pinned_mempool = cp.get_default_pinned_memory_pool()
+        # print('MEMPOOL USED BYTES', mempool.used_bytes())
+        # print('MEMPOOL TOTAL BYTES', mempool.total_bytes())
+        # print('PINNED MEMPOOL FREE BLOCKS', pinned_mempool.n_free_blocks())
+
         # * Perform SFFT Subtraction
         if ConvdSide == 'REF':
             PixA_mI, PixA_mJ = PixA_mREF, PixA_mSCI
@@ -160,11 +187,26 @@ def CP(FITS_REF, FITS_SCI, FITS_mREF, FITS_mSCI, ForceConv, GKerHW, \
         if VERBOSE_LEVEL in [0, 1, 2]:
             print('MeLOn CheckPoint: TRIGGER SFFT-SUBTRACTION!')
 
+        # print('MEMPOOL USED BYTES', mempool.used_bytes())
+        # print('MEMPOOL TOTAL BYTES', mempool.total_bytes())
+        # print('PINNED MEMPOOL FREE BLOCKS', pinned_mempool.n_free_blocks())
+
         Tsub_start = time.time()
+
+        # tracemalloc.start()
+
         _tmp = GeneralSFFTSubtract.GSS(PixA_I=PixA_I, PixA_J=PixA_J, PixA_mI=PixA_mI, PixA_mJ=PixA_mJ, \
             SFFTConfig=SFFTConfig, ContamMask_I=None, BACKEND_4SUBTRACT=BACKEND_4SUBTRACT, \
             NUM_CPU_THREADS_4SUBTRACT=NUM_CPU_THREADS_4SUBTRACT, VERBOSE_LEVEL=VERBOSE_LEVEL)
 
+        # print('MEMPOOL USED BYTES', mempool.used_bytes())
+        # print('MEMPOOL TOTAL BYTES', mempool.total_bytes())
+        # print('PINNED MEMPOOL FREE BLOCKS', pinned_mempool.n_free_blocks())
+
+        # size, peak = tracemalloc.get_traced_memory()
+        # logger.debug(f'MEMORY IN CUSTOMIZEDPACKET FROM GENERALSFFTSUBTRACT size={size}, peak={peak}')
+        # tracemalloc.clear_traces()
+
         Solution, PixA_DIFF = _tmp[:2]
         if VERBOSE_LEVEL in [1, 2]:
             _message = 'SFFT-SUBTRACTION TAKES [%.3f s]' %(time.time() - Tsub_start)

sfft/utils/ConvKernelConvertion.py

@@ -1,4 +1,5 @@
 import numpy as np
+import cupy
 
 __author__ = "Lei Hu <leihu@andrew.cmu.edu>"
 __version__ = "v1.0"
@@ -17,16 +18,16 @@ def CSZ(ConvKernel, N0, N1):
         L0, L1 = ConvKernel.shape  
         w0, w1 = (L0-1) // 2, (L1-1) // 2      
         pd0, pd1 = N0 - L0, N1 - L1
-        TailZP = np.lib.pad(ConvKernel, ((0, pd0), (0, pd1)), 'constant', constant_values=(0, 0))    # Tail-Zero-Padding
-        KIMG_CSZ = np.roll(np.roll(TailZP, -w0, axis=0), -w1, axis=1)    # Circular-Shift
+        TailZP = cupy.pad(ConvKernel, ((0, pd0), (0, pd1)), 'constant', constant_values=(0, 0))    # Tail-Zero-Padding
+        KIMG_CSZ = cupy.roll(np.roll(TailZP, -w0, axis=0), -w1, axis=1)    # Circular-Shift
         return KIMG_CSZ
 
     def iCSZ(KIMG, L0, L1):
         N0, N1 = KIMG.shape
         w0, w1 = (L0-1) // 2, (L1-1) // 2
-        KIMG_iCSZ = np.roll(np.roll(KIMG, w1, axis=1), w0, axis=0)    # inverse Circular-Shift
+        KIMG_iCSZ = cupy.roll(cupy.roll(KIMG, w1, axis=1), w0, axis=0)    # inverse Circular-Shift
         ConvKernel = KIMG_iCSZ[:L0, :L1]    # Tail-Truncation 
-        lost_weight = 1.0 - np.sum(np.abs(ConvKernel)) / np.sum(np.abs(KIMG_iCSZ))
+        lost_weight = 1.0 - cupy.sum(cupy.abs(ConvKernel)) / cupy.sum(cupy.abs(KIMG_iCSZ))
         print('MeLOn CheckPoint: Tail-Truncation Lost-Weight [%.4f %s] (Absolute Percentage Error) ' %(lost_weight*100, '%'))
         return ConvKernel