psfmodel.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 27 17:31:18 2019

@author: rfetick
"""

import numpy as _np
from scipy.optimize import least_squares as _least_squares
from astropy.io import fits as _fits
import time as _time
import numpy.fft as _fft # import fft2, fftshift, ifft2
import sys as _sys
from maoppy.utils import binning as _binning
#from functools import lru_cache as _lru_cache

# Value to compute finite differences
_EPSILON = _np.sqrt(_sys.float_info.epsilon)

#%% FUNCTIONS
def oversample(samp):
    """Find the minimal integer that allows oversampling"""
    k = int(_np.ceil(2.0/samp))
    return (k*samp,k)

#%% FITTING
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)
    """
    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)

    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),
           positive_bck=False,fixed=None,npixfit=None,**kwargs):
    """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`)
        Inverse of the noise's variance
        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)
    npixfit : int (default=None)
        Increased pixel size for improved fitting accuracy
    **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
    """
    
    if weights is None: weights = _np.ones_like(psf)
    elif len(psf)!=len(weights): raise ValueError("Keyword `weights` must have same number of elements as `psf`")
    
    # Increase array size for improved fitting accuracy
    if npixfit is not None:
        sx,sy = _np.shape(psf)
        if (sx>npixfit) or (sy>npixfit): raise ValueError('npixfit must be greater or equal to both psf dimensions')
        psfbig = _np.zeros((npixfit,npixfit))
        wbig = _np.zeros((npixfit,npixfit))
        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
    
    sqW = _np.sqrt(weights)
    Model_inst = Model(_np.shape(psf),**kwargs)
    
    class CostClass(object):
        def __init__(self):
            self.iter = 0
        def __call__(self,y):
            if (self.iter%3)==0: print("-",end="")
            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)
            return _np.reshape(sqW * (amp * m + bck - psf), _np.size(psf))
    
    cost = CostClass()
    
    if fixed is not None:
        if len(fixed)!=len(x0): raise ValueError("When defined, `fixed` must be same size as `x0`")
        FREE = [not fixed[i] for i in range(len(fixed))]
        INDFREE = _np.where(FREE)[0]
    
    def input2mini(x,dxdy):
        # Transform user parameters to minimizer parameters
        if fixed is None: xfree = x
        else: xfree = _np.take(x,INDFREE)
        return _np.concatenate((xfree,dxdy))
    
    def mini2input(y):
        # Transform minimizer parameters to user parameters
        if fixed is None:
            xall = y[0:-2]
        else:
            xall = _np.copy(x0)
            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]
        if fixed is not None: b_low = _np.take(b_low,INDFREE)
        b_low = _np.concatenate((b_low,[-_np.inf,-_np.inf]))
        b_up = inst.bounds[1]
        if fixed is not None: b_up = _np.take(b_up,INDFREE)
        b_up = _np.concatenate((b_up,[_np.inf,_np.inf]))
        return (b_low,b_up)
    
    result = _least_squares(cost, input2mini(x0,dxdy), bounds=get_bound(Model_inst))
    
    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

#%% BASIC FUNCTIONS
def reduced_coord(XY,ax,ay,theta,cx,cy,grad=False):
    """Create an array of reduced coordinated with elongation and rotation"""
    c = _np.cos(theta)
    s = _np.sin(theta)
    s2 = _np.sin(2.0 * theta)

    Rxx = (c/ax)**2 + (s/ay)**2
    Rxy =  s2/ay**2 -  s2/ax**2
    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
    
    return u    

def moffat(XY, param, norm=None, removeInside=0, grad=False):
    """Compute a Moffat function on a meshgrid
    ```moff = E * (1+u)^(-beta)```
    with `u` the quadratic coordinates in the shifted and rotated frame
    and `E` the energy normalization factor
    
    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
    --------
    norm : None, np.inf, float (>0), int (>0)
        Radius for energy normalization
        None      - No energy normalization (maximum=1.0)
                    E = 1.0
        np.inf    - Total energy normalization (on the whole X-Y plane)
                    E = (beta-1)/(pi*ax*ay)
        float,int - Energy normalization up to the radius defined by this value
                    E = (beta-1)/(pi*ax*ay)*(1-(1+(R**2)/(ax*ay))**(1-beta))    
    
    removeInside: float (default=0)
        Used to remove the central pixel in energy computation
    """
    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
    
    if grad:
        u,du = reduced_coord(XY,ax,ay,theta,cx,cy,grad=True)
    else:
        u = reduced_coord(XY,ax,ay,theta,cx,cy,grad=False)
    
    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)
    
    if norm is None:
        if grad:
            return V,dV
    else: # norm can be float or np.inf
        if (beta<=1) and (norm==_np.inf): raise ValueError("Cannot compute Moffat energy for beta<=1")
        if beta==1: raise ValueError("Energy computation for beta=1.0 not implemented yet. Sorry!")
        E = (beta-1) / (_np.pi*ax*ay)
        Fout = (1 +         (norm**2)/(ax*ay))**(1-beta)
        Fin  = (1 + (removeInside**2)/(ax*ay))**(1-beta)
        F = 1/(Fin-Fout)
        if grad:
            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

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))
    ax = _np.sqrt(2)*param[0]
    ay = _np.sqrt(2)*param[1]
    u = reduced_coord(XY,ax,ay,param[2],param[3],param[4])
    return 1.0/(2*_np.pi*param[0]*param[1])*_np.exp(-u)

#%% 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
    """
    
    def __init__(self,npix):
        """
        Parameters
        ----------
        npix : tuple of two elements
            Model X and Y pixel size when called
        """
        if type(npix)!=tuple: raise TypeError("Argument `npix` must be a tuple")
        self.npix = npix
        self.bounds = (-_np.inf,_np.inf)
    
    def __repr__(self):
        return "ParametricPSF of size (%u,%u)"%self.npix
    
    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)
        return _fft.fftshift(_fft.fft2(psf))
    
    def mtf(self,*args,**kwargs):
        """Return the Modulation Transfer Function (MTF)"""
        return _np.abs(self.otf(args,kwargs))
    
    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)
        hdu = _fits.PrimaryHDU(psf, hdr)
        hdu.writeto(filename, overwrite=True)
        
    def _getfitshdr(self,param,keys=None,keys_comment=None):
        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`")
        if keys_comment is not None:
            if len(keys_comment)!=len(param): raise ValueError("When defined, `keys_comment` must be same size as `param`")
        hdr = _fits.Header()
        
        hdr["HIERARCH ORIGIN"] = "MAOPPY automatic header"
        hdr["HIERARCH CREATION"] = (_time.ctime(),"Date of file creation")
        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):
        super().__init__(_np.shape(image_psf))
        self.image_psf = image_psf
        self.bounds = ()
        
    def __call__(self,*args,**kwargs):
        return self.image_psf


class Moffat(ParametricPSF):
    def __init__(self,npix,norm=_np.inf):
        self.npix = npix
        self.norm = norm
        bounds_down = [_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
        bounds_up = [_np.inf for i in range(4)]
        self.bounds = (bounds_down,bounds_up)
        
        self._XY = _np.mgrid[0:npix[0],0:npix[1]]
        self._XY[0] -= npix[0]//2
        self._XY[1] -= npix[1]//2
    
    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
        """
        y = _np.concatenate((x,[dx,dy]))
        return moffat(self._XY,y,norm=self.norm)
    
    def tofits(self,param,filename,*args,keys=["ALPHA_X","ALPHA_Y","THETA","BETA"],**kwargs):
        super(Moffat,self).tofits(param,filename,*args,keys=keys,**kwargs)


class Gaussian(ParametricPSF):
    def __init__(self,npix):
        self.npix = npix
        bounds_down = [_EPSILON,_EPSILON,-_np.inf]
        bounds_up = [_np.inf for i in range(3)]
        self.bounds = (bounds_down,bounds_up)
        
        self._XY = _np.mgrid[0:npix[0],0:npix[1]]
        self._XY[1] -= npix[0]/2
        self._XY[0] -= npix[1]/2
    
    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
        """
        y = _np.concatenate((x,[dx,dy]))
        return gauss(self._XY,y)
    
    def tofits(self,param,filename,*args,keys=["SIGMA_X","SIGMA_Y","THETA"],**kwargs):
        super(Gaussian,self).tofits(param,filename,*args,keys=keys,**kwargs)

#%% MASTER CLASS
class ParametricPSFfromPSD(ParametricPSF):
    """This class is NOT to be instantiated"""
    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")
        if system is None: raise ValueError("Keyword `system` must be defined")
        if samp is None: raise ValueError("Keyword `samp` must be defined")
        self._npix = npix # "_" to bypass the _xyarray update, that will be made with samp setter
        self.samp = samp # also init _xyarray
        self.system = system
        self._nparam = nparam
        
    @property
    def npix(self):
        return self._npix
    
    @npix.setter
    def npix(self,value):
        self._npix = value
        self._computeXYarray()
        
    @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)
        self._computeXYarray()
        
    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
    
    #@_lru_cache(maxsize=2)
    @property
    def _pix2freq(self):
        """
        pixel to frequency conversion in the PSD plane
        """
        return 1.0/(self.system.D*self._samp_over)
    
    #@_lru_cache(maxsize=2)
    @property
    def _nxnyOver(self):
        """
        Return the number of pixels for the correctly sampled array (at least at Shannon-Nyquist)
        """
        return self.npix[0]*self._k, self.npix[1]*self._k
    
    #@_lru_cache(maxsize=2)
    @property
    def _nullFreqIndex(self):
        nx,ny = self._nxnyOver
        return nx//2, ny//2
    
    #@_lru_cache(maxsize=2)
    @property
    def _fxfy(self):
        """
        Return the array of frequencies (correctly sampled above Shannon-Nyquist)
        """
        return self._xyarray * self._pix2freq
    
    def strehlOTF(self,x0):
        """Compute the Strehl ratio based on the sum of the OTF"""
        otf = self.otf(x0,_caller='self')
        otf_diff = self._otfDiffraction()
        return _np.real(_np.sum(otf)/_np.sum(otf_diff))
    
    def strehlMarechal(self,x0):
        """Compute the Strehl ratio based on the Maréchal approximation"""
        _,sig2 = self.psd(x0)
        return _np.exp(-sig2)
    
    def psd(self,x0,grad=False):
        raise ValueError("ParametricPSFfromPSD is not to be instantiated. the `psd` method must be override in the subclasses")
    
    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',grad=False):
        """
        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)")
        
        otfDiffr = self._otfDiffraction()
        otfShift = self._otfShift(dx,dy)
        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)
        
        otf = otfTurb * otfDiffr * otfShift
        if grad: return otf,g
        return otf
    
    def _otfTurbulent(self,x0,grad=False):
        L = self.system.D * self._samp_over
        if grad:
            PSD,integral,g,integral_g = self.psd(x0,grad=True)
        else:
            PSD,integral = self.psd(x0,grad=False)
            
        Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
        #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:
            g2 = _np.zeros(g.shape,dtype=complex) # I cannot override 'g' here due to float to complex type
            for i in range(len(x0)):
                Bg = _fft.fft2(_fft.fftshift(g[i,...])) / L**2
                #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
                g2[i,...] = -Bphi*otf
            return otf, g2
        return otf
    
    #@_lru_cache(maxsize=2)
    def _otfDiffraction(self):
        nx, ny = self._nxnyOver
        NpupX = _np.ceil(nx/self._samp_over)
        NpupY = _np.ceil(ny/self._samp_over)
        tab = _np.zeros((nx, ny), dtype=complex)
        tab[0:int(NpupX), 0:int(NpupY)] = self.system.pupil((NpupX,NpupY),samp=self._samp_over)
        return _fft.fftshift(_np.abs(_fft.ifft2(_np.abs(_fft.fft2(tab)) ** 2)) / _np.sum(tab))
    
    #@_lru_cache(maxsize=2)
    def _otfShift(self,dx,dy):
        Y, X = self._xyarray
        ny,nx = X.shape
        return _np.exp(-2j*_np.pi*self._k*(dx*X/nx + dy*Y/ny))
    
    def __call__(self,x0,dx=0,dy=0,grad=False):
        """
        Parameters
        ----------
        x0 : numpy.array (dim=1), tuple, list
            Array of parameters for the PSD (see __doc__)
        dx : float
            PSF X shifting [pix] (default = 0)
        dy : float
            PSF Y shifting [pix] (default = 0)
        """
        
        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)
        
        psf = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(otf))))
            
        k = int(self._k)
        if k==1:
            if grad: return psf,g
            return psf
        else:
            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,...].astype(float),k)
                return _binning(psf,k), g2
            return _binning(psf,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]
    """
    def __init__(self,npix,system=None,samp=None):
        """
        Parameters
        ----------
        npix : tuple
            Size of output PSF
        system : OpticalSystem
            Optical system for this PSF
        samp : float
            Sampling at the observation wavelength
        """
        super().__init__(2,npix,system=system,samp=samp)
        
        # r0,L0
        bounds_down = [_EPSILON, _EPSILON]
        bounds_up = [_np.inf, _np.inf]
        self.bounds = (bounds_down,bounds_up)
    
    def psd(self,x0,grad=False):
        """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)
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
        
        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
        fmax = _np.min([nx0,ny0])*self._pix2freq
        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)
        integral = integral_in + integral_out
        
        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)
            
            # 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
        
        if grad: return PSD,integral,g,integral_g
        return PSD, integral

#%% PSFAO MODEL        
class Psfao(ParametricPSFfromPSD):
    """
    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
    """
    def __init__(self,npix,system=None,Lext=10.,samp=None):
        """
        Parameters
        ----------
        npix : tuple
            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
        """
        super().__init__(7,npix,system=system,samp=samp)
        
        self.Lext = Lext
        # 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]
        self.bounds = (bounds_down,bounds_up)
    
    def psd(self,x0,grad=False):
        """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
        grad : bool (default=False)
            Return both (psd,integral,gradient) if set to True
            Warning: not finished yet, do not use!
            
        Returns
        -------
        psd : numpy.array (dim=2)
        integral : float : the integral of the `psd` up to infinity
            
        """
        self.check_parameters(x0)
        nx0,ny0 = self._nullFreqIndex
        f2D = self._fxfy
        pix = self._pix2freq
        F2 = f2D[0]**2. + f2D[1]**2.
        Fao = self.system.Nact/(2.0*self.system.D)
        maskin = (F2 < Fao**2.)
        
        r0,C,A,alpha,ratio,theta,beta = x0
        
        PSD = 0.0229* r0**(-5./3.) * ((1. / self.Lext**2.) + F2)**(-11./6.)
        PSD *= (1-maskin)
        
        ax = alpha*ratio
        ay = alpha/ratio
        param = (ax,ay,theta,beta,0,0)
        
        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)
        moff[nx0,ny0] = 0.0 # Set Moffat PSD = 0 at null frequency
        #moff = moff / (_np.sum(moff)*pix**2)  # normalize moffat numerically to get correct A=sigma² in the AO zone
        # Warning: Moffat numerical normalization generates strehlOTF jump when "_k" is changed
        
        PSD += maskin * (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
        
        # Compute PSD integral up to infinity
        fmax = _np.min([nx0,ny0])*pix
        integral_in = _np.sum(PSD*(F2<(fmax**2))) * pix**2 # numerical sum
        integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum
        integral = integral_in + integral_out
        
        # TODO: finish gradient computation!
        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
            g[3,...] = A*maskin*(dm[0,...]*ratio + dm[1,...]/ratio)
            # derivative towards ratio
            g[4,...] = A*maskin*(dm[0,...]*alpha - dm[1,...]*ay/ratio)
            # derivative towards theta
            g[5,...] = A*maskin*dm[2,...]
            # derivative towards beta
            g[6,...] = A*maskin*dm[3,...]
            # Remove central freq from all derivative
            g[:,nx0,ny0] = 0
        
            # 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
        return PSD, integral
        
    def tofits(self,param,filename,*args,keys=None,overwrite=False,**kwargs):
        if keys is None:
            keys = ["R0","CST","SIGMA2","ALPHA","RATIO","THETA","BETA"]
            keys_comment = ["Fried parameter [m]",
                            "PSD AO area constant C [rad2]",
                            "PSD AO area Moffat variance A [rad2]",
                            "PSD AO area Moffat alpha [1/m]",
                            "PSD AO area Moffat sqrt(ax/ay) ratio",
                            "PSD AO area Moffat theta [rad]",
                            "PSD AO area Moffat beta"]
        
        # 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")
        hdr["HIERARCH SYSTEM D"] = (self.system.D,"Primary mirror diameter")
        hdr["HIERARCH SYSTEM NACT"] = (self.system.Nact,"Linear number of AO actuators")
        hdr["HIERARCH SAMP"] = (self.samp,"Sampling (eg. 2 for Shannon)")
        hdr["HIERARCH LEXT"] = (self.Lext,"Von-Karman outer scale")
        
        hdu = _fits.PrimaryHDU(psf, hdr)
        hdu.writeto(filename, overwrite=overwrite)