stageSupervisor.py

 
 
from shesha.supervisor.compassSupervisor import CompassSupervisor
from shesha.supervisor.components import AtmosCompass, DmCompass, RtcCompass, TargetCompass, TelescopeCompass, WfsCompass
from shesha.supervisor.optimizers import ModalBasis, Calibration, ModalGains
import numpy as np
import time
 
import shesha.constants as scons
 
from typing import Iterable
 
 
class StageSupervisor(CompassSupervisor):
    """ This class implements a single stage (e.g. first stage, second stage) supervisor 
    to handle compass simulations of cascaded AO. The main supervision will be handled by another 
    supervisor (manager). 
 
    Attributes:
        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
 
        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 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, stack_wfs_image: bool = False,
             do_centroids: bool = True, compute_corono: bool=True) -> None:
        """Iterates the AO loop, with optional parameters, considering it is a single 
        stage and may be called in the middle of WFS frames. 
        Overload the CompassSupervisor next() method to arrange tasks orders and allow cascaded 
        simulation.
 
        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
 
            stack_wfs_image : (bool) : If False (default), the Wfs image is computed as 
                                        usual. Otherwise, a newly computed WFS image is accumulated 
                                        to the previous one.
 
            do_centroids : (bool) : If True (default), the last WFS image is stacked and 
                                    centroids computation is done. WFS image must be reset before 
                                    next loop (in the manager).
 
            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 apply_control:
                        self.rtcrtc.apply_control(nControl)
 
                    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 self.cacaocacao:
                        self.rtcrtc.publish()
                        
        else:
            # start updating the DM shape
            if apply_control:
                for ncontrol in nControl :
                    # command buffer is updated and commands voltages update is applied
                    self.rtcrtc.apply_control(ncontrol)
                    # Note: clipping is always made by apply_control (CBE. 2023.01.27)
 
            # start the propagations
            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)
 
                    if stack_wfs_image:
                        # accumulate image during sub-integration frames
                        wfs_image = self.wfswfs.get_wfs_image(w)
                        self.wfswfs.compute_wfs_image(w)
                        self.wfswfs.set_wfs_image(w, self.wfswfs.get_wfs_image(w) + wfs_image)
                    else:
                        self.wfswfs.compute_wfs_image(w)
                    
            if self.rtcrtc is not None:
                for ncontrol in nControl : # range(len(self.config.p_controllers)):                                                          
                    # modified to allow do_centroids when the WFS exposure is over.
                    # Also useful for calibration. (CBE 2023.01.30)
                    if do_centroids:
                        self.rtcrtc.do_centroids(ncontrol)
 
                    if do_control:
                        self.rtcrtc.do_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 reset(self):
        """ 
        Reset the simulation to return to its original state.
        Overwrites the compassSupervisor reset function, reseting explicitely the WFS image, 
        to force a new integration of the frame.
        """
        self.atmosatmos.reset_turbu()
        self.wfswfs.reset_noise()
        self.wfswfs.reset_image()
        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()