feat: incorporate 2D-DFT to obtain 2D-Frank image

Author: mariajmellado

frank/fitExample/AS209_1mm_FrankFit.py renamed to frank/Examples/AS209_1mm_FrankFit.py

@@ -0,0 +1,77 @@
+
+import numpy as np
+
+class DiscreteFourierTransform2D(object):
+    def __init__(self, Rmax, N, nu=0):
+        # Remember that now N is to create N**2 points in image plane.
+        self.Xmax  = 2*Rmax # [Rmax] = rad.
+        self.Ymax = 2*Rmax
+        self.Nx = N 
+        self.Ny = N
+        self.N = N**2 # Number of points we want to use in the 2D-DFT.
+
+        # Real space collocation points
+        x1n = np.linspace(0, self.Xmax, self.Nx) # rad
+        x2n = np.linspace(0, self.Ymax, self.Ny) # rad
+        x1n, x2n = np.meshgrid(x1n, x1n, indexing='ij')
+        # x1n.shape = N**2 X 1, so now, we have N**2 collocation points in the image plane.
+        x1n, x2n = x1n.reshape(-1), x2n.reshape(-1) # x1n = x_array and x2n = y_array
+        dx = 2*self.Xmax/self.N
+        dy = 2*self.Ymax/self.N
+
+        # Frequency space collocation points.
+        # The [1:] is because to not consider the 0 baseline. But we're missing points. 
+        q1n = np.fft.fftfreq(self.Nx+1, d = dx)[1:]
+        q2n = np.fft.fftfreq(self.Ny+1, d = dy)[1:]
+        q1n, q2n = np.meshgrid(q1n, q2n, indexing='ij') 
+        # q1n.shape = N**2 X 1, so now, we have N**2 collocation points.
+        q1n, q2n = q1n.reshape(-1), q2n.reshape(-1) # q1n = u_array and q2n = v_array
+
+        self.Xn = x1n
+        self.Yn = x2n
+        self.Un = q1n
+        self.Vn = q2n
+
+    def get_collocation_points(self):        
+        return np.array([self.Xn, self.Yn]), np.array([self.Un, self.Vn])
+
+    def coefficients(self, u = None, v = None, direction="forward"):
+        if direction == 'forward':
+            norm = 1
+            factor = -2j*np.pi/self.Nx
+        elif direction == 'backward':
+            norm = 1 / self.N
+            factor = 2j*np.pi/self.Nx
+        else:
+            raise AttributeError("direction must be one of {}"
+                                 "".format(['forward', 'backward']))
+        if u is None:
+            u = self.Un
+            v = self.Vn
+        
+        if direction == "forward":
+            H = norm * np.exp(factor*(np.outer(u, self.Xn) + np.outer(v, self.Yn)))
+        else:
+            H = norm * np.exp(factor*(np.outer(u, self.Xn) + np.outer(v, self.Yn)))
+        return H
+
+
+    def transform(self, f, u, v, direction="forward"):
+        Y = self.coefficients(u, v, direction=direction)
+        return np.dot(Y, f)
+          
+          
+    @property
+    def size(self):
+        """Number of points used in the 2D-DFT"""
+        return self.N
+    
+    @property
+    def uv_points(self):
+        """u and v  collocation points"""
+        return self.Un, self.Vn
+
+    @property
+    def q(self):
+        """Frequency points"""
+        return np.hypot(self.Un, self.Vn)

frank/Examples/DFT_2D_test.ipynb

@@ -28,6 +28,7 @@
 from frank.constants import rad_to_arcsec
 from frank.filter import CriticalFilter
 from frank.hankel import DiscreteHankelTransform
+from frank.fourier2d import DiscreteFourierTransform2D
 from frank.statistical_models import (
     GaussianModel, LogNormalMAPModel, VisibilityMapping
 )
@@ -443,6 +444,7 @@ def __init__(self, Rmax, N, geometry, nu=0, block_data=True,
         self._geometry = geometry
 
         self._DHT = DiscreteHankelTransform(Rmax, N, nu)
+        self._DFT = DiscreteFourierTransform2D(Rmax, N)
 
         if assume_optically_thick:
             if scale_height is not None:
@@ -457,7 +459,8 @@ def __init__(self, Rmax, N, geometry, nu=0, block_data=True,
         self._vis_map = VisibilityMapping(self._DHT, geometry, 
                                           model, scale_height=scale_height,
                                           block_data=block_data, block_size=block_size,
-                                          check_qbounds=False, verbose=verbose)
+                                          check_qbounds=False, verbose=verbose,
+                                          DFT = self._DFT)
 
         self._info  = {'Rmax' : self._DHT.Rmax * rad_to_arcsec,
                        'N' : self._DHT.size
@@ -577,7 +580,7 @@ def _fit(self):
         """Fit step. Computes the best fit given the pre-processed data"""
         fit = GaussianModel(self._DHT, self._M, self._j,
                             noise_likelihood=self._H0,
-                            Wvalues= self._Wvalues, V = self._V)
+                            Wvalues= self._Wvalues, V = self._V, DFT = self._DFT)
 
         self._sol = FrankGaussianFit(self._vis_map, fit, self._info,
                                      geometry=self._geometry.clone())

frank/fourier2d.py

@@ -66,7 +66,8 @@ class VisibilityMapping:
     """
     def __init__(self, DHT, geometry,  
                  vis_model='opt_thick', scale_height=None, block_data=True,
-                 block_size=10 ** 5, check_qbounds=True, verbose=True):
+                 block_size=10 ** 5, check_qbounds=True, verbose=True,
+                 DFT = None):
 
         _vis_models = ['opt_thick', 'opt_thin', 'debris']
         if vis_model not in _vis_models:
@@ -81,6 +82,7 @@ def __init__(self, DHT, geometry,
         self._chunk_size = block_size
 
         self._DHT = DHT
+        self._DFT = DFT
         self._geometry = geometry
 
         # Check for consistency and report the model choice.
@@ -178,8 +180,8 @@ def map_visibilities(self, u, v, V, weights, frequencies=None, geometry=None):
             frequencies = np.ones_like(V)
 
         channels = np.unique(frequencies)
-        Ms = np.zeros([len(channels), self.size, self.size], dtype='f8')
-        js = np.zeros([len(channels), self.size], dtype='f8')
+        Ms = np.zeros([len(channels), self.size, self.size], dtype='c8')
+        js = np.zeros([len(channels), self.size], dtype='c8')
         for i, f in enumerate(channels):
             idx = frequencies == f
 
@@ -199,11 +201,13 @@ def map_visibilities(self, u, v, V, weights, frequencies=None, geometry=None):
             Ndata = len(Vi)
             while start < Ndata:
                 qs = qi[start:end]
+                us = u[start:end]
+                vs = v[start:end]
                 ks = ki[start:end]
                 ws = wi[start:end]
                 Vs = Vi[start:end]
 
-                X = self._get_mapping_coefficients(qs, ks)
+                X = self._get_mapping_coefficients(qs, ks, us, vs)
 
                 wXT = np.array(X.T * ws, order='C')
 
@@ -462,7 +466,7 @@ def interpolate(self, f, r, space='Real'):
         return self._DHT.interpolate(f, r, space)
 
 
-    def _get_mapping_coefficients(self, qs, ks, geometry=None, inverse=False):
+    def _get_mapping_coefficients(self, qs, ks, u, v, geometry=None, inverse=False):
         """Get :math:`H(q)`, such that :math:`V(q) = H(q) I_\nu`"""
 
         if self._vis_model == 'opt_thick':
@@ -486,7 +490,8 @@ def _get_mapping_coefficients(self, qs, ks, geometry=None, inverse=False):
         else:
             direction='forward'
 
-        H = self._DHT.coefficients(qs, direction=direction) * scale
+        #H = self._DHT.coefficients(qs, direction=direction) * scale
+        H = self._DFT.coefficients(u, v, direction=direction) * scale
 
         return H
 
@@ -529,7 +534,8 @@ def Rmax(self):
     @property
     def q(self):
         r"""Frequency points, unit = :math:`\lambda`"""
-        return self._DHT.q
+        #return self._DHT.q
+        return self._DFT.q
 
     @property
     def Qmax(self):
@@ -539,7 +545,8 @@ def Qmax(self):
     @property
     def size(self):
         """Number of points in reconstruction"""
-        return self._DHT.size
+        #return self._DHT.size
+        return self._DFT.size
 
     @property
     def scale_height(self):
@@ -631,9 +638,10 @@ class GaussianModel:
 
     def __init__(self, DHT, M, j, p=None, scale=None, guess=None,
                  Nfields=None, noise_likelihood=0,
-                 Wvalues = None, V = None):
+                 Wvalues = None, V = None, DFT = None):
 
         self._DHT = DHT
+        self._DFT = DFT
         self._Wvalues = Wvalues
         self._V = V
 
@@ -691,24 +699,13 @@ def __init__(self, DHT, M, j, p=None, scale=None, guess=None,
                 en = (n+1)*Nr 
 
                 self._Sinv[sn:en, sn:en] += Sj[n]
-        else:
+        else: # p is None, so we are interested in this case.
+            " New GP Kernel below"
             #self._Sinv =  None
-            q_array = self._DHT.q
-            
-            def true_squared_exponential_kernel(q, p, l):
-
-                q1, q2 = np.meshgrid(q, q)
-                p1, p2 = np.meshgrid(p, p)
-                SE_Kernel = np.sqrt(p1 * p2) * np.exp(-0.5*(q1-q2)**2 / l**2)
-                return SE_Kernel
-            
-            Ykm = self._DHT.coefficients(direction="backward")
-            # We continue after set M matrix because is needed to calculate
-            # best parameters for S matrix.
 
         # Compute the design matrix
-        self._M = np.zeros([Nr*Nfields, Nr*Nfields], dtype='f8')
-        self._j = np.zeros(Nr*Nfields, dtype='f8')
+        self._M = np.zeros([Nr*Nfields, Nr*Nfields], dtype='c8')
+        self._j = np.zeros(Nr*Nfields, dtype='c8')
         for si, Mi, ji in zip(self._scale, M, j):
 
             for n in range(0, Nfields):
@@ -724,35 +721,42 @@ def true_squared_exponential_kernel(q, p, l):
 
         self._like_noise = noise_likelihood
 
-        # M is already defined, so we find best parameters for S matrix and use it.
-        m, c , l = self.minimizeS()
-        pI =  np.exp(m*np.log(q_array) + c)
-        S_fspace = true_squared_exponential_kernel(q_array, pI, l)
-        S_real = np.dot(np.transpose(Ykm), np.dot(S_fspace, Ykm))
+        " New GP "
+        self.u, self.v = self._DFT.uv_points
+        self.Ykm = self._DFT.coefficients(direction="backward")
+        self.Ykm_f = self._DFT.coefficients(direction="forward")
+        m, c , l = -4.8, 59.21, 1.21e5 
+        #m, c, l = self.minimizeS() # Finding best parameters to S matrix.
+        S_real = self.calculate_S(self.u, self.v, l, m, c)
         S_real_inv = np.linalg.inv(S_real)
         self._Sinv =  S_real_inv
 
         self._fit()
 
+    def true_squared_exponential_kernel(self, u, v, l, m, c):      
+        u1, u2 = np.meshgrid(u, u)
+        v1, v2 = np.meshgrid(v, v)
+        q1 = np.hypot(u1, v1)
+        q2 = np.hypot(u2, v2)
+
+        def power_spectrum(q, m, c):
+            return np.exp(m*np.log(q) + c)
+        p1 = power_spectrum(q1, m, c)
+        p2 = power_spectrum(q2, m, c)
+
+        SE_Kernel = np.sqrt(p1 * p2) * np.exp(-0.5*((u1-u2)**2 + (v1-v2)**2)/ l**2)
+        return SE_Kernel
+
+    def calculate_S(self, u, v, l, m, c):
+        S_fspace = self.true_squared_exponential_kernel(u, v, l, m, c)
+        S_real = np.matmul(self.Ykm, np.matmul(S_fspace, self.Ykm_f))
+        return S_real
+        
     def minimizeS(self):
         from scipy.optimize import minimize
         from scipy.special import gamma
         V = self._V
 
-        def calculate_S(m, c, l):
-            q_array = self._DHT.q
-            p_array = c*(q_array**m)
-            def true_squared_exponential_kernel(q, p, l):
-                q1, q2 = np.meshgrid(q, q)
-                p1, p2 = np.meshgrid(p, p)
-                SE_Kernel = np.sqrt(p1 * p2) * np.exp(-0.5*(q1-q2)**2 / l**2)
-                return SE_Kernel
-            
-            Ykm = self._DHT.coefficients(direction="backward")
-            S_fspace = true_squared_exponential_kernel(q_array, p_array, l)
-            S_real = np.dot(np.transpose(Ykm), np.dot(S_fspace, Ykm))
-            return S_real
-
         def calculate_D(S):
             S_real_inv = np.linalg.inv(S)
             Dinv = self._M + S_real_inv
@@ -782,7 +786,7 @@ def likelihood(param, data):
             def inv_gamma_function(l, alpha, beta):
                 return ((gamma(alpha)*beta)**(-1))*((beta/l)**(alpha + 1))*np.exp(-beta/l) 
 
-            S  = calculate_S(m,c, l)
+            S  = self.calculate_S(self.u, self.v, l, m, c)
             [Dinv, D] = calculate_D(S)
             mu = calculate_mu(Dinv)
             logdetS = np.linalg.slogdet(S)[1]
@@ -796,18 +800,17 @@ def inv_gamma_function(l, alpha, beta):
             - 0.5*(factor + logdetS) \
             + 0.5*(factor + logdetD) \
             + 0.5*np.dot(np.transpose(self._j), mu) \
-            - 0.5*np.dot(np.transpose(data), np.dot(np.diag(Wvalues), data)) 
+            - 0.5*np.dot(np.transpose(np.conjugate(data)), np.dot(np.diag(Wvalues), data)) 
             return -log_likelihood
 
         result = minimize(likelihood, x0=np.array([-5, 60, 1e5]), args=(V,),
                            bounds=[(-6, 6), (1, 70), (1e4, 1e6)],
                            method="Nelder-Mead", tol=1e-7,
                            )
         m, c, l = result.x
-        print("Result: ", "m: ", m, "c: ", c, "l: ", "{:e}".format(l))
+        print("Best parameters: ", "m: ", m, "c: ", c, "l: ", "{:e}".format(l))
         return [m, c, l]
 
-
     def _fit(self):
         """Compute the mean and variance"""
         # Compute the inverse prior covariance, S(p)^-1
@@ -980,7 +983,8 @@ def num_fields(self):
     @property
     def size(self):
         """Number of points in reconstruction"""
-        return self._DHT.size
+        #return self._DHT.size
+        return self._DFT.size
 
 
 class LogNormalMAPModel: