effects.py 28.8 KB
Newer Older
Fang Yuedong's avatar
Fang Yuedong committed
1
2
3
4
5
import galsim
from matplotlib.pyplot import flag
import numpy as np
from numpy.core.fromnumeric import mean, size
from numpy.random import Generator, PCG64
Wei Chengliang's avatar
Wei Chengliang committed
6
7
import math
import copy
Fang Yuedong's avatar
Fang Yuedong committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from numba import jit
from astropy import stats


def AddOverscan(GSImage, overscan=1000, gain=1, widthl=27, widthr=27, widtht=8, widthb=8, read_noise=5):
    """
    Add overscan/gain; gain=e-/ADU
    widthl: left pre-scan width
    widthr: right pre-scan width
    widtht: top over-scan width (the top of nd-array with small row-index)
    widthb: bottom over-scan width (the bottom of nd-array with large row-index)
    """
    imgshape = GSImage.array.shape
    newimg = galsim.Image(imgshape[1]+widthl+widthr, imgshape[0]+widtht+widthb, init_value=0)
    rng = galsim.UniformDeviate()
    NoiseOS = galsim.GaussianNoise(rng, sigma=read_noise)
    newimg.addNoise(NoiseOS)
    newimg = (newimg+overscan)/gain
Wei Chengliang's avatar
Wei Chengliang committed
26
    newimg.array[widtht:(widtht+imgshape[0]), widthl:(widthl+imgshape[1])] = GSImage.array
Fang Yuedong's avatar
Fang Yuedong committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
    newimg.wcs = GSImage.wcs
    # if GSImage.wcs is not None:
    #     newwcs = GSImage.wcs.withOrigin(galsim.PositionD(widthl,widtht))
    #     newimg.wcs = newwcs
    # else:
    #     pass
    return newimg


def DefectivePixels(GSImage, IfHotPix=True, IfDeadPix=True, fraction=1E-4, seed=20210304, biaslevel=0):
    # Also called bad pixels, including hot pixels and dead pixels
    # Hot Pixel > 20e-/s
    # Dead Pixel < 70%*Mean
    rgf = Generator(PCG64(int(seed*1.1)))
Wei Chengliang's avatar
Wei Chengliang committed
41
    if IfHotPix is True and IfDeadPix is True:
Fang Yuedong's avatar
Fang Yuedong committed
42
        HotFraction = rgf.random()             # fraction in total bad pixels
Wei Chengliang's avatar
Wei Chengliang committed
43
    elif IfHotPix is False and IfDeadPix is False:
Fang Yuedong's avatar
Fang Yuedong committed
44
        return GSImage
Wei Chengliang's avatar
Wei Chengliang committed
45
    elif IfHotPix is True:
Fang Yuedong's avatar
Fang Yuedong committed
46
47
48
49
50
51
52
53
        HotFraction = 1
    else:
        HotFraction = 0

    NPix = GSImage.array.size
    NPixBad = int(NPix*fraction)
    NPixHot = int(NPix*fraction*HotFraction)
    NPixDead = NPixBad-NPixHot
Wei Chengliang's avatar
Wei Chengliang committed
54
55

    NPix_y, NPix_x = GSImage.array.shape
Fang Yuedong's avatar
Fang Yuedong committed
56
57
    mean = np.mean(GSImage.array)
    rgp = Generator(PCG64(int(seed)))
Wei Chengliang's avatar
Wei Chengliang committed
58
    yxposfrac = rgp.random((NPixBad, 2))
Wei Chengliang's avatar
Wei Chengliang committed
59
60
61
62
    YPositHot = np.array(NPix_y*yxposfrac[0:NPixHot, 0]).astype(np.int32)
    XPositHot = np.array(NPix_x*yxposfrac[0:NPixHot, 1]).astype(np.int32)
    YPositDead = np.array(NPix_y*yxposfrac[NPixHot:, 0]).astype(np.int32)
    XPositDead = np.array(NPix_x*yxposfrac[NPixHot:, 1]).astype(np.int32)
Fang Yuedong's avatar
Fang Yuedong committed
63
64
65

    rgh = Generator(PCG64(int(seed*1.2)))
    rgd = Generator(PCG64(int(seed*1.3)))
Wei Chengliang's avatar
Wei Chengliang committed
66
67
68
69
    if IfHotPix is True:
        GSImage.array[YPositHot, XPositHot] += rgh.gamma(2, 25*150, size=NPixHot)
    if IfDeadPix is True:
        GSImage.array[YPositDead, XPositDead] = rgd.random(NPixDead)*(mean-biaslevel)*0.7+biaslevel+rgp.standard_normal()*5
Fang Yuedong's avatar
Fang Yuedong committed
70
71
72
73
74
    return GSImage


def BadColumns(GSImage, seed=20240309, chipid=1, logger=None):
    # Set bad column values
Wei Chengliang's avatar
Wei Chengliang committed
75
    ysize, xsize = GSImage.array.shape
Fang Yuedong's avatar
Fang Yuedong committed
76
77
78
79
80
81
82
83
84
85
    subarr = GSImage.array[int(ysize*0.1):int(ysize*0.12), int(xsize*0.1):int(xsize*0.12)]
    subarr = stats.sigma_clip(subarr, sigma=4, cenfunc='median', maxiters=3, masked=False)
    meanimg = np.median(subarr)
    stdimg = np.std(subarr)
    seed += chipid
    rgn = Generator(PCG64(int(seed)))
    rgcollen = Generator(PCG64(int(seed*1.1)))
    rgxpos = Generator(PCG64(int(seed*1.2)))
    rgdn = Generator(PCG64(int(seed*1.3)))

Wei Chengliang's avatar
Wei Chengliang committed
86
87
    nbadsecA, nbadsecD = rgn.integers(low=1, high=5, size=2)
    collen = rgcollen.integers(low=int(ysize*0.1), high=int(ysize*0.7), size=(nbadsecA+nbadsecD))
Fang Yuedong's avatar
Fang Yuedong committed
88
89
90
91
92
93
94
    xposit = rgxpos.integers(low=int(xsize*0.05), high=int(xsize*0.95), size=(nbadsecA+nbadsecD))
    if logger is not None:
        logger.info(xposit+1)
    else:
        print(xposit+1)
    # signs = 2*rgdn.integers(0,2,size=(nbadsecA+nbadsecD))-1
    # if meanimg>0:
Wei Chengliang's avatar
Wei Chengliang committed
95
    dn = rgdn.integers(low=np.abs(meanimg)*1.3+50, high=np.abs(meanimg)*2+150, size=(nbadsecA+nbadsecD))  # *signs
Fang Yuedong's avatar
Fang Yuedong committed
96
97
98
    # elif meanimg<0:
    #     dn = rgdn.integers(low=meanimg*2-150, high=meanimg*1.3-50, size=(nbadsecA+nbadsecD)) #*signs
    for badcoli in range(nbadsecA):
Wei Chengliang's avatar
Wei Chengliang committed
99
        GSImage.array[(ysize-collen[badcoli]):ysize, xposit[badcoli]:(xposit[badcoli]+1)] = (np.abs(np.random.normal(0, stdimg*2, (collen[badcoli], 1)))+dn[badcoli])
Fang Yuedong's avatar
Fang Yuedong committed
100
    for badcoli in range(nbadsecD):
Wei Chengliang's avatar
Wei Chengliang committed
101
        GSImage.array[0:collen[badcoli+nbadsecA], xposit[badcoli+nbadsecA]:(xposit[badcoli+nbadsecA]+1)] = (np.abs(np.random.normal(0, stdimg*2, (collen[badcoli+nbadsecA], 1)))+dn[badcoli+nbadsecA])
Fang Yuedong's avatar
Fang Yuedong committed
102
103
104
    return GSImage


Wei Chengliang's avatar
Wei Chengliang committed
105
def AddBiasNonUniform16(GSImage, bias_level=500, nsecy=2, nsecx=8, seed=202102, logger=None):
Fang Yuedong's avatar
Fang Yuedong committed
106
107
108
    # Generate Bias and its non-uniformity, and add the 16 bias values to the GS-Image
    rg = Generator(PCG64(int(seed)))
    Random16 = (rg.random(nsecy*nsecx)-0.5)*20
Wei Chengliang's avatar
Wei Chengliang committed
109
110
111
    if int(bias_level) == 0:
        BiasLevel = np.zeros((nsecy, nsecx))
    elif bias_level > 0:
Wei Chengliang's avatar
Wei Chengliang committed
112
        BiasLevel = Random16.reshape((nsecy, nsecx)) + bias_level
Fang Yuedong's avatar
Fang Yuedong committed
113
114
115
116
    if logger is not None:
        msg = str(" Biases of 16 channels: " + str(BiasLevel))
        logger.info(msg)
    else:
Wei Chengliang's avatar
Wei Chengliang committed
117
        print(" Biases of 16 channels:\n", BiasLevel)
Fang Yuedong's avatar
Fang Yuedong committed
118
119
120
121
122
    arrshape = GSImage.array.shape
    secsize_x = int(arrshape[1]/nsecx)
    secsize_y = int(arrshape[0]/nsecy)
    for rowi in range(nsecy):
        for coli in range(nsecx):
Wei Chengliang's avatar
Wei Chengliang committed
123
            GSImage.array[rowi*secsize_y:(rowi+1)*secsize_y, coli*secsize_x:(coli+1)*secsize_x] += BiasLevel[rowi, coli]
Fang Yuedong's avatar
Fang Yuedong committed
124
125
126
127
128
    return GSImage


def MakeBiasNcomb(npix_x, npix_y, bias_level=500, ncombine=1, read_noise=5, gain=1, seed=202102, logger=None):
    # Start with 0 value bias GS-Image
Wei Chengliang's avatar
Wei Chengliang committed
129
    ncombine = int(ncombine)
Fang Yuedong's avatar
Fang Yuedong committed
130
    BiasSngImg0 = galsim.Image(npix_x, npix_y, init_value=0)
Wei Chengliang's avatar
Wei Chengliang committed
131
132
    BiasSngImg = AddBiasNonUniform16(BiasSngImg0,
                                     bias_level=bias_level,
Wei Chengliang's avatar
Wei Chengliang committed
133
                                     nsecy=2, nsecx=8,
Wei Chengliang's avatar
Wei Chengliang committed
134
135
                                     seed=int(seed),
                                     logger=logger)
Fang Yuedong's avatar
Fang Yuedong committed
136
137
138
139
140
141
142
    BiasCombImg = BiasSngImg*ncombine
    rng = galsim.UniformDeviate()
    NoiseBias = galsim.GaussianNoise(rng=rng, sigma=read_noise*ncombine**0.5)
    BiasCombImg.addNoise(NoiseBias)
    if ncombine == 1:
        BiasTag = 'Single'
        pass
Wei Chengliang's avatar
Wei Chengliang committed
143
    elif ncombine > 1:
Fang Yuedong's avatar
Fang Yuedong committed
144
145
146
147
148
149
150
        BiasCombImg /= ncombine
        BiasTag = 'Combine'
    # BiasCombImg.replaceNegative(replace_value=0)
    # BiasCombImg.quantize()
    return BiasCombImg, BiasTag


Wei Chengliang's avatar
Wei Chengliang committed
151
def ApplyGainNonUniform16(GSImage, gain=1, nsecy=2, nsecx=8, seed=202102, logger=None):
Fang Yuedong's avatar
Fang Yuedong committed
152
153
154
    # Generate Gain non-uniformity, and multipy the different factors (mean~1 with sigma~1%) to the GS-Image
    rg = Generator(PCG64(int(seed)))
    Random16 = (rg.random(nsecy*nsecx)-0.5)*0.04+1   # sigma~1%
Wei Chengliang's avatar
Wei Chengliang committed
155
    Gain16 = Random16.reshape((nsecy, nsecx))/gain
Fang Yuedong's avatar
Fang Yuedong committed
156
157
158
159
160
    gain_array = np.ones(nsecy*nsecx)*gain
    if logger is not None:
        msg = str("Gain of 16 channels: " + str(Gain16))
        logger.info(msg)
    else:
Wei Chengliang's avatar
Wei Chengliang committed
161
        print("Gain of 16 channels: ", Gain16)
Fang Yuedong's avatar
Fang Yuedong committed
162
163
164
165
166
    arrshape = GSImage.array.shape
    secsize_x = int(arrshape[1]/nsecx)
    secsize_y = int(arrshape[0]/nsecy)
    for rowi in range(nsecy):
        for coli in range(nsecx):
Wei Chengliang's avatar
Wei Chengliang committed
167
168
            GSImage.array[rowi*secsize_y:(rowi+1)*secsize_y, coli*secsize_x:(coli+1)*secsize_x] *= Gain16[rowi, coli]
            gain_array[rowi*nsecx+coli] = 1/Gain16[rowi, coli]
Fang Yuedong's avatar
Fang Yuedong committed
169
170
171
    return GSImage, gain_array


Wei Chengliang's avatar
Wei Chengliang committed
172
def GainsNonUniform16(GSImage, gain=1, nsecy=2, nsecx=8, seed=202102, logger=None):
Fang Yuedong's avatar
Fang Yuedong committed
173
174
175
    # Generate Gain non-uniformity, and multipy the different factors (mean~1 with sigma~1%) to the GS-Image
    rg = Generator(PCG64(int(seed)))
    Random16 = (rg.random(nsecy*nsecx)-0.5)*0.04+1   # sigma~1%
Wei Chengliang's avatar
Wei Chengliang committed
176
    Gain16 = Random16.reshape((nsecy, nsecx))/gain
Fang Yuedong's avatar
Fang Yuedong committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
    if logger is not None:
        msg = str(seed-20210202, "Gains of 16 channels: " + str(Gain16))
        logger.info(msg)
    else:
        print(seed-20210202, "Gains of 16 channels:\n", Gain16)
    # arrshape = GSImage.array.shape
    # secsize_x = int(arrshape[1]/nsecx)
    # secsize_y = int(arrshape[0]/nsecy)
    # for rowi in range(nsecy):
    #     for coli in range(nsecx):
    #         GSImage.array[rowi*secsize_y:(rowi+1)*secsize_y,coli*secsize_x:(coli+1)*secsize_x] *= Gain16[rowi,coli]
    # return GSImage
    return Gain16


def MakeFlatSmooth(GSBounds, seed):
    rg = Generator(PCG64(int(seed)))
Wei Chengliang's avatar
Wei Chengliang committed
194
    r1, r2, r3, r4 = rg.random(4)
Fang Yuedong's avatar
Fang Yuedong committed
195
196
197
198
    a1 = -0.5 + 0.2*r1
    a2 = -0.5 + 0.2*r2
    a3 = r3+5
    a4 = r4+5
Wei Chengliang's avatar
Wei Chengliang committed
199
    xmin, xmax, ymin, ymax = GSBounds.getXMin(), GSBounds.getXMax(), GSBounds.getYMin(), GSBounds.getYMax()
Fang Yuedong's avatar
Fang Yuedong committed
200
201
    Flty, Fltx = np.mgrid[ymin:(ymax+1), xmin:(xmax+1)]
    rg = Generator(PCG64(int(seed)))
Wei Chengliang's avatar
Wei Chengliang committed
202
    p1, p2, bg = rg.poisson(1000, 3)
Fang Yuedong's avatar
Fang Yuedong committed
203
204
205
206
207
208
    Fltz = 0.6*1e-7*(a1 * (Fltx-p1) ** 2 + a2 * (Flty-p2) ** 2 - a3*Fltx - a4*Flty) + bg*20
    FlatImg = galsim.ImageF(Fltz)
    return FlatImg


def MakeFlatNcomb(flat_single_image, ncombine=1, read_noise=5, gain=1, overscan=500, biaslevel=500, seed_bias=20210311, logger=None):
Wei Chengliang's avatar
Wei Chengliang committed
209
    ncombine = int(ncombine)
Fang Yuedong's avatar
Fang Yuedong committed
210
211
212
213
214
215
216
217
218
    FlatCombImg = flat_single_image*ncombine
    rng = galsim.UniformDeviate()
    NoiseFlatPoi = galsim.PoissonNoise(rng=rng, sky_level=0)
    FlatCombImg.addNoise(NoiseFlatPoi)
    NoiseFlatReadN = galsim.GaussianNoise(rng=rng, sigma=read_noise*ncombine**0.5)
    FlatCombImg.addNoise(NoiseFlatReadN)
    # NoiseFlat = galsim.CCDNoise(rng, gain=gain, read_noise=read_noise*ncombine**0.5, sky_level=0)
    for i in range(ncombine):
        FlatCombImg = AddBiasNonUniform16(
Wei Chengliang's avatar
Wei Chengliang committed
219
220
221
            FlatCombImg,
            bias_level=biaslevel,
            nsecy=2, nsecx=8,
Fang Yuedong's avatar
Fang Yuedong committed
222
223
224
225
226
            seed=seed_bias,
            logger=logger)
    if ncombine == 1:
        FlatTag = 'Single'
        pass
Wei Chengliang's avatar
Wei Chengliang committed
227
    elif ncombine > 1:
Fang Yuedong's avatar
Fang Yuedong committed
228
229
230
231
232
233
234
235
        FlatCombImg /= ncombine
        FlatTag = 'Combine'
    # FlatCombImg.replaceNegative(replace_value=0)
    # FlatCombImg.quantize()
    return FlatCombImg, FlatTag


def MakeDarkNcomb(npix_x, npix_y, overscan=500, bias_level=500, seed_bias=202102, darkpsec=0.02, exptime=150, ncombine=10, read_noise=5, gain=1, logger=None):
Wei Chengliang's avatar
Wei Chengliang committed
236
    ncombine = int(ncombine)
Fang Yuedong's avatar
Fang Yuedong committed
237
238
239
240
241
242
243
244
245
246
    darkpix = darkpsec*exptime
    DarkSngImg = galsim.Image(npix_x, npix_y, init_value=darkpix)
    rng = galsim.UniformDeviate()
    NoiseDarkPoi = galsim.PoissonNoise(rng=rng, sky_level=0)
    NoiseReadN = galsim.GaussianNoise(rng=rng, sigma=read_noise*ncombine**0.5)
    DarkCombImg = DarkSngImg*ncombine
    DarkCombImg.addNoise(NoiseDarkPoi)
    DarkCombImg.addNoise(NoiseReadN)
    for i in range(ncombine):
        DarkCombImg = AddBiasNonUniform16(
Wei Chengliang's avatar
Wei Chengliang committed
247
248
            DarkCombImg,
            bias_level=bias_level,
Wei Chengliang's avatar
Wei Chengliang committed
249
            nsecy=2, nsecx=8,
Fang Yuedong's avatar
Fang Yuedong committed
250
251
252
253
254
            seed=int(seed_bias),
            logger=logger)
    if ncombine == 1:
        DarkTag = 'Single'
        pass
Wei Chengliang's avatar
Wei Chengliang committed
255
    elif ncombine > 1:
Fang Yuedong's avatar
Fang Yuedong committed
256
257
258
259
260
261
262
263
264
        DarkCombImg /= ncombine
        DarkTag = 'Combine'
    # DarkCombImg.replaceNegative(replace_value=0)
    # DarkCombImg.quantize()
    return DarkCombImg, DarkTag


def PRNU_Img(xsize, ysize, sigma=0.01, seed=202101):
    rg = Generator(PCG64(int(seed)))
Wei Chengliang's avatar
Wei Chengliang committed
265
    prnuarr = rg.normal(1, sigma, (ysize, xsize))
Fang Yuedong's avatar
Fang Yuedong committed
266
267
268
269
    prnuimg = galsim.ImageF(prnuarr)
    return prnuimg


Wei Chengliang's avatar
Wei Chengliang committed
270
def NonLinear_f(x, beta_1, beta_2):
Wei Chengliang's avatar
Wei Chengliang committed
271
272
273
    return x - beta_1 * x * x + beta_2 * x * x * x


Fang Yuedong's avatar
Fang Yuedong committed
274
def NonLinearity(GSImage, beta1=5E-7, beta2=0):
Wei Chengliang's avatar
Wei Chengliang committed
275
    # NonLinear_f = lambda x, beta_1, beta_2: x - beta_1*x*x + beta_2*x*x*x
Fang Yuedong's avatar
Fang Yuedong committed
276
277
278
279
    GSImage.applyNonlinearity(NonLinear_f, beta1, beta2)
    return GSImage


Wei Chengliang's avatar
Wei Chengliang committed
280
# Saturation & Bleeding Start#
Fang Yuedong's avatar
Fang Yuedong committed
281
def BleedingTrail(aa, yy):
Wei Chengliang's avatar
Wei Chengliang committed
282
283
    if aa < 0.2:
        aa = 0.2
Fang Yuedong's avatar
Fang Yuedong committed
284
285
286
287
    else:
        pass
    try:
        fy = 0.5*(math.exp(math.log(yy+1)**3/aa)+np.exp(-1*math.log(yy+1)**3/aa))
Wei Chengliang's avatar
Wei Chengliang committed
288
        faa = 0.5*(math.e+1/math.e)
Fang Yuedong's avatar
Fang Yuedong committed
289
290
291
292
293
294
295
        trail_frac = 1-0.1*(fy-1)/(faa-1)
    except Exception as e:
        print(e)
        trail_frac = 1

    return trail_frac

Wei Chengliang's avatar
Wei Chengliang committed
296

Fang Yuedong's avatar
Fang Yuedong committed
297
298
299
300
def MakeTrail(imgarr, satuyxtuple, charge, fullwell=9e4, direction='up', trailcutfrac=0.9):
    '''
    direction: "up" or "down". For "up", bleeds along Y-decreasing direction; for "down", bleeds along Y-increasing direction.
    '''
Wei Chengliang's avatar
Wei Chengliang committed
301
    yi, xi = satuyxtuple
Fang Yuedong's avatar
Fang Yuedong committed
302
303
304
    aa = np.log(charge/fullwell)**3              # scale length of the bleeding trail
    yy = 1

Wei Chengliang's avatar
Wei Chengliang committed
305
306
    while charge > 0:
        if yi < 0 or yi > imgarr.shape[0]-1:
Fang Yuedong's avatar
Fang Yuedong committed
307
            break
Wei Chengliang's avatar
Wei Chengliang committed
308
309
        if yi == 0 or yi == imgarr.shape[0]-1:
            imgarr[yi, xi] = fullwell
Fang Yuedong's avatar
Fang Yuedong committed
310
            break
Wei Chengliang's avatar
Wei Chengliang committed
311
312
313
        if direction == 'up':
            if imgarr[yi-1, xi] >= fullwell:
                imgarr[yi, xi] = fullwell
Wei Chengliang's avatar
Wei Chengliang committed
314
                yi -= 1
Fang Yuedong's avatar
Fang Yuedong committed
315
                continue
Wei Chengliang's avatar
Wei Chengliang committed
316
317
318
319
        elif direction == 'down':
            if imgarr[yi+1, xi] >= fullwell:
                imgarr[yi, xi] = fullwell
                yi += 1
Fang Yuedong's avatar
Fang Yuedong committed
320
                continue
Wei Chengliang's avatar
Wei Chengliang committed
321
        if aa <= 1:
Wei Chengliang's avatar
Wei Chengliang committed
322
323
324
325
326
327
328
            while imgarr[yi, xi] >= fullwell:
                imgarr[yi, xi] = fullwell
                if direction == 'up':
                    imgarr[yi-1, xi] += charge
                    charge = imgarr[yi-1, xi]-fullwell
                    yi -= 1
                    if yi < 0:
Fang Yuedong's avatar
Fang Yuedong committed
329
                        break
Wei Chengliang's avatar
Wei Chengliang committed
330
331
332
333
334
                elif direction == 'down':
                    imgarr[yi+1, xi] += charge
                    charge = imgarr[yi+1, xi]-fullwell
                    yi += 1
                    if yi > imgarr.shape[0]:
Fang Yuedong's avatar
Fang Yuedong committed
335
336
337
                        break
        else:
            # calculate bleeding trail:
Wei Chengliang's avatar
Wei Chengliang committed
338
            trail_frac = BleedingTrail(aa, yy)
Fang Yuedong's avatar
Fang Yuedong committed
339
340

            # put charge upwards
Wei Chengliang's avatar
Wei Chengliang committed
341
342
343
344
345
346
            if trail_frac >= 0.99:
                imgarr[yi, xi] = fullwell
                if direction == 'up':
                    yi -= 1
                elif direction == 'down':
                    yi += 1
Fang Yuedong's avatar
Fang Yuedong committed
347
348
                yy += 1
            else:
Wei Chengliang's avatar
Wei Chengliang committed
349
                if trail_frac < trailcutfrac:
Fang Yuedong's avatar
Fang Yuedong committed
350
351
                    break
                charge = fullwell*trail_frac
Wei Chengliang's avatar
Wei Chengliang committed
352
353
354
355
356
357
358
359
                imgarr[yi, xi] += charge
                if imgarr[yi, xi] > fullwell:
                    imgarr[yi, xi] = fullwell

                if direction == 'up':
                    yi -= 1
                elif direction == 'down':
                    yi += 1
Fang Yuedong's avatar
Fang Yuedong committed
360
361
362
363
364
365
                yy += 1

    return imgarr


def ChargeFlow(imgarr, fullwell=9E4):
Wei Chengliang's avatar
Wei Chengliang committed
366
    size_y, size_x = imgarr.shape
Wei Chengliang's avatar
Wei Chengliang committed
367
    satupos_y, satupos_x = np.where(imgarr > fullwell)
Fang Yuedong's avatar
Fang Yuedong committed
368

Wei Chengliang's avatar
Wei Chengliang committed
369
    if satupos_y.shape[0] == 0:
Fang Yuedong's avatar
Fang Yuedong committed
370
371
372
373
374
375
376
377
378
        # make no change for the image array
        return imgarr
    elif satupos_y.shape[0]/imgarr.size > 0.5:
        imgarr.fill(fullwell)
        return imgarr

    chargedict = {}
    imgarrorig = copy.deepcopy(imgarr)

Wei Chengliang's avatar
Wei Chengliang committed
379
380
381
    for yi, xi in zip(satupos_y, satupos_x):
        yxidx = ''.join([str(yi), str(xi)])
        chargedict[yxidx] = imgarrorig[yi, xi]-fullwell
Fang Yuedong's avatar
Fang Yuedong committed
382

Wei Chengliang's avatar
Wei Chengliang committed
383
384
    for yi, xi in zip(satupos_y, satupos_x):
        yxidx = ''.join([str(yi), str(xi)])
Fang Yuedong's avatar
Fang Yuedong committed
385
386
387
388
389
390
        satcharge = chargedict[yxidx]
        chargeup = ((np.random.random()-0.5)*0.05+0.5)*satcharge
        chargedn = satcharge - chargeup

        try:
            # Charge Clump moves up
Wei Chengliang's avatar
Wei Chengliang committed
391
            if yi >= 0 and yi < imgarr.shape[0]:
Wei Chengliang's avatar
Wei Chengliang committed
392
                imgarr = MakeTrail(imgarr, (yi, xi), chargeup, fullwell=9e4, direction='up', trailcutfrac=0.9)
Fang Yuedong's avatar
Fang Yuedong committed
393
                # Charge Clump moves down
Wei Chengliang's avatar
Wei Chengliang committed
394
                imgarr = MakeTrail(imgarr, (yi, xi), chargedn, fullwell=9e4, direction='down', trailcutfrac=0.9)
Fang Yuedong's avatar
Fang Yuedong committed
395
        except Exception as e:
Wei Chengliang's avatar
Wei Chengliang committed
396
            print(e, '@pix ', (yi+1, xi+1))
Fang Yuedong's avatar
Fang Yuedong committed
397
398
399
            return imgarr
    return imgarr

Wei Chengliang's avatar
Wei Chengliang committed
400

Fang Yuedong's avatar
Fang Yuedong committed
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
def SaturBloom(GSImage, nsect_x=1, nsect_y=1, fullwell=9e4):
    """
    To simulate digital detector's saturation and blooming effect. The blooming is along the read-out direction, perpendicular to the charge transfer direction. Charge clumpy overflows the pixel well will flow to two oposite directions with nearly same charges.
    Work together with chargeflow() function.
    Parameters:
      GSImage: a GalSim Image of a chip;
      nsect_x: number of sections divided along x direction;
      nsect_y: number of sections divided along y direction.
    Return: a saturated image array with the same size of input.
    """
    imgarr = GSImage.array
    size_y, size_x = imgarr.shape
    subsize_y = int(size_y/nsect_y)
    subsize_x = int(size_x/nsect_x)

    for i in range(nsect_y):
        for j in range(nsect_x):
            subimg = imgarr[subsize_y*i:subsize_y*(i+1), subsize_x*j:subsize_x*(j+1)]

            subimg = ChargeFlow(subimg, fullwell=fullwell)

            imgarr[subsize_y*i:subsize_y*(i+1), subsize_x*j:subsize_x*(j+1)] = subimg

    return GSImage

Wei Chengliang's avatar
Wei Chengliang committed
426
#    Saturation & Bleeding End    #
Fang Yuedong's avatar
Fang Yuedong committed
427
428
429
430
431


def readout16(GSImage, rowi=0, coli=0, overscan_value=0):
    # readout image as 16 outputs of sub-images plus prescan & overscan.
    # assuming image width and height sizes are both even.
Wei Chengliang's avatar
Wei Chengliang committed
432
    # assuming image has 2 columns and 8 rows of output channels.
Fang Yuedong's avatar
Fang Yuedong committed
433
434
435
436
437
    # 00  01
    # 10  11
    # 20  21
    # ...
    # return: GS Image Object
Wei Chengliang's avatar
Wei Chengliang committed
438
    npix_y, npix_x = GSImage.array.shape
Fang Yuedong's avatar
Fang Yuedong committed
439
440
441
    subheight = int(8+npix_y/2+8)
    subwidth = int(16+npix_x/8+27)
    OutputSubimg = galsim.ImageUS(subwidth, subheight, init_value=overscan_value)
Wei Chengliang's avatar
Wei Chengliang committed
442
    if rowi < 4 and coli == 0:
Fang Yuedong's avatar
Fang Yuedong committed
443
444
445
        subbounds = galsim.BoundsI(1, int(npix_x/2),  int(npix_y/8*rowi+1), int(npix_y/8*(rowi+1)))
        subbounds = subbounds.shift(galsim.PositionI(GSImage.bounds.getXMin()-1, GSImage.bounds.getYMin()-1))
        subimg = GSImage[subbounds]
Wei Chengliang's avatar
Wei Chengliang committed
446
447
        OutputSubimg.array[27:int(npix_y/8)+27, 8:int(npix_x/2)+8] = subimg.array
    elif rowi < 4 and coli == 1:
Fang Yuedong's avatar
Fang Yuedong committed
448
449
450
        subbounds = galsim.BoundsI(npix_x/2+1, npix_x,  npix_y/8*rowi+1, npix_y/8*(rowi+1))
        subbounds = subbounds.shift(galsim.PositionI(GSImage.bounds.getXMin()-1, GSImage.bounds.getYMin()-1))
        subimg = GSImage[subbounds]
Wei Chengliang's avatar
Wei Chengliang committed
451
452
        OutputSubimg.array[27:int(npix_y/8)+27, 8:int(npix_x/2)+8] = subimg.array
    elif rowi >= 4 and rowi < 8 and coli == 0:
Fang Yuedong's avatar
Fang Yuedong committed
453
454
455
        subbounds = galsim.BoundsI(1, npix_x/2,  npix_y/8*rowi+1, npix_y/8*(rowi+1))
        subbounds = subbounds.shift(galsim.PositionI(GSImage.bounds.getXMin()-1, GSImage.bounds.getYMin()-1))
        subimg = GSImage[subbounds]
Wei Chengliang's avatar
Wei Chengliang committed
456
457
        OutputSubimg.array[16:int(npix_y/8)+16, 8:int(npix_x/2)+8] = subimg.array
    elif rowi >= 4 and rowi < 8 and coli == 1:
Fang Yuedong's avatar
Fang Yuedong committed
458
459
460
        subbounds = galsim.BoundsI(npix_x/2+1, npix_x,  npix_y/8*rowi+1, npix_y/8*(rowi+1))
        subbounds = subbounds.shift(galsim.PositionI(GSImage.bounds.getXMin()-1, GSImage.bounds.getYMin()-1))
        subimg = GSImage[subbounds]
Wei Chengliang's avatar
Wei Chengliang committed
461
        OutputSubimg.array[16:int(npix_y/8)+16, 8:int(npix_x/2)+8] = subimg.array
Fang Yuedong's avatar
Fang Yuedong committed
462
463
464
465
466
467
468
469
470
    else:
        print("\n\033[31mError: "+"Wrong rowi or coli assignment. Permitted: 0<=rowi<=7, 0<=coli<=1."+"\033[0m\n")
        return OutputSubimg
    return OutputSubimg


def CTE_Effect(GSImage, threshold=27, direction='column'):
    # Devide the image into 4 sections and apply CTE effect with different trail directions.
    # GSImage: a GalSim Image object.
Wei Chengliang's avatar
Wei Chengliang committed
471
    size_y, size_x = GSImage.array.shape
Fang Yuedong's avatar
Fang Yuedong committed
472
473
474
475
    size_sect_y = int(size_y/2)
    size_sect_x = int(size_x/2)
    imgarr = GSImage.array
    if direction == 'column':
Wei Chengliang's avatar
Wei Chengliang committed
476
        imgarr[0:size_sect_y, :] = CTEModelColRow(imgarr[0:size_sect_y, :], trail_direction='down', direction='column', threshold=threshold)
Wei Chengliang's avatar
Wei Chengliang committed
477
        imgarr[size_sect_y:size_y, :] = CTEModelColRow(imgarr[size_sect_y:size_y, :], trail_direction='up', direction='column', threshold=threshold)
Fang Yuedong's avatar
Fang Yuedong committed
478
    elif direction == 'row':
Wei Chengliang's avatar
Wei Chengliang committed
479
480
        imgarr[:, 0:size_sect_x] = CTEModelColRow(imgarr[:, 0:size_sect_x], trail_direction='right', direction='row', threshold=threshold)
        imgarr[:, size_sect_x:size_x] = CTEModelColRow(imgarr[:, size_sect_x:size_x], trail_direction='left', direction='row', threshold=threshold)
Fang Yuedong's avatar
Fang Yuedong committed
481
482
483
484
    return GSImage


@jit()
Wei Chengliang's avatar
Wei Chengliang committed
485
def CTEModelColRow(img, trail_direction='up', direction='column', threshold=27):
Fang Yuedong's avatar
Fang Yuedong committed
486

Wei Chengliang's avatar
Wei Chengliang committed
487
488
489
    # total trail flux vs (pixel flux)^1/2 is approximately linear
    # total trail flux = trail_a * (pixel flux)^1/2 + trail_b
    # trail pixel flux = pow(0.5,x)/0.5, normalize to 1
Wei Chengliang's avatar
Wei Chengliang committed
490
    trail_a = 5.651803799619966
Fang Yuedong's avatar
Fang Yuedong committed
491
492
493
494
495
    trail_b = -2.671933068990729

    sh1 = img.shape[0]
    sh2 = img.shape[1]
    n_img = img*0
Wei Chengliang's avatar
Wei Chengliang committed
496
    idx = np.where(img < threshold)
Fang Yuedong's avatar
Fang Yuedong committed
497
498
    if len(idx[0]) == 0:
        pass
Wei Chengliang's avatar
Wei Chengliang committed
499
    elif len(idx[0]) > 0:
Fang Yuedong's avatar
Fang Yuedong committed
500
501
        n_img[idx] = img[idx]

Wei Chengliang's avatar
Wei Chengliang committed
502
    yidx, xidx = np.where(img >= threshold)
Fang Yuedong's avatar
Fang Yuedong committed
503
504
    if len(yidx) == 0:
        pass
Wei Chengliang's avatar
Wei Chengliang committed
505
    elif len(yidx) > 0:
Fang Yuedong's avatar
Fang Yuedong committed
506
        # print(index)
Wei Chengliang's avatar
Wei Chengliang committed
507
508
        for i, j in zip(yidx, xidx):
            f = img[i, j]
Fang Yuedong's avatar
Fang Yuedong committed
509
510
511
512
            trail_f = (np.sqrt(f)*trail_a + trail_b)*0.5
            # trail_f=5E-5*f**1.5
            xy_num = 10
            all_trail = np.zeros(xy_num)
Wei Chengliang's avatar
Wei Chengliang committed
513

Wei Chengliang's avatar
Wei Chengliang committed
514
            xy_upstr = np.arange(1, xy_num, 1)
Fang Yuedong's avatar
Fang Yuedong committed
515
516
517
518

            # all_trail_pix = np.sum(pow(0.5,xy_upstr)/0.5)
            all_trail_pix = 0
            for m in xy_upstr:
Wei Chengliang's avatar
Wei Chengliang committed
519
                a1 = 12.97059491
Wei Chengliang's avatar
Wei Chengliang committed
520
521
522
523
524
                b1 = 0.54286652
                c1 = 0.69093105
                a2 = 2.77298856
                b2 = 0.11231055
                c2 = -0.01038675
Fang Yuedong's avatar
Fang Yuedong committed
525
                # t_pow = 0
Wei Chengliang's avatar
Wei Chengliang committed
526
527
                am = 1
                bm = 1
Fang Yuedong's avatar
Fang Yuedong committed
528
529
530
531
532
                t_pow = am*np.exp(-bm*m)
                # if m < 5:
                #     t_pow = a1*np.exp(-b1*m)+c1
                # else:
                #     t_pow = a2*np.exp(-b2*m)+c2
Wei Chengliang's avatar
Wei Chengliang committed
533
                if t_pow < 0:
Fang Yuedong's avatar
Fang Yuedong committed
534
535
536
537
538
539
540
                    t_pow = 0

                all_trail_pix += t_pow
                all_trail[m] = t_pow
            trail_pix_eff = trail_f/all_trail_pix
            all_trail = trail_pix_eff*all_trail
            all_trail[0] = f - trail_f
Wei Chengliang's avatar
Wei Chengliang committed
541

Wei Chengliang's avatar
Wei Chengliang committed
542
            for m in np.arange(0, xy_num, 1):
Fang Yuedong's avatar
Fang Yuedong committed
543
544
545
546
547
                if direction == 'column':
                    if trail_direction == 'down':
                        y_pos = i + m
                    elif trail_direction == 'up':
                        y_pos = i - m
Wei Chengliang's avatar
Wei Chengliang committed
548
                    if y_pos < 0 or y_pos >= sh1:
Fang Yuedong's avatar
Fang Yuedong committed
549
                        break
Wei Chengliang's avatar
Wei Chengliang committed
550
                    n_img[y_pos, j] = n_img[y_pos, j] + all_trail[m]
Fang Yuedong's avatar
Fang Yuedong committed
551
552
553
554
555
                elif direction == 'row':
                    if trail_direction == 'left':
                        x_pos = j - m
                    elif trail_direction == 'right':
                        x_pos = j + m
Wei Chengliang's avatar
Wei Chengliang committed
556
                    if x_pos < 0 or x_pos >= sh2:
Fang Yuedong's avatar
Fang Yuedong committed
557
                        break
Wei Chengliang's avatar
Wei Chengliang committed
558
                    n_img[i, x_pos] = n_img[i, x_pos] + all_trail[m]
Fang Yuedong's avatar
Fang Yuedong committed
559
560
561
562

    return n_img


Wei Chengliang's avatar
Wei Chengliang committed
563
564
# ---------- For Cosmic-Ray Simulation ------------
# ---------- Zhang Xin ----------------------------
Fang Yuedong's avatar
Fang Yuedong committed
565
def getYValue(collection, x):
Wei Chengliang's avatar
Wei Chengliang committed
566
    index = 0
Fang Yuedong's avatar
Fang Yuedong committed
567
    if (collection.shape[1] == 2):
Wei Chengliang's avatar
Wei Chengliang committed
568
        while (x > collection[index, 0] and index < collection.shape[0]):
Wei Chengliang's avatar
Wei Chengliang committed
569
            index = index + 1
Fang Yuedong's avatar
Fang Yuedong committed
570
        if (index == collection.shape[0] or index == 0):
Wei Chengliang's avatar
Wei Chengliang committed
571
            return 0
Fang Yuedong's avatar
Fang Yuedong committed
572

Wei Chengliang's avatar
Wei Chengliang committed
573
574
        deltX = collection[index, 0] - collection[index-1, 0]
        deltY = collection[index, 1] - collection[index-1, 1]
Fang Yuedong's avatar
Fang Yuedong committed
575
576
577
578

        if deltX == 0:
            return (collection[index, 1] + collection[index-1, 1])/2.0
        else:
Wei Chengliang's avatar
Wei Chengliang committed
579
580
581
            a = deltY/deltX
            return a * (x - collection[index-1, 0]) + collection[index-1, 1]
    return 0
Fang Yuedong's avatar
Fang Yuedong committed
582
583


Wei Chengliang's avatar
Wei Chengliang committed
584
def selectCosmicRayCollection(attachedSizes, xLen, yLen, cr_pixelRatio, CR_max_size):
Fang Yuedong's avatar
Fang Yuedong committed
585
586
587
588

    normalRay = 0.90
    nnormalRay = 1-normalRay
    max_nrayLen = 100
Wei Chengliang's avatar
Wei Chengliang committed
589
590
591
592
    pixelNum = int(xLen * yLen * cr_pixelRatio * normalRay)
    pixelNum_n = int(xLen * yLen * cr_pixelRatio * nnormalRay)
    CRPixelNum = 0

Wei Chengliang's avatar
Wei Chengliang committed
593
    maxValue = max(attachedSizes[:, 1])
Wei Chengliang's avatar
Wei Chengliang committed
594
    maxValue += 0.1
Fang Yuedong's avatar
Fang Yuedong committed
595

Wei Chengliang's avatar
Wei Chengliang committed
596
597
    cr_event_num = 0
    CRs = np.zeros(pixelNum)
Fang Yuedong's avatar
Fang Yuedong committed
598
    while (CRPixelNum < pixelNum):
Wei Chengliang's avatar
Wei Chengliang committed
599
600
        x = CR_max_size * np.random.random()
        y = maxValue * np.random.random()
Fang Yuedong's avatar
Fang Yuedong committed
601
        if (y <= getYValue(attachedSizes, x)):
Wei Chengliang's avatar
Wei Chengliang committed
602
603
604
            CRs[cr_event_num] = np.ceil(x)
            cr_event_num = cr_event_num + 1
            CRPixelNum = CRPixelNum + round(x)
Fang Yuedong's avatar
Fang Yuedong committed
605
606
607

    while (CRPixelNum < pixelNum + pixelNum_n):
        nx = np.random.random()*(max_nrayLen-CR_max_size)+CR_max_size
Wei Chengliang's avatar
Wei Chengliang committed
608
609
610
        CRs[cr_event_num] = np.ceil(nx)
        cr_event_num = cr_event_num + 1
        CRPixelNum = CRPixelNum + np.ceil(nx)
Fang Yuedong's avatar
Fang Yuedong committed
611

Wei Chengliang's avatar
Wei Chengliang committed
612
    return CRs[0:cr_event_num]
Fang Yuedong's avatar
Fang Yuedong committed
613
614


Wei Chengliang's avatar
Wei Chengliang committed
615
def defineEnergyForCR(cr_event_size, seed=12345):
Fang Yuedong's avatar
Fang Yuedong committed
616
    import random
Wei Chengliang's avatar
Wei Chengliang committed
617
618
    sigma = 0.6 / 2.355
    mean = 3.3
Fang Yuedong's avatar
Fang Yuedong committed
619
    random.seed(seed)
Wei Chengliang's avatar
Wei Chengliang committed
620
    energys = np.zeros(cr_event_size)
Fang Yuedong's avatar
Fang Yuedong committed
621
    for i in np.arange(cr_event_size):
Wei Chengliang's avatar
Wei Chengliang committed
622
        energy_index = random.normalvariate(mean, sigma)
Wei Chengliang's avatar
Wei Chengliang committed
623
624
        energys[i] = pow(10, energy_index)
    return energys
Fang Yuedong's avatar
Fang Yuedong committed
625

Wei Chengliang's avatar
Wei Chengliang committed
626

Wei Chengliang's avatar
Wei Chengliang committed
627
def convCR(CRmap=None, addPSF=None, sp_n=4):
Fang Yuedong's avatar
Fang Yuedong committed
628
629
630
631
632
633
634
635
    sh = CRmap.shape

    # sp_n = 4
    subCRmap = np.zeros(np.array(sh)*sp_n)
    pix_v0 = 1/(sp_n*sp_n)
    for i in np.arange(sh[0]):
        i_st = sp_n*i
        for j in np.arange(sh[1]):
Wei Chengliang's avatar
Wei Chengliang committed
636
            if CRmap[i, j] == 0:
Fang Yuedong's avatar
Fang Yuedong committed
637
638
                continue
            j_st = sp_n*j
Wei Chengliang's avatar
Wei Chengliang committed
639
            pix_v1 = CRmap[i, j]*pix_v0
Fang Yuedong's avatar
Fang Yuedong committed
640
641
642
643
644
645
            for m in np.arange(sp_n):
                for n in np.arange(sp_n):
                    subCRmap[i_st+m, j_st + n] = pix_v1

    m_size = addPSF.shape[0]

Wei Chengliang's avatar
Wei Chengliang committed
646
    subCRmap_n = np.zeros(np.array(subCRmap.shape) + m_size - 1)
Fang Yuedong's avatar
Fang Yuedong committed
647
648
649

    for i in np.arange(subCRmap.shape[0]):
        for j in np.arange(subCRmap.shape[1]):
Wei Chengliang's avatar
Wei Chengliang committed
650
            if subCRmap[i, j] > 0:
Wei Chengliang's avatar
Wei Chengliang committed
651
652
                convPix = addPSF*subCRmap[i, j]
                subCRmap_n[i:i+m_size, j:j+m_size] += convPix
Fang Yuedong's avatar
Fang Yuedong committed
653
654
655
656
657
658
659
660

    CRmap_n = np.zeros((np.array(subCRmap_n.shape)/sp_n).astype(np.int32))
    sh_n = CRmap_n.shape

    for i in np.arange(sh_n[0]):
        i_st = sp_n*i
        for j in np.arange(sh_n[1]):
            p_v = 0
Wei Chengliang's avatar
Wei Chengliang committed
661
            j_st = sp_n*j
Fang Yuedong's avatar
Fang Yuedong committed
662
663
664
665
            for m in np.arange(sp_n):
                for n in np.arange(sp_n):
                    p_v += subCRmap_n[i_st+m, j_st + n]

Wei Chengliang's avatar
Wei Chengliang committed
666
            CRmap_n[i, j] = p_v
Fang Yuedong's avatar
Fang Yuedong committed
667
668
669
670
671
672
673
674

    return CRmap_n


def produceCR_Map(xLen, yLen, exTime, cr_pixelRatio, gain, attachedSizes, seed=20210317):
    # Return: an 2-D numpy array
    # attachedSizes = np.loadtxt('./wfc-cr-attachpixel.dat');
    np.random.seed(seed)
Wei Chengliang's avatar
Wei Chengliang committed
675
676
    CR_max_size = 20.0
    cr_size = selectCosmicRayCollection(attachedSizes, xLen, yLen, cr_pixelRatio, CR_max_size)
Fang Yuedong's avatar
Fang Yuedong committed
677

Wei Chengliang's avatar
Wei Chengliang committed
678
679
    cr_event_size = cr_size.shape[0]
    cr_energys = defineEnergyForCR(cr_event_size, seed=seed)
Fang Yuedong's avatar
Fang Yuedong committed
680

Wei Chengliang's avatar
Wei Chengliang committed
681
    CRmap = np.zeros([yLen, xLen])
Fang Yuedong's avatar
Fang Yuedong committed
682

Wei Chengliang's avatar
Wei Chengliang committed
683
    # produce conv kernel
Fang Yuedong's avatar
Fang Yuedong committed
684
685
686
687
688
689
690
691
692
693
694
695
    from astropy.modeling.models import Gaussian2D
    o_size = 4
    sp_n = 8

    m_size = o_size*sp_n+1
    m_cen = o_size*sp_n/2
    sigma_psf = 0.2*sp_n
    addPSF_ = Gaussian2D(1, m_cen, m_cen, sigma_psf, sigma_psf)
    yp, xp = np.mgrid[0:m_size, 0:m_size]
    addPSF = addPSF_(xp, yp)
    convKernel = addPSF/addPSF.sum()

Wei Chengliang's avatar
Wei Chengliang committed
696
    # ---------------------------------
Fang Yuedong's avatar
Fang Yuedong committed
697
    for i in np.arange(cr_event_size):
Wei Chengliang's avatar
Wei Chengliang committed
698
699
        xPos = round((xLen - 1) * np.random.random())
        yPos = round((yLen - 1) * np.random.random())
Wei Chengliang's avatar
Wei Chengliang committed
700
        cr_lens = int(cr_size[i])
Wei Chengliang's avatar
Wei Chengliang committed
701
702
703
704
        if cr_lens == 0:
            continue
        pix_energy = cr_energys[i]/gain/cr_lens
        pos_angle = 1/2*math.pi*np.random.random()
Fang Yuedong's avatar
Fang Yuedong committed
705
706
707
708

        crMatrix = np.zeros([cr_lens+1, cr_lens + 1])

        for j in np.arange(cr_lens):
Wei Chengliang's avatar
Wei Chengliang committed
709
            x_n = int(np.cos(pos_angle)*j - np.sin(pos_angle)*0)
Fang Yuedong's avatar
Fang Yuedong committed
710
711
            if x_n < 0:
                x_n = x_n + cr_lens+1
Wei Chengliang's avatar
Wei Chengliang committed
712
            y_n = int(np.sin(pos_angle)*j + np.cos(pos_angle)*0)
Wei Chengliang's avatar
Wei Chengliang committed
713
            if x_n < 0 or x_n > cr_lens or y_n < 0 or y_n > cr_lens:
Wei Chengliang's avatar
Wei Chengliang committed
714
715
                continue
            crMatrix[y_n, x_n] = pix_energy
Fang Yuedong's avatar
Fang Yuedong committed
716
717
718
719
720
721
722
723
724
725
726
727
728

        crMatrix_n = convCR(crMatrix, convKernel, sp_n)
        # crMatrix_n = crMatrix

        xpix = np.arange(crMatrix_n.shape[0]) + int(xPos)
        ypix = np.arange(crMatrix_n.shape[1]) + int(yPos)

        sh = CRmap.shape
        okx = (xpix >= 0) & (xpix < sh[1])
        oky = (ypix >= 0) & (ypix < sh[0])

        sly = slice(ypix[oky].min(), ypix[oky].max()+1)
        slx = slice(xpix[okx].min(), xpix[okx].max()+1)
Wei Chengliang's avatar
Wei Chengliang committed
729
        CRmap[sly, slx] += crMatrix_n[oky, :][:, okx]
Fang Yuedong's avatar
Fang Yuedong committed
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
    return CRmap.astype(np.int32), cr_event_size


def ShutterEffectArr(GSImage, t_exp=150, t_shutter=1.3, dist_bearing=735, dt=1E-3):
    # Generate Shutter-Effect normalized image
    # t_shutter: time of shutter movement
    # dist_bearing: distance between two bearings of shutter leaves
    # dt: delta_t of sampling

    from scipy import interpolate

    SampleNumb = int(t_shutter/dt+1)
    DistHalf = dist_bearing/2

    t = np.arange(SampleNumb)*dt
    a_arr = 5.84*np.sin(2*math.pi/t_shutter*t)
    v = np.zeros(SampleNumb)
    theta = np.zeros(SampleNumb)
    x = np.arange(SampleNumb)/(SampleNumb-1)*dist_bearing
    s = np.zeros(SampleNumb)
    s1 = np.zeros(SampleNumb)
    s2 = np.zeros(SampleNumb)
Wei Chengliang's avatar
Wei Chengliang committed
752
    brt = np.zeros(SampleNumb)
Fang Yuedong's avatar
Fang Yuedong committed
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
    idx = np.arange(SampleNumb)
    sidx = np.zeros(SampleNumb)
    s1idx = np.zeros(SampleNumb)
    s2idx = np.zeros(SampleNumb)

    v[0] = 0
    theta[0] = 0

    for i in range(SampleNumb-1):
        v[i+1] = v[i]+a_arr[i]*dt
        theta[i+1] = theta[i]+v[i]*dt
        s1[i] = DistHalf*np.cos(theta[i])
        s2[i] = dist_bearing-DistHalf*np.cos(theta[i])
        s1idx[i] = int(s1[i]/dist_bearing*(SampleNumb))
        s2idx[i] = int(s2[i]/dist_bearing*(SampleNumb))
Wei Chengliang's avatar
Wei Chengliang committed
768
        brt[(idx > s1idx[i]) & (idx < s2idx[i])] += dt
Fang Yuedong's avatar
Fang Yuedong committed
769

Wei Chengliang's avatar
Wei Chengliang committed
770
    if t_exp > t_shutter*2:
Fang Yuedong's avatar
Fang Yuedong committed
771
772
773
774
775
776
777
778
779
780
781
782
        brt = brt*2+(t_exp-t_shutter*2)
    else:
        brt = brt*2

    x = (x-dist_bearing/2)*100

    intp = interpolate.splrep(x, brt, s=0)

    xmin = GSImage.bounds.getXMin()
    xmax = GSImage.bounds.getXMax()
    ymin = GSImage.bounds.getYMin()
    ymax = GSImage.bounds.getYMax()
Wei Chengliang's avatar
Wei Chengliang committed
783
    if xmin < np.min(x) or xmax > np.max(x):
Fang Yuedong's avatar
Fang Yuedong committed
784
        raise LookupError("Out of focal-plane bounds in X-direction.")
Wei Chengliang's avatar
Wei Chengliang committed
785
    if ymin < -25331 or ymax > 25331:
Fang Yuedong's avatar
Fang Yuedong committed
786
787
788
789
790
        raise LookupError("Out of focal-plane bounds in Y-direction.")
    sizex = xmax-xmin+1
    sizey = ymax-ymin+1
    xnewgrid = np.mgrid[xmin:(xmin+sizex)]
    expeffect = interpolate.splev(xnewgrid, intp, der=0)
Zhang Xin's avatar
Zhang Xin committed
791
    expeffect /= t_exp
Wei Chengliang's avatar
Wei Chengliang committed
792
    exparrnormal = np.tile(expeffect, (sizey, 1))
Fang Yuedong's avatar
Fang Yuedong committed
793
794
795
    # GSImage *= exparrnormal

    return exparrnormal