psfmodel.py 26.1 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
#%% FUNCTIONS
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def oversample(samp, fixed_k = None):
    """
        Find the minimal integer that allows oversampling

    Args:
        samp (float): input sampling
        fixed_k (int, optional): Oversampling factor to be fixed. Defaults to None.

    Returns:
        (k*samp, k): the oversampling (>=2) and the corresponding oversampling factor
    """
    if fixed_k == None: 
        k = int(_np.ceil(2.0/samp))
    else: 
        k = fixed_k
rfetick's avatar
rfetick committed
37
38
39
    return (k*samp,k)

#%% FITTING
FETICK Romain's avatar
FETICK Romain committed
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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
57
58
59
60
61
    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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

    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),
78
           positive_bck=False,fixed=None,npixfit=None,**kwargs):
FETICK Romain's avatar
FETICK Romain committed
79
80
81
82
83
84
85
86
87
88
89
90
91
    """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
92
        Inverse of the noise's variance
FETICK Romain's avatar
FETICK Romain committed
93
94
95
96
97
98
99
100
101
102
        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)
103
104
    npixfit : int (default=None)
        Increased pixel size for improved fitting accuracy
FETICK Romain's avatar
FETICK Romain committed
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
    **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
    """
131
    
rfetick's avatar
rfetick committed
132
    if weights is None: weights = _np.ones_like(psf)
133
    elif len(psf)!=len(weights): raise ValueError("Keyword `weights` must have same number of elements as `psf`")
134
135
136
    
    # Increase array size for improved fitting accuracy
    if npixfit is not None:
rfetick's avatar
rfetick committed
137
        sx,sy = _np.shape(psf)
138
        if (sx>npixfit) or (sy>npixfit): raise ValueError('npixfit must be greater or equal to both psf dimensions')
rfetick's avatar
rfetick committed
139
140
        psfbig = _np.zeros((npixfit,npixfit))
        wbig = _np.zeros((npixfit,npixfit))
141
142
143
144
145
        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
146
147
    sqW = _np.sqrt(weights)
    Model_inst = Model(_np.shape(psf),**kwargs)
FETICK Romain's avatar
FETICK Romain committed
148
149
150
151
152
    
    class CostClass(object):
        def __init__(self):
            self.iter = 0
        def __call__(self,y):
153
            if (self.iter%3)==0: print("-",end="")
FETICK Romain's avatar
FETICK Romain committed
154
155
156
157
            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
158
            return _np.reshape(sqW * (amp * m + bck - psf), _np.size(psf))
FETICK Romain's avatar
FETICK Romain committed
159
160
161
162
    
    cost = CostClass()
    
    if fixed is not None:
163
        if len(fixed)!=len(x0): raise ValueError("When defined, `fixed` must be same size as `x0`")
FETICK Romain's avatar
FETICK Romain committed
164
        FREE = [not fixed[i] for i in range(len(fixed))]
rfetick's avatar
rfetick committed
165
        INDFREE = _np.where(FREE)[0]
FETICK Romain's avatar
FETICK Romain committed
166
167
168
    
    def input2mini(x,dxdy):
        # Transform user parameters to minimizer parameters
169
        if fixed is None: xfree = x
rfetick's avatar
rfetick committed
170
171
        else: xfree = _np.take(x,INDFREE)
        return _np.concatenate((xfree,dxdy))
FETICK Romain's avatar
FETICK Romain committed
172
173
174
175
176
177
    
    def mini2input(y):
        # Transform minimizer parameters to user parameters
        if fixed is None:
            xall = y[0:-2]
        else:
rfetick's avatar
rfetick committed
178
            xall = _np.copy(x0)
FETICK Romain's avatar
FETICK Romain committed
179
180
181
182
183
184
            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
185
186
        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
187
        b_up = inst.bounds[1]
rfetick's avatar
rfetick committed
188
189
        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
190
191
        return (b_low,b_up)
    
rfetick's avatar
rfetick committed
192
    result = _least_squares(cost, input2mini(x0,dxdy), bounds=get_bound(Model_inst))
FETICK Romain's avatar
FETICK Romain committed
193
194
195
196
197
198
199
200
201
202
203
204
    
    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

205
#%% BASIC FUNCTIONS
206
def reduced_coord(XY,ax,ay,theta,cx,cy):
rfetick's avatar
rfetick committed
207
208
209
    c = _np.cos(theta)
    s = _np.sin(theta)
    s2 = _np.sin(2.0 * theta)
210
211
212
213
214
215
216
217

    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    

218
def moffat(XY, param, norm=None, removeInside=0):
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
    """
    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
238
    norm : None, _np.inf, float (>0), int (>0)
239
240
241
        Radius for energy normalization
        None      - No energy normalization (maximum=1.0)
                    Enorm = 1.0
rfetick's avatar
rfetick committed
242
        _np.inf    - Total energy normalization (on the whole X-Y plane)
243
244
245
                    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))    
246
247
248
    
    removeInside: float (default=0)
        Used to remove the central pixel in energy computation
249
250
251
252
253
254
255
256
    """
    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
257
258
    else: # norm can be float or np.inf
        if (beta<=1) and (norm==_np.inf): raise ValueError("Cannot compute Moffat energy for beta<=1")
259
        if beta==1: raise ValueError("Energy computation for beta=1.0 not implemented yet. Sorry!")
rfetick's avatar
rfetick committed
260
        Enorm = (beta-1) / (_np.pi*ax*ay)
261
262
263
        Fout = (1 - (1 + (norm**2) / (ax*ay))**(1-beta))
        Fin = (1 - (1 + (removeInside**2) / (ax*ay))**(1-beta))
        Enorm = Enorm / (Fout-Fin)
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
    
    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
283
284
    ax = _np.sqrt(2)*param[0]
    ay = _np.sqrt(2)*param[1]
285
    u = reduced_coord(XY,ax,ay,param[2],param[3],param[4])
rfetick's avatar
rfetick committed
286
    return 1.0/(2*_np.pi*param[0]*param[1])*_np.exp(-u)
287

FETICK Romain's avatar
FETICK Romain committed
288
289
290
291
292
293
#%% 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
    """
    
294
    def __init__(self,npix):
FETICK Romain's avatar
FETICK Romain committed
295
296
297
        """
        Parameters
        ----------
298
        npix : tuple of two elements
FETICK Romain's avatar
FETICK Romain committed
299
300
            Model X and Y pixel size when called
        """
301
302
        if type(npix)!=tuple: raise TypeError("Argument `npix` must be a tuple")
        self.npix = npix
rfetick's avatar
rfetick committed
303
        self.bounds = (-_np.inf,_np.inf)
FETICK Romain's avatar
FETICK Romain committed
304
305
    
    def __repr__(self):
306
        return "ParametricPSF of size (%u,%u)"%self.npix
FETICK Romain's avatar
FETICK Romain committed
307
308
309
310
311
312
313
    
    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
314
        return _fft.fftshift(_fft.fft2(psf))
FETICK Romain's avatar
FETICK Romain committed
315
316
317
    
    def mtf(self,*args,**kwargs):
        """Return the Modulation Transfer Function (MTF)"""
rfetick's avatar
rfetick committed
318
        return _np.abs(self.otf(args,kwargs))
FETICK Romain's avatar
FETICK Romain committed
319
320
321
322
    
    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
323
        hdu = _fits.PrimaryHDU(psf, hdr)
FETICK Romain's avatar
FETICK Romain committed
324
325
326
        hdu.writeto(filename, overwrite=True)
        
    def _getfitshdr(self,param,keys=None,keys_comment=None):
327
328
        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
329
        if keys_comment is not None:
330
            if len(keys_comment)!=len(param): raise ValueError("When defined, `keys_comment` must be same size as `param`")
rfetick's avatar
rfetick committed
331
        hdr = _fits.Header()
FETICK Romain's avatar
FETICK Romain committed
332
        
FETICK Romain's avatar
FETICK Romain committed
333
        hdr["HIERARCH ORIGIN"] = "MAOPPY automatic header"
rfetick's avatar
rfetick committed
334
        hdr["HIERARCH CREATION"] = (_time.ctime(),"Date of file creation")
FETICK Romain's avatar
FETICK Romain committed
335
336
337
338
339
340
341
342
343
344
345
346
347
        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
348
        super().__init__(_np.shape(image_psf))
FETICK Romain's avatar
FETICK Romain committed
349
350
351
352
353
354
355
356
        self.image_psf = image_psf
        self.bounds = ()
        
    def __call__(self,*args,**kwargs):
        return self.image_psf


class Moffat(ParametricPSF):
357
358
    def __init__(self,npix,norm=_np.inf):
        self.npix = npix
FETICK Romain's avatar
FETICK Romain committed
359
        self.norm = norm
rfetick's avatar
rfetick committed
360
361
        bounds_down = [_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
        bounds_up = [_np.inf for i in range(4)]
FETICK Romain's avatar
FETICK Romain committed
362
        self.bounds = (bounds_down,bounds_up)
FETICK Romain's avatar
FETICK Romain committed
363
        
364
365
366
        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
367
368
369
370
371
372
373
374
375
376
377
    
    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
378
        y = _np.concatenate((x,[dx,dy]))
FETICK Romain's avatar
FETICK Romain committed
379
        return moffat(self._XY,y,norm=self.norm)
FETICK Romain's avatar
FETICK Romain committed
380
381
382
383
384
    
    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
385
class Gaussian(ParametricPSF):
386
387
    def __init__(self,npix):
        self.npix = npix
rfetick's avatar
rfetick committed
388
389
        bounds_down = [_EPSILON,_EPSILON,-_np.inf]
        bounds_up = [_np.inf for i in range(3)]
FETICK Romain's avatar
FETICK Romain committed
390
391
        self.bounds = (bounds_down,bounds_up)
        
392
393
394
        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
395
396
397
398
399
400
401
402
403
404
    
    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
405
        y = _np.concatenate((x,[dx,dy]))
FETICK Romain's avatar
FETICK Romain committed
406
407
408
        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
409
        super(Gaussian,self).tofits(param,filename,*args,keys=keys,**kwargs)
410

411
412
413
#%% MASTER CLASS
class ParametricPSFfromPSD(ParametricPSF):
    """This class is NOT to be instantiated"""
414
    def __init__(self,nparam,npix,system=None,samp=None, fixed_k = None):
415
416
417
        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
418
419
        if system is None: raise ValueError("Keyword `system` must be defined")
        if samp is None: raise ValueError("Keyword `samp` must be defined")
420
421
        
        self.fixed_k = fixed_k        
rfetick's avatar
rfetick committed
422
423
        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
424
        self.system = system
425
426
        self._nparam = nparam
        
427
        
rfetick's avatar
rfetick committed
428
429
430
431
432
433
434
435
436
    @property
    def npix(self):
        return self._npix
    
    @npix.setter
    def npix(self,value):
        self._npix = value
        self._computeXYarray()
        
rfetick's avatar
rfetick committed
437
438
439
440
441
    @property
    def samp(self):
        return self._samp_over/self._k
    
    @samp.setter
442
    def samp(self, value):
rfetick's avatar
rfetick committed
443
        # Manage cases of undersampling
444
        self._samp_over, self._k = oversample(value, fixed_k = self.fixed_k)
rfetick's avatar
rfetick committed
445
        self._computeXYarray()
rfetick's avatar
rfetick committed
446
        
rfetick's avatar
rfetick committed
447
448
449
450
451
    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
452
    
453
    #@_lru_cache(maxsize=2)
454
455
    @property
    def _pix2freq(self):
456
457
458
459
460
        """
        pixel to frequency conversion in the PSD plane
        """
        return 1.0/(self.system.D*self._samp_over)
    
461
    #@_lru_cache(maxsize=2)
462
463
    @property
    def _nxnyOver(self):
464
465
466
        """
        Return the number of pixels for the correctly sampled array (at least at Shannon-Nyquist)
        """
467
        return self.npix[0]*self._k, self.npix[1]*self._k
468
    
469
    #@_lru_cache(maxsize=2)
470
471
472
    @property
    def _nullFreqIndex(self):
        nx,ny = self._nxnyOver
473
474
        return nx//2, ny//2
    
475
    #@_lru_cache(maxsize=2)
476
477
    @property
    def _fxfy(self):
478
479
480
        """
        Return the array of frequencies (correctly sampled above Shannon-Nyquist)
        """
rfetick's avatar
rfetick committed
481
        return self._xyarray * self._pix2freq
482
    
rfetick's avatar
rfetick committed
483
    def strehlOTF(self,x0):
484
        """Compute the Strehl ratio based on the sum of the OTF"""
rfetick's avatar
rfetick committed
485
        otf = self.otf(x0,_caller='self')
486
        otf_diff = self._otfDiffraction()
rfetick's avatar
rfetick committed
487
488
489
        return _np.real(_np.sum(otf)/_np.sum(otf_diff))
    
    def strehlMarechal(self,x0):
490
        """Compute the Strehl ratio based on the Maréchal approximation"""
rfetick's avatar
rfetick committed
491
492
493
        _,sig2 = self.psd(x0)
        return _np.exp(-sig2)
    
494
495
496
    def psd(self,x0):
        raise ValueError("ParametricPSFfromPSD is not to be instantiated. the `psd` method must be override in the subclasses")
    
rfetick's avatar
rfetick committed
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
    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')
        
    def otf(self,x0,dx=0,dy=0,_caller='user'):
        """
        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)")
        
517
518
519
        otfTurbulent = self._otfTurbulent(x0)
        otfDiffraction = self._otfDiffraction()
        otfShift = self._otfShift(dx,dy)
rfetick's avatar
rfetick committed
520
        
521
        return otfTurbulent * otfDiffraction * otfShift
rfetick's avatar
rfetick committed
522
    
523
    def _otfTurbulent(self,x0):
524
        PSD,integral = self.psd(x0)
rfetick's avatar
rfetick committed
525
        L = self.system.D * self._samp_over
rfetick's avatar
rfetick committed
526
        Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
527
528
529
        #Dphi = _fft.fftshift(_np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV 
        Dphi = _fft.fftshift(_np.real(2 * (integral - Bg))) # normalized up to infinity
        return _np.exp(-Dphi/2.)
rfetick's avatar
rfetick committed
530
    
531
    #@_lru_cache(maxsize=2)
532
    def _otfDiffraction(self):
rfetick's avatar
rfetick committed
533
534
535
536
        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
537
        tab[0:int(NpupX), 0:int(NpupY)] = self.system.pupil((NpupX,NpupY),samp=self._samp_over)
rfetick's avatar
rfetick committed
538
        return _fft.fftshift(_np.abs(_fft.ifft2(_np.abs(_fft.fft2(tab)) ** 2)) / _np.sum(tab))
rfetick's avatar
rfetick committed
539
    
540
    #@_lru_cache(maxsize=2)
541
    def _otfShift(self,dx,dy):
rfetick's avatar
rfetick committed
542
543
544
        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
545
546
547
548
549
550
    
    def __call__(self,x0,dx=0,dy=0):
        """
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
551
            Array of parameters for the PSD (see __doc__)
rfetick's avatar
rfetick committed
552
553
554
555
556
557
        dx : float
            PSF X shifting [pix] (default = 0)
        dy : float
            PSF Y shifting [pix] (default = 0)
        """
        
rfetick's avatar
rfetick committed
558
        out = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(self.otf(x0,dx=dx,dy=dy,_caller='self')))))
rfetick's avatar
rfetick committed
559
560
561
562
        
        if self._k==1:
            return out
        else:
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
            return _binning(out,int(self._k)) # undersample PSF if needed (if it was oversampled for computations)

#%% 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]
    """
579
    def __init__(self,npix,system=None,samp=None):
580
581
582
        """
        Parameters
        ----------
583
        npix : tuple
584
585
586
587
588
589
            Size of output PSF
        system : OpticalSystem
            Optical system for this PSF
        samp : float
            Sampling at the observation wavelength
        """
590
        super().__init__(2,npix,system=system,samp=samp)
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
        
        # r0,L0
        bounds_down = [_EPSILON, _EPSILON]
        bounds_up = [_np.inf, _np.inf]
        self.bounds = (bounds_down,bounds_up)
    
    def psd(self,x0):
        """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)
613
614
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
615
616
617
618
619
620
621
622
623
624
        
        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
625
        fmax = _np.min([nx0,ny0])*self._pix2freq
626
627
        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)
628
629
630
        integral = integral_in + integral_out
        
        return PSD, integral
631

632
#%% PSFAO MODEL        
633
class Psfao(ParametricPSFfromPSD):
FETICK Romain's avatar
FETICK Romain committed
634
    """
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
    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
654
    """
655
    def __init__(self,npix,system=None,Lext=10.,samp=None, fixed_k = None):
FETICK Romain's avatar
FETICK Romain committed
656
657
658
        """
        Parameters
        ----------
659
        npix : tuple
FETICK Romain's avatar
FETICK Romain committed
660
661
662
663
664
665
666
667
668
            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
        """
669
        super().__init__(7,npix,system=system,samp=samp,fixed_k=fixed_k)
670
        
FETICK Romain's avatar
FETICK Romain committed
671
        self.Lext = Lext
rfetick's avatar
rfetick committed
672
        # r0,C,A,alpha,ratio,theta,beta 
rfetick's avatar
rfetick committed
673
674
        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
675
676
677
678
679
680
681
682
683
        self.bounds = (bounds_down,bounds_up)
    
    def psd(self,x0):
        """Compute the PSD model from parameters
        PSD is given in [rad²/f²] = [rad² m²]
        
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
684
            See __doc__ for more details
FETICK Romain's avatar
FETICK Romain committed
685
686
687
688
            
        Returns
        -------
        psd : numpy.array (dim=2)
689
        integral : float : the integral of the `psd` up to infinity
FETICK Romain's avatar
FETICK Romain committed
690
691
            
        """
692
        self.check_parameters(x0)
693
694
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
FETICK Romain's avatar
FETICK Romain committed
695
        
696
697
        F2 = f2D[0]**2. + f2D[1]**2.
        
FETICK Romain's avatar
FETICK Romain committed
698
699
        Fao = self.system.Nact/(2.0*self.system.D)
        
rfetick's avatar
rfetick committed
700
        r0,C,A,alpha,ratio,theta,beta = x0
701
702
        
        PSD = 0.0229* r0**(-5./3.) * ((1. / self.Lext**2.) + F2)**(-11./6.)
FETICK Romain's avatar
FETICK Romain committed
703
704
        PSD *= (F2 >= Fao**2.)
        
rfetick's avatar
rfetick committed
705
706
        ax = alpha*ratio
        ay = alpha/ratio
707
        param = (ax,ay,theta,beta,0,0)
708
        
709
710
        removeInside = (1+_np.sqrt(2))/2 * self._pix2freq/2 # remove central pixel in energy computation
        moff = moffat(f2D,param,norm=Fao,removeInside=removeInside) * (F2 < Fao**2.)
711
        moff[nx0,ny0] = 0.0 # Set Moffat PSD = 0 at null frequency
712
        # Newly added for the PSFAO19 model
713
        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
714
        # Warning: Moffat numerical normalization generates strehlOTF jump when "_k" is changed
715
716
717
        
        PSD += (F2 < Fao**2.) * _np.abs(C + A*moff) # before I wrote "moffat(f2D,param,norm=Fao)" instead of "moff", so no numerical normalization
        PSD[nx0,ny0] = 0.0 # Set PSD = 0 at null frequency
718
719
        
        # Compute PSD integral up to infinity
720
721
        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
722
        integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum
723
724
725
        integral = integral_in + integral_out
        
        return PSD, integral
FETICK Romain's avatar
FETICK Romain committed
726
        
rfetick's avatar
rfetick committed
727
    def tofits(self,param,filename,*args,keys=None,overwrite=False,**kwargs):
FETICK Romain's avatar
FETICK Romain committed
728
        if keys is None:
rfetick's avatar
rfetick committed
729
            keys = ["R0","CST","SIGMA2","ALPHA","RATIO","THETA","BETA"]
FETICK Romain's avatar
FETICK Romain committed
730
731
732
            keys_comment = ["Fried parameter [m]",
                            "PSD AO area constant C [rad2]",
                            "PSD AO area Moffat variance A [rad2]",
733
                            "PSD AO area Moffat alpha [1/m]",
rfetick's avatar
rfetick committed
734
                            "PSD AO area Moffat sqrt(ax/ay) ratio",
FETICK Romain's avatar
FETICK Romain committed
735
736
                            "PSD AO area Moffat theta [rad]",
                            "PSD AO area Moffat beta"]
FETICK Romain's avatar
FETICK Romain committed
737
738
739
740
741
742
        
        # 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")
743
744
        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
745
746
747
        hdr["HIERARCH SAMP"] = (self.samp,"Sampling (eg. 2 for Shannon)")
        hdr["HIERARCH LEXT"] = (self.Lext,"Von-Karman outer scale")
        
rfetick's avatar
rfetick committed
748
        hdu = _fits.PrimaryHDU(psf, hdr)
rfetick's avatar
rfetick committed
749
        hdu.writeto(filename, overwrite=overwrite)