Skip to content

Commit 367ef3a

Browse files
yuedong0607yuedong0607
committed
reformat to meet pep8
1 parent b26288e commit 367ef3a

20 files changed

Lines changed: 387 additions & 1576 deletions

‎roman_imsim/__init__.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,4 @@
1818
from .wcs import *
1919
from .skycat import *
2020
from .photonOps import *
21-
from .bandpass import *
22-
# from .detector_physics import *
23-
from .detector_effects import *
21+
from .bandpass import *

‎roman_imsim/detector_effects.py‎

Lines changed: 0 additions & 1031 deletions
This file was deleted.

‎roman_imsim/effects/background.py‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,26 @@
11
import os
2-
import numpy as np
3-
import fitsio as fio
42
import galsim
53
from roman_imsim.effects import roman_effects
64
import galsim.roman as roman
7-
from .utils import sca_number_to_file
5+
86

97
class background(roman_effects):
108
def __init__(self, params, base, logger, rng, rng_iter=None):
119
super().__init__(params, base, logger, rng, rng_iter)
12-
self.thermal_background = self.params['thermal_background'] if 'thermal_background' in self.params else False
10+
self.thermal_background = self.params['thermal_background'] if 'thermal_background' \
11+
in self.params else False
1312
self.stray_light = self.params['stray_light'] if 'stray_light' in self.params else False
14-
13+
1514
self.model = getattr(self, self.params['model'])
1615
if self.model is None:
17-
self.logger.warning("%s hasn't been implemented yet, the simple model will be applied for %s"%(str(self.params['model']), str(self.__class__.__name__)))
16+
self.logger.warning("%s hasn't been implemented yet, the simple model will be applied for %s"%(
17+
str(self.params['model']), str(self.__class__.__name__)))
1818
self.model = self.simple_model
19-
19+
2020
def simple_model(self, image):
2121
if self.save_diff:
2222
orig = image.copy()
23-
23+
2424
self.logger.warning("Simple model will be applied for background.")
2525
pointing = self.pointing
2626
# Build current specification sky level if sky level not given
@@ -29,28 +29,29 @@ def simple_model(self, image):
2929
else:
3030
radec = pointing.radec
3131
sky_level = roman.getSkyLevel(pointing.bpass, world_pos=radec, date=pointing.date)
32-
self.logger.debug('Adding sky_level = %s',sky_level)
33-
32+
self.logger.debug('Adding sky_level = %s', sky_level)
33+
3434
if self.stray_light:
35-
self.logger.debug('Stray light fraction = %s',roman.stray_light_fraction)
35+
self.logger.debug('Stray light fraction = %s', roman.stray_light_fraction)
3636
sky_level *= (1.0 + roman.stray_light_fraction)
3737
# Create sky image
3838
self.sky = galsim.Image(bounds=image.bounds, wcs=pointing.WCS)
3939
pointing.WCS.makeSkyImage(self.sky, sky_level)
4040
if self.thermal_background:
4141
tb = roman.thermal_backgrounds[pointing.filter] * pointing.exptime
42-
self.logger.debug('Adding thermal background: %s',tb)
42+
self.logger.debug('Adding thermal background: %s', tb)
4343
self.sky += tb
4444
self.sky.addNoise(self.noise)
45-
46-
# [TODO] Not entirely sure about this block, since the 'auto' option is meant to let the software choose which drawing method to use based on the total flux.
45+
46+
# [TODO] Not entirely sure about this block, since the 'auto' option is meant to
47+
# let the software choose which drawing method to use based on the total flux.
4748
if self.base['image']['draw_method'] not in ['phot', 'auto']:
4849
image.addNoise(self.noise)
49-
50+
5051
# Adding sky level to the image.
5152
image += self.sky[self.sky.bounds & image.bounds]
5253
if self.save_diff:
5354
prev = image.copy()
5455
diff = prev - orig
5556
diff.write(os.path.join(self.diff_dir, 'sky_a.fits'))
56-
return image
57+
return image

‎roman_imsim/effects/bias.py‎

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,31 @@
11
import os
2-
import numpy as np
32
import fitsio as fio
4-
import galsim
5-
import galsim.roman as roman
63
from roman_imsim.effects import roman_effects
7-
import roman_imsim.effects as effects
8-
from galsim.config import ParseValue
9-
from .utils import sca_number_to_file, get_pointing
4+
from .utils import sca_number_to_file
5+
106

117
class bias(roman_effects):
128
def __init__(self, params, base, logger, rng, rng_iter=None):
139
super().__init__(params, base, logger, rng, rng_iter)
14-
10+
1511
self.model = getattr(self, self.params['model'])
1612
if self.model is None:
17-
self.logger.warning("%s hasn't been implemented yet, the simple model will be applied for %s"%(str(self.params['model']), str(self.__class__.__name__)))
13+
self.logger.warning("%s hasn't been implemented yet, the simple model will be applied for %s"%(
14+
str(self.params['model']), str(self.__class__.__name__)))
1815
self.model = self.simple_model
1916

2017
def simple_model(self, image):
2118
self.logger.warning("No bias will be applied.")
2219
return image
23-
20+
2421
def lab_model(self, image):
2522
if self.sca_filepath is None:
2623
self.logger.warning("No bias data provided; no bias will be applied.")
2724
return image
2825
self.df = fio.FITS(os.path.join(self.sca_filepath, sca_number_to_file[self.sca]))
29-
26+
3027
self.logger.warning("Lab measured model will be applied for bias.")
31-
bias = self.df['BIAS'][:,:] #4096x4096 img
28+
bias = self.df['BIAS'][:, :] # 4096x4096 img
3229

33-
image.array[:,:] += bias
34-
return image
30+
image.array[:, :] += bias
31+
return image

‎roman_imsim/effects/brighter_fatter.py‎

Lines changed: 92 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -4,20 +4,21 @@
44
from roman_imsim.effects import roman_effects
55
from .utils import sca_number_to_file
66

7+
78
class brighter_fatter(roman_effects):
89
def __init__(self, params, base, logger, rng, rng_iter=None):
910
super().__init__(params, base, logger, rng, rng_iter)
10-
# self.saturation_level = self.params['saturation_level'] if 'saturation_level' in self.params else 100000
11-
11+
1212
self.model = getattr(self, self.params['model'])
1313
if self.model is None:
14-
self.logger.warning("%s hasn't been implemented yet, the simple model will be applied for %s"%(str(self.params['model']), str(self.__class__.__name__)))
14+
self.logger.warning("%s hasn't been implemented yet, the simple model will be applied for %s"%(
15+
str(self.params['model']), str(self.__class__.__name__)))
1516
self.model = self.simple_model
16-
17+
1718
def simple_model(self, image):
1819
self.logger.warning("No bfe effect will be applied.")
1920
return image
20-
21+
2122
def lab_model(self, image):
2223
"""
2324
Apply brighter-fatter effect.
@@ -41,132 +42,136 @@ def lab_model(self, image):
4142
self.logger.warning("No BFE kernel data file provided; no bfe effect will be applied.")
4243
return image
4344
self.df = fio.FITS(os.path.join(self.sca_filepath, sca_number_to_file[self.sca]))
44-
45+
4546
self.logger.warning("Lab measured model will be applied for brighter-fatter effect.")
4647

47-
nbfe = 2 ## kernel of bfe in shape (2 x nbfe+1)*(2 x nbfe+1)
48+
nbfe = 2 # kernel of bfe in shape (2 x nbfe+1)*(2 x nbfe+1)
4849
bin_size = 128
4950
n_max = 32
5051
m_max = 32
5152
num_grids = 4
5253
n_sub = n_max//num_grids
5354
m_sub = m_max//num_grids
5455

55-
##=======================================================================
56-
## solve boundary shfit kernel aX components
57-
##=======================================================================
58-
a_area = self.df['BFE'][:,:,:,:] #5x5x32x32
59-
a_components = np.zeros( (4, 2*nbfe+1, 2*nbfe+1, n_max, m_max) ) #4x5x5x32x32
56+
# =======================================================================
57+
# solve boundary shfit kernel aX components
58+
# =======================================================================
59+
a_area = self.df['BFE'][:, :, :, :] # 5x5x32x32
60+
a_components = np.zeros((4, 2*nbfe+1, 2*nbfe+1, n_max, m_max)) # 4x5x5x32x32
6061

61-
##solve aR aT aL aB for each a
62-
for n in range(n_max): #m_max and n_max = 32 (binned in 128x128)
62+
# solve aR aT aL aB for each a
63+
for n in range(n_max): # m_max and n_max = 32 (binned in 128x128)
6364
for m in range(m_max):
64-
a = a_area[:,:, n, m] ## a in (2 x nbfe+1)*(2 x nbfe+1)
65-
66-
## assume two parity symmetries
67-
a = ( a + np.fliplr(a) + np.flipud(a) + np.flip(a) )/4.
68-
69-
r = 0.5* ( 3.25/4.25 )**(1.5) / 1.5 ## source-boundary projection
70-
B = (a[2,2], a[3,2], a[2,3], a[3,3],
71-
a[4,2], a[2,4], a[3,4], a[4,4] )
72-
73-
A = np.array( [ [ -2 , -2 , 0 , 0 , 0 , 0 , 0 ],
74-
[ 0 , 1 , 0 , -1 , -2 , 0 , 0 ],
75-
[ 1 , 0 , -1 , 0 , -2 , 0 , 0 ],
76-
[ 0 , 0 , 0 , 0 , 2 , -2 , 0 ],
77-
[ 0 , 0 , 0 , 1 , 0 ,-2*r, 0 ],
78-
[ 0 , 0 , 1 , 0 , 0 ,-2*r, 0 ],
79-
[ 0 , 0 , 0 , 0 , 0 , 1+r, -1 ],
80-
[ 0 , 0 , 0 , 0 , 0 , 0 , 2 ] ])
81-
82-
83-
s1,s2,s3,s4,s5,s6,s7 = np.linalg.lstsq(A, B, rcond=None)[0]
84-
85-
aR = np.array( [[ 0. , -s7 ,-r*s6 , r*s6 , s7 ],
86-
[ 0. , -s6 , -s5 , s5 , s6 ],
87-
[ 0. , -s3 , -s1 , s1 , s3 ],
88-
[ 0. , -s6 , -s5 , s5 , s6 ],
89-
[ 0. , -s7 ,-r*s6 , r*s6 , s7 ],])
90-
91-
92-
aT = np.array( [[ 0. , 0. , 0. , 0. , 0. ],
93-
[ -s7 , -s6 , -s4 , -s6 , -s7 ],
94-
[ -r*s6 , -s5 , -s2 , -s5 , -r*s6 ],
95-
[ r*s6 , s5 , s2 , s5 , r*s6 ],
96-
[ s7 , s6 , s4 , s6 , s7 ],])
97-
65+
a = a_area[:, :, n, m] # a in (2 x nbfe+1)*(2 x nbfe+1)
66+
67+
# assume two parity symmetries
68+
a = (a + np.fliplr(a) + np.flipud(a) + np.flip(a))/4.
69+
70+
r = 0.5 * (3.25/4.25)**(1.5) / 1.5 # source-boundary projection
71+
B = (a[2, 2], a[3, 2], a[2, 3], a[3, 3],
72+
a[4, 2], a[2, 4], a[3, 4], a[4, 4])
73+
74+
A = np.array([[-2 , -2 , 0 , 0 , 0 , 0 , 0],
75+
[0 , 1 , 0 , -1 , -2 , 0 , 0],
76+
[1 , 0 , -1 , 0 , -2 , 0 , 0],
77+
[0 , 0 , 0 , 0 , 2 , -2 , 0],
78+
[0 , 0 , 0 , 1 , 0 , -2*r, 0],
79+
[0 , 0 , 1 , 0 , 0 , -2*r, 0],
80+
[0 , 0 , 0 , 0 , 0 , 1+r, -1],
81+
[0 , 0 , 0 , 0 , 0 , 0 , 2]])
82+
83+
s1, s2, s3, s4, s5, s6, s7 = np.linalg.lstsq(A, B, rcond=None)[0]
84+
85+
aR = np.array([[0. , -s7 , -r*s6 , r*s6 , s7],
86+
[0. , -s6 , -s5 , s5 , s6],
87+
[0. , -s3 , -s1 , s1 , s3],
88+
[0. , -s6 , -s5 , s5 , s6],
89+
[0. , -s7 , -r*s6 , r*s6 , s7],])
90+
91+
aT = np.array([[0. , 0. , 0. , 0. , 0.],
92+
[-s7 , -s6 , -s4 , -s6 , -s7],
93+
[-r*s6 , -s5 , -s2 , -s5 , -r*s6],
94+
[r*s6 , s5 , s2 , s5 , r*s6],
95+
[s7 , s6 , s4 , s6 , s7],])
9896

9997
aL = aR[::-1, ::-1]
10098
aB = aT[::-1, ::-1]
10199

100+
a_components[0, :, :, n, m] = aR[:, :]
101+
a_components[1, :, :, n, m] = aT[:, :]
102+
a_components[2, :, :, n, m] = aL[:, :]
103+
a_components[3, :, :, n, m] = aB[:, :]
102104

105+
# =============================
106+
# Apply bfe to image
107+
# =============================
103108

104-
105-
a_components[0, :,:, n, m] = aR[:,:]
106-
a_components[1, :,:, n, m] = aT[:,:]
107-
a_components[2, :,:, n, m] = aL[:,:]
108-
a_components[3, :,:, n, m] = aB[:,:]
109-
110-
##=============================
111-
## Apply bfe to image
112-
##=============================
113-
114-
## pad and expand kernels
115-
## The img is clipped by the saturation level here to cap the brighter fatter effect and avoid unphysical behavior
109+
# pad and expand kernels
110+
# The img is clipped by the saturation level here to cap the brighter fatter effect
111+
# and avoid unphysical behavior
116112

117113
# array_pad = image.copy().array
118114
# saturation_array = np.ones_like(array_pad) * self.saturation_level
119115
# where_sat = np.where(array_pad > saturation_array)
120116
# array_pad[ where_sat ] = saturation_array[ where_sat ]
121117
# array_pad = array_pad[4:-4,4:-4]
122118
saturate = self.cross_refer('saturate')
123-
array_pad = saturate.apply(image = image.copy()).array[4:-4,4:-4] # img of interest 4088x4088
124-
array_pad = np.pad(array_pad, [(4+nbfe,4+nbfe),(4+nbfe,4+nbfe)], mode='symmetric') #4100x4100 array
125-
119+
array_pad = saturate.apply(image=image.copy()).array[4:-4, 4:-4] # img of interest 4088x4088
120+
array_pad = np.pad(array_pad, [(4+nbfe, 4+nbfe), (4+nbfe, 4+nbfe)],
121+
mode='symmetric') # 4100x4100 array
126122

127-
dQ_components = np.zeros( (4, bin_size*n_max, bin_size*m_max) ) #(4, 4096, 4096) in order of [aR, aT, aL, aB]
123+
# (4, 4096, 4096) in order of [aR, aT, aL, aB]
124+
dQ_components = np.zeros((4, bin_size*n_max, bin_size*m_max))
128125

126+
# run in sub grids to reduce memory
129127

130-
### run in sub grids to reduce memory
131-
132-
## pad and expand kernels
128+
# pad and expand kernels
133129
t = np.zeros((bin_size*n_sub, n_sub))
134130
for row in range(t.shape[0]):
135-
t[row, row//(bin_size) ] =1
136-
137-
131+
t[row, row//(bin_size)] = 1
138132

139133
for gj in range(num_grids):
140134
for gi in range(num_grids):
141135

142-
a_components_pad = np.zeros( (4, 2*nbfe+1, 2*nbfe+1, bin_size*n_sub+2*nbfe, bin_size*m_sub+2*nbfe) ) #(4,5,5,sub_grid,sub_grid)
143-
136+
# (4,5,5,sub_grid,sub_grid)
137+
a_components_pad = np.zeros((4, 2*nbfe+1, 2*nbfe+1, bin_size
138+
* n_sub+2*nbfe, bin_size*m_sub+2*nbfe))
144139

145140
for comp in range(4):
146141
for j in range(2*nbfe+1):
147142
for i in range(2*nbfe+1):
148-
tmp = (t.dot( a_components[comp,j,i,gj*n_sub:(gj+1)*n_sub,gi*m_sub:(gi+1)*m_sub] ) ).dot(t.T) #sub_grid*sub_grid
149-
a_components_pad[comp, j, i, :, :] = np.pad(tmp, [(nbfe,nbfe),(nbfe,nbfe)], mode='symmetric')
143+
# sub_grid*sub_grid
144+
tmp = (t.dot(a_components[comp, j, i, gj*n_sub:(gj+1)
145+
* n_sub, gi*m_sub:(gi+1)*m_sub])).dot(t.T)
146+
a_components_pad[comp, j, i, :, :] = np.pad(
147+
tmp, [(nbfe, nbfe), (nbfe, nbfe)], mode='symmetric')
150148

151-
#convolve aX_ij with Q_ij
149+
# convolve aX_ij with Q_ij
152150
for comp in range(4):
153151
for dy in range(-nbfe, nbfe+1):
154152
for dx in range(-nbfe, nbfe+1):
155-
dQ_components[comp, gj*bin_size*n_sub : (gj+1)*bin_size*n_sub , gi*bin_size*m_sub : (gi+1)*bin_size*m_sub]\
156-
+= a_components_pad[comp, nbfe+dy, nbfe+dx, nbfe-dy:nbfe-dy+bin_size*n_sub, nbfe-dx:nbfe-dx+bin_size*m_sub ]\
157-
*array_pad[ -dy + nbfe + gj*bin_size*n_sub : -dy + nbfe+ (gj+1)*bin_size*n_sub , -dx + nbfe + gi*bin_size*m_sub : -dx + nbfe + (gi+1)*bin_size*m_sub ]
153+
dQ_components[comp, gj*bin_size*n_sub : (gj+1)*bin_size*n_sub,
154+
gi*bin_size*m_sub : (gi+1)*bin_size*m_sub]\
155+
+= a_components_pad[comp, nbfe+dy, nbfe+dx, nbfe-dy:nbfe-dy+bin_size*n_sub,
156+
nbfe-dx:nbfe-dx+bin_size*m_sub]\
157+
* array_pad[-dy + nbfe + gj*bin_size*n_sub : -dy + nbfe
158+
+ (gj+1)*bin_size*n_sub, -dx + nbfe + gi*bin_size*m_sub : -dx
159+
+ nbfe + (gi+1)*bin_size*m_sub]
158160

159161
dj = int(np.sin(comp*np.pi/2))
160162
di = int(np.cos(comp*np.pi/2))
161163

162-
dQ_components[comp, gj*bin_size*n_sub : (gj+1)*bin_size*n_sub , gi*bin_size*m_sub : (gi+1)*bin_size*m_sub]\
163-
*= 0.5*(array_pad[ nbfe + gj*bin_size*n_sub : nbfe+ (gj+1)*bin_size*n_sub , nbfe + gi*bin_size*m_sub : nbfe + (gi+1)*bin_size*m_sub ] +\
164-
array_pad[dj+nbfe + gj*bin_size*n_sub : dj+nbfe+ (gj+1)*bin_size*n_sub , di+nbfe + gi*bin_size*m_sub : di+nbfe + (gi+1)*bin_size*m_sub] )
164+
dQ_components[comp, gj*bin_size*n_sub : (gj+1)*bin_size*n_sub ,
165+
gi*bin_size*m_sub : (gi+1)*bin_size*m_sub]\
166+
*= 0.5*(array_pad[nbfe + gj*bin_size*n_sub : nbfe + (gj+1)*bin_size*n_sub,
167+
nbfe + gi*bin_size*m_sub : nbfe + (gi+1)*bin_size*m_sub]
168+
+ array_pad[dj+nbfe + gj*bin_size*n_sub : dj+nbfe + (gj+1)*bin_size*n_sub
169+
, di+nbfe + gi*bin_size*m_sub : di+nbfe + (gi+1)*bin_size*m_sub])
165170

166-
image.array[:,:] -= dQ_components.sum(axis=0)
167-
image.array[:,1:] += dQ_components[0][:,:-1]
168-
image.array[1:,:] += dQ_components[1][:-1,:]
169-
image.array[:,:-1] += dQ_components[2][:,1:]
170-
image.array[:-1,:] += dQ_components[3][1:,:]
171+
image.array[:, :] -= dQ_components.sum(axis=0)
172+
image.array[:, 1:] += dQ_components[0][:, :-1]
173+
image.array[1:, :] += dQ_components[1][:-1, :]
174+
image.array[:, :-1] += dQ_components[2][:, 1:]
175+
image.array[:-1, :] += dQ_components[3][1:, :]
171176

172-
return image
177+
return image

0 commit comments

Comments
 (0)