Fix a bug in multislice where it doesn't work when prevent_alias is set to True

CDTools/models/base.py

@@ -39,7 +39,7 @@ import numpy as np
 import threading
 import queue
 import time
-import pytorch_warmup 
+#import pytorch_warmup
 
 __all__ = ['CDIModel']
 
@@ -282,7 +282,8 @@ class CDIModel(t.nn.Module):
             scheduler = None
 
         if warmup:
-            warmup_scheduler = pytorch_warmup.UntunedLinearWarmup(optimizer)
+            print('Warmup is not currently implemented, sorry!')
+            #warmup_scheduler = pytorch_warmup.UntunedLinearWarmup(optimizer)
         else:
             warmup_scheduler = None
             

CDTools/models/fancy_ptycho.py

@@ -122,7 +122,7 @@ class FancyPtycho(CDIModel):
 
         
     @classmethod
-    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, loss='amplitude mse', units='um'):
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'):
         
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
@@ -496,7 +496,7 @@ class FancyPtycho(CDIModel):
         
 
         weights = cmath.torch_to_complex(self.weights.detach().cpu())
-
+        
         probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)*
                                    weights.conj(),axis=2)
 

CDTools/models/multislice_2d_ptycho.py

@@ -329,7 +329,10 @@ class Multislice2DPtycho(CDIModel):
                     exit_waves = tools.interactions.ptycho_2D_round(
                         exit_waves, obj, pix_trans,
                         multiple_modes=True)
-                
+            #if self.iteration_count >= 10:
+            #    plt.figure()
+            #    plt.imshow(np.abs(cmath.torch_to_complex(exit_waves[0,0].detach().cpu())))
+            #    plt.show()
 
                     
             elif self.obj.dim() == 4:
@@ -346,17 +349,19 @@ class Multislice2DPtycho(CDIModel):
             if i < self.nz-1: #on all but the last iteration
                 exit_waves = tools.propagators.near_field(
                     exit_waves,self.as_prop)
-                    
         return exit_waves
     
         
     def forward_propagator(self, wavefields):
-        left = [self.probe.shape[-3]//2,self.probe.shape[-2]//2]
-        right = [self.probe.shape[-3]//2+self.probe.shape[-3],
-                 self.probe.shape[-2]//2+self.probe.shape[-2]]
-
-        return tools.propagators.far_field(wavefields)[...,left[0]:right[0],
-                                                       left[1]:right[1],:]
+        if self.prevent_aliasing:
+            left = [self.probe.shape[-3]//2,self.probe.shape[-2]//2]
+            right = [self.probe.shape[-3]//2+self.probe.shape[-3],
+                     self.probe.shape[-2]//2+self.probe.shape[-2]]
+            
+            return tools.propagators.far_field(wavefields)[...,left[0]:right[0],
+                                                           left[1]:right[1],:]
+        else:
+            return tools.propagators.far_field(wavefields)
 
     
     def measurement(self, wavefields):

CDTools/models/rpi.py

@@ -266,7 +266,7 @@ class RPI(CDIModel):
         try:
             # will fail if index has no length, for example when index
             # is just an int. In this case, we just do nothing instead
-            output = output.unsqueeze(1).repeat(1,len(index),1,1,1)
+            output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
         except TypeError:
             pass
         

CDTools/tools/data/data.py

@@ -354,7 +354,11 @@ def get_data(cxi_file, cut_zeroes = True):
         data[data < 0] = 0
 
     if 'axes' in cxi_file[pull_from].attrs:
-        axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
+        try:
+            axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
+        except AttributeError as e: # Weird string vs bytes thing, ehhh
+            axes = str(cxi_file[pull_from].attrs['axes']).split(':')
+
         axes = [axis.strip().lower() for axis in axes]
     else:
         axes = None

examples/transmission_RPI.py

@@ -34,7 +34,7 @@ dataset.get_as(device='cuda')
 # The regularization is an L2 regularizer that empirically helps accelerate
 # convergence
 for i, loss in enumerate(model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05])):#0.1)):
-    #model.inspect(dataset)
+    model.inspect(dataset)
     print(i,loss)
     
 #model.inspect(dataset)