Add and build tests for a Fourier Ring Correlation method

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')
 

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

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)):

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())