From cea860bdf91a518651250744cc37d0f9557d5df8 Mon Sep 17 00:00:00 2001
From: Abe Levitan <alevitan@mit.edu>
Date: Fri, 5 Feb 2021 12:30:29 -0500
Subject: [PATCH] Fix an issue with unified_mode_ptycho not working in
 model.compare, and fix colorbars for model.compare

---
 CDTools/models/base.py                | 23 ++++++++++++++++++----
 CDTools/models/unified_mode_ptycho.py | 28 ++-------------------------
 2 files changed, 21 insertions(+), 30 deletions(-)

diff --git a/CDTools/models/base.py b/CDTools/models/base.py
index 138485b..166749e 100644
--- a/CDTools/models/base.py
+++ b/CDTools/models/base.py
@@ -53,6 +53,10 @@ class CDIModel(t.nn.Module):
     functions.
     """
 
+    def __init__(self):
+        super(CDIModel,self).__init__()
+        self.iteration_count = 0
+        
     def from_dataset(self, dataset):
         raise NotImplementedError()
 
@@ -135,6 +139,7 @@ class CDIModel(t.nn.Module):
         def run_iteration(stop_event=None):
             loss = 0
             N = 0
+            t0 = time.time()
             for inputs, patterns in data_loader:
                 N += 1
                 def closure():
@@ -175,6 +180,9 @@ class CDIModel(t.nn.Module):
             if scheduler is not None:
                 scheduler.step(loss)
 
+            self.latest_loss = loss
+            self.latest_iteration_time = time.time() - t0
+            self.iteration_count += 1
             return loss
 
         if thread:
@@ -377,6 +385,15 @@ class CDIModel(t.nn.Module):
                                 calculation_width=calculation_width)
 
 
+    def report(self):
+        """Returns a string informing on the latest reconstruction iteration"""
+        if hasattr(self, 'latest_loss'):
+            return 'Iteration ' + str(self.iteration_count) + \
+                  ' completed in %0.2f s with loss ' %\
+                  self.latest_iteration_time + str(self.latest_loss)
+        else:
+            return 'No reconstruction iterations performed yet!'
+        
     # By default, the plot_list is empty
     plot_list = []
     
@@ -490,10 +507,7 @@ class CDIModel(t.nn.Module):
                 return
             
             im.norecurse=True
-            im.colorbar.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
-            im.colorbar.ax.set_ylim(0,1)
-            im.colorbar.set_ticks(ticker.LinearLocator(numticks=5))
-            im.colorbar.draw_all()
+            im.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
 
         
         def update(idx):
@@ -530,6 +544,7 @@ class CDIModel(t.nn.Module):
                 cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff))
                 
             else:
+
                 sim = axes[0].images[-1]
                 sim.set_data(sim_data)
                 update_colorbar(sim)
diff --git a/CDTools/models/unified_mode_ptycho.py b/CDTools/models/unified_mode_ptycho.py
index 60921e9..2ec1e3b 100644
--- a/CDTools/models/unified_mode_ptycho.py
+++ b/CDTools/models/unified_mode_ptycho.py
@@ -242,36 +242,12 @@ class UnifiedModePtycho(CDIModel):
             pix_trans += self.translation_scale * self.translation_offsets[index]
 
         Ws = self.Ws[index]
-
-        # This probably needs to be fixed, I doubt it will really still work
-        # to deal with single-pattern sims.
-        #if type(index) == type(0):
-        #    print('hi')
-        #    index = [index]
-        #    Ws = [Ws]
-        #    pix_trans = [pix_trans]
-        #    single_frame = True
-        #else:
-        #    single_frame = False
-
-        probes = []
-        all_exit_waves = []
-
-        # This is the case if a purely stable, incoherent model is defined.
-        #if len(Ws[0].shape) == 0:
-            # What we do here is generate an identity matrix, and multiply
-            # that identity matrix by the per-frame weight.
-            #Ws = [W * t.stack([t.eye(self.probe.shape[0]),
-            #                   t.zeros([self.probe.shape[0]]*2)],dim=-1).to(
-            #                       dtype=W.dtype, device=W.device)
-            #      for W in Ws]
-            #print(Ws)
-            
+             
         # This restricts the basis probes with the probe support
         basis_prs = self.probe * self.probe_support[...,:,:]    
 
         
-        if len(Ws[0].shape) == 0:
+        if len(self.Ws[0].shape) == 0:
             # If a purely stable coherent illumination is defined
             prs = cmath.cmult(Ws[...,None,None,None,:],basis_prs)
         else: