diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py index bef729b..dd3e8fb 100644 --- a/CDTools/tools/analysis/analysis.py +++ b/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) dm_2 sqrt(dm_1)))**2 + + where dm_n refers to the density matrix encoded by fields_n such + that dm_n = fields_n fields_.conjtranspose(), sqrt + refers to the matrix square root, and 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_ parameter: + + fidelity = sum(svdvals(fields_1 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) + diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py index fc124b7..bca164e 100644 --- a/CDTools/tools/propagators/propagators.py +++ b/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) diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index b27bca5..86597cd 100644 --- a/tests/tools/test_analysis.py +++ b/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])) +