diff --git a/CDTools/datasets/__init__.py b/CDTools/datasets/__init__.py index 121b152..30b0322 100644 --- a/CDTools/datasets/__init__.py +++ b/CDTools/datasets/__init__.py @@ -1,3 +1,32 @@ +""" 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 +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 +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: + +* __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 @@ -13,6 +42,8 @@ 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 @@ -21,12 +52,46 @@ from matplotlib import ticker # 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) @@ -45,6 +110,13 @@ class CDataset(torchdata.Dataset): 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: @@ -59,9 +131,25 @@ class CDataset(torchdata.Dataset): 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 @@ -80,13 +168,46 @@ class CDataset(torchdata.Dataset): def _load(self, index): - # Internal function to load data + """ 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: @@ -109,6 +230,23 @@ class CDataset(torchdata.Dataset): 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: @@ -129,5 +267,14 @@ class CDataset(torchdata.Dataset): 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.ptycho_2d_dataset import Ptycho2DDataset diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index 7ae3ad5..6e224b0 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -15,15 +15,47 @@ from matplotlib import ticker __all__ = ['Ptycho2DDataset'] -# -# This is the standard dataset for a 2D ptychography experiment, -# which saves and loads files compatible with most reconstruction -# programs (only tested against SHARP) -# class Ptycho2DDataset(CDataset): + """The standard dataset for a 2D ptychography scan + Subclasses datasets.CDataset + + This class loads and saves 2D ptychography scan data from .cxi files. + It should save and load files compatible with most reconstruction + programs, although it is only tested against SHARP. + """ def __init__(self, translations, patterns, axes=None, *args, **kwargs): + """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 + ---------- + translations : array + An nx3 array containing the probe translations at each scan point + patterns : array + An nxmxl array containing the full stack of measured diffraction patterns + 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 + """ + super(Ptycho2DDataset,self).__init__(*args, **kwargs) self.axes = copy(axes) @@ -40,10 +72,41 @@ class Ptycho2DDataset(CDataset): return self.patterns.shape[0] def _load(self, index): + """ Internal function to load data + + 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. + + The inputs for a 2D ptychogaphy data set are: + + 1) The indices of the patterns to use + 2) The recorded probe positions associated with those points + + Parameters + ---------- + index : int or slice + The index or indices of the scan points to use + + Returns + ------- + inputs : tuple + A tuple of the inputs to the related forward models + outputs : tuple + The output pattern or stack of output patterns + """ return (index, self.translations[index]), self.patterns[index] def to(self, *args, **kwargs): + """Sends the relevant data to the given device and dtype + + This function sends the stored translations, patterns, + mask and background to the specified device and dtype + + Accepts the same parameters as torch.Tensor.to + """ super(Ptycho2DDataset,self).to(*args,**kwargs) self.translations = self.translations.to(*args, **kwargs) self.patterns = self.patterns.to(*args, **kwargs) @@ -53,6 +116,21 @@ class Ptycho2DDataset(CDataset): # perhaps there is a way but I couldn't figure it out. @classmethod def from_cxi(cls, cxi_file): + """Generates a new CDataset from a .cxi file directly + + This generates a new Ptycho2DDataset from a .cxi file storing + a 2D ptychography scan. + + Parameters + ---------- + file : str, pathlib.Path, or h5py.File + The .cxi file to load from + + Returns + ------- + dataset : Ptycho2DDataset + 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: @@ -80,12 +158,38 @@ class Ptycho2DDataset(CDataset): def to_cxi(self, cxi_file): + """Saves out a Ptycho2DDataset as a .cxi file + + This function saves all the compatible information in a + Ptycho2DDataset object into a .cxi file. This saved .cxi file + should be compatible with any standard .cxi file based + reconstruction tool, such as SHARP. + + 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) + super(Ptycho2DDataset,self).to_cxi(cxi_file) cdtdata.add_data(cxi_file, self.patterns, axes=self.axes) cdtdata.add_ptycho_translations(cxi_file, self.translations) def inspect(self): + """Launches an interactive plot for perusing the data + + This launches an interactive plotting tool in matplotlib that + shows the spatial map constructed from the integrated intensity + at each position on the left, next to a panel on the right that + can display a base-10 log plot of the detector readout at each + position. + """ 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]) @@ -146,7 +250,7 @@ class Ptycho2DDataset(CDataset): cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.15,fraction=0.1) cb1.ax.tick_params(labelrotation=20) - meas = axes[1].imshow(meas_data * mask) + meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask) cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.15,fraction=0.1) cb2.ax.tick_params(labelrotation=20) @@ -173,7 +277,7 @@ class Ptycho2DDataset(CDataset): meas = axes[1].images[-1] - meas.set_data(meas_data * mask) + meas.set_data(np.log(meas_data) / np.log(10) * mask) update_colorbar(meas) diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 252d005..1eb7802 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -1,3 +1,48 @@ +"""This module contains all the models for different CDI Reconstructions + +All the reconstructions are coordinated through the ptychography models +defined here. The models are, at their core, just subclasses of the +:code:`torch.nn.model` class, so they contain the same structure of +parameters, etc. Their central functionality is as a simulation that maps +some input (usually, the index number of a scan point) to an output that +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. + + +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 + +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 @@ -7,56 +52,19 @@ from matplotlib.widgets import Slider from matplotlib import ticker import numpy as np -# -# This is unrelated, but it will then be important to be able to save and load -# models easily from a predefined format. To be honest, this could just -# literally be by pickling the model. They could also be saved out as -# state_dicts or via torch.save. I think it's best to just save the whole -# model - I lose out on the modularity of just saving the state_dict, but -# I gain in it being easy to reload the non-learned aspects of the model, -# like the wavelength and sample geometry. Remember that it's important -# that the final outputs of the reconstructions are transferrable to other -# places -# - - -# -# For now, just save/load model via the built-in t.save() and t.load() -# functions -# - +__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho'] 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. The functions required can be - split into several subsections: - - Creation: - from_dataset : a function to create a CDIModel from an appropriate CDataset - - Simulation: - interaction : a function to simulate exit waves from experimental parameters - forward_propagator : the propagator from the experiment plane to the detector plane - backward_propagator : the propagator from the detector plane to the experiment plane - measurement : a function to simulate the detector readout from a detector plane wavefront - forward : predefined, the entire stacked forward model - loss : the loss function to report and use for automatic differentiation - simulation : predefined, simulates a stack of detector images from the forward model - simulate_to_dataset : a function to create a CDataset from the simulation defined in the model - - Reconstruction: - AD_optimize : predefined, a generic automatic differentiation reconstruction - Adam_optimize : predefined, sensible automatic differentiation reconstruction using ADAM - - The work of defining the various subclasses boils down to creating an - appropriate implementation for this set of functions. + 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() @@ -78,27 +86,51 @@ class CDIModel(t.nn.Module): 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() - # I know this is silly but it makes it clear this should be explicitly - # overwritten def to(self, *args, **kwargs): super(CDIModel,self).to(*args,**kwargs) - def simulate(self, args_list): - return t.Tensor([self.forward(*args) for args in args_list]) - - 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): + """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 + """ for it in range(iterations): loss = 0 @@ -127,7 +159,26 @@ class CDIModel(t.nn.Module): def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False): + """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 + """ # Make a dataloader data_loader = torchdata.DataLoader(dataset, batch_size=batch_size, shuffle=True) @@ -147,7 +198,27 @@ class CDIModel(t.nn.Module): def LBFGS_optimize(self, iterations, dataset, batch_size=None, lr=0.1,history_size=2): + """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. + """ + # Make a dataloader if batch_size is not None: data_loader = torchdata.DataLoader(dataset, batch_size=batch_size, @@ -178,9 +249,20 @@ class CDIModel(t.nn.Module): registered plots which need to incorporate some information from the dataset (such as geometry or a comparison with measured data). - Args: - dataset (CDataset): Optional, a dataset matched to the model type - update (bool) : Whether to update existing plots or plot new ones + 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 @@ -239,8 +321,10 @@ class CDIModel(t.nn.Module): def compare(self, dataset): """Opens a tool for comparing simulated and measured diffraction patterns - Args: - dataset (CDataset) : A dataset containing the simulated diffraction patterns to compare agains + Parameters + ---------- + dataset : CDataset + A dataset containing the simulated diffraction patterns to compare against """ fig, axes = plt.subplots(1,3,figsize=(12,5.3)) diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 54a39a4..5e5deb1 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -12,7 +12,11 @@ from datetime import datetime import numpy as np class SimplePtycho(CDIModel): + """A simple ptychography model for exploring ideas and extensions + + + """ def __init__(self, wavelength, detector_geometry, probe_basis, detector_slice, probe_guess, obj_guess, min_translation = [0,0], diff --git a/docs/source/conf.py b/docs/source/conf.py index ac2b4af..b5dd381 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -47,6 +47,13 @@ extensions = [ ] +# One only works in 1.8+, the other is depricated in >1.8 +autodoc_default_options = { + 'show-inheritance': True +} +autodoc_default_flags = [ 'show-inheritance'] + + # Add any paths that contain templates here, relative to this directory. templates_path = ['_templates'] diff --git a/docs/source/datasets.rst b/docs/source/datasets.rst index 1368c14..fade617 100644 --- a/docs/source/datasets.rst +++ b/docs/source/datasets.rst @@ -3,7 +3,7 @@ Datasets .. automodule:: CDTools.datasets :members: - + :private-members: diff --git a/docs/source/models.rst b/docs/source/models.rst index 3ed0c8a..5291f60 100644 --- a/docs/source/models.rst +++ b/docs/source/models.rst @@ -1,4 +1,7 @@ Models ====== +.. automodule:: CDTools.models + :members: + :private-members: