cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-09 13:02:41 +02:00
Code Issues Packages Projects Releases Wiki Activity

rewrite the probe orthogonalization for fancy_ptycho, and introduce a new probe orthogonalization tool in analysis. Also, update some tracking plots in fancy_ptycho

Browse Source
This commit is contained in:
levita_a
2024-07-25 15:46:53 +02:00
parent a02a0466ae
commit 5141320c6f
4 changed files with 411 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/models/fancy_ptycho.py
+76 -80
View File
@@ -79,7 +79,7 @@ class FancyPtycho(CDIModel):
self.register_buffer('mask',
t.tensor(mask, dtype=t.bool))
probe_guess = t.tensor(probe_guess, dtype=t.complex4)
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
@@ -533,18 +533,7 @@ class FancyPtycho(CDIModel):
else:
return translations
def get_rhos(self):
# If this is the general unified mode model
if self.weights.dim() >= 2:
Ws = self.weights.detach().cpu().numpy()
rhos_out = np.matmul(np.swapaxes(Ws, 1, 2), Ws.conj())
return rhos_out
# This is the purely incoherent case
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def center_probes(self, iterations=4):
"""Centers the probes
@@ -557,95 +546,97 @@ class FancyPtycho(CDIModel):
self.probe.data = centered_probe.to(device=self.probe.data.device)
def tidy_probes(self, normalize=False, tidy_each_frame=False):
def tidy_probes(self):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
As a first step, we calculate the state of the light field across the
full experiment, using the weight matrices and basis probes. Then, we
use an SVD to update the basis probes so they form an eigenbasis of
the implied density matrix for the full experiment.
Next, the weight matrices for each shot are recalculated so that the
probes generated by weights * basis_probes for each shot are themselves
an eigenbasis for that individual shot's density matrix.
"""
# First we treat the purely incoherent case
# I don't love this pattern of using an if statement with a return
# to catch this case, but because it's so much simpler than the
# unified mode case I think it's appropriate
# First we treat the incoherent but stable case, where the weights are
# just one per-shot overall weight
if self.weights.dim() == 1:
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes(probe)
self.probe.data = t.as_tensor(
ortho_probes,
device=self.probe.device,
dtype=self.probe.dtype)
ortho_probes = analysis.orthogonalize_probes_t(self.probe.detach())
self.probe.data = ortho_probes
return
# This is for the unified mode case
# What follows is for the unified OPRP and incoherent multi-mode model,
# where each shot has it's own matrix of weights such that the probe
# state for each shot is self.weights @ self.probe
# We concatenate all the weight matrices
all_weights = t.cat(t.unbind(self.weights.detach().cpu(), dim=0), dim=0)
# We use that to calculate the density matrix of the full experiment,
# normalized by number of exposures
overall_rho = (t.mm(all_weights.transpose(0,1), all_weights.conj())
/ self.weights.shape[0])
# We concatenate all the weight matrices, to come up with a state
# corresponding to the summed light field across all the exposures.
# This state will have a large number of modes, but all built from
# the same small number of basis modes
all_weights = t.cat(t.unbind(self.weights.detach(), dim=0), dim=0)
# We generate the orthogonal probes based on this full-experiment
# density matrix. We also keep the transform matrix A
probe = self.probe.detach().cpu().numpy()
ortho_probes, A = analysis.orthogonalize_probes(
probe, density_matrix=overall_rho,
keep_transform=True, normalize=normalize)
# representation of the light field.
ortho_probes, reexpressed_weights = \
analysis.orthogonalize_probes_t(
self.probe.detach(),
weight_matrix=all_weights,
return_reexpressed_weights=True
)
# We apply A to the weight matrices to update them along with the
# probes
new_weights = t.matmul(
t.as_tensor(A).transpose(0,1),
self.weights.detach().cpu().transpose(-2,-1)).transpose(-2,-1)
self.probe.data = t.as_tensor(
ortho_probes, device=self.probe.device, dtype=self.probe.dtype)
# We just orthogonalized the incoherent sum of all the exposures
# across the full experiment, so the output probes are normalized so
# that their intensity matches the summed intensity across the full
# experiment. We divide their amplitudes by the square root of the
# number of shots so that we now have a set of probes corresponding
# to the mean shot
ortho_probes /= np.sqrt(self.weights.shape[0])
reexpressed_weights *= np.sqrt(self.weights.shape[0])
self.weights.data = new_weights.to(device=self.weights.device,
dtype=self.weights.dtype)
# We now replace the shot-to-shot weights with the versions that have
# been re-expressed in the new basis.
new_weights = t.stack(t.split(reexpressed_weights,
self.weights.shape[1]), dim=0)
# At this point, we now have a new set of basis probes, which are the
# eigenbasis for the full-experiment density matrix, and we have
# re-expressed all the shot-to-shot weight matrices in that basis.
# But, the shot-to-shot probes (self.weights <matmul> self.probe)
# are still exactly the same as they were before.
#
# Oftentimes, we also want the shot-to-shot weights to be re-expressed
# so that the shot-to-shot probes are the eigenbasis for each
# individual shot's density matrix. That's what we do below.
#
if tidy_each_frame:
# And we save it back to the model
self.probe.data = ortho_probes.to(
device=self.probe.device, dtype=self.probe.dtype)
self.weights.data = new_weights.to(
device=self.weights.device, dtype=self.weights.dtype)
# TODO: I need to check that this really works
dm_rank = self.weights.shape[-2]
for idx in range(self.weights.shape[0]):
weights = self.weights.data[idx].detach().cpu()
rho = t.mm(weights.transpose(0,1), weights.conj())
ortho_probes, A = analysis.orthogonalize_probes(
self.probe.detach().cpu(), density_matrix=rho,
normalize=False, keep_transform=True)
# NOTE: I used to have this part as an option, with "tidy_each_frame",
# because it took such a long time. Now that I've rewritten it properly,
# it's quite fast and so I removed the kwarg because there's really
# no situation where you woudn't want to do this.
new_weights = t.linalg.pinv(A)[:dm_rank, :].conj()
self.weights.data[idx] = new_weights.to(
device=self.weights.device,
dtype=self.weights.dtype)
# Now, we seek to edit the shot-to-shot weight matrices such that
# self.weights[i] @ self.probes will be properly orthogonalized for
# all i.
# All we need to know about the probes is that they are orthogonalized
# and the intensity within each probe mode
probe_sqrt_intensities = t.linalg.norm(self.probe.data, dim=(-2,-1))
# This does a super fast batched computation
U, S, Vh = t.linalg.svd(self.weights.data * probe_sqrt_intensities,
full_matrices=False)
# We discard the U matrix and re-multiply S & Vh
self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities)
def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs):
def get_probes(idx):
basis_prs = self.probe * self.probe_support[..., :, :]
prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs,
axis=-4)
ortho_probes = analysis.orthogonalize_probes(prs)
ortho_probes = analysis.orthogonalize_probes_t(prs)
if mode.lower() == 'amplitude':
return np.abs(ortho_probes.detach().cpu().numpy())
@@ -737,15 +728,20 @@ class FancyPtycho(CDIModel):
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Average Density Matrix Amplitudes',
('Average Weight Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.get_rhos()), axis=0),
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% Power in Top Mode (only accurate after tidy_probes)',
('% of Power in Top Mode',
lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
analysis.calc_top_mode_fraction(self.get_rhos()),
100 * t.stack([
analysis.calc_mode_power_fractions(
self.probe.data,
weight_matrix=self.weights.data[i])[0]
for i in range(self.weights.shape[0])
], dim=0),
fig=fig,
units=self.units),
lambda self: len(self.weights.shape) >= 2),
src/cdtools/tools/analysis/analysis.py
+183 -36
View File
@@ -13,69 +13,164 @@ from scipy import linalg as sla
from scipy import special
from scipy import optimize as opt
__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
__all__ = ['product_svd', 'orthogonalize_probes_t',
'orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction',
'calc_mode_power_fractions',
'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error',
'remove_phase_ramp', 'remove_amplitude_exponent',
'standardize_reconstruction_set']
def product_svd(A, B):
""" Computes the SVD of A @ B
This function uses a method which uses a QR decomposition of
A and B to calculate the final reduced SVD, without explicitly
calculating the full matrix. The output is defined such that
A B = U S Vh, and as a reduced SVD
Parameters
----------
A : array
An nxr matrix
B : array
An rxm matrix
Returns
-------
U : array
An nxr matrix of left singular vectors
S : array
An length-r array, t.diag(S) is the diagonal matrix of singular values
Vh : array
And rxm matrix of the conjugate-transposed right singular vectors
"""
# Handle the case of numpy input
return_np = False
if isinstance(A, np.ndarray):
A = t.as_tensor(A)
return_np = True
if isinstance(B, np.ndarray):
B = t.as_tensor(B)
return_np = True
# We take a QR decomposition of the two matrices
Qa, Ra = t.linalg.qr(A)
Qb, Rb = t.linalg.qr(B.conj().transpose(0,1))
# And now we take the SVD of the product of the two R matrices
U, S, Vh = t.linalg.svd(t.matmul(Ra, Rb.conj().transpose(0,1)),
full_matrices=False)
# And build back the final SVD of the product matrix
U_final = t.matmul(Qa, U)
Vh_final = t.matmul(Vh, Qb.conj().transpose(0,1))
if return_np:
U_final = U_final.numpy()
S = S.numpy()
Vh_final = Vh_final.numpy()
return U_final, S, Vh_final
def orthogonalize_probes_t(
probes,
density_matrix=None,
keep_transform=False,
normalize=False,
weight_matrix=None,
n_probe_dims=2,
return_reexpressed_weights=False,
):
""" Orthogonalizes a set of incoherently mixing probes
TODO: actually make this, and replace ortho_probes with a fully
pytorch-based function
## TODO fully replace orthogonalize_probes with orthogonalize_probes_t
Any set of probe modes defines a density matrix (a.k.a mutual coherence
function) that is the ultimate description of the state of the light
field. This function takes any set of probe modes - not necessarily
orthogonalized - and returns an orthogonalized set of probe modes.
Formally, it returns the eigenbasis of the density matrix, ordered
from largest to smallest eigenvalue.
This function takes any set of probe modes for mixed-mode ptychography,
which are considered to define a mutual coherence function, and returns
an orthogonalized set probe modes which refer to the same mutual coherence
function. The orthogonalized modes are extracted via a singular value
decomposition and are unique up to a global per-mode phase factor.
If normalize is set to True, then it will return the normalized
eigenbasis. Otherwise, it will return a scaled version of the eigenbasis,
so that the returned probes can be used directly for multi-mode
ptychography.
If a weight matrix is explicitly given, the function will instead
orthogonalize the light field defined by weight_matrix @ probes. It
accomplishes this via a method which avoids explicitly constructing this
potentially large matrix. This can be useful for Orthogonal Probe
Relaxation ptychography. In this case, one may have a large stacked
matrix of shot-to-shot weights but a small basis set of probes.
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
In addition to returning the orthogonalized probe modes, this function
also returns a re-expression of the original weight matrix in the basis
of the orthogonalized probe modes, such that:
reexpressed_weight_matrix @ orthogonalized_probes = weight_matrix @ probes.
This re-expressed weight matrix is guaranteed to have orthonormalized rows,
such that:
reexpressed_weight_matrix^\\dagger @ reexpressed_weight_matrix = I.
There is usually no reason to use the re-expressed weight matrix.
However, it can be useful in situations where the individual rows in the
weight matrix have a specific meaning, such as an exposure number, which
should be preserved.
If the parameter "keep_transform" is set, the function will additionally
return the matrix A such that A * ortho_probes^dagger = probes^dagger
TODO: is the above right, or are ortho_probes and probes flipped?
Warning! The shape of the output orthogonalized_probes may not be equal to
the shape of input probes, when a weight matrix is used. If the input
weight matrix has m < l rows, where l is the number of probe modes,
then the output orthogonalized probes will have length m, not length l.
Parameters
----------
probes : array
An l x (<n_probe_dims>) complex array representing a stack of probes
density_matrix : 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
An l x (<n_probe_pix>) array representing a stack of l probes
weight_matrix : array
Optional, an m x l weight matrix further elaborating on the state
n_probe_dims : int
Default 2, the number of trailing dimensions defining each probe
Default is 2, the number of trailing dimensions for each probe state
Returns
-------
ortho_probes: array
An l x (<n_probe_dims>) complex array representing a stack of probes
orthogonalized_probes : array
A min(m,l) x (<n_probe_dims>) array representing a stack of probes
reexpressed_weight_matrix : array
A the original weight matrix, re-expressed to work with the new probes
"""
pass
return_np = False
if isinstance(probes, np.ndarray):
probes = t.as_tensor(probes)
return_np = True
if weight_matrix is not None and isinstance(weight_matrix, np.ndarray):
weight_matrix = t.as_tensor(weight_matrix)
return_np = True
n_probe_pix = np.prod(np.array(probes.shape[-n_probe_dims:]))
probes_mat = probes.reshape(probes.shape[:-n_probe_dims] +
(n_probe_pix,))
if weight_matrix is None:
# We just calculate a straight up SVD of the probes
U, S, Vh = t.linalg.svd(probes_mat, full_matrices=False)
else:
U, S, Vh = product_svd(weight_matrix, probes_mat)
output_shape = (-1,) + tuple(n for n in probes.shape[1:])
orthogonalized_probes = (S[:,None] * Vh).reshape(output_shape)
reexpressed_weight_matrix = U
if return_np:
orthogonalized_probes = orthogonalized_probes.numpy()
reexpressed_weight_matrix = reexpressed_weight_matrix.numpy()
to_return = (orthogonalized_probes,)
if return_reexpressed_weights:
to_return += (reexpressed_weight_matrix,)
return to_return
def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
"""Orthogonalizes a set of incoherently mixing probes
@@ -668,7 +763,58 @@ def calc_vn_entropy(matrix):
entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig)
return entropy
def calc_mode_power_fractions(
probes,
weight_matrix=None,
n_probe_dims=2,
assume_preorthogonalized=False,
):
"""Calculates the fraction of total power in each orthogonalized mode
This code first orthogonalizes the probe modes, so the result of this
function are independent of the particular way that the multi-mode
breakdown is expressed.
Parameters
----------
probes : array
An l x (<n_probe_pix>) array representing a stack of l probes
weight_matrix : array
Optional, an m x l weight matrix further elaborating on the state
n_probe_dims : int
Default is 2, the number of trailing dimensions for each probe state
assume_preorthogonalized : bool
Default is False. If True, will not orthogonalize the probes
Returns
-------
power_fractions : array
The fraction of the total power in each mode
"""
if not assume_preorthogonalized:
ortho_probes, reexpressed_weights = \
orthogonalize_probes_t(
probes,
weight_matrix=weight_matrix,
n_probe_dims=n_probe_dims,
)
else:
weight_slice = np.s_[...,] + np.s_[None,] * n_probe_dims
if weight_matrix is None:
ortho_probes = probes
else:
ortho_probes = t.sum(weight_matrix[weight_slice] * probes,
axis=-(n_probe_dims + 1))
dims = [-d-1 for d in range(n_probe_dims)]
power = t.sum(t.abs(ortho_probes)**2, dim=dims)
power_fractions = power / t.sum(power)
return power_fractions
def calc_top_mode_fraction(matrix):
"""Calculates the fraction of total power in the top mode of a density matrix
@@ -687,6 +833,7 @@ def calc_top_mode_fraction(matrix):
entropy: float or np.array
The fraction of power in the top mode of each matrix
"""
# TODO depricate
if len(matrix.shape) == 3:
# Get the eigenvalues
src/cdtools/tools/plotting/plotting.py
+6 -4
View File
@@ -780,11 +780,13 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
im = im.reshape(-1,im.shape[-2],im.shape[-1])
im_idx = axes[1].image_idx
if event.key == 'up' or event.button == 'up' \
or event.key == 'left':
if (event.key == 'up'
or (hasattr(event, 'button') and event.button == 'up')
or event.key == 'left'):
im_idx = (im_idx - 1) % im.shape[0]
if event.key == 'down' or event.button == 'down' \
or event.key == 'right':
if (event.key == 'down'
or (hasattr(event, 'button') and event.button == 'down')
or event.key == 'right'):
im_idx = (im_idx + 1) % im.shape[0]
axes[1].image_idx=im_idx
tests/tools/test_analysis.py
+146
View File
@@ -1,11 +1,157 @@
import numpy as np
from scipy import linalg as la
from scipy.sparse import linalg as spla
import torch as t
from itertools import combinations
from cdtools.tools import analysis, initializers
def test_product_svd():
rank = 4
shape_A = (12, rank)
shape_B = (rank, 9)
A = np.random.rand(*shape_A) + 1j * np.random.rand(*shape_A)
B = np.random.rand(*shape_B) + 1j * np.random.rand(*shape_B)
AB = np.matmul(A,B)
U_1, S_1, Vh_1 = t.linalg.svd(t.as_tensor(AB), full_matrices=False)
U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A),t.as_tensor(B))
check_AB = U_2 @ t.diag_embed(S_2).to(dtype=Vh_2.dtype) @ Vh_2
# So, at a minimum, U S Vh = AB
assert np.allclose(AB, check_AB.numpy())
# SVD is only defined up to an arbitrary complex valued phase per
# singular vector, so all we can ask for in the comparison is that the
# magnitudes here are
assert np.allclose(S_1[:rank].numpy(), S_2.numpy())
prod_U = U_1[:,:rank].transpose(0,1).conj() @ U_2
prod_Vh = Vh_1[:rank,:] @ Vh_2.transpose(0,1).conj()
assert np.allclose(t.abs(prod_U).numpy(), np.eye(rank))
assert np.allclose(t.abs(prod_Vh).numpy(), np.eye(rank))
# Confirms that the phases are consistent between the two, I think
# it's redundant with the first check but I'm not sure
assert np.allclose(prod_Vh.numpy(), prod_U.numpy())
# test with numpy
U_3, S_3, Vh_3 = analysis.product_svd(A,B)
assert isinstance(U_3, np.ndarray)
assert isinstance(S_3, np.ndarray)
assert isinstance(Vh_3, np.ndarray)
assert np.allclose(S_1[:rank].numpy(), S_3)
prod_U = U_1[:,:rank].transpose(0,1).numpy().conj() @ U_3
prod_Vh = Vh_1[:rank,:].numpy() @ Vh_3.transpose().conj()
assert np.allclose(np.abs(prod_U), np.eye(rank))
assert np.allclose(np.abs(prod_Vh), np.eye(rank))
# Confirms that the phases are consistent between the two, I think
# it's redundant with the first check but I'm not sure
assert np.allclose(prod_Vh, prod_U)
def test_orthogonalize_probes_t():
op = analysis.orthogonalize_probes_t
probe_xs = np.arange(64) - 32
probe_ys = np.arange(76) - 38
probe_Ys, probe_Xs = np.meshgrid(probe_ys, probe_xs)
probe_Rs = np.sqrt(probe_Xs**2 + probe_Ys**2)
probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2 + 1j)),
3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)),
1*np.exp(-probe_Rs**2 / (2 * 15**2))])
weight_matrix_none = None
weight_matrix_single = np.random.randn(1,3) + 1j * np.random.randn(1,3)
weight_matrix_small = np.random.randn(2,3) + 1j * np.random.randn(2,3)
weight_matrix_medium = np.random.randn(3,3) + 1j * np.random.randn(3,3)
weight_matrix_large = np.random.randn(7,3) + 1j * np.random.randn(7,3)
weight_matrices = [
weight_matrix_none,
weight_matrix_single,
weight_matrix_small,
weight_matrix_medium,
weight_matrix_large
]
for weight_matrix in weight_matrices:
ortho_probes_np, rwm_np = op(probes, weight_matrix=weight_matrix)
assert isinstance(ortho_probes_np, np.ndarray)
assert isinstance(rwm_np, np.ndarray)
probes_t = t.as_tensor(probes)
wm_t = (t.as_tensor(weight_matrix) if weight_matrix is not None
else weight_matrix)
ortho_probes_t, rwm_t = op(probes_t, weight_matrix=wm_t)
assert t.is_tensor(ortho_probes_t)
assert t.is_tensor(rwm_t)
assert np.allclose(ortho_probes_np, ortho_probes_t.numpy())
assert np.allclose(rwm_np, rwm_t.numpy())
# Now we test a if wm @ ortho_probes is actually the original
# input
if weight_matrix is not None:
realized_probes = np.tensordot(weight_matrix, probes, axes=1)
else:
realized_probes = probes
calculated_probes = np.tensordot(rwm_np, ortho_probes_np, axes=1)
assert np.allclose(realized_probes, calculated_probes)
# Now we test if the orthogonalized probes are orthogonalized
reshaped_probes = ortho_probes_np.reshape(
(ortho_probes_np.shape[0],
ortho_probes_np.shape[1] * ortho_probes_np.shape[2]))
products = np.matmul(reshaped_probes,
reshaped_probes.conj().transpose())
if weight_matrix is not None:
output_nmodes = min(weight_matrix.shape[0], probes.shape[0])
else:
output_nmodes = probes.shape[0]
for i in range(output_nmodes):
for j in range(i+1, output_nmodes):
assert np.isclose(products[i,j], 0)
# And now we test if they multiply to the same density matrix as
# the original probes + weight matrix
reshaped_realized_probes = realized_probes.reshape(
(realized_probes.shape[0],
realized_probes.shape[1] * realized_probes.shape[2]))
dm_original = np.matmul(
reshaped_realized_probes.conj().transpose(),
reshaped_realized_probes
)
dm_output = np.matmul(
reshaped_probes.conj().transpose(),
reshaped_probes
)
assert np.allclose(dm_original, dm_output)
# And finally, we confirm that what we have are the eigenvectors/values
# of that density matrix
w, v = spla.eigsh(dm_original, k=output_nmodes)
assert np.allclose(w, np.diag(products))
cross_products = np.matmul(reshaped_probes, np.sqrt(w) * v)
# The abs accounts for the fact that the phase of the eigenvectors
# is undefined
assert np.allclose(np.abs(cross_products), np.abs(products))
def test_orthogonalize_probes():
# The test strategy should be to define a few non-orthogonal probes