def single_analytical_ft_convolution(without_analytic, with_analytic): """Convolves two convolvable objects utilizing their analytic fourier transforms. Parameters ---------- without_analytic : `Convolvable` A Convolvable object which lacks an analytic fourier transform with_analytic : `Convolvable` A Convolvable object which has an analytic fourier transform Returns ------- `ConvolutionResult` A convolution result """ fourier_data = m.fftshift(m.fft2(m.fftshift(without_analytic.data))) fourier_unit_x = forward_ft_unit(without_analytic.sample_spacing, without_analytic.samples_x) fourier_unit_y = forward_ft_unit(without_analytic.sample_spacing, without_analytic.samples_y) a_ft = with_analytic.analytic_ft(fourier_unit_x, fourier_unit_y) result = _conv_result_core(fourier_data, a_ft) return Convolvable(result, without_analytic.unit_x, without_analytic.unit_y, False)
def _unequal_spacing_conv_core(psf1, psf2): '''Interpolates psf2 before using fft-based convolution Args: psf1 (prysm.PSF): PSF. This one defines the sampling of the output. psf2 (prysm.PSF): PSF. This one will have its frequency response truncated. Returns: PSF: a new `PSF` that is the convolution of psf1 and psf2. ''' # map psf1 into the fourier domain ft1 = fft2(fftshift(psf1.data)) unit1x = forward_ft_unit(psf1.sample_spacing, psf1.samples_x) unit1y = forward_ft_unit(psf1.sample_spacing, psf1.samples_y) # map psf2 into the fourier domain ft2 = fft2(fftshift(psf2.data)) unit2x = forward_ft_unit(psf2.sample_spacing, psf2.samples_x) unit2y = forward_ft_unit(psf2.sample_spacing, psf2.samples_y) ft3 = ifftshift(resample_2d_complex(fftshift(ft2), (unit2y, unit2x), (unit1y, unit1x))) psf3 = PSF(data=abs(ifftshift(ifft2(ft1 * ft3))), sample_spacing=psf1.sample_spacing) return psf3._renorm()
def prop_pupil_plane_to_psf_plane(wavefunction, Q, incoherent=True, norm=None): """Propagate a pupil plane to a PSF plane and compute the grid along which the PSF exists. Parameters ---------- wavefunction : `numpy.ndarray` the pupil wavefunction Q : `float` oversampling / padding factor incoherent : `bool`, optional whether to return the incoherent (real valued) PSF, or the coherent (complex-valued) PSF. Incoherent = |coherent|^2 norm : `str`, {None, 'ortho'} normalization parameter passed directly to numpy/cupy fft Returns ------- psf : `numpy.ndarray` incoherent point spread function """ padded_wavefront = pad2d(wavefunction, Q) impulse_response = m.ifftshift( m.fft2(m.fftshift(padded_wavefront), norm=norm)) if incoherent: return abs(impulse_response)**2 else: return impulse_response
def conv(self, psf2): '''Convolves this PSF with another Args: psf2 (`PSF`): PSf to convolve with this one. Returns: PSF: A new `PSF` that is the convolution of these two PSFs. Notes: output PSF has equal sampling to whichever PSF has a lower nyquist frequency. ''' try: psf_ft = fftshift(fft2(self.data)) psf_unit_x = forward_ft_unit(self.sample_spacing, self.samples_x) psf_unit_y = forward_ft_unit(self.sample_spacing, self.samples_y) psf2_ft = psf2.analytic_ft(psf_unit_x, psf_unit_y) psf3 = PSF(data=abs(ifft2(psf_ft * psf2_ft)), sample_spacing=self.sample_spacing) return psf3._renorm() except AttributeError: # no analytic FT on the PSF/subclass print('No analytic FT, falling back to numerical approach.') return convpsf(self, psf2)
def forward_ft_unit(sample_spacing, samples, shift=True): """Compute the units resulting from a fourier transform. Parameters ---------- sample_spacing : `float` center-to-center spacing of samples in an array samples : `int` number of samples in the data shift : `bool`, optional whether to shift the output. If True, first element is a negative freq if False, first element is 0 freq. Returns ------- `numpy.ndarray` array of sample frequencies in the output of an fft """ unit = m.fftfreq(samples, sample_spacing) if shift: return m.fftshift(unit) else: return unit
def from_pupil(pupil, efl, padding=1): ''' Uses scalar diffraction propogation to generate a PSF from a pupil. Args: pupil (prysm.Pupil): Pupil, with OPD data and wavefunction. efl (float): effective focal length of the optical system. padding (number): number of pupil widths to pad each side of the pupil with during computation. Returns: PSF. A new PSF instance. ''' # padded pupil contains 1 pupil width on each side for a width of 3 psf_samples = (pupil.samples * padding) * 2 + pupil.samples sample_spacing = pupil_sample_to_psf_sample(pupil_sample=pupil.sample_spacing * 1000, num_samples=psf_samples, wavelength=pupil.wavelength, efl=efl) padded_wavefront = pad2d(pupil.fcn, padding) impulse_response = ifftshift(fft2(fftshift(padded_wavefront))) psf = abs(impulse_response)**2 return PSF(psf / np.max(psf), sample_spacing)
def _numerical_ft_convolution_core_equalspacing_unequalsamplecount( more_samples, less_samples): # compute the ordinate axes of the input and output in_x = forward_ft_unit(less_samples.sample_spacing, less_samples.data.shape[0]) in_y = forward_ft_unit(less_samples.sample_spacing, less_samples.data.shape[1]) output_x = forward_ft_unit(more_samples.sample_spacing, more_samples.data.shape[0]) output_y = forward_ft_unit(more_samples.sample_spacing, more_samples.data.shape[1]) # FFT the less sampled one and map it onto the denser grid less_fourier = m.fftshift(m.fft2(m.fftshift(less_samples.data))) interpf = interp2d(in_x, in_y, less_fourier, kind='linear') resampled_less = interpf(output_x, output_y) # FFT convolve the two convolvables more_fourier = m.fftshift(m.fft2(m.fftshift(more_samples.data))) data = _conv_result_core(resampled_less, more_fourier) return Convolvable(data, more_samples.unit_x, more_samples.unit_y, False)
def bandreject_filter(array, sample_spacing, wllow, wlhigh): sy, sx = array.shape # compute the bandpass in sample coordinates ux, uy = forward_ft_unit(sample_spacing, sx), forward_ft_unit(sample_spacing, sy) fhigh, flow = 1 / wllow, 1 / wlhigh # make an ordinate array in frequency space and use it to make a mask uxx, uyy = m.meshgrid(ux, uy) highpass = ((uxx < -fhigh) | (uxx > fhigh)) | ((uyy < -fhigh) | (uyy > fhigh)) lowpass = ((uxx > -flow) & (uxx < flow)) & ((uyy > -flow) & (uyy < flow)) mask = highpass | lowpass # adjust NaNs and FFT work = array.copy() work[~m.isfinite(work)] = 0 fourier = m.fftshift(m.fft2(m.ifftshift(work))) fourier[mask] = 0 out = m.fftshift(m.ifft2(m.ifftshift(fourier))) return out.real
def from_psf(psf): ''' Generates an MTF from a PSF. Args: psf (:class:`PSF`): PSF to compute an MTF from. Returns: :class:`MTF`: A new MTF instance. ''' dat = abs(fftshift(fft2(psf.data))) unit_x = forward_ft_unit(psf.sample_spacing, psf.samples_x) unit_y = forward_ft_unit(psf.sample_spacing, psf.samples_y) return MTF(dat / dat[psf.center_x, psf.center_y], unit_x, unit_y)
def _numerical_ft_convolution_core_unequalspacing(finer_sampled, coarser_sampled): # compute the ordinate axes of the input of each in_x_more = forward_ft_unit(finer_sampled.sample_spacing, finer_sampled.data.shape[0]) in_y_more = forward_ft_unit(finer_sampled.sample_spacing, finer_sampled.data.shape[1]) in_x_less = forward_ft_unit(coarser_sampled.sample_spacing, coarser_sampled.data.shape[0]) in_y_less = forward_ft_unit(coarser_sampled.sample_spacing, coarser_sampled.data.shape[1]) # fourier-space interpolate the larger bandwidth signal onto the grid defined by the lower # bandwidth signal. This assumes the lower bandwidth signal is Nyquist sampled, which is # not necessarily the case. The accuracy of this method depends on the quality of the input. more_fourier = m.fftshift(m.fft2(m.fftshift(finer_sampled.data))) interpf = interp2d(in_x_more, in_y_more, more_fourier, kind='linear') resampled_more = interpf(in_x_less, in_y_less) # FFT the less well sampled input and perform the Fourier based convolution. less_fourier = m.fftshift(m.fft2(m.fftshift(coarser_sampled.data))) data = _conv_result_core(resampled_more, less_fourier) return Convolvable(data, in_x_less, in_y_less, False)
def forward_ft_unit(sample_spacing, samples): """Compute the units resulting from a fourier transform. Parameters ---------- sample_spacing : `float` center-to-center spacing of samples in an array samples : `int` number of samples in the data Returns ------- `numpy.ndarray` array of sample frequencies in the output of an fft """ return m.fftshift(m.fftfreq(samples, sample_spacing))
def prop_pupil_plane_to_psf_plane(wavefunction, Q, norm=None): """Propagate a pupil plane to a PSF plane and compute the grid along which the PSF exists. Parameters ---------- wavefunction : `numpy.ndarray` the pupil wavefunction Q : `float` oversampling / padding factor norm : `str`, {None, 'ortho'} normalization parameter passed directly to numpy/cupy fft Returns ------- psf : `numpy.ndarray` incoherent point spread function """ padded_wavefront = pad2d(wavefunction, Q) impulse_response = m.ifftshift(m.fft2(m.fftshift(padded_wavefront), norm=norm)) return abs(impulse_response) ** 2
def from_psf(psf): """Generate an MTF from a PSF. Parameters ---------- psf : `PSF` PSF to compute an MTF from Returns ------- `MTF` A new MTF instance """ if getattr(psf, '_mtf', None) is not None: return psf._mtf else: dat = abs(m.fftshift(m.fft2(psf.data))) unit_x = forward_ft_unit(psf.sample_spacing / 1e3, psf.samples_x) # 1e3 for microns => mm unit_y = forward_ft_unit(psf.sample_spacing / 1e3, psf.samples_y) return MTF(dat / dat[psf.center_y, psf.center_x], unit_x, unit_y)
def psd(height, sample_spacing, window=None): """Compute the power spectral density of a signal. Parameters ---------- height : `numpy.ndarray` height or phase data sample_spacing : `float` spacing of samples in the input data window : {'welch', 'hann'} or ndarray, optional Returns ------- unit_x : `numpy.ndarray` ordinate x frequency axis unit_y : `numpy.ndarray` ordinate y frequency axis psd : `numpy.ndarray` power spectral density Notes ----- See GH_FFT for a rigorous treatment of FFT scalings https://holometer.fnal.gov/GH_FFT.pdf """ window = make_window(height, sample_spacing, window) fft = m.ifftshift(m.fft2(m.fftshift(height * window))) psd = abs(fft)**2 # mag squared first as per GH_FFT fs = 1 / sample_spacing S2 = (window**2).sum() coef = S2 * fs * fs psd /= coef ux = forward_ft_unit(sample_spacing, height.shape[1]) uy = forward_ft_unit(sample_spacing, height.shape[0]) return ux, uy, psd
def show_fourier(self, interp_method='lanczos', fig=None, ax=None): ''' Displays the fourier transform of the image. Args: interp_method (`string`): method used to interpolate the data for display. fig (`matplotlib.figure`): figure to plot in. ax (`matplotlib.axis`): axis to plot in. Returns: `tuple` containing: `matplotlib.figure`: figure containing the plot. `matplotlib.axis`: axis containing the plot. ''' dat = abs(fftshift(fft2(pad2d(self.data)))) dat /= dat.max() unit_x = forward_ft_unit(self.sample_spacing, self.samples_x) unit_y = forward_ft_unit(self.sample_spacing, self.samples_y) xmin, xmax = unit_x[0], unit_x[-1] ymin, ymax = unit_y[0], unit_y[-1] fig, ax = share_fig_ax(fig, ax) im = ax.imshow(dat**0.1, extent=[xmin, xmax, ymin, ymax], cmap='Greys_r', interpolation=interp_method, origin='lower') fig.colorbar(im) ax.set(xlabel='Spatial Frequency X [cy/mm]', ylabel='Spatial Frequency Y [cy/mm]', title='Normalized FT of image, to 0.1 power') return fig, ax
def synthesize_surface_from_psd(psd, nu_x, nu_y): """Synthesize a surface height map from PSD data. Parameters ---------- psd : `numpy.ndarray` PSD data, units nm²/(cy/mm)² nu_x : `numpy.ndarray` x spatial frequency, cy/mm nu_y : `numpy.ndarray` y spatial frequency, cy_mm """ # generate a random phase to be matched to the PSD randnums = m.rand(*psd.shape) randfft = m.fft2(randnums) phase = m.angle(randfft) # calculate the output window # the 0th element of nu_y has the greatest frequency in magnitude because of # the convention to put the nyquist sample at -fs instead of +fs for even-size arrays fs = -2 * nu_y[0] dx = dy = 1 / fs ny, nx = psd.shape x, y = m.arange(nx) * dx, m.arange(ny) * dy # calculate the area of the output window, "S2" in GH_FFT notation A = x[-1] * y[-1] # use ifft to compute the PSD signal = m.exp(1j * phase) * m.sqrt(A * psd) coef = 1 / dx / dy out = m.ifftshift(m.ifft2(m.fftshift(signal))) * coef out = out.real return x, y, out
def show_fourier(self, freq_x=None, freq_y=None, interp_method='lanczos', fig=None, ax=None): '''Display the fourier transform of the image. Parameters ---------- interp_method : `string` method used to interpolate the data for display. freq_x : iterable x frequencies to use for convolvable with analytical FT and no data freq_y : iterable y frequencies to use for convolvable with analytic FT and no data fig : `matplotlib.figure.Figure` Figure containing the plot ax : `matplotlib.axes.Axis` Axis containing the plot Returns ------- fig : `matplotlib.figure.Figure` Figure containing the plot ax : `matplotlib.axes.Axis` Axis containing the plot Notes ----- freq_x and freq_y are unused when the convolvable has a .data field. ''' if self.has_analytic_ft: if self.data is None: if freq_x is None or freq_y is None: raise ValueError( 'Convolvable has analytic FT and no data, must provide x and y coordinates' ) else: lx, ly, ss = len(self.unit_x), len( self.unit_y), self.sample_spacing freq_x, freq_y = forward_ft_unit(ss, lx), forward_ft_unit(ss, ly) data = self.analytic_ft(freq_x, freq_y) else: data = abs(m.fftshift(m.fft2(pad2d(self.data)))) data /= data.max() freq_x = forward_ft_unit(self.sample_spacing, self.samples_x) freq_y = forward_ft_unit(self.sample_spacing, self.samples_y) xmin, xmax = freq_x[0], freq_x[-1] ymin, ymax = freq_y[0], freq_y[-1] fig, ax = share_fig_ax(fig, ax) im = ax.imshow(data, extent=[xmin, xmax, ymin, ymax], cmap='Greys_r', interpolation=interp_method, origin='lower') fig.colorbar(im, ax=ax, label='Normalized Spectral Intensity [a.u.]') ax.set(xlim=(xmin, xmax), xlabel=r' $\nu_x$ [cy/mm]', ylim=(ymin, ymax), ylabel=r' $\nu_y$ [cy/mm]') return fig, ax
def _conv_result_core(data1, data2): dat = abs(m.fftshift(m.ifft2(data1 * data2))) return dat / dat.max()
def _numerical_ft_convolution_core_equalspacing(convolvable1, convolvable2): # two are identically sampled; just FFT convolve them without modification ft1 = m.fftshift(m.fft2(m.fftshift(convolvable1.data))) ft2 = m.fftshift(m.fft2(m.fftshift(convolvable2.data))) data = _conv_result_core(ft1, ft2) return Convolvable(data, convolvable1.unit_x, convolvable1.unit_y, False)