diff --git a/CDTools/datasets.py b/CDTools/datasets.py index f6be09d..c18d801 100644 --- a/CDTools/datasets.py +++ b/CDTools/datasets.py @@ -221,7 +221,6 @@ class Ptycho_2D_Dataset(CDataset): mask = 1 if not updating: - print('hi') axes[0].set_title('Nanomap') axes[1].set_title('Pattern') diff --git a/CDTools/tools/analysis.py b/CDTools/tools/analysis.py index d859f50..e7e7dc4 100644 --- a/CDTools/tools/analysis.py +++ b/CDTools/tools/analysis.py @@ -7,7 +7,8 @@ from CDTools.tools import image_processing as ip from scipy import fftpack __all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions', - 'calc_consistency_prtf'] + 'calc_consistency_prtf', 'calc_deconvolved_cross_correlation', + 'calc_frc'] from matplotlib import pyplot as plt @@ -315,5 +316,465 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None): prtfs = t.Tensor(prtfs) return bins[:-1], np.mean(prtfs,axis=0) - + + +def calc_deconvolved_cross_correlation(im1, im2): + """Calculates a cross-correlation between two images with their autocorrelations deconvolved. + + This can also be thought of as the inverse Fourier transform of the + object from which the Fourier Ring Correlation is defined. + + Args: + im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + + Returns: + (t.Tensor) : The deconvolved cross-correlation, in real space + + """ + # + # Here's my approach, perhaps it's a little unconventional. I will first + # calculate the phase correlation function as found in ____ (cite a paper + # defining it). This is strongly peaked, so I can take a small window + # of say, 10x10 pixels, and then do a sinc interpolation of that area + # using an FFT with upsampling by a factor of resolution in reciprocal + # space + # + + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + + cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2))) + + # Not sure if this is more or less stable than just the correlation + # maximum - requires some testing + cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2) + + if im_np: + cor = cmath.torch_to_complex(cor) + + return cor + + +def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1): + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + + if im_slice is None: + im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, + (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] + im_slice = np.s_[300:-300,300:-300] + + if nbins is None: + nbins = np.max(synth_obj[im_slice].shape) // 4 + + + cor_fft = cmath.cmult(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy() + + + di = np.linalg.norm(basis[:,0]) + dj = np.linalg.norm(basis[:,1]) + + i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di)) + j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj)) + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + + + synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft) + + prtfs = [] + for obj in objects: + obj = obj[obj_slice] + single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy() + single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft) + + prtfs.append(synth_ints/single_ints) + + + if not obj_np: + bins = t.Tensor(bins) + prtfs = t.Tensor(prtfs) + + return bins[:-1], np.mean(prtfs,axis=0) + + + +def calc_deconvolved_cross_correlation(im1, im2): + """Calculates a cross-correlation between two images with their autocorrelations deconvolved. + + This can also be thought of as the inverse Fourier transform of the + object from which the Fourier Ring Correlation is defined. + + Args: + im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + + Returns: + (t.Tensor) : The deconvolved cross-correlation, in real space + + """ + # + # Here's my approach, perhaps it's a little unconventional. I will first + # calculate the phase correlation function as found in ____ (cite a paper + # defining it). This is strongly peaked, so I can take a small window + # of say, 10x10 pixels, and then do a sinc interpolation of that area + # using an FFT with upsampling by a factor of resolution in reciprocal + # space + # + + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + + cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2))) + + # Not sure if this is more or less stable than just the correlation + # maximum - requires some testing + cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2) + + if im_np: + cor = cmath.torch_to_complex(cor) + + return cor + + +def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1): + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + + if im_slice is None: + im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, + (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] + im_slice = np.s_[300:-300,300:-300] + + if nbins is None: + nbins = np.max(synth_obj[im_slice].shape) // 4 + + + cor_fft = cmath.cmult(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy() + + + di = np.linalg.norm(basis[:,0]) + dj = np.linalg.norm(basis[:,1]) + + i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di)) + j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj)) + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + + + synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft) + + prtfs = [] + for obj in objects: + obj = obj[obj_slice] + single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy() + single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft) + + prtfs.append(synth_ints/single_ints) + + + if not obj_np: + bins = t.Tensor(bins) + prtfs = t.Tensor(prtfs) + + return bins[:-1], np.mean(prtfs,axis=0) + + + +def calc_deconvolved_cross_correlation(im1, im2): + """Calculates a cross-correlation between two images with their autocorrelations deconvolved. + + This can also be thought of as the inverse Fourier transform of the + object from which the Fourier Ring Correlation is defined. + + Args: + im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + + Returns: + (t.Tensor) : The deconvolved cross-correlation, in real space + + """ + # + # Here's my approach, perhaps it's a little unconventional. I will first + # calculate the phase correlation function as found in ____ (cite a paper + # defining it). This is strongly peaked, so I can take a small window + # of say, 10x10 pixels, and then do a sinc interpolation of that area + # using an FFT with upsampling by a factor of resolution in reciprocal + # space + # + + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + + cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2))) + + # Not sure if this is more or less stable than just the correlation + # maximum - requires some testing + cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2) + + if im_np: + cor = cmath.torch_to_complex(cor) + + return cor + + +def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1): + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + + if im_slice is None: + im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, + (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] + im_slice = np.s_[300:-300,300:-300] + + if nbins is None: + nbins = np.max(synth_obj[im_slice].shape) // 4 + + + cor_fft = cmath.cmult(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy() + + + di = np.linalg.norm(basis[:,0]) + dj = np.linalg.norm(basis[:,1]) + + i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di)) + j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj)) + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + + + synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft) + + prtfs = [] + for obj in objects: + obj = obj[obj_slice] + single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy() + single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft) + + prtfs.append(synth_ints/single_ints) + + + if not obj_np: + bins = t.Tensor(bins) + prtfs = t.Tensor(prtfs) + + return bins[:-1], np.mean(prtfs,axis=0) + + + +def calc_deconvolved_cross_correlation(im1, im2, im_slice=None): + """Calculates a cross-correlation between two images with their autocorrelations deconvolved. + + This can also be thought of as the inverse Fourier transform of the + object from which the Fourier Ring Correlation is defined. + + Args: + im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + im_slice (slice) : Default is from 3/8 to 5/8 across the image, a slice to use in the processing. + + Returns: + (t.Tensor) : The deconvolved cross-correlation, in real space + + """ + + + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + if im_slice is None: + im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, + (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] + + + cor_fft = cmath.cmult(t.fft(im1[im_slice],2), + cmath.cconj(t.fft(im2[im_slice],2))) + + # Not sure if this is more or less stable than just the correlation + # maximum - requires some testing + cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2) + + if im_np: + cor = cmath.torch_to_complex(cor) + + return cor + + +def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): + """Calculates a Fourier ring correlation between two images + + This function requires an input of a basis to allow for FRC calculations + to be related to physical units. + + Like other analysis functions, this can take input in numpy or pytorch, + and will return output in the respective format. + + Args: + im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array + basis (t.Tensor) : The basis for the images, defined as is standard for datasets + im_slice (slice) : Default is from 3/8 to 5/8 across the image, a slice to use in the processing. + nbins (int) : Number of bins to break the FRC up into + snr (float) : The signal to noise ratio (for the combined information in both images) to return a threshold curve for. + + Returns: + (t.Tensor) : The frequencies associated with each FRC value + (t.Tensor) : The FRC values + (t.Tensor) : The threshold curve for comparison + + """ + + + + im_np = False + if isinstance(im1, np.ndarray): + im1 = cmath.complex_to_torch(im1) + im_np = True + if isinstance(im2, np.ndarray): + im2 = cmath.complex_to_torch(im2) + im_np = True + + if isinstance(basis, np.ndarray): + basis = t.tensor(basis) + + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + if im2.shape[-1] != 2: + im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) + + + if im_slice is None: + im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, + (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] + + if nbins is None: + nbins = np.max(synth_obj[im_slice].shape) // 4 + + + cor_fft = cmath.cmult(cmath.fftshift(t.fft(im1[im_slice],2)), + cmath.fftshift(cmath.cconj(t.fft(im2[im_slice],2)))) + + F1 = cmath.cabs(cmath.fftshift(t.fft(im1[im_slice],2)))**2 + F2 = cmath.cabs(cmath.fftshift(t.fft(im2[im_slice],2)))**2 + + + di = np.linalg.norm(basis[:,0]) + dj = np.linalg.norm(basis[:,1]) + + i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di)) + j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj)) + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + + + + numerator, bins = np.histogram(Rs,bins=nbins,weights=cmath.torch_to_complex(cor_fft)) + denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy()) + denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy()) + n_pix, bins = np.histogram(Rs,bins=nbins) + + frc = np.abs(numerator / np.sqrt(denominator_F1*denominator_F2)) + + # This moves from combined-image SNR to single-image SNR + snr /= 2 + + threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \ + (1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix)) + + if not im_np: + bins = t.tensor(bins) + frc = t.tensor(frc) + threshold = t.tensor(threshold) + + return bins[:-1], frc, threshold diff --git a/examples/MIT_BNL_logo.py b/examples/MIT_BNL_logo.py index c2df1d7..581b904 100644 --- a/examples/MIT_BNL_logo.py +++ b/examples/MIT_BNL_logo.py @@ -60,6 +60,7 @@ for i, loss in enumerate(model.Adam_optimize(15, dataset, batch_size=15)): for i, loss in enumerate(model.Adam_optimize(15, dataset, batch_size=15, lr=0.0005)): print(i,loss) model.inspect(dataset) + # The third phase lowers the rate further for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.00005)): diff --git a/examples/example_reconstructions/gold_balls.pickle b/examples/example_reconstructions/gold_balls.pickle index 86a4f04..8db6ff1 100644 Binary files a/examples/example_reconstructions/gold_balls.pickle and b/examples/example_reconstructions/gold_balls.pickle differ diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 131c3e7..c0b5ada 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -2,6 +2,7 @@ from __future__ import division, print_function, absolute_import import pytest import numpy as np +from scipy import fftpack as ffts import torch as t from itertools import combinations @@ -179,3 +180,87 @@ def test_calc_consistency_prtf(): # Check that the maximum frequency is correct for the basis assert np.isclose(freqs[-1]-freqs[-2] + freqs[-1], np.sqrt(1/4**2 + 1/6**2)) + +def test_calc_deconvolved_cross_correlation(): + + obj1 = np.random.rand(200,300) + 1j * np.random.rand(200,300) + obj2 = np.random.rand(200,300) + 1j * np.random.rand(200,300) + + cor_fft = np.fft.fft2(obj1) * np.conj(np.fft.fft2(obj2)) + + # Not sure if this is more or less stable than just the correlation + # maximum - requires some testing + np_cor = np.fft.ifft2(cor_fft / np.abs(cor_fft)) + # test with numpy inputs + test_cor = analysis.calc_deconvolved_cross_correlation(obj1,obj2, im_slice=np.s_[:,:]) + assert np.allclose(test_cor, np_cor) + + # test with pytorch inputs + obj1_t = cmath.complex_to_torch(obj1) + obj2_t = cmath.complex_to_torch(obj2) + test_cor_t = analysis.calc_deconvolved_cross_correlation(obj1_t,obj2_t, im_slice=np.s_[:,:]) + + assert np.allclose(cmath.torch_to_complex(test_cor_t), np_cor) + + + +def test_calc_frc(): + + obj1 = np.random.rand(270,230) + 1j * np.random.rand(270,230) + obj2 = np.random.rand(270,230) + 1j * np.random.rand(270,230) + + basis = np.array([[0,2,0], + [3,0,0]]) + + nbins = 100 + snr = 2 + + cor_fft = ffts.fftshift(ffts.fft2(obj1[10:-10,20:-20])) * \ + ffts.fftshift(np.conj(ffts.fft2(obj2[10:-10,20:-20]))) + + F1 = np.abs(ffts.fftshift(ffts.fft2(obj1[10:-10,20:-20])))**2 + F2 = np.abs(ffts.fftshift(ffts.fft2(obj2[10:-10,20:-20])))**2 + + + di = np.linalg.norm(basis[:,0]) + dj = np.linalg.norm(basis[:,1]) + + i_freqs = ffts.fftshift(ffts.fftfreq(cor_fft.shape[0],d=di)) + j_freqs = ffts.fftshift(ffts.fftfreq(cor_fft.shape[1],d=dj)) + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + + numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft) + denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1) + denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2) + n_pix, bins = np.histogram(Rs,bins=nbins) + bins = bins[:-1] + + frc = np.abs(numerator / np.sqrt(denominator_F1*denominator_F2)) + # This moves from combined-image SNR to single-image SNR + snr /= 2 + + threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \ + (1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix)) + + test_bins, test_frc, test_threshold = analysis.calc_frc(obj1, obj2, + basis, im_slice=np.s_[10:-10,20:-20], nbins=100, snr=2) + + assert np.allclose(bins, test_bins) + assert np.allclose(frc, test_frc) + assert np.allclose(threshold, test_threshold) + + # try again with complex + obj1_torch = cmath.complex_to_torch(obj1) + obj2_torch = cmath.complex_to_torch(obj2) + basis_torch = t.tensor(basis) + + test_bins_t, test_frc_t, test_threshold_t = analysis.calc_frc(obj1_torch, + obj2_torch, + basis_torch, + im_slice=np.s_[10:-10,20:-20], nbins=100, snr=2) + + assert np.allclose(bins, test_bins_t.numpy()) + assert np.allclose(frc, test_frc_t.numpy()) + assert np.allclose(threshold, test_threshold_t.numpy())