diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index bfe5c15..7f71fb9 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -26,18 +26,18 @@ class PolarizedFancyPtycho(FancyPtycho): polarizer_offsets=None, analyzer_offsets=None, polarizer_scale=1, analyzer_scale=1, mask=None, weights = None, translation_scale = 1, saturation=None, - probe_support = None, obj_support=None, oversampling=1, + probe_support = None, oversampling=1, loss='amplitude mse',units='um'): - super(FancyPtycho, self).__init__(wavelength, detector_geometry, + super(PolarizedFancyPtycho, self).__init__(wavelength, detector_geometry, probe_basis, probe_guess, obj_guess, detector_slice=None, surface_normal=np.array([0.,0.,1.]), min_translation = t.Tensor([0,0]), background = None, translation_offsets=None, mask=None, - weights = None, translation_scale = 1, saturation=None, - probe_support = None, obj_support=None, oversampling=1, + weights = weights, translation_scale = 1, saturation=None, + probe_support = None, oversampling=1, loss='amplitude mse',units='um') if polarizer_offsets is None: @@ -49,7 +49,7 @@ class PolarizedFancyPtycho(FancyPtycho): self.analyzer_offsets = None else: self.analyzer_offsets = t.nn.Parameter(t.tensor(analyzer_offsets).to(dtype=t.float32)) / analyzer_scale - + self.polarizer = polarizer self.analyzer = analyzer probe_guess = t.tensor(probe_guess, dtype=t.complex64) @@ -63,7 +63,8 @@ class PolarizedFancyPtycho(FancyPtycho): @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=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True): - model = FancyPtycho.from_dataset(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, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True) + # When using this method, remember to pass through the inputs + model = FancyPtycho.from_dataset(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, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um') # Mutate the class to its subclass @@ -81,10 +82,13 @@ class PolarizedFancyPtycho(FancyPtycho): probe = t.stack([probe, ] + probe_stack) print('probe', type(probe), probe.shape) model.probe.data = probe + print(model.probe.shape) # obj = t.stack((model.obj.data, model.obj.data), dim=-3) # model.obj.data = t.stack((obj.data, obj.data), dim=-4) # obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) obj = model.obj.detach() + # Abe - Probably something identity matrix-like would be a better + # initialization (e.g. ((obj,0*obj),(0*obj,obj)) obj = t.stack((obj, obj), dim=-3) obj = t.stack((obj, obj), dim=-4) print('object', type(obj), obj.shape) @@ -275,12 +279,12 @@ class PolarizedFancyPtycho(FancyPtycho): raise NotImplementedError('Unstable Modes not Implemented for polarized light') pol_probes = polarization.apply_linear_polarizer(prs, polarizer) + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( - prs, self.obj, pix_trans, + pol_probes, self.obj, pix_trans, shift_probe=True, multiple_modes=True, polarized=True) analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer) - return analyzed_exit_waves diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index 7f60942..246698a 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/CDTools/tools/interactions/interactions.py @@ -423,9 +423,10 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi exit_waves : torch.Tensor An (N)x(P)xMxL tensor of the calculated exit waves """ + single_translation = False if translations.dim() == 1: - translations = translations[None,:] + translations = translations[None, :] single_translation = True # Separate the translations into a part that chooses the window @@ -438,8 +439,8 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi tr[1]:tr[1]+probe.shape[-1]] for tr in integer_translations]) else: - selections = t.stack([obj[:, :,tr[0]:tr[0]+probe.shape[-2], - tr[1]:tr[1]+probe.shape[-1]] + 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 @@ -480,7 +481,6 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi else: raise NotImplementedError('Object shift not yet implemented') - print('ptyvho 2d sinc', output.shape) if single_translation: return output[0] else: diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index 1f77d6f..f1451b2 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -17,7 +17,11 @@ from matplotlib import ticker, patheffects __all__ = ['colorize', 'plot_amplitude', 'plot_phase', 'plot_colorized', 'plot_translations', 'get_units_factor', 'plot_nanomap', 'plot_real', 'plot_imag', - 'plot_nanomap_with_images'] + 'plot_nanomap_with_images', + 'polarized_plot_component_amplitudes', + 'polarized_plot_phase_ret', + 'polarized_plot_global_phases', + 'polarized_plot_ellipses'] def colorize(z): diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index f3bc04c..01c9525 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -27,6 +27,7 @@ def generate_linear_polarizer(pol_angle): if pol_angle.dim() == 0: pol_angle = t.unsqueeze(pol_angle,0) single_angle = True + pol_angle_rad = t.deg2rad(pol_angle) a = t.cos(pol_angle_rad) ** 2 b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad)