diff --git a/src/cdtools/models/__init__.py b/src/cdtools/models/__init__.py index 9a10e91..12643df 100644 --- a/src/cdtools/models/__init__.py +++ b/src/cdtools/models/__init__.py @@ -20,24 +20,22 @@ defining a new ptychography model before attempting to do so. """ -# I don't believe that __all__ really needed, but it's nice to define it -# to be explicit that import * is safe -__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'PolarizationSweptPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI', 'TimeResolvedPtychoCalibration', 'TimeResolvedRPI'] +# We define __all__ to be sure that import * only imports what we want +__all__ = [ + 'CDIModel', + 'SimplePtycho', + 'FancyPtycho', + 'Bragg2DPtycho', + 'Multislice2DPtycho', + 'MultislicePtycho', + 'RPI', +] from cdtools.models.base import CDIModel from cdtools.models.simple_ptycho import SimplePtycho from cdtools.models.fancy_ptycho import FancyPtycho -from cdtools.models.polarized_fancy_ptycho import PolarizedFancyPtycho -from cdtools.models.polarization_swept_ptycho import PolarizationSweptPtycho from cdtools.models.bragg_2d_ptycho import Bragg2DPtycho from cdtools.models.multislice_2d_ptycho import Multislice2DPtycho from cdtools.models.multislice_ptycho import MultislicePtycho 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/base.py b/src/cdtools/models/base.py index 754fe77..efbcb30 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -39,7 +39,6 @@ import queue import time from scipy import io from contextlib import contextmanager -from .complex_lbfgs import MyLBFGS from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch __all__ = ['CDIModel'] @@ -629,8 +628,6 @@ class CDIModel(t.nn.Module): optimizer = t.optim.LBFGS(self.parameters(), lr = lr, history_size=history_size, line_search_fn=line_search_fn) - #optimizer = MyLBFGS(self.parameters(), - # lr = lr, history_size=history_size) return self.AD_optimize(iterations, data_loader, optimizer, regularization_factor=regularization_factor, diff --git a/src/cdtools/models/complex_lbfgs.py b/src/cdtools/models/complex_lbfgs.py deleted file mode 100644 index 2a03696..0000000 --- a/src/cdtools/models/complex_lbfgs.py +++ /dev/null @@ -1,485 +0,0 @@ -import torch -from functools import reduce -from torch.optim.optimizer import Optimizer - - -def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None): - # ported from https://github.com/torch/optim/blob/master/polyinterp.lua - # Compute bounds of interpolation area - if bounds is not None: - xmin_bound, xmax_bound = bounds - else: - xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1) - - # Code for most common case: cubic interpolation of 2 points - # w/ function and derivative values for both - # Solution in this case (where x2 is the farthest point): - # d1 = g1 + g2 - 3*(f1-f2)/(x1-x2); - # d2 = sqrt(d1^2 - g1*g2); - # min_pos = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2)); - # t_new = min(max(min_pos,xmin_bound),xmax_bound); - d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2) - d2_square = d1**2 - g1 * g2 - if d2_square >= 0: - d2 = d2_square.sqrt() - if x1 <= x2: - min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2)) - else: - min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2)) - return min(max(min_pos, xmin_bound), xmax_bound) - else: - return (xmin_bound + xmax_bound) / 2. - - -def _strong_wolfe(obj_func, - x, - t, - d, - f, - g, - gtd, - c1=1e-4, - c2=0.9, - tolerance_change=1e-9, - max_ls=25): - # ported from https://github.com/torch/optim/blob/master/lswolfe.lua - d_norm = d.abs().max() - g = g.clone(memory_format=torch.contiguous_format) - # evaluate objective and gradient using initial step - f_new, g_new = obj_func(x, t, d) - ls_func_evals = 1 - gtd_new = g_new.dot(d) - - # bracket an interval containing a point satisfying the Wolfe criteria - t_prev, f_prev, g_prev, gtd_prev = 0, f, g, gtd - done = False - ls_iter = 0 - while ls_iter < max_ls: - # check conditions - if f_new > (f + c1 * t * gtd) or (ls_iter > 1 and f_new >= f_prev): - bracket = [t_prev, t] - bracket_f = [f_prev, f_new] - bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)] - bracket_gtd = [gtd_prev, gtd_new] - break - - if abs(gtd_new) <= -c2 * gtd: - bracket = [t] - bracket_f = [f_new] - bracket_g = [g_new] - done = True - break - - if gtd_new >= 0: - bracket = [t_prev, t] - bracket_f = [f_prev, f_new] - bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)] - bracket_gtd = [gtd_prev, gtd_new] - break - - # interpolate - min_step = t + 0.01 * (t - t_prev) - max_step = t * 10 - tmp = t - t = _cubic_interpolate( - t_prev, - f_prev, - gtd_prev, - t, - f_new, - gtd_new, - bounds=(min_step, max_step)) - - # next step - t_prev = tmp - f_prev = f_new - g_prev = g_new.clone(memory_format=torch.contiguous_format) - gtd_prev = gtd_new - f_new, g_new = obj_func(x, t, d) - ls_func_evals += 1 - gtd_new = g_new.dot(d) - ls_iter += 1 - - # reached max number of iterations? - if ls_iter == max_ls: - bracket = [0, t] - bracket_f = [f, f_new] - bracket_g = [g, g_new] - - # zoom phase: we now have a point satisfying the criteria, or - # a bracket around it. We refine the bracket until we find the - # exact point satisfying the criteria - insuf_progress = False - # find high and low points in bracket - low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0) - while not done and ls_iter < max_ls: - # line-search bracket is so small - if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change: - break - - # compute new trial value - t = _cubic_interpolate(bracket[0], bracket_f[0], bracket_gtd[0], - bracket[1], bracket_f[1], bracket_gtd[1]) - - # test that we are making sufficient progress: - # in case `t` is so close to boundary, we mark that we are making - # insufficient progress, and if - # + we have made insufficient progress in the last step, or - # + `t` is at one of the boundary, - # we will move `t` to a position which is `0.1 * len(bracket)` - # away from the nearest boundary point. - eps = 0.1 * (max(bracket) - min(bracket)) - if min(max(bracket) - t, t - min(bracket)) < eps: - # interpolation close to boundary - if insuf_progress or t >= max(bracket) or t <= min(bracket): - # evaluate at 0.1 away from boundary - if abs(t - max(bracket)) < abs(t - min(bracket)): - t = max(bracket) - eps - else: - t = min(bracket) + eps - insuf_progress = False - else: - insuf_progress = True - else: - insuf_progress = False - - # Evaluate new point - f_new, g_new = obj_func(x, t, d) - ls_func_evals += 1 - gtd_new = g_new.dot(d) - ls_iter += 1 - - if f_new > (f + c1 * t * gtd) or f_new >= bracket_f[low_pos]: - # Armijo condition not satisfied or not lower than lowest point - bracket[high_pos] = t - bracket_f[high_pos] = f_new - bracket_g[high_pos] = g_new.clone(memory_format=torch.contiguous_format) - bracket_gtd[high_pos] = gtd_new - low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0) - else: - if abs(gtd_new) <= -c2 * gtd: - # Wolfe conditions satisfied - done = True - elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0: - # old high becomes new low - bracket[high_pos] = bracket[low_pos] - bracket_f[high_pos] = bracket_f[low_pos] - bracket_g[high_pos] = bracket_g[low_pos] - bracket_gtd[high_pos] = bracket_gtd[low_pos] - - # new point becomes new low - bracket[low_pos] = t - bracket_f[low_pos] = f_new - bracket_g[low_pos] = g_new.clone(memory_format=torch.contiguous_format) - bracket_gtd[low_pos] = gtd_new - - # return stuff - t = bracket[low_pos] - f_new = bracket_f[low_pos] - g_new = bracket_g[low_pos] - return f_new, g_new, t, ls_func_evals - - -class MyLBFGS(Optimizer): - """Implements L-BFGS algorithm, heavily inspired by `minFunc - `_. - .. warning:: - This optimizer doesn't support per-parameter options and parameter - groups (there can be only one). - .. warning:: - Right now all parameters have to be on a single device. This will be - improved in the future. - .. note:: - This is a very memory intensive optimizer (it requires additional - ``param_bytes * (history_size + 1)`` bytes). If it doesn't fit in memory - try reducing the history size, or use a different algorithm. - Args: - lr (float): learning rate (default: 1) - max_iter (int): maximal number of iterations per optimization step - (default: 20) - max_eval (int): maximal number of function evaluations per optimization - step (default: max_iter * 1.25). - tolerance_grad (float): termination tolerance on first order optimality - (default: 1e-5). - tolerance_change (float): termination tolerance on function - value/parameter changes (default: 1e-9). - history_size (int): update history size (default: 100). - line_search_fn (str): either 'strong_wolfe' or None (default: None). - """ - - def __init__(self, - params, - lr=1, - max_iter=20, - max_eval=None, - tolerance_grad=1e-7, - tolerance_change=1e-9, - history_size=100, - line_search_fn=None): - if max_eval is None: - max_eval = max_iter * 5 // 4 - defaults = dict( - lr=lr, - max_iter=max_iter, - max_eval=max_eval, - tolerance_grad=tolerance_grad, - tolerance_change=tolerance_change, - history_size=history_size, - line_search_fn=line_search_fn) - super(MyLBFGS, self).__init__(params, defaults) - - if len(self.param_groups) != 1: - raise ValueError("LBFGS doesn't support per-parameter options " - "(parameter groups)") - - self._params = self.param_groups[0]['params'] - self._numel_cache = None - - def _numel(self): - if self._numel_cache is None: - self._numel_cache = reduce(lambda total, p: total + p.numel(), self._params, 0) - return self._numel_cache - - def _gather_flat_grad(self): - views = [] - for p in self._params: - if p.grad is None: - view = p.new(p.numel()).zero_() - elif p.grad.is_sparse: - view = p.grad.to_dense().view(-1) - else: - view = p.grad.view(-1) - views.append(view) - return torch.cat(views, 0) - - def _add_grad(self, step_size, update): - offset = 0 - for p in self._params: - numel = p.numel() - if ((update.dtype==torch.complex64 or update.dtype==torch.complex128) - and (p.dtype==torch.float32 or p.dtype==torch.float64)): - p.add_(update[offset:offset + numel].real.view_as(p), alpha=step_size) - else: - p.add_(update[offset:offset + numel].view_as(p), alpha=step_size) - # view as to avoid deprecated pointwise semantics - #try: - # print(p.dtype) - # print(update.dtype) - # p.add_(update[offset:offset + numel].view_as(p), alpha=step_size) - # print('Worked fine') - # print(update[offset:offset + numel].view_as(p)) - #except: - # print('Failed') - # print(update[offset:offset + numel].view_as(p)) - # exit() - offset += numel - assert offset == self._numel() - - def _clone_param(self): - return [p.clone(memory_format=torch.contiguous_format) for p in self._params] - - def _set_param(self, params_data): - for p, pdata in zip(self._params, params_data): - p.copy_(pdata) - - def _directional_evaluate(self, closure, x, t, d): - self._add_grad(t, d) - loss = float(closure()) - flat_grad = self._gather_flat_grad() - self._set_param(x) - return loss, flat_grad - - @torch.no_grad() - def step(self, closure): - """Performs a single optimization step. - Args: - closure (callable): A closure that reevaluates the model - and returns the loss. - """ - assert len(self.param_groups) == 1 - - # Make sure the closure is always called with grad enabled - closure = torch.enable_grad()(closure) - - group = self.param_groups[0] - lr = group['lr'] - max_iter = group['max_iter'] - max_eval = group['max_eval'] - tolerance_grad = group['tolerance_grad'] - tolerance_change = group['tolerance_change'] - line_search_fn = group['line_search_fn'] - history_size = group['history_size'] - - # NOTE: LBFGS has only global state, but we register it as state for - # the first param, because this helps with casting in load_state_dict - state = self.state[self._params[0]] - state.setdefault('func_evals', 0) - state.setdefault('n_iter', 0) - - # evaluate initial f(x) and df/dx - orig_loss = closure() - loss = float(orig_loss) - current_evals = 1 - state['func_evals'] += 1 - - flat_grad = self._gather_flat_grad() - opt_cond = flat_grad.abs().max() <= tolerance_grad - - # optimal condition - if opt_cond: - return orig_loss - - # tensors cached in state (for tracing) - d = state.get('d') - t = state.get('t') - old_dirs = state.get('old_dirs') - old_stps = state.get('old_stps') - ro = state.get('ro') - H_diag = state.get('H_diag') - prev_flat_grad = state.get('prev_flat_grad') - prev_loss = state.get('prev_loss') - - n_iter = 0 - # optimize for a max of max_iter iterations - while n_iter < max_iter: - # keep track of nb of iterations - n_iter += 1 - state['n_iter'] += 1 - - ############################################################ - # compute gradient descent direction - ############################################################ - if state['n_iter'] == 1: - d = flat_grad.neg() - old_dirs = [] - old_stps = [] - ro = [] - H_diag = 1 - else: - # do lbfgs update (update memory) - y = flat_grad.sub(prev_flat_grad) - s = d.mul(t) - ys = y.dot(s) # y*s - if ys.abs() > 1e-10: - # updating memory - if len(old_dirs) == history_size: - # shift history by one (limited-memory) - old_dirs.pop(0) - old_stps.pop(0) - ro.pop(0) - - # store new direction/step - old_dirs.append(y) - old_stps.append(s) - ro.append(1. / ys) - - # update scale of initial Hessian approximation - H_diag = ys / y.dot(y) # (y*y) - - # compute the approximate (L-BFGS) inverse Hessian - # multiplied by the gradient - num_old = len(old_dirs) - - if 'al' not in state: - state['al'] = [None] * history_size - al = state['al'] - - # iteration in L-BFGS loop collapsed to use just one buffer - q = flat_grad.neg() - for i in range(num_old - 1, -1, -1): - al[i] = old_stps[i].dot(q) * ro[i] - q.add_(old_dirs[i], alpha=-al[i]) - - # multiply by initial Hessian - # r/d is the final direction - d = r = torch.mul(q, H_diag) - for i in range(num_old): - be_i = old_dirs[i].dot(r) * ro[i] - r.add_(old_stps[i], alpha=al[i] - be_i) - - if prev_flat_grad is None: - prev_flat_grad = flat_grad.clone(memory_format=torch.contiguous_format) - else: - prev_flat_grad.copy_(flat_grad) - prev_loss = loss - - ############################################################ - # compute step length - ############################################################ - # reset initial guess for step size - if state['n_iter'] == 1: - t = min(1., 1. / flat_grad.abs().sum()) * lr - else: - t = lr - - # directional derivative - gtd = flat_grad.dot(d) # g * d - # Why did this used to be gtd > -tol_change? - - # directional derivative is below tolerance - if gtd.abs() < tolerance_change: - break - - # optional line search: user function - ls_func_evals = 0 - if line_search_fn is not None: - # perform line search, using user function - if line_search_fn != "strong_wolfe": - raise RuntimeError("only 'strong_wolfe' is supported") - else: - x_init = self._clone_param() - - def obj_func(x, t, d): - return self._directional_evaluate(closure, x, t, d) - - loss, flat_grad, t, ls_func_evals = _strong_wolfe( - obj_func, x_init, t, d, loss, flat_grad, gtd) - self._add_grad(t, d) - opt_cond = flat_grad.abs().max() <= tolerance_grad - else: - # no line search, simply move with fixed-step - self._add_grad(t, d) - if n_iter != max_iter: - # re-evaluate function only if not in last iteration - # the reason we do this: in a stochastic setting, - # no use to re-evaluate that function here - with torch.enable_grad(): - loss = float(closure()) - flat_grad = self._gather_flat_grad() - opt_cond = flat_grad.abs().max() <= tolerance_grad - ls_func_evals = 1 - - # update func eval - current_evals += ls_func_evals - state['func_evals'] += ls_func_evals - - ############################################################ - # check conditions - ############################################################ - if n_iter == max_iter: - break - - if current_evals >= max_eval: - break - - # optimal condition - if opt_cond: - break - - # lack of progress - if d.mul(t).abs().max() <= tolerance_change: - break - - if abs(loss - prev_loss) < tolerance_change: - break - - state['d'] = d - state['t'] = t - state['old_dirs'] = old_dirs - state['old_stps'] = old_stps - state['ro'] = ro - state['H_diag'] = H_diag - state['prev_flat_grad'] = prev_flat_grad - state['prev_loss'] = prev_loss - - return orig_loss diff --git a/src/cdtools/models/fastccd_ptycho.py b/src/cdtools/models/fastccd_ptycho.py deleted file mode 100644 index 59f8e3b..0000000 --- a/src/cdtools/models/fastccd_ptycho.py +++ /dev/null @@ -1,691 +0,0 @@ -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, - 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, - 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): - """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 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(self.probe.detach()) - self.probe.data = ortho_probes - return - - # 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) - - # 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 - ) - - # 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) - - # 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) - - - 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=-3) - 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 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), - ('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/multimode_rpi.py b/src/cdtools/models/multimode_rpi.py deleted file mode 100644 index c022154..0000000 --- a/src/cdtools/models/multimode_rpi.py +++ /dev/null @@ -1,354 +0,0 @@ -import torch as t -from cdtools.models import CDIModel -from cdtools import tools -from cdtools.tools import plotting as p -from cdtools.tools.interactions import RPI_interaction -from cdtools.tools import initializers -from scipy.ndimage import binary_dilation -import numpy as np -from copy import copy - -__all__ = ['MultimodeRPI'] - -class MultimodeRPI(CDIModel): - - @property - def obj(self): - return t.complex(self.obj_real, self.obj_imag) - - @property - def weights(self): - ws = t.complex(self.weights_real, self.weights_imag) - return ws / 10# / self.obj_real.size().numel() - - def __init__(self, wavelength, detector_geometry, probe_basis, - probe, obj_guess, detector_slice=None, - background=None, mask=None, saturation=None, - obj_support=None, oversampling=1, weight_matrix=False): - - super(MultimodeRPI, self).__init__() - - self.wavelength = t.tensor(wavelength) - self.detector_geometry = copy(detector_geometry) - - det_geo = self.detector_geometry - if hasattr(det_geo, 'distance'): - det_geo['distance'] = t.tensor(det_geo['distance']) - if hasattr(det_geo, 'basis'): - det_geo['basis'] = t.tensor(det_geo['basis']) - if hasattr(det_geo, 'corner'): - det_geo['corner'] = t.tensor(det_geo['corner']) - - self.probe_basis = t.tensor(probe_basis) - - scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1], - probe.shape[-2]/obj_guess.shape[-2]]) - self.obj_basis = self.probe_basis * scale_factor - self.detector_slice = detector_slice - - # Maybe something to include in a bit - # self.surface_normal = t.tensor(surface_normal) - - self.saturation = saturation - - if mask is None: - self.mask = mask - else: - self.mask = t.tensor(mask, dtype=t.bool) - - - self.probe = t.tensor(probe, dtype=t.complex64) - - obj_guess = t.tensor(obj_guess, dtype=t.complex64) - - self.obj_real = t.nn.Parameter(obj_guess.real) - self.obj_imag = t.nn.Parameter(obj_guess.imag) - - self.weights_real = t.nn.Parameter(t.eye(probe.shape[0])* 10)# * self.obj_real.size().numel()) - self.weights_imag = t.nn.Parameter(t.zeros(probe.shape[0])) - - if not weight_matrix: - self.weights_real.requires_grad=False - self.weights_imag.requires_grad=False - - # Wait for LBFGS to be updated for complex-valued parameters - # self.obj = t.nn.Parameter(obj_guess.to(t.float32)) - - if background is None: - if detector_slice is not None: - background = 1e-6 * t.ones( - self.probe[0][self.detector_slice].shape, - dtype=t.float32) - else: - background = 1e-6 * t.ones(self.probe[0].shape, - dtype=t.float32) - - self.background = t.tensor(background, dtype=t.float32) - - if obj_support is not None: - self.obj_support = obj_support - self.obj.data = self.obj * obj_support[None, ...] - else: - self.obj_support = t.ones_like(self.obj[0, ...]) - - self.oversampling = oversampling - - - @classmethod - def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', weight_matrix=False, probe_threshold=0): - raise NotImplementedError() - - 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 only need the patterns here, not the inputs associated with them. - _, 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, - oversampling=oversampling) - - if not isinstance(probe,t.Tensor): - probe = t.as_tensor(probe) - - # Potentially need all of this orientation stuff later - - #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) - - - if background is None and hasattr(dataset, 'background') \ - and dataset.background is not None: - background = t.sqrt(dataset.background) - elif background is not None: - background = t.sqrt(t.Tensor(background).to(dtype=t.float32)) - - det_geo = dataset.detector_geometry - - # If no mask is given, but one exists in the dataset, load it. - if mask is None and hasattr(dataset, 'mask') \ - and dataset.mask is not None: - mask = dataset.mask.to(t.bool) - - # Now we initialize the object - if obj_size is None: - # This is a standard size for a well-matched probe and detector - obj_size = (np.array(probe_shape) // 2).astype(int) - - if initialization.lower().strip() == 'random': - # I think something to do with the fact that the object is defined - # on a coarser grid needs to be accounted for here that is not - # accounted for yet - scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2) - obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size)) - elif initialization.lower().strip() == 'spectral': - if background is not None: - obj_guess = initializers.RPI_spectral_init( - patterns[0], probe, obj_size, mask=mask, - background=background**2, n_modes=n_modes) - else: - obj_guess = initializers.RPI_spectral_init( - patterns[0], probe, obj_size, mask=mask, - n_modes=n_modes) - - else: - raise KeyError('Initialization "' + str(initialization) + \ - '" invalid - use "spectral" or "random"') - - probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0)) - probe_fft = tools.propagators.far_field(probe_intensity) - pad0l = (probe.shape[-2] - obj_size[-2])//2 - pad0r = probe.shape[-2] - obj_size[-2] - pad0l - pad1l = (probe.shape[-1] - obj_size[-1])//2 - pad1r = probe.shape[-1] - obj_size[-1] - pad1l - probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r] - probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft)) - - obj_support = probe_lr > t.max(probe_lr) * probe_threshold - obj_support = t.as_tensor(binary_dilation(obj_support)) - - return cls(wavelength, det_geo, probe_basis, - probe, obj_guess, detector_slice=det_slice, - background=background, mask=mask, saturation=saturation, - obj_support=obj_support, oversampling=oversampling, - weight_matrix=weight_matrix) - - - def random_init(self, pattern): - scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2) - self.obj.data = scale * t.exp( - 2j * np.pi * t.rand(self.obj.shape)).to( - dtype=self.obj.dtype, device=self.obj.device) - - def spectral_init(self, pattern): - if self.background is not None: - self.obj.data = initializers.RPI_spectral_init( - pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask, - background=self.background**2, n_modes=self.obj.shape[0]).to( - dtype=self.obj.dtype, device=self.obj.device) - else: - self.obj.data = initializers.RPI_spectral_init( - pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask, - n_modes=self.obj.shape[0]).to( - dtype=self.obj.dtype, device=self.obj.device) - - # Needs work - def interaction(self, index, *args): - # including *args allows this to work with all sorts of datasets - # that might include other information in with the index in their - # "input" parameters (such as translations for a ptychography dataset). - # This makes it seamless to use such a dataset even though those - # extra arguments will not be used. - - - all_exit_waves = [] - - # Mix the probes with the weight matrix - prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3) - - for i in range(self.probe.shape[0]): - pr = prs[i] - # Here we have a 3D probe (one single mode) - # and a 4D object (multiple modes mixing incoherently) - exit_waves = RPI_interaction(pr, - self.obj_support * self.obj[i]) - all_exit_waves.append(exit_waves.unsqueeze(0)) - - # This creates a bunch of modes generated from all possible combos - # of the probe and object modes all strung out along the first index - - output = t.cat(all_exit_waves) - - # If we have multiple indexes input, we unsqueeze and repeat the stack - # of wavefields enough times to simulate each requested index. This - # seems silly, but it enables (for example) one to do a reconstruction - # from a set of diffraction patterns that are all known to be from the - # same object. - try: - # will fail if index has no length, for example when index - # is just an int. In this case, we just do nothing instead - output = output.unsqueeze(0).repeat(1,len(index),1,1,1) - except TypeError: - pass - return output - - - 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): - # Here I'm taking advantage of an undocumented feature in the - # incoherent_sum measurement function where it will work with - # a 4D wavefield array as well as a 5D array. - return tools.measurements.quadratic_background(wavefields, - self.background, - detector_slice=self.detector_slice, - measurement=tools.measurements.incoherent_sum, - saturation=self.saturation, - oversampling=self.oversampling) - - def loss(self, sim_data, real_data, mask=None): - return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) - #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) - - def regularizer(self, factors): - return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \ - + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2) - - def to(self, *args, **kwargs): - super(MultimodeRPI, self).to(*args, **kwargs) - self.wavelength = self.wavelength.to(*args,**kwargs) - # move the detector geometry too - det_geo = self.detector_geometry - if hasattr(det_geo, 'distance'): - det_geo['distance'] = det_geo['distance'].to(*args,**kwargs) - if hasattr(det_geo, 'basis'): - det_geo['basis'] = det_geo['basis'].to(*args,**kwargs) - if hasattr(det_geo, 'corner'): - det_geo['corner'] = det_geo['corner'].to(*args,**kwargs) - - if self.mask is not None: - self.mask = self.mask.to(*args, **kwargs) - - self.probe = self.probe.to(*args,**kwargs) - self.probe_basis = self.probe_basis.to(*args,**kwargs) - self.obj_basis = self.obj_basis.to(*args,**kwargs) - self.obj_support = self.obj_support.to(*args,**kwargs) - self.background = self.background.to(*args, **kwargs) - - # Maybe include in a bit - #self.surface_normal = self.surface_normal.to(*args, **kwargs) - - def sim_to_dataset(self, args_list): - raise NotImplementedError('No sim to dataset yet, sorry!') - - plot_list = [ - ('Root Sum Squared Amplitude of all Probes', - lambda self, fig: p.plot_amplitude( - np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)), - fig=fig, basis=self.probe_basis)), - ('Object Amplitudes', - lambda self, fig: p.plot_amplitude(self.obj, fig=fig, - basis=self.obj_basis)), - ('Object Phases', - lambda self, fig: p.plot_phase(self.obj, fig=fig, - basis=self.obj_basis)) - ] - - - def save_results(self, dataset=None, full_obj=False): - # dataset is set as a kwarg here because it isn't needed, but the - # common pattern is to pass a dataset. This makes it okay if one - # continues to use that standard pattern - probe_basis = self.probe_basis.detach().cpu().numpy() - obj_basis = self.obj_basis.detach().cpu().numpy() - probe = self.probe.detach().cpu().numpy() - # Provide the option to save out the subdominant objects or - # just the dominant one - if full_obj: - obj = self.obj.detach().cpu().numpy() - else: - obj = self.obj[0].detach().cpu().numpy() - background = self.background.detach().cpu().numpy()**2 - - return {'probe_basis': probe_basis, 'obj_basis': obj_basis, - 'probe': probe,'obj': obj, - 'background': background} - diff --git a/src/cdtools/models/polarization_swept_ptycho.py b/src/cdtools/models/polarization_swept_ptycho.py deleted file mode 100644 index 8e9594b..0000000 --- a/src/cdtools/models/polarization_swept_ptycho.py +++ /dev/null @@ -1,605 +0,0 @@ -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__ = ['FancyPtycho'] - -class PolarizationSweptPtycho(CDIModel): - - def __init__(self, wavelength, detector_geometry, - probe_basis, - probe_guess, - obj_guess, - polarization_states, - detector_slice=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, - dtype=t.float32, - obj_view_crop=0 - ): - - super(PolarizationSweptPtycho, self).__init__() - self.register_buffer('wavelength', - t.tensor(wavelength, dtype=dtype)) - self.store_detector_geometry(detector_geometry, - dtype=dtype) - - self.register_buffer('min_translation', - t.tensor(min_translation, dtype=dtype)) - - self.register_buffer('probe_basis', - t.tensor(probe_basis, dtype=dtype)) - - self.detector_slice = copy(detector_slice) - self.register_buffer('surface_normal', - t.tensor(surface_normal, dtype=dtype)) - if saturation is None: - self.saturation = None - else: - self.register_buffer('saturation', - t.tensor(saturation, dtype=dtype)) - # Not sure how to make this a buffer... - self.units = units - - self.fourier_probe = fourier_probe - - if mask is None: - self.mask = None - else: - self.register_buffer('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: - probe_norm = 1 * t.max(t.abs(probe_guess[0])) - else: - probe_norm = 1 * t.max(t.abs(probe_guess)) - self.register_buffer('probe_norm', probe_norm.to(dtype)) - - self.probe = t.nn.Parameter(probe_guess / self.probe_norm) - self.obj = t.nn.Parameter(obj_guess) - - - self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop, - obj_view_crop:-obj_view_crop] - - if background is None: - if detector_slice is not None: - dummy_det = t.empty([s//oversampling - for s in self.probe.shape[-2:]]) - shape = dummy_det[self.detector_slice].shape - #shape = self.probe[0][self.detector_slice].shape - else: - shape = [s//oversampling for s in self.probe.shape[-2:]] - background = 1e-6 * t.ones(shape, dtype=t.float32) - - self.background = t.nn.Parameter(background) - - if weights is None: - self.weights = None - else: - self.weights = t.nn.Parameter(t.tensor(weights, - dtype=t.float32)) - - 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.register_buffer('translation_scale', - t.tensor(translation_scale, dtype=dtype)) - - if probe_support is None: - probe_support = t.ones(self.probe.shape[-2:], dtype=t.bool) - self.register_buffer('probe_support', - t.tensor(probe_support, 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=dtype) - Js = t.arange(self.probe.shape[-1], dtype=dtype) - Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js)) - - I_phase = 2 * np.pi* Is * self.oversampling - J_phase = 2 * np.pi* Js * self.oversampling - self.register_buffer('I_phase', I_phase) - self.register_buffer('J_phase', J_phase) - - - self.simulate_finite_pixels = simulate_finite_pixels - - self.polarization_states = t.nn.Parameter( - t.tensor(polarization_states, dtype=t.complex64)) - # by default, don't optimize this, but I think it might be - # interesting to try it because the polarization states - # are not very pure - self.polarization_states.requires_grad = False - - - # 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, - translation_scale=1, - saturation=None, - probe_support_radius=None, - probe_fourier_crop=None, - propagation_distance=None, - scattering_mode=None, - oversampling=1, - auto_center=False, - fourier_probe=False, - loss='amplitude mse', - units='um', - simulate_probe_translation=False, - simulate_finite_pixels=False, - obj_view_crop=None, - obj_padding=200 - ): - - 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[:] - polarization_states = dataset.polarization_states - - 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, - 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=obj_padding) - - 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) - - if probe_fourier_crop is not None: - probe = tools.propagators.far_field(probe) - probe = probe[probe_fourier_crop:-probe_fourier_crop, - probe_fourier_crop:-probe_fourier_crop] - probe = tools.propagators.inverse_far_field(probe) - - # 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) - # The probe gets one extra dimension, so we can simulate one probe - # per polarization state - n_states = polarization_states.shape[0] - probe = t.stack([probe] * n_states, dim=0) - - # Looks like an identity matrix - obj_base = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) - obj_top = t.stack([obj_base, obj_base*0], dim=0) - obj_bottom = t.stack([obj_base*0, obj_base], dim=0) - obj = t.stack([obj_top, obj_bottom], dim=0) - - if obj_view_crop is None: - obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2 - if obj_view_crop < 0: - obj_view_crop += min(probe.shape[-2], probe.shape[-1]) // 2 - obj_view_crop += obj_padding - - det_geo = dataset.detector_geometry - - translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5) - - 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.shape[-2:], 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, - polarization_states=polarization_states, - 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, - obj_view_crop=obj_view_crop) - - - def interaction(self, index, translations, polarization_indices, *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[..., :, :] - # Now we expand the polarization states explicitly. - # pol_basis_prs has dimensions n_states x 2 x n_modes x m x l - pol_basis_prs = ( - self.polarization_states[...,None,None,None] - * basis_prs[...,None,:,:,:]) - # For a Fourier-space probe, we take an IFT - if self.fourier_probe: - basis_prs = tools.propagators.inverse_far_field(basis_prs) - - # So, first we get the appropriate probe for each shot. - # prs is now n_frames x 2 x n_modes x m x l - prs = pol_basis_prs[polarization_indices] - - # Now we construct the probes for each shot from the basis probes - if self.weights is not None: - Ws = self.weights[index] - # And then we multiply by the weights, along the 0th dimension - prs = Ws[..., None, None, None, None] * prs#basis_prs - - if self.simulate_probe_translation: - det_pix_trans = tools.interactions.translations_to_pixel( - self.det_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,None,:,:] - - # We automatically rescale the probe to match the background size, - # which allows us to do stuff like let the object be super-resolution, - # while restricting the probe to the detector resolution but still - # doing an explicit real-space limitation of the probe - padding = [self.oversampling * self.background.shape[-2] - prs.shape[-2], - self.oversampling * self.background.shape[-1] - prs.shape[-1]] - - if any([p != 0 for p in padding]): # For probe_fourier_crop != 0. - padding = [padding[-1]//2, padding[-1]-padding[-1]//2, - padding[-2]//2, padding[-2]-padding[-2]//2] - prs = tools.propagators.far_field(prs) - prs = t.nn.functional.pad(prs, padding) - prs = tools.propagators.inverse_far_field(prs) - # Now we actually do the interaction, using the sinc subpixel - # translation model as per usual - - # After the object, no point in treating the polarization modes - # any differently from the object/probe modes - # So, this stacks up all the output modes in lexicographic order - exit_waves = t.cat( - [self.probe_norm * tools.interactions.ptycho_2D_sinc( - prs[...,0,:,:,:], self.obj[0,0], pix_trans, - shift_probe=True, multiple_modes=True), - self.probe_norm * tools.interactions.ptycho_2D_sinc( - prs[...,0,:,:,:], self.obj[1,0], pix_trans, - shift_probe=True, multiple_modes=True)], - dim=1) - exit_waves = exit_waves + t.cat( - [self.probe_norm * tools.interactions.ptycho_2D_sinc( - prs[...,1,:,:,:], self.obj[0,1], pix_trans, - shift_probe=True, multiple_modes=True), - self.probe_norm * tools.interactions.ptycho_2D_sinc( - prs[...,1,:,:,:], self.obj[1,1], pix_trans, - shift_probe=True, multiple_modes=True)], - dim=1) - - 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): - return 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) - - - # Note: No "loss" function is defined here, because it is added - # dynamically during object creation in __init__ - - 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} - - - 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=self.get_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 - """ - # 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]) - - - - plot_list = [ - ('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[np.s_[:,:] + self.obj_view_slice], - fig=fig, - basis=self.probe_basis, - units=self.units)), - ('Object Phase', - lambda self, fig: p.plot_phase( - self.obj[np.s_[:,:] + self.obj_view_slice], - 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)) - ] - -# def plot_errors(self, dataset): - - - - def save_results(self, dataset): - 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 - wavelength = self.wavelength.cpu().numpy() - - return {'basis': basis, 'translation': translations, - 'probe': probe, 'obj': obj, - 'background': background, - 'oversampling': oversampling, - 'weights': weights, 'wavelength': wavelength} diff --git a/src/cdtools/models/polarized_fancy_ptycho.py b/src/cdtools/models/polarized_fancy_ptycho.py deleted file mode 100644 index a9ad37b..0000000 --- a/src/cdtools/models/polarized_fancy_ptycho.py +++ /dev/null @@ -1,556 +0,0 @@ -import torch as t -from cdtools.models import CDIModel, FancyPtycho -from cdtools.datasets import Ptycho2DDataset -from cdtools import tools -from cdtools.tools import plotting as p -# from cdtools.tools import polarized_plotting as pp -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 -from cdtools.tools import polarization - -__all__ = ['PolarizedFancyPtycho'] - -class PolarizedFancyPtycho(FancyPtycho): - - def __init__(self, wavelength, detector_geometry, - probe_basis, - probe_guess, obj_guess, polarizer, analyzer, - detector_slice=None, - surface_normal=np.array([0.,0.,1.]), - min_translation = t.Tensor([0,0]), - background = None, translation_offsets=None, - polarizer_offsets=None, analyzer_offsets=None, - polarizer_scale=1, analyzer_scale=1, mask=None, - weights = None, translation_scale = 1, saturation=None, - probe_support = None, oversampling=1, - loss='amplitude mse',units='um'): - - super(PolarizedFancyPtycho, self).__init__(wavelength, detector_geometry, - probe_basis, - probe_guess, obj_guess, - detector_slice=None, - surface_normal=np.array([0.,0.,1.]), - min_translation = t.Tensor([0,0]), - background = None, translation_offsets=None, mask=None, - weights = weights, translation_scale = 1, saturation=None, - probe_support = None, oversampling=1, - loss='amplitude mse',units='um') - - if polarizer_offsets is None: - self.polarizer_offsets = None - else: - self.polarizer_offsets = t.nn.Parameter(t.tensor(polarizer_offsets).to(dtype=t.float32)) / polarizer_scale - - if analyzer_offsets is None: - self.analyzer_offsets = None - else: - self.analyzer_offsets = t.nn.Parameter(t.tensor(analyzer_offsets).to(dtype=t.float32)) / analyzer_scale - - self.polarizer = polarizer - self.analyzer = analyzer - probe_guess = t.tensor(probe_guess, dtype=t.complex64) - if probe_guess.dim() > 4: - 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) - - @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, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', left_polarized=True): - - # When using this method, remember to pass through the inputs - model = FancyPtycho.from_dataset( - dataset, - probe_size=probe_size, - randomize_ang=randomize_ang, - padding=padding, - n_modes=n_modes, - dm_rank=dm_rank, - translation_scale=translation_scale, - saturation=saturation, - probe_support_radius=probe_support_radius, - propagation_distance=propagation_distance, - scattering_mode=scattering_mode, - oversampling=oversampling, - auto_center=auto_center, - loss=loss, - units=units) - - - # Mutate the class to its subclass - model.__class__ = cls - - if left_polarized: - x = 1j - else: - x = -1j - - probe = model.probe.detach() - probe = t.cat((probe, probe * x), dim=-3) - 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)] - probe = t.stack([probe, ] + probe_stack) - #print('probe', type(probe), probe.shape) - model.probe.data = probe - #print(model.probe.shape) - # obj = t.stack((model.obj.data, model.obj.data), dim=-3) - # model.obj.data = t.stack((obj.data, obj.data), dim=-4) - # obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) - obj = model.obj.detach() - # Abe - Probably something identity matrix-like would be a better - # initialization (e.g. ((obj,0*obj),(0*obj,obj)) - obj = t.stack((obj, obj), dim=-3) - obj = t.stack((obj, obj), dim=-4) - #print('object', type(obj), obj.shape) - model.obj.data = obj - #print('polarized fancy ptycho from datset obj') - a = obj.detach() - #plt.imshow(np.real(a[0, 0, :, :])) - #plt.figure() - #plt.imshow(np.real(a[0, 1, :, :])) - #plt.show() - - # tensor vs tensor.data - return model - - polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)] - - @classmethod - def from_dataset2(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, loss='amplitude mse', units='um'): - - 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, polarizer, analyzer), 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 - - if left_polarized: - x = 1j - else: - x = -1j - - - # 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, - oversampling=oversampling) - - probe_shape = t.stack((2, probe_shape), dim=-3) - 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)] - probe = t.stack([probe, ] + probe_stack) - # probe = t.stack([tools.propagators.far_field(probe),] + probe_stack) - probe_x, probe_y = probe, probe * x - probe = t.stact((probe_x, probe_y), dim=-3) - - a = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) - b = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) - c = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) - d = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) - - ab = t.stack((a, b), dim=-3) - cd = t.stack((c, d), dim=-3) - obj = t.stack((ab, cd), dim=-4) - 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, '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, - oversampling=oversampling, - loss=loss, units=units) - - - - def interaction(self, index, translations, polarizer, analyzer, test=False): - - # 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[...,:,:] # This makes no sense - # self.probe is an Nx2xXxY stach of probes - - # Now we construct the probes for each shot from the basis probes - Ws = self.weights[index] - if len(self.weights[0].shape) == 0: - # If a purely stable coherent illumination is defined - # Ws is a tensor of length M, M is the number of frames to be processed - prs = Ws[...,None,None,None,None] * basis_prs - else: - raise NotImplementedError('Unstable Modes not Implemented for polarized light') - - pol_probes = polarization.apply_linear_polarizer(prs, polarizer) - exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( - pol_probes, self.obj, pix_trans, - shift_probe=True, multiple_modes=True, polarized=True) - # We're losing some efficiency here, because we only need to keep - # around the scalar wavefield after analyzing the waves. - # But I think it's not a huge issue - Abe - analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer) - return analyzed_exit_waves - - - def vectorial_wavefields(wavefields, func, *args, **kwargs): - wavefields_x = wavefields[..., 0, :, :, :] - wavefields_y = wavefields[..., 1, :, :, :] - out_x = func(wavefields_x, *args, **kwargs) - out_y = func(wavefields_y, *args, **kwargs) - out = t.stack((out_x, out_y), dim=-4) - return out[..., None, :, :] - - 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): - wavefields_x = wavefields[..., 0, :, :] - wavefields_y = wavefields[..., 1, :, :] - out_x = tools.measurements.quadratic_background(wavefields_x, - self.background, - detector_slice=self.detector_slice, - measurement=tools.measurements.incoherent_sum, - saturation=self.saturation, - oversampling=self.oversampling) - # now, set bckgr to 0 since t shouldn't be calculated twice - out_y = tools.measurements.quadratic_background(wavefields_y, - 0, - detector_slice=self.detector_slice, - measurement=tools.measurements.incoherent_sum, - saturation=self.saturation, - oversampling=self.oversampling) - - return out_x + out_y - - - # Note: No "loss" function is defined here, because it is added - # dynamically during object creation in __init__ - - def to(self, *args, **kwargs): - super(PolarizedFancyPtycho, self).to(*args, **kwargs) - - - def sim_to_dataset(self, args_list): - # 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 - - # Then we simulate the results - data = self.forward(indices, translations) - - # 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) - 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 - - - - 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 - - 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 - # 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 Amplitudes (scroll to view modes)', - lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)), - ('Basis Probe Phases (scroll to view modes)', - lambda self, fig: p.plot_phase(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)) - ] - - - def save_results(self, dataset): - 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() - - return {'basis':basis, 'translation':translations, - 'probe':probe,'obj':obj, - 'background':background, - 'weights':weights} diff --git a/src/cdtools/models/s_matrix_ptycho.py b/src/cdtools/models/s_matrix_ptycho.py deleted file mode 100644 index adb7b46..0000000 --- a/src/cdtools/models/s_matrix_ptycho.py +++ /dev/null @@ -1,429 +0,0 @@ -import torch as t -from cdtools.models import CDIModel -from cdtools.datasets import Ptycho2DDataset -from cdtools import tools -from cdtools.tools import cmath -from cdtools.tools import plotting as p -from matplotlib import pyplot as plt -from datetime import datetime -import numpy as np -from copy import copy - -__all__ = ['SMatrixPtycho'] - -class SMatrixPtycho(CDIModel): - - def __init__(self, wavelength, detector_geometry, - probe_basis, probe_guess, probe_fourier_support, - s_matrix_guess, - detector_slice=None, - surface_normal=np.array([0.,0.,1.]), - min_translation = t.Tensor([0,0]), - background = None, translation_offsets=None, - probe_planes = None, mask=None, - weights = None, translation_scale = 1, saturation=None, - oversampling=1): - - super(SMatrixPtycho,self).__init__() - self.wavelength = t.Tensor([wavelength]) - self.detector_geometry = copy(detector_geometry) - det_geo = self.detector_geometry - if hasattr(det_geo, 'distance'): - det_geo['distance'] = t.Tensor(det_geo['distance']) - if hasattr(det_geo, 'basis'): - det_geo['basis'] = t.Tensor(det_geo['basis']) - if hasattr(det_geo, 'corner'): - det_geo['corner'] = t.Tensor(det_geo['corner']) - - self.min_translation = t.Tensor(min_translation) - - self.probe_basis = t.Tensor(probe_basis) - self.detector_slice = detector_slice - self.surface_normal = t.Tensor(surface_normal) - - self.saturation = saturation - - if mask is None: - self.mask = mask - else: - self.mask = t.BoolTensor(mask) - - # We rescale the probe here so it learns at the same rate as the - # object - # Remember that for S-matrix we have several probes for different - # planes - if probe_guess.dim() > 4: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0,0].to(t.float32))) - else: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) - - self.probe = t.nn.Parameter(probe_guess.to(t.float32) - / self.probe_norm) - - self.s_matrix = t.nn.Parameter(s_matrix_guess.to(t.float32)) - - if background is None: - ew_shape = [s_matrix_guess.shape[0] - 1 + probe_guess.shape[-3], - s_matrix_guess.shape[1] - 1 + probe_guess.shape[-2]] - if detector_slice is not None: - background = 1e-6 * t.ones(t.ones(ew_shape)[self.detector_slice].shape).to(t.float32) - else: - background = 1e-6 * t.ones(ew_shape).to(t.float32) - - - self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) - - if weights is None: - self.weights = None - else: - self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32)) - - if translation_offsets is None: - self.translation_offsets = None - else: - self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale) - - # This maps indices to probe planes to be used. If none, it defaults - # to always being plane 0 - if probe_planes is None: - self.probe_planes = None - else: - self.probe_planes = t.LongTensor(probe_planes) - - self.translation_scale = translation_scale - - self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32) - - self.oversampling = oversampling - - - @classmethod - def from_dataset(cls, dataset, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1): - - datasets = [dataset] - propagation_distances = [propagation_distance] - - # We only return the 0th element because in the general case, the - # constructor needs to return a stacked datset in addition to - # a model, but for the case of one dataset we only need to return - # the model. - return cls.from_datasets(datasets, probe_convergence_radius, - locality_radius=locality_radius, - probe_size=probe_size, - randomize_ang=randomize_ang, - padding=padding, - n_modes=n_modes, - translation_scale=translation_scale, - saturation=saturation, - propagation_distances=propagation_distances, - scattering_mode=scattering_mode, - oversampling=oversampling)[0] - - # This is for the multi-focal-plane case, where each dataset will correspond - # to a different focal plane. The guess propagation distance for each - # dataset can be set individually but otherwise the probes are - # reconstructed entirely separately. All datasets are assumed to have - # the same basic parameters (wavelength, detector geometry, etc) and share - # the same origin in the x-y plane. - @classmethod - def from_datasets(cls, datasets, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distances=None, scattering_mode=None, oversampling=1): - - wavelength = datasets[0].wavelength - det_basis = datasets[0].detector_geometry['basis'] - det_shape = datasets[0][0][1].shape - distance = datasets[0].detector_geometry['distance'] - - # Then, generate the probe geometry from the dataset - ewg = tools.initializers.exit_wave_geometry - probe_basis, ew_shape, det_slice = ewg(det_basis, - det_shape, - wavelength, - distance, - padding=padding, - oversampling=oversampling) - - if propagation_distances is None: - propagation_distances = [None] * len(datasets) - - # This shrinks the probe to ensure that the output wavefield - # is the correct shape - probe_shape = t.Size(np.array(ew_shape) - (2*locality_radius)) - - # always do this on the cpu - probe_planes = [] - translations = [] - patterns = [] - probes = [] - for i, dataset in enumerate(datasets): - get_as_args = dataset.get_as_args - - dataset.get_as(device='cpu') - (indices, tx), pats = dataset[:] - dataset.get_as(*get_as_args[0],**get_as_args[1]) - translations.append(tx) - probe_planes.extend([i]*tx.shape[0]) - patterns.append(pats) - - # Finally, initialize the probe and object using this information - if locality_radius != 0: - probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling)[locality_radius:-locality_radius,locality_radius:-locality_radius] - else: - probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling) - - # Now we initialize all the subdominant probe modes - probe_max = t.max(cmath.cabs(probe)) - probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)] - probe = t.stack([tools.propagators.inverse_far_field(probe),] + probe_stack) - probes.append(probe) - - translations = t.cat(translations) - patterns = t.cat(patterns) - probes = t.stack(probes) - - - if hasattr(datasets[0], 'sample_info') and \ - datasets[0].sample_info is not None and \ - 'orientation' in datasets[0].sample_info: - surface_normal = datasets[0].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(outgoing_dir) - - - # 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) - - # The locality radius correction is probably not needed because - # it will always be way less than 200, but it ensures that there - # is no wrapping in the s-matrix - obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200+2*locality_radius) - - if hasattr(dataset, 'background') and dataset.background is not None: - background = t.sqrt(datasets[0].background) - else: - background = None - - - - s_matrix = t.zeros([2*locality_radius+1,2*locality_radius+1,obj_size[0], - obj_size[1],2]) - s_matrix[locality_radius,locality_radius,:,:,:] = \ - tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) - - - - det_geo = dataset.detector_geometry - - translation_offsets = 0 * (t.rand((translations.shape[0],2)) - 0.5) - - weights = t.ones(translations.shape[0]) - - if hasattr(datasets[0], 'mask') and datasets[0].mask is not None: - mask = datasets[0].mask.to(t.bool) - else: - mask = None - - - probe_support = t.zeros_like(probes[0,0].to(dtype=t.float32)) - - xs, ys = np.mgrid[:probes.shape[-3],:probes.shape[-2]] - xs = xs - np.mean(xs) - ys = ys - np.mean(ys) - Rs = np.sqrt(xs**2 + ys**2) - - probe_support[Rs=2), - ('Second Dominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[1,0], fig=fig, basis=self.probe_basis), - lambda self: self.probe.shape[0] >=2), - ('Exit Wave Amplitude under Uniform Illumination', - lambda self, fig: p.plot_amplitude(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)), - ('Exit Wave Phase under Uniform Illumination', - lambda self, fig: p.plot_phase(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)), - ('Corrected Translations', - lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)), - ('Background', - lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)) - ] - - - def save_results(self, dataset): - basis = self.probe_basis.detach().cpu().numpy() - translations = self.corrected_translations(dataset).detach().cpu().numpy() - probe = cmath.torch_to_complex(self.probe.detach().cpu()) - probe = probe * self.probe_norm.detach().cpu().numpy() - s_matrix = cmath.torch_to_complex(self.s_matrix.detach().cpu()) - background = self.background.detach().cpu().numpy()**2 - weights = self.weights.detach().cpu().numpy() - wavelength = self.wavelength.cpu().numpy() - - return {'basis':basis, 'translation':translations, - 'probe':probe,'s_matrix':s_matrix, - 'background':background, - 'weights':weights, 'wavelength':wavelength} diff --git a/src/cdtools/models/time_resolved_ptycho_calibration.py b/src/cdtools/models/time_resolved_ptycho_calibration.py deleted file mode 100644 index 1461050..0000000 --- a/src/cdtools/models/time_resolved_ptycho_calibration.py +++ /dev/null @@ -1,503 +0,0 @@ -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 - -# -# Basic points: -# Just one probe mode, no need to overcomplicate things. -# Weights is only a list of numbers, no matrices or anything like that -# Mandatory probe support in Fourier space -# -# When loading from a dataset, we need information on the zone plate geometry -# - -__all__ = ['TimeResolvedPtychoCalibration'] - - -class TimeResolvedPtychoCalibration(CDIModel): - - def __init__(self, wavelength, detector_geometry, - probe_basis, - probe_guess, obj_guess, - fourier_times, probe_fourier_support, - times, time_dependence, frame_delays, - detector_slice=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, - oversampling=1, - loss='amplitude mse', units='um', - simulate_probe_translation=False): - - super(TimeResolvedPtychoCalibration, 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: - 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 - - 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) - - 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: - background = 1e-6 * t.ones( - self.probe[self.detector_slice].shape, - dtype=t.float32) - else: - background = 1e-6 * t.ones(self.probe[0].shape, - dtype=t.float32) - - self.background = t.nn.Parameter(background) - - if weights is None: - self.weights = None - else: - self.weights = t.nn.Parameter(t.tensor(weights, - dtype=t.float32)) - - 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 - - self.probe_fourier_support = probe_fourier_support - self.fourier_times = fourier_times - self.times = times - self.time_dependence = t.nn.Parameter(time_dependence) - self.frame_delays = frame_delays - - 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 - - # 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, zp_geometry, time_window, n_times, n_frames, randomize_ang=0, padding=0, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', simulate_probe_translation=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, - 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 - probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) - - probe = tools.propagators.far_field(probe) - - 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) - # we define a set of weights which only has one index - Ws = t.ones(len(dataset)) - - if hasattr(dataset, 'mask') and dataset.mask is not None: - mask = dataset.mask.to(t.bool) - else: - mask = None - - # What do we know about the zp? - # delta_r, N, and beamstop ratio. It's probably best, though, - # to just read diameter, beamstop_diameter and focal_length directly - # because that is the most general even if the optic isn't truly - # a zone plate. - zp_distance = zp_geometry['focal_length'] - probe_size = probe_basis * t.as_tensor(probe_shape, dtype=t.float32) - pinv_basis = t.tensor(np.linalg.pinv(probe_size).transpose()).to(t.float32) - zone_plate_basis = pinv_basis * wavelength * zp_distance - zone_plate_steps = t.sum(zone_plate_basis,axis=1) - # This may very well mix up x & y and fail on non-square detectors - probe_fourier_support = t.zeros(probe.shape, dtype=t.bool) - xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]] - xs = zone_plate_steps[1] * (xs - np.mean(xs)) - ys = zone_plate_steps[0] * (ys - np.mean(ys)) - Rs = np.sqrt(xs**2 + ys**2) - - distances = np.sqrt(zp_distance**2 + xs**2 + ys**2) - times = (distances - t.min(distances)) / 2.99792e8 - - # This sets the support of the probe and also restricts the - # timing matrix so it only considers times that are actually - # in the window defined by the probe fourier support - probe_fourier_support[Rs < zp_geometry['diameter']/2] = 1 - times[Rs > zp_geometry['diameter']/2] = 0 - times[Rs > zp_geometry['diameter']/2] = t.max(times) - - probe_fourier_support[Rs < zp_geometry['beamstop_diameter']/2] = 0 - times[Rs < zp_geometry['beamstop_diameter']/2] = t.max(times) - times[Rs < zp_geometry['beamstop_diameter']/2] = t.min(times) - - fourier_times = times - t.min(times) - probe = probe * probe_fourier_support - - # This is now the time axis for the probe's envelope - times = t.linspace(0, time_window, n_times+1) - time_dependence = t.ones(n_times, dtype=t.complex64) - frame_delays = t.linspace(0, t.max(fourier_times) + time_window, n_frames+2) - frame_delays -= time_window - frame_delays = frame_delays[1:-1] - - return cls(wavelength, det_geo, probe_basis, probe, obj, - fourier_times, probe_fourier_support, - times, time_dependence, frame_delays, - 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, - oversampling=oversampling, - loss=loss, units=units, - simulate_probe_translation=simulate_probe_translation) - - - 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]) - - probes = self.get_probes(space='real') - Ws = self.weights[index] - - # This might not work well - prs = Ws[...,None,None,None] * probes - #prs = t.sum(Ws[..., None, None, None] * probes, axis=-3) - #print(prs.shape) - - 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): - return tools.measurements.quadratic_background( - wavefields, - self.background, - detector_slice=self.detector_slice, - measurement=tools.measurements.incoherent_sum, - saturation=self.saturation, - oversampling=self.oversampling) - - - # Note: No "loss" function is defined here, because it is added - # dynamically during object creation in __init__ - - def to(self, *args, **kwargs): - super(TimeResolvedPtychoCalibration, 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: - 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_fourier_support = self.probe_fourier_support.to(*args, - **kwargs) - self.fourier_times = self.fourier_times.to(*args, **kwargs) - self.times = self.times.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_probes(self, space='real'): - # This is the part where I need to create the probe modes using the - # time-dependent stuff - - # This restricts the basis probes to stay within the probe support - optic_mask = self.probe * self.probe_fourier_support - - probes = t.zeros([len(self.frame_delays)] + list(self.probe.shape), - dtype=self.probe.dtype, device=self.probe.device) - for i, delay in enumerate(self.frame_delays): - indices = t.bucketize(self.fourier_times, self.times + delay) - clamped_indices = t.clamp(indices-1, max=len(self.time_dependence)-1) - illumination = t.take(self.time_dependence, clamped_indices) - illumination[indices==0] = 0 - illumination[indices==len(self.time_dependence)+1] = 0 - probes[i] = illumination * optic_mask - - if space.lower()=='real': - return tools.propagators.inverse_far_field(probes) - elif space.lower()=='fourier' or space.lower()=='reciprocal': - return probes - - - plot_list = [ - ('Probe Amplitudes (scroll to view modes)', - lambda self, fig: p.plot_amplitude(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)), - ('Probe Phases (scroll to view modes)', - lambda self, fig: p.plot_phase(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)), - ('Fourier Probe Amplitudes (scroll to view modes)', - lambda self, fig: p.plot_amplitude(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)), - ('Fourier Probe Phases (scroll to view modes)', - lambda self, fig: p.plot_phase(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)), - ('Optic Amplitude', - lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)), - ('Optic Phase', - lambda self, fig: p.plot_phase(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)), - ('Time Structure', - lambda self, fig: (plt.figure(fig.number) and (plt.clf() or True) and plt.plot(self.time_dependence.real.detach().cpu().numpy()) and plt.plot(self.time_dependence.imag.detach().cpu().numpy()))) - ] - -# def plot_errors(self, dataset): - - - - def save_results(self, dataset): - basis = self.probe_basis.detach().cpu().numpy() - translations = self.corrected_translations(dataset).detach().cpu().numpy() - optic = self.probe.detach().cpu().numpy() - optic = optic * self.probe_norm.detach().cpu().numpy() - time_dependence = self.time_dependence.detach().cpu().numpy() - times = self.times.detach().cpu().numpy() - fourier_times = self.fourier_times.detach().cpu().numpy() - probes = self.get_probes().detach().cpu().numpy() - obj = self.obj.detach().cpu().numpy() - background = self.background.detach().cpu().numpy()**2 - weights = self.weights.detach().cpu().numpy() - - return {'basis': basis, 'translation': translations, - 'probes': probes, 'optic': optic, - 'times': times, 'time_dependence': time_dependence, - 'fourier_times': fourier_times, - 'obj': obj, - 'background': background, - 'weights': weights, - } diff --git a/src/cdtools/models/time_resolved_rpi.py b/src/cdtools/models/time_resolved_rpi.py deleted file mode 100644 index 146826b..0000000 --- a/src/cdtools/models/time_resolved_rpi.py +++ /dev/null @@ -1,298 +0,0 @@ -import torch as t -from cdtools.models import CDIModel -from cdtools import tools -from cdtools.tools import plotting as p -from cdtools.tools.interactions import RPI_interaction -from cdtools.tools import initializers -from scipy.ndimage import binary_dilation -import numpy as np -from copy import copy - - -__all__ = ['TimeResolvedRPI'] - -class TimeResolvedRPI(CDIModel): - - @property - def obj(self): - return t.complex(self.obj_real, self.obj_imag) - - - def __init__(self, wavelength, detector_geometry, probe_basis, - probe, obj_guess, framerate, detector_slice=None, - background=None, mask=None, saturation=None, - obj_support=None, oversampling=1): - - super(TimeResolvedRPI, self).__init__() - - self.wavelength = t.tensor(wavelength) - self.framerate = framerate - self.detector_geometry = copy(detector_geometry) - - det_geo = self.detector_geometry - if hasattr(det_geo, 'distance'): - det_geo['distance'] = t.tensor(det_geo['distance']) - if hasattr(det_geo, 'basis'): - det_geo['basis'] = t.tensor(det_geo['basis']) - if hasattr(det_geo, 'corner'): - det_geo['corner'] = t.tensor(det_geo['corner']) - - self.probe_basis = t.tensor(probe_basis) - - scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1], - probe.shape[-2]/obj_guess.shape[-2]]) - self.obj_basis = self.probe_basis * scale_factor - self.detector_slice = detector_slice - - # Maybe something to include in a bit - # self.surface_normal = t.tensor(surface_normal) - - self.saturation = saturation - - if mask is None: - self.mask = mask - else: - self.mask = t.tensor(mask, dtype=t.bool) - - - self.probe = t.tensor(probe, dtype=t.complex64) - - obj_guess = t.tensor(obj_guess, dtype=t.complex64) - - self.obj_real = t.nn.Parameter(obj_guess.real) - self.obj_imag = t.nn.Parameter(obj_guess.imag) - - # Wait for LBFGS to be updated for complex-valued parameters - # self.obj = t.nn.Parameter(obj_guess.to(t.float32)) - - if background is None: - if detector_slice is not None: - background = 1e-6 * t.ones( - self.probe[0][self.detector_slice].shape, - dtype=t.float32) - else: - background = 1e-6 * t.ones(self.probe[0].shape, - dtype=t.float32) - - self.background = t.tensor(background, dtype=t.float32) - - if obj_support is not None: - self.obj_support = obj_support - self.obj.data = self.obj * obj_support[None, ...] - else: - self.obj_support = t.ones_like(self.obj[0, ...]) - - self.oversampling = oversampling - - - @classmethod - def from_dataset(cls, dataset, probe, framerate, obj_size=None, background=None, mask=None, padding=0, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', probe_threshold=0): - - 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 only need the patterns here, not the inputs associated with them. - _, 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, - oversampling=oversampling) - - if not isinstance(probe,t.Tensor): - probe = t.as_tensor(probe) - - if background is None and hasattr(dataset, 'background') \ - and dataset.background is not None: - background = t.sqrt(dataset.background) - elif background is not None: - background = t.sqrt(t.Tensor(background).to(dtype=t.float32)) - - det_geo = dataset.detector_geometry - - # If no mask is given, but one exists in the dataset, load it. - if mask is None and hasattr(dataset, 'mask') \ - and dataset.mask is not None: - mask = dataset.mask.to(t.bool) - - # Now we initialize the object - if obj_size is None: - # This is a standard size for a well-matched probe and detector - obj_size = list((np.array(probe_shape) // 2).astype(int)) - - # I think something to do with the fact that the object is defined - # on a coarser grid needs to be accounted for here that is not - # accounted for yet - scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2) - n_modes = (probe.shape[0] - 1) // framerate + 1 - - obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size)) - - - probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0)) - probe_fft = tools.propagators.far_field(probe_intensity) - pad0l = (probe.shape[-2] - obj_size[-2])//2 - pad0r = probe.shape[-2] - obj_size[-2] - pad0l - pad1l = (probe.shape[-1] - obj_size[-1])//2 - pad1r = probe.shape[-1] - obj_size[-1] - pad1l - probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r] - probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft)) - - obj_support = probe_lr > t.max(probe_lr) * probe_threshold - obj_support = t.as_tensor(binary_dilation(obj_support)) - - return cls(wavelength, det_geo, probe_basis, - probe, obj_guess, framerate, detector_slice=det_slice, - background=background, mask=mask, saturation=saturation, - obj_support=obj_support, oversampling=oversampling) - - - def random_init(self, pattern): - scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2) - self.obj.data = scale * t.exp( - 2j * np.pi * t.rand(self.obj.shape)).to( - dtype=self.obj.dtype, device=self.obj.device) - - - # Needs work - def interaction(self, index, *args): - # including *args allows this to work with all sorts of datasets - # that might include other information in with the index in their - # "input" parameters (such as translations for a ptychography dataset). - # This makes it seamless to use such a dataset even though those - # extra arguments will not be used. - - - all_exit_waves = [] - - # Mix the probes with the weight matrix - prs = self.probe - - for i in range(self.probe.shape[0]): - obj_frame = i // self.framerate - pr = prs[i] - exit_waves = RPI_interaction(pr, - self.obj_support * self.obj[obj_frame]) - all_exit_waves.append(exit_waves.unsqueeze(0)) - - # This creates a bunch of modes generated from all possible combos - # of the probe and object modes all strung out along the first index - - output = t.cat(all_exit_waves) - - # If we have multiple indexes input, we unsqueeze and repeat the stack - # of wavefields enough times to simulate each requested index. This - # seems silly, but it enables (for example) one to do a reconstruction - # from a set of diffraction patterns that are all known to be from the - # same object. - try: - # will fail if index has no length, for example when index - # is just an int. In this case, we just do nothing instead - output = output.unsqueeze(0).repeat(1,len(index),1,1,1) - except TypeError: - pass - return output - - - 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): - # Here I'm taking advantage of an undocumented feature in the - # incoherent_sum measurement function where it will work with - # a 4D wavefield array as well as a 5D array. - return tools.measurements.quadratic_background(wavefields, - self.background, - detector_slice=self.detector_slice, - measurement=tools.measurements.incoherent_sum, - saturation=self.saturation, - oversampling=self.oversampling) - - def loss(self, sim_data, real_data, mask=None): - return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) - #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) - - def regularizer(self, factors): - return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \ - + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2) - - def to(self, *args, **kwargs): - super(TimeResolvedRPI, self).to(*args, **kwargs) - self.wavelength = self.wavelength.to(*args,**kwargs) - # move the detector geometry too - det_geo = self.detector_geometry - if hasattr(det_geo, 'distance'): - det_geo['distance'] = det_geo['distance'].to(*args,**kwargs) - if hasattr(det_geo, 'basis'): - det_geo['basis'] = det_geo['basis'].to(*args,**kwargs) - if hasattr(det_geo, 'corner'): - det_geo['corner'] = det_geo['corner'].to(*args,**kwargs) - - if self.mask is not None: - self.mask = self.mask.to(*args, **kwargs) - - self.probe = self.probe.to(*args,**kwargs) - self.probe_basis = self.probe_basis.to(*args,**kwargs) - self.obj_basis = self.obj_basis.to(*args,**kwargs) - self.obj_support = self.obj_support.to(*args,**kwargs) - self.background = self.background.to(*args, **kwargs) - - # Maybe include in a bit - #self.surface_normal = self.surface_normal.to(*args, **kwargs) - - def sim_to_dataset(self, args_list): - raise NotImplementedError('No sim to dataset yet, sorry!') - - plot_list = [ - ('Probe Amplitudes', - lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)), - ('Object Amplitudes', - lambda self, fig: p.plot_amplitude(self.obj, fig=fig, - basis=self.obj_basis)), - ('Object Phases', - lambda self, fig: p.plot_phase(self.obj, fig=fig, - basis=self.obj_basis)) - ] - - - def save_results(self, dataset=None, full_obj=False): - # dataset is set as a kwarg here because it isn't needed, but the - # common pattern is to pass a dataset. This makes it okay if one - # continues to use that standard pattern - probe_basis = self.probe_basis.detach().cpu().numpy() - obj_basis = self.obj_basis.detach().cpu().numpy() - probe = self.probe.detach().cpu().numpy() - # Provide the option to save out the subdominant objects or - # just the dominant one - if full_obj: - obj = self.obj.detach().cpu().numpy() - else: - obj = self.obj[0].detach().cpu().numpy() - background = self.background.detach().cpu().numpy()**2 - - return {'probe_basis': probe_basis, 'obj_basis': obj_basis, - 'probe': probe,'obj': obj, - 'background': background} - diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py index 91e1597..7558cd4 100644 --- a/src/cdtools/tools/analysis/analysis.py +++ b/src/cdtools/tools/analysis/analysis.py @@ -775,8 +775,8 @@ def calc_rms_error(field_1, field_2, align_phases=True, normalize=False, The RMS error, or tensor of RMS errors, depending on the dim argument """ - fields_1 = t.as_tensor(fields_1) - fields_2 = t.as_tensor(fields_2) + field_1 = t.as_tensor(field_1) + field_2 = t.as_tensor(field_2) sumdims = tuple(d - dims for d in range(dims)) @@ -922,6 +922,8 @@ def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2): The generalized RMS error, or tensor of generalized RMS errors, depending on the dim argument """ + # TODO either make this work correctly or just make everything work + # only for tensors fields_1 = t.as_tensor(fields_1) fields_2 = t.as_tensor(fields_2) @@ -1378,14 +1380,6 @@ def standardize_reconstruction_pair( limit=frc_limit, ) - - #probe_freqs, probe_frc, probe_frc_threshold = calc_generalized_frc( - # probe_1, - # probe_2, - # half_1['probe_basis'], - # nbins=probe_nbins, - # limit=frc_limit, - #) probe_freqs, probe_frc, probe_frc_threshold = calc_generalized_frc( probe_1, probe_2, diff --git a/src/cdtools/tools/interactions/interactions.py b/src/cdtools/tools/interactions/interactions.py index 3cd244c..5f9b453 100644 --- a/src/cdtools/tools/interactions/interactions.py +++ b/src/cdtools/tools/interactions/interactions.py @@ -625,6 +625,7 @@ def RPI_interaction(probe, obj): # The far-field propagator is just a 2D FFT but with an fftshift fftobj = propagators.far_field(obj) fftobj_npix = fftobj.shape[-2] * fftobj.shape[-1] + # We calculate the padding that we need to do the upsampling # This is carefully set up to keep the zero-frequency pixel in the correct # location as the overall shape changes. Don't mess with this without @@ -635,8 +636,12 @@ def RPI_interaction(probe, obj): pad1r = probe.shape[-1] - obj.shape[-1] - pad1l fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad2l, pad2r)) + + # This keeps the mean intensity equal, instead of spreading out the + # intensity over the upsampled region + # TODO: Keeping this in probably isn't the most efficient fftobj_npix_new = fftobj.shape[-2] * fftobj.shape[-1] - scale_factor =np.sqrt(fftobj_npix_new / fftobj_npix) + scale_factor = np.sqrt(fftobj_npix_new / fftobj_npix) # Again, just an inverse FFT but with an fftshift upsampled_obj = scale_factor * propagators.inverse_far_field(fftobj) diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index f29abd4..ae3d6bd 100644 --- a/tests/tools/test_interactions.py +++ b/tests/tools/test_interactions.py @@ -253,7 +253,6 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj): def test_RPI_interaction(random_probe, random_obj): - random_obj1 = random_obj[:79,:68] * 0 + 1 random_probe1 = random_probe * 0 + 1 t_random_obj1 = t.as_tensor(random_obj1) @@ -269,6 +268,9 @@ def test_RPI_interaction(random_probe, random_obj): output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups), norm='ortho')) + output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1] + / (random_obj1.shape[-2] * random_obj1.shape[-1])) + assert np.allclose(t_output1, output1) random_obj2 = np.stack([random_obj[:64,:89]]*3) @@ -286,5 +288,9 @@ def test_RPI_interaction(random_probe, random_obj): output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups), norm='ortho')) + output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1] + / (random_obj2.shape[-2] * random_obj2.shape[-1])) + assert np.allclose(t_output2, output2) +