Merge pull request #4 from ojdf/dev

Optimisations

fast/ao_power_spectra.py

@@ -159,7 +159,7 @@ def Jol_noise_openloop(freq, Dsubap, noise_variance, lf_mask):
 
     return lf_mask * powerspec
 
-def Jol_alias_openloop(freq, Dsubap, p, lf_mask, v=None, Delta_t=None, wvl=None, lmax=10, kmax=10, L0=numpy.inf, l0=1e-6):
+def Jol_alias_openloop(freq, Dsubap, p, lf_mask, v=None, Delta_t=None, wvl=None, lmax=3, kmax=3, L0=numpy.inf, l0=1e-6):
     fx = freq.fx
     fy = freq.fy
     fabs = freq.fabs
@@ -183,12 +183,14 @@ def Jol_alias_openloop(freq, Dsubap, p, lf_mask, v=None, Delta_t=None, wvl=None,
     else:
         v_dot_kappa = 0
 
-    sinc_term = numpy.sinc(Delta_t * v_dot_kappa / (2*numpy.pi))**2
 
     #TODO: the central pixel may still be wrong
     with warnings.catch_warnings(): 
         warnings.filterwarnings("ignore", category=RuntimeWarning) # avoid annoying RuntimeWarnings
 
+        sinc_term = numpy.sinc(Delta_t * v_dot_kappa / (2*numpy.pi))**2
+        term_0 = freq.fx**2 * freq.fy**2 / freq.fabs**4
+
         for l in ls:
             for k in ks:
                 if l == 0 and k == 0:
@@ -201,7 +203,7 @@ def Jol_alias_openloop(freq, Dsubap, p, lf_mask, v=None, Delta_t=None, wvl=None,
 
                 term_1 = (freq.fx/(freq_shift.fy) + freq.fy/(freq_shift.fx))**2
                 term_2 = funcs.turb_powerspectrum_vonKarman(freq_shift, p, L0=L0, l0=l0)
-                mult = term_1 * term_2 *  freq.fx**2 * freq.fy**2 / freq.fabs**4
+                mult = term_1 * term_2 * term_0
                 mult[...,midpt_x,midpt_y] = 0.
                 if l == 0:
                     mult[...,midpt_x,:] = term_2[...,midpt_x,:]

fast/fast.py

@@ -381,8 +381,8 @@ def init_pupil_mask(self):
             N_req = int(2*numpy.ceil(2*numpy.pi/(self.freq.main.df * dx_req)/2)) # ensure even
 
             pupil_temporal = funcs.compute_pupil(N_req, dx_req, self.D_ground,
-                self.W0, Tx_obsc=self.obsc_ground, ptype=ptype)
-            self.freq.make_logamp_freqs(N=N_req, dx=dx_req)
+                self.W0, Tx_obsc=self.obsc_ground, ptype=ptype, Ny=2*self.Npxls_pup)
+            self.freq.make_logamp_freqs(Nx=N_req, dx=dx_req, Ny=2*self.Npxls_pup, dy=self.dx)
             self.pupil_filter_temporal = funcs.pupil_filter(self.freq.logamp, pupil_temporal, spline=True)
 
         self.smf = self.params['SMF']
@@ -836,14 +836,16 @@ def make_temporal_freqs(self, nlayer, Ny, Nx, wind_speed, wind_dir, dt):
 
         self.temporal = SpatialFrequencyStruct(numpy.array(fx_axes), numpy.array(fy_axes), rot=numpy.radians(wind_dir), freq_per_layer=True)
 
-    def make_logamp_freqs(self, N=None, dx=None):
-        if N is None and dx is None:
+    def make_logamp_freqs(self, Nx=None, dx=None, Ny=None, dy=None):
+        if Nx is None and dx is None:
             self.logamp = self.main
 
         else:
-            df = 2*numpy.pi/(N*dx)
-            fx_axis = numpy.arange(-N/2., N/2.) * df
-            self.logamp = SpatialFrequencyStruct(fx_axis)
+            dfx = 2*numpy.pi/(Nx*dx)
+            fx_axis = numpy.arange(-Nx/2., Nx/2.) * dfx
+            dfy = 2*numpy.pi/(Ny*dy)
+            fy_axis = numpy.arange(-Ny/2., Ny/2.) * dfy
+            self.logamp = SpatialFrequencyStruct(fx_axis, fy_axis)
 
 class SpatialFrequencyStruct():
 

fast/funcs.py

@@ -257,9 +257,22 @@ def make_phase_subharm(rand, freq, N, dx, double=False):
     else:
         return phs_lo.real
 
-def compute_pupil(N, dx, Tx, W0=None, Tx_obsc=0, Raxicon=None, ptype='gauss'):
+def compute_pupil(N, dx, Tx, W0=None, Tx_obsc=0, Raxicon=None, ptype='gauss', Ny=None):
+
     circ_ap = circle(Tx/dx/2, N) - circle(Tx_obsc/dx/2, N)
 
+    if Ny != None:
+        Nx = N
+        assert ((Ny-Nx)%2)==0, f"(Nx-Ny)/2 must be even"
+        if Ny > Nx:
+            Npad = (Ny - Nx)//2
+            circ_ap = numpy.pad(circ_ap, [(0,0),(Npad,Npad)])
+        if Ny < Nx:
+            Ncut = (Nx - Ny)//2
+            circ_ap = circ_ap[:,Ncut:-Ncut]
+    else:
+        Nx = Ny = N
+
     if ptype == 'circ':
         return circ_ap / numpy.sqrt(circ_ap.sum()*dx**2)
 
@@ -270,13 +283,14 @@ def compute_pupil(N, dx, Tx, W0=None, Tx_obsc=0, Raxicon=None, ptype='gauss'):
             return g * circ_ap, opt
         else:
             I0 = 2 / (numpy.pi * W0**2)
-            return gaussian2d(N, W0/dx/numpy.sqrt(2)) * circ_ap * numpy.sqrt(I0)
+            return gaussian2d((Nx, Ny), W0/dx/numpy.sqrt(2)) * circ_ap * numpy.sqrt(I0)
 
     elif ptype == 'axicon':
         if W0 == "opt":
             raise TypeError("Using 'axicon' and W0='opt' not supported, please set a value for W0")
-        x = numpy.arange(-N/2, N/2, 1) * dx
-        xx, yy = numpy.meshgrid(x,x)
+        x = numpy.arange(-Nx/2, Nx/2, 1) * dx
+        y = numpy.arange(-Ny/2, Ny/2, 1) * dx
+        xx, yy = numpy.meshgrid(x,y)
         r = numpy.sqrt(xx**2 + yy**2)
         if Raxicon == None:
             midpt = Tx_obsc/2 + (Tx/2-Tx_obsc/2)/2
@@ -299,16 +313,16 @@ def pupil_filter(freq, pupil, spline=False):
         return P
 
 def optimize_fibre(pupil, dx, size_min=None, size_max=None, return_size=False):
-    N = pupil.shape[-1]
+    Nx, Ny = pupil.shape
 
     if size_max is None:
-        size_max = N * dx
+        size_max = max(Ny,Nx) * dx
 
     if size_min is None:
         size_min = dx
 
     def _opt_func(W):
-        return coupling_loss(W, N, pupil, dx)
+        return coupling_loss(W, (Nx, Ny), pupil, dx)
 
     opt = minimize_scalar(_opt_func, bracket=[size_min, size_max]).x
 
@@ -321,7 +335,7 @@ def _opt_func(W):
         if abs(opt) < dx:
             raise Exception("Cannot optimise gaussian mode, try changing DX?")
 
-    g = gaussian2d(N, opt/dx/numpy.sqrt(2)) * numpy.sqrt(2./(numpy.pi * opt**2))
+    g = gaussian2d((Nx,Ny), opt/dx/numpy.sqrt(2)) * numpy.sqrt(2./(numpy.pi * opt**2))
 
     if return_size:
         return g, numpy.abs(opt)