mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-11-12 17:23:54 +01:00
Clean up namespaces and also add a test of RPI with incoherent probes
This commit is contained in:
Abe Levitan committed
1 parent
4be2af8609
commit
6b735a431f
39 files changed
+871
-764
No files matched your search
@@ -3,3 +3,4 @@ from __future__ import division, print_function, absolute_import
|
||||
from CDTools import tools
|
||||
from CDTools import datasets
|
||||
from CDTools import models
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
+10
-494
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from datetime import datetime
|
||||
import numpy as np
|
||||
from copy import copy
|
||||
|
||||
__all__ = ['Bragg2DPtycho']
|
||||
|
||||
#
|
||||
# Key ideas:
|
||||
|
||||
@@ -11,6 +11,7 @@ from datetime import datetime
|
||||
import numpy as np
|
||||
from copy import copy
|
||||
|
||||
__all__ = ['FancyPtycho']
|
||||
|
||||
class FancyPtycho(CDIModel):
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ from datetime import datetime
|
||||
import numpy as np
|
||||
from copy import copy
|
||||
|
||||
__all__ = ['Multislice2DPtycho']
|
||||
|
||||
class Multislice2DPtycho(CDIModel):
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ from datetime import datetime
|
||||
import numpy as np
|
||||
from copy import copy
|
||||
|
||||
__all__ = ['PinholePlanePtycho']
|
||||
|
||||
class PinholePlanePtycho(CDIModel):
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -11,6 +11,7 @@ from datetime import datetime
|
||||
import numpy as np
|
||||
from copy import copy
|
||||
|
||||
__all__ = ['SMatrixPtycho']
|
||||
|
||||
class SMatrixPtycho(CDIModel):
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.analysis.analysis import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.atoms.atoms import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.cmath.cmath import *
|
||||
@@ -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']
|
||||
|
||||
|
||||
#
|
||||
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.data.data import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.image_processing.image_processing import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.initializers.initializers import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.interactions.interactions import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.losses.losses import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.measurements.measurements import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.plotting.plotting import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.projectors.projectors import *
|
||||
File renamed without changes.
@@ -0,0 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools.tools.propagators.propagators import *
|
||||
@@ -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):
|
||||
Binary file not shown.
@@ -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()
|
||||
|
||||
Reference in new issue
Block a user