From 3c8de5dc192072fdc5e0753d00183e3c31cca41f Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 16 Oct 2025 14:16:12 +0200 Subject: [PATCH] Update the example codes --- examples/fancy_ptycho.py | 20 ++++++++++++++++---- examples/gold_ball_ptycho.py | 11 +++++++---- examples/gold_ball_split.py | 9 ++++++--- examples/simple_ptycho.py | 11 ++++++++++- 4 files changed, 39 insertions(+), 12 deletions(-) diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index 1cf5855..067f78d 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -19,19 +19,31 @@ device = 'cuda' model.to(device=device) dataset.get_as(device=device) +# Now, to do the reconstruction we will use the more flexible pattern of +# creating an explicit reconstructor. This is what is used behind the +# scenes by the convenience functions model._optimize. For a long +# reconstruction script with multiple steps, it is better to create the +# reconstructor explicitly. +# +# The reconstructor will store the model and dataset and create an appropriate +# optimizer. This allows the optimizer to persist, along with e.g. estimates +# of the moments of individual parameters between loops +recon = cdtools.reconstructors.AdamReconstructor(model, dataset) + # The learning rate parameter sets the alpha for Adam. # The beta parameters are (0.9, 0.999) by default # The batch size sets the minibatch size -for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10): +for loss in recon.optimize(50, lr=0.02, batch_size=10): print(model.report()) # Plotting is expensive, so we only do it every tenth epoch if model.epoch % 10 == 0: model.inspect(dataset) # It's common to chain several different reconstruction loops. Here, we -# started with an aggressive refinement to find the probe, and now we -# polish the reconstruction with a lower learning rate and larger minibatch -for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50): +# started with an aggressive refinement to find the probe in the previous +# loop, and now we polish the reconstruction with a lower learning rate +# and larger minibatch +for loss in recon.optimize(50, lr=0.005, batch_size=50): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 4747666..4971975 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -29,6 +29,7 @@ model = cdtools.models.FancyPtycho.from_dataset( probe_fourier_crop=pad ) + # This is a trick that my grandmother taught me, to combat the raster grid # pathology: we randomze the our initial guess of the probe positions. # The units here are pixels in the object array. @@ -42,17 +43,20 @@ device = 'cuda' model.to(device=device) dataset.get_as(device=device) +# Create the reconstructor +recon = cdtools.reconstructors.AdamReconstructor(model, dataset) + # This will save out the intermediate results if an exception is thrown # during the reconstruction with model.save_on_exception( 'example_reconstructions/gold_balls_earlyexit.h5', dataset): - for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50): + for loss in recon.optimize(20, lr=0.005, batch_size=50): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) - for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100): + for loss in recon.optimize(50, lr=0.002, batch_size=100): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) @@ -64,8 +68,7 @@ with model.save_on_exception( # Setting schedule=True automatically lowers the learning rate if # the loss fails to improve after 10 epochs - for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100, - schedule=True): + for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) diff --git a/examples/gold_ball_split.py b/examples/gold_ball_split.py index 8624b9f..9fc5b08 100644 --- a/examples/gold_ball_split.py +++ b/examples/gold_ball_split.py @@ -36,15 +36,18 @@ for label, dataset in zip(labels, datasets): model.to(device=device) dataset.get_as(device=device) + # Create the reconstructor + recon = cdtools.reconstructors.AdamReconstructor(model, dataset) + # For batched reconstructions like this, there's no need to live-plot # the progress - for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50): + for loss in recon.optimize(20, lr=0.005, batch_size=50): print(model.report()) - for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100): + for loss in recon.optimize(50, lr=0.002, batch_size=100): print(model.report()) - for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100, + for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True): print(model.report()) diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 726af94..217d96b 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -1,3 +1,12 @@ +""" +Runs a very simple reconstruction using the SimplePtycho model, which was +designed to be an easy introduction to show how the models are made and used. + +For a more realistic example of how to use cdtools for real-world data, +look at fancy_ptycho.py and gold_ball_ptycho.py, both of which use the +more powerful FancyPtycho model and include more information on how to +correct for common sources of error. +""" import cdtools from matplotlib import pyplot as plt @@ -13,7 +22,7 @@ device = 'cuda' model.to(device=device) dataset.get_as(device=device) -# We run the actual reconstruction +# We run the reconstruction for loss in model.Adam_optimize(100, dataset, batch_size=10): # We print a quick report of the optimization status print(model.report())