hide imports with underscore

maoppy/psfmodel.py

@@ -10,10 +10,10 @@ import numpy as _np
 from scipy.optimize import least_squares as _least_squares
 from astropy.io import fits as _fits
 import time as _time
-from numpy.fft import fft2, fftshift, ifft2
+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
+from functools import lru_cache as _lru_cache
 
 # Value to compute finite differences
 _EPSILON = _np.sqrt(_sys.float_info.epsilon)
@@ -297,7 +297,7 @@ class ParametricPSF(object):
     def otf(self,*args,**kwargs):
         """Return the Optical Transfer Function (OTF)"""
         psf = self.__call__(args,kwargs)
-        return fftshift(fft2(psf))
+        return _fft.fftshift(_fft.fft2(psf))
     
     def mtf(self,*args,**kwargs):
         """Return the Modulation Transfer Function (MTF)"""
@@ -503,12 +503,12 @@ class Turbulent(ParametricPSF):
     def _otf_turbulent(self,x0):
         PSD = self.psd(x0)
         L = self.system.D * self._samp_over
-        Bg = fft2(fftshift(PSD)) / L**2
-        Dphi = fftshift(_np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV 
+        Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
+        Dphi = _fft.fftshift(_np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV 
         #Dphi = fftshift(_np.real(2 * (integral - Bg))) # normalized up to infinity
         return _np.exp(-Dphi/2.) 
     
-    @lru_cache(maxsize=2)
+    @_lru_cache(maxsize=2)
     def _otf_diffraction(self):
         Nx_over = self.Npix[0]*self._k
         Ny_over = self.Npix[1]*self._k
@@ -517,7 +517,7 @@ class Turbulent(ParametricPSF):
         NpupY = _np.ceil(Ny_over/self._samp_over)
         tab = _np.zeros((Nx_over, Ny_over), dtype=_np.complex)
         tab[0:int(NpupX), 0:int(NpupY)] = self.system.pupil((NpupX,NpupY),samp=self._samp_over)
-        return fftshift(_np.abs(ifft2(_np.abs(fft2(tab)) ** 2)) / _np.sum(tab))
+        return _fft.fftshift(_np.abs(_fft.ifft2(_np.abs(_fft.fft2(tab)) ** 2)) / _np.sum(tab))
     
     def _otf_shift(self,dx,dy):
         Nx_over = self.Npix[0]*self._k
@@ -541,7 +541,7 @@ class Turbulent(ParametricPSF):
             PSF Y shifting [pix] (default = 0)
         """
         
-        out = _np.real(fftshift(ifft2(fftshift(self.otf(x0,dx=dx,dy=dy,_caller='self')))))
+        out = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(self.otf(x0,dx=dx,dy=dy,_caller='self')))))
         
         if self._k==1:
             return out
@@ -673,12 +673,12 @@ class Psfao(ParametricPSF):
     def _otf_turbulent(self,x0):
         PSD,integral = self.psd(x0)
         L = self.system.D * self._samp_over
-        Bg = fft2(fftshift(PSD)) / L**2
-        #Dphi = fftshift(_np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV 
-        Dphi = fftshift(_np.real(2 * (integral - Bg))) # normalized up to infinity
+        Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
+        #Dphi = _fft.fftshift(_np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV 
+        Dphi = _fft.fftshift(_np.real(2 * (integral - Bg))) # normalized up to infinity
         return _np.exp(-Dphi/2.) 
     
-    @lru_cache(maxsize=2)
+    @_lru_cache(maxsize=2)
     def _otf_diffraction(self):
         Nx_over = self.Npix[0]*self._k
         Ny_over = self.Npix[1]*self._k
@@ -687,7 +687,7 @@ class Psfao(ParametricPSF):
         NpupY = _np.ceil(Ny_over/self._samp_over)
         tab = _np.zeros((Nx_over, Ny_over), dtype=_np.complex)
         tab[0:int(NpupX), 0:int(NpupY)] = self.system.pupil((NpupX,NpupY),samp=self._samp_over)
-        return fftshift(_np.abs(ifft2(_np.abs(fft2(tab)) ** 2)) / _np.sum(tab))
+        return _fft.fftshift(_np.abs(_fft.ifft2(_np.abs(_fft.fft2(tab)) ** 2)) / _np.sum(tab))
     
     def _otf_shift(self,dx,dy):
         Nx_over = self.Npix[0]*self._k
@@ -716,7 +716,7 @@ class Psfao(ParametricPSF):
             PSF Y shifting [pix] (default = 0)
         """
         
-        out = _np.real(fftshift(ifft2(fftshift(self.otf(x0,dx=dx,dy=dy,_caller='self')))))
+        out = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(self.otf(x0,dx=dx,dy=dy,_caller='self')))))
         
         if self._k==1:
             return out