From e80bf05f253c4795cbbdeb4f63e1ac2266c96c4b Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 4 Sep 2019 16:27:55 -0400 Subject: [PATCH] Further improve docs layout and add unfinished ptycho models to repo --- CDTools/models/bragg_2d_ptycho.py | 36 +++ CDTools/models/pinhole_plane_ptycho.py | 383 +++++++++++++++++++++++++ docs/source/examples.rst | 16 ++ docs/source/index.rst | 37 +-- docs/source/indices_tables.rst | 6 + docs/source/installation.rst | 43 +-- examples/simple_ptycho.py | 18 +- 7 files changed, 487 insertions(+), 52 deletions(-) create mode 100644 CDTools/models/bragg_2d_ptycho.py create mode 100644 CDTools/models/pinhole_plane_ptycho.py create mode 100644 docs/source/indices_tables.rst diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py new file mode 100644 index 0000000..daf7fd2 --- /dev/null +++ b/CDTools/models/bragg_2d_ptycho.py @@ -0,0 +1,36 @@ +from __future__ import division, print_function, absolute_import + +import torch as t +from CDTools.models import CDIModel +from CDTools.datasets import Ptycho_2D_Dataset +from CDTools import tools +from CDTools.tools import cmath +from CDTools.tools import plotting as p +from matplotlib import pyplot as plt +from datetime import datetime +import numpy as np +from copy import copy + + +# +# Key ideas: +# 1) To a first approximation, do the reconstruction on a parallelogram +# shaped grid on the sample which is conjugate to the detector coordinates +# 2) Simulate a probe in those same coordinates, but propagate it back and +# forth (along with the translations), using the angular spectrum method +# 3) Apply a correction to the simulated data to account for the tilt of the +# sample with respect to the detector +# 4) Include a correction for the thickness of the sample +# + + +class Bragg2DPtycho(CDIModel): + + def __init__(self, wavelength, detector_geometry, + probe_basis, probe_guess, obj_guess, + detector_slice=None, + surface_normal=np.array([0.,0.,1.]), + min_translation = t.Tensor([0,0]), + background = None, translation_offsets=None, mask=None, + weights = None, translation_scale = 1, saturation=None, + probe_support = None, obj_support=None, oversampling=1): diff --git a/CDTools/models/pinhole_plane_ptycho.py b/CDTools/models/pinhole_plane_ptycho.py new file mode 100644 index 0000000..58485b1 --- /dev/null +++ b/CDTools/models/pinhole_plane_ptycho.py @@ -0,0 +1,383 @@ +from __future__ import division, print_function, absolute_import + +import torch as t +from CDTools.models import CDIModel +from CDTools.datasets import Ptycho_2D_Dataset +from CDTools import tools +from CDTools.tools import cmath +from CDTools.tools import plotting as p +from matplotlib import pyplot as plt +from datetime import datetime +import numpy as np +from copy import copy + + +class PinholePlanePtycho(CDIModel): + + def __init__(self, wavelength, detector_geometry, + probe_basis, + pinhole_guess, obj_guess, illumination_guess, + pinhole_propagation, + detector_slice=None, + surface_normal=np.array([0.,0.,1.]), + min_translation = t.Tensor([0,0]), + background = None, translation_offsets=None, mask=None, + weights = None, translation_scale = 1, saturation=None, + pinhole_support = None, obj_support=None, oversampling=1): + + super(PinholePlanePtycho,self).__init__() + self.wavelength = t.Tensor([wavelength]) + self.detector_geometry = copy(detector_geometry) + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = t.Tensor(det_geo['distance']) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = t.Tensor(det_geo['basis']) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = t.Tensor(det_geo['corner']) + + self.min_translation = t.Tensor(min_translation) + + self.probe_basis = t.Tensor(probe_basis) + self.detector_slice = detector_slice + self.surface_normal = t.Tensor(surface_normal) + + self.saturation = saturation + + if mask is None: + self.mask = mask + else: + self.mask = t.ByteTensor(mask) + + # We rescale the pinhole here so it learns at the same rate as the + # object + self.pinhole_norm = 1 * t.max(tools.cmath.cabs(pinhole_guess.to(t.float32))) + + self.pinhole = t.nn.Parameter(pinhole_guess.to(t.float32) + / self.pinhole_norm) + + self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + self.illumination = t.nn.Parameter(illumination_guess.to(t.float32)) + + spacing = np.abs([self.probe_basis.numpy()[0,1], self.probe_basis.numpy()[1,0]]) + + + self.pinhole_propagator = tools.propagators.generate_angular_spectrum_propagator(self.pinhole.shape[:-1], spacing, self.wavelength, pinhole_propagation).to(dtype=self.pinhole.dtype) + + + if background is None: + if detector_slice is not None: + background = 1e-6 * t.ones(self.pinhole[self.detector_slice].shape[:-1]) + else: + background = 1e-6 * t.ones(self.pinhole.shape[:-1]) + + + self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) + + if weights is None: + self.weights = None + else: + self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32)) + + if translation_offsets is None: + self.translation_offsets = None + else: + self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale) + + self.translation_scale = translation_scale + + if pinhole_support is not None: + self.pinhole_support = pinhole_support + else: + self.pinhole_support = t.ones_like(self.pinhole) + + if obj_support is not None: + self.obj_support = obj_support + self.obj.data = self.obj * obj_support + else: + self.obj_support = t.ones_like(self.obj) + + self.oversampling = oversampling + + + @classmethod + def from_dataset(cls, dataset, pinhole_propagation, pinhole_size=None, randomize_ang=0, padding=0, translation_scale = 1, saturation=None, pinhole_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1): + + wavelength = dataset.wavelength + det_basis = dataset.detector_geometry['basis'] + det_shape = dataset[0][1].shape + distance = dataset.detector_geometry['distance'] + + # always do this on the cpu + get_as_args = dataset.get_as_args + dataset.get_as(device='cpu') + (indices, translations), patterns = dataset[:] + dataset.get_as(*get_as_args[0],**get_as_args[1]) + + # Set to none to avoid issues with things outside the detector + center = tools.image_processing.centroid(t.sum(patterns,dim=0)) + + # Then, generate the pinhole geometry from the dataset + ewg = tools.initializers.exit_wave_geometry + probe_basis, pinhole_shape, det_slice = ewg(det_basis, + det_shape, + wavelength, + distance, + center=center, + padding=padding, + opt_for_fft=False, + oversampling=oversampling) + + + if hasattr(dataset, 'sample_info') and \ + dataset.sample_info is not None and \ + 'orientation' in dataset.sample_info: + surface_normal = dataset.sample_info['orientation'][2] + else: + surface_normal = np.array([0.,0.,1.]) + + + # If this information is supplied when the function is called, + # then we override the information in the .cxi file + if scattering_mode in {'t', 'transmission'}: + surface_normal = np.array([0.,0.,1.]) + elif scattering_mode in {'r', 'reflection'}: + outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) + outgoing_dir /= np.linalg.norm(outgoing_dir) + surface_normal = outgoing_dir + np.array([0.,0.,1.]) + surface_normal /= np.linalg.norm(outgoing_dir) + + + # Next generate the object geometry from the probe geometry and + # the translations + pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal) + + obj_size, min_translation = tools.initializers.calc_object_setup(pinhole_shape, pix_translations, padding=200) + + if hasattr(dataset, 'background') and dataset.background is not None: + background = t.sqrt(dataset.background) + else: + background = None + + # Finally, initialize the pinhole and object using this information + if pinhole_size is None: + pinhole = tools.initializers.SHARP_style_probe(dataset, pinhole_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) + else: + pinhole = tools.initializers.gaussian_probe(dataset, probe_basis, pinhole_shape, pinhole_size, propagation_distance=propagation_distance) + + + illumination = tools.cmath.expi(t.zeros(obj_size)) + obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) + + det_geo = dataset.detector_geometry + + translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) + + weights = t.ones(len(dataset)) + + if hasattr(dataset, 'mask') and dataset.mask is not None: + mask = dataset.mask.to(t.uint8) + else: + mask = None + + if pinhole_support_radius is not None: + pinhole_support = t.zeros_like(pinhole.to(dtype=t.float32)) + p_cent = np.array(pinhole.shape[:2]).astype(int) // 2 + psr = int(pinhole_support_radius) + pinhole_support[p_cent[0]-psr:p_cent[0]+psr, + p_cent[1]-psr:p_cent[1]+psr] = 1 + pinhole = pinhole_support * pinhole + else: + pinhole_support = None; + + if restrict_obj != -1: + ro = restrict_obj + os = np.array(obj_size) + ps = np.array(pinhole_shape) + obj_support = t.zeros_like(obj.to(dtype=t.float32)) + obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2, + ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1 + else: + obj_support = None + + return cls(wavelength, det_geo, probe_basis, pinhole, obj, + illumination, pinhole_propagation, + detector_slice=det_slice, + surface_normal=surface_normal, + min_translation=min_translation, + translation_offsets = translation_offsets, + weights=weights, mask=mask, background=background, + translation_scale=translation_scale, + saturation=saturation, + pinhole_support=pinhole_support, + obj_support=obj_support, + oversampling=oversampling) + + + def interaction(self, index, translations): + pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, + translations, + surface_normal=self.surface_normal) + pix_trans -= self.min_translation + + if self.translation_offsets is not None: + pix_trans += self.translation_scale * self.translation_offsets[index] + + pr = self.pinhole * self.pinhole_support + pinhole_exit_waves = self.pinhole_norm * \ + tools.interactions.ptycho_2D_sinc(pr, + self.illumination, + pix_trans, + shift_probe=True) + sample_illumination = tools.propagators.near_field(pinhole_exit_waves, + self.pinhole_propagator) + exit_waves = [] + if pinhole_exit_waves.dim() == 4: + for pix_tran, illum in zip(pix_trans, sample_illumination): + exit_waves.append(tools.interactions.ptycho_2D_sinc(illum, + self.obj_support * self.obj, + pix_tran, + shift_probe=True)) + exit_waves = t.stack(exit_waves) + else: + exit_waves = tools.interactions.ptycho_2D_sinc(sample_illumination, + self.obj_support * self.obj, + pix_trans, + shift_probe=True) + + + + if exit_waves.dim() == 4: + exit_waves = self.weights[index][:,None,None,None] * exit_waves + else: + exit_waves = self.weights[index] * exit_waves + + + return exit_waves + + + def forward_propagator(self, wavefields): + return tools.propagators.far_field(wavefields) + + + def backward_propagator(self, wavefields): + return tools.propagators.inverse_far_field(wavefields) + + + def measurement(self, wavefields): + return tools.measurements.quadratic_background(wavefields, + self.background, + detector_slice=self.detector_slice, + saturation=self.saturation, + oversampling=self.oversampling) + + + def loss(self, sim_data, real_data, mask=None): + return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) + + + def to(self, *args, **kwargs): + super(PinholePlanePtycho, self).to(*args, **kwargs) + self.wavelength = self.wavelength.to(*args,**kwargs) + # move the detector geometry too + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = det_geo['distance'].to(*args,**kwargs) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = det_geo['basis'].to(*args,**kwargs) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = det_geo['corner'].to(*args,**kwargs) + + if self.mask is not None: + self.mask = self.mask.to(*args, **kwargs) + + + self.min_translation = self.min_translation.to(*args,**kwargs) + self.probe_basis = self.probe_basis.to(*args,**kwargs) + self.pinhole_norm = self.pinhole_norm.to(*args,**kwargs) + self.pinhole_support = self.pinhole_support.to(*args,**kwargs) + self.obj_support = self.obj_support.to(*args,**kwargs) + self.surface_normal = self.surface_normal.to(*args, **kwargs) + self.pinhole_propagator = self.pinhole_propagator.to(*args, **kwargs) + + + def sim_to_dataset(self, args_list): + # In the future, potentially add more control + # over what metadata is saved (names, etc.) + + # First, I need to gather all the relevant data + # that needs to be added to the dataset + entry_info = {'program_name': 'CDTools', + 'instrument_n': 'Simulated Data', + 'start_time': datetime.now()} + + surface_normal = self.surface_normal.detach().cpu().numpy() + xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal) + xsurfacevec /= np.linalg.norm(xsurfacevec) + ysurfacevec = np.cross(surface_normal, xsurfacevec) + ysurfacevec /= np.linalg.norm(ysurfacevec) + orientation = np.array([xsurfacevec, ysurfacevec, surface_normal]) + + sample_info = {'description': 'A simulated sample', + 'orientation': orientation} + + + detector_geometry = self.detector_geometry + mask = self.mask + wavelength = self.wavelength + indices, translations = args_list + + # Then we simulate the results + data = self.forward(indices, translations) + + # And finally, we make the dataset + return Ptycho_2D_Dataset(translations, data, + entry_info = entry_info, + sample_info = sample_info, + wavelength=wavelength, + detector_geometry=detector_geometry, + mask=mask) + + + def corrected_translations(self,dataset): + translations = dataset.translations.to(dtype=self.pinhole.dtype,device=self.pinhole.device) + t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal) + return translations + t_offset + + + # Needs to be updated to allow for plotting to an existing figure + plot_list = [ + ('Pinhole Amplitude', + lambda self, fig: p.plot_amplitude(self.pinhole, fig=fig, basis=self.probe_basis)), + ('Pinhole Phase', + lambda self, fig: p.plot_phase(self.pinhole, fig=fig, basis=self.probe_basis)), + ('Illumination Amplitude', + lambda self, fig: p.plot_amplitude(self.illumination, fig=fig, basis=self.probe_basis), + lambda self: len(self.pinhole) >=2), + ('Illumination Phase', + lambda self, fig: p.plot_phase(self.illumination, fig=fig, basis=self.probe_basis), + lambda self: len(self.pinhole) >=2), + ('Object Amplitude', + lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), + ('Object Phase', + lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)), + ('Corrected Translations', + lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)), + ('Background', + lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)) + ] + + + def save_results(self, dataset): + basis = self.probe_basis.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + probe = cmath.torch_to_complex(self.pinhole.detach().cpu()) + probe = probe * self.pinhole_norm.detach().cpu().numpy() + obj = cmath.torch_to_complex(self.obj.detach().cpu()) + background = self.background.detach().cpu().numpy()**2 + weights = self.weights.detach().cpu().numpy() + + return {'basis':basis, 'translation':translations, + 'probe':probe,'obj':obj, + 'background':background, + 'weights':weights} diff --git a/docs/source/examples.rst b/docs/source/examples.rst index ce68152..b67d694 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -1,3 +1,19 @@ Examples ======== +Included with the repository are a number of example scripts that demonstrate how various aspects of CDTools work. It is recommended to read through several of them before continuing to the tutorial, to get a feel for how the ptychography scripts work. + + +Simple Ptycho +------------- + +This script runs a ptychography reconstruction using the SimplePtycho model. This model implements a straightforward transmission ptychography geometry, with no position refinement, incoherent modes, or other doodads (only a background model). + +.. literalinclude:: ../../examples/simple_ptycho.py + +Running this script leads to a very poor reconstruction on the example data, however depending on the dataset the model can produce reasonable results. When reading this script, note the basic workflow: First, data is read into a dataset object. Then, a model is created to match the geometry stored in the dataset, with a default initialization for all the parameters. Next, a reconstruction is run (and the progress reported out). Finally, the finished reconstruction is inspected. + + +Gold Balls Ptycho +----------------- + diff --git a/docs/source/index.rst b/docs/source/index.rst index 42c19cc..c2f1f36 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,3 +1,16 @@ +.. toctree:: + :hidden: + :maxdepth: 1 + + self + installation + examples + tutorial + general + datasets + models + tools/index + indices_tables Introduction to CDTools ======================= @@ -32,27 +45,3 @@ CDTools makes it simple to load and inspect from data stored in .cxi files using The high-level interface to CDTools is built on a lower level "three-legged stool". This consists of tools to access stored data, tools to visualize data and reconstructions, and tools that implement basic operations relevant to coherent diffraction. All of these tools can be used directly alongside the high-level interface, when needed. Enough blabber. If you're interested, read the docs! - - -Documentation Overview -====================== - -.. toctree:: - :maxdepth: 1 - - installation - examples - tutorial - general - datasets - models - tools/index - - - -Indices and tables -================== - -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` diff --git a/docs/source/indices_tables.rst b/docs/source/indices_tables.rst new file mode 100644 index 0000000..d46b839 --- /dev/null +++ b/docs/source/indices_tables.rst @@ -0,0 +1,6 @@ +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/docs/source/installation.rst b/docs/source/installation.rst index f7be677..b80e6fa 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -1,17 +1,19 @@ Installation ============ -CDTools can be downloaded from it's `MIT github page`_, and the relevant prerequisites can be downloaded via pip, conda, or most python package managers. +Step 1: Download +---------------- -.. _`MIT github page`: https://github.mit.edu/Scattering/ADCD +The source code for CDTools is hosted on it's `MIT github page`_. Access to the MIT github page can be granted to any member of the MIT community - if you are not a member, access can be arranged through a guest account on MIT's GitHub enterprise server. +.. _`MIT github page`: https://github.mit.edu/Scattering/CDTools -Access to the MIT github page can be granted to any member of the MIT community - if you are not a member, access can be arranged through a guest account on the github enterprise server +It is recommended that you clone the repository, rather than just downloading the contents, as it remains under heavy development. Cloning the repository will allow you to get access to new updates. -Prerequisites -------------- +Step 2: Install Dependencies +---------------------------- -CDTools has the following prerequisites: +CDTools depends on the following packages: * `numpy `_ * `scipy `_ @@ -20,30 +22,39 @@ CDTools has the following prerequisites: * `python-dateutil `_ * `h5py `_ -All of these can be installed via pip or conda. It is required that pytorch is ilt with MKL, as that enables FFTs. Additionally, installing pytorch with CUDA support is recommended for running any serious reconstructions with the package. The code is written to be python 2.7+ compatible, although it is only tested in python 3. +And has optional dependencies on -Finally, to run the tests, pytest is required, and to build the docs, sphinx and sphinx-argparse are required. + * `pytest `_ + * `sphinx `_ + * `sphinx-argparse `_ + +All of these can be installed via pip or conda. Finally, CDTools is written to be python 2.7+ compatible, but is only actively tested on python 3. + +*It is required that pytorch is built with MKL*, as that enables FFTs. Additionally, installing pytorch *with CUDA support* is recommended, if you intend to run any serious reconstructions with the package. + +Finally, the optional depencency on pytest enables the tests to be run to confirm a successful installation. Sphinx and sphinx-argparse are only required if you plan on building the documentation. -Installation ------------- +Step 3: Install +--------------- -CDTools can be installed via pip, although it is recommended to install it in development mode as changes to the code are pushed to the MIT github quite frequently. - -To install in developer mode, run the following command from the top level directory (the directory including the setup.py file) +To install in CDTools in developer mode (recommended, to allow any updates to be pushed immediately through your system), run the following command from the top level directory (the directory including the setup.py file) .. code:: bash $ pip install -e . +If you prefer to use a tool other than pip, CDTools can be installed via any other package management tool that works with a setup.py file. -Run The Tests -------------- + +Step 4: Run The Tests +--------------------- -To ensure that the installation has worked correctly, it is recommended to run the unit tests. After ensuring that `pytest `_ is installed, run the following command from the top level directory: +To ensure that the installation has worked correctly, it is recommended, although not required, that you run the unit tests. After ensuring that pytest is installed, run the following command from the top level directory: .. code:: bash $ pytest +If any tests fail, make sure that you have all the noted dependencies properly installed (specifically that pytorch is installed with MKL support). If so, `shoot me an email `_ and we'll get to the bottom of it. diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 77cf05a..9334ab0 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -1,27 +1,21 @@ from __future__ import division, print_function, absolute_import import CDTools -from CDTools.tools import cmath -from CDTools.tools.plotting import * from matplotlib import pyplot as plt -import pickle +# First, we load an example dataset from a .cxi file filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' - dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename) +# Next, we create a ptychography model from the dataset model = CDTools.models.SimplePtycho.from_dataset(dataset) +# Now, we run a short reconstruction from the dataset! for i, loss in enumerate(model.Adam_optimize(10, dataset)): + model.inspect(dataset) print(i, loss) -for i, loss in enumerate(model.Adam_optimize(10, dataset, lr=0.001)): - print(i, loss) - -model.inspect(update=False) +# Finally, we plot the results +model.inspect(dataset) model.compare(dataset) - -#with open('test_results.pickle', 'wb') as f: -# pickle.dump(model.save_results(),f) - plt.show()