optics.py

import os
import yaml

import numpy as np
from .config import cpism_refdata, which_focalplane, S  # S is synphot
from .config import optics_config
from .utils import region_replace
from .io import log


FILTERS = {
    "f565": S.FileBandpass(f"{cpism_refdata}/throughtput/f565_total.fits"),
    "f661": S.FileBandpass(f"{cpism_refdata}/throughtput/f661_total.fits"),
    "f743": S.FileBandpass(f"{cpism_refdata}/throughtput/f743_total.fits"),
    "f883": S.FileBandpass(f"{cpism_refdata}/throughtput/f883_total.fits"),
    "f940": S.FileBandpass(f"{cpism_refdata}/throughtput/f940_total.fits"),
    "f1265": S.FileBandpass(f"{cpism_refdata}/throughtput/f1265_total.fits"),
    "f1425": S.FileBandpass(f"{cpism_refdata}/throughtput/f1425_total.fits"),
    "f1542": S.FileBandpass(f"{cpism_refdata}/throughtput/f1542_total.fits")
}


def filter_throughput(filter_name):
    """
    Totally throughput of 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 = "f661" if filter_name == "default" else filter_name
    if filter_name not in FILTERS.keys():
        log.warning(f"滤光片名称错误({filter_name})，返回默认滤光片(f661)透过率")
        filter_name = "f661"

    return FILTERS[filter_name]


def example_psf_func(band, spectrum, frame_size, error=0.1):
    """
    Example psf generating function.

    Parameters
    -------------
    band: str
        The name of the band.
    spectrum: synphot.Spectrum or synphot.SourceSpectrum
        The spectrum of the target.
    frame_size: int
        The size of the frame.
    error: float
        Phase RMS error.

    Returns
    ---------------
    2D array
        psf image with shape of `frame_size`

    """
    pass


def make_focus_image(
    band: str,
    targets: list,
    psf_function: callable,
    init_shifts: list = [0, 0],
    rotation: float = 0,
    platesize: list = [1024, 1024],
) -> np.ndarray:
    """
    Make the focus image of the targets.

    Parameters
    -----------
    band: str
        The name of the band.
    targets: list
        The list of the targets.
        Each element of the list is a tuple of (x, y, spectrum).

        - x, y: float
            - The position of the target in the focal plane.
        - spectrum: synphot.Spectrum or synphot.SourceSpectrum
            - The spectrum of the target.
    psf_function: callable
        The function to generate the PSF, with same parameters and return as `example_psf_func`.

    init_shifts: list
        The initial shifts of the center targets. Unit: arcsec.
        The default is [0, 0].
    rotation: float
        The rotation of the focal plane. Unit: degree.
        The default is 0 degree.
    platesize: list
        The size of the focal plane. Unit: pixel.
        The default is [1024, 1024].

    Returns
    --------
    np.ndarray
        The focus image of the targets.
        2D array with the shape of platesize.
    """

    config = optics_config[which_focalplane(band)]
    platescale = config["platescale"]

    focal_image = np.zeros(platesize)
    if not targets:
        return focal_image

    def rotate_and_shift(shift):
        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)

    cstar_x, cstar_y, cstar_spectrum = targets[0]
    cstar_shift = rotate_and_shift([cstar_x, cstar_y]) / platescale

    error_value = 0  # nm

    cstar_psf = psf_function(
        band, cstar_spectrum, config["cstar_frame_size"], error=error_value
    )

    platesize = np.array(platesize)[::-1]
    psf_shape = np.array(cstar_psf.shape)[::-1]
    cstar_shift += (platesize - 1) / 2 - (psf_shape - 1) / 2

    focal_image = region_replace(
        focal_image,
        cstar_psf,
        cstar_shift,
        padded_in=False,
        padded_out=False,
        subpix=True,
    )

    for i_target in range(1, len(targets)):
        sub_x, sub_y, sub_spectrum = targets[i_target]
        pdout = False if i_target == len(targets) - 1 else True
        pdin = False if i_target == 1 else True
        log.debug(f"input target {sub_x=:}, {sub_y=:}")
        sub_shift = rotate_and_shift([sub_x, sub_y]) / platescale
        log.debug(f"after rotate and shift {sub_shift=:}")
        sub_psf = psf_function(
            band, sub_spectrum, config["substellar_frame_size"], error=error_value
        )
        psf_shape = np.array(sub_psf.shape)[::-1]
        sub_shift += (platesize - 1) / 2 - (psf_shape - 1) / 2
        log.debug(f"input shift of region_replace: {sub_shift=:}")
        focal_image = region_replace(
            focal_image,
            sub_psf,
            sub_shift,
            padded_in=pdin,
            padded_out=pdout,
            subpix=True,
        )

    return focal_image


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