From 8cf217f5d8a93e4cec7bb0eb72713c41baf34905 Mon Sep 17 00:00:00 2001
From: Pengju Sheng <pengju.sheng@psi.ch>
Date: Thu, 23 Jul 2026 11:50:01 +0200
Subject: [PATCH] Add feature adjust learning rate

---
 examples/fancy_ptycho_adjust_lr.py | 54 ++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 examples/fancy_ptycho_adjust_lr.py

diff --git a/examples/fancy_ptycho_adjust_lr.py b/examples/fancy_ptycho_adjust_lr.py
new file mode 100644
index 0000000..c23c077
--- /dev/null
+++ b/examples/fancy_ptycho_adjust_lr.py
@@ -0,0 +1,54 @@
+import cdtools
+import torch as t
+from matplotlib import pyplot as plt
+
+filename = 'example_data/lab_ptycho_data.cxi'
+dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
+
+# FancyPtycho is the workhorse model
+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')
+
+# For this script, we use a slightly different pattern where we explicitly
+# create a `Reconstructor` class to orchestrate the reconstruction. The
+# reconstructor will store the model and dataset and create an appropriate
+# optimizer. This allows the optimizer to persist between loops, along with
+# e.g. estimates of the moments of individual parameters
+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
+lr = {'translation_offsets':0.04,'background':0.01}
+for loss in recon.optimize(50, lr=lr, batch_size=10, default_lr = 0.03):
+    print(model.report())
+    # Because plotting can be expensive, setting a minimum plotting interval
+    # (in seconds) can avoid excessive replots. 
+    model.inspect(min_interval=10)
+
+# It's common to chain several different reconstruction loops. Here, we
+# 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())
+    model.inspect(min_interval=10)
+
+# This orthogonalizes the recovered probe modes
+model.tidy_probes()
+
+# Setting replot_all will reopen any windows which were closed earlier
+model.inspect(replot_all=True)
+model.compare(dataset)
+plt.show()