diff --git a/CDTools/__init__.py b/CDTools/__init__.py index d0f9e36..ae7ca77 100644 --- a/CDTools/__init__.py +++ b/CDTools/__init__.py @@ -3,3 +3,4 @@ from __future__ import division, print_function, absolute_import from CDTools import tools from CDTools import datasets from CDTools import models + diff --git a/CDTools/datasets/__init__.py b/CDTools/datasets/__init__.py index 504647a..16655e7 100644 --- a/CDTools/datasets/__init__.py +++ b/CDTools/datasets/__init__.py @@ -1,7 +1,7 @@ """ This module contains all the datasets for interacting with ptychography data All the access to data from standard ptychography and CDI experiments is -coordinated through the various datasets defined in this file. They make use +coordinated through the various datasets defined in this module. They make use of the lower-level data reading and writing functions defined in tools.data, but critically all of these datasets subclass torch.Dataset. This allows them to be used as standard torch datasets during reconstructions, which @@ -9,8 +9,11 @@ helps make it easy to use the various data-handling strategies that are implemented by default in pytorch (such as drawing data in a random order, drawing minibatches, etc.) -Subclasses of CDataset are required to define their own implementations -of the following functions: +New Datasets can be defined a subclass of the main CDataset class defined +in the base.py file. Example implementations of all these functions +can be found in the code for the Ptycho2DDataset class. In addition, it is +recommended to read through the tutorial section on defining a new CDI +dataset before attempting to do so * __init__ * __len__ @@ -20,267 +23,11 @@ of the following functions: * to_cxi * inspect -Example implementations of all these functions can be found in the code -for the Ptycho2DDataset class. - -In addition, it is recommended to read through the tutorial section on -defining a new ptychography dataset before attempting to do so """ from __future__ import division, print_function, absolute_import -import numpy as np -import torch as t -from copy import copy -import h5py -try: - import pathlib -except ImportError: - import pathlib2 as pathlib - -from CDTools.tools import data as cdtdata -from CDTools.tools import plotting -from torch.utils import data as torchdata -from matplotlib import pyplot as plt -from matplotlib.widgets import Slider -from matplotlib import ticker - __all__ = ['CDataset','Ptycho2DDataset'] -# -# This loads and stores all the kinds of metadata that are common to -# All different kinds of diffraction experiments -# Other datasets can subclass this and not worry about loading and -# saving that metadata. -# - -class CDataset(torchdata.Dataset): - """ The base dataset class which all other datasets subclass - - Subclasses torch.utils.data.Dataset - - This base dataset class defines the functionality which should be - common to all subclassed datasets. This includes the loading and - storage of the metadata portions of .cxi files, as well as the tools - needed to allow for easy mixing of data on the CPU and GPU. - """ - - def __init__(self, entry_info=None, sample_info=None, - wavelength=None, - detector_geometry=None, mask=None, - background=None): - - """The __init__ function allows construction from python objects. - - The detector_geometry dictionary is defined to have the - entries defined by the outputs of data.get_detector_geometry. - - - Parameters - ---------- - entry_info : dict - A dictionary containing the entry_info metadata - sample_info : dict - A dictionary containing the sample_info metadata - wavelength : float - The wavelength of light used in the experiment - detector_geometry : dict - A dictionary containing the various detector geometry - parameters - mask : array - A mask for the detector, defined as 1 for live pixels, 0 - for dead - background : array - An initial guess for the not-previously-subtracted - detector background - """ - - # Force pass-by-value-like behavior to stop strangeness - self.entry_info = copy(entry_info) - self.sample_info = copy(sample_info) - self.wavelength = wavelength - self.detector_geometry = copy(detector_geometry) - if mask is not None: - if isinstance(mask, t.Tensor): - self.mask = mask.detach().to(dtype=t.bool) - else: - self.mask = t.BoolTensor(mask) - else: - self.mask = None - if background is not None: - self.background = t.Tensor(background) - else: - self.background = None - - self.get_as(device='cpu') - - - def to(self,*args,**kwargs): - """Sends the relevant data to the given device and dtype - - This function sends the stored mask and background to the - specified device and dtype - - Accepts the same parameters as torch.Tensor.to - """ - # The mask should always stay a uint8, but it should switch devices - mask_kwargs = copy(kwargs) - try: - mask_kwargs.pop('dtype') - except KeyError as r: - pass - - if self.mask is not None: - 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 - - Oftentimes there isn't room to store an entire dataset on a GPU, - but it is still worth running the calculation on the GPU even with - the overhead incurred by transferring data back and forth. In that - case, get_as can be used instead of to, to declare a set of - device and dtype that the data should be returned as, whenever it - is accessed through the __getitem__ function (as it would be in - any reconstructions). - - Parameters - ---------- - Accepts the same parameters as torch.Tensor.to - """ - self.get_as_args = (args, kwargs) - - def __len__(self): - raise NotImplementedError() - - def __getitem__(self, index): - # Deals with loading to appropriate device/dtype, if - # specified via a call to get_as - inputs, outputs = self._load(index) - if hasattr(self, 'get_as_args'): - outputs = outputs.to(*self.get_as_args[0],**self.get_as_args[1]) - moved_inputs = [] - for inp in inputs: - try: - moved_inputs.append(inp.to(*self.get_as_args[0],**self.get_as_args[1]) ) - except: - moved_inputs.append(inp) - else: - moved_inputs = inputs - return moved_inputs, outputs - - - def _load(self, index): - """ Internal function to load data - - In all subclasses of CDataset, a _load function should be defined. - This function is used internally by the global __getitem__ function - defined in the base class, which handles moving data around when - the dataset is (for example) storing the data on the CPU but - getting data as GPU tensors. - - It should accept an index or slice, and return output as a tuple. - The first item of the tuple is a tuple containing the inputs to - the forward model for the related ptychography model. The second - item of the tuple should be the set of diffraction patterns - associated with the returned inputs. - - Since there is no kind of data stored in a CDataset, this - function is defined as returing a NotImplemented Error - """ - raise NotImplementedError() - - - @classmethod - def from_cxi(cls, cxi_file): - """Generates a new CDataset from a .cxi file directly - - This is the most commonly used constructor for CDatasets and - subclasses thereof. It populates the dataset using the information - in a .cxi file. It can either take an h5py.File object directly, - or a filename or pathlib object pointing to the file - - Parameters - ---------- - file : str, pathlib.Path, or h5py.File - The .cxi file to load from - - Returns - ------- - dataset : CDataset - The constructed dataset object - """ - - # If a bare string is passed - if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): - with h5py.File(cxi_file,'r') as f: - return cls.from_cxi(f) - - entry_info = cdtdata.get_entry_info(cxi_file) - sample_info = cdtdata.get_sample_info(cxi_file) - wavelength = cdtdata.get_wavelength(cxi_file) - distance, basis, corner = cdtdata.get_detector_geometry(cxi_file) - detector_geometry = {'distance' : distance, - 'basis' : basis, - 'corner' : corner} - mask = cdtdata.get_mask(cxi_file) - dark = cdtdata.get_dark(cxi_file) - return cls(entry_info = entry_info, - sample_info = sample_info, - wavelength=wavelength, - detector_geometry=detector_geometry, - mask=mask, background=dark) - - - def to_cxi(self, cxi_file): - """Saves out a CDataset as a .cxi file - - This function saves all the compatible information in a CDataset - object into a .cxi file. This is useful for saving out modified - or simulated datasets - - Parameters - ---------- - cxi_file : str, pathlib.Path, or h5py.File - The .cxi file to write to - """ - - # If a bare string is passed - if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): - with h5py.File(cxi_file,'w') as f: - return self.to_cxi(f) - - if self.entry_info is not None: - cdtdata.add_entry_info(cxi_file, self.entry_info) - if self.sample_info is not None: - cdtdata.add_sample_info(cxi_file, self.sample_info) - if self.wavelength is not None: - cdtdata.add_source(cxi_file, self.wavelength) - if self.detector_geometry is not None: - if 'corner' in self.detector_geometry: - corner = self.detector_geometry['corner'] - else: - corner = None - cdtdata.add_detector(cxi_file, - self.detector_geometry['distance'], - self.detector_geometry['basis'], - corner = corner) - if self.mask is not None: - cdtdata.add_mask(cxi_file, self.mask) - if self.background is not None: - cdtdata.add_dark(cxi_file, self.background) - - def inspect(self): - """The prototype for the inspect function - - In all subclasses of CDataset, an inspect function should be - defined which opens a tool that shows the data in a natural - layout for that kind of experiment. In the base class, no actual - data is stored, so this is defined to raise a NotImplementedError - """ - raise NotImplementedError - +from CDTools.datasets.base import CDataset from CDTools.datasets.ptycho_2d_dataset import Ptycho2DDataset diff --git a/CDTools/datasets/base.py b/CDTools/datasets/base.py new file mode 100644 index 0000000..4b177eb --- /dev/null +++ b/CDTools/datasets/base.py @@ -0,0 +1,278 @@ +""" This module contains the base CDataset class for handling CDI data + +Subclasses of CDataset are required to define their own implementations +of the following functions: + +* __init__ +* __len__ +* _load +* to +* from_cxi +* to_cxi +* inspect + +Example implementations of all these functions can be found in the code +for the Ptycho2DDataset class. + +In addition, it is recommended to read through the tutorial section on +defining a new ptychography dataset before attempting to do so +""" + +from __future__ import division, print_function, absolute_import + +import numpy as np +import torch as t +from copy import copy +import h5py +try: + import pathlib +except ImportError: + import pathlib2 as pathlib + +from CDTools.tools import data as cdtdata +from CDTools.tools import plotting +from torch.utils import data as torchdata +from matplotlib import pyplot as plt +from matplotlib.widgets import Slider +from matplotlib import ticker + + +__all__ = ['CDataset'] + +# +# This loads and stores all the kinds of metadata that are common to +# All different kinds of diffraction experiments +# Other datasets can subclass this and not worry about loading and +# saving that metadata. +# + +class CDataset(torchdata.Dataset): + """ The base dataset class which all other datasets subclass + + Subclasses torch.utils.data.Dataset + + This base dataset class defines the functionality which should be + common to all subclassed datasets. This includes the loading and + storage of the metadata portions of .cxi files, as well as the tools + needed to allow for easy mixing of data on the CPU and GPU. + """ + + def __init__(self, entry_info=None, sample_info=None, + wavelength=None, + detector_geometry=None, mask=None, + background=None): + + """The __init__ function allows construction from python objects. + + The detector_geometry dictionary is defined to have the + entries defined by the outputs of data.get_detector_geometry. + + + Parameters + ---------- + entry_info : dict + A dictionary containing the entry_info metadata + sample_info : dict + A dictionary containing the sample_info metadata + wavelength : float + The wavelength of light used in the experiment + detector_geometry : dict + A dictionary containing the various detector geometry + parameters + mask : array + A mask for the detector, defined as 1 for live pixels, 0 + for dead + background : array + An initial guess for the not-previously-subtracted + detector background + """ + + # Force pass-by-value-like behavior to stop strangeness + self.entry_info = copy(entry_info) + self.sample_info = copy(sample_info) + self.wavelength = wavelength + self.detector_geometry = copy(detector_geometry) + if mask is not None: + if isinstance(mask, t.Tensor): + self.mask = mask.detach().to(dtype=t.bool) + else: + self.mask = t.BoolTensor(mask) + else: + self.mask = None + if background is not None: + self.background = t.Tensor(background) + else: + self.background = None + + self.get_as(device='cpu') + + + def to(self,*args,**kwargs): + """Sends the relevant data to the given device and dtype + + This function sends the stored mask and background to the + specified device and dtype + + Accepts the same parameters as torch.Tensor.to + """ + # The mask should always stay a uint8, but it should switch devices + mask_kwargs = copy(kwargs) + try: + mask_kwargs.pop('dtype') + except KeyError as r: + pass + + if self.mask is not None: + 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 + + Oftentimes there isn't room to store an entire dataset on a GPU, + but it is still worth running the calculation on the GPU even with + the overhead incurred by transferring data back and forth. In that + case, get_as can be used instead of to, to declare a set of + device and dtype that the data should be returned as, whenever it + is accessed through the __getitem__ function (as it would be in + any reconstructions). + + Parameters + ---------- + Accepts the same parameters as torch.Tensor.to + """ + self.get_as_args = (args, kwargs) + + def __len__(self): + raise NotImplementedError() + + def __getitem__(self, index): + # Deals with loading to appropriate device/dtype, if + # specified via a call to get_as + inputs, outputs = self._load(index) + if hasattr(self, 'get_as_args'): + outputs = outputs.to(*self.get_as_args[0],**self.get_as_args[1]) + moved_inputs = [] + for inp in inputs: + try: + moved_inputs.append(inp.to(*self.get_as_args[0],**self.get_as_args[1]) ) + except: + moved_inputs.append(inp) + else: + moved_inputs = inputs + return moved_inputs, outputs + + + def _load(self, index): + """ Internal function to load data + + In all subclasses of CDataset, a _load function should be defined. + This function is used internally by the global __getitem__ function + defined in the base class, which handles moving data around when + the dataset is (for example) storing the data on the CPU but + getting data as GPU tensors. + + It should accept an index or slice, and return output as a tuple. + The first item of the tuple is a tuple containing the inputs to + the forward model for the related ptychography model. The second + item of the tuple should be the set of diffraction patterns + associated with the returned inputs. + + Since there is no kind of data stored in a CDataset, this + function is defined as returing a NotImplemented Error + """ + raise NotImplementedError() + + + @classmethod + def from_cxi(cls, cxi_file): + """Generates a new CDataset from a .cxi file directly + + This is the most commonly used constructor for CDatasets and + subclasses thereof. It populates the dataset using the information + in a .cxi file. It can either take an h5py.File object directly, + or a filename or pathlib object pointing to the file + + Parameters + ---------- + file : str, pathlib.Path, or h5py.File + The .cxi file to load from + + Returns + ------- + dataset : CDataset + The constructed dataset object + """ + + # If a bare string is passed + if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): + with h5py.File(cxi_file,'r') as f: + return cls.from_cxi(f) + + entry_info = cdtdata.get_entry_info(cxi_file) + sample_info = cdtdata.get_sample_info(cxi_file) + wavelength = cdtdata.get_wavelength(cxi_file) + distance, basis, corner = cdtdata.get_detector_geometry(cxi_file) + detector_geometry = {'distance' : distance, + 'basis' : basis, + 'corner' : corner} + mask = cdtdata.get_mask(cxi_file) + dark = cdtdata.get_dark(cxi_file) + return cls(entry_info = entry_info, + sample_info = sample_info, + wavelength=wavelength, + detector_geometry=detector_geometry, + mask=mask, background=dark) + + + def to_cxi(self, cxi_file): + """Saves out a CDataset as a .cxi file + + This function saves all the compatible information in a CDataset + object into a .cxi file. This is useful for saving out modified + or simulated datasets + + Parameters + ---------- + cxi_file : str, pathlib.Path, or h5py.File + The .cxi file to write to + """ + + # If a bare string is passed + if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): + with h5py.File(cxi_file,'w') as f: + return self.to_cxi(f) + + if self.entry_info is not None: + cdtdata.add_entry_info(cxi_file, self.entry_info) + if self.sample_info is not None: + cdtdata.add_sample_info(cxi_file, self.sample_info) + if self.wavelength is not None: + cdtdata.add_source(cxi_file, self.wavelength) + if self.detector_geometry is not None: + if 'corner' in self.detector_geometry: + corner = self.detector_geometry['corner'] + else: + corner = None + cdtdata.add_detector(cxi_file, + self.detector_geometry['distance'], + self.detector_geometry['basis'], + corner = corner) + if self.mask is not None: + cdtdata.add_mask(cxi_file, self.mask) + if self.background is not None: + cdtdata.add_dark(cxi_file, self.background) + + def inspect(self): + """The prototype for the inspect function + + In all subclasses of CDataset, an inspect function should be + defined which opens a tool that shows the data in a natural + layout for that kind of experiment. In the base class, no actual + data is stored, so this is defined to raise a NotImplementedError + """ + raise NotImplementedError + + diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index fbe6807..58af8be 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -9,508 +9,24 @@ corresponds to the measured data (usually, a diffraction pattern). This model can then be used as the heart of an automatic differentiation reconstruction which retrieves the parameters that were used in the model. +A main CDIModel class is defined in the base.py file, and models for +various CDI geometries can be defined as subclasses of this base model. +The subclasses of the main CDIModel class are required to implement a set of +functions defined in the base.py file. Example implementations of +these functions can be found in the code for the SimplePtycho class. -The subclasses of the main CDIModel class are required to define their -own implementations of the following functions: +Finally, it is recommended to read through the tutorial section on +defining a new ptychography model before attempting to do so. -Loading and Saving ------------------- -from_dataset - Creates a CDIModel from an appropriate CDataset -simulate_to_dataset - Creates a CDataset from the simulation defined in the model -save_results - Saves out a dictionary with the recovered parameters - - -Simulation ----------- -interaction - Simulates exit waves from experimental parameters -forward_propagator - The propagator from the experiment plane to the detector plane -backward_propagator - Optional, the propagator from the detector plane to the experiment plane -measurement - Simulates the detector readout from a detector plane wavefront -loss - the loss function to report and use for automatic differentiation - -Example implementations of all these functions can be found in the code -for the SimplePtycho class. - -In addition, it is recommended to read through the tutorial section on -defining a new ptychography model before attempting to do so """ from __future__ import division, print_function, absolute_import -import torch as t -from torch.utils import data as torchdata -from matplotlib import pyplot as plt -from matplotlib.widgets import Slider -from matplotlib import ticker -import numpy as np - +# 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'] - -class CDIModel(t.nn.Module): - """This base model defines all the functions that must be exposed for a valid CDIModel subclass - - Most of the functions only raise a NotImplementedError at this level and - must be explicitly defined by any subclass - these are noted explocitly - in the module-level intro. The work of defining the various subclasses - boils down to creating an appropriate implementation for this set of - functions. - """ - - def from_dataset(self, dataset): - raise NotImplementedError() - - - def interaction(self, *args): - raise NotImplementedError() - - - def forward_propagator(self, exit_wave): - raise NotImplementedError() - - - def backward_propagator(self, detector_wave): - raise NotImplementedError() - - - def measurement(self, detector_wave): - raise NotImplementedError() - - - def forward(self, *args): - """The complete forward model - - This model relies on composing the interaction, forward propagator, - and measurement functions which are required to be defined by all - subclasses. It therefore should not be redefined by the subclasses. - - The arguments to this function, for any given subclass, will be - the same as the arguments to the interaction function. - """ - return self.measurement(self.forward_propagator(self.interaction(*args))) - - def loss(self, sim_data, real_data): - raise NotImplementedError() - - - def to(self, *args, **kwargs): - super(CDIModel,self).to(*args,**kwargs) - - - def simulate_to_dataset(self, args_list): - raise NotImplementedError() - - def save_results(self): - raise NotImplementedError() - - def AD_optimize(self, iterations, data_loader, optimizer,\ - scheduler=None, regularization_factor=None): - """Runs a round of reconstruction using the provided optimizer - - This is the basic automatic differentiation reconstruction tool - which all the other, algorithm-specific tools, use. - - Like all the other optimization routines, it is defined as a - generator function which yields the average loss each epoch. - - Parameters - ---------- - iterations : int - How many epochs of the algorithm to run - dataset : CDataset - The dataset to reconstruct against - optimizer : torch.optim.Optimizer - The optimizer to run the reconstruction with - scheduler : torch.optim.lr_scheduler._LRScheduler - Optional, a learning rate scheduler to use - regularization_factor : float or list(float) - Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method - """ - # First, calculate the normalization - normalization = 0 - for inputs, patterns in data_loader: - normalization += t.sum(patterns).cpu().numpy() - - for it in range(iterations): - loss = 0 - N = 0 - for inputs, patterns in data_loader: - N += 1 - def closure(): - 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) - - if regularization_factor is not None \ - and hasattr(self, 'regularizer'): - loss += self.regularizer(regularization_factor) - - loss.backward() - return loss - - loss += optimizer.step(closure).detach().cpu().numpy() - - loss /= normalization - if scheduler is not None: - scheduler.step(loss) - - yield loss - - - def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, - schedule=False, amsgrad=False, subset=None, - regularization_factor=None): - """Runs a round of reconstruction using the Adam optimizer - - This is generally accepted to be the most robust algorithm for use - with ptychography. Like all the other optimization routines, - it is defined as a generator function, which yields the average - loss each epoch. - - Parameters - ---------- - iterations : int - 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 (alpha) to use - schedule : float - Optional, whether to use the ReduceLROnPlateau scheduler - subset : list(int) or int - Optional, a pattern index or list of pattern indices to use - regularization_factor : float or list(float) - Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method - """ - - if subset is not None: - # if just one pattern, turn into a list for convenience - if type(subset) == type(1): - subset = [subset] - dataset = torchdata.Subset(dataset, subset) - - # Make a dataloader - data_loader = torchdata.DataLoader(dataset, batch_size=batch_size, - shuffle=True) - - # Define the optimizer - optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad) - - - # Define the scheduler - if schedule: - scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9) - else: - scheduler = None - - return self.AD_optimize(iterations, data_loader, optimizer, - scheduler=scheduler, - regularization_factor=regularization_factor) - - - def LBFGS_optimize(self, iterations, dataset, batch_size=None, - lr=0.1,history_size=2, subset=None, - regularization_factor=None): - """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. - - Parameters - ---------- - iterations : int - 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 - Optional, the length of the history to use. - subset : list(int) or int - Optional, a pattern index or list of pattern indices to ues - regularization_factor : float or list(float) - Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method - """ - if subset is not None: - # if just one pattern, turn into a list for convenience - if type(subset) == type(1): - 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) - - - # Define the optimizer - optimizer = t.optim.LBFGS(self.parameters(), - lr = lr, history_size=history_size) - - return self.AD_optimize(iterations, data_loader, optimizer, - regularization_factor=regularization_factor) - - - 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): - """Runs a round of reconstruction using the SGDoptimizer - - This algorithm is often less stable that Adam, but it is simpler - and is the basic workhorse of gradience descent. - - Parameters - ---------- - iterations : int - 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 - momentum : float - Optional, the length of the history to use. - subset : list(int) or int - Optional, a pattern index or list of pattern indices to use - regularization_factor : float or list(float) - Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method - """ - - if subset is not None: - # if just one pattern, turn into a list for convenience - if type(subset) == type(1): - 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) - - - # Define the optimizer - optimizer = t.optim.SGD(self.parameters(), - lr = lr, momentum=momentum, - dampening=dampening, - weight_decay=weight_decay, - nesterov=nesterov) - - return self.AD_optimize(iterations, data_loader, optimizer, - regularization_factor=regularization_factor) - - - # By default, the plot_list is empty - plot_list = [] - - - def inspect(self, dataset=None, update=True): - """Plots all the plots defined in the model's plot_list attribute - - If update is set to True, it will update any previously plotted set - of plots, if one exists, and then redraw them. Otherwise, it will - plot a new set, and any subsequent updates will update the new set - - Optionally, a dataset can be passed, which then will plot any - registered plots which need to incorporate some information from - the dataset (such as geometry or a comparison with measured data). - - Plots can be registered in any subclass by defining the plot_list - attribute. This should be a list of tuples in the following format: - ( 'Plot Title', function_to_generate_plot(self), - function_to_determine_whether_to_plot(self)) - - Where the third element in the tuple (a function that returns - True if the plot is relevant) is not required. - - Parameters - ---------- - dataset : CDataset - Optional, a dataset matched to the model type - update : bool - Whether to update existing plots or plot new ones - - """ - first_update = False - if update and hasattr(self, 'figs') and self.figs: - figs = self.figs - elif update: - figs = None - self.figs = [] - first_update = True - else: - figs = None - self.figs = [] - - idx = 0 - for plots in self.plot_list: - # If a conditional is included in the plot - try: - if len(plots) >=3 and not plots[2](self): - continue - except TypeError as e: - if len(plots) >= 3 and not plots[2](self, dataset): - continue - - name = plots[0] - plotter = plots[1] - - if figs is None: - fig = plt.figure() - self.figs.append(fig) - else: - fig = figs[idx] - - try: - plotter(self,fig) - plt.title(name) - except TypeError as e: - if dataset is not None: - try: - plotter(self, fig, dataset) - plt.title(name) - except (IndexError, KeyError, AttributeError) as e: - pass - - except (IndexError, KeyError, AttributeError) as e: - pass - - idx += 1 - - if update: - plt.draw() - fig.canvas.start_event_loop(0.001) - - if first_update: - plt.pause(0.05 * len(self.figs)) - - - def compare(self, dataset): - """Opens a tool for comparing simulated and measured diffraction patterns - - Parameters - ---------- - dataset : CDataset - A dataset containing the simulated diffraction patterns to compare against - """ - - 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 - # (say, from clicking back in the navbar) - # Holy fuck this was annoying. Sorry future for how - # crappy this solution is. - #if not np.allclose(im.colorbar.ax.get_xlim(), - # (np.min(im.get_array()), - # np.max(im.get_array()))): - if hasattr(im, 'norecurse') and im.norecurse: - im.norecurse=False - 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() - - - def update(idx): - idx = int(idx) % len(dataset) - fig.pattern_idx = idx - updating = True if len(axes[0].images) >= 1 else False - - inputs, output = dataset[idx] - sim_data = self.forward(*inputs).detach().cpu().numpy() - sim_data = sim_data - 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 - - if not updating: - axes[0].set_title('Simulated') - axes[1].set_title('Measured') - axes[2].set_title('Difference') - - sim = axes[0].imshow(sim_data) - meas = axes[1].imshow(meas_data * mask) - diff = axes[2].imshow((sim_data-meas_data) * mask) - - cb1 = plt.colorbar(sim, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) - cb1.ax.tick_params(labelrotation=20) - cb1.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(sim)) - cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) - cb2.ax.tick_params(labelrotation=20) - cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) - cb3 = plt.colorbar(diff, ax=axes[2], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) - cb3.ax.tick_params(labelrotation=20) - 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) - - meas = axes[1].images[-1] - meas.set_data(meas_data * mask) - update_colorbar(meas) - - diff = axes[2].images[-1] - diff.set_data((sim_data-meas_data) * mask) - update_colorbar(diff) - - - # This is dumb but the slider doesn't work unless a reference to it is - # kept somewhere... - self.slider = Slider(axslider, 'Pattern #', 0, len(dataset)-1, valstep=1, valfmt="%d") - self.slider.on_changed(update) - - 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.pattern_idx - 1) - elif event.key == 'down' or event.button == 'down': - update(fig.pattern_idx + 1) - self.slider.set_val(fig.pattern_idx) - plt.draw() - - fig.canvas.mpl_connect('key_press_event',on_action) - fig.canvas.mpl_connect('scroll_event',on_action) - update(0) - - - - +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 diff --git a/CDTools/models/base.py b/CDTools/models/base.py new file mode 100644 index 0000000..e5fb9c3 --- /dev/null +++ b/CDTools/models/base.py @@ -0,0 +1,496 @@ +"""This module contains the base CDIModel class for CDI Models. + +The subclasses of the main CDIModel class are required to define their +own implementations of the following functions: + +Loading and Saving +------------------ +from_dataset + Creates a CDIModel from an appropriate CDataset +simulate_to_dataset + Creates a CDataset from the simulation defined in the model +save_results + Saves out a dictionary with the recovered parameters + + +Simulation +---------- +interaction + Simulates exit waves from experimental parameters +forward_propagator + The propagator from the experiment plane to the detector plane +backward_propagator + Optional, the propagator from the detector plane to the experiment plane +measurement + Simulates the detector readout from a detector plane wavefront +loss + the loss function to report and use for automatic differentiation + +""" + +from __future__ import division, print_function, absolute_import + +import torch as t +from torch.utils import data as torchdata +from matplotlib import pyplot as plt +from matplotlib.widgets import Slider +from matplotlib import ticker +import numpy as np + +__all__ = ['CDIModel'] + + +class CDIModel(t.nn.Module): + """This base model defines all the functions that must be exposed for a valid CDIModel subclass + + Most of the functions only raise a NotImplementedError at this level and + must be explicitly defined by any subclass - these are noted explocitly + in the module-level intro. The work of defining the various subclasses + boils down to creating an appropriate implementation for this set of + functions. + """ + + def from_dataset(self, dataset): + raise NotImplementedError() + + + def interaction(self, *args): + raise NotImplementedError() + + + def forward_propagator(self, exit_wave): + raise NotImplementedError() + + + def backward_propagator(self, detector_wave): + raise NotImplementedError() + + + def measurement(self, detector_wave): + raise NotImplementedError() + + + def forward(self, *args): + """The complete forward model + + This model relies on composing the interaction, forward propagator, + and measurement functions which are required to be defined by all + subclasses. It therefore should not be redefined by the subclasses. + + The arguments to this function, for any given subclass, will be + the same as the arguments to the interaction function. + """ + return self.measurement(self.forward_propagator(self.interaction(*args))) + + def loss(self, sim_data, real_data): + raise NotImplementedError() + + + def to(self, *args, **kwargs): + super(CDIModel,self).to(*args,**kwargs) + + + def simulate_to_dataset(self, args_list): + raise NotImplementedError() + + def save_results(self): + raise NotImplementedError() + + def AD_optimize(self, iterations, data_loader, optimizer,\ + scheduler=None, regularization_factor=None): + """Runs a round of reconstruction using the provided optimizer + + This is the basic automatic differentiation reconstruction tool + which all the other, algorithm-specific tools, use. + + Like all the other optimization routines, it is defined as a + generator function which yields the average loss each epoch. + + Parameters + ---------- + iterations : int + How many epochs of the algorithm to run + dataset : CDataset + The dataset to reconstruct against + optimizer : torch.optim.Optimizer + The optimizer to run the reconstruction with + scheduler : torch.optim.lr_scheduler._LRScheduler + Optional, a learning rate scheduler to use + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + """ + # First, calculate the normalization + normalization = 0 + for inputs, patterns in data_loader: + normalization += t.sum(patterns).cpu().numpy() + + for it in range(iterations): + loss = 0 + N = 0 + for inputs, patterns in data_loader: + N += 1 + def closure(): + 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) + + if regularization_factor is not None \ + and hasattr(self, 'regularizer'): + loss += self.regularizer(regularization_factor) + + loss.backward() + return loss + + loss += optimizer.step(closure).detach().cpu().numpy() + + loss /= normalization + if scheduler is not None: + scheduler.step(loss) + + yield loss + + + def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, + schedule=False, amsgrad=False, subset=None, + regularization_factor=None): + """Runs a round of reconstruction using the Adam optimizer + + This is generally accepted to be the most robust algorithm for use + with ptychography. Like all the other optimization routines, + it is defined as a generator function, which yields the average + loss each epoch. + + Parameters + ---------- + iterations : int + 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 (alpha) to use + schedule : float + Optional, whether to use the ReduceLROnPlateau scheduler + subset : list(int) or int + Optional, a pattern index or list of pattern indices to use + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + """ + + if subset is not None: + # if just one pattern, turn into a list for convenience + if type(subset) == type(1): + subset = [subset] + dataset = torchdata.Subset(dataset, subset) + + # Make a dataloader + data_loader = torchdata.DataLoader(dataset, batch_size=batch_size, + shuffle=True) + + # Define the optimizer + optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad) + + + # Define the scheduler + if schedule: + scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9) + else: + scheduler = None + + return self.AD_optimize(iterations, data_loader, optimizer, + scheduler=scheduler, + regularization_factor=regularization_factor) + + + def LBFGS_optimize(self, iterations, dataset, batch_size=None, + lr=0.1,history_size=2, subset=None, + regularization_factor=None): + """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. + + Parameters + ---------- + iterations : int + 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 + Optional, the length of the history to use. + subset : list(int) or int + Optional, a pattern index or list of pattern indices to ues + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + """ + if subset is not None: + # if just one pattern, turn into a list for convenience + if type(subset) == type(1): + 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) + + + # Define the optimizer + optimizer = t.optim.LBFGS(self.parameters(), + lr = lr, history_size=history_size) + + return self.AD_optimize(iterations, data_loader, optimizer, + regularization_factor=regularization_factor) + + + 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): + """Runs a round of reconstruction using the SGDoptimizer + + This algorithm is often less stable that Adam, but it is simpler + and is the basic workhorse of gradience descent. + + Parameters + ---------- + iterations : int + 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 + momentum : float + Optional, the length of the history to use. + subset : list(int) or int + Optional, a pattern index or list of pattern indices to use + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + """ + + if subset is not None: + # if just one pattern, turn into a list for convenience + if type(subset) == type(1): + 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) + + + # Define the optimizer + optimizer = t.optim.SGD(self.parameters(), + lr = lr, momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov) + + return self.AD_optimize(iterations, data_loader, optimizer, + regularization_factor=regularization_factor) + + + # By default, the plot_list is empty + plot_list = [] + + + def inspect(self, dataset=None, update=True): + """Plots all the plots defined in the model's plot_list attribute + + If update is set to True, it will update any previously plotted set + of plots, if one exists, and then redraw them. Otherwise, it will + plot a new set, and any subsequent updates will update the new set + + Optionally, a dataset can be passed, which then will plot any + registered plots which need to incorporate some information from + the dataset (such as geometry or a comparison with measured data). + + Plots can be registered in any subclass by defining the plot_list + attribute. This should be a list of tuples in the following format: + ( 'Plot Title', function_to_generate_plot(self), + function_to_determine_whether_to_plot(self)) + + Where the third element in the tuple (a function that returns + True if the plot is relevant) is not required. + + Parameters + ---------- + dataset : CDataset + Optional, a dataset matched to the model type + update : bool + Whether to update existing plots or plot new ones + + """ + first_update = False + if update and hasattr(self, 'figs') and self.figs: + figs = self.figs + elif update: + figs = None + self.figs = [] + first_update = True + else: + figs = None + self.figs = [] + + idx = 0 + for plots in self.plot_list: + # If a conditional is included in the plot + try: + if len(plots) >=3 and not plots[2](self): + continue + except TypeError as e: + if len(plots) >= 3 and not plots[2](self, dataset): + continue + + name = plots[0] + plotter = plots[1] + + if figs is None: + fig = plt.figure() + self.figs.append(fig) + else: + fig = figs[idx] + + try: + plotter(self,fig) + plt.title(name) + except TypeError as e: + if dataset is not None: + try: + plotter(self, fig, dataset) + plt.title(name) + except (IndexError, KeyError, AttributeError) as e: + pass + + except (IndexError, KeyError, AttributeError) as e: + pass + + idx += 1 + + if update: + plt.draw() + fig.canvas.start_event_loop(0.001) + + if first_update: + plt.pause(0.05 * len(self.figs)) + + + def compare(self, dataset): + """Opens a tool for comparing simulated and measured diffraction patterns + + Parameters + ---------- + dataset : CDataset + A dataset containing the simulated diffraction patterns to compare against + """ + + 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 + # (say, from clicking back in the navbar) + # Holy fuck this was annoying. Sorry future for how + # crappy this solution is. + #if not np.allclose(im.colorbar.ax.get_xlim(), + # (np.min(im.get_array()), + # np.max(im.get_array()))): + if hasattr(im, 'norecurse') and im.norecurse: + im.norecurse=False + 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() + + + def update(idx): + idx = int(idx) % len(dataset) + fig.pattern_idx = idx + updating = True if len(axes[0].images) >= 1 else False + + inputs, output = dataset[idx] + sim_data = self.forward(*inputs).detach().cpu().numpy() + sim_data = sim_data + 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 + + if not updating: + axes[0].set_title('Simulated') + axes[1].set_title('Measured') + axes[2].set_title('Difference') + + sim = axes[0].imshow(sim_data) + meas = axes[1].imshow(meas_data * mask) + diff = axes[2].imshow((sim_data-meas_data) * mask) + + cb1 = plt.colorbar(sim, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) + cb1.ax.tick_params(labelrotation=20) + cb1.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(sim)) + cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) + cb2.ax.tick_params(labelrotation=20) + cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) + cb3 = plt.colorbar(diff, ax=axes[2], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) + cb3.ax.tick_params(labelrotation=20) + 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) + + meas = axes[1].images[-1] + meas.set_data(meas_data * mask) + update_colorbar(meas) + + diff = axes[2].images[-1] + diff.set_data((sim_data-meas_data) * mask) + update_colorbar(diff) + + + # This is dumb but the slider doesn't work unless a reference to it is + # kept somewhere... + self.slider = Slider(axslider, 'Pattern #', 0, len(dataset)-1, valstep=1, valfmt="%d") + self.slider.on_changed(update) + + 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.pattern_idx - 1) + elif event.key == 'down' or event.button == 'down': + update(fig.pattern_idx + 1) + self.slider.set_val(fig.pattern_idx) + plt.draw() + + fig.canvas.mpl_connect('key_press_event',on_action) + fig.canvas.mpl_connect('scroll_event',on_action) + update(0) + + diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index 28dc806..2dfec4c 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/CDTools/models/bragg_2d_ptycho.py @@ -12,6 +12,7 @@ from datetime import datetime import numpy as np from copy import copy +__all__ = ['Bragg2DPtycho'] # # Key ideas: diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index e517008..1bfb3d3 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -11,6 +11,7 @@ from datetime import datetime import numpy as np from copy import copy +__all__ = ['FancyPtycho'] class FancyPtycho(CDIModel): diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index dd79cd1..f529189 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -11,6 +11,7 @@ from datetime import datetime import numpy as np from copy import copy +__all__ = ['Multislice2DPtycho'] class Multislice2DPtycho(CDIModel): diff --git a/CDTools/models/pinhole_plane_ptycho.py b/CDTools/models/pinhole_plane_ptycho.py index d2b067b..9cb9d3d 100644 --- a/CDTools/models/pinhole_plane_ptycho.py +++ b/CDTools/models/pinhole_plane_ptycho.py @@ -11,6 +11,7 @@ from datetime import datetime import numpy as np from copy import copy +__all__ = ['PinholePlanePtycho'] class PinholePlanePtycho(CDIModel): diff --git a/CDTools/models/rpi.py b/CDTools/models/rpi.py index a36d1ef..7bc35bf 100644 --- a/CDTools/models/rpi.py +++ b/CDTools/models/rpi.py @@ -13,6 +13,8 @@ from datetime import datetime import numpy as np from copy import copy +__all__ = ['RPI'] + # # This model has to work a bit differently from a ptychography model # because a typical RPI dataset will have lots of images, each of which diff --git a/CDTools/models/s_matrix_ptycho.py b/CDTools/models/s_matrix_ptycho.py index 6c30537..b3ba1bd 100644 --- a/CDTools/models/s_matrix_ptycho.py +++ b/CDTools/models/s_matrix_ptycho.py @@ -11,6 +11,7 @@ from datetime import datetime import numpy as np from copy import copy +__all__ = ['SMatrixPtycho'] class SMatrixPtycho(CDIModel): diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 98c5756..baa662c 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -11,6 +11,8 @@ from matplotlib import pyplot as plt from datetime import datetime import numpy as np +__all__ = ['SimplePtycho'] + class SimplePtycho(CDIModel): """A simple ptychography model for exploring ideas and extensions diff --git a/CDTools/tools/__init__.py b/CDTools/tools/__init__.py index ec25428..a94f74b 100644 --- a/CDTools/tools/__init__.py +++ b/CDTools/tools/__init__.py @@ -1,13 +1,31 @@ +""" This module contains various packages of tools supporting common needs + +It at times feels a bit silly to keep all the tools siloed under a particular +subpackage such as "projectors" or "measurements", especially given that +"flat is better than nested", but I find something comforting about the +organization and the ability to only pull in the particular set of tools +that one needs for a specific application. + +The submodules are all structured as modules with their own __init__ files, +which use an import * statement to import from a file defining the various +functions. This is done to prevent leakage of imported packages into the +namespace of CDTools. I know, we're all consenting adults, but I just hate +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 losses +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 from CDTools.tools import analysis +from CDTools.tools import atoms diff --git a/CDTools/tools/analysis/__init__.py b/CDTools/tools/analysis/__init__.py new file mode 100644 index 0000000..5366b02 --- /dev/null +++ b/CDTools/tools/analysis/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.analysis.analysis import * diff --git a/CDTools/tools/analysis.py b/CDTools/tools/analysis/analysis.py similarity index 100% rename from CDTools/tools/analysis.py rename to CDTools/tools/analysis/analysis.py diff --git a/CDTools/tools/atoms/__init__.py b/CDTools/tools/atoms/__init__.py new file mode 100644 index 0000000..f5e7d52 --- /dev/null +++ b/CDTools/tools/atoms/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.atoms.atoms import * diff --git a/CDTools/tools/atoms.py b/CDTools/tools/atoms/atoms.py similarity index 100% rename from CDTools/tools/atoms.py rename to CDTools/tools/atoms/atoms.py diff --git a/CDTools/tools/cmath/__init__.py b/CDTools/tools/cmath/__init__.py new file mode 100644 index 0000000..828fd09 --- /dev/null +++ b/CDTools/tools/cmath/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.cmath.cmath import * diff --git a/CDTools/tools/cmath.py b/CDTools/tools/cmath/cmath.py similarity index 98% rename from CDTools/tools/cmath.py rename to CDTools/tools/cmath/cmath.py index 38e1f38..d969f53 100644 --- a/CDTools/tools/cmath.py +++ b/CDTools/tools/cmath/cmath.py @@ -12,8 +12,8 @@ import numpy as np import torch as t -__all__ = ['complex_to_torch','torch_to_complex','cabssq','cabs','cconj', - 'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift','expi'] +__all__ = ['complex_to_torch', 'torch_to_complex', 'cabssq', 'cabs', 'cconj', + 'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi'] # diff --git a/CDTools/tools/data/__init__.py b/CDTools/tools/data/__init__.py new file mode 100644 index 0000000..5e91e48 --- /dev/null +++ b/CDTools/tools/data/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.data.data import * diff --git a/CDTools/tools/data.py b/CDTools/tools/data/data.py similarity index 100% rename from CDTools/tools/data.py rename to CDTools/tools/data/data.py diff --git a/CDTools/tools/image_processing/__init__.py b/CDTools/tools/image_processing/__init__.py new file mode 100644 index 0000000..583a3f8 --- /dev/null +++ b/CDTools/tools/image_processing/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.image_processing.image_processing import * diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing/image_processing.py similarity index 100% rename from CDTools/tools/image_processing.py rename to CDTools/tools/image_processing/image_processing.py diff --git a/CDTools/tools/initializers/__init__.py b/CDTools/tools/initializers/__init__.py new file mode 100644 index 0000000..04928b4 --- /dev/null +++ b/CDTools/tools/initializers/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.initializers.initializers import * diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers/initializers.py similarity index 100% rename from CDTools/tools/initializers.py rename to CDTools/tools/initializers/initializers.py diff --git a/CDTools/tools/interactions/__init__.py b/CDTools/tools/interactions/__init__.py new file mode 100644 index 0000000..eef5f9b --- /dev/null +++ b/CDTools/tools/interactions/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.interactions.interactions import * diff --git a/CDTools/tools/interactions.py b/CDTools/tools/interactions/interactions.py similarity index 100% rename from CDTools/tools/interactions.py rename to CDTools/tools/interactions/interactions.py diff --git a/CDTools/tools/losses/__init__.py b/CDTools/tools/losses/__init__.py new file mode 100644 index 0000000..8114811 --- /dev/null +++ b/CDTools/tools/losses/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.losses.losses import * diff --git a/CDTools/tools/losses.py b/CDTools/tools/losses/losses.py similarity index 100% rename from CDTools/tools/losses.py rename to CDTools/tools/losses/losses.py diff --git a/CDTools/tools/measurements/__init__.py b/CDTools/tools/measurements/__init__.py new file mode 100644 index 0000000..544a9e7 --- /dev/null +++ b/CDTools/tools/measurements/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.measurements.measurements import * diff --git a/CDTools/tools/measurements.py b/CDTools/tools/measurements/measurements.py similarity index 100% rename from CDTools/tools/measurements.py rename to CDTools/tools/measurements/measurements.py diff --git a/CDTools/tools/plotting/__init__.py b/CDTools/tools/plotting/__init__.py new file mode 100644 index 0000000..4cc7e0b --- /dev/null +++ b/CDTools/tools/plotting/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.plotting.plotting import * diff --git a/CDTools/tools/plotting.py b/CDTools/tools/plotting/plotting.py similarity index 100% rename from CDTools/tools/plotting.py rename to CDTools/tools/plotting/plotting.py diff --git a/CDTools/tools/projectors/__init__.py b/CDTools/tools/projectors/__init__.py new file mode 100644 index 0000000..6234ec0 --- /dev/null +++ b/CDTools/tools/projectors/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.projectors.projectors import * diff --git a/CDTools/tools/projectors.py b/CDTools/tools/projectors/projectors.py similarity index 100% rename from CDTools/tools/projectors.py rename to CDTools/tools/projectors/projectors.py diff --git a/CDTools/tools/propagators/__init__.py b/CDTools/tools/propagators/__init__.py new file mode 100644 index 0000000..91285c0 --- /dev/null +++ b/CDTools/tools/propagators/__init__.py @@ -0,0 +1,3 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools.propagators.propagators import * diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators/propagators.py similarity index 99% rename from CDTools/tools/propagators.py rename to CDTools/tools/propagators/propagators.py index 2ab0bad..3e8ab1d 100644 --- a/CDTools/tools/propagators.py +++ b/CDTools/tools/propagators/propagators.py @@ -14,7 +14,10 @@ from matplotlib import pyplot as plt __all__ = ['far_field', 'near_field', 'generate_angular_spectrum_propagator', - 'inverse_far_field', 'inverse_near_field'] + 'inverse_far_field', 'inverse_near_field', + 'generate_high_NA_k_intensity_map', + 'high_NA_far_field', + 'generate_generalized_angular_spectrum_propagator'] def far_field(wavefront): diff --git a/examples/example_data/Optical_ptycho_incoherent.pickle b/examples/example_data/Optical_ptycho_incoherent.pickle new file mode 100644 index 0000000..13964cf Binary files /dev/null and b/examples/example_data/Optical_ptycho_incoherent.pickle differ diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index 81afc57..c28ee26 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -9,7 +9,8 @@ from torch.utils.data import Subset ss_filename = 'example_data/Optical_Data_ss.cxi' -with open('example_data/Optical_ptycho.pickle', 'rb') as f: +#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) probe = ptycho_results['probe'] @@ -19,7 +20,7 @@ dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(ss_filename) # Next, we create a ptychography model from the dataset # Note that we explicitly as for two incoherent probe modes -model = CDTools.models.RPI.from_dataset(dataset, probe, [800,800], +model = CDTools.models.RPI.from_dataset(dataset, probe, [900,900], background=background, n_modes=2) @@ -31,12 +32,14 @@ dataset.get_as(device='cuda') # 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) + #model.inspect(dataset) print(i,loss) - + +model.inspect(dataset) + # Now we use the regularizer to damp all but the top modes -for i, loss in enumerate(model.LBFGS_optimize(20, dataset, lr=0.4, regularization_factor=[0.001,0.1])): - model.inspect(dataset) +for i, loss in enumerate(model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1])): + #model.inspect(dataset) print(i,loss) results = model.save_results()