COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
lgs_init.py
1 
37 
38 import os
39 try:
40  shesha_db = os.environ['SHESHA_DB_ROOT']
41 except KeyError as err:
42  import warnings
43  shesha_db = os.environ['SHESHA_ROOT'] + "/data/"
44  # warnings.warn("'SHESHA_DB_ROOT' not defined, using default one: " + shesha_db)
45 finally:
46  shesha_savepath = shesha_db
47 # print("shesha_savepath:", shesha_savepath)
48 
49 import shesha.config as conf
50 import shesha.constants as scons
51 from shesha.util import utilities as util
52 import numpy as np
53 
54 from shesha.sutra_wrap import Sensors
55 import scipy.ndimage.interpolation as sci
56 
57 
58 def make_lgs_prof1d(p_wfs: conf.Param_wfs, p_tel: conf.Param_tel, prof: np.ndarray,
59  h: np.ndarray, beam: float, center=""):
60  """same as prep_lgs_prof but cpu only. original routine from rico
61 
62  :parameters:
63  p_tel: (Param_tel) : telescope settings
64 
65  prof: (np.ndarray[dtype=np.float32]) : Na profile intensity, in arbitrary units
66 
67  h: (np.ndarray[dtype=np.float32]) : altitude, in meters. h MUST be an array with EQUALLY spaced elements.
68 
69  beam: (float) : size in arcsec of the laser beam
70 
71  center: (string) : either "image" or "fourier" depending on where the centre should be.
72  """
73 
74  p_wfs._prof1d = prof
75  p_wfs._profcum = np.zeros(h.shape[0] + 1, dtype=np.float32)
76  p_wfs._profcum[1:] = prof.cumsum()
77 
78  subapdiam = p_tel.diam / p_wfs.nxsub # diam of subap
79  if (p_wfs.nxsub > 1):
80  xsubs = np.linspace((subapdiam - p_tel.diam) / 2, (p_tel.diam - subapdiam) / 2,
81  p_wfs.nxsub).astype(np.float32)
82  else:
83  xsubs = np.zeros(1, dtype=np.float32)
84  ysubs = xsubs.copy().astype(np.float32)
85 
86  # cdef int nP=prof.shape[0] #UNUSED
87  hG = np.sum(h * prof) / np.sum(prof)
88  x = np.arange(p_wfs._Ntot).astype(np.float32) - p_wfs._Ntot / 2
89  # x expressed in pixels. (0,0) is in the fourier-center.
90  x = x * p_wfs._qpixsize # x expressed in arcseconds
91  # cdef float dx=x[1]-x[0] #UNUSED
92  # cdef float dh=h[1]-h[0] #UNUSED
93 
94  if (p_wfs.nxsub > 1):
95  dOffAxis = np.sqrt((xsubs[p_wfs._validsubsy // p_wfs.npix] - p_wfs.lltx)**2 +
96  (ysubs[p_wfs._validsubsx // p_wfs.npix] - p_wfs.llty)**2)
97  else:
98  dOffAxis = np.sqrt((xsubs - p_wfs.lltx)**2 + (ysubs - p_wfs.llty)**2)
99 
100  profi = np.zeros((p_wfs._Ntot, p_wfs._nvalid), dtype=np.float32)
101 
102  subsdone = np.ones(p_wfs._nvalid, dtype=np.int32)
103  dif2do = np.zeros(p_wfs._nvalid, dtype=np.int32)
104 
105  while (np.any(subsdone)):
106  tmp = dOffAxis[np.where(subsdone)][0]
107  inds = np.where(dOffAxis == tmp)[0]
108  # height, translated in arcsec due to perspective effect
109  zhc = (h - hG) * (206265. * tmp / hG**2)
110  dzhc = zhc[1] - zhc[0]
111 
112  if (p_wfs._qpixsize > dzhc):
113  avg_zhc = np.zeros(zhc.size + 1, dtype=np.float32)
114  avg_zhc[0] = zhc[0]
115  avg_zhc[avg_zhc.size - 1] = zhc[zhc.size - 1]
116  avg_zhc[1:-1] = 0.5 * (zhc[1:] + zhc[:-1])
117  avg_x = np.zeros(x.size + 1, dtype=np.float32)
118  avg_x[0] = x[0]
119  avg_x[avg_x.size - 1] = x[x.size - 1]
120  avg_x[1:-1] = 0.5 * (x[1:] + x[:-1])
121 
122  for i in range(inds.size):
123  profi[:, inds[i]] = np.diff(np.interp(avg_x, avg_zhc,
124  p_wfs._profcum)).astype(np.float32)
125 
126  else:
127  for i in range(inds.size):
128  profi[:, inds[i]] = np.interp(x, zhc, prof)
129  subsdone[inds] = 0
130 
131  w = beam / 2.35482005
132  if (w == 0):
133  # TODO what is n
134  n = 1
135  g = np.zeros(n, dtype=np.float32)
136  if (center == "image"):
137  g[n / 2 - 1] = 0.5
138  g[n / 2] = 0.5
139  else:
140  g[n / 2] = 1
141 
142  else:
143  if (center == "image"):
144  if ((p_wfs.npix * p_wfs._nrebin) % 2 != p_wfs._Nfft % 2):
145  g = np.exp(-(x + p_wfs._qpixsize)**2 / (2 * w**2.))
146  else:
147  g = np.exp(-(x + p_wfs._qpixsize / 2)**2 / (2 * w**2.))
148 
149  else:
150  g = np.exp(-x**2 / (2 * w**2.))
151 
152  p_wfs._ftbeam = np.fft.fft(g).astype(np.complex64)
153  p_wfs._beam = g.astype(np.float32)
154  # convolved profile in 1D.
155 
156  g_extended = np.tile(g, (p_wfs._nvalid, 1)).T
157 
158  p1d = np.fft.ifft(
159  np.fft.fft(profi, axis=0) * np.fft.fft(g_extended, axis=0),
160  axis=0).real.astype(np.float32)
161  p1d = p1d * p1d.shape[0]
162  p1d = np.roll(p1d, int(p_wfs._Ntot / 2. - 0.5), axis=0)
163  p1d = np.abs(p1d)
164 
165  im = np.zeros((p1d.shape[1], p1d.shape[0], p1d.shape[0]), dtype=np.float32)
166  for i in range(p1d.shape[0]):
167  im[:, i, :] = g[i] * p1d.T
168 
169  if (ysubs.size > 1):
170  azimuth = np.arctan2(ysubs[p_wfs._validsubsy // p_wfs.npix] - p_wfs.llty,
171  xsubs[p_wfs._validsubsx // p_wfs.npix] - p_wfs.lltx)
172  else:
173  azimuth = np.arctan2(ysubs - p_wfs.llty, xsubs - p_wfs.lltx)
174 
175  p_wfs._azimuth = azimuth
176 
177  if (center == "image"):
178  xcent = p_wfs._Ntot / 2. - 0.5
179  ycent = xcent
180  else:
181  xcent = p_wfs._Ntot / 2. #+ 1
182  ycent = xcent
183 
184  if (ysubs.size > 0):
185  # TODO rotate
186  # im = util.rotate3d(im, azimuth * 180 / np.pi, xcent, ycent) --> ça marche pas !!
187  # max_im = np.max(im, axis=(1, 2))
188  # im = (im.T / max_im).T
189  for k in range(im.shape[0]):
190  img = im[k, :, :] / im[k, :, :].max()
191  im[k, :, :] = sci.rotate(img, azimuth[k] * 180 / np.pi, reshape=False)
192 
193  else:
194  # im = util.rotate(im, azimuth * 180 / np.pi, xcent, ycent)
195  # im = im / np.max(im)
196  im = sci.rotate(img, azimuth * 180 / np.pi, reshape=False)
197 
198  p_wfs._lgskern = im.T
199 
200 
201 def prep_lgs_prof(p_wfs: conf.Param_wfs, nsensors: int, p_tel: conf.Param_tel,
202  sensors: Sensors, center="", imat=0):
203  """The function returns an image array(double,n,n) of a laser beacon elongated by perpective
204  effect. It is obtaind by convolution of a gaussian of width "lgsWidth" arcseconds, with the
205  line of the sodium profile "prof". The altitude of the profile is the array "h".
206 
207  :parameters:
208  p_wfs: (Param_wfs) : WFS settings
209 
210  nsensors: (int) : wfs index
211 
212  p_tel: (Param_tel) : telescope settings
213 
214  Sensors: (Sensors) : WFS object
215 
216  center: (string) : either "image" or "fourier" depending on where the centre should be.
217 
218  Computation of LGS spot from the sodium profile:
219  Everything is done here in 1D, because the Na profile is the result of the convolution of a function
220  P(x,y) = profile(x) . dirac(y)
221  by a gaussian function, for which variables x and y can be split :
222  exp(-(x^2+y^2)/2.s^2) = exp(-x^2/2.s^2) * exp(-y^2/2.s^2)
223  The convolution is (symbol $ denotes integral)
224  C(X,Y) = $$ exp(-x^2/2.s^2) * exp(-y^2/2.s^2) * profile(x-X) * dirac(y-Y) dx dy
225  First one performs the integration along y
226  C(X,Y) = exp(-Y^2/2.s^2) $ exp(-x^2/2.s^2) * profile(x-X) dx
227  which shows that the profile can be computed by
228  - convolving the 1-D profile
229  - multiplying it in the 2nd dimension by a gaussian function
230 
231  If one has to undersample the inital profile, then some structures may be "lost". In this case,
232  it's better to try to "save" those structures by re-sampling the integral of the profile, and
233  then derivating it afterwards.
234  Now, if the initial profile is a coarse one, and that one has to oversample it, then a
235  simple re-sampling of the profile is adequate.
236  """
237  if (p_wfs.proftype is None or p_wfs.proftype == ""):
238  p_wfs.set_proftype(scons.ProfType.GAUSS1)
239 
240  profilename = scons.ProfType.FILES[p_wfs.proftype]
241 
242  profile_path = shesha_savepath + profilename
243  print("reading Na profile from", profile_path)
244  prof = np.load(profile_path)
245  make_lgs_prof1d(p_wfs, p_tel, np.mean(prof[1:, :], axis=0), prof[0, :],
246  p_wfs.beamsize, center="image")
247  p_wfs.set_altna(prof[0, :].astype(np.float32))
248  p_wfs.set_profna(np.mean(prof[1:, :], axis=0).astype(np.float32))
249 
250  p_wfs._prof1d = p_wfs._profna
251  p_wfs._profcum = np.zeros(p_wfs._profna.size + 1, dtype=np.float32)
252  p_wfs._profcum[1:] = p_wfs._profna.cumsum()
253  subapdiam = p_tel.diam / p_wfs.nxsub # diam of subap
254 
255  if (p_wfs.nxsub > 1):
256  xsubs = np.linspace((subapdiam - p_tel.diam) / 2, (p_tel.diam - subapdiam) / 2,
257  p_wfs.nxsub).astype(np.float32)
258  else:
259  xsubs = np.zeros(1, dtype=np.float32)
260  ysubs = xsubs.copy().astype(np.float32)
261 
262  # center of gravity of the profile
263  hG = np.sum(p_wfs._altna * p_wfs._profna) / np.sum(p_wfs._profna)
264  x = np.arange(p_wfs._Ntot).astype(np.float32) - p_wfs._Ntot / 2
265  # x expressed in pixels. (0,0) is in the fourier-center
266 
267  x = x * p_wfs._qpixsize # x expressed in arcseconds
268  # cdef float dx=x[1]-x[0] #UNUSED
269  dh = p_wfs._altna[1] - p_wfs._altna[0]
270 
271  if (p_wfs.nxsub > 1):
272  dOffAxis = np.sqrt((xsubs[p_wfs._validsubsx // p_wfs.npix] - p_wfs.lltx)**2 +
273  (ysubs[p_wfs._validsubsy // p_wfs.npix] - p_wfs.llty)**2)
274  else:
275  dOffAxis = np.sqrt((xsubs - p_wfs.lltx)**2 + (ysubs - p_wfs.llty)**2)
276 
277  if (imat > 0):
278  dOffAxis *= 0.
279 
280  w = p_wfs.beamsize / 2.35482005 # TODO: FIXME
281  if (w == 0):
282  # TODO what is n
283  n = 1
284  g = np.zeros(n, dtype=np.float32)
285  if (center == "image"):
286  g[n / 2 - 1] = 0.5
287  g[n / 2] = 0.5
288  else:
289  g[n / 2] = 1
290 
291  else:
292  if (center == "image"):
293  g = np.exp(-(x + p_wfs._qpixsize / 2)**2 / (2 * w**2.))
294  else:
295  g = np.exp(-x**2 / (2 * w**2.))
296 
297  p_wfs._ftbeam = np.fft.fft(g, axis=0).astype(np.complex64)
298  p_wfs._beam = g
299  # convolved profile in 1D.
300 
301  if (xsubs.size > 1):
302  azimuth = np.arctan2(ysubs[p_wfs._validsubsy // p_wfs.npix] - p_wfs.llty,
303  xsubs[p_wfs._validsubsx // p_wfs.npix] - p_wfs.lltx)
304  else:
305  azimuth = np.arctan2(ysubs - p_wfs.llty, xsubs - p_wfs.lltx)
306 
307  p_wfs._azimuth = azimuth
308 
309  sensors.d_wfs[nsensors].d_gs.d_lgs.lgs_init(
310  p_wfs._prof1d.size, hG, p_wfs._altna[0], dh, p_wfs._qpixsize, dOffAxis,
311  p_wfs._prof1d, p_wfs._profcum, p_wfs._beam, p_wfs._ftbeam, p_wfs._azimuth)
shesha.init.lgs_init.prep_lgs_prof
def prep_lgs_prof(conf.Param_wfs p_wfs, int nsensors, conf.Param_tel p_tel, Sensors sensors, center="", imat=0)
The function returns an image array(double,n,n) of a laser beacon elongated by perpective effect.
Definition: lgs_init.py:235
shesha.util
Utilities functions.
Definition: shesha/shesha/util/__init__.py:1
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.lgs_init.make_lgs_prof1d
def make_lgs_prof1d(conf.Param_wfs p_wfs, conf.Param_tel p_tel, np.ndarray prof, np.ndarray h, float beam, center="")
same as prep_lgs_prof but cpu only.
Definition: lgs_init.py:71