Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
csst-sims
csst_sim_docs
Commits
7c99b572
Commit
7c99b572
authored
May 17, 2024
by
GZhao
Browse files
update cpic sourcecode
parent
7283405d
Changes
9
Show whitespace changes
Inline
Side-by-side
source/csst_cpic_sim/__init__.py
0 → 100644
View file @
7c99b572
from
.main
import
quick_run_v2
,
vis_observation
from
.optics
import
focal_mask
from
.target
import
star_photlam
,
planet_contrast
,
extract_target_x_y
,
spectrum_generator
from
.camera
import
CosmicRayFrameMaker
,
sky_frame_maker
from
.config
import
__version__
__all__
=
[
"CosmicRayFrameMaker"
,
"sky_frame_maker"
,
"star_photlam"
,
"planet_contrast"
,
"extract_target_x_y"
,
"spectrum_generator"
,
"focal_mask"
,
"quick_run_v2"
,
"vis_observation"
,
"__version__"
]
\ No newline at end of file
source/csst_cpic_sim/camera.py
0 → 100644
View file @
7c99b572
import
math
import
numpy
as
np
import
scipy.ndimage
as
nd
from
astropy.io
import
fits
import
matplotlib.pyplot
as
plt
from
.config
import
config
,
S
from
.utils
import
region_replace
,
random_seed_select
from
.io
import
log
from
.optics
import
filter_throughput
cpism_refdata
=
config
[
'cpism_refdata'
]
MAG_SYSTEM
=
config
[
'mag_system'
]
solar_spectrum
=
S
.
FileSpectrum
(
config
[
'solar_spectrum'
])
solar_spectrum
.
convert
(
'photlam'
)
def
sky_frame_maker
(
band
,
skybg
,
platescale
,
shape
):
"""
generate a sky background frame.
Parameters
----------
band : str
The band of the sky background.
skybg : str
The sky background file name.
platescale : float
The platescale of the camera in arcsec/pixel.
shape : tuple
The shape of the output frame. (y, x)
Returns
-------
sky_bkg_frame : numpy.ndarray
The sky background frame.
"""
filter
=
filter_throughput
(
band
)
sk_spec
=
solar_spectrum
.
renorm
(
skybg
,
MAG_SYSTEM
,
filter
)
sky_bkg_frame
=
np
.
zeros
(
shape
)
sky_bkg_frame
+=
(
sk_spec
*
filter
).
integrate
()
*
platescale
**
2
return
sky_bkg_frame
class
CRobj
(
object
):
"""
Cosmic ray object.
Attributes
----------
flux : float
The flux of the cosmic ray.
angle : float
The angle of the cosmic ray.
sigma : float
The width of the cosmic ray.
length : int
The length of the cosmic ray.
"""
def
__init__
(
self
,
flux
,
angle
,
sigma
,
length
)
->
None
:
self
.
flux
=
flux
self
.
angle
=
angle
self
.
sigma
=
sigma
self
.
length
=
length
class
CosmicRayFrameMaker
(
object
):
"""
Cosmic ray frame maker.
Parameters
----------
depth : float
The depth of the camera pixel in um.
pitch : float
The pitch of the camera pixel in um.
cr_rate : float
The cosmic ray rate per second per cm2.
"""
def
__init__
(
self
)
->
None
:
self
.
tmp_size
=
[
7
,
101
]
self
.
freq_power
=
-
0.9
self
.
trail_std
=
0.1
self
.
depth
=
10
# um
self
.
pitch
=
13
# um
self
.
cr_rate
=
1
# particle per s per cm2 from Miles et al. 2021
def
make_CR
(
self
,
length
,
sigma
,
seed
=-
1
):
"""
make a image of cosmic ray with given length and sigma.
Parameters
----------
length : int
The length of the cosmic ray in pixel.
sigma : float
The width of the cosmic ray in pixel.
Returns
-------
output : numpy.ndarray
The image of cosmic ray.
"""
h
=
self
.
tmp_size
[
0
]
w
=
self
.
tmp_size
[
1
]
freq
=
((
w
-
1
)
/
2
-
np
.
abs
(
np
.
arange
(
w
)
-
(
w
-
1
)
/
2
)
+
1
)
**
(
self
.
freq_power
)
x
=
np
.
arange
(
w
)
-
(
w
-
1
)
/
2
hl
=
(
length
-
1
)
/
2
x_wing
=
np
.
exp
(
-
(
np
.
abs
(
x
)
-
hl
)
**
2
/
sigma
/
sigma
/
2
)
x_wing
[
np
.
abs
(
x
)
<
hl
]
=
1
cr
=
np
.
zeros
([
h
,
w
])
center
=
(
h
-
1
)
/
2
for
i
in
range
(
h
):
phase
=
np
.
random
.
rand
(
w
)
*
2
*
np
.
pi
trail0
=
abs
(
np
.
fft
.
fft
(
freq
*
np
.
sin
(
phase
)
+
1j
*
x
*
np
.
cos
(
phase
)))
# TODO maybe somthing wrong
trail_norm
=
(
trail0
-
trail0
.
mean
())
/
trail0
.
std
()
cr
[
i
,
:]
=
np
.
exp
(
-
(
i
-
center
)
**
2
/
sigma
/
sigma
/
2
)
\
*
(
trail_norm
*
self
.
trail_std
+
1
)
*
x_wing
output
=
np
.
zeros
([
w
,
w
])
d
=
(
w
-
h
)
//
2
output
[
d
:
d
+
h
,
:]
=
cr
return
output
def
_length_rand
(
self
,
N
,
seed
=-
1
):
"""
randomly generate N cosmic ray length.
"""
len_out
=
[]
seed
=
random_seed_select
(
seed
=
seed
)
log
.
debug
(
f
"cr length seed:
{
seed
}
"
)
for
i
in
range
(
N
):
x2y2
=
2
while
x2y2
>
1
:
lx
=
1
-
2
*
np
.
random
.
rand
()
ly
=
1
-
2
*
np
.
random
.
rand
()
x2y2
=
lx
*
lx
+
ly
*
ly
z
=
1
-
2
*
x2y2
r
=
2
*
np
.
sqrt
(
x2y2
*
(
1
-
x2y2
))
length
=
abs
(
r
/
z
*
self
.
depth
)
pitch
=
self
.
pitch
len_out
.
append
(
int
(
length
/
pitch
))
return
np
.
array
(
len_out
)
def
_number_rand
(
self
,
expt
,
pixsize
,
random
=
False
,
seed
=-
1
):
"""
randomly generate the number of cosmic rays.
"""
area
=
(
self
.
pitch
/
1e4
)
**
2
*
pixsize
[
0
]
*
pixsize
[
1
]
ncr
=
area
*
expt
*
self
.
cr_rate
if
random
:
seed
=
random_seed_select
(
seed
=
seed
)
log
.
debug
(
f
"cr count:
{
seed
}
"
)
ncr
=
np
.
random
.
poisson
(
ncr
)
else
:
ncr
=
int
(
ncr
)
self
.
ncr
=
ncr
return
ncr
def
_sigma_rand
(
self
,
N
,
seed
=-
1
):
"""
randomly generate N cosmic ray sigma.
"""
sig_sig
=
0.5
# asuming the sigma of size of cosmic ray is 0.5
seed
=
random_seed_select
(
seed
=
seed
)
log
.
debug
(
f
"cr width seed:
{
seed
}
"
)
sig
=
abs
(
np
.
random
.
randn
(
N
))
*
sig_sig
+
1
/
np
.
sqrt
(
12
)
*
1.2
# assume sigma is 1.2 times of pictch sig
return
sig
def
_flux_rand
(
self
,
N
,
seed
=-
1
):
"""
randomly generate N cosmic ray flux.
"""
seed
=
random_seed_select
(
seed
=
seed
)
log
.
debug
(
f
"cr flux seed:
{
seed
}
"
)
u
=
np
.
random
.
rand
(
N
)
S0
=
800
lam
=
0.57
S
=
(
-
np
.
log
(
1
-
u
)
/
lam
+
S0
**
0.25
)
**
4
return
S
def
random_CR_parameter
(
self
,
expt
,
pixsize
):
"""
randomly generate cosmic ray parameters, including number, length, flux, sigma and angle.
Parameters
----------
expt : float
The exposure time in second.
pixsize : list
The size of the image in pixel.
Returns
-------
CRs : list
A list of cosmic ray objects.
"""
N
=
self
.
_number_rand
(
expt
,
pixsize
)
log
.
debug
(
f
"cr count:
{
N
}
"
)
length
=
self
.
_length_rand
(
N
)
if
N
>
0
:
log
.
debug
(
f
"cr length, max:
{
length
.
max
()
}
, min:
{
length
.
min
()
}
"
)
flux
=
self
.
_flux_rand
(
N
)
log
.
debug
(
f
"cr flux, max:
{
flux
.
max
()
}
, min:
{
flux
.
min
()
}
"
)
sig
=
self
.
_sigma_rand
(
N
)
log
.
debug
(
f
"cr width, max:
{
sig
.
max
()
}
, min:
{
sig
.
min
()
}
"
)
seed
=
random_seed_select
(
seed
=-
1
)
log
.
debug
(
f
"cr angle seed:
{
seed
}
"
)
angle
=
np
.
random
.
rand
(
N
)
*
180
CRs
=
[]
for
i
in
range
(
N
):
CRs
.
append
(
CRobj
(
flux
[
i
],
angle
[
i
],
sig
[
i
],
length
[
i
]))
return
CRs
def
make_cr_frame
(
self
,
shape
,
expt
,
seed
=-
1
):
"""
make a cosmic ray frame.
Parameters
----------
shape : list
The size of the image in pixel.
expt : float
The exposure time in second.
seed : int, optional
The random seed. The default is -1. If seed is -1, the seed will be randomly selected.
Returns
-------
image : numpy.ndarray
The cosmic ray frame.
"""
image
=
np
.
zeros
(
shape
)
sz
=
shape
cr_array
=
self
.
random_CR_parameter
(
expt
,
shape
)
cr_center
=
(
self
.
tmp_size
[
1
]
-
1
)
/
2
seed
=
random_seed_select
(
seed
=
seed
)
log
.
debug
(
f
"cr position seed:
{
seed
}
"
)
for
i
in
range
(
len
(
cr_array
)):
cr
=
cr_array
[
i
]
x
=
np
.
random
.
rand
()
*
sz
[
1
]
y
=
np
.
random
.
rand
()
*
sz
[
0
]
cr_img
=
self
.
make_CR
(
cr
.
length
,
cr
.
sigma
)
cr_img
*=
cr
.
flux
cr_img
=
abs
(
nd
.
rotate
(
cr_img
,
cr
.
angle
,
reshape
=
False
))
if
i
==
0
:
pdin
=
False
else
:
pdin
=
True
if
i
==
len
(
cr_array
)
-
1
:
pdout
=
False
else
:
pdout
=
True
image
=
region_replace
(
image
,
cr_img
,
[
y
-
cr_center
,
x
-
cr_center
],
padded_in
=
pdin
,
padded_out
=
pdout
,
subpix
=
True
)
image
=
np
.
maximum
(
image
,
0
)
log
.
debug
(
f
"cr image max:
{
image
.
max
()
}
, min:
{
image
.
min
()
}
"
)
return
image
class
CpicVisEmccd
(
object
):
"""
The class for Cpic EMCCD camera.
Attributes
-----------
switch : dict
A dictionary to switch on/off the camera effects.
"""
def
__init__
(
self
,
config_dict
=
None
):
"""Initialize the camera.
Parameters
----------
config_dict : dict, optional
The configuration dictionary.
Returns
-------
None
"""
self
.
_defaut_config
()
if
config_dict
is
not
None
:
old_switch
=
self
.
switch
self
.
__dict__
.
update
(
config_dict
)
#not safe, be careful
old_switch
.
update
(
self
.
switch
)
self
.
switch
=
old_switch
self
.
config_init
()
def
_defaut_config
(
self
):
"""set up defaut config for the camera."""
self
.
plszx
=
1024
self
.
plszy
=
1024
self
.
pscan1
=
8
self
.
pscan2
=
0
self
.
oscan1
=
16
self
.
oscan2
=
18
self
.
udark
=
6
self
.
bdark
=
2
self
.
ldark
=
16
self
.
rdark
=
16
self
.
max_adu
=
16_383
self
.
switch
=
{
'bias_vp'
:
True
,
'bias_hp'
:
True
,
'bias_ci'
:
True
,
'bias_shift'
:
True
,
'cic'
:
True
,
'dark'
:
True
,
'flat'
:
True
,
'badcolumn'
:
True
,
'shutter'
:
True
,
'cte'
:
True
,
'nonlinear'
:
True
,
'cosmicray'
:
True
,
'blooming'
:
True
,
'em_blooming'
:
True
,
}
self
.
dark_file
=
cpism_refdata
+
'/camera/emccd_dark_current.fits'
self
.
flat_file
=
cpism_refdata
+
'/camera/emccd_flat_field.fits'
self
.
cic_file
=
cpism_refdata
+
'/camera/emccd_cic2.fits'
self
.
bad_col_file
=
cpism_refdata
+
'/camera/emccd_bad_columns.fits'
self
.
cic
=
None
self
.
dark
=
None
self
.
flat
=
None
self
.
fullwell
=
80_000
self
.
em_fullwell
=
500_000
#780_000
self
.
em_cte
=
0.9996
self
.
emreg_cal_num
=
10
# 用来加速计算
self
.
emreg_num
=
604
self
.
readout_speed
=
6.25e6
# Hz
self
.
readout_time
=
0.365
# s
self
.
heat_speed
=
1
/
1000
# voltage / 1000 degree per frame
self
.
temper_speed
=
0.05
# degree per second
self
.
cooler_temp
=
-
80
self
.
readout_noise
=
160
self
.
ph_per_adu
=
59
self
.
bias_level
=
200
self
.
vertical_param
=
[
0
,
14
]
self
.
horizontal1_param
=
[
5.3
/
2
,
12
,
5
/
2
,
50
]
self
.
horizontal2_param
=
[
5.3
/
4
,
16.17
,
5
/
4
,
76.65
]
self
.
bias_hp_resd
=
2
self
.
cooler_interfence
=
20
self
.
bias_level_std
=
3
self
.
bias_shift_per_volt
=
-
1.34
self
.
shift_time
=
1
/
1000
self
.
nonlinear_coefficient
=
-
0.1
self
.
detector_name
=
'EMCCD'
self
.
ccd_label
=
'CCD201-20'
self
.
pitch_size
=
13
def
config_init
(
self
):
"""initialize the camera.
If the config is set, call this function to update the config.
"""
self
.
_img_index
=
0
self
.
_ccd_temp
=
self
.
cooler_temp
self
.
_vertical_i0
=
np
.
random
.
randint
(
0
,
2
)
self
.
_frame_read_time
=
2080
*
1055
/
6.25e6
self
.
flat_shape
=
[
self
.
plszy
,
self
.
plszx
]
darksz_x
=
self
.
plszx
+
self
.
rdark
+
self
.
ldark
darksz_y
=
self
.
plszy
+
self
.
udark
+
self
.
bdark
self
.
dark_shape
=
[
darksz_y
,
darksz_x
]
biassz_x
=
darksz_x
+
self
.
pscan1
+
self
.
oscan1
biassz_y
=
darksz_y
+
self
.
pscan2
+
self
.
oscan2
self
.
bias_shape
=
[
biassz_y
,
biassz_x
]
if
self
.
flat
is
None
:
self
.
flat
=
fits
.
getdata
(
self
.
flat_file
)
if
self
.
cic
is
None
:
self
.
cic
=
fits
.
getdata
(
self
.
cic_file
)
if
self
.
dark
is
None
:
self
.
dark
=
fits
.
getdata
(
self
.
dark_file
)
self
.
bad_col
=
fits
.
getdata
(
self
.
bad_col_file
)
self
.
ccd_temp
=
self
.
cooler_temp
self
.
system_time
=
0
self
.
time_syn
(
0
,
initial
=
True
)
self
.
emgain_set
(
1024
,
-
80
)
def
CTE_cal
(
n
,
N_p
,
CTE
,
S_0
):
'''
CTE_cal(order of trail pixels, number of pixel transfers, CTE, initial intensity of target pixel)
'''
CTI
=
1
-
CTE
S_0_n
=
S_0
*
((
N_p
*
CTI
)
**
n
)
/
math
.
factorial
(
n
)
*
math
.
exp
(
-
N_p
*
CTI
)
return
S_0_n
def
cte_201
(
cte
,
start
=
0
,
length
=
10
):
N_p
=
604
S_0
=
1
res
=
[
0
]
*
start
for
n
in
range
(
length
):
s
=
CTE_cal
(
n
,
N_p
,
cte
,
S_0
)
res
.
append
(
s
)
return
np
.
array
(
res
)
cti_trail
=
cte_201
(
self
.
em_cte
,
start
=
0
,
length
=
10
)
self
.
cti_trail
=
cti_trail
/
cti_trail
.
sum
()
def
em_fix_fuc_fit
(
self
,
emgain
):
"""Calculate the emgain fix coeficient to fix the gamma distribution.
The coeficient is from fixing of ideal emgain distribution.
Parameters
----------
emgain : float
The emgain.
Returns
-------
float
The coeficient.
"""
emgain
=
np
.
array
([
emgain
]).
flatten
()
p
=
[
0.01014486
,
-
0.00712984
,
-
0.17163414
,
0.09523666
,
-
0.53926089
]
def
kernel
(
em
):
log_em
=
np
.
log10
(
em
)
loglog_g
=
np
.
log10
(
log_em
)
loglog_t
=
np
.
polyval
(
p
,
loglog_g
)
log_t
=
10
**
loglog_t
t
=
10
**
log_t
-
1
return
t
output
=
[]
for
em
in
emgain
:
if
em
<=
1
:
output
.
append
(
0
)
elif
em
>
80
:
output
.
append
(
kernel
(
80
))
else
:
output
.
append
(
kernel
(
em
))
return
np
.
array
(
output
)
def
bias_frame
(
self
):
"""Generate bias frame
The bias frame contains vertical, horizontal, peper-salt noise, bias drift effect.
Can be configurable using self.switch.
Returns
-------
np.ndarray
bias frame
"""
shape
=
self
.
bias_shape
TPI
=
np
.
pi
*
2
# vertical pattern
# 使用一维的曲线描述竖条纹的截面
vp_1d
=
np
.
zeros
(
shape
[
1
])
# 以下代码用于模拟垂直间隔的竖条纹在奇偶幅图像时的不同表现
# 后续相机更新过程中,已经将改特性进行修改,因此不再使用此代码
# vp_1d[0::2] = self.vertical_param[self._vertical_i0]
# self._vertical_i0 = 1 - self._vertical_i0
# vp_1d[1::2] = self.vertical_param[self._vertical_i0]
vp_1d
[
0
::
2
]
=
self
.
vertical_param
[
self
.
_vertical_i0
]
vp_1d
[
1
::
2
]
=
self
.
vertical_param
[
1
-
self
.
_vertical_i0
]
vp_frame
=
np
.
zeros
(
shape
)
if
self
.
switch
[
'bias_vp'
]:
vp_frame
+=
vp_1d
# if show: # pragma: no cover
# plt.figure(figsize=(10, 3))
# plt.plot(vp_1d)
# plt.xlim([0, 100])
# plt.xlabel('x-axis')
# plt.ylabel('ADU')
# plt.title('vertical pattern')
# horizontal pattern
# 图像上的横条纹是梳状,分为两个部分,左边大约77个像素是周期小一点,其余的会大一点
boundary
=
77
# boundary between left and width
boundary_width
=
5
# 左右需要平滑过度一下
y
=
np
.
arange
(
self
.
bias_shape
[
0
])
hp_left_param
=
self
.
horizontal1_param
# 实测数据拟合得到的
hp_left_1d
=
hp_left_param
[
0
]
*
np
.
sin
(
TPI
*
(
y
/
hp_left_param
[
1
]
+
np
.
random
.
rand
()))
hp_left_1d
+=
hp_left_param
[
2
]
*
np
.
sin
(
TPI
*
(
y
/
hp_left_param
[
3
]
+
np
.
random
.
rand
()))
hp_left_frame
=
np
.
broadcast_to
(
hp_left_1d
,
[
boundary
+
boundary_width
,
len
(
hp_left_1d
),]).
T
hp_right_param
=
self
.
horizontal2_param
hp_right_1d
=
hp_right_param
[
0
]
*
np
.
sin
(
TPI
*
(
y
/
hp_right_param
[
1
]
+
np
.
random
.
rand
()))
hp_right_1d
+=
hp_right_param
[
2
]
*
np
.
sin
(
TPI
*
(
y
/
hp_right_param
[
3
]
+
np
.
random
.
rand
()))
hp_right_frame
=
np
.
broadcast_to
(
hp_right_1d
,
[
shape
[
1
]
-
boundary
,
len
(
hp_right_1d
)]).
T
combine_profile_left
=
np
.
ones
(
boundary
+
boundary_width
)
combine_profile_left
[
-
boundary_width
:]
=
(
boundary_width
-
np
.
arange
(
boundary_width
)
-
1
)
/
boundary_width
combine_profile_right
=
np
.
ones
(
shape
[
1
]
-
boundary
)
combine_profile_right
[:
boundary_width
]
=
np
.
arange
(
boundary_width
)
/
boundary_width
hp_frame
=
np
.
zeros
(
shape
)
hp_frame
[:,
:
boundary
+
boundary_width
]
=
hp_left_frame
*
combine_profile_left
hp_frame
[:,
boundary
:]
=
hp_right_frame
*
combine_profile_right
# residual frame 横条纹外,还有一个垂直方向的梯度,根据测试数据,使用直线加指数函数的方式来拟合
exp_a
,
exp_b
,
lin_a
,
lin_b
=
(
-
1.92377463e+00
,
1.32698365e-01
,
8.39509583e-04
,
4.25384480e-01
)
y_trend
=
exp_a
*
np
.
exp
(
-
(
y
+
1
)
*
exp_b
)
+
y
*
lin_a
-
lin_b
# random horizontal pattern generated in frequency domain
# 除了规则横条纹外,还有随机的横条纹,随机的横条纹在频域空间生成,具有相同的频率谱和随机的相位
rsd_freq_len
=
len
(
y_trend
)
*
4
red_freq
=
np
.
arange
(
rsd_freq_len
)
red_freq
=
red_freq
-
(
len
(
red_freq
)
-
1
)
/
2
red_freq
=
np
.
exp
(
-
red_freq
**
2
/
230
**
2
)
*
240
+
np
.
random
.
randn
(
rsd_freq_len
)
*
30
red_freq
=
np
.
fft
.
fftshift
(
red_freq
)
phase
=
np
.
random
.
rand
(
rsd_freq_len
)
*
TPI
hp_rsd_1d
=
np
.
fft
.
ifft
(
np
.
exp
(
1j
*
phase
)
*
red_freq
)
hp_rsd_1d
=
np
.
abs
(
hp_rsd_1d
)
*
self
.
bias_hp_resd
hp_rsd_1d
=
hp_rsd_1d
[:
rsd_freq_len
//
4
]
hp_rsd_1d
=
hp_rsd_1d
-
np
.
mean
(
hp_rsd_1d
)
hp_rsd_1d
=
y_trend
+
hp_rsd_1d
hp_frame
=
(
hp_frame
.
T
+
hp_rsd_1d
).
T
if
not
self
.
switch
[
'bias_hp'
]:
hp_frame
*=
0
# if show: # pragma: no cover
# plt.figure(figsize=(10, 3))
# # plt.plot(hp_right_1d)
# plt.plot(hp_rsd_1d)
# # plt.xlim([0, 200])
# plt.xlabel('y-axis')
# plt.ylabel('ADU')
# plt.title('vertical pattern')
# 接上制冷机后,会有亮暗点
#cooler interfence effect
ci_position
=
10
ci_sub_struct
=
80
ci_sub_exp
=
2.5
ci_x_shft
=
3
ci_interval
=
250
# 6.25MHz readout / 2.5KHz interfence
ci_dn
=
self
.
cooler_interfence
npix
=
shape
[
0
]
*
shape
[
1
]
n_ci_event
=
npix
//
ci_interval
ci_align
=
np
.
zeros
((
n_ci_event
,
ci_interval
))
ci_align
[:,
ci_position
]
=
np
.
random
.
randn
(
n_ci_event
)
*
ci_dn
ci_align
[:,
ci_position
+
1
]
=
np
.
random
.
randn
(
n_ci_event
)
*
ci_dn
yi0
=
np
.
random
.
randint
(
0
,
ci_sub_struct
)
xs0
=
(
ci_interval
-
ci_position
)
/
(
ci_sub_struct
/
2
)
**
ci_sub_exp
for
yi
in
range
(
n_ci_event
):
sub_yi
=
(
yi
-
yi0
)
%
ci_sub_struct
sub_yi
=
abs
(
sub_yi
-
ci_sub_struct
/
2
)
shiftx
=
int
(
sub_yi
**
ci_sub_exp
*
xs0
)
ci_align
[
yi
,
:]
=
np
.
roll
(
ci_align
[
yi
,
:],
shiftx
)
ci_align
=
np
.
pad
(
ci_align
.
flatten
(),
(
0
,
npix
-
n_ci_event
*
ci_interval
))
ci_frame
=
ci_align
.
reshape
(
shape
[
0
],
shape
[
1
])
for
yi
in
range
(
shape
[
0
]):
ci_frame
[
yi
,
:]
=
np
.
roll
(
ci_frame
[
yi
,
:],
yi
*
ci_x_shft
)
if
not
self
.
switch
[
'bias_ci'
]:
ci_frame
*=
0
bias_shift
=
0
if
self
.
switch
[
'bias_shift'
]:
bias_shift
=
(
self
.
volt
-
25
)
*
self
.
bias_shift_per_volt
# 混合在一起
rn_adu
=
self
.
readout_noise
/
self
.
ph_per_adu
bias_level
=
self
.
bias_level
+
np
.
random
.
randn
()
*
self
.
bias_level_std
bias_frame
=
vp_frame
+
ci_frame
+
hp_frame
+
bias_level
bias_frame
+=
rn_adu
*
np
.
random
.
randn
(
shape
[
0
],
shape
[
1
])
+
bias_shift
return
bias_frame
def
nonlinear_effect
(
self
,
image
):
"""
nonlinear effect
"""
fullwell
=
self
.
fullwell
nonlinear_coefficient
=
self
.
nonlinear_coefficient
log
.
debug
(
f
"nonlinear effect added with coefficient
{
nonlinear_coefficient
}
"
)
image
+=
(
image
/
fullwell
)
**
2
*
nonlinear_coefficient
*
fullwell
return
image
def
time_syn
(
self
,
t
,
readout
=
True
,
initial
=
False
):
"""
time synchronization and update the system time and ccd temperature
Parameters
----------
t : float
relative time
readout : bool, optional
If the camera is readout before time synchronization, set readout to True
if True, the ccd temperature will increase, otherwise, it will decrease
initial : bool, optional
If inital is True, the ccd will be intialized to the cooler temperature
"""
if
initial
:
self
.
ccd_temp
=
self
.
cooler_temp
self
.
system_time
=
t
return
dt
=
np
.
maximum
(
t
,
0
)
heat
=
0
if
readout
:
heat
=
self
.
volt
*
self
.
heat_speed
self
.
ccd_temp
=
heat
+
self
.
cooler_temp
+
(
self
.
ccd_temp
-
self
.
cooler_temp
)
*
np
.
exp
(
-
dt
*
self
.
temper_speed
)
if
self
.
ccd_temp
<
self
.
cooler_temp
:
#
self
.
ccd_temp
=
self
.
cooler_temp
self
.
system_time
+=
dt
# def em_cte(self, img):
# i_shift = 1
# cte_coe = 0.1
# img_shift_i = np.zeros_like(img)
# img_shift_i = np.random.poisson(img * cte_coe)
# pass
def
emgain_set
(
self
,
em_set
,
ccd_temp
=
None
,
self_update
=
True
):
"""Set emgain from em set value.
Parameters
----------
em_set : int
em set value. 3FF is about 1.17×, 200 is about 1000×.
ccd_temp : float, optional
CCD temperature. If not given, use the current ccd temperature.
self_update : bool, optional
if True, update the emgain and emset. Default is True.
if False, only return the emgain.
"""
if
ccd_temp
is
None
:
ccd_temp
=
self
.
ccd_temp
volt_coe_a
=
-
0.01828
volt_coe_b
=
43.61
volt_func
=
lambda
es
:
volt_coe_a
*
es
+
volt_coe_b
self
.
volt
=
volt_func
(
em_set
)
volt_3ff
=
volt_func
(
int
(
'3ff'
,
16
))
volt_190
=
volt_func
(
int
(
'190'
,
16
))
em_coe_c
=
0.24
# using the expression of em_b = ln(g199) / constant to make fitting easier
constant
=
(
np
.
exp
(
em_coe_c
*
volt_190
)
-
np
.
exp
(
em_coe_c
*
volt_3ff
))
# fitting from the ccd test result
ln_g190
=
(
-
ccd_temp
-
7
)
*
0.0325
em_coe_b
=
ln_g190
/
constant
emgain
=
np
.
exp
(
em_coe_b
*
np
.
exp
(
em_coe_c
*
self
.
volt
))
emgain
=
np
.
maximum
(
1
,
emgain
)
# print(emgain, em_coe_b, em_coe_c * self.volt, self.volt, np.exp(em_coe_c * self.volt))
if
self_update
:
self
.
emgain
=
emgain
self
.
emset
=
em_set
return
emgain
def
vertical_blooming
(
self
,
image
):
"""
vertical blooming effect
"""
fullwell
=
self
.
fullwell
line
=
np
.
arange
(
image
.
shape
[
0
])
yp
,
xp
=
np
.
where
(
image
>
fullwell
)
n_saturated
=
len
(
xp
)
log
.
debug
(
f
"
{
len
(
xp
)
}
pixels are saturated!"
)
if
n_saturated
>
5000
:
log
.
warning
(
f
"More than 5000(
{
len
(
xp
)
}
) pixels are saturated!"
)
img0
=
image
.
copy
()
for
x
,
y
in
zip
(
xp
,
yp
):
image
[:,
x
]
+=
np
.
exp
(
-
(
line
-
y
)
**
2
/
20
**
2
)
*
img0
[
y
,
x
]
*
0.2
return
np
.
minimum
(
image
,
fullwell
)
def
emregester_blooming
(
self
,
image
,
max_iteration
=
5
):
"""
emregester blooming effect
"""
line
=
image
.
flatten
().
copy
()
curve_x
=
np
.
arange
(
1300
)
+
2
curve_y
=
np
.
exp
(
11
*
curve_x
**
(
-
0.19
)
-
11
)
curve_y
[
0
]
=
0
curve_y
/=
curve_y
.
sum
()
over_limit_coe
=
0.999
saturated
=
image
>
self
.
em_fullwell
n_saturated
=
saturated
.
sum
()
if
n_saturated
>
0
:
log
.
debug
(
f
"
{
n_saturated
}
pixels are saturated during EM process."
)
if
n_saturated
>
2000
:
log
.
warning
(
f
"More than 2000 (
{
n_saturated
}
) pixels are saturated during EM process!"
)
for
index
in
range
(
max_iteration
):
over_limit
=
np
.
maximum
(
line
-
self
.
em_fullwell
*
over_limit_coe
,
0
)
line
=
np
.
minimum
(
line
,
self
.
em_fullwell
*
over_limit_coe
)
blooming
=
np
.
convolve
(
over_limit
,
curve_y
,
mode
=
'full'
)[
:
len
(
line
)]
line
=
line
+
blooming
n_over
=
(
line
>
self
.
em_fullwell
).
sum
()
if
n_over
<=
0
:
break
log
.
debug
(
f
'
{
index
}
/
{
max_iteration
}
loop: saturated pixel number:
{
n_over
}
'
)
line
=
np
.
minimum
(
line
,
self
.
em_fullwell
)
return
line
.
reshape
(
image
.
shape
)
# def add_em_cte_conv(self, image):
# shape = image.shape
# img_line = np.convolve(image.flatten(), self.cti_trail, mode='full')
# return img_line[:shape[0]*shape[1]].reshape(shape)
def
readout
(
self
,
image_focal
,
em_set
,
expt_set
,
image_cosmic_ray
=
False
,
emgain
=
None
):
"""From focal planet image to ccd output.
Interface function for emccd. Simulate the readout process.
Parameters
----------
image_focal : np.ndarray
focal planet image. Unit: photon-electron/s/pixel
em_set : int
emgain set value.
expt_set: float
exposure time set value. Unit: s (0 mains 1ms)
image_cosmic_ray : np.ndarray, optional
cosmic ray image. Unit: photon-electron/pixel
emgain: float, optional
if not None, use the given emgain. Else, calculate the emgain using em_set
Returns
-------
np.ndarray
ccd output image. Unit: ADU
"""
expt
=
expt_set
if
expt_set
==
0
:
expt
=
0.001
dt
=
self
.
readout_time
+
expt
self
.
time_syn
(
dt
,
readout
=
True
)
if
emgain
is
None
:
self
.
emgain_set
(
em_set
,
self
.
ccd_temp
)
else
:
self
.
emgain
=
emgain
self
.
emset
=
0
emgain
=
self
.
emgain
image
=
image_focal
.
astype
(
float
)
log
.
debug
(
f
"image total photon:
{
image
.
sum
()
}
"
)
if
self
.
switch
[
'flat'
]:
image
=
image
*
self
.
flat
img_dark
=
np
.
zeros
(
self
.
dark_shape
)
img_dark
[
self
.
bdark
:
self
.
plszy
+
self
.
bdark
,
self
.
ldark
:
self
.
ldark
+
self
.
plszx
]
=
image
image
=
img_dark
if
self
.
switch
[
'dark'
]:
image
=
image
+
self
.
dark
image
*=
expt
if
image_cosmic_ray
is
not
None
and
self
.
switch
[
'cosmicray'
]:
image
+=
image_cosmic_ray
if
self
.
switch
[
'nonlinear'
]:
image
=
self
.
nonlinear_effect
(
image
)
if
self
.
switch
[
'blooming'
]:
image
=
self
.
vertical_blooming
(
image
)
if
self
.
switch
[
'badcolumn'
]:
for
i
in
range
(
self
.
bad_col
.
shape
[
1
]):
deadpix_x
=
self
.
bad_col
[
0
,
i
]
deadpix_y
=
self
.
bad_col
[
1
,
i
]
image
[
deadpix_y
:,
deadpix_x
]
=
0
img_bias
=
np
.
zeros
(
self
.
bias_shape
,
dtype
=
int
)
img_bias
[
self
.
pscan2
:
-
self
.
oscan2
,
self
.
pscan1
:
-
self
.
oscan1
]
=
np
.
random
.
poisson
(
image
)
image
=
img_bias
if
self
.
switch
[
'shutter'
]:
img_line
=
image_focal
.
sum
(
axis
=
0
)
image_shutter
=
np
.
broadcast_to
(
img_line
,
[
self
.
bias_shape
[
0
],
self
.
flat_shape
[
1
]])
image_shutter
=
image_shutter
/
self
.
flat_shape
[
1
]
*
self
.
shift_time
image_shutter
=
np
.
random
.
poisson
(
image_shutter
)
image
[:,
self
.
pscan1
+
self
.
ldark
:
-
self
.
oscan1
-
self
.
rdark
]
+=
image_shutter
if
self
.
switch
[
'cic'
]:
cic_frame
=
np
.
zeros
((
self
.
dark_shape
[
0
],
self
.
bias_shape
[
1
]))
+
self
.
cic
image
[
self
.
pscan2
:
-
self
.
oscan2
,
:]
+=
np
.
random
.
poisson
(
cic_frame
)
# em_fix found to fitting the gamma distribution to theoritical one.
# here is the theoritical distribution. See Robbins and Hadwen 2003 for more details.
# >>> pEM = np.exp(np.log(emgain)/self.emreg_num) - 1
# >>> for _ in range(self.emreg_num):
# >>> image += np.random.binomial(image, pEM)
# This code is too slow, so we used a modified gamma
em_fix
=
self
.
em_fix_fuc_fit
(
emgain
)
*
emgain
image
=
np
.
random
.
gamma
(
image
,
em_fix
)
+
image
*
(
emgain
-
em_fix
)
if
self
.
switch
[
'em_blooming'
]:
image
=
self
.
emregester_blooming
(
image
)
image
=
np
.
maximum
(
image
,
0
)
image
=
image
.
astype
(
int
)
if
self
.
switch
[
'cte'
]:
big_cte
=
self
.
em_cte
**
(
self
.
emreg_num
/
self
.
emreg_cal_num
)
for
_
in
range
(
self
.
emreg_cal_num
):
residual
=
np
.
random
.
binomial
(
image
,
1
-
big_cte
)
image
[:,
1
:]
+=
residual
[:,
:
-
1
]
-
residual
[:,
1
:]
bias
=
self
.
bias_frame
()
image
=
image
/
self
.
ph_per_adu
+
bias
image
=
np
.
minimum
(
image
,
self
.
max_adu
)
image
=
np
.
maximum
(
image
,
0
)
log
.
debug
(
f
"emccd paramters:
\
emset:
{
em_set
}
\
emgain:
{
emgain
}
\
expt:
{
expt
}
\
readout time:
{
dt
}
"
)
return
image
.
astype
(
dtype
=
np
.
uint16
)
# if __name__ == '__main__':
# import matplotlib.pyplot as plt
# emccd = EMCCD()
# image_focal = np.zeros((emccd.plszy, emccd.plszx)) + 1000
# image_focal[100:105, 100:105] = 10_000_000
# after_cte = emccd.emregester_blooming(image_focal, max_iteration=100)
# print(after_cte.sum(), image_focal.sum())
# fits.writeto('after_cte.fits', after_cte, overwrite=True)
# # darksz_x = emccd.plszx + emccd.rdark + emccd.ldark
# # darksz_y = emccd.plszy + emccd.udark + emccd.bdark
# # iamge_cosmic_ray = np.zeros((darksz_y, darksz_x))
# # emgain = 10
# # expt = 10
# # image = emccd.readout(image_focal, emgain, expt, iamge_cosmic_ray)
# # fits.writeto('test.fits', image, overwrite=True)
# image = np.zeros((1000, 1000))
# make_cosmic_ray_frame = CosmicRayFrameMaker()
# crimage = make_cosmic_ray_frame(image.shape, 3000)
# fits.writeto('crimage.fits', crimage, overwrite=True)
# # emccd.add_stripe_effect(image)
source/csst_cpic_sim/config.py
0 → 100644
View file @
7c99b572
import
os
,
yaml
import
warnings
from
datetime
import
datetime
import
numpy
as
np
config_aim
=
os
.
path
.
dirname
(
os
.
path
.
dirname
(
__file__
))
config_aim
=
os
.
path
.
join
(
config_aim
,
'data/refdata_path.yaml'
)
# def set_config(refdata_path=None):
# if refdata_path is None:
# print("input cpism refencence data folder")
# refdata_path = input()
# refdata_path = os.path.abspath(refdata_path)
# with open(config_aim, 'w') as f:
# yaml.dump(refdata_path, f)
# return refdata_path
# try:
# with open(config_aim, 'r') as f:
# cpism_refdata = yaml.load(f, Loader=yaml.FullLoader)
# if not os.path.isdir(cpism_refdata):
# raise FileNotFoundError('cpism refdata path not found')
# config_set = True
# except FileNotFoundError:
# warnings.warn(f'refdata not setup yet, set it before use')
# cpism_refdata = set_config()
def
load_refdata_path
(
config_aim
):
"""Load refdata path.
The refdata path is stored in config_aim file. If not found, set it.
Parameters
----------
config_aim : str
config_aim file path
"""
with
open
(
config_aim
,
'r'
)
as
f
:
refdata_list
=
yaml
.
load
(
f
,
Loader
=
yaml
.
FullLoader
)
for
refdata
in
refdata_list
:
if
os
.
path
.
isdir
(
refdata
):
return
refdata
print
(
"csst_cpic_sim refdata folder not found, please input cpism refencence data folder"
)
refdata
=
input
()
refdata
=
os
.
path
.
abspath
(
refdata
)
if
os
.
path
.
isdir
(
refdata
):
refdata_list
.
append
(
refdata
)
with
open
(
config_aim
,
'w'
)
as
f
:
yaml
.
dump
(
refdata_list
,
f
)
exit
()
cpism_refdata
=
load_refdata_path
(
config_aim
)
config
=
{}
config
[
'cpism_refdata'
]
=
cpism_refdata
config
[
'utc0'
]
=
'2024-05-01T00:00:00'
config
[
'hybrid_model'
]
=
f
'
{
cpism_refdata
}
/target_model/hybrid_model.fits'
config
[
'bcc_model'
]
=
f
'
{
cpism_refdata
}
/target_model/bccmodels'
config
[
'mag_system'
]
=
'abmag'
config
[
'apm_file'
]
=
f
'
{
cpism_refdata
}
/optics/apm.fits'
config
[
'actuator_file'
]
=
f
'
{
cpism_refdata
}
/optics/actuator.fits'
config
[
'aberration'
]
=
f
'
{
cpism_refdata
}
/optics/initial_phase_aberration.fits'
config
[
'mask_width'
]
=
0.4
config
[
'check_fits_header'
]
=
False
config
[
'bands'
]
=
{
'f661'
:
f
'
{
cpism_refdata
}
/throughtput/f661_total.fits'
,
'f743'
:
f
'
{
cpism_refdata
}
/throughtput/f743_total.fits'
,
'f883'
:
f
'
{
cpism_refdata
}
/throughtput/f883_total.fits'
,
'f565'
:
f
'
{
cpism_refdata
}
/throughtput/f565_total.fits'
,
'f520'
:
f
'
{
cpism_refdata
}
/throughtput/f520.fits'
,
'f662'
:
f
'
{
cpism_refdata
}
/throughtput/f662.fits'
,
'f850'
:
f
'
{
cpism_refdata
}
/throughtput/f850.fits'
,
'f720'
:
f
'
{
cpism_refdata
}
/throughtput/f720.fits'
,
}
config
[
'diameter'
]
=
2
# in meters
config
[
'platescale'
]
=
0.016153
config
[
'datamodel'
]
=
f
'
{
cpism_refdata
}
/io/csst-cpic-l0.yaml'
config
[
'log_dir'
]
=
f
'
{
cpism_refdata
}
/log'
config
[
'log_level'
]
=
f
'info'
config
[
'output'
]
=
f
'./'
config
[
'sp2teff_model'
]
=
f
'
{
cpism_refdata
}
/target_model/sptype2teff_lut.json'
config
[
'dm_pickle'
]
=
f
'
{
cpism_refdata
}
/optics/dm_model.pkl'
config
[
'pysyn_refdata'
]
=
f
'
{
cpism_refdata
}
/starmodel/grp/redcat/trds'
config
[
'catalog_folder'
]
=
f
'
{
cpism_refdata
}
/demo_catalog'
config
[
'csst_format'
]
=
True
config
[
'nsample'
]
=
5
update_able_keys
=
[
'apm_file'
,
'actuator_file'
,
'aberration'
,
'log_dir'
,
'log_level'
,
'catalog_folder'
,
'nsample'
,
'csst_format'
,
'output'
,
'check_fits_header'
]
def
replace_cpism_refdata
(
config
:
dict
,
output
:
str
=
'$'
)
->
None
:
"""Replace the cpism_refdata in the config.
In the config file, we use ${cpism_refdata} to indicate the cpism_refdata.
This function is used to replace the cpism_refdata in the config, or replace back.
Parameters
----------
config: dict
config dict.
output: str
'$' or 'other'. If output is '$', then replace the cpism_refdata in the config with ${cpism_refdata}.
If output is 'other', then replace the ${cpism_refdata} in the config file with the real path.
'$' is used meanly to generate a demo config file.
"""
aim
=
cpism_refdata
target
=
'${cpism_refdata}'
if
output
!=
'$'
:
aim
,
target
=
target
,
aim
for
key
,
value
in
config
.
items
():
if
isinstance
(
value
,
str
):
config
[
key
]
=
value
.
replace
(
aim
,
target
)
if
isinstance
(
value
,
dict
):
replace_cpism_refdata
(
value
,
output
)
with
open
(
cpism_refdata
+
'/cpism_config.yaml'
,
'r'
)
as
f
:
new_config
=
yaml
.
load
(
f
,
Loader
=
yaml
.
FullLoader
)
replace_cpism_refdata
(
new_config
,
None
)
config
.
update
(
new_config
)
if
os
.
environ
.
get
(
'PYSYN_CDBS'
)
is
None
:
os
.
environ
[
'PYSYN_CDBS'
]
=
config
[
'pysyn_refdata'
]
__version__
=
'2.0.0'
# we need to ignore the warning from pysynphot, because we only use the star models.
with
warnings
.
catch_warnings
():
# pragma: no cover
warnings
.
filterwarnings
(
"ignore"
)
import
pysynphot
as
S
def
setup_config
(
new_config
):
"""Set up config from a dict.
Some of the configs need to calcuate.
Parameters
----------
new_config: dict
new config dict.
Returns
-------
None
"""
config
.
update
(
new_config
)
config
[
'utc0_float'
]
=
datetime
.
timestamp
(
datetime
.
fromisoformat
(
config
[
'utc0'
]))
config
[
'solar_spectrum'
]
=
f
"
{
os
.
environ
[
'PYSYN_CDBS'
]
}
/grid/solsys/solar_spec.fits"
config
[
'aperature_area'
]
=
(
config
[
'diameter'
]
*
50
)
**
2
*
np
.
pi
# cm^2
config
[
'default_band'
]
=
list
(
config
[
'bands'
].
keys
())[
0
]
config
[
'default_filter'
]
=
config
[
'bands'
][
config
[
'default_band'
]]
setup_config
({})
def
which_focalplane
(
band
):
"""
Return the name of the focalplane which the band belongs to.
right now only support vis
Parameters
-----------
band: str
The name of the band.
Returns
--------
str
The name of the focalplane.
'vis' or 'nir' or 'wfs'
Raises
-------
ValueError
If the band is not in ['f565', 'f661', 'f743', 'f883', 'f940', 'f1265', 'f1425', 'f1542', 'wfs']
"""
# band = band.lower()
# if band in ['f565', 'f661', 'f743', 'f883']:
# return 'vis'
# if band in ['f940', 'f1265', 'f1425', 'f1542']:
# return 'nir'
# if band in ['wfs']:
# return 'wfs'
return
'vis'
# raise ValueError(f"未知的波段{band}")
def
iso_time
(
time
):
"""Transfer relative time to iso time format
Parameters
----------
time: str or float
The relative time in seconds.
Returns
-------
str
The iso time format.
"""
if
isinstance
(
time
,
str
):
_
=
datetime
.
fromisoformat
(
time
)
return
time
utc0
=
config
[
'utc0'
]
time0
=
datetime
.
timestamp
(
datetime
.
fromisoformat
(
utc0
))
time
=
datetime
.
fromtimestamp
(
time0
+
time
)
return
time
.
isoformat
()
def
relative_time
(
time
):
"""Transfer iso time format to relative time in seconds
Parameters
----------
time: str or float
The iso time format.
Returns
-------
float
The relative time in seconds.
"""
if
isinstance
(
time
,
float
):
return
time
if
isinstance
(
time
,
int
):
return
float
(
time
)
utc0
=
config
[
'utc0'
]
time0
=
datetime
.
timestamp
(
datetime
.
fromisoformat
(
utc0
))
return
datetime
.
timestamp
(
datetime
.
fromisoformat
(
time
))
-
time0
\ No newline at end of file
source/csst_cpic_sim/io.py
0 → 100644
View file @
7c99b572
import
yaml
,
os
,
re
from
datetime
import
datetime
import
numpy
as
np
import
pandas
as
pd
from
astropy.io
import
fits
from
astropy.coordinates
import
SkyCoord
from
.config
import
__version__
,
which_focalplane
from
.utils
import
Logger
from
.config
import
config
,
iso_time
default_output_dir
=
config
[
'output'
]
log_level
=
config
[
'log_level'
]
header_check
=
config
[
'check_fits_header'
]
def
set_up_logger
(
log_dir
):
if
not
os
.
path
.
exists
(
log_dir
):
os
.
makedirs
(
log_dir
)
return
Logger
(
log_dir
+
'/cpism_pack.log'
,
log_level
).
logger
log
=
set_up_logger
(
config
[
'log_dir'
])
def
check_and_update_fits_header
(
header
):
"""
Check the header keywords and update the description according to the data model.
Parameters
-----------
header: astropy.io.fits.header.Header
The header to be checked.
Returns
--------
None
"""
# model_file = f"{cpism_refdata}/io/image_header.json"
# if hdu == 'primary':
# model_file = f"{cpism_refdata}/io/primary_header.json"
model_file
=
config
[
'datamodel'
]
with
open
(
model_file
,
'r'
,
encoding
=
'utf-8'
)
as
fid
:
data_model
=
yaml
.
load
(
fid
,
Loader
=
yaml
.
FullLoader
)
if
'FILETYPE'
in
header
.
keys
():
header_model
=
data_model
[
'HDU0'
]
hdu
=
'hdu0'
else
:
header_model
=
data_model
[
'HDU1'
]
hdu
=
'hdu1'
dm_comment
=
{}
def
print_warning
(
info
):
if
header_check
:
log
.
warning
(
info
)
# check existance and format of keyword in fits header
for
_
,
content
in
header_model
.
items
():
comment
=
content
[
'comment'
]
keyword
=
content
[
'key'
]
dtype
=
content
[
'dtype'
]
if
pd
.
isnull
(
comment
):
comment
=
''
if
len
(
comment
)
>
47
:
# comment = comment[:46]
print_warning
(
f
"Keyword
{
keyword
}
has a comment longer than 47 characters."
)
dm_comment
[
keyword
]
=
comment
if
keyword
not
in
header
.
keys
():
print_warning
(
f
"Keyword
{
keyword
}
not found in header."
)
elif
not
pd
.
isnull
(
dtype
):
value
=
header
[
keyword
]
# check the type of the value, I for int, R for float, C for str
if
isinstance
(
value
,
str
):
key_type
=
'str'
elif
isinstance
(
value
,
float
):
key_type
=
'flo'
elif
isinstance
(
value
,
bool
):
key_type
=
'boo'
elif
isinstance
(
value
,
int
):
key_type
=
'int'
elif
isinstance
(
value
,
type
(
header
[
'COMMENT'
])):
key_type
=
'str'
else
:
key_type
=
'ukn'
if
key_type
!=
dtype
[
0
:
3
]:
print_warning
(
f
"Keyword
{
keyword
}
has wrong type in [
{
hdu
}
].
{
dtype
}
expected,
{
key_type
}
found."
)
# check if there are extral keyword in header, and update the comment
for
keyword
in
header
.
keys
():
# print(keyword)
if
keyword
not
in
dm_comment
.
keys
():
print_warning
(
f
"Keyword
{
keyword
}
not found in the [
{
hdu
}
] data model."
)
elif
keyword
!=
'COMMENT'
:
# comment keyword is not allowed to be updated
header
[
keyword
]
=
(
header
[
keyword
],
dm_comment
[
keyword
])
return
header
def
obsid_parser
(
obsid
:
int
):
"""
Parse the obsid to get the obstype.
Parameters
----------
obsid: str
The obsid of the observation.
Obsid must be 11 digits and start with 5 for CPIC.
Returns
-------
str
The obstype of the observation.
Raises
------
ValueError
If the obsid is not 11 digits or does not start with 5.
"""
obsid
=
str
(
obsid
)
if
len
(
obsid
)
!=
11
:
raise
ValueError
(
'Obsid must be 11 digits.'
)
if
obsid
[
0
]
!=
'4'
:
raise
ValueError
(
'Obsid must start with 4 for CPIC'
)
obstype_dict
=
{
'20'
:
'BIAS'
,
'21'
:
'DARK'
,
'22'
:
'FLAT'
,
'23'
:
'BKG'
,
'24'
:
'LASER'
,
'01'
:
'SCI'
,
'02'
:
'DSF'
,
'10'
:
'CALS'
,
}
obstype
=
obstype_dict
.
get
(
obsid
[
1
:
3
],
'DEFT'
)
return
obstype
def
datetime_obj_to_mjd
(
time_obj
):
"""
transfer datetime object to mean julian date (MJD).
Parameters
----------
time_obj: datetime.datetime
The datetime object.
Returns
-------
float
The mean julian date (MJD).
"""
return
(
time_obj
-
datetime
(
1858
,
11
,
17
)).
total_seconds
()
/
86400
def
primary_hdu
(
obs_info
:
dict
,
gnc_info
:
dict
,
filename_output
=
False
):
"""
Generate the primary hdu of the fits file.
Parameters
----------
obs_info: dict
The parameters of the observation. See `save_fits` function.
gnc_info: dict
The gnc information of the observation.
filename_output: bool (optional)
If True, the folder and the filename will be returned.
Returns
-------
fits.PrimaryHDU
The primary hdu of the fits file.
str, str
The folder and filename of the fits file. Only returned if filename_output is True.
Notes
-----
The gnc_info dict should contain the information of orbit and observation.
these informations are used to genrated the hdu header. Refer to the data model for more information.
"""
# camera_config, _ = load_camera_and_optics_config(obs_info['band'])
obsid
=
obs_info
[
'obsid'
]
exp_start
=
obs_info
[
'EXPSTART'
]
exp_start
=
datetime
.
fromisoformat
(
exp_start
)
exp_end
=
obs_info
[
'EXPEND'
]
exp_end
=
datetime
.
fromisoformat
(
exp_end
)
filename
=
"CSST_CPIC"
filename
+=
"_"
+
which_focalplane
(
obs_info
[
'band'
]).
upper
()
filename
+=
"_"
+
obsid_parser
(
obsid
)
filename
+=
"_"
+
exp_start
.
strftime
(
"%Y%m%d%H%M%S"
)
filename
+=
"_"
+
exp_end
.
strftime
(
"%Y%m%d%H%M%S"
)
filename
+=
f
"_
{
obsid
}
_X_L0_V01.fits"
type_dir
=
'CAL'
if
int
(
f
'
{
obsid
}
'
[
1
:
3
])
<=
10
:
type_dir
=
'SCI'
mjd_dir
=
f
"
{
datetime_obj_to_mjd
(
exp_start
)
:
.
0
f
}
"
folder
=
f
"
{
type_dir
}
/
{
mjd_dir
}
"
header
=
fits
.
Header
()
# General keywords
header
[
'SIMPLE'
]
=
True
header
[
'BITPIX'
]
=
8
header
[
'NAXIS'
]
=
0
header
[
'EXTEND'
]
=
True
header
[
'NEXTEND'
]
=
1
# + parameters['nframe']
# header['GROUPS'] = False
header
[
'DATE'
]
=
datetime
.
now
().
isoformat
(
timespec
=
'seconds'
)
heaer_filename
=
filename
[:
-
4
]
if
len
(
heaer_filename
)
>
68
:
heaer_filename
=
heaer_filename
[:
68
]
header
[
'FILENAME'
]
=
heaer_filename
header
[
'FILETYPE'
]
=
obsid_parser
(
obsid
)
header
[
'TELESCOP'
]
=
'CSST'
header
[
'INSTRUME'
]
=
'CPIC'
header
[
'RADECSYS'
]
=
'ICRS'
header
[
'EQUINOX'
]
=
2000.0
header
[
'FITSSWV'
]
=
f
'CPISM V
{
__version__
}
'
header
[
'COMMENT'
]
=
''
cstar
=
{
'ra'
:
'0d'
,
'dec'
:
'0d'
}
if
obs_info
[
'target'
]
!=
{}:
cstar
=
obs_info
[
'target'
][
'cstar'
]
radec
=
SkyCoord
(
cstar
[
'ra'
],
cstar
[
'dec'
])
target_name
=
radec
.
to_string
(
'hmsdms'
)
target_name
=
re
.
sub
(
R
'[hdms\s]'
,
''
,
target_name
)
header
[
'OBJECT'
]
=
cstar
.
get
(
'name'
,
target_name
)
header
[
'TARGET'
]
=
target_name
header
[
'OBSID'
]
=
str
(
obsid
)
header
[
'RA_OBJ'
]
=
radec
.
ra
.
degree
header
[
'DEC_OBJ'
]
=
radec
.
dec
.
degree
# telescope information
header
[
'REFFRAME'
]
=
'CSSTGSC-1.0'
header
[
'DATE-OBS'
]
=
exp_start
.
isoformat
(
timespec
=
'seconds'
)
header
[
'SATESWV'
]
=
'1'
header
[
'EXPSTART'
]
=
datetime_obj_to_mjd
(
exp_start
)
cabstart
=
gnc_info
.
get
(
'CABSTART'
,
exp_start
.
isoformat
(
timespec
=
'seconds'
))
cabstart
=
iso_time
(
cabstart
)
cabstart_mjd
=
datetime_obj_to_mjd
(
datetime
.
fromisoformat
(
cabstart
))
header
[
'CABSTART'
]
=
cabstart_mjd
header
[
'SUNANGL0'
]
=
gnc_info
.
get
(
'SUNANGL0'
,
-
1.0
)
header
[
'MOONANG0'
]
=
gnc_info
.
get
(
'MOONANG0'
,
-
1.0
)
header
[
'TEL_ALT0'
]
=
gnc_info
.
get
(
'TEL_ALT0'
,
-
1.0
)
header
[
'POS_ANG0'
]
=
gnc_info
.
get
(
'POS_ANG0'
,
float
(
obs_info
[
'rotation'
]))
header
[
'POSI0_X'
]
=
gnc_info
.
get
(
'POSI0_X'
,
-
1.0
)
header
[
'POSI0_Y'
]
=
gnc_info
.
get
(
'POSI0_Y'
,
-
1.0
)
header
[
'POSI0_Z'
]
=
gnc_info
.
get
(
'POSI0_Z'
,
-
1.0
)
header
[
'VELO0_X'
]
=
gnc_info
.
get
(
'VELO0_X'
,
-
1.0
)
header
[
'VELO0_Y'
]
=
gnc_info
.
get
(
'VELO0_Y'
,
-
1.0
)
header
[
'VELO0_Z'
]
=
gnc_info
.
get
(
'VELO0_Z'
,
-
1.0
)
header
[
'EULER0_1'
]
=
gnc_info
.
get
(
'EULER0_1'
,
-
1.0
)
header
[
'EULER0_2'
]
=
gnc_info
.
get
(
'EULER0_2'
,
-
1.0
)
header
[
'EULER0_3'
]
=
gnc_info
.
get
(
'EULER0_3'
,
-
1.0
)
header
[
'RA_PNT0'
]
=
gnc_info
.
get
(
'RA_PNT0'
,
header
[
'RA_OBJ'
])
header
[
'DEC_PNT0'
]
=
gnc_info
.
get
(
'DEC_PNT0'
,
header
[
'DEC_OBJ'
])
header
[
'EXPEND'
]
=
datetime_obj_to_mjd
(
exp_end
)
cabend
=
gnc_info
.
get
(
'CABEND'
,
exp_end
.
isoformat
(
timespec
=
'seconds'
))
cabend
=
iso_time
(
cabend
)
cabend_mjd
=
datetime_obj_to_mjd
(
datetime
.
fromisoformat
(
cabend
))
header
[
'CABEND'
]
=
cabend_mjd
header
[
'SUNANGL1'
]
=
gnc_info
.
get
(
'SUNANGL1'
,
header
[
'SUNANGL0'
])
header
[
'MOONANG1'
]
=
gnc_info
.
get
(
'MOONANG1'
,
header
[
'MOONANG0'
])
header
[
'TEL_ALT1'
]
=
gnc_info
.
get
(
'TEL_ALT1'
,
header
[
'TEL_ALT0'
])
header
[
'POS_ANG1'
]
=
gnc_info
.
get
(
'POS_ANG1'
,
header
[
'POS_ANG0'
])
header
[
'POSI1_X'
]
=
gnc_info
.
get
(
'POSI1_X'
,
header
[
'POSI0_x'
])
header
[
'POSI1_Y'
]
=
gnc_info
.
get
(
'POSI1_Y'
,
header
[
'POSI0_y'
])
header
[
'POSI1_Z'
]
=
gnc_info
.
get
(
'POSI1_Z'
,
header
[
'POSI0_z'
])
header
[
'VELO1_X'
]
=
gnc_info
.
get
(
'VELO1_X'
,
header
[
'VELO0_x'
])
header
[
'VELO1_Y'
]
=
gnc_info
.
get
(
'VELO1_Y'
,
header
[
'VELO0_y'
])
header
[
'VELO1_Z'
]
=
gnc_info
.
get
(
'VELO1_Z'
,
header
[
'VELO0_z'
])
header
[
'EULER1_1'
]
=
gnc_info
.
get
(
'EULER1_1'
,
header
[
'EULER0_1'
])
header
[
'EULER1_2'
]
=
gnc_info
.
get
(
'EULER1_2'
,
header
[
'EULER0_2'
])
header
[
'EULER1_3'
]
=
gnc_info
.
get
(
'EULER1_3'
,
header
[
'EULER0_3'
])
header
[
'RA_PNT1'
]
=
gnc_info
.
get
(
'RA_PNT1'
,
header
[
'RA_PNT0'
])
header
[
'DEC_PNT1'
]
=
gnc_info
.
get
(
'DEC_PNT1'
,
header
[
'DEC_PNT0'
])
header
[
'EPOCH'
]
=
float
(
exp_start
.
year
)
header
[
'EXPTIME'
]
=
(
exp_end
-
exp_start
).
total_seconds
()
header
[
'CHECKSUM'
]
=
'0000000000000000'
header
[
'DATASUM'
]
=
'0000000000'
check_and_update_fits_header
(
header
)
# other information
hdu
=
fits
.
PrimaryHDU
(
header
=
header
)
hdu
.
add_checksum
()
if
filename_output
:
return
hdu
,
folder
,
filename
else
:
return
hdu
def
frame_header
(
obs_info
,
index
,
primary_header
,
camera_dict
):
"""
Generate the header for a single frame.
Parameters
----------
obs_info : dict
Dictionary of parameters. See `save_fits` function.
index : int
Frame index.
primary_header : dict
Primary header
Returns
-------
astropy.io.fits.Header
"""
header
=
fits
.
Header
()
# camera_config, optics_config = load_camera_and_optics_config(
# obs_info['band'])
camera_config
=
camera_dict
plszx
=
camera_config
[
'plszx'
]
plszy
=
camera_config
[
'plszy'
]
pscan1
=
camera_config
[
'pscan1'
]
pscan2
=
camera_config
[
'pscan2'
]
oscan1
=
camera_config
[
'oscan1'
]
oscan2
=
camera_config
[
'oscan2'
]
udark
=
camera_config
[
'udark'
]
bdark
=
camera_config
[
'bdark'
]
ldark
=
camera_config
[
'ldark'
]
rdark
=
camera_config
[
'rdark'
]
imgszx
=
plszx
+
pscan1
+
oscan1
+
ldark
+
rdark
imgszy
=
plszy
+
pscan2
+
oscan2
+
udark
+
bdark
header
[
'XTENSION'
]
=
'IMAGE'
header
[
'BITPIX'
]
=
16
header
[
'NAXIS'
]
=
2
header
[
'NAXIS1'
]
=
1080
header
[
'NAXIS2'
]
=
1050
header
[
'PCOUNT'
]
=
0
header
[
'GCOUNT'
]
=
1
header
[
'BSCALE'
]
=
1
header
[
'BZERO'
]
=
32768
header
[
'EXTNAME'
]
=
'IMAGE'
header
[
'EXTVER'
]
=
1
header
[
'BUNIT'
]
=
'ADU'
header
[
'FILTER'
]
=
obs_info
[
'band'
]
header
[
'DETSN'
]
=
'0'
header
[
'DETNAME'
]
=
camera_config
[
'detector_name'
]
header
[
'CHIPLAB'
]
=
camera_config
[
'ccd_label'
]
header
[
'DEWTEMP'
]
=
float
(
camera_config
[
'cooler_temp'
])
frame_info
=
obs_info
[
'frame_info'
][
index
]
header
[
'CHIPTEMP'
]
=
float
(
frame_info
[
'chiptemp'
])
header
[
'DETSIZE'
]
=
f
"
{
imgszx
}
*
{
imgszy
}
"
header
[
'IMGINDEX'
]
=
index
+
1
utc0_float
=
config
[
'utc0_float'
]
ra0
=
primary_header
.
get
(
'RA_PNT0'
,
-
1.0
)
dec0
=
primary_header
.
get
(
'DEC_PNT0'
,
-
1.0
)
pa0
=
primary_header
.
get
(
'POS_ANG0'
,
-
1.0
)
cab0
=
primary_header
.
get
(
'CABSTART'
,
obs_info
[
'frame_info'
][
0
][
'expt_start'
])
ra1
=
primary_header
.
get
(
'RA_PNT1'
,
ra0
)
dec1
=
primary_header
.
get
(
'DEC_PNT1'
,
dec0
)
pa1
=
primary_header
.
get
(
'POS_ANG1'
,
pa0
)
cab1
=
primary_header
.
get
(
'CABEND'
,
cab0
)
frame_stamp
=
frame_info
[
'expt_start'
]
+
utc0_float
frame_mjd
=
datetime_obj_to_mjd
(
datetime
.
fromtimestamp
(
frame_stamp
))
cab0
=
iso_time
(
cab0
)
cab1
=
iso_time
(
cab1
)
cab0
=
datetime_obj_to_mjd
(
datetime
.
fromisoformat
(
cab0
))
cab1
=
datetime_obj_to_mjd
(
datetime
.
fromisoformat
(
cab1
))
ratio
=
(
frame_mjd
-
cab0
)
/
(
cab1
-
cab0
)
ra
=
ra0
+
(
ra1
-
ra0
)
*
ratio
dec
=
dec0
+
(
dec1
-
dec0
)
*
ratio
pa
=
pa0
+
(
pa1
-
pa0
)
*
ratio
header
[
'IMG_EXPT'
]
=
datetime
.
fromtimestamp
(
frame_stamp
).
isoformat
(
timespec
=
'seconds'
)
header
[
'IMG_CABT'
]
=
header
[
'IMG_EXPT'
]
header
[
'IMG_DUR'
]
=
float
(
obs_info
[
'expt'
])
header
[
'IMG_PA'
]
=
ra
header
[
'IMG_RA'
]
=
dec
header
[
'IMG_DEC'
]
=
pa
header
[
'DATASECT'
]
=
f
"
{
plszx
}
*
{
plszy
}
"
header
[
'PIXSCAL'
]
=
frame_info
[
'platescale'
]
header
[
'PIXSIZE'
]
=
float
(
camera_config
[
'pitch_size'
])
header
[
'NCHAN'
]
=
1
header
[
'PSCAN1'
]
=
pscan1
header
[
'PSCAN2'
]
=
pscan2
header
[
'OSCAN1'
]
=
oscan1
header
[
'OSCAN2'
]
=
oscan2
header
[
'UDARK'
]
=
udark
header
[
'BDARK'
]
=
bdark
header
[
'LDARK'
]
=
ldark
header
[
'RDARK'
]
=
rdark
# WCS
cstar
=
{
'ra'
:
'0d'
,
'dec'
:
'0d'
}
if
obs_info
[
'target'
]
!=
{}:
cstar
=
obs_info
[
'target'
][
'cstar'
]
radec
=
SkyCoord
(
cstar
[
'ra'
],
cstar
[
'dec'
])
shift
=
obs_info
[
'shift'
]
platescale
=
frame_info
[
'platescale'
]
rotation
=
np
.
radians
(
obs_info
[
'rotation'
])
header
[
'WCSAXES'
]
=
2
header
[
'CRPIX1'
]
=
(
plszx
+
1
)
/
2
+
pscan1
+
ldark
+
shift
[
0
]
/
platescale
header
[
'CRPIX2'
]
=
(
plszy
+
1
)
/
2
+
pscan2
+
udark
+
shift
[
0
]
/
platescale
header
[
'CRVAL1'
]
=
radec
.
ra
.
degree
header
[
'CRVAL2'
]
=
radec
.
dec
.
degree
header
[
'CTYPE1'
]
=
'RA---TAN'
header
[
'CTYPE2'
]
=
'DEC--TAN'
header
[
'CD1_1'
]
=
np
.
cos
(
rotation
)
header
[
'CD1_2'
]
=
-
np
.
sin
(
rotation
)
header
[
'CD2_1'
]
=
np
.
sin
(
rotation
)
header
[
'CD2_2'
]
=
np
.
cos
(
rotation
)
# Readout information
header
[
'EMGAIN'
]
=
float
(
obs_info
[
'emgain'
])
header
[
'GAIN'
]
=
float
(
camera_config
[
'ph_per_adu'
])
header
[
'DET_BIAS'
]
=
float
(
camera_config
[
'bias_level'
])
header
[
'RON'
]
=
float
(
camera_config
[
'readout_noise'
])
header
[
'READTIME'
]
=
float
(
camera_config
[
'readout_time'
])
header
[
'ROSPEED'
]
=
float
(
camera_config
[
'readout_speed'
])
# CPIC information
header
[
'LS_STAT'
]
=
'OFF'
header
[
'IWA'
]
=
frame_info
[
'iwa'
]
header
[
'WFSINFO1'
]
=
-
1.0
header
[
'WFSINFO2'
]
=
-
1.0
header
[
'CHECKSUM'
]
=
'0000000000000000'
header
[
'DATASUM'
]
=
'0000000000'
header
=
check_and_update_fits_header
(
header
)
return
header
def
save_fits_simple
(
images
,
obs_info
,
output_folder
=
'./'
):
"""
Save the image to a fits file with a simple header to TMP directory.
Parameters
----------
images : numpy.ndarray
Image array to be written.
obs_info : dict
Dictionary of observation informations. See `save_fits` function.
Returns
----------
Filename of the saved fits file.
"""
target
=
obs_info
[
'target'
]
target_info
=
'NO_TARGET'
if
'cstar'
in
target
.
keys
():
target_info
=
''
target_info
=
f
"S
{
target
[
'cstar'
][
'magnitude'
]
:
.
1
f
}
"
target_info
+=
f
"_P
{
len
(
target
.
get
(
'planets'
,
'[]'
))
}
"
name
=
target_info
if
'name'
in
target
.
keys
():
name
=
target
[
'name'
]
name
=
name
.
replace
(
'/'
,
'_'
)
name
=
name
.
replace
(
','
,
'_'
)
now
=
datetime
.
now
()
time
=
now
.
strftime
(
"%Y%m%d%H%M%S"
)
filename
=
f
"
{
name
}
_
{
time
}
.fits"
header
=
fits
.
Header
()
header
[
'skybg'
]
=
obs_info
[
'skybg'
]
header
[
'name'
]
=
name
header
[
'exptime'
]
=
obs_info
[
'expt'
]
header
[
'nframe'
]
=
obs_info
[
'nframe'
]
header
[
'band'
]
=
obs_info
[
'band'
]
header
[
'emgain'
]
=
obs_info
[
'emgain'
]
header
[
'obsid'
]
=
obs_info
[
'obsid'
]
header
[
'rotation'
]
=
obs_info
[
'rotation'
]
shift
=
obs_info
[
'shift'
]
header
[
'shift'
]
=
f
"x:
{
shift
[
0
]
}
,y:
{
shift
[
1
]
}
"
fullname
=
os
.
path
.
join
(
output_folder
,
filename
)
print
(
fullname
)
if
not
os
.
path
.
exists
(
output_folder
):
os
.
makedirs
(
output_folder
)
log
.
debug
(
f
"Output folder
{
output_folder
}
is created."
)
log
.
debug
(
f
"save fits file to
{
fullname
}
"
)
fits
.
writeto
(
fullname
,
images
,
overwrite
=
True
,
header
=
header
)
return
fullname
def
save_fits
(
images
,
obs_info
,
gnc_info
,
camera_dict
=
{},
csst_format
=
True
,
output_folder
=
'./'
):
"""
Save the image to a fits file.
Parameters
----------
images : numpy.ndarray
Image array to be saved.
obs_info : dict
Dictionary of observation informations.
Must contain the following keys
- band: str
- Band of the image.
- expt: float
- Exposure time of the each image.
- nframe: int
- Number of frames in the image.
- emgain: int
- EM gain of the camera.
- obsid: str
- Observation ID. Obsid must be 11 digits and start with 5 for CPIC. See pharse_obsid() for details.
- rotation: float
- Rotation angle of the image.
- shift: list
- Shift of the image.
gnc_info : dict
Dictionary of GNC information.
Contains the keywords in the primary header. See primary_hdu() for details.
csst_format : bool, optional
Whether to save the fits file in CSST format, by default True.
"""
if
not
csst_format
:
save_fits_simple
(
images
,
obs_info
,
output_folder
=
output_folder
)
return
hdu_header
,
folder
,
filename
=
primary_hdu
(
obs_info
,
gnc_info
,
True
)
hdu_list
=
fits
.
HDUList
([
hdu_header
])
if
len
(
images
.
shape
)
==
2
:
images
=
np
.
array
([
images
])
for
index
in
range
(
images
.
shape
[
0
]):
header
=
frame_header
(
obs_info
,
index
,
hdu_header
.
header
,
camera_dict
=
camera_dict
,
)
frame_hdu
=
fits
.
ImageHDU
(
images
[
index
,
:,
:],
header
=
header
)
frame_hdu
.
add_checksum
()
hdu_list
.
append
(
frame_hdu
)
folder
=
f
"
{
output_folder
}
/
{
folder
}
"
if
not
os
.
path
.
exists
(
folder
):
os
.
makedirs
(
folder
)
log
.
debug
(
f
'make new folder
{
folder
}
'
)
full_path
=
f
"
{
folder
}
/
{
filename
}
"
log
.
debug
(
f
'save fits file:
{
full_path
}
'
)
hdu_list
.
writeto
(
full_path
,
overwrite
=
True
)
source/csst_cpic_sim/main.py
0 → 100644
View file @
7c99b572
import
argparse
,
sys
,
tqdm
,
time
,
os
,
yaml
from
glob
import
glob
from
datetime
import
datetime
import
traceback
import
numpy
as
np
from
.target
import
spectrum_generator
,
target_file_load
from
.optics
import
focal_mask
,
focal_convolve
from
.camera
import
CosmicRayFrameMaker
,
sky_frame_maker
,
CpicVisEmccd
from
.io
import
save_fits
,
log
from
.config
import
update_able_keys
,
relative_time
from
.config
import
config
as
default_config
def
vis_observation
(
target
:
dict
,
skybg
:
float
,
expt
:
float
,
nframe
:
int
,
band
:
str
,
emset
:
int
,
obsid
:
int
=
41000000000
,
rotation
:
float
=
0
,
shift
:
list
=
[
0
,
0
],
gnc_info
:
dict
=
{},
csst_format
:
bool
=
True
,
camera
=
CpicVisEmccd
(),
crmaker
=
CosmicRayFrameMaker
(),
nsample
:
int
=
1
,
emgain
=
None
,
prograss_bar
=
None
,
output
=
'./'
)
->
np
.
ndarray
:
"""Observation simulation main process on visable band using EMCCD.
Parameters
-----------
target : dict
target dictionary. See the input of `target.spectrum_generator` for more details.
skybg : float
sky background in unit of magnitude/arcsec^2
expt: float
exposure time in unit of second
nframe: int
number of frames
band: str
band name
emset: int
EM gain setting value. 1023(0x3FF) for ~1.0× EM gain.
obsid: int
observation ID. Start from 4 for CPIC, 01 for science observation. See the input of io.obsid_parser for more details.
rotation: float
rotation of the telescope. in unit of degree. 0 means North is up.
shift: list
target shift in unit of arcsec
gnc_info: dict
gnc_info dictionary. See the input of io.primary_hdu for more details.
csst_format: bool
if True, the output fits file will be in the csst format.
crmaker: CosmicRayFrameMaker
CosmicRayFrameMaker object. See the input of camera.CosmicRayFrameMaker for more details.
nsample: int
number of samples for wide bandpass.
emgain: float or None
if None, emgain are set using emset parameter. Else, emgain are set using emgain parameter.
prograss_bar: bool
if True, a prograss_bar will be shown.
output: str
the output directory.
Returns
-------
image_cube: np.ndarray
"""
start_time
=
time
.
time
()
platescale
=
default_config
[
'platescale'
]
iwa
=
default_config
[
'mask_width'
]
/
2
area
=
default_config
[
'aperature_area'
]
expt_start
=
camera
.
system_time
target_list
=
[]
if
'cstar'
in
target
.
keys
():
target_list
=
spectrum_generator
(
target
)
image_cube
=
[]
if
emgain
is
None
:
emgain_value
=
camera
.
emgain_set
(
emset
,
self_update
=
False
)
else
:
emgain_value
,
emset
=
emgain
,
-
emgain
params
=
{
'target'
:
target
,
'skybg'
:
skybg
,
'expt'
:
expt
,
'nframe'
:
nframe
,
'band'
:
band
,
'emset'
:
emset
,
'emgain'
:
emgain_value
,
'obsid'
:
obsid
,
'rotation'
:
rotation
,
'shift'
:
shift
,
}
all_frame_info
=
[]
if
prograss_bar
:
pg_bar
=
tqdm
.
tqdm
(
total
=
nframe
,
leave
=
False
)
for
i
in
range
(
nframe
):
if
prograss_bar
:
pg_bar
.
update
(
1
)
else
:
print
(
f
'
\r
Simulation Running: Frame
{
i
+
1
}
/
{
nframe
}
'
,
end
=
''
)
frame_info
=
{}
frame_info
[
'expt_start'
]
=
camera
.
system_time
focal_frame
=
focal_convolve
(
band
,
target_list
,
init_shifts
=
shift
,
rotation
=
rotation
,
nsample
=
nsample
,
platesize
=
camera
.
flat_shape
)
if
skybg
is
None
or
skybg
>
100
:
sky_bkg_frame
=
0
else
:
sky_bkg_frame
=
sky_frame_maker
(
band
,
skybg
,
platescale
,
camera
.
flat_shape
)
sky_bkg_frame
=
sky_bkg_frame
*
area
*
expt
sky_bkg_frame
=
focal_mask
(
sky_bkg_frame
,
iwa
,
platescale
)
cr_frame
=
crmaker
.
make_cr_frame
(
camera
.
dark_shape
,
expt
)
camera
.
time_syn
(
expt
,
readout
=
True
)
image
=
camera
.
readout
(
focal_frame
+
sky_bkg_frame
,
emset
,
expt
,
image_cosmic_ray
=
cr_frame
,
emgain
=
emgain
)
image_cube
.
append
(
image
)
frame_info
[
'expt_end'
]
=
camera
.
system_time
frame_info
[
'chiptemp'
]
=
camera
.
ccd_temp
frame_info
[
'platescale'
]
=
platescale
frame_info
[
'iwa'
]
=
iwa
all_frame_info
.
append
(
frame_info
)
image_cube
=
np
.
array
(
image_cube
)
expt_end
=
camera
.
system_time
utc0
=
default_config
[
'utc0'
]
utc0_stamp
=
datetime
.
timestamp
(
datetime
.
fromisoformat
(
utc0
))
expt_start_iso
=
datetime
.
fromtimestamp
(
utc0_stamp
+
expt_start
)
expt_end_iso
=
datetime
.
fromtimestamp
(
utc0_stamp
+
expt_end
)
params
[
'EXPSTART'
]
=
expt_start_iso
.
isoformat
()
params
[
'EXPEND'
]
=
expt_end_iso
.
isoformat
()
params
[
'frame_info'
]
=
all_frame_info
save_fits
(
image_cube
,
params
,
gnc_info
,
camera
.
__dict__
.
copy
(),
csst_format
=
csst_format
,
output_folder
=
output
)
if
prograss_bar
:
pg_bar
.
close
()
print
(
f
' Done [
{
time
.
time
()
-
start_time
:
.
1
f
}
s] '
)
else
:
print
(
f
'
\r
Done [
{
time
.
time
()
-
start_time
:
.
1
f
}
s] '
)
return
image_cube
def
quick_run_v2
(
target_str
:
str
,
band
:
str
,
expt
:
float
,
emgain
:
float
,
nframe
:
int
,
skybg
:
float
=
None
,
rotation
:
float
=
0
,
shift
:
list
=
[
0
,
0
],
emset_input
:
bool
=
False
,
cr_frame
:
bool
=
True
,
camera_effect
:
bool
=
True
,
prograss_bar
=
False
,
output
=
'./'
)
->
np
.
ndarray
:
"""A quick run function for CPIC.
Parameters
----------
target_str: str
target string. See the input of `target.target_file_load` for more details.
band: str
band name
expt: float
exposure time in unit of second
emgain: float
EM gain value. Note that emgain is not the same as emset, controled by emset_input keyword.
nframe: int
number of frames
skybg: float or None
sky background in unit of magnitude. If None, sky background will not be set.
rotation: float
rotation of the telescope. in unit of degree. 0 means North is up.
shift: list
target shift in unit of arcsec
emset_input: bool
if True, emgain paramter is emset. Else, emgain is set using emgain parameter.
cr_frame: bool
if True, cosmic ray frame will be added to the image.
camera_effect: bool
if True, camera effect will be added to the image.
prograss_bar: bool
if True, a prograss_bar will be shown.
output: str
the output directory.
Returns
-------
image_cube: np.ndarray
"""
print
(
f
'Quick Run:
{
target_str
}
'
)
log
.
debug
(
f
"""input parameters:
target_str:
{
target_str
}
skybg:
{
skybg
}
band:
{
band
}
expt:
{
expt
}
nframe:
{
nframe
}
emgain:
{
emgain
}
rotation:
{
rotation
}
shift:
{
shift
}
emset_input:
{
emset_input
}
cr_frame:
{
cr_frame
}
camera_effect:
{
camera_effect
}
prograss_bar:
{
prograss_bar
}
output:
{
output
}
"""
)
target_dict
=
target_file_load
(
target_str
)
if
target_dict
==
{}:
target_dict
[
'name'
]
=
'blank'
emgain_value
=
emgain
if
emset_input
:
emgain_value
=
None
camera
=
CpicVisEmccd
()
if
not
camera_effect
:
camera
.
switch
[
'bias_vp'
]
=
False
camera
.
switch
[
'bias_hp'
]
=
False
camera
.
switch
[
'bias_ci'
]
=
False
camera
.
switch
[
'bias_shift'
]
=
False
camera
.
switch
[
'badcolumn'
]
=
False
camera
.
switch
[
'flat'
]
=
False
camera
.
switch
[
'cte'
]
=
False
camera
.
switch
[
'shutter'
]
=
False
camera
.
switch
[
'nonlinear'
]
=
False
if
not
cr_frame
:
camera
.
switch
[
'cosmicray'
]
=
False
return
vis_observation
(
target
=
target_dict
,
skybg
=
skybg
,
expt
=
expt
,
nframe
=
nframe
,
band
=
band
,
emset
=
emgain
,
emgain
=
emgain_value
,
csst_format
=
False
,
shift
=
shift
,
rotation
=
rotation
,
camera
=
camera
,
prograss_bar
=
prograss_bar
,
output
=
output
)
def
deduplicate_names_add_count
(
names
:
list
):
"""remove duplicate names and add count"""
for
i
in
range
(
len
(
names
)
-
1
,
-
1
,
-
1
):
if
names
.
count
(
names
[
i
])
>
1
:
names
[
i
]
=
names
[
i
]
+
'_'
+
str
(
names
.
count
(
names
[
i
]))
def
observation_simulation_from_config
(
obs_file
,
config_file
):
"""Run observation simulation from config file
Parameters
-----------
obs_file: str
observation file or folder.
config_file: str
config file.
Examples:
see examples in `example` folder.
"""
config
=
{}
if
config_file
:
with
open
(
config_file
)
as
fid
:
config
=
yaml
.
load
(
fid
,
Loader
=
yaml
.
FullLoader
)
camera_config
=
config
.
get
(
'camera'
,
[{}])
all_camera
=
[]
all_camera_name
=
[]
for
c_config
in
camera_config
:
all_camera
.
append
(
CpicVisEmccd
(
c_config
))
all_camera_name
.
append
(
c_config
.
get
(
'name'
,
'camera'
))
deduplicate_names_add_count
(
all_camera_name
)
for
key
in
config
.
keys
():
if
key
in
update_able_keys
:
default_config
[
key
]
=
config
[
key
]
output_folder
=
default_config
[
'output'
]
csst_format
=
default_config
[
'csst_format'
]
nsample
=
default_config
[
'nsample'
]
obs_file
=
os
.
path
.
abspath
(
obs_file
)
file_list
=
[]
if
os
.
path
.
isdir
(
obs_file
):
file_list
=
glob
(
f
"
{
obs_file
}
/*.yaml"
)
elif
'.yaml'
==
obs_file
[
-
5
:]:
file_list
=
[
obs_file
]
if
not
file_list
:
log
.
warning
(
f
"No observation file found in
{
obs_file
}
"
)
for
ind_target
,
file
in
enumerate
(
file_list
):
try
:
with
open
(
file
,
'r'
)
as
fid
:
obs_info
=
yaml
.
load
(
fid
,
Loader
=
yaml
.
FullLoader
)
target
=
target_file_load
(
obs_info
.
get
(
'target'
,
{}))
skybg
=
obs_info
.
get
(
'skybg'
,
None
)
expt
=
obs_info
[
'expt'
]
band
=
obs_info
[
'band'
]
emset
=
obs_info
[
'emset'
]
nframe
=
obs_info
[
'nframe'
]
obsid
=
obs_info
[
'obsid'
]
rotation
=
obs_info
.
get
(
'rotation'
,
0
)
shift
=
obs_info
.
get
(
'shift'
,
[
0
,
0
])
gnc_info
=
obs_info
.
get
(
'gnc_info'
,
{})
time
=
obs_info
.
get
(
'time'
,
0
)
emgain
=
obs_info
.
get
(
'emgain'
,
None
)
time
=
relative_time
(
time
)
except
Exception
as
e
:
log
.
error
(
f
"
{
file
}
is not a valid yaml file."
)
log
.
error
(
f
"Failed with
{
type
(
e
).
__name__
}{
e
}
.
\n\n
{
traceback
.
format_exc
()
}
"
)
continue
ind_camera
=
0
for
camera_name
,
camera
in
zip
(
all_camera_name
,
all_camera
):
ind_camera
+=
1
ind_run
=
ind_target
*
len
(
all_camera
)
+
ind_camera
all_run
=
len
(
all_camera
)
*
len
(
file_list
)
info_text
=
f
"(
{
ind_run
}
/
{
all_run
}
) obsid[
{
obsid
}
] with
{
camera_name
}
"
log
.
info
(
info_text
)
if
time
==
0
:
camera
.
time_syn
(
time
,
initial
=
True
)
else
:
dt
=
time
-
camera
.
system_time
if
dt
<
0
:
log
.
warning
(
f
'Time is not synced.
{
dt
}
seconds are added.'
)
dt
=
0
camera
.
time_syn
(
dt
,
readout
=
False
)
if
len
(
all_camera
)
>
1
:
output
=
os
.
path
.
join
(
output_folder
,
camera_name
)
else
:
output
=
output_folder
try
:
vis_observation
(
target
,
skybg
,
expt
,
nframe
,
band
,
emset
,
obsid
,
emgain
=
emgain
,
rotation
=
rotation
,
shift
=
shift
,
gnc_info
=
gnc_info
,
camera
=
camera
,
output
=
output
,
nsample
=
nsample
,
csst_format
=
csst_format
,
prograss_bar
=
True
)
except
Exception
as
e
:
log
.
error
(
f
"
{
info_text
}
failed with
{
type
(
e
).
__name__
}{
e
}
.
\n\n
{
traceback
.
format_exc
()
}
"
)
def
main
(
argv
=
None
):
"""
Command line interface of csst_cpic_sim
Parameters
-----------
argv: list
input arguments. Default is None. If None, sys.argv is used.
"""
parser
=
argparse
.
ArgumentParser
(
description
=
'Cpic obsevation image simulation'
)
parser
.
set_defaults
(
func
=
lambda
_
:
parser
.
print_usage
())
subparsers
=
parser
.
add_subparsers
(
help
=
'type of runs'
)
parser_quickrun
=
subparsers
.
add_parser
(
'quickrun'
,
help
=
'a quick observation with no configration file'
)
parser_quickrun
.
add_argument
(
'target_string'
,
type
=
str
,
help
=
'example: *5.1/25.3(1.3,1.5)/22.1(2.3,-4.5)'
)
parser_quickrun
.
add_argument
(
'expt'
,
type
=
float
,
help
=
'exposure time [ms]'
)
parser_quickrun
.
add_argument
(
'emgain'
,
type
=
float
,
help
=
'emgain or emgain set value if emgain_input is False'
)
parser_quickrun
.
add_argument
(
'nframe'
,
type
=
int
,
help
=
'number of frames'
)
parser_quickrun
.
add_argument
(
'-b'
,
'--band'
,
type
=
str
,
default
=
'f661'
,
help
=
'band, one of f565/f661/f743/f883'
)
parser_quickrun
.
add_argument
(
'-r'
,
'--rotation'
,
type
=
float
,
default
=
0
,
help
=
'rotation angle [degree]'
)
parser_quickrun
.
add_argument
(
'-s'
,
'--skybk'
,
type
=
float
,
default
=
21
,
help
=
'magnitude of sky background [mag/arcsec^2]'
)
parser_quickrun
.
add_argument
(
'-f'
,
'--cr_frame'
,
action
=
'store_true'
,
help
=
'if True, cosmic ray frame will be added'
)
parser_quickrun
.
add_argument
(
'-e'
,
'--emset'
,
action
=
'store_true'
,
help
=
'if True, emgain set value will be used as input'
)
parser_quickrun
.
add_argument
(
'-c'
,
'--camera_effect'
,
action
=
'store_true'
,
help
=
'if True, camera effect will be added'
)
parser_quickrun
.
add_argument
(
'-o'
,
'--output'
,
type
=
str
,
default
=
'./'
,
help
=
'output folder'
)
def
quick_run_call
(
args
):
quick_run_v2
(
target_str
=
args
.
target_string
,
expt
=
args
.
expt
,
emgain
=
args
.
emgain
,
nframe
=
args
.
nframe
,
band
=
args
.
band
,
rotation
=
args
.
rotation
,
skybg
=
args
.
skybk
,
cr_frame
=
args
.
cr_frame
,
emset_input
=
args
.
emset
,
camera_effect
=
args
.
camera_effect
,
output
=
os
.
path
.
abspath
(
args
.
output
),
prograss_bar
=
True
)
parser_quickrun
.
set_defaults
(
func
=
quick_run_call
)
parser_run
=
subparsers
.
add_parser
(
'run'
,
help
=
'observation with configration files or folder'
)
parser_run
.
add_argument
(
'target_config'
,
type
=
str
,
help
=
'configration file or folder'
)
parser_run
.
add_argument
(
'-c'
,
'--config'
,
type
=
str
,
help
=
'instrument configration file'
)
def
run_all
(
args
):
observation_simulation_from_config
(
args
.
target_config
,
args
.
config
,
)
parser_run
.
set_defaults
(
func
=
run_all
)
args
=
parser
.
parse_args
(
argv
)
args
.
func
(
args
)
# if __name__ == '__main__': # pragma: no cover
# sys.exit(main())
# target_example = {
# 'cstar': {
# 'magnitude': 1,
# 'ra': '120d',
# 'dec': '40d',
# 'distance': 10,
# 'sptype': 'F0III',
# },
# 'stars': [
# {
# 'magnitude': 20,
# 'pangle': 60,
# 'separation': 1,
# 'sptype': 'F0III'
# }
# ]
# }
# # quick_run('', 10, 'f661', 1, 1, 30)
# # quick_run('*2.4/10(3,5)/15(-4,2)', 21, 'f661', 1, 1, 30)
# # # normal target
# observation_simulation(
# target=target_example,
# skybg=21,
# expt=1,
# nframe=2,
# band='f661',
# emgain=30,
# obsid=51012345678,
# )
# # bias
# observation_simulation(
# target=target_example,
# skybg=999,
# expt=1,
# nframe=2,
# band='f661',
# emgain=1,
# obsid=51012345678,
# shift=[3, 3],
# rotation=60
# )
# # bias-gain
# observation_simulation(
# target={},
# skybg=999,
# expt=0.01,
# nframe=2,
# band='f661',
# emgain=1000,
# obsid=50012345678,
# )
# # dark
# observation_simulation(
# target={},
# skybg=999,
# expt=100,
# nframe=2,
# band='f661',
# emgain=30,
# obsid=50112345678,
# )
# # flat
# observation_simulation(
# target={},
# skybg=15,
# expt=10,
# nframe=2,
# band='f661',
# emgain=30,
# obsid=50112345678,
# )
source/csst_cpic_sim/optics.py
0 → 100644
View file @
7c99b572
import
numpy
as
np
from
scipy.signal
import
fftconvolve
from
scipy.ndimage
import
rotate
from
.config
import
config
,
S
# S is synphot
from
.utils
import
region_replace
from
.io
import
log
from
.psf_simulation
import
single_band_masked_psf
,
single_band_psf
FILTERS
=
{}
for
key
,
value
in
config
[
'bands'
].
items
():
FILTERS
[
key
]
=
S
.
FileBandpass
(
value
)
default_band
=
config
[
'default_band'
]
def
filter_throughput
(
filter_name
):
"""
Totally throughput of the each CPIC band.
Including the throughput of the filter, telescope, cpic, and camera QE.
If the filter_name is not supported, return the throughput of the default filter(f661).
Parameters
-----------
filter_name: str
The name of the filter.
One of ['f565', 'f661'(default), 'f743', 'f883', 'f940', 'f1265', 'f1425', 'f1542']
Returns
--------
synphot.Bandpass
The throughput of the filter.
"""
filter_name
=
filter_name
.
lower
()
filter_name
=
default_band
if
filter_name
==
'default'
else
filter_name
if
filter_name
not
in
FILTERS
.
keys
():
log
.
warning
(
f
"滤光片名称错误(
{
filter_name
}
),返回默认滤光片(
{
default_band
}
)透过率"
)
filter_name
=
default_band
return
FILTERS
[
filter_name
]
def
_rotate_and_shift
(
shift
,
rotation
,
init_shifts
):
rotation_rad
=
rotation
/
180
*
np
.
pi
return
np
.
array
([
shift
[
0
]
*
np
.
cos
(
rotation_rad
)
+
shift
[
1
]
*
np
.
sin
(
rotation_rad
),
-
shift
[
0
]
*
np
.
sin
(
rotation_rad
)
+
shift
[
1
]
*
np
.
cos
(
rotation_rad
)
])
+
np
.
array
(
init_shifts
)
def
ideal_focus_image
(
bandpass
:
S
.
spectrum
.
SpectralElement
,
targets
:
list
,
platescale
,
platesize
:
list
=
[
1024
,
1024
],
init_shifts
:
list
=
[
0
,
0
],
rotation
:
float
=
0
)
->
np
.
ndarray
:
"""Ideal focus image of the targets.
Each star is a little point of 1pixel.
Parameters
-----------
bandpass: synphot.SpectralElement
The bandpass of the filter.
targets: list
The list of the targets. See the output of `spectrum_generator` for details.
platescale: float
The platescale of the camera. Unit: arcsec/pixel
platesize: list
The size of the image. Unit: pixel
init_shifts: list
The shifts of the targets to simulate the miss alignment. Unit: arcsec
rotation: float
The rotation of the image. Unit: degree
Returns
--------
np.ndarray
The ideal focus image.
"""
focal_image
=
np
.
zeros
(
platesize
)
focal_shape
=
np
.
array
(
platesize
)[::
-
1
]
# x, y
if
not
targets
:
return
focal_image
for
target
in
targets
:
sub_x
,
sub_y
,
sub_spectrum
,
sub_image
=
target
sub_shift
=
_rotate_and_shift
([
sub_x
,
sub_y
],
rotation
,
init_shifts
)
/
platescale
sed
=
(
sub_spectrum
*
bandpass
).
integrate
()
if
sub_image
is
None
:
x
=
(
focal_shape
[
0
]
-
1
)
/
2
+
sub_shift
[
0
]
y
=
(
focal_shape
[
1
]
-
1
)
/
2
+
sub_shift
[
1
]
int_x
=
int
(
x
)
int_y
=
int
(
y
)
if
int_x
<
0
or
int_x
>=
focal_shape
[
0
]
-
1
or
int_y
<
0
or
int_y
>=
focal_shape
[
1
]
-
1
:
continue
dx1
=
x
-
int_x
dx0
=
1
-
dx1
dy1
=
y
-
int_y
dy0
=
1
-
dy1
sub
=
np
.
array
([
[
dx0
*
dy0
,
dx1
*
dy0
],
[
dx0
*
dy1
,
dx1
*
dy1
]])
*
sed
focal_image
[
int_y
:
int_y
+
2
,
int_x
:
int_x
+
2
]
+=
sub
else
:
# sub_image = sub_image
sub_image
=
np
.
abs
(
rotate
(
sub_image
,
rotation
,
reshape
=
False
))
sub_image
=
sub_image
/
sub_image
.
sum
()
sub_img_shape
=
np
.
array
(
sub_image
.
shape
)[::
-
1
]
sub_shift
+=
(
focal_shape
-
1
)
/
2
-
(
sub_img_shape
-
1
)
/
2
focal_image
=
region_replace
(
focal_image
,
sub_image
*
sed
,
sub_shift
,
subpix
=
True
)
return
focal_image
def
focal_convolve
(
band
:
str
,
targets
:
list
,
init_shifts
:
list
=
[
0
,
0
],
rotation
:
float
=
0
,
nsample
:
int
=
5
,
error
:
float
=
0
,
platesize
:
list
=
[
1024
,
1024
])
->
np
.
ndarray
:
"""PSF convolution of the ideal focus image.
Parameters
----------
band: str
The name of the band.
target: list
The list of thetargets. See the output of `spectrum_generator` for details.
init_shifts: list
The shifts of the targets to simulate the miss alignment. Unit: arcsec
rotation: float
The rotation of the image. Unit: degree
error: float
The error of the DM acceleration. Unit: nm
platesize: list
The size of the image. Unit: pixel
Returns
--------
np.ndarray
"""
# config = optics_config[which_focalplane(band)]
platescale
=
config
[
'platescale'
]
# telescope_config = optics_config['telescope']
area
=
config
[
'aperature_area'
]
filter
=
filter_throughput
(
band
)
throughput
=
filter
.
throughput
wave
=
filter
.
wave
throughput_criterion
=
throughput
.
max
()
*
0.1
wave_criterion
=
wave
[
throughput
>
throughput_criterion
]
min_wave
=
wave_criterion
[
0
]
max_wave
=
wave_criterion
[
-
1
]
# print(min_wave, max_wave)
platescale
=
config
[
'platescale'
]
iwa
=
config
[
'mask_width'
]
/
2
if
abs
(
init_shifts
[
0
])
>
4
or
abs
(
init_shifts
[
1
])
>
4
:
print
(
'Input shifts are too large, and are set to zero'
)
init_shifts
=
[
0
,
0
]
all_fp_image
=
[]
if
not
targets
:
return
np
.
zeros
((
platesize
[
1
],
platesize
[
0
]))
for
i_wave
in
range
(
nsample
):
d_wave
=
(
max_wave
-
min_wave
)
/
nsample
wave0
=
min_wave
+
i_wave
*
d_wave
wave1
=
min_wave
+
(
i_wave
+
1
)
*
d_wave
center_wavelength
=
(
wave0
+
wave1
)
/
2
*
1e-10
i_throughput
=
throughput
.
copy
()
i_throughput
[(
wave
>
wave1
)
|
(
wave
<
wave0
)]
=
0
i_band
=
S
.
ArrayBandpass
(
wave
,
i_throughput
,
waveunits
=
'angstrom'
)
i_fp_image
=
ideal_focus_image
(
i_band
,
targets
[
1
:],
platescale
,
platesize
,
init_shifts
,
rotation
)
psf
=
single_band_psf
(
center_wavelength
,
error
=
error
)
_
,
_
,
cstar_sp
,
_
=
targets
[
0
]
cstar_flux
=
(
cstar_sp
*
i_band
).
integrate
()
cstar_psf
=
single_band_masked_psf
(
center_wavelength
,
error
=
error
,
shift
=
init_shifts
)
c_fp_image
=
fftconvolve
(
i_fp_image
,
psf
,
mode
=
'same'
)
c_fp_image
=
focal_mask
(
c_fp_image
,
iwa
,
platescale
)
c_fp_image
=
c_fp_image
+
cstar_flux
*
cstar_psf
all_fp_image
.
append
(
c_fp_image
*
area
)
# trans to photon/second
return
np
.
array
(
all_fp_image
).
sum
(
axis
=
0
)
def
focal_mask
(
image
,
iwa
,
platescale
,
throughtput
=
1e-6
):
"""
Mask the image outside the inner working angle.
Parameters
-----------
image: np.ndarray
The image to be masked.
iwa: float
The inner working angle. Unit: arcsec.
platescale: float
The platescale of the image. Unit: arcsec/pixel.
throughtput: float
The throughtput of the mask. The default is 1e-6.
Returns
--------
np.ndarray
The masked image.
"""
xx
,
yy
=
np
.
mgrid
[
0
:
image
.
shape
[
0
],
0
:
image
.
shape
[
1
]]
center
=
np
.
array
([(
image
.
shape
[
0
]
-
1
)
/
2
,
(
image
.
shape
[
1
]
-
1
)
/
2
])
mask
=
(
abs
(
xx
-
center
[
0
])
<
iwa
/
platescale
)
|
(
abs
(
yy
-
center
[
1
])
<
iwa
/
platescale
)
image_out
=
image
.
copy
()
image_out
[
mask
]
*=
throughtput
return
image_out
source/csst_cpic_sim/psf_simulation.py
0 → 100644
View file @
7c99b572
import
os
,
pickle
import
numpy
as
np
from
astropy.io
import
fits
from
hcipy
import
Field
,
Wavefront
,
FraunhoferPropagator
from
hcipy
import
SurfaceApodizer
,
SurfaceAberrationAtDistance
from
hcipy
import
TipTiltMirror
,
DeformableMirror
from
hcipy
import
ComplexSurfaceApodizer
from
hcipy
import
make_pupil_grid
,
make_circular_aperture
,
make_obstructed_circular_aperture
from
hcipy
import
make_xinetics_influence_functions
,
make_uniform_grid
from
.config
import
config
DMCACHE
=
True
ARCSEC2RAD
=
np
.
radians
(
1
/
3600
)
# initial psf simulation
apm
,
apm_header
=
fits
.
getdata
(
config
[
'apm_file'
],
header
=
True
)
actuator
=
fits
.
getdata
(
config
[
'actuator_file'
])
surface_aberration
=
fits
.
getdata
(
config
[
'aberration'
])
pupil_diameter
=
2
# m
F_number
=
83
focal_length
=
pupil_diameter
*
F_number
pupil_shape
=
apm
.
shape
pupil_rate
=
apm_header
[
'PHRATE'
]
# meter/pixel
pupil_grid
=
make_pupil_grid
(
pupil_shape
[
0
],
pupil_shape
[
0
]
*
pupil_rate
)
aperture
=
make_circular_aperture
(
pupil_diameter
)(
pupil_grid
)
aperture
=
aperture
*
Field
(
apm
.
flatten
(),
pupil_grid
)
second_pupil_size
=
pupil_diameter
*
2
# just gauss a number
second_aperture
=
make_circular_aperture
(
second_pupil_size
)(
pupil_grid
)
lyot_stop
=
ComplexSurfaceApodizer
(
second_aperture
,
second_aperture
*
0
,
lambda
_
:
2
)
emccd_pixel_size
=
13e-6
# m
platescale
=
emccd_pixel_size
/
(
ARCSEC2RAD
*
focal_length
)
# arcsec / pixel
focal_shape
=
np
.
array
([
1024
,
1024
])
focal_size
=
focal_shape
*
emccd_pixel_size
# platescale = vis_focal_config['platescale']
mask_width
=
config
[
'mask_width'
]
focal_full_frame
=
make_uniform_grid
(
focal_shape
,
focal_size
,
has_center
=
True
)
prop_full_frame
=
FraunhoferPropagator
(
pupil_grid
,
focal_full_frame
,
focal_length
)
spider_width
=
mask_width
/
platescale
*
emccd_pixel_size
spider
=
make_obstructed_circular_aperture
(
focal_size
[
0
]
*
2
,
0
,
4
,
spider_width
)(
focal_full_frame
)
pupil_cross_mask
=
ComplexSurfaceApodizer
(
spider
,
spider
*
0
,
lambda
_
:
2
)
num_actuators_across
=
32
# dm spacing is little smaller than pupil
actuator_spacing
=
0.95
*
pupil_diameter
/
num_actuators_across
pickle_file
=
config
[
'dm_pickle'
]
if
os
.
path
.
exists
(
pickle_file
)
and
DMCACHE
:
with
open
(
pickle_file
,
'rb'
)
as
fid
:
influence_functions
=
pickle
.
load
(
fid
)
else
:
influence_functions
=
make_xinetics_influence_functions
(
pupil_grid
,
num_actuators_across
,
actuator_spacing
)
with
open
(
pickle_file
,
'wb'
)
as
fid
:
pickle
.
dump
(
influence_functions
,
fid
)
deformable_mirror
=
DeformableMirror
(
influence_functions
)
tiptilt_mirror
=
TipTiltMirror
(
pupil_grid
)
aberration
=
SurfaceApodizer
(
surface_sag
=
surface_aberration
.
flatten
(),
refractive_index
=-
1
)
# arbitrary distance for the aberration to propagate
aberration_distance
=
80
*
focal_length
aberration
=
SurfaceAberrationAtDistance
(
aberration
,
aberration_distance
)
def
single_band_masked_psf
(
wavelength
:
float
,
error
:
float
=
0
,
shift
:
list
=
[
0
,
0
])
->
np
.
ndarray
:
"""CPIC PSF considering the focal plane mask.
Parameters
-----------
wavelength : float
observation wavelength in meter
error : float
deformable mirror control error in nm
shift : list
angular shift of the target in arcsec.
Returns
----------
psf : np.ndarray
psf in the focal plane. Normalized as the input flux is 1.
(Note that total flux of the psf is not 1, because it is masked)
"""
error
=
np
.
random
.
normal
(
0
,
error
*
1e-9
,
actuator
.
shape
)
deformable_mirror
.
actuators
=
actuator
+
error
wf
=
Wavefront
(
aperture
,
wavelength
)
shift
=
np
.
array
(
shift
)
*
ARCSEC2RAD
/
2
tiptilt_mirror
.
actuators
=
shift
wf
=
tiptilt_mirror
(
wf
)
wf
=
aberration
(
wf
)
first_focal
=
prop_full_frame
(
deformable_mirror
(
wf
))
strength
=
first_focal
.
intensity
.
shaped
.
sum
()
second_pupil
=
prop_full_frame
.
backward
(
pupil_cross_mask
(
first_focal
))
second_focal
=
prop_full_frame
(
lyot_stop
(
second_pupil
))
psf
=
second_focal
.
intensity
.
shaped
return
psf
/
strength
def
single_band_psf
(
wavelength
:
float
,
error
:
float
=
0
)
->
np
.
ndarray
:
"""CPIC PSF without considering the focal plane mask.
Used to simulate the off-axis PSF.
Parameters
-----------
wavelength : float
observation wavelength in meter
error : float
deformable mirror control error in nm
Returns
----------
psf : np.ndarray
psf in the focal plane. Normalized. The total flux is 1.
"""
error
=
np
.
random
.
normal
(
0
,
error
*
1e-9
,
actuator
.
shape
)
deformable_mirror
.
actuators
=
actuator
+
error
wf
=
Wavefront
(
aperture
,
wavelength
)
wf
=
aberration
(
wf
)
first_focal
=
prop_full_frame
(
deformable_mirror
(
wf
))
image
=
np
.
array
(
first_focal
.
intensity
.
shaped
)
return
image
/
image
.
sum
()
# def single_band_psf(wavelength, error=0, aber_phase=None):
# error = np.random.normal(0, error*1e-9, actuator.shape)
# deformable_mirror.actuators = actuator + error
# wf = Wavefront(aperture, wavelength)
# wf = aberration(wf)
# if aber_phase is not None:
# other_error = SurfaceApodizer(
# surface_sag=aber_phase.flatten(), refractive_index=-1)
# wf = other_error(wf)
# first_focal = prop_full_frame(deformable_mirror(wf))
# image = np.array(first_focal.intensity.shaped)
# return image / image.sum()
\ No newline at end of file
source/csst_cpic_sim/target.py
0 → 100644
View file @
7c99b572
import
os
,
re
,
json
,
yaml
from
typing
import
Union
import
numpy
as
np
from
scipy
import
constants
from
astropy.io
import
fits
from
astropy.coordinates
import
SkyCoord
from
.config
import
S
# pysynphot
from
.config
import
config
from
.optics
import
filter_throughput
from
.io
import
log
from
pysynphot.renorm
import
StdRenorm
PLANK_H
=
constants
.
h
# ~6.62607015e-34
LIGHT_SPEED
=
constants
.
c
# ~3e8
DEG2RAD
=
np
.
pi
/
180
R_JUPITER_METER
=
6.99e4
AU_METER
=
1.49e8
DEFAULT_FILTER_FILE
=
config
[
'default_band'
]
HYBRID_MODEL_FILE
=
config
[
'hybrid_model'
]
BCC_MODEL_FOLDER
=
config
[
'bcc_model'
]
MAG_SYSTEM
=
config
[
'mag_system'
]
CATALOG_CACHE
=
{}
class
AlbedoCat
(
S
.
spectrum
.
TabularSpectralElement
,
S
.
catalog
.
Icat
):
"""Generate albedo spectrum from the planet reflection model in Batalha et al. 2018
Parameters
----------
phase : float
[degree], The phase angle of the planet, range from 0 to 180.
metallicity : float
The metallicity of the planet. log(Z/Zsun), range from 0 to 2.
f_sed : float
The sedimentation efficiency of the cloud. log(f_sed), range from -2 to 2. 2 is for cloud free case.
Notes
-----
f_sed is only reliable between -2 and log10(6) and the cloud free case (2).
Values between log10(6) and 2 are interpolated from the cloud free case (2) and log10(6).
Reference
---------
Color Classification of Extrasolar Giant Planets: Prospects and Cautions
Natasha E. Batalha et al 2018 AJ 156 158
"""
def
__init__
(
self
,
phase
:
float
,
metallicity
:
float
,
f_sed
:
float
,
):
catname
=
'BCCalbedo'
self
.
isAnalytic
=
False
self
.
name
=
f
"
{
catname
}
_
{
phase
}
_
{
metallicity
}
_
{
f_sed
}
"
self
.
parameter_names
=
[
'phase'
,
'metallicity'
,
'f_sed'
]
self
.
catalog_folder
=
BCC_MODEL_FOLDER
filename
=
os
.
path
.
join
(
self
.
catalog_folder
,
'catalog.fits'
)
if
filename
in
CATALOG_CACHE
:
indices
=
CATALOG_CACHE
[
filename
]
else
:
with
fits
.
open
(
filename
)
as
table
:
indexList
=
table
[
1
].
data
.
field
(
'INDEX'
)
filenameList
=
table
[
1
].
data
.
field
(
'FILENAME'
)
indices
=
self
.
_getArgs
(
indexList
,
filenameList
)
CATALOG_CACHE
[
filename
]
=
indices
list0
,
list1
=
self
.
_breakList
(
indices
,
0
,
phase
)
list2
,
list3
=
self
.
_breakList
(
list0
,
1
,
metallicity
)
list4
,
list5
=
self
.
_breakList
(
list1
,
1
,
metallicity
)
list6
,
list7
=
self
.
_breakList
(
list2
,
2
,
f_sed
)
list8
,
list9
=
self
.
_breakList
(
list3
,
2
,
f_sed
)
list10
,
list11
=
self
.
_breakList
(
list4
,
2
,
f_sed
)
list12
,
list13
=
self
.
_breakList
(
list5
,
2
,
f_sed
)
sp1
,
wave
,
waveunits
=
self
.
_getSpectrum
(
list6
[
0
],
catname
,
wave_output
=
True
)
sp2
=
self
.
_getSpectrum
(
list7
[
0
],
catname
)
sp3
=
self
.
_getSpectrum
(
list8
[
0
],
catname
)
sp4
=
self
.
_getSpectrum
(
list9
[
0
],
catname
)
sp5
=
self
.
_getSpectrum
(
list10
[
0
],
catname
)
sp6
=
self
.
_getSpectrum
(
list11
[
0
],
catname
)
sp7
=
self
.
_getSpectrum
(
list12
[
0
],
catname
)
sp8
=
self
.
_getSpectrum
(
list13
[
0
],
catname
)
spa1
=
self
.
_interpolateSpectrum
(
sp1
,
sp2
,
f_sed
)
spa2
=
self
.
_interpolateSpectrum
(
sp3
,
sp4
,
f_sed
)
spa3
=
self
.
_interpolateSpectrum
(
sp5
,
sp6
,
f_sed
)
spa4
=
self
.
_interpolateSpectrum
(
sp7
,
sp8
,
f_sed
)
spa5
=
self
.
_interpolateSpectrum
(
spa1
,
spa2
,
metallicity
)
spa6
=
self
.
_interpolateSpectrum
(
spa3
,
spa4
,
metallicity
)
spa7
=
self
.
_interpolateSpectrum
(
spa5
,
spa6
,
phase
)
sp
=
spa7
[
0
]
self
.
_wavetable
=
wave
*
1e4
self
.
_throughputtable
=
sp
self
.
waveunits
=
S
.
units
.
Units
(
waveunits
.
lower
())
self
.
warnings
=
{}
def
_getSpectrum
(
self
,
parlist
,
catdir
,
wave_output
=
False
):
name
=
parlist
[
3
]
filename
=
name
.
split
(
'['
)[
0
]
column
=
name
.
split
(
'['
)[
1
][:
-
1
]
filename
=
f
"
{
self
.
catalog_folder
}
/
{
filename
}
"
# sp = S.spectrum.TabularSpectralElement(filename, thrucol=column)
with
fits
.
open
(
filename
)
as
td
:
sp
=
td
[
1
].
data
.
field
(
column
)
wave
=
td
[
1
].
data
.
field
(
'wavelength'
)
# waveunits = td[1].header['tunit1']
waveunits
=
'micron'
result
=
[]
for
member
in
parlist
:
result
.
append
(
member
)
result
.
pop
()
result
.
append
(
sp
)
if
wave_output
:
return
result
,
wave
,
waveunits
return
result
def
_sptype2num
(
spectral_type
:
str
)
->
tuple
:
"""
convert spectral type string to number, for interpretation
case0: normal case
- M1V: 6.1, 5
- O5IV: 0.5, 4
- F3V: 3.3, 5
- K4.5II: 5.45, 2
case 1: star type or subtype missing
zero or V will return
- M1: 6.1, 5
- M: 6.0, 5
case 2: spectral type + subtype
subtype will be ignored
- K3Vvar: 5.3, 5
- F6Vbwvar: 3.6, 5
- K0IV SB: 5.0, 4
- F5V+: 3.5, 5
case 3: multi spectral type
only the first sptype is used
- G5IV/V +K1IV: 4.5, 4
- F7IV-V: 3.7, 4
- O4/O5IV: 0.4, 0
- G3/G5V: 4.3, 0
case 4: illegal input
ValueError will be raised
"""
obafgkm
=
'OBAFGKML'
spectral_type
=
spectral_type
.
upper
()
# match spectral type such as M1V, O5IV, F3V, K4.5II
matched
=
re
.
match
(
RF
'^([
{
obafgkm
}
])([0-9]\d*\.?\d*)*([IV]*)'
,
spectral_type
)
if
not
matched
:
raise
ValueError
(
f
"illegal spectral type input:
{
spectral_type
}
"
)
shorttype
=
obafgkm
.
find
(
matched
.
group
(
1
))
subtype
=
0.0
if
matched
.
group
(
2
):
subtype
=
float
(
matched
.
group
(
2
))
stlist
=
[
'O'
,
'I'
,
'II'
,
'III'
,
'IV'
,
'V'
]
startype
=
dict
(
zip
(
stlist
,
range
(
len
(
stlist
)))).
get
(
matched
.
group
(
3
),
5
)
return
shorttype
+
subtype
/
10
,
startype
def
_spnum2teff
(
spn
:
int
,
stn
:
int
)
->
tuple
:
"""
interpret of spectral number (by __sptype2num) to get t_eff and log_g
look up table from the document of ck04model
"""
with
open
(
config
[
'sp2teff_model'
],
'r'
)
as
fid
:
sptype_teff_lut
=
json
.
load
(
fid
)
def
_interp
(
spn
,
stn
):
lut
=
sptype_teff_lut
[
f
'
{
stn
}
'
]
teff
=
np
.
interp
(
spn
,
lut
[
0
],
lut
[
1
])
logg
=
np
.
interp
(
spn
,
lut
[
0
],
lut
[
2
])
return
teff
,
logg
stn
=
5
if
stn
not
in
[
1
,
2
,
3
,
4
,
5
]
else
stn
if
stn
in
[
1
,
3
,
5
]:
return
_interp
(
spn
,
stn
)
else
:
teff_low
,
logg_low
=
_interp
(
spn
,
stn
-
1
)
teff_high
,
logg_high
=
_interp
(
spn
,
stn
+
1
)
return
(
teff_high
+
teff_low
)
/
2
,
(
logg_high
+
logg_low
)
/
2
def
star_photlam
(
magnitude
:
float
,
sptype
:
str
,
is_blackbody
:
bool
=
False
,
mag_input_band
:
str
=
'f661'
)
->
S
.
ArraySpectrum
:
"""
genrate flux spectrum of a star by its observal magnitude and spectral type
Parameters
----------
magnitude: float
magnitude of the star
sptype: str
spectral type of the star
is_blackbody: bool
if True, use blackbody spectrum instead of ck04model
mag_input_band: str
bandpass of the input magnitude, default is f661
Returns
-------
pysynphot.spectrum.ArraySpectrum
spectrum of the star in photlam unit
"""
spn
,
stn
=
_sptype2num
(
sptype
)
t_eff
,
log_g
=
_spnum2teff
(
spn
,
stn
)
log
.
debug
(
f
"
{
sptype
}
star => [
{
t_eff
=
:}
,
{
log_g
=
:}
];
{
is_blackbody
=
:}
"
)
filter
=
filter_throughput
(
mag_input_band
)
if
not
is_blackbody
:
METALLICITY
=
0
spectrum
=
S
.
Icat
(
'ck04models'
,
t_eff
,
METALLICITY
,
log_g
)
else
:
spectrum
=
S
.
BlackBody
(
t_eff
)
star_sp
=
spectrum
.
renorm
(
magnitude
,
MAG_SYSTEM
,
filter
)
star_sp
.
convert
(
'photlam'
)
return
star_sp
def
standard_spectrum
(
magnitude
:
float
)
->
S
.
ArraySpectrum
:
"""Standard spectrum of magnitude system.
For example, the standard_spectrum of vega megnitude is vega spectrum.
The standard spectrum of ab magnitude is a flat spectrum.
Parameters
-----------
magnitude : float
magnitude of the standard spectrum
Returns
----------
star_sp : S.spectrum.Spectrum
"""
inner_std
=
S
.
units
.
Units
(
MAG_SYSTEM
).
StdSpectrum
std_spectrum
=
S
.
ArraySpectrum
(
inner_std
.
wave
,
inner_std
.
flux
,
inner_std
.
waveunits
,
inner_std
.
fluxunits
)
filter
=
filter_throughput
(
DEFAULT_FILTER_FILE
)
std_spectrum
=
StdRenorm
(
std_spectrum
,
filter
,
magnitude
,
MAG_SYSTEM
)
std_spectrum
.
convert
(
'photlam'
)
return
std_spectrum
def
bcc_spectrum
(
coe_cloud
:
float
,
coe_metalicity
:
float
)
->
S
.
spectrum
.
ArraySpectralElement
:
"""Albedo spectrum of planet using BCC model (Batalha et al. 2018),
Parameters
----------
coe_cloud: float
The sedimentation efficiency of the cloud. log(f_sed), range from -2 to 2. 2 is for cloud free case.
coe_metalicity: float
The metallicity of the planet. log(Z/Zsun), range from 0 to 2.
Returns
-------
pysynphot.spectrum.ArrayBandpass
albedo spectrum of the planet
"""
spec
=
AlbedoCat
(
0
,
coe_metalicity
,
coe_cloud
)
spec
.
convert
(
'angstrom'
)
return
spec
def
hybrid_albedo_spectrum
(
coe_b
:
float
,
coe_r
:
float
)
->
S
.
spectrum
.
ArraySpectralElement
:
"""Albedo spectrum of planet using hybrid-jupiter-neptune model (Lacy et al. 2018)
jupiter and neptune spectrum is from Karkoschka’s 1994
Parameters
----------
coe_b: float
coefficient of blue spectrum, 1 for jupiter, 0 for neptune
coe_r: float
coefficient of red spectrum, 1 for jupiter, 0 for neptune
Returns
-------
pysynphot.spectrum.ArrayBandpass
albedo spectrum of the planet
"""
log
.
debug
(
f
"planet hybrid spectrum with
{
coe_b
=
:}
,
{
coe_r
=
:}
"
)
model
=
fits
.
getdata
(
HYBRID_MODEL_FILE
)
spec
=
model
[
1
,
:]
*
coe_r
spec
+=
model
[
2
,
:]
*
coe_b
spec
+=
model
[
3
,
:]
*
(
1
-
coe_r
)
spec
+=
model
[
4
,
:]
*
(
1
-
coe_b
)
albedo
=
S
.
ArrayBandpass
(
wave
=
model
[
0
,
:],
throughput
=
spec
,
waveunits
=
'nm'
)
albedo
.
convert
(
'angstrom'
)
return
albedo
def
extract_target_x_y
(
target
:
dict
,
ra0
:
str
=
None
,
dec0
:
str
=
None
)
->
tuple
:
"""
extract x, y of target from target dict
Parameters
----------
target: dict
target dict. must contain either (ra, dec) or (pangle, spearation)
ra0: str
ra of center star. must be provided if (ra, dec) of target is used
dec0: str
dec of center star. must be provided if (ra, dec) of target is used
Returns
-------
x, y: float
x, y of target in arcsec
Raises
------
ValueError
if (ra, dec) of target is used but (ra, dec) of center star is not provided.
ValueError
one of (ra, dec) or (pangle, spearation) is not provided.
"""
def
_pa2xy
(
p_angle
,
separation
):
p_angle_rad
=
p_angle
*
DEG2RAD
x
=
separation
*
np
.
sin
(
p_angle_rad
)
y
=
separation
*
np
.
cos
(
p_angle_rad
)
log
.
debug
(
f
"(
{
p_angle
=
:}
,
{
separation
=
:}
) => (
{
x
=
:}
,
{
y
=
:}
)"
)
return
x
,
y
if
'pangle'
in
target
.
keys
()
and
'separation'
in
target
.
keys
():
return
_pa2xy
(
target
[
'pangle'
],
target
[
'separation'
])
if
'ra'
not
in
target
.
keys
()
or
'dec'
not
in
target
.
keys
():
raise
ValueError
(
'either (ra, dec) or (pangle, separation) needed in target dict'
)
if
ra0
is
None
or
dec0
is
None
:
raise
ValueError
(
'(ra, dec) of center star must be provided if (ra, dec) of bkstar is used'
)
ra
,
dec
=
target
[
'ra'
],
target
[
'dec'
]
log
.
debug
(
f
"target:
{
ra
=
:}
,
{
dec
=
:}
, center star:
{
ra0
=
:}
,
{
dec0
=
:}
"
)
cstar
=
SkyCoord
(
ra0
,
dec0
)
bkstar
=
SkyCoord
(
ra
,
dec
)
separation
=
cstar
.
separation
(
bkstar
).
arcsec
p_angle
=
cstar
.
position_angle
(
bkstar
).
degree
x
,
y
=
_pa2xy
(
p_angle
,
separation
)
return
x
,
y
def
detect_template_path
(
template
:
str
)
->
str
:
"""Find template file in catalog folder or current folder.
Parameters
----------
template: str
template file name
Returns
-------
str
absolute path of template file
"""
if
os
.
path
.
isfile
(
template
):
return
os
.
path
.
abspath
(
template
)
catalog
=
config
[
'catalog_folder'
]
cat_temp
=
os
.
path
.
join
(
catalog
,
template
)
if
os
.
path
.
isfile
(
cat_temp
):
return
cat_temp
raise
FileExistsError
(
f
'cant find
{
template
}
in ./ or catalog folder'
)
class
TargetOjbect
(
object
):
"""A helper class to generate target spectrum and albedo spectrum
Attributes
----------
x: float
x of target in arcsec
y: fload
y of target in arcsec
ra: str
ra string for center star, such as '15d23m05s'
dec: str
dec string for center star
distance: float
distance of center star in pc
image: 2d np.array
image of the target
spectrum: pysynphot.spectrum.Spectrum
spectrum of the target
"""
def
__init__
(
self
,
info
,
cstar
=
None
):
"""Initialize a target object
Parameters
----------
info: dict
target info, see Example for more details
cstar: TargetOjbect or None
center star object bounded, if None, means the target is the center star itself
center star object is used to calculate the projection x, y of each target according to its ra and dec
also center star's distance is used to calculate seperation of planet
Examples:
--------
cstar = {
'magnitude': 0,
'ra': '120d',
'dec': '40d',
'distance': 10,
'sptype': 'G0III'
'sp_model': 'star'
}
stars = {
'magnitude': 15,
'ra': '120.001d',
'dec': '40.001d',
'sptype': 'F0III',
'sp_model': 'blackbody'
}
planets = {
'radius': 2,
'pangle': 60,
'coe_cloud': 0.3,
'coe_metal': 0.7,
'separation': 0.5,
'phase_angle': 90,
'sp_model': 'hybrid_planet',
'image': 'extend_planet.fits'
}
# planet using input costum albedo spectrum!
# Note that the albedo spectrum is not normalized!
# so the contrast of the planet sould be considered in the input file by user!
# The input file is in pysynphot.FileSpectralElement format.
# See the documentation of pysynphot for details.
planets = {
'pangle': 60,
'separation': 0.5,
'sp_spectrum': 'template_planet',
'template': 'planet_albedo.fits'
}
"""
self
.
sp_model
=
info
.
get
(
'sp_model'
,
'star'
)
if
cstar
is
None
:
self
.
x
,
self
.
y
=
0
,
0
self
.
ra
=
info
[
'ra'
]
self
.
dec
=
info
[
'dec'
]
self
.
distance
=
info
.
get
(
'distance'
,
None
)
else
:
self
.
x
,
self
.
y
=
extract_target_x_y
(
info
,
cstar
.
ra
,
cstar
.
dec
)
self
.
image
=
None
if
'image'
in
info
.
keys
():
self
.
image
=
fits
.
getdata
(
detect_template_path
(
info
[
'image'
]))
if
self
.
sp_model
==
'blackbody'
:
self
.
spectrum
=
star_photlam
(
info
[
'magnitude'
],
info
[
'sptype'
],
is_blackbody
=
True
,
mag_input_band
=
info
.
get
(
'mag_input_band'
,
'f661'
)
)
if
self
.
sp_model
==
'reference'
:
self
.
spectrum
=
standard_spectrum
(
info
[
'magnitude'
])
if
self
.
sp_model
==
'template_star'
:
self
.
spectrum
=
S
.
FileSpectrum
(
detect_template_path
(
info
[
'template'
]))
if
self
.
sp_model
==
'star'
:
self
.
spectrum
=
star_photlam
(
info
[
'magnitude'
],
info
[
'sptype'
],
is_blackbody
=
False
,
mag_input_band
=
info
.
get
(
'mag_input_band'
,
'f661'
)
)
if
self
.
sp_model
in
[
'hybrid_planet'
,
'bcc_planet'
,
'template_planet'
]:
planet
=
info
phase_angle
=
planet
.
get
(
'phase_angle'
,
90
)
sp_model
=
self
.
sp_model
radius
=
planet
.
get
(
'radius'
,
1
)
if
cstar
.
distance
is
None
:
raise
ValueError
(
'distance of center star must be provided if planet is added'
)
if
planet
.
get
(
'contrast'
,
None
)
is
not
None
:
contrast
=
planet
[
'contrast'
]
else
:
contrast
=
planet_contrast
(
self
.
x
*
cstar
.
distance
,
self
.
y
*
cstar
.
distance
,
phase_angle
,
radius
,
)
if
sp_model
==
'hybrid_planet'
:
coe_blue
,
coe_red
=
planet
.
get
(
'coe_b'
,
1
),
planet
.
get
(
'coe_r'
,
1
)
albedo_spect
=
hybrid_albedo_spectrum
(
coe_blue
,
coe_red
)
elif
sp_model
==
'bcc_planet'
:
coe_c
,
coe_m
=
planet
.
get
(
'coe_cloud'
,
2
),
planet
.
get
(
'coe_metal'
,
0
)
albedo_spect
=
bcc_spectrum
(
coe_c
,
coe_m
)
else
:
#sp_model == 'template_planet'
albedo_spect
=
S
.
FileBandpass
(
detect_template_path
(
planet
[
'template'
]))
contrast
=
1
self
.
spectrum
=
cstar
.
spectrum
*
albedo_spect
*
contrast
def
planet_contrast
(
planet_x_au
:
float
,
planet_y_au
:
float
,
phase_angle
:
float
,
radius
:
float
)
->
float
:
"""
calculate the contrast of a planet
Parameters
----------
planet_x_au: float
x position of the planet in au
planet_y_au: float
y position of the planet in au
phase_angle: float
phase angle of the planet in degree
radius: float
radius of the planet in jupiter radius
Returns
-------
contrast: float
contrast of the planet
"""
separation
=
np
.
sqrt
(
planet_x_au
**
2
+
planet_y_au
**
2
)
phase_angle
=
phase_angle
*
DEG2RAD
if
np
.
sin
(
phase_angle
)
<
1e-9
:
raise
ValueError
(
'sin(phase_angle) can not be 0'
)
sep_3d
=
separation
/
np
.
sin
(
phase_angle
)
# Lambert Scattering phase function
# from Madhusudhan and Burrows 2012 equation 33.
phase
=
(
np
.
sin
(
phase_angle
)
+
(
np
.
pi
-
phase_angle
)
*
np
.
cos
(
phase_angle
))
/
np
.
pi
log
.
debug
(
f
'alpha:
{
phase_angle
/
np
.
pi
*
180
}
{
phase
=
}
'
)
contrast
=
(
radius
/
sep_3d
*
R_JUPITER_METER
/
AU_METER
)
**
2
*
phase
return
contrast
def
spectrum_generator
(
targets
:
dict
)
->
list
:
"""
generate the spectrum due to the input target list
Parameters
----------
targets: dict
target dictionary which contains keyword 'cstar' (necessary), 'stars'(optional), 'planets' (optional).
The values are:
- cstar: dict
- center star information. must contain keywords ra, dec, distance, magnitude, sptype
- objects: list of dict, optional, recommended! V2.0 new!
- list of targets. each dict must contain ra, dec, magnitude, sptype
Returns
-------
obj_sp_list: list
list of [x, y, spectrum, image] of each target
"""
cstar
=
targets
[
'cstar'
]
objects
=
targets
.
get
(
'objects'
,
[])
obj_sp_list
=
[]
cstar_obj
=
TargetOjbect
(
cstar
)
obj_sp_list
.
append
([
cstar_obj
.
x
,
cstar_obj
.
y
,
cstar_obj
.
spectrum
,
cstar_obj
.
image
])
for
target
in
objects
:
target_obj
=
TargetOjbect
(
target
,
cstar
=
cstar_obj
)
obj_sp_list
.
append
([
target_obj
.
x
,
target_obj
.
y
,
target_obj
.
spectrum
,
target_obj
.
image
])
return
obj_sp_list
def
target_file_load
(
target
:
Union
[
dict
,
str
])
->
dict
:
"""Generate target dict from file, string or dict.
Parameters
----------
target: Union[dict, str]
target file name, formated string or target dict.
Outputs
--------
target: dict
dictionary of target. Format of input
Note
-------
If target is a string start with *, it will be treated as a formatted string.
e.g. "*5.1/25.3(1.3,1.5)/22.1(2.3,-4.5)" which means a central object
with magnitude 5.1, and two substellar with magnitude 25.3 and 22.1, respectively.
The spectrum of each object is standard refernece spectrum of ABmag.
The first number in the parenthesis is the x position in arcsec, and the second is the y position.
If target is a string without *, it will be treated as file name. And find the file in catalog folder.
The file need to be in yaml format.
And end with .yaml (note .yml not work). If not .yaml will be added.
If target is a dict, it will be returned directly.
If all the above conditions are not met, an empty dict will be returned.
"""
if
isinstance
(
target
,
dict
):
return
target
if
not
target
:
# None or empty string
return
{}
if
isinstance
(
target
,
str
):
#filename or formatted string
target
=
target
.
strip
()
if
not
target
:
return
{}
catalog_folder
=
config
[
'catalog_folder'
]
target_file
=
target
target_file
+=
'.yaml'
if
target_file
[
-
5
:].
lower
()
!=
'.yaml'
else
""
target_name
=
os
.
path
.
basename
(
target_file
)[:
-
5
]
file_search
=
[
target_file
,
os
.
path
.
join
(
catalog_folder
,
target_file
)]
for
file
in
file_search
:
if
os
.
path
.
isfile
(
file
):
with
open
(
file
)
as
fid
:
target
=
yaml
.
load
(
fid
,
Loader
=
yaml
.
FullLoader
)
target
[
'name'
]
=
target_name
return
target
target_str
=
target
if
(
target_str
[
0
]
==
'*'
):
objects
=
target_str
[
1
:].
split
(
'/'
)
cstar_mag
=
float
(
objects
[
0
])
cstar
=
{
'magnitude'
:
cstar_mag
,
'ra'
:
'0d'
,
'dec'
:
'0d'
,
'sp_model'
:
'reference'
,
'distance'
:
10
,
}
stars
=
[]
for
sub_stellar
in
objects
[
1
:]:
float_regex
=
R
"[+-]?\d+(?:\.\d+)?"
match
=
re
.
match
(
rf
"(
{
float_regex
}
)\((
{
float_regex
}
),(
{
float_regex
}
)\)"
,
sub_stellar
)
if
not
match
:
log
.
error
(
f
'Wrong format for sub stellar:
{
sub_stellar
}
, Skip it'
)
continue
mag
=
float
(
match
.
group
(
1
))
x
=
float
(
match
.
group
(
2
))
y
=
float
(
match
.
group
(
3
))
pangle
=
np
.
arctan2
(
x
,
y
)
*
180
/
np
.
pi
separation
=
np
.
sqrt
(
x
**
2
+
y
**
2
)
stars
.
append
({
'magnitude'
:
mag
,
'pangle'
:
pangle
,
'separation'
:
separation
,
'sp_model'
:
'reference'
,
})
target_dict
=
{
'name'
:
target_str
[
1
:],
'cstar'
:
cstar
,
'objects'
:
stars
,
}
return
target_dict
log
.
error
(
f
'Wrong format for target:
{
target
}
, using blank target instead'
)
return
{}
source/csst_cpic_sim/utils.py
0 → 100644
View file @
7c99b572
import
numpy
as
np
import
scipy.ndimage
as
nd
import
logging
import
random
import
matplotlib.pyplot
as
plt
# DO NOT IMPORT CPICIMGSIM MODULES HERE
class
Logger
(
object
):
def
__init__
(
self
,
filename
,
level
=
'INFO'
):
self
.
logger
=
logging
.
getLogger
(
'cpism_log'
)
self
.
logger
.
setLevel
(
logging
.
DEBUG
)
shinfo
=
logging
.
StreamHandler
()
onlyinfo
=
logging
.
Filter
()
onlyinfo
.
filter
=
lambda
record
:
(
record
.
levelno
<
logging
.
WARNING
)
fmtstr
=
'%(message)s'
shinfo
.
setFormatter
(
logging
.
Formatter
(
fmtstr
))
# 设置屏幕上显示的格式
shinfo
.
setLevel
(
logging
.
INFO
)
shinfo
.
addFilter
(
onlyinfo
)
sh
=
logging
.
StreamHandler
()
fmtstr
=
'!%(levelname)s!: %(message)s [%(filename)s - %(funcName)s (line: %(lineno)d)]: '
sh
.
setFormatter
(
logging
.
Formatter
(
fmtstr
))
# 设置屏幕上显示的格式
sh
.
setLevel
(
logging
.
WARNING
)
th
=
logging
.
FileHandler
(
filename
)
# 往文件里写入#指定间隔时间自动生成文件的处理器
fmtstr
=
'%(asctime)s %(filename)s [%(funcName)s] - %(levelname)s: %(message)s'
th
.
setFormatter
(
logging
.
Formatter
(
fmtstr
))
# 设置文件里写入的格式
th
.
setLevel
(
logging
.
__dict__
.
get
(
level
.
upper
()))
self
.
logger
.
addHandler
(
shinfo
)
self
.
logger
.
addHandler
(
sh
)
self
.
logger
.
addHandler
(
th
)
def
random_seed_select
(
seed
=-
1
):
"""
Select a random seed for numpy.random and return it.
"""
if
seed
==
-
1
:
seed
=
random
.
randint
(
0
,
2
**
32
-
1
)
np
.
random
.
seed
(
seed
)
return
seed
def
region_replace
(
background
:
np
.
ndarray
,
front
:
np
.
ndarray
,
shift
:
list
,
bmask
:
float
=
1.0
,
fmask
:
float
=
1.0
,
padded_in
:
bool
=
False
,
padded_out
:
bool
=
False
,
subpix
:
bool
=
False
):
"""
replace a region of the background with the front image.
Parameters
----------
background: np.ndarray
The background image.
front: np.ndarray
The front image.
shift: list
The [x, y] shift of the front image. Unit: pixel.
Relative to the lower-left corner of the background image.
[0, 0] means the lower-left corner of the front image is at the lower-left corner of the background image.
bmask: float
The mask of the background image. Default: 1.0
0.0 means the background image is masked.
1.0 means the background image is fully added.
fmask: float
The mask of the front image. Default: 1.0
0.0 means the front image is masked (not added).
1.0 means the front image is fully added.
padded_in: bool
Whether the input background image is padded. Default: False
In the function, the background image is padded by the size of the front image.
If True, means the background image is padded.
padded_out: bool
Whether the output image is padded. Default: False
In the function, the background image is padded by the size of the front image.
If True, means the output image is padded.
padded_in and padded_out are designed for the case that replace_region fuction is called multiple times.
subpix: bool
Whether the shift is subpixel. Default: False
If True, the shift is subpixel, using scipy.ndimage.shift to shift the front image.
If False, the shift is integer, using numpy slicing to shift the front image.
Returns
-------
np.ndarray
The output image.
shape = background.shape if padded_out = False
shape = background.shape + 2 * front.shape if padded_out = True
"""
int_shift
=
np
.
array
(
shift
).
astype
(
int
)
b_sz
=
np
.
array
(
background
.
shape
)
f_sz
=
np
.
array
(
front
.
shape
)
if
padded_in
:
padded
=
background
b_sz
=
b_sz
-
f_sz
*
2
else
:
padded
=
np
.
pad
(
background
,
((
f_sz
[
0
],
f_sz
[
0
]),
(
f_sz
[
1
],
f_sz
[
1
])))
if
np
.
any
((
int_shift
<
-
b_sz
)
|
(
int_shift
>
b_sz
)):
if
padded_out
:
return
padded
return
background
if
subpix
:
subs
=
np
.
array
(
shift
)
-
int_shift
front
=
nd
.
shift
(
front
,
(
subs
[
0
],
subs
[
1
]))
int_shift
+=
f_sz
roi_y
=
int_shift
[
1
]
roi_x
=
int_shift
[
0
]
padded
[
roi_y
:
roi_y
+
f_sz
[
0
],
roi_x
:
roi_x
+
f_sz
[
1
]]
*=
bmask
padded
[
roi_y
:
roi_y
+
f_sz
[
0
],
roi_x
:
roi_x
+
f_sz
[
1
]]
+=
fmask
*
front
if
padded_out
:
return
padded
return
padded
[
f_sz
[
0
]:
b_sz
[
0
]
+
f_sz
[
0
],
f_sz
[
1
]:
b_sz
[
1
]
+
f_sz
[
1
]]
def
psf_imshow
(
psf
,
vmin
=
1e-8
,
vmax
=
0.1
,
log
=
True
,
region
=
1
):
focal_img
=
psf
.
copy
()
focal_img
=
(
focal_img
-
focal_img
.
min
())
/
(
focal_img
.
max
()
-
focal_img
.
min
())
if
log
:
focal_img
=
np
.
log10
(
focal_img
*
9
+
1
)
plt
.
imshow
(
focal_img
,
origin
=
'lower'
,
cmap
=
'gray'
,
vmin
=
vmin
,
vmax
=
vmax
)
shape
=
psf
.
shape
plt
.
xlim
(
shape
[
1
]
*
(
1
-
region
)
/
2
,
shape
[
1
]
*
(
1
+
region
)
/
2
)
plt
.
ylim
(
shape
[
0
]
*
(
1
-
region
)
/
2
,
shape
[
0
]
*
(
1
+
region
)
/
2
)
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment