├── LICENSE ├── PD.py ├── PD.pyc ├── README.md ├── __init__.py ├── __init__.pyc ├── __pycache__ ├── PD.cpython-38.pyc ├── aperture.cpython-38.pyc ├── cost_func.cpython-38.pyc ├── deconvolution.cpython-38.pyc ├── logging.cpython-38.pyc ├── noise.cpython-38.pyc ├── processing.cpython-38.pyc ├── telescope.cpython-38.pyc ├── tools.cpython-38.pyc ├── wavefront.cpython-38.pyc └── zernike.cpython-38.pyc ├── aperture.py ├── aperture.pyc ├── cost_func.py ├── deconvolution.py ├── deconvolution.pyc ├── minimization.py ├── noise.py ├── noise.pyc ├── patch_pd.py ├── processing.py ├── pypd ├── setup.py ├── telescope.py ├── telescope.pyc ├── tools.py ├── tools.pyc ├── wavefront.py ├── wavefront.pyc ├── zernike.py └── zernike.pyc /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2021 Fatima Kahil 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /PD.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import scipy 4 | from scipy.fftpack import fftshift, ifftshift, fft2, ifft2 5 | from functools import partial 6 | 7 | 8 | 9 | 10 | ## function to compute the FT of focused and defocused image 11 | def FT(im0, imk): 12 | d0 = fft2(im0) 13 | dk = fft2(imk) 14 | return d0, dk 15 | 16 | 17 | ## function to define the Q matrix (Eq. 6 of L\"odfahl & scharmer. 1994) 18 | def Q_matrix(t0,tk,reg,gamma): 19 | tmp = np.abs(t0)**2 + gamma*np.abs(tk)**2 +reg 20 | q = 1./(np.sqrt(tmp)) 21 | q2 = q*q 22 | return q, q2 23 | 24 | 25 | 26 | # 27 | ## function to compute the optimized object (Eq. 5 of L\"odfahl & scharmer.1994) 28 | def F_M(q2,d0,dk,t0,tk,filter,gamma): 29 | F_M = filter*q2*(d0*np.conj(t0) + gamma*dk*np.conj(tk)) 30 | return F_M 31 | 32 | 33 | 34 | 35 | ## function to define the error metric to be minimized (Eq. 9 of L\"odfahl & scharmer.1994) 36 | def Error_metric(t0,tk,d0,dk,q,filter): 37 | ef = filter*(dk*t0 - d0*tk) 38 | ef = q*ef 39 | EF = fft2((ef)) 40 | EF = EF.real 41 | EF = EF-EF.mean() 42 | return EF 43 | 44 | 45 | 46 | ## 47 | def L_M(EF,size): 48 | L_m = np.sum(np.abs(EF)**2)/(size**2) 49 | return L_m 50 | 51 | -------------------------------------------------------------------------------- /PD.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/PD.pyc -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # PyPD 3 | 4 | This repo contains a Python package that I have written for applying Phase Diversity Reconstruction on Solar Orbiter data (or any other telescope with a circular aperture). It could be applied on any astronomical data as well. 5 | 6 | It is based on the algorithm described in [Löfdahl and Scharmer. 1994](http://adsabs.harvard.edu/full/1994A&AS..107..243L) 7 | 8 | # Classes explained 9 | An important note here is to use a unified unit for all parameters (nm or m) 10 | 11 | ## The `Telescope` class: 12 | The parameters needed for this class are: 13 | - `lam`: the wavelength of study 14 | - `diameter`: the diameter of the telescope's aperture 15 | - `focal_length`: the effective focal length of the telescope 16 | - `name`: the name of your telescope 17 | - `platescale`: the pixel size in arcseconds 18 | - `size`: the size of the detector in pixels 19 | 20 | This class allows the user to compute: 21 | 1. `calc_ps`: platescale in arcsec/pixel with an additional input of `pix_size` in units the same as the focal length 22 | 2. `spatial_resolution`: the spatial resolution (Rayleigh limit) 23 | 3. `pupil_size`: the size of the exit pupil of the system 24 | 25 | ## The `minimization` class: 26 | 27 | This class allows the user to: 28 | 29 | 1. Minimize the "Error Metric" and get the Zernike Polynomials describing the wavefront error of the telescope 30 | 2. Plot the wavefront error, the point spread function (PSF), the modulation transfer function (MTF), the 1D azimuthal MTF of the telescope 31 | 3. Construct the PSF and MTF from a set of Zernike Polynomials (from wavefront to PSF to OTF to MTF) 32 | 4. Restore blurred images and correct for the PSF choosing between the Wiener filter and the Richardson-Lucy filter 33 | 34 | The parameters of this class: 35 | 1. the pair of focused and defocused images, the input image should have a format of `2xsize_xxsize_y`. 36 | 2. the parameters of the `Telescope` class 37 | 3. `cutoff` frequency for the noise filtering 38 | 4. `reg`: regularization parameter for the Wiener filter 39 | 5. `ap`: the amount of the apodization rim in pixels 40 | 6. `x1, x2, y1, y2`: the delimeters of the region of interest 41 | 7. `co_num`: number of Zernike polynomials to fit the wavefront 42 | 8. `del_z`: the amount of defocus in lambda 43 | 44 | The modules of this class: 45 | 1. `fit`: takes in the method of fitting and returns the best-fit Zernike coefficients 46 | 2. `plot_results`: takes in the returned Zernike coefficients and plots the wavefront error, the 2D MTF, the 1D MTF 47 | 3. `restored_scene`: perform the deconvolution of a blurred image with an input of Zernike Coefficients. The user can use between two filter `Wiener` and `richardsonLucy` 48 | 49 | ## The `patch_pd` class: 50 | For fitting the wavefront error is sub-regions of the full FOV of the PD dataset. The user can enter the size of the subregion (it has to be quadratic) plus the number of Zernike polynomials to be fit and the names of the output files (one for the wavefront error and one for the 2D MTF). 51 | 52 | The class offers the option to run parallel computation by setting the parser -p to True (see below). 53 | # How to use the code? 54 | The `minimization` class returns the best-fit zernike polynomials, a visualisation of the results (wavefront error+MTF), and the restored scene (in this case the focused image of the PD dataset). To get these specific results, type in the shell terminal: 55 | ``` 56 | python3 minimization.py -i 'path/input.fits' -w 617.3e-6 -a 140 -f 4125.3 -p 0.5 -c 0.5 -r 1e-10 -ap 10 -x1 500 -x2 650 -y1 500 -y2 650 -z 10 -del 0.5 -o path/reduced.fits -fl 'Wiener' 57 | ``` 58 | 59 | The specific description of the parsers and input to the class can be found inside [the main code](https://github.com/fakahil/PyPD/blob/master/minimization.py). The values given above are for an example PD dataset taken by the PHI/HRT telescope. You can change the values according to your telescope. 60 | 61 | To use the `patch_pd` class: 62 | 63 | ``` 64 | python3 patch_pd.py -i 'path/input.fits' -z 10 -d 265 -ow 'path/output_wf.fits' -om 'path/output_wf.fits -p True 65 | 66 | ``` 67 | The parsers description can be found in the [main code](https://github.com/fakahil/PyPD/blob/master/patch_pd.py) 68 | 69 | 70 | 71 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from . import aperture, deconvolution, minimization 2 | from . import telescope, noise, aperture, PD 3 | from . import zernike, tools, wavefront 4 | -------------------------------------------------------------------------------- /__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__init__.pyc -------------------------------------------------------------------------------- /__pycache__/PD.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/PD.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/aperture.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/aperture.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/cost_func.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/cost_func.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/deconvolution.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/deconvolution.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/logging.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/logging.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/noise.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/noise.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/processing.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/processing.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/telescope.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/telescope.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/tools.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/tools.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/wavefront.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/wavefront.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/zernike.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/__pycache__/zernike.cpython-38.pyc -------------------------------------------------------------------------------- /aperture.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | 5 | 6 | 7 | def mask_pupil(rpupil, size): 8 | r = 1 9 | xx = np.linspace(-r, r, 2*rpupil) 10 | yy = np.linspace(-r, r, 2*rpupil) 11 | [X,Y] = np.meshgrid(xx,yy) 12 | R = np.sqrt(X**2+Y**2) 13 | theta = np.arctan2(Y, X) 14 | M = 1*(np.cos(theta)**2+np.sin(theta)**2) 15 | M[R>1] = 0 16 | Mask = np.zeros([size,size]) 17 | Mask[size//2-rpupil+1:size//2+rpupil+1,size//2-rpupil+1:size//2+rpupil+1]= M 18 | return Mask 19 | 20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /aperture.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/aperture.pyc -------------------------------------------------------------------------------- /cost_func.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | import scipy 4 | from scipy.fftpack import fftshift, ifftshift, fft2, ifft2 5 | 6 | 7 | import tools 8 | import aperture 9 | import PD 10 | import noise 11 | import wavefront 12 | import telescope 13 | from telescope import * 14 | from noise import * 15 | from wavefront import * 16 | from tools import * 17 | from PD import * 18 | 19 | 20 | class cost_func: 21 | def __init__(self,size,cut_off,reg,ap,co_num,del_z,lam, diameter,focal_length,platescale): 22 | 23 | self.cut_off = cut_off 24 | self.reg = reg 25 | self.size = size 26 | 27 | self.co_num = co_num 28 | 29 | self.del_z = del_z 30 | self.lam = lam 31 | self.diameter = diameter 32 | self.focal_length = focal_length 33 | self.platescale = platescale 34 | self.telescope = Telescope(self.lam,self.diameter,self.focal_length,'HRT',self.platescale,self.size) 35 | 36 | 37 | 38 | 39 | 40 | def Minimize_res(self,coefficients): 41 | 42 | noise_filter = fftshift(noise_mask_high(self.size,self.cut_off)) 43 | 44 | ph = wavefront.phase_embed(coefficients,self.telescope.pupil_size(),self.size,self.co_num) 45 | A_f = wavefront.pupil_foc(coefficients,self.size,self.telescope.pupil_size(),self.co_num) 46 | 47 | 48 | psf_foc = wavefront.PSF( aperture.mask_pupil(self.telescope.pupil_size(),self.size),A_f,False) 49 | 50 | t0 = wavefront.OTF(psf_foc) 51 | 52 | 53 | return t0,ph 54 | 55 | 56 | 57 | -------------------------------------------------------------------------------- /deconvolution.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | import scipy 5 | from scipy.fftpack import fftshift, ifftshift, fft2, ifft2 6 | from astropy.io import fits 7 | import imreg_dft 8 | 9 | import tools 10 | import aperture 11 | import PD 12 | import noise 13 | import wavefront 14 | import wavefront 15 | from wavefront import * 16 | import tools 17 | from tools import * 18 | from noise import * 19 | 20 | def ft2(g): 21 | G = fftshift(fft2(ifftshift(g))) 22 | return G 23 | 24 | def ift2(G): 25 | 26 | numPixels = G.shape[0] 27 | 28 | g = fftshift(ifft2(ifftshift(G))) 29 | return g 30 | 31 | 32 | def conv2(g1, g2): 33 | 34 | G1 = ft2(g1) 35 | G2 = ft2(g2) 36 | G_out = G1*G2 37 | 38 | numPixels = g1.shape[0] 39 | 40 | g_out = ift2(G_out) 41 | return g_out 42 | 43 | 44 | def richardsonLucy(img, psf, iterations=100): 45 | f = img 46 | blur_psf = np.matrix(psf) 47 | 48 | psf_mirror = blur_psf.T 49 | 50 | for _ in range(iterations): 51 | 52 | f = f * conv2(img / conv2(f, blur_psf), psf_mirror) 53 | 54 | return np.abs(f) 55 | 56 | def Wienerfilter(img,t0,reg,cut_off,size,ap): 57 | temp = noise_mask_high(size,cut_off) 58 | noise_filter = fftshift(temp) 59 | 60 | im0 = tools.apo2d(img,ap) 61 | d0 = fft2(im0) 62 | scene = noise_filter*d0*(np.conj(t0)/(np.abs(t0)**2+reg)) 63 | scene2 = ifft2(scene).real 64 | return scene2 65 | 66 | 67 | 68 | -------------------------------------------------------------------------------- /deconvolution.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/deconvolution.pyc -------------------------------------------------------------------------------- /minimization.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import scipy 5 | from scipy.fftpack import fftshift, ifftshift, fft2, ifft2 6 | from mpl_toolkits.axes_grid1 import make_axes_locatable, axes_size 7 | from scipy.signal import correlate2d as correlate 8 | from scipy.signal import general_gaussian 9 | from astropy.io import fits 10 | from scipy import ndimage 11 | from functools import partial 12 | import time 13 | import imreg_dft 14 | import pyfits 15 | import tools 16 | import aperture 17 | import PD 18 | import noise 19 | import wavefront 20 | import deconvolution 21 | from scipy.optimize import minimize, minimize_scalar 22 | 23 | import argparse 24 | #from . import detector 25 | #from .detector import * 26 | from telescope import * 27 | 28 | from tools import imreg, apo2d 29 | from aperture import * 30 | from PD import * 31 | from noise import * 32 | from wavefront import * 33 | from deconvolution import * 34 | 35 | class minimization(object): 36 | 37 | def __init__(self,foc_defoc,lam, diameter,focal_length,platescale,cut_off,reg,ap,x1,x2,y1,y2,co_num,del_z,output,filterr): 38 | self.data = fits.getdata(foc_defoc) 39 | 40 | 41 | self.foc = self.data[0,:,:] 42 | self.defoc = self.data[1,:,:] 43 | self.cut_off = cut_off 44 | self.reg = reg 45 | self.size =self.y2-self.y1 46 | self.ap = ap 47 | self.co_num = co_num 48 | self.x1 = x1 49 | self.x2 = x2 50 | self.y1 = y1 51 | self.y2 = y2 52 | self.del_z = del_z 53 | self.lam = lam 54 | self.diameter = diameter 55 | self.focal_length = focal_length 56 | self.platescale = platescale 57 | self.reg = reg 58 | self.telescope = Telescope(self.lam,self.diameter,self.focal_length,'HRT',self.platescale,self.size) 59 | self.output = output 60 | self.filterr = filterr 61 | 62 | 63 | 64 | def fit(self,m = 'L-BFGS-B'): 65 | im0 = self.foc[self.y1:self.y2,self.x1:self.x2] 66 | imk = self.defoc[self.y1:self.y2,self.x1:self.x2] 67 | imk = tools.imreg(im0,imk) 68 | im0 = tools.apo2d(im0,self.ap) 69 | imk = tools.apo2d(imk, self.ap) 70 | rpupil = self.telescope.pupil_size() 71 | Mask = aperture.mask_pupil(rpupil,self.size) 72 | 73 | d0,dk = PD.FT(im0,imk) 74 | noise_temp = noise_mask_high(self.size,self.cut_off) 75 | noise_filter = fftshift(noise_temp) 76 | 77 | M_gamma = noise_mask(self.size,self.cut_off) 78 | gam = Gamma(d0,dk,M_gamma) 79 | p0 = np.zeros(self.co_num) 80 | 81 | def Minimise(coefficients): 82 | A_f = wavefront.pupil_foc(coefficients,self.size,self.telescope.pupil_size(),self.co_num) 83 | A_def = wavefront.pupil_defocus(coefficients,self.size,self.del_z,self.telescope.pupil_size(),self.co_num) 84 | psf_foc = wavefront.PSF(Mask,A_f,False) 85 | psf_defoc = wavefront.PSF(Mask,A_def,False) 86 | t0 = wavefront.OTF(psf_foc) 87 | tk = wavefront.OTF(psf_defoc) 88 | 89 | q,q2 = PD.Q_matrix(t0,tk,self.reg,gam) 90 | 91 | #noise_filter = sch_filter(N0, t0, tk, d0, dk,reg) 92 | F_m = PD.F_M(q2,d0, dk,t0,tk,noise_filter,gam) 93 | 94 | E_metric = PD.Error_metric(t0,tk,d0,dk,q,noise_filter) 95 | L_m = PD.L_M(E_metric,self.size) 96 | return L_m 97 | 98 | Minimise_partial = partial(Minimise) 99 | mini = scipy.optimize.minimize(Minimise_partial,p0,method=m) 100 | return mini.x 101 | 102 | 103 | def plot_results(self,Z): 104 | 105 | 106 | 107 | A_f = wavefront.pupil_foc(Z,self.size,self.telescope.pupil_size(),self.co_num) 108 | rpupil = self.telescope.pupil_size() 109 | Mask = aperture.mask_pupil(rpupil,self.size) 110 | psf_foc = wavefront.PSF(Mask,A_f,False) 111 | t0 = wavefront.OTF(psf_foc) 112 | ph =wavefront.phase(Z, self.telescope.pupil_size(),self.co_num) 113 | 114 | 115 | fig = plt.figure(figsize=(20,20)) 116 | ax1 = fig.add_subplot(1,3,1) 117 | im1 = ax1.imshow(ph/(2*np.pi), origin='lower',cmap='gray') 118 | ax1.set_xlabel('[Pixels]',fontsize=18) 119 | ax1.set_ylabel('[Pixels]',fontsize=18) 120 | divider1 = make_axes_locatable(ax1) 121 | cax1 = divider1.append_axes("right", size=0.15, pad=0.05) 122 | cbar1 = plt.colorbar(im1, cax=cax1,orientation='vertical')#,ticks=np.arange(0.4,0.9,0.1)) 123 | ax1.set_title('WF error HRT [$\lambda$]',fontsize=20) 124 | 125 | ax2 = fig.add_subplot(1,3,2) 126 | im2 = ax2.imshow(wavefront.MTF(fftshift(t0)), origin='lower',cmap='gray') 127 | ax2.set_xlabel('[Pixels]',fontsize=18) 128 | ax2.set_ylabel('[Pixels]',fontsize=18) 129 | divider2 = make_axes_locatable(ax2) 130 | cax2 = divider2.append_axes("right", size=0.15, pad=0.05) 131 | cbar2 = plt.colorbar(im2, cax=cax2,orientation='vertical')#,ticks=np.arange(0.4,0.9,0.1)) 132 | ax2.set_title('MTF',fontsize=20) 133 | 134 | 135 | az = tools.GetPSD1D(wavefront.MTF(fftshift(t0))) 136 | freq=np.linspace(0,0.5,int(self.size/2)) 137 | freq_c_hrt = self.diameter/(self.lam*self.focal_length*100) 138 | phi_hrt = np.arccos(freq/freq_c_hrt) 139 | MTF_p_hrt = (2/np.pi)*(phi_hrt - (np.cos(phi_hrt))*np.sin(phi_hrt)) 140 | ax3 = fig.add_subplot(1,3,3) 141 | ax3.set_xlabel('Freq (1/pixel)',fontsize=22) 142 | ax3.set_ylabel('MTF',fontsize=22) 143 | ax3.plot(freq,MTF_p_hrt,label='Theoretical MTF') 144 | ax3.plot(freq,az,label='Observed MTF') 145 | ax3.set_xlim(0,0.5) 146 | plt.legend() 147 | plt.subplots_adjust(hspace=0.1, wspace=0.6) 148 | 149 | plt.savefig('resutls.png') 150 | 151 | 152 | def restored_scene(self,Z,iterations_RL=10): 153 | to_clean = self.foc[self.y1:self.y2,self.x1:self.x2] 154 | A_f = wavefront.pupil_foc(Z,self.size,self.telescope.pupil_size(),self.co_num) 155 | rpupil = self.telescope.pupil_size() 156 | Mask = aperture.mask_pupil(rpupil,self.size) 157 | psf_foc = wavefront.PSF(Mask,A_f,False) 158 | t0 = wavefront.OTF(psf_foc) 159 | if to_clean.shape != t0.shape: 160 | raise ValueError('Image and PSF do not have the same dimensions') 161 | elif to_clean.shape == t0.shape: 162 | 163 | if self.filterr == 'Wiener': 164 | restored =Wienerfilter(to_clean,t0,self.reg,self.cut_off,self.size,self.ap) 165 | 166 | 167 | if self.filterr == 'richardsonLucy': 168 | restored = richardsonLucy(to_clean, psf_foc, iterations_RL) 169 | 170 | print("Saving data...") 171 | hdu = fits.PrimaryHDU(restored) 172 | hdu.writeto(self.output,overwrite=True) 173 | ''' 174 | import os.path 175 | if os.path.exists(self.output): 176 | os.system('rm {0}'.format(self.output)) 177 | print('Overwriting...') 178 | hdu.writeto(self.output,overwrite=True) 179 | 180 | ''' 181 | 182 | 183 | if (__name__ == '__main__'): 184 | 185 | 186 | 187 | parser = argparse.ArgumentParser(description='Retrieving wavefront error') 188 | parser.add_argument('-i','--input', help='input') 189 | parser.add_argument('-o','--out', help='out') 190 | parser.add_argument('-w','--wavelength', help='wavelength',default=617.3e-6) 191 | parser.add_argument('-a','--aperture', help='aperture', default=140) 192 | parser.add_argument('-f','--focal_length', help='focal_length',default=4125.3) 193 | parser.add_argument('-p','--plate_scale', help='plate_scale',default=0.5) 194 | parser.add_argument('-c','--cut_off', help='cut_off',default=0.5) 195 | parser.add_argument('-r','--reg', help='reg',default=1e-10) 196 | parser.add_argument('-ap','--apod', help='apod',default=10) 197 | parser.add_argument('-x1','--x1', help='x1') 198 | parser.add_argument('-x2','--x2', help='x2') 199 | parser.add_argument('-y1','--y1', help='y1') 200 | parser.add_argument('-y2','--y2', help='y2') 201 | parser.add_argument('-z','--Z', help='Z',default=10) 202 | parser.add_argument('-del','--del', help='del',default=0.5) 203 | parser.add_argument('-fl','--filter', help='filter', choices=['richardsonLucy', 'Wiener'], default='Wiener') 204 | 205 | parsed = vars(parser.parse_args()) 206 | 207 | 208 | 209 | res = minimization(foc_defoc='{0}'.format(parsed['input']), lam=float(parsed['wavelength']),diameter=float(parsed['aperture']),focal_length=float(parsed['focal_length']), 210 | platescale=float(parsed['plate_scale']),cut_off=float(parsed['cut_off']),reg=float(parsed['reg']),ap=int(parsed['apod']),x1=int(parsed['x1']),x2=int(parsed['x2']),y1=int(parsed['y1']),y2=int(parsed['y2']), 211 | co_num=int(parsed['Z']),del_z=float(parsed['del']),output='{0}'.format(parsed['out']),filterr=parsed['filter']) 212 | 213 | Z = res.fit() 214 | print(Z) 215 | tools.plot_zernike(Z) 216 | res.plot_results(Z) 217 | res.restored_scene(Z,10) 218 | 219 | 220 | 221 | -------------------------------------------------------------------------------- /noise.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import scipy 4 | from scipy.fftpack import fftshift, ifftshift, fft2, ifft2 5 | 6 | 7 | 8 | 9 | 10 | # compute the gamma parameter 11 | def Gamma(d0,dk,mask): 12 | d0 = np.abs(d0)**2 13 | dk = np.abs(dk)**2 14 | d0 = fftshift(d0) 15 | dk = fftshift(dk) 16 | denom = np.sum(mask) 17 | sigma_0 = np.sum(mask*d0)/denom 18 | sigma_k = np.sum(mask*dk)/denom 19 | gamma = sigma_0/sigma_k 20 | return gamma 21 | 22 | 23 | # mask frequencies below a frequency threshold (to make the noise image) 24 | def noise_mask(size,cut_off): 25 | X = np.linspace(-0.5,0.5,size) 26 | x,y = np.meshgrid(X,X) 27 | mask = np.ones((size,size)) 28 | m = x * x + y * y <= cut_off**2 29 | mask[m] = 0 30 | return mask 31 | 32 | 33 | 34 | # to mask frequencies above a frequency threshold 35 | def noise_mask_high(size,cut_off): 36 | X = np.linspace(-0.5,0.5,size) 37 | x,y = np.meshgrid(X,X) 38 | mask = np.zeros((size,size)) 39 | m = x * x + y * y <= cut_off**2 40 | mask[m] = 1 41 | return mask 42 | 43 | 44 | 45 | 46 | def Noise_d0(im0, freq_mask): 47 | pd = fft2(im0) 48 | pd = np.abs(pd)**2 49 | pd = fftshift(pd) 50 | pd_m = pd*freq_mask 51 | pd_m = ifftshift(pd_m) 52 | return pd_m 53 | 54 | 55 | 56 | 57 | ## Noise filter (Loedfahl and Berger 1998) 58 | def sch_filter(noise_im,t0,tk,d0,dk,reg): 59 | 60 | ab_noise = np.abs(noise_im)**2 61 | filter = ab_noise*(np.abs(t0)**2 + np.abs(tk)**2)/(np.abs(d0*np.conj(t0) + dk*np.conj(tk))**2 + reg) 62 | filter = 1- filter 63 | filter_2 = filter 64 | filter_2[filter<0.2] = 0 65 | filter_2[filter>1] = 1 66 | return filter_2 67 | 68 | -------------------------------------------------------------------------------- /noise.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/noise.pyc -------------------------------------------------------------------------------- /patch_pd.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import tools 4 | import aperture 5 | import PD 6 | import noise 7 | import wavefront 8 | import deconvolution 9 | import cost_func 10 | from scipy.optimize import minimize, minimize_scalar 11 | import astropy 12 | from astropy.io import fits 13 | import argparse 14 | import telescope 15 | from telescope import * 16 | from tools import * 17 | from aperture import * 18 | from PD import * 19 | from noise import * 20 | from wavefront import * 21 | from deconvolution import * 22 | from cost_func import * 23 | import os 24 | import processing 25 | 26 | 27 | class patch_pd(object): 28 | def __init__(self,pd_data,Del,co_num,output_wf,output_mtf,parallel=True): 29 | 30 | 31 | self.data= fits.getdata(pd_data) 32 | self.Del = Del 33 | self.co_num = co_num 34 | self.output_wf = output_wf 35 | self.output_mtf = output_mtf 36 | 37 | 38 | self.mean_im0 = self.data[0,500:1000,500:1000].mean() 39 | self.mean_imk = self.data[1,500:1000,500:1000].mean() 40 | 41 | self.Im0 = self.data[0,:,:]/self.mean_im0 42 | self.Imk = self.data[1,:,:]/self.mean_imk 43 | 44 | self.output_WF = np.zeros((2048,2048)) 45 | self.output_MTF = np.zeros((2048,2048)) 46 | self.parallel = parallel 47 | #@staticmethod 48 | def run_pd(self,k): 49 | 50 | im0 = self.patch[k,:,:,0] 51 | imk =self.patch[k,:,:,1] 52 | im0 = im0/im0.mean() 53 | imk = imk/imk.mean() 54 | imk = imreg(im0,imk) 55 | 56 | 57 | im0 = apo2d(im0,10) 58 | imk = apo2d(imk,10) 59 | 60 | d0,dk = FT(im0,imk) 61 | gam =1# Gamma(d0,dk,M_gamma) 62 | p0 = np.zeros(self.co_num) 63 | fit = cost_func(self.Del,0.5,1e-10,10,self.co_num,0.5,617.3e-6, 140,4125.3,0.5) 64 | Mask = aperture.mask_pupil(fit.telescope.pupil_size(),fit.size) 65 | noise_temp = noise_mask_high(fit.size,fit.cut_off) 66 | noise_filter = fftshift(noise_temp) 67 | def Minimise(coefficients): 68 | A_f = wavefront.pupil_foc(coefficients,fit.size,fit.telescope.pupil_size(),self.co_num) 69 | A_def = wavefront.pupil_defocus(coefficients,fit.size,fit.del_z,fit.telescope.pupil_size(),self.co_num) 70 | psf_foc = wavefront.PSF(Mask,A_f,False) 71 | psf_defoc = wavefront.PSF(Mask,A_def,False) 72 | t0 = wavefront.OTF(psf_foc) 73 | tk = wavefront.OTF(psf_defoc) 74 | q,q2 = PD.Q_matrix(t0,tk,fit.reg,gam) 75 | F_m = PD.F_M(q2,d0, dk,t0,tk,noise_filter,gam) 76 | E_metric = PD.Error_metric(t0,tk,d0,dk,q,noise_filter) 77 | L_m = PD.L_M(E_metric,fit.size) 78 | return L_m 79 | 80 | 81 | 82 | Minimize_partial = partial(Minimise) 83 | mini = minimize(Minimize_partial,p0,method='L-BFGS-B') 84 | result = fit.Minimize_res(mini.x) 85 | patch_wfe = result[1] 86 | patch_mtf = MTF(fftshift(result[0])) 87 | return patch_wfe,patch_mtf 88 | 89 | def fit_patch(self): 90 | 91 | upper = 1700 92 | Nx = np.arange(300,upper,self.Del) 93 | Ny = np.arange(300,upper,self.Del) 94 | if not self.parallel: 95 | for n1 in Nx : 96 | for n2 in Ny: 97 | 98 | print('fitting the area in:'+str(n1)+','+str(n2)) 99 | im0 = self.Im0[n2:n2+self.Del,n1:n1+self.Del] 100 | imk = self.Imk[n2:n2+self.Del,n1:n1+self.Del] 101 | im0 = im0/im0.mean() 102 | imk = imk/imk.mean() 103 | imk = imreg(im0,imk) 104 | 105 | im0 = apo2d(im0,10) 106 | imk = apo2d(imk,10) 107 | 108 | d0,dk = FT(im0,imk) 109 | gam =1# Gamma(d0,dk,M_gamma) 110 | 111 | p0 = np.zeros(self.co_num) 112 | fit = cost_func(self.Del,0.5,1e-10,10,self.co_num,0.5,617.3e-6, 140,4125.3,0.5) 113 | Mask = aperture.mask_pupil(fit.telescope.pupil_size(),fit.size) 114 | noise_temp = noise_mask_high(fit.size,fit.cut_off) 115 | noise_filter = fftshift(noise_temp) 116 | def Minimise(coefficients): 117 | A_f = wavefront.pupil_foc(coefficients,fit.size,fit.telescope.pupil_size(),self.co_num) 118 | A_def = wavefront.pupil_defocus(coefficients,fit.size,fit.del_z,fit.telescope.pupil_size(),self.co_num) 119 | psf_foc = wavefront.PSF(Mask,A_f,False) 120 | psf_defoc = wavefront.PSF(Mask,A_def,False) 121 | t0 = wavefront.OTF(psf_foc) 122 | tk = wavefront.OTF(psf_defoc) 123 | 124 | q,q2 = PD.Q_matrix(t0,tk,fit.reg,gam) 125 | 126 | 127 | F_m = PD.F_M(q2,d0, dk,t0,tk,noise_filter,gam) 128 | 129 | E_metric = PD.Error_metric(t0,tk,d0,dk,q,noise_filter) 130 | L_m = PD.L_M(E_metric,fit.size) 131 | return L_m 132 | 133 | 134 | Minimise_partial = partial(Minimise) 135 | mini = scipy.optimize.minimize(Minimise_partial,p0,method= 'L-BFGS-B') 136 | 137 | result = fit.Minimize_res(mini.x) 138 | self.output_WF[n2:n2+self.Del,n1:n1+self.Del] = result[1] 139 | self.output_MTF[n2:n2+self.Del,n1:n1+self.Del] = MTF(fftshift(result[0])) 140 | hdu = fits.PrimaryHDU(self.output_WF) 141 | hdu.writeto(self.output_wf,overwrite=True) 142 | hdu = fits.PrimaryHDU(self.output_mtf) 143 | hdu.writeto(self.output_mtf,overwrite=True) 144 | else: 145 | print('Initialising parallel computation') 146 | t0 = time.time() 147 | self.patch = tools.prepare_patches(self.data,self.Del,self.Im0,self.Imk) 148 | n_workers = min(6, os.cpu_count()) 149 | print(f'number of workers is {n_workers}') 150 | self.args_list = [i for i in range(len(self.patch))] 151 | self.results_parallel = list(processing.MP.simultaneous(self.run_pd, self.args_list, workers=n_workers)) 152 | dt = (time.time() - t0)/60. 153 | print(f'Time spent in fitting the wavefront error is: {dt: .3f}min') 154 | 155 | def plot_results(self): 156 | 157 | if not self.parallel: 158 | data_mtf = self.output_MTF 159 | data_wfe = self.output_WF 160 | 161 | if self.parallel: 162 | data_mtf,data_wfe = tools.stitch_patches(self.results_parallel,self.Del) 163 | hdu = fits.PrimaryHDU(data_mtf) 164 | hdu.writeto(self.output_mtf,overwrite=True) 165 | hdu = fits.PrimaryHDU(data_wfe) 166 | hdu.writeto(self.output_wf,overwrite=True) 167 | tools.plot_mtf_wf(data_wfe,data_mtf) 168 | 169 | 170 | if (__name__ == '__main__'): 171 | parser = argparse.ArgumentParser(description='PD on sub-fields') 172 | parser.add_argument('-i','--input', help='input') 173 | parser.add_argument('-z','--Z', help='Zernike',default=10) 174 | parser.add_argument('-d','--Del', help='Del',default=265) 175 | parser.add_argument('-ow','--ow', help='output_WFE') 176 | parser.add_argument('-om','--om', help='output MTF') 177 | 178 | parser.add_argument('-p','--parallel',choices=['True','False'],default=True) 179 | parsed = vars(parser.parse_args()) 180 | st = patch_pd(pd_data='{0}'.format(parsed['input']),Del=int(parsed['Del']),co_num=int(parsed['Z']),output_wf='{0}'.format(parsed['ow']),output_mtf='{0}'.format(parsed['om']),parallel=bool(parsed['parallel'])) 181 | st.fit_patch() 182 | st.plot_results() 183 | 184 | -------------------------------------------------------------------------------- /processing.py: -------------------------------------------------------------------------------- 1 | import os 2 | from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, Executor 3 | class Thread: 4 | """ 5 | # Thread 6 | The threads holds the information on the function to execute in a thread or process. 7 | Provides an interface to the `future` object once submitted to an executer. 8 | """ 9 | 10 | def __init__(self, func, args): 11 | self.function = func 12 | self.arguments = args 13 | self.future = None 14 | 15 | def submit(self, executor: Executor): 16 | """Start execution via executor""" 17 | if not self.is_submitted(): 18 | self.future = executor.submit(self.function, self.arguments) 19 | return self 20 | 21 | def is_submitted(self) -> bool: 22 | return self.future is not None 23 | 24 | def is_done(self): 25 | return self.is_submitted() and self.future.done() 26 | 27 | def exception(self): 28 | if not self.is_done(): 29 | return None 30 | return self.future.exception() 31 | 32 | def result(self): 33 | if not self.is_submitted(): 34 | return None 35 | return self.future.result() 36 | 37 | class MP: 38 | """ 39 | ## MP Multi-Processing 40 | Class provides housekeeping / setup methods to reduce the programming overhead of 41 | spawning threads or processes. 42 | """ 43 | 44 | #: Number of CPUs of the current machine 45 | NUM_CPUs = round(os.cpu_count() * 0.8) 46 | 47 | @staticmethod 48 | def threaded(func, args, workers=10, raise_exception=True): 49 | """ 50 | Calls the given function in multiple threads for the set of given arguments 51 | Note that this does not spawn processes, but threads. Use this for non CPU 52 | CPU dependent tasks, i.e. I/O 53 | Method returns once all calls are done. 54 | 55 | ### Params 56 | - func: [Function] the function to call 57 | - args: [Iterable] the 'list' of arguments for each call 58 | - workers: [Integer] the number of concurrent threads to use 59 | - raise_exception: [Bool] Flag if an exception in a thread shall be raised or just logged 60 | 61 | ### Returns 62 | Results from all `Threads` as list 63 | """ 64 | if len(args) == 1: 65 | return list(func(arg) for arg in args) 66 | 67 | with ThreadPoolExecutor(workers) as ex: 68 | threads = [Thread(func, arg).submit(ex) for arg in args] 69 | return MP.collect_results(threads, raise_exception) 70 | 71 | @staticmethod 72 | def simultaneous(func, args, workers=None, raise_exception=True): 73 | """ 74 | Calls the given function in multiple processes for the set of given arguments 75 | Note that this does spawn processes, not threads. Use this for task that 76 | depend heavily on CPU and can be done in parallel. 77 | Method returns once all calls are done. 78 | 79 | ### Params 80 | - func: [Function] the function to call 81 | - args: [Iterable] the 'list' of arguments for each call 82 | - workers: [Integer] the number of concurrent threads to use (Default: NUM_CPUs) 83 | - raise_exception: [Bool] Flag if an exception in a thread shall be raised or just logged 84 | 85 | ### Returns 86 | Results from all `Threads` as list 87 | """ 88 | if len(args) == 1: 89 | return list(func(arg) for arg in args) 90 | 91 | if workers is None: 92 | workers = MP.NUM_CPUs 93 | with ProcessPoolExecutor(workers) as ex: 94 | threads = [Thread(func, arg).submit(ex) for arg in args] 95 | return MP.collect_results(threads, raise_exception) 96 | 97 | @staticmethod 98 | def collect_results(threads: list, raise_exception: bool = True) -> list: 99 | """ 100 | Takes a list of threads and waits for them to be executed. Collects results. 101 | 102 | ### Params 103 | - threads: [List] a list of submitted threads 104 | - raise_exception: [Bool] Flag if an exception in a thread shall be raised or just logged 105 | 106 | ### Returns 107 | Results from all `Threads` as list 108 | """ 109 | result = [] 110 | while len(threads) > 0: 111 | for thread in threads: 112 | if not thread.is_submitted(): 113 | threads.remove(thread) 114 | if not thread.is_done(): 115 | continue 116 | 117 | if thread.exception() is not None: 118 | MP.__exception_handling(threads, thread, raise_exception) 119 | else: 120 | result.append(thread.result()) 121 | threads.remove(thread) 122 | return result 123 | 124 | 125 | -------------------------------------------------------------------------------- /pypd: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | echo "Phase diversity package" 4 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | setup(name='PyPD', 3 | version='0.2.1', 4 | keywords = ('optics', 'telescope', 'zernike','phase diversity'), 5 | description='Module for Phase diversity analysis', 6 | license = 'MPS License', 7 | install_requires = ['numpy>=1.9.3','matplotlib>=1.4.3','unwrap'], 8 | author='Fatima Kahil', 9 | author_email='kahil@mps.mpg.de', 10 | url= 'https://github.com/fakahil/PyPD', 11 | packages = setuptools.find_packages(), 12 | classifiers=[ 13 | "Programming Language :: Python :: 3", 14 | "License :: OSI Approved :: MIT License", 15 | "Operating System :: OS Independent",],) 16 | -------------------------------------------------------------------------------- /telescope.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | 5 | 6 | 7 | #from . import detector 8 | #from .detector import * 9 | 10 | class Telescope(object): 11 | 12 | def __init__(self,lam,diameter,focal_length,name,platescale,size): 13 | self.lam = lam 14 | self.diameter = diameter 15 | self.focal_length = focal_length 16 | self.name = name 17 | self.nu_cutoff = self.diameter/self.lam 18 | self.size = size 19 | self.platescale = platescale 20 | #self.camera = Detector() 21 | 22 | def calc_ps(self, pix_size): 23 | return 206265*pix_size/self.focal_length 24 | 25 | def show_lam(self): 26 | return self.lam 27 | 28 | def show_diameter(self): 29 | return self.diameter 30 | 31 | def pupil_size(self): 32 | pix = self.platescale 33 | pixrad = pix*np.pi/(180*3600) 34 | 35 | deltanu = 1./(self.size*pixrad) 36 | rpupil = self.nu_cutoff/(2*deltanu) 37 | return np.int(rpupil) 38 | 39 | def spatial_resolution(self): 40 | return self.diameter/self.lam 41 | 42 | ''' 43 | def pupil_size(D,lam,pix,size): 44 | pixrad = pix*np.pi/(180*3600) # Pixel-size in radians 45 | nu_cutoff = D/lam # Cutoff frequency in rad^-1 46 | deltanu = 1./(size*pixrad) # Sampling interval in rad^-1 47 | rpupil = nu_cutoff/(2*deltanu) #pupil size in pixels 48 | return np.int(rpupil) 49 | 50 | def calc_plate_scale(pix_size, focal_length): 51 | return 206265*pix_size/focal_length 52 | ''' 53 | 54 | 55 | 56 | -------------------------------------------------------------------------------- /telescope.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/telescope.pyc -------------------------------------------------------------------------------- /tools.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pylab 3 | import matplotlib.pyplot as plt 4 | import scipy 5 | from scipy.fftpack import fftshift, ifftshift, fft2, ifft2 6 | from mpl_toolkits.axes_grid1 import make_axes_locatable, axes_size 7 | from scipy.signal import correlate2d as correlate 8 | from scipy.signal import general_gaussian 9 | from scipy import ndimage 10 | import imreg_dft as ird 11 | from image_registration import chi2_shift 12 | from image_registration.fft_tools import shift 13 | 14 | 15 | def GetPSD1D(psd2D): 16 | h = psd2D.shape[0] 17 | w = psd2D.shape[1] 18 | wc = w//2 19 | hc = h//2 20 | 21 | # create an array of integer radial distances from the center 22 | Y, X = np.ogrid[0:h, 0:w] 23 | r = np.hypot(X - wc, Y - hc).astype(np.int) 24 | 25 | # SUM all psd2D pixels with label 'r' for 0<=r<=wc 26 | # NOTE: this will miss power contributions in 'corners' r>wc 27 | psd1D = ndimage.mean(psd2D, r, index=np.arange(0, wc)) 28 | 29 | return psd1D 30 | 31 | def RMS_WF(array): 32 | rms = np.sqrt(np.sum((1/(2*np.pi)*array)**2)) 33 | rms = 1./(rms) 34 | return rms 35 | 36 | 37 | def Apod(im, size,power, sigma): 38 | av = im.mean() 39 | im = im-av 40 | gen_gaussian = general_gaussian(size, power, sigma) 41 | window = np.outer(gen_gaussian, gen_gaussian) 42 | apod = im*window 43 | apod = apod+av 44 | return apod 45 | 46 | def apo2d(masi,perc): 47 | s = masi.shape 48 | edge = 100./perc 49 | mean = np.mean(masi) 50 | masi = masi-mean 51 | xmask = np.ones(s[1]) 52 | ymask = np.ones(s[0]) 53 | smooth_x = np.int(s[1]/edge) 54 | smooth_y = np.int(s[0]/edge) 55 | 56 | for i in range(0,smooth_x): 57 | xmask[i] = (1.-np.cos(np.pi*np.float(i)/np.float(smooth_x)))/2. 58 | ymask[i] = (1.-np.cos(np.pi*np.float(i)/np.float(smooth_y)))/2. 59 | 60 | xmask[s[1] - smooth_x:s[1]] = (xmask[0:smooth_x])[::-1] 61 | ymask[s[0] - smooth_y:s[0]] = (ymask[0:smooth_y])[::-1] 62 | 63 | #mask_x = np.outer(xmask,xmask) 64 | #mask_y = np.outer(ymask,ymask) 65 | for i in range(0,s[1]): 66 | masi[:,i] = masi[:,i]*xmask[i] 67 | for i in range(0,s[0]): 68 | masi[i,:] = masi[i,:]*ymask[i] 69 | masi = masi+mean 70 | return masi 71 | 72 | 73 | def strehl(rms): 74 | return np.exp(-2*(np.pi*rms**2)) 75 | 76 | def imreg(im0,imk): 77 | 78 | xoff, yoff, exoff, eyoff = chi2_shift(im0,imk) 79 | timg = ird.transform_img(imk, tvec=np.array([-yoff,-xoff])) 80 | return timg 81 | 82 | def RMS_cont(data): 83 | return data.std()/data.mean() 84 | 85 | def noise(im): 86 | from skimage.restoration import estimate_sigma 87 | s = estimate_sigma(im) 88 | return s 89 | 90 | def plot_zernike(coeff): 91 | n = coeff.shape[0] 92 | index = np.arange(n) 93 | fig = plt.figure(figsize=(9, 6), dpi=80) 94 | width = 0.4 95 | for i in index: 96 | #xticklist.append('Z'+str(i+4)) 97 | barfigure = plt.bar(index, coeff/(2*np.pi), width,color = '#2E9AFE',edgecolor = '#2E9AFE') 98 | plt.xticks(np.arange(1, 11, step=1)) 99 | plt.xlabel('Zernike Polynomials',fontsize=18) 100 | plt.ylabel('Coefficient [$\lambda$]',fontsize=18) 101 | plt.title('Zernike Polynomials Coefficients',fontsize=18) 102 | plt.savefig('Zernikes.png',dpi=300) 103 | 104 | def prepare_patches(d,Del,Im0,Imk): 105 | n = d.shape[0] 106 | upper = 1700 107 | lower = 300 108 | Nx = np.arange(lower,upper,Del) 109 | Ny = np.arange(lower,upper,Del) 110 | i_max = np.floor((upper-lower)/Del)+1 111 | patches = np.zeros((int(i_max**2),Del,Del,n)) 112 | #output_WF =np.zeros((int(i_max**2),Del,Del)) 113 | #output_mtf = np.zeros((int(i_max**2),Del,Del)) 114 | k=0 115 | for n1 in Nx : 116 | for n2 in Ny: 117 | patches[k,:,:,0]=Im0[n2:n2+Del,n1:n1+Del] 118 | patches[k,:,:,1]=Imk[n2:n2+Del,n1:n1+Del] 119 | k = k+1 120 | return patches 121 | 122 | 123 | def stitch_patches(results,Del): 124 | data1 = [r[0] for r in results] 125 | data2 = [r[1] for r in results] 126 | upper = 1700-Del 127 | lower = 300 128 | Nx = np.arange(lower,upper,Del) 129 | Ny = np.arange(lower,upper,Del) 130 | i_max = np.floor((upper-lower)/Del)+1 131 | k = 0 132 | if len(data1)==(np.floor((upper-lower)/Del)+1)**2: 133 | 134 | st_wf = np.zeros((2048,2048)) 135 | st_mtf = np.zeros((2048,2048)) 136 | else: 137 | raise TypeError('Check dimensions!') 138 | 139 | for n1 in Nx : 140 | for n2 in Ny: 141 | st_wf[n2:n2+Del,n1:n1+Del] = data1[k] 142 | st_mtf[n2:n2+Del,n1:n1+Del] = data2[k] 143 | k=k+1 144 | return st_mtf,st_wf 145 | 146 | 147 | 148 | def plot_mtf_wf(ph,mtf): 149 | 150 | fig=plt.figure(figsize=(20,8)) 151 | aspect = 5 152 | pad_fraction = 0.5 153 | ax = fig.add_subplot(1,2,1) 154 | im=ax.imshow(ph/(2*np.pi), cmap=pylab.gray(),origin='lower',vmin=-1.2,vmax=1.2) 155 | 156 | ax.set_xlabel('[Pixels]',fontsize=18) 157 | ax.set_ylabel('[Pixels]',fontsize=18) 158 | divider = make_axes_locatable(ax) 159 | cax = divider.append_axes("right", size=0.15, pad=0.05) 160 | cbar = plt.colorbar(im, cax=cax,orientation='vertical') 161 | cbar.set_label('WF error HRT [$\lambda$]',fontsize=20) 162 | cax.tick_params(labelsize=14) 163 | ax2 = fig.add_subplot(1,2,2) 164 | 165 | im2=ax2.imshow(mtf,cmap=pylab.gray(),origin='lower',vmin=0,vmax=1) 166 | ax2.set_xlabel('[Pixels]',fontsize=18) 167 | divider = make_axes_locatable(ax2) 168 | cax2 = divider.append_axes("right", size=0.15, pad=0.05) 169 | cbar2 = plt.colorbar(im2, cax=cax2,orientation='vertical') 170 | cax2.tick_params(labelsize=14) 171 | cbar2.set_label('MTF',fontsize=16) 172 | 173 | plt.subplots_adjust(wspace=.2, hspace=None) 174 | plt.savefig('WFE+MTF.png',dpi=300) 175 | 176 | 177 | def compute_residual_shifts(pd_pair,Del): 178 | 179 | d = fits.getdata(pd_pair) 180 | xoff, yoff, exoff, eyoff = chi2_shift(d[0,500:1000,500:1000],d[1,500:1000,500:1000]) 181 | Imk = ird.transform_img(d[1,:,:], tvec=np.array([-yoff,-xoff])) 182 | Nx = np.arange(200,1800,Del) 183 | Ny = np.arange(200,1800,Del) 184 | shifts_x = np.zeros((2048,2048)) 185 | shifts_y = np.zeros((2048,2048)) 186 | S_x = [] 187 | S_y = [] 188 | 189 | for n1 in Nx : 190 | for n2 in Ny: 191 | 192 | im0 = Im0[n2:n2+Del,n1:n1+Del] 193 | imk = Imk[n2:n2+Del,n1:n1+Del] 194 | xoff, yoff, exoff, eyoff = chi2_shift(im0,imk) 195 | print(xoff, yoff) 196 | shifts_x[n2:n2+Del,n1:n1+Del] = xoff 197 | shifts_y[n2:n2+Del,n1:n1+Del] = yoff 198 | return shifts_x, shifts_y 199 | -------------------------------------------------------------------------------- /tools.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/tools.pyc -------------------------------------------------------------------------------- /wavefront.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import scipy 4 | from scipy.fftpack import fftshift, ifftshift, fft2, ifft2 5 | from scipy.signal import correlate2d as correlate 6 | from scipy import ndimage 7 | import zernike 8 | 9 | def phase(coefficients,rpupil, co_num): 10 | r = 1 11 | x = np.linspace(-r, r, 2*rpupil) 12 | y = np.linspace(-r, r, 2*rpupil) 13 | 14 | [X,Y] = np.meshgrid(x,y) 15 | R = np.sqrt(X**2+Y**2) 16 | theta = np.arctan2(Y, X) 17 | 18 | Z = zernike.Zernike_polar(coefficients,R,theta, co_num) 19 | Z[R>1] = 0 20 | return Z 21 | 22 | 23 | def phase_aberr(del_z,rpupil): 24 | alpha_4 = (del_z*np.pi)/np.sqrt(3) 25 | r = 1 26 | x = np.linspace(-r, r, 2*rpupil) 27 | y = np.linspace(-r, r, 2*rpupil) 28 | [X,Y] = np.meshgrid(x,y) 29 | #ph_defocus = -(np.pi*del_z*(X**2+Y**2)*D**2)/(4*lam*f**2) 30 | R = np.sqrt(X**2 + Y**2) 31 | ph_defocus = alpha_4 * np.sqrt(3)*(2*R**2-1) 32 | ph_defocus[R>1] = 0 33 | return ph_defocus 34 | 35 | 36 | # Embed the phase in the center of a quadratic map of dimensions sizexsize 37 | def phase_embed(coefficients,rpupil,size,co_num): 38 | ph = phase(coefficients,rpupil,co_num) 39 | A = np.zeros([size,size]) 40 | A[size//2-rpupil+1:size//2+rpupil+1,size//2-rpupil+1:size//2+rpupil+1]= phase(coefficients,rpupil,co_num) 41 | return A 42 | 43 | 44 | ## function for making the complex pupil function from Zernike 45 | def pupil_foc(coefficients,size,rpupil,co_num): 46 | 47 | A = np.zeros([size,size]) 48 | A[size//2-rpupil+1:size//2+rpupil+1,size//2-rpupil+1:size//2+rpupil+1]= phase(coefficients,rpupil,co_num) 49 | aberr = np.exp(1j*A) 50 | return aberr 51 | 52 | 53 | ## function for making the complex pupil function for the defocused wavefront 54 | def pupil_defocus(coefficients,size,del_z,rpupil, co_num): 55 | A = np.zeros([size,size]) 56 | A[size//2-rpupil+1:size//2+rpupil+1,size//2-rpupil+1:size//2+rpupil+1] = phase_aberr(del_z,rpupil)+phase(coefficients,rpupil,co_num) 57 | 58 | aberr_defocus = np.exp(1j*A) 59 | return aberr_defocus 60 | 61 | 62 | 63 | 64 | def PSF(mask,abbe,norm): 65 | ## making zero where the aberration is equal to 1 (the zero background) 66 | 67 | abbe_z = np.zeros((len(abbe),len(abbe)),dtype=np.complex) 68 | abbe_z = mask*abbe 69 | PSF = ifftshift(fft2(fftshift(abbe_z))) #from brandon 70 | PSF = (np.abs(PSF))**2 #or PSF*PSF.conjugate() 71 | if norm=='True': 72 | PSF = PSF/PSF.sum() 73 | return PSF 74 | else: 75 | return PSF 76 | 77 | def OTF(psf): 78 | otf = ifftshift(psf) 79 | otf = fft2(otf) 80 | otf = otf/np.real(otf[0,0]) 81 | 82 | #otf = otf/otf.max() # or otf_max = otf[size/2,size/2] if max is shifted to center 83 | 84 | return otf 85 | 86 | 87 | def MTF(otf): 88 | mtf = np.abs(otf) 89 | return mtf 90 | 91 | 92 | 93 | -------------------------------------------------------------------------------- /wavefront.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/wavefront.pyc -------------------------------------------------------------------------------- /zernike.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | ''' 5 | # Zernike polynomials in cartesian coordinates 6 | def Zernike_cartesien(coefficients, x,y): 7 | Z = [0]+coefficients 8 | r = np.sqrt(x**2 + y**2) 9 | Z1 = Z[1] * 1 10 | Z2 = Z[2] * 2*x 11 | Z3 = Z[3] * 2*y 12 | 13 | Z4 = Z[4] * np.sqrt(3)*(2*r**2-1) 14 | Z5 = Z[5] * 2*np.sqrt(6)*x*y 15 | Z6 = Z[6] * np.sqrt(6)*(x**2-y**2) 16 | Z7 = Z[7] * np.sqrt(8)*y*(3*r**2-2) 17 | Z8 = Z[8] * np.sqrt(8)*x*(3*r**2-2) 18 | Z9 = Z[9] * np.sqrt(8)*y*(3*x**2-y**2) 19 | Z10 = Z[10] * np.sqrt(8)*x*(x**2-3*y**2) 20 | Z11 = Z[11] * np.sqrt(5)*(6*r**4-6*r**2+1) 21 | Z12 = Z[12] * np.sqrt(10)*(x**2-y**2)*(4*r**2-3) 22 | Z13 = Z[13] * 2*np.sqrt(10)*x*y*(4*r**2-3) 23 | Z14 = Z[14] * np.sqrt(10)*(r**4-8*x**2*y**2) 24 | Z15 = Z[15] * 4*np.sqrt(10)*x*y*(x**2-y**2) 25 | Z16 = Z[16] * np.sqrt(12)*x*(10*r**4-12*r**2+3) 26 | Z17 = Z[17] * np.sqrt(12)*y*(10*r**4-12*r**2+3) 27 | Z18 = Z[18] * np.sqrt(12)*x*(x**2-3*y**2)*(5*r**2-4) 28 | Z19 = Z[19] * np.sqrt(12)*y*(3*x**2-y**2)*(5*r**2-4) 29 | Z20 = Z[20] * np.sqrt(12)*x*(16*x**4-20*x**2*r**2+5*r**4) 30 | Z21 = Z[21] * np.sqrt(12)*y*(16*y**4-20*y**2*r**2+5*r**4) 31 | Z22 = Z[22] * np.sqrt(7)*(20*r**6-30*r**4+12*r**2-1) 32 | Z23 = Z[23] * 2*np.sqrt(14)*x*y*(15*r**4-20*r**2+6) 33 | Z24 = Z[24] * np.sqrt(14)*(x**2-y**2)*(15*r**4-20*r**2+6) 34 | Z25 = Z[25] * 4*np.sqrt(14)*x*y*(x**2-y**2)*(6*r**2-5) 35 | Z26 = Z[26] * np.sqrt(14)*(8*x**4-8*x**2*r**2+r**4)*(6*r**2-5) 36 | Z27 = Z[27] * np.sqrt(14)*x*y*(32*x**4-32*x**2*r**2+6*r**4) 37 | Z28 = Z[28] * np.sqrt(14)*(32*x**6-48*x**4*r**2+18*x**2*r**4-r**6) 38 | Z29 = Z[29] * 4*y*(35*r**6-60*r**4+30*r**2-4) 39 | Z30 = Z[30] * 4*x*(35*r**6-60*r**4+30*r**2-4) 40 | Z31 = Z[31] * 4*y*(3*x**2-y**2)*(21*r**4-30*r**2+10) 41 | Z32 = Z[32] * 4*x*(x**2-3*y**2)*(21*r**4-30*r**2+10) 42 | Z33 = Z[33] * 4*(7*r**2-6)*(4*x**2*y*(x**2-y**2)+y*(r**4-8*x**2*y**2)) 43 | Z34 = Z[34] * (4*(7*r**2-6)*(x*(r**4-8*x**2*y**2)-4*x*y**2*(x**2-y**2))) 44 | Z35 = Z[35] * (8*x**2*y*(3*r**4-16*x**2*y**2)+4*y*(x**2-y**2)*(r**4-16*x**2*y**2)) 45 | Z36 = Z[36] * (4*x*(x**2-y**2)*(r**4-16*x**2*y**2)-8*x*y**2*(3*r**4-16*x**2*y**2)) 46 | Z37 = Z[37] * 3*(70*r**8-140*r**6+90*r**4-20*r**2+1) 47 | 48 | ZW = Z1 + Z2 + Z3 49 | #+ Z4+ Z5+ Z6+ Z7+ Z8+ Z9+ Z10+ Z11+ Z12+ Z13+ Z14+ Z15+ Z16+ Z17+ Z18+ Z19+Z20+ Z21+ Z22+ Z23+ Z24+ Z25+ Z26+ Z27+ Z28+ Z29+Z30+ Z31+ Z32+ Z33+ Z34+ Z35+ Z36+ Z37 50 | return Zw 51 | ''' 52 | 53 | #https://oeis.org/A176988 54 | 55 | 56 | # Zernike polynomials in polar coordinates 57 | def Zernike_polar(coefficients, r, u, co_num): 58 | #Z= np.insert(np.array([0,0,0]),3,coefficients) 59 | Z = np.zeros(37) 60 | Z[:co_num] = coefficients 61 | 62 | #Z1 = Z[0] * 1*(np.cos(u)**2+np.sin(u)**2) 63 | #Z2 = Z[1] * 2*r*np.cos(u) 64 | #Z3 = Z[2] * 2*r*np.sin(u) 65 | 66 | Z4 = Z[0] * np.sqrt(3)*(2*r**2-1) #defocus 67 | 68 | Z5 = Z[1] * np.sqrt(6)*r**2*np.sin(2*u) #astigma 69 | Z6 = Z[2] * np.sqrt(6)*r**2*np.cos(2*u) 70 | 71 | Z7 = Z[3] * np.sqrt(8)*(3*r**2-2)*r*np.sin(u) #coma 72 | Z8 = Z[4] * np.sqrt(8)*(3*r**2-2)*r*np.cos(u) 73 | 74 | Z9 = Z[5] * np.sqrt(8)*r**3*np.sin(3*u) #trefoil 75 | 76 | Z10= Z[6] * np.sqrt(8)*r**3*np.cos(3*u) 77 | 78 | Z11 = Z[7] * np.sqrt(5)*(1-6*r**2+6*r**4) #secondary spherical 79 | 80 | Z12 = Z[8] * np.sqrt(10)*(4*r**2-3)*r**2*np.cos(2*u) #2 astigma 81 | Z13 = Z[9] * np.sqrt(10)*(4*r**2-3)*r**2*np.sin(2*u) 82 | 83 | Z14 = Z[10] * np.sqrt(10)*r**4*np.cos(4*u) #tetrafoil 84 | Z15 = Z[11] * np.sqrt(10)*r**4*np.sin(4*u) 85 | 86 | Z16 = Z[12] * np.sqrt(12)*(10*r**4-12*r**2+3)*r*np.cos(u) #secondary coma 87 | Z17 = Z[13] * np.sqrt(12)*(10*r**4-12*r**2+3)*r*np.sin(u) 88 | 89 | Z18 = Z[14] * np.sqrt(12)*(5*r**2-4)*r**3*np.cos(3*u) #secondary trefoil 90 | Z19 = Z[15] * np.sqrt(12)*(5*r**2-4)*r**3*np.sin(3*u) 91 | 92 | Z20 = Z[16] * np.sqrt(12)*r**5*np.cos(5*u) #pentafoil 93 | Z21 = Z[17] * np.sqrt(12)*r**5*np.sin(5*u) 94 | 95 | Z22 = Z[18] * np.sqrt(7)*(20*r**6-30*r**4+12*r**2-1) #spherical 96 | 97 | Z23 = Z[19] * np.sqrt(14)*(15*r**4-20*r**2+6)*r**2*np.sin(2*u) #astigma 98 | Z24 = Z[20] * np.sqrt(14)*(15*r**4-20*r**2+6)*r**2*np.cos(2*u) 99 | 100 | Z25 = Z[21] * np.sqrt(14)*(6*r**2-5)*r**4*np.sin(4*u)#trefoil 101 | Z26 = Z[22] * np.sqrt(14)*(6*r**2-5)*r**4*np.cos(4*u) 102 | 103 | Z27 = Z[23] * np.sqrt(14)*r**6*np.sin(6*u) #hexafoil 104 | Z28 = Z[24] * np.sqrt(14)*r**6*np.cos(6*u) 105 | 106 | Z29 = Z[25] * 4*(35*r**6-60*r**4+30*r**2-4)*r*np.sin(u) #coma 107 | Z30 = Z[26] * 4*(35*r**6-60*r**4+30*r**2-4)*r*np.cos(u) 108 | 109 | Z31 = Z[27] * 4*(21*r**4-30*r**2+10)*r**3*np.sin(3*u)#trefoil 110 | Z32 = Z[28] * 4*(21*r**4-30*r**2+10)*r**3*np.cos(3*u) 111 | 112 | Z33 = Z[29] * 4*(7*r**2-6)*r**5*np.sin(5*u) #pentafoil 113 | Z34 = Z[30] * 4*(7*r**2-6)*r**5*np.cos(5*u) 114 | 115 | Z35 = Z[31] * 4*r**7*np.sin(7*u) #heptafoil 116 | Z36 = Z[32] * 4*r**7*np.cos(7*u) 117 | 118 | Z37 = Z[33] * 3*(70*r**8-140*r**6+90*r**4-20*r**2+1) #spherical 119 | 120 | #Z1+Z2+Z3+ 121 | ZW = Z4+Z5+Z6+Z7+Z8+Z9+Z10+Z11+Z12+Z13+Z14+Z15+Z16+ Z17+Z18+Z19+Z20+Z21+Z22+Z23+ Z24+Z25+Z26+Z27+Z28+ Z29+ Z30+ Z31+ Z32+ Z33+ Z34+ Z35+ Z36+ Z37 122 | return ZW 123 | 124 | -------------------------------------------------------------------------------- /zernike.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fakahil/PyPD/366f6e31a961f32d187e6f22609966732bd252d3/zernike.pyc --------------------------------------------------------------------------------