# Path: pyqstem/wave.py
# class Wave(BaseArray):
# 
#     def __init__(self, array, energy, sampling=None, periodic_xy=True):
#         BaseArray.__init__(self, array, sampling)
#         self.energy = energy
#     
#     @property
#     def shape(self):
#         return self.array.shape
#     
#     @property
#     def wavelength(self):
#         return 0.38783/np.sqrt(self.energy+9.78476*10**(-4)*self.energy**2)
# 
#     def z_slice(self,ind=-1):
# 
#         if len(self.array.shape)==2:
#             raise RuntimeError('z_slice() only works for 3d wavefunctions')
# 
#         return Wavefunction(self.array[:,:,ind],self.energy,self.sampling[:2],self.offset)
# 
#     def apply_ctf(self,ctf):
#         return ctf.apply(self)
# 
#     def resample(self,sampling):
# 
#         if len(self.array.shape)==3:
#             raise RuntimeError('resample() only works for 2d wavefunctions')
# 
#         if not isinstance(sampling, (list, tuple)):
#             sampling=(sampling,)*2
# 
#         zoom=(self.sampling[0]/sampling[0],self.sampling[1]/sampling[1])
# 
#         real = scipy.ndimage.interpolation.zoom(np.real(self.array), zoom)
#         imag = scipy.ndimage.interpolation.zoom(np.imag(self.array), zoom)
# 
#         sampling=(self.array.shape[0]*self.sampling[0]/real.shape[0],
#                     self.array.shape[1]*self.sampling[1]/real.shape[1])
# 
#         return Wavefunction(real+1.j*imag,self.energy,sampling,self.offset)
# 
#     def detect(self,dose=None,MTF_param=None,MTF_func=None,blur=None,resample=None):
#         sampling=self.sampling
#         img=np.abs(self.array)**2
#         return detect(img,sampling,dose=dose,MTF_param=MTF_param,MTF_func=MTF_func,blur=blur,resample=resample)
# 
#     def save(self,name):
#         np.savez(name,self.array,self.energy,self.sampling)
# 
#     def view(self,method='intensity',nav_axis=2,ind=-1,slider=False,ax=None,**kwargs):
# 
#         self.refs += view(self,method=method,nav_axis=nav_axis,ind=ind,slider=slider,ax=ax,**kwargs)
# 
# Path: pyqstem/util.py
# def energy2wavelength(v0):
#     m0=0.5109989461*10**3 # keV / c**2
#     h=4.135667662*10**(-15)*10**(-3) # eV * s
#     c=2.99792458*10**8 # m / s
#     return h*c/np.sqrt(v0*(2*m0+v0))*10**10
# 
# def spatial_frequencies(shape,sampling,return_polar=False,return_nyquist=False,wavelength=None):
# 
#     if not isinstance(shape, collections.Iterable):
#         shape = (shape,)*2
# 
#     if not isinstance(sampling, collections.Iterable):
#         sampling = (sampling,)*2
#     
#     dkx=1/(shape[0]*sampling[0])
#     dky=1/(shape[1]*sampling[1])
# 
#     if shape[0]%2==0:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2,shape[0]/2,1))
#     else:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2-.5,shape[0]/2-.5,1))
# 
#     if shape[1]%2==0:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2,shape[1]/2,1))
#     else:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2-.5,shape[1]/2-.5,1))
# 
#     ky,kx = np.meshgrid(ky,kx)
# 
#     k2 = kx**2+ky**2
# 
#     ret = (kx,ky,k2)
#     if return_nyquist:
#         knx = 1/(2*sampling[0])
#         kny = 1/(2*sampling[1])
#         Kx=kx/knx
#         Ky=ky/knx
#         K2 = Kx**2+Ky**2
#         ret += (Kx,Ky,K2)
#     if return_polar:
#         theta = np.sqrt(k2*wavelength**2)
#         phi = np.arctan2(ky,kx)
#         ret += (theta,phi)
# 
#     return ret


from .wave import Wave
from .util import energy2wavelength, spatial_frequencies
        from scipy.interpolate import spline,interp1d
import matplotlib.pyplot as plt
import numpy as np
    def check_recalculate(self,shape,sampling,wavelength,tol=1e-12):
        
        if np.any([None is i for i in [self.array,self.sampling,self.wavelength]]):
            return True
        elif ((not np.isclose(self.wavelength,wavelength,rtol=tol))|
              (not np.isclose(self.sampling,sampling,rtol=tol).all())|
              (not (self.array.shape == shape))):
            return True
        else:
            return False

    def apply(self,wave,keep=True):
        
        shape=wave.array.shape
        sampling=wave.sampling
        wavelength=wave.wavelength
        
        if self.check_recalculate(shape,sampling,wavelength):
            ctf = self.calculate(shape,sampling,wavelength)
        
        new_wave_array=np.fft.ifft2(np.fft.fft2(wave.array)*ctf)
        
        if keep:
            self.sampling=sampling
            self.wavelength=wavelength
            self.array=ctf
        
        if len(new_wave_array.shape) > 2:
            return [Wave(a,wave.energy,wave.sampling) for a in new_wave_array]
        else:

# Path: pyqstem/wave.py
# class Wave(BaseArray):
# 
#     def __init__(self, array, energy, sampling=None, periodic_xy=True):
#         BaseArray.__init__(self, array, sampling)
#         self.energy = energy
#     
#     @property
#     def shape(self):
#         return self.array.shape
#     
#     @property
#     def wavelength(self):
#         return 0.38783/np.sqrt(self.energy+9.78476*10**(-4)*self.energy**2)
# 
#     def z_slice(self,ind=-1):
# 
#         if len(self.array.shape)==2:
#             raise RuntimeError('z_slice() only works for 3d wavefunctions')
# 
#         return Wavefunction(self.array[:,:,ind],self.energy,self.sampling[:2],self.offset)
# 
#     def apply_ctf(self,ctf):
#         return ctf.apply(self)
# 
#     def resample(self,sampling):
# 
#         if len(self.array.shape)==3:
#             raise RuntimeError('resample() only works for 2d wavefunctions')
# 
#         if not isinstance(sampling, (list, tuple)):
#             sampling=(sampling,)*2
# 
#         zoom=(self.sampling[0]/sampling[0],self.sampling[1]/sampling[1])
# 
#         real = scipy.ndimage.interpolation.zoom(np.real(self.array), zoom)
#         imag = scipy.ndimage.interpolation.zoom(np.imag(self.array), zoom)
# 
#         sampling=(self.array.shape[0]*self.sampling[0]/real.shape[0],
#                     self.array.shape[1]*self.sampling[1]/real.shape[1])
# 
#         return Wavefunction(real+1.j*imag,self.energy,sampling,self.offset)
# 
#     def detect(self,dose=None,MTF_param=None,MTF_func=None,blur=None,resample=None):
#         sampling=self.sampling
#         img=np.abs(self.array)**2
#         return detect(img,sampling,dose=dose,MTF_param=MTF_param,MTF_func=MTF_func,blur=blur,resample=resample)
# 
#     def save(self,name):
#         np.savez(name,self.array,self.energy,self.sampling)
# 
#     def view(self,method='intensity',nav_axis=2,ind=-1,slider=False,ax=None,**kwargs):
# 
#         self.refs += view(self,method=method,nav_axis=nav_axis,ind=ind,slider=slider,ax=ax,**kwargs)
# 
# Path: pyqstem/util.py
# def energy2wavelength(v0):
#     m0=0.5109989461*10**3 # keV / c**2
#     h=4.135667662*10**(-15)*10**(-3) # eV * s
#     c=2.99792458*10**8 # m / s
#     return h*c/np.sqrt(v0*(2*m0+v0))*10**10
# 
# def spatial_frequencies(shape,sampling,return_polar=False,return_nyquist=False,wavelength=None):
# 
#     if not isinstance(shape, collections.Iterable):
#         shape = (shape,)*2
# 
#     if not isinstance(sampling, collections.Iterable):
#         sampling = (sampling,)*2
#     
#     dkx=1/(shape[0]*sampling[0])
#     dky=1/(shape[1]*sampling[1])
# 
#     if shape[0]%2==0:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2,shape[0]/2,1))
#     else:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2-.5,shape[0]/2-.5,1))
# 
#     if shape[1]%2==0:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2,shape[1]/2,1))
#     else:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2-.5,shape[1]/2-.5,1))
# 
#     ky,kx = np.meshgrid(ky,kx)
# 
#     k2 = kx**2+ky**2
# 
#     ret = (kx,ky,k2)
#     if return_nyquist:
#         knx = 1/(2*sampling[0])
#         kny = 1/(2*sampling[1])
#         Kx=kx/knx
#         Ky=ky/knx
#         K2 = Kx**2+Ky**2
#         ret += (Kx,Ky,K2)
#     if return_polar:
#         theta = np.sqrt(k2*wavelength**2)
#         phi = np.arctan2(ky,kx)
#         ret += (theta,phi)
# 
#     return ret


# Path: pyqstem/imaging.py
from .wave import Wave
from .util import energy2wavelength, spatial_frequencies
        from scipy.interpolate import spline,interp1d
import matplotlib.pyplot as plt
import numpy as np

    def check_recalculate(self,shape,sampling,wavelength,tol=1e-12):
        
        if np.any([None is i for i in [self.array,self.sampling,self.wavelength]]):
            return True
        elif ((not np.isclose(self.wavelength,wavelength,rtol=tol))|
              (not np.isclose(self.sampling,sampling,rtol=tol).all())|
              (not (self.array.shape == shape))):
            return True
        else:
            return False

    def apply(self,wave,keep=True):
        
        shape=wave.array.shape
        sampling=wave.sampling
        wavelength=wave.wavelength
        
        if self.check_recalculate(shape,sampling,wavelength):
            ctf = self.calculate(shape,sampling,wavelength)
        
        new_wave_array=np.fft.ifft2(np.fft.fft2(wave.array)*ctf)
        
        if keep:
            self.sampling=sampling
            self.wavelength=wavelength
            self.array=ctf
        
        if len(new_wave_array.shape) > 2:
            return [Wave(a,wave.energy,wave.sampling) for a in new_wave_array]
        else:

            return Wave(new_wave_array,wave.energy,wave.sampling)