From 35f140d3486fc414391c6dd6cbebb0f900934e9a Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 25 Jul 2024 17:35:48 +0200 Subject: [PATCH] Fully replace orthogonalize_probes and update most models --- src/cdtools/models/fancy_ptycho.py | 6 +- src/cdtools/models/fastccd_ptycho.py | 136 ++++++++------ src/cdtools/models/multislice_2d_ptycho.py | 142 +++++++++------ src/cdtools/models/multislice_ptycho.py | 150 ++++++++-------- .../models/polarization_swept_ptycho.py | 13 +- src/cdtools/models/polarized_fancy_ptycho.py | 4 +- src/cdtools/tools/analysis/analysis.py | 170 ++---------------- tests/tools/test_analysis.py | 74 +------- 8 files changed, 270 insertions(+), 425 deletions(-) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index bf3728f..09e2718 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -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()) diff --git a/src/cdtools/models/fastccd_ptycho.py b/src/cdtools/models/fastccd_ptycho.py index a887bd6..59f8e3b 100644 --- a/src/cdtools/models/fastccd_ptycho.py +++ b/src/cdtools/models/fastccd_ptycho.py @@ -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)), diff --git a/src/cdtools/models/multislice_2d_ptycho.py b/src/cdtools/models/multislice_2d_ptycho.py index b1455f6..1d6b3cd 100644 --- a/src/cdtools/models/multislice_2d_ptycho.py +++ b/src/cdtools/models/multislice_2d_ptycho.py @@ -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), diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index 9a1c660..f9eab3e 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -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 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), diff --git a/src/cdtools/models/polarization_swept_ptycho.py b/src/cdtools/models/polarization_swept_ptycho.py index 7deb572..8e9594b 100644 --- a/src/cdtools/models/polarization_swept_ptycho.py +++ b/src/cdtools/models/polarization_swept_ptycho.py @@ -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 = [ diff --git a/src/cdtools/models/polarized_fancy_ptycho.py b/src/cdtools/models/polarized_fancy_ptycho.py index af6999f..a9ad37b 100644 --- a/src/cdtools/models/polarized_fancy_ptycho.py +++ b/src/cdtools/models/polarized_fancy_ptycho.py @@ -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 diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py index 754892a..5da67e5 100644 --- a/src/cdtools/tools/analysis/analysis.py +++ b/src/cdtools/tools/analysis/analysis.py @@ -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 diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 61323da..27c6a65 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -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():