From 530a7717acb382d7b68434e62c3b6ccdf8b21bc0 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Wed, 28 Jul 2021 00:40:15 -0400 Subject: [PATCH] . --- .../datasets/polarized_ptycho_2d_dataset.py | 4 +- CDTools/tools/polarization.py | 29 ++++ CDTools/tools/polarization/polarization.py | 163 ++++++++++++++++++ 3 files changed, 194 insertions(+), 2 deletions(-) create mode 100644 CDTools/tools/polarization.py create mode 100644 CDTools/tools/polarization/polarization.py diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py index 95c426c..1c59cdd 100644 --- a/CDTools/datasets/polarized_ptycho_2d_dataset.py +++ b/CDTools/datasets/polarized_ptycho_2d_dataset.py @@ -138,7 +138,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): # If a bare string is passed if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): with h5py.File(cxi_file, 'r') as f: - return cls.from_cxi(f, polarized) + return cls.from_cxi(f) # Generate a base dataset dataset = Ptycho2DDataset.from_cxi(cxi_file) @@ -172,7 +172,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): # If a bare string is passed if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): with cdtdata.create_cxi(cxi_file) as f: - return self.to_cxi(f, polarized) + return self.to_cxi(f) # This saves the translations, patterns, etc. super(PolarizedPtycho2DDataset, self).to_cxi(cxi_file) diff --git a/CDTools/tools/polarization.py b/CDTools/tools/polarization.py new file mode 100644 index 0000000..b47ecf8 --- /dev/null +++ b/CDTools/tools/polarization.py @@ -0,0 +1,29 @@ +import numpy as numpy +import torch as t + +def polarize_probe_linear(probe, polar_angle): + """ + Applies a linear polarizer to the probe + + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe + polar_angle: float + The angle between the fast-axis of the linear polarizer and the x-axis + """ + + polarizer = t.tensor([[(math.cos(theta)) ** 2, math.sin(2 * theta) / 2], [math.sin(2 * theta) / 2, (math.sin(theta)) ** 2]]) + return t.tensordot(polarizer, probe, dims=([-3], [-4])) + +def apply_jones_matrix(probe, jones_matrix): + """ + Applies the given Jones matrix to the probe + + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe + jones_matrix: t.tensor + 2x2 + """ \ No newline at end of file diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py new file mode 100644 index 0000000..87b551a --- /dev/null +++ b/CDTools/tools/polarization/polarization.py @@ -0,0 +1,163 @@ +import numpy as numpy +import torch as t +import math +from math import sin +from math import cos + +def apply_linear_polarizer(probe, polar_angle): + """ + Applies a linear polarizer to the probe + + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe + polar_angle: float + The angle between the fast-axis of the linear polarizer and the horizontal axis + + Returns: + -------- + linearly polarized probe: t.Tensor + (...)x2x1xMxL + """ + probe = probe.to(dtype=t.cfloat) + theta = math.radians(polar_angle) + polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat) + # I haven't figured out how to multiply tensors using tensordot yet, + # so we'll be temporarily using matmul on the previously tranposed vector + # (since it returns the matrix multiplication product over the last two dimensions + + #Swap the dimensions for the prober to be (...)xMxLx2x1 to perform matmul on it + probe = probe.transpose(-1, -3).transpose(-2, -4) + polarized_probe = t.matmul(polarizer, probe) + + # Transpose it back + return polarized_probe.transpose(-1, -3).transpose(-2, -4) + +def apply_jones_matrix(probe, jones_matrix): + """ + Applies a given Jones matrix to the probe + + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe + jones_matrix: t.tensor + (...)x2x2 + + Returns: + -------- + linearly polarized probe: t.Tensor + (...)x2x1xMxL + """ + probe = probe.to(dtype=t.cfloat) + probe = probe.transpose(-1, -3).transpose(-2, -4) + polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe) + + # Transpose it back + return polarized_probe.transpose(-1, -3).transpose(-2, -4) + +def apply_phase_retardance(probe, phase_shift): + """ + Shifts the y-component of the field wrt the x-component by a given phase shift + + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe + phase_shift: float + phase shift in degrees + + Returns: + -------- + probe: t.Tensor + (...)x2x1xMxL + """ + probe = probe.to(dtype=t.cfloat) + jones_matrix = t.tensor([[1, 0], [0, phase_shift]]) + probe = probe.transpose(-1, -3).transpose(-2, -4) + polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe) + + # Transpose it back + return polarized_probe.transpose(-1, -3).transpose(-2, -4) + +def apply_circular_polarizer(probe, left_polarized=True): + """ + Applies a circular polarizer to the probe + + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe + left_polarizd: bool + True for the left-polarization, False for the right + + Returns: + -------- + circularly polarized probe: t.Tensor + (...)x2x1xMxL + """ + probe = probe.to(dtype=t.cfloat) + if left_polarized: + jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]])) + else: + jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]]) + probe = probe.transpose(-1, -3).transpose(-2, -4) + polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe) + + # Transpose it back + return polarized_probe.transpose(-1, -3).transpose(-2, -4) + +def apply_quarter_wave_plate(probe, fast_axis_angle): + """ + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe + fast_axis_angle: float + The angle between the fast-axis of the polarizer and the horizontal axis + + Returns: + -------- + polarized probe: t.Tensor + (...)x2x1xMxL + """ + probe = probe.to(dtype=t.cfloat) + 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]]) + probe = probe.transpose(-1, -3).transpose(-2, -4) + polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe) + # Transpose it back + return polarized_probe.transpose(-1, -3).transpose(-2, -4) + + +def apply_half_wave_plate(probe, fast_axis_angle): + """ + Parameters: + ---------- + probe: t.Tensor + A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe + fast_axis_angle: float + The angle between the fast-axis of the polarizer and the horizontal axis + + Returns: + -------- + polarized probe: t.Tensor + (...)x2x1xMxL + """ + probe = probe.to(dtype=t.cfloat) + 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]]) + probe = probe.transpose(-1, -3).transpose(-2, -4) + polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe) + # Transpose it back + return polarized_probe.transpose(-1, -3).transpose(-2, -4) + +# probe = t.rand(3, 2, 1, 5, 6) +# print(apply_linear_polarizer(probe, 30).shape) +# print(apply_circular_polarizer(probe).shape) +# print(apply_phase_retardance(probe, 29).shape) +# print(apply_half_wave_plate(probe, 29).shape) +# print(apply_quarter_wave_plate(probe, 29).shape)