psf_simulation.py

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