mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-19 09:02:09 +02:00
More work on documentation
This commit is contained in:
Abe Levitan
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
+137
-53
@@ -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))
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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']
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ Datasets
|
||||
|
||||
.. automodule:: CDTools.datasets
|
||||
:members:
|
||||
|
||||
:private-members:
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
Models
|
||||
======
|
||||
|
||||
.. automodule:: CDTools.models
|
||||
:members:
|
||||
:private-members:
|
||||
|
||||
|
||||
Reference in New Issue
Block a user