From 899317ae0f1a3325af9505a9a54db9fead0f1609 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Wed, 28 Jul 2021 06:59:18 -0400 Subject: [PATCH] not tested/debugged yet --- .../datasets/polarized_ptycho_2d_dataset.py | 27 ++++++++-- CDTools/tools/initializers/initializers.py | 52 ++++++++++++++++--- CDTools/tools/interactions/interactions.py | 43 ++++++++++++--- CDTools/tools/polarization/polarization.py | 7 +++ 4 files changed, 111 insertions(+), 18 deletions(-) diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py index 1c59cdd..93deb31 100644 --- a/CDTools/datasets/polarized_ptycho_2d_dataset.py +++ b/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) diff --git a/CDTools/tools/initializers/initializers.py b/CDTools/tools/initializers/initializers.py index b61b5d5..5d6791c 100644 --- a/CDTools/tools/initializers/initializers.py +++ b/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 diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index e7fd7cc..2be19a1 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/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') diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index 87b551a..b687b22 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/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