effects.py 29.2 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
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
342
                    # if yi > imgarr.shape[0]:
                    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
353
354
            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
355
356
                yy += 1
            else:
Wei Chengliang's avatar
Wei Chengliang committed
357
                if trail_frac < trailcutfrac:
Fang Yuedong's avatar
Fang Yuedong committed
358
359
                    break
                charge = fullwell*trail_frac
Wei Chengliang's avatar
Wei Chengliang committed
360
361
362
363
364
365
366
367
                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
368
369
370
371
372
373
                yy += 1

    return imgarr


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

Wei Chengliang's avatar
Wei Chengliang committed
377
    if satupos_y.shape[0] == 0:
Fang Yuedong's avatar
Fang Yuedong committed
378
379
380
381
382
383
384
385
386
        # 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
387
388
389
    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
390

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

Wei Chengliang's avatar
Wei Chengliang committed
408

Fang Yuedong's avatar
Fang Yuedong committed
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
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
434
#    Saturation & Bleeding End    #
Fang Yuedong's avatar
Fang Yuedong committed
435
436
437
438
439


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


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

Wei Chengliang's avatar
Wei Chengliang committed
495
496
497
    # 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
498
    trail_a = 5.651803799619966
Fang Yuedong's avatar
Fang Yuedong committed
499
500
501
502
503
    trail_b = -2.671933068990729

    sh1 = img.shape[0]
    sh2 = img.shape[1]
    n_img = img*0
Wei Chengliang's avatar
Wei Chengliang committed
504
    idx = np.where(img < threshold)
Fang Yuedong's avatar
Fang Yuedong committed
505
506
    if len(idx[0]) == 0:
        pass
Wei Chengliang's avatar
Wei Chengliang committed
507
    elif len(idx[0]) > 0:
Fang Yuedong's avatar
Fang Yuedong committed
508
509
        n_img[idx] = img[idx]

Wei Chengliang's avatar
Wei Chengliang committed
510
    yidx, xidx = np.where(img >= threshold)
Fang Yuedong's avatar
Fang Yuedong committed
511
512
    if len(yidx) == 0:
        pass
Wei Chengliang's avatar
Wei Chengliang committed
513
    elif len(yidx) > 0:
Fang Yuedong's avatar
Fang Yuedong committed
514
        # print(index)
Wei Chengliang's avatar
Wei Chengliang committed
515
516
        for i, j in zip(yidx, xidx):
            f = img[i, j]
Fang Yuedong's avatar
Fang Yuedong committed
517
518
519
520
            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
521

Wei Chengliang's avatar
Wei Chengliang committed
522
            xy_upstr = np.arange(1, xy_num, 1)
Fang Yuedong's avatar
Fang Yuedong committed
523
524
525
526

            # 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
527
                a1 = 12.97059491
Wei Chengliang's avatar
Wei Chengliang committed
528
529
530
531
532
                b1 = 0.54286652
                c1 = 0.69093105
                a2 = 2.77298856
                b2 = 0.11231055
                c2 = -0.01038675
Fang Yuedong's avatar
Fang Yuedong committed
533
                # t_pow = 0
Wei Chengliang's avatar
Wei Chengliang committed
534
535
                am = 1
                bm = 1
Fang Yuedong's avatar
Fang Yuedong committed
536
537
538
539
540
                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
541
                if t_pow < 0:
Fang Yuedong's avatar
Fang Yuedong committed
542
543
544
545
546
547
548
                    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
549

Wei Chengliang's avatar
Wei Chengliang committed
550
            for m in np.arange(0, xy_num, 1):
Fang Yuedong's avatar
Fang Yuedong committed
551
552
553
554
555
                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
556
                    if y_pos < 0 or y_pos >= sh1:
Fang Yuedong's avatar
Fang Yuedong committed
557
                        break
Wei Chengliang's avatar
Wei Chengliang committed
558
                    n_img[y_pos, j] = n_img[y_pos, j] + all_trail[m]
Fang Yuedong's avatar
Fang Yuedong committed
559
560
561
562
563
                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
564
                    if x_pos < 0 or x_pos >= sh2:
Fang Yuedong's avatar
Fang Yuedong committed
565
                        break
Wei Chengliang's avatar
Wei Chengliang committed
566
                    n_img[i, x_pos] = n_img[i, x_pos] + all_trail[m]
Fang Yuedong's avatar
Fang Yuedong committed
567
568
569
570

    return n_img


Wei Chengliang's avatar
Wei Chengliang committed
571
572
# ---------- For Cosmic-Ray Simulation ------------
# ---------- Zhang Xin ----------------------------
Fang Yuedong's avatar
Fang Yuedong committed
573
def getYValue(collection, x):
Wei Chengliang's avatar
Wei Chengliang committed
574
    index = 0
Fang Yuedong's avatar
Fang Yuedong committed
575
    if (collection.shape[1] == 2):
Wei Chengliang's avatar
Wei Chengliang committed
576
        while (x > collection[index, 0] and index < collection.shape[0]):
Wei Chengliang's avatar
Wei Chengliang committed
577
            index = index + 1
Fang Yuedong's avatar
Fang Yuedong committed
578
        if (index == collection.shape[0] or index == 0):
Wei Chengliang's avatar
Wei Chengliang committed
579
            return 0
Fang Yuedong's avatar
Fang Yuedong committed
580

Wei Chengliang's avatar
Wei Chengliang committed
581
582
        deltX = collection[index, 0] - collection[index-1, 0]
        deltY = collection[index, 1] - collection[index-1, 1]
Fang Yuedong's avatar
Fang Yuedong committed
583
584
585
586

        if deltX == 0:
            return (collection[index, 1] + collection[index-1, 1])/2.0
        else:
Wei Chengliang's avatar
Wei Chengliang committed
587
588
589
            a = deltY/deltX
            return a * (x - collection[index-1, 0]) + collection[index-1, 1]
    return 0
Fang Yuedong's avatar
Fang Yuedong committed
590
591


Wei Chengliang's avatar
Wei Chengliang committed
592
def selectCosmicRayCollection(attachedSizes, xLen, yLen, cr_pixelRatio, CR_max_size):
Fang Yuedong's avatar
Fang Yuedong committed
593
594
595
596

    normalRay = 0.90
    nnormalRay = 1-normalRay
    max_nrayLen = 100
Wei Chengliang's avatar
Wei Chengliang committed
597
598
599
600
    pixelNum = int(xLen * yLen * cr_pixelRatio * normalRay)
    pixelNum_n = int(xLen * yLen * cr_pixelRatio * nnormalRay)
    CRPixelNum = 0

Wei Chengliang's avatar
Wei Chengliang committed
601
    maxValue = max(attachedSizes[:, 1])
Wei Chengliang's avatar
Wei Chengliang committed
602
    maxValue += 0.1
Fang Yuedong's avatar
Fang Yuedong committed
603

Wei Chengliang's avatar
Wei Chengliang committed
604
605
    cr_event_num = 0
    CRs = np.zeros(pixelNum)
Fang Yuedong's avatar
Fang Yuedong committed
606
    while (CRPixelNum < pixelNum):
Wei Chengliang's avatar
Wei Chengliang committed
607
608
        x = CR_max_size * np.random.random()
        y = maxValue * np.random.random()
Fang Yuedong's avatar
Fang Yuedong committed
609
        if (y <= getYValue(attachedSizes, x)):
Wei Chengliang's avatar
Wei Chengliang committed
610
611
612
            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
613
614
615

    while (CRPixelNum < pixelNum + pixelNum_n):
        nx = np.random.random()*(max_nrayLen-CR_max_size)+CR_max_size
Wei Chengliang's avatar
Wei Chengliang committed
616
617
618
        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
619

Wei Chengliang's avatar
Wei Chengliang committed
620
    return CRs[0:cr_event_num]
Fang Yuedong's avatar
Fang Yuedong committed
621
622


Wei Chengliang's avatar
Wei Chengliang committed
623
def defineEnergyForCR(cr_event_size, seed=12345):
Fang Yuedong's avatar
Fang Yuedong committed
624
    import random
Wei Chengliang's avatar
Wei Chengliang committed
625
626
    sigma = 0.6 / 2.355
    mean = 3.3
Fang Yuedong's avatar
Fang Yuedong committed
627
    random.seed(seed)
Wei Chengliang's avatar
Wei Chengliang committed
628
    energys = np.zeros(cr_event_size)
Fang Yuedong's avatar
Fang Yuedong committed
629
    for i in np.arange(cr_event_size):
Wei Chengliang's avatar
Wei Chengliang committed
630
        energy_index = random.normalvariate(mean, sigma)
Wei Chengliang's avatar
Wei Chengliang committed
631
632
        energys[i] = pow(10, energy_index)
    return energys
Fang Yuedong's avatar
Fang Yuedong committed
633

Wei Chengliang's avatar
Wei Chengliang committed
634

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

    for i in np.arange(subCRmap.shape[0]):
        for j in np.arange(subCRmap.shape[1]):
Wei Chengliang's avatar
Wei Chengliang committed
658
            if subCRmap[i, j] > 0:
Wei Chengliang's avatar
Wei Chengliang committed
659
660
                convPix = addPSF*subCRmap[i, j]
                subCRmap_n[i:i+m_size, j:j+m_size] += convPix
Fang Yuedong's avatar
Fang Yuedong committed
661
662
663
664
665
666
667
668

    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
669
            j_st = sp_n*j
Fang Yuedong's avatar
Fang Yuedong committed
670
671
672
673
            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
674
            CRmap_n[i, j] = p_v
Fang Yuedong's avatar
Fang Yuedong committed
675
676
677
678
679
680
681
682

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

Wei Chengliang's avatar
Wei Chengliang committed
686
687
    cr_event_size = cr_size.shape[0]
    cr_energys = defineEnergyForCR(cr_event_size, seed=seed)
Fang Yuedong's avatar
Fang Yuedong committed
688

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

Wei Chengliang's avatar
Wei Chengliang committed
691
    # produce conv kernel
Fang Yuedong's avatar
Fang Yuedong committed
692
693
694
695
696
697
698
699
700
701
702
703
    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
704
    # ---------------------------------
Fang Yuedong's avatar
Fang Yuedong committed
705
    for i in np.arange(cr_event_size):
Wei Chengliang's avatar
Wei Chengliang committed
706
707
        xPos = round((xLen - 1) * np.random.random())
        yPos = round((yLen - 1) * np.random.random())
Wei Chengliang's avatar
Wei Chengliang committed
708
        cr_lens = int(cr_size[i])
Wei Chengliang's avatar
Wei Chengliang committed
709
710
711
712
        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
713
714
715
716

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

        for j in np.arange(cr_lens):
Wei Chengliang's avatar
Wei Chengliang committed
717
            x_n = int(np.cos(pos_angle)*j - np.sin(pos_angle)*0)
Fang Yuedong's avatar
Fang Yuedong committed
718
719
            if x_n < 0:
                x_n = x_n + cr_lens+1
Wei Chengliang's avatar
Wei Chengliang committed
720
            y_n = int(np.sin(pos_angle)*j + np.cos(pos_angle)*0)
Wei Chengliang's avatar
Wei Chengliang committed
721
            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
722
723
                continue
            crMatrix[y_n, x_n] = pix_energy
Fang Yuedong's avatar
Fang Yuedong committed
724
725
726
727
728
729
730
731
732
733
734
735
736

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

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

    return exparrnormal