COMPASS  5.4.4
End-to-end AO simulation tool using GPU acceleration
atmos_init.py
1 
37 
38 import shesha.config as conf
39 from shesha.constants import CONST
40 import shesha.util.iterkolmo as itK
41 import shesha.util.hdf5_util as h5u
42 from shesha.sutra_wrap import carmaWrap_context, Atmos
43 from rich.progress import track
44 import numpy as np
45 
46 
47 def atmos_init(context: carmaWrap_context, p_atmos: conf.Param_atmos,
48  p_tel: conf.Param_tel, p_geom: conf.Param_geom, ittime=None, p_wfss=None,
49  p_targets=None, dataBase={}, use_DB=False):
50  """
51  Initializes an Atmos object
52 
53  Args:
54  context: (carmaWrap_context): GPU device context
55 
56  p_atmos: (Param_atmos): Atmosphere parameters
57 
58  p_tel: (Param_tel): Telescope parameters
59 
60  p_geom: (Param_geom): Geometry parameters
61 
62  Kwargs:
63  ittime: (float): exposition time [s]
64 
65  p_wfss: (list of Param_wfs): WFS parameters
66 
67  p_targets: (list of Param_target): target parameters
68 
69  dataBase: (dict): dictionary for data base
70 
71  use_DB: (bool): flag for using the dataBase system
72 
73  :return:
74  atm : (Atmos): Atmos object
75  """
76  if not p_geom.is_init:
77  raise RuntimeError("Cannot init atmosphere with uninitialized p_geom.")
78 
79  # Deleted carmaWrap_context : get the singleton
80 
81  if p_atmos.r0 is None:
82  p_atmos.r0 = 0.
83 
84  if ittime is None:
85  ittime = 1.
86  # Adjust layers alt using zenith angle
87  p_atmos.alt = p_atmos.alt / np.cos(p_geom.zenithangle * CONST.DEG2RAD)
88  # Pixel size in meters
89  p_atmos.pupixsize = p_tel.diam / p_geom.pupdiam
90 
91  # Off-axis wavefront sensors and targets
92  # Note : p_wfss is a list of single-WFS configs
93  # but p_target groups several targets
94  # hence different xpos, ypos syntaxes
95  norms = [0.]
96  if p_wfss is not None:
97  norms += [(w.xpos**2 + w.ypos**2)**0.5 for w in p_wfss]
98  if p_targets is not None:
99  norms += [(p_target.xpos**2 + p_target.ypos**2)**0.5 for p_target in p_targets]
100 
101  max_size = max(norms)
102 
103  # Meta-pupil diameter for all layers depending on altitude
104  patch_diam = (p_geom._n + 2 * (max_size * CONST.ARCSEC2RAD * p_atmos.alt) /
105  p_atmos.pupixsize + 4).astype(np.int64)
106  p_atmos.dim_screens = (patch_diam + patch_diam % 2)
107 
108  # Phase screen speeds
109  lin_delta = p_geom.pupdiam / p_tel.diam * p_atmos.windspeed * \
110  np.cos(CONST.DEG2RAD * p_geom.zenithangle) * ittime
111  p_atmos._deltax = lin_delta * np.sin(CONST.DEG2RAD * p_atmos.winddir + np.pi)
112  p_atmos._deltay = lin_delta * np.cos(CONST.DEG2RAD * p_atmos.winddir + np.pi)
113 
114  # Fraction normalization
115  p_atmos.frac /= np.sum(p_atmos.frac)
116 
117  if p_atmos.L0 is None:
118  # Set almost infinite L0
119  p_atmos.L0 = np.ones(p_atmos.nscreens, dtype=np.float32) * 1e5
120  L0_pix = p_atmos.L0 * p_geom.pupdiam / p_tel.diam
121 
122  if p_atmos.seeds is None:
123  p_atmos.seeds = np.arange(p_atmos.nscreens, dtype=np.int64) + 1234
124 
125  r0_layers = p_atmos.r0 / (p_atmos.frac**(3. / 5.) * p_atmos.pupixsize)
126  stencil_size = itK.stencil_size_array(p_atmos.dim_screens)
127 
128  atm = Atmos(context, p_atmos.nscreens, p_atmos.r0, r0_layers, p_atmos.dim_screens,
129  stencil_size, p_atmos.alt, p_atmos.windspeed, p_atmos.winddir,
130  p_atmos._deltax, p_atmos._deltay, context.active_device)
131 
132  print("Creating turbulent layers :")
133  for i in track(range(p_atmos.nscreens)):
134  if "A" in dataBase:
135  A, B, istx, isty = h5u.load_AB_from_dataBase(dataBase, i)
136  else:
137  A, B, istx, isty = itK.AB(p_atmos.dim_screens[i], L0_pix[i],
138  p_atmos._deltax[i], p_atmos._deltay[i], 0)
139  if use_DB:
140  h5u.save_AB_in_database(i, A, B, istx, isty)
141 
142  atm.init_screen(i, A, B, istx, isty, p_atmos.seeds[i])
143 
144  return atm
Parameter classes for COMPASS.
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
def atmos_init(carmaWrap_context context, conf.Param_atmos p_atmos, conf.Param_tel p_tel, conf.Param_geom p_geom, ittime=None, p_wfss=None, p_targets=None, dataBase={}, use_DB=False)
Initializes an Atmos object.
Definition: atmos_init.py:75
Functions for handling the database system.
Definition: hdf5_util.py:1
Stencil and matrices computation for the creation of a turbulent screen.
Definition: iterkolmo.py:1