From eb64770258398a6c9c74116b7e2abed9c2cfa654 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 3 Aug 2021 02:25:29 -0400 Subject: [PATCH 01/15] . --- CDTools/tools/polarization/polarization.py | 9 +- tests/test_polarization.py | 438 +++++++++++++++++++++ 2 files changed, 445 insertions(+), 2 deletions(-) create mode 100644 tests/test_polarization.py diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index d8dbe06..af5b0b5 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -61,9 +61,13 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True) if transpose: if len(jones_matrix.shape) >= 4: jones_matrix = jones_matrix[..., None, :, :, :, :] + 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) @@ -82,8 +86,9 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True) if len(jones_matrix.shape) < 4: jones_matrix = jones_matrix[..., None, None] probe = probe[..., None, :, :] - print('probs', probe.shape) - print('joness', jones_matrix.shape) + # 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) 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 + From 332541c9515ca9376c98b45449fa99771408d5f2 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 3 Aug 2021 03:14:00 -0400 Subject: [PATCH 02/15] '.' --- CDTools/datasets/polarized_ptycho_2d_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py index 93deb31..62c3a25 100644 --- a/CDTools/datasets/polarized_ptycho_2d_dataset.py +++ b/CDTools/datasets/polarized_ptycho_2d_dataset.py @@ -68,7 +68,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): polarizer = [] analyzer = [] - for k in t.tensor(translations).shape[0]: + for k in range(t.tensor(translations).shape[0]): polarizer.append((k//3)%3) analyzer.append((k%3)) self.polarizer = t.tensor(polarizer) From 8c22c99d66fdeb1c0ebc704677de80398f960814 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 3 Aug 2021 13:57:25 -0400 Subject: [PATCH 03/15] . --- CDTools/models/base.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 5ceca97..0da45eb 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 +pp = PdfPages('multipage.pdf') __all__ = ['CDIModel'] @@ -514,6 +516,7 @@ class CDIModel(t.nn.Module): if update: plt.draw() fig.canvas.start_event_loop(0.001) + pp.savefig() if first_update: plt.pause(0.05 * len(self.figs)) From 061b7f647aa88cee4135fbe91ea95f5f19be4003 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 3 Aug 2021 19:34:08 -0400 Subject: [PATCH 04/15] '.' --- CDTools/tools/polarization/polarization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index af5b0b5..7294c80 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -29,7 +29,7 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True (N)(P)x2x1xMxL """ - if len(polarizer) == 1: + if len(polarizer.shape) == 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: From 5d4fba7819892b0dc7864e7b05066749831712f6 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 3 Aug 2021 19:54:18 -0400 Subject: [PATCH 05/15] '.' --- CDTools/tools/polarization/polarization.py | 1 + 1 file changed, 1 insertion(+) diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index 7294c80..3090b2e 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -30,6 +30,7 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True """ if len(polarizer.shape) == 1: + print('polarizer', polarizer) 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: From 24475dbdac5b1e0f687d6c588cee33b9d37d2b07 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 3 Aug 2021 20:14:15 -0400 Subject: [PATCH 06/15] '.' --- CDTools/datasets/polarized_ptycho_2d_dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py index 62c3a25..b1e22a7 100644 --- a/CDTools/datasets/polarized_ptycho_2d_dataset.py +++ b/CDTools/datasets/polarized_ptycho_2d_dataset.py @@ -69,8 +69,8 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): polarizer = [] analyzer = [] for k in range(t.tensor(translations).shape[0]): - polarizer.append((k//3)%3) - analyzer.append((k%3)) + polarizer.append((k//3)%3 * 45) + analyzer.append((k%3 * 45)) self.polarizer = t.tensor(polarizer) self.analyzer = t.tensor(analyzer) From 3bc3622910180562c460c9a52a5a6b419b96f05a Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 3 Aug 2021 20:24:43 -0400 Subject: [PATCH 07/15] '.' --- CDTools/tools/polarization/polarization.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index 3090b2e..c6764d6 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -10,6 +10,7 @@ __all__ = ['apply_linear_polarizer', 'apply_quarter_wave_plate', 'apply_circular_polarizer', 'apply_jones_matrix'] +print() def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True): """ @@ -29,14 +30,9 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True (N)(P)x2x1xMxL """ - if len(polarizer.shape) == 1: - print('polarizer', polarizer) - 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))])) + 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) From 8f9cefac0500cbafbc4cad3cdb8340f1e366d596 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 3 Aug 2021 22:06:53 -0400 Subject: [PATCH 08/15] . --- CDTools/tools/polarization/polarization.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index c6764d6..1efc5e9 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -29,7 +29,8 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True linearly polarized probe: t.Tensor (N)(P)x2x1xMxL """ - + 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))])) From 6adce3a575c3e3021e99a832834650c1524cc110 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 3 Aug 2021 22:33:32 -0400 Subject: [PATCH 09/15] . --- CDTools/tools/polarization/polarization.py | 324 +++++++++++---------- 1 file changed, 163 insertions(+), 161 deletions(-) diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index 1efc5e9..d164fa6 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -5,199 +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 - """ - 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))])) + Returns: + -------- + linearly polarized probe: t.Tensor + (N)(P)x2x1xMxL + """ + # if len(polarizer.shape) == 0: + # polarizer = t.tensor([polarizer]) + if transpose: + print('a') + 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, :, :, :, :] + 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, :, :] - # 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: + 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 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 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 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) From 1925739174d3f89b0a11d931d4b05601cf59678b Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 3 Aug 2021 22:54:39 -0400 Subject: [PATCH 10/15] . --- CDTools/tools/polarization/polarization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index d164fa6..d9b42f3 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -31,8 +31,8 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True """ # if len(polarizer.shape) == 0: # polarizer = t.tensor([polarizer]) - if transpose: - print('a') + 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))])) From 21036f625aff3f78857690038a5f2e4afa5ebcbb Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 3 Aug 2021 23:52:19 -0400 Subject: [PATCH 11/15] . --- CDTools/models/base.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 0da45eb..3ca232e 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -43,7 +43,7 @@ import time from .complex_adam import MyAdam from .complex_lbfgs import MyLBFGS from matplotlib.backends.backend_pdf import PdfPages -pp = PdfPages('multipage.pdf') + __all__ = ['CDIModel'] @@ -478,7 +478,7 @@ class CDIModel(t.nn.Module): self.figs = [] idx = 0 - for plots in self.plot_list: + for plots, c in in enumerate(self.plot_list): # If a conditional is included in the plot try: if len(plots) >=3 and not plots[2](self): @@ -505,6 +505,7 @@ class CDIModel(t.nn.Module): try: plotter(self, fig, dataset) plt.title(name) + plt.savefig('img-{0}.pdf'.format(index), bbox_inches='tight') except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e: pass @@ -516,7 +517,6 @@ class CDIModel(t.nn.Module): if update: plt.draw() fig.canvas.start_event_loop(0.001) - pp.savefig() if first_update: plt.pause(0.05 * len(self.figs)) From bf4db86ad49e75f69016d6e282e359d0c52c0139 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 3 Aug 2021 23:59:23 -0400 Subject: [PATCH 12/15] . --- CDTools/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 3ca232e..726b17c 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -478,7 +478,7 @@ class CDIModel(t.nn.Module): self.figs = [] idx = 0 - for plots, c in in enumerate(self.plot_list): + for plots, c in enumerate(self.plot_list): # If a conditional is included in the plot try: if len(plots) >=3 and not plots[2](self): From 35c1643721c72f2dc50fbfc308a0ae3b8c584b66 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Wed, 4 Aug 2021 00:04:29 -0400 Subject: [PATCH 13/15] . --- CDTools/models/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 726b17c..1b3405d 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -478,7 +478,7 @@ class CDIModel(t.nn.Module): self.figs = [] idx = 0 - for plots, c in enumerate(self.plot_list): + for plots in self.plot_list: # If a conditional is included in the plot try: if len(plots) >=3 and not plots[2](self): @@ -505,7 +505,7 @@ class CDIModel(t.nn.Module): try: plotter(self, fig, dataset) plt.title(name) - plt.savefig('img-{0}.pdf'.format(index), bbox_inches='tight') + plt.savefig('img-{0}.pdf'.format(idx), bbox_inches='tight') except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e: pass From 9f580102baabdb7129357f384d04d39273c5e462 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Wed, 4 Aug 2021 00:09:54 -0400 Subject: [PATCH 14/15] . --- CDTools/models/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 1b3405d..85c5fbe 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -505,6 +505,7 @@ class CDIModel(t.nn.Module): 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 From daeac7e983fbbd39e476282a674006c9c85efd03 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Wed, 4 Aug 2021 09:39:49 -0400 Subject: [PATCH 15/15] . --- CDTools/models/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 85c5fbe..3c26947 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -499,7 +499,9 @@ 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: