cdtools/examples/fancy_ptycho_interactive.ipynb

In [ ]:

%matplotlib widget
import cdtools
import torch as t
from matplotlib import pyplot as plt

In [ ]:

# Load and inspect a dataset

filename = 'example_data/lab_ptycho_data.cxi'
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)

dataset.inspect();

In [ ]:

# Initialize a model from the dataset and move it to the GPU.

model = cdtools.models.FancyPtycho.from_dataset(
    dataset,
    n_modes=3, # Use 3 incoherently mixing probe modes
    oversampling=2, # Simulate the probe on a 2xlarger real-space array
    probe_support_radius=120, # Force the probe to 0 outside a radius of 120 pix
    propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm
    units='mm', # Set the units for the live plots
    obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix,
)

if t.cuda.is_available():
    model.to(device='cuda')
    dataset.get_as(device='cuda')

# Then, create a reconstructor object and view the initialized model

recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
model.inspect(dataset);

In [ ]:

# Workaround reconstruction pattern for interactive plotting in jupyter:

# With this pattern, it is safe to interrupt the kernel. Doing so will
# trigger an update of the plots, at which point the current state can
# be viewed. Then this cell can be re-run to continue the reconstruction
while model.epoch < 50:
    for loss in recon.optimize(1, lr=0.02, batch_size=10):
        print(model.report())
        model.inspect(dataset, min_interval=10)

while model.epoch < 100:
    for loss in recon.optimize(1, lr=0.005, batch_size=10):
        print(model.report())
        model.inspect(dataset, min_interval=10)

In [ ]:

# Save out the results

model.save_to_h5('lab_ptycho_reconstruction.h5', dataset)

In [ ]:

# Finalize the plotting and create the comparison plot

# This orthogonalizes the recovered probe modes. It is best to do so
# after saving the results, if you intend to initialize any further
# reconstructions with the probe.
model.tidy_probes()

# Final plotting
model.inspect(dataset)
model.compare(dataset);

In [ ]: