atmosCompass.py

 
from shesha.init.atmos_init import atmos_init
from shesha.constants import CONST
import numpy as np
from typing import List
 
class AtmosCompass(object):
    """ Atmosphere handler for compass simulation
 
    Attributes:
        atmos : (sutraWrap.Atmos) : Sutra atmos instance
 
        is_enable : (bool) : Flag to enable/disable atmophere
 
        context : (carmaContext) : CarmaContext instance
 
        config : (config module) : Parameters configuration structure module
    """
    def __init__(self, context, config):
        """ Initialize an AtmosCompass component for atmosphere related supervision
 
        Parameters:
            context : (carmaContext) : CarmaContext instance
 
            config : (config module) : Parameters configuration structure module
        """
        self.context = context
        self.is_enable = True # Flag to enable/disable atmophere
        self.config = config # Parameters configuration coming from supervisor init
        print("->atmosphere init")
        self.atmos = atmos_init(self.context, self.config.p_atmos, self.config.p_tel,
                        self.config.p_geom, self.config.p_loop.ittime, p_wfss=self.config.p_wfss,
                        p_targets=self.config.p_targets)
 
    
    def enable_atmos(self, enable : bool) -> None:
        """ Set or unset whether atmos is enabled when running loop
 
        Parameters:
            enable : (bool) : True to enable, False to disable
        """
        self.is_enable = enable
 
    def set_r0(self, r0 : float, *, reset_seed : int=-1) -> None:
        """ Change the current r0 for all layers
 
        Parameters:
            r0 : (float) : r0 @ 0.5 µm
 
            reset_seed : (int, optional): if -1 (default), keep same seed and same screen
                                if 0 random seed is applied and refresh screens
                                if (value) set the given seed and refresh screens
        """
        self.atmos.set_r0(r0)
        if reset_seed != -1:
            if reset_seed == 0:
                ilayer = np.random.randint(1e4)
            else:
                ilayer = reset_seed
            for k in range(self.atmos.nscreens):
                self.atmos.set_seed(k, 1234 + ilayer)
                self.atmos.refresh_screen(k)
                ilayer += 1
        self.config.p_atmos.set_r0(r0)
 
    def set_wind(self, screen_index : int, *, windspeed : float = None, winddir : float = None) -> None:
        """ Set new wind information for the given screen
 
        Parameters:
            screen_index : (int) : Atmos screen to change
 
            windspeed : (float) [m/s] : new wind speed of the screen. If None, the wind speed is unchanged
 
            winddir : (float) [deg]: new wind direction of the screen. If None, the wind direction is unchanged
        """
        if windspeed is not None:
            self.config.p_atmos.windspeed[screen_index] = windspeed
        if winddir is not None:
            self.config.p_atmos.winddir[screen_index] = winddir
        
        lin_delta = self.config.p_geom.pupdiam / self.config.p_tel.diam * self.config.p_atmos.windspeed[screen_index] * \
                    np.cos(CONST.DEG2RAD * self.config.p_geom.zenithangle) * self.config.p_loop.ittime
        oldx = self.config.p_atmos._deltax[screen_index]
        oldy = self.config.p_atmos._deltay[screen_index]
        self.config.p_atmos._deltax[screen_index] = lin_delta * np.sin(CONST.DEG2RAD * self.config.p_atmos.winddir[screen_index] + np.pi)
        self.config.p_atmos._deltay[screen_index] = lin_delta * np.cos(CONST.DEG2RAD * self.config.p_atmos.winddir[screen_index] + np.pi)
        self.atmos.d_screens[screen_index].set_deltax(self.config.p_atmos._deltax[screen_index])
        self.atmos.d_screens[screen_index].set_deltay(self.config.p_atmos._deltay[screen_index])
        if(oldx * self.config.p_atmos._deltax[screen_index] < 0): #Sign has changed, must change the stencil
            stencilx = np.array(self.atmos.d_screens[screen_index].d_istencilx)
            n = self.config.p_atmos.dim_screens[screen_index]
            stencilx = (n * n - 1) - stencilx
            self.atmos.d_screens[screen_index].set_istencilx(stencilx)
        if(oldy * self.config.p_atmos._deltay[screen_index] < 0): #Sign has changed, must change the stencil
            stencily = np.array(self.atmos.d_screens[screen_index].d_istencily)
            n = self.config.p_atmos.dim_screens[screen_index]
            stencily = (n * n - 1) - stencily
            self.atmos.d_screens[screen_index].set_istencily(stencily)
 
    def reset_turbu(self) -> None:
        """ Reset the turbulence layers to their original state
        """
        ilayer = 0
        for k in range(self.atmos.nscreens):
            self.atmos.set_seed(k, 1234 + ilayer)
            self.atmos.refresh_screen(k)
            ilayer += 1
 
    def get_atmos_layer(self, indx: int) -> np.ndarray:
        """ Return the selected atmos screen
 
        Parameters:
            indx : (int) : Index of the turbulent layer to return
 
        Return:
            layer : (np.ndarray) : turbulent layer phase screen
        """
        return np.array(self.atmos.d_screens[indx].d_screen)
 
    def move_atmos(self) -> None:
        """ Move the turbulent layers according to wind speed and direction for a single iteration
        """
        self.atmos.move_atmos()