continue gradient computation

maoppy/psfmodel.py

@@ -192,6 +192,7 @@ def psffit(psf,Model,x0,weights=None,dxdy=(0.,0.),flux_bck=(True,True),
 
 #%% 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)
@@ -221,11 +222,11 @@ def reduced_coord(XY,ax,ay,theta,cx,cy,grad=False):
     
     return u    
 
-def moffat(XY, param, norm=None, removeInside=0):
-    """
-    Compute a Moffat function on a meshgrid
-    moff = Enorm * (1+u)^(-beta)
+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
     ----------
@@ -241,14 +242,14 @@ def moffat(XY, param, norm=None, removeInside=0):
         
     Keywords
     --------
-    norm : None, _np.inf, float (>0), int (>0)
+    norm : None, np.inf, float (>0), int (>0)
         Radius for energy normalization
         None      - No energy normalization (maximum=1.0)
-                    Enorm = 1.0
-        _np.inf    - Total energy normalization (on the whole X-Y plane)
-                    Enorm = (beta-1)/(pi*ax*ay)
+                    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
-                    Enorm = (beta-1)/(pi*ax*ay)*(1-(1+(R**2)/(ax*ay))**(1-beta))    
+                    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
@@ -256,19 +257,35 @@ def moffat(XY, param, norm=None, removeInside=0):
     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 grad:
+        u,ug = reduced_coord(XY,ax,ay,theta,cx,cy,grad=True)
+        g = _np.zeros((4,)+u.shape)
+    else:
+        u = reduced_coord(XY,ax,ay,theta,cx,cy,grad=False)
+    
+    V = (1.0+u)**(-beta)
+    E = 1.0
     
     if norm is None:
-        Enorm = 1
+        if grad:
+            dVdu = -beta*V/(1.0+u)
+            g[0,...] = dVdu*ug[0,...]
+            g[1,...] = dVdu*ug[1,...]
+            g[2,...] = dVdu*ug[2,...]
+            g[3,...] = -V*_np.log(1.0+u)
+            return V,g
     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!")
-        Enorm = (beta-1) / (_np.pi*ax*ay)
-        Fout = (1 - (1 + (norm**2) / (ax*ay))**(1-beta))
-        Fin = (1 - (1 + (removeInside**2) / (ax*ay))**(1-beta))
-        Enorm = Enorm / (Fout-Fin)
+        E = (beta-1) / (_np.pi*ax*ay)
+        Fout = 1 - (1 + (norm**2)         / (ax*ay))**(1-beta)
+        Fin  = 1 - (1 + (removeInside**2) / (ax*ay))**(1-beta)
+        E = E / (Fout-Fin)
+        if grad:
+            # TODO: finish gradient computation
+            raise ValueError("Gradient not implemented yet!")
     
-    return Enorm * (1.0+u)**(-beta)
+    return V*E
 
 def gauss(XY,param):
     """

maoppy/test/test_psfmodel.py

@@ -7,16 +7,20 @@ Created on Tue Jun  1 16:30:27 2021
 
 import pytest
 import numpy as np
+import copy
 from maoppy.psfmodel import (reduced_coord, lsq_flux_bck, moffat, gauss,
                              ParametricPSFfromPSD, Psfao)
 from maoppy.instrument import muse_nfm, zimpol
 
+#%% DEFINITIONS
 def errL1(a,b):
+    """L1-norm error between two arrays"""
     avg = np.sum(np.abs(a)+np.abs(b))/2
     diff = np.sum(np.abs(a-b))
     if avg==0: return 0 # if the norms of a and b are both null!
     return diff/avg
 
+#%% TEST FUNCTIONS
 def test_reduced_coord_grad():
     ax = 3
     ay = 4
@@ -57,7 +61,7 @@ def test_fluxbck_positive():
     w = np.ones_like(image)
     f,b = lsq_flux_bck(model,image,w,positive_bck=True)
     assert b==0
-    
+
 def test_moffat_energy():
     npix = 512
     XY = np.mgrid[0:npix,0:npix] - npix/2.0
@@ -74,6 +78,23 @@ def test_gauss_energy():
     g = gauss(XY,param)
     assert np.sum(g)==pytest.approx(1.0,abs=1e-2)
     
+def test_moffat_grad():
+    npix = 512
+    XY = np.mgrid[0:npix,0:npix]-npix//2
+    param = [35,25,1.3,1.5,0,0]
+    # for norm=None
+    m,g = moffat(XY, param, norm=None, grad=True)
+    gnum = g*0
+    eps = 1e-6
+    for i in range(gnum.shape[0]):
+        dp = copy.copy(param)
+        dp[i] += eps
+        m2 = moffat(XY, dp, norm=None)
+        gnum[i,...] = (m2-m)/eps
+    assert errL1(g,gnum)==pytest.approx(0,abs=1e-5)
+    # TODO: similar gradient test for norm=np.inf or any float
+    
+#%% TEST PSFfromPSD and PSFAO MODELS
 def test_psffrompsd_oversampling():
     npix = 100
     samp = 0.3
@@ -124,7 +145,7 @@ def test_psfao_psd_integral():
     int_ana = A
     assert int_num==pytest.approx(int_ana,abs=1e-2)
     
-def test_otf_max():
+def test_psfao_otf_max():
     npix = 1024
     samp = 2.0
     P = Psfao((npix,npix),system=zimpol,samp=samp)
@@ -140,5 +161,6 @@ def test_otf_max():
     assert mtf==pytest.approx(1.0,abs=1e-2)
     assert mtf<=1.0
     
+#%% RUN FILE
 if __name__=="__main__":
     pytest.main()
\ No newline at end of file