Finish fleshing out the examples

This commit is contained in:
Abe Levitan
2019-09-04 18:23:34 -04:00
parent 0377ea90be
commit 99d02f95cf
10 changed files with 158 additions and 105 deletions
+8 -8
View File
@@ -103,8 +103,8 @@ def plot_amplitude(im, fig = None, basis=None, units='um', cmap='viridis', **kwa
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
Returns
-------
@@ -169,8 +169,8 @@ def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
Returns
-------
@@ -244,8 +244,8 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
Optional, the 3x2 probe basis
units : str
The length units to mark on the plot, default is um
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
Returns
-------
@@ -299,8 +299,8 @@ def plot_translations(translations, fig=None, units='um', lines=True, **kwargs):
Default is um, units to report in (assuming input in m)
lines : bool
Whether to plot lines indicating the path taken
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
Returns
+85 -2
View File
@@ -1,9 +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.
Included with the repository are a number of example scripts that demonstrate how various aspects of CDTools work. It is recommended to read through at least a few of them before continuing to the tutorial, to get a feel for how the ptychography scripts work.
Inspect Dataset
---------------
The most important thing to be able to do is look at the raw data and confirm that it looks like you expect, hasn't been corrupted, etc. This script demonstrates how to load and look at your data, if it's stored as a compliant .cxi file.
.. literalinclude:: ../../examples/inspect_dataset.py
First, data is read into a dataset object. The dataset object knows how to generate a set of plots describing the basic information (diffraction patterns and nanomaps), which can be generated by running :code:`dataset.inspect`.
Simple Ptycho
-------------
@@ -11,9 +21,82 @@ This script runs a ptychography reconstruction using the SimplePtycho model. Thi
.. 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.
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. After the data is loaded, 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 plotted.First, data is read
Gold Balls Ptycho
-----------------
This script uses the FancyPtycho model to perform a reconstruction from the classic `gold balls <http://www.cxidb.org/id-65.html>`_ dataset. The FancyPtycho model is the workhorse model for 2D Ptychographic reconstructions, and has the option to enable correction for a variety of potential sources of error.
.. literalinclude:: ../../examples/gold_ball_ptycho.py
A number of things are changed about this script. The first difference is the section where we transfer the data and the model to the GPU. This is an important step if we would like to get a reconstruction anytime soon!
Next, note that we can liveplot the results by calling :code:`model.inspect` within the reconstruction loop. In fact, we could put any code we wanted here to look at the state of the model, and it would be executed after each epoch.
Finally, note that we take the time to save the results of this reconstruction. This is handled by pickling the dictionary produced by :code:`model.save_results`, which contains all the parameters which were reconstructed by the model.
There are a number of things going on behind the scenes as well, just because we decided to use the FancyPtycho model. This reconstruction, in addition to using two incoherently mixing modes, is also performing position annealing and reconstructing any instabilities in the probe flux that might have occured.
MIT BNL Logo
------------
Next, we look at an example showing a more involved reconstruction from a dataset that has some serious errors in it. This is a dataset we collected at the CSX beamline of NSLS-II in 2016 from a target that said "MIT BNL" on it, using a defocused zone plate.
.. literalinclude:: ../../examples/MIT_BNL_logo.py
The first new thing that pops out at us is the choice to remove part of the diffraction patterns stored in the dataset. In this case, that data is noisy, and in any case, it's masked off with the mask stored in the .cxi file. We can simply edit all the patterns like we would any other pytorch or numpy array, taking care to crop the mask as well to match.
The next thing we note is that a few more options are being taken advantage of during the model construction. First of all, we introduce the :code:`translation_scale` option, which sets the aggressiveness of the position reconstruction. A small number like 1 (the default) is often too mild when there are large errors, whereas numbers above 10 are often too agressive to lead to convergence.
In this case, the :code:`propagation_distance` option tells us how far to propagate the naive initial guess of the probe. This data was taken from a very defocused zone plate, which is a classic difficult-to-reconstruct probe. The standard probe initialization, however, is generally quite good at generating a good guess of the focal point of any zone-plate-like optics. This option lets us start with a guess of that probe, which has already been propagated from it's focal point by a known distance.
Next, note that we can start to play a bit with the reconstruction procedure in response to the kinds of error that exist in the data. In this case, we actually perform three rounds of reconstruction, in serial, with different parameters! First, we we turn off the position reconstruction, and run the reconstruction with an explicitly chosen batch size which was found to work well with this dataset. Next, we turn position reconstruction back on, and continue the reconstruction. Finally, we improve the reconstruction quality by running ten iterations at a smaller learning rate (the default is set to a sensible 0.005).
Specular Ptycho
---------------
This file demonstrates a reconstruction from an experiment in the specular reflection geometry.
.. literalinclude:: ../../examples/specular_ptycho.py
While the reconstruction itself is simple this time, a few new parameters have been set in the model definition. First, the :code:`randomize_ang` parameter was set. This seeds the object with phase noise spanning a given angular range, rather than a constant. Second, the :code:`scattering_mode` parameter is set. When set, it will override any information in the .cxi file about the sample orientation, and define the orientation based on the detector geometry.
In this case, the 'reflection' geometry knows to set the sample normal parallel to the scattering vector defined by an outgoing ray which intersects the detector at a perpendicular, and an incoming ray along the positive z-direction (as defined in the cxi specification). The other options, 'transmision', assumes that the surface normal of your sample is parallel to the incoming light ray. If this is left unset, it will use any information stored in the .cxi file, and will default to a transmission geometry if no data exists.
Ensemble Reconstruction
-----------------------
One important check when working with ptychography is performing multiple independently seeded reconstructions from the same data, to understand to what extent they converge to the same solution. This script demonstrates how to run an ensemble of reconstructions
.. literalinclude:: ../../examples/ensemble_reconstruction.py
Note that essentially all that's different here is that the reconstructions are performed within a loop, with each reconstruction added to a list. Importantly, the :code:`randomize_ang` parameter is set, so that the reconstructions are seeded with a different pattern of noise.
Ensemble Analysis
-----------------
This file inspects the results from the ensemble reconstruction
.. literalinclude:: ../../examples/ensemble_analysis.py
After loading the results and manipulating them a bit using regular python, two analysis functions from CDTools are used to glean information from the ensemble.
First is :code:`analysis.synthesize_reconstructions`. This function takes a stack of reconstructed probes and objects, corrects each reconstruction to account for the various ambiguities (phase ramps, translation offsets, etc.), and sums them into a single synthesized probe and object. Once the reconstructions are synthesized, this final synthesized probe and object can be plotted
Second is :code:`analysis.calc_consistency_prtf`. This function compares the power into various spatial frequencies between a stack of individual reconstructions and a synthesized reconstruction, in analogy with the phase retrieval transfer function. It then outputs a curve analogous to the PRTF, that shows at what spatial frequencies the reconstructions have converged consistently. This result is then plotted as well.
Much of this functionality is duplicated by the analysis script, which will perform much of the same analysis on any saved reconstruction ensemble saved in this format.
.. code-block:: bash
$ python -m CDTools.scripts.synthesize example_reconstructions/gold_balls_ensemble.pickle
+1 -1
View File
@@ -30,7 +30,7 @@ And has optional dependencies on
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.
**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.
+16 -34
View File
@@ -1,71 +1,53 @@
from __future__ import division, print_function, absolute_import
import CDTools
from CDTools.tools import cmath
from CDTools import tools
from CDTools.tools.plotting import *
from matplotlib import pyplot as plt
import pickle
from time import time
import datetime
import torch as t
import numpy as np
# This file is too large to be distributed via Github.
# Please contact Abe Levitan (alevitan@mit) if you would like access
filename = '/media/Data Bank/CSX_6_17/Processed_CXIs/79511_p.cxi'
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
# In this dataset, the edges of the patterns are too noisy and are
# masked off anyway. We can easily just remove this data instead of
# leaving it to float.
# In this dataset, the edges of the patterns are masked off anyway
# We can easily just remove this data instead of leaving it to float.
dataset.patterns = dataset.patterns[:,70:-70,70:-70]
dataset.mask = dataset.mask[70:-70,70:-70]
# This model definition includes lots of parameters.
# In this case:
# This model definition includes lots of tweaks, described below.
#
# randomize_ang defines the initial random phase noise's extent
# translations_scale defines how aggressive the position reconstruction is
# n_modes is the number of incoherent modes
# propagation_distance is the distance to propagate from the SHARP-style guess of the probe's focal spot (in this case, the value comes from knowledge of the experimental geometry).
model = CDTools.models.FancyPtycho.from_dataset(dataset,
randomize_ang = np.pi/4,
translation_scale = 4,
n_modes=2,
propagation_distance=-73e-6)
# Uncomment these to use on the CPU
# default is CPU with 32-bit floats
# Move to the GPU
model.to(device='cuda')
dataset.get_as(device='cuda')
# We can run the first phase of phase retrieval while leaving
# the probe positions fixed (whether this is good is debatable)
# model.translation_offsets.requires_grad = False
for i, loss in enumerate(model.Adam_optimize(15, dataset, batch_size=15)):
print(i,loss)
# We turn off position reconstruction for the first phase
model.translation_offsets.requires_grad = False
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15)):
model.inspect(dataset)
# And we turn it on for the second phase, as we also lower the learning rate
# model.translation_offsets.requires_grad = True
for i, loss in enumerate(model.Adam_optimize(15, dataset, batch_size=15, lr=0.0005)):
print(i,loss)
# And we turn it on for the second phase
model.translation_offsets.requires_grad = True
for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=15)):
model.inspect(dataset)
print(i,loss)
# The third phase lowers the rate further
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.00005)):
print(i,loss)
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005)):
model.inspect(dataset)
print(i,loss)
model.inspect(dataset)
model.compare(dataset)
plt.show()
+9 -18
View File
@@ -1,23 +1,12 @@
from __future__ import division, print_function, absolute_import
import numpy as np
from matplotlib import pyplot as plt
import pickle
from CDTools.tools import cmath, plotting
from CDTools.tools.analysis import *
from CDTools.tools import plotting
from CDTools.tools import analysis
#
# Note that much of this functionality is duplicated by the convenience
# script. Try running:
#
# python -m CDTools.scripts.synthesize example_reconstructions/gold_balls_ensemble.pickle
#
# Which will perform much of the same analysis on any saved reconstruction
# ensemble
#
with open('example_reconstructions/gold_balls_ensemble.pickle', 'rb') as f:
dataset = pickle.load(f)
@@ -29,12 +18,13 @@ if type(dataset) == type([]):
for key in dataset[0]}
# Now we synthesize the object using the tool from CDTools
synth_probe, synth_obj, aligned_objs = synthesize_reconstructions(
# Now we synthesize an average reconstruction
synth_probe, synth_obj, aligned_objs = analysis.synthesize_reconstructions(
dataset['probe'], dataset['obj'])
# And then we calculate the consistency PRTF from this
freqs, prtf = calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'][0])
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'][0])
# Plot the first mode in detail
plotting.plot_phase(synth_probe[0],basis=dataset['basis'][0])
@@ -44,12 +34,13 @@ plotting.plot_colorized(synth_probe[0],basis=dataset['basis'][0])
# Just plot the colorized version of the subdominant modes
plotting.plot_colorized(synth_probe[1],basis=dataset['basis'][0])
plotting.plot_colorized(synth_probe[2],basis=dataset['basis'][0])
# And now we plot the object
plotting.plot_amplitude(synth_obj,basis=dataset['basis'][0])
plotting.plot_colorized(synth_obj,basis=dataset['basis'][0])
plotting.plot_phase(synth_obj,basis=dataset['basis'][0])
# Now plot the PRTF
# Finally, plot the consistency PRTF
plt.figure()
plt.plot(freqs*1e-6, prtf)
plt.xlabel('Spatial Frequency (cycles/um)')
@@ -1,12 +1,10 @@
from __future__ import division, print_function, absolute_import
import CDTools
import numpy as np
import pickle
from matplotlib import pyplot as plt
# Load the data
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
results = []
@@ -14,23 +12,23 @@ results = []
for idx in range(25):
print('Starting Reconstruction', idx)
model = CDTools.models.FancyPtycho.from_dataset(dataset,n_modes=3,randomize_ang=0.1*np.pi)
# Create a new model each time
model = CDTools.models.FancyPtycho.from_dataset(dataset,n_modes=3,
randomize_ang=0.1*np.pi)
# default is CPU with 32-bit floats
# Work on the GPU
model.to(device='cuda')
dataset.get_as(device='cuda')
# Run the reconstruction
for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)):
print(i,loss)
# Here we see how to liveplot the results - this call will create
# or update a readout of the various parameters being reconstructed
# And add the results to the ensemble
results.append(model.save_results(dataset))
# Save out the ensemble
with open('example_reconstructions/gold_balls_ensemble.pickle', 'wb') as f:
pickle.dump(results,f)
model.inspect(dataset)
model.compare(dataset)
plt.show()
+9 -14
View File
@@ -1,36 +1,31 @@
from __future__ import division, print_function, absolute_import
import CDTools
import numpy as np
import pickle
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
# Note that we explicitly as for two incoherent probe modes
model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
# default is CPU with 32-bit floats
# Let's do this reconstruction on the GPU, shall we?
model.to(device='cuda')
dataset.get_as(device='cuda')
for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)):
print(i,loss)
# Here we see how to liveplot the results - this call will create
# or update a readout of the various parameters being reconstructed
# And we liveplot the updates to the model as they happen
model.inspect(dataset)
print(i,loss)
# And we save the reconstruction out to a file
with open('example_reconstructions/gold_balls.pickle', 'wb') as f:
pickle.dump(model.save_results(dataset),f)
# Finally, we plot the results
model.inspect(dataset)
dataset.inspect()
model.compare(dataset)
plt.show()
+12
View File
@@ -0,0 +1,12 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
# 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)
# And we take a look at the data
dataset.inspect()
plt.show()
-1
View File
@@ -12,7 +12,6 @@ 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)
# Finally, we plot the results
@@ -1,44 +1,37 @@
from __future__ import division, print_function, absolute_import
import CDTools
from CDTools.tools import cmath
from CDTools import tools
from CDTools.tools.plotting import *
from matplotlib import pyplot as plt
import pickle
from time import time
import datetime
import h5py
import torch as t
import numpy as np
# This file is too large to be distributed via Github.
# Please contact Abe Levitan (alevitan@mit) if you would like access
filename = '/media/Data Bank/CSX_10_18/Processed_CXIs/110531_p.cxi'
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
# This model definition includes lots of tweaks, described below.
#
# randomize_ang defines the initial random phase noise's extent
# translations_scale defines how aggressive the position reconstruction is
# scattering_mode overrides any sample normal information stored in the .cxi file
model = CDTools.models.FancyPtycho.from_dataset(dataset,
randomize_ang = np.pi/4,
padding=0,
translation_scale=10,
scattering_mode='reflection')
# Uncomment these to use on the CPU
# default is CPU with 32-bit floats
# Move to the GPU
model.to(device='cuda')
dataset.get_as(device='cuda')
# Run the reconstruction
for i, loss in enumerate(model.Adam_optimize(250, dataset,batch_size=5)):
print(i,loss)
model.inspect(dataset)
model.inspect(dataset)
dataset.inspect()
model.compare(dataset)
plt.show()