add Straylight unittest code

import unittest

from ObservationSim.Straylight import Straylight

import numpy as np

import math

import astropy.constants as cons

import galsim

from astropy.table import Table

from scipy import interpolate

import matplotlib.pyplot as plt

hubbleAverZodiacal = {'nuv':0.0035,'u':0.0163,'g':0.1109,'r':0.1471,'i':0.1568,'z':0.0953,'y':0.0283}

hubbleAverEarthShine = {'nuv':0.00024,'u':0.0051,'g':0.0506,'r':0.0591,'i':0.0568,'z':0.0315,'y':0.0090}

def transRaDec2D(ra, dec):

    x1 = np.cos(dec / 57.2957795) * np.cos(ra / 57.2957795);

    y1 = np.cos(dec / 57.2957795) * np.sin(ra / 57.2957795);

    z1 = np.sin(dec / 57.2957795);

    return np.array([x1, y1, z1])

def getAngle132(x1=0, y1=0, z1=0, x2=0, y2=0, z2=0, x3=0, y3=0, z3=0):

    cosValue = 0;

    angle = 0;

    x11 = x1 - x3;

    y11 = y1 - y3;

    z11 = z1 - z3;

    x22 = x2 - x3;

    y22 = y2 - y3;

    z22 = z2 - z3;

    tt = np.sqrt((x11 * x11 + y11 * y11 + z11 * z11) * (x22 * x22 + y22 * y22 + z22 * z22));

    if (tt == 0):

        return 0;

    cosValue = (x11 * x22 + y11 * y22 + z11 * z22) / tt;

    if (cosValue > 1):

        cosValue = 1;

    if (cosValue < -1):

        cosValue = -1;

    angle = math.acos(cosValue);

    return angle * 360 / (2 * math.pi);

def calculateAnglePwithEarth(sat = np.array([0,0,0]), pointing = np.array([0,0,0]), sun = np.array([0,0,0])):

    modSat = np.sqrt(sat[0]*sat[0] + sat[1]*sat[1]+sat[2]*sat[2])

    modPoint = np.sqrt(pointing[0]*pointing[0] + pointing[1]*pointing[1] + pointing[2]*pointing[2])

    withLocalZenithAngle = (pointing[0] * sat[0] + pointing[1] * sat[1] + pointing[2] * sat[2]) / (modPoint*modSat)

    innerM_sat_sun = sat[0] * sun[0] + sat[1] * sun[1] + sat[2] * sun[2]

    cosAngle = innerM_sat_sun / (modSat * cons.au.value/1000)

    isInSunSide = 1

    if (cosAngle < -0.3385737): #cos109.79

        isInSunSide = -1;

    elif cosAngle >= -0.3385737 and cosAngle <= 0.3385737:

        isInSunSide = 0;

    return math.acos(withLocalZenithAngle)*180/math.pi,isInSunSide

class TestStraylight(unittest.TestCase):

    def __init__(self, methodName='runTest',datFn = '', filter = 'i', grating = "GI"):

        super(TestStraylight,self).__init__(methodName)

        self.pointingData = np.loadtxt(datFn, dtype=np.double)

        self.filter = filter

        self.grating = grating

    def test_EarthShineFilter(self):

        d_sh = self.pointingData.shape

        sl_e_pix = np.zeros([d_sh[0],3],dtype=np.double)

        for i in np.arange(d_sh[0]):

            # if i > 50:

            #     continue

            ju = self.pointingData[i, 5]

            # pointing = transRaDec2D(self.pointingData[i, 0], self.pointingData[i, 1])

            # print(ju, pointing, surveylist[i,3:9])

            sl = Straylight(jtime=ju, sat_pos=self.pointingData[i, 6:9], pointing_radec=np.array([self.pointingData[i, 0], self.pointingData[i, 1]]),sun_pos=self.pointingData[i,9:12])

            e1, py = sl.calculateEarthShineFilter(filter=self.filter)

            earthZenithAngle, isInSunSide = calculateAnglePwithEarth(sat=self.pointingData[i, 6:9], pointing= sl.pointing, sun=self.pointingData[i,9:12])

            # e2, _ = sl.calculateZodiacalFilter2(filter='i', sun_pos=sl.sun_pos)

            # e3 = sl.calculateStarLightFilter(filter='i', pointYaxis=py)

            # e_all = sl.calculateStrayLightFilter(filter='i')

            # s_pix, spec = sl.calculateStrayLightGrating(grating='GI')

            sl_e_pix[i,0] = e1

            sl_e_pix[i, 1] = earthZenithAngle

            sl_e_pix[i, 2] = isInSunSide

        median  = np.median(sl_e_pix[:,0])

        print(' average Earthshine %s: %e' % (self.filter, median))

        self.assertTrue(median-hubbleAverEarthShine[self.filter] < 0.1)

        plt.figure()

        ids1 = sl_e_pix[:, 2] == 1

        ids2 = sl_e_pix[:, 2] != 1

        plt.plot(sl_e_pix[ids1, 0], sl_e_pix[ids1, 1], 'r.')

        plt.plot(sl_e_pix[ids2, 0], sl_e_pix[ids2, 1], 'b.')

        plt.legend(['In Sun Side', 'In Earths shadow'])

        plt.xlabel('straylight-earthshine(e-/pixel/s)')

        plt.ylabel('Angle with local zenith(degree)')

        plt.show()

    def test_ZodiacalFilter(self):

        d_sh = self.pointingData.shape

        sl_e_pix = np.zeros([d_sh[0],2],dtype=np.double)

        for i in np.arange(d_sh[0]):

            ju = self.pointingData[i, 5]

            sl = Straylight(jtime=ju, sat_pos=self.pointingData[i, 6:9], pointing_radec=np.array([self.pointingData[i, 0], self.pointingData[i, 1]]),sun_pos=self.pointingData[i,9:12])

            e1, _ = sl.calculateZodiacalFilter2(filter=self.filter, sun_pos=sl.sun_pos)

            sl_e_pix[i,0] = e1

            sl_e_pix[i,1] = getAngle132(x1=self.pointingData[i,9], y1=self.pointingData[i,10], z1=self.pointingData[i,11], x2=sl.pointing[0],

                                        y2=sl.pointing[1], z2=sl.pointing[2], x3=0, y3=0, z3=0)

        plt.figure()

        plt.plot(sl_e_pix[:, 0], sl_e_pix[:, 1], 'r.')

        plt.xlabel('straylight-zodiacal(e-/pixel/s)')

        plt.ylabel('Angle between pointing and sun(degree)')

        plt.show()

        median  = np.median(sl_e_pix[:,0])

        print(' average Zodiacal %s: %f' % (self.filter, median))

        self.assertTrue(median-hubbleAverZodiacal[self.filter] < 0.1)

    def test_StarFilter(self):

        d_sh = self.pointingData.shape

        sl_e_pix = np.zeros(d_sh[0],dtype=np.double)

        tnum = 10

        for i in np.arange(tnum):

            # if i > 50:

            #     continue

            ju = self.pointingData[i, 5]

            # pointing = transRaDec2D(self.pointingData[i, 0], self.pointingData[i, 1])

            # print(ju, pointing, surveylist[i,3:9])

            sl = Straylight(jtime=ju, sat_pos=self.pointingData[i, 6:9], pointing_radec=np.array([self.pointingData[i, 0], self.pointingData[i, 1]]),sun_pos=self.pointingData[i,9:12])

            e1, py = sl.calculateEarthShineFilter(filter=self.filter)

            # e2, _ = sl.calculateZodiacalFilter2(filter='i', sun_pos=sl.sun_pos)

            e3 = sl.calculateStarLightFilter(filter=self.filter, pointYaxis=py)

            # e_all = sl.calculateStrayLightFilter(filter='i')

            # s_pix, spec = sl.calculateStrayLightGrating(grating='GI')

            sl_e_pix[i] = e3

        median  = np.median(sl_e_pix[0:tnum])

        print(' average Earthshine %s: %e' % (self.filter, median))

        self.assertTrue(median-hubbleAverEarthShine[self.filter] < 0.2)

    def test_GratingStraylight(self):

        d_sh = self.pointingData.shape

        sl_e_pix = np.zeros(d_sh[0],dtype=np.double)

        tnum = 10

        for i in np.arange(tnum):

            # if i > 50:

            #     continue

            ju = self.pointingData[i, 5]

            # pointing = transRaDec2D(self.pointingData[i, 0], self.pointingData[i, 1])

            # print(ju, pointing, surveylist[i,3:9])

            sl = Straylight(jtime=ju, sat_pos=self.pointingData[i, 6:9], pointing_radec=np.array([self.pointingData[i, 0], self.pointingData[i, 1]]),sun_pos=self.pointingData[i,9:12])

            # e1, py = sl.calculateEarthShineFilter(filter=self.filter)

            # e2, _ = sl.calculateZodiacalFilter2(filter='i', sun_pos=sl.sun_pos)

            # e3 = sl.calculateStarLightFilter(filter=self.filter, pointYaxis=py)

            # e_all = sl.calculateStrayLightFilter(filter='i')

            s_pix, spec = sl.calculateStrayLightGrating(grating=self.grating)

            sl_e_pix[i] = s_pix

        plt.figure()

        plt.plot(spec['WAVELENGTH'], spec['FLUX'], 'r')

        plt.xlabel('WAVELENGTH')

        plt.ylabel('F$\lambda$(erg/s/cm2/A/arcsec2)')

        plt.xlim(2000,10000)

        plt.show()

        median  = np.median(sl_e_pix[0:tnum])

        print(' average Earthshine %s: %e' % (self.grating, median))

        self.assertTrue(median < 0.8)

if __name__ == '__main__':

    suit = unittest.TestSuite()

    # case1 = TestStraylight('test_EarthShineFilter',datFn = 'Straylight_test.dat', filter = 'i')

    # suit.addTest(case1)

    # case2 = TestStraylight('test_ZodiacalFilter',datFn = 'Straylight_test.dat',filter = 'i')

    # suit.addTest(case2)

    # case3 = TestStraylight('test_StarFilter', datFn='Straylight_test.dat', filter='i')

    # suit.addTest(case3)

    case4 = TestStraylight('test_GratingStraylight', datFn='Straylight_test.dat', grating = 'GI')

    suit.addTest(case4)

    unittest.TextTestRunner(verbosity=2).run(suit)

 No newline at end of file