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

Fully replace orthogonalize_probes and update most models

Browse Source
This commit is contained in:
levita_a
2024-07-25 17:35:48 +02:00
parent 943dde13e6
commit 35f140d348
8 changed files with 270 additions and 425 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/models/fancy_ptycho.py
+3 -3
View File
@@ -567,7 +567,7 @@ class FancyPtycho(CDIModel):
# just one per-shot overall weight
if self.weights.dim() == 1:
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes_t(self.probe.detach())
ortho_probes = analysis.orthogonalize_probes(self.probe.detach())
self.probe.data = ortho_probes
return
@@ -584,7 +584,7 @@ class FancyPtycho(CDIModel):
# We generate the orthogonal probes based on this full-experiment
# representation of the light field.
ortho_probes, reexpressed_weights = \
analysis.orthogonalize_probes_t(
analysis.orthogonalize_probes(
self.probe.detach(),
weight_matrix=all_weights,
return_reexpressed_weights=True
@@ -636,7 +636,7 @@ class FancyPtycho(CDIModel):
basis_prs = self.probe * self.probe_support[..., :, :]
prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs,
axis=-3)
ortho_probes = analysis.orthogonalize_probes_t(prs)
ortho_probes = analysis.orthogonalize_probes(prs)
if mode.lower() == 'amplitude':
return np.abs(ortho_probes.detach().cpu().numpy())
src/cdtools/models/fastccd_ptycho.py
+81 -55
View File
@@ -498,81 +498,96 @@ class FastCCDPtycho(CDIModel):
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1, normalize=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(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, 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)
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix.
# We generate the orthogonal probes based on this full-experiment
# representation of the light field.
ortho_probes, reexpressed_weights = \
analysis.orthogonalize_probes(
self.probe.detach(),
weight_matrix=all_weights,
return_reexpressed_weights=True
)
rhos = self.get_rhos()
overall_rho = np.mean(rhos, axis=0)
probe = self.probe.detach().cpu().numpy()
ortho_probes, A = analysis.orthogonalize_probes(
probe, density_matrix=overall_rho,
keep_transform=True, normalize=normalize)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans, np.matmul(rhos, Aconj))
# 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])
# 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)
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
# 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)
dm_rank = self.weights.shape[1]
# 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_Ws = []
for rho in new_rhos:
# These are returned from smallest to largest - we want to keep
# the largest ones
w, v = sla.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:, ::-1][:, :dm_rank]
# For situations where the rank of the density matrix is not
# full in reality, but we keep more modes around than needed,
# some ws can go negative due to numerical error! This is
# extremely rare, but comon enough to cause crashes occasionally
# when there are thousands of individual matrices to transform
# every time this is called.
w = np.maximum(w, 0)
# 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.
new_Ws.append(np.dot(np.diag(np.sqrt(w)), v.transpose()))
# 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))
new_Ws = np.array(new_Ws)
# This does a super fast batched computation
U, S, Vh = t.linalg.svd(self.weights.data * probe_sqrt_intensities,
full_matrices=False)
self.weights.data = t.as_tensor(
new_Ws, dtype=self.weights.dtype, device=self.weights.device)
self.probe.data = t.as_tensor(
ortho_probes, device=self.probe.device, dtype=self.probe.dtype)
# 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)
axis=-3)
ortho_probes = analysis.orthogonalize_probes(prs)
if mode.lower() == 'amplitude':
@@ -625,11 +640,22 @@ class FastCCDPtycho(CDIModel):
lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Basis Probe Real Space Phases',
lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
('Average Weight Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(
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)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig, units=self.units),
('% of Power in Top Mode',
lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
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),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
src/cdtools/models/multislice_2d_ptycho.py
+84 -58
View File
@@ -495,74 +495,89 @@ class Multislice2DPtycho(CDIModel):
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1, normalize=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(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
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix.
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = self.probe.detach().cpu().numpy()
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=normalize)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
# 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)
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
# We generate the orthogonal probes based on this full-experiment
# representation of the light field.
ortho_probes, reexpressed_weights = \
analysis.orthogonalize_probes(
self.probe.detach(),
weight_matrix=all_weights,
return_reexpressed_weights=True
)
dm_rank = self.weights.shape[1]
# 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])
new_Ws = []
for rho in new_rhos:
# These are returned from smallest to largest - we want to keep
# the largest ones
w,v = sla.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
# For situations where the rank of the density matrix is not
# full in reality, but we keep more modes around than needed,
# some ws can go negative due to numerical error! This is
# extremely rare, but comon enough to cause crashes occasionally
# when there are thousands of individual matrices to transform
# every time this is called.
w = np.maximum(w,0)
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
# 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)
self.weights.data = t.as_tensor(new_Ws,
dtype=self.weights.dtype,device=self.weights.device)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
# 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)
# 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.
# 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)
# Needs to be updated to allow for plotting to an existing figure
@@ -575,12 +590,23 @@ class Multislice2DPtycho(CDIModel):
lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Probe Real Space Phase',
lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()),axis=0), fig=fig),
lambda self: len(self.weights.shape) >=2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig,units=self.units),
lambda self: len(self.weights.shape) >=2),
('Average Weight Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% of Power in Top Mode',
lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
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),
('Slice by Slice Real Part of T',
lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: self.exponentiate_obj),
src/cdtools/models/multislice_ptycho.py
+74 -76
View File
@@ -160,14 +160,6 @@ class MultislicePtycho(CDIModel):
self.register_buffer('I_phase', I_phase)
self.register_buffer('J_phase', J_phase)
#from matplotlib import pyplot as plt
#p.plot_real(
# self.obj[(np.s_[:],) + self.obj_view_slice],
# fig=fig,
# basis=self.obj_basis,
# units=self.units)
#plt.show()
self.register_buffer('simulate_finite_pixels',
t.tensor(simulate_finite_pixels, dtype=bool))
@@ -623,95 +615,96 @@ class MultislicePtycho(CDIModel):
self.probe.data.cpu(), iterations=iterations)
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(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(
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)
axis=-3)
ortho_probes = analysis.orthogonalize_probes(prs)
if mode.lower() == 'amplitude':
@@ -804,15 +797,20 @@ class MultislicePtycho(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/models/polarization_swept_ptycho.py
+4 -9
View File
@@ -521,16 +521,11 @@ class PolarizationSweptPtycho(CDIModel):
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# NOTE: untested, edited when ortho_probes was updated
for idx in range(self.probe.shape[0]):
oself.probe.data[idx] = analysis.orthogonalize_probes(
self.probe.data[idx])
probe = self.probe.detach().cpu().numpy()
ortho_probes = np.empty(probe.shape, dtype=probe.dtype)
for idx in range(probe.shape[0]):
ortho_probes[idx] = analysis.orthogonalize_probes(probe[idx])
self.probe.data = t.as_tensor(
ortho_probes,
device=self.probe.device,
dtype=self.probe.dtype)
plot_list = [
src/cdtools/models/polarized_fancy_ptycho.py
+3 -1
View File
@@ -408,8 +408,10 @@ class PolarizedFancyPtycho(FancyPtycho):
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
NOTE: Not updated with the update to orthogonalize_probes
"""
raise NotImplementedError('Function is known to be incorrect! Contact Abe Levitan at abraham.levitan@psi.ch to ask him to fix it.')
# First we treat the purely incoherent case
# I don't love this pattern of using an if statement with a return
src/cdtools/tools/analysis/analysis.py
+19 -151
View File
@@ -13,14 +13,23 @@ from scipy import linalg as sla
from scipy import special
from scipy import optimize as opt
__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']
__all__ = [
'product_svd',
'orthogonalize_probes',
'standardize',
'synthesize_reconstructions',
'calc_consistency_prtf',
'calc_deconvolved_cross_correlation',
'calc_frc',
'calc_vn_entropy',
'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):
@@ -76,15 +85,13 @@ def product_svd(A, B):
return U_final, S, Vh_final
def orthogonalize_probes_t(
def orthogonalize_probes(
probes,
weight_matrix=None,
n_probe_dims=2,
return_reexpressed_weights=False,
):
""" Orthogonalizes a set of incoherently mixing probes
## TODO fully replace orthogonalize_probes with orthogonalize_probes_t
This function takes any set of probe modes for mixed-mode ptychography,
which are considered to define a mutual coherence function, and returns
@@ -170,113 +177,6 @@ def orthogonalize_probes_t(
return orthogonalized_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
all of the retrieved probes, and then build the density matrix
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
TODO: is the above right, or are ortho_probes and probes flipped?
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 = probes.detach().cpu().numpy()
send_to_torch = True
except:
send_to_torch = False
# 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 = w[::-1]
v = v[:,::-1]
# We do this just to avoid total failure when the density
# matrix is not positive definite.
# In most cases (such as when rho is generated directly from some other
# matrix A such that rho=A A^dagger), w should never have any negative
# entries.
w = np.maximum(w,0)
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)
if normalize:
ortho_probes = 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,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:
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 = t.as_tensor(np.stack(ortho_probes))
A = t.as_tensor(A)
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
@@ -794,7 +694,7 @@ def calc_mode_power_fractions(
"""
if not assume_preorthogonalized:
ortho_probes = orthogonalize_probes_t(
ortho_probes = orthogonalize_probes(
probes,
weight_matrix=weight_matrix,
n_probe_dims=n_probe_dims,
@@ -814,38 +714,6 @@ def calc_mode_power_fractions(
return power_fractions
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
"""
# TODO depricate
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
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
tests/tools/test_analysis.py
+2 -72
View File
@@ -52,9 +52,9 @@ def test_product_svd():
assert np.allclose(prod_Vh, prod_U)
def test_orthogonalize_probes_t():
def test_orthogonalize_probes():
op = analysis.orthogonalize_probes_t
op = analysis.orthogonalize_probes
probe_xs = np.arange(64) - 32
probe_ys = np.arange(76) - 38
@@ -150,76 +150,6 @@ def test_orthogonalize_probes_t():
# 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
# and orthogonalize them. Then we can test two features of the results:
# 1) Are they orthogonal?
# 2) Is the total intensity at each point the same as it was originally?
probe_xs = np.arange(128) - 64
probe_ys = np.arange(150) - 75
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))]).astype(np.complex64)
# test that it works on numpy arrays
ortho_probes = analysis.orthogonalize_probes(probes)
# test that it also works on torch tensors
ortho_probes_t = analysis.orthogonalize_probes(t.as_tensor(probes)).numpy()
# This tests for orthogonality
for p1,p2 in combinations(ortho_probes,2):
assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
for p1,p2 in combinations(ortho_probes_t,2):
assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
probe_intensity = np.sum(np.abs(probes)**2,axis=0)
ortho_probe_intensity = np.sum(np.abs(ortho_probes)**2,axis=0)
ortho_probe_t_intensity = np.sum(np.abs(ortho_probes_t)**2,axis=0)
assert np.allclose(probe_intensity,ortho_probe_intensity)
assert np.allclose(probe_intensity,ortho_probe_t_intensity)
# Check that it returns normalized probes if we ask
ortho_probes = analysis.orthogonalize_probes(probes, normalize=True)
assert np.allclose([1,1,1],np.sum(np.abs(ortho_probes)**2,axis=(1,2)))
# And now we check that the A matrices actually work
ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True,
normalize=False)
assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes)
# And now we check that the A matrices actually work
ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True,
normalize=False)
assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes)
# The big problem here is that we haven't tested if it returns the
# probes or their complex conjugate.. we can test that by sending in
# one probe and checking that we get the right thing out
probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 20**2 + 100j))])
ortho_probes = analysis.orthogonalize_probes(probes)
# We need to correct for the arbitrary phase offset to be able to
# compare them
correction = ortho_probes[0]/probes[0]
correction = correction / np.abs(correction)
ortho_probes *= correction
assert np.allclose(ortho_probes,probes)
def test_standardize():