Start to make the lab bragg ptycho code sane, and fix the bug with probe bases when using probe fourier crop

examples/gold_ball_ptycho.py

@@ -1,6 +1,7 @@
 import cdtools
 from matplotlib import pyplot as plt
 from scipy import io
+import torch as t
 
 # First, we load an example dataset from a .cxi file
 filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
@@ -8,15 +9,18 @@ dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 # Next, we create a ptychography model from the dataset
 # Note that we explicitly ask for two incoherent probe modes
-model = cdtools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
+model = cdtools.models.FancyPtycho.from_dataset(
+    dataset,
+    n_modes=2,
+)
 
 # Let's do this reconstruction on the GPU, shall we? 
 #model.to(device='cuda')
 #dataset.get_as(device='cuda')
 
-with model.save_on_exit('example_reconstructions/gold_balls.mat', dataset):
+with model.save_on_exit('example_reconstructions/gold_balls.h5', dataset):
     # Now, we run a short reconstruction from the dataset
-    for loss in model.Adam_optimize(10, dataset, batch_size=50):
+    for loss in model.Adam_optimize(1, dataset, batch_size=50):
         # And we liveplot the updates to the model as they happen
         print(model.report())
         model.inspect(dataset)

examples/lab_bragg_2d_ptycho.py

@@ -10,14 +10,19 @@ from scipy import io
 
 # This file is too large to be distributed via Github.
 # Please contact Abe Levitan (alevitan@mit) if you would like access
-filename = 'Zone Plate Bragg 633.cxi'
+filename = 'example_data/Zone Plate Bragg 633.cxi'
 dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
 
 #dataset.inspect()
 
-model = cdtools.models.Bragg2DPtycho.from_dataset(dataset,probe_support_radius=60,correct_tilt=False)
-model.to(device='cuda')
-dataset.get_as(device='cuda')
+model = cdtools.models.Bragg2DPtycho.from_dataset(
+    dataset,
+    probe_support_radius=60,
+    correct_tilt=False
+)
+
+#model.to(device='cuda')
+#dataset.get_as(device='cuda')
 
 
 model.translation_offsets.requires_grad = False
@@ -25,8 +30,8 @@ for loss in model.Adam_optimize(100, dataset):
     model.inspect(dataset)
     print(model.report())
 
-io.savemat('example_reconstructions/lab_bragg_2d_ptycho.mat',
-           model.save_results(dataset))
+#io.savemat('example_reconstructions/lab_bragg_2d_ptycho.mat',
+#           model.save_results(dataset))
 
 model.compare(dataset)
 plt.show()