From d0a0d125219b9ed60ccfd596913094c258f131ea Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 4 Jan 2024 16:50:10 +0100 Subject: [PATCH] Add a normalize reconstruction set function --- src/cdtools/models/fancy_ptycho.py | 42 +++- src/cdtools/tools/analysis/analysis.py | 231 +++++++++++++++++- .../image_processing/image_processing.py | 73 +++++- tests/tools/test_image_processing.py | 2 + 4 files changed, 337 insertions(+), 11 deletions(-) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 5a3384e..72d8039 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -517,6 +517,18 @@ class FancyPtycho(CDIModel): 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 + + Note that this does not compensate for the centering by adjusting + the object, so it's a good idea to reset the object after centering + the probes + """ + centered_probe = tools.image_processing.center( + self.probe.data.cpu(), iterations=iterations) + self.probe.data = centered_probe.to(device=self.probe.data.device) + + def tidy_probes(self, normalization=1, normalize=False): """Tidies up the probes @@ -703,14 +715,18 @@ class FancyPtycho(CDIModel): ('Background', lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)) ] - -# def plot_errors(self, dataset): - def save_results(self, dataset): + # This will save out everything needed to recreate the object + # in the same state, but it's not the best formatted. For example, + # "background" stores the square root of the background, etc. + state_dict = super().save_results() + + # So, we also save out the main results in a more readable format basis = self.probe_basis.detach().cpu().numpy() - translations = self.corrected_translations(dataset).detach().cpu().numpy() + translations=self.corrected_translations(dataset).detach().cpu().numpy() + original_translations = dataset.translations.detach().cpu().numpy() probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() obj = self.obj.detach().cpu().numpy() @@ -719,8 +735,16 @@ class FancyPtycho(CDIModel): oversampling = self.oversampling wavelength = self.wavelength.cpu().numpy() - return {'basis': basis, 'translation': translations, - 'probe': probe, 'obj': obj, - 'background': background, - 'oversampling': oversampling, - 'weights': weights, 'wavelength': wavelength} + return { + 'basis': basis, + 'translations': translations, + 'original_translations': original_translations, + 'probe': probe, + 'obj': obj, + 'background': background, + 'oversampling': oversampling, + 'weights': weights, + 'wavelength': wavelength, + 'state_dict': state_dict, + } + diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py index 5d61874..75048ba 100644 --- a/src/cdtools/tools/analysis/analysis.py +++ b/src/cdtools/tools/analysis/analysis.py @@ -15,7 +15,9 @@ from scipy import special __all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions', 'calc_consistency_prtf', 'calc_deconvolved_cross_correlation', 'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction', - 'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error'] + 'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error', + 'remove_phase_ramp', 'remove_amplitude_exponent', + 'standardize_reconstruction_set'] def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False): @@ -960,3 +962,230 @@ def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, s return bins[:-1], frc, threshold +def remove_phase_ramp(im, window, probe=None): + + window = im[window] + + Is, Js = np.mgrid[:window.shape[0],:window.shape[1]] + def zero_freq_component(freq): + phase_ramp = np.exp(2j * np.pi * (freq[0] * Is + freq[1] * Js)) + return -np.abs(np.sum(phase_ramp * window))**2 + + x0 = np.array([0,0]) + result = opt.minimize(zero_freq_component, x0) + center_freq = result['x'] + + Is, Js = np.mgrid[:im.shape[0],:im.shape[1]] + phase_ramp = np.exp(2j * np.pi * (center_freq[0] * Is + center_freq[1] * Js)) + im = im * phase_ramp + + if probe is not None: + Is, Js = np.mgrid[:probe.shape[-2],:probe.shape[-1]] + phase_ramp = np.exp(-2j * np.pi * (center_freq[0] * Is + center_freq[1] * Js)) + probe = probe * phase_ramp + return im, probe + else: + return im + + +def remove_amplitude_exponent(im, window, probe=None, weights=None, translations=None, basis=None): + window = np.abs(im[window]) + + Is, Js = np.mgrid[:window.shape[0],:window.shape[1]] + def rms_error(x): + constant = x[0] + growth_rate = x[1:] + exponential_decay = constant * np.exp((growth_rate[0] * Is + growth_rate[1] * Js)) + return np.sum((window - exponential_decay)**2) + + x0 = np.array([1,0,0]) + result = opt.minimize(rms_error, x0, method='Nelder-Mead') + growth_rate = result['x'][1:] + + Is, Js = np.mgrid[:im.shape[0],:im.shape[1]] + exponential_decay = np.exp(-(growth_rate[0] * Is + growth_rate[1] * Js)) + im = im * exponential_decay + to_return = (im,) + + if probe is not None: + Is, Js = np.mgrid[:probe.shape[-2],:probe.shape[-1]] + exponential_decay = np.exp((growth_rate[0] * Is + growth_rate[1] * Js)) + probe = probe * exponential_decay + to_return = to_return + (probe,) + + if weights is not None: + pix_translations = cdtools.tools.interactions.translations_to_pixel(t.as_tensor(basis), t.as_tensor(translations)).numpy() + pix_translations -= np.min(pix_translations,axis=0) + weights = weights * np.exp(growth_rate[0] * pix_translations[:,0] + growth_rate[1] * pix_translations[:,1]) + to_return = to_return + (weights,) + + if len(to_return) == 1: + return to_return[0] + else: + return to_return + +def make_illumination_map(results, total_probe_intensity=None, padding=200): + + probe_intensity = np.sum(np.abs(results['probe'])**2, axis=0) + if total_probe_intensity is not None: + probe_intensity = probe_intensity / np.sum(probe_intensity) * total_probe_intensity + + illumination_map = np.zeros_like(results['obj'], dtype=np.float32) + translations = t.as_tensor(results['translations']) + pix_translations = cdtools.tools.interactions.translations_to_pixel( + t.as_tensor(results['basis']), translations) + + # We use the min_translation stored in the model, which was calculated + # from the uncorrected translations and therefore may be different from + # what it is if we recalculate it on the corrected translations + min_translation = t.as_tensor(results['state_dict']['min_translation']) + pix_translations = np.round((pix_translations - min_translation).numpy()).astype(int) + + for translation in pix_translations: + illumination_map[translation[0]:translation[0]+probe_intensity.shape[0], + translation[1]:translation[1]+probe_intensity.shape[1]] += probe_intensity + + return illumination_map + + +def standardize_reconstruction_set( + half_1, + half_2, + full, + correct_phase_offset=True, + remove_phase_ramp=True, + remove_amplitude_exponent=False, + window=np.s_[:,:], + nbins=50 +): + """Normalizes and analyses a set of 50/50/100% reconstructions + + It's very common to split a ptychography dataset into two sub-datasets, + each with 50% of the exposures, so that the difference between the two + sub-datasets can be used to estimate the quality and resolution of the + final, full reconstruction. But to do that analysis, first the + reconstructions need to be aligned with respect to each other and + normalized in a few ways. + + This function takes the results (as output by model.save_results) of + a set of 50/50/100% reconstructions and: + + - Aligns the object reconstructions with one another + - Corrects for the global phase offset (by default) + - Sets a sensible value for the object/probe phase ramp (by default) + - Sets a sensible value for the object/probe exponential decay (off by + default) + - Calculates the FRC and derived SSNR. + + Then, these results are packaged into an output dictionary. The output + does not retain all the information from the inputs, so if full traceability + is desired, do not delete the files containing the individual + reconstructions + + Parameters + ---------- + half_1 : dict + The result of the first half dataset, as returned by model.save_results + half_2 : dict + The result of the second half dataset, as returned by model.save_results + full : dict + The result of the full dataset, as returned by model.save_results + + Returns + ------- + results : dict + A dictionary containing the synthesized results + """ + # We get the two half-data reconstructions + obj_1, probe_1, weights_1 = half_1['obj'],half_1['probe'],half_1['weights'] + obj_2, probe_2, weights_2 = half_2['obj'],half_2['probe'],half_2['weights'] + obj, probe, weights = full['obj'], full['probe'], full['weights'] + + + if remove_phase_ramp: + obj_1, probe_1 = remove_phase_ramp( + half_1['obj'], window, probe=half_1['probe']) + obj_2, probe_2 = remove_phase_ramp( + half_2['obj'], window, probe=half_2['probe']) + obj, probe = remove_phase_ramp( + full['obj'], window, probe=full['probe']) + + if remove_amplitude_exponent: + obj_1, probe_1, weights_1 = remove_amplitude_exponent( + obj_1, window, probe=probe_1, + weights=half_1['weights'], + basis=half_1['basis'], + translations=half_1['translations']) + obj_2, probe_2, weights_2 = remove_amplitude_exponent( + obj_2, window, probe=probe_2, + weights=half_2['weights'], + basis=half_2['basis'], + translations=half_2['translations']) + obj, probe, weights = remove_amplitude_exponent( + obj, window, probe=probe, + weights=full['weights'], + basis=full['basis'], + translations=full['translations']) + + if correct_phase_offset: + obj_1 = np.exp(-1j* np.angle(np.sum(obj_1[window]))) * obj_1 + obj_2 = np.exp(-1j* np.angle(np.sum(obj_2))) * obj_2 + obj = np.exp(-1j* np.angle(np.sum(obj))) * obj + + + # TODO calculate the illumination map using the corrected probes and + # the appropriately shifted translations, so that all the maps are + # directly comparable + illumination_map_1 = make_illumination_map(half_1, total_probe_intensity=total_probe_intensity) + illumination_map_2 = make_illumination_map(half_2, total_probe_intensity=total_probe_intensity) + illumination_map = make_illumination_map(full, total_probe_intensity=total_probe_intensity) + + # Todo update the translations to account for the determined shift + shift_1 = cdtools.tools.image_processing.find_shift( + t.as_tensor(ip.hann_window(np.abs(obj[window]))), + t.as_tensor(ip.hann_window(np.abs(obj_1[window])))) + obj_1 = cdtools.tools.image_processing.sinc_subpixel_shift( + t.as_tensor(obj_1), shift_1).numpy() + + shift_2 = cdtools.tools.image_processing.find_shift( + t.as_tensor(ip.hann_window(np.abs(obj[window]))), + t.as_tensor(ip.hann_window(np.abs(obj_2[window])))) + obj_2 = cdtools.tools.image_processing.sinc_subpixel_shift( + t.as_tensor(obj_2), shift_2).numpy() + + freqs, frc, threshold = cdtools.tools.analysis.calc_frc( + ip.hann_window(obj_1[window]), + ip.hann_window(obj_2[window]), + half_1['basis'], nbins=nbins, limit='corner') + + # The correct formulation when the final output is the full reconstruction + ssnr = 2 * frc / (1 - frc) + + results = { + 'obj_half_1': obj_1, + 'probe_half_1': probe_1, + 'weights_half_1': weights_1, + 'translations_half_1': half_1['translations'], + 'background_1': half_1['background'], + 'illumination_map_1': illumination_map_1, + 'obj_half_2': obj_2, + 'probe_half_2': probe_2, + 'weights_half_2': weights_2, + 'translations_half_2': half_2['translations'], + 'background_2': half_2['background'], + 'illumination_map_2': illumination_map_2, + 'obj_full': obj, + 'probe_full': probe, + 'weights_full': weights, + 'translations_full': full['translations'], + 'background_full': full['background'], + 'illumination_map_full': illumination_map, + 'wavelength': full['wavelength'], + 'basis': full['basis'], + 'frc_freqs': freqs, + 'frc': frc, + 'frc_threshold': threshold, + 'ssnr': ssnr} + + return results + diff --git a/src/cdtools/tools/image_processing/image_processing.py b/src/cdtools/tools/image_processing/image_processing.py index c7f9437..8a9eb4b 100644 --- a/src/cdtools/tools/image_processing/image_processing.py +++ b/src/cdtools/tools/image_processing/image_processing.py @@ -13,7 +13,7 @@ from cdtools.tools import propagators __all__ = ['hann_window', 'centroid', 'centroid_sq', 'sinc_subpixel_shift', 'find_subpixel_shift', 'find_pixel_shift', 'find_shift', - 'convolve_1d', 'fourier_upsample'] + 'convolve_1d', 'fourier_upsample', 'center'] def hann_window(im): """ Applies a hann window to a 2D image to apodize it @@ -337,3 +337,74 @@ def fourier_upsample(ims, preserve_mean=False): return propagators.inverse_far_field(upsampled) +def center(image, image_dims=2, use_power=True, iterations=4): + """Automatically centers an image or stack of images + + This function is designed with probes for ptychography in mind, so the + default centering method is to place the centroid of the magnitude-squared + of the input image on the central pixel. + + # TODO I should place the centroid actually at the zero-frequency + # pixel, rather than the center of the array, so this can be used in + # Fourier space also. + + Because it's intended for probes, the function will sum over all extra + dimensions beyond when calculating the centroid, and shift + all the images by the same amount. It does *not* calculate a separate + shift for each image in the stack + + It also centers by using circular shifts. This means that, after calculating + the centroid position and shifting by that amount, the centroid will not + be perfectly centered. To counteract this, multiple iterations are run, + by default 4 + + + Parameters + ---------- + image : torch.Tensor + The ... x N x M image to center + image_dims : int + Default 2, the number of dimensions to center along + use_power : bool + Default True, whether to use the square of the magnitude + iterations : int + Default 4, the number of iterations to do + + Returns + ------- + centered_im : torch.Tensor + The centered image + + """ + # Make sure we dont screw with the input image + image = t.clone(image) + + if image_dims !=2: + raise NotImplementedError('Implementing centerings with dimension != ' + '2 requires modifying some other functions ' + 'and is not yet implemented.') + + if image_dims > image.ndim: + raise IndexError('Number of image dimensions cannot exceed the ' + 'dimensionality of the input') + + im_shape = image.shape + # This adds an extra dimension if the image dimensionality is equal to + # image_dims, and ravels any extra dimensions + reshaped_im = image.reshape([-1, ] + list(im_shape[-image_dims:])) + + for i in range(iterations): + if use_power: + to_center = t.sum(t.abs(reshaped_im)**2, dim=0) + else: + to_center = t.sum(t.abs(reshaped_im), dim=0) + + im_centroid = centroid(to_center) + + for i in range(reshaped_im.shape[0]): + reshaped_im[i] = sinc_subpixel_shift( + reshaped_im[i], + (-im_centroid[0] + im_shape[-2] / 2, + -im_centroid[1] + im_shape[-1] / 2)) + + return reshaped_im.reshape(im_shape) diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index 57795fb..573f901 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -139,3 +139,5 @@ def test_convolve_1d(): assert np.allclose(convolved,np_result) +def test_center(): + raise Exception('Not yet tested')