psfmodel.py 30.8 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,grad=False):
rfetick's avatar
rfetick committed
195
    """Create an array of reduced coordinated with elongation and rotation"""
rfetick's avatar
rfetick committed
196
197
198
    c = _np.cos(theta)
    s = _np.sin(theta)
    s2 = _np.sin(2.0 * theta)
199
200
201

    Rxx = (c/ax)**2 + (s/ay)**2
    Rxy =  s2/ay**2 -  s2/ax**2
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
    Ryy = (c/ay)**2 + (s/ax)**2
    
    uxx = (XY[0]-cx)**2
    uxy = (XY[0]-cx)*(XY[1]-cy)
    uyy = (XY[1]-cy)**2
    
    u = Rxx*uxx + Rxy*uxy + Ryy*uyy
    
    if grad:
        g = _np.zeros((3,)+u.shape)
        # du/dax
        g[0,...] = -2/ax * (uxx*(c/ax)**2 - uxy*s2/ax**2 + uyy*(s/ax)**2)
        # du/day
        g[1,...] = -2/ay * (uxx*(s/ay)**2 + uxy*s2/ay**2 + uyy*(c/ay)**2)
        # du/dtheta
        drxx = -2*c*s/ax**2 + 2*s*c/ay**2
        drxy = 2*_np.cos(2*theta)*(1/ay**2 - 1/ax**2)
        dryy = -2*c*s/ay**2 + 2*s*c/ax**2
        g[2,...] = drxx*uxx + drxy*uxy + dryy*uyy
        return u,g
222
223
224
    
    return u    

rfetick's avatar
rfetick committed
225
226
227
def moffat(XY, param, norm=None, removeInside=0, grad=False):
    """Compute a Moffat function on a meshgrid
    ```moff = E * (1+u)^(-beta)```
228
    with `u` the quadratic coordinates in the shifted and rotated frame
rfetick's avatar
rfetick committed
229
    and `E` the energy normalization factor
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
    
    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
245
    norm : None, np.inf, float (>0), int (>0)
246
247
        Radius for energy normalization
        None      - No energy normalization (maximum=1.0)
rfetick's avatar
rfetick committed
248
249
250
                    E = 1.0
        np.inf    - Total energy normalization (on the whole X-Y plane)
                    E = (beta-1)/(pi*ax*ay)
251
        float,int - Energy normalization up to the radius defined by this value
rfetick's avatar
rfetick committed
252
                    E = (beta-1)/(pi*ax*ay)*(1-(1+(R**2)/(ax*ay))**(1-beta))    
253
254
255
    
    removeInside: float (default=0)
        Used to remove the central pixel in energy computation
256
257
258
259
    """
    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
    
rfetick's avatar
rfetick committed
260
    if grad:
rfetick's avatar
rfetick committed
261
        u,du = reduced_coord(XY,ax,ay,theta,cx,cy,grad=True)
rfetick's avatar
rfetick committed
262
263
264
    else:
        u = reduced_coord(XY,ax,ay,theta,cx,cy,grad=False)
    
rfetick's avatar
rfetick committed
265
266
267
268
269
270
271
272
273
    V = (1.0+u)**(-beta) # Moffat shape
    E = 1.0 # normalization factor (eventually up to infinity)
    F = 1.0 # normalization factor (eventually up to a limited radius)
    
    if grad:
        dVdu = -beta*V/(1.0+u)
        dV = _np.zeros((4,)+u.shape)
        for i in range(3): dV[i,...] = dVdu*du[i,...]
        dV[3,...] = -V*_np.log(1.0+u)
274
275
    
    if norm is None:
rfetick's avatar
rfetick committed
276
        if grad:
rfetick's avatar
rfetick committed
277
            return V,dV
278
279
    else: # norm can be float or np.inf
        if (beta<=1) and (norm==_np.inf): raise ValueError("Cannot compute Moffat energy for beta<=1")
280
        if beta==1: raise ValueError("Energy computation for beta=1.0 not implemented yet. Sorry!")
rfetick's avatar
rfetick committed
281
        E = (beta-1) / (_np.pi*ax*ay)
rfetick's avatar
rfetick committed
282
283
284
        Fout = (1 +         (norm**2)/(ax*ay))**(1-beta)
        Fin  = (1 + (removeInside**2)/(ax*ay))**(1-beta)
        F = 1/(Fin-Fout)
285
        if grad:
rfetick's avatar
rfetick committed
286
287
288
289
290
291
292
293
294
295
296
            dE = [-E/ax,-E/ay,0,E/(beta-1)]
            k = (1-beta)*Fout/(1 + (norm**2)/(ax*ay))*(norm**2)/(ax*ay)
            dFout = _np.array([-k/ax,-k/ay,0,-Fout*_np.log(1 + (norm**2)/(ax*ay))]) if norm<_np.inf else _np.zeros(4)
            k = (1-beta)*Fin/(1 + (removeInside**2)/(ax*ay))*(removeInside**2)/(ax*ay)
            dFin  = _np.array([-k/ax,-k/ay,0,-Fin*_np.log(1 + (removeInside**2)/(ax*ay))])
            dF = -(dFin-dFout)/(Fin-Fout)**2
            dm = 0*dV
            for i in range(4): dm[i,...] = V*E*dF[i] + V*dE[i]*F + dV[i,...]*E*F
            return V*E*F, dm
            
    return V*E*F
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313

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
314
315
    ax = _np.sqrt(2)*param[0]
    ay = _np.sqrt(2)*param[1]
316
    u = reduced_coord(XY,ax,ay,param[2],param[3],param[4])
rfetick's avatar
rfetick committed
317
    return 1.0/(2*_np.pi*param[0]*param[1])*_np.exp(-u)
318

FETICK Romain's avatar
FETICK Romain committed
319
320
321
322
323
324
#%% 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
    """
    
325
    def __init__(self,npix):
FETICK Romain's avatar
FETICK Romain committed
326
327
328
        """
        Parameters
        ----------
329
        npix : tuple of two elements
FETICK Romain's avatar
FETICK Romain committed
330
331
            Model X and Y pixel size when called
        """
332
333
        if type(npix)!=tuple: raise TypeError("Argument `npix` must be a tuple")
        self.npix = npix
rfetick's avatar
rfetick committed
334
        self.bounds = (-_np.inf,_np.inf)
FETICK Romain's avatar
FETICK Romain committed
335
336
    
    def __repr__(self):
337
        return "ParametricPSF of size (%u,%u)"%self.npix
FETICK Romain's avatar
FETICK Romain committed
338
339
340
341
342
343
344
    
    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
345
        return _fft.fftshift(_fft.fft2(psf))
FETICK Romain's avatar
FETICK Romain committed
346
347
348
    
    def mtf(self,*args,**kwargs):
        """Return the Modulation Transfer Function (MTF)"""
rfetick's avatar
rfetick committed
349
        return _np.abs(self.otf(args,kwargs))
FETICK Romain's avatar
FETICK Romain committed
350
351
352
353
    
    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
354
        hdu = _fits.PrimaryHDU(psf, hdr)
FETICK Romain's avatar
FETICK Romain committed
355
356
357
        hdu.writeto(filename, overwrite=True)
        
    def _getfitshdr(self,param,keys=None,keys_comment=None):
358
359
        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
360
        if keys_comment is not None:
361
            if len(keys_comment)!=len(param): raise ValueError("When defined, `keys_comment` must be same size as `param`")
rfetick's avatar
rfetick committed
362
        hdr = _fits.Header()
FETICK Romain's avatar
FETICK Romain committed
363
        
FETICK Romain's avatar
FETICK Romain committed
364
        hdr["HIERARCH ORIGIN"] = "MAOPPY automatic header"
rfetick's avatar
rfetick committed
365
        hdr["HIERARCH CREATION"] = (_time.ctime(),"Date of file creation")
FETICK Romain's avatar
FETICK Romain committed
366
367
368
369
370
371
372
373
374
375
376
377
378
        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
379
        super().__init__(_np.shape(image_psf))
FETICK Romain's avatar
FETICK Romain committed
380
381
382
383
384
385
386
387
        self.image_psf = image_psf
        self.bounds = ()
        
    def __call__(self,*args,**kwargs):
        return self.image_psf


class Moffat(ParametricPSF):
388
389
    def __init__(self,npix,norm=_np.inf):
        self.npix = npix
FETICK Romain's avatar
FETICK Romain committed
390
        self.norm = norm
rfetick's avatar
rfetick committed
391
392
        bounds_down = [_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
        bounds_up = [_np.inf for i in range(4)]
FETICK Romain's avatar
FETICK Romain committed
393
        self.bounds = (bounds_down,bounds_up)
FETICK Romain's avatar
FETICK Romain committed
394
        
395
396
397
        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
398
399
400
401
402
403
404
405
406
407
408
    
    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
409
        y = _np.concatenate((x,[dx,dy]))
FETICK Romain's avatar
FETICK Romain committed
410
        return moffat(self._XY,y,norm=self.norm)
FETICK Romain's avatar
FETICK Romain committed
411
412
413
414
415
    
    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
416
class Gaussian(ParametricPSF):
417
418
    def __init__(self,npix):
        self.npix = npix
rfetick's avatar
rfetick committed
419
420
        bounds_down = [_EPSILON,_EPSILON,-_np.inf]
        bounds_up = [_np.inf for i in range(3)]
FETICK Romain's avatar
FETICK Romain committed
421
422
        self.bounds = (bounds_down,bounds_up)
        
423
424
425
        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
426
427
428
429
430
431
432
433
434
435
    
    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
436
        y = _np.concatenate((x,[dx,dy]))
FETICK Romain's avatar
FETICK Romain committed
437
438
439
        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
440
        super(Gaussian,self).tofits(param,filename,*args,keys=keys,**kwargs)
441

442
443
444
#%% MASTER CLASS
class ParametricPSFfromPSD(ParametricPSF):
    """This class is NOT to be instantiated"""
445
446
447
448
    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
449
450
        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
451
452
        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
453
        self.system = system
454
455
        self._nparam = nparam
        
rfetick's avatar
rfetick committed
456
457
458
459
460
461
462
463
464
    @property
    def npix(self):
        return self._npix
    
    @npix.setter
    def npix(self,value):
        self._npix = value
        self._computeXYarray()
        
rfetick's avatar
rfetick committed
465
466
467
468
469
470
471
472
    @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
473
        self._computeXYarray()
rfetick's avatar
rfetick committed
474
        
rfetick's avatar
rfetick committed
475
476
477
478
479
    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
480
    
481
    #@_lru_cache(maxsize=2)
482
483
    @property
    def _pix2freq(self):
484
485
486
487
488
        """
        pixel to frequency conversion in the PSD plane
        """
        return 1.0/(self.system.D*self._samp_over)
    
489
    #@_lru_cache(maxsize=2)
490
491
    @property
    def _nxnyOver(self):
492
493
494
        """
        Return the number of pixels for the correctly sampled array (at least at Shannon-Nyquist)
        """
495
        return self.npix[0]*self._k, self.npix[1]*self._k
496
    
497
    #@_lru_cache(maxsize=2)
498
499
500
    @property
    def _nullFreqIndex(self):
        nx,ny = self._nxnyOver
501
502
        return nx//2, ny//2
    
503
    #@_lru_cache(maxsize=2)
504
505
    @property
    def _fxfy(self):
506
507
508
        """
        Return the array of frequencies (correctly sampled above Shannon-Nyquist)
        """
rfetick's avatar
rfetick committed
509
        return self._xyarray * self._pix2freq
510
    
rfetick's avatar
rfetick committed
511
    def strehlOTF(self,x0):
512
        """Compute the Strehl ratio based on the sum of the OTF"""
rfetick's avatar
rfetick committed
513
        otf = self.otf(x0,_caller='self')
514
        otf_diff = self._otfDiffraction()
rfetick's avatar
rfetick committed
515
516
517
        return _np.real(_np.sum(otf)/_np.sum(otf_diff))
    
    def strehlMarechal(self,x0):
518
        """Compute the Strehl ratio based on the Maréchal approximation"""
rfetick's avatar
rfetick committed
519
520
521
        _,sig2 = self.psd(x0)
        return _np.exp(-sig2)
    
rfetick's avatar
rfetick committed
522
    def psd(self,x0,grad=False):
523
524
        raise ValueError("ParametricPSFfromPSD is not to be instantiated. the `psd` method must be override in the subclasses")
    
rfetick's avatar
rfetick committed
525
526
527
528
529
530
531
532
533
    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
534
    def otf(self,x0,dx=0,dy=0,_caller='user',grad=False):
rfetick's avatar
rfetick committed
535
536
537
538
539
540
541
542
543
544
        """
        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
545
        otfDiffr = self._otfDiffraction()
546
        otfShift = self._otfShift(dx,dy)
rfetick's avatar
rfetick committed
547
548
549
550
551
        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
552
        
rfetick's avatar
rfetick committed
553
554
555
        otf = otfTurb * otfDiffr * otfShift
        if grad: return otf,g
        return otf
rfetick's avatar
rfetick committed
556
    
rfetick's avatar
rfetick committed
557
    def _otfTurbulent(self,x0,grad=False):
rfetick's avatar
rfetick committed
558
        L = self.system.D * self._samp_over
rfetick's avatar
rfetick committed
559
        if grad:
rfetick's avatar
rfetick committed
560
            PSD,integral,g,integral_g = self.psd(x0,grad=True)
rfetick's avatar
rfetick committed
561
562
563
        else:
            PSD,integral = self.psd(x0,grad=False)
            
rfetick's avatar
rfetick committed
564
        Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
rfetick's avatar
rfetick committed
565
566
567
568
569
        #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
570
            g2 = _np.zeros(g.shape,dtype=complex) # I cannot override 'g' here due to float to complex type
rfetick's avatar
rfetick committed
571
572
            for i in range(len(x0)):
                Bg = _fft.fft2(_fft.fftshift(g[i,...])) / L**2
rfetick's avatar
rfetick committed
573
574
                #Bphi = _fft.fftshift(_np.real(Bg[0, 0] - Bg)) # normalized on the numerical FoV
                Bphi = _fft.fftshift(_np.real(integral_g[i] - Bg)) # normalized up to infinity
rfetick's avatar
rfetick committed
575
576
577
                g2[i,...] = -Bphi*otf
            return otf, g2
        return otf
rfetick's avatar
rfetick committed
578
    
579
    #@_lru_cache(maxsize=2)
580
    def _otfDiffraction(self):
rfetick's avatar
rfetick committed
581
582
583
        nx, ny = self._nxnyOver
        NpupX = _np.ceil(nx/self._samp_over)
        NpupY = _np.ceil(ny/self._samp_over)
584
        tab = _np.zeros((nx, ny), dtype=complex)
rfetick's avatar
rfetick committed
585
        tab[0:int(NpupX), 0:int(NpupY)] = self.system.pupil((NpupX,NpupY),samp=self._samp_over)
rfetick's avatar
rfetick committed
586
        return _fft.fftshift(_np.abs(_fft.ifft2(_np.abs(_fft.fft2(tab)) ** 2)) / _np.sum(tab))
rfetick's avatar
rfetick committed
587
    
588
    #@_lru_cache(maxsize=2)
589
    def _otfShift(self,dx,dy):
rfetick's avatar
rfetick committed
590
591
592
        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
593
    
rfetick's avatar
rfetick committed
594
    def __call__(self,x0,dx=0,dy=0,grad=False):
rfetick's avatar
rfetick committed
595
596
597
598
        """
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
599
            Array of parameters for the PSD (see __doc__)
rfetick's avatar
rfetick committed
600
601
602
603
604
605
        dx : float
            PSF X shifting [pix] (default = 0)
        dy : float
            PSF Y shifting [pix] (default = 0)
        """
        
rfetick's avatar
rfetick committed
606
607
608
609
610
611
        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
612
        
rfetick's avatar
rfetick committed
613
614
615
616
617
618
        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
619
        else:
rfetick's avatar
rfetick committed
620
621
622
            if grad:
                g2 = _np.zeros((len(x0),psf.shape[0]//k,psf.shape[1]//k))
                for i in range(len(x0)):
rfetick's avatar
rfetick committed
623
                    g2[i,...] = _binning(g[i,...].astype(float),k)
rfetick's avatar
rfetick committed
624
625
                return _binning(psf,k), g2
            return _binning(psf,k) # undersample PSF if needed (if it was oversampled for computations)
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640

#%% 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]
    """
641
    def __init__(self,npix,system=None,samp=None):
642
643
644
        """
        Parameters
        ----------
645
        npix : tuple
646
647
648
649
650
651
            Size of output PSF
        system : OpticalSystem
            Optical system for this PSF
        samp : float
            Sampling at the observation wavelength
        """
652
        super().__init__(2,npix,system=system,samp=samp)
653
654
655
656
657
658
        
        # r0,L0
        bounds_down = [_EPSILON, _EPSILON]
        bounds_up = [_np.inf, _np.inf]
        self.bounds = (bounds_down,bounds_up)
    
rfetick's avatar
rfetick committed
659
    def psd(self,x0,grad=False):
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
        """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)
675
676
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
677
678
679
680
681
682
683
684
685
686
        
        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
687
        fmax = _np.min([nx0,ny0])*self._pix2freq
688
689
        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)
690
691
        integral = integral_in + integral_out
        
rfetick's avatar
rfetick committed
692
693
694
695
        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)
rfetick's avatar
rfetick committed
696
697
698
699
700
            
            # compute integral gradient
            integral_g = [0,0]
            for i in range(2): integral_g[i] = _np.sum(g[i,...]*(F2<(fmax**2))) * self._pix2freq**2 # numerical sum (in the array's tangent circle)
            integral_g[0] += integral_out*(-5./3)/r0
rfetick's avatar
rfetick committed
701
        
rfetick's avatar
rfetick committed
702
        if grad: return PSD,integral,g,integral_g
703
        return PSD, integral
704

705
#%% PSFAO MODEL        
706
class Psfao(ParametricPSFfromPSD):
FETICK Romain's avatar
FETICK Romain committed
707
    """
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
    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
727
    """
728
    def __init__(self,npix,system=None,Lext=10.,samp=None):
FETICK Romain's avatar
FETICK Romain committed
729
730
731
        """
        Parameters
        ----------
732
        npix : tuple
FETICK Romain's avatar
FETICK Romain committed
733
734
735
736
737
738
739
740
741
            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
        """
742
        super().__init__(7,npix,system=system,samp=samp)
743
        
FETICK Romain's avatar
FETICK Romain committed
744
        self.Lext = Lext
rfetick's avatar
rfetick committed
745
746
747
748
749
750
751
        # r0,C,A,alpha,ratio,theta,beta
        ### Mathematical bounds
        #bounds_down = [_EPSILON,0,0,_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
        #bounds_up = [_np.inf for i in range(7)]
        ### Physical bounds
        bounds_down = [1e-3,0,0,1e-3,1e-2,-_np.inf,1.01]
        bounds_up = [_np.inf]*4 + [1e2,_np.inf,5]
FETICK Romain's avatar
FETICK Romain committed
752
753
        self.bounds = (bounds_down,bounds_up)
    
rfetick's avatar
rfetick committed
754
    def psd(self,x0,grad=False):
FETICK Romain's avatar
FETICK Romain committed
755
756
757
758
759
760
        """Compute the PSD model from parameters
        PSD is given in [rad²/f²] = [rad² m²]
        
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
761
            See __doc__ for more details
rfetick's avatar
rfetick committed
762
763
764
        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
765
766
767
768
            
        Returns
        -------
        psd : numpy.array (dim=2)
769
        integral : float : the integral of the `psd` up to infinity
FETICK Romain's avatar
FETICK Romain committed
770
771
            
        """
772
        self.check_parameters(x0)
773
774
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
775
        pix = self._pix2freq
776
        F2 = f2D[0]**2. + f2D[1]**2.
FETICK Romain's avatar
FETICK Romain committed
777
        Fao = self.system.Nact/(2.0*self.system.D)
rfetick's avatar
rfetick committed
778
        maskin = (F2 < Fao**2.)
FETICK Romain's avatar
FETICK Romain committed
779
        
rfetick's avatar
rfetick committed
780
        r0,C,A,alpha,ratio,theta,beta = x0
781
782
        
        PSD = 0.0229* r0**(-5./3.) * ((1. / self.Lext**2.) + F2)**(-11./6.)
rfetick's avatar
rfetick committed
783
        PSD *= (1-maskin)
FETICK Romain's avatar
FETICK Romain committed
784
        
rfetick's avatar
rfetick committed
785
786
        ax = alpha*ratio
        ay = alpha/ratio
787
        param = (ax,ay,theta,beta,0,0)
788
        
789
790
791
792
793
794
795
796
        removeInside = (1+_np.sqrt(2))/2 * pix/2 # remove central pixel in energy computation
        if grad:
            moff,dm = moffat(f2D,param,norm=Fao,removeInside=removeInside,grad=True)
        else:
            moff = moffat(f2D,param,norm=Fao,removeInside=removeInside)
        moff *= maskin
        
        # Warning: Numerical normalization! (not compatible with analytical gradient)
797
        moff[nx0,ny0] = 0.0 # Set Moffat PSD = 0 at null frequency
798
        #moff = moff / (_np.sum(moff)*pix**2)  # normalize moffat numerically to get correct A=sigma² in the AO zone
rfetick's avatar
rfetick committed
799
        # Warning: Moffat numerical normalization generates strehlOTF jump when "_k" is changed
800
        
rfetick's avatar
rfetick committed
801
        PSD += maskin * (C + A*moff) # before I wrote "moffat(f2D,param,norm=Fao)" instead of "moff", so no numerical normalization
802
        PSD[nx0,ny0] = 0.0 # Set PSD = 0 at null frequency
803
804
        
        # Compute PSD integral up to infinity
805
806
        fmax = _np.min([nx0,ny0])*pix
        integral_in = _np.sum(PSD*(F2<(fmax**2))) * pix**2 # numerical sum
rfetick's avatar
rfetick committed
807
        integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum
808
809
        integral = integral_in + integral_out
        
rfetick's avatar
rfetick committed
810
        # TODO: finish gradient computation!
rfetick's avatar
rfetick committed
811
812
813
814
815
816
817
818
819
        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
820
            g[3,...] = A*maskin*(dm[0,...]*ratio + dm[1,...]/ratio)
rfetick's avatar
rfetick committed
821
            # derivative towards ratio
822
            g[4,...] = A*maskin*(dm[0,...]*alpha - dm[1,...]*ay/ratio)
rfetick's avatar
rfetick committed
823
            # derivative towards theta
824
            g[5,...] = A*maskin*dm[2,...]
rfetick's avatar
rfetick committed
825
            # derivative towards beta
826
            g[6,...] = A*maskin*dm[3,...]
rfetick's avatar
rfetick committed
827
828
829
            # Remove central freq from all derivative
            g[:,nx0,ny0] = 0
        
rfetick's avatar
rfetick committed
830
831
832
833
834
835
            # Compute integral gradient
            integral_g = _np.zeros(len(x0))
            for i in range(len(x0)): integral_g[i] = _np.sum(g[i,...]*(F2<(fmax**2))) * self._pix2freq**2 # numerical sum
            integral_g[0] += integral_out*(-5./3)/r0
            
        if grad: return PSD,integral,g,integral_g
836
        return PSD, integral
FETICK Romain's avatar
FETICK Romain committed
837
        
rfetick's avatar
rfetick committed
838
    def tofits(self,param,filename,*args,keys=None,overwrite=False,**kwargs):
FETICK Romain's avatar
FETICK Romain committed
839
        if keys is None:
rfetick's avatar
rfetick committed
840
            keys = ["R0","CST","SIGMA2","ALPHA","RATIO","THETA","BETA"]
FETICK Romain's avatar
FETICK Romain committed
841
842
843
            keys_comment = ["Fried parameter [m]",
                            "PSD AO area constant C [rad2]",
                            "PSD AO area Moffat variance A [rad2]",
844
                            "PSD AO area Moffat alpha [1/m]",
rfetick's avatar
rfetick committed
845
                            "PSD AO area Moffat sqrt(ax/ay) ratio",
FETICK Romain's avatar
FETICK Romain committed
846
847
                            "PSD AO area Moffat theta [rad]",
                            "PSD AO area Moffat beta"]
FETICK Romain's avatar
FETICK Romain committed
848
849
850
851
852
853
        
        # 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")
854
855
        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
856
857
858
        hdr["HIERARCH SAMP"] = (self.samp,"Sampling (eg. 2 for Shannon)")
        hdr["HIERARCH LEXT"] = (self.Lext,"Von-Karman outer scale")
        
rfetick's avatar
rfetick committed
859
        hdu = _fits.PrimaryHDU(psf, hdr)
rfetick's avatar
rfetick committed
860
        hdu.writeto(filename, overwrite=overwrite)