compassSupervisor.py

 
 
from shesha.supervisor.genericSupervisor import GenericSupervisor
from shesha.supervisor.components import AtmosCompass, DmCompass, RtcCompass, TargetCompass, TelescopeCompass, WfsCompass, CoronagraphCompass
from shesha.supervisor.optimizers import ModalBasis, Calibration, ModalGains
import numpy as np
import time
 
import shesha.constants as scons
 
from typing import Iterable
 
 
class CompassSupervisor(GenericSupervisor):
    """ This class implements generic supervisor to handle compass simulation
 
    Attributes inherited from GenericSupervisor:
        context : (CarmaContext) : a CarmaContext instance
 
        config : (config) : Parameters structure
 
        is_init : (bool) : Flag equals to True if the supervisor has already been initialized
 
        iter : (int) : Frame counter
 
    Attributes:
        telescope : (TelescopeComponent) : a TelescopeComponent instance
 
        atmos : (AtmosComponent) : An AtmosComponent instance
 
        target : (targetComponent) : A TargetComponent instance
 
        wfs : (WfsComponent) : A WfsComponent instance
 
        dms : (DmComponent) : A DmComponent instance
 
        rtc : (RtcComponent) : A Rtc component instance
 
        cacao : (bool) : CACAO features enabled in the RTC
 
        basis : (ModalBasis) : a ModalBasis instance (optimizer)
 
        calibration : (Calibration) : a Calibration instance (optimizer)
 
        modalgains : (ModalGains) : a ModalGain instance (optimizer) using CLOSE algorithm
 
        close_modal_gains : (list of floats) : list of the previous values of the modal gains
    """
 
    def __init__(self, config, *, cacao: bool = False):
        """ Instantiates a CompassSupervisor object
 
        Args:
            config: (config module) : Configuration module
 
        Kwargs:
            cacao : (bool) : If True, enables CACAO features in RTC (Default is False)
                                      Requires OCTOPUS to be installed
        """
        self.cacaocacao = cacao
        self.teltel = None
        self.atmosatmos = None
        self.targettarget = None
        self.wfswfs = None
        self.dmsdms = None
        self.rtcrtc = None
        self.coronocorono = None
        
        GenericSupervisor.__init__(self, config)
        self.basisbasis = ModalBasis(self.configconfig, self.dmsdms, self.targettarget)
        self.calibrationcalibration = Calibration(self.configconfig, self.teltel, self.atmosatmos, self.dmsdms,
                                       self.targettarget, self.rtcrtc, self.wfswfs)
        if config.p_controllers is not None:
            self.modalgainsmodalgains = ModalGains(self.configconfig, self.rtcrtc)
        self.close_modal_gainsclose_modal_gains = []
 
#     ___                  _      __  __     _   _            _
#    / __|___ _ _  ___ _ _(_)__  |  \/  |___| |_| |_  ___  __| |___
#   | (_ / -_) ' \/ -_) '_| / _| | |\/| / -_)  _| ' \/ _ \/ _` (_-<
#    \___\___|_||_\___|_| |_\__| |_|  |_\___|\__|_||_\___/\__,_/__/
 
    def _init_tel(self):
        """Initialize the telescope component of the supervisor as a TelescopeCompass
        """
        self.teltel = TelescopeCompass(self.contextcontext, self.configconfig)
 
    def _init_atmos(self):
        """Initialize the atmosphere component of the supervisor as a AtmosCompass
        """
        self.atmosatmos = AtmosCompass(self.contextcontext, self.configconfig)
 
    def _init_dms(self):
        """Initialize the DM component of the supervisor as a DmCompass
        """
        self.dmsdms = DmCompass(self.contextcontext, self.configconfig)
 
    def _init_target(self):
        """Initialize the target component of the supervisor as a TargetCompass
        """
        if self.teltel is not None:
            self.targettarget = TargetCompass(self.contextcontext, self.configconfig, self.teltel)
        else:
            raise ValueError("Configuration not loaded or Telescope not initilaized")
 
    def _init_wfs(self):
        """Initialize the wfs component of the supervisor as a WfsCompass
        """
        if self.teltel is not None:
            self.wfswfs = WfsCompass(self.contextcontext, self.configconfig, self.teltel)
        else:
            raise ValueError("Configuration not loaded or Telescope not initilaized")
 
    def _init_rtc(self):
        """Initialize the rtc component of the supervisor as a RtcCompass
        """
        if self.wfswfs is not None:
            self.rtcrtc = RtcCompass(self.contextcontext, self.configconfig, self.teltel, self.wfswfs,
                                  self.dmsdms, self.atmosatmos, cacao=self.cacaocacao)
        else:
            raise ValueError("Configuration not loaded or Telescope not initilaized")
 
    def _init_components(self) -> None:
        """ Initialize all the components
        """
 
        if self.configconfig.p_tel is None or self.configconfig.p_geom is None:
            raise ValueError("Telescope geometry must be defined (p_geom and p_tel)")
        self._init_tel_init_tel()
 
        if self.configconfig.p_atmos is not None:
            self._init_atmos_init_atmos()
        if self.configconfig.p_dms is not None:
            self._init_dms_init_dms()
        if self.configconfig.p_targets is not None:
            self._init_target_init_target()
        if self.configconfig.p_wfss is not None:
            self._init_wfs_init_wfs()
        if self.configconfig.p_controllers is not None or self.configconfig.p_centroiders is not None:
            self._init_rtc_init_rtc()
        if self.configconfig.p_coronos is not None:
            self._init_coronagraph_init_coronagraph()
 
        GenericSupervisor._init_components(self)
 
    def _init_coronagraph(self):
        """ Initialize the coronagraph
        """
        self.coronocorono = CoronagraphCompass()
        for p_corono in self.configconfig.p_coronos:
            self.coronocorono.add_corono(self.contextcontext, p_corono, self.configconfig.p_geom, self.targettarget)
 
    def next(self, *, move_atmos: bool = True, nControl: int = 0,
             tar_trace: Iterable[int] = None, wfs_trace: Iterable[int] = None,
             do_control: bool = True, apply_control: bool = True,
             compute_tar_psf: bool = True, compute_corono: bool=True) -> None:
        """Iterates the AO loop, with optional parameters.
 
        Overload the GenericSupervisor next() method to handle the GEO controller
        specific raytrace order operations
 
        Kwargs:
            move_atmos: (bool): move the atmosphere for this iteration. Default is True
 
            nControl: (int): Controller number to use. Default is 0 (single control configuration)
 
            tar_trace: (List): list of targets to trace. None is equivalent to all (default)
 
            wfs_trace: (List): list of WFS to trace. None is equivalent to all (default)
 
            do_control : (bool) : Performs RTC operations if True (Default)
 
            apply_control: (bool): if True (default), apply control on DMs
 
            compute_tar_psf : (bool) : If True (default), computes the PSF at the end of the iteration
 
            compute_corono: (bool): If True (default), computes the coronagraphic image
        """
        try:
            iter(nControl)
        except TypeError:
            # nControl is not an iterable creating a list
            nControl = [nControl]
 
        #get the index of the first GEO controller (-1 if there is no GEO controller)
        geo_index = next(( i for i,c in enumerate(self.configconfig.p_controllers)
            if c.type== scons.ControllerType.GEO ), -1)
 
        if tar_trace is None and self.targettarget is not None:
            tar_trace = range(len(self.configconfig.p_targets))
        if wfs_trace is None and self.wfswfs is not None:
            wfs_trace = range(len(self.configconfig.p_wfss))
 
        if move_atmos and self.atmosatmos is not None:
            self.atmosatmos.move_atmos()
        # in case there is at least 1 controller GEO in the controller list : use this one only
        self.teltel.update_input_phase()
        if ( geo_index > -1):
            nControl = geo_index
            if tar_trace is not None:
                for t in tar_trace:
                    if self.atmosatmos.is_enable:
                        self.targettarget.raytrace(t, tel=self.teltel, atm=self.atmosatmos, ncpa=False)
                    else:
                        self.targettarget.raytrace(t, tel=self.teltel, ncpa=False)
 
                    if do_control and self.rtcrtc is not None:
                        self.rtcrtc.do_control(nControl, sources=self.targettarget.sources)
                        self.targettarget.raytrace(t, dms=self.dmsdms, ncpa=True, reset=False)
                        if apply_control:
                            self.rtcrtc.apply_control(nControl)
                        if self.cacaocacao:
                            self.rtcrtc.publish()
        else:
            if tar_trace is not None: # already checked at line 213?
                for t in tar_trace:
                    if self.atmosatmos.is_enable:
                        self.targettarget.raytrace(t, tel=self.teltel, atm=self.atmosatmos,
                                             dms=self.dmsdms)
                    else:
                        self.targettarget.raytrace(t, tel=self.teltel, dms=self.dmsdms)
 
            if wfs_trace is not None: # already checked at line 215?
                for w in wfs_trace:
                    if self.atmosatmos.is_enable:
                        self.wfswfs.raytrace(w, tel=self.teltel, atm=self.atmosatmos)
                    else:
                        self.wfswfs.raytrace(w, tel=self.teltel)
 
                    if not self.configconfig.p_wfss[w].open_loop and self.dmsdms is not None:
                        self.wfswfs.raytrace(w, dms=self.dmsdms, ncpa=False, reset=False)
                    self.wfswfs.compute_wfs_image(w)
            if do_control and self.rtcrtc is not None:
                for ncontrol in nControl : # range(len(self.config.p_controllers)):
                    self.rtcrtc.do_centroids(ncontrol)
                    self.rtcrtc.do_control(ncontrol)
                    self.rtcrtc.do_clipping(ncontrol)
 
            if apply_control:
                for ncontrol in nControl :
                    self.rtcrtc.apply_control(ncontrol)
 
            if self.cacaocacao:
                self.rtcrtc.publish()
 
        if compute_tar_psf:
            for tar_index in tar_trace:
                self.targettarget.comp_tar_image(tar_index)
                self.targettarget.comp_strehl(tar_index)
 
        if self.coronocorono is not None and compute_corono:
            for coro_index in range(len(self.configconfig.p_coronos)):
                self.coronocorono.compute_image(coro_index)
 
        if self.configconfig.p_controllers[0].close_opti and (not self.rtcrtc._rtc.d_control[0].open_loop):
            self.modalgainsmodalgains.update_mgains()
            self.close_modal_gainsclose_modal_gains.append(self.modalgainsmodalgains.get_modal_gains())
 
        self.iteriter += 1
 
    def _print_strehl(self, monitoring_freq: int, iters_time: float, total_iters: int, *,
                      tar_index: int = 0):
        """ Print the Strehl ratio SE and LE from a target on the terminal, the estimated remaining time and framerate
 
        Args:
            monitoring_freq : (int) : Number of frames between two prints
 
            iters_time : (float) : time elapsed between two prints
 
            total_iters : (int) : Total number of iterations
 
        Kwargs:
            tar_index : (int) : Index of the target. Default is 0
        """
        framerate = monitoring_freq / iters_time
        strehl = self.targettarget.get_strehl(tar_index)
        etr = (total_iters - self.iteriter) / framerate
        print("%d \t %.3f \t  %.3f\t     %.1f \t %.1f" % (self.iteriter + 1, strehl[0],
                                                          strehl[1], etr, framerate))
 
    def loop(self, number_of_iter: int, *, monitoring_freq: int = 100,
             compute_tar_psf: bool = True, **kwargs):
        """ Perform the AO loop for <number_of_iter> iterations
 
        Args:
            number_of_iter: (int) : Number of iteration that will be done
 
        Kwargs:
            monitoring_freq: (int) : Monitoring frequency [frames]. Default is 100
 
            compute_tar_psf : (bool) : If True (default), computes the PSF at each iteration
                                                 Else, only computes it each <monitoring_freq> frames
        """
        if not compute_tar_psf:
            print("WARNING: Target PSF will be computed (& accumulated) only during monitoring"
                  )
 
        print("----------------------------------------------------")
        print("iter# | S.E. SR | L.E. SR | ETR (s) | Framerate (Hz)")
        print("----------------------------------------------------")
        # self.next(**kwargs)
        t0 = time.time()
        t1 = time.time()
        if number_of_iter == -1:  # Infinite loop
            while (True):
                self.nextnext(compute_tar_psf=compute_tar_psf, **kwargs)
                if ((self.iteriter + 1) % monitoring_freq == 0):
                    if not compute_tar_psf:
                        self.targettarget.comp_tar_image(0)
                        self.targettarget.comp_strehl(0)
                    self._print_strehl_print_strehl(monitoring_freq, time.time() - t1, number_of_iter)
                    t1 = time.time()
 
        for _ in range(number_of_iter):
            self.nextnext(compute_tar_psf=compute_tar_psf, **kwargs)
            if ((self.iteriter + 1) % monitoring_freq == 0):
                if not compute_tar_psf:
                    self.targettarget.comp_tar_image(0)
                    self.targettarget.comp_strehl(0)
                self._print_strehl_print_strehl(monitoring_freq, time.time() - t1, number_of_iter)
                t1 = time.time()
        t1 = time.time()
        print(" loop execution time:", t1 - t0, "  (", number_of_iter, "iterations), ",
              (t1 - t0) / number_of_iter, "(mean)  ", number_of_iter / (t1 - t0), "Hz")
 
    def reset(self):
        """ Reset the simulation to return to its original state
        """
        self.atmosatmos.reset_turbu()
        self.wfswfs.reset_noise()
        for tar_index in range(len(self.configconfig.p_targets)):
            self.targettarget.reset_strehl(tar_index)
        self.dmsdms.reset_dm()
        self.rtcrtc.open_loop()
        self.rtcrtc.close_loop()
 
 
#    ___              _  __ _      __  __     _   _            _
#   / __|_ __  ___ __(_)/ _(_)__  |  \/  |___| |_| |_  ___  __| |___
#   \__ \ '_ \/ -_) _| |  _| / _| | |\/| / -_)  _| ' \/ _ \/ _` (_-<
#   |___/ .__/\___\__|_|_| |_\__| |_|  |_\___|\__|_||_\___/\__,_/__/
#       |_|
 
    def record_ao_circular_buffer(
            self, cb_count: int, sub_sample: int = 1, controller_index: int = 0,
            tar_index: int = 0, see_atmos: bool = True, cube_data_type: str = None,
            cube_data_file_path: str = "", ncpa: int = 0, ncpa_wfs: np.ndarray = None,
            ref_slopes: np.ndarray = None, ditch_strehl: bool = True,
            projection_matrix: np.ndarray = None):
        """ Used to record a synchronized circular buffer AO loop data.
 
        Args:
            cb_count: (int) : the number of iterations to record.
 
            sub_sample:  (int) : sub sampling of the data (default=1, I.e no subsampling)
 
            controller_index:  (int) :
 
            tar_index:  (int) : target number
 
            see_atmos:  (int) : used for the next function to enable or not the Atmos
 
            cube_data_type:   (int) : if  specified ("tarPhase" or "psfse") returns the target phase or short exposure PSF data cube in the output variable
 
            cube_data_file_path:  (int) : if specified it will also save the target phase cube data (full path on the server)
 
            ncpa:  (int) : !!experimental!!!: Used only in the context of PYRWFS + NCPA compensation on the fly (with optical gain)
            defines how many iters the NCPA refslopes are updates with the proper optical gain. Ex: if NCPA=10 refslopes will be updates every 10 iters.
 
            ncpa_wfs:  (int) : the ncpa phase as seen from the wfs array with dims = size of Mpupil
 
            ref_slopes:  (int) : the reference slopes to use.
 
            ditch_strehl:  (int) : resets the long exposure SR computation at the beginning of the Circular buffer (default= True)
 
            projection_matrix : (np.ndarray) : projection matrix on modal basis to compute residual coefficients
 
        Returns:
            slopes:  (int) : the slopes CB
 
            volts:  (int) : the volts applied to the DM(s) CB
 
            ai:  (int) : the modal coefficient of the residual phase projected on the currently used modal Basis
 
            psf_le:  (int) : Long exposure PSF over the <cb_count> iterations (I.e SR is reset at the begining of the CB if ditch_strehl=True)
 
            strehl_se_list:  (int) : The SR short exposure evolution during CB recording
 
            strehl_le_list:  (int) : The SR long exposure evolution during CB recording
 
            g_ncpa_list:  (int) : the gain applied to the NCPA (PYRWFS CASE) if NCPA is set to True
 
            cube_data:  (int) : the tarPhase or psfse cube data (see cube_data_type)
        """
        slopes_data = None
        volts_data = None
        cube_data = None
        ai_data = None
        k = 0
        sthrel_se_list = []
        sthrel_le_list = []
        g_ncpa_list = []
 
        # Resets the target so that the PSF LE is synchro with the data
        # Doesn't reset it if Ditch_strehl == False (used for real time gain computation)
        if ditch_strehl:
            for i in range(len(self.configconfig.p_targets)):
                self.targettarget.reset_strehl(i)
 
        # Starting CB loop...
        for j in range(cb_count):
            print(j, end="\r")
            if (ncpa):
                if (j % ncpa == 0):
                    ncpa_diff = ref_slopes[None, :]
                    ncpa_turbu = self.calibrationcalibration.do_imat_phase(
                            controller_index, -ncpa_wfs[None, :, :], noise=False,
                            with_turbu=True)
                    g_ncpa = float(
                            np.sqrt(
                                    np.dot(ncpa_diff, ncpa_diff.T) / np.dot(
                                            ncpa_turbu, ncpa_turbu.T)))
                    if (g_ncpa > 1e18):
                        g_ncpa = 0
                        print('Warning NCPA ref slopes gain too high!')
                        g_ncpa_list.append(g_ncpa)
                        self.rtcrtc.set_ref_slopes(-ref_slopes * g_ncpa)
                    else:
                        g_ncpa_list.append(g_ncpa)
                        print('NCPA ref slopes gain: %4.3f' % g_ncpa)
                        self.rtcrtc.set_ref_slopes(-ref_slopes / g_ncpa)
 
            self.atmosatmos.enable_atmos(see_atmos)
            self.nextnext()
            for t in range(len(self.configconfig.p_targets)):
                self.targettarget.comp_tar_image(t)
 
            srse, srle, _, _ = self.targettarget.get_strehl(tar_index)
            sthrel_se_list.append(srse)
            sthrel_le_list.append(srle)
            if (j % sub_sample == 0):
                if (projection_matrix is not None):
                    ai_vector = self.calibrationcalibration.compute_modal_residuals(
                            projection_matrix, selected_actus=self.basisbasis.selected_actus)
                    if (ai_data is None):
                        ai_data = np.zeros((len(ai_vector), int(cb_count / sub_sample)))
                    ai_data[:, k] = ai_vector
 
                slopes_vector = self.rtcrtc.get_slopes(controller_index)
                if (slopes_data is None):
                    slopes_data = np.zeros((len(slopes_vector),
                                            int(cb_count / sub_sample)))
                slopes_data[:, k] = slopes_vector
 
                volts_vector = self.rtcrtc.get_command(
                        controller_index)  # get_command or get_voltages ?
                if (volts_data is None):
                    volts_data = np.zeros((len(volts_vector),
                                           int(cb_count / sub_sample)))
                volts_data[:, k] = volts_vector
 
                if (cube_data_type):
                    if (cube_data_type == "tarPhase"):
                        dataArray = self.targettarget.get_tar_phase(tar_index, pupil=True)
                    elif (cube_data_type == "psfse"):
                        dataArray = self.targettarget.get_tar_image(tar_index, expo_type="se")
                    else:
                        raise ValueError("unknown dataData" % cube_data_type)
                    if (cube_data is None):
                        cube_data = np.zeros((*dataArray.shape,
                                              int(cb_count / sub_sample)))
                    cube_data[:, :, k] = dataArray
                k += 1
        if (cube_data_file_path != ""):
            print("Saving tarPhase cube at: ", cube_data_file_path)
            from astropy.io import fits as pf
            pf.writeto(cube_data_file_path, cube_data, overwrite=True)
 
        psf_le = self.targettarget.get_tar_image(tar_index, expo_type="le")
        return slopes_data, volts_data, ai_data, psf_le, sthrel_se_list, sthrel_le_list, g_ncpa_list, cube_data
 
    def export_config(self):
        """
        Extract and convert compass supervisor configuration parameters
        into 2 dictionnaries containing relevant AO parameters
 
        Args:
            root: (object), COMPASS supervisor object to be parsed
 
        Returns :
            2 dictionaries... See F. Vidal :)
        """
        return self.configconfig.export_config()
 
    def get_s_pupil(self):
        """
        Returns the so called S Pupil of COMPASS
 
        Return:
            s_pupil: (np.array) : S Pupil of COMPASS
        """
        return self.configconfig.p_geom.get_spupil()
 
    def get_i_pupil(self):
        """
        Returns the so called I Pupil of COMPASS
 
        Return:
            i_pupil: (np.array) : I Pupil of COMPASS
        """
        return self.configconfig.p_geom.get_ipupil()
 
    def get_m_pupil(self):
        """
        Returns the so called M Pupil of COMPASS
 
        Return:
            m_pupil: (np.array) : M Pupil of COMPASS
        """
        return self.configconfig.p_geom.get_mpupil()