test_effect_unit.py

import unittest
import numpy as np
from ObservationSim.Instrument.Chip import Effects
import galsim
import matplotlib.pyplot as plt
import os,sys,math,copy
from numpy.random import Generator, PCG64
import warnings
from astropy.io import fits

warnings.filterwarnings("ignore", '.*Numba.*',)

width = 9216
height = 9232

class DetTest(unittest.TestCase):

    def __init__(self, methodName='runTest'):
        super(DetTest,self).__init__(methodName)
        self.filePath('csst_msc_sim/test_sls_and_straylight')

    def filePath(self, file_name):
        self.datafn = os.path.join(os.getenv('UNIT_TEST_DATA_ROOT'), file_name)
        self.outDataFn = os.path.join(self.datafn,'output')
        if os.path.isdir(self.outDataFn):
            pass
        else:
            os.mkdir(self.outDataFn)
        

    def test_prnu(self):
        '''
        Unit test for PRNU. Expected result: a randomized GS image contains PRNU with sigma=0.01, mean=1.
        '''
        print('PRNU Test:')
        sigma = 0.01
        seed = 20210911
        prnuimg = Effects.PRNU_Img(width, height, sigma=sigma, seed=seed)
        meanval, stdval = np.mean(prnuimg.array), np.std(prnuimg.array)
        print('  Mean & STDDEV of PRNU image are %6.4f & %6.4f.' % (meanval, stdval))
        print('  PRNU Image Array:')
        print('  ',prnuimg.array)
        self.assertTrue(np.abs(meanval-1)<1e-6)
        self.assertTrue(np.abs(stdval-sigma)<0.002)
        print('\nUnit test for PRNU has been passed.')
        del prnuimg

    def test_dark(self):
        '''
        Test add dark current to image. Expected result: an image with dark current 3.4 e- and noise=1.844 e-.
        '''
        rng_poisson = galsim.BaseDeviate(20210911)
        dark_noise = galsim.DeviateNoise(galsim.PoissonDeviate(rng_poisson, 0.02*(150+0.5*40)))
        img = galsim.Image(200,200,dtype=np.float32, init_value=0)
        print('Initial Mean & STD = %6.3f & %6.3f' % (np.mean(img.array), np.std(img.array)))
        img.addNoise(dark_noise)
        meanval = np.mean(img.array)
        stdval = np.std(img.array)
        print('Dark added Mean & STD = %6.3f & %6.3f' % (meanval, stdval))
        self.assertTrue(np.abs(meanval-3.4)<0.05)
        self.assertTrue(np.abs(stdval-1.844)<0.02)
        print('\nUnit test for dark current has been passed.')
        del img

    def test_satu(self):
        '''
        Test saturation and bleeding. Expected result: an image with bleeding effect.
        '''
        img = galsim.Image(500,500,dtype=np.float32)
        star = galsim.Gaussian(flux=60e5,fwhm=3)
        img = star.drawImage(image=img,center=(150,200))
        # gal = galsim.Sersic(n=1, half_light_radius=3,flux=50e5)
        # img = gal.drawImage(image=img,center=(350,300))
        img.addNoise(galsim.GaussianNoise(sigma=7))
        # plt.imshow(img.array)
        # plt.show()
        filename1 = os.path.join(self.outDataFn,'test_satu_initimg.fits')
        img.write(filename1)
        newimg = Effects.SaturBloom(img, fullwell=9e4)
        # plt.imshow(newimg.array)
        # plt.show()
        filename2 = os.path.join(self.outDataFn,'test_satu_bleedimg.fits')
        newimg.write(filename2)
        del img,newimg, star

    def test_nonlinear(self):
        '''
        Test non-linear effect. Expected result: an image with non-linearity effect.
        '''
        imgarr = np.arange(1,9e4,4).reshape((150,150))
        img = galsim.Image(copy.deepcopy(imgarr))
        filename1 = os.path.join(self.outDataFn,'test_nonlinear_initimg.fits')
        img.write(filename1)
        newimg = Effects.NonLinearity(img, beta1=5E-7, beta2=0)
        filename2 = os.path.join(self.outDataFn,'test_nonlinear_finalimg.fits')
        newimg.write(filename2)
        plt.scatter(imgarr.flatten(), newimg.array.flatten(), s=2, alpha=0.5)
        plt.plot([-1e3,9e4],[-1e3,9e4],color='black', lw=1, ls='--')
        plt.xlabel('input (e-)')
        plt.ylabel('output (e-)')
        plt.savefig(os.path.join(self.outDataFn,'test_nonlinearity.png'), dpi=200)
        plt.show()
        del img,newimg,imgarr

    def test_badpixel_HtrDtr(self):
        img = galsim.Image(500,500,init_value=1000)
        rgbadpix = Generator(PCG64(20210911))
        badfraction = 5E-5*(rgbadpix.random()*0.5+0.7)
        img = Effects.DefectivePixels(img, IfHotPix=True, IfDeadPix=True, fraction=badfraction, seed=20210911, biaslevel=0)
        img.write(os.path.join(self.outDataFn,'test_badpixel_HtrDtr.fits'))
        del img
    def test_badpixel_HfsDtr(self):
        img = galsim.Image(500,500,init_value=1000)
        rgbadpix = Generator(PCG64(20210911))
        badfraction = 5E-5*(rgbadpix.random()*0.5+0.7)
        img = Effects.DefectivePixels(img, IfHotPix=False, IfDeadPix=True, fraction=badfraction, seed=20210911, biaslevel=0)
        img.write(os.path.join(self.outDataFn,'test_badpixel_HfsDtr.fits'))
        del img
    def test_badpixel_HtrDfs(self):
        img = galsim.Image(500,500,init_value=1000)
        rgbadpix = Generator(PCG64(20210911))
        badfraction = 5E-5*(rgbadpix.random()*0.5+0.7)
        img = Effects.DefectivePixels(img, IfHotPix=True, IfDeadPix=False, fraction=badfraction, seed=20210911, biaslevel=0)
        img.write(os.path.join(self.outDataFn,'test_badpixel_HtrDfs.fits'))
        del img
    def test_badpixel_HfsDfs(self):
        img = galsim.Image(500,500,init_value=1000)
        rgbadpix = Generator(PCG64(20210911))
        badfraction = 5E-5*(rgbadpix.random()*0.5+0.7)
        img = Effects.DefectivePixels(img, IfHotPix=False, IfDeadPix=False, fraction=badfraction, seed=20210911, biaslevel=0)
        img.write(os.path.join(self.outDataFn,'test_badpixel_HfsDfs.fits'))
        del img

    def test_badlines(self):
        img = galsim.Image(500,500,init_value=-1000)
        img.addNoise(galsim.GaussianNoise(sigma=7))
        newimg = Effects.BadColumns(copy.deepcopy(img), seed=20210911)
        newimg.write(os.path.join(self.outDataFn,'test_badlines.fits'))
        del newimg,img

    # def test_cte(self):
    #     img = galsim.Image(200,200,init_value=1000)
    #     img.array[50,80] = 1e4
    #     img.array[150,150] = 3e4
    #     newimgcol = Effects.CTE_Effect(copy.deepcopy(img),direction='column')
    #     newimgrow = Effects.CTE_Effect(copy.deepcopy(img),direction='row')
    #     newimgcol.write(os.path.join(self.outDataFn,'test_ctecol.fits'))
    #     newimgrow.write(os.path.join(self.outDataFn,'test_cterow.fits'))
    #     del img,newimgcol,newimgrow

    def test_readnoise(self):
        img = galsim.Image(200,200,init_value=1000)
        seed = 20210911
        rng_readout = galsim.BaseDeviate(seed)
        readout_noise = galsim.GaussianNoise(rng=rng_readout, sigma=5)
        img.addNoise(readout_noise)
        img.write(os.path.join(self.outDataFn,'test_readnoise.fits'))
        stdval = np.std(img.array)
        self.assertTrue(np.abs(stdval-5)<0.01*5)
        print('\nUnit test for readout noise has been passed.')
        del img

    def test_addbias(self):
        img = galsim.Image(200,200,init_value=0)
        img = Effects.AddBiasNonUniform16(img,bias_level=500, nsecy = 2, nsecx=8,seed=20210911)
        img.write('./output/test_addbias.fits')
        del img

    def test_apply16gains(self):
        img = galsim.Image(500,500,init_value=100)
        img,_ = Effects.ApplyGainNonUniform16(img, gain=1.5, nsecy=2, nsecx=8, seed=202102)
        img.write(os.path.join(self.outDataFn,'test_apply16gains.fits'))
        rightedge = int(500/8)*8
        print('gain=%6.2f' % 1.5)
        meanimg = np.mean(img.array[:,:rightedge])
        sigmaimg = np.std(img.array[:,:rightedge])
        print('mean, sigma = %6.2f, %6.2f' % (meanimg,sigmaimg))
        self.assertTrue(np.abs(meanimg-100/1.5)<1)
        self.assertTrue(np.abs(sigmaimg/meanimg-0.01)<0.001)
        print('\nUnit test for applying 16 channel gains has been passed.')
        del img


    def test_cosmicray(self):
        attachedSizes = np.loadtxt(os.path.join(self.datafn,'wfc-cr-attachpixel.dat'))
        cr_map,_ = Effects.produceCR_Map(
                xLen=500, yLen=500, exTime=150+0.5*40, 
                cr_pixelRatio=0.003*(1+0.5*40/150), 
                gain=1, attachedSizes=attachedSizes, seed=20210911)
        crimg = galsim.Image(cr_map)
        crimg.write(os.path.join(self.outDataFn,'test_cosmicray.fits'))
        del cr_map,crimg

    def test_shutter(self):
        img = galsim.Image(5000,5000,init_value=1000)
        shuttimg = Effects.ShutterEffectArr(img, t_exp=150, t_shutter=1.3, dist_bearing=735, dt=1E-3)    # shutter effect normalized image for this chip
        img *= shuttimg
        img.write(os.path.join(self.outDataFn,'test_shutter.fits'))
        del img

    def test_vignette(self):
        img = galsim.Image(2000,2000,init_value=1000)
        print(img.bounds)
        # # img.bounds = galsim.BoundsI(1, width, 1, height)
        img.setOrigin(10000,10000)
        flat_img = Effects.MakeFlatSmooth(img.bounds,20210911)
        flat_normal = flat_img / np.mean(flat_img.array)
        flat_normal.write(os.path.join(self.outDataFn,'test_vignette.fits'))
        del flat_img,img,flat_normal



if __name__ == '__main__':
    unittest.main()