Modified Initial Probe

Learning how to modify the simulated initial probe in the source plane.

This example is using the same data and parameters as in Loading HDF5 Data and shows how it is possible in PtyPy to set up a reconstruction, execute with level=4, then modify the model of the initial probe and run the reconstruction via P.run() using the newly generated initial probe

---

import ptypy, os
import ptypy.utils as u

# This will import the HDF5Loader class
ptypy.load_ptyscan_module("hdf5_loader")

# Root directory of tutorial data
tutorial_data_home = "../../data/"

# Dataset for this tutorial
dataset = "small_data/dls_i08_nanogold_small.h5"

# Absolute path to HDF5 file with raw data
path_to_data = os.path.join(tutorial_data_home, dataset)

# Create parameter tree
p = u.Param()

# Set verbose level to interactive
p.verbose_level = "interactive"

# Set io settings (no files saved)
p.io = u.Param()
p.io.rfile = None
p.io.autosave = u.Param(active=False)
p.io.interaction = u.Param(active=False)

# Live-plotting during the reconstruction
p.io.autoplot = u.Param()
p.io.autoplot.active=True
p.io.autoplot.threaded = False
p.io.autoplot.layout = "jupyter"
p.io.autoplot.interval = 10

# Define the scan model
p.scans = u.Param()
p.scans.scan_00 = u.Param()
p.scans.scan_00.name = 'Full'

# Initial illumination (based on simulated optics)
p.scans.scan_00.illumination = u.Param()
p.scans.scan_00.illumination.model = None
p.scans.scan_00.illumination.photons = None
p.scans.scan_00.illumination.aperture = u.Param()
p.scans.scan_00.illumination.aperture.form = "circ"
p.scans.scan_00.illumination.aperture.size = 333e-6
p.scans.scan_00.illumination.propagation = u.Param()
p.scans.scan_00.illumination.propagation.focussed = 13.725e-3
p.scans.scan_00.illumination.propagation.parallel = 45e-6

# Data loader
p.scans.scan_00.data = u.Param()
p.scans.scan_00.data.name = 'Hdf5Loader'

# Read diffraction data
p.scans.scan_00.data.intensities = u.Param()
p.scans.scan_00.data.intensities.file = path_to_data
p.scans.scan_00.data.intensities.key = "data"

# Read positions data
p.scans.scan_00.data.positions = u.Param()
p.scans.scan_00.data.positions.file = path_to_data
p.scans.scan_00.data.positions.slow_key = "posy_mm"
p.scans.scan_00.data.positions.slow_multiplier = 1e-3
p.scans.scan_00.data.positions.fast_key = "posx_mm"
p.scans.scan_00.data.positions.fast_multiplier = 1e-3

# Read meta data: photon energy
p.scans.scan_00.data.recorded_energy = u.Param()
p.scans.scan_00.data.recorded_energy.file = path_to_data
p.scans.scan_00.data.recorded_energy.key = "energy_ev"
p.scans.scan_00.data.recorded_energy.multiplier = 1e-3

# Read meta data: detector distance
p.scans.scan_00.data.recorded_distance = u.Param()
p.scans.scan_00.data.recorded_distance.file = path_to_data
p.scans.scan_00.data.recorded_distance.key = "det_distance_mm"
p.scans.scan_00.data.recorded_distance.multiplier = 1e-3

# Read meta data: detector pixelsize
p.scans.scan_00.data.recorded_psize = u.Param()
p.scans.scan_00.data.recorded_psize.file = path_to_data
p.scans.scan_00.data.recorded_psize.key = "det_pixelsize_um"
p.scans.scan_00.data.recorded_psize.multiplier = 1e-6

# Define reconstruction engine (using DM)
p.engines = u.Param()
p.engines.engine = u.Param()
p.engines.engine.name = "DM"
p.engines.engine.numiter = 200
p.engines.engine.numiter_contiguous = 10
p.engines.engine.alpha = 0.9
p.engines.engine.probe_support = None
p.engines.engine.probe_update_start = 0

# Run reconstruction
P = ptypy.core.Ptycho(p,level=4)

from ptypy.core.illumination import _propagation, illumination_desc, aperture
import numpy as np

# gather parameters
prop_pars = illumination_desc.make_default(99).propagation
prop_pars.update(p.scans.scan_00.illumination.propagation)
aperture_pars = p.scans.scan_00.illumination.aperture

# get current probe model
prb_storage = P.probe.S["Sscan_00G00"]
prb_initial = np.copy(prb_storage.data[0])

# Create new model in source plane
ap_size, grids, prop = _propagation(prop_pars, prb_storage.shape[1:], prb_storage.psize, prb_storage.views[0].pod.geometry.energy)
model_source = aperture(np.zeros(prb_storage.shape[1:]), grids, aperture_pars)

# modify model in source place
model_source[60:] = 0

# propagate model into sample plane
model = prop(model_source)

# rescale probe power
model *= np.sqrt(P.diff.S["S0000"].max_power / u.norm2(model))

# Store probe 
P.probe.S["Sscan_00G00"].data[0] = model

import matplotlib.pyplot as plt
fig, axes = plt.subplots(ncols=5, figsize=(12,3))
axes[0].set_title("Original simulated probe")
axes[0].imshow(np.abs(prb_initial))
axes[1].set_title("Initial aperture")
axes[1].imshow(aperture(np.zeros(prb_storage.shape[1:]), grids, aperture_pars))
axes[2].set_title("Modified aperture")
axes[2].imshow(model_source)
axes[3].set_title("New simulated probe")
axes[3].imshow(np.abs(model))
axes[4].set_title("Probe in detector plane")
axes[4].imshow(np.abs(prb_storage.views[0].pod.geometry.propagator.fw(model)))
plt.show()

# Run the reconstruction using the modified initial probe
P.run()