Skip to content

Commit f84d985

Browse files
yuedong0607yuedong0607
committed
black reformated
1 parent 78ad325 commit f84d985

19 files changed

Lines changed: 351 additions & 264 deletions

‎roman_imsim/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,4 +18,4 @@
1818
from .wcs import *
1919
from .skycat import *
2020
from .photonOps import *
21-
from .bandpass import *
21+
from .bandpass import *

‎roman_imsim/effects/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,4 +13,4 @@
1313
from .read_noise import read_noise
1414
from .gain import gain
1515
from .bias import bias
16-
from .setup_sky import setup_sky
16+
from .setup_sky import setup_sky

‎roman_imsim/effects/background.py‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,14 +7,17 @@
77
class background(roman_effects):
88
def __init__(self, params, base, logger, rng, rng_iter=None):
99
super().__init__(params, base, logger, rng, rng_iter)
10-
self.thermal_background = self.params['thermal_background'] if 'thermal_background' \
11-
in self.params else False
12-
self.stray_light = self.params['stray_light'] if 'stray_light' in self.params else False
10+
self.thermal_background = (
11+
self.params["thermal_background"] if "thermal_background" in self.params else False
12+
)
13+
self.stray_light = self.params["stray_light"] if "stray_light" in self.params else False
1314

14-
self.model = getattr(self, self.params['model'])
15+
self.model = getattr(self, self.params["model"])
1516
if self.model is None:
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__)))
17+
self.logger.warning(
18+
"%s hasn't been implemented yet, the simple model will be applied for %s"
19+
% (str(self.params["model"]), str(self.__class__.__name__))
20+
)
1821
self.model = self.simple_model
1922

2023
def simple_model(self, image):
@@ -29,29 +32,29 @@ def simple_model(self, image):
2932
else:
3033
radec = pointing.radec
3134
sky_level = roman.getSkyLevel(pointing.bpass, world_pos=radec, date=pointing.date)
32-
self.logger.debug('Adding sky_level = %s', sky_level)
35+
self.logger.debug("Adding sky_level = %s", sky_level)
3336

3437
if self.stray_light:
35-
self.logger.debug('Stray light fraction = %s', roman.stray_light_fraction)
36-
sky_level *= (1.0 + roman.stray_light_fraction)
38+
self.logger.debug("Stray light fraction = %s", roman.stray_light_fraction)
39+
sky_level *= 1.0 + roman.stray_light_fraction
3740
# Create sky image
3841
self.sky = galsim.Image(bounds=image.bounds, wcs=pointing.WCS)
3942
pointing.WCS.makeSkyImage(self.sky, sky_level)
4043
if self.thermal_background:
4144
tb = roman.thermal_backgrounds[pointing.filter] * pointing.exptime
42-
self.logger.debug('Adding thermal background: %s', tb)
45+
self.logger.debug("Adding thermal background: %s", tb)
4346
self.sky += tb
4447
self.sky.addNoise(self.noise)
4548

4649
# [TODO] Not entirely sure about this block, since the 'auto' option is meant to
4750
# let the software choose which drawing method to use based on the total flux.
48-
if self.base['image']['draw_method'] not in ['phot', 'auto']:
51+
if self.base["image"]["draw_method"] not in ["phot", "auto"]:
4952
image.addNoise(self.noise)
5053

5154
# Adding sky level to the image.
5255
image += self.sky[self.sky.bounds & image.bounds]
5356
if self.save_diff:
5457
prev = image.copy()
5558
diff = prev - orig
56-
diff.write(os.path.join(self.diff_dir, 'sky_a.fits'))
59+
diff.write(os.path.join(self.diff_dir, "sky_a.fits"))
5760
return image

‎roman_imsim/effects/bias.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,12 @@ class bias(roman_effects):
88
def __init__(self, params, base, logger, rng, rng_iter=None):
99
super().__init__(params, base, logger, rng, rng_iter)
1010

11-
self.model = getattr(self, self.params['model'])
11+
self.model = getattr(self, self.params["model"])
1212
if self.model is None:
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__)))
13+
self.logger.warning(
14+
"%s hasn't been implemented yet, the simple model will be applied for %s"
15+
% (str(self.params["model"]), str(self.__class__.__name__))
16+
)
1517
self.model = self.simple_model
1618

1719
def simple_model(self, image):
@@ -25,7 +27,7 @@ def lab_model(self, image):
2527
self.df = fio.FITS(os.path.join(self.sca_filepath, sca_number_to_file[self.sca]))
2628

2729
self.logger.warning("Lab measured model will be applied for bias.")
28-
bias = self.df['BIAS'][:, :] # 4096x4096 img
30+
bias = self.df["BIAS"][:, :] # 4096x4096 img
2931

3032
image.array[:, :] += bias
3133
return image

‎roman_imsim/effects/brighter_fatter.py‎

Lines changed: 115 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,12 @@ class brighter_fatter(roman_effects):
99
def __init__(self, params, base, logger, rng, rng_iter=None):
1010
super().__init__(params, base, logger, rng, rng_iter)
1111

12-
self.model = getattr(self, self.params['model'])
12+
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"%(
15-
str(self.params['model']), str(self.__class__.__name__)))
14+
self.logger.warning(
15+
"%s hasn't been implemented yet, the simple model will be applied for %s"
16+
% (str(self.params["model"]), str(self.__class__.__name__))
17+
)
1618
self.model = self.simple_model
1719

1820
def simple_model(self, image):
@@ -50,49 +52,60 @@ def lab_model(self, image):
5052
n_max = 32
5153
m_max = 32
5254
num_grids = 4
53-
n_sub = n_max//num_grids
54-
m_sub = m_max//num_grids
55+
n_sub = n_max // num_grids
56+
m_sub = m_max // num_grids
5557

5658
# =======================================================================
5759
# solve boundary shfit kernel aX components
5860
# =======================================================================
59-
a_area = self.df['BFE'][:, :, :, :] # 5x5x32x32
60-
a_components = np.zeros((4, 2*nbfe+1, 2*nbfe+1, n_max, m_max)) # 4x5x5x32x32
61+
a_area = self.df["BFE"][:, :, :, :] # 5x5x32x32
62+
a_components = np.zeros((4, 2 * nbfe + 1, 2 * nbfe + 1, n_max, m_max)) # 4x5x5x32x32
6163

6264
# solve aR aT aL aB for each a
6365
for n in range(n_max): # m_max and n_max = 32 (binned in 128x128)
6466
for m in range(m_max):
6567
a = a_area[:, :, n, m] # a in (2 x nbfe+1)*(2 x nbfe+1)
6668

6769
# 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]])
70+
a = (a + np.fliplr(a) + np.flipud(a) + np.flip(a)) / 4.0
71+
72+
r = 0.5 * (3.25 / 4.25) ** (1.5) / 1.5 # source-boundary projection
73+
B = (a[2, 2], a[3, 2], a[2, 3], a[3, 3], a[4, 2], a[2, 4], a[3, 4], a[4, 4])
74+
75+
A = np.array(
76+
[
77+
[-2, -2, 0, 0, 0, 0, 0],
78+
[0, 1, 0, -1, -2, 0, 0],
79+
[1, 0, -1, 0, -2, 0, 0],
80+
[0, 0, 0, 0, 2, -2, 0],
81+
[0, 0, 0, 1, 0, -2 * r, 0],
82+
[0, 0, 1, 0, 0, -2 * r, 0],
83+
[0, 0, 0, 0, 0, 1 + r, -1],
84+
[0, 0, 0, 0, 0, 0, 2],
85+
]
86+
)
8287

8388
s1, s2, s3, s4, s5, s6, s7 = np.linalg.lstsq(A, B, rcond=None)[0]
8489

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],])
90+
aR = np.array(
91+
[
92+
[0.0, -s7, -r * s6, r * s6, s7],
93+
[0.0, -s6, -s5, s5, s6],
94+
[0.0, -s3, -s1, s1, s3],
95+
[0.0, -s6, -s5, s5, s6],
96+
[0.0, -s7, -r * s6, r * s6, s7],
97+
]
98+
)
99+
100+
aT = np.array(
101+
[
102+
[0.0, 0.0, 0.0, 0.0, 0.0],
103+
[-s7, -s6, -s4, -s6, -s7],
104+
[-r * s6, -s5, -s2, -s5, -r * s6],
105+
[r * s6, s5, s2, s5, r * s6],
106+
[s7, s6, s4, s6, s7],
107+
]
108+
)
96109

97110
aL = aR[::-1, ::-1]
98111
aB = aT[::-1, ::-1]
@@ -115,58 +128,96 @@ def lab_model(self, image):
115128
# where_sat = np.where(array_pad > saturation_array)
116129
# array_pad[ where_sat ] = saturation_array[ where_sat ]
117130
# array_pad = array_pad[4:-4,4:-4]
118-
saturate = self.cross_refer('saturate')
131+
saturate = self.cross_refer("saturate")
119132
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
133+
array_pad = np.pad(
134+
array_pad, [(4 + nbfe, 4 + nbfe), (4 + nbfe, 4 + nbfe)], mode="symmetric"
135+
) # 4100x4100 array
122136

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

126140
# run in sub grids to reduce memory
127141

128142
# pad and expand kernels
129-
t = np.zeros((bin_size*n_sub, n_sub))
143+
t = np.zeros((bin_size * n_sub, n_sub))
130144
for row in range(t.shape[0]):
131-
t[row, row//(bin_size)] = 1
145+
t[row, row // (bin_size)] = 1
132146

133147
for gj in range(num_grids):
134148
for gi in range(num_grids):
135149

136150
# (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))
151+
a_components_pad = np.zeros(
152+
(4, 2 * nbfe + 1, 2 * nbfe + 1, bin_size * n_sub + 2 * nbfe, bin_size * m_sub + 2 * nbfe)
153+
)
139154

140155
for comp in range(4):
141-
for j in range(2*nbfe+1):
142-
for i in range(2*nbfe+1):
156+
for j in range(2 * nbfe + 1):
157+
for i in range(2 * nbfe + 1):
143158
# 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)
159+
tmp = (
160+
t.dot(
161+
a_components[
162+
comp,
163+
j,
164+
i,
165+
gj * n_sub : (gj + 1) * n_sub,
166+
gi * m_sub : (gi + 1) * m_sub,
167+
]
168+
)
169+
).dot(t.T)
146170
a_components_pad[comp, j, i, :, :] = np.pad(
147-
tmp, [(nbfe, nbfe), (nbfe, nbfe)], mode='symmetric')
171+
tmp, [(nbfe, nbfe), (nbfe, nbfe)], mode="symmetric"
172+
)
148173

149174
# convolve aX_ij with Q_ij
150175
for comp in range(4):
151-
for dy in range(-nbfe, nbfe+1):
152-
for dx in range(-nbfe, nbfe+1):
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]
160-
161-
dj = int(np.sin(comp*np.pi/2))
162-
di = int(np.cos(comp*np.pi/2))
163-
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])
176+
for dy in range(-nbfe, nbfe + 1):
177+
for dx in range(-nbfe, nbfe + 1):
178+
dQ_components[
179+
comp,
180+
gj * bin_size * n_sub : (gj + 1) * bin_size * n_sub,
181+
gi * bin_size * m_sub : (gi + 1) * bin_size * m_sub,
182+
] += (
183+
a_components_pad[
184+
comp,
185+
nbfe + dy,
186+
nbfe + dx,
187+
nbfe - dy : nbfe - dy + bin_size * n_sub,
188+
nbfe - dx : nbfe - dx + bin_size * m_sub,
189+
]
190+
* array_pad[
191+
-dy
192+
+ nbfe
193+
+ gj * bin_size * n_sub : -dy
194+
+ nbfe
195+
+ (gj + 1) * bin_size * n_sub,
196+
-dx
197+
+ nbfe
198+
+ gi * bin_size * m_sub : -dx
199+
+ nbfe
200+
+ (gi + 1) * bin_size * m_sub,
201+
]
202+
)
203+
204+
dj = int(np.sin(comp * np.pi / 2))
205+
di = int(np.cos(comp * np.pi / 2))
206+
207+
dQ_components[
208+
comp,
209+
gj * bin_size * n_sub : (gj + 1) * bin_size * n_sub,
210+
gi * bin_size * m_sub : (gi + 1) * bin_size * m_sub,
211+
] *= 0.5 * (
212+
array_pad[
213+
nbfe + gj * bin_size * n_sub : nbfe + (gj + 1) * bin_size * n_sub,
214+
nbfe + gi * bin_size * m_sub : nbfe + (gi + 1) * bin_size * m_sub,
215+
]
216+
+ array_pad[
217+
dj + nbfe + gj * bin_size * n_sub : dj + nbfe + (gj + 1) * bin_size * n_sub,
218+
di + nbfe + gi * bin_size * m_sub : di + nbfe + (gi + 1) * bin_size * m_sub,
219+
]
220+
)
170221

171222
image.array[:, :] -= dQ_components.sum(axis=0)
172223
image.array[:, 1:] += dQ_components[0][:, :-1]

‎roman_imsim/effects/dark_current.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,12 @@ class dark_current(roman_effects):
1010
def __init__(self, params, base, logger, rng, rng_iter=None):
1111
super().__init__(params, base, logger, rng, rng_iter)
1212

13-
self.model = getattr(self, self.params['model'])
13+
self.model = getattr(self, self.params["model"])
1414
if self.model is None:
15-
self.logger.warning("%s hasn't been implemented yet, the simple model will be applied for %s"%(
16-
str(self.params['model']), str(self.__class__.__name__)))
15+
self.logger.warning(
16+
"%s hasn't been implemented yet, the simple model will be applied for %s"
17+
% (str(self.params["model"]), str(self.__class__.__name__))
18+
)
1719
self.model = self.simple_model
1820

1921
def simple_model(self, image):
@@ -35,7 +37,7 @@ def lab_model(self, image):
3537

3638
exptime = self.pointing.exptime
3739
self.logger.warning("Lab measured model will be applied for dark current.")
38-
self.dark_current_ = roman.dark_current * exptime + self.df['DARK'][:, :].flatten() * exptime
40+
self.dark_current_ = roman.dark_current * exptime + self.df["DARK"][:, :].flatten() * exptime
3941
dark_current_ = self.dark_current_.clip(0)
4042
# opt for numpy random geneator instead for speed
4143
self.im_dark = self.rng_np.poisson(dark_current_).reshape(image.array.shape).astype(image.dtype)

‎roman_imsim/effects/dead_pix.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,12 @@ class dead_pix(roman_effects):
88
def __init__(self, params, base, logger, rng, rng_iter=None):
99
super().__init__(params, base, logger, rng, rng_iter)
1010

11-
self.model = getattr(self, self.params['model'])
11+
self.model = getattr(self, self.params["model"])
1212
if self.model is None:
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__)))
13+
self.logger.warning(
14+
"%s hasn't been implemented yet, the simple model will be applied for %s"
15+
% (str(self.params["model"]), str(self.__class__.__name__))
16+
)
1517
self.model = self.simple_model
1618

1719
def simple_model(self, image):
@@ -25,7 +27,7 @@ def lab_model(self, image):
2527
self.df = fio.FITS(os.path.join(self.sca_filepath, sca_number_to_file[self.sca]))
2628

2729
self.logger.warning("Lab measured model will be applied for dead pixel.")
28-
dead_mask = self.df['BADPIX'][:, :] & 1 # 4096x4096 array
30+
dead_mask = self.df["BADPIX"][:, :] & 1 # 4096x4096 array
2931
image.array[dead_mask > 0] = 0
3032

3133
return image

0 commit comments

Comments
 (0)