mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-18 16:42:09 +02:00
3.9 KiB
3.9 KiB
In [ ]:
%matplotlib widget
import cdtools
import torch as t
from matplotlib import pyplot as pltIn [ ]:
# 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();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(min_interval=10)
while model.epoch < 100:
for loss in recon.optimize(1, lr=0.005, batch_size=10):
print(model.report())
model.inspect(min_interval=10)In [ ]:
# Save out the results
model.save_to_h5('lab_ptycho_reconstruction.h5')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()
model.compare(dataset);In [ ]: