not tested/debugged yet

CDTools/datasets/polarized_ptycho_2d_dataset.py

@@ -60,9 +60,19 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
         
         super(PolarizedPtycho2DDataset,self).__init__(translations, patterns,
                                              *args, **kwargs)
-        self.polarizer = t.tensor(polarizer_angles, dtype=t.float32)
-        self.analyzer = t.tensor(analyzer_angles, dtype=t.float32)
+        
+        
+        # self.polarizer = t.tensor(polarizer_angles, dtype=t.float32)
+        # self.analyzer = t.tensor(analyzer_angles, dtype=t.float32)
+        
 
+        polarizer = []
+        analyzer = []
+        for k in t.tensor(translations).shape[0]:
+            polarizer.append((k//3)%3)
+            analyzer.append((k%3))
+        self.polarizer = t.tensor(polarizer)
+        self.analyzer = t.tensor(analyzer)
 
     def _load(self, index):
         """ Internal function to load data
@@ -147,8 +157,17 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
         dataset.__class__ = cls
 
         # Now, we save out the polarizer and analyzer states
-        polarizer = cdtdata.get_shot_to_shot_info(cxi_file, 'polarizer_angle')
-        analyzer = cdtdata.get_shot_to_shot_info(cxi_file, 'analyzer_angle')
+        # polarizer = cdtdata.get_shot_to_shot_info(cxi_file, 'polarizer_angle')
+        # analyzer = cdtdata.get_shot_to_shot_info(cxi_file, 'analyzer_angle')
+
+        polarizer = []
+        analyzer = []
+        for k in range(dataset.translations.shape[0]):
+            polarizer.append((k//3)%3)
+            analyzer.append((k%3))
+        dataset.polarizer = t.tensor(polarizer)
+        dataset.analyzer = t.tensor(analyzer)
+
 
         dataset.analyzer = t.tensor(analyzer, dtype=t.float32)
         dataset.polarizer = t.tensor(polarizer, dtype=t.float32)

CDTools/tools/initializers/initializers.py

@@ -5,6 +5,7 @@ initialize the reconstrucions, and the heuristic calculations for
 geierating sensible initializations for the probe guess.
 """
 
+import math
 import numpy as np
 import torch as t
 from CDTools.tools.propagators import *
@@ -210,7 +211,7 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
 
 
 
-def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
+def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0, polarized=False, left_polarized=True):
     """Initializes a gaussian probe based on experimental parameters
 
     This function generates a gaussian probe initialization which has a
@@ -272,13 +273,33 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
         
     # Finally, we should calculate the average pattern intensity from the
     # dataset and normalize the gaussian probe. This should be done by
-    avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))]
+    if not polarized:
+        avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))]
+
+    else:
+        polarizer = dataset.polarizer.tolist()
+        analyzer = dataset.analyzer.tolist()
+        factors = [(math.cos(math.radians(polarizer[idx] - analyzer[idx])))**2 for idx in range(len(dataset)) if (abs(polarizer[idx] - analyzer[idx]) > 5)]
+        avg_intensities = [t.sum(dataset[idx][1]) / factor[idx] for idx in range(len(dataset))]
+
     avg_intensity = t.mean(t.tensor(avg_intensities))
     probe_intensity = t.sum(t.abs(probe)**2)
-    return avg_intensity / probe_intensity * probe
-    
+    probe = t.sqrt(avg_intensity / probe_intensity) * probe
 
-def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, oversampling=1):
+    if polarized:
+        if left_polarized == True:
+            x = 1j
+        else:
+            x = -1j
+        probe = t.stack((probe.to(dtype=t.cfloat), x * probe.to(dtype=t.cfloat)), dim=-4)
+        return probe[..., None, :, :]
+
+    return probe
+
+
+
+
+def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, oversampling=1, polarized=False, left_polarized=True):
     """Generates a SHARP style probe guess from a dataset
 
     What we call the "SHARP" style probe guess is to take a mean of all
@@ -320,13 +341,22 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
 
     # to use the mask or not?
     intensities = np.zeros([dim // oversampling for dim in shape])
+    
+    if polarized:
+        factors = [(math.cos(math.radians(polarizer[idx] - analyzer[idx])))**2 for idx in range(len(dataset)) if abs(polarizer[idx] - analyzer[idx]) > 5]
+    else:
+        factors = [1 for idx in range(len(dataset))]
+
+
     for params, im in dataset:
         if hasattr(dataset,'mask') and dataset.mask is not None:
-            intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy()
+            intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy() / factors[params[0]]
         else:
-            intensities[det_slice] += im.cpu().numpy()         
+            intensities[det_slice] += im.cpu().numpy() / params[factors[0]]       
     intensities /= len(dataset)
 
+
+
     # Subtract off a known background if it's stored
     if hasattr(dataset, 'background') and dataset.background is not None:
         intensities[det_slice] = np.clip(intensities[det_slice] - dataset.background.cpu().numpy(), a_min=0,a_max=None)
@@ -374,6 +404,14 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
     top = shape[1]//2 - probe_guess.shape[1] // 2 
     final_probe[left:left+probe_guess.shape[0],
                 top:top+probe_guess.shape[1]] = probe_guess
+
+    if polarized:
+        if left_polarized:
+            x = 1j
+        else:
+            x = -1j
+        final_probe = t.stack((final_probe.to(dtype=t.cfloat), final_probe.to(dtype=t.cfloat)), dim=-4)
+        final_probe = final_probe[..., None, :, :]
     
     return final_probe
 

CDTools/tools/interactions/interactions.py

@@ -7,7 +7,7 @@ for ptychographic reconstruction.
 
 import torch as t
 import numpy as np
-from CDTools.tools import propagators, image_processing
+from CDTools.tools import propagators, image_processing, polarization
 
 __all__ = ['translations_to_pixel', 'pixel_to_translations',
            'project_translations_to_sample',
@@ -385,7 +385,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
         return t.stack(exit_waves)
 
 
-def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True):
+def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, polarized=False, polarizer=None, analyzer=None):
     """Returns a stack of exit waves accounting for subpixel shifts
  
     This function returns a collection of exit waves, with the first
@@ -437,6 +437,13 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
     selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
                               tr[1]:tr[1]+probe.shape[-1]]
                           for tr in integer_translations])
+    if polarized:
+        polarized_probes = t.stack([polarization.apply_linear_polarizer(probe, polarizer[idx]) for idx in range(polarizer.shape)])
+        # Nx(P)x2x1xMxL tensor
+        selections = t.stack([obj[:, :,tr[0]:tr[0]+probe.shape[-2],
+                          tr[1]:tr[1]+probe.shape[-1]]
+                      for tr in integer_translations])
+        # Nx2x2xMxL tensor
     
     exit_waves = []
     if shift_probe:
@@ -449,22 +456,44 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
         J = 2 * np.pi * J / probe.shape[-1]
         phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
                                 -subpixel_translations[:,1,None,None]*J))
+
+        if polarized:
+            phase_masks = t.stack((phase_masks, phase_masks), dim=-4)
+            phase_masks = t.unsqueeze(phase_masks, dim=-3)
+            # Nx2x1xMxL tensor
+            # probe is (P)x2x1xMxL tensor
         
         fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
 
+
         if multiple_modes: # Multi-mode probe
-            shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
+            if polarized:
+                shifted_fft_probe = fft_probe * phase_masks[...,None,:,:,:,:]
+            else:
+                shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
         else:
             shifted_fft_probe = fft_probe * phase_masks
 
+
         shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
                                                     dim=(-1,-2)))
 
-        if multiple_modes: # Multi-mode probe
-            output = shifted_probe * selections[...,None,:,:]
+        if not polarized:
+            if multiple_modes: # Multi-mode probe
+                output = shifted_probe * selections[...,None,:,:]
+            else:
+                output = shifted_probe * selections
+        # selections: Nx2x2xMxL
+        # probe: Nx(P)x2x1xMxL
         else:
-            output = shifted_probe * selections
-        
+            if multiple_modes:
+                selections = selections[..., None, :,:,:,:]
+            shift_probe = shifted_probe.transpose(-1, -3).transpose(-2, -4)
+            selections = selections.transpose(-1, -3).transpose(-2, -4)
+            output = t.matmul(selections, shift_probe).transpose(-1, -3).transpose(-2, -4)
+            
+            output = t.stack([polarization.apply_linear_polarizer(output[idx, ...], analyzer[idx]) for  idx in range(analyzer.shape)])
+    
     else:
         raise NotImplementedError('Object shift not yet implemented')
 

CDTools/tools/polarization/polarization.py

@@ -4,6 +4,13 @@ import math
 from math import sin
 from math import cos
 
+__all__ = ['apply_linear_polarizer',
+		   'apply_phase_retardance',
+		   'apply_half_wave_plate',
+		   'apply_quarter_wave_plate',
+		   'apply_circular_polarizer',
+		   'apply_jones_matrix']
+		   
 def apply_linear_polarizer(probe, polar_angle):
 	"""
 	Applies a linear polarizer to the probe