-
Notifications
You must be signed in to change notification settings - Fork 0
/
extra_proc.py
49 lines (36 loc) · 1.55 KB
/
extra_proc.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#############################################
# extra_proc.py
#
# Aritra Ghosh
#
# This script is meant to do extra processing to all the simulated images -- Convolve with the suitable PSF--add the noise.
############################################
import numpy as np
import pylab as plt
import tqdm
from astropy.io import fits
from multiprocessing import Pool
from astropy.convolution import convolve
dataReadPath = ''
dataWritePath = ''
psfFilePath = ''
noiseFilePath = ''
NUM_FILES_TOTAL = 100000
NUM_THREADS = 15
IMG_WIDTH = IMG_HEIGHT = 83 #No of pixels length/width -wise in each cutout
noise = np.load(noiseFilePath)
def psf_and_noise(i):
img_data = fits.getdata(dataReadPath+"output_img_"+str(i)+".fits",memmap=False)
psf_data = fits.getdata(psfFilePath, memmap=False)
convolved_data = convolve(img_data,psf_data,boundary='extend')
final_data = convolved_data + np.random.choice(noise,size=(IMG_WIDTH,IMG_HEIGHT))
hdu = fits.PrimaryHDU(final_data)
hdu.writeto(dataWritePath+"output_img_"+str(i)+".fits",overwrite=True)
# !!!! --- if you want to replace the exisiting file -- replace the name with the same name as that of the existing file and also add overwrite=True
del hdu #to avoid garbage collection issues
pl = Pool(NUM_THREADS)
pl.map(psf_and_noise,range(0,NUM_FILES_TOTAL))
#if you don't want this just replace this with the previous line
#this next line does the same thing as the pool call above, but now with a progress bar
#for _ in tqdm.tqdm(pl.imap_unordered(psf_and_noise,range(0,NUM_FILES_TOTAL)), total=NUM_FILES_TOTAL):
# pass