psfmodel.py 28 KB
Newer Older
FETICK Romain's avatar
FETICK Romain committed
1
2
3
4
5
6
7
8
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 27 17:31:18 2019

@author: rfetick
"""

rfetick's avatar
rfetick committed
9
10
11
12
import numpy as _np
from scipy.optimize import least_squares as _least_squares
from astropy.io import fits as _fits
import time as _time
rfetick's avatar
rfetick committed
13
import numpy.fft as _fft # import fft2, fftshift, ifft2
rfetick's avatar
rfetick committed
14
15
import sys as _sys
from maoppy.utils import binning as _binning
16
#from functools import lru_cache as _lru_cache
FETICK Romain's avatar
FETICK Romain committed
17

18
# Value to compute finite differences
rfetick's avatar
rfetick committed
19
_EPSILON = _np.sqrt(_sys.float_info.epsilon)
20

rfetick's avatar
rfetick committed
21
22
23
#%% FUNCTIONS
def oversample(samp):
    """Find the minimal integer that allows oversampling"""
rfetick's avatar
rfetick committed
24
    k = int(_np.ceil(2.0/samp))
rfetick's avatar
rfetick committed
25
26
27
    return (k*samp,k)

#%% FITTING
FETICK Romain's avatar
FETICK Romain committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def lsq_flux_bck(model, data, weights, background=True, positive_bck=False):
    """Compute the analytical least-square solution for flux and background
    LS = SUM_pixels { weights*(flux*model + bck - data)² }
    
    Parameters
    ----------
    model: numpy.ndarray
    data: numpy.ndarray
    weights: numpy.ndarray
    
    Keywords
    --------
    background: bool
        Activate/inactivate background (activated by default:True)
    positive_bck : bool
        Makes background positive (default:False)
    """
rfetick's avatar
rfetick committed
45
46
47
48
49
    ws = _np.sum(weights)
    mws = _np.sum(model * weights)
    mwds = _np.sum(model * weights * data)
    m2ws = _np.sum(weights * (model ** 2))
    wds = _np.sum(weights * data)
FETICK Romain's avatar
FETICK Romain committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

    if background:
        delta = mws ** 2 - ws * m2ws
        amp = 1. / delta * (mws * wds - ws * mwds)
        bck = 1. / delta * (-m2ws * wds + mws * mwds)
    else:
        amp = mwds / m2ws
        bck = 0.0
        
    if bck<0 and positive_bck: #re-implement above equation
        amp = mwds / m2ws
        bck = 0.0

    return amp, bck

def psffit(psf,Model,x0,weights=None,dxdy=(0.,0.),flux_bck=(True,True),
66
           positive_bck=False,fixed=None,npixfit=None,**kwargs):
FETICK Romain's avatar
FETICK Romain committed
67
68
69
70
71
72
73
74
75
76
77
78
79
    """Fit a PSF with a parametric model solving the least-square problem
       epsilon(x) = SUM_pixel { weights * (amp * Model(x) + bck - psf)² }
    
    Parameters
    ----------
    psf : numpy.ndarray
        The experimental image to be fitted
    Model : class
        The class representing the fitting model
    x0 : tuple, list, numpy.ndarray
        Initial guess for parameters
    weights : numpy.ndarray
        Least-square weighting matrix (same size as `psf`)
FETICK Romain's avatar
FETICK Romain committed
80
        Inverse of the noise's variance
FETICK Romain's avatar
FETICK Romain committed
81
82
83
84
85
86
87
88
89
90
        Default: uniform weighting
    dxdy : tuple of two floats
        Eventual guess on PSF shifting
    flux_bck : tuple of two bool
        Only background can be activate/inactivated
        Flux is always activated (sorry!)
    positive_bck : bool
        Force background to be positive or null
    fixed : numpy.ndarray
        Fix some parameters to their initial value (default: None)
91
92
    npixfit : int (default=None)
        Increased pixel size for improved fitting accuracy
FETICK Romain's avatar
FETICK Romain committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
    **kwargs :
        All keywords used to instantiate your `Model`
    
    Returns
    -------
    out.x : numpy.array
            Parameters at optimum
       .dxdy : tuple of 2 floats
           PSF shift at optimum
       .flux_bck : tuple of two floats
           Estimated image flux and background
       .psf : numpy.ndarray (dim=2)
           Image of the PSF model at optimum
       .success : bool
           Minimization success
       .status : int
           Minimization status (see scipy doc)
       .message : string
           Human readable minimization status
       .active_mask : numpy.array
           Saturated bounds
       .nfev : int
           Number of function evaluations
       .cost : float
           Value of cost function at optimum
    """
119
    
rfetick's avatar
rfetick committed
120
    if weights is None: weights = _np.ones_like(psf)
121
    elif len(psf)!=len(weights): raise ValueError("Keyword `weights` must have same number of elements as `psf`")
122
123
124
    
    # Increase array size for improved fitting accuracy
    if npixfit is not None:
rfetick's avatar
rfetick committed
125
        sx,sy = _np.shape(psf)
126
        if (sx>npixfit) or (sy>npixfit): raise ValueError('npixfit must be greater or equal to both psf dimensions')
rfetick's avatar
rfetick committed
127
128
        psfbig = _np.zeros((npixfit,npixfit))
        wbig = _np.zeros((npixfit,npixfit))
129
130
131
132
133
        psfbig[npixfit//2-sx//2:npixfit//2+sx//2,npixfit//2-sy//2:npixfit//2+sy//2] = psf
        wbig[npixfit//2-sx//2:npixfit//2+sx//2,npixfit//2-sy//2:npixfit//2+sy//2] = weights
        psf = psfbig
        weights = wbig
    
rfetick's avatar
rfetick committed
134
135
    sqW = _np.sqrt(weights)
    Model_inst = Model(_np.shape(psf),**kwargs)
FETICK Romain's avatar
FETICK Romain committed
136
137
138
139
140
    
    class CostClass(object):
        def __init__(self):
            self.iter = 0
        def __call__(self,y):
141
            if (self.iter%3)==0: print("-",end="")
FETICK Romain's avatar
FETICK Romain committed
142
143
144
145
            self.iter += 1
            x, dxdy = mini2input(y)
            m = Model_inst(x,dx=dxdy[0],dy=dxdy[1])
            amp, bck = lsq_flux_bck(m, psf, weights, background=flux_bck[1], positive_bck=positive_bck)
rfetick's avatar
rfetick committed
146
            return _np.reshape(sqW * (amp * m + bck - psf), _np.size(psf))
FETICK Romain's avatar
FETICK Romain committed
147
148
149
150
    
    cost = CostClass()
    
    if fixed is not None:
151
        if len(fixed)!=len(x0): raise ValueError("When defined, `fixed` must be same size as `x0`")
FETICK Romain's avatar
FETICK Romain committed
152
        FREE = [not fixed[i] for i in range(len(fixed))]
rfetick's avatar
rfetick committed
153
        INDFREE = _np.where(FREE)[0]
FETICK Romain's avatar
FETICK Romain committed
154
155
156
    
    def input2mini(x,dxdy):
        # Transform user parameters to minimizer parameters
157
        if fixed is None: xfree = x
rfetick's avatar
rfetick committed
158
159
        else: xfree = _np.take(x,INDFREE)
        return _np.concatenate((xfree,dxdy))
FETICK Romain's avatar
FETICK Romain committed
160
161
162
163
164
165
    
    def mini2input(y):
        # Transform minimizer parameters to user parameters
        if fixed is None:
            xall = y[0:-2]
        else:
rfetick's avatar
rfetick committed
166
            xall = _np.copy(x0)
FETICK Romain's avatar
FETICK Romain committed
167
168
169
170
171
172
            for i in range(len(y)-2):
                xall[INDFREE[i]] = y[i]
        return (xall,y[-2:])
    
    def get_bound(inst):
        b_low = inst.bounds[0]
rfetick's avatar
rfetick committed
173
174
        if fixed is not None: b_low = _np.take(b_low,INDFREE)
        b_low = _np.concatenate((b_low,[-_np.inf,-_np.inf]))
FETICK Romain's avatar
FETICK Romain committed
175
        b_up = inst.bounds[1]
rfetick's avatar
rfetick committed
176
177
        if fixed is not None: b_up = _np.take(b_up,INDFREE)
        b_up = _np.concatenate((b_up,[_np.inf,_np.inf]))
FETICK Romain's avatar
FETICK Romain committed
178
179
        return (b_low,b_up)
    
rfetick's avatar
rfetick committed
180
    result = _least_squares(cost, input2mini(x0,dxdy), bounds=get_bound(Model_inst))
FETICK Romain's avatar
FETICK Romain committed
181
182
183
184
185
186
187
188
189
190
191
192
    
    print("") #finish line of "-"
    
    result.x, result.dxdy = mini2input(result.x) #split output between x and dxdy

    m = Model_inst(result.x,dx=result.dxdy[0],dy=result.dxdy[1])
    amp, bck = lsq_flux_bck(m, psf, weights, background=flux_bck[1], positive_bck=positive_bck)
    
    result.flux_bck = (amp,bck)
    result.psf = m    
    return result

193
#%% BASIC FUNCTIONS
194
def reduced_coord(XY,ax,ay,theta,cx,cy):
rfetick's avatar
rfetick committed
195
196
197
    c = _np.cos(theta)
    s = _np.sin(theta)
    s2 = _np.sin(2.0 * theta)
198
199
200
201
202
203
204
205

    Rxx = (c/ax)**2 + (s/ay)**2
    Ryy = (c/ay)**2 + (s/ax)**2
    Rxy =  s2/ay**2 -  s2/ax**2
    
    u = Rxx*(XY[0]-cx)**2 + Rxy*(XY[0]-cx)*(XY[1]-cy) + Ryy*(XY[1]-cy)**2
    return u    

206
def moffat(XY, param, norm=None, removeInside=0):
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
    """
    Compute a Moffat function on a meshgrid
    moff = Enorm * (1+u)^(-beta)
    with `u` the quadratic coordinates in the shifted and rotated frame
    
    Parameters
    ----------
    XY : numpy.ndarray (dim=2)
        The (X,Y) meshgrid with X = XY[0] and Y = XY[1]
    param : list, tuple, numpy.ndarray (len=6)
        param[0] - Alpha X
        param[1] - Alpha Y
        param[2] - Theta
        param[3] - Beta
        param[4] - Center X
        param[5] - Center Y
        
    Keywords
    --------
rfetick's avatar
rfetick committed
226
    norm : None, _np.inf, float (>0), int (>0)
227
228
229
        Radius for energy normalization
        None      - No energy normalization (maximum=1.0)
                    Enorm = 1.0
rfetick's avatar
rfetick committed
230
        _np.inf    - Total energy normalization (on the whole X-Y plane)
231
232
233
                    Enorm = (beta-1)/(pi*ax*ay)
        float,int - Energy normalization up to the radius defined by this value
                    Enorm = (beta-1)/(pi*ax*ay)*(1-(1+(R**2)/(ax*ay))**(1-beta))    
234
235
236
    
    removeInside: float (default=0)
        Used to remove the central pixel in energy computation
237
238
239
240
241
242
243
244
    """
    if len(param)!=6: raise ValueError("Parameter `param` must contain exactly 6 elements, but input has %u elements"%len(param))
    ax,ay,theta,beta,cx,cy = param
    
    u = reduced_coord(XY,ax,ay,theta,cx,cy)
    
    if norm is None:
        Enorm = 1
245
246
    else: # norm can be float or np.inf
        if (beta<=1) and (norm==_np.inf): raise ValueError("Cannot compute Moffat energy for beta<=1")
247
        if beta==1: raise ValueError("Energy computation for beta=1.0 not implemented yet. Sorry!")
rfetick's avatar
rfetick committed
248
        Enorm = (beta-1) / (_np.pi*ax*ay)
249
250
251
        Fout = (1 - (1 + (norm**2) / (ax*ay))**(1-beta))
        Fin = (1 - (1 + (removeInside**2) / (ax*ay))**(1-beta))
        Enorm = Enorm / (Fout-Fin)
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
    
    return Enorm * (1.0+u)**(-beta)

def gauss(XY,param):
    """
    Compute a Gaussian function on a meshgrid

    Parameters
    ----------
    XY : numpy.ndarray (dim=2)
        The (X,Y) meshgrid with X = XY[0] and Y = XY[1]
    param : list, tuple, numpy.ndarray (len=5)
        param[0] - Sigma X
        param[1] - Sigma Y
        param[2] - Theta
        param[3] - Center X
        param[4] - Center Y  
    """
    if len(param)!=5: raise ValueError("Parameter `param` must contain exactly 5 elements, but input has %u elements"%len(param))
rfetick's avatar
rfetick committed
271
272
    ax = _np.sqrt(2)*param[0]
    ay = _np.sqrt(2)*param[1]
273
    u = reduced_coord(XY,ax,ay,param[2],param[3],param[4])
rfetick's avatar
rfetick committed
274
    return 1.0/(2*_np.pi*param[0]*param[1])*_np.exp(-u)
275

FETICK Romain's avatar
FETICK Romain committed
276
277
278
279
280
281
#%% CLASS PARAMETRIC PSF AND ITS SUBCLASSES
class ParametricPSF(object):
    """Super-class defining parametric PSFs
    Not to be instantiated, only serves as a referent for subclasses
    """
    
282
    def __init__(self,npix):
FETICK Romain's avatar
FETICK Romain committed
283
284
285
        """
        Parameters
        ----------
286
        npix : tuple of two elements
FETICK Romain's avatar
FETICK Romain committed
287
288
            Model X and Y pixel size when called
        """
289
290
        if type(npix)!=tuple: raise TypeError("Argument `npix` must be a tuple")
        self.npix = npix
rfetick's avatar
rfetick committed
291
        self.bounds = (-_np.inf,_np.inf)
FETICK Romain's avatar
FETICK Romain committed
292
293
    
    def __repr__(self):
294
        return "ParametricPSF of size (%u,%u)"%self.npix
FETICK Romain's avatar
FETICK Romain committed
295
296
297
298
299
300
301
    
    def __call__(self,*args,**kwargs):
        raise ValueError("ParametricPSF is not made to be instantiated. Better use its subclasses")
    
    def otf(self,*args,**kwargs):
        """Return the Optical Transfer Function (OTF)"""
        psf = self.__call__(args,kwargs)
rfetick's avatar
rfetick committed
302
        return _fft.fftshift(_fft.fft2(psf))
FETICK Romain's avatar
FETICK Romain committed
303
304
305
    
    def mtf(self,*args,**kwargs):
        """Return the Modulation Transfer Function (MTF)"""
rfetick's avatar
rfetick committed
306
        return _np.abs(self.otf(args,kwargs))
FETICK Romain's avatar
FETICK Romain committed
307
308
309
310
    
    def tofits(self,param,filename,*args,keys=None,keys_comment=None,**kwargs):
        psf = self.__call__(param,*args,**kwargs)
        hdr = self._getfitshdr(param,keys=keys,keys_comment=keys_comment)
rfetick's avatar
rfetick committed
311
        hdu = _fits.PrimaryHDU(psf, hdr)
FETICK Romain's avatar
FETICK Romain committed
312
313
314
        hdu.writeto(filename, overwrite=True)
        
    def _getfitshdr(self,param,keys=None,keys_comment=None):
315
316
        if keys is None: keys = ["PARAM %u"%(i+1) for i in range(len(param))]
        if len(keys)!=len(param): raise ValueError("When defined, `keys` must be same size as `param`")
FETICK Romain's avatar
FETICK Romain committed
317
        if keys_comment is not None:
318
            if len(keys_comment)!=len(param): raise ValueError("When defined, `keys_comment` must be same size as `param`")
rfetick's avatar
rfetick committed
319
        hdr = _fits.Header()
FETICK Romain's avatar
FETICK Romain committed
320
        
FETICK Romain's avatar
FETICK Romain committed
321
        hdr["HIERARCH ORIGIN"] = "MAOPPY automatic header"
rfetick's avatar
rfetick committed
322
        hdr["HIERARCH CREATION"] = (_time.ctime(),"Date of file creation")
FETICK Romain's avatar
FETICK Romain committed
323
324
325
326
327
328
329
330
331
332
333
334
335
        for i in range(len(param)):
            if keys_comment is None:
                hdr["HIERARCH PARAM "+keys[i]] = param[i]
            else:
                hdr["HIERARCH PARAM "+keys[i]] = (param[i],keys_comment[i])
        return hdr


class ConstantPSF(ParametricPSF):
    """Create a constant PSF, given as a 2D image, using ParametricPSF formalism
    With such a formalism, a constant PSF is just a particular case of a parametric PSF
    """
    def __init__(self,image_psf):
rfetick's avatar
rfetick committed
336
        super().__init__(_np.shape(image_psf))
FETICK Romain's avatar
FETICK Romain committed
337
338
339
340
341
342
343
344
        self.image_psf = image_psf
        self.bounds = ()
        
    def __call__(self,*args,**kwargs):
        return self.image_psf


class Moffat(ParametricPSF):
345
346
    def __init__(self,npix,norm=_np.inf):
        self.npix = npix
FETICK Romain's avatar
FETICK Romain committed
347
        self.norm = norm
rfetick's avatar
rfetick committed
348
349
        bounds_down = [_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
        bounds_up = [_np.inf for i in range(4)]
FETICK Romain's avatar
FETICK Romain committed
350
        self.bounds = (bounds_down,bounds_up)
FETICK Romain's avatar
FETICK Romain committed
351
        
352
353
354
        self._XY = _np.mgrid[0:npix[0],0:npix[1]]
        self._XY[0] -= npix[0]//2
        self._XY[1] -= npix[1]//2
FETICK Romain's avatar
FETICK Romain committed
355
356
357
358
359
360
361
362
363
364
365
    
    def __call__(self,x,dx=0,dy=0):
        """
        Parameters
        ----------
        x : list, tuple, numpy.ndarray (len=4)
            x[0] - Alpha X
            x[1] - Alpha Y
            x[2] - Theta
            x[3] - Beta
        """
rfetick's avatar
rfetick committed
366
        y = _np.concatenate((x,[dx,dy]))
FETICK Romain's avatar
FETICK Romain committed
367
        return moffat(self._XY,y,norm=self.norm)
FETICK Romain's avatar
FETICK Romain committed
368
369
370
371
372
    
    def tofits(self,param,filename,*args,keys=["ALPHA_X","ALPHA_Y","THETA","BETA"],**kwargs):
        super(Moffat,self).tofits(param,filename,*args,keys=keys,**kwargs)


FETICK Romain's avatar
FETICK Romain committed
373
class Gaussian(ParametricPSF):
374
375
    def __init__(self,npix):
        self.npix = npix
rfetick's avatar
rfetick committed
376
377
        bounds_down = [_EPSILON,_EPSILON,-_np.inf]
        bounds_up = [_np.inf for i in range(3)]
FETICK Romain's avatar
FETICK Romain committed
378
379
        self.bounds = (bounds_down,bounds_up)
        
380
381
382
        self._XY = _np.mgrid[0:npix[0],0:npix[1]]
        self._XY[1] -= npix[0]/2
        self._XY[0] -= npix[1]/2
FETICK Romain's avatar
FETICK Romain committed
383
384
385
386
387
388
389
390
391
392
    
    def __call__(self,x,dx=0,dy=0):
        """
        Parameters
        ----------
        x : list, tuple, numpy.ndarray (len=4)
            x[0] - Sigma X
            x[1] - Sigma Y
            x[2] - Theta
        """
rfetick's avatar
rfetick committed
393
        y = _np.concatenate((x,[dx,dy]))
FETICK Romain's avatar
FETICK Romain committed
394
395
396
        return gauss(self._XY,y)
    
    def tofits(self,param,filename,*args,keys=["SIGMA_X","SIGMA_Y","THETA"],**kwargs):
FETICK Romain's avatar
FETICK Romain committed
397
        super(Gaussian,self).tofits(param,filename,*args,keys=keys,**kwargs)
398

399
400
401
#%% MASTER CLASS
class ParametricPSFfromPSD(ParametricPSF):
    """This class is NOT to be instantiated"""
402
403
404
405
    def __init__(self,nparam,npix,system=None,samp=None):
        if not (type(npix) in [tuple,list,_np.ndarray]): raise ValueError("npix must be a tuple, list or numpy.ndarray")
        if len(npix)!=2: raise ValueError("npix must be of length = 2")
        if (npix[0]%2) or (npix[1]%2): raise ValueError("Each npix component must be even")
rfetick's avatar
rfetick committed
406
407
        if system is None: raise ValueError("Keyword `system` must be defined")
        if samp is None: raise ValueError("Keyword `samp` must be defined")
rfetick's avatar
rfetick committed
408
409
        self._npix = npix # "_" to bypass the _xyarray update, that will be made with samp setter
        self.samp = samp # also init _xyarray
rfetick's avatar
rfetick committed
410
        self.system = system
411
412
        self._nparam = nparam
        
rfetick's avatar
rfetick committed
413
414
415
416
417
418
419
420
421
    @property
    def npix(self):
        return self._npix
    
    @npix.setter
    def npix(self,value):
        self._npix = value
        self._computeXYarray()
        
rfetick's avatar
rfetick committed
422
423
424
425
426
427
428
429
    @property
    def samp(self):
        return self._samp_over/self._k
    
    @samp.setter
    def samp(self,value):
        # Manage cases of undersampling
        self._samp_over, self._k = oversample(value)
rfetick's avatar
rfetick committed
430
        self._computeXYarray()
rfetick's avatar
rfetick committed
431
        
rfetick's avatar
rfetick committed
432
433
434
435
436
    def _computeXYarray(self):
        nx,ny = self._nxnyOver
        self._xyarray = _np.mgrid[0:nx, 0:ny].astype(float)
        self._xyarray[0] -= nx//2
        self._xyarray[1] -= ny//2
rfetick's avatar
rfetick committed
437
    
438
    #@_lru_cache(maxsize=2)
439
440
    @property
    def _pix2freq(self):
441
442
443
444
445
        """
        pixel to frequency conversion in the PSD plane
        """
        return 1.0/(self.system.D*self._samp_over)
    
446
    #@_lru_cache(maxsize=2)
447
448
    @property
    def _nxnyOver(self):
449
450
451
        """
        Return the number of pixels for the correctly sampled array (at least at Shannon-Nyquist)
        """
452
        return self.npix[0]*self._k, self.npix[1]*self._k
453
    
454
    #@_lru_cache(maxsize=2)
455
456
457
    @property
    def _nullFreqIndex(self):
        nx,ny = self._nxnyOver
458
459
        return nx//2, ny//2
    
460
    #@_lru_cache(maxsize=2)
461
462
    @property
    def _fxfy(self):
463
464
465
        """
        Return the array of frequencies (correctly sampled above Shannon-Nyquist)
        """
rfetick's avatar
rfetick committed
466
        return self._xyarray * self._pix2freq
467
    
rfetick's avatar
rfetick committed
468
    def strehlOTF(self,x0):
469
        """Compute the Strehl ratio based on the sum of the OTF"""
rfetick's avatar
rfetick committed
470
        otf = self.otf(x0,_caller='self')
471
        otf_diff = self._otfDiffraction()
rfetick's avatar
rfetick committed
472
473
474
        return _np.real(_np.sum(otf)/_np.sum(otf_diff))
    
    def strehlMarechal(self,x0):
475
        """Compute the Strehl ratio based on the Maréchal approximation"""
rfetick's avatar
rfetick committed
476
477
478
        _,sig2 = self.psd(x0)
        return _np.exp(-sig2)
    
rfetick's avatar
rfetick committed
479
    def psd(self,x0,grad=False):
480
481
        raise ValueError("ParametricPSFfromPSD is not to be instantiated. the `psd` method must be override in the subclasses")
    
rfetick's avatar
rfetick committed
482
483
484
485
486
487
488
489
490
    def check_parameters(self,x0):
        bd, bu = self.bounds
        x0 = _np.array(x0)
        bd = _np.array(bd)
        bu = _np.array(bu)
        if len(x0)!=self._nparam: raise ValueError('len(x0) is different from length of bounds')
        if _np.any(x0<bd): raise ValueError('Lower bounds are not respected')
        if _np.any(x0>bu): raise ValueError('Upper bounds are not respected')
        
rfetick's avatar
rfetick committed
491
    def otf(self,x0,dx=0,dy=0,_caller='user',grad=False):
rfetick's avatar
rfetick committed
492
493
494
495
496
497
498
499
500
501
        """
        See __call__ for input arguments
        Warning: result of otf will be unconsistent if undersampled!!!
        This issue is solved with oversampling + binning in __call__ but not here
        For the moment, the `_caller` keyword prevents user to misuse otf
        """
        
        if (self._k > 1) and (_caller != 'self'):
            raise ValueError("Cannot call `otf(...)` when undersampled (functionality not implemented yet)")
        
rfetick's avatar
rfetick committed
502
        otfDiffr = self._otfDiffraction()
503
        otfShift = self._otfShift(dx,dy)
rfetick's avatar
rfetick committed
504
505
506
507
508
        if grad:
            otfTurb,g = self._otfTurbulent(x0,grad=True)
            for i in range(g.shape[0]): g[i,:] *= otfDiffr*otfShift
        else:
            otfTurb = self._otfTurbulent(x0,grad=False)
rfetick's avatar
rfetick committed
509
        
rfetick's avatar
rfetick committed
510
511
512
        otf = otfTurb * otfDiffr * otfShift
        if grad: return otf,g
        return otf
rfetick's avatar
rfetick committed
513
    
rfetick's avatar
rfetick committed
514
    def _otfTurbulent(self,x0,grad=False):
rfetick's avatar
rfetick committed
515
        L = self.system.D * self._samp_over
rfetick's avatar
rfetick committed
516
517
518
519
520
        if grad:
            PSD,integral,g = self.psd(x0,grad=True)
        else:
            PSD,integral = self.psd(x0,grad=False)
            
rfetick's avatar
rfetick committed
521
        Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
rfetick's avatar
rfetick committed
522
523
524
525
526
        #Bphi = _fft.fftshift(_np.real(Bg[0, 0] - Bg)) # normalized on the numerical FoV 
        Bphi = _fft.fftshift(_np.real(integral - Bg)) # normalized up to infinity
        otf = _np.exp(-Bphi)
        
        if grad:
rfetick's avatar
rfetick committed
527
            g2 = _np.zeros(g.shape,dtype=complex) # I cannot override 'g' here due to float to complex type
rfetick's avatar
rfetick committed
528
529
            for i in range(len(x0)):
                Bg = _fft.fft2(_fft.fftshift(g[i,...])) / L**2
rfetick's avatar
rfetick committed
530
                # BUG: grad is normalized on the FoV, that is different of the otf!
rfetick's avatar
rfetick committed
531
532
533
534
                Bphi = _fft.fftshift(_np.real(Bg[0, 0] - Bg)) # normalized on the numerical FoV
                g2[i,...] = -Bphi*otf
            return otf, g2
        return otf
rfetick's avatar
rfetick committed
535
    
536
    #@_lru_cache(maxsize=2)
537
    def _otfDiffraction(self):
rfetick's avatar
rfetick committed
538
539
540
541
        nx, ny = self._nxnyOver
        NpupX = _np.ceil(nx/self._samp_over)
        NpupY = _np.ceil(ny/self._samp_over)
        tab = _np.zeros((nx, ny), dtype=_np.complex)
rfetick's avatar
rfetick committed
542
        tab[0:int(NpupX), 0:int(NpupY)] = self.system.pupil((NpupX,NpupY),samp=self._samp_over)
rfetick's avatar
rfetick committed
543
        return _fft.fftshift(_np.abs(_fft.ifft2(_np.abs(_fft.fft2(tab)) ** 2)) / _np.sum(tab))
rfetick's avatar
rfetick committed
544
    
545
    #@_lru_cache(maxsize=2)
546
    def _otfShift(self,dx,dy):
rfetick's avatar
rfetick committed
547
548
549
        Y, X = self._xyarray
        ny,nx = X.shape
        return _np.exp(-2j*_np.pi*self._k*(dx*X/nx + dy*Y/ny))
rfetick's avatar
rfetick committed
550
    
rfetick's avatar
rfetick committed
551
    def __call__(self,x0,dx=0,dy=0,grad=False):
rfetick's avatar
rfetick committed
552
553
554
555
        """
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
556
            Array of parameters for the PSD (see __doc__)
rfetick's avatar
rfetick committed
557
558
559
560
561
562
        dx : float
            PSF X shifting [pix] (default = 0)
        dy : float
            PSF Y shifting [pix] (default = 0)
        """
        
rfetick's avatar
rfetick committed
563
564
565
566
567
568
        if grad:
            otf,g = self.otf(x0,dx=dx,dy=dy,_caller='self',grad=True)
            for i in range(len(x0)):
                g[i,...] = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(g[i,...]))))
        else:
            otf = self.otf(x0,dx=dx,dy=dy,_caller='self',grad=False)
rfetick's avatar
rfetick committed
569
        
rfetick's avatar
rfetick committed
570
571
572
573
574
575
        psf = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(otf))))
            
        k = int(self._k)
        if k==1:
            if grad: return psf,g
            return psf
rfetick's avatar
rfetick committed
576
        else:
rfetick's avatar
rfetick committed
577
578
579
580
581
582
            if grad:
                g2 = _np.zeros((len(x0),psf.shape[0]//k,psf.shape[1]//k))
                for i in range(len(x0)):
                    g2[i,...] = _binning(g[i,...],k)
                return _binning(psf,k), g2
            return _binning(psf,k) # undersample PSF if needed (if it was oversampled for computations)
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597

#%% TURBULENT PSF
class Turbulent(ParametricPSFfromPSD):
    """
    Summary
    -------
    PSF model dedicated to long-exposure imaging with turbulence
    p = Turbulent((npix,npix),system=system,samp=samp)
    
    Description
    -----------
    the PSF is parametrised through the PSD of the electromagnetic phase
        x[0] - Fried parameter r0 [m]
        x[1] - Von Karman external length [m]
    """
598
    def __init__(self,npix,system=None,samp=None):
599
600
601
        """
        Parameters
        ----------
602
        npix : tuple
603
604
605
606
607
608
            Size of output PSF
        system : OpticalSystem
            Optical system for this PSF
        samp : float
            Sampling at the observation wavelength
        """
609
        super().__init__(2,npix,system=system,samp=samp)
610
611
612
613
614
615
        
        # r0,L0
        bounds_down = [_EPSILON, _EPSILON]
        bounds_up = [_np.inf, _np.inf]
        self.bounds = (bounds_down,bounds_up)
    
rfetick's avatar
rfetick committed
616
    def psd(self,x0,grad=False):
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
        """Compute the PSD model from parameters
        PSD is given in [rad²/f²] = [rad² m²]
        
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
            See __doc__ for more details
            
        Returns
        -------
        psd : numpy.array (dim=2)
        integral : float : the integral of the `psd` up to infinity
            
        """
        self.check_parameters(x0)
632
633
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
634
635
636
637
638
639
640
641
642
643
        
        F2 = f2D[0]**2. + f2D[1]**2.
        
        r0,Lext = x0
        PSD = 0.0229* r0**(-5./3.) * ((1./Lext**2.) + F2)**(-11./6.)

        # Set PSD = 0 at null frequency
        PSD[nx0,ny0] = 0.0
        
        # Compute PSD integral up to infinity
644
        fmax = _np.min([nx0,ny0])*self._pix2freq
645
646
        integral_in = _np.sum(PSD*(F2<(fmax**2))) * self._pix2freq**2 # numerical sum (in the array's tangent circle)
        integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum (outside the array's tangent circle)
647
648
        integral = integral_in + integral_out
        
rfetick's avatar
rfetick committed
649
650
        # TODO: I can also compute the integral gradient
        
rfetick's avatar
rfetick committed
651
652
653
654
655
656
        if grad:
            g = _np.zeros((len(x0),)+F2.shape)
            g[0,...] = PSD*(-5./3)/r0
            g[1,...] = PSD*(-11./6)/((1./Lext**2.) + F2)*(-2/Lext**3)
        
        if grad: return PSD,integral,g
657
        return PSD, integral
658

659
#%% PSFAO MODEL        
660
class Psfao(ParametricPSFfromPSD):
FETICK Romain's avatar
FETICK Romain committed
661
    """
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
    Summary
    -------
    PSF model dedicated to long-exposure imaging with adaptive optics
    p = Psfao((npix,npix),system=system,samp=samp)
    
    Description
    -----------
    the PSF is parametrised through the PSD of the electromagnetic phase
        x[0] - Fried parameter           : r0 [m]
        x[1] - Corrected area background : C [rad² m²]
        x[2] - Moffat phase variance     : A [rad²]
        x[3] - Moffat width              : alpha [1/m]
        x[4] - Moffat elongation ratio   : sqrt(ax/ay)
        x[5] - Moffat angle              : theta [rad]
        x[6] - Moffat power law          : beta
        
    Reference
    ---------
    Fétick et al., August 2019, A&A, Vol.628
FETICK Romain's avatar
FETICK Romain committed
681
    """
682
    def __init__(self,npix,system=None,Lext=10.,samp=None):
FETICK Romain's avatar
FETICK Romain committed
683
684
685
        """
        Parameters
        ----------
686
        npix : tuple
FETICK Romain's avatar
FETICK Romain committed
687
688
689
690
691
692
693
694
695
            Size of output PSF
        system : OpticalSystem
            Optical system for this PSF
        samp : float
            Sampling at the observation wavelength
        Lext : float
            Von-Karman external scale (default = 10 m)
            Useless if Fao >> 1/Lext
        """
696
        super().__init__(7,npix,system=system,samp=samp)
697
        
FETICK Romain's avatar
FETICK Romain committed
698
        self.Lext = Lext
rfetick's avatar
rfetick committed
699
        # r0,C,A,alpha,ratio,theta,beta 
rfetick's avatar
rfetick committed
700
701
        bounds_down = [_EPSILON,0,0,_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
        bounds_up = [_np.inf for i in range(7)]
FETICK Romain's avatar
FETICK Romain committed
702
703
        self.bounds = (bounds_down,bounds_up)
    
rfetick's avatar
rfetick committed
704
    def psd(self,x0,grad=False):
FETICK Romain's avatar
FETICK Romain committed
705
706
707
708
709
710
        """Compute the PSD model from parameters
        PSD is given in [rad²/f²] = [rad² m²]
        
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
711
            See __doc__ for more details
rfetick's avatar
rfetick committed
712
713
714
        grad : bool (default=False)
            Return both (psd,integral,gradient) if set to True
            Warning: not finished yet, do not use!
FETICK Romain's avatar
FETICK Romain committed
715
716
717
718
            
        Returns
        -------
        psd : numpy.array (dim=2)
719
        integral : float : the integral of the `psd` up to infinity
FETICK Romain's avatar
FETICK Romain committed
720
721
            
        """
722
        self.check_parameters(x0)
723
724
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
725
        F2 = f2D[0]**2. + f2D[1]**2.
FETICK Romain's avatar
FETICK Romain committed
726
        Fao = self.system.Nact/(2.0*self.system.D)
rfetick's avatar
rfetick committed
727
        maskin = (F2 < Fao**2.)
FETICK Romain's avatar
FETICK Romain committed
728
        
rfetick's avatar
rfetick committed
729
        r0,C,A,alpha,ratio,theta,beta = x0
730
731
        
        PSD = 0.0229* r0**(-5./3.) * ((1. / self.Lext**2.) + F2)**(-11./6.)
rfetick's avatar
rfetick committed
732
        PSD *= (1-maskin)
FETICK Romain's avatar
FETICK Romain committed
733
        
rfetick's avatar
rfetick committed
734
735
        ax = alpha*ratio
        ay = alpha/ratio
736
        param = (ax,ay,theta,beta,0,0)
737
        
738
        removeInside = (1+_np.sqrt(2))/2 * self._pix2freq/2 # remove central pixel in energy computation
rfetick's avatar
rfetick committed
739
        moff = moffat(f2D,param,norm=Fao,removeInside=removeInside) * maskin
740
        moff[nx0,ny0] = 0.0 # Set Moffat PSD = 0 at null frequency
741
        moff = moff / (_np.sum(moff)*self._pix2freq**2)  # normalize moffat numerically to get correct A=sigma² in the AO zone
rfetick's avatar
rfetick committed
742
        # Warning: Moffat numerical normalization generates strehlOTF jump when "_k" is changed
743
        
rfetick's avatar
rfetick committed
744
        PSD += maskin * (C + A*moff) # before I wrote "moffat(f2D,param,norm=Fao)" instead of "moff", so no numerical normalization
745
        PSD[nx0,ny0] = 0.0 # Set PSD = 0 at null frequency
746
747
        
        # Compute PSD integral up to infinity
748
749
        fmax = _np.min([nx0,ny0])*self._pix2freq
        integral_in = _np.sum(PSD*(F2<(fmax**2))) * self._pix2freq**2 # numerical sum
rfetick's avatar
rfetick committed
750
        integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum
751
752
        integral = integral_in + integral_out
        
rfetick's avatar
rfetick committed
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
        if grad:
            g = _np.zeros((len(x0),)+F2.shape)
            # derivative towards r0
            g[0,...] = PSD*(1-maskin)*(-5./3)/r0
            # derivative towards C
            g[1,...] = maskin
            # derivative towards A
            g[2,...] = maskin*moff
            # derivative towards alpha
            # derivative towards ratio
            # derivative towards theta
            # derivative towards beta
            # Remove central freq from all derivative
            g[:,nx0,ny0] = 0
        
rfetick's avatar
rfetick committed
768
769
        #TODO: I can also compute the integral gradient towards the parameters!
        
rfetick's avatar
rfetick committed
770
        if grad: return PSD,integral,g
771
        return PSD, integral
FETICK Romain's avatar
FETICK Romain committed
772
        
rfetick's avatar
rfetick committed
773
    def tofits(self,param,filename,*args,keys=None,overwrite=False,**kwargs):
FETICK Romain's avatar
FETICK Romain committed
774
        if keys is None:
rfetick's avatar
rfetick committed
775
            keys = ["R0","CST","SIGMA2","ALPHA","RATIO","THETA","BETA"]
FETICK Romain's avatar
FETICK Romain committed
776
777
778
            keys_comment = ["Fried parameter [m]",
                            "PSD AO area constant C [rad2]",
                            "PSD AO area Moffat variance A [rad2]",
779
                            "PSD AO area Moffat alpha [1/m]",
rfetick's avatar
rfetick committed
780
                            "PSD AO area Moffat sqrt(ax/ay) ratio",
FETICK Romain's avatar
FETICK Romain committed
781
782
                            "PSD AO area Moffat theta [rad]",
                            "PSD AO area Moffat beta"]
FETICK Romain's avatar
FETICK Romain committed
783
784
785
786
787
788
        
        # redefine tofits() because extra hdr is required
        psf = self.__call__(param,*args,**kwargs)
        hdr = self._getfitshdr(param,keys=keys,keys_comment=keys_comment)
        
        hdr["HIERARCH SYSTEM"] = (self.system.name,"System name")
789
790
        hdr["HIERARCH SYSTEM D"] = (self.system.D,"Primary mirror diameter")
        hdr["HIERARCH SYSTEM NACT"] = (self.system.Nact,"Linear number of AO actuators")
FETICK Romain's avatar
FETICK Romain committed
791
792
793
        hdr["HIERARCH SAMP"] = (self.samp,"Sampling (eg. 2 for Shannon)")
        hdr["HIERARCH LEXT"] = (self.Lext,"Von-Karman outer scale")
        
rfetick's avatar
rfetick committed
794
        hdu = _fits.PrimaryHDU(psf, hdr)
rfetick's avatar
rfetick committed
795
        hdu.writeto(filename, overwrite=overwrite)