.

CDTools/models/polarized_fancy_ptycho.py

@@ -118,15 +118,15 @@ class PolarizedFancyPtycho(FancyPtycho):
         return out[..., None, :, :]
 
     def forward_propagator(self, wavefields):
-        return tools.propagators.far_field
+        return tools.propagators.far_field(wavefields)
 
     def backward_propagator(self, wavefields):
-        return tools.propagators.inverse_far_field
+        return tools.propagators.inverse_far_field(wavefields)
 
     
     def measurement(self, wavefields):
-        wavefields_x = wavefields[..., 0, :, :, :]
-        wavefields_y = wavefields[..., 1, :, :, :]
+        wavefields_x = wavefields[..., 0, :, :]
+        wavefields_y = wavefields[..., 1, :, :]
         out_x = tools.measurements.quadratic_background(wavefields_x,
                             self.background,
                             detector_slice=self.detector_slice,
@@ -135,7 +135,7 @@ class PolarizedFancyPtycho(FancyPtycho):
                             oversampling=self.oversampling)
         # now, set bckgr to 0 since t shouldn't be calculated twice
         out_y = tools.measurements.quadratic_background(wavefields_y,
-                            None,
+                            0,
                             detector_slice=self.detector_slice,
                             measurement=tools.measurements.incoherent_sum,
                             saturation=self.saturation,

CDTools/tools/polarization/polarization.py

@@ -48,12 +48,14 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
 	probe: t.Tensor
 		A (N)(P)x2xMxL tensor representing the probe
 	jones_matrix: t.tensor
-		(N)x2x2 
+		(N)x2x2x(M)x(L) 
 
 	Returns:
 	--------
 	a probe with the jones matrix applied: t.Tensor
 		(N)(P)x2xMxL 
+
+	Assume that if the probe has a dimension (N), so does the jones matrix
 	"""
 	if multiple_modes:
 		if transpose:
@@ -80,10 +82,13 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
 			if len(jones_matrix.shape) < 4:
 				jones_matrix = jones_matrix[..., None, None]
 			probe = probe[..., None, :, :]
+			print('probs', probe.shape)
+			print('joness', jones_matrix.shape)
 			probe = probe.transpose(-1, -3).transpose(-2, -4)
 			jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
 			output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
 
+
 		else:
 			if len(jones_matrix.shape) < 4:
 				jones_matrix = jones_matrix[..., None, None]