From 998245ae4aabf3fca1800fc75c6f01c92feff27f Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 13 Aug 2021 14:17:12 -0400 Subject: [PATCH] '.' --- CDTools/models/polarized_fancy_ptycho.py | 9 ++---- CDTools/tools/interactions/interactions.py | 24 +++++--------- CDTools/tools/polarization/polarization.py | 37 +++++++++++++--------- 3 files changed, 33 insertions(+), 37 deletions(-) diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index d2483f3..ed21fa1 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -52,7 +52,7 @@ 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, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', 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') # Mutate the class to its subclass @@ -97,13 +97,10 @@ class PolarizedFancyPtycho(FancyPtycho): prs = Ws[...,None,None,None,None] * basis_prs else: raise NotImplementedError('Unstable Modes not Implemented for polarized light') - - polarizer = tools.polarization.generate_linear_polarizer(polarizer) - analyzer = tools.polarization.generate_linear_polarizer(analyzer) - pol_probes = polarization.apply_linear_polarizer(prs, polarizer) + pol_probes = polarization.apply_linear_polarizer(prs, polarizer) exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( - prs, self.obj_support * self.obj,pix_trans, + prs, self.obj, pix_trans, shift_probe=True, multiple_modes=True, polarized=True) analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer) diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index 1a70400..d2d8667 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/CDTools/tools/interactions/interactions.py @@ -43,7 +43,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1. pixel_translations : torch.Tensor A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation """ - + projection_1 = t.Tensor([[1,0,0], [0,1,0], [0,0,0]]).to(device=translations.device,dtype=translations.dtype) @@ -433,16 +433,16 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi integer_translations = t.floor(translations) subpixel_translations = translations - integer_translations integer_translations = integer_translations.to(dtype=t.int32) - - 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: + if not polarized: + selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2], + 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]] for tr in integer_translations]) # Nx2x2xMxL tensor - + exit_waves = [] if shift_probe: i = t.arange(probe.shape[-2],device=probe.device,dtype=t.float32) \ @@ -454,15 +454,11 @@ 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 = phase_masks[..., None, :, :] # Nx2x1xMxL tensor # probe is (N)(P)x2xMxL tensor - fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2)) - - if multiple_modes: # Multi-mode probe if polarized: shifted_fft_probe = fft_probe * phase_masks[...,None,:,:,:] @@ -470,11 +466,8 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi 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 not polarized: if multiple_modes: # Multi-mode probe output = shifted_probe * selections[...,None,:,:] @@ -483,8 +476,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi # selections: Nx2x2xMxL # probe: Nx(P)x2x1xMxL else: - # print('F', shift_probe) - output = polarization.apply_jones_matrix(shifted_probe, selections) + output = polarization.apply_jones_matrix(shifted_probe, selections, multiple_modes=multiple_modes) else: raise NotImplementedError('Object shift not yet implemented') diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index 81d2edf..cb6da63 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -22,15 +22,18 @@ __all__ = ['apply_linear_polarizer', def generate_linear_polarizer(pol_angle): single_angle = False - pol_angle = t.as_tensor(pol_angle) + pol_angle = t.as_tensor(pol_angle).to(dtype=t.float32) if pol_angle.dim() == 0: pol_angle = t.unsqueeze(pol_angle,0) single_angle = True - pol_angle_rad = t.deg2rad(pol_angle) - jones_matrices = t.stack([t.tensor([[(t.cos(p)) ** 2, t.sin(p) * t.cos(p)], - [t.sin(p) * t.cos(p), (t.sin(p)) ** 2]]) - for p in pol_angle_rad]) + a = t.cos(pol_angle_rad) ** 2 + b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad) + c = b + d = t.sin(pol_angle_rad) ** 2 + ab = t.stack((a, b), dim=-1) + cd = t.stack((c, d), dim=-1) + jones_matrices = t.stack((ab, cd), dim=-2) if single_angle: return jones_matrices[0].to(dtype=t.cfloat) else: @@ -54,11 +57,15 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True linearly polarized probe: t.Tensor (N)(P)x2x1xMxL """ - jones_matrices = generate_linear_polarizer(polarization) + jones_matrices = generate_linear_polarizer(polarizer) return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes) def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True): + # print('probe', probe.shape, 'jones matrix', jones_matrix.shape) + # if jones_matrix.shape == t.Size([5, 2, 2, 2, 2]): + # print('probe', probe.shape, 'jones', jones_matrix.shape) + # print('JONES', jones_matrix) """ Applies a given Jones matrix to the probe @@ -87,6 +94,7 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True) # vice versa elif jones_matrix.dim() > probe.dim(): probe = probe.unsqueeze(0) + # print('apply jonesmatrix: probe', probe.shape, 'matrix:', jones_matrix) jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4) probe = probe.transpose(-1, -3).transpose(-2, -4) output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3) @@ -97,7 +105,7 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True) return output -def apply_phase_retardance(probe, phase_shift): +def apply_phase_retardance(probe, phase_shift, multiple_modes=True): """ Shifts the y-component of the field wrt the x-component by a given phase shift @@ -115,11 +123,11 @@ def apply_phase_retardance(probe, phase_shift): """ probe = probe.to(dtype=t.cfloat) jones_matrix = t.tensor([[1, 0], [0, phase_shift]]).to(dtype=t.cfloat) - polarized = apply_jones_matrix(probe, jones_matrix) + polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) return polarized -def apply_circular_polarizer(probe, left_polarized=True): +def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True): """ Applies a circular polarizer to the probe @@ -140,11 +148,10 @@ def apply_circular_polarizer(probe, left_polarized=True): jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]])).to(dtype=t.cfloat) else: jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]]).to(dtype=t.cfloat) - polarized = apply_jones_matrix(probe, jones_matrix) - + polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) return polarized -def apply_quarter_wave_plate(probe, fast_axis_angle): +def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True): """ Parameters: ---------- @@ -162,11 +169,11 @@ def apply_quarter_wave_plate(probe, fast_axis_angle): theta = math.radians(fast_axis_angle) exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2)) jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]]).to(dtype=t.cfloat) - out = apply_jones_matrix(probe, jones_matrix) + out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) return out -def apply_half_wave_plate(probe, fast_axis_angle): +def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True): """ Parameters: ---------- @@ -184,7 +191,7 @@ def apply_half_wave_plate(probe, fast_axis_angle): theta = math.radians(fast_axis_angle) exponent = t.exp(-1j * math.pi / 2 * t.ones(2, 2)) jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]]).to(dtype=t.cfloat) - out = apply_jones_matrix(probe, jones_matrix) + out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) return out \ No newline at end of file