COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
wfs_init.py
1 
37 
38 import shesha.config as conf
39 import shesha.constants as scons
40 from shesha.constants import CONST
41 
42 from . import lgs_init as LGS
43 from shesha.sutra_wrap import carmaWrap_context, Sensors, Telescope
44 import numpy as np
45 
46 
47 def wfs_init(context: carmaWrap_context, telescope: Telescope, p_wfss: list,
48  p_tel: conf.Param_tel, p_geom: conf.Param_geom, p_dms=None, p_atmos=None):
49  """
50  Create and initialise a Sensors object
51 
52  :parameters:
53  context : (carmaWrap_context)
54  telescope: (Telescope) : Telescope object
55  p_wfss: (list of Param_wfs) : wfs settings
56  p_tel: (Param_tel) : telescope settings
57  p_geom: (Param_geom) : geom settings
58  p_dms : (list of Param_dm) : (optional) dms settings
59  p_atmos: (Param_atmos) : (optional) atmos settings
60  :return:
61  g_wfs: (Sensors): Sensors object
62  """
63  # create sensor object on gpu
64  # and init sensor gs object on gpu
65  nsensors = len(p_wfss)
66  # arrays needed to call Sensors constructor
67  t_wfs = [
68  'shlo' if (o.type == scons.WFSType.SH and o.is_low_order) else
69  'pyrhr' if o.type == scons.WFSType.PYRLR else o.type for o in p_wfss
70  ]
71 
72  # cdef np.ndarray t_wfs = np.array([o.type for o in
73  # wfs],dtype=np.str)
74  nxsub = np.array([o.nxsub for o in p_wfss], dtype=np.int64)
75  nvalid = np.array([o._nvalid for o in p_wfss], dtype=np.int64)
76  nPupils = np.array([o.nPupils for o in p_wfss], dtype=np.int64)
77  nphase = np.array([o._pdiam for o in p_wfss], dtype=np.int64)
78  pdiam = np.array([o._subapd for o in p_wfss], dtype=np.float32)
79  npix = np.array([o.npix for o in p_wfss], dtype=np.int64)
80  nrebin = np.array([o._nrebin for o in p_wfss], dtype=np.int64)
81  nfft = np.array([o._Nfft for o in p_wfss], dtype=np.int64)
82  ntota = np.array([o._Ntot for o in p_wfss], dtype=np.int64)
83  nphot = np.array([o._nphotons for o in p_wfss], dtype=np.float32)
84  nphot4imat = np.array([o.nphotons4imat for o in p_wfss], dtype=np.float32)
85  lgs = np.array([o.gsalt > 0 for o in p_wfss], dtype=np.int32)
86  fakecam = np.array([o.fakecam for o in p_wfss], dtype=np.bool)
87  maxFlux = np.array([o.max_flux_per_pix for o in p_wfss], dtype=np.int32)
88  max_pix_value = np.array([o.max_pix_value for o in p_wfss], dtype=np.int32)
89 
90  # arrays needed to call initgs
91  xpos = np.array([o.xpos for o in p_wfss], dtype=np.float32)
92  ypos = np.array([o.ypos for o in p_wfss], dtype=np.float32)
93  Lambda = np.array([o.Lambda for o in p_wfss], dtype=np.float32)
94  zerop = p_wfss[0].zerop
95  size = np.zeros(nsensors, dtype=np.int64) + p_geom._n
96  seed = np.arange(nsensors, dtype=np.int64) + 1234
97  npup = (np.zeros((nsensors)) + p_geom._n).astype(np.int64)
98 
99  G = np.array([o.G for o in p_wfss], dtype=np.float32)
100  thetaML = np.array([o.thetaML for o in p_wfss], dtype=np.float32)
101  dx = np.array([o.dx for o in p_wfss], dtype=np.float32)
102  dy = np.array([o.dy for o in p_wfss], dtype=np.float32)
103 
104  roket_flag = any([w.roket for w in p_wfss])
105 
106  if (p_wfss[0].type == scons.WFSType.SH):
107  g_wfs = Sensors(context, telescope, t_wfs, nsensors, nxsub, nvalid, nPupils,
108  npix, nphase, nrebin, nfft, ntota, npup, pdiam, nphot,
109  nphot4imat, lgs, fakecam, maxFlux, max_pix_value,
110  context.active_device, roket_flag)
111 
112  mag = np.array([o.gsmag for o in p_wfss], dtype=np.float32)
113  noise = np.array([o.noise for o in p_wfss], dtype=np.float32)
114 
115  g_wfs.initgs(xpos, ypos, Lambda, mag, zerop, size, noise, seed, G, thetaML, dx,
116  dy)
117 
118  elif (p_wfss[0].type == scons.WFSType.PYRHR or
119  p_wfss[0].type == scons.WFSType.PYRLR):
120  npup = np.array([o.pyr_npts for o in p_wfss])
121  npix = np.array([o._validsubsx.size for o in p_wfss])
122  G = np.array([o.G for o in p_wfss], dtype=np.float32)
123  thetaML = np.array([o.thetaML for o in p_wfss], dtype=np.float32)
124  dx = np.array([o.dx for o in p_wfss], dtype=np.float32)
125  dy = np.array([o.dy for o in p_wfss], dtype=np.float32)
126 
127  g_wfs = Sensors(context, telescope, t_wfs, nsensors, nxsub, nvalid, nPupils,
128  npix, nphase, nrebin, nfft, ntota, npup, pdiam, nphot,
129  nphot4imat, lgs, fakecam, maxFlux, max_pix_value,
130  context.active_device, roket_flag)
131 
132  mag = np.array([o.gsmag for o in p_wfss], dtype=np.float32)
133  noise = np.array([o.noise for o in p_wfss], dtype=np.float32)
134  g_wfs.initgs(xpos, ypos, Lambda, mag, zerop, size, noise, seed, G, thetaML, dx,
135  dy)
136 
137  else:
138  raise Exception("WFS type unknown")
139 
140  # fill sensor object with data
141 
142  for i in range(nsensors):
143  p_wfs = p_wfss[i]
144  wfs = g_wfs.d_wfs[i]
145  fluxPerSub = p_wfs._fluxPerSub.T[np.where(p_wfs._isvalid.T > 0)].copy()
146  if p_wfs.type == scons.WFSType.PYRHR or p_wfs.type == scons.WFSType.PYRLR:
147  halfxy = np.exp(1j * p_wfs._halfxy).astype(np.complex64).T.copy()
148  if (p_wfs._pyr_weights is None):
149  p_wfs.set_pyr_weights(np.ones(p_wfs._pyr_cx.size))
150  wfs.compute_pyrfocalplane = p_wfs.pyr_compute_focalplane
151  wfs.load_arrays(halfxy, p_wfs._pyr_cx, p_wfs._pyr_cy, p_wfs._pyr_weights,
152  p_wfs._sincar, p_wfs._submask, p_wfs._validsubsx,
153  p_wfs._validsubsy, p_wfs._phasemap, fluxPerSub)
154  else:
155  wfs.load_arrays(p_wfs._phasemap, p_wfs._hrmap, p_wfs._binmap, p_wfs._halfxy,
156  fluxPerSub, p_wfs._validsubsx, p_wfs._validsubsy,
157  p_wfs._validpuppixx, p_wfs._validpuppixy, p_wfs._ftkernel)
158 
159  # lgs case
160  for i in range(nsensors):
161  if (p_wfss[i].gsalt > 0):
162  # lgs mode requested
163  # init sensor lgs object with necessary data
164  LGS.prep_lgs_prof(p_wfss[i], i, p_tel, g_wfs)
165 
166  type_target = "atmos" # FIXME
167 
168  for i in range(len(p_wfss)):
169  p_wfs = p_wfss[i]
170  if p_wfs.gsalt > 0:
171  gsalt = 1. / p_wfs.gsalt
172  else:
173  gsalt = 0
174 
175  if p_wfs.atmos_seen is not None and p_atmos is not None:
176  for j in range(p_atmos.nscreens):
177  xoff = (gsalt * p_atmos.alt[j] * p_tel.diam / 2. +
178  p_wfs.xpos * CONST.ARCSEC2RAD * p_atmos.alt[j]) / \
179  p_atmos.pupixsize
180  yoff = (gsalt * p_atmos.alt[j] * p_tel.diam / 2. +
181  p_wfs.ypos * CONST.ARCSEC2RAD * p_atmos.alt[j]) / \
182  p_atmos.pupixsize
183  xoff = xoff + (p_atmos.dim_screens[j] - p_geom._n) / 2.
184  yoff = yoff + (p_atmos.dim_screens[j] - p_geom._n) / 2.
185  g_wfs.d_wfs[i].d_gs.add_layer(type_target, j, xoff, yoff)
186 
187  if (not p_wfs.open_loop and p_dms is not None):
188  if (p_wfs.dms_seen is None):
189  p_wfs.dms_seen = np.arange(len(p_dms)).astype(np.int32)
190  for j in range(p_wfs.dms_seen.size):
191  k = p_wfs.dms_seen[j]
192  dims = p_dms[k]._n2 - p_dms[k]._n1 + 1
193  dim = p_geom._mpupil.shape[0]
194  if (dim < dims):
195  dim = dims
196  xoff = (gsalt * p_dms[k].alt * p_tel.diam / 2. + \
197  p_wfs.xpos * CONST.ARCSEC2RAD * p_dms[k].alt ) * \
198  p_geom.pupdiam / p_tel.diam
199  yoff = (gsalt * p_dms[k].alt * p_tel.diam / 2. + \
200  p_wfs.ypos * CONST.ARCSEC2RAD * p_dms[k].alt ) * \
201  p_geom.pupdiam / p_tel.diam
202  xoff = xoff + (dim - p_geom._n) / 2
203  yoff = yoff + (dim - p_geom._n) / 2
204  g_wfs.d_wfs[i].d_gs.add_layer(p_dms[k].type, k, xoff, yoff)
205 
206  return g_wfs
shesha.sutra_wrap
Definition: sutra_wrap.py:1
shesha.constants
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
shesha.config
Parameter classes for COMPASS.
Definition: shesha/shesha/config/__init__.py:1
shesha.init.wfs_init.wfs_init
def wfs_init(carmaWrap_context context, Telescope telescope, list p_wfss, conf.Param_tel p_tel, conf.Param_geom p_geom, p_dms=None, p_atmos=None)
Create and initialise a Sensors object.
Definition: wfs_init.py:61