diff --git a/csst_cpic_sim/camera.py b/csst_cpic_sim/camera.py index 185875969a49558b87d1f0f7961e40ebd82a5538..7e91b17960bd8524c687e6b6016937a6e0a3d5ed 100644 --- a/csst_cpic_sim/camera.py +++ b/csst_cpic_sim/camera.py @@ -278,7 +278,26 @@ class CosmicRayFrameMaker(object): 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 @@ -288,6 +307,7 @@ class CpicVisEmccd(object): self.config_init() def _defaut_config(self): + """set up defaut config for the camera.""" self.plszx = 1024 self.plszy = 1024 self.pscan1 = 8 @@ -359,6 +379,10 @@ class CpicVisEmccd(object): 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) @@ -409,6 +433,19 @@ class CpicVisEmccd(object): 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): @@ -430,6 +467,16 @@ class CpicVisEmccd(object): 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 @@ -574,6 +621,20 @@ class CpicVisEmccd(object): 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 @@ -599,6 +660,18 @@ class CpicVisEmccd(object): # 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 @@ -689,6 +762,27 @@ class CpicVisEmccd(object): # 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: diff --git a/csst_cpic_sim/config.py b/csst_cpic_sim/config.py index 77c29f825e3ab5b978eaf7248abdce4738d643c8..ca822e8e8e2e693944d9689a68ad51d412b4acfd 100644 --- a/csst_cpic_sim/config.py +++ b/csst_cpic_sim/config.py @@ -30,6 +30,16 @@ config_aim = os.path.join(config_aim, 'data/refdata_path.yaml') 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) @@ -84,7 +94,22 @@ 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, output='$'): +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 != '$': @@ -111,6 +136,18 @@ with warnings.catch_warnings(): # pragma: no cover 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" @@ -123,12 +160,13 @@ 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. - from ['f565', 'f661', 'f743', 'f883', 'f940', 'f1265', 'f1425', 'f1542', 'wfs'] + Returns -------- @@ -152,6 +190,20 @@ def which_focalplane(band): # 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 @@ -162,6 +214,19 @@ def iso_time(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): diff --git a/csst_cpic_sim/main.py b/csst_cpic_sim/main.py index 7e0f613a588bdcda7d0abe019c56cd6fa45482d8..f8006854bf1c7e8258811798da8048ac95527fe5 100644 --- a/csst_cpic_sim/main.py +++ b/csst_cpic_sim/main.py @@ -31,7 +31,48 @@ def vis_observation( 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 @@ -148,6 +189,41 @@ def quick_run_v2( 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}') @@ -209,12 +285,25 @@ def quick_run_v2( 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: @@ -317,6 +406,15 @@ def observation_simulation_from_config(obs_file, config_file): 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()) diff --git a/csst_cpic_sim/optics.py b/csst_cpic_sim/optics.py index 8b03d3c42e416c22079a4c99111aae98df4cabdd..26d7456616cc9d8003ff9e4e29fe80e433ce912b 100644 --- a/csst_cpic_sim/optics.py +++ b/csst_cpic_sim/optics.py @@ -41,7 +41,7 @@ def filter_throughput(filter_name): return FILTERS[filter_name] -def rotate_and_shift(shift, rotation, init_shifts): +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), @@ -55,7 +55,30 @@ def ideal_focus_image( platescale, platesize: list = [1024, 1024], init_shifts: list = [0, 0], - rotation: float = 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 @@ -65,7 +88,7 @@ def ideal_focus_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 + sub_shift = _rotate_and_shift([sub_x, sub_y], rotation, init_shifts) / platescale sed = (sub_spectrum * bandpass).integrate() if sub_image is None: @@ -111,6 +134,27 @@ def focal_convolve( 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'] diff --git a/csst_cpic_sim/psf_simulation.py b/csst_cpic_sim/psf_simulation.py index 7c36fe40e8f8497317ec4f053a1c462fbcfd3b71..b8164214f217c3362af44cced78e0b949ae1cbce 100644 --- a/csst_cpic_sim/psf_simulation.py +++ b/csst_cpic_sim/psf_simulation.py @@ -73,7 +73,27 @@ aberration = SurfaceApodizer( aberration_distance = 80 * focal_length aberration = SurfaceAberrationAtDistance(aberration, aberration_distance) -def single_band_masked_psf(wavelength, error=0, shift=[0, 0]): +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) @@ -89,15 +109,43 @@ def single_band_masked_psf(wavelength, error=0, shift=[0, 0]): psf = second_focal.intensity.shaped return psf / strength -def single_band_psf(wavelength, error=0, aber_phase=None, intensity=True): +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) - 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 + 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 diff --git a/csst_cpic_sim/target.py b/csst_cpic_sim/target.py index 78e765f46b1e2411631e0fc3251f5915816767af..b1c44b86d67c9ad01b09c4139e50cb6e763eb0c2 100644 --- a/csst_cpic_sim/target.py +++ b/csst_cpic_sim/target.py @@ -255,8 +255,11 @@ def star_photlam( star_sp.convert('photlam') return star_sp -def standard_spectrum(magnitude: float): +def standard_spectrum( + magnitude: float) -> S.spectrum.Spectrum: """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 ----------- @@ -282,8 +285,7 @@ def standard_spectrum(magnitude: float): def bcc_spectrum( coe_cloud: float, - coe_metalicity: float -): + coe_metalicity: float) -> S.spectrum.ArraySpectralElement: """Albedo spectrum of planet using BCC model (Batalha et al. 2018), Parameters @@ -306,7 +308,7 @@ def bcc_spectrum( def hybrid_albedo_spectrum( coe_b: float, - coe_r: float): + coe_r: float) -> S.spectrum.ArrayBandpass: """Albedo spectrum of planet using hybrid-jupiter-neptune model (Lacy et al. 2018) jupiter and neptune spectrum is from Karkoschka’s 1994 @@ -341,7 +343,7 @@ def hybrid_albedo_spectrum( def extract_target_x_y( target: dict, ra0: str = None, - dec0: str = None): + dec0: str = None) -> tuple: """ extract x, y of target from target dict @@ -454,19 +456,8 @@ class TargetOjbect(object): 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 x, y of each target according to its ra and dec + 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 - sp_model: str - keyword to compatible with V1.0 input. See Notes - spectral type of the target, can be 'star' or 'planet' - - Notes - ----- - keyargv 中的sp_model 完全用来用来兼容V1.0版的输入,执行逻辑如下: - 1. 如果没有设置sp_model, 对应了V2.0的输入, 则sp_model从info中获取 - 2. 如果设置了sp_model, 则区分为star和planet - 3. 如果sp_model为star, 则判断是否设置了isblackbody, 如果是的话, 则sp_model改为blackbody - 4. 如果sp_model为planet, 则进一步区分为hybrid_planet 和 bcc_planet 和 template_planet Examples: -------- @@ -477,6 +468,7 @@ class TargetOjbect(object): 'dec': '40d', 'distance': 10, 'sptype': 'G0III' + 'sp_model': 'star' } stars = { @@ -484,7 +476,7 @@ class TargetOjbect(object): 'ra': '120.001d', 'dec': '40.001d', 'sptype': 'F0III', - 'sp_spectrum': 'blackbody', + 'sp_model': 'blackbody' } planets = { @@ -494,7 +486,7 @@ class TargetOjbect(object): 'coe_metal': 0.7, 'separation': 0.5, 'phase_angle': 90, - 'sp_spectrum': 'hybrid_planet', + 'sp_model': 'hybrid_planet', 'image': 'extend_planet.fits' } @@ -581,7 +573,7 @@ def planet_contrast( planet_x_au: float, planet_y_au: float, phase_angle: float, - radius: float): + radius: float) -> float: """ calculate the contrast of a planet @@ -619,7 +611,7 @@ def planet_contrast( return contrast def spectrum_generator( - targets: dict): + targets: dict) -> list: """ generate the spectrum due to the input target list