Maybe last time it didn't actually commit?

CDTools/tools/analysis/analysis.py

@@ -12,14 +12,14 @@ import numpy as np
 from CDTools.tools import cmath
 from CDTools.tools import image_processing as ip
 from scipy import fftpack
+from scipy import linalg as sla
 
-__all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions',
+__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
            'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
-           'calc_frc']
+           'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction']
 
 
-from matplotlib import pyplot as plt
-def orthogonalize_probes(probes):
+def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
     """Orthogonalizes a set of incoherently mixing probes
 
     The strategy is to define a reduced orthogonal basis that spans
@@ -27,55 +27,104 @@ def orthogonalize_probes(probes):
     defined by the probes in that basis. After diagonalization, the
     eigenvectors can be recast into the original basis and returned
 
+    By default, it assumes that the set of probes are defined just as
+    standard incoherently mixing probe modes, and orthogonalizes them.
+    However, if a density matrix is explicitly given, it will instead
+    consider the problem of extracting the eigenbasis of the matrix
+    probes * denstity_matrix * probes^dagger, where probes is the
+    column matrix of the given probe functions. This latter problem arises
+    in the generalization of the probe mixing model, and reduces to the
+    simpler case when the density matrix is equal to the identity matrix
+    
+    If the parameter "keep_transform" is set, the function will additionally
+    return the matrix A such that A * ortho_probes^dagger = probes^dagger
+
+    If the parameter "normalize" is False (as is the default), the variation
+    in intensities in the probe modes will be kept in the probe modes, as is
+    natural for a purely incoherent model. If it is set to "True", the
+    returned probe modes will all be normalized instead.
+
     Parameters
     ----------
     probes : array
         An l x n x m complex array representing  a stack of probes
+    density_matrix : np.array
+        An optional l x l density matrix further elaborating on the state
+    keep_transform : bool
+        Default False, whether to return the map from probes to ortho_probes
+    normalize : bool
+        Default False, whether to normalize the probe modes
 
     Returns
     -------
     ortho_probes: array
         An l x n x m complex array representing a stack of probes
     """
-
+    
     try:
         probes = cmath.torch_to_complex(probes.detach().cpu())
         send_to_torch = True
     except:
         send_to_torch = False
 
-    bases = []
-    coefficients = np.zeros((probes.shape[0],probes.shape[0]), dtype=np.complex64)
-    for i, probe in enumerate(probes):
-        ortho_probe = np.copy(probe)
-        for j, basis in enumerate(bases):
-            coefficients[j,i] = np.sum(basis.conj()*ortho_probe)
-            ortho_probe -= basis * coefficients[j,i]
+    from matplotlib import pyplot as plt
+    # We can do the orthogonalization with an SVD, so first we have to
+    # reshape the final two dimensions (the image shape) into a single
+    # vectorized dimension. This matrix is probes^dagger, hence the
+    # conjugation
+    probes_mat = probes.reshape(probes.shape[0],
+                                probes.shape[1]*probes.shape[2])
 
+    if density_matrix is None:
+        density_matrix = np.eye(probes.shape[0])
+    
+        
+    # next we want to extract the eigendecomposition of the density matrix
+    # itself
+    w,v = sla.eigh(density_matrix)
+    w,v = np.linalg.eigh(density_matrix)
+    w = w[::-1]
+    v = v[:,::-1]
 
-        coefficients[i,i] = np.sqrt(np.sum(np.abs(ortho_probe)**2))
-        bases.append(ortho_probe / coefficients[i,i])
+    # We do this just to avoid total failure when the density
+    # matrix is not positive definite.
+    w = np.abs(w)
+    # Note: this will fail if any w are less than zero
+    # Probably should figure out a good way to deal with that
+    B_dagger = np.dot(np.diag(np.sqrt(w)), v.conj().transpose())
 
+    #u,s,vh = np.linalg.svd(np.dot(B_dagger,probes_mat), full_matrices=False)
+    u,s,vh = sla.svd(np.dot(B_dagger,probes_mat), full_matrices=False)
+    
 
-    density_mat = coefficients.dot(np.conj(coefficients).transpose())
-    eigvals, eigvecs = np.linalg.eigh(density_mat)
+    if normalize:
+        ortho_probes = vh.reshape(probes.shape[0],
+                                  probes.shape[1],
+                                  probes.shape[2])
 
-    ortho_probes = []
-    for i in range(len(eigvals)):
-        coefficients = np.sqrt(eigvals[i]) * eigvecs[:,i]
-        probe = np.zeros(bases[0].shape, dtype=np.complex64)
-        for coefficient, basis in zip(coefficients, bases):
-            probe += basis * coefficient
-        ortho_probes.append(probe)
-
-
-    if send_to_torch:
-        return cmath.complex_to_torch(np.stack(ortho_probes[::-1]))
+        B_dagger_inv = np.linalg.pinv(B_dagger)
+        A = np.dot(B_dagger_inv,np.dot(u,np.diag(s)))
+        #A_dagger = np.dot(np.linalg.pinv(np.diag(s)),
+        #                  np.dot(np.transpose(u).conj(),B_dagger))
     else:
-        return np.stack(ortho_probes[::-1])
-
+        ortho_probes = np.dot(np.diag(s),vh).reshape(probes.shape[0],
+                                                     probes.shape[1],
+                                                     probes.shape[2])
+        B_dagger_inv = np.linalg.pinv(B_dagger)
+        A = np.dot(B_dagger_inv,u)
+        #A_dagger = np.dot(np.transpose(u).conj(),B_dagger)
+    
+    if send_to_torch:
+        ortho_probes = cmath.complex_to_torch(np.stack(ortho_probes))
+        A = cmath.complex_to_torch(A)
+        #A_dagger = cmath.complex_to_torch(A_dagger)
 
+    if keep_transform:
+        return ortho_probes, A#_dagger
+    else:
+        return ortho_probes 
 
+    
 def standardize(probe, obj, obj_slice=None, correct_ramp=False):
     """Standardizes a probe and object to prepare them for comparison
 
@@ -512,3 +561,66 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
         threshold = t.tensor(threshold)
 
     return bins[:-1], frc, threshold
+
+
+def calc_vn_entropy(matrix):
+    """Calculates the Von Neumann entropy of a density matrix
+    
+    Will either accept a single matrix, or a stack of matrices. Matrices
+    are assumed to be Hermetian and positive definite, to be well-formed
+    density matrices
+    
+    Parameters
+    ----------
+    matrix : np.array
+        The nxn matrix or lxnxn stack of matrices to calculate the entropy of
+
+    
+    Returns
+    -------
+    entropy: float or np.array
+        The entropy or entropies of the arrays
+    """
+
+    if len(matrix.shape) == 3:
+        # Get the eigenvalues
+        eigs = [np.linalg.eigh(mat)[0] for mat in matrix]
+        # Normalize them to match standard density matrix form
+        eigs = [eig / np.sum(eig) for eig in eigs]
+        # And calculate the VN entropy!
+        entropies = [-np.sum(eig*np.log(eig)) for eig in eigs]
+        return np.array(entropies)
+    else:
+        eig = np.linalg.eigh(matrix)[0]
+        entropy = -np.sum(eig*np.log(eig))/np.sum(eig)
+        return -np.trace(np.dot(matrix,sla.logm(matrix))) 
+
+def calc_top_mode_fraction(matrix):
+    """Calculates the fraction of total power in the top mode of a density matrix
+    
+    Will either accept a single matrix, or a stack of matrices. Matrices
+    are assumed to be Hermetian and positive definite, to be well-formed
+    density matrices
+    
+    Parameters
+    ----------
+    matrix : np.array
+        The nxn matrix or lxnxn stack of matrices to work from
+
+    
+    Returns
+    -------
+    entropy: float or np.array
+        The fraction of power in the top mode of each matrix
+    """
+
+    if len(matrix.shape) == 3:
+        # Get the eigenvalues
+        eigs = [np.linalg.eigh(mat)[0] for mat in matrix]
+        # Normalize them to match standard density matrix form
+        fractions = [np.max(eig) / np.sum(eig) for eig in eigs]
+        return np.array(fractions)
+    else:
+        eig = np.linalg.eigh(matrix)[0]
+        fraction = np.max(eig) / np.sum(eig)
+        return fraction

CDTools/tools/cmath/cmath.py

@@ -13,7 +13,7 @@ import torch as t
 
 
 __all__ = ['complex_to_torch', 'torch_to_complex', 'cabssq', 'cabs', 'cconj',
-           'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi']
+           'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi', 'cexpi']
 
 
 #

CDTools/tools/interactions/interactions.py

@@ -140,7 +140,15 @@ def project_translations_to_sample(sample_basis, translations):
     describing the location of the probe's intersection with the sample
     plane, assuming the basis found in sample_basis is used. Second, an
     amount that the probe must be propagated along the z-axis to reach
-    that location on the sample plane.
+    the sample plane at the given location. This includes both the effect of
+    the tilted sample plane and any explicitly defined motion along the z axis
+    of the probe-forming optic as included in the input translations.
+
+
+    Note that because of the sign convention (that this function returns th
+    relative amount the probe needs to be propagated to reach any given
+    location), a positive motion along the z-axis of the probe forming optics
+    will lead to a negative propagation distance.
     
     The assumed geometry is incoming radiation with a wavevector parallel
     to the +z axis, [0,0,1].
@@ -177,6 +185,8 @@ def project_translations_to_sample(sample_basis, translations):
         dtype=surface_normal.dtype)
     
     # Here we're setting up a matrix-vector equation mat*answer=input
+    # At some point ger will need to be replaced by outer, but for now
+    # outer many places still don't have new enough versions of torch.
     mat = t.cat((I - t.ger(propagation_dir,propagation_dir),
                  surface_normal.unsqueeze(0)))
     
@@ -205,8 +215,9 @@ def project_translations_to_sample(sample_basis, translations):
         single_translation = True
 
     pixel_translations = t.mm(translations, sample_projection)
-    propagations = t.mm(translations, prop_projection)
-    
+    propagations = t.mm(translations, prop_projection) \
+        - t.mm(translations,propagation_dir[:,None])
+
     if single_translation:
         return pixel_translations[0], propagations[0]
     else:

CDTools/tools/measurements/measurements.py

@@ -17,7 +17,8 @@ from torch.nn.functional import avg_pool2d
 # intensity pattern on a detector
 #
 
-__all__ = ['intensity', 'incoherent_sum', 'quadratic_background']
+__all__ = ['intensity', 'incoherent_sum', 'density_matrix',
+           'quadratic_background']
 
 
 def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
@@ -66,6 +67,101 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
         return t.clamp(output + epsilon,0,saturation)
 
 
+def density_matrix(wavefields, density_matrix, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
+    """Returns the intensities associated with a given density matrix state
+
+    The essential idea is that the most general description of a light field
+    at the detector plane will consist of a density matrix state. Here, that
+    low rank density matrix state is encoded as a set of basis wavefields
+    and a density matrix in that basis.
+
+    For computational efficiency, the density matrix is coded in an unusual
+    format. The density matrix formally is a complex Hermetian matrix,
+    which also happens to be positive definite. Here, we store it as a
+    real-valued matrix, where the upper triangle corresponds to the real
+    part of the elements in the upper triangle, and the lower triangle
+    corresponds to the imaginary parts. The elements on the diagonal are
+    purely real, and are stored as they are.
+
+    As with other multi-mode measurement functions, the modes are stored in
+    the first index, and the index of the diffraction pattern in the stack
+    of diffraction patterns is the second index. The stack-direction index
+    can be omitted if only a single pattern needs to be simulated
+    
+    It is important to note that this method does not inforce the positive
+    definiteness of the density matrix, this it is possible for negative
+    values of intensity to appear if the underlying density matrix passed
+    to this method is not positive definite
+
+    Parameters
+    ----------
+    wavefields : torch.Tensor
+        An Lx(Jx)MxNx2 stack of complex wavefields
+    density_matrix : torch.Tensor
+        A (Jx)LxL stack of real-valued representations of density matrices, as per above
+    saturation : float
+        Optional, a maximum saturation value to clamp the resulting intensities to
+    oversampling : int
+        Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
+
+    Returns
+    -------
+    sim_patterns : torch.Tensor 
+        A real Lx(Jx)MxN array storing the incoherently summed intensities
+
+    """
+    
+    #if wavefields.dim() == 4:
+    #    wavefields.unsqueeze(1)
+    #    single_frame = True
+    #elif wavefields.dim() == 5:
+    #    single_frame=False
+        
+    output = t.zeros(wavefields.shape[1:-1],
+                     dtype=wavefields.dtype,
+                     device=wavefields.device)
+    
+    # flat is better than nested, but simple is better than complex...
+    for (i,j) in ((i,j) for i in range(density_matrix.shape[-2])
+                  for j in range(density_matrix.shape[-1])):
+        if i == j: # diagonal
+            output += density_matrix[...,i,j,None,None] \
+                * cmath.cabssq(wavefields[i])
+        if i < j: # upper triangle, real part 
+            output += 2 * density_matrix[...,i,j,None,None] \
+                * (wavefields[i,...,0] * wavefields[j,...,0]
+                   +  wavefields[i,...,1] * wavefields[j,...,1])
+        if i > j: # lower triangle, imaginary part
+            # We pull the i,jth element from the density matrix,
+            # but this correponds to wavefield j and wavefield i,
+            # unlike above where it was wavefield i and j (swapped order).
+            # We also get the one negative sign because this is the imaginary
+            # part
+            output += 2 * density_matrix[...,i,j,None,None] \
+                * (wavefields[j,...,0] * wavefields[i,...,1]
+                   -  wavefields[j,...,1] * wavefields[i,...,0])
+        
+    # Now we apply oversampling
+    if oversampling != 1:
+        if wavefields.dim() == 4:
+            output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
+        else:
+            output = avg_pool2d(output, oversampling)
+            
+    # Then we grab the detector slice
+    if detector_slice is not None:
+        if wavefields.dim() == 4:
+            output = output[detector_slice]
+        else:
+            output = output[(np.s_[:],) + detector_slice]
+
+    if saturation is None:
+        return t.clamp(output,min=0) + epsilon
+    else:
+        return t.clamp(output + epsilon,0,saturation)
+    
+            
+            
 def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
     """Returns the incoherent sum of the intensities of the wavefields
     
@@ -94,7 +190,6 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
     sim_patterns : torch.Tensor 
         A real JXMxN array storing the incoherently summed intensities
     """
-    # This syntax just adds an axis to the slice to preserve the J direction
 
     output = t.sum(cmath.cabssq(wavefields),dim=0) 
 
@@ -118,7 +213,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
         return t.clamp(output + epsilon,0,saturation)
 
 
-def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1):
+def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1):
     """Returns the intensity of a wavefield plus a background
     
     The intensity is calculated via the given measurment function 
@@ -148,10 +243,10 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement
     """
     
     if detector_slice is None:
-        output = measurement(wavefield, epsilon=epsilon,
+        output = measurement(wavefield, *args, epsilon=epsilon,
                              oversampling=oversampling) + background**2
     else:
-        output = measurement(wavefield, detector_slice,
+        output = measurement(wavefield, *args, detector_slice=detector_slice,
                              epsilon=epsilon, oversampling=oversampling) \
                              + background**2
 

CDTools/tools/plotting/plotting.py

@@ -17,7 +17,7 @@ from matplotlib.colors import hsv_to_rgb
 
 __all__ = ['colorize', 'plot_amplitude', 'plot_phase',
            'plot_colorized', 'plot_translations', 'get_units_factor',
-           'plot_nanomap']
+           'plot_nanomap', 'plot_real', 'plot_imag']
 
 
 def colorize(z):
@@ -82,6 +82,140 @@ def get_units_factor(units):
     return factor
 
 
+def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
+    """Plots the real part of a complex array with dimensions NxM
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    if fig is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, **kwargs)
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+
+    if isinstance(im, t.Tensor):
+        real = im[...,0].detach().cpu().numpy()
+    else:
+        real = np.real(im)
+
+    #Plot in a basis if it exists, otherwise dont
+    if basis is not None:
+        if isinstance(basis,t.Tensor):
+            basis = basis.detach().cpu().numpy()
+        # This fails if the 
+        basis_norm = np.linalg.norm(basis, axis = 0)
+        basis_norm = basis_norm * get_units_factor(units)
+
+        extent = [0, real.shape[-1]*basis_norm[1], 0, real.shape[-2]*basis_norm[0]]
+    else:
+        extent=None
+
+    plt.imshow(real, cmap = cmap, extent = extent)
+    cbar = plt.colorbar()
+    cbar.set_label('Real Part (a.u.)')
+
+    if basis is not None:
+        plt.xlabel('X (' + units + ')')
+        plt.ylabel('Y (' + units + ')')
+    else:
+        plt.xlabel('j (pixels)')
+        plt.ylabel('i (pixels)')
+
+    return fig
+
+
+def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
+    """Plots the imaginary part of a complex array with dimensions NxM
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    if fig is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, **kwargs)
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+
+    if isinstance(im, t.Tensor):
+        imag = im[...,1].detach().cpu().numpy()
+    else:
+        imag = np.imag(im)
+
+    #Plot in a basis if it exists, otherwise dont
+    if basis is not None:
+        if isinstance(basis,t.Tensor):
+            basis = basis.detach().cpu().numpy()
+        # This fails if the 
+        basis_norm = np.linalg.norm(basis, axis = 0)
+        basis_norm = basis_norm * get_units_factor(units)
+
+        extent = [0, imag.shape[-1]*basis_norm[1], 0, imag.shape[-2]*basis_norm[0]]
+    else:
+        extent=None
+
+    plt.imshow(imag, cmap = cmap, extent = extent)
+    cbar = plt.colorbar()
+    cbar.set_label('Imaginary Part (a.u.)')
+
+    if basis is not None:
+        plt.xlabel('X (' + units + ')')
+        plt.ylabel('Y (' + units + ')')
+    else:
+        plt.xlabel('j (pixels)')
+        plt.ylabel('i (pixels)')
+
+    return fig
+
+
 def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
     """Plots the amplitude of a complex array with dimensions NxM
 

CDTools/tools/propagators/propagators.py

@@ -205,7 +205,7 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     
     intensity_map = t.Tensor(intensity_map).to(*args, **kwargs)
 
-
+    
     # This accounts for the implied phase ramp along the exit wave direction
     # In other words, it prevents the diffraction pattern from sliding off the
     # detector when the sample is tilted but represented by an object with