diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py index 93deb31..b1e22a7 100644 --- a/CDTools/datasets/polarized_ptycho_2d_dataset.py +++ b/CDTools/datasets/polarized_ptycho_2d_dataset.py @@ -68,9 +68,9 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): polarizer = [] analyzer = [] - for k in t.tensor(translations).shape[0]: - polarizer.append((k//3)%3) - analyzer.append((k%3)) + for k in range(t.tensor(translations).shape[0]): + polarizer.append((k//3)%3 * 45) + analyzer.append((k%3 * 45)) self.polarizer = t.tensor(polarizer) self.analyzer = t.tensor(analyzer) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 5ceca97..3c26947 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -42,6 +42,8 @@ import time #import pytorch_warmup from .complex_adam import MyAdam from .complex_lbfgs import MyLBFGS +from matplotlib.backends.backend_pdf import PdfPages + __all__ = ['CDIModel'] @@ -497,12 +499,16 @@ class CDIModel(t.nn.Module): try: plotter(self,fig) plt.title(name) - + print('f', idx) + plt.savefig('img-{0}.pdf'.format(idx), bbox_inches='tight') + except TypeError as e: if dataset is not None: try: plotter(self, fig, dataset) plt.title(name) + print('f', idx) + plt.savefig('img-{0}.pdf'.format(idx), bbox_inches='tight') except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e: pass diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index dd5a58b..1f072a5 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -5,196 +5,201 @@ 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'] - + 'apply_phase_retardance', + 'apply_half_wave_plate', + 'apply_quarter_wave_plate', + 'apply_circular_polarizer', + 'apply_jones_matrix'] +print() + def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True): - """ - Applies a linear polarizer to the probe + """ + Applies a linear polarizer to the probe - Parameters: - ---------- - probe: t.Tensor - A (N)(P)x2xMxL tensor representing the probe, MxL - the size of the probe - The angle between the fast-axis of the linear polarizer and the horizontal axis - polarizer: t.Tensor - A 1D tensor (N) representing the polarizer angles for each of the patterns (or a single tensor of shape (1)) + Parameters: + ---------- + probe: t.Tensor + A (N)(P)x2xMxL tensor representing the probe, MxL - the size of the probe + The angle between the fast-axis of the linear polarizer and the horizontal axis + polarizer: t.Tensor + A 1D tensor (N) representing the polarizer angles for each of the patterns (or a single tensor of shape (1)) - Returns: - -------- - linearly polarized probe: t.Tensor - (N)(P)x2x1xMxL - """ + Returns: + -------- + linearly polarized probe: t.Tensor + (N)(P)x2x1xMxL + """ + # if len(polarizer.shape) == 0: + # polarizer = t.tensor([polarizer]) + if len(polarizer,shape) == 0: + polarizer = t.tensor([polarizer]) + pol_cos = lambda idx: cos(math.radians(polarizer[idx])) + pol_sin = lambda idx: sin(math.radians(polarizer[idx])) + jones_matrices = t.stack(([t.tensor([[(pol_cos(idx)) ** 2, pol_sin(idx) * pol_cos(idx)], [pol_sin(idx) * pol_cos(idx), (pol_sin(idx)) ** 2]]).to(dtype=t.cfloat) for idx in range(len(polarizer))])) - if len(polarizer) == 1: - theta = math.radians(polarizer) - jones_matrices = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat) - else: - pol_cos = lambda idx: cos(math.radians(polarizer[idx])) - pol_sin = lambda idx: sin(math.radians(polarizer[idx])) - jones_matrices = t.stack(([t.tensor([[(pol_cos(idx)) ** 2, pol_sin(idx) * pol_cos(idx)], [pol_sin(idx) * pol_cos(idx), (pol_sin(idx)) ** 2]]).to(dtype=t.cfloat) for idx in range(len(polarizer))])) - - return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes) + 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): - """ - Applies a given Jones matrix to the probe + """ + Applies a given Jones matrix to the probe - Parameters: - ---------- - probe: t.Tensor - A (N)(P)x2xMxL tensor representing the probe - jones_matrix: t.tensor - (N)x2x2x(M)x(L) + Parameters: + ---------- + probe: t.Tensor + A (N)(P)x2xMxL tensor representing the probe + jones_matrix: t.tensor + (N)x2x2x(M)x(L) - Returns: - -------- - a probe with the jones matrix applied: t.Tensor - (N)(P)x2xMxL + 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: - if len(jones_matrix.shape) >= 4: - jones_matrix = jones_matrix[..., None, :, :, :, :] - else: - jones_matrix = jones_matrix[..., None, :, :, None, None] - probe = probe[..., None, :, :] - jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4) - # (N)1xMxLx2x2 or (N)1x1x1x2x2 - probe = probe.transpose(-1, -3).transpose(-2, -4) - output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3) - # (N)Px2xMxL + Assume that if the probe has a dimension (N), so does the jones matrix + """ + if multiple_modes: + if transpose: + if len(jones_matrix.shape) >= 4: + jones_matrix = jones_matrix[..., None, :, :, :, :] - else: - if len(jones_matrix.shape) < 4: - jones_matrix = jones_matrix[..., None, :, :, :, :] - probe = t.stack((probe, probe), dim=-4) - output = t.sum(jones_matrix * probe, dim=-3) - #(N)x2xMxL + else: + jones_matrix = jones_matrix[..., None, :, :, None, None] + probe = probe[..., None, :, :] + # if jones matrices do not differ from pattern to pattern + if len(probe.shape) > len(jones_matrix.shape): + jones_matrix = jones_matrix[None, ...] + jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4) + # (N)1xMxLx2x2 or (N)1x1x1x2x2 + probe = probe.transpose(-1, -3).transpose(-2, -4) + output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3) + # (N)Px2xMxL - else: - if transpose: - 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, :, :, :, :] + probe = t.stack((probe, probe), dim=-4) + output = t.sum(jones_matrix * probe, dim=-3) + #(N)x2xMxL + + else: + if transpose: + if len(jones_matrix.shape) < 4: + jones_matrix = jones_matrix[..., None, None] + probe = probe[..., None, :, :] + # if jones matrices do not differ from pattern to pattern + if len(probe.shape) > len(jones_matrix.shape): + jones_matrix = jones_matrix[None, ...] + 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] - probe = t.stack((probe, probe), dim=-4) - output = t.sum(jones_matrix * probe, dim=-3) - - return output + else: + if len(jones_matrix.shape) < 4: + jones_matrix = jones_matrix[..., None, None] + probe = t.stack((probe, probe), dim=-4) + output = t.sum(jones_matrix * probe, dim=-3) + + return output def apply_phase_retardance(probe, phase_shift): - """ - Shifts the y-component of the field wrt the x-component by a given 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 + 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) + 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) + # 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 + """ + 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) + 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) + # 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 + """ + 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) + 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 + """ + 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) + 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(17, 7, 2, 6, 4) # polarizer = t.rand(7) diff --git a/tests/test_polarization.py b/tests/test_polarization.py new file mode 100644 index 0000000..f4ac70c --- /dev/null +++ b/tests/test_polarization.py @@ -0,0 +1,438 @@ +from __future__ import division, print_function, absolute_import +import numpy as np +import torch as t +from copy import copy +import h5py +import pathlib +from CDTools.datasets import CDataset, Ptycho2DDataset, PolarizedPtycho2DDataset +from CDTools.tools import data as cdtdata, initializers, polarization, interactions +from CDTools.tools.polarization import apply_jones_matrix as jones +from CDTools.tools import plotting +from torch.utils import data as torchdata +from matplotlib import pyplot as plt +from matplotlib.widgets import Slider +from matplotlib import ticker +from CDTools.models import FancyPtycho, PolarizedFancyPtycho +from torch.utils import data as torchdata +from math import cos as cos, sin as sin +import math + + + + +angle = 87 +angle_2 = angle - 45 + +def polarizer(angle): + theta = math.radians(angle) + polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat) + return polarizer +exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2)) +theta2 = math.radians(angle_2) +quarter_plate = t.tensor([[(cos(theta2))**2 + 1j * (sin(theta2))**2, (1 - 1j) * sin(theta2) * cos(theta2)], + [(1 - 1j) * sin(theta2) * cos(theta2), (sin(theta2))**2 + 1j * (cos(theta2))**2]]) + +def build_from_quarters(jones1, jones2, jones3, jones4): + x = t.cat((t.stack((jones1, jones1), dim=-1), t.stack((jones2, jones2), dim=-1)), dim=-1) + y = t.cat((t.stack((jones3, jones3), dim=-1), t.stack((jones4, jones4), dim=-1)), dim=-1) + x = t.stack((x, x), dim=-2) + y = t.stack((y, y), dim=-2) + return t.cat((x, y), dim=-2).to(dtype=t.cfloat) + +jones_plate = t.matmul(quarter_plate, polarizer(angle)) +jones0 = polarizer(0) +jones90 = polarizer(90) +jones45 = polarizer(45) + + +''' +after applying the polarizer and the quarter_plate, the probe should get circularly polarized +probe: no multiple modes, 1 diffr pattern + 2xMxL +jones_matrix: same jones matrix applied to all the pixels + 2x2 +''' +transpose = True +def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr(): + probe = t.rand(2, 4, 4, dtype=t.cfloat) + print(polarizer) + out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose), + quarter_plate, multiple_modes=False, transpose=transpose) + print('expected shape:(2, 3, 4)') + print('actual:', out.shape) + print('simulated:', out) + + assert np.allclose(np.real(out[0]), np.imag(out[1])) + + +'''probe: no multiple modes, 1 diffr pattern + 2xMxL = 2x4x4 +jones_matrix: jones matrices differ from pixel to pixel + 2x2xMxL = 2x2x4x4 +4 quarters: + 1: [:, :, :-2, :-2] - circular_polarizer, + 2: [:, :, :-2, -2:] - 0 + 3: [:, :, -2:, :-2] - 90 + 4: [:, :, -2:, -2:] - 45 +''' +def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr(): + jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45) + probe = t.ones(2, 4, 4).to(dtype=t.cfloat) + print('jones:', jones_matr) + # print('jones:', jones_matr) + out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose) + + # jones -> (2,2,x,y), probe, output probe + # interaction -> + print('expected shape:(2, 4, 4)') + print('simulated:', out.shape) + print('simulated:', out) + + assert (np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2])) + and t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, -2:, -2:], out[0, -2:, -2:])) + + +''' +probe: no multiple modes, multiple diffr patterns + Nx2xMxL = 3x2x4x4 +jones_matrix: same jones matrix applied to all the pixels + Nx2x2 = 3x2x2 + 3 different matrices for each probe: + 1: 0 + 2: 45 + 3: 90 +''' +def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr(): + probe = t.ones(2, 4, 4, dtype=t.cfloat) + probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0) + print('probe shape:', probe.shape, 'jones shape', jones_matr.shape) + out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose) + + print('expected shape: (3, 2, 4, 4)') + print('actual:', out.shape) + print('simulated:', out) + + assert (t.allclose(out[0, 0, :, :], t.ones(4, 4, dtype=t.cfloat)) + and t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat)) + and t.allclose(out[1, 0, :, :], out[1, 1]) + and t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat)) + and t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat))) + + +''' +probe: no multiple modes, multiple diffr pattern + Nx2xMxL = 3x2x4x4 +jones_matrix: jones matrices differ from pixel to pixel + Nx2x2xMxL = 3x2x2x4x4 + jones matrix for the 1st pattern: + 1: quat plate, 2: 90, 3: 0, 4: 45 + jones matrix for the 2nd pattern: + 1: 0, 2: 45, 3: 90, 4: quat plate + jones matrix for the 3rd pattern: + 1: 90, 2: 0, 3: 45, 4: quat plate +''' +def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr(): + jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), + build_from_quarters(jones0, jones45, jones90, jones_plate), + build_from_quarters(jones90, jones0, jones45, jones_plate)] + jones_matr = t.stack(([i for i in jones_m]), dim=0) + probe = t.ones(2, 4, 4, dtype=t.cfloat) + probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose) + + print('expected shape: (3, 2, 4, 4)') + print('actual shape:', out.shape) + print('simulated:', out) + + + assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2])) + and t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:]) + + and t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:]) + and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat)) + and np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:])) + + and t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2]) + and np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:]))) + + +''' +probe: multiple modes, 1 diffr pattern + Px2xMxL = 2x2x3x4 +jones_matrix: same jones matrix applied to all the pixels + 2x2 - quarter plate +''' +def test_apply_jones_matrix_mult_modes_1_pattern_one_jones_matr(): + probe = t.rand(2, 2, 3, 4, dtype=t.cfloat) + out = jones(jones(probe, polarizer(angle), multiple_modes=True, transpose=transpose), quarter_plate, multiple_modes=True, transpose=transpose) + + print('expected shape: (2, 2, 3, 4)') + print('actual:', out.shape) + print('simulated:', out) + + assert (np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :])) + and np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :]))) + + +''' +probe: multiple modes, 1 diffr pattern + Px2xMxL = 3x2x4x4 +jones_matrix: jones matrices differ from pixel to pixel + 2xMxL = 2x4x4 +4 quarters: + 1: [:, :, :-2, :-2] - circular_polarizer, + 2: [:, :, :-2, -2:] - 0 + 3: [:, :, -2:, :-2] - 90 + 4: [:, :, -2:, -2:] - 45 +''' +def test_apply_jones_matrix_mult_modes_1_pattern_diff_jones_matr(): + jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45) + probe = t.ones(2, 4, 4, dtype=t.cfloat) + probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose) + + print('expected shape: (3, 2, 4, 4)') + print('actual shape:', out.shape) + print('simulated:', out) + + assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2])) + and t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:]) + + and np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2])) + and t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:]) + + and np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2])) + and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])) + + +''' +probe: multiple modes, multiple diffr patterns + NxPx2xMxL = 3x7x2x3x4 +jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern) + Nx2x2 = 3x2x2 + 3 different matrices for each probe in one mode: + 1: 0 + 2: 45 + 3: 90 +''' +def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr(): + probe = t.ones(2, 3, 4, dtype=t.cfloat) + # 1st mode + probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + # 2nd mode + probe_mode2 = 10 * t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1) + jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0) + print('probe shape:', probe.shape, 'jones shape', jones_matr.shape) + out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose) + + print('expected shape: (3, 7, 2, 3, 4)') + print('actual:', out.shape) + print('simulated (patterns in one mode):', out[:, 6, :, :, :]) + + # we'll be checking only one mode + assert (t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat)) + and t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat)) + and t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :]) + and t.allclose(out[2, 6, 0, :, :], 3 * (10**6) * t.zeros(3, 4, dtype=t.cfloat)) + and t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat))) + + +''' +probe: multiple modes, multiple diffr pattern + NxPx2xMxL = 3x4x2x4x4 +jones_matrix: jones matrices differ from pixel to pixel + Nx2x2xMxL = 3x2x2x4x4 + (differs across the patterns in each mode) + jones matrix for the 1st pattern: + 1: quat plate, 2: 90, 3: 0, 4: 45 + jones matrix for the 2nd pattern: + 1: 0, 2: 45, 3: 90, 4: quat plate + jones matrix for the 3rd pattern: + 1: 90, 2: 0, 3: 45, 4: quat plate +''' +def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr(): + jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), + build_from_quarters(jones0, jones45, jones90, jones_plate), + build_from_quarters(jones90, jones0, jones45, jones_plate)] + jones_matr = t.stack(([i for i in jones_m]), dim=0) + probe = t.ones(2, 4, 4, dtype=t.cfloat) + # one mode + probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1) + print('probe:', probe.shape) + print('jones:', jones_matr.shape) + out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose) + + print('expected shape: (3, 2, 2, 4, 4)') + print('actual shape:', out.shape) + print('simulated patterns in one mode:', out[:, 3, :, :, :]) + + o = 10 ** 3 + # we'll be checking only one mode (4th) + assert (np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2])) + and t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:]) + + and t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:]) + and t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat)) + and np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:])) + + and t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2]) + and np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:]))) + + + + + + ''' + I REALIZED THAT ANOTHER POSSIBLE SISTUATION WE HAVEN'T CONSIDERED BEFORE IS WHEN THE JONES MATRIX DOESN'T DIFFER ACROSS THE PATTERNS + IN OTHER WORDS, PROBE.SHAPE CONTAINS N (NUM OF PATTERNS) BUT JONES_MATRIX.SHAPE DOESN'T + + + NOW, WE'LL TEST THE CASES WHEN THE PROBE DOESN'T DIFFER FROM PATTERN TO PATTERN + ''' + + + ''' +probe: no multiple modes, multiple diffr patterns + Nx2xMxL = 3x2x4x4 +jones_matrix: same jones matrix applied to all the pixels + 2x2 = 2x2 - a quarter waveplate +''' +def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr_1(): + probe = t.ones(2, 4, 4, dtype=t.cfloat) + probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + jones_matr = jones_plate + print('probe shape:', probe.shape, 'jones shape', jones_matr.shape) + out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose) + + print('expected shape: (3, 2, 4, 4)') + print('actual:', out.shape) + print('simulated:', out) + + assert np.allclose(np.real(out[:, 0, :, :]), np.imag(out[:, 1, :, :])) + + +''' +probe: no multiple modes, multiple diffr pattern + Nx2xMxL = 3x2x4x4 +jones_matrix: jones matrices differ from pixel to pixel + 2x2xMxL = 2x2x4x4 + 1: quat plate, 2: 90, 3: 0, 4: 45 + +''' +def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr_1(): + jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45) + probe = t.ones(2, 4, 4, dtype=t.cfloat) + probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose) + + print('expected shape: (3, 2, 4, 4)') + print('actual shape:', out.shape) + print('simulated:', out) + +# check the 3rd pattern + assert (np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2])) + and t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])) + + +''' +probe: multiple modes, multiple diffr patterns + NxPx2xMxL = 3x7x2x3x4 +jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern) + 2x2 = 2x2 + quarter wave plate +''' +def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr_1(): + probe = t.ones(2, 3, 4, dtype=t.cfloat) + # 1st mode + probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1) + jones_matr = jones_plate + print('probe shape:', probe.shape, 'jones shape', jones_matr.shape) + out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose) + + print('expected shape: (3, 7, 2, 3, 4)') + print('actual:', out.shape) + print('simulated (patterns in one mode):', out[:, 6, :, :, :]) + + # we'll be checking only one mode + assert np.allclose(np.real(out[:, :, 0, :, :]), np.imag(out[:, :, 1, :, :])) + + +''' +probe: multiple modes, multiple diffr pattern + NxPx2xMxL = 3x4x2x4x4 +jones_matrix: jones matrices differ from pixel to pixel + 2x2xMxL = 2x2x4x4 + jones matrix: + 1: quat plate, 2: 90, 3: 0, 4: 45 +''' +def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr_1(): + jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45) + probe = t.ones(2, 4, 4, dtype=t.cfloat) + # one mode + probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0) + probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1) + print('probe:', probe.shape) + print('jones:', jones_matr.shape) + out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose) + + print('expected shape: (3, 2, 2, 4, 4)') + print('actual shape:', out.shape) + print('simulated patterns in one mode:', out[:, 3, :, :, :]) + + o = 10 ** 2 + # we'll be checking only one mode (3th) + assert (np.allclose(np.real(out[0, 2, 0, :-2, :-2]), np.imag(out[0, 2, 1, :-2, :-2])) + and t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + and t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:])) + + return None +