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 1/6] 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()

From d0122b1161d92ab6419c74f1d69604d7b09fa953 Mon Sep 17 00:00:00 2001
From: Pengju Sheng <pengju.sheng@psi.ch>
Date: Thu, 23 Jul 2026 11:54:14 +0200
Subject: [PATCH 2/6] adjust lr

---
 src/cdtools/reconstructors/adam.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
index e298bdb..b335166 100644
--- a/src/cdtools/reconstructors/adam.py
+++ b/src/cdtools/reconstructors/adam.py
@@ -52,16 +52,21 @@ def __init__(self,
                  subset: List[int] = None):
 
         # Define the optimizer for use in this subclass
-        optimizer = t.optim.Adam(model.parameters())
+        param_groups = []
+        for name, param in model.named_parameters():
+            param_groups.append({'params':[param], 'name':name})
+
+        optimizer = t.optim.Adam(param_groups)
 
         super().__init__(model, dataset, optimizer, subset=subset)
 
 
 
     def adjust_optimizer(self,
-                         lr: int = 0.005,
+                         lr: int | dict = 0.005,
                          betas: Tuple[float] = (0.9, 0.999),
-                         amsgrad: bool = False):
+                         amsgrad: bool = False,
+                         default_lr : int = 0.005):
         """
         Change hyperparameters for the utilized optimizer.
 
@@ -77,16 +82,25 @@ def adjust_optimizer(self,
             Optional, whether to use the AMSGrad variant of this algorithm.
         """
         for param_group in self.optimizer.param_groups:
-            param_group['lr'] = lr
             param_group['betas'] = betas
             param_group['amsgrad'] = amsgrad
-
+            param_name = param_group['name']
+            if isinstance(lr, dict):
+                if param_name not in lr:
+                    param_group['lr'] = default_lr
+                else:
+                    param_group['lr'] = lr[param_name]
+            else:
+                param_group['lr'] = lr
+
+            print(f"Found paramter {param_name}, learning rate : {param_group['lr']}")
         
     def optimize(self,
                  iterations: int,
                  batch_size: int = 15,
-                 lr: float = 0.005,
+                 lr: int | dict = 0.005,
                  betas: Tuple[float] = (0.9, 0.999),
+                 default_lr : int = 0.005,
                  custom_data_loader: t.utils.data.DataLoader = None,
                  schedule: bool = False,
                  amsgrad: bool = False,
@@ -155,7 +169,8 @@ def optimize(self,
         # hyperparameters need to be set up with self.adjust_optimizer
         self.adjust_optimizer(lr=lr,
                               betas=betas,
-                              amsgrad=amsgrad)
+                              amsgrad=amsgrad,
+                              default_lr = default_lr)
 
         # Set up the scheduler
         if schedule:

From a6dca51e3ef44becc31be17a5698524d5ae4b200 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Thu, 23 Jul 2026 14:11:03 +0200
Subject: [PATCH 3/6] First working version

---
 examples/fancy_ptycho_adjust_lr.py         |  54 --------
 examples/fancy_ptycho_per_parameter_lrs.py |  50 +++++++
 src/cdtools/reconstructors/adam.py         | 144 +++++++++++++++------
 3 files changed, 157 insertions(+), 91 deletions(-)
 delete mode 100644 examples/fancy_ptycho_adjust_lr.py
 create mode 100644 examples/fancy_ptycho_per_parameter_lrs.py

diff --git a/examples/fancy_ptycho_adjust_lr.py b/examples/fancy_ptycho_adjust_lr.py
deleted file mode 100644
index c23c077..0000000
--- a/examples/fancy_ptycho_adjust_lr.py
+++ /dev/null
@@ -1,54 +0,0 @@
-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()
diff --git a/examples/fancy_ptycho_per_parameter_lrs.py b/examples/fancy_ptycho_per_parameter_lrs.py
new file mode 100644
index 0000000..4c9a0a1
--- /dev/null
+++ b/examples/fancy_ptycho_per_parameter_lrs.py
@@ -0,0 +1,50 @@
+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)
+
+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')
+
+# Here, we tune the learning rates of individual parameters. The default
+# learning rate factor is 1. Any learning rate factor set here will multiply
+# the learning rate for each recon.optimize loop. The dictionary can be passed
+# to the reconstructor object at creation time, as done here. It can also be
+# updated later with the call to recon.optimize(..., lr_factors=lr_factors).
+lr_factors = {
+    'translation_offsets' : 1.2,
+    'weights' : 0.2,
+    'background' : 0.3,
+}
+
+recon = cdtools.reconstructors.AdamReconstructor(
+    model, dataset, lr_factors=lr_factors)
+
+# For example, background will get a lr of 0.03 * 0.3 (lr * lr_factor).
+for loss in recon.optimize(50, lr=0.03, batch_size=10, lr_factors=lr_factors, verbose=True):
+    print(model.report())
+    model.inspect(min_interval=10)
+
+# And here background will get a lr of 0.005 * 0.3 (lr * lr_factor).
+for loss in recon.optimize(50, lr=0.005, batch_size=50):
+    print(model.report())
+    model.inspect(min_interval=10)
+
+model.tidy_probes()
+
+model.inspect(replot_all=True)
+model.compare(dataset)
+plt.show()
diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
index b335166..477e9ef 100644
--- a/src/cdtools/reconstructors/adam.py
+++ b/src/cdtools/reconstructors/adam.py
@@ -8,6 +8,7 @@
 """
 from __future__ import annotations
 from typing import TYPE_CHECKING
+import warnings
 
 import torch as t
 from typing import Tuple, List, Union
@@ -46,10 +47,13 @@ class AdamReconstructor(Reconstructor):
     - **data_loader** -- A torch.utils.data.DataLoader that is defined by
         calling the `setup_dataloader` method.
     """
-    def __init__(self,
-                 model: CDIModel,
-                 dataset: Ptycho2DDataset,
-                 subset: List[int] = None):
+    def __init__(
+            self,
+            model: CDIModel,
+            dataset: Ptycho2DDataset,
+            subset: List[int] = None,
+            lr_factors: dict = {}
+    ):
 
         # Define the optimizer for use in this subclass
         param_groups = []
@@ -60,54 +64,109 @@ def __init__(self,
 
         super().__init__(model, dataset, optimizer, subset=subset)
 
+        self._set_lr_factors(lr_factors)
 
+        
+    def _set_lr_factors(self, lr_factors):
+        """Sets the learning rate factors from a provided dictionary
+
+        This is broken out into it's own function to avoid replicating the
+        code to emit a warning, and to enable easy future changes as the logic
+        may need to become more complicated.
+
+        Parameters
+        ----------
+        lr_factors : dict
+            A dictionary mapping optimizer parameters to adjustment factors for the learning rate.
+        """
+        self.lr_factors = lr_factors
+        param_group_names = {p['name'] for p in self.optimizer.param_groups}
+        unused_lr_factors = self.lr_factors.keys() - param_group_names
+        
+        if len(unused_lr_factors) != 0:
+            warnings.warn(
+                'The lr_factor dictionary defines some entries ' +
+                'which are unused. Check the following entries for typos:' +
+                str(unused_lr_factors),
+                stacklevel=3,
+            )
 
-    def adjust_optimizer(self,
-                         lr: int | dict = 0.005,
-                         betas: Tuple[float] = (0.9, 0.999),
-                         amsgrad: bool = False,
-                         default_lr : int = 0.005):
+    def print_lrs(self):
+        """Prints the current per-parameter learning rates.
+        """
+
+        for param_group in self.optimizer.param_groups:
+            print(
+                f"Paramter {param_group['name']} has learning rate "
+                f"{param_group['lr']}."
+            )
+
+    def adjust_optimizer(
+            self,
+            lr: int = 0.005,
+            betas: Tuple[float] = (0.9, 0.999),
+            amsgrad: bool = False,
+            lr_factors: dict = None,
+            verbose: bool = False,
+    ):
         """
         Change hyperparameters for the utilized optimizer.
 
         Parameters
         ----------
         lr : float
-            Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
+            Optional, the learning rate (alpha) to use. Default is 0.005. 0.05
             is typically the highest possible value with any chance of being
             stable.
         betas : tuple
             Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
         amsgrad : bool
             Optional, whether to use the AMSGrad variant of this algorithm.
+        lr_factors : dict
+            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
+        verbose : bool
+            Default False, whether to print out setup information after adjustment.
         """
+
+        # Update the learning rate factors if explicitly given. Otherwise,
+        # persist the existing dictionary. A common pattern is to set the
+        # factors once at the start, and then adjust only the learning rate
+        # afterward.
+        if lr_factors is not None:
+            self._set_lr_factors(lr_factors)
+                    
         for param_group in self.optimizer.param_groups:
             param_group['betas'] = betas
             param_group['amsgrad'] = amsgrad
             param_name = param_group['name']
-            if isinstance(lr, dict):
-                if param_name not in lr:
-                    param_group['lr'] = default_lr
+            if isinstance(self.lr_factors, dict):
+                if param_name not in self.lr_factors:
+                    param_group['lr'] = lr
                 else:
-                    param_group['lr'] = lr[param_name]
+                    param_group['lr'] = lr * self.lr_factors[param_name]
             else:
                 param_group['lr'] = lr
 
-            print(f"Found paramter {param_name}, learning rate : {param_group['lr']}")
-        
-    def optimize(self,
-                 iterations: int,
-                 batch_size: int = 15,
-                 lr: int | dict = 0.005,
-                 betas: Tuple[float] = (0.9, 0.999),
-                 default_lr : int = 0.005,
-                 custom_data_loader: t.utils.data.DataLoader = None,
-                 schedule: bool = False,
-                 amsgrad: bool = False,
-                 regularization_factor: Union[float, List[float]] = None,
-                 thread: bool = True,
-                 calculation_width: int = 10,
-                 shuffle: bool = True):
+        if verbose:
+            self.print_lrs()
+
+                
+    def optimize(
+            self,
+            iterations: int,
+            batch_size: int = 15,
+            lr: int = 0.005,
+            betas: Tuple[float] = (0.9, 0.999),
+            lr_factors : dict = None,
+            custom_data_loader: t.utils.data.DataLoader = None,
+            schedule: bool = False,
+            amsgrad: bool = False,
+            regularization_factor: Union[float, List[float]] = None,
+            thread: bool = True,
+            calculation_width: int = 10,
+            shuffle: bool = True,
+            verbose: bool = False,
+    ):
         """
         Runs a round of reconstruction using the Adam optimizer
 
@@ -157,21 +216,32 @@ def optimize(self,
         shuffle : bool
             Optional, enable/disable shuffling of the dataset. This option
             is intended for diagnostic purposes and should be left as True.
+        lr_factors : dict
+            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
+        verbose : bool
+            Default False, whether to print out setup information about the planned run.
         """
+
+        # The optimizer is created in self.__init__, but the
+        # hyperparameters need to be set up with self.adjust_optimizer
+        self.adjust_optimizer(
+            lr=lr,
+            betas=betas,
+            amsgrad=amsgrad,
+            lr_factors=lr_factors,
+            verbose=verbose,
+        )
+
         # Update the training history
         self.model.training_history += (
             f'Planning {iterations} epochs of Adam, with a learning rate = '
             f'{lr}, batch size = {batch_size}, regularization_factor = '
             f'{regularization_factor}, and schedule = {schedule}.\n'
         )
-
-        # The optimizer is created in self.__init__, but the
-        # hyperparameters need to be set up with self.adjust_optimizer
-        self.adjust_optimizer(lr=lr,
-                              betas=betas,
-                              amsgrad=amsgrad,
-                              default_lr = default_lr)
-
+        self.model.training_history += (
+            f'The learning rate factors are {self.lr_factors}, default = 1.\n'
+        )
+        
         # Set up the scheduler
         if schedule:
             self.scheduler = \

From f6e0654445a866d986df6338b819f0e0f9368132 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Thu, 23 Jul 2026 14:18:24 +0200
Subject: [PATCH 4/6] Update model.Adam_optimize()

---
 src/cdtools/models/base.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index ebaaf85..1beb560 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -405,6 +405,7 @@ def Adam_optimize(
             dataset: CDataset,
             batch_size: int = 15,
             lr: float = 0.005,
+            lr_factors: dict = {},
             betas: Tuple[float] = (0.9, 0.999),
             schedule: bool = False,
             amsgrad: bool = False,
@@ -434,6 +435,8 @@ def Adam_optimize(
             Optional, The learning rate (alpha) to use. Defaultis 0.005. 
             0.05 is typically the highest possible value with any chance 
             of being stable.
+        lr_factors : dict
+            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
         betas : tuple(float)
             Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
         schedule : bool
@@ -458,6 +461,7 @@ def Adam_optimize(
             model=self,
             dataset=dataset,
             subset=subset,
+            lr_factors=lr_factors,
         )
         
         # Run some reconstructions

From 3252a0dd8f795676063b5f75686f36a3d576ccc7 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Thu, 23 Jul 2026 15:00:33 +0200
Subject: [PATCH 5/6] Update SGD to allow per-parameter learning rates and add
 test coverage

---
 examples/fancy_ptycho_per_parameter_lrs.py |   2 +-
 src/cdtools/models/base.py                 |   4 +
 src/cdtools/reconstructors/adam.py         |  31 ++---
 src/cdtools/reconstructors/sgd.py          | 149 +++++++++++++++++----
 tests/test_reconstructors.py               |  56 ++++++--
 5 files changed, 182 insertions(+), 60 deletions(-)

diff --git a/examples/fancy_ptycho_per_parameter_lrs.py b/examples/fancy_ptycho_per_parameter_lrs.py
index 4c9a0a1..39cf20f 100644
--- a/examples/fancy_ptycho_per_parameter_lrs.py
+++ b/examples/fancy_ptycho_per_parameter_lrs.py
@@ -34,7 +34,7 @@
     model, dataset, lr_factors=lr_factors)
 
 # For example, background will get a lr of 0.03 * 0.3 (lr * lr_factor).
-for loss in recon.optimize(50, lr=0.03, batch_size=10, lr_factors=lr_factors, verbose=True):
+for loss in recon.optimize(50, lr=0.03, batch_size=10):
     print(model.report())
     model.inspect(min_interval=10)
 
diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index 1beb560..943f5ed 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -544,6 +544,7 @@ def SGD_optimize(self,
                      dataset: CDataset,
                      batch_size: int = None,
                      lr: float = 2e-7,
+                     lr_factors : dict = {},
                      momentum: float = 0,
                      dampening: float = 0,
                      weight_decay: float = 0,
@@ -569,6 +570,8 @@ def SGD_optimize(self,
             Optional, the size of the minibatches to use.
         lr : float
             Optional, the learning rate to use.
+        lr_factors : dict
+            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
         momentum : float
             Optional, the length of the history to use.
         dampening : float
@@ -595,6 +598,7 @@ def SGD_optimize(self,
             model=self, 
             dataset=dataset, 
             subset=subset,
+            lr_factors=lr_factors,
         )
 
         # Run some reconstructions
diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
index 477e9ef..8eaf016 100644
--- a/src/cdtools/reconstructors/adam.py
+++ b/src/cdtools/reconstructors/adam.py
@@ -34,14 +34,15 @@ class AdamReconstructor(Reconstructor):
         The dataset to reconstruct against.
     subset : list(int) or int
         Optional, a pattern index or list of pattern indices to use.
-    schedule : bool
-        Optional, create a learning rate scheduler
-        (torch.optim.lr_scheduler._LRScheduler).
+    lr_factors : dict
+        Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
 
     Important attributes:
     - **model** -- Always points to the core model used.
     - **optimizer** -- This class by default uses `torch.optim.Adam` to perform
         optimizations.
+    - **lr_factors** -- A map from optimizer parameters to learning rate
+        factors.
     - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the
         `optimize` method.
     - **data_loader** -- A torch.utils.data.DataLoader that is defined by
@@ -90,6 +91,7 @@ def _set_lr_factors(self, lr_factors):
                 str(unused_lr_factors),
                 stacklevel=3,
             )
+            
 
     def print_lrs(self):
         """Prints the current per-parameter learning rates.
@@ -101,13 +103,13 @@ def print_lrs(self):
                 f"{param_group['lr']}."
             )
 
+            
     def adjust_optimizer(
             self,
             lr: int = 0.005,
             betas: Tuple[float] = (0.9, 0.999),
             amsgrad: bool = False,
             lr_factors: dict = None,
-            verbose: bool = False,
     ):
         """
         Change hyperparameters for the utilized optimizer.
@@ -124,8 +126,6 @@ def adjust_optimizer(
             Optional, whether to use the AMSGrad variant of this algorithm.
         lr_factors : dict
             Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
-        verbose : bool
-            Default False, whether to print out setup information after adjustment.
         """
 
         # Update the learning rate factors if explicitly given. Otherwise,
@@ -139,17 +139,12 @@ def adjust_optimizer(
             param_group['betas'] = betas
             param_group['amsgrad'] = amsgrad
             param_name = param_group['name']
-            if isinstance(self.lr_factors, dict):
-                if param_name not in self.lr_factors:
-                    param_group['lr'] = lr
-                else:
-                    param_group['lr'] = lr * self.lr_factors[param_name]
+            if isinstance(self.lr_factors, dict) and \
+               param_name in self.lr_factors:
+                param_group['lr'] = lr * self.lr_factors[param_name]
             else:
                 param_group['lr'] = lr
 
-        if verbose:
-            self.print_lrs()
-
                 
     def optimize(
             self,
@@ -165,7 +160,6 @@ def optimize(
             thread: bool = True,
             calculation_width: int = 10,
             shuffle: bool = True,
-            verbose: bool = False,
     ):
         """
         Runs a round of reconstruction using the Adam optimizer
@@ -195,6 +189,8 @@ def optimize(
             stable.
         betas : tuple
             Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
+        lr_factors : dict
+            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
         schedule : bool
             Optional, create a learning rate scheduler
             (torch.optim.lr_scheduler._LRScheduler).
@@ -216,10 +212,6 @@ def optimize(
         shuffle : bool
             Optional, enable/disable shuffling of the dataset. This option
             is intended for diagnostic purposes and should be left as True.
-        lr_factors : dict
-            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
-        verbose : bool
-            Default False, whether to print out setup information about the planned run.
         """
 
         # The optimizer is created in self.__init__, but the
@@ -229,7 +221,6 @@ def optimize(
             betas=betas,
             amsgrad=amsgrad,
             lr_factors=lr_factors,
-            verbose=verbose,
         )
 
         # Update the training history
diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py
index cb1e26b..c5c495a 100644
--- a/src/cdtools/reconstructors/sgd.py
+++ b/src/cdtools/reconstructors/sgd.py
@@ -34,23 +34,34 @@ class SGDReconstructor(Reconstructor):
         The dataset to reconstruct against.
     subset : list(int) or int
         Optional, a pattern index or list of pattern indices to use.
+    lr_factors : dict
+        Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
 
     Important attributes:
     - **model** -- Always points to the core model used.
     - **optimizer** -- This class by default uses `torch.optim.Adam` to perform
         optimizations.
+    - **lr_factors** -- A map from optimizer parameters to learning rate
+        factors.
     - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the
         `optimize` method.
     - **data_loader** -- A torch.utils.data.DataLoader that is defined by
         calling the `setup_dataloader` method.
     """
-    def __init__(self,
-                 model: CDIModel,
-                 dataset: Ptycho2DDataset,
-                 subset: List[int] = None):
+    def __init__(
+            self,
+            model: CDIModel,
+            dataset: Ptycho2DDataset,
+            subset: List[int] = None,
+            lr_factors: dict = {}
+    ):
 
         # Define the optimizer for use in this subclass
-        optimizer = t.optim.SGD(model.parameters())
+        param_groups = []
+        for name, param in model.named_parameters():
+            param_groups.append({'params':[param], 'name':name})
+
+        optimizer = t.optim.SGD(param_groups)
 
         super().__init__(
             model,
@@ -59,13 +70,54 @@ def __init__(self,
             subset=subset,
         )
 
+        self._set_lr_factors(lr_factors)
+
+
+    def _set_lr_factors(self, lr_factors):
+        """Sets the learning rate factors from a provided dictionary
+
+        This is broken out into it's own function to avoid replicating the
+        code to emit a warning, and to enable easy future changes as the logic
+        may need to become more complicated.
+
+        Parameters
+        ----------
+        lr_factors : dict
+            A dictionary mapping optimizer parameters to adjustment factors for the learning rate.
+        """
+        self.lr_factors = lr_factors
+        param_group_names = {p['name'] for p in self.optimizer.param_groups}
+        unused_lr_factors = self.lr_factors.keys() - param_group_names
+        
+        if len(unused_lr_factors) != 0:
+            warnings.warn(
+                'The lr_factor dictionary defines some entries ' +
+                'which are unused. Check the following entries for typos:' +
+                str(unused_lr_factors),
+                stacklevel=3,
+            )
+            
+
+    def print_lrs(self):
+        """Prints the current per-parameter learning rates.
+        """
+
+        for param_group in self.optimizer.param_groups:
+            print(
+                f"Paramter {param_group['name']} has learning rate "
+                f"{param_group['lr']}."
+            )
+
         
-    def adjust_optimizer(self,
-                         lr: int = 0.005,
-                         momentum: float = 0,
-                         dampening: float = 0,
-                         weight_decay: float = 0,
-                         nesterov: bool = False):
+    def adjust_optimizer(
+            self,
+            lr: int = 0.005,
+            momentum: float = 0,
+            dampening: float = 0,
+            weight_decay: float = 0,
+            nesterov: bool = False,
+            lr_factors: dict = None,
+    ):
         """
         Change hyperparameters for the utilized optimizer.
 
@@ -84,26 +136,48 @@ def adjust_optimizer(self,
         nesterov : bool
             Optional, enables Nesterov momentum. Only applicable when momentum
             is non-zero.
+        lr_factors : dict
+            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
+
         """
+
+        # Update the learning rate factors if explicitly given. Otherwise,
+        # persist the existing dictionary. A common pattern is to set the
+        # factors once at the start, and then adjust only the learning rate
+        # afterward.
+        if lr_factors is not None:
+            self._set_lr_factors(lr_factors)
+                    
+
         for param_group in self.optimizer.param_groups:
-            param_group['lr'] = lr
             param_group['momentum'] = momentum
             param_group['dampening'] = dampening
             param_group['weight_decay'] = weight_decay
             param_group['nesterov'] = nesterov
 
-    def optimize(self,
-                 iterations: int,
-                 batch_size: int = 15,
-                 lr: float = 2e-7,
-                 momentum: float = 0,
-                 dampening: float = 0,
-                 weight_decay: float = 0,
-                 nesterov: bool = False,
-                 regularization_factor: Union[float, List[float]] = None,
-                 thread: bool = True,
-                 calculation_width: int = 10,
-                 shuffle: bool = True):
+            param_name = param_group['name']
+            if isinstance(self.lr_factors, dict) and \
+               param_name in self.lr_factors:
+                param_group['lr'] = lr * self.lr_factors[param_name]
+            else:
+                param_group['lr'] = lr
+
+
+    def optimize(
+            self,
+            iterations: int,
+            batch_size: int = 15,
+            lr: float = 2e-7,
+            momentum: float = 0,
+            dampening: float = 0,
+            weight_decay: float = 0,
+            nesterov: bool = False,
+            lr_factors : dict = None,
+            regularization_factor: Union[float, List[float]] = None,
+            thread: bool = True,
+            calculation_width: int = 10,
+            shuffle: bool = True,
+    ):
         """
         Runs a round of reconstruction using the Adam optimizer
 
@@ -131,6 +205,8 @@ def optimize(self,
         nesterov : bool
             Optional, enables Nesterov momentum. Only applicable when momentum
             is non-zero.
+        lr_factors : dict
+            Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate.
         regularization_factor : float or list(float)
             Optional, if the model has a regularizer defined, the set of
             parameters to pass the regularizer method.
@@ -148,12 +224,27 @@ def optimize(self,
 
         # The optimizer is created in self.__init__, but the
         # hyperparameters need to be set up with self.adjust_optimizer
-        self.adjust_optimizer(lr=lr,
-                              momentum=momentum,
-                              dampening=dampening,
-                              weight_decay=weight_decay,
-                              nesterov=nesterov)
+        self.adjust_optimizer(
+            lr=lr,
+            momentum=momentum,
+            dampening=dampening,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            lr_factors=lr_factors,
+        )
 
+        # Update the training history
+        self.model.training_history += (
+            f'Planning {iterations} epochs of SGD, with a learning rate = '
+            f'{lr}, batch size = {batch_size}, regularization_factor = '
+            f'{regularization_factor}, momentum history length = {momentum},'
+            f'momemntum dampening = {dampening}, weight_decay = {weight_decay},'
+            f' and nesterov = {nesterov}.\n'
+        )
+        self.model.training_history += (
+            f'The learning rate factors are {self.lr_factors}, default = 1.\n'
+        )
+        
         # Now, we run the optimize routine defined in the base class
         return super(SGDReconstructor, self).optimize(
             iterations,
diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py
index ff34a7a..1ac6731 100644
--- a/tests/test_reconstructors.py
+++ b/tests/test_reconstructors.py
@@ -20,9 +20,10 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
         4) Reconstructions performed by `Adam.optimize` and
            `model.Adam_optimize` calls produce identical results when
            run over one round of optimization.
-        5) The quality of the reconstruction remains below a specified
+        5) Checks that the per-parameter learning rates work in both cases
+        6) The quality of the reconstruction remains below a specified
            threshold.
-        5) Ensure that the FancyPtycho model works fine and dandy with the
+        7) Ensure that the FancyPtycho model works fine and dandy with the
            Reconstructors.
     """
 
@@ -53,12 +54,20 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     model_recon.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
 
+    lr_factors = {
+        'obj' : 1.1,
+        'weights' : 0.5,
+    }
+    
     # ******* Reconstructions with AdamReconstructor.optimize *******
     print('Running reconstruction using AdamReconstructor.optimize' +
           ' on provided reconstruction_device,', reconstruction_device)
 
-    recon = cdtools.reconstructors.AdamReconstructor(model=model_recon,
-                                                     dataset=dataset)
+    recon = cdtools.reconstructors.AdamReconstructor(
+        model=model_recon,
+        dataset=dataset,
+        lr_factors=lr_factors,
+    )
     t.manual_seed(0)
 
     # Run a reconstruction
@@ -100,10 +109,13 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
 
     # We only need to test the first loop to ensure it's identical
     for i, iterations in enumerate(epoch_tup[:1]):
-        for loss in model.Adam_optimize(iterations,
-                                        dataset,
-                                        lr=lr_tup[i],
-                                        batch_size=batch_size_tup[i]):
+        for loss in model.Adam_optimize(
+                iterations,
+                dataset,
+                lr=lr_tup[i],
+                lr_factors=lr_factors,
+                batch_size=batch_size_tup[i],
+        ):
             print(model.report())
             if show_plot:
                 model.inspect(dataset, min_interval=10)
@@ -161,8 +173,16 @@ def test_intensity_MSE(gold_ball_cxi, reconstruction_device, show_plot):
     for loss in recon.optimize(5, lr=.05, batch_size=10):
         print(model.report())
 
-    # Threshold to be updated after running on a GPU machine
-    assert model.loss_history[-1] < 1e7
+    # Test that Adam optimizer post-creation update of lr_factors works
+    lr_factors = {
+        'background' : 0.3,
+        'translation_offsets': 1.2,
+    }
+    
+    for loss in recon.optimize(3, lr=.05, batch_size=10, lr_factors=lr_factors):
+        print(model.report())
+
+    assert model.loss_history[-1] < 6.5e6
 
 
 @pytest.mark.slow
@@ -371,3 +391,19 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     # a threshold of 7.2e-4 for the tested loss. If this value has been
     # exceeded, the reconstructions have gotten worse.
     assert model.loss_history[-1] < 0.95
+
+    print('Testing per-parameter learning rates')
+
+    lr_factors = {
+        'background': 0.4,
+    }
+    
+    for loss in model.SGD_optimize(epochs,
+                                   dataset,
+                                   lr=lr,
+                                   lr_factors=lr_factors,
+                                   batch_size=batch_size):
+        print(model.report())
+        if show_plot:
+            model.inspect(dataset, min_interval=10)
+    

From 289182f3f3904b3d5eca76ef11c4cf665203f195 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Thu, 23 Jul 2026 15:05:23 +0200
Subject: [PATCH 6/6] Update the documentation to include the example for
 adjusting learning rates per-parameter

---
 docs/source/examples.rst                             | 12 ++++++++++++
 ...ameter_lrs.py => per_parameter_learning_rates.py} |  0
 2 files changed, 12 insertions(+)
 rename examples/{fancy_ptycho_per_parameter_lrs.py => per_parameter_learning_rates.py} (100%)

diff --git a/docs/source/examples.rst b/docs/source/examples.rst
index 6f356cd..879f86f 100644
--- a/docs/source/examples.rst
+++ b/docs/source/examples.rst
@@ -110,6 +110,18 @@ If :code:`propagation_distance` is set, it will assume a Fresnel scaling theorem
 Finally, note the addition of the :code:`panel_plot_mode=True` argument. This is the default mode, and returns the plots in a panel format, good for easily monitoring the progress of a reconstruction. If individual plots are needed for use in presentations, papers, or otherwise, setting :code:`panel_plot_mode=False` will plot each output in it's own window.
 
 
+Per-Parameter Learning Rates
+----------------------------
+
+This script shows how the learning rates can be adjusted per parameter, which can sometimes accelerate convergence substantially.
+
+.. literalinclude:: ../../examples/per_parameter_learning_rates.py
+
+The major addition is the inclusion of a dictionary, :code:`lr_factors`, which multiplies the main learning rate for each individual parameter. If a specific parameter is not being updated aggressively enough, increase this value from the default of 1. If it is being updated too aggressively and preventing convergence, lower the value.
+
+This dictionary will persist through all further :code:`recon.optimize()` calls, unles explicitly reset.
+
+
 Gold Ball Split
 ---------------
 
diff --git a/examples/fancy_ptycho_per_parameter_lrs.py b/examples/per_parameter_learning_rates.py
similarity index 100%
rename from examples/fancy_ptycho_per_parameter_lrs.py
rename to examples/per_parameter_learning_rates.py