From a260168d9d327fbd804953b6bf4e4182281c29b5 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 10 Jan 2023 16:30:35 -0800 Subject: [PATCH] Honestly I don't know, I must have forgotten to commit something --- examples/simple_ptycho.py | 13 +- src/cdtools/models/__init__.py | 1 + src/cdtools/models/fancy_ptycho.py | 2 + src/cdtools/models/fastccd_ptycho.py | 667 ++++++++++++++++++ src/cdtools/models/simple_ptycho.py | 17 +- .../image_processing/image_processing.py | 22 +- 6 files changed, 704 insertions(+), 18 deletions(-) create mode 100644 src/cdtools/models/fastccd_ptycho.py diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 04f4218..313e29f 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -36,7 +36,7 @@ normalization=0 for inputs, patterns in data_loader: normalization += t.sum(patterns).cpu().numpy() -def run_iteration(stop_event=None): +for it in range(20): loss = 0 N = 0 t0 = time.time() @@ -50,22 +50,19 @@ def run_iteration(stop_event=None): sim_patterns = model.forward(*inp) if hasattr(model, 'mask'): - loss = model.module.loss(pats,sim_patterns, mask=model.module.mask) + loss = model.loss(pats,sim_patterns, mask=model.mask) else: - loss = model.module.loss(pats,sim_patterns) + loss = model.loss(pats,sim_patterns) loss.backward() return loss.detach() loss += optimizer.step(closure).detach().cpu().numpy() print('time', time.time()-t0) - return loss / normalization + print(loss / normalization) -for it in range(20): - print(run_iteration()) - -print(model.module.save_results().keys()) +print(model.save_results().keys()) exit() # Finally, we plot the results model.module.inspect(dataset) diff --git a/src/cdtools/models/__init__.py b/src/cdtools/models/__init__.py index 9b2286b..df61f96 100644 --- a/src/cdtools/models/__init__.py +++ b/src/cdtools/models/__init__.py @@ -34,6 +34,7 @@ from cdtools.models.rpi import RPI from cdtools.models.multimode_rpi import MultimodeRPI from cdtools.models.time_resolved_ptycho_calibration import TimeResolvedPtychoCalibration from cdtools.models.time_resolved_rpi import TimeResolvedRPI +from cdtools.models.fastccd_ptycho import FastCCDPtycho # Still needs to be updated for the new complex numbers #from cdtools.models.s_matrix_ptycho import SMatrixPtycho diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 7812b4c..93edd02 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -647,8 +647,10 @@ class FancyPtycho(CDIModel): obj = self.obj.detach().cpu().numpy() background = self.background.detach().cpu().numpy()**2 weights = self.weights.detach().cpu().numpy() + oversampling = self.oversampling return {'basis': basis, 'translation': translations, 'probe': probe, 'obj': obj, 'background': background, + 'oversampling': oversampling, 'weights': weights} diff --git a/src/cdtools/models/fastccd_ptycho.py b/src/cdtools/models/fastccd_ptycho.py new file mode 100644 index 0000000..88eca59 --- /dev/null +++ b/src/cdtools/models/fastccd_ptycho.py @@ -0,0 +1,667 @@ +import torch as t +from cdtools.models import CDIModel +from cdtools.datasets import Ptycho2DDataset +from cdtools import tools +from cdtools.tools import plotting as p +from cdtools.tools import analysis +from matplotlib import pyplot as plt +from datetime import datetime +import numpy as np +from scipy import linalg as sla +from copy import copy + +__all__ = ['FastCCDPtycho'] + + +class FastCCDPtycho(CDIModel): + + def __init__(self, wavelength, detector_geometry, + probe_basis, + probe_guess, + obj_guess, + detector_slice=None, + threshold=None, + surface_normal=t.tensor([0., 0., 1.], dtype=t.float32), + min_translation=t.tensor([0, 0], dtype=t.float32), + background=None, + translation_offsets=None, + mask=None, + weights=None, + translation_scale=1, + saturation=None, + probe_support=None, + oversampling=1, + fourier_probe=False, + loss='amplitude mse', + units='um', + simulate_probe_translation=False, + simulate_finite_pixels=False, + ): + + super(FastCCDPtycho, self).__init__() + self.wavelength = t.tensor(wavelength) + self.detector_geometry = copy(detector_geometry) + det_geo = self.detector_geometry + if 'distance' in det_geo: + det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32) + if 'basis' in det_geo: + det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32) + if 'corner' in det_geo and det_geo['corner'] is not None: + det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32) + + self.min_translation = t.tensor(min_translation) + + self.probe_basis = t.tensor(probe_basis) + self.detector_slice = copy(detector_slice) + self.surface_normal = t.tensor(surface_normal) + + self.saturation = saturation + self.units = units + self.fourier_probe = fourier_probe + + if mask is None: + self.mask = mask + else: + self.mask = t.tensor(mask, dtype=t.bool) + + 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 + # object + if probe_guess.dim() > 2: + self.probe_norm = 1 * t.max(t.abs(probe_guess[0])) + else: + self.probe_norm = 1 * t.max(t.abs(probe_guess)) + + self.probe = t.nn.Parameter(probe_guess / self.probe_norm) + self.obj = t.nn.Parameter(obj_guess) + + if background is None: + if detector_slice is not None: + dummy_det = t.empty([s//oversampling + for s in self.probe[0].shape]) + shape = dummy_det[self.detector_slice].shape + #shape = self.probe[0][self.detector_slice].shape + else: + shape = [s//oversampling for s in self.probe[0]] + background = 1e-6 * t.ones(shape, dtype=t.float32) + + self.background = t.nn.Parameter(background) + + if threshold is None: + threshold = t.zeros([960,2], dtype=t.float32) + + self.threshold = t.nn.Parameter(threshold) + + if weights is None: + self.weights = None + else: + # We now need to distinguish between real-valued per-image + # weights and complex-valued per-mode weight matrices + if len(weights.shape) == 1: + # This is if it's just a list of numbers + self.weights = t.nn.Parameter(t.tensor(weights, + dtype=t.float32)) + else: + # Now this is a matrix of weights, so it needs to be complex + self.weights = t.nn.Parameter(t.tensor(weights, + dtype=t.complex64)) + + if translation_offsets is None: + self.translation_offsets = None + else: + t_o = t.tensor(translation_offsets, dtype=t.float32) + t_o = t_o / translation_scale + self.translation_offsets = t.nn.Parameter(t_o) + + self.translation_scale = translation_scale + + if probe_support is not None: + self.probe_support = probe_support + else: + self.probe_support = t.ones_like(self.probe[0], dtype=t.bool) + + self.oversampling = oversampling + + self.simulate_probe_translation = simulate_probe_translation + if simulate_probe_translation: + Is = t.arange(self.probe.shape[-2], dtype=t.float32) + Js = t.arange(self.probe.shape[-1], dtype=t.float32) + Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js)) + self.I_phase = 2 * np.pi* Is + self.J_phase = 2 * np.pi* Js + + self.simulate_finite_pixels = simulate_finite_pixels + + # Here we set the appropriate loss function + if (loss.lower().strip() == 'amplitude mse' + or loss.lower().strip() == 'amplitude_mse'): + self.loss = tools.losses.amplitude_mse + elif (loss.lower().strip() == 'poisson nll' + or loss.lower().strip() == 'poisson_nll'): + self.loss = tools.losses.poisson_nll + else: + raise KeyError('Specified loss function not supported') + + + @classmethod + def from_dataset(cls, + dataset, + probe_size=None, + randomize_ang=0, + padding=0, + n_modes=1, + dm_rank=None, + translation_scale=1, + saturation=None, + probe_support_radius=None, + propagation_distance=None, + scattering_mode=None, + oversampling=1, + auto_center=False, + opt_for_fft=False, + fourier_probe=False, + loss='amplitude mse', + units='um', + simulate_probe_translation=False, + simulate_finite_pixels=False, + ): + + wavelength = dataset.wavelength + det_basis = dataset.detector_geometry['basis'] + det_shape = dataset[0][1].shape + distance = dataset.detector_geometry['distance'] + + # always do this on the cpu + get_as_args = dataset.get_as_args + dataset.get_as(device='cpu') + + # We include the *extras to make this work even with datasets, like + # polarization dependent datasets, that might toss out extra inputs + (indices, translations, *extras), patterns = dataset[:] + + dataset.get_as(*get_as_args[0], **get_as_args[1]) + + # Set to none to avoid issues with things outside the detector + if auto_center: + center = tools.image_processing.centroid(t.sum(patterns, dim=0)) + else: + center = None + + # Then, generate the probe geometry from the dataset + ewg = tools.initializers.exit_wave_geometry + probe_basis, probe_shape, det_slice = ewg(det_basis, + det_shape, + wavelength, + distance, + center=center, + padding=padding, + opt_for_fft=opt_for_fft, + oversampling=oversampling) + + if hasattr(dataset, 'sample_info') and \ + dataset.sample_info is not None and \ + 'orientation' in dataset.sample_info: + surface_normal = dataset.sample_info['orientation'][2] + else: + surface_normal = np.array([0., 0., 1.]) + + # If this information is supplied when the function is called, + # then we override the information in the .cxi file + if scattering_mode in {'t', 'transmission'}: + surface_normal = np.array([0., 0., 1.]) + elif scattering_mode in {'r', 'reflection'}: + outgoing_dir = np.cross(det_basis[:, 0], det_basis[:, 1]) + outgoing_dir /= np.linalg.norm(outgoing_dir) + surface_normal = outgoing_dir + np.array([0., 0., 1.]) + surface_normal /= -np.linalg.norm(surface_normal) + + # Next generate the object geometry from the probe geometry and + # the translations + + pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal) + + obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) + + if hasattr(dataset, 'background') and dataset.background is not None: + background = t.sqrt(dataset.background) + else: + background = None + + # Finally, initialize the probe and object using this information + if probe_size is None: + probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) + else: + probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance) + + # Now we initialize all the subdominant probe modes + probe_max = t.max(t.abs(probe)) + probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)] + + # For a Fourier space probe + if fourier_probe: + probe = tools.propagators.far_field(probe) + + probe = t.stack([probe, ] + probe_stack) + + obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + + det_geo = dataset.detector_geometry + + translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5) + + if dm_rank is not None and dm_rank != 0: + if dm_rank > n_modes: + raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.') + elif dm_rank == -1: + # dm_rank == -1 is defined to mean full-rank + dm_rank = n_modes + + Ws = t.zeros(len(dataset), dm_rank, n_modes, dtype=t.complex64) + # Start with as close to the identity matrix as possible, + # cutting of when we hit the specified maximum rank + for i in range(0, dm_rank): + Ws[:, i, i] = 1 + else: + # dm_rank == None or dm_rank = 0 triggers a special case where + # a standard incoherent multi-mode model is used. This is the + # default, because it is so common. + # In this case, we define a set of weights which only has one index + Ws = t.ones(len(dataset)) + + if hasattr(dataset, 'intensities') and dataset.intensities is not None: + Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...] + / t.mean(dataset.intensities)) + + if hasattr(dataset, 'mask') and dataset.mask is not None: + mask = dataset.mask.to(t.bool) + else: + mask = None + + if probe_support_radius is not None: + probe_support = t.zeros(probe[0].shape, dtype=t.bool) + xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]] + xs = xs - np.mean(xs) + ys = ys - np.mean(ys) + Rs = np.sqrt(xs**2 + ys**2) + + probe_support[Rs < probe_support_radius] = 1 + probe = probe * probe_support[None, :, :] + + else: + probe_support = None + + return cls(wavelength, det_geo, probe_basis, probe, obj, + detector_slice=det_slice, + surface_normal=surface_normal, + min_translation=min_translation, + translation_offsets=translation_offsets, + weights=Ws, mask=mask, background=background, + translation_scale=translation_scale, + saturation=saturation, + probe_support=probe_support, + fourier_probe=fourier_probe, + oversampling=oversampling, + loss=loss, units=units, + simulate_probe_translation=simulate_probe_translation, + simulate_finite_pixels=simulate_finite_pixels) + + + def interaction(self, index, translations, *args): + + # The *args is included so that this can work even when given, say, + # a polarized ptycho dataset that might spit out more inputs. + + # Step 1 is to convert the translations for each position into a + # value in pixels + pix_trans = tools.interactions.translations_to_pixel( + self.probe_basis, + translations, + surface_normal=self.surface_normal) + pix_trans -= self.min_translation + # We then add on any recovered translation offset, if they exist + if self.translation_offsets is not None: + pix_trans += (self.translation_scale * + self.translation_offsets[index]) + + # This restricts the basis probes to stay within the probe support + basis_prs = self.probe * self.probe_support[..., :, :] + + # For a Fourier-space probe, we take an IFT + if self.fourier_probe: + basis_prs = tools.propagators.inverse_far_field(basis_prs) + + # Now we construct the probes for each shot from the basis probes + if self.weights is not None: + Ws = self.weights[index] + else: + try: + Ws = t.ones(len(index)) # I'm positive this introduced a bug + except: + Ws = 1 + + if self.weights is None or len(self.weights[0].shape) == 0: + # If a purely stable coherent illumination is defined + prs = Ws[..., None, None, None] * basis_prs + else: + # If a frame-by-frame weight matrix is defined + # This takes the dot product of all the weight matrices with + # the probes. The output has dimensions of translation, then + # coherent mode index, then x,y, and then complex index + # Maybe this can be done with a matmul now? + prs = t.sum(Ws[..., None, None] * basis_prs, axis=-3) + + if self.simulate_probe_translation: + det_pix_trans = tools.interactions.translations_to_pixel( + self.detector_geometry['basis'], + translations, + surface_normal=self.surface_normal) + + probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] * + self.I_phase[None,...] + + det_pix_trans[:,1,None,None] * + self.J_phase[None,...])) + prs = prs * probe_masks[...,None,:,:] + + + # Now we actually do the interaction, using the sinc subpixel + # translation model as per usual + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs, self.obj, pix_trans, + shift_probe=True, multiple_modes=True) + return exit_waves + + + def forward_propagator(self, wavefields): + return tools.propagators.far_field(wavefields) + + + def backward_propagator(self, wavefields): + return tools.propagators.inverse_far_field(wavefields) + + + def measurement(self, wavefields): + measured = tools.measurements.quadratic_background( + wavefields, + self.background, + detector_slice=self.detector_slice, + measurement=tools.measurements.incoherent_sum, + saturation=self.saturation, + oversampling=self.oversampling, + simulate_finite_pixels=self.simulate_finite_pixels) + thresholds = t.repeat_interleave(self.threshold, 480,dim=-1) + return t.clamp(measured - thresholds, min=0.001) + + + # Note: No "loss" function is defined here, because it is added + # dynamically during object creation in __init__ + + def to(self, *args, **kwargs): + super(FastCCDPtycho, self).to(*args, **kwargs) + self.wavelength = self.wavelength.to(*args, **kwargs) + # move the detector geometry too + det_geo = self.detector_geometry + if 'distance' in det_geo: + det_geo['distance'] = det_geo['distance'].to(*args, **kwargs) + if 'basis' in det_geo: + det_geo['basis'] = det_geo['basis'].to(*args, **kwargs) + if 'corner' in det_geo and det_geo['corner'] is not None: + det_geo['corner'] = det_geo['corner'].to(*args, **kwargs) + + if self.mask is not None: + self.mask = self.mask.to(*args, **kwargs) + + if self.simulate_probe_translation: + self.I_phase = self.I_phase.to(*args, **kwargs) + self.J_phase = self.J_phase.to(*args, **kwargs) + + self.min_translation = self.min_translation.to(*args, **kwargs) + self.probe_basis = self.probe_basis.to(*args, **kwargs) + self.probe_norm = self.probe_norm.to(*args, **kwargs) + self.probe_support = self.probe_support.to(*args, **kwargs) + self.surface_normal = self.surface_normal.to(*args, **kwargs) + + + def sim_to_dataset(self, args_list, calculation_width=None): + # In the future, potentially add more control + # over what metadata is saved (names, etc.) + + # First, I need to gather all the relevant data + # that needs to be added to the dataset + entry_info = {'program_name': 'cdtools', + 'instrument_n': 'Simulated Data', + 'start_time': datetime.now()} + + surface_normal = self.surface_normal.detach().cpu().numpy() + xsurfacevec = np.cross(np.array([0., 1., 0.]), surface_normal) + xsurfacevec /= np.linalg.norm(xsurfacevec) + ysurfacevec = np.cross(surface_normal, xsurfacevec) + ysurfacevec /= np.linalg.norm(ysurfacevec) + orientation = np.array([xsurfacevec, ysurfacevec, surface_normal]) + + sample_info = {'description': 'A simulated sample', + 'orientation': orientation} + + + detector_geometry = self.detector_geometry + mask = self.mask + wavelength = self.wavelength + indices, translations = args_list + + data = [] + len(indices) + if calculation_width is None: + calculation_width = len(indices) + index_chunks = [indices[i:i + calculation_width] + for i in range(0, len(indices), + calculation_width)] + translation_chunks = [translations[i:i + calculation_width] + for i in range(0, len(indices), + calculation_width)] + + + # Then we simulate the results + data = [self.forward(idx, trans).detach() + for idx, trans in zip(index_chunks, translation_chunks)] + + data = t.cat(data, dim=0) + # And finally, we make the dataset + return Ptycho2DDataset( + translations, data, + entry_info=entry_info, + sample_info=sample_info, + wavelength=wavelength, + detector_geometry=detector_geometry, + mask=mask) + + + def corrected_translations(self, dataset): + translations = dataset.translations.to( + dtype=t.float32, device=self.probe.device) + if (hasattr(self, 'translation_offsets') and + self.translation_offsets is not None): + t_offset = tools.interactions.pixel_to_translations( + self.probe_basis, + self.translation_offsets * self.translation_scale, + surface_normal=self.surface_normal) + return translations + t_offset + 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 tidy_probes(self, normalization=1, normalize=False): + """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 + """ + + # 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 + 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) + return + + # This is for the unified mode case + + # 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)) + + new_rhos /= normalization + ortho_probes *= np.sqrt(normalization) + + dm_rank = self.weights.shape[1] + + 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) + + 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) + + + 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) + + if mode.lower() == 'amplitude': + return np.abs(ortho_probes.detach().cpu().numpy()) + if mode.lower() == 'root_sum_intensity': + return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2, + axis=0) + if mode.lower() == 'phase': + return np.angle(ortho_probes.detach().cpu().numpy()) + + probe_matrix = np.zeros([self.probe.shape[0]]*2, + dtype=np.complex64) + np_probes = self.probe.detach().cpu().numpy() + for i in range(probe_matrix.shape[0]): + for j in range(probe_matrix.shape[0]): + probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj()) + + weights = self.weights.detach().cpu().numpy() + + probe_intensities = np.sum(np.tensordot(weights, probe_matrix, axes=1) + * weights.conj(), axis=2) + + # Imaginary part is already essentially zero up to rounding error + probe_intensities = np.real(probe_intensities) + + values = np.sum(probe_intensities, axis=1) + if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity': + cmap = 'viridis' + else: + cmap = 'twilight' + + p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity', cmap=cmap, **kwargs), + + + plot_list = [ + ('', + lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'), + lambda self: len(self.weights.shape) >= 2), + ('', + lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'), + lambda self: len(self.weights.shape) >= 2), + ('', + lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'), + lambda self: len(self.weights.shape) >= 2), + ('Basis Probe Fourier Space Amplitudes', + lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Basis Probe Fourier Space Phases', + lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Basis Probe Real Space Amplitudes', + 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), + 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), + ('Object Amplitude', + lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)), + ('Object Phase', + lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)), + ('Corrected Translations', + lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), + ('Background', + lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)), + ('Thresholds', + lambda self, fig: p.plot_real(t.repeat_interleave(self.threshold, 480,dim=-1), fig=fig)) + ] + +# def plot_errors(self, dataset): + + + + def save_results(self, dataset): + thresholds = self.thresholds.detach().cpu().numpy() + basis = self.probe_basis.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + probe = self.probe.detach().cpu().numpy() + probe = probe * self.probe_norm.detach().cpu().numpy() + obj = self.obj.detach().cpu().numpy() + background = self.background.detach().cpu().numpy()**2 + weights = self.weights.detach().cpu().numpy() + oversampling = self.oversampling + + return {'basis': basis, 'translation': translations, + 'probe': probe, 'obj': obj, + 'background': background, + 'oversampling': oversampling, + 'weights': weights, 'thresholds':thresholds} diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index 79c0df4..91303e6 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -8,12 +8,8 @@ from torch.utils import data as torchdata from datetime import datetime import numpy as np -__all__ = ['SimplePtycho'] +__all__ = ['SimplePtycho'] -class complexWrapper(t.Tensor): - def __new__(base_tensor): - return t.view_as_complex(base_tensor) - class SimplePtycho(CDIModel): """A simple ptychography model for exploring ideas and extensions @@ -49,14 +45,17 @@ class SimplePtycho(CDIModel): # object self.register_buffer('probe_norm', t.max(t.abs(probe_guess))) - self.probe_data = complexParameter(probe_guess/self.probe_norm) - #self.probe_data = t.nn.Parameter(t.view_as_real(probe_guess / self.probe_norm)) + #self.probe_data = complexWrapper(probe_guess/self.probe_norm) + self.probe_data = t.nn.Parameter(t.view_as_real(probe_guess / self.probe_norm)) self.obj_data = t.nn.Parameter(t.view_as_real(obj_guess)) + # Do this for all the complex-valued parameters which are stored as real + # valued for compatibility reasons. Note that updating model.probe.data + # won't update the underlying storage, but something like: + # model.probe.data[:] = ... + # will update the underlying storage @property def probe(self): - probe = t.view_as_complex(self.probe_data) - return t.view_as_complex(self.probe_data) @property diff --git a/src/cdtools/tools/image_processing/image_processing.py b/src/cdtools/tools/image_processing/image_processing.py index d202319..c7f9437 100644 --- a/src/cdtools/tools/image_processing/image_processing.py +++ b/src/cdtools/tools/image_processing/image_processing.py @@ -11,10 +11,30 @@ import numpy as np import torch as t from cdtools.tools import propagators -__all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift', +__all__ = ['hann_window', 'centroid', 'centroid_sq', 'sinc_subpixel_shift', 'find_subpixel_shift', 'find_pixel_shift', 'find_shift', 'convolve_1d', 'fourier_upsample'] +def hann_window(im): + """ Applies a hann window to a 2D image to apodize it + + TODO: update for pytorch + + Parameters + ---------- + im: np.array + A numpy array to apply a hann window to + + Returns + ------- + apodidzed : np.array + The image, apodized by a hann window + """ + Xs, Ys = np.mgrid[:im.shape[0],:im.shape[1]] + Xhann = np.sin(np.pi*Xs/(im.shape[1]-1))**2 + Yhann = np.sin(np.pi*Ys/(im.shape[0]-1))**2 + Hann = Xhann * Yhann + return im * Hann def centroid(im, dims=2): """Returns the centroid of an image or a stack of images