Add analysis tools for comparing incoherent wavefields

CDTools/tools/analysis/analysis.py

@@ -15,7 +15,8 @@ from scipy import special
 
 __all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
            'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
-           'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction']
+           'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction',
+           'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error']
 
 
 def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
@@ -581,6 +582,7 @@ def calc_vn_entropy(matrix):
         entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig)
         return entropy
 
+    
 def calc_top_mode_fraction(matrix):
     """Calculates the fraction of total power in the top mode of a density matrix
     
@@ -610,3 +612,213 @@ def calc_top_mode_fraction(matrix):
         eig = np.linalg.eigh(matrix)[0]
         fraction = np.max(eig) / np.sum(eig)
         return fraction
+    
+
+def calc_rms_error(field_1, field_2, align_phases=True, normalize=False,
+                   dims=2):
+    """Calculates the root-mean-squared error between two complex wavefields
+
+    The formal definition of this function is:
+
+    output = norm * sqrt(mean(abs(field_1 - gamma * field_2)**2))
+    
+    Where norm is an optional normalization factor, and gamma is an
+    optional phase factor which is appropriate when the wavefields suffer
+    from a global phase degeneracy as is often the case in diffractive
+    imaging.
+
+    The normalization is defined as the square root of the total intensity
+    contained in field_1, which is appropriate when field_1 represents a
+    known ground truth:
+
+    norm = sqrt(mean(abs(field_1)**2))
+    
+    The phase offset is an analytic expression for the phase offset which
+    will minimize the RMS error between the two wavefields:
+
+    gamma = exp(1j * angle(sum(field_1 * conj(field_2))))
+    
+    This implementation is stable even in cases where field_1 and field_2
+    are completely orthogonal.
+
+    In the definitions above, the field_n are n-dimensional wavefields. The
+    dimensionality of the wavefields can be altered via the dims argument,
+    but the default is 2 for a 2D wavefield.
+    
+    Parameters
+    ----------
+    field_1 : array
+        The first complex-valued field
+    field_2 : array
+        The second complex-valued field
+    align_phases : bool
+        Default is True, whether to account for a global phase offset
+    normalize : bool
+        Default is False, whether to normalize to the intensity of field_1
+    dims : (int or tuple of python:ints)
+        Default is 2, the number of final dimensions to reduce over.
+
+    
+    Returns
+    -------
+    rms_error : float or t.Tensor
+        The RMS error, or tensor of RMS errors, depending on the dim argument
+
+    """
+
+    sumdims = tuple(d - dims for d in range(dims))
+        
+    if align_phases:
+        # Keepdim allows us to broadcast the result correctly when we
+        # multiply by the fields
+        gamma = t.exp(1j * t.angle(t.sum(field_1 * t.conj(field_2), dim=sumdims,
+                                         keepdim=True)))
+    else:
+        gamma = 1
+
+    if normalize:
+        norm = 1 / t.mean(t.abs(field_1)**2, dim=sumdims)
+    else:
+        norm = 1
+
+    difference = field_1 - gamma * field_2
+    
+    return t.sqrt(norm * t.mean(t.abs(difference)**2, dim=sumdims))
+
+
+def calc_fidelity(fields_1, fields_2, dims=2):
+    """Calculates the fidelity between two density matrices
+
+    The fidelity is a comparison metric between two density matrices
+    (i.e. mutual coherence functions) that extends the idea of the
+    overlap to incoherent light. As a reminder, the overlap between two
+    fields is:
+
+    overlap = abs(sum(field_1 * field_2))**2
+    
+    Whereas the fidelity is defined as:
+    
+    fidelity = trace(sqrt(sqrt(dm_1) <dot> dm_2 <dot> sqrt(dm_1)))**2
+
+    where dm_n refers to the density matrix encoded by fields_n such
+    that dm_n = fields_n <dot> fields_<n>.conjtranspose(), sqrt
+    refers to the matrix square root, and <dot> is the matrix product.
+    
+    This is not a practical implementation, however, as it is not feasible
+    to explicitly construct the matrices dm_1 and dm_2 in memory. Therefore,
+    we take advantage of the alternate definition based directly on the
+    fields_<n> parameter:
+
+    fidelity = sum(svdvals(fields_1 <dot> fields_2.conjtranspose()))**2
+    
+    In the definitions above, the fields_n are regarded as collections of
+    wavefields, where each wavefield is by default 2-dimensional. The
+    dimensionality of the wavefields can be altered via the dims argument,
+    but the fields_n arguments must always have at least one more dimension
+    than the dims argument. Any additional dimensions are treated as batch
+    dimensions.
+    
+    Parameters
+    ----------
+    fields_1 : array
+        The first set of complex-valued field modes
+    fields_2 : array
+        The second set of complex-valued field modes
+    dims : int
+        Default is 2, the number of final dimensions to reduce over.
+
+    
+    Returns
+    -------
+    fidelity : float or t.Tensor
+        The fidelity, or tensor of fidelities, depending on the dim argument
+
+    """
+
+    fields_1 = t.as_tensor(fields_1)
+    fields_2 = t.as_tensor(fields_2)
+    
+    mult = fields_1.unsqueeze(-dims-2) * fields_2.unsqueeze(-dims-1).conj()
+    sumdims = tuple(d - dims for d in range(dims))
+    mat = t.sum(mult,dim=sumdims)
+    
+    svdvals = t.linalg.svdvals(mat)
+    return t.sum(svdvals, dim=-1)**2
+
+
+def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2):
+    """Calculates a generalization of the root-mean-squared error between two complex wavefields
+
+    This function calculates an generalization of the RMS error which uses the
+    concept of fidelity to extend it to capture the error between
+    incoherent wavefields, defined as a mode decomposition. The extension has
+    several nice properties, in particular:
+
+    1) For coherent wavefields, it precisely matches the RMS error including
+       a correction for the global phase degeneracy (align_phases=True)
+    2) All mode decompositions of either field that correspond to the same
+       density matrix / mutual coherence function will produce the same 
+       output
+    3) The error will only be zero when comparing mode decompositions that
+       correspond to the same density matrix.
+    4) Due to (2), one need not worry about the ordering of the modes,
+       properly orthogonalizing the modes, and it is even possible to
+       compare mode decompositions with different numbers of modes.    
+    
+    The formal definition of this function is:
+
+    output = norm * sqrt(mean(abs(fields_1)**2)
+                         + mean(abs(fields_2)**2)
+                         - 2 * sqrt(fidelity(fields_1,fields_2)))
+    
+    Where norm is an optional normalization factor, and the fidelity is
+    defined based on the mean, rather than the sum, to match the convention
+    for the root *mean* squared error.
+
+    The normalization is defined as the square root of the total intensity
+    contained in fields_1, which is appropriate when fields_1 represents a
+    known ground truth:
+
+    norm = sqrt(mean(abs(fields_1)**2))
+
+    In the definitions above, the fields_n are regarded as collections of
+    wavefields, where each wavefield is by default 2-dimensional. The
+    dimensionality of the wavefields can be altered via the dims argument,
+    but the fields_n arguments must always have at least one more dimension
+    than the dims argument. Any additional dimensions are treated as batch
+    dimensions.
+    
+    Parameters
+    ----------
+    fields_1 : array
+        The first set of complex-valued field modes
+    fields_2 : array
+        The second set of complex-valued field modes
+    normalize : bool
+        Default is False, whether to normalize to the intensity of fields_1
+    dims : (int or tuple of python:ints)
+        Default is 2, the number of final dimensions to reduce over.
+
+    Returns
+    -------
+    rms_error : float or t.Tensor
+        The generalized RMS error, or tensor of generalized RMS errors, depending on the dim argument
+
+    """
+    fields_1 = t.as_tensor(fields_1)
+    fields_2 = t.as_tensor(fields_2)
+
+    npix = t.prod(t.as_tensor(fields_1.shape[-dims:],dtype=t.int32))
+    
+    sumdims = tuple(d - dims - 1 for d in range(dims+1))
+    fields_1_intensity = t.sum(t.abs(fields_1)**2,dim=sumdims) / npix
+    fields_2_intensity = t.sum(t.abs(fields_2)**2,dim=sumdims) / npix
+    fidelity = calc_fidelity(fields_1, fields_2, dims=dims) / npix**2
+
+    result = fields_1_intensity + fields_2_intensity - 2 * t.sqrt(fidelity)
+    
+    if normalize:
+        result /= fields_1_intensity
+
+    return t.sqrt(result)
+    

CDTools/tools/propagators/propagators.py

@@ -382,6 +382,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
     # Bandlimiting is not implemented in the generalized function, because it
     # has a less clear meaning in that setting, so we apply it here instead
     if bandlimit is not None:
+        # No need to multiply by 2pi
         ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0])
         kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1])
         Ki, Kj = t.meshgrid(ki,kj)

tests/tools/test_analysis.py

@@ -1,5 +1,6 @@
 import numpy as np
 from scipy import fftpack as ffts
+from scipy import linalg as la
 import torch as t
 from itertools import combinations
 
@@ -239,7 +240,6 @@ def test_calc_deconvolved_cross_correlation():
     assert np.allclose(test_cor_t.numpy(), np_cor)
 
     
-
 def test_calc_frc():
 
     obj1 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
@@ -300,3 +300,141 @@ def test_calc_frc():
     assert np.allclose(bins, test_bins_t.numpy())
     assert np.allclose(frc, test_frc_t.numpy())
     assert np.allclose(threshold, test_threshold_t.numpy())
+
+    
+def test_calc_rms_error():
+    field_1 = t.rand(14,19, dtype=t.complex64)
+    field_2 = t.rand(14,19, dtype=t.complex64)
+
+    # Check that the calculation is insensitive to phase
+    assert t.allclose(analysis.calc_rms_error(field_1, field_2),
+                      analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2))
+
+    # And that it is sensitive to phase if we turn off the 
+    assert not t.allclose(
+        analysis.calc_rms_error(field_1, field_2, align_phases=False),
+        analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2,
+                                align_phases=False))
+
+    # Check that the result is positive
+    assert analysis.calc_rms_error(field_1, field_2) > 0
+
+    # And that it is a smaller number with align_phases on
+    assert (analysis.calc_rms_error(field_1, field_2) <=
+            analysis.calc_rms_error(field_1, field_2, align_phases=False))
+
+    # Now we check against an explicit implementation:
+    gamma = field_1 * t.conj(field_2)
+    gamma /= t.abs(gamma)
+
+    # This is an alternate way of doing the calculation. Actually, would this
+    # be a better implementation anyway? Probably no difference tbh.
+    rms_error_nophase = t.sqrt((t.mean(t.abs(field_1)**2) +
+                                t.mean(t.abs(field_2)**2) -
+                                2 * t.abs(t.mean(field_1 * t.conj(field_2)))))
+    assert t.allclose(rms_error_nophase,
+                      analysis.calc_rms_error(field_1, field_2))
+
+    rms_error_phase = t.sqrt((t.mean(t.abs(field_1)**2) +
+                              t.mean(t.abs(field_2)**2) -
+                              2 * t.real(t.mean(field_1 * t.conj(field_2)))))
+
+    assert t.allclose(rms_error_phase,
+                      analysis.calc_rms_error(field_1, field_2,
+                                              align_phases=False))
+
+    # Now let's test that it works along a dimension:
+    
+    field_1 = t.rand(3,14,19, dtype=t.complex64)
+    field_2 = t.rand(3,14,19, dtype=t.complex64)
+    result = analysis.calc_rms_error(field_1, field_2, normalize=True)
+    assert (result.shape == t.Size([3]))
+    
+    for i in range(3):
+        assert t.allclose(analysis.calc_rms_error(field_1[i],
+                                                  field_2[i],
+                                                  normalize=True),
+                          result[i])
+    
+
+def test_calc_fidelity():
+
+    fields_1 = t.rand(2,30,17, dtype=t.complex128)
+    fields_2 = t.rand(3,30,17, dtype=t.complex128)
+
+    dm_1 = t.reshape(fields_1, (2,-1))
+    dm_1 = t.tensordot(dm_1.transpose(0,1), dm_1.conj(), dims=1).numpy()
+    dm_2 = t.reshape(fields_2, (3,-1))
+    dm_2 = t.tensordot(dm_2.transpose(0,1), dm_2.conj(), dims=1).numpy()
+    
+    inner_mat = la.sqrtm(np.dot(np.dot(la.sqrtm(dm_1),dm_2),la.sqrtm(dm_1)))
+    fidelity = t.as_tensor(np.abs(np.trace(inner_mat))**2)
+    
+    assert t.isclose(fidelity, analysis.calc_fidelity(fields_1, fields_2))
+
+    fields_1 = t.rand(1,30,17, dtype=t.complex128)
+    fields_2 = t.rand(1,30,17, dtype=t.complex128)
+
+    assert t.isclose(t.abs(t.sum(fields_1*fields_2.conj()))**2,
+                     analysis.calc_fidelity(fields_1, fields_2))
+
+    # Checking that it works with extra dimensions
+    fields_1 = t.rand(3,3,30,17, dtype=t.complex128)
+    fields_2 = t.rand(3,1,30,17, dtype=t.complex128)
+    field_3 = t.rand(1,30,17, dtype=t.complex128)
+
+    fidelities = analysis.calc_fidelity(fields_1, fields_2)
+    fidelities_2 = analysis.calc_fidelity(fields_1, field_3)
+    for i in range(3):
+        assert t.isclose(analysis.calc_fidelity(fields_1[i], fields_2[i]),
+                         fidelities[i])
+        assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3),
+                         fidelities_2[i])
+
+    # Check that the diensionality argument works
+    fields_1 = t.rand(3,2,12, dtype=t.complex128)
+    fields_2 = t.rand(3,2,12, dtype=t.complex128)
+
+    assert (analysis.calc_fidelity(fields_1, fields_2, dims=1).shape
+            == t.Size([3]))
+    
+    fields_1 = t.rand(3,2,12,4,5, dtype=t.complex128)
+    fields_2 = t.rand(3,2,12,4,5, dtype=t.complex128)
+
+    assert (analysis.calc_fidelity(fields_1, fields_2, dims=3).shape
+            == t.Size([3]))
+
+def test_calc_generalized_rms_error():
+
+    # Test that it matches the rms error for coherent fields
+    
+    fields_1 = t.rand(1,30,17, dtype=t.complex128)
+    fields_2 = t.rand(1,30,17, dtype=t.complex128)
+    
+    assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
+                     analysis.calc_rms_error(fields_1[0], fields_2[0],
+                                             align_phases=True))
+
+    # Test that it is independent of field order
+    fields_1 = t.rand(5,30,17, dtype=t.complex128)
+    fields_2 = t.rand(3,30,17, dtype=t.complex128)
+    fields_3 = fields_2.flip(0)
+    
+    assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
+                     analysis.calc_generalized_rms_error(fields_1, fields_3))
+
+    # Test with leading dimensions
+    fields_1 = t.rand(3,4,2,10,17, dtype=t.complex128)
+    fields_2 = t.rand(3,4,3,10,17, dtype=t.complex128)
+    
+    assert (analysis.calc_generalized_rms_error(fields_1, fields_2).shape
+            == t.Size([3,4]))
+
+    # And test with different number of dimensions dims
+    # Test that it is independent of field order
+    fields_1 = t.rand(3,6,17, dtype=t.complex128)
+    fields_2 = t.rand(3,1,17, dtype=t.complex128)
+    fields_3 = fields_2.flip(0)
+
+    assert (analysis.calc_generalized_rms_error(fields_1, fields_2, dims=1).shape == t.Size([3]))
+