Notebook: Ptychography reconstruction with positions update from a CXI dataset#

This notebook uses the the runner API (pynx.ptycho.runner) to load the data and prepare the optimisation.

This is what is used for the command-line scripts, here we just grap the Ptycho object from the runner.

This uses a dataset recorded on id01@ESRF, which exhibits some position distortions near the ceneter of the spiral scan.

[1]:
# Select language and/or GPU name or rank through environment variable (optional)
#import os
#os.environ['PYNX_PU'] = 'cuda.0'

%matplotlib widget
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 *

# Import CXI runner
from pynx.ptycho.runner.cxi import PtychoRunnerScanCXI, default_params as params

# This can be used to have a wide screen for the notebook
#from IPython.core.display import display, HTML
#display(HTML("<style>.container { width:90% !important; }</style>"))

Load CXI data if necessary#

this is available from the PyNX ESRF public folder

[2]:
if not os.path.exists('ptycho-siemens-star-id01.cxi'):
    os.system('curl -O http://ftp.esrf.fr/pub/scisoft/PyNX/data/ptycho-siemens-star-id01.cxi')

Load the data & setup the runner parameters#

As we are using the pynx.ptycho.runner API, we can setup the parameters exactly as for the command-line scripts.

[3]:
params['cxifile']='ptycho-siemens-star-id01.cxi'

# Initial probe description
params['probe']='gauss,150e-9x200e-9'  # Starting from a simulated probe, here described as gaussian
params['defocus'] = 200e-6

# Initial object
params['object'] = 'random,0.8,1,0,0.5'

# Update plots and LLK metrics every N cycles (10 is wasteful, slows down a lot calculations, but adequate for a tutorial)
params['verbose'] = 10

# It would be also possible to supply an algorithm e.g. "ML**50,DM**100,probe=1"
# Here "manual" means the runner will not perform any optimisation when ws.run() is executed
params['algorithm'] = "manual"

# Limit the total number of frames imported (faster
params['maxframe'] = 200
# params['moduloframe'] = 2

ws = PtychoRunnerScanCXI(params, 0)
ws.load_scan()
ws.load_data()  # Load all frames from a maxipix detector using CXI/HDF5 data
ws.prepare_processing_unit()

ws.center_crop_data()  # Auto-crop the data
ws.prepare()  # This will prepare the initial object

MAXFRAME: only using first 200 frames
CXI: read scan number=13
Reading 200 frames from CXI-HDF5 file: 0.20.40.60.80.100.120.140.160.180.200
Time to read all frames:  0.6s [84.38 Mpixel/s]
Loaded mask from CXI data: /entry_1/instrument_1/detector_1/mask
Initialized 2D mask with 6156 ( 2.312%) bad pixels
Ptycho runner: preparing processing unit
Searching available OpenCL GPU [ranking by global memory bandwidth]:
                                        Apple M1 Pro [Apple]:   10 Gb,   574 Gbytes/s
Using OpenCL GPU: Apple M1 Pro
Center of diffraction: X=206.75 Y=209.19
Largest prime number acceptable for FFT size: 13
Interpolating data for 4764 (masked) pixels
Interpolating data for 4764 (masked) pixels - finished
Final iobs data size after cropping / centering / rebinning: (200, 400, 400)
E= 7.99071keV, zdetector= 1.386m, pixel size= 55.00um, pixel size(object)=   9.8nm
Using random object type with amplitude range:  0.80- 1.00 and phase range:  0.00- 0.50

Setup optimisation#

We set the number of probe modes to 3.

Sub-pixel interpolation can also be activated but does not change significantly the final result.

The Ptycho object can be accessed as ws.p

[4]:
# Initial optimisation
ws.run()
ws.run_algorithm('nbprobe=3')

# Get Ptycho object
p = ws.p

# Enable bilinear interpolation (lowers resolution, but smoother position update)
p._interpolation = False

# Keep copies of object and probe to restore later
obj = p.get_obj()
probe = p.get_probe()

 ####################################################################################################
#
# Scan:  13 Run: 1
#
 ####################################################################################################
Simulating object: random
Making obj: (592, 596) 588 600
Simulating probe: gauss
ScaleObjProbe: 1.3848974 317549.97 196573.2503602902 212304.6260141112 0.9259018610595349

 ####################################################################################################
#
# Scan:  13 Run: 1 , Algorithm: nbprobe=3
#
 ####################################################################################################

Total elapsed time for algorithms:     0.02s

 ####################################################################################################

Probe statistics at sample position:
  FWHM (peak intensity):   167.49nm(H) x  171.61nm(V)
  FW @20% intensity    :   255.33nm(H) x  261.31nm(V)
  FWHM (statistical)   :   167.22nm(H) x  171.28nm(V), average=  169.24nm

Optimise object and probe#

[5]:
fig = plt.figure()  # Open the figure in the cell before the computing to see live changes
[6]:
p = DM(update_object=True, update_probe=True, calc_llk=20, show_obj_probe=20)**100 * p
p = ML(update_object=True, update_probe=True, calc_llk=20, show_obj_probe=20)**100 * p
DM/o/3p    #  0 LLK=   487.80(p)  1687.94(g)   542.84(e), nb photons=8.866051e+09, dt/cycle= 7.4600s
DM/o/3p    # 20 LLK=     8.26(p)    31.32(g)    13.80(e), nb photons=8.322520e+09, dt/cycle= 0.1224s
DM/o/3p    # 40 LLK=     7.39(p)    24.67(g)    12.22(e), nb photons=8.297144e+09, dt/cycle= 0.1215s
DM/o/3p    # 60 LLK=     8.39(p)    15.39(g)    13.57(e), nb photons=8.282478e+09, dt/cycle= 0.1235s
DM/o/3p    # 80 LLK=     7.42(p)    17.86(g)    12.19(e), nb photons=8.304001e+09, dt/cycle= 0.1250s
DM/o/3p    # 99 LLK=     6.70(p)    30.49(g)    11.09(e), nb photons=8.332631e+09, dt/cycle= 0.1279s
ML/o/3p    #100 LLK=     6.17(p)    31.09(g)    10.35(e), nb photons=8.856857e+09, dt/cycle= 0.8669s
ML/o/3p    #120 LLK=     3.82(p)    35.20(g)     7.75(e), nb photons=8.879905e+09, dt/cycle= 0.2586s
ML/o/3p    #140 LLK=     3.39(p)    32.00(g)     7.02(e), nb photons=8.879475e+09, dt/cycle= 0.2605s
ML/o/3p    #160 LLK=     3.18(p)    29.94(g)     6.70(e), nb photons=8.860971e+09, dt/cycle= 0.2639s
ML/o/3p    #180 LLK=     3.08(p)    28.49(g)     6.59(e), nb photons=8.848262e+09, dt/cycle= 0.2640s
ML/o/3p    #199 LLK=     3.05(p)    28.20(g)     6.63(e), nb photons=8.859660e+09, dt/cycle= 0.2667s

Optimise object, probe and positions#

This uses the AP and ML algorithms, which are more stable than DM for positions optimisations

Note the final improvement of the object near the center.

We use pos_history=True which will be used to plot the history of positions at the end, but this slows down the calculation as the values have to be grabbed from the GPU for each update. It’s only useful to study how the algorithm works, in production it should not be used.

You can play with different parameters to tweak the position update and see how it affects the convergence: * update_pos=5: update positions every N cycle. It’s better for object and probe convergence to not update positions every cycle * pos_mult=5: amplify the positions update. This speeds up the convergence, but can lead to unstabilities * pos_threshold=0.2: this inhibits the positions update when the local object gradient *probe is smaller than 0.2 times the average gradient for all positions. This prevents updating positions in areas of the object which are too smooth * min_shift: don’t update shifts when change is lower than a certain number of pixels (value has to be small as this is changes for a given cycle) * max_shift: prevent too large changes per cycle. Should not be necessary, but can be safe

[7]:
fig = plt.figure()
[8]:
p = AP(update_object=True, update_probe=True, update_pos=5, pos_mult=5,
       pos_history=True, calc_llk=20, show_obj_probe=40)**200 * p
p = ML(update_object=True, update_probe=True, update_pos=5, pos_mult=5,
       pos_history=True, calc_llk=10, show_obj_probe=20)**100 * p

AP/o/3p/t  #200 LLK=     3.05(p)    28.23(g)     6.63(e), nb photons=8.860580e+09, dt/cycle=13.9576s
AP/o/3p/t  #220 LLK=     2.68(p)    26.16(g)     5.28(e), nb photons=8.828312e+09, dt/cycle= 0.1158s
AP/o/3p/t  #240 LLK=     2.47(p)    26.84(g)     4.88(e), nb photons=8.822498e+09, dt/cycle= 0.1058s
AP/o/3p/t  #260 LLK=     2.36(p)    27.77(g)     4.67(e), nb photons=8.821308e+09, dt/cycle= 0.1169s
AP/o/3p/t  #280 LLK=     2.30(p)    28.02(g)     4.55(e), nb photons=8.821497e+09, dt/cycle= 0.1059s
AP/o/3p/t  #300 LLK=     2.27(p)    28.77(g)     4.49(e), nb photons=8.823389e+09, dt/cycle= 0.1185s
AP/o/3p/t  #320 LLK=     2.23(p)    28.56(g)     4.42(e), nb photons=8.824457e+09, dt/cycle= 0.1058s
AP/o/3p/t  #340 LLK=     2.21(p)    28.42(g)     4.37(e), nb photons=8.824548e+09, dt/cycle= 0.1198s
AP/o/3p/t  #360 LLK=     2.21(p)    28.30(g)     4.37(e), nb photons=8.823866e+09, dt/cycle= 0.1060s
AP/o/3p/t  #380 LLK=     2.20(p)    28.85(g)     4.36(e), nb photons=8.823943e+09, dt/cycle= 0.1213s
AP/o/3p/t  #399 LLK=     2.22(p)    28.69(g)     4.38(e), nb photons=8.820827e+09, dt/cycle= 0.1024s
ML/o/3p/t  #400 LLK=     2.22(p)    28.78(g)     4.39(e), nb photons=8.856237e+09, dt/cycle= 0.9523s
ML/o/3p/t  #410 LLK=     2.24(p)    30.00(g)     4.65(e), nb photons=8.857693e+09, dt/cycle= 0.3071s
ML/o/3p/t  #420 LLK=     2.24(p)    29.10(g)     4.68(e), nb photons=8.853694e+09, dt/cycle= 0.2740s
ML/o/3p/t  #430 LLK=     2.22(p)    29.80(g)     4.63(e), nb photons=8.852134e+09, dt/cycle= 0.3160s
ML/o/3p/t  #440 LLK=     2.23(p)    30.19(g)     4.66(e), nb photons=8.854716e+09, dt/cycle= 0.2726s
ML/o/3p/t  #450 LLK=     2.21(p)    31.00(g)     4.64(e), nb photons=8.859780e+09, dt/cycle= 0.3192s
ML/o/3p/t  #460 LLK=     2.23(p)    31.83(g)     4.68(e), nb photons=8.861436e+09, dt/cycle= 0.2724s
ML/o/3p/t  #470 LLK=     2.21(p)    30.81(g)     4.64(e), nb photons=8.861177e+09, dt/cycle= 0.3167s
ML/o/3p/t  #480 LLK=     2.24(p)    31.27(g)     4.71(e), nb photons=8.862857e+09, dt/cycle= 0.2750s
ML/o/3p/t  #490 LLK=     2.21(p)    30.55(g)     4.66(e), nb photons=8.861356e+09, dt/cycle= 0.3192s
ML/o/3p/t  #499 LLK=     2.16(p)    31.35(g)     4.56(e), nb photons=8.848157e+09, dt/cycle= 0.2660s

Plot the position shifts#

This standard plot is also created when using the analysis step in the algorithm with the command-line scripts.

Shifts are represented both with arrows (scaled if necessary) and a heat map relative to the maximum shift (the color indicates the direction of the shift, the intensity the relative amplitude).

In this case the piezo motors had some trouble in the center, either because of the slow start or the quick change of direction.

[9]:
p = PlotPositions(show_plot=True, save_prefix=None, fig_size=(10,4)) * p
Average distance between points: 0.091490 µm
Shift in positions: mean=0.021617 µm, max=0.148069µm
Scale: 0.6178886185770237

Plot the history of position shifts near the center#

We recorded the position history, so plot them for the center positions vs the cycle #.

The position update parameters (e.g. pos_mult) can be modified to see how that affects convergence

[10]:
ipos = range(15)
fig = plt.figure(figsize=(9.5,6))
#ax = plt.axes((0.04,0.3,0.22,0.3))
plt.subplot(441)
#for i in range(default_processing_unit.get_stack_size()):
for i in range(20):
    x = [v[1] for v in p.position_history[i]]
    y = [v[2] for v in p.position_history[i]]
    plt.scatter(x,y, 1)
    plt.text(x[0], y[0], '%d' % i)
    #print("%3d  dr = %5.3f" % (i, np.sqrt((x[0]-x[-1])**2 + (y[0]-y[-1])**2)))
plt.gca().set_aspect(1)

for i in range(len(ipos)):
    #plt.subplot(4,4,2 + i // 3 * 4 + i % 3)
    plt.subplot(4,4,i+2)
    ix, x, y = [v[0] for v in p.position_history[ipos[i]]], [v[1] for v in p.position_history[ipos[i]]], \
               [v[2] for v in p.position_history[ipos[i]]]
    plt.plot(ix,x,'b.', label='x[%d]'%ipos[i])
    plt.twinx()
    plt.plot(ix,y,'r.', label='y[%d]'%ipos[i])
    plt.text(0.46, 1.03,'x', color='blue', alpha=0.9,transform=plt.gca().transAxes)
    plt.text(0.54, 1.03,'y', color='red', alpha=0.9,transform=plt.gca().transAxes)
    plt.text(0.62, 1.03,'[%d]'%ipos[i], alpha=0.9,transform=plt.gca().transAxes)
    plt.text(0.9, -.14, 'cycle',transform=plt.gca().transAxes)
plt.tight_layout()

Free GPU memory#

It is good practice in a notebook to use FreePU() to release memory from the GPU

[11]:
p = FreePU() * p
[ ]: