Notebook example of ptychograpic reconstruction on simulated data¶
This is a basic example, showing how to: * simulate data * compute the object shape from the data * use the Ptycho operators for reconstruction
[1]:
# Optional: select language and/or GPU name or rank through environment variable
#import os
#os.environ['PYNX_PU'] = 'cuda'
%matplotlib notebook
import matplotlib.pyplot as plt
from pynx.ptycho import simulation, shape
# Import Ptycho, PtychoData and operators (automatically selecting OpenCL or CUDA)
from pynx.ptycho import *
Simulate the Ptychography dataset¶
[2]:
# 2D detector size (square)
nxy = 256
# Pixel size in meters
pixel_size_detector = 55e-6
# Wavelength in meters
wavelength = 1.5e-10
# Detector distance in meters
detector_distance = 1
# Object options 'siemens' simulates Siemens star (with a few holes)
# 'logo' simulates PyNX logo
# obj_info = {'type': 'logo', 'phase_stretch': 1.57, 'alpha_win': .2}
obj_info = {'type': 'siemens', 'phase_stretch': 1.57, 'alpha_win': .2}
# Probe description, either as a Gaussian, or as a focused aperture
probe_info = {'type': 'focus', 'aperture': (150e-6, 150e-6), 'focal_length': .08,
'defocus': 350e-6, 'shape': (nxy, nxy)}
# probe_info = {'type': 'gauss', 'sigma_pix': (20, 20), 'shape': (nxy, nxy)}
# Spiral scan: 50 positions = 4 turns, 78 = 5 turns, 113 = 6 turns
scan_info = {'type': 'spiral', 'scan_step_pix': 20, 'n_scans': 200}
if False:
# Use last frame without sample (direct beam ; serves as an absolute reference)
s = simulation.Simulation(obj_info=None, probe_info=None, scan_info=scan_info, data_info=None, verbose=False)
s.make_scan()
posx, posy = s.scan.values
posx[-1] = 1e20
posy[-1] = 1e20
posx = np.ma.masked_array(posx, posx>=1e10)
posy = np.ma.masked_array(posy, posy>=1e10)
scan_info = {'type': 'custom', 'x': posx, 'y': posy}
# Data info, with the different parameters and using Poisson noise
# nb_photons_per_frame is the average number of photons per frame
data_info = {'nb_photons_per_frame': 1e9, 'bg': 0, 'wavelength': wavelength,
'detector_distance': detector_distance,
'detector_pixel_size': pixel_size_detector,
'noise': 'poisson'}
# Initialisation of the simulation
s = simulation.Simulation(obj_info=obj_info, probe_info=probe_info, scan_info=scan_info, data_info=data_info)
s.make_data()
posx, posy = s.scan.values
pixel_size_object = wavelength * detector_distance / pixel_size_detector / nxy
ampl = s.amplitude.values # square root of the measured diffraction pattern intensity
Simulating object: siemens
Simulating probe: focus
Simulating scan: spiral
Simulating ptychographic data [200 frames].
Parameters of the simulation:
Data info: {'pix_size_direct_nm': 10, 'num_phot_max': None, 'nb_photons_per_frame': 1000000000.0, 'bg': 0, 'beam_stop_transparency': 0, 'noise': 'poisson', 'wavelength': 1.5e-10, 'detector_distance': 1, 'detector_pixel_size': 5.5e-05}
Scan info: {'type': 'spiral', 'scan_step_pix': 20, 'n_scans': 200, 'integer_values': False}
Object info: {'type': 'Custom', 'phase_stretch': 1.57, 'alpha_win': 0.2}
Probe info: {'type': 'focus', 'shape': (256, 256), 'sigma_pix': (50, 50), 'rotation': 0, 'aperture': (0.00015, 0.00015), 'focal_length': 0.08, 'defocus': 0.00035}
Prepare the initial object and probe¶
This uses the pynx.ptycho.simulation module for an explicit simulation of object and probe.
Note that if the initial object array is not supplied to the Ptychoobject (obj=None), its size will be automatically calculated, and the object initialised to an homogeneous object (array of 1)
[3]:
# Size of the reconstructed object
nyo, nxo = shape.calc_obj_shape(posx, posy, ampl.shape[1:])
# Initial object
obj_init_info = {'type': 'random', 'range': (0.9, 1, 0, 0.5), 'shape': (nyo, nxo)}
# Initial probe
probe_init_info = {'type': 'focus', 'aperture': (150e-6, 150e-6), 'focal_length': .08,
'defocus': 250e-6, 'shape': (nxy, nxy)}
# Basic data info, used to compute the object pixel size
data_info = {'wavelength': wavelength, 'detector_distance': detector_distance,
'detector_pixel_size': pixel_size_detector}
# Perform the actual simulation
init = simulation.Simulation(obj_info=obj_init_info, probe_info=probe_init_info, data_info=data_info)
init.make_obj()
init.make_probe()
Simulating object: random
Simulating probe: focus
Create the PtychoData and Ptycho objects¶
[4]:
data = PtychoData(iobs=ampl ** 2, positions=(posx * pixel_size_object, posy * pixel_size_object),
detector_distance=1, mask=None, pixel_size_detector=55e-6, wavelength=1.5e-10)
# Random object start
p = Ptycho(probe=init.probe.values, obj=init.obj.values, data=data, background=None)
# Initial scaling of object and probe
p = ScaleObjProbe(verbose=True) * p
ScaleObjProbe: 3191.7012 328421.5 3428.621194668073 593.9698936479831 5.772382035549923
Optimise the Ptycho object¶
This can use different algorithms: * Difference Map * Alternating Projections * Maximum Likelihood conjugate gradient
For each algorithm it is possible to update object, probe, positions, and to display the result.
Each algorithm operator is elevated to the the number of cycles, e.g. DM()**40 will perform 40 cycles
[5]:
plt.figure()
p = DM(update_object=True, update_probe=True, calc_llk=20, show_obj_probe=20)**100 * p
#p = ShowObjProbe() * AP(update_object=True, update_probe=False, calc_llk=10)**40 * p
p = ML(update_object=True, update_probe=True, calc_llk=20, show_obj_probe=20)**100 * p
DM/o/p # 0 LLK= 44435.33(p) 208742.22(g) 69049.22(e), nb photons=2.332854e+12, dt/cycle=0.156s
DM/o/p # 20 LLK= 946.86(p) 1065.00(g) 1391.22(e), nb photons=2.584728e+12, dt/cycle=0.035s
DM/o/p # 40 LLK= 11.21(p) 22.30(g) 22.28(e), nb photons=2.600215e+12, dt/cycle=0.036s
DM/o/p # 60 LLK= 3.42(p) 6.81(g) 6.82(e), nb photons=2.600465e+12, dt/cycle=0.037s
DM/o/p # 80 LLK= 2.09(p) 4.15(g) 4.17(e), nb photons=2.600445e+12, dt/cycle=0.039s
DM/o/p # 99 LLK= 1.30(p) 2.56(g) 2.58(e), nb photons=2.600384e+12, dt/cycle=0.037s
ML/o/p #101 LLK= 1.15(p) 2.30(g) 2.30(e), nb photons=2.600946e+12, dt/cycle=0.137s
ML/o/p #121 LLK= 1.10(p) 2.19(g) 2.19(e), nb photons=2.600728e+12, dt/cycle=0.056s
ML/o/p #141 LLK= 1.09(p) 2.17(g) 2.17(e), nb photons=2.600670e+12, dt/cycle=0.057s
ML/o/p #161 LLK= 1.07(p) 2.14(g) 2.13(e), nb photons=2.600712e+12, dt/cycle=0.056s
ML/o/p #181 LLK= 1.05(p) 2.10(g) 2.09(e), nb photons=2.600689e+12, dt/cycle=0.056s
ML/o/p #200 LLK= 1.02(p) 2.05(g) 2.04(e), nb photons=2.600697e+12, dt/cycle=0.058s
Use DM and ML options to smooth the object and/or probe¶
For DM or AP, use the *_smooth_sigma and *_inertia parameters For ML, use the reg_fac_* parameters to
[6]:
plt.figure() # Create a new figure instead of using the one in the previous cell
p = DM(update_object=True, update_probe=True, calc_llk=20, show_obj_probe=20,
obj_smooth_sigma=0.1, obj_inertia=0.01, probe_smooth_sigma=0.05, probe_inertia=0.001)**40 * p
p = ML(update_object=True, update_probe=True, show_obj_probe=20,
calc_llk=20, reg_fac_obj=0.01, reg_fac_probe=0)**100 * p
DM/o/p #201 LLK= 0.90(p) 1.80(g) 1.80(e), nb photons=2.600667e+12, dt/cycle=0.242s
DM/o/p #221 LLK= 0.62(p) 1.24(g) 1.24(e), nb photons=2.600650e+12, dt/cycle=0.040s
DM/o/p #240 LLK= 0.61(p) 1.21(g) 1.22(e), nb photons=2.600595e+12, dt/cycle=0.041s
ML/o/p #242 LLK= 0.58(p) 1.14(g) 1.15(e), nb photons=2.600743e+12, dt/cycle=0.134s
ML/o/p #262 LLK= 0.57(p) 1.13(g) 1.14(e), nb photons=2.600670e+12, dt/cycle=0.057s
ML/o/p #282 LLK= 0.57(p) 1.13(g) 1.14(e), nb photons=2.600676e+12, dt/cycle=0.058s
ML/o/p #302 LLK= 0.57(p) 1.13(g) 1.14(e), nb photons=2.600676e+12, dt/cycle=0.057s
ML/o/p #322 LLK= 0.57(p) 1.13(g) 1.14(e), nb photons=2.600676e+12, dt/cycle=0.057s
ML/o/p #341 LLK= 0.57(p) 1.13(g) 1.14(e), nb photons=2.600676e+12, dt/cycle=0.059s
Add probe modes and continue optimising¶
The DM/o/3p indicates: * the algorithm (DM or AP or ML) * the parts which are optimised (o for object, p for probe, t for translations) * the number of modes (when >1)
(in this case with simulated data, the additional probe modes are not useful)
[7]:
pr = p.get_probe()
nb_probe, ny, nx = pr.shape
# New number of probe modes
nb_probe = 3
pr1 = np.empty((nb_probe, ny, nx), dtype=np.complex64)
pr1[0] = pr[0]
for i in range(1, nb_probe):
n = abs(pr).mean() / 10
pr1[i] = np.random.uniform(0, n, (ny, nx)) * np.exp(1j * np.random.uniform(0,2*np.pi, (ny,nx)))
p.set_probe(pr1)
plt.figure()
p = DM(update_object=True, update_probe=True, calc_llk=20, show_obj_probe=20)**40 * p
p = ML(update_object=True, update_probe=True, calc_llk=20, show_obj_probe=20)**100 * p
DM/o/3p #342 LLK= 14.06(p) 108.07(g) 32.55(e), nb photons=2.600547e+12, dt/cycle=0.293s
DM/o/3p #362 LLK= 0.52(p) 1.04(g) 1.04(e), nb photons=2.600664e+12, dt/cycle=0.066s
DM/o/3p #381 LLK= 0.56(p) 1.10(g) 1.11(e), nb photons=2.600607e+12, dt/cycle=0.059s
ML/o/3p #383 LLK= 0.53(p) 1.05(g) 1.06(e), nb photons=2.600730e+12, dt/cycle=0.170s
ML/o/3p #403 LLK= 0.53(p) 1.05(g) 1.05(e), nb photons=2.600681e+12, dt/cycle=0.101s
ML/o/3p #423 LLK= 0.53(p) 1.05(g) 1.05(e), nb photons=2.600675e+12, dt/cycle=0.101s
ML/o/3p #443 LLK= 0.53(p) 1.05(g) 1.05(e), nb photons=2.600675e+12, dt/cycle=0.101s
ML/o/3p #463 LLK= 0.53(p) 1.05(g) 1.05(e), nb photons=2.600675e+12, dt/cycle=0.101s
ML/o/3p #482 LLK= 0.53(p) 1.05(g) 1.05(e), nb photons=2.600675e+12, dt/cycle=0.102s
[8]:
# Manual decompositon of algorithms
#p = Psi2Obj() * PropagateApplyAmplitude()* ObjProbe2Psi() * SelectStack(0) * p
#p = Psi2ObjMerge() * LoopStack(Psi2Obj() * PropagateApplyAmplitude() * ObjProbe2Psi()) * p
Export data and/or result object & probe to CXI (hdf5) files¶
[9]:
if False:
#
p.save_obj_probe_cxi('obj_probe.cxi')
save_ptycho_data_cxi('data.cxi', ampl ** 2, pixel_size_detector, wavelength, detector_distance,
posx * pixel_size_object, posy * pixel_size_object, z=None, monitor=None,
mask=None, instrument='simulation', overwrite=True)