Commit 63b26f3b authored by rfetick's avatar rfetick
Browse files

precompute fx2, fy2 and fxy arrays for moffat

parent b8fe833d
......@@ -5,8 +5,12 @@ Created on Tue Aug 31 17:16:53 2021
Run profiler on this script to find computation bottleneck
Results:
31 Aug 2021 - precompute _otfDiffraction - time drops from 160 to 122 ms/iter (100% -> 76%)
09 Sep 2021 - remove some fftshift, use scipy.fft - time drops from 110 to 97 ms/iter (100% -> 88%)
31 Aug 2021 - precompute _otfDiffraction
time drops from 160 to 122 ms/iter (100% -> 76%)
09 Sep 2021 - remove some fftshift, use scipy.fft
time drops from 110 to 97 ms/iter (100% -> 88%)
17 Nov 2021 - precompute fx2, fy2 and fxy for moffat
time drops from 95 to 83 ms/iter (100% -> 87%)
@author: rfetick
"""
......@@ -15,13 +19,13 @@ from time import time
from maoppy.psfmodel import Psfao
from maoppy.instrument import muse_nfm
Npix = 128 # pixel size of PSF
npix = 128 # pixel size of PSF
wvl = 600*1e-9 # wavelength [m]
niter = 20 # number of PSF to generate
niter = 50 # number of PSF to generate
#%% Initialize PSF model
samp = muse_nfm.samp(wvl) # sampling (2.0 for Shannon-Nyquist)
Pmodel = Psfao((Npix,Npix),system=muse_nfm,samp=samp)
Pmodel = Psfao((npix,npix), system=muse_nfm, samp=samp)
#%% Choose parameters and compute PSF
r0 = 0.15 # Fried parameter [m]
......@@ -36,6 +40,6 @@ param = [r0,bck,amp,alpha,ratio,theta,beta]
t0 = time()
for i in range(niter):
psf = Pmodel(param,dx=0,dy=0)
psf = Pmodel(param, dx=0, dy=0)
t1 = time()
print("Elapsed: %u ms/iter"%round((t1-t0)*1e3/niter))
\ No newline at end of file
......@@ -16,8 +16,9 @@ from maoppy.utils import binning
__all__ = ["oversample", "lsq_flux_bck", "psffit", "reduced_coord",
"moffat", "gauss", "ParametricPSF", "ConstantPSF", "Moffat",
"Gaussian", "ParametricPSFfromPSD", "Turbulent", "Psfao"]
"reduced_center_coord", "moffat", "moffat_center", "gauss",
"ParametricPSF", "ConstantPSF", "Moffat", "Gaussian",
"ParametricPSFfromPSD", "Turbulent", "Psfao"]
# Value to compute finite differences
_EPSILON = _np.sqrt(sys.float_info.epsilon)
......@@ -218,7 +219,15 @@ def psffit(psf,Model,x0,weights=None,dxdy=(0.,0.),flux_bck=(True,True),
#%% BASIC FUNCTIONS
def reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=False):
def reduced_coord(xxyy, ax, ay, theta, cx, cy, grad=False):
"""Create an array of reduced coordinated with elongation and rotation"""
uxx = (xxyy[0]-cx)**2
uxy = (xxyy[0]-cx)*(xxyy[1]-cy)
uyy = (xxyy[1]-cy)**2
return reduced_center_coord(uxx, uxy, uyy, ax, ay, theta, grad=grad)
def reduced_center_coord(uxx, uxy, uyy, ax, ay, theta, grad=False):
"""Create an array of reduced coordinated with elongation and rotation"""
c = _np.cos(theta)
s = _np.sin(theta)
......@@ -228,10 +237,6 @@ def reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=False):
rxy = s2/ay**2 - s2/ax**2
ryy = (c/ay)**2 + (s/ax)**2
uxx = (xxyy[0]-cx)**2
uxy = (xxyy[0]-cx)*(xxyy[1]-cy)
uyy = (xxyy[1]-cy)**2
uu = rxx*uxx + rxy*uxy + ryy*uyy
if grad:
......@@ -247,13 +252,13 @@ def reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=False):
gg[2,...] = drxx*uxx + drxy*uxy + dryy*uyy
return uu, gg
return uu
return uu
def moffat(xxyy, param, norm=None, removeInside=0, grad=False):
"""Compute a Moffat function on a meshgrid
"""Compute a Moffat function on a meshgrid [xx,yy]
```moff = E * (1+u)^(-beta)```
with `u` the quadratic coordinates in the shifted and rotated frame
with `u` the reduced quadratic coordinates in the shifted and rotated frame
and `E` the energy normalization factor
Parameters
......@@ -284,12 +289,53 @@ def moffat(xxyy, param, norm=None, removeInside=0, grad=False):
"""
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
_,_,_,_,cx,cy = param
uxx = (xxyy[0]-cx)**2
uxy = (xxyy[0]-cx)*(xxyy[1]-cy)
uyy = (xxyy[1]-cy)**2
return moffat_center(uxx, uxy, uyy, param[:-2], norm=norm, removeInside=removeInside, grad=grad)
def moffat_center(uxx, uxy, uyy, param, norm=None, removeInside=0, grad=False):
"""Compute a Moffat function, given the quadratic coordinates
```moff = E * (1+u)^(-beta)```
with `u` the reduced quadratic coordinates in the shifted and rotated frame
and `E` the energy normalization factor
Parameters
----------
xxyy : numpy.ndarray (dim=2)
The (xx,yy) meshgrid with xx = xxyy[0] and yy = xxyy[1]
param : list, tuple, numpy.ndarray (len=6)
param[0] - Alpha X
param[1] - Alpha Y
param[2] - Theta
param[3] - Beta
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)!=4:
raise ValueError("Parameter `param` must contain exactly 4 elements, but input has %u elements"%len(param))
ax,ay,theta,beta = param
if grad:
uu,du = reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=True)
uu,du = reduced_center_coord(uxx, uxy, uyy, ax, ay, theta, grad=grad)
else:
uu = reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=False)
uu = reduced_center_coord(uxx, uxy, uyy, ax, ay, theta, grad=grad)
V = (1.0+uu)**(-beta) # Moffat shape
E = 1.0 # normalization factor (eventually up to infinity)
......@@ -566,8 +612,13 @@ class ParametricPSFfromPSD(ParametricPSF):
xyarray = _np.mgrid[0:nx, 0:ny].astype(float)
xyarray[0] -= nx//2
xyarray[1] -= ny//2
self._fxfy = xyarray * self.pix2freq
self._f2 = self._fxfy[0]**2. + self._fxfy[1]**2.
self._fx_fy = xyarray * self.pix2freq
self._xx = xyarray[1]/ny
self._yy = xyarray[0]/nx
self._fx2 = self._fx_fy[0]**2.
self._fy2 = self._fx_fy[1]**2.
self._fxy = self._fx_fy[0]*self._fx_fy[1]
self._f2 = self._fx2 + self._fy2
@property
......@@ -697,15 +748,13 @@ class ParametricPSFfromPSD(ParametricPSF):
def _otfShift(self, dx, dy):
"""Shifting part of the OTF"""
yy, xx = self._fxfy/self.pix2freq # so it is just an array of pixels
ny,nx = xx.shape
# Compensate oversampling shift
dx += (self._k-1)/(2*self._k)
dy += (self._k-1)/(2*self._k)
# Compensate odd pixel shift
dx -= (self.npix[1]%2)/2
dy -= (self.npix[0]%2)/2
return ifftshift(_np.exp(-2j*_np.pi*self._k*(dx*xx/nx + dy*yy/ny)))
return ifftshift(_np.exp(-2j*_np.pi*self._k*(dx*self._xx + dy*self._yy)))
def __call__(self, parampsd, dx=0, dy=0, grad=False):
......@@ -926,13 +975,15 @@ class Psfao(ParametricPSFfromPSD):
# Moffat
ax = alpha*ratio
ay = alpha/ratio
param = (ax,ay,theta,beta,0,0)
param = (ax,ay,theta,beta)
remove_inside = (1+_np.sqrt(2))/2 * pix/2 # remove central pixel in energy computation
if grad:
moff,dm = moffat(self._fxfy,param,norm=self.cutoffAOfreq,removeInside=remove_inside,grad=True)
moff,dm = moffat_center(self._fx2, self._fxy, self._fy2,
param, norm=self.cutoffAOfreq, removeInside=remove_inside, grad=True)
else:
moff = moffat(self._fxfy,param,norm=self.cutoffAOfreq,removeInside=remove_inside)
moff = moffat_center(self._fx2, self._fxy, self._fy2,
param, norm=self.cutoffAOfreq, removeInside=remove_inside, grad=False)
moff *= self._maskin
numerical_norm = False # set to true in order to activate code below
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment