update doc string

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:

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):

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())
 

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']

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()
+
+
+# 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

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