Updating LBFGS line search type

CDTools/datasets/ptycho_2d_dataset.py

@@ -60,9 +60,9 @@ class Ptycho2DDataset(CDataset):
 
         super(Ptycho2DDataset,self).__init__(*args, **kwargs)
         self.axes = copy(axes)
-        self.translations = t.tensor(translations, dtype=t.float32)
+        self.translations = t.tensor(translations)
         
-        self.patterns = t.as_tensor(patterns, dtype=t.float32)
+        self.patterns = t.as_tensor(patterns)
 
         if self.mask is None:
             self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.bool)

CDTools/models/base.py

@@ -293,7 +293,7 @@ class CDIModel(t.nn.Module):
     def LBFGS_optimize(self, iterations, dataset,
                        lr=0.1,history_size=2, subset=None,
                        regularization_factor=None, thread=True,
-                       calculation_width=10):
+                       calculation_width=10, line_search_fn=None):
         """Runs a round of reconstruction using the L-BFGS optimizer
 
         This algorithm is often less stable that Adam, however in certain
@@ -335,7 +335,8 @@ class CDIModel(t.nn.Module):
 
         # Define the optimizer
         optimizer = t.optim.LBFGS(self.parameters(),
-                                  lr = lr, history_size=history_size)
+                                  lr = lr, history_size=history_size,
+                                  line_search_fn=line_search_fn)
         #optimizer = MyLBFGS(self.parameters(),
         #                    lr = lr, history_size=history_size)
 

CDTools/models/fancy_ptycho.py

@@ -224,6 +224,10 @@ class FancyPtycho(CDIModel):
             # In this case, we define a set of weights which only has one index
             Ws = t.ones(len(dataset))
 
+        if hasattr(dataset, 'intensities') and dataset.intensities is not None:
+            Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...]
+                   / t.mean(dataset.intensities))
+
         if hasattr(dataset, 'mask') and dataset.mask is not None:
             mask = dataset.mask.to(t.bool)
         else:

CDTools/tools/polarization/polarization.py

@@ -17,9 +17,6 @@ __all__ = ['apply_linear_polarizer',
            'generate_birefringent_obj']
 
 
-# Abe - split these into two functions
-
-# Note for the future: this function should
 def generate_linear_polarizer(pol_angle):
     single_angle = False