diff --git a/CDTools/datasets/base.py b/CDTools/datasets/base.py index 469a978..f4d3067 100644 --- a/CDTools/datasets/base.py +++ b/CDTools/datasets/base.py @@ -121,7 +121,7 @@ class CDataset(torchdata.Dataset): self.mask = self.mask.to(*args,**mask_kwargs) if self.background is not None: self.background = self.background.to(*args,**kwargs) - + def get_as(self, *args, **kwargs): """Sets the dataset to return data on the given device and dtype diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index ce46fc4..0b84d54 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -198,169 +198,27 @@ class Ptycho2DDataset(CDataset): position. """ - # We start by making the figure and axes - fig, axes = plt.subplots(1,2,figsize=(8,5.3)) - fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96]) - axslider = plt.axes([0.15,0.06,0.75,0.03]) - - # - # Then we define some helper functions for getting the right data - # that are used both in the initial setup and the updates - # - - def get_data(idx): + def get_images(idx): inputs, output = self[idx] meas_data = output.detach().cpu().numpy() if hasattr(self, 'mask') and self.mask is not None: mask = self.mask.detach().cpu().numpy() else: mask = 1 - - return mask, meas_data - - - def calculate_sizes(idx): - bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted()) - s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch - s0 /= 4 # A rough value to make the size work out - s = np.ones(len(self)) * s0 - - s[idx] *= 4 - return s - - def update_colorbar(im): - # - # This solves the problem of the colorbar being changed - # when the forward and back buttons are used!!! - # - if hasattr(im, 'norecurse') and im.norecurse: - im.norecurse=False - return - - im.norecurse=True - # This is needed to update the colorbar - # only change limits if array contains multiple values - if np.min(im.get_array()) != np.max(im.get_array()): - im.set_clim(vmin=np.min(im.get_array()), - vmax=np.max(im.get_array())) - - # - # The meatiest part of this program, here we just go through and - # set up the plot how we want it - # - - # First we set up the left-hand plot, which shows an overview map - axes[0].set_title('Relative Displacement Map') + + if logarithmic: + return np.log(meas_data) / np.log(10) * mask + else: + return meas_data * mask translations = self.translations.detach().cpu().numpy() nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy() - s = calculate_sizes(0) - units_factor = plotting.get_units_factor(units) - nanomap = axes[0].scatter(units_factor * translations[:,0],units_factor * translations[:,1],s=s,c=nanomap_values, picker=True) - - axes[0].invert_xaxis() - axes[0].set_facecolor('k') - axes[0].set_xlabel('Translation x ('+units+')', labelpad=1) - axes[0].set_ylabel('Translation y ('+units+')', labelpad=1) - cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal', - format='%.2e', - ticks=ticker.LinearLocator(numticks=5), - pad=0.17,fraction=0.1) - cb1.ax.set_title('Integrated Intensity', size="medium", pad=5) - cb1.ax.tick_params(labelrotation=20) - - - # Now we set up the second plot, which shows the individual - # diffraction patterns - axes[1].set_title('Diffraction Pattern') - mask, meas_data = get_data(0) if logarithmic: - meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask) + cbar_title='Log Base 10 of Diffraction Intensity' else: - meas = axes[1].imshow(meas_data * mask) + cbar_title='Diffraction Intensity' + + plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title) - cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal', - format='%.2e', - ticks=ticker.LinearLocator(numticks=5), - pad=0.17,fraction=0.1) - cb2.ax.tick_params(labelrotation=20) - cb2.ax.set_title('Pixel Intensity', size="medium", pad=5) - cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) - - # This function handles all the updating, except for moving the - # slider value. This is done because the slider widget is - # ultimately responsible for triggering an update, so all other - # updates are done by changing the slider widget value which - # then triggers this - def update(idx): - # We have to explicitly make it an integer because the slider will - # output floats (even if they are still integer-valued) - idx = int(idx) - - # Get the new data for this index - mask, meas_data = get_data(idx) - - # Now we resize the nanomap to show the new selection - axes[0].collections[0].set_sizes(calculate_sizes(idx)) - - # And we update the data in the image as well - meas = axes[1].images[-1] - if logarithmic: - meas.set_data(np.log(meas_data) / np.log(10) * mask) - else: - meas.set_data(meas_data * mask) - - update_colorbar(meas) - - - # - # Now we define the functions to handle various kinds of events - # that can be thrown our way - # - - # We start by creating the slider here, so it can be used - # by the update hooks. - slider = Slider(axslider, 'Pattern #', 0, len(self)-1, valstep=1, valfmt="%d") - - # This handles scroll wheel and keypress events - def on_action(event): - # Otherwise the if statements can throw errors when the - # event type isn't right, this way they just don't trigger - if not hasattr(event, 'button'): - event.button = None - if not hasattr(event, 'key'): - event.key = None - - if event.key == 'up' or event.button == 'up' or event.key == 'right': - idx = slider.val - 1 - elif event.key == 'down' or event.button == 'down' or event.key == 'left': - idx = slider.val + 1 - else: - # This prevents errors from being thrown on irrelevant key - # or mouse input - return - - # Handle the wraparound and trigger the update - idx = int(idx) % len(self) - slider.set_val(idx) - - # This handles "pick" events in the nanomap - def on_pick(event): - # If we don't filter on type of event, this will also capture, - # for example, scroll events that happen over the nanomap - if event.mouseevent.button == 1: - slider.set_val(event.ind[0]) - - - # Here we connect the various update functions - fig.canvas.mpl_connect('pick_event',on_pick) - fig.canvas.mpl_connect('key_press_event',on_action) - fig.canvas.mpl_connect('scroll_event',on_action) - slider.on_changed(update) - - # Throw an extra update into the mix just to get rid of any things - # (like the nanomap dot sizes) that otherwise would change on the - # first update - update(0) diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 8604d7b..1e991fb 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -20,19 +20,15 @@ defining a new ptychography model before attempting to do so. """ -from __future__ import division, print_function, absolute_import - # 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', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI', 'UnifiedModePtycho'] +#__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI'] from CDTools.models.base import CDIModel from CDTools.models.simple_ptycho import SimplePtycho from CDTools.models.fancy_ptycho import FancyPtycho -from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho -from CDTools.models.s_matrix_ptycho import SMatrixPtycho from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho from CDTools.models.rpi import RPI -from CDTools.models.unified_mode_ptycho import UnifiedModePtycho -from CDTools.models.unified_mode_ptycho2 import UnifiedModePtycho2 +#from CDTools.models.s_matrix_ptycho import SMatrixPtycho + diff --git a/CDTools/models/base.py b/CDTools/models/base.py index cefa169..5ceca97 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -39,6 +39,9 @@ import numpy as np import threading import queue import time +#import pytorch_warmup +from .complex_adam import MyAdam +from .complex_lbfgs import MyLBFGS __all__ = ['CDIModel'] @@ -53,6 +56,10 @@ class CDIModel(t.nn.Module): functions. """ + def __init__(self): + super(CDIModel,self).__init__() + self.iteration_count = 0 + def from_dataset(self, dataset): raise NotImplementedError() @@ -100,7 +107,8 @@ class CDIModel(t.nn.Module): raise NotImplementedError() def AD_optimize(self, iterations, data_loader, optimizer,\ - scheduler=None, regularization_factor=None, thread=True): + scheduler=None, regularization_factor=None, thread=True, + calculation_width=10, warmup_scheduler=None): """Runs a round of reconstruction using the provided optimizer This is the basic automatic differentiation reconstruction tool @@ -123,41 +131,85 @@ class CDIModel(t.nn.Module): Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method thread : bool Default True, whether to run the computation in a separate thread to allow interaction with plots during computation + calculation_width : int + Default 10, how many translations to pass through at once for each round of gradient accumulation """ # First, calculate the normalization normalization = 0 for inputs, patterns in data_loader: normalization += t.sum(patterns).cpu().numpy() - + def run_iteration(stop_event=None): loss = 0 N = 0 + t0 = time.time() for inputs, patterns in data_loader: N += 1 def closure(): - # This is just used to allow graceful exit when threading - if stop_event is not None and stop_event.is_set(): - exit() optimizer.zero_grad() - sim_patterns = self.forward(*inputs) - if hasattr(self, 'mask'): - loss = self.loss(patterns,sim_patterns, mask=self.mask) - else: - loss = self.loss(patterns,sim_patterns) + + input_chunks = [[inp[i:i + calculation_width] + for inp in inputs] + for i in range(0, len(inputs[0]), + calculation_width)] + pattern_chunks = [patterns[i:i + calculation_width] + for i in range(0, len(inputs[0]), + calculation_width)] + + total_loss = 0 + for inp, pats in zip(input_chunks, pattern_chunks): + # This is just used to allow graceful exit when threading + if stop_event is not None and stop_event.is_set(): + exit() + + sim_patterns = self.forward(*inp) + #sim_patterns.retain_grad() + if hasattr(self, 'mask'): + loss = self.loss(pats,sim_patterns, mask=self.mask) + else: + loss = self.loss(pats,sim_patterns) + + loss.backward()#retain_variables=True) + #plt.figure() + #plt.imshow(sim_patterns.grad.cpu()[0]) + #plt.colorbar() + #plt.show() + + total_loss += loss.detach() + + #print('probe grad') + #print(t.mean(self.obj.grad)) + if regularization_factor is not None \ and hasattr(self, 'regularizer'): - loss += self.regularizer(regularization_factor) - #print(loss) - loss.backward() - return loss - + loss = self.regularizer(regularization_factor) + loss.backward() + return total_loss + + + + if warmup_scheduler is not None: + old_lrs = [group['lr'] for group in + optimizer.param_groups] + warmup_scheduler.dampen() + #print([group['lr'] for group in + # optimizer.param_groups], end='\r') + loss += optimizer.step(closure).detach().cpu().numpy() + + if warmup_scheduler is not None: + for old_lr, group in zip(old_lrs, optimizer.param_groups): + group['lr'] = old_lr + loss /= normalization if scheduler is not None: scheduler.step(loss) + self.latest_loss = loss + self.latest_iteration_time = time.time() - t0 + self.iteration_count += 1 return loss if thread: @@ -176,6 +228,7 @@ class CDIModel(t.nn.Module): self.figs[0].canvas.start_event_loop(0.01) else: calc.join() + except KeyboardInterrupt as e: stop_event.set() print('\nAsking execution thread to stop cleanly - please be patient.') @@ -189,7 +242,8 @@ class CDIModel(t.nn.Module): def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False, amsgrad=False, subset=None, - regularization_factor=None, thread=True): + regularization_factor=None, thread=True, + calculation_width=10, warmup=False): """Runs a round of reconstruction using the Adam optimizer This is generally accepted to be the most robust algorithm for use @@ -215,6 +269,9 @@ class CDIModel(t.nn.Module): Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method thread : bool Default True, whether to run the computation in a separate thread to allow interaction with plots during computation + calculation_width : int + Default 1, how many translations to pass through at once for each round of gradient accumulation + """ if subset is not None: @@ -228,7 +285,8 @@ class CDIModel(t.nn.Module): shuffle=True) # Define the optimizer - optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad) + #optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad) + optimizer = MyAdam(self.parameters(), lr = lr, amsgrad=amsgrad) # Define the scheduler @@ -237,21 +295,33 @@ class CDIModel(t.nn.Module): else: scheduler = None + if warmup: + print('Warmup is not currently implemented, sorry!') + #warmup_scheduler = pytorch_warmup.UntunedLinearWarmup(optimizer) + else: + warmup_scheduler = None + return self.AD_optimize(iterations, data_loader, optimizer, scheduler=scheduler, + warmup_scheduler=warmup_scheduler, regularization_factor=regularization_factor, - thread=thread) + thread=thread, + calculation_width=calculation_width) - def LBFGS_optimize(self, iterations, dataset, batch_size=None, + def LBFGS_optimize(self, iterations, dataset, lr=0.1,history_size=2, subset=None, - regularization_factor=None, thread=True): + regularization_factor=None, thread=True, + calculation_width=10): """Runs a round of reconstruction using the L-BFGS optimizer This algorithm is often less stable that Adam, however in certain situations or geometries it can be shockingly efficient. Like all the other optimization routines, it is defined as a generator function which yields the average loss each epoch. + + Note: There is no batch size, because it is a usually a bad idea to use + LBFGS on anything but all the data at onece Parameters ---------- @@ -259,8 +329,6 @@ class CDIModel(t.nn.Module): How many epochs of the algorithm to run dataset : CDataset The dataset to reconstruct against - batch_size : int - Optional, the size of the minibatches to use lr : float Optional, the learning rate to use history_size : int @@ -271,6 +339,7 @@ class CDIModel(t.nn.Module): Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method thread : bool Default True, whether to run the computation in a separate thread to allow interaction with plots during computation + """ if subset is not None: # if just one pattern, turn into a list for convenience @@ -278,27 +347,27 @@ class CDIModel(t.nn.Module): subset = [subset] dataset = torchdata.Subset(dataset, subset) - # Make a dataloader - if batch_size is not None: - data_loader = torchdata.DataLoader(dataset, batch_size=batch_size, - shuffle=True) - else: - data_loader = torchdata.DataLoader(dataset) + # Make a dataloader. This basically does nothing but load all the + # data at once + data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset)) # Define the optimizer optimizer = t.optim.LBFGS(self.parameters(), lr = lr, history_size=history_size) - + #optimizer = MyLBFGS(self.parameters(), + # lr = lr, history_size=history_size) + return self.AD_optimize(iterations, data_loader, optimizer, regularization_factor=regularization_factor, - thread=thread) + thread=thread, + calculation_width=calculation_width) def SGD_optimize(self, iterations, dataset, batch_size=None, lr=0.01, momentum=0, dampening=0, weight_decay=0, nesterov=False, subset=None, regularization_factor=None, - thread=True): + thread=True, calculation_width=10): """Runs a round of reconstruction using the SGDoptimizer This algorithm is often less stable that Adam, but it is simpler @@ -322,6 +391,9 @@ class CDIModel(t.nn.Module): Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method thread : bool Default True, whether to run the computation in a separate thread to allow interaction with plots during computation + calculation_width : int + Default 1, how many translations to pass through at once for each round of gradient accumulation + """ if subset is not None: @@ -347,9 +419,19 @@ class CDIModel(t.nn.Module): return self.AD_optimize(iterations, data_loader, optimizer, regularization_factor=regularization_factor, - thread=thread) + thread=thread, + calculation_width=calculation_width) + def report(self): + """Returns a string informing on the latest reconstruction iteration""" + if hasattr(self, 'latest_loss'): + return 'Iteration ' + str(self.iteration_count) + \ + ' completed in %0.2f s with loss ' %\ + self.latest_iteration_time + str(self.latest_loss) + else: + return 'No reconstruction iterations performed yet!' + # By default, the plot_list is empty plot_list = [] @@ -407,7 +489,7 @@ class CDIModel(t.nn.Module): plotter = plots[1] if figs is None: - fig = plt.figure() + fig = plt.figure() self.figs.append(fig) else: fig = figs[idx] @@ -415,6 +497,7 @@ class CDIModel(t.nn.Module): try: plotter(self,fig) plt.title(name) + except TypeError as e: if dataset is not None: try: @@ -448,7 +531,7 @@ class CDIModel(t.nn.Module): fig, axes = plt.subplots(1,3,figsize=(12,5.3)) fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96]) axslider = plt.axes([0.15,0.06,0.75,0.03]) - + def update_colorbar(im): # If the update brought the colorbar out of whack @@ -463,12 +546,8 @@ class CDIModel(t.nn.Module): return im.norecurse=True - im.colorbar.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array())) - im.colorbar.ax.set_ylim(0,1) - im.colorbar.set_ticks(ticker.LinearLocator(numticks=5)) - im.colorbar.draw_all() + im.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array())) - def update(idx): idx = int(idx) % len(dataset) fig.pattern_idx = idx @@ -503,6 +582,7 @@ class CDIModel(t.nn.Module): cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff)) else: + sim = axes[0].images[-1] sim.set_data(sim_data) update_colorbar(sim) diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index e2a00fe..650b5ab 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/CDTools/models/bragg_2d_ptycho.py @@ -4,7 +4,6 @@ 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 CDTools.tools.propagators import generate_generalized_angular_spectrum_propagator as ggasp from matplotlib import pyplot as plt @@ -110,9 +109,9 @@ class Bragg2DPtycho(CDIModel): # We rescale the probe here so it learns at the same rate as the # object if probe_guess.dim() > 3: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) + self.probe_norm = 1 * t.max(t.abs(probe_guess[0]).to(t.float32)) else: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32))) + self.probe_norm = 1 * t.max(t.abs(probe_guess).to(t.float32)) # Not strictly necessary but otherwise it will return # a probe with the stuff outside of the support unchanged after @@ -124,42 +123,44 @@ class Bragg2DPtycho(CDIModel): # you look at the phase map probe_guess[probe_guess == 0] = 0 - self.probe = t.nn.Parameter(probe_guess.to(t.float32) + self.probe = t.nn.Parameter(probe_guess.to(t.complex64) / self.probe_norm) - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + self.obj = t.nn.Parameter(obj_guess.to(t.complex64)) if background is None: if detector_slice is not None: - background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1]) + background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape) else: - background = 1e-6 * t.ones(self.probe[0].shape[:-1]) + background = 1e-6 * t.ones(self.probe[0].shape) - self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) + self.background = t.nn.Parameter(t.as_tensor(background,dtype=t.float32)) if weights is None: self.weights = None else: - self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32)) + self.weights = t.nn.Parameter(t.as_tensor(weights,dtype=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) + self.translation_offsets = t.nn.Parameter( + t.as_tensor(translation_offsets,dtype=t.float32) / + translation_scale) 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]) + self.probe_support = t.ones(self.probe[0].shape,dtype=t.bool) if obj_support is not None: self.obj_support = obj_support self.obj.data = self.obj * obj_support else: - self.obj_support = t.ones_like(self.obj) + self.obj_support = t.ones_like(self.obj, dtype=t.bool) self.oversampling = oversampling @@ -170,7 +171,7 @@ class Bragg2DPtycho(CDIModel): # recall that here we always want the shape of the detector # before it's cut down by the detector slice to match the # physical detector region - probe_shape = self.probe[0].shape[:-1] + probe_shape = self.probe[0] self.k_map, self.intensity_map = \ tools.propagators.generate_high_NA_k_intensity_map( @@ -188,11 +189,11 @@ class Bragg2DPtycho(CDIModel): # This propagator should be able to be multiplied by the propagation # distance each time to get a propagator - self.universal_propagator = cmath.cphase(ggasp(self.probe.shape[1:], + self.universal_propagator = t.angle(ggasp(self.probe.shape[1:], self.probe_basis, self.wavelength, t.Tensor([0,0,self.wavelength/(2*np.pi)]), propagation_vector=self.prop_dir, - dtype=t.float32, + dtype=t.complex64, propagate_along_offset=True)) @@ -299,11 +300,11 @@ class Bragg2DPtycho(CDIModel): # Now we initialize all the subdominant probe modes - probe_max = t.max(cmath.cabs(probe)) + 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) - obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) + obj = t.exp(1j*(randomize_ang * (t.rand(obj_size)-0.5))) det_geo = dataset.detector_geometry @@ -317,8 +318,8 @@ class Bragg2DPtycho(CDIModel): mask = None if probe_support_radius is not None: - probe_support = t.zeros_like(probe[0].to(dtype=t.float32)) - p_cent = np.array(probe.shape[1:3]).astype(int) // 2 + probe_support = t.zeros(probe[0].shape, dtype=t.bool) + p_cent = np.array(probe[0].shape).astype(int) // 2 psr = int(probe_support_radius) probe_support[p_cent[0]-psr:p_cent[0]+psr, p_cent[1]-psr:p_cent[1]+psr] = 1 @@ -329,7 +330,7 @@ class Bragg2DPtycho(CDIModel): ro = restrict_obj os = np.array(obj_size) ps = np.array(probe_shape) - obj_support = t.zeros_like(obj.to(dtype=t.float32)) + obj_support = t.zeros(obj.shape,dtype=t.bool) obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2, ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1 else: @@ -368,53 +369,29 @@ class Bragg2DPtycho(CDIModel): pix_trans -= self.min_translation props -= self.median_propagation - + if self.translation_offsets is not None: pix_trans += self.translation_scale * self.translation_offsets[index] - single_translation = False - if translations.dim() == 1: - translations = translations[None,:] - pix_trans = pix_trans[None,:] - single_translation = True - - all_exit_waves = [] - - for i in range(self.probe.shape[0]): - pr = self.probe[i] * self.probe_support - - exit_waves = [] - for j in range(translations.size()[0]): - if self.propagate_probe: - # Minus sign is empirical - propagator = cmath.expi( - (-1*props[j]*(2*np.pi)/self.wavelength) - * self.universal_propagator) - prop_pr = tools.propagators.near_field(pr, propagator) - #plt.close('all') - #plt.imshow(np.abs(cmath.torch_to_complex(prop_pr.detach().cpu()))) - #plt.show() - else: - prop_pr = pr - - exit_waves.append(self.probe_norm * - tools.interactions.ptycho_2D_sinc(prop_pr, - self.obj_support * self.obj, - pix_trans[j], - shift_probe=True)) - exit_waves = t.stack(exit_waves) - - - if not single_translation: - exit_waves = self.weights[index][:,None,None,None] * exit_waves - else: - exit_waves = self.weights[index] * exit_waves[0,:,:,:] - - all_exit_waves.append(exit_waves) + Ws = self.weights[index] + prs = Ws[...,None,None,None] * self.probe + # Now we need to propagate each of the probes - return t.stack(all_exit_waves) - + for j in range(prs.shape[0]): + # I believe this -1 sign is in error, but I need a dataset with + # well understood geometry to figure it out + propagator = t.exp( + 1j*(props[j]*(2*np.pi)/self.wavelength) + * self.universal_propagator) + prs[j] = tools.propagators.near_field(prs[j], propagator) + + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs, self.obj_support * self.obj,pix_trans, + shift_probe=True, multiple_modes=True) + + return exit_waves + def forward_propagator(self, wavefields): if self.correct_tilt: @@ -514,32 +491,12 @@ class Bragg2DPtycho(CDIModel): def corrected_translations(self,dataset): - translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device) + translations = dataset.translations.to(dtype=self.probe.real.dtype, + 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 - # Needs to be updated to allow for plotting to an existing figure - # plot_list = [ - # ('Dominant Probe Amplitude', - # lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), - # ('Dominant Probe Phase', - # lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), - # ('Subdominant Probe Amplitude', - # lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), - # lambda self: len(self.probe) >=2), - # ('Subdominant Probe Phase', - # lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), - # lambda self: len(self.probe) >=2), - # ('Object Amplitude', - # lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), - # ('Object Phase', - # lambda self, fig: p.plot_phase(self.obj, 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)) - # ] plot_list = [ ('Dominant Probe Amplitude', lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig)), @@ -565,9 +522,9 @@ class Bragg2DPtycho(CDIModel): 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 = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() - obj = cmath.torch_to_complex(self.obj.detach().cpu()) + obj = self.obj.detach().cpu().numpy() background = self.background.detach().cpu().numpy()**2 weights = self.weights.detach().cpu().numpy() diff --git a/CDTools/models/complex_adam.py b/CDTools/models/complex_adam.py new file mode 100644 index 0000000..1ded7d9 --- /dev/null +++ b/CDTools/models/complex_adam.py @@ -0,0 +1,163 @@ +import math +import torch +from torch import Tensor +from torch.optim.optimizer import Optimizer +from typing import List, Optional + +class MyAdam(Optimizer): + r"""Implements Adam algorithm. + It has been proposed in `Adam: A Method for Stochastic Optimization`_. + The implementation of the L2 penalty follows changes proposed in + `Decoupled Weight Decay Regularization`_. + Args: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups + lr (float, optional): learning rate (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) + .. _Adam\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _Decoupled Weight Decay Regularization: + https://arxiv.org/abs/1711.05101 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, + weight_decay=0, amsgrad=False): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + if not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + if not 0.0 <= weight_decay: + raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) + defaults = dict(lr=lr, betas=betas, eps=eps, + weight_decay=weight_decay, amsgrad=amsgrad) + super(MyAdam, self).__init__(params, defaults) + + def __setstate__(self, state): + super(MyAdam, self).__setstate__(state) + for group in self.param_groups: + group.setdefault('amsgrad', False) + + @torch.no_grad() + def step(self, closure=None): + """Performs a single optimization step. + Args: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad = [] + grads = [] + exp_avgs = [] + exp_avg_sqs = [] + max_exp_avg_sqs = [] + state_steps = [] + beta1, beta2 = group['betas'] + + for p in group['params']: + if p.grad is not None: + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead') + grads.append(p.grad) + + state = self.state[p] + # Lazy state initialization + if len(state) == 0: + state['step'] = 0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format) + if group['amsgrad']: + # Maintains max of all exp. moving avg. of sq. grad. values + state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format) + + exp_avgs.append(state['exp_avg']) + exp_avg_sqs.append(state['exp_avg_sq']) + + if group['amsgrad']: + max_exp_avg_sqs.append(state['max_exp_avg_sq']) + + # update the steps for each param group update + state['step'] += 1 + # record the step after step update + state_steps.append(state['step']) + + adam(params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=group['amsgrad'], + beta1=beta1, + beta2=beta2, + lr=group['lr'], + weight_decay=group['weight_decay'], + eps=group['eps']) + return loss + + + +def adam(params: List[Tensor], + grads: List[Tensor], + exp_avgs: List[Tensor], + exp_avg_sqs: List[Tensor], + max_exp_avg_sqs: List[Tensor], + state_steps: List[int], + *, + amsgrad: bool, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float): + r"""Functional API that performs Adam algorithm computation. + See :class:`~torch.optim.Adam` for details. + """ + + for i, param in enumerate(params): + + grad = grads[i] + exp_avg = exp_avgs[i] + exp_avg_sq = exp_avg_sqs[i] + step = state_steps[i] + + bias_correction1 = 1 - beta1 ** step + bias_correction2 = 1 - beta2 ** step + + if weight_decay != 0: + grad = grad.add(param, alpha=weight_decay) + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2) + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i]) + # Use the max. for normalizing running avg. of gradient + denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps) + else: + denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps) + + step_size = lr / bias_correction1 + + param.addcdiv_(exp_avg, denom, value=-step_size) diff --git a/CDTools/models/complex_lbfgs.py b/CDTools/models/complex_lbfgs.py new file mode 100644 index 0000000..2a03696 --- /dev/null +++ b/CDTools/models/complex_lbfgs.py @@ -0,0 +1,485 @@ +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/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 936ef1d..89a3750 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -4,11 +4,12 @@ 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 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'] @@ -23,7 +24,8 @@ class FancyPtycho(CDIModel): min_translation = t.Tensor([0,0]), background = None, translation_offsets=None, mask=None, weights = None, translation_scale = 1, saturation=None, - probe_support = None, obj_support=None, oversampling=1): + probe_support = None, obj_support=None, oversampling=1, + loss='amplitude mse',units='um'): super(FancyPtycho,self).__init__() self.wavelength = t.Tensor([wavelength]) @@ -43,6 +45,7 @@ class FancyPtycho(CDIModel): self.surface_normal = t.Tensor(surface_normal) self.saturation = saturation + self.units = units if mask is None: self.mask = mask @@ -51,21 +54,21 @@ class FancyPtycho(CDIModel): # We rescale the probe here so it learns at the same rate as the # object - if probe_guess.dim() > 3: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) + if probe_guess.dim() > 2: + self.probe_norm = 1 * t.max(t.abs(probe_guess[0].to(t.complex64))) else: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32))) + self.probe_norm = 1 * t.max(t.abs(probe_guess.to(t.complex64))) - self.probe = t.nn.Parameter(probe_guess.to(t.float32) + self.probe = t.nn.Parameter(probe_guess.to(t.complex64) / self.probe_norm) - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + self.obj = t.nn.Parameter(obj_guess.to(t.complex64)) if background is None: if detector_slice is not None: - background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1]) + background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape) else: - background = 1e-6 * t.ones(self.probe[0].shape[:-1]) + background = 1e-6 * t.ones(self.probe[0].shape) self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) @@ -73,8 +76,20 @@ class FancyPtycho(CDIModel): if weights is None: self.weights = None else: - self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32)) - + # 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).to(t.float32)) + else: + # Now this is a matrix of weights, so we + if type(weights) == type(t.zeros(1)): + self.weights = t.nn.Parameter(weights.to(t.complex64)) + else: + # There is a good chance that this doesn't work + self.weights = t.nn.Parameter(t.Tensor(weights).to(t.complex64)) + + if translation_offsets is None: self.translation_offsets = None else: @@ -96,8 +111,19 @@ class FancyPtycho(CDIModel): self.oversampling = oversampling + # 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, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False): + 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, opt_for_fft=False, loss='amplitude mse', units='um'): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -153,6 +179,7 @@ class FancyPtycho(CDIModel): 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: @@ -166,18 +193,35 @@ class FancyPtycho(CDIModel): # Now we initialize all the subdominant probe modes - probe_max = t.max(cmath.cabs(probe)) + 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) - obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) + 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) - weights = t.ones(len(dataset)) + 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) @@ -185,14 +229,17 @@ class FancyPtycho(CDIModel): mask = None if probe_support_radius is not None: - probe_support = t.zeros_like(probe[0].to(dtype=t.float32)) - p_cent = np.array(probe.shape[1:3]).astype(int) // 2 - psr = int(probe_support_radius) - probe_support[p_cent[0]-psr:p_cent[0]+psr, - p_cent[1]-psr:p_cent[1]+psr] = 1 + 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= 2: + Ws = self.weights.detach().cpu().numpy() + rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj()) + return rhos_out + # This is the purely incoherent case + else: + return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0], + dtype=np.complex64) + + def tidy_probes(self, normalization=1, normalize=False): + """Tidies up the probes + + What we want to do here is use all the information on all the probes + to calculate a natural basis for the experiment, and update all the + density matrices to operate in that updated basis + + """ + + # First we treat the purely incoherent case + + # I don't love this pattern of using an if statement with a return + # to catch this case, but because it's so much simpler than the + # unified mode case I think it's appropriate + if self.weights.dim() == 1: + probe = self.probe.detach().cpu().numpy() + ortho_probes = analysis.orthogonalize_probes(probe) + self.probe.data = t.as_tensor(ortho_probes, + device=self.probe.device,dtype=self.probe.dtype) + return - # Needs to be updated to allow for plotting to an existing figure + # 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 = [ - ('Dominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Dominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Subdominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Subdominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), + ('', + 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)), + 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)), + 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)), + 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)) ] @@ -366,9 +547,9 @@ class FancyPtycho(CDIModel): 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 = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() - obj = cmath.torch_to_complex(self.obj.detach().cpu()) + obj = self.obj.detach().cpu().numpy() background = self.background.detach().cpu().numpy()**2 weights = self.weights.detach().cpu().numpy() diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index 0032044..a996f0b 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -1,20 +1,29 @@ -from __future__ import division, print_function, absolute_import - 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 analysis, image_processing 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 +from functools import reduce __all__ = ['Multislice2DPtycho'] + + class Multislice2DPtycho(CDIModel): + @property + def probe(self): + return t.complex(self.probe_real,self.probe_imag) + + @property + def obj(self): + return t.complex(self.obj_real,self.obj_imag) + def __init__(self, wavelength, detector_geometry, probe_basis, probe_guess, obj_guess, dz, nz, @@ -23,15 +32,20 @@ class Multislice2DPtycho(CDIModel): min_translation = t.Tensor([0,0]), background = None, translation_offsets=None, mask=None, weights = None, translation_scale = 1, saturation=None, - #probe_support = None, - probe_fourier_support=None, - obj_support=None, oversampling=1, - bandlimit=4/5): + probe_support = None, + oversampling=1, + bandlimit=None, + subpixel=True, + exponentiate_obj=True, + fourier_probe=False, + prevent_aliasing=True, + phase_only=False, + units='um'): super(Multislice2DPtycho,self).__init__() self.wavelength = t.Tensor([wavelength]) self.detector_geometry = copy(detector_geometry) - self.dz = -dz + self.dz = dz self.nz = nz det_geo = self.detector_geometry if hasattr(det_geo, 'distance'): @@ -48,6 +62,12 @@ class Multislice2DPtycho(CDIModel): self.surface_normal = t.Tensor(surface_normal) self.saturation = saturation + self.subpixel = subpixel + self.exponentiate_obj = exponentiate_obj + self.fourier_probe = fourier_probe + self.units = units + self.phase_only=phase_only + self.prevent_aliasing=prevent_aliasing if mask is None: self.mask = mask @@ -57,66 +77,80 @@ class Multislice2DPtycho(CDIModel): # We rescale the probe here so it learns at the same rate as the # object if probe_guess.dim() > 3: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) + self.probe_norm = 1 * t.max(t.abs(probe_guess[0]).to(t.float32)) else: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32))) - - self.probe = t.nn.Parameter(probe_guess.to(t.float32) - / self.probe_norm) + self.probe_norm = 1 * t.max(t.abs(probe_guess).to(t.float32)) + + pg = probe_guess.to(t.complex64)/self.probe_norm + self.probe_real = t.nn.Parameter(pg.real) + self.probe_imag = t.nn.Parameter(pg.imag) + + og = obj_guess.to(t.complex64) + self.obj_real = t.nn.Parameter(og.real) + self.obj_imag = t.nn.Parameter(og.imag) - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + #self.probe = t.nn.Parameter(probe_guess.to(t.complex64) + # / self.probe_norm) + + #self.obj = t.nn.Parameter(obj_guess.to(t.complex64)) + if background is None: if detector_slice is not None: - background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1]) + background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape) else: - background = 1e-6 * t.ones(self.probe[0].shape[:-1]) + background = 1e-6 * t.ones(self.probe[0].shape) - self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) + self.background = t.nn.Parameter(t.as_tensor(background,dtype=t.float32)) if weights is None: self.weights = None else: - self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32)) + # 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.as_tensor(weights, + dtype=t.float32)) + else: + # Now this is a matrix of weights, so we + self.weights = t.nn.Parameter(t.as_tensor(weights, + dtype=t.complex64)) 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) + self.translation_offsets = t.nn.Parameter(t.as_tensor(translation_offsets,dtype=t.float32)/ translation_scale) self.translation_scale = translation_scale - self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32) - #if probe_support is not None: - # self.probe_support = probe_support - #else: - # self.probe_support = t.ones_like(self.probe[0]) - - if obj_support is not None: - self.obj_support = obj_support - if self.obj.dim() == 3: - self.obj.data = self.obj * obj_support - elif self.obj.dim() == 4: - self.obj.data = self.obj * obj_support[None,...] + if probe_support is not None: + self.probe_support = t.as_tensor(probe_support,dtype=t.bool) else: - if self.obj.dim() == 3: - self.obj_support = t.ones_like(self.obj) - elif self.obj.dim() == 4: - self.obj_support = t.ones_like(self.obj[0]) + self.probe_support = None#t.ones_like(self.probe,dtype=t.bool)#None + self.oversampling = oversampling spacing = np.linalg.norm(self.probe_basis,axis=0) - shape = np.array(self.probe.shape[1:-1]) - + shape = np.array(self.probe.shape[1:]) + if prevent_aliasing: + shape *= 2 + spacing /= 2 + self.bandlimit = bandlimit - # Big question: Should there be a minus sign before self.dz, or not? - self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=self.bandlimit) + + self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=1/np.sqrt(2))#self.bandlimit) + #plt.imshow(t.abs(self.as_prop)) + #plt.figure() + #plt.imshow(t.abs(t.fft.fftshift(t.fft.ifft2(self.as_prop)))) + #plt.show() + #exit() @classmethod - def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=4/5, replicate_slice=False): + def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -165,12 +199,11 @@ class Multislice2DPtycho(CDIModel): 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) - - obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) + + obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=100) if hasattr(dataset, 'background') and dataset.background is not None: background = t.sqrt(dataset.background) @@ -178,32 +211,53 @@ class Multislice2DPtycho(CDIModel): 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) + probe = tools.initializers.STEM_style_probe(dataset, probe_shape, det_slice, probe_convergence_semiangle, propagation_distance=propagation_distance, oversampling=oversampling) + #probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) + + # Now we initialize all the subdominant probe modes + probe_max = t.max(t.abs(probe)) + if n_modes >=2: + probe_stack = list(0.01*tools.initializers.generate_subdominant_modes(probe,n_modes-1,circular=False)) + #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) else: - probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance) + probe = t.unsqueeze(probe,0) # For a Fourier space probe - probe = tools.propagators.inverse_far_field(probe) - - # 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([probe,] + probe_stack) + if fourier_probe: + probe = tools.propagators.far_field(probe) + # Consider a different start - obj = t.zeros(obj_size+(2,)) - #obj = tools.cmath.expi(t.zeros(obj_size)) + if exponentiate_obj: + obj = t.zeros(obj_size, dtype=t.complex64) + else: + obj = t.exp(1j*t.zeros(obj_size)) # If we will use a separate object per slice if not replicate_slice: obj = t.stack([obj]*nz) - det_geo = dataset.detector_geometry translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) - weights = t.ones(len(dataset)) + 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,2) + # 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,0] = 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) @@ -211,51 +265,36 @@ class Multislice2DPtycho(CDIModel): mask = None if probe_support_radius is not None: - probe_support = t.zeros_like(probe[0].to(dtype=t.float32)) - p_cent = np.array(probe.shape[1:3]).astype(int) // 2 - psr = int(probe_support_radius) - probe_support[p_cent[0]-psr:p_cent[0]+psr, - p_cent[1]-psr:p_cent[1]+psr] = 1 - probe = probe * probe_support[None,:,:] - else: - probe_support = None; - - if restrict_obj != -1: - ro = restrict_obj - os = np.array(obj_size) - ps = np.array(probe_shape) - if replicate_slice: - obj_support = t.zeros_like(obj.to(dtype=t.float32)) - else: - obj_support = t.zeros_like(obj[0].to(dtype=t.float32)) - obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2, - ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1 - else: - obj_support = None - - probe_support = t.zeros_like(probe[0].to(dtype=t.float32)) - - xs, ys = np.mgrid[:probe.shape[-3],:probe.shape[-2]] - xs = xs - np.mean(xs) - ys = ys - np.mean(ys) - Rs = np.sqrt(xs**2 + ys**2) + 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= 1: + # plt.imshow(t.abs(tools.propagators.far_field( + # exit_waves[0,0].detach()).cpu())) + # plt.colorbar() + # plt.show() + if i < self.nz-1: #on all but the last iteration + exit_waves = tools.propagators.near_field( + exit_waves,self.as_prop) + #plt.imshow(t.abs(exit_waves[0,0].detach().cpu())) + #plt.show() + + 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) + if self.prevent_aliasing: + left = [self.probe.shape[-2]//2,self.probe.shape[-1]//2] + right = [self.probe.shape[-2]//2+self.probe.shape[-2], + self.probe.shape[-1]//2+self.probe.shape[-1]] + + return tools.propagators.far_field(wavefields)[...,left[0]:right[0], + left[1]:right[1]] + else: + return tools.propagators.far_field(wavefields) def measurement(self, wavefields): @@ -341,11 +409,8 @@ class Multislice2DPtycho(CDIModel): def loss(self, sim_data, real_data, mask=None): - regularizer = t.mean(t.abs(self.obj)) - loss = tools.losses.amplitude_mse(real_data, sim_data, mask=mask) - lambd = 1 - return loss + lambd * regularizer - #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) + return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) + #return tools.losses.poisson_nll(real_data, sim_data, mask=mask,eps=0.5) def to(self, *args, **kwargs): @@ -362,14 +427,14 @@ class Multislice2DPtycho(CDIModel): if self.mask is not None: self.mask = self.mask.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.probe_fourier_support = self.probe_fourier_support.to(*args,**kwargs) - self.obj_support = self.obj_support.to(*args,**kwargs) + if self.probe_support is not None: + self.probe_support = self.probe_support.to(*args,**kwargs) + + self.surface_normal = self.surface_normal.to(*args, **kwargs) self.as_prop = self.as_prop.to(*args, **kwargs) @@ -412,29 +477,134 @@ class Multislice2DPtycho(CDIModel): def corrected_translations(self,dataset): - translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device) + translations = dataset.translations.to(dtype=self.probe.real.dtype,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 + + """ + + # 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) + # Needs to be updated to allow for plotting to an existing figure plot_list = [ - ('Dominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Dominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Subdominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Subdominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), + ('Probe Fourier Space Amplitude', + lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Probe Fourier Space Phase', + lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Probe Real Space Amplitude', + lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), + ('Probe Real Space Phase', + lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), + ('Average Density Matrix Amplitudes', + lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()),axis=0), fig=fig), + lambda self: len(self.weights.shape) >=2), + ('% Power in Top Mode (only accurate after tidy_probes)', + lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig,units=self.units), + lambda self: len(self.weights.shape) >=2), + ('Slice by Slice Real Part of T', + lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), + lambda self: self.exponentiate_obj), + ('Slice by Slice Imaginary Part of T', + lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), + lambda self: self.exponentiate_obj), ('Integrated Real Part of T', - lambda self, fig: p.plot_real(t.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)), + lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), + lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3), ('Integrated Imaginary Part of T', - lambda self, fig: p.plot_imag(t.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)), + lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units), + lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3), + ('Slice by Slice Amplitude of Object Function', + lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), + lambda self: not self.exponentiate_obj), + ('Slice by Slice Phase of Object Function', + lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'), + lambda self: not self.exponentiate_obj), + ('Amplitude of Stacked Object Function', + lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units), + lambda self: (not self.exponentiate_obj) and self.obj.dim() >=3), + ('Phase of Stacked Object Function', + lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), + lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3), ('Corrected Translations', - lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)), + 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)) ] @@ -443,9 +613,14 @@ class Multislice2DPtycho(CDIModel): 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()) + if self.fourier_probe: + probe = tools.propagators.inverse_far_field(self.probe) + else: + probe = self.probe + + probe = probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() - obj = cmath.torch_to_complex(self.obj.detach().cpu()) + obj = self.obj.detach().cpu().numpy() background = self.background.detach().cpu().numpy()**2 weights = self.weights.detach().cpu().numpy() dz = self.dz diff --git a/CDTools/models/pinhole_plane_ptycho.py b/CDTools/models/pinhole_plane_ptycho.py deleted file mode 100644 index 9cb9d3d..0000000 --- a/CDTools/models/pinhole_plane_ptycho.py +++ /dev/null @@ -1,384 +0,0 @@ -from __future__ import division, print_function, absolute_import - -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__ = ['PinholePlanePtycho'] - -class PinholePlanePtycho(CDIModel): - - def __init__(self, wavelength, detector_geometry, - probe_basis, - pinhole_guess, obj_guess, illumination_guess, - pinhole_propagation, - detector_slice=None, - surface_normal=np.array([0.,0.,1.]), - min_translation = t.Tensor([0,0]), - background = None, translation_offsets=None, mask=None, - weights = None, translation_scale = 1, saturation=None, - pinhole_support = None, obj_support=None, oversampling=1): - - super(PinholePlanePtycho,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.ByteTensor(mask) - - # We rescale the pinhole here so it learns at the same rate as the - # object - self.pinhole_norm = 1 * t.max(tools.cmath.cabs(pinhole_guess.to(t.float32))) - - self.pinhole = t.nn.Parameter(pinhole_guess.to(t.float32) - / self.pinhole_norm) - - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) - self.illumination = t.nn.Parameter(illumination_guess.to(t.float32)) - - spacing = np.abs([self.probe_basis.numpy()[0,1], self.probe_basis.numpy()[1,0]]) - - - self.pinhole_propagator = tools.propagators.generate_angular_spectrum_propagator(self.pinhole.shape[:-1], spacing, self.wavelength, pinhole_propagation).to(dtype=self.pinhole.dtype) - - - if background is None: - if detector_slice is not None: - background = 1e-6 * t.ones(self.pinhole[self.detector_slice].shape[:-1]) - else: - background = 1e-6 * t.ones(self.pinhole.shape[:-1]) - - - 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) - - self.translation_scale = translation_scale - - if pinhole_support is not None: - self.pinhole_support = pinhole_support - else: - self.pinhole_support = t.ones_like(self.pinhole) - - if obj_support is not None: - self.obj_support = obj_support - self.obj.data = self.obj * obj_support - else: - self.obj_support = t.ones_like(self.obj) - - self.oversampling = oversampling - - - @classmethod - def from_dataset(cls, dataset, pinhole_propagation, pinhole_size=None, randomize_ang=0, padding=0, translation_scale = 1, saturation=None, pinhole_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1): - - 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') - (indices, translations), patterns = dataset[:] - dataset.get_as(*get_as_args[0],**get_as_args[1]) - - # Set to none to avoid issues with things outside the detector - center = tools.image_processing.centroid(t.sum(patterns,dim=0)) - - # Then, generate the pinhole geometry from the dataset - ewg = tools.initializers.exit_wave_geometry - probe_basis, pinhole_shape, det_slice = ewg(det_basis, - det_shape, - wavelength, - distance, - center=center, - padding=padding, - opt_for_fft=False, - 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(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) - - obj_size, min_translation = tools.initializers.calc_object_setup(pinhole_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 pinhole and object using this information - if pinhole_size is None: - pinhole = tools.initializers.SHARP_style_probe(dataset, pinhole_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) - else: - pinhole = tools.initializers.gaussian_probe(dataset, probe_basis, pinhole_shape, pinhole_size, propagation_distance=propagation_distance) - - - illumination = tools.cmath.expi(t.zeros(obj_size)) - obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) - - det_geo = dataset.detector_geometry - - translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) - - weights = t.ones(len(dataset)) - - if hasattr(dataset, 'mask') and dataset.mask is not None: - mask = dataset.mask.to(t.uint8) - else: - mask = None - - if pinhole_support_radius is not None: - pinhole_support = t.zeros_like(pinhole.to(dtype=t.float32)) - p_cent = np.array(pinhole.shape[:2]).astype(int) // 2 - psr = int(pinhole_support_radius) - pinhole_support[p_cent[0]-psr:p_cent[0]+psr, - p_cent[1]-psr:p_cent[1]+psr] = 1 - pinhole = pinhole_support * pinhole - else: - pinhole_support = None; - - if restrict_obj != -1: - ro = restrict_obj - os = np.array(obj_size) - ps = np.array(pinhole_shape) - obj_support = t.zeros_like(obj.to(dtype=t.float32)) - obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2, - ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1 - else: - obj_support = None - - return cls(wavelength, det_geo, probe_basis, pinhole, obj, - illumination, pinhole_propagation, - detector_slice=det_slice, - surface_normal=surface_normal, - min_translation=min_translation, - translation_offsets = translation_offsets, - weights=weights, mask=mask, background=background, - translation_scale=translation_scale, - saturation=saturation, - pinhole_support=pinhole_support, - obj_support=obj_support, - oversampling=oversampling) - - - def interaction(self, index, translations): - pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, - translations, - surface_normal=self.surface_normal) - pix_trans -= self.min_translation - - if self.translation_offsets is not None: - pix_trans += self.translation_scale * self.translation_offsets[index] - - pr = self.pinhole * self.pinhole_support - pinhole_exit_waves = self.pinhole_norm * \ - tools.interactions.ptycho_2D_sinc(pr, - self.illumination, - pix_trans, - shift_probe=True) - sample_illumination = tools.propagators.near_field(pinhole_exit_waves, - self.pinhole_propagator) - exit_waves = [] - if pinhole_exit_waves.dim() == 4: - for pix_tran, illum in zip(pix_trans, sample_illumination): - exit_waves.append(tools.interactions.ptycho_2D_sinc(illum, - self.obj_support * self.obj, - pix_tran, - shift_probe=True)) - exit_waves = t.stack(exit_waves) - else: - exit_waves = tools.interactions.ptycho_2D_sinc(sample_illumination, - self.obj_support * self.obj, - pix_trans, - shift_probe=True) - - - - if exit_waves.dim() == 4: - exit_waves = self.weights[index][:,None,None,None] * exit_waves - else: - exit_waves = self.weights[index] * exit_waves - - - 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, - 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) - - - def to(self, *args, **kwargs): - super(PinholePlanePtycho, 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.min_translation = self.min_translation.to(*args,**kwargs) - self.probe_basis = self.probe_basis.to(*args,**kwargs) - self.pinhole_norm = self.pinhole_norm.to(*args,**kwargs) - self.pinhole_support = self.pinhole_support.to(*args,**kwargs) - self.obj_support = self.obj_support.to(*args,**kwargs) - self.surface_normal = self.surface_normal.to(*args, **kwargs) - self.pinhole_propagator = self.pinhole_propagator.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=self.pinhole.dtype,device=self.pinhole.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 - - - # Needs to be updated to allow for plotting to an existing figure - plot_list = [ - ('Pinhole Amplitude', - lambda self, fig: p.plot_amplitude(self.pinhole, fig=fig, basis=self.probe_basis)), - ('Pinhole Phase', - lambda self, fig: p.plot_phase(self.pinhole, fig=fig, basis=self.probe_basis)), - ('Illumination Amplitude', - lambda self, fig: p.plot_amplitude(self.illumination, fig=fig, basis=self.probe_basis), - lambda self: len(self.pinhole) >=2), - ('Illumination Phase', - lambda self, fig: p.plot_phase(self.illumination, fig=fig, basis=self.probe_basis), - lambda self: len(self.pinhole) >=2), - ('Object Amplitude', - lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), - ('Object Phase', - lambda self, fig: p.plot_phase(self.obj, 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.pinhole.detach().cpu()) - probe = probe * self.pinhole_norm.detach().cpu().numpy() - obj = cmath.torch_to_complex(self.obj.detach().cpu()) - 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/CDTools/models/rpi.py b/CDTools/models/rpi.py index 7bc35bf..3bd9e20 100644 --- a/CDTools/models/rpi.py +++ b/CDTools/models/rpi.py @@ -4,7 +4,6 @@ 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 CDTools.tools.interactions import RPI_interaction from CDTools.tools import initializers @@ -51,6 +50,10 @@ __all__ = ['RPI'] __all__ = ['RPI'] class RPI(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, detector_slice=None, @@ -73,8 +76,8 @@ class RPI(CDIModel): self.probe_basis = t.Tensor(probe_basis) - scale_factor = t.Tensor([probe.shape[-2]/obj_guess.shape[-2], - probe.shape[-3]/obj_guess.shape[-3]]) + 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 @@ -89,12 +92,17 @@ class RPI(CDIModel): self.mask = t.BoolTensor(mask) - self.probe = probe.to(t.float32) + self.probe = probe.to(t.complex64) - if obj_guess.dim() == 3: - obj_guess = obj_guess[None,:,:,:] + if obj_guess.dim() == 2: + obj_guess = obj_guess[None,:,:] - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + + self.obj_real = t.nn.Parameter(obj_guess.real.to(t.float32)) + self.obj_imag = t.nn.Parameter(obj_guess.imag.to(t.float32)) + + # 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: @@ -146,7 +154,7 @@ class RPI(CDIModel): oversampling=oversampling) if not isinstance(probe,t.Tensor): - probe = cmath.complex_to_torch(probe) + probe = t.as_tensor(probe) # Potentially need all of this orientation stuff later @@ -190,8 +198,8 @@ class RPI(CDIModel): # 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(cmath.cabssq(probe)) - obj_guess = scale * cmath.expi(2 * np.pi * t.rand([n_modes,]+obj_size)) + 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( @@ -218,9 +226,9 @@ class RPI(CDIModel): def random_init(self, pattern): - scale = t.sum(pattern) / t.sum(cmath.cabssq(self.probe)) - self.obj.data = scale * cmath.expi( - 2 * np.pi * t.rand(self.obj.shape[:-1])).to( + 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): @@ -249,8 +257,8 @@ class RPI(CDIModel): pr = self.probe[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[None,:,:] * self.obj) + exit_waves = RPI_interaction(pr, + self.obj_support * self.obj) all_exit_waves.append(exit_waves) @@ -266,7 +274,7 @@ class RPI(CDIModel): 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(1).repeat(1,len(index),1,1,1) + output = output.unsqueeze(0).repeat(1,len(index),1,1,1) except TypeError: pass @@ -297,8 +305,8 @@ class RPI(CDIModel): #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) def regularizer(self, factors): - return factors[0] * t.sum(cmath.cabssq(self.obj[0,:,:,:])) \ - + factors[1] * t.sum(cmath.cabssq(self.obj[1:,:,:,:])) + 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(RPI, self).to(*args, **kwargs) @@ -330,7 +338,7 @@ class RPI(CDIModel): plot_list = [ ('Root Sum Squared Amplitude of all Probes', lambda self, fig: p.plot_amplitude( - np.sqrt(np.sum(cmath.cabssq(self.probe).cpu().numpy(),axis=0)), + np.sqrt(np.sum((t.abs(self.probe)**2).cpu().numpy(),axis=0)), fig=fig, basis=self.probe_basis)), ('Dominant Object Amplitude', lambda self, fig: p.plot_amplitude(self.obj[0], fig=fig, @@ -355,13 +363,13 @@ class RPI(CDIModel): # continues to use that standard pattern probe_basis = self.probe_basis.detach().cpu().numpy() obj_basis = self.obj_basis.detach().cpu().numpy() - probe = cmath.torch_to_complex(self.probe.detach().cpu()) + probe = self.probe.detach().cpu().numpy() # Provide the option to save out the subdominant objects or # just the dominant one if full_obj: - obj = cmath.torch_to_complex(self.obj.detach().cpu()) + obj = self.obj.detach().cpu().numpy() else: - obj = cmath.torch_to_complex(self.obj[0].detach().cpu()) + obj = self.obj[0].detach().cpu().numpy() background = self.background.detach().cpu().numpy()**2 return {'probe_basis': probe_basis, 'obj_basis': obj_basis, diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index baa662c..a178d78 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -10,6 +10,7 @@ from torch.utils import data as torchdata from matplotlib import pyplot as plt from datetime import datetime import numpy as np +from .complex_adam import MyAdam __all__ = ['SimplePtycho'] @@ -25,22 +26,22 @@ class SimplePtycho(CDIModel): surface_normal=np.array([0.,0.,1.]), mask=None): super(SimplePtycho,self).__init__() - self.wavelength = t.Tensor([wavelength]) + 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']) + det_geo['distance'] = t.tensor(det_geo['distance']) if hasattr(det_geo, 'basis'): - det_geo['basis'] = t.Tensor(det_geo['basis']) + det_geo['basis'] = t.tensor(det_geo['basis']) if hasattr(det_geo, 'corner'): - det_geo['corner'] = t.Tensor(det_geo['corner']) + det_geo['corner'] = t.tensor(det_geo['corner']) - self.min_translation = t.Tensor(min_translation) + self.min_translation = t.tensor(min_translation) - self.probe_basis = t.Tensor(probe_basis) + self.probe_basis = t.tensor(probe_basis) self.detector_slice = detector_slice - self.surface_normal = t.Tensor(surface_normal) + self.surface_normal = t.tensor(surface_normal) if mask is None: self.mask = None @@ -49,11 +50,11 @@ class SimplePtycho(CDIModel): # We rescale the probe here so it learns at the same rate as the # object - self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32))) + self.probe_norm = t.max(t.abs(probe_guess.to(t.complex64))) - self.probe = t.nn.Parameter(probe_guess.to(t.float32) + self.probe = t.nn.Parameter(probe_guess.to(t.complex64) / self.probe_norm) - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + self.obj = t.nn.Parameter(obj_guess.to(t.complex64)) @@ -94,8 +95,8 @@ class SimplePtycho(CDIModel): # Finally, initialize the probe and object using this information probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice) - - obj = t.ones(obj_size+(2,)) + + obj = t.ones(obj_size).to(dtype=t.complex64) det_geo = dataset.detector_geometry @@ -112,6 +113,7 @@ class SimplePtycho(CDIModel): translations, surface_normal=self.surface_normal) pix_trans -= self.min_translation + return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe, self.obj, pix_trans) @@ -130,7 +132,7 @@ class SimplePtycho(CDIModel): detector_slice=self.detector_slice) - def loss(self, sim_data, real_data, mask=None): + def loss(self, real_data, sim_data, mask=None): return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) @@ -206,9 +208,9 @@ class SimplePtycho(CDIModel): def save_results(self): - probe = tools.cmath.torch_to_complex(self.probe.detach().cpu()) + probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() - obj = tools.cmath.torch_to_complex(self.obj.detach().cpu()) + obj = self.obj.detach().cpu().numpy() return {'probe':probe,'obj':obj} @@ -272,4 +274,4 @@ class SimplePtycho(CDIModel): # Calculate loss loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns)) - yield t.mean(t.Tensor(loss)).cpu().numpy() + yield t.mean(t.tensor(loss)).cpu().numpy() diff --git a/CDTools/models/unified_mode_ptycho.py b/CDTools/models/unified_mode_ptycho.py deleted file mode 100644 index afaa19d..0000000 --- a/CDTools/models/unified_mode_ptycho.py +++ /dev/null @@ -1,449 +0,0 @@ -from __future__ import division, print_function, absolute_import - -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 analysis -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__ = ['UnifiedModePtycho'] - -class UnifiedModePtycho(CDIModel): - - def __init__(self, wavelength, detector_geometry, - probe_basis, - probe_guess, obj_guess, rhos_guess, - detector_slice=None, - surface_normal=np.array([0.,0.,1.]), - min_translation = t.Tensor([0,0]), - background = None, translation_offsets=None, mask=None, - translation_scale = 1, saturation=None, - probe_support = None, obj_support=None, oversampling=1): - - super(UnifiedModePtycho,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 - if probe_guess.dim() > 3: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) - else: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32))) - - self.probe = t.nn.Parameter(probe_guess.to(t.float32) - / self.probe_norm) - - 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[:-1]) - else: - background = 1e-6 * t.ones(self.probe[0].shape[:-1]) - - - self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) - - self.rhos = t.nn.Parameter(t.Tensor(rhos_guess).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) - - 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]) - - if obj_support is not None: - self.obj_support = obj_support - self.obj.data = self.obj * obj_support - else: - self.obj_support = t.ones_like(self.obj) - - self.oversampling = oversampling - - - @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, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, mixing_mode='unified'): - - 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') - (indices, translations), patterns = dataset[:] - dataset.get_as(*get_as_args[0],**get_as_args[1]) - - # Set to none to avoid issues with things outside the detector - if auto_center: - center = tools.image_processing.centroid(t.sum(patterns,dim=0)) - else: - center = None - - # Then, generate the probe geometry from the dataset - ewg = tools.initializers.exit_wave_geometry - probe_basis, probe_shape, det_slice = ewg(det_basis, - det_shape, - wavelength, - distance, - center=center, - padding=padding, - opt_for_fft=opt_for_fft, - oversampling=oversampling) - - - if hasattr(dataset, 'sample_info') and \ - dataset.sample_info is not None and \ - 'orientation' in dataset.sample_info: - surface_normal = dataset.sample_info['orientation'][2] - else: - surface_normal = np.array([0.,0.,1.]) - - - # If this information is supplied when the function is called, - # then we override the information in the .cxi file - if scattering_mode in {'t', 'transmission'}: - surface_normal = np.array([0.,0.,1.]) - elif scattering_mode in {'r', 'reflection'}: - outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) - outgoing_dir /= np.linalg.norm(outgoing_dir) - surface_normal = outgoing_dir + np.array([0.,0.,1.]) - surface_normal /= -np.linalg.norm(surface_normal) - - - # Next generate the object geometry from the probe geometry and - # the translations - pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal) - - obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) - - if hasattr(dataset, 'background') and dataset.background is not None: - background = t.sqrt(dataset.background) - else: - background = None - - # Finally, initialize the probe and object using this information - if probe_size is None: - probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) - else: - probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance) - - - # Now we initialize all the subdominant probe modes - probe_max = t.max(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([probe,] + probe_stack) - #probe = t.stack([tools.propagators.far_field(probe),] + probe_stack) - - obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) - - det_geo = dataset.detector_geometry - - translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) - - # - if mixing_mode.lower().strip() == 'unified': - rhos = t.zeros(len(dataset),n_modes,n_modes) - rhos[:,0,0] = 1 - for i in range(1,n_modes): - rhos[:,i,i] = 1/n_modes - - 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_like(probe[0].to(dtype=t.float32)) - p_cent = np.array(probe.shape[1:3]).astype(int) // 2 - psr = int(probe_support_radius) - probe_support[p_cent[0]-psr:p_cent[0]+psr, - p_cent[1]-psr:p_cent[1]+psr] = 1 - probe = probe * probe_support[None,:,:] - else: - probe_support = None; - - if restrict_obj != -1: - ro = restrict_obj - os = np.array(obj_size) - ps = np.array(probe_shape) - obj_support = t.zeros_like(obj.to(dtype=t.float32)) - obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2, - ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1 - else: - obj_support = None - - return cls(wavelength, det_geo, probe_basis, probe, obj, rhos, - detector_slice=det_slice, - surface_normal=surface_normal, - min_translation=min_translation, - translation_offsets = translation_offsets, - mask=mask, background=background, - translation_scale=translation_scale, - saturation=saturation, - probe_support=probe_support, - obj_support=obj_support, - oversampling=oversampling) - - - def interaction(self, index, translations): - pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, - translations, - surface_normal=self.surface_normal) - pix_trans -= self.min_translation - - if self.translation_offsets is not None: - pix_trans += self.translation_scale * self.translation_offsets[index] - - probes = [] - all_exit_waves = [] - for i in range(self.probe.shape[0]): - # from storing the probe in Fourier space - #pr = tools.propagators.inverse_far_field(self.probe[i]) * self.probe_support - pr = self.probe[i] * self.probe_support - exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr, - self.obj_support * self.obj, - pix_trans, - shift_probe=True) - exit_waves = exit_waves * self.probe_support[...,:,:] - - all_exit_waves.append(exit_waves) - - - return t.stack(all_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, indices): - #return tools.measurements.density_matrix(wavefields,self.rhos[indices], - # detector_slice=self.detector_slice, - # saturation=self.saturation, - # oversampling=self.oversampling) - return tools.measurements.quadratic_background(wavefields, - self.background, self.rhos[indices], - detector_slice=self.detector_slice, - measurement=tools.measurements.density_matrix, - saturation=self.saturation, - oversampling=self.oversampling) - - - def forward(self, *args): - """The complete forward model - - We need to override this to enable the wavefield mixing at the - level of the measurement function - """ - indices = args[0] - return self.measurement(self.forward_propagator(self.interaction(*args)),indices) - - 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 to(self, *args, **kwargs): - super(UnifiedModePtycho, 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.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.obj_support = self.obj_support.to(*args,**kwargs) - self.surface_normal = self.surface_normal.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 get_rhos(self): - - rhos_out = np.zeros([self.rhos.shape[0], - self.rhos.shape[1],self.rhos.shape[2]], - dtype=np.complex64) - for (i,j) in ((i,j) for i in range(rhos_out.shape[1]) - for j in range(rhos_out.shape[2])): - if i == j: - rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy() - if i < j: # upper triangle, real part - rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy() - rhos_out[:,j,i] += self.rhos.data[:,i,j].cpu().detach().numpy() - if i > j: # upper triangle, real part - rhos_out[:,j,i] += 1j * self.rhos.data[:,i,j].cpu().detach().numpy() - rhos_out[:,i,j] -= 1j * self.rhos.data[:,i,j].cpu().detach().numpy() - return rhos_out - - def tidy_probes(self, normalization=1): - """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 - - """ - rhos = self.get_rhos() - overall_rho = np.mean(rhos,axis=0) - probe = cmath.torch_to_complex(self.probe.detach().cpu()) - ortho_probes, A = analysis.orthogonalize_probes(probe, - density_matrix=overall_rho, - keep_transform=True, - normalize=True) - Aconj = A.conj() - Atrans = np.transpose(A) - new_rhos = np.swapaxes(np.dot(Atrans,np.dot(rhos,Aconj)),0,1) - - new_rhos /= normalization - ortho_probes *= np.sqrt(normalization) - #print(np.dot(Aconjinv,rhos).shape) - #print(np.dot(rhos,Aconj).shape) - new_rhos = cmath.complex_to_torch(new_rhos).to( - dtype=self.rhos.dtype,device=self.rhos.device) - - # This repacks the data into the format used internally - for (i,j) in ((i,j) for i in range(new_rhos.shape[1]) - for j in range(new_rhos.shape[2])): - if i == j: - self.rhos.data[:,i,j] = new_rhos[:,i,j,0] - if i < j: # upper triangle, real part - self.rhos.data[:,i,j] = new_rhos[:,i,j,0] - self.rhos.data[:,j,i] = new_rhos[:,i,j,1] - - - self.probe.data = cmath.complex_to_torch(ortho_probes).to( - device=self.probe.device,dtype=self.probe.dtype) - - - def corrected_translations(self,dataset): - translations = dataset.translations.to(dtype=self.probe.dtype,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 - - - # Needs to be updated to allow for plotting to an existing figure - plot_list = [ - ('Dominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Dominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Subdominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Subdominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Average Density Matrix Amplitudes', - lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig), - lambda self: len(self.probe) >=2), - ('Von Neumann Entropy (only accurate after tidy_probes)', - lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_vn_entropy(self.get_rhos()), fig=fig)), - ('Object Amplitude', - lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), - ('Object Phase', - lambda self, fig: p.plot_phase(self.obj, 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() - obj = cmath.torch_to_complex(self.obj.detach().cpu()) - 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/CDTools/models/unified_mode_ptycho2.py b/CDTools/models/unified_mode_ptycho2.py deleted file mode 100644 index dbab1f5..0000000 --- a/CDTools/models/unified_mode_ptycho2.py +++ /dev/null @@ -1,490 +0,0 @@ -from __future__ import division, print_function, absolute_import - -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 analysis -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__ = ['UnifiedModePtycho2'] - -class UnifiedModePtycho2(CDIModel): - - def __init__(self, wavelength, detector_geometry, - probe_basis, - probe_guess, obj_guess, Ws_guess, - detector_slice=None, - surface_normal=np.array([0.,0.,1.]), - min_translation = t.Tensor([0,0]), - background = None, translation_offsets=None, mask=None, - translation_scale = 1, saturation=None, - probe_support = None, obj_support=None, oversampling=1): - - super(UnifiedModePtycho2,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 - if probe_guess.dim() > 3: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) - else: - self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32))) - - self.probe = t.nn.Parameter(probe_guess.to(t.float32) - / self.probe_norm) - - 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[:-1]) - else: - background = 1e-6 * t.ones(self.probe[0].shape[:-1]) - - - self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) - - if type(Ws_guess) == type(t.zeros(1)): - self.Ws = t.nn.Parameter(Ws_guess.to(t.float32)) - else: - self.Ws = t.nn.Parameter(cmath.complex_to_torch(Ws_guess).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) - - 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]) - - if obj_support is not None: - self.obj_support = obj_support - self.obj.data = self.obj * obj_support - else: - self.obj_support = t.ones_like(self.obj) - - self.oversampling = oversampling - - - @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, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, dm_rank=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') - (indices, translations), patterns = dataset[:] - dataset.get_as(*get_as_args[0],**get_as_args[1]) - - # Set to none to avoid issues with things outside the detector - if auto_center: - center = tools.image_processing.centroid(t.sum(patterns,dim=0)) - else: - center = None - - # Then, generate the probe geometry from the dataset - ewg = tools.initializers.exit_wave_geometry - probe_basis, probe_shape, det_slice = ewg(det_basis, - det_shape, - wavelength, - distance, - center=center, - padding=padding, - opt_for_fft=opt_for_fft, - oversampling=oversampling) - - - if hasattr(dataset, 'sample_info') and \ - dataset.sample_info is not None and \ - 'orientation' in dataset.sample_info: - surface_normal = dataset.sample_info['orientation'][2] - else: - surface_normal = np.array([0.,0.,1.]) - - - # If this information is supplied when the function is called, - # then we override the information in the .cxi file - if scattering_mode in {'t', 'transmission'}: - surface_normal = np.array([0.,0.,1.]) - elif scattering_mode in {'r', 'reflection'}: - outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) - outgoing_dir /= np.linalg.norm(outgoing_dir) - surface_normal = outgoing_dir + np.array([0.,0.,1.]) - surface_normal /= -np.linalg.norm(surface_normal) - - - # Next generate the object geometry from the probe geometry and - # the translations - pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal) - - obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) - - if hasattr(dataset, 'background') and dataset.background is not None: - background = t.sqrt(dataset.background) - else: - background = None - - # Finally, initialize the probe and object using this information - if probe_size is None: - probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) - else: - probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance) - - - # Now we initialize all the subdominant probe modes - probe_max = t.max(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([probe,] + probe_stack) - #probe = t.stack([tools.propagators.far_field(probe),] + probe_stack) - - obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) - - det_geo = dataset.detector_geometry - - translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) - - # dm_rank defines the rank of the shot-by-shot density matrices - if dm_rank > n_modes: - raise KeyError('Density matrix rank cannot be greater than the number of modes') - elif dm_rank != 0: - if dm_rank == -1: - dm_rank = n_modes - Ws = t.zeros(len(dataset),dm_rank,n_modes,2) - Ws[:,0,0,0] = 1 - for i in range(1,dm_rank): - Ws[:,i,i,0] = 1/np.sqrt(n_modes) - else: - # dm_rank=0 is a special case defining a purely stable, incoherent - # mode mixing model. This is passed on by defining 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_like(probe[0].to(dtype=t.float32)) - p_cent = np.array(probe.shape[1:3]).astype(int) // 2 - psr = int(probe_support_radius) - probe_support[p_cent[0]-psr:p_cent[0]+psr, - p_cent[1]-psr:p_cent[1]+psr] = 1 - probe = probe * probe_support[None,:,:] - else: - probe_support = None; - - if restrict_obj != -1: - ro = restrict_obj - os = np.array(obj_size) - ps = np.array(probe_shape) - obj_support = t.zeros_like(obj.to(dtype=t.float32)) - obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2, - ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1 - else: - obj_support = None - - return cls(wavelength, det_geo, probe_basis, probe, obj, Ws, - detector_slice=det_slice, - surface_normal=surface_normal, - min_translation=min_translation, - translation_offsets = translation_offsets, - mask=mask, background=background, - translation_scale=translation_scale, - saturation=saturation, - probe_support=probe_support, - obj_support=obj_support, - oversampling=oversampling) - - - def interaction(self, index, translations): - pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, - translations, - surface_normal=self.surface_normal) - pix_trans -= self.min_translation - - if self.translation_offsets is not None: - pix_trans += self.translation_scale * self.translation_offsets[index] - - Ws = self.Ws[index] - - if type(index) == type(0): - index = [index] - Ws = [Ws] - pix_trans = [pix_trans] - single_frame = True - else: - single_frame = False - - probes = [] - all_exit_waves = [] - - # This is the case if a purely stable, incoherent model is defined. - if len(Ws[0].shape) == 0: - # What we do here is generate an identity matrix, and multiply - # that identity matrix by the per-frame weight. - Ws = [W * t.stack([t.eye(self.probe.shape[0]), - t.zeros([self.probe.shape[0]]*2)],dim=-1).to( - dtype=W.dtype, device=W.device) - for W in Ws] - - - # Outer iteration is the mode index iteration - for i in range(Ws[0].shape[0]): - # Now we need to separately treat each mode - exit_waves = [] - for W, pix_tran in zip(Ws, pix_trans): - # from storing the probe in Fourier space - pr = [cmath.cmult(W[i,j,:], self.probe[j] * self.probe_support) - for j in range(self.probe.shape[0])] - - pr = t.sum(t.stack(pr), axis=0) - - exit_waves.append(self.probe_norm * - tools.interactions.ptycho_2D_sinc( - pr, self.obj_support * self.obj, - pix_tran, shift_probe=True)) - - exit_waves = t.stack(exit_waves) - - if single_frame: - exit_waves = exit_waves[0] - # Multiply again by probe support to suppress the fringes from the - # sinc-interpolated shift - exit_waves = exit_waves * self.probe_support[...,:,:] - - all_exit_waves.append(exit_waves) - - - return t.stack(all_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) - - 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 to(self, *args, **kwargs): - super(UnifiedModePtycho2, 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.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.obj_support = self.obj_support.to(*args,**kwargs) - self.surface_normal = self.surface_normal.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 get_rhos(self): - # If this is not a purely stable model - if len(self.Ws.shape) >= 2: - Ws = cmath.torch_to_complex(self.Ws.detach().cpu()) - rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj()) - return rhos_out - else: - return np.array([np.eye(self.probe.shape[0])]*self.Ws.shape[0], - dtype=np.complex64) - - def tidy_probes(self, normalization=1): - """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 - - """ - - # Must also implement a version that works appropriately with - # a purely incoherent model - - # - # 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. This could avoid potential stability issues - # due to the existence of zero eigenvalues in the full rho matrix - # when dm_rank < n_modes - # - - rhos = self.get_rhos() - overall_rho = np.mean(rhos,axis=0) - probe = cmath.torch_to_complex(self.probe.detach().cpu()) - ortho_probes, A = analysis.orthogonalize_probes(probe, - density_matrix=overall_rho, - keep_transform=True, - normalize=True) - 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.Ws.shape[1] - - new_Ws = [] - for rho in new_rhos: - # These are returned from smalles to largest - we want to keep - # the largest ones - w,v = np.linalg.eigh(rho) - w = w[::-1][:dm_rank] - v = v[:,::-1][:,:dm_rank] - new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose())) - - new_Ws = np.array(new_Ws) - - self.Ws.data = cmath.complex_to_torch(new_Ws).to( - dtype=self.Ws.dtype,device=self.Ws.device) - - self.probe.data = cmath.complex_to_torch(ortho_probes).to( - device=self.probe.device,dtype=self.probe.dtype) - - - def corrected_translations(self,dataset): - translations = dataset.translations.to(dtype=self.probe.dtype,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 - - - # Needs to be updated to allow for plotting to an existing figure - plot_list = [ - ('Dominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Dominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Subdominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Subdominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Average Density Matrix Amplitudes', - lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig), - lambda self: len(self.Ws.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), - lambda self: len(self.Ws.shape) >=2), - ('Object Amplitude', - lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), - ('Object Phase', - lambda self, fig: p.plot_phase(self.obj, 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() - obj = cmath.torch_to_complex(self.obj.detach().cpu()) - background = self.background.detach().cpu().numpy()**2 - Ws = cmath.torch_to_complex(self.Ws.detach().cpu()) - - return {'basis':basis, 'translation':translations, - 'probe':probe,'obj':obj, - 'background':background, - 'Ws':Ws} diff --git a/CDTools/tools/__init__.py b/CDTools/tools/__init__.py index a94f74b..ff1cc7f 100644 --- a/CDTools/tools/__init__.py +++ b/CDTools/tools/__init__.py @@ -17,13 +17,11 @@ having numpy and torch defined under CDTools.tools.cmath, you know? from __future__ import division, print_function, absolute_import -from CDTools.tools import cmath from CDTools.tools import losses from CDTools.tools import data from CDTools.tools import image_processing from CDTools.tools import initializers from CDTools.tools import plotting -from CDTools.tools import projectors from CDTools.tools import interactions from CDTools.tools import propagators from CDTools.tools import measurements diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py index 538d6e8..2bc67da 100644 --- a/CDTools/tools/analysis/analysis.py +++ b/CDTools/tools/analysis/analysis.py @@ -5,14 +5,13 @@ or numpy arrays, so they can be used either directly after reconstructions on the attributes of the models themselves, or after-the-fact once the data has been stored in numpy arrays. """ -from __future__ import division, print_function import torch as t import numpy as np -from CDTools.tools import cmath from CDTools.tools import image_processing as ip from scipy import fftpack from scipy import linalg as sla +from scipy import special __all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions', 'calc_consistency_prtf', 'calc_deconvolved_cross_correlation', @@ -62,12 +61,11 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm """ try: - probes = cmath.torch_to_complex(probes.detach().cpu()) + probes = probes.detach().cpu().numpy() send_to_torch = True except: send_to_torch = False - from matplotlib import pyplot as plt # We can do the orthogonalization with an SVD, so first we have to # reshape the final two dimensions (the image shape) into a single # vectorized dimension. This matrix is probes^dagger, hence the @@ -82,15 +80,16 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm # next we want to extract the eigendecomposition of the density matrix # itself w,v = sla.eigh(density_matrix) - w,v = np.linalg.eigh(density_matrix) w = w[::-1] v = v[:,::-1] # We do this just to avoid total failure when the density # matrix is not positive definite. - w = np.abs(w) - # Note: this will fail if any w are less than zero - # Probably should figure out a good way to deal with that + # In most cases (such as when rho is generated directly from some other + # matrix A such that rho=A A^dagger), w should never have any negative + # entries. + w = np.maximum(w,0) + B_dagger = np.dot(np.diag(np.sqrt(w)), v.conj().transpose()) #u,s,vh = np.linalg.svd(np.dot(B_dagger,probes_mat), full_matrices=False) @@ -115,9 +114,8 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm #A_dagger = np.dot(np.transpose(u).conj(),B_dagger) if send_to_torch: - ortho_probes = cmath.complex_to_torch(np.stack(ortho_probes)) - A = cmath.complex_to_torch(A) - #A_dagger = cmath.complex_to_torch(A_dagger) + ortho_probes = t.as_tensor(np.stack(ortho_probes)) + A = t.as_tensor(A) if keep_transform: return ortho_probes, A#_dagger @@ -171,22 +169,21 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False): # First, we normalize the probe intensity to a fixed value. probe_np = False if isinstance(probe, np.ndarray): - probe = cmath.complex_to_torch(probe).to(t.float32) + probe = t.as_tensor(probe, dtype=t.complex64) probe_np = True obj_np = False if isinstance(obj, np.ndarray): - obj = cmath.complex_to_torch(obj).to(t.float32) + obj = t.as_tensor(obj,dtype=t.complex64) obj_np = True # If this is a single probe and not a stack of probes - if len(probe.shape) == 3: + if len(probe.shape) == 2: probe = probe[None,...] single_probe = True else: single_probe = False - - normalization = t.sqrt(t.sum(cmath.cabssq(probe[0])) / (len(probe[0].view(-1))/2)) + normalization = t.sqrt(t.sum(t.abs(probe[0])**2) / (len(probe[0].view(-1)))) probe = probe / normalization obj = obj * normalization @@ -198,37 +195,38 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False): if correct_ramp: # Need to check if this is actually working and, if not, why not - center_freq = ip.centroid(cmath.cabssq(cmath.fftshift(t.fft(probe[0],2)))) - center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32) - center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32) + center_freq = ip.centroid(t.abs(t.fft.fftshift(t.fft.fft2(probe[0]), + dim=(-1,-2)))**2) + center_freq -= t.div(t.tensor(probe[0].shape,dtype=t.float32),2,rounding_mode='floor') + center_freq /= t.as_tensor(probe[0].shape,dtype=t.float32) Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]] - probe_phase_ramp = cmath.expi(2 * np.pi * - (center_freq[0] * t.tensor(Is).to(t.float32) + - center_freq[1] * t.tensor(Js).to(t.float32))) - probe = cmath.cmult(probe, cmath.cconj(probe_phase_ramp)) + probe_phase_ramp = t.exp(2j * np.pi * + (center_freq[0] * t.tensor(Is).to(t.float32) + + center_freq[1] * t.tensor(Js).to(t.float32))) + probe = probe * t.conj(probe_phase_ramp) Is, Js = np.mgrid[:obj.shape[0],:obj.shape[1]] - obj_phase_ramp = cmath.expi(2*np.pi * - (center_freq[0] * t.tensor(Is).to(t.float32) + - center_freq[1] * t.tensor(Js).to(t.float32))) - obj = cmath.cmult(obj, obj_phase_ramp) + obj_phase_ramp = t.exp(2j*np.pi * + (center_freq[0] * t.tensor(Is).to(t.float32) + + center_freq[1] * t.tensor(Js).to(t.float32))) + obj = obj * obj_phase_ramp # Then, we set them to consistent absolute phases - obj_angle = cmath.cphase(t.sum(obj[obj_slice],dim=(0,1))) - obj = cmath.cmult(obj, cmath.expi(-obj_angle)) + obj_angle = t.angle(t.sum(obj[obj_slice])) + obj = obj * t.exp(-1j*obj_angle) for i in range(probe.shape[0]): - probe_angle = cmath.cphase(t.sum(probe[i],dim=(0,1))) - probe[i] = cmath.cmult(probe[i], cmath.expi(-probe_angle)) + probe_angle = t.angle(t.sum(probe[i])) + probe[i] = probe[i] * t.exp(-1j*probe_angle) if single_probe: probe = probe[0] if probe_np: - probe = cmath.torch_to_complex(probe.detach().cpu()) + probe = probe.detach().cpu().numpy() if obj_np: - obj = cmath.torch_to_complex(obj.detach().cpu()) + obj = obj.detach().cpu().numpy() return probe, obj @@ -266,16 +264,17 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None, A list of standardized objects, for further processing """ + # This should be cleaned up so it accepts anything array_like probe_np = False if isinstance(probes[0], np.ndarray): - probes = [cmath.complex_to_torch(probe).to(t.float32) for probe in probes] + probes = [t.as_tensor(probe,dtype=t.complex64) for probe in probes] probe_np = True obj_np = False if isinstance(objects[0], np.ndarray): - objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects] + objects = [t.as_tensor(obj,dtype=t.complex64) for obj in objects] obj_np = True - obj_shape = np.min(np.array([obj.shape[:-1] for obj in objects]),axis=0) + obj_shape = np.min(np.array([obj.shape for obj in objects]),axis=0) objects = [obj[:obj_shape[0],:obj_shape[1]] for obj in objects] if obj_slice is None: @@ -298,8 +297,7 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None, obj = ip.sinc_subpixel_shift(obj,np.array(shift)) - - if len(probe.shape) == 4: + if len(probe.shape) == 3: probe = t.stack([ip.sinc_subpixel_shift(p,tuple(shift)) for p in probe],dim=0) else: @@ -317,10 +315,10 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None, i = -1 if probe_np: - synth_probe = cmath.torch_to_complex(synth_probe) + synth_probe = synth_probe.numpy() if obj_np: - synth_obj = cmath.torch_to_complex(synth_obj) - obj_stack = [cmath.torch_to_complex(obj) for obj in obj_stack] + synth_obj = synth_obj.numpy() + obj_stack = [obj.numpy() for obj in obj_stack] return synth_probe/(i+2), synth_obj/(i+2), obj_stack @@ -358,10 +356,10 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None): obj_np = False if isinstance(objects[0], np.ndarray): - objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects] + objects = [t.as_tensor(obj, dtype=t.complex64) for obj in objects] obj_np = True if isinstance(synth_obj, np.ndarray): - synth_obj = cmath.complex_to_torch(synth_obj).to(t.float32) + synth_obj = t.as_tensor(synth_obj, dtype=t.complex64) if isinstance(basis, t.Tensor): basis = basis.detach().cpu().numpy() @@ -373,7 +371,7 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None): if nbins is None: nbins = np.max(synth_obj[obj_slice].shape) // 4 - synth_fft = cmath.cabssq(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy() + synth_fft = (t.abs(t.fft.fftshift(t.fft.fft2(synth_obj[obj_slice]), dim=(-1,-2)))**2).numpy() di = np.linalg.norm(basis[:,0]) @@ -391,7 +389,8 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None): prtfs = [] for obj in objects: obj = obj[obj_slice] - single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy() + single_fft = (t.abs(t.fft.fftshift(t.fft.fft2(obj), + dim=(-1,-2)))**2).numpy() single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft) prtfs.append(synth_ints/single_ints) @@ -433,32 +432,26 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None): im_np = False if isinstance(im1, np.ndarray): - im1 = cmath.complex_to_torch(im1) + im1 = t.as_tensor(im1) im_np = True if isinstance(im2, np.ndarray): - im2 = cmath.complex_to_torch(im2) + im2 = t.as_tensor(im2) im_np = True - # If last dimension is not 2, then convert to a complex tensor now - if im1.shape[-1] != 2: - im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) - if im2.shape[-1] != 2: - im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) - if im_slice is None: im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] - cor_fft = cmath.cmult(t.fft(im1[im_slice],2), - cmath.cconj(t.fft(im2[im_slice],2))) + cor_fft = t.fft.fft2(im1[im_slice]) * \ + t.conj(t.fft.fft2(im2[im_slice])) # Not sure if this is more or less stable than just the correlation # maximum - requires some testing - cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2) + cor = t.fft.ifft2(cor_fft / t.abs(cor_fft)) if im_np: - cor = cmath.torch_to_complex(cor) + cor = cor.numpy() return cor @@ -500,22 +493,15 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): im_np = False if isinstance(im1, np.ndarray): - im1 = cmath.complex_to_torch(im1) + im1 = t.as_tensor(im1) im_np = True if isinstance(im2, np.ndarray): - im2 = cmath.complex_to_torch(im2) + im2 = t.as_tensor(im2) im_np = True if isinstance(basis, np.ndarray): basis = t.tensor(basis) - # If last dimension is not 2, then convert to a complex tensor now - if im1.shape[-1] != 2: - im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) - if im2.shape[-1] != 2: - im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) - - if im_slice is None: im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] @@ -524,11 +510,11 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): nbins = np.max(im1[im_slice].shape) // 4 - cor_fft = cmath.cmult(cmath.fftshift(t.fft(im1[im_slice],2)), - cmath.fftshift(cmath.cconj(t.fft(im2[im_slice],2)))) + cor_fft = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) * \ + t.fft.fftshift(t.conj(t.fft.fft2(im2[im_slice])),dim=(-1,-2)) - F1 = cmath.cabs(cmath.fftshift(t.fft(im1[im_slice],2)))**2 - F2 = cmath.cabs(cmath.fftshift(t.fft(im2[im_slice],2)))**2 + F1 = t.abs(t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)))**2 + F2 = t.abs(t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)))**2 di = np.linalg.norm(basis[:,0]) @@ -542,7 +528,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): - numerator, bins = np.histogram(Rs,bins=nbins,weights=cmath.torch_to_complex(cor_fft)) + numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft.numpy()) denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy()) denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy()) n_pix, bins = np.histogram(Rs,bins=nbins) @@ -588,12 +574,12 @@ def calc_vn_entropy(matrix): # Normalize them to match standard density matrix form eigs = [eig / np.sum(eig) for eig in eigs] # And calculate the VN entropy! - entropies = [-np.sum(eig*np.log(eig)) for eig in eigs] + entropies = [-np.sum(special.xlogy(eig,eig)) for eig in eigs] return np.array(entropies) else: eig = np.linalg.eigh(matrix)[0] - entropy = -np.sum(eig*np.log(eig))/np.sum(eig) - return -np.trace(np.dot(matrix,sla.logm(matrix))) + entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig) + return entropy def calc_top_mode_fraction(matrix): """Calculates the fraction of total power in the top mode of a density matrix diff --git a/CDTools/tools/atoms/atoms.py b/CDTools/tools/atoms/atoms.py index db4212b..9493214 100644 --- a/CDTools/tools/atoms/atoms.py +++ b/CDTools/tools/atoms/atoms.py @@ -7,7 +7,7 @@ Appendix C. doi: 10.1007/978-1-4419-6533-2_11 They are actually taken from https://github.com/LeBeauGroup/pyMultislicer written by Jim LeBeau, but above is the original reference. """ -from __future__ import division, print_function, absolute_import + import numpy as np import torch as t from scipy import fftpack diff --git a/CDTools/tools/cmath/__init__.py b/CDTools/tools/cmath/__init__.py deleted file mode 100644 index 828fd09..0000000 --- a/CDTools/tools/cmath/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from __future__ import division, print_function, absolute_import - -from CDTools.tools.cmath.cmath import * diff --git a/CDTools/tools/cmath/cmath.py b/CDTools/tools/cmath/cmath.py deleted file mode 100644 index 282b840..0000000 --- a/CDTools/tools/cmath/cmath.py +++ /dev/null @@ -1,327 +0,0 @@ -"""Contains basic functions for dealing with complex numbers in pytorch. - -Since pytorch doesn't have built-in support for complex numbers, but the -fast fourier transforms in pytorch assume a specific format for complex -arrays, this module uses that format to store complex numbers. It exposes -functions for converting between complex numpy arrays and torch tensors -stored in that format, as well as basic complex math operations implemented -on the torch tensors -""" -from __future__ import division, print_function, absolute_import -import numpy as np -import torch as t - - -__all__ = ['complex_to_torch', 'torch_to_complex', 'cabssq', 'cabs', 'cconj', - 'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi', 'cexpi'] - - -# -# These define the conversions to and from this format -# - -def complex_to_torch(x): - """Maps a complex numpy array to a torch tensor - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This maps a complex type numpy array to a torch - tensor following this convention - - Parameters - ---------- - x : np.ndarray - A numpy array to convert - - Returns - ------- - torch.Tensor - A torch tensor representation of that array - - """ - return t.from_numpy(np.stack((np.real(x),np.imag(x)),axis=-1)) - - -def torch_to_complex(x): - """Maps a torch tensor to the a complex numpy array - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This maps a torch tensor following that convention - to the appropriate numpy complex array. Note that, in order for this - function to work, the tensor must be detached from any parameters and - living on the CPU. - - Parameters - ---------- - x : torch.Tensor - A tensor to convert - - Returns - ------- - np.array - A complex typed numpy array corresponding to the input - - """ - x = np.array(x) - x = x[...,0] + x[...,1] * 1j - return x - - -# -# And these define the basic operations on these arrays. Note that -# multiplication between a complex valued and real valued pytorch -# tensor will proceed as expected because of torch's broadcasting -# and thus doesn't need it's own function -# - -def cabssq(x): - """Returns the square of the absolute value of a complex torch tensor - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This calculates the elementwise absolute value - squared of any toch tensor following that standard. - - Parameters - ---------- - x : torch.Tensor - An input tensor - - Returns - ------- - torch.Tensor - A tensor storing the elementwise absolute value squared - - """ - return x[...,0]**2 + x[...,1]**2 - - -def cabs(x): - """Returns the absolute value of a complex torch tensor - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This calculates the elementwise absolute value - of any torch tensor following that standard. - - Parameters - ---------- - x : torch.Tensor - An input tensor - - Returns - ------- - torch.Tensor - A tensor storing the elementwise absolute value - - """ - return t.sqrt(cabssq(x)) - - -def cphase(x): - """Returns the phase of a complex torch tensor - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This calculates the elementwise complex phase - of any torch tensor following that standard. - - Parameters - ---------- - x : torch.Tensor - An input tensor - - Returns - ------- - torch.Tensor - A tensor storing the elementwise phase - - """ - return t.atan2(x[...,1],x[...,0]) - - -def cconj(x): - """Returns the complex conjugate of a complex torch tensor - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This calculates the elementwise complex conjugate - of any torch tensor following that standard. - - Parameters - ---------- - x : torch.Tensor - An input tensor - - Returns - ------- - torch.Tensor - A tensor storing the elementwise complex conjugate - - """ - return t.stack((x[...,0],-x[...,1]),dim=-1) - - - -def cmult(a,b): - """Returns the complex product of two torch tensors - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This calculates the elementwise product - of two torch tensors following that standard. - - Parameters - ---------- - a : torch.Tensor - An input tensor - b : torch.Tensor - A second input tensor - - Returns - ------- - torch.Tensor - A tensor storing the elementwise product - - """ - - real = a[...,0] * b[...,0] - a[...,1] * b[...,1] - imag = a[...,0] * b[...,1] + a[...,1] * b[...,0] - return t.stack((real,imag),dim=-1) - - -def cdiv(a,b): - """Returns the complex quotient of two torch tensors - - Pytorch uses tensors with a final dimension of 2 to represent - complex numbers. This calculates the elementwise quotient - of two torch tensors following that standard. - - Parameters - ---------- - a : torch.Tensor - An input tensor - b : torch.Tensor - A second input tensor - - Returns - ------- - torch.Tensor - A tensor storing the elementwise complex quotient - - """ - return cmult(a, cconj(b)) / t.unsqueeze(cabssq(b),-1) - - - -# -# Not entirely sure if these belong here, but heck with it. -# We just need the ability to do fftshifts -# - - -def fftshift(array,dims=None): - """Drop-in torch replacement for scipy.fftpack.fftshift - - This maps a tensor, assumed to be the output of a fast Fourier - transform, into a tensor whose zero-frequency element is at the - center of the tensor instead of the start. It will by default shift - every dimension in the tensor but the last (which is assumed to - represent the complex number and be of dimension 2), but can shift - any arbitrary set of dimensions. - - Parameters - ---------- - array : torch.Tensor - An array of data to be fftshifted - dims : iterable - A list of all dimensions to shift - - Returns - ------- - torch.Tensor - The fftshifted tensor - - """ - - if dims is None: - dims=list(range(array.dim()))[:-1] - for dim in dims: - length = array.size()[dim] - cut_to = (length + 1) // 2 - cut_len = length - cut_to - array = t.cat((array.narrow(dim,cut_to,cut_len), - array.narrow(dim,0,cut_to)), dim) - return array - - - -def ifftshift(array,dims=None): - """Drop-in torch replacement for scipy.fftpack.iftshift - - This maps a tensor, assumed to be the shifted output of a fast - Fourier transform, into a tensor whose zero-frequency element is - back at the start of the tensor instead of the center. It is the - inverse of the fftshift operator. It will by default shift - every dimension in the tensor but the last (which is assumed to - represent the complex number and be of dimension 2), but can shift - any arbitrary set of dimensions. - - Parameters - ---------- - array : torch.Tensor - An array of data to be ifftshifted - dims : list(int) - A list of all dimensions to shift - - Returns - ------- - torch.Tensor - The ifftshifted tensor - - """ - - if dims is None: - dims=list(range(array.dim()))[:-1] - for dim in dims: - length = array.size()[dim] - cut_to = length // 2 - cut_len = length - cut_to - - array = t.cat((array.narrow(dim,cut_to,cut_len), - array.narrow(dim,0,cut_to)), dim) - return array - - -def expi(x): - """Returns a complex-format tensor for exp(i* (x)) - - Expects the input to be in the form of a real-valued tensor - - Parameters - ---------- - x : torch.Tensor - An array to be exponentiated - - Returns - ------- - torch.Tensor - A complex-format tensor - - """ - return t.stack((t.cos(x),t.sin(x)),dim=-1) - - -def cexpi(z): - """Returns a complex-format tensor for exp(i* (z)) - - Expects the input to be in the form of a complex-valued tensor - - Parameters - ---------- - x : torch.Tensor - An array to be exponentiated - - Returns - ------- - torch.Tensor - A complex-format tensor - - """ - real = t.cos(z[...,0]) * t.exp(-z[...,1]) - imag = t.sin(z[...,0]) * t.exp(-z[...,1]) - return t.stack((real, imag),dim=-1) diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py index 0dfe6ed..b3f05c4 100644 --- a/CDTools/tools/data/data.py +++ b/CDTools/tools/data/data.py @@ -5,7 +5,6 @@ specific desired information from a .cxi file. These functions should handle all the needed conversions between standard formats (for example, transposes of the basis arrays, shifting from object to probe motion, etc). """ -from __future__ import division, print_function, absolute_import import h5py import numpy as np @@ -277,7 +276,7 @@ def get_mask(cxi_file): mask = np.array(i1['detector_1/mask']).astype(np.uint32) mask_on = np.equal(mask,np.uint32(0)) mask_has_signal = np.equal(mask,np.uint32(0x00001000)) - return np.logical_or(mask_on,mask_has_signal).astype(np.bool) + return np.logical_or(mask_on,mask_has_signal).astype(bool) else: return None @@ -354,7 +353,11 @@ def get_data(cxi_file, cut_zeroes = True): data[data < 0] = 0 if 'axes' in cxi_file[pull_from].attrs: - axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':') + try: + axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':') + except AttributeError as e: # Weird string vs bytes thing, ehhh + axes = str(cxi_file[pull_from].attrs['axes']).split(':') + axes = [axis.strip().lower() for axis in axes] else: axes = None diff --git a/CDTools/tools/image_processing/image_processing.py b/CDTools/tools/image_processing/image_processing.py index 4ea5d3c..8e57afc 100644 --- a/CDTools/tools/image_processing/image_processing.py +++ b/CDTools/tools/image_processing/image_processing.py @@ -7,14 +7,13 @@ kind of tools perform common image manipulations on torch tensors, in such a way that it is safe to include them in automatic differentiation models. """ -from __future__ import division, print_function, absolute_import import numpy as np import torch as t -from CDTools.tools import cmath +from CDTools.tools import propagators __all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift', 'find_subpixel_shift', 'find_pixel_shift', 'find_shift', - 'convolve_1d'] + 'convolve_1d', 'fourier_upsample'] def centroid(im, dims=2): @@ -75,7 +74,7 @@ def centroid_sq(im, dims=2, comp=False): An (i,j) index or stack of indices """ if comp: - im_sq = cmath.cabssq(im) + im_sq = t.abs(im)**2 else: im_sq = im**2 @@ -109,9 +108,9 @@ def sinc_subpixel_shift(im, shift): I = I.to(dtype=im.dtype,device=im.device) J = J.to(dtype=im.dtype,device=im.device) - fft_im = cmath.fftshift(t.fft(im, 2)) - shifted_fft_im = cmath.cmult(fft_im, cmath.expi(-shift[0]*I - shift[1]*J)) - return t.ifft(cmath.ifftshift(shifted_fft_im),2) + fft_im = t.fft.fftshift(t.fft.fft2(im),dim=(-2,-1)) + shifted_fft_im = fft_im * t.exp(1j * (-shift[0]*I - shift[1]*J)) + return t.fft.ifft2(t.fft.ifftshift(shifted_fft_im, dim=(-2,-1))) @@ -139,6 +138,7 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): shift : torch.Tensor The relative shift (i,j) needed to best map im1 onto im2 """ + # # Here's my approach, perhaps it's a little unconventional. I will first # calculate the phase correlation function as found in ____ (cite a paper @@ -147,19 +147,11 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): # using an FFT with upsampling by a factor of resolution in reciprocal # space # - - # If last dimension is not 2, then convert to a complex tensor now - if im1.shape[-1] != 2: - im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) - if im2.shape[-1] != 2: - im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) - - - cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2))) + cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2)) # Not sure if this is more or less stable than just the correlation # maximum - requires some testing - cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2) + cor = t.fft.ifft2(cor_fft / t.abs(cor_fft)) # Now, I need to shift the array to pull out a contiguous window @@ -167,27 +159,33 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): try: search_around = search_around.cpu() except: - search_around = t.tensor(search_around) + search_around = t.as_tensor(search_around) window_size = 15 shift_zero = tuple(-search_around + t.tensor([window_size,window_size])) - cor_window = t.roll(cor, shift_zero, dims=(0,1))[:2*window_size,:2*window_size] + cor_window = t.roll(cor, shift_zero, dims=(-2,-1))[...,:2*window_size,:2*window_size] # Now we upsample this window - cor_window_fft = cmath.fftshift(t.fft(cor_window,2)) - upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape)[:-1] * resolution) + (2,), + cor_window_fft = t.fft.fftshift(t.fft.fft2(cor_window),dim=(-2,-1)) + upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape) * resolution), dtype=cor.dtype,device=cor.device) - upsampled[:2*window_size,:2*window_size] = cor_window_fft + upsampled[...,:2*window_size,:2*window_size] = cor_window_fft upsampled = t.roll(upsampled,(-window_size,-window_size),dims=(0,1)) - upsampled = t.roll(cmath.cabssq(t.ifft(upsampled, 2)),(-window_size*resolution,-window_size*resolution), dims=(0,1)) + upsampled = t.roll(t.abs(t.fft.ifft2(upsampled))**2, + (-window_size*resolution,-window_size*resolution), + dims=(0,1)) # And we extract the shift from the window - sh = t.tensor(upsampled.shape).to(device=upsampled.device) - cormax = t.tensor([t.argmax(upsampled) // sh[1], - t.argmax(upsampled) % sh[1]]).to(device=upsampled.device) - subpixel_shift = ((cormax + sh // 2) % sh - sh//2).to(dtype=upsampled.dtype) + sh = t.as_tensor(upsampled.shape, device=upsampled.device) + cormax = t.as_tensor([t.div(t.argmax(upsampled), sh[1], + rounding_mode='floor'), + t.argmax(upsampled) % sh[1]], + device=upsampled.device) + + sh_over_2 = t.div(sh,2,rounding_mode='floor') + subpixel_shift = ((cormax + sh_over_2) % sh - sh_over_2).to(dtype=upsampled.dtype) return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \ subpixel_shift / resolution @@ -213,24 +211,19 @@ def find_pixel_shift(im1, im2): shift : torch.Tensor The integer-valued shift (i,j) that best maps im1 onto im2 """ - # If last dimension is not 2, then convert to a complex tensor now - if im1.shape[-1] != 2: - im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) - if im2.shape[-1] != 2: - im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) - - - cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2))) + cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2)) # Not sure if this is more or less stable than just the correlation # maximum - requires some testing - cor = cmath.cabs(t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)) + cor = t.abs(t.fft.ifft2(cor_fft / t.abs(cor_fft))) - sh = t.tensor(cor.shape).to(device=im1.device) - cormax = t.tensor([t.argmax(cor) // sh[1], + sh = t.as_tensor(cor.shape,device=im1.device) + cormax = t.tensor([t.div(t.argmax(cor),sh[1],rounding_mode='floor'), t.argmax(cor) % sh[1]]).to(device=im1.device) - return (cormax + sh // 2) % sh - sh//2 + + sh_over_2 = t.div(sh,2,rounding_mode='floor') + return (cormax + sh_over_2) % sh - sh_over_2 @@ -290,38 +283,37 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True): The convolved image """ - complex_things = 2 - im_complex = True - if image.shape[-1] != 2: - image = t.stack((image,t.zeros_like(image)),dim=-1) - complex_things -= 1 - im_complex = False - - if kernel.shape[-1] != 2: - kernel = t.stack((kernel,t.zeros_like(kernel)),dim=-1) - complex_things -= 1 - + if fftshift_kernel: - kernel = cmath.ifftshift(kernel) + kernel = t.fft.ifftshift(kernel,dim=(-1,)) - # If the image wasn't originally complex, and the dimension - # was passed with the nexative-indexing convention - if not im_complex and dim < 0: - dim = dim-1 - - # We have to transpose the relevant dimension to -2 before using the fft, - # which expects to operate on the final non-complex dimension - trans_im = t.transpose(image, dim, -2) + # We have to transpose the relevant dimension to -1 before using the fft, + # which expects to operate on the final dimension + trans_im = t.transpose(image, dim, -1) # Take a correlation - fft_im = t.fft(trans_im, 1) - fft_kernel = t.fft(kernel, 1) - trans_conv = t.ifft(cmath.cmult(fft_im,fft_kernel), 1) + fft_im = t.fft.fft(trans_im) + fft_kernel = t.fft.fft(kernel) + trans_conv = t.fft.ifft(fft_im * fft_kernel) + + conv_im = t.transpose(trans_conv, dim, -1) + + return conv_im + + +def fourier_upsample(ims, preserve_mean=False): + # If preserve_mean is true, it preserves the mean pixel intensity + # otherwise, it preserves the total summed intensity + upsampled = t.zeros(ims.shape[:-2]+(2*ims.shape[-2],2*ims.shape[-1]), + dtype=ims.dtype, + device=ims.device) + left = [ims.shape[-2]//2,ims.shape[-1]//2] + right = [ims.shape[-2]//2+ims.shape[-2], + ims.shape[-1]//2+ims.shape[-1]] + + upsampled[...,left[0]:right[0],left[1]:right[1]] = propagators.far_field(ims) + if preserve_mean: + upsampled *= 2 + return propagators.inverse_far_field(upsampled) - conv_im = t.transpose(trans_conv, dim, -2) - # If nothing was input as complex, the result should be returned as real - if complex_things == 0: - return conv_im[...,0] - else: - return conv_im diff --git a/CDTools/tools/initializers/initializers.py b/CDTools/tools/initializers/initializers.py index 68487ef..b61b5d5 100644 --- a/CDTools/tools/initializers/initializers.py +++ b/CDTools/tools/initializers/initializers.py @@ -4,21 +4,22 @@ The functions in this module both do the geometric calculations needed to initialize the reconstrucions, and the heuristic calculations for geierating sensible initializations for the probe guess. """ -from __future__ import division, print_function, absolute_import + import numpy as np import torch as t - -__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian', - 'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init'] - -from CDTools.tools import cmath from CDTools.tools.propagators import * from CDTools.tools.analysis import orthogonalize_probes +from CDTools.tools import image_processing from scipy.fftpack import next_fast_len from scipy.sparse import linalg as spla from torch.nn.functional import pad import numpy as np +from functools import * +__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian', + 'gaussian_probe', 'SHARP_style_probe', 'STEM_style_probe', + 'RPI_spectral_init', + 'generate_subdominant_modes'] def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1): """Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry @@ -59,14 +60,17 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, The slice corresponding to the physical detector """ - det_shape = t.Tensor(tuple(det_shape)).to(t.int32) - det_basis = t.Tensor(det_basis) + det_shape = t.as_tensor(tuple(det_shape), dtype=t.int32) + det_basis = t.as_tensor(det_basis) # First, set the center if it's not already specified # This definition matches the center pixel of an fftshifted array if center is None: - center = det_shape // 2 + # this is center//2, but in pytorch 1.9.0 it throws a warning + # if you just do that + + center = t.div(det_shape,2,rounding_mode='floor')# // 2 else: - center = t.Tensor(center).to(t.int32) + center = t.as_tensor(center, dtype=t.int32) # Then, calculate the required detector size from the centering # This is a bit opaque but was worth doing accurately @@ -79,11 +83,14 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, if opt_for_fft: - full_shape = t.Tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32) + full_shape = t.as_tensor([next_fast_len(dim) for dim in full_shape], + dtype=t.int32) # Then, generate a slice that pops the actual detector from the full # detector shape - full_center = full_shape // 2 + # this is full_center//2, but in pytorch 1.9.0 it throws a warning + # if you just do that + full_center = t.div(full_shape,2, rounding_mode='floor')# // 2 det_slice = np.s_[int(full_center[0]-center[0]): int(full_center[0]-center[0]+det_shape[0]), int(full_center[1]-center[1]): @@ -95,7 +102,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, # This method should work for a general parallelogram # shaped detector det_shape = det_basis * full_shape.to(t.float32) - pinv_basis = t.Tensor(np.linalg.pinv(det_shape).transpose()).to(t.float32) + pinv_basis = t.tensor(np.linalg.pinv(det_shape).transpose()).to(t.float32) real_space_basis = pinv_basis * wavelength * distance # This is definitely correct, but less simple. Included here @@ -103,7 +110,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, #oop_dir = np.cross(det_basis[:,0],det_basis[:,1]) #oop_dir /= np.linalg.norm(oop_dir) #full_basis = np.array([np.array(det_basis[:,0]),np.array(det_basis[:,1]),oop_dir]).transpose() - #inv_basis = t.Tensor(np.linalg.inv(full_basis)[:2,:].transpose()).to(t.float32) + #inv_basis = t.tensor(np.linalg.inv(full_basis)[:2,:].transpose()).to(t.float32) #real_space_basis = inv_basis*wavelength * distance / \ # full_shape.to(t.float32) @@ -199,7 +206,7 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]): jsq = (j - center[1])**2 result = np.exp((1j*curvature[0] / 2 - 1 / (2 * sigma[0]**2)) * isq + \ (1j*curvature[1] / 2 - 1 / (2 * sigma[1]**2)) * jsq) - return cmath.complex_to_torch(amplitude*result) + return t.as_tensor(amplitude*result,dtype=t.complex64) @@ -266,8 +273,8 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0): # Finally, we should calculate the average pattern intensity from the # dataset and normalize the gaussian probe. This should be done by avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))] - avg_intensity = t.mean(t.Tensor(avg_intensities)) - probe_intensity = t.sum(cmath.cabssq(probe)) + avg_intensity = t.mean(t.tensor(avg_intensities)) + probe_intensity = t.sum(t.abs(probe)**2) return avg_intensity / probe_intensity * probe @@ -308,6 +315,8 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over The complex-style tensor storing the probe guess """ + # NOTE: I don't love the way np and torch are mixed here, I think this + # function deserves some love. # to use the mask or not? intensities = np.zeros([dim // oversampling for dim in shape]) @@ -322,30 +331,29 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over if hasattr(dataset, 'background') and dataset.background is not None: intensities[det_slice] = np.clip(intensities[det_slice] - dataset.background.cpu().numpy(), a_min=0,a_max=None) - probe_fft = cmath.complex_to_torch(np.sqrt(intensities)) - - probe_guess = cmath.torch_to_complex(inverse_far_field(probe_fft)) - - # Now we remove the central pixel - center = np.array(probe_guess.shape) // 2 - - # I'm always divided on whether to use this modification: - - probe_guess[center[0], center[1]]=np.mean([ - probe_guess[center[0]-1, center[1]], - probe_guess[center[0]+1, center[1]], - probe_guess[center[0], center[1]-1], - probe_guess[center[0], center[1]+1]]) + probe_fft = t.tensor(np.sqrt(intensities)).to(dtype=t.complex64) - probe_guess = cmath.complex_to_torch(probe_guess) + probe_guess = inverse_far_field(probe_fft).numpy() + # Now we remove the central pixel + #center = np.array(probe_guess.shape) // 2 + # I'm always unsure whether to use this modification: + + #probe_guess[center[0], center[1]]=np.mean([ + # probe_guess[center[0]-1, center[1]], + # probe_guess[center[0]+1, center[1]], + # probe_guess[center[0], center[1]-1], + # probe_guess[center[0], center[1]+1]]) + + probe_guess = t.as_tensor(probe_guess, dtype=t.complex64) + if propagation_distance is not None: # First generate the propagation array - probe_shape = t.Tensor(tuple(probe_guess.shape))[:-1] + probe_shape = t.as_tensor(tuple(probe_guess.shape)) # Start by recalculating the probe basis from the given information - det_basis = t.Tensor(dataset.detector_geometry['basis']) + det_basis = t.as_tensor(dataset.detector_geometry['basis']) basis_dirs = det_basis / t.norm(det_basis, dim=0) distance = dataset.detector_geometry['distance'] probe_basis = basis_dirs * dataset.wavelength * distance / \ @@ -355,19 +363,117 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over probe_spacing = t.norm(probe_basis,dim=0).numpy() probe_shape = probe_shape.numpy().astype(np.int32) - #assert 0 # And generate the propagator AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance) probe_guess = near_field(probe_guess,AS_prop) - # Finally, place this probe in a full-sized array if there is oversampling - final_probe = t.zeros([dim for dim in shape] + [2]) + final_probe = t.zeros(shape,dtype=t.complex64) left = shape[0]//2 - probe_guess.shape[0] // 2 top = shape[1]//2 - probe_guess.shape[1] // 2 final_probe[left:left+probe_guess.shape[0], - top:top+probe_guess.shape[1],:] = probe_guess + top:top+probe_guess.shape[1]] = probe_guess + + return final_probe + + +def STEM_style_probe(dataset, shape, det_slice, convergence_semiangle, propagation_distance=None, oversampling=1): + """Generates a STEM style probe guess from a dataset + + What we call the "STEM" style probe guess is a probe generated by + a uniform aperture in Fourier space, with no optical aberrations. + This is the kind of probe than an ideal, aberration-corrected STEM + with perfect coherence would produce, so it is a good starting + guess for STEM datasets. + + We set the initial intensity of the probe with relation to the + diffraction patterns so that a typical object will have an intensity + of around 1. We also set the initial phase ramp of the probe such that + the undiffracted probe has a centroid on the detector which matches the + centroid of the diffraction patterns + + Parameters + ---------- + dataset : Ptycho_2D_Dataset + The dataset to work from + shape : torch.Size + The size of the probe array to simulate + det_slice : slice + A slice or tuple of slices corresponding to the detector region in Fourier space + convergence_angle : float + The convergence angle of the probe, in mrad. + propagation_distance : float + Default is no propagation, an amount to propagate the guessed probe from it's focal point + oversampling : int + Default 1, the width of the region of pixels in the wavefield to bin into a single detector pixel + + Returns + ------- + torch.Tensor + The complex-style tensor storing the probe guess + """ + + # The basis in the dataset should describe the basis of the probe's + # Fourier transform, even if the probe is simulated on larger stage than + # the detector. The only issue is if the probe is oversampled in + # Fourier space (simulated on a larger stage in real space). That factor + # is defined by oversampling. + + probe_basis = (t.as_tensor(dataset.detector_geometry['basis'], + dtype=t.float32) + / oversampling) + + mean_im = t.mean(dataset.patterns,dim=0) + center = image_processing.centroid(mean_im) + + Is = t.arange(shape[0], dtype=t.float32) - center[0] + Js = t.arange(shape[1], dtype=t.float32) - center[1] + Is,Js = t.meshgrid(Is,Js) + + Rs = t.tensordot(probe_basis,t.stack([Is,Js]),dims=1) + forward = t.Tensor([0,0,1]) + Rs = (forward * dataset.detector_geometry['distance'])[:,None,None] + Rs + dirs = Rs / t.linalg.norm(Rs,dim=(0,)) + angles = t.acos(t.tensordot(forward,dirs,dims=1)) * 1000 #in mrad + + probe_fft = t.zeros(shape,dtype=t.complex64) + probe_fft[angles1024: + raise KeyError('Are you sure you want this many orders?') + i = 1 + n = 0 + total = 0 + while True: + if total >= n_orders: + break + + yield (n,i-n) + total += 1 + if n j: # lower triangle, imaginary part - # We pull the i,jth element from the density matrix, - # but this correponds to wavefield j and wavefield i, - # unlike above where it was wavefield i and j (swapped order). - # We also get the one negative sign because this is the imaginary - # part - output += 2 * density_matrix[...,i,j,None,None] \ - * (wavefields[j,...,0] * wavefields[i,...,1] - - wavefields[j,...,1] * wavefields[i,...,0]) - - # Now we apply oversampling - if oversampling != 1: - if wavefields.dim() == 4: - output = avg_pool2d(output.unsqueeze(0), oversampling)[0] - else: - output = avg_pool2d(output, oversampling) - - # Then we grab the detector slice - if detector_slice is not None: - if wavefields.dim() == 4: - output = output[detector_slice] - else: - output = output[(np.s_[:],) + detector_slice] - - if saturation is None: - return t.clamp(output,min=0) + epsilon - else: - return t.clamp(output + epsilon,0,saturation) - def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1): @@ -169,15 +70,15 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non the wavefields. If a detector slice is given, the returned array will only include that slice from the simulated wavefronts. - The first index is the set of incoherently adding patterns, and - the second index is the index of the diffraction pattern to measure. - The next two indices index the wavefield. The final index is the complex - index. + The (-3) index is the set of incoherently adding patterns, and any + indexes further to the front correspond to the set of diffraction patterns + to measure. The (-2) and (-1) indices are the wavefield, and the final + index is the complex index Parameters ---------- wavefields : torch.Tensor - An LxJxMxNx2 stack of complex wavefields + An LxJxMxNx stack of complex wavefields detector_slice : slice Optional, a slice or tuple of slices defining a section of the simulation to return saturation : float @@ -188,21 +89,21 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non Returns ------- sim_patterns : torch.Tensor - A real JXMxN array storing the incoherently summed intensities + A real LXMxN array storing the incoherently summed intensities """ - output = t.sum(cmath.cabssq(wavefields),dim=0) + output = t.sum(t.abs(wavefields)**2,dim=-3) # Now we apply oversampling if oversampling != 1: - if wavefields.dim() == 4: + if wavefields.dim() == 3: output = avg_pool2d(output.unsqueeze(0), oversampling)[0] else: output = avg_pool2d(output, oversampling) # Then we grab the detector slice if detector_slice is not None: - if wavefields.dim() == 4: + if wavefields.dim() == 3: output = output[detector_slice] else: output = output[(np.s_[:],) + detector_slice] @@ -224,7 +125,7 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas Parameters ---------- wavefield : torch.Tensor - A JxMxNx2 stack of complex wavefields + A JxMxN stack of complex-valued wavefields background : torch.Tensor An tensor storing the square root of the detector background detector_slice : slice diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index 14c9a4d..b5c65a8 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -6,18 +6,18 @@ and for after-the-fact analysis. Utilities for plotting complex valued images exist, as well as plotting scan patterns and nanomaps """ -from __future__ import division, print_function, absolute_import - -from CDTools.tools import cmath import torch as t import numpy as np import matplotlib.pyplot as plt from matplotlib.colors import hsv_to_rgb +from matplotlib.widgets import Slider +from matplotlib import ticker, patheffects __all__ = ['colorize', 'plot_amplitude', 'plot_phase', 'plot_colorized', 'plot_translations', 'get_units_factor', - 'plot_nanomap', 'plot_real', 'plot_imag'] + 'plot_nanomap', 'plot_real', 'plot_imag', + 'plot_nanomap_with_images'] def colorize(z): @@ -82,7 +82,146 @@ def get_units_factor(units): return factor -def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): +def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, **kwargs): + """Plots an image with a colorbar and on an appropriate spatial grid + + If a figure is given explicitly, it will clear that existing figure and + plot over it. Otherwise, it will generate a new figure. + + If a basis is explicitly passed, the image will be plotted in real-space + coordinates + + Finally, if a function is passed to the plot_func argument, this function + will be called on each slice of data before it is plotted. This is used + internally to enable the plot_real, plot_image, plot_phase, etc. functions. + + + Parameters + ---------- + im : array + An complex array with dimensions NxM + plot_func : callable + A function which maps numpy arrays to the image to be plotted + fig : matplotlib.figure.Figure + Default is a new figure, a matplotlib figure to use to plot + basis : np.array + Optional, the 3x2 probe basis + units : str + The length units to mark on the plot, default is um + cmap : str + Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting + \\**kwargs + All other args are passed to fig.add_subplot(111, \\**kwargs) + + Returns + ------- + used_fig : matplotlib.figure.Figure + The figure object that was actually plotted to. + """ + + # convert to numpy + if isinstance(im, t.Tensor): + # If final dimension is 2, assume it is a complex array. If not, + # assume it represents a real array + if im.shape[-1] == 2: + im = im.detach().cpu().numpy() + else: + im = im.detach().cpu().numpy() + + if fig is None: + fig = plt.figure() + ax = fig.add_subplot(111, **kwargs) + + # This nukes everything and updates either the appropriate image from the + # stack of images, or the only image if only a single image has been + # given + def make_plot(idx): + plt.figure(fig.number) + title = plt.gca().get_title() + fig.clear() + + + # If im only has two dimensions, this reshape will add a leading + # dimension, and update will be called on index 0. If it has 3 or more + # dimensions, then all the leading dimensions will be compressed into + # one long dimension which can be scrolled through. + s = im.shape + reshaped_im = im.reshape(-1,s[-2],s[-1]) + num_images = reshaped_im.shape[0] + fig.plot_idx = idx % num_images + + to_plot = plot_func(reshaped_im[fig.plot_idx]) + + #Plot in a basis if it exists, otherwise dont + if basis is not None: + if isinstance(basis,t.Tensor): + np_basis = basis.detach().cpu().numpy() + else: + np_basis = basis + # This fails if the basis is not rectangular + basis_norm = np.linalg.norm(np_basis, axis = 0) + basis_norm = basis_norm * get_units_factor(units) + + extent = [0, im.shape[-1]*basis_norm[1], 0, + im.shape[-2]*basis_norm[0]] + else: + extent=None + + plt.imshow(to_plot, cmap = cmap, extent = extent) + cbar = plt.colorbar() + if cmap_label is not None: + cbar.set_label(cmap_label) + + if basis is not None: + plt.xlabel('X (' + units + ')') + plt.ylabel('Y (' + units + ')') + else: + plt.xlabel('j (pixels)') + plt.ylabel('i (pixels)') + + + plt.title(title) + + if len(im.shape) >= 3: + plt.text(0.03, 0.03, str(fig.plot_idx), fontsize=14, transform=plt.gcf().transFigure) + return fig + + if hasattr(fig, 'plot_idx'): + result = make_plot(fig.plot_idx) + else: + result = make_plot(0) + + update = make_plot + + + def on_action(event): + if not hasattr(event, 'button'): + event.button = None + if not hasattr(event, 'key'): + event.key = None + + if event.key == 'up' or event.button == 'up': + update(fig.plot_idx - 1) + elif event.key == 'down' or event.button == 'down': + update(fig.plot_idx + 1) + plt.draw() + + if len(im.shape) >=3: + if not hasattr(fig,'my_callbacks'): + fig.my_callbacks = [] + + for cid in fig.my_callbacks: + fig.canvas.mpl_disconnect(cid) + fig.my_callbacks = [] + fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action)) + fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action)) + + return result + + +def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs): """Plots the real part of a complex array with dimensions NxM If a figure is given explicitly, it will clear that existing figure and @@ -103,6 +242,8 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -111,45 +252,14 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa used_fig : matplotlib.figure.Figure The figure object that was actually plotted to. """ - if fig is None: - fig = plt.figure() - ax = fig.add_subplot(111, **kwargs) - else: - plt.figure(fig.number) - plt.gcf().clear() - - if isinstance(im, t.Tensor): - real = im[...,0].detach().cpu().numpy() - else: - real = np.real(im) - - #Plot in a basis if it exists, otherwise dont - if basis is not None: - if isinstance(basis,t.Tensor): - basis = basis.detach().cpu().numpy() - # This fails if the - basis_norm = np.linalg.norm(basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - - extent = [0, real.shape[-1]*basis_norm[1], 0, real.shape[-2]*basis_norm[0]] - else: - extent=None - - plt.imshow(real, cmap = cmap, extent = extent) - cbar = plt.colorbar() - cbar.set_label('Real Part (a.u.)') - - if basis is not None: - plt.xlabel('X (' + units + ')') - plt.ylabel('Y (' + units + ')') - else: - plt.xlabel('j (pixels)') - plt.ylabel('i (pixels)') - - return fig + plot_func = lambda x: np.real(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) + -def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): +def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs): """Plots the imaginary part of a complex array with dimensions NxM If a figure is given explicitly, it will clear that existing figure and @@ -170,6 +280,8 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -178,53 +290,21 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa used_fig : matplotlib.figure.Figure The figure object that was actually plotted to. """ - if fig is None: - fig = plt.figure() - ax = fig.add_subplot(111, **kwargs) - else: - plt.figure(fig.number) - plt.gcf().clear() - - if isinstance(im, t.Tensor): - imag = im[...,1].detach().cpu().numpy() - else: - imag = np.imag(im) - - #Plot in a basis if it exists, otherwise dont - if basis is not None: - if isinstance(basis,t.Tensor): - basis = basis.detach().cpu().numpy() - # This fails if the - basis_norm = np.linalg.norm(basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - - extent = [0, imag.shape[-1]*basis_norm[1], 0, imag.shape[-2]*basis_norm[0]] - else: - extent=None - - plt.imshow(imag, cmap = cmap, extent = extent) - cbar = plt.colorbar() - cbar.set_label('Imaginary Part (a.u.)') - - if basis is not None: - plt.xlabel('X (' + units + ')') - plt.ylabel('Y (' + units + ')') - else: - plt.xlabel('j (pixels)') - plt.ylabel('i (pixels)') - - return fig + plot_func = lambda x: np.imag(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) -def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): +def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', **kwargs): """Plots the amplitude of a complex array with dimensions NxM If a figure is given explicitly, it will clear that existing figure and plot over it. Otherwise, it will generate a new figure. If a basis is explicitly passed, the image will be plotted in real-space - coordinates - + coordinates. + Parameters ---------- im : array @@ -237,6 +317,8 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -245,45 +327,13 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', used_fig : matplotlib.figure.Figure The figure object that was actually plotted to. """ - if fig is None: - fig = plt.figure() - ax = fig.add_subplot(111, **kwargs) - else: - plt.figure(fig.number) - plt.gcf().clear() - - if isinstance(im, t.Tensor): - absolute = cmath.cabs(im).detach().cpu().numpy() - else: - absolute = np.absolute(im) - - #Plot in a basis if it exists, otherwise dont - if basis is not None: - if isinstance(basis,t.Tensor): - basis = basis.detach().cpu().numpy() - # This fails if the - basis_norm = np.linalg.norm(basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - - extent = [0, absolute.shape[-1]*basis_norm[1], 0, absolute.shape[-2]*basis_norm[0]] - else: - extent=None - - plt.imshow(absolute, cmap = cmap, extent = extent) - cbar = plt.colorbar() - cbar.set_label('Amplitude (a.u.)') - - if basis is not None: - plt.xlabel('X (' + units + ')') - plt.ylabel('Y (' + units + ')') - else: - plt.xlabel('j (pixels)') - plt.ylabel('i (pixels)') - - return fig + plot_func = lambda x: np.absolute(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) -def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs): +def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', cmap_label='Phase (rad)', **kwargs): """ Plots the phase of a complex array with dimensions NxMx2 If a figure is given explicitly, it will clear that existing figure and @@ -304,6 +354,8 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs) The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -312,49 +364,18 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs) used_fig : matplotlib.figure.Figure The figure object that was actually plotted to. """ - if fig is None: - fig = plt.figure() - ax = fig.add_subplot(111, **kwargs) - else: - plt.figure(fig.number) - plt.gcf().clear() - - if isinstance(im, t.Tensor): - phase = cmath.cphase(im).detach().cpu().numpy() - else: - phase = np.angle(im) - - if basis is not None: - if isinstance(basis,t.Tensor): - basis = basis.detach().cpu().numpy() - basis_norm = np.linalg.norm(basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - - extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]] - else: - extent=None - - - # If the user has matplotlib >=3.0, use the preferred colormap if cmap == 'auto': - try: - plt.imshow(phase, cmap = 'twilight', extent=extent) - except: - plt.imshow(phase, cmap = 'hsv', extent=extent) - else: - plt.imshow(phase, cmap = cmap, extent=extent) + if 'twilight' in plt.colormaps(): + cmap = 'twilight' + elif 'hsv' in plt.colormaps(): + cmap = 'hsv' + else: + raise AttributeError('Neither twilight or hsv colormap exists in this screwed up matplotlib install') - cbar = plt.colorbar() - cbar.set_label('Phase (rad)') - - if basis is not None: - plt.xlabel('X (' + units + ')') - plt.ylabel('Y (' + units + ')') - else: - plt.xlabel('j (pixels)') - plt.ylabel('i (pixels)') - - return fig + plot_func = lambda x: np.angle(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) def plot_amplitude_surfacenorm(): @@ -390,38 +411,9 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs): used_fig : matplotlib.figure.Figure The figure object that was actually plotted to. """ - if fig is None: - fig = plt.figure() - ax = fig.add_subplot(111, **kwargs) - else: - plt.figure(fig.number) - plt.gcf().clear() - - if isinstance(im, t.Tensor): - im = cmath.torch_to_complex(im.detach().cpu()) - - if basis is not None: - if isinstance(basis,t.Tensor): - basis = basis.detach().cpu().numpy() - basis_norm = np.linalg.norm(basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - - extent = [0, im.shape[-1]*basis_norm[1], 0, im.shape[-2]*basis_norm[0]] - else: - extent=None - - colorized = colorize(im) - plt.imshow(colorized, extent=extent) - - if basis is not None: - plt.xlabel('X (' + units + ')') - plt.ylabel('Y (' + units + ')') - else: - plt.xlabel('j (pixels)') - plt.ylabel('i (pixels)') - - return fig - + plot_func = lambda x: colorize(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, **kwargs) def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs): @@ -524,3 +516,254 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr plt.colorbar() return fig + + +def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs): + """Plots a nanomap, with an image or stack of images for each point + + In many situations, ptychography data or the output of ptychography + reconstructions is formatted as a set of images associated with various + points in real space. This function is designed to allow for browsing + through this kind of data, by making it possible to visualize a + + """ + + # This should pull heavily from the dataset.inspect function + # In fact, I should be able to replace most of that function with a + # call to this function once it's built + # We start by making the figure and axes + + # The key will be writing this so it works okay when called in "update" + # mode, i.e. on a figure that already has this thing showing. + + if fig is None: + fig = plt.figure(figsize=(8,5.3)) + else: + plt.figure(fig.number) + plt.gcf().clear() + if hasattr(fig, 'nanomap_cids'): + for cid in fig.nanomap_cids: + fig.canvas.mpl_disconnect(cid) + + # Does figsize work with the fig.subplots, or just for plt.subplots? + axes = fig.subplots(1,2) + + fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96]) + plt.subplots_adjust(wspace=0.25) #avoids overlap of labels with plots + axslider = plt.axes([0.15,0.06,0.75,0.03]) + + # This gets the set of sizes for the points in the nanomap + def calculate_sizes(idx): + bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted()) + s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch + s0 /= 4 # A rough value to make the size work out + s = np.ones(translations.shape[0]) * s0 + + s[idx] *= 4 + return s + + + def update_colorbar(im): + # + # This solves the problem of the colorbar being changed + # when the forward and back buttons are used!!! + # + if hasattr(im, 'norecurse') and im.norecurse: + im.norecurse=False + return + + im.norecurse=True + # This is needed to update the colorbar + # only change limits if array contains multiple values + if np.min(im.get_array()) != np.max(im.get_array()): + im.set_clim(vmin=np.min(im.get_array()), + vmax=np.max(im.get_array())) + + # + # The meatiest part of this program, here we just go through and + # set up the plot how we want it + # + + # First we set up the left-hand plot, which shows an overview map + axes[0].set_title('Relative Displacement Map') + + translations = translations.detach().cpu().numpy() + + if convention.lower() != 'probe': + translations = translations * -1 + + s = calculate_sizes(0) + + nanomap_units_factor = get_units_factor(nanomap_units) + nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0], + nanomap_units_factor * translations[:,1], + s=s,c=values, picker=True) + + axes[0].invert_xaxis() + axes[0].set_facecolor('k') + axes[0].set_xlabel('Translation x ('+nanomap_units+')', labelpad=1) + axes[0].set_ylabel('Translation y ('+nanomap_units+')', labelpad=1) + cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal', + format='%.2e', + ticks=ticker.LinearLocator(numticks=5), + pad=0.17,fraction=0.1) + cb1.ax.set_title(nanomap_colorbar_title, size="medium", pad=5) + cb1.ax.tick_params(labelrotation=20) + if values is None: + # This seems to do a good job of leaving the appropriate space + # where the colorbar should have been to avoid stretching the + # nanomap plot, while still not showing the (now useless) colorbar. + cb1.remove() + + # Now we set up the second plot, which shows the individual + # diffraction patterns + axes[1].set_title(image_title) + #Plot in a basis if it exists, otherwise dont + if basis is not None: + axes[1].set_xlabel('X (' + image_units + ')') + axes[1].set_ylabel('Y (' + image_units + ')') + + example_im = get_image_func(0) + if isinstance(example_im, t.Tensor): + example_im = example_im.cpu().numpy() + + if isinstance(basis,t.Tensor): + np_basis = basis.detach().cpu().numpy() + else: + np_basis = basis + # This fails if the basis is not rectangular + basis_norm = np.linalg.norm(np_basis, axis = 0) + basis_norm = basis_norm * get_units_factor(image_units) + + extent = [0, example_im.shape[-1]*basis_norm[1], 0, + example_im.shape[-2]*basis_norm[0]] + else: + axes[1].set_xlabel('j (pixels)') + axes[1].set_ylabel('i (pixels)') + extent=None + + im=get_image_func(0) + if len(im.shape) >= 3: + im_idx=0 + axes[1].image_idx = im_idx + im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx] + axes[1].text_box = axes[1].text(0.98, 0.98, str(im_idx), color='w', + fontsize=14, + horizontalalignment='right', + verticalalignment='top', + transform=axes[1].transAxes) + axes[1].text_box.set_path_effects( + [patheffects.Stroke(linewidth=2, foreground='black'), + patheffects.Normal()]) + + meas = axes[1].imshow(im, extent=extent, cmap=cmap) + + cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal', + format='%.2e', + ticks=ticker.LinearLocator(numticks=5), + pad=0.17,fraction=0.1) + cb2.ax.tick_params(labelrotation=20) + cb2.ax.set_title(image_colorbar_title, size="medium", pad=5) + cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) + + # This function handles all the updating, except for moving the + # slider value. This is done because the slider widget is + # ultimately responsible for triggering an update, so all other + # updates are done by changing the slider widget value which + # then triggers this + def update(idx, im_idx=None): + # We have to explicitly make it an integer because the slider will + # output floats (even if they are still integer-valued) + idx = int(idx) + + # Get the new data for this index + im = get_image_func(idx) + if len(im.shape) >= 3: + if im_idx == None and hasattr(axes[1],'image_idx'): + im_idx = axes[1].image_idx + elif im_idx == None: + im_idx=0 + axes[1].image_idx = im_idx + axes[1].text_box.set_text(str(im_idx)) + im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx] + + # Now we resize the nanomap to show the new selection + axes[0].collections[0].set_sizes(calculate_sizes(idx)) + + # And we update the data in the image as well + + ax_im = axes[1].images[-1] + ax_im.set_data(im) + update_colorbar(ax_im) + + + # + # Now we define the functions to handle various kinds of events + # that can be thrown our way + # + + # We start by creating the slider here, so it can be used + # by the update hooks. + slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d") + + # This handles scroll wheel and keypress events + def on_action(event): + im = get_image_func(0) + if event.inaxes is axes[1] and len(im.shape) >=3: + # This is triggered if the data to display has more than 2 + # dimensions (i.e. is an image stack) and the event originates + # while the mouse is within the image display + im = im.reshape(-1,im.shape[-2],im.shape[-1]) + im_idx = axes[1].image_idx + + if event.key == 'up' or event.button == 'up' \ + or event.key == 'left': + im_idx = (im_idx - 1) % im.shape[0] + if event.key == 'down' or event.button == 'down' \ + or event.key == 'right': + im_idx = (im_idx + 1) % im.shape[0] + + axes[1].image_idx=im_idx + slider.set_val(slider.val)#update(slider.val,im_idx=im_idx) + return # This prevents the rest from also happening + + # Otherwise the if statements can throw errors when the + # event type isn't right, this way they just don't trigger + if not hasattr(event, 'button'): + event.button = None + if not hasattr(event, 'key'): + event.key = None + + if event.key == 'up' or event.button == 'up' or event.key == 'left': + idx = slider.val - 1 + elif event.key == 'down' or event.button == 'down' or event.key == 'right': + idx = slider.val + 1 + else: + # This prevents errors from being thrown on irrelevant key + # or mouse input + return + + # Handle the wraparound and trigger the update + idx = int(idx) % translations.shape[0] + slider.set_val(idx) + + # This handles "pick" events in the nanomap + def on_pick(event): + # If we don't filter on type of event, this will also capture, + # for example, scroll events that happen over the nanomap + if event.mouseevent.button == 1: + slider.set_val(event.ind[0]) + + + # Here we connect the various update functions + cid1 = fig.canvas.mpl_connect('pick_event',on_pick) + cid2 = fig.canvas.mpl_connect('key_press_event',on_action) + cid3 = fig.canvas.mpl_connect('scroll_event',on_action) + # It's so dumb that matplotlib doesn't automatically track this for you + fig.nanomap_cids = [cid1,cid2,cid3] + slider.on_changed(update) + + # Throw an extra update into the mix just to get rid of any things + # (like the nanomap dot sizes) that otherwise would change on the + # first update + update(0) diff --git a/CDTools/tools/projectors/__init__.py b/CDTools/tools/projectors/__init__.py deleted file mode 100644 index 6234ec0..0000000 --- a/CDTools/tools/projectors/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from __future__ import division, print_function, absolute_import - -from CDTools.tools.projectors.projectors import * diff --git a/CDTools/tools/projectors/projectors.py b/CDTools/tools/projectors/projectors.py deleted file mode 100644 index eef5461..0000000 --- a/CDTools/tools/projectors/projectors.py +++ /dev/null @@ -1,80 +0,0 @@ -"""This module contains various projection functions - -These functions are useful when defining declarative algorithms to run -alongside the automatic differentiation ones, for comparison or in a -situation where they might be needed. -""" -from __future__ import division, print_function, absolute_import -from CDTools.tools.cmath import * -import torch as t - -__all__ = ['modulus', 'support'] - - -def modulus(wavefront, intensities, mask = None): - """Implements the modulus constraint in torch - - This accepts a tensor representing the propagated simulated wavefront(s), - where the last dimension represents the real and imaginary components of - the propagated wavefield(s). It projects the modulus of the diffraction - pattern onto the modulus of the simulated wavefield. - - It assumes that the wavefront is stored in an array - [i,j] where i corresponds to the y-axis and j corresponds to the - x-axis, with the origin following the CS standard of being in the - upper right. - - Parameters - ---------- - wavefront : torch.Tensor - The JxNxMx2 stack of complex propagated wavefronts - intensities : torch.Tensor - The measured diffraction pattern(s) stored as an JxNxM stack of real tensors - mask : torch.Tensor - A mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1 - - Returns - ------- - projected : torch.Tensor - The JxNxMx2 projected wavefield with corrected intensities - """ - # Calculate amplitudes from intensities - amplitudes = t.sqrt(intensities) - # Normalize wavefront so the complex elements have modulus one - wavefront_mag = cabs(wavefront) - projected = wavefront * (amplitudes / wavefront_mag)[...,None] - # Replace amplitude of wavefront with measured amplitude - if mask is not None: - selection = mask == 0 - # Apply the mask to replace unmasked pixels in the original wavefront - projected = projected.masked_scatter(selection, wavefront.masked_select(selection)) - - return projected - - -def support(wavefront, support): - """Implements the support constraint in torch - - This accepts a torch tensor representing (a) simulated wavefield(s), - where the last dimension represents the real and imaginary components of - the propagated wavefield(s). It projects the support of the imaged object - onto the simulated wavefront via a support mask. - - It assumes that the wavefront is stored in an array - [i,j] where i corresponds to the y-axis and j corresponds to the - x-axis, with the origin following the CS standard of being in the - upper right. - - Parameters - ---------- - wavefront : torch.Tensor - The JxNxMx2 stack of complex propagated wavefronts - support : torch.Tensor - An NxM support, with 1s within the support and 0s outside - - Returns - ------- - projected : torch.Tensor - The JxNxMx2 wavefield with the support mask applied - """ - return wavefront * support.to(wavefront.dtype)[...,None] diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py index 2eca125..2d517ab 100644 --- a/CDTools/tools/propagators/propagators.py +++ b/CDTools/tools/propagators/propagators.py @@ -3,9 +3,7 @@ All the functions here are designed for use in an automatic differentiation ptychography model. Each function implements a different propagator. """ -from __future__ import division, print_function, absolute_import -from CDTools.tools.cmath import * import torch as t from torch.nn.functional import grid_sample from scipy import fftpack @@ -19,7 +17,6 @@ __all__ = ['far_field', 'near_field', 'high_NA_far_field', 'generate_generalized_angular_spectrum_propagator'] - def far_field(wavefront): """Implements a far-field propagator in torch @@ -47,8 +44,10 @@ def far_field(wavefront): propagated : torch.Tensor The JxNxMx2 propagated wavefield """ - - return fftshift(t.fft(ifftshift(wavefront), 2, normalized=True)) + + shifted = t.fft.ifftshift(wavefront, dim=(-1,-2)) + propagated = t.fft.fft2(shifted, norm='ortho') + return t.fft.fftshift(propagated, dim=(-1,-2)) def inverse_far_field(wavefront): @@ -74,7 +73,9 @@ def inverse_far_field(wavefront): propagated : torch.Tensor The JxNxMx2 exit wavefield """ - return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True)) + shifted = t.fft.ifftshift(wavefront, dim=(-1,-2)) + propagated = t.fft.ifft2(shifted, norm='ortho') + return t.fft.fftshift(propagated, dim=(-1,-2)) def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, lens=False, **kwargs): @@ -286,38 +287,37 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None): # np.ones_like(k_map[0,:-1,:-1,0].cpu().numpy())) #plt.show() def process_wavefield_stack(low_NA_wavefield): - real_output = grid_sample(low_NA_wavefield[None,:,:,:,0],k_map,mode='bilinear',padding_mode='zeros', align_corners=False) - imag_output = grid_sample(low_NA_wavefield[None,:,:,:,1],k_map,mode='bilinear',padding_mode='zeros', align_corners=False) + # grid_sample doesn't work on complex-valued wavefields + real_output = grid_sample(low_NA_wavefield[None,:,:,:].real,k_map,mode='bilinear',padding_mode='zeros', align_corners=False) + imag_output = grid_sample(low_NA_wavefield[None,:,:,:].imag,k_map,mode='bilinear',padding_mode='zeros', align_corners=False) - result = t.stack((real_output[0,:,:,:],imag_output[0,:,:,:]),dim=3) + result = real_output[0,:,:,:] + 1j * imag_output[0,:,:,:] if intensity_map is not None: - result = result * intensity_map[None,:,:,None] + result = result * intensity_map[None,:,:] return result original_dim = wavefront.dim() + if original_dim == 2: + result = process_wavefield_stack(low_NA_wavefield[None,:,:]) + return result[0,:,:] if original_dim == 3: - result = process_wavefield_stack(low_NA_wavefield[None,:,:,:]) - return result[0,:,:,:] - if original_dim == 4: result = process_wavefield_stack(low_NA_wavefield) return result - if original_dim == 5: + if original_dim == 4: result = [] for i in range(low_NA_wavefield.size()[0]): - result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:,:])) + result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:])) return t.stack(result) else: raise IndexError('Wavefield had incorrect number of dimensions') - - def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, bandlimit=None, **kwargs): """Generates an angular-spectrum based near-field propagator from experimental quantities - + This function generates an angular-spectrum based near field propagator that will work on torch Tensors. The function is structured this way - to generate the propagator first - because the @@ -325,15 +325,14 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r propagator is used in a reconstruction program, then it will be best to calculate this mask once and reuse it. - Formally, this propagator is the complex conjugate of the fourier - transform of the convolution kernel for light propagation in free - space + Formally, this propagator is the the fourier transform of the convolution + kernel for light propagation in free space If the optional bandlimit parameter is set, the propagator will be set to zero beyond an explicit bandlimiting frequency. This is helpful if the propagator will be used in a repeated multiply/propagate framework such as a multislice algorithm, where it helps to prevent aliasing. - + Parameters ---------- shape : array @@ -354,34 +353,43 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r propagator : torch.Tensor A phase mask which accounts for the phase change that each plane wave will undergo. """ - - ki = 2 * np.pi * fftpack.fftfreq(shape[0],spacing[0]) - kj = 2 * np.pi * fftpack.fftfreq(shape[1],spacing[1]) - Kj, Ki = np.meshgrid(kj,ki) - - # Define this as complex so the square root properly gives - # k>k0 components imaginary frequencies - k0 = np.complex128((2*np.pi/wavelength)) - - # Properly accuount for evanescent waves - if z >=0: - propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z) + # Internally, the generalized propagation function is used, so we start + # by creating an appropriate basis + # This creates a real-valued tensor which matches the kind of complex + # number dtype requested in **kwargs + if 'dtype' in kwargs: + basis = t.real(t.zeros([3,2], **kwargs)) else: - propagator = np.exp(1j*np.conj(np.sqrt(k0**2 - Ki**2 - Kj**2)) * z) - - if remove_z_phase: - propagator *= np.exp(-1j * k0 * z) + basis = t.zeros([3,2], dtype=t.float32) + spacing = t.as_tensor(spacing, dtype=basis.dtype) + basis[1,0] = -spacing[0] + basis[0,1] = -spacing[1] + # And similarly, the offset is just z along the z direction + offset = t.zeros([3], dtype=basis.dtype) + offset[2] = z + + + # And we call the generalized function! + propagator = generate_generalized_angular_spectrum_propagator(shape, basis, + wavelength, offset, + propagate_along_offset=remove_z_phase, **kwargs) + if z < 0: + propagator = t.conj(propagator) + + + # Bandlimiting is not implemented in the generalized function, because it + # has a less clear meaning in that setting, so we apply it here instead if bandlimit is not None: - Rs = np.sqrt((Ki / np.max(ki))**2 + (Kj / np.max(kj))**2) + ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0]) + kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1]) + Ki, Kj = t.meshgrid(ki,kj) + min_radius = min(t.max(ki),t.max(kj)) + Rs = t.sqrt((Ki/t.max(ki))**2 + (Kj/t.max(kj))**2) propagator = propagator * (Rs < bandlimit) - # Take the conjugate explicitly here instead of negating - # the previous expression to ensure that complex frequencies - # get mapped to values <1 instead of >1 - propagator = complex_to_torch(np.conj(propagator)) - - return propagator.to(*args, **kwargs) + return propagator + def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, offset_vector, *args, propagation_vector=None, propagate_along_offset=False, **kwargs): @@ -391,16 +399,15 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o propagator that will work on torch Tensors. The function is structured this way - to generate the propagator first - because the generation of the propagation mask is a bit expensive and if this - propagator is used in a reconstruction program, then it will be best - to calculate this mask once and close over it. + propagator is used in a reconstruction program, it will be best + to calculate this mask once and then reuse it it. - Formally, this propagator is the complex conjugate of the fourier - transform of the convolution kernel for light propagation in free - space. It will map a ligh field at an input plane, with the size - and shape defined by the shape and basis inputs, and map it to a - plane of the same size and shape offset by the offset vector. It - is designed to work on any wavefield defined on an array of - parallelograms. + Formally, this propagator is the fourier transform of the convolution + kernel for light propagation in free space. It will map a light field + at an input plane, with the size and shape defined by the shape and basis + inputs, and map it to a plane of the same size and shape offset by the + offset vector. It is designed to work on any wavefield defined on an + array of parallelograms. In addition, if the propagation_vector is set, there is an assumed @@ -418,33 +425,27 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o propagation_vector option Note that, unlike in the case of the simple angular spectrum propagator, - the direction of "forward propagation" is defined by the offset vector. - Therefore, in the simple case of a perpendicular offset, there will be - no difference between using an offset vector or the negative of the - offset vector. This is because, for the light propagation problem to - be well posed, the assumption must be made that light only passes through - the plane of the known wavefield in one direction. Mathematically, this - corresponds to a choice of uniform phase objects either accumulating - positive or negative phase. In the simple propagation case, there is - no ambiguity introduced by always choosing the light field to propagate - along the positive z direction. In the general case, there is no equivalent - obvious choice - thus, the light is always assumed to pass through the - initial plane travelling in the direction of the final plane. + the direction of "forward propagation" is defined by the propagation + vector, or (if the propagation vector is not defined), the offset + vector. Therefore, in the simple case of a perpendicular offset, + there will be no difference between using an offset vector or the + negative of the offset vector. This is because, for the light + propagation problem to be well posed, the assumption must be made that + light only passes through the plane of the known wavefield in one + direction. We always assume that light passes through the initial plane + travelling in the direction of the final plane. - Practically, if one wants to simulate inverse propagation, there are then + Practically, if one wants to simulate inverse propagation, there are two possible approaches. First, one can use the inverse_near_field function, which simulates the inverse propagation problem and therefore will naturally simulate propagation in the opposite direction. Second, - one can explicitly include a propagation_vector argument, which overrides - the offset vector in defining the direction in which light passes through - the input plane. However, in this case, the resulting light field will have - the overall phase accumulation due to propagation along the propagation - vector removed, which may not be the intended behavior. However, this is - not recommended, as inverse propagation will tend to magnify evanescent - waves - it is therefore preferable (unless there is a specific need to - account for evanescent waves properly) to use the inverse near field - propagator - + one can explicitly include a propagation_vector argument, in the direction + opposite to the offset vector. Note for both cases that this function + will only return propagators which suppress evanescent waves. Thus, + propagating forward and then backward by either of these two methods + will lead to a supression of evanescant waves. If you need a propagator + that will cause evanescent waves to undergo exponential growth, good + for you, but this function will not provide it for you. Parameters ---------- @@ -467,126 +468,101 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o A phase mask which accounts for the phase change that each plane wave will undergo. """ - # These check for any pytorch inputs and convert them - try: - basis = basis.detach().cpu().numpy() - except: - pass + # make sure everything is in pytorch, and set the propagation vector + # appropriately if propagate_along_offset is chosen + basis = t.as_tensor(basis) + offset_vector = t.as_tensor(offset_vector, dtype=basis.dtype) - shape = np.array(tuple(shape)) - - try: - offset_vector = offset_vector.detach().cpu().numpy() - except: - pass - - if propagation_vector is not None: - try: - propagation_vector = propagation_vector.detach().cpu().numpy() - except: - pass - - try: - wavelength = wavelength.detach().cpu().numpy()[0] - except: - pass - - - # First we calculate a dual basis for the real space grid - inv_basis = np.linalg.pinv(basis).transpose() - - # Then we calculate the frequencies in (i,j) space - ki = 2 * np.pi * fftpack.fftfreq(shape[0]) - kj = 2 * np.pi * fftpack.fftfreq(shape[1]) - K_ij = np.stack(np.meshgrid(ki,kj, indexing='ij')) - - # Now we convert these to frequencies in reciprocal space - # These frequencies span the 2D plane of the input wavefield. - K_xyz = np.tensordot(inv_basis, K_ij, axes=1) - - # Now we need to apply two corrections to the standard AS method. - # First, we calculate a phase mask which corresponds to the - # shift of the final plane away from the perpendicular direction - # from the input plane. We don't need to extract the perpendicular - # component of the shift because the K_xyz vectors are naturally in the - # input plane. - - # This may have a sign error - must be checked - phase_mask = np.exp(-1j * np.tensordot(offset_vector,K_xyz,axes=1)) - - # Next, we apply a shift to the k-space vectors which sets up - # propagation such that a uniform phase object will propagate along the - # propagation axis. This is not modeling a physical effect, but simply is - # the clearest way to do a rigorous simulation while preventing - # aliasing-related challenges. If used (as is by default), be aware - # and prepare the input wavefields appropriately. if propagate_along_offset: propagation_vector = offset_vector - perpendicular_dir = np.cross(basis[:,1],basis[:,0]) - perpendicular_dir /= np.linalg.norm(perpendicular_dir) - offset_perpendicular = np.dot(perpendicular_dir, offset_vector) - - - k0 = 2*np.pi/wavelength - - sign_correction = 1 - - # Only implement the shift if the flag is set to True if propagation_vector is not None: - - prop_perpendicular = np.dot(perpendicular_dir, propagation_vector) - prop_parallel = propagation_vector - perpendicular_dir \ - * prop_perpendicular - - - if np.linalg.norm(propagation_vector) == 0: - # for numerical stability if this is exactly zero we need - # a special case - k_offset = np.array([0,0,0]) - else: - k_offset = prop_parallel * k0 / np.linalg.norm(propagation_vector) + propagation_vector = t.as_tensor(propagation_vector, dtype=basis.dtype) - K_xyz = K_xyz - k_offset[:,None,None] - - # There apparently is a sign correction that I need to apply - #sign_correction = np.sign(np.dot(perpendicular_dir,propagation_vector)) - sign_correction = np.sign(np.dot(offset_vector,propagation_vector)) - - # we also need to remove the z-dependence on the phase - # This time, though, the z-dependence actually has to do with - # the out of plane component of k at the central offset. Normally - # this is 0, so the z-component is just k0, but not in this case - # We only need one case here, unlike with the propagator, because - # k_offset will always be less than k0 - phase_mask *= np.exp(-1j * np.sqrt(k0**2 - np.linalg.norm(k_offset)**2) - * sign_correction - * np.abs(offset_perpendicular)) - + # + # In this section, we calculate the wavevectors associated with each + # pixel in Fourier space. This is the meat of the function + # - # Redefine this as complex so the square root properly gives - # k>k0 components imaginary frequencies - k0 = np.complex128(k0) + # First we calculate a dual basis for the real space grid + inv_basis = t.linalg.pinv(basis).transpose(0,1) + + # Then we calculate the frequencies in (i,j) space + ki = 2 * np.pi * t.fft.fftfreq(shape[0], dtype=inv_basis.dtype) + kj = 2 * np.pi * t.fft.fftfreq(shape[1], dtype=inv_basis.dtype) + K_ij = t.stack(t.meshgrid(ki,kj)) - # Finally, generate the propagator! - # Must have cases to ensure that evanescent waves decay instead of grow - if sign_correction > 0: - propagator = np.exp(1j*np.sqrt(k0**2 - np.linalg.norm(K_xyz,axis=0)**2) - * sign_correction * np.abs(offset_perpendicular)) + # Now we convert these to frequencies in reciprocal space + # These frequencies span the 2D plane of the input wavefield, + # hence K_ip for "in-plane" + K_ip = t.tensordot(inv_basis, K_ij, dims=1) + + + # Now, we need to generate the out-of-plane direction, so we can + # expand these Ks to the full Ks in 3D reciprocal space. + + # THis is broken down into steps to avoid floating point underflow + # which was a real problem that showed up for electron ptycho + b1_dir = basis[:,0] / t.linalg.norm(basis[:,0]) + b2_dir = basis[:,1] / t.linalg.norm(basis[:,1]) + perpendicular_dir = t.cross(b1_dir,b2_dir) + perpendicular_dir /= t.linalg.norm(perpendicular_dir) + + # We set the sign of the propagation direction appropriately + if propagation_vector is not None: + perpendicular_dir *= t.sign(t.dot(perpendicular_dir,propagation_vector)) else: - propagator = np.exp(-1j * np.conj(np.sqrt(k0**2 - - np.linalg.norm(K_xyz,axis=0)**2)) - * np.abs(offset_perpendicular)) - propagator *= phase_mask + pass + perpendicular_dir *= t.sign(t.dot(perpendicular_dir,offset_vector)) - - # Take the conjugate explicitly here instead of negating - # the previous expression to ensure that complex frequencies - # get mapped to values <1 instead of >1 - propagator = complex_to_torch(np.conj(propagator)) - - return propagator.to(**kwargs) + # Then, if we have a propagation vector, we shift the in-plane + # components of all the pixels to be centered around the in-plane + # component of the propagation vector. + if propagation_vector is not None: + prop_dir = (propagation_vector / + t.linalg.norm(propagation_vector)) + K_0 = 2*np.pi / wavelength * prop_dir + K_0_ip = K_0 - t.dot(perpendicular_dir,K_0) * perpendicular_dir + K_ip += K_0_ip[:,None,None] + else: + K_0 = t.zeros_like(offset_vector) + # Now, we have accurate in-plane values for K, so we can calculate the + # out-of-plane part. We start by calculating it's squared magnitude + K_oop_squared = (2*np.pi/wavelength)**2 - t.linalg.norm(K_ip,dim=0)**2 + + # Then, we take the square root and assign it the appropriate direction, + # adding to get the full 3D wavevectors. Note that we convert to complex + # before the square root to appropriately map negative numbers to + # complex frequencies + K = K_ip + perpendicular_dir[:,None,None] \ + * t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared))) + + # + # In this section, we take the inner product of the calcualted + # wavevectors with the offset vector, to get the phase shift + # experienced by each plane wave. + # + + # We need to convert to complex because K is complex + offset_vector = t.complex(offset_vector,t.zeros_like(offset_vector)) + + # We then subtract off K_0, essentially setting the phase offset + # experienced by K_0 to 0. K_0 will already have been set to 0 if + # there was no propagation vector set. + K_m_K_0 = K - K_0[:,None,None] + # We actually calculate the phase mask + phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1) + + # And we take the conjugate if needed, which makes sure that the + # imaginary part is always positive (supressing evanescent waves) + if propagation_vector is not None \ + and t.sign(t.dot(propagation_vector,offset_vector.real)) == -1: + phase_mask = t.conj(phase_mask) + + # We return the final mask, in the form asked for! + return t.exp(1j*phase_mask).to(*args,**kwargs) + def near_field(wavefront, angular_spectrum_propagator): """ Propagates a wavefront via the angular spectrum method @@ -595,13 +571,14 @@ def near_field(wavefront, angular_spectrum_propagator): represents the real and imaginary components of the wavefield, and returns the near-field propagated version of it. It does this using the supplied angular spectrum propagator, which is a premade - phase mask. + phase mask representing the Fourier transform of the kernel for + light propagation in the desired geometry. Parameters ---------- wavefront : torch.Tensor - The JxNxMx2 stack of complex wavefronts to be propagated + The (Leading Dims)xNxM stack of complex wavefronts to be propagated angular_spectrum_propagator : torch.Tensor The NxM phase mask to be applied during propagation @@ -610,7 +587,7 @@ def near_field(wavefront, angular_spectrum_propagator): propagated : torch.Tensor The propagated wavefront """ - return t.ifft(cmult(angular_spectrum_propagator,t.fft(wavefront,2)), 2) + return t.fft.ifft2(angular_spectrum_propagator * t.fft.fft2(wavefront)) @@ -651,7 +628,8 @@ def inverse_near_field(wavefront, angular_spectrum_propagator): propagated : torch.Tensor The inverse propagated wavefront """ - return t.ifft(cmult(t.fft(wavefront,2), cconj(angular_spectrum_propagator)), 2) + return t.fft.ifft2(t.fft.fft2(wavefront) + * t.conj(angular_spectrum_propagator)) diff --git a/README.md b/README.md index 711e646..6ada9df 100644 --- a/README.md +++ b/README.md @@ -6,21 +6,21 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A # imports from matplotlib import pyplot as plt from CDTools.datasets import Ptycho_2D_Dataset -from CDTools.models import SimplePtycho +from CDTools.models import FancyPtycho # Load the file dataset = Ptycho_2D_Dataset.from_cxi('ptycho_data.cxi') # Generate a model from the data -model = SimplePtycho.from_dataset(dataset) +model = FancyPtycho.from_dataset(dataset) # Run a reconstruction -for i, loss in enumerate(model.Adam_optimize(10, dataset)): - print(i, loss) +for loss in model.Adam_optimize(10, dataset): + print(model.report()) # And look at the results! -model.inspect(dataset) -model.compare(dataset) +model.inspect(dataset) # See the reconstructed object, probe, etc. +model.compare(dataset) # See how the simulated and measured patterns compare plt.show() ``` diff --git a/conda_requirements.txt b/conda_requirements.txt index da4dcb2..4726895 100644 --- a/conda_requirements.txt +++ b/conda_requirements.txt @@ -2,7 +2,7 @@ numpy>=1.0 scipy>=1.0 matplotlib>=2.0 python-dateutil -pytorch>=1.3.0 +pytorch>=1.9.0 h5py>=2.1 pytest sphinx diff --git a/converters/NSLS2_HXN_hdf5_to_CXI.py b/converters/NSLS2_HXN_hdf5_to_CXI.py deleted file mode 100755 index ed800f4..0000000 --- a/converters/NSLS2_HXN_hdf5_to_CXI.py +++ /dev/null @@ -1,131 +0,0 @@ -""" -Purpose: Convert NSLSII HXN hdf5 files to CXI files for analysis with CDTools. -Author: David Rower -Date: December 2019 -""" - -import numpy as np -import pickle -import h5py -import os -import CDTools -from CDTools.tools import data as cdtdata -from matplotlib import pyplot as plt -from scipy.spatial.transform import Rotation -from datetime import datetime - -def create_cxi_from_NSLS2_HXN_2DFly(data_dir, save_str, scan_number, - wavelength, theta, ROI_corner_xy, metadata): - """Converts NSLS2 HXN 2D Fly scan data (from pickle and hdf5) to CXI format - - Assumes scan files will live in data_dir with naming convention - pickle: /scan_.pickle, - hdf5: /scan_.hdf5, - and will create the file /scan_.cxi. - - Parameters - ---------- - data_dir : str - Input data directory - save_str : str - Output data name - scan_number : int - A scan index number - theta : float - Rotation angle of sample in HXN convention, in degrees - ROI_corner_xy : np.array - 1x2 array containing x, y corner of detector ROI - metadata : dict - Contains metadata relevant to the experiment - """ - - ## Load in pickle and hdf5 files - scan_str = "scan_" + scan_number - - # Load pickle (includes useful data about scan not in .hdf5 file) - with open(os.path.join(data_dir, scan_str+".pickle"), 'rb') as f: - scan_pickle = pickle.load(f) - - assert scan_pickle['plan_type'] == "FlyPlan2D", "Code only for FlyPlan2D." - - # Load hdf5 file - scan_hdf5 = h5py.File(os.path.join(data_dir, scan_str+".h5"), 'r') - - - ## Let's attempt to convert this bad boy - scan_cxi = cdtdata.create_cxi(save_str) - - - ## Add source - cdtdata.add_source(scan_cxi, wavelength=wavelength) - scan_cxi['entry_1/instrument_1/source_1']['name'] = scan_pickle['beamline_id'] - - - ## Add sample - theta = np.radians(theta) - sample_unit_vecs = Rotation.from_rotvec(-theta * np.array([0,1,0])).as_dcm() - orientation = np.hstack((sample_unit_vecs[:,0], sample_unit_vecs[:,1])) - translation = np.zeros(3) - sample_info_dict = { - "name" : "TaTe4", - "orientation" : orientation, - "translation" : translation - } - cdtdata.add_sample_info(scan_cxi, sample_info_dict) - - ## Add other metadata for experiment - metadata['start_time'] = datetime.fromtimestamp(scan_pickle['time']) - cdtdata.add_entry_info(scan_cxi, metadata) - - - ## Add detector - - # Constant detector parameters - detector_pixel_size = 55e-6 # meters - detector_height_px = 515 # px ### WARNING: NEED TO CHECK THIS - detector_width_px = 515 # px - - # Geometry parameters from scan files - distance = scan_pickle['dist_detector'] * 1e-3 # assuming mm, almost sure - gamma = np.radians(scan_pickle['gamma_detector']) - delta = np.radians(scan_pickle['delta_detector']) - Rg = Rotation.from_rotvec(-gamma * np.array([0,1,0])).as_dcm() # cw about y - Rd = Rotation.from_rotvec(-delta * Rg[:,0]).as_dcm() # cw about rotated x - RdRg = np.matmul(Rd, Rg) - - # Define detector basis: row vectors for y and x detector axes - basis = detector_pixel_size * np.array([[0.,-1.,0.],[-1.,0.,0.]]) - basis = np.matmul(RdRg,basis.T).T - - # Define corner posiiton: first find center, then offset it - corner_pos = np.dot(RdRg, distance * np.array([0.,0.,1.])) - if ROI_corner_xy[0] is None: - ROI_corner_xy[0] = 0. - if ROI_corner_xy[1] is None: - ROI_corner_xy[1] = 0. - corner_pos -= basis[0,:] * (detector_width_px/2. - ROI_corner_xy[0]) - corner_pos -= basis[1,:] * (detector_height_px/2. - ROI_corner_xy[1]) - - # Add detector data finally - cdtdata.add_detector(scan_cxi, distance, basis.T, corner=corner_pos) - - - ## Add data - axes = ['translation'] + scan_pickle['axes'] # THIS IS ONLY FOR FLY2D - data = np.copy(scan_hdf5['entry']['instrument']['detector']['data']) - data[data == 0] = 1 # to prevent divide by zero in log error - cdtdata.add_data(scan_cxi, data, axes) - - - ## Add translations - x_bounds = scan_pickle['scan_range'][0] - y_bounds = scan_pickle['scan_range'][1] - xx, yy = np.meshgrid(np.linspace(*x_bounds, scan_pickle['num1']), - np.linspace(*y_bounds, scan_pickle['num2'])) - translations = (1e-6 * - np.stack((xx.ravel(), yy.ravel(), np.zeros_like(xx.ravel())), axis=1)) - cdtdata.add_ptycho_translations(scan_cxi, translations) - - - ## Close hdf5 file - scan_hdf5.close() diff --git a/converters/TITAN_stem.py b/converters/TITAN_stem.py deleted file mode 100644 index 611d6c4..0000000 --- a/converters/TITAN_stem.py +++ /dev/null @@ -1,208 +0,0 @@ -""" -Purpose: Convert file collection from Jim Lebeau's TITAN microscope to .CXI -Author: Abe Levitan -Date: January 2019 -""" - -import numpy as np -import pickle -import h5py -import os -import CDTools -from CDTools.tools import data as cdtdata -from CDTools.datasets import Ptycho2DDataset -from matplotlib import pyplot as plt -from scipy.spatial.transform import Rotation -from datetime import datetime -import xml.etree.ElementTree as ET - - - -def load_raw_image_stack(filename): - # The resulting data is an array of (exposure, image-i, image-j), - # with image0i corresponding to y and image-j corresponding to x - # Note that the real-space scanning is done from the bottom right - # corner, first heading left (in x) then scanning up. - rawdata = np.fromfile(filename,dtype='`_ * `sphinx-argparse `_ * `sphinx_rtd_theme `_ - * `pathlib2 `_ + -All of these can be installed via pip or conda. Finally, CDTools is written to be python 2.7+ compatible, but is only actively tested on python 3. +All of these can be installed via pip or conda. Finally, CDTools is NOT python 2 compatible. **It is required that pytorch is built with MKL**, as that enables FFTs. Additionally, installing pytorch **with CUDA support** is recommended, if you intend to run any serious reconstructions with the package. diff --git a/docs/source/intro.rst b/docs/source/intro.rst index 5c981e8..b550ff0 100644 --- a/docs/source/intro.rst +++ b/docs/source/intro.rst @@ -7,24 +7,23 @@ Introduction to CDTools Introduction to CDTools ----------------------- -CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation based approach. +CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation (AD) based approach. .. code-block:: python # imports + import CDTools from matplotlib import pyplot as plt - from CDTools.datasets import Ptycho2DDataset - from CDTools.models import SimplePtycho # Load the file - dataset = Ptycho2DDataset.from_cxi('ptycho_data.cxi') + dataset = CDTools.datasets.Ptycho2DDataset.from_cxi('ptycho_data.cxi') # Generate a model from the data - model = SimplePtycho.from_dataset(dataset) + model = CDTools.models.SimplePtycho.from_dataset(dataset) # Run a reconstruction - for i, loss in enumerate(model.Adam_optimize(10, dataset)): - print(i, loss) + for loss in model.Adam_optimize(20, dataset): + print(model.report()) # And look at the results! model.inspect(dataset) @@ -32,7 +31,7 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A plt.show() -CDTools makes it simple to load and inspect data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a bunch of modular functions for AD ptychography, which can then be used right away from the same scripting framework. +CDTools makes it simple to load and inspect data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a modular functions for AD ptychography. These can can be used to construct new forward models. The high-level interface to CDTools is built on a lower level "three-legged stool". This consists of tools to access stored data, tools to visualize data and reconstructions, and tools that implement basic operations relevant to coherent diffraction. All of these tools can be used directly alongside the high-level interface, when needed. diff --git a/docs/source/tools/cmath.rst b/docs/source/tools/cmath.rst deleted file mode 100644 index 763bb0a..0000000 --- a/docs/source/tools/cmath.rst +++ /dev/null @@ -1,5 +0,0 @@ -Cmath -===== - -.. automodule:: CDTools.tools.cmath - :members: diff --git a/docs/source/tools/index.rst b/docs/source/tools/index.rst index 5d8bea6..3c17dc6 100644 --- a/docs/source/tools/index.rst +++ b/docs/source/tools/index.rst @@ -4,7 +4,6 @@ Tools .. toctree:: :maxdepth: 1 - cmath image_processing data initializers @@ -14,4 +13,3 @@ Tools losses plotting analysis - projectors diff --git a/docs/source/tools/projectors.rst b/docs/source/tools/projectors.rst deleted file mode 100644 index 60b042f..0000000 --- a/docs/source/tools/projectors.rst +++ /dev/null @@ -1,5 +0,0 @@ -Projectors -========== - -.. automodule:: CDTools.tools.projectors - :members: diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst index 04528ce..1c16313 100644 --- a/docs/source/tutorial.rst +++ b/docs/source/tutorial.rst @@ -13,11 +13,9 @@ Our first step will be creating the file and filling out the boilerplate: All th .. code-block:: python - from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt - import pickle + from scipy import io You can always import more libraries, like numpy, or pytorch, or pandas, or what have you, as needed. Next, we load the dataset and give it a look-over @@ -46,9 +44,9 @@ We then try a basic Adam reconstruction with this model, with no changes to the .. code-block:: python - for i, loss in enumerate(model.Adam_optimize(50, dataset)): + for loss in model.Adam_optimize(50, dataset): model.inspect(dataset) - print(i,loss) + print(model.report()) model.compare(dataset) plt.show() @@ -58,7 +56,7 @@ It is worth noting here exactly how this code is working. The reconstruction met In CDTools, every reconstruction method will return a generator. Whenever the generator is asked for the next item, it runs a single epoch of the reconstructionalgorithm, and then returns the average loss over that epoch as that next item. This allows the execution of the reconstruction algorithm to pause once every epoch, allowing some time for the user to run a small snippet of code to inspect how the reconstruction is coming along. -From the end user perspective, all this means is: follow the format above, or more generally put the :code:`model.Adam_optimize(n, dataset)` call anywhere that you would feel comfortable putting a call to :code:`range(n)` - list comprehensions, for loops, etc. +From the end user perspective, all this means is: follow the format above, or more generally put the :code:`model.Adam_optimize(n, dataset)` call anywhere that you would feel comfortable putting a call to :code:`range(n)` - list comprehensions, for loops, etc. In this case, we have called a function to plot out the current state of the reconstruction, and a function to print out the current loss and iteration time. Once we run this, we can take a look at the result. What we see is pretty good, but we can see that there are some issues with the reconstruction near the edge, and the probe itself seems to be larger than the "stage" on which we're reconstructing it. So, we can make two tweaks to this code in response. First, we increase the oversampling ratio, which doubles the size of the stage (this often can cause other issues as well, but generally works well in situations like this where the probe is honestly too large. @@ -94,8 +92,8 @@ Now we expect to get a nice reconstruction, so we can save the data. You can sav .. code-block:: python - with open('example_reconstructions/lab_ptycho.pickle', 'wb') as f: - pickle.dump(model.save_results(dataset),f) + io.savemat('example_reconstructions/lab_ptycho.pickle', + model.save_results(dataset)) This is usually placed before the call to :code:`plt.show()`, to make sure that if the user manually exits the program once all the plot windows are opened, the data will still have been saved. @@ -127,7 +125,7 @@ We can start with the basic skeleton for this file. In addition to our standard .. code-block:: python - from __future__ import division, print_function, absolute_import + import numpy as np import torch as t from matplotlib import pyplot as plt @@ -288,7 +286,6 @@ Once again, we start with the basic skeleton .. code-block:: python - from __future__ import division, print_function, absolute_import import numpy as np import torch as t from CDTools.models import CDIModel @@ -323,10 +320,10 @@ It's important to note that there's not requirement for what the arguments to th self.probe_basis = t.Tensor(probe_basis) # We rescale the probe so it learns at the same rate as the object - self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32))) - self.probe = t.nn.Parameter(probe_guess.to(t.float32) + self.probe_norm = t.max(t.abs(probe_guess).to(t.float32)) + self.probe = t.nn.Parameter(probe_guess.to(t.complex64) / self.probe_norm) - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + self.obj = t.nn.Parameter(obj_guess.to(t.complex64)) Here, we chose to define the model based on a basis matrix describing the probe array, an initial guess at the probe, and an initial object. In addition, an optional offset for the translations is included. @@ -375,7 +372,7 @@ To initialize the object from a dataset, we need to start by extracting the rele probe_shape, det_slice) - obj = t.ones(obj_size+(2,)) + obj = t.ones(obj_size, dtype=t.complex64) return cls(probe_basis, probe, obj, min_translation=min_translation) @@ -472,9 +469,9 @@ At the moment, there is no consistent way to save out the results across the boa .. code-block:: python def save_results(self): - probe = tools.cmath.torch_to_complex(self.probe.detach().cpu()) + probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() - obj = tools.cmath.torch_to_complex(self.obj.detach().cpu()) + obj = self.obj.detach().cpu().numpy() return {'probe':probe,'obj':obj} @@ -499,9 +496,9 @@ We can test this model with a simple script, shown below. By filling in the back model.to(device='cuda') dataset.get_as(device='cuda') - for i, loss in enumerate(model.Adam_optimize(100, dataset)): + for loss in model.Adam_optimize(100, dataset): model.inspect(dataset) - print(i,loss) + print(model.report()) model.compare(dataset) plt.show() diff --git a/examples/MIT_BNL_logo.py b/examples/MIT_BNL_logo.py index f6c31cf..07d7ad8 100644 --- a/examples/MIT_BNL_logo.py +++ b/examples/MIT_BNL_logo.py @@ -1,5 +1,3 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt @@ -23,7 +21,7 @@ dataset.mask = dataset.mask[70:-70,70:-70] model = CDTools.models.FancyPtycho.from_dataset(dataset, translation_scale = 4, n_modes=2, - propagation_distance=-73e-6) + propagation_distance=73e-6) # Move to the GPU model.to(device='cuda') @@ -32,20 +30,20 @@ dataset.get_as(device='cuda') # We turn off position reconstruction for the first phase model.translation_offsets.requires_grad = False -for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15)): +for loss in model.Adam_optimize(10, dataset, batch_size=15): model.inspect(dataset) - print(i,loss) + print(model.report()) # And we turn it on for the second phase model.translation_offsets.requires_grad = True -for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=15)): +for loss in model.Adam_optimize(20, dataset, batch_size=15): model.inspect(dataset) - print(i,loss) + print(model.report()) # The third phase lowers the rate further -for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005)): +for loss in model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005): model.inspect(dataset) - print(i,loss) + print(model.report()) model.inspect(dataset) diff --git a/examples/basic_ptycho_dataset.py b/examples/basic_ptycho_dataset.py index 1848873..1dd5e26 100644 --- a/examples/basic_ptycho_dataset.py +++ b/examples/basic_ptycho_dataset.py @@ -1,4 +1,3 @@ -from __future__ import division, print_function, absolute_import import numpy as np import torch as t from matplotlib import pyplot as plt diff --git a/examples/debug_3d.py b/examples/debug_3d.py deleted file mode 100644 index c39ab90..0000000 --- a/examples/debug_3d.py +++ /dev/null @@ -1,17 +0,0 @@ -from __future__ import division, print_function, absolute_import -import CDTools -from matplotlib import pyplot as plt -# First, we load an example dataset from a .cxi file -filename = '/media/Data Bank/APS_HXN_07_19/CXIs/scan_94513_cxi.h5' - -#filename = 'data/scan_94361_cxi.h5' -dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) -# Next, we create a ptychography model from the dataset -model = CDTools.models.Simple3DBPP.from_dataset(dataset) -#import pdb; pdb.set_trace() -model.inspect(dataset) -plt.show() - -# # Now, we run a short reconstruction from the dataset! -# for i, loss in enumerate(model.Adam_optimize(1, dataset)): -# print(i, loss) diff --git a/examples/ensemble_analysis.py b/examples/ensemble_analysis.py index bfdfe37..bc60a5b 100644 --- a/examples/ensemble_analysis.py +++ b/examples/ensemble_analysis.py @@ -1,21 +1,12 @@ -from __future__ import division, print_function, absolute_import - from matplotlib import pyplot as plt -import pickle +from scipy import io from CDTools.tools import plotting from CDTools.tools import analysis -with open('example_reconstructions/gold_balls_ensemble.pickle', 'rb') as f: - dataset = pickle.load(f) - -# This converts from a list of dictionaries to a dictionary of lists -# It's safe to assume that all elements have the same set of keys -if type(dataset) == type([]): - dataset = {key: [element[key] for element in dataset] - for key in dataset[0]} +dataset = io.loadmat('example_reconstructions/gold_balls_ensemble.mat') # Now we synthesize an average reconstruction @@ -25,7 +16,6 @@ synth_probe, synth_obj, aligned_objs = analysis.synthesize_reconstructions( # And then we calculate the consistency PRTF from this freqs, prtf = analysis.calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'][0]) - # Plot the first mode in detail plotting.plot_phase(synth_probe[0],basis=dataset['basis'][0]) plotting.plot_amplitude(synth_probe[0],basis=dataset['basis'][0]) diff --git a/examples/ensemble_reconstruction.py b/examples/ensemble_reconstruction.py index ea85994..b57a33b 100644 --- a/examples/ensemble_reconstruction.py +++ b/examples/ensemble_reconstruction.py @@ -1,16 +1,16 @@ -from __future__ import division, print_function, absolute_import - import CDTools -import pickle +import numpy as np +from scipy import io # Load the data filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) results = [] +n = 25 -for idx in range(25): - print('Starting Reconstruction', idx) +for idx in range(n): + print('Starting Reconstruction', idx+1, 'of',n) # Create a new model each time model = CDTools.models.FancyPtycho.from_dataset(dataset,n_modes=3, @@ -21,14 +21,21 @@ for idx in range(25): dataset.get_as(device='cuda') # Run the reconstruction - for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)): - print(i,loss) + for loss in model.Adam_optimize(30, dataset, batch_size=100): + print(model.report(), end='\r') - # And add the results to the ensemble + # Print a summary that won't be overwritten by the next line + print('Finished:',model.report()) + + # And add the results to the ensemble dictionary results.append(model.save_results(dataset)) - +# This converts from a list of dictionaries to a dictionary of lists +# It's safe to assume that all elements have the same set of keys +# After this, e.g. dataset['probe'] will return a list of all probes. +results = {key: np.array([element[key] for element in results]) + for key in results[0].keys()} + +print(results['probe'].shape) # Save out the ensemble -with open('example_reconstructions/gold_balls_ensemble.pickle', 'wb') as f: - pickle.dump(results,f) - +io.savemat('example_reconstructions/gold_balls_ensemble.mat', results) diff --git a/examples/example_data/Optical_ptycho.pickle b/examples/example_data/Optical_ptycho.pickle deleted file mode 100644 index 712cce4..0000000 Binary files a/examples/example_data/Optical_ptycho.pickle and /dev/null differ diff --git a/examples/example_reconstructions/gold_balls.mat b/examples/example_reconstructions/gold_balls.mat new file mode 100644 index 0000000..c760136 Binary files /dev/null and b/examples/example_reconstructions/gold_balls.mat differ diff --git a/examples/example_reconstructions/gold_balls.pickle b/examples/example_reconstructions/gold_balls.pickle deleted file mode 100644 index 9115f28..0000000 Binary files a/examples/example_reconstructions/gold_balls.pickle and /dev/null differ diff --git a/examples/example_reconstructions/gold_balls_ensemble.mat b/examples/example_reconstructions/gold_balls_ensemble.mat new file mode 100644 index 0000000..1b812cd Binary files /dev/null and b/examples/example_reconstructions/gold_balls_ensemble.mat differ diff --git a/examples/example_reconstructions/gold_balls_ensemble.pickle b/examples/example_reconstructions/gold_balls_ensemble.pickle deleted file mode 100644 index 946cb05..0000000 Binary files a/examples/example_reconstructions/gold_balls_ensemble.pickle and /dev/null differ diff --git a/examples/example_reconstructions/lab_ptycho.mat b/examples/example_reconstructions/lab_ptycho.mat new file mode 100644 index 0000000..491cfa7 Binary files /dev/null and b/examples/example_reconstructions/lab_ptycho.mat differ diff --git a/examples/example_reconstructions/lab_ptycho.pickle b/examples/example_reconstructions/lab_ptycho.pickle deleted file mode 100644 index 745ca6d..0000000 Binary files a/examples/example_reconstructions/lab_ptycho.pickle and /dev/null differ diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 6917c0e..5149bfd 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -1,30 +1,31 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt -import pickle -import time +from scipy import io # First, we load an example dataset from a .cxi file filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) # Next, we create a ptychography model from the dataset -# Note that we explicitly as for two incoherent probe modes +# Note that we explicitly ask for two incoherent probe modes model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2) # Let's do this reconstruction on the GPU, shall we? model.to(device='cuda') dataset.get_as(device='cuda') -for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)): +# Now, we run a short reconstruction from the dataset +for loss in model.Adam_optimize(100, dataset, batch_size=50, schedule=True): # And we liveplot the updates to the model as they happen - print(i,loss) + print(model.report()) model.inspect(dataset) -# And we save the reconstruction out to a file -with open('example_reconstructions/gold_balls.pickle', 'wb') as f: - pickle.dump(model.save_results(dataset),f) +# This orthogonalizes the incoherent probe modes +model.tidy_probes() + +# And we save out the results as a .mat file +io.savemat('example_reconstructions/gold_balls.mat', + model.save_results(dataset)) # Finally, we plot the results model.inspect(dataset) diff --git a/examples/inspect_dataset.py b/examples/inspect_dataset.py index 0b71d13..10e4de5 100644 --- a/examples/inspect_dataset.py +++ b/examples/inspect_dataset.py @@ -1,5 +1,3 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt diff --git a/examples/lab_bragg_2d_ptycho.py b/examples/lab_bragg_2d_ptycho.py index fe6f6f5..d7c5f94 100644 --- a/examples/lab_bragg_2d_ptycho.py +++ b/examples/lab_bragg_2d_ptycho.py @@ -1,36 +1,32 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt -import pickle - +from scipy import io # -# This is the testing file for the Bragg ptychography code that includes -# corrections for probe propagation and for off-axis far field diffraction -# I'm starting to build it up using the lab ptychography forward data so I -# can initially test that it works on "softball" data before moving to the -# actual Bragg geometry data +# This is the testing dataset for the Bragg ptychography code, that was +# collected with our optical setup using a low NA zone plate illuminating +# a reflective diffraction grating with writing on it. # -#filename = 'example_data/lab_ptycho_data.cxi' +# This file is too large to be distributed via Github. +# Please contact Abe Levitan (alevitan@mit) if you would like access filename = '/media/Data Bank/Lab Ptycho/Zone Plate Bragg 633.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) #dataset.inspect() -model = CDTools.models.Bragg2DPtycho.from_dataset(dataset,probe_support_radius=60,correct_tilt=False)#propagate_probe=False)#, n_modes=8) +model = CDTools.models.Bragg2DPtycho.from_dataset(dataset,probe_support_radius=60,correct_tilt=False) model.to(device='cuda') dataset.get_as(device='cuda') model.translation_offsets.requires_grad = False -for i, loss in enumerate(model.Adam_optimize(100, dataset)): +for loss in model.Adam_optimize(100, dataset): model.inspect(dataset) - print(i,loss) - -#with open('example_reconstructions/lab_bragg_2d_ptycho.pickle', 'wb') as f: -# pickle.dump(model.save_results(dataset),f) + print(model.report()) -#model.compare(dataset) +io.savemat('example_reconstructions/lab_bragg_2d_ptycho.mat', + model.save_results(dataset)) + +model.compare(dataset) plt.show() diff --git a/examples/lab_ptycho_data.py b/examples/lab_ptycho_data.py index a4cadfe..746ceff 100644 --- a/examples/lab_ptycho_data.py +++ b/examples/lab_ptycho_data.py @@ -1,28 +1,32 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt -import pickle +from scipy import io + +# A simple dataset collected from our optical setup filename = 'example_data/lab_ptycho_data.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) -# dataset.inspect() -# plt.show() +dataset.inspect(units='mm') +plt.show() model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=2, - probe_support_radius=90) + probe_support_radius=90, + n_modes=2, + dm_rank=0,units='mm') model.to(device='cuda') dataset.get_as(device='cuda') model.translation_offsets.requires_grad = False -for i, loss in enumerate(model.Adam_optimize(50, dataset)): +for loss in model.Adam_optimize(50, dataset): model.inspect(dataset) - print(i,loss) - -with open('example_reconstructions/lab_ptycho.pickle', 'wb') as f: - pickle.dump(model.save_results(dataset),f) + print(model.report()) + +model.tidy_probes() + +io.savemat('example_reconstructions/lab_ptycho.mat', + model.save_results(dataset)) model.compare(dataset) plt.show() diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 866a7a0..9ee4933 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -1,18 +1,16 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt # First, we load an example dataset from a .cxi file filename = 'example_data/lab_ptycho_data.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) -plt.ion() + # Next, we create a ptychography model from the dataset model = CDTools.models.SimplePtycho.from_dataset(dataset) # Now, we run a short reconstruction from the dataset! -for i, loss in enumerate(model.Adam_optimize(10, dataset)): - print(i, loss) +for loss in model.Adam_optimize(20, dataset): + print(model.report()) # Finally, we plot the results model.inspect(dataset) diff --git a/examples/simple_ptycho_model.py b/examples/simple_ptycho_model.py index 012af2c..a5047f7 100644 --- a/examples/simple_ptycho_model.py +++ b/examples/simple_ptycho_model.py @@ -1,4 +1,3 @@ -from __future__ import division, print_function, absolute_import import numpy as np import torch as t from CDTools.models import CDIModel diff --git a/examples/specular_ptycho.py b/examples/specular_ptycho.py index f501e26..cf8f24b 100644 --- a/examples/specular_ptycho.py +++ b/examples/specular_ptycho.py @@ -1,5 +1,3 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt import numpy as np diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index 9c49978..6d1bdba 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -1,15 +1,10 @@ -from __future__ import division, print_function, absolute_import - import CDTools -from matplotlib import pyplot as plt import pickle -from torch.utils.data import Subset +from matplotlib import pyplot as plt # First, we load an example dataset from a .cxi file - ss_filename = 'example_data/Optical_Data_ss.cxi' -#with open('example_data/Optical_ptycho.pickle', 'rb') as f: with open('example_data/Optical_ptycho_incoherent.pickle', 'rb') as f: ptycho_results = pickle.load(f) @@ -19,9 +14,9 @@ background = ptycho_results['background'] dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(ss_filename) -# Next, we create a ptychography model from the dataset +# Next, we create an RPI model from the dataset # Note that we explicitly as for two incoherent probe modes -model = CDTools.models.RPI.from_dataset(dataset, probe, [900,900], +model = CDTools.models.RPI.from_dataset(dataset, probe, [500,500], background=background, n_modes=2, initialization='random') @@ -33,20 +28,19 @@ dataset.get_as(device='cuda') # Note that the inspect step takes the vast majority of the time # The regularization is an L2 regularizer that empirically helps accelerate # convergence -for i, loss in enumerate(model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05])):#0.1)): - #model.inspect(dataset) - print(i,loss) - -#model.inspect(dataset) +for loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]): + model.inspect(dataset) + print(model.report()) + # Now we use the regularizer to damp all but the top modes -for i, loss in enumerate(model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1])): +for loss in model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1]): #model.inspect(dataset) - print(i,loss) + print(model.report()) +# Save results to a python dictionary results = model.save_results() - # Finally, we plot the results model.inspect(dataset) model.compare(dataset) diff --git a/examples/unified_modes.py b/examples/unified_modes.py index 05dbc27..666e57f 100644 --- a/examples/unified_modes.py +++ b/examples/unified_modes.py @@ -1,39 +1,37 @@ -from __future__ import division, print_function, absolute_import - import CDTools from matplotlib import pyplot as plt -import pickle +from scipy import io + filename = 'example_data/lab_ptycho_data.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) -# dataset.inspect() -# plt.show() +#dataset.inspect(units='mm') +#plt.show() + +# dm_rank=-1 tells it to use a full-rank unified mode approximation +model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1,n_modes=3, dm_rank=-1) -#model = CDTools.models.UnifiedModePtycho.from_dataset(dataset, oversampling=2,n_modes=3)#, probe_support_radius=90) -model = CDTools.models.UnifiedModePtycho2.from_dataset(dataset, oversampling=1,n_modes=3)#, probe_support_radius=90) -#model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1) model.to(device='cuda') dataset.get_as(device='cuda') - model.translation_offsets.requires_grad = False -for i, loss in enumerate(model.Adam_optimize(100, dataset)): +for loss in model.Adam_optimize(200, dataset): model.inspect(dataset) - print(i,loss) + print(model.report()) model.tidy_probes() -for i, loss in enumerate(model.Adam_optimize(20, dataset, lr=0.0001)): +for loss in model.Adam_optimize(200, dataset, lr=0.0001): model.inspect(dataset) - print(i,loss) + print(model.report()) -model.tidy_probes() +model.tidy_probes(normalize=True) model.inspect(dataset) -with open('example_reconstructions/unified_modes.pickle', 'wb') as f: - pickle.dump(model.save_results(dataset),f) +io.savemat('example_reconstructions/unified_modes.mat', + model.save_results(dataset)) model.compare(dataset) plt.show() diff --git a/kiara_debugging_code/debugging_mod.py b/kiara_debugging_code/debugging_mod.py deleted file mode 100644 index 231f12f..0000000 --- a/kiara_debugging_code/debugging_mod.py +++ /dev/null @@ -1,112 +0,0 @@ - -import numpy as np -from matplotlib import pyplot as plt -import imageio - -#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -class ImageSeries: - - def __init__(self, images, crop,fps=None): - self.images = images #list of images (numpy arrays), saved itself as a numpy array - self.L = len(self.images) - self.x1, self.x2, self.y1, self.y2 = crop #(x1,y1) and (x2,y2) are cropping coords - self.crimages = [] - for i in range(self.L): - self.crimages.append(self.images[i][self.x1:self.x2, self.y1:self.y2]) - self.crimages = np.array(self.crimages) - self.fps = fps - - def show_frame(self,t): #t = time to show - if 0 <= t <= self.L: - plt.imshow(self.crimages[t],interpolation="none") - plt.show() - return - - def save_vid(self, filename, secspersec): - imageio.mimwrite(filename, self.crimages, fps=self.fps*secspersec) - - def waterfall(self,x,y): - if x == None: - #plot row y=y over time - return self.crimages[:,:,y] - elif y == None: - return np.transpose(self.crimages[:,x,:]) - - def plot_waterfall(self,x,y,tint=None): - waterfall = self.waterfall(x,y) - if x == None: - #plot col y over time - fig = plt.figure() - W = fig.add_subplot(111) - W.imshow(waterfall,interpolation="none") - if tint: - plt.yticks(np.arange(0,self.L,self.fps*tint),np.arange(0,self.L/self.fps,tint)) - W.set_title("waterfall plot of row y = " + str(y)) - W.set_ylabel("time [s]") - plt.show() - return - elif y == None: - #plot col x=x over time - fig = plt.figure() - W = fig.add_subplot(111) - W.imshow(waterfall,interpolation="none") - if tint: - plt.xticks(np.arange(0,self.L,self.fps*tint),np.arange(0,self.L/self.fps,tint)) - W.set_title("waterfall plot of col x = " + str(x)) - W.set_xlabel("time [s]") - plt.show() - return - return - - def Ipixel(self,pixel): - px,py = pixel - return self.crimages[:,px,py] - - def plot_Ipixel(self,pixel,description=""): - I = self.Ipixel(pixel) - fig = plt.figure() - f1 = fig.add_subplot(111) - f1.set_title("I(t) for pixel" + str(pixel) + " (" + description + ")") - f1.set_xlabel("time [s]") - f1.plot(np.arange(0,self.L/self.fps,1/self.fps),I) - plt.show() - return - - def fftIpixel(self,pixel): - I = self.Ipixel(pixel) - IfreqA = np.fft.fft(I)/self.L - Ifreq = np.fft.fftfreq(self.L,d=(1/self.fps)) - return(Ifreq, IfreqA) - - def plot_fftIpixel(self,pixel): - Ifreq, IfreqA = self.fftIpixel(pixel) - fig = plt.figure() - f1 = fig.add_subplot(111) - f1.set_title("FFT for pixel" + str(pixel)) - f1.set_xlabel("frequency [1/s]") - f1.plot(Ifreq, abs(IfreqA)) - plt.show() - return - - def g2(self,pixel): - I = self.Ipixel(pixel) - g2 = [] - avgsq = np.mean(I)**2 - for tau in range(len(I)-1): - if tau == 0: - dotp = 0 - for t in range(len(I)): - dotp += I[t]*I[t] - g2.append(dotp/(len(I)*avgsq) ) - elif tau != 0: - dotp = 0 - for t in range(len(I)-tau): - dotp += I[t]*I[t+tau] - g2.append( dotp / (len(I[:-tau])*avgsq)) - g2 = np.array(g2) - return g2 - - - - diff --git a/kiara_debugging_code/vid_read.py b/kiara_debugging_code/vid_read.py deleted file mode 100644 index 26bf618..0000000 --- a/kiara_debugging_code/vid_read.py +++ /dev/null @@ -1,50 +0,0 @@ - -import numpy as np -from matplotlib import pyplot as plt -import imageio -from PIL import Image, ImageSequence -from debugging_mod import ImageSeries - - -vid = Image.open('data_10_4/kiara_20fps_redlaser_vid.tif') - -vidarray = [] -for i, page in enumerate(ImageSequence.Iterator(vid)): - pg = np.array(page) - vidarray.append(pg) -vidarray = np.array(vidarray) - -RedLaserExp = ImageSeries(vidarray, (430,606,590,766),fps=20) #cropping x1:x2, y1:y2 - -#with np.load('data_9_28/kiara_data_300sec_green') as data: -# GreenLaserExp = ImageSeries(data['arr_0'], (500,676,580,756), fps=5) - #crop to: x1=500, x2=676, y1=580, y2=756 - -#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - -#Note: human eye can see at c. 150 fps - -#inner top right corner of disk (120,70) -#inner top left corner of disk (66,63) -#inner bottom left corner of disk (66,107) -#inner bottom (87,118) -#inner right (125, 86) - -#pixel intensity plot: --------------------------------------- -#RedLaserExp.plot_Ipixel((125,86),description="inner right") -#RedLaserExp.plot_Ipixel((87,118),description="inner bottom") -#RedLaserExp.plot_Ipixel((66,63),description="inner top left") - -#pixel fft plot: -------------------------------------------- -#RedLaserExp.plot_fftIpixel((125,86)) -#RedLaserExp.plot_fftIpixel((87,118)) -#RedLaserExp.plot_fftIpixel((66,63)) - -#waterfall plot for row: ------------------------------------- -#RedLaserExp.plot_waterfall(None,118,tint=1) - -#waterfall plot for col: ------------------------------------- -#RedLaserExp.plot_waterfall(66,None,tint=1) - -#save a video: ----------------------------------------------- -#RedLaserExp.save_vid("redlaser_10-4_20fps_1x.mp4",1) diff --git a/setup.py b/setup.py index 44ee5c6..10cca68 100644 --- a/setup.py +++ b/setup.py @@ -5,8 +5,9 @@ with open("README.md", "r") as fh: setuptools.setup( name="CDTools", - version="0.0.1", - author="Abe Levitan, Madelyn Cain", + version="0.1.1", + python_requires='>3.4', # we use pathlib, introduced in 3.4 + author="Abe Levitan", author_email="alevitan@mit.edu", description="Coherent Diffraction Tools", long_description=long_description, @@ -15,15 +16,13 @@ setuptools.setup( install_requires=[ "numpy>=1.0", "scipy>=1.0", - "matplotlib>=2.0", + "matplotlib>=2.0", # Matplotlib 2.0 introduces better colormaps and no I'm not sorry "python-dateutil", - "torch>=1.3.0", #1.3.0 adds the align_corners option for grid_sample which is used for high NA far-field propagation. 1.2.0 introduced boolean tensors in a breaking way, we use the boolean tensors here for masking. - "h5py>=2.1", - "pathlib2 ; python_version<'3.4'"], + "torch>=1.9.0", #1.9.0 implements support for autograd on indexed complex tensors, which we need in order to use complex tensors in the forward models + "h5py>=2.1"], extras_require={ 'tests': ["pytest"], - 'docs': ["sphinx","sphinx-argparse","sphinx_rtd_theme"], - ":python_version<'3.4'": ["pathlib2"], + 'docs': ["sphinx","sphinx-argparse","sphinx_rtd_theme"] }, packages=setuptools.find_packages(), classifiers=[ diff --git a/tests/conftest.py b/tests/conftest.py index 6b31585..bbf4ecb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,3 @@ -from __future__ import division, print_function, absolute_import import numpy as np import h5py import pytest @@ -104,7 +103,7 @@ def ptycho_cxi_1(): # Remember the format for the CXI file differs from the format used # internally mask = np.zeros((256,256)).astype(np.int32) - expected['mask'] = np.ones((256,256)).astype(np.bool) + expected['mask'] = np.ones((256,256)).astype(bool) d1f.create_dataset('mask',data=mask) # Create an initial background @@ -272,7 +271,7 @@ def ptycho_cxi_3(): # Remember the format for the CXI file differs from the format used # internally mask = np.ones((256,256)).astype(np.uint32) * 0x00001000 - expected['mask'] = np.ones((256,256)).astype(np.bool) + expected['mask'] = np.ones((256,256)).astype(bool) d1f.create_dataset('mask',data=mask) expected['dark'] = None diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 7347ad7..30404f3 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -28,7 +28,7 @@ def test_CDataset_init(): dataset = CDataset(entry_info, sample_info, wavelength, detector_geometry, mask) - assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(np.bool)))) + assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(bool)))) assert dataset.entry_info == entry_info assert dataset.sample_info == sample_info assert dataset.wavelength == wavelength diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 12e897f..e7857c4 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -6,7 +6,7 @@ from scipy import fftpack as ffts import torch as t from itertools import combinations -from CDTools.tools import analysis, cmath, initializers +from CDTools.tools import analysis, initializers def test_orthogonalize_probes(): @@ -30,7 +30,7 @@ def test_orthogonalize_probes(): ortho_probes = analysis.orthogonalize_probes(probes) # test that it also works on torch tensors - ortho_probes_t = cmath.torch_to_complex(analysis.orthogonalize_probes(cmath.complex_to_torch(probes))) + ortho_probes_t = analysis.orthogonalize_probes(t.as_tensor(probes)).numpy() # This tests for orthogonality for p1,p2 in combinations(ortho_probes,2): @@ -80,13 +80,12 @@ def test_standardize(): # Start by making a probe and object that should meet the standardization # conditions - probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)) - probe = cmath.torch_to_complex(probe) + probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy() probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2)) probe = probe * np.exp(-1j * np.angle(np.sum(probe))) assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel())) - assert np.isclose(0,np.angle(np.sum(probe))) + assert np.angle(np.sum(probe)) < 1e-7 obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5)) obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5, @@ -105,9 +104,9 @@ def test_standardize(): assert np.allclose(obj, s_obj) # Test that it works on torch tensors - s_probe, s_obj = analysis.standardize(cmath.complex_to_torch(test_probe).to(t.float32), cmath.complex_to_torch(test_obj).to(t.float32)) - s_probe = cmath.torch_to_complex(s_probe) - s_obj = cmath.torch_to_complex(s_obj) + s_probe, s_obj = analysis.standardize(t.as_tensor(test_probe,dtype=t.complex64), t.as_tensor(test_obj,dtype=t.complex64)) + s_probe = s_probe.numpy() + s_obj = s_obj.numpy() assert np.allclose(probe, s_probe) assert np.allclose(obj, s_obj) @@ -146,7 +145,6 @@ def test_standardize(): -from matplotlib import pyplot as plt def test_synthesize_reconstructions(): # I can only really test for a lack of failures, so I think my plan # will be to create a dataset that just needs to be added and see that @@ -154,13 +152,12 @@ def test_synthesize_reconstructions(): # Start by making a probe and object that should meet the standardization # conditions - probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)) - probe = cmath.torch_to_complex(probe) + probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy() probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2)) probe = probe * np.exp(-1j * np.angle(np.sum(probe))) assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel())) - assert np.isclose(0,np.angle(np.sum(probe))) + assert np.abs(np.angle(np.sum(probe))) < 1e-7 obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5)) obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5, @@ -203,14 +200,14 @@ def test_calc_consistency_prtf(): assert np.allclose(prtf, 0.7) # Check that it also works with torch input - t_synth_obj = cmath.complex_to_torch(synth_obj) - t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack] + t_synth_obj = t.as_tensor(synth_obj) + t_obj_stack = [t.as_tensor(obj) for obj in obj_stack] freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, basis, nbins=30) assert np.allclose(prtf.numpy(), 0.7) # And also when the basis is in torch - t_synth_obj = cmath.complex_to_torch(synth_obj) - t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack] + t_synth_obj = t.as_tensor(synth_obj) + t_obj_stack = [t.as_tensor(obj) for obj in obj_stack] freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30) assert np.allclose(prtf.numpy(), 0.7) @@ -238,11 +235,11 @@ def test_calc_deconvolved_cross_correlation(): assert np.allclose(test_cor, np_cor) # test with pytorch inputs - obj1_t = cmath.complex_to_torch(obj1) - obj2_t = cmath.complex_to_torch(obj2) + obj1_t = t.as_tensor(obj1) + obj2_t = t.as_tensor(obj2) test_cor_t = analysis.calc_deconvolved_cross_correlation(obj1_t,obj2_t, im_slice=np.s_[:,:]) - assert np.allclose(cmath.torch_to_complex(test_cor_t), np_cor) + assert np.allclose(test_cor_t.numpy(), np_cor) @@ -294,8 +291,8 @@ def test_calc_frc(): assert np.allclose(threshold, test_threshold) # try again with complex - obj1_torch = cmath.complex_to_torch(obj1) - obj2_torch = cmath.complex_to_torch(obj2) + obj1_torch = t.as_tensor(obj1) + obj2_torch = t.as_tensor(obj2) basis_torch = t.tensor(basis) test_bins_t, test_frc_t, test_threshold_t = analysis.calc_frc(obj1_torch, diff --git a/tests/tools/test_cmath.py b/tests/tools/test_cmath.py deleted file mode 100644 index e39d821..0000000 --- a/tests/tools/test_cmath.py +++ /dev/null @@ -1,135 +0,0 @@ -from __future__ import division, print_function, absolute_import - -from CDTools.tools import cmath -import numpy as np -import torch as t -from scipy.fftpack import fftshift, ifftshift - -def test_complex_to_torch(): - arr = np.random.rand(100,4) + 1j * np.random.rand(100,4) - tensor = cmath.complex_to_torch(arr) - assert np.allclose(tensor[:,:,0].numpy(),np.real(arr)) - assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr)) - - -def test_torch_to_complex(): - tensor = t.rand(100,4,2) - arr = cmath.torch_to_complex(tensor) - assert np.allclose(tensor[:,:,0].numpy(),np.real(arr)) - assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr)) - - -def test_cabssq(): - arr = np.random.rand(56,23,2) + 1j * np.random.rand(56,23,2) - cabssq = cmath.cabssq(cmath.complex_to_torch(arr)) - assert np.allclose(cabssq.numpy(),np.abs(arr)**2) - - -def test_cabs(): - arr = np.random.rand(2,3,4,5) + 1j * np.random.rand(2,3,4,5) - cabs = cmath.cabs(cmath.complex_to_torch(arr)) - assert np.allclose(cabs.numpy(),np.abs(arr)) - - -def test_cconj(): - arr = np.random.rand(50) + 1j * np.random.rand(50) - cconj = cmath.cconj(cmath.complex_to_torch(arr)) - assert np.allclose(cmath.torch_to_complex(cconj), np.conj(arr)) - - -def test_cmult(): - arr1 = np.random.rand(50) + 1j * np.random.rand(50) - arr2 = np.random.rand(50) + 1j * np.random.rand(50) - mult = cmath.cmult(cmath.complex_to_torch(arr1), - cmath.complex_to_torch(arr2)) - assert np.allclose(cmath.torch_to_complex(mult),arr1*arr2) - - -def test_cdiv(): - arr1 = np.random.rand(50) + 1j * np.random.rand(50) - arr2 = np.random.rand(50) + 1j * np.random.rand(50) - div = cmath.cdiv(cmath.complex_to_torch(arr1), - cmath.complex_to_torch(arr2)) - assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2) - - -def test_cphase(): - arr = np.random.rand(50) + 1j * np.random.rand(50) - cphase = cmath.cphase(cmath.complex_to_torch(arr)) - assert np.allclose(cphase.numpy(), np.angle(arr)) - - -def test_scalars(): - arr1 = np.random.rand(1) + 1j * np.random.rand(1) - arr2 = np.random.rand(1) + 1j * np.random.rand(1) - div = cmath.cdiv(cmath.complex_to_torch(arr1), - cmath.complex_to_torch(arr2)) - assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2) - -def test_fftshift(): - #1D, even - arr = np.random.rand(300) + 1j * np.random.rand(300) - shifted = cmath.fftshift(cmath.complex_to_torch(arr)) - assert np.allclose(fftshift(arr), - cmath.torch_to_complex(shifted)) - #1D, odd - arr = np.random.rand(301) + 1j * np.random.rand(301) - shifted = cmath.fftshift(cmath.complex_to_torch(arr)) - assert np.allclose(fftshift(arr), - cmath.torch_to_complex(shifted)) - - #2D - arr = np.random.rand(20,21) + 1j * np.random.rand(20,21) - shifted = cmath.fftshift(cmath.complex_to_torch(arr)) - assert np.allclose(fftshift(arr), - cmath.torch_to_complex(shifted)) - #3D - arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) - shifted = cmath.fftshift(cmath.complex_to_torch(arr)) - assert np.allclose(fftshift(arr), - cmath.torch_to_complex(shifted)) - - #3D, choosing specific axes - arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) - shifted = cmath.fftshift(cmath.complex_to_torch(arr),dims=(0,1)) - assert np.allclose(fftshift(arr,axes=(0,1)), - cmath.torch_to_complex(shifted)) - - -def test_ifftshift(): - #1D, even - arr = np.random.rand(300) + 1j * np.random.rand(300) - shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) - assert np.allclose(ifftshift(arr), - cmath.torch_to_complex(shifted)) - #1D, odd - arr = np.random.rand(301) + 1j * np.random.rand(301) - shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) - assert np.allclose(ifftshift(arr), - cmath.torch_to_complex(shifted)) - - #2D - arr = np.random.rand(20,21) + 1j * np.random.rand(20,21) - shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) - assert np.allclose(ifftshift(arr), - cmath.torch_to_complex(shifted)) - #3D - arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) - shifted = cmath.ifftshift(cmath.complex_to_torch(arr)) - assert np.allclose(ifftshift(arr), - cmath.torch_to_complex(shifted)) - - #3D, choosing specific axes - arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17) - shifted = cmath.ifftshift(cmath.complex_to_torch(arr),dims=(0,1)) - assert np.allclose(ifftshift(arr,axes=(0,1)), - cmath.torch_to_complex(shifted)) - - -def test_expi(): - phases = np.random.rand(20,70) * 2 * np.pi - np_result = np.exp(1j * phases) - - torch_result = cmath.expi(t.Tensor(phases)) - - assert np.allclose(np_result, cmath.torch_to_complex(torch_result)) diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index 5e0b7d7..a39b15d 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -1,5 +1,3 @@ -from __future__ import division, print_function, absolute_import - from CDTools.tools import data import numpy as np import torch as t @@ -8,10 +6,7 @@ import pytest import os import datetime import numbers -try: - import pathlib -except ImportError: - import pathlib2 as pathlib +import pathlib diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index f37581c..2c2402a 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -5,7 +5,7 @@ import pytest import numpy as np import torch as t -from CDTools.tools import image_processing, cmath, initializers, interactions +from CDTools.tools import image_processing, initializers, interactions from scipy import ndimage from scipy.signal import fftconvolve @@ -32,8 +32,8 @@ def test_centroid_sq(): assert t.allclose(centroid, t.Tensor(sp_centroid)) # Test complex with multiple ims - ims = t.rand((5,30,40,2)) - np_ims = cmath.torch_to_complex(ims) + ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40)) + np_ims = ims.numpy() sp_centroids = [ndimage.measurements.center_of_mass(np.abs(im)**2) for im in np_ims] centroids = image_processing.centroid_sq(ims, comp=True) @@ -51,12 +51,12 @@ def test_sinc_subpixel_shift(): Ys,Xs = np.meshgrid(xs,xs) sinc_im = np.sinc(Xs-0.3) * np.sinc(Ys-0.6) - torch_im = cmath.complex_to_torch(im) + torch_im = t.as_tensor(im) test_im = image_processing.sinc_subpixel_shift(torch_im,(0.3,0.6)) # The fidelity isn't great due to the FFT-based approach, so we need # a pretty relaxed condition - assert np.max(np.abs(sinc_im - cmath.torch_to_complex(test_im))) < 0.005 + assert np.max(np.abs(sinc_im - test_im.numpy())) < 0.005 def test_find_pixel_shift(): @@ -69,13 +69,13 @@ def test_find_pixel_shift(): # Test a real and complex im big_im = t.rand((30,70)) - im1 = t.stack((big_im[:-5,10:],t.zeros_like(big_im[:-5,10:])),dim=-1) + im1 = big_im[:-5,10:].to(dtype=t.complex64) im2 = big_im[5:,:-10] assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10])) assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10])) # Test two complex ims - big_im = t.rand((45,45,2)) + big_im = t.rand((45,45)) + 1j * t.rand((45,45)) im1 = big_im[:-5,:-4] im2 = big_im[5:,4:] assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4])) @@ -83,13 +83,13 @@ def test_find_pixel_shift(): def test_find_subpixel_shift(): # We can do this by creating a test probe and a test object - test_probe = t.rand((70,70,2)) - test_obj = t.ones((300,300,2)) + test_probe = t.rand((70,70)) + 1j * t.rand((70,70)) + test_obj = t.ones((300,300)) + 1j * t.rand((300,300)) shift = t.tensor((0.8,0.75)) - im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift) - + im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, multiple_modes=False) + retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50) # tolerance of 0.03 on this measurement assert t.all(t.abs(shift - retrieved_shift) < 0.03) @@ -98,12 +98,13 @@ def test_find_subpixel_shift(): def test_find_shift(): # We can do this by creating a test probe and a test object - test_probe = t.rand((200,200,2)) - test_obj = t.ones((300,300,2)) + test_probe = t.rand((200,200)) + 1j * t.rand((200,200)) + test_obj = t.ones((300,300)) + 1j * t.rand((300,300)) shift = t.tensor((0.8,0.75)) - im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift)[:-40,:-6] + im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, + multiple_modes=False)[:-40,:-6] retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50) # tolerance of 0.03 on this measurement @@ -111,14 +112,14 @@ def test_find_shift(): def test_convolve_1d(): - from matplotlib import pyplot as plt test_image = np.random.rand(400,300) #test_image = np.hstack((np.ones((400,150)),np.zeros((400,150)))) xs = np.linspace(-100,100,300) kernel = 1/(1+xs**2) # First, we test with everything real, dim=1 - convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(kernel),dim=1) + convolved = image_processing.convolve_1d(t.as_tensor(test_image), + t.as_tensor(kernel),dim=1) np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1)) assert np.allclose(convolved.numpy(),np_result) @@ -128,13 +129,16 @@ def test_convolve_1d(): kernel = 1/(1+xs**2) # Then with dim=0, and a non-fftshifted kernel - convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(np.fft.ifftshift(kernel)), fftshift_kernel=False) + convolved = image_processing.convolve_1d(t.as_tensor(test_image), + t.as_tensor(np.fft.ifftshift(kernel)), + fftshift_kernel=False) np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)) assert np.allclose(convolved.numpy(),np_result) # And finally with complex input - convolved = cmath.torch_to_complex(image_processing.convolve_1d(cmath.complex_to_torch(test_image),cmath.complex_to_torch(kernel))) + convolved = image_processing.convolve_1d(t.as_tensor(test_image,dtype=t.complex64), + t.as_tensor(kernel,dtype=t.complex64)).numpy() np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0) assert np.allclose(convolved,np_result) diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index 4f6871c..e340315 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -1,7 +1,6 @@ from __future__ import division, print_function, absolute_import from CDTools.tools import initializers -from CDTools.tools import cmath from CDTools.datasets import Ptycho2DDataset import numpy as np import torch as t @@ -78,11 +77,12 @@ def test_gaussian(): # Generate gaussian as a numpy array (square array) shape = [10, 10] sigma = [2.5, 2.5] + center = ((shape[0]-1)/2, (shape[1]-1)/2) y, x = np.mgrid[:shape[0], :shape[1]] np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2 -0.5*((y-center[0])/sigma[0])**2) - init_result = cmath.torch_to_complex(initializers.gaussian([10, 10], [2.5, 2.5], amplitude=10)) + init_result = initializers.gaussian(shape, sigma, amplitude=10).numpy() assert np.allclose(init_result, np_result) # Generate gaussian as a numpy array (rectangular array) @@ -92,7 +92,7 @@ def test_gaussian(): y, x = np.mgrid[:shape[0], :shape[1]] np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2 -0.5*((y-center[0])/sigma[0])**2) - init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma)) + init_result = initializers.gaussian(shape, sigma).numpy() assert np.allclose(init_result, np_result) # Generate gaussian with curvature @@ -105,8 +105,8 @@ def test_gaussian(): -0.5*((y-center[0])/sigma[0])**2) np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2 +0.5j*curvature[0]*(y-center[0])**2) - init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma, - center=center, curvature=curvature, amplitude=10)) + init_result = initializers.gaussian(shape, sigma, center=center, + curvature=curvature, amplitude=10).numpy() assert np.allclose(init_result, np_result) @@ -147,8 +147,7 @@ def test_gaussian_probe(ptycho_cxi_1): normalization_1 = normalization / np.sum(np.abs(np_probe)**2) - probe = initializers.gaussian_probe(dataset, basis, shape, sigma) - probe = cmath.torch_to_complex(probe) + probe = initializers.gaussian_probe(dataset, basis, shape, sigma).numpy() assert np.allclose(probe, normalization_1*np_probe) # And then a propagated probe @@ -163,8 +162,7 @@ def test_gaussian_probe(ptycho_cxi_1): normalization_2 = normalization / np.sum(np.abs(np_probe)**2) probe = initializers.gaussian_probe(dataset, basis, shape, sigma, - propagation_distance=z) - probe = cmath.torch_to_complex(probe) + propagation_distance=z).numpy() assert np.allclose(probe, normalization_2*np_probe) @@ -185,10 +183,10 @@ def test_SHARP_style_probe(ptycho_cxi_1): distance) probe = initializers.SHARP_style_probe(dataset, shape, det_slice) - assert probe.shape == t.Size([256,256,2]) + assert probe.shape == t.Size([256,256]) probe = initializers.SHARP_style_probe(dataset, shape, det_slice, propagation_distance=20e-6) - assert probe.shape == t.Size([256,256,2]) + assert probe.shape == t.Size([256,256]) def test_RPI_spectral_init(): @@ -200,24 +198,24 @@ def test_RPI_spectral_init(): probe = np.random.rand(230,253).astype(np.complex64) obj_shape = [37,53] mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04) - background = t.Tensor(np.random.rand(*pattern.shape) .astype(np.float32)* 0.05) + background = t.as_tensor(np.random.rand(*pattern.shape),dtype=t.float32) * 0.05 - probe = cmath.complex_to_torch(probe) - pattern = t.Tensor(pattern) + probe = t.as_tensor(probe) + pattern = t.as_tensor(pattern) obj = initializers.RPI_spectral_init(pattern, probe, obj_shape) - assert list(obj.shape) == [1]+obj_shape+[2] + assert list(obj.shape) == [1]+obj_shape obj = initializers.RPI_spectral_init(pattern, probe, obj_shape, n_modes=2, mask=mask) - assert list(obj.shape) == [2]+obj_shape+[2] + assert list(obj.shape) == [2]+obj_shape obj = initializers.RPI_spectral_init(pattern, probe, obj_shape, n_modes=2, background=background) - assert list(obj.shape) == [2]+obj_shape+[2] + assert list(obj.shape) == [2]+obj_shape obj = initializers.RPI_spectral_init(pattern, probe, obj_shape, n_modes=2, mask=mask, background=background) - assert list(obj.shape) == [2]+obj_shape+[2] + assert list(obj.shape) == [2]+obj_shape diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index 79e2bfe..b0c10f8 100644 --- a/tests/tools/test_interactions.py +++ b/tests/tools/test_interactions.py @@ -1,6 +1,5 @@ from __future__ import division, print_function, absolute_import -from CDTools.tools import cmath from CDTools.tools import interactions import numpy as np import torch as t @@ -137,16 +136,16 @@ def test_ptycho_2D_round(random_probe, random_obj): random_obj[tr[0]:tr[0]+random_probe.shape[0], tr[1]:tr[1]+random_probe.shape[1]] for tr in np.round(translations).astype(int)] - exit_waves_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe), - cmath.complex_to_torch(random_obj), - t.tensor(translations)) - assert np.allclose(cmath.torch_to_complex(exit_waves_t), exit_waves_np) + exit_waves_t = interactions.ptycho_2D_round(t.as_tensor(random_probe), + t.as_tensor(random_obj), + t.as_tensor(translations)) + assert np.allclose(exit_waves_t.numpy(), exit_waves_np) # Test the single wave case - exit_wave_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe), - cmath.complex_to_torch(random_obj), - t.tensor(translations[0])) - assert np.allclose(cmath.torch_to_complex(exit_wave_t), exit_waves_np[0]) + exit_wave_t = interactions.ptycho_2D_round(t.as_tensor(random_probe), + t.as_tensor(random_obj), + t.as_tensor(translations[0])) + assert np.allclose(exit_wave_t.numpy(), exit_waves_np[0]) @@ -158,14 +157,14 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj): translation = np.array([46.7,53.2]) exit_waves_probe = interactions.ptycho_2D_linear( - cmath.complex_to_torch(single_pixel_probe), - cmath.complex_to_torch(random_obj), - t.tensor(translations), + t.as_tensor(single_pixel_probe), + t.as_tensor(random_obj), + t.as_tensor(translations), shift_probe=True) exit_wave_probe = interactions.ptycho_2D_linear( - cmath.complex_to_torch(single_pixel_probe), - cmath.complex_to_torch(random_obj), + t.as_tensor(single_pixel_probe), + t.as_tensor(random_obj), t.tensor(translation), shift_probe=True) @@ -174,14 +173,14 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj): exit_waves_obj = interactions.ptycho_2D_linear( - cmath.complex_to_torch(single_pixel_probe), - cmath.complex_to_torch(random_obj), + t.as_tensor(single_pixel_probe), + t.as_tensor(random_obj), t.tensor(translations), shift_probe=False) exit_wave_obj = interactions.ptycho_2D_linear( - cmath.complex_to_torch(single_pixel_probe), - cmath.complex_to_torch(random_obj), + t.as_tensor(single_pixel_probe), + t.as_tensor(random_obj), t.tensor(translation), shift_probe=False) @@ -189,7 +188,7 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj): assert t.allclose(exit_waves_obj[0],exit_wave_obj) # For the shifted probe, we should find 4 pixels with intensity - exit_waves_probe = cmath.torch_to_complex(exit_waves_probe)[0] + exit_waves_probe = t.as_tensor(exit_waves_probe)[0] probe_shift = np.array([[0.3*0.8,0.3*0.2], [0.7*0.8,0.7*0.2]]) @@ -199,7 +198,7 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj): assert np.allclose(probe_shift * obj_section, exit_section) # For the shifted obj, we should find one pixel with intensity - exit_waves_obj = cmath.torch_to_complex(exit_waves_obj)[0] + exit_waves_obj = t.as_tensor(exit_waves_obj)[0] obj_shift = np.array([[0.3*0.8,0.3*0.2], [0.7*0.8,0.7*0.2]]) obj_section = random_obj[128+46:128+48, @@ -219,15 +218,15 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj): translation = np.array([46.7,53.2]) exit_waves_probe = interactions.ptycho_2D_sinc( - cmath.complex_to_torch(single_pixel_probe), - cmath.complex_to_torch(random_obj), - t.tensor(translations), + t.as_tensor(single_pixel_probe), + t.as_tensor(random_obj), + t.as_tensor(translations), shift_probe=True) exit_wave_probe = interactions.ptycho_2D_sinc( - cmath.complex_to_torch(single_pixel_probe), - cmath.complex_to_torch(random_obj), - t.tensor(translation), + t.as_tensor(single_pixel_probe), + t.as_tensor(random_obj), + t.as_tensor(translation), shift_probe=True) # Check that the outputs match @@ -247,7 +246,7 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj): 53:53+256] exit_wave_np = sinc_shifted_probe * obj_section - exit_wave_torch = cmath.torch_to_complex(exit_wave_probe) + exit_wave_torch = exit_wave_probe.numpy() # The fidelity isn't great due to the FFT-based approach, so we need # a pretty relaxed condition @@ -259,8 +258,8 @@ def test_RPI_interaction(random_probe, random_obj): random_obj1 = random_obj[:79,:68] random_probe1 = random_probe - t_random_obj1 = cmath.complex_to_torch(random_obj1) - t_random_probe1 = cmath.complex_to_torch(random_probe1) + t_random_obj1 = t.as_tensor(random_obj1) + t_random_probe1 = t.as_tensor(random_probe1) t_output1 = interactions.RPI_interaction(t_random_probe1, t_random_obj1) obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho')) @@ -272,13 +271,13 @@ def test_RPI_interaction(random_probe, random_obj): output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups), norm='ortho')) - assert np.allclose(cmath.torch_to_complex(t_output1), output1) + assert np.allclose(t.as_tensor(t_output1), output1) random_obj2 = np.stack([random_obj[:64,:89]]*3) random_probe2 = random_probe[3:,5:] - t_random_obj2 = cmath.complex_to_torch(random_obj2) - t_random_probe2 = cmath.complex_to_torch(random_probe2) + t_random_obj2 = t.as_tensor(random_obj2) + t_random_probe2 = t.as_tensor(random_probe2) t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2) obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho')) diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index 5c372cf..0c7df38 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -1,7 +1,6 @@ from __future__ import division, print_function, absolute_import from CDTools.tools import losses -from CDTools.tools import cmath import numpy as np import torch as t @@ -16,7 +15,7 @@ def test_amplitude_mse(): # And add some noise to it sim = data + 0.1 * np.random.rand(10,100,100) # and define a simple mask that needs to be broadcast - mask = (np.random.rand(100,100) > 0.1).astype(np.bool) + mask = (np.random.rand(100,100) > 0.1).astype(bool) # First, test without a mask np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2) @@ -38,7 +37,7 @@ def test_intensity_mse(): # And add some noise to it sim = data + 0.1 * np.random.rand(10,100,100) # and define a simple mask that needs to be broadcast - mask = (np.random.rand(100,100) > 0.1).astype(np.bool) + mask = (np.random.rand(100,100) > 0.1).astype(bool) # First, test without a mask @@ -61,7 +60,7 @@ def test_poisson_nll(): # And add some noise to it sim = data + 0.1 * np.random.rand(10,100,100) # and define a simple mask that needs to be broadcast - mask = (np.random.rand(100,100) > 0.1).astype(np.bool) + mask = (np.random.rand(100,100) > 0.1).astype(bool) # First, test without a mask diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py index 36375c4..7407b97 100644 --- a/tests/tools/test_measurements.py +++ b/tests/tools/test_measurements.py @@ -1,31 +1,30 @@ from __future__ import division, print_function, absolute_import from CDTools.tools import measurements -from CDTools.tools import cmath import torch as t import numpy as np import pytest def test_intensity(): - wavefields = t.rand((5,10,10,2)) + wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10)) epsilon=1e-6 - np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + epsilon + np_result = np.abs(t.as_tensor(wavefields))**2 + epsilon assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon), - t.tensor(np_result)) + t.as_tensor(np_result)) # Test single field case assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon), - t.tensor(np_result[0])) + t.as_tensor(np_result[0])) det_slice = np.s_[3:,5:8] assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon), - t.tensor(np_result[(np.s_[:],)+det_slice])) + t.as_tensor(np_result[(np.s_[:],)+det_slice])) # Test single field case assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon), - t.tensor(np_result[0][det_slice])) + t.as_tensor(np_result[0][det_slice])) # With oversampling on @@ -36,34 +35,34 @@ def test_intensity(): # With multiple fields assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon, oversampling=2), - t.tensor(np_oversampling_result,)) + t.as_tensor(np_oversampling_result,)) # With a single field assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon, oversampling=2), - t.tensor(np_oversampling_result[0],)) + t.as_tensor(np_oversampling_result[0],)) def test_incoherent_sum(): # With no explicit slice given - wavefields = t.rand((5,4,10,10,2)) + wavefields = t.rand((5,4,10,10)) + 1j * t.rand((5,4,10,10)) epsilon=1e-6 - np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) + epsilon + np_result = np.sum(np.abs(wavefields.numpy())**2,axis=-3) + epsilon assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon), - t.tensor(np_result)) + t.as_tensor(np_result)) # Test single field case - assert t.allclose(measurements.incoherent_sum(wavefields[:,0],epsilon=epsilon), - t.tensor(np_result[0])) + assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon), + t.as_tensor(np_result[0])) # With a slice given det_slice = np.s_[3:,5:8] assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon), - t.tensor(np_result[(np.s_[:],)+det_slice])) + t.as_tensor(np_result[(np.s_[:],)+det_slice])) # Test single field case - assert t.allclose(measurements.incoherent_sum(wavefields[:,0],det_slice,epsilon=epsilon), - t.tensor(np_result[0][det_slice])) + assert t.allclose(measurements.incoherent_sum(wavefields[0,:],det_slice,epsilon=epsilon), + t.as_tensor(np_result[0][det_slice])) # With oversampling on np_oversampling_result = (np_result[:,::2,::2] + \ @@ -73,19 +72,19 @@ def test_incoherent_sum(): # With multiple fields assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon, oversampling=2), - t.tensor(np_oversampling_result,)) + t.as_tensor(np_oversampling_result,)) # With a single field - assert t.allclose(measurements.incoherent_sum(wavefields[:,0],epsilon=epsilon, oversampling=2), - t.tensor(np_oversampling_result[0],)) + assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon, oversampling=2), + t.as_tensor(np_oversampling_result[0],)) def test_quadratic_background(): # test with intensity - wavefields = t.rand((5,10,10,2)) + wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10)) epsilon=1e-6 background = t.rand((10,10)) - np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + background.numpy()**2 + epsilon + np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon det_slice = np.s_[3:,5:8] result = measurements.quadratic_background(wavefields,background[det_slice], @@ -96,8 +95,8 @@ def test_quadratic_background(): # test with incoherent sum but no slice and no stack - wavefields = t.rand((4,10,10,2)) - np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) + wavefields = t.rand((4,10,10)) + 1j * t.rand((4,10,10)) + np_result = np.sum(np.abs(wavefields.numpy())**2,axis=0) np_result += background.numpy()**2 result = measurements.quadratic_background(wavefields, background, epsilon=epsilon, diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py index a6dbe47..009457a 100644 --- a/tests/tools/test_plotting.py +++ b/tests/tools/test_plotting.py @@ -1,6 +1,5 @@ from __future__ import division, print_function, absolute_import -from CDTools.tools import cmath from CDTools.tools import plotting from CDTools.tools import initializers import numpy as np @@ -11,7 +10,7 @@ import matplotlib.pyplot as plt def test_plot_amplitude(show_plot): # Test with tensor - im = cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64)) + im = t.as_tensor(scipy.misc.ascent(),dtype=t.complex128) plotting.plot_amplitude(im, basis = np.array([[1,1], [1,1], [0,0]]), title = 'Test Amplitude') if show_plot: plt.show() @@ -31,7 +30,7 @@ def test_plot_phase(show_plot): plt.show() # Test with numpy array - im = cmath.torch_to_complex(initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])) + im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy() plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[1,1], [1,1], [0,0]])) if show_plot: plt.show() @@ -39,14 +38,13 @@ def test_plot_phase(show_plot): def test_plot_colorize(show_plot): # Test with tensor gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]) - im = cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64))) + im = gaussian * t.as_tensor(scipy.misc.ascent(), dtype=t.complex64) plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[1,1], [1,1], [0,0]])) if show_plot: plt.show() # Test with numpy array - gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]) - im = cmath.torch_to_complex(cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64)))) + im = im.numpy() plotting.plot_colorized(im, title = 'Test Colorize') if show_plot: plt.show() diff --git a/tests/tools/test_projectors.py b/tests/tools/test_projectors.py deleted file mode 100644 index 060b70a..0000000 --- a/tests/tools/test_projectors.py +++ /dev/null @@ -1,39 +0,0 @@ -from __future__ import division, print_function, absolute_import - -from CDTools.tools import cmath -from CDTools.tools import projectors -import numpy as np -import torch as t - - -def test_modulus(): - # Create a complex array with random modulus and known phase - np_result = np.sqrt(6) * (1 + 1j) * np.random.rand(10,10) - projection_intensity = t.from_numpy(np.abs(np_result)**2).to(t.float32) - original_wavefront = cmath.complex_to_torch((1+1j) * np.random.rand(10,10)).to(t.float32) - # Test without masks - torch_result = projectors.modulus(original_wavefront,projection_intensity) - assert np.allclose(cmath.torch_to_complex(torch_result),np_result) - - # Test with mask - mask = t.ones((10,10,2), dtype = t.uint8) - mask[5]*=0 - np_result[5] = cmath.torch_to_complex(original_wavefront[5]) - torch_result = projectors.modulus(original_wavefront,projection_intensity, mask=mask) - print(np_result[5]) - print(cmath.torch_to_complex(torch_result)[5]) - assert np.allclose(cmath.torch_to_complex(torch_result),np_result) - - - -def test_support(): - # Define a mask as uint8 to make sure it works when support is not - # the same type as the wavefield - support = t.zeros((10,10)).to(t.uint8) - # some masked, some unmasked - support[:3,:3] = 1 - - np_result = np.zeros((10,10)).astype(np.complex128) - np_result[:3,:3] = 1 + 1j - - assert(np.allclose(cmath.torch_to_complex(projectors.support(t.ones((10,10,2)), support)), np_result)) diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index e4ff41c..abe5e9c 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -1,8 +1,8 @@ from __future__ import division, print_function, absolute_import -from CDTools.tools import cmath from CDTools.tools import initializers from CDTools.tools import propagators +from CDTools.tools import image_processing import numpy as np import torch as t @@ -19,26 +19,26 @@ def exit_waves_1(): obj = scipy.misc.ascent()[0:64,0:64].astype(np.complex128) arr = np.random.random_sample((64,64)) obj *= (arr+(1-arr**2)**.5*1j) - obj = cmath.complex_to_torch(obj) + obj = t.as_tensor(obj) # Construct wavefront from image probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3) - return cmath.cmult(probe,obj) + return probe * obj def test_far_field(exit_waves_1): # Far field diffraction patterns calculated by numpy with zero frequency in center - np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho')) + np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho')) - assert(np.allclose(np_result, cmath.torch_to_complex(propagators.far_field(exit_waves_1)))) + assert(np.allclose(np_result, propagators.far_field(exit_waves_1).numpy())) def test_inverse_far_field(exit_waves_1): # We want the inverse far field to map back to the exit waves with no intensity corrections # Far field result for exit waves calculated with numpy - far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho'))) + far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho'))) assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result))) @@ -67,14 +67,14 @@ def test_generate_high_NA_k_intensity_map(): j = (np.arange(573) - 270) Is,Js = np.meshgrid(i,j,indexing='ij') wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128) - t_wavefield = cmath.complex_to_torch(wavefield).to(dtype=t.float32) + t_wavefield = t.as_tensor(wavefield, dtype=t.complex64) high_NA_propagated = propagators.high_NA_far_field( t_wavefield, k_map, intensity_map=intensity_map) low_NA_propagated = propagators.far_field(t_wavefield) - low_NA = cmath.torch_to_complex(low_NA_propagated) - high_NA = cmath.torch_to_complex(high_NA_propagated) + low_NA = low_NA_propagated.numpy() + high_NA = high_NA_propagated.numpy() # Checking first that for a low-NA propagation they give the same result # 1e-4 tolerance seems to be reasonable in this comparison given my @@ -100,8 +100,8 @@ def test_generate_high_NA_k_intensity_map(): t_wavefield, k_map, intensity_map=intensity_map) low_NA_propagated = propagators.far_field(t_wavefield) - low_NA = cmath.torch_to_complex(low_NA_propagated) - high_NA = cmath.torch_to_complex(high_NA_propagated) + low_NA = low_NA_propagated.numpy() + high_NA = high_NA_propagated.numpy() #plt.close('all') #plt.imshow(np.abs(low_NA)) @@ -121,11 +121,54 @@ def test_generate_high_NA_k_intensity_map(): # I believe this works, but I still would like to get a second method for # simulating at least one diffraction pattern as an independent check - assert 0 - - + #assert 0 +def test_near_field_direction(): + # + # The idea is to make sure that near field propagation is consistent + # with far-field propagation. Because we consider the Fourier + # transform to be the far-field propagator, we should find that + # near-field propagation becomes far-field propagation in the limit + # that z goes to positive infinity. This consistency requirement + # forces a specific choice of sign for near-field propagation - choosing + # it incorrectly will make the correct limit obtain at negative + # infinity. + # + # We can check consistency by setting up a gaussian wavefield with an + # offset in the far field. Then, we can do some near-field propagation, + # and we will ensure that the centroid of the beam moves in the correct + # direction. If the sign is incorrect, it will move in the wrong direction. + + x = (np.arange(901) - 400) + y = (np.arange(1200) - 500) + Ys,Xs = np.meshgrid(y,x) + Rs = np.sqrt(Xs**2+Ys**2) + E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)),dtype=t.complex64) + E0_real = propagators.inverse_far_field(E0_fourier) + # This is in the top-left corner in Fourier space + + wavelength = 3e-9 #nm + z = 1000e-9 + asp = propagators.generate_angular_spectrum_propagator( + E0_real.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex64) + + Ez_real = propagators.near_field(E0_real,asp) + + centroid = image_processing.centroid(t.abs(Ez_real)) + + # Assert it's in top half + assert centroid[0] < Ez_real.shape[0] // 2 + # Assert it's in left half + assert centroid[1] < Ez_real.shape[1] // 2 + + #plt.imshow(t.abs(E0_fourier)) + #plt.figure() + #plt.imshow(t.abs(E0_real)) + #plt.figure() + #plt.imshow(t.abs(Ez_real)) + #plt.show() + def test_near_field(): @@ -151,17 +194,37 @@ def test_near_field(): # The analytical expression for propagation of a gaussian beam in the # paraxial approx - Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * ( z + Rs**2 / (2 * Rz)) + 1j * np.arctan(z / zr)) - Ez_nozphase = Ez * np.exp(1j * k * z) + # The sign convention I'm using here is opposite from what is on + # wikipedia (as of June 2021), but it is I believe the most sensible + # choice. By choosing e^(ikx) to be the plane wave that's propagating + # in the x direction, we preserve two nice properties: + # 1) The time-dependence is e^(-iwt+ikx), which meshes with the choice + # we like to make as physicists of e^(-iwt) for time-dependence + # 2) This makes it so that the Fourier transform does far-field + # propagation in direction of light propagation. The logic is as + # follows: + # a) We must assume that light only propagates one direction through + # the plane, in order to have well-defined propagation + # b) Fixing that direction to be the positive direciton of propagation, + # we can consider the in-plane image of a plane wave propagating + # through the plane of interest. + # c) We can ask whether the 2D Fourier transform of that in-plane image + # produces a spot on the same side of zero as the in-plane component + # of k. + # If we choose e^(-ikx) to represent light propagating along K, the + # answer is no, and we find we have to use the inverse FT instead. + # Thus, e^(ikx) is the right choice here. + + Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * ( z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr)) + Ez_nozphase = Ez * np.exp(-1j * k * z) # First we check it normally asp = propagators.generate_angular_spectrum_propagator( - E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.float64) + E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128) - Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp) - Ez_t = cmath.torch_to_complex(Ez_t) + Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy() # Check for at least 10^-3 relative accuracy in this scenario assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)) @@ -169,8 +232,7 @@ def test_near_field(): Emz = np.conj(Ez) - Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp) - Emz_t = cmath.torch_to_complex(Emz_t) + Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy() # Again, 10^-3 is about all the accuracy we can expect assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz)) @@ -178,11 +240,10 @@ def test_near_field(): # Then, we check it with the phase correction asp = propagators.generate_angular_spectrum_propagator( E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True, - dtype=t.float64) + dtype=t.complex128) - Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp) - Ez_t = cmath.torch_to_complex(Ez_t) + Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy() # Check for at least 10^-3 relative accuracy in this scenario assert np.max(np.abs(Ez_nozphase-Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase)) @@ -190,12 +251,20 @@ def test_near_field(): Emz = np.conj(Ez_nozphase) - Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp) - Emz_t = cmath.torch_to_complex(Emz_t) + Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy() # Again, 10^-3 is about all the accuracy we can expect assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz)) + # Finally, we check that the bandlimiting at least does something + asp = propagators.generate_angular_spectrum_propagator( + E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True, + dtype=t.complex128, bandlimit=0.3) + + assert asp[140,0] == 0 + assert asp[0,180] == 0 + assert asp[130,0] != 0 + assert asp[0,175] != 0 def test_generalized_near_field(): @@ -236,14 +305,18 @@ def test_generalized_near_field(): return Zs / (Zs**2 + zr**2) def get_E(Xs, Ys, Zs, correct=False): - # if correct is True, remove the e^(-ikz) dependence + # Again, this follows the convention opposite from Wikipedia. See + # the note in test_angular_spectrum_propagator. + + # if correct is True, remove the e^(ikz) dependence + Rs_sq = Xs**2 + Ys**2 Wzs = get_w(Zs) E = w0 / Wzs * np.exp(-Rs_sq / Wzs**2) *\ - np.exp(-1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \ - 1j * np.arctan(Zs / zr)) + np.exp(1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \ + - 1j * np.arctan(Zs / zr)) if correct: - E = E * np.exp(1j * k * Zs) + E = E * np.exp(-1j * k * Zs) return E @@ -338,22 +411,22 @@ def test_generalized_near_field(): E0 = get_E(Xs,Ys,Zs_0, correct=False) Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=False) asp = propagators.generate_generalized_angular_spectrum_propagator( - E0.shape,new_basis,wavelength,offset_vec,dtype=t.float64) + E0.shape,new_basis,wavelength,offset_vec,dtype=t.complex128) elif str(propagation_vec) == 'offset': E0 = get_E(Xs,Ys,Zs_0, correct=True) Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True) asp = propagators.generate_generalized_angular_spectrum_propagator( E0.shape,new_basis,wavelength,offset_vec, - dtype=t.float64, propagate_along_offset=True) + dtype=t.complex128, propagate_along_offset=True) else: E0 = get_E(Xs,Ys,Zs_0, correct=True) Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True) asp = propagators.generate_generalized_angular_spectrum_propagator( E0.shape,new_basis,wavelength,offset_vec, - dtype=t.float64, propagation_vector=propagation_vec) + dtype=t.complex128, propagation_vector=propagation_vec) - Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp) - Ez_t = cmath.torch_to_complex(Ez_t) + Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy() + # Check for at least 10^-3 relative accuracy in this scenario if not np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)): #if True: @@ -362,14 +435,17 @@ def test_generalized_near_field(): plt.figure() plt.imshow(np.angle(Ez)) plt.figure() + plt.imshow(np.abs(Ez_t)) + plt.figure() + plt.imshow(np.angle(Ez_t)) + plt.figure() plt.imshow(np.abs(Ez-Ez_t)/np.max(np.abs(Ez))) plt.show() assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)) - Em0_t = propagators.inverse_near_field(cmath.complex_to_torch(Ez),asp) - Em0_t = cmath.torch_to_complex(Em0_t) + Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy() # Again, 10^-3 is about all the accuracy we can expect assert np.max(np.abs(E0-Em0_t)) < 1e-3 * np.max(np.abs(E0)) @@ -388,9 +464,8 @@ def test_generalized_near_field(): E0 = get_E(Xs,Ys,Zs_0, correct=True) asp = propagators.generate_generalized_angular_spectrum_propagator( E0.shape, new_basis, wavelength,offset_vec, - dtype=t.float64, propagation_vector=propagation_vec) - Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp) - Ez_t = cmath.torch_to_complex(Ez_t) + dtype=t.complex128, propagation_vector=propagation_vec) + Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy() for Rrand in Rrands: @@ -400,9 +475,9 @@ def test_generalized_near_field(): rot_prop = np.dot(Rrand, propagation_vec) asp = propagators.generate_generalized_angular_spectrum_propagator( E0.shape, rot_basis, wavelength, rot_offset, - dtype=t.float64, propagation_vector=rot_prop) - Ez_rot_t = propagators.near_field(cmath.complex_to_torch(E0),asp) - Ez_rot_t = cmath.torch_to_complex(Ez_rot_t) + dtype=t.complex128, propagation_vector=rot_prop) + Ez_rot_t = propagators.near_field(t.as_tensor(E0),asp).numpy() + assert np.max(np.abs(Ez_t-Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t)) @@ -421,9 +496,10 @@ def test_inverse_near_field(): E0 = np.exp(-Rs**2 / w0**2) asp = propagators.generate_angular_spectrum_propagator( - E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.float64) + E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128) - E0 = cmath.complex_to_torch(E0) + + E0 = t.as_tensor(E0,dtype=t.complex128) E_prop = propagators.near_field(E0,asp) E_backprop = propagators.inverse_near_field(E_prop, asp)