Upload New File

CpicImgSim/main.py

+import numpy as np
+import re
+
+from .target import spectrum_generator
+from .optics import make_focus_image, focal_mask, optics_config
+from .psf_simulation import simulate_psf
+from .camera import EMCCD, CosmicRayFrameMaker, sky_frame_maker
+from .io import save_fits, log
+from .config import which_focalplane
+
+
+def psf_function(band, cstar_spectrum, shape, error=0.1):
+    cstar = True
+    if shape < 300:
+        cstar = False
+    return simulate_psf(error, band, cstar_spectrum, nsample=1, cstar=cstar)
+
+
+def observation_simulation(
+        target: dict,
+        skybg: float,
+        expt: float,
+        nframe: int,
+        band: str,
+        emgain: float,
+        obsid: int = 51900000000,
+        rotation: float = 0,
+        shift: list = [0, 0],
+        gnc_info: dict = {},
+        csst_format: bool = True,
+        psf_function: callable = psf_function):
+    """
+    Simulate the observation. All-In-One function of the package.
+
+    Parameters
+    -----------
+    target: dict
+        The target information. See target.py for details.
+    skybg: float
+        magnitude of the skybackground at the input b and. (abmag system)
+    expt: float
+        exposure time in second.
+    nframe: int
+        number of frames to be simulated.
+    band: str
+        the band of the observation. (e.g. 'f661')
+    emgain: float
+        the EM gain of the camera.
+    obsid: int
+        the observation id. Default is 51900000000.
+    rotation: float
+        the rotation angle of the target in degree. Default is 0.
+    shift: list
+        the shift of the target in arcsec. Default is [0, 0].
+    gnc_info: dict
+        the gnc information. Default is {}. See io.py for details.
+    csst_format: bool
+        whether to save the fits file in CSST format. Default is True.
+    psf_function: callable
+        the function to generate the psf. See optics.py for details.
+
+    Returns
+    -----------
+    np.ndarray of the simulated images with shape (nframe, 1088, 1050).
+
+    """
+    target_list = []
+    if 'cstar' in target.keys():
+        target_list = spectrum_generator(target)
+
+    focal_name = which_focalplane(band)
+    this_focal_config = optics_config[focal_name]
+    telescope_config = optics_config['telescope']
+    area = telescope_config['aperature_area']
+
+    if focal_name == 'vis':
+        camera = EMCCD()
+    else:
+        raise ValueError('Only VIS focal plane is supported.')
+
+    platescale = this_focal_config['platescale']
+    iwa = this_focal_config['mask_width'] / 2
+
+    crmaker = CosmicRayFrameMaker()
+    images = []
+
+    params = {
+        'target': target,
+        'skybg': skybg,
+        'expt': expt,
+        'nframe': nframe,
+        'band': band,
+        'emgain': emgain,
+        'obsid': obsid,
+        'rotation': rotation,
+        'shift': shift,
+    }
+    paramstr = ', '.join([f'{k}={v}' for k, v in params.items()])
+    log.debug(f"parameters: {paramstr}")
+
+    for i in range(nframe):
+
+        log.info(f'Simulation Running: Frame {i+1}/{nframe}')
+
+        focal_frame = make_focus_image(
+            band,
+            target_list,
+            psf_function,
+            platesize=camera.flat_shape,
+            rotation=rotation,
+            init_shifts=shift,
+        )
+
+        if skybg is None or skybg > 100:
+            sky_bkg_frame = 0
+        else:
+            sky_bkg_frame = sky_frame_maker(
+                band,
+                skybg,
+                platescale,
+                camera.flat_shape
+            )
+
+        focal_frame = (focal_frame + sky_bkg_frame) * expt * area
+        focal_frame = focal_mask(focal_frame, iwa, platescale)
+
+        cr_frame = crmaker.make_cr_frame(camera.dark_shape, expt)
+
+        img = camera.readout(
+            focal_frame,
+            emgain,
+            expt,
+            image_cosmic_ray=cr_frame
+        )
+
+        images.append(img)
+
+    images = np.array(images)
+
+    save_fits(images, params, gnc_info, csst_format=csst_format)
+
+    return images
+
+
+def quick_run(
+        target_str: str,
+        skymag: float,
+        band: str,
+        expt: float,
+        nframe: int,
+        emgain: float,
+        rotation: float = 0,
+        shift: list = [0, 0]) -> np.ndarray:
+    """
+    A quick run function to simulate the observation.
+
+    Parameters
+    -----------
+    target_str: str
+        The target information in string format.
+        In the format of "\*5.1/25.3(1.3,1.5)/22.1(2.3,-4.5)" which means a central star
+        with magnitude 5.1, and two substellar with magnitude 25.3 and 22.1, respectively.
+        The first number in the parenthesis is the x position in arcsec, and the second is the y position.
+    skybg: float
+        magnitude of the skybackground at the input band. (abmag system)
+    band: str
+        the band of the observation. (e.g. 'f661')
+    expt: float
+        exposure time in second.
+    nframe: int
+        number of frames to be simulated.
+    emgain: float
+        the EM gain of the camera.
+    rotation: float (optional)
+        the rotation angle of the target in degree. Default is 0.
+    shift: list (optional)
+        the shift of the target in arcsec. Default is [0, 0].
+
+    Returns
+    -----------
+    np.ndarray of the simulated images, with shape (nframe, 1088, 1050)
+
+    Notes
+    -----------
+    1. stars are assumed to be G0III with distance 10pc.
+    2. magnitude of the star and substellar are assumed to be in the same band.
+    3. Csst format is not supported.
+    4. The psf is assumed to be the default one.
+    5. fits file will be saved in the current directory.
+
+
+    """
+    log.info(f'Quick Run: {target_str}')
+    target_dict = {
+        'name': 'cal',
+    }
+
+    if (target_str != '') and (target_str[0] == '*'):
+        objects = target_str[1:].split('/')
+        cstar_mag = float(objects[0])
+
+        cstar = {
+            'magnitude': cstar_mag,
+            'ra': '0d',
+            'dec': '0d',
+            'sptype': 'G0III',
+            'distance': 10,
+            'mag_input_band': band
+        }
+
+        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:
+                raise ValueError('Wrong format for sub stellar.')
+            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,
+                'sptype': 'G0III',
+                'mag_input_band': band
+            })
+        target_dict = {
+            'name': target_str[1:],
+            'cstar': cstar,
+            'stars': stars,
+        }
+
+    return observation_simulation(
+        target=target_dict,
+        skybg=skymag,
+        expt=expt,
+        nframe=nframe,
+        band=band,
+        emgain=emgain,
+        csst_format=False,
+        shift=shift,
+        rotation=rotation,
+    )
+
+
+# observation_simulation(
+#     target={},
+#     skybg=15,
+#     expt=10,
+#     nframe=2,
+#     band='f661',
+#     emgain=30,
+#     obsid=50112345678,
+# )
+
+# quick_run('*5.1/25.3(0.8,0.8)', None, 'f661', 10, 1, 10)
+# quick_run('*5/20(0.8,0.8)', None, 'f883', 10, 1, 10)
+
+# # quick *5.1/25.3(1.3,1.5) expt nframe emgain band rotation shift
+# # quick target_name expt nframe emgain band rotation shift
+# # plan plan_file_or_folder
+if __name__ == '__main__':  # pragma: no cover
+    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,
+#     )