effects.py 29.5 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
    nbadsecA, nbadsecD = rgn.integers(low=1, high=5, size=2)
87
    collen = rgcollen.integers(low=int(ysize*0.1), high=int(ysize*0.5), 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
Fang Yuedong's avatar
Fang Yuedong committed
311

Wei Chengliang's avatar
Wei Chengliang committed
312
313
314
        if direction == 'up':
            if imgarr[yi-1, xi] >= fullwell:
                imgarr[yi, xi] = fullwell
Wei Chengliang's avatar
Wei Chengliang committed
315
                yi -= 1
Fang Yuedong's avatar
Fang Yuedong committed
316
317
318
                # [TEST] charge in the middle
                if yi == (imgarr.shape[0] // 2 - 1):
                    break
Fang Yuedong's avatar
Fang Yuedong committed
319
                continue
Wei Chengliang's avatar
Wei Chengliang committed
320
321
322
323
        elif direction == 'down':
            if imgarr[yi+1, xi] >= fullwell:
                imgarr[yi, xi] = fullwell
                yi += 1
Fang Yuedong's avatar
Fang Yuedong committed
324
325
                if yi == (imgarr.shape[0] // 2):
                    break
Fang Yuedong's avatar
Fang Yuedong committed
326
                continue
Wei Chengliang's avatar
Wei Chengliang committed
327
        if aa <= 1:
Wei Chengliang's avatar
Wei Chengliang committed
328
329
330
331
332
333
            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
Fang Yuedong's avatar
Fang Yuedong committed
334
335
                    # if yi < 0:
                    if yi < 0 or yi == (imgarr.shape[0]//2 - 1):
Fang Yuedong's avatar
Fang Yuedong committed
336
                        break
Wei Chengliang's avatar
Wei Chengliang committed
337
338
339
340
                elif direction == 'down':
                    imgarr[yi+1, xi] += charge
                    charge = imgarr[yi+1, xi]-fullwell
                    yi += 1
Fang Yuedong's avatar
Fang Yuedong committed
341
                    # if yi > imgarr.shape[0]:
Wei Chengliang's avatar
Wei Chengliang committed
342
                    if yi > imgarr.shape[0] or yi == (imgarr.shape[0]//2):
Fang Yuedong's avatar
Fang Yuedong committed
343
344
345
                        break
        else:
            # calculate bleeding trail:
Wei Chengliang's avatar
Wei Chengliang committed
346
            trail_frac = BleedingTrail(aa, yy)
Fang Yuedong's avatar
Fang Yuedong committed
347
348

            # put charge upwards
Wei Chengliang's avatar
Wei Chengliang committed
349
350
351
352
            if trail_frac >= 0.99:
                imgarr[yi, xi] = fullwell
                if direction == 'up':
                    yi -= 1
353
354
                    if yi == (imgarr.shape[0]//2 - 1):
                        break
Wei Chengliang's avatar
Wei Chengliang committed
355
356
                elif direction == 'down':
                    yi += 1
357
358
                    if yi == (imgarr.shape[0]//2):
                        break
Fang Yuedong's avatar
Fang Yuedong committed
359
360
                yy += 1
            else:
Wei Chengliang's avatar
Wei Chengliang committed
361
                if trail_frac < trailcutfrac:
Fang Yuedong's avatar
Fang Yuedong committed
362
363
                    break
                charge = fullwell*trail_frac
Wei Chengliang's avatar
Wei Chengliang committed
364
365
366
367
368
369
                imgarr[yi, xi] += charge
                if imgarr[yi, xi] > fullwell:
                    imgarr[yi, xi] = fullwell

                if direction == 'up':
                    yi -= 1
370
371
                    if yi == (imgarr.shape[0]//2 - 1):
                        break
Wei Chengliang's avatar
Wei Chengliang committed
372
373
                elif direction == 'down':
                    yi += 1
374
375
                    if yi == (imgarr.shape[0]//2):
                        break
Fang Yuedong's avatar
Fang Yuedong committed
376
377
378
379
380
381
                yy += 1

    return imgarr


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

Wei Chengliang's avatar
Wei Chengliang committed
385
    if satupos_y.shape[0] == 0:
Fang Yuedong's avatar
Fang Yuedong committed
386
387
388
389
390
391
392
393
394
        # 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
395
396
397
    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
398

Wei Chengliang's avatar
Wei Chengliang committed
399
400
    for yi, xi in zip(satupos_y, satupos_x):
        yxidx = ''.join([str(yi), str(xi)])
Fang Yuedong's avatar
Fang Yuedong committed
401
402
403
404
405
406
        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
407
            if yi >= 0 and yi < imgarr.shape[0]:
Wei Chengliang's avatar
Wei Chengliang committed
408
                imgarr = MakeTrail(imgarr, (yi, xi), chargeup, fullwell=9e4, direction='up', trailcutfrac=0.9)
Fang Yuedong's avatar
Fang Yuedong committed
409
                # Charge Clump moves down
Wei Chengliang's avatar
Wei Chengliang committed
410
                imgarr = MakeTrail(imgarr, (yi, xi), chargedn, fullwell=9e4, direction='down', trailcutfrac=0.9)
Fang Yuedong's avatar
Fang Yuedong committed
411
        except Exception as e:
Wei Chengliang's avatar
Wei Chengliang committed
412
            print(e, '@pix ', (yi+1, xi+1))
Fang Yuedong's avatar
Fang Yuedong committed
413
414
415
            return imgarr
    return imgarr

Wei Chengliang's avatar
Wei Chengliang committed
416

Fang Yuedong's avatar
Fang Yuedong committed
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
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
442
#    Saturation & Bleeding End    #
Fang Yuedong's avatar
Fang Yuedong committed
443
444
445
446
447


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
448
    # assuming image has 2 columns and 8 rows of output channels.
Fang Yuedong's avatar
Fang Yuedong committed
449
450
451
452
453
    # 00  01
    # 10  11
    # 20  21
    # ...
    # return: GS Image Object
Wei Chengliang's avatar
Wei Chengliang committed
454
    npix_y, npix_x = GSImage.array.shape
Fang Yuedong's avatar
Fang Yuedong committed
455
456
457
    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
458
    if rowi < 4 and coli == 0:
Fang Yuedong's avatar
Fang Yuedong committed
459
460
461
        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
462
463
        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
464
465
466
        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
467
468
        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
469
470
471
        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
472
473
        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
474
475
476
        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
477
        OutputSubimg.array[16:int(npix_y/8)+16, 8:int(npix_x/2)+8] = subimg.array
Fang Yuedong's avatar
Fang Yuedong committed
478
479
480
481
482
483
484
485
486
    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
487
    size_y, size_x = GSImage.array.shape
Fang Yuedong's avatar
Fang Yuedong committed
488
489
490
491
    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
492
        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
493
        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
494
    elif direction == 'row':
Wei Chengliang's avatar
Wei Chengliang committed
495
496
        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
497
498
499
500
    return GSImage


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

Wei Chengliang's avatar
Wei Chengliang committed
503
504
505
    # 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
506
    trail_a = 5.651803799619966
Fang Yuedong's avatar
Fang Yuedong committed
507
508
509
510
511
    trail_b = -2.671933068990729

    sh1 = img.shape[0]
    sh2 = img.shape[1]
    n_img = img*0
Wei Chengliang's avatar
Wei Chengliang committed
512
    idx = np.where(img < threshold)
Fang Yuedong's avatar
Fang Yuedong committed
513
514
    if len(idx[0]) == 0:
        pass
Wei Chengliang's avatar
Wei Chengliang committed
515
    elif len(idx[0]) > 0:
Fang Yuedong's avatar
Fang Yuedong committed
516
517
        n_img[idx] = img[idx]

Wei Chengliang's avatar
Wei Chengliang committed
518
    yidx, xidx = np.where(img >= threshold)
Fang Yuedong's avatar
Fang Yuedong committed
519
520
    if len(yidx) == 0:
        pass
Wei Chengliang's avatar
Wei Chengliang committed
521
    elif len(yidx) > 0:
Fang Yuedong's avatar
Fang Yuedong committed
522
        # print(index)
Wei Chengliang's avatar
Wei Chengliang committed
523
524
        for i, j in zip(yidx, xidx):
            f = img[i, j]
Fang Yuedong's avatar
Fang Yuedong committed
525
526
527
528
            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
529

Wei Chengliang's avatar
Wei Chengliang committed
530
            xy_upstr = np.arange(1, xy_num, 1)
Fang Yuedong's avatar
Fang Yuedong committed
531
532
533
534

            # 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
535
                a1 = 12.97059491
Wei Chengliang's avatar
Wei Chengliang committed
536
537
538
539
540
                b1 = 0.54286652
                c1 = 0.69093105
                a2 = 2.77298856
                b2 = 0.11231055
                c2 = -0.01038675
Fang Yuedong's avatar
Fang Yuedong committed
541
                # t_pow = 0
Wei Chengliang's avatar
Wei Chengliang committed
542
543
                am = 1
                bm = 1
Fang Yuedong's avatar
Fang Yuedong committed
544
545
546
547
548
                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
549
                if t_pow < 0:
Fang Yuedong's avatar
Fang Yuedong committed
550
551
552
553
554
555
556
                    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
557

Wei Chengliang's avatar
Wei Chengliang committed
558
            for m in np.arange(0, xy_num, 1):
Fang Yuedong's avatar
Fang Yuedong committed
559
560
561
562
563
                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
564
                    if y_pos < 0 or y_pos >= sh1:
Fang Yuedong's avatar
Fang Yuedong committed
565
                        break
Wei Chengliang's avatar
Wei Chengliang committed
566
                    n_img[y_pos, j] = n_img[y_pos, j] + all_trail[m]
Fang Yuedong's avatar
Fang Yuedong committed
567
568
569
570
571
                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
572
                    if x_pos < 0 or x_pos >= sh2:
Fang Yuedong's avatar
Fang Yuedong committed
573
                        break
Wei Chengliang's avatar
Wei Chengliang committed
574
                    n_img[i, x_pos] = n_img[i, x_pos] + all_trail[m]
Fang Yuedong's avatar
Fang Yuedong committed
575
576
577
578

    return n_img


Wei Chengliang's avatar
Wei Chengliang committed
579
580
# ---------- For Cosmic-Ray Simulation ------------
# ---------- Zhang Xin ----------------------------
Fang Yuedong's avatar
Fang Yuedong committed
581
def getYValue(collection, x):
Wei Chengliang's avatar
Wei Chengliang committed
582
    index = 0
Fang Yuedong's avatar
Fang Yuedong committed
583
    if (collection.shape[1] == 2):
Wei Chengliang's avatar
Wei Chengliang committed
584
        while (x > collection[index, 0] and index < collection.shape[0]):
Wei Chengliang's avatar
Wei Chengliang committed
585
            index = index + 1
Fang Yuedong's avatar
Fang Yuedong committed
586
        if (index == collection.shape[0] or index == 0):
Wei Chengliang's avatar
Wei Chengliang committed
587
            return 0
Fang Yuedong's avatar
Fang Yuedong committed
588

Wei Chengliang's avatar
Wei Chengliang committed
589
590
        deltX = collection[index, 0] - collection[index-1, 0]
        deltY = collection[index, 1] - collection[index-1, 1]
Fang Yuedong's avatar
Fang Yuedong committed
591
592
593
594

        if deltX == 0:
            return (collection[index, 1] + collection[index-1, 1])/2.0
        else:
Wei Chengliang's avatar
Wei Chengliang committed
595
596
597
            a = deltY/deltX
            return a * (x - collection[index-1, 0]) + collection[index-1, 1]
    return 0
Fang Yuedong's avatar
Fang Yuedong committed
598
599


Wei Chengliang's avatar
Wei Chengliang committed
600
def selectCosmicRayCollection(attachedSizes, xLen, yLen, cr_pixelRatio, CR_max_size):
Fang Yuedong's avatar
Fang Yuedong committed
601
602
603
604

    normalRay = 0.90
    nnormalRay = 1-normalRay
    max_nrayLen = 100
Wei Chengliang's avatar
Wei Chengliang committed
605
606
607
608
    pixelNum = int(xLen * yLen * cr_pixelRatio * normalRay)
    pixelNum_n = int(xLen * yLen * cr_pixelRatio * nnormalRay)
    CRPixelNum = 0

Wei Chengliang's avatar
Wei Chengliang committed
609
    maxValue = max(attachedSizes[:, 1])
Wei Chengliang's avatar
Wei Chengliang committed
610
    maxValue += 0.1
Fang Yuedong's avatar
Fang Yuedong committed
611

Wei Chengliang's avatar
Wei Chengliang committed
612
613
    cr_event_num = 0
    CRs = np.zeros(pixelNum)
Fang Yuedong's avatar
Fang Yuedong committed
614
    while (CRPixelNum < pixelNum):
Wei Chengliang's avatar
Wei Chengliang committed
615
616
        x = CR_max_size * np.random.random()
        y = maxValue * np.random.random()
Fang Yuedong's avatar
Fang Yuedong committed
617
        if (y <= getYValue(attachedSizes, x)):
Wei Chengliang's avatar
Wei Chengliang committed
618
619
620
            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
621
622
623

    while (CRPixelNum < pixelNum + pixelNum_n):
        nx = np.random.random()*(max_nrayLen-CR_max_size)+CR_max_size
Wei Chengliang's avatar
Wei Chengliang committed
624
625
626
        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
627

Wei Chengliang's avatar
Wei Chengliang committed
628
    return CRs[0:cr_event_num]
Fang Yuedong's avatar
Fang Yuedong committed
629
630


Wei Chengliang's avatar
Wei Chengliang committed
631
def defineEnergyForCR(cr_event_size, seed=12345):
Fang Yuedong's avatar
Fang Yuedong committed
632
    import random
Wei Chengliang's avatar
Wei Chengliang committed
633
634
    sigma = 0.6 / 2.355
    mean = 3.3
Fang Yuedong's avatar
Fang Yuedong committed
635
    random.seed(seed)
Wei Chengliang's avatar
Wei Chengliang committed
636
    energys = np.zeros(cr_event_size)
Fang Yuedong's avatar
Fang Yuedong committed
637
    for i in np.arange(cr_event_size):
Wei Chengliang's avatar
Wei Chengliang committed
638
        energy_index = random.normalvariate(mean, sigma)
Wei Chengliang's avatar
Wei Chengliang committed
639
640
        energys[i] = pow(10, energy_index)
    return energys
Fang Yuedong's avatar
Fang Yuedong committed
641

Wei Chengliang's avatar
Wei Chengliang committed
642

Wei Chengliang's avatar
Wei Chengliang committed
643
def convCR(CRmap=None, addPSF=None, sp_n=4):
Fang Yuedong's avatar
Fang Yuedong committed
644
645
646
647
648
649
650
651
    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
652
            if CRmap[i, j] == 0:
Fang Yuedong's avatar
Fang Yuedong committed
653
654
                continue
            j_st = sp_n*j
Wei Chengliang's avatar
Wei Chengliang committed
655
            pix_v1 = CRmap[i, j]*pix_v0
Fang Yuedong's avatar
Fang Yuedong committed
656
657
658
659
660
661
            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
662
    subCRmap_n = np.zeros(np.array(subCRmap.shape) + m_size - 1)
Fang Yuedong's avatar
Fang Yuedong committed
663
664
665

    for i in np.arange(subCRmap.shape[0]):
        for j in np.arange(subCRmap.shape[1]):
Wei Chengliang's avatar
Wei Chengliang committed
666
            if subCRmap[i, j] > 0:
Wei Chengliang's avatar
Wei Chengliang committed
667
668
                convPix = addPSF*subCRmap[i, j]
                subCRmap_n[i:i+m_size, j:j+m_size] += convPix
Fang Yuedong's avatar
Fang Yuedong committed
669
670
671
672
673
674
675
676

    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
677
            j_st = sp_n*j
Fang Yuedong's avatar
Fang Yuedong committed
678
679
680
681
            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
682
            CRmap_n[i, j] = p_v
Fang Yuedong's avatar
Fang Yuedong committed
683
684
685
686
687
688
689
690

    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
691
692
    CR_max_size = 20.0
    cr_size = selectCosmicRayCollection(attachedSizes, xLen, yLen, cr_pixelRatio, CR_max_size)
Fang Yuedong's avatar
Fang Yuedong committed
693

Wei Chengliang's avatar
Wei Chengliang committed
694
695
    cr_event_size = cr_size.shape[0]
    cr_energys = defineEnergyForCR(cr_event_size, seed=seed)
Fang Yuedong's avatar
Fang Yuedong committed
696

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

Wei Chengliang's avatar
Wei Chengliang committed
699
    # produce conv kernel
Fang Yuedong's avatar
Fang Yuedong committed
700
701
702
703
704
705
706
707
708
709
710
711
    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
712
    # ---------------------------------
Fang Yuedong's avatar
Fang Yuedong committed
713
    for i in np.arange(cr_event_size):
Wei Chengliang's avatar
Wei Chengliang committed
714
715
        xPos = round((xLen - 1) * np.random.random())
        yPos = round((yLen - 1) * np.random.random())
Wei Chengliang's avatar
Wei Chengliang committed
716
        cr_lens = int(cr_size[i])
Wei Chengliang's avatar
Wei Chengliang committed
717
718
719
720
        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
721
722
723
724

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

        for j in np.arange(cr_lens):
Wei Chengliang's avatar
Wei Chengliang committed
725
            x_n = int(np.cos(pos_angle)*j - np.sin(pos_angle)*0)
Fang Yuedong's avatar
Fang Yuedong committed
726
727
            if x_n < 0:
                x_n = x_n + cr_lens+1
Wei Chengliang's avatar
Wei Chengliang committed
728
            y_n = int(np.sin(pos_angle)*j + np.cos(pos_angle)*0)
Wei Chengliang's avatar
Wei Chengliang committed
729
            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
730
731
                continue
            crMatrix[y_n, x_n] = pix_energy
Fang Yuedong's avatar
Fang Yuedong committed
732
733
734
735
736
737
738
739
740
741
742
743
744

        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
745
        CRmap[sly, slx] += crMatrix_n[oky, :][:, okx]
Fang Yuedong's avatar
Fang Yuedong committed
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
    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
768
    brt = np.zeros(SampleNumb)
Fang Yuedong's avatar
Fang Yuedong committed
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
    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
784
        brt[(idx > s1idx[i]) & (idx < s2idx[i])] += dt
Fang Yuedong's avatar
Fang Yuedong committed
785

Wei Chengliang's avatar
Wei Chengliang committed
786
    if t_exp > t_shutter*2:
Fang Yuedong's avatar
Fang Yuedong committed
787
788
789
790
791
792
793
794
795
796
797
798
        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
799
    if xmin < np.min(x) or xmax > np.max(x):
Fang Yuedong's avatar
Fang Yuedong committed
800
        raise LookupError("Out of focal-plane bounds in X-direction.")
Wei Chengliang's avatar
Wei Chengliang committed
801
    if ymin < -25331 or ymax > 25331:
Fang Yuedong's avatar
Fang Yuedong committed
802
803
804
805
806
        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
807
    expeffect /= t_exp
Wei Chengliang's avatar
Wei Chengliang committed
808
    exparrnormal = np.tile(expeffect, (sizey, 1))
Fang Yuedong's avatar
Fang Yuedong committed
809
810
811
    # GSImage *= exparrnormal

    return exparrnormal