test_effect_unit.py

import unittest
import numpy as np
from observation_sim.instruments.chip import effects
import galsim
import matplotlib.pyplot as plt
import os
import sys
import math
import 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)
        # shutter effect normalized image for this chip
        shuttimg = effects.ShutterEffectArr(
            img, t_exp=150, t_shutter=1.3, dist_bearing=735, dt=1E-3)
        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()