From 960740f39f53b2b7c9a773ef51e5d683f3680991 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 5 Aug 2021 10:44:32 -0400 Subject: [PATCH] Go over Anastasia's polarization tests and clean up repo a bit --- CDTools/tools/polarization.py | 29 --- CDTools/tools/polarization/polarization.py | 12 +- CDTools/tools/polarization/test.py | 48 ---- .../polarization/testing_interactions.py | 246 ------------------ tests/{ => tools}/test_polarization.py | 69 ++--- tests/tools/test_propagators.py | 2 +- 6 files changed, 47 insertions(+), 359 deletions(-) delete mode 100644 CDTools/tools/polarization.py delete mode 100644 CDTools/tools/polarization/test.py delete mode 100644 CDTools/tools/polarization/testing_interactions.py rename tests/{ => tools}/test_polarization.py (91%) diff --git a/CDTools/tools/polarization.py b/CDTools/tools/polarization.py deleted file mode 100644 index b47ecf8..0000000 --- a/CDTools/tools/polarization.py +++ /dev/null @@ -1,29 +0,0 @@ -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 index 1f072a5..0728185 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -1,5 +1,8 @@ import numpy as numpy import torch as t +# Abe - again, we don't need math here. replace with torch-native functions +# in all the definitions here. We have to use torch here if we want to be +# able to calculate derivatives w.r.t. the polarizer angle, say. import math from math import sin from math import cos @@ -10,8 +13,10 @@ __all__ = ['apply_linear_polarizer', 'apply_quarter_wave_plate', 'apply_circular_polarizer', 'apply_jones_matrix'] -print() - + + +# Abe - + def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True): """ Applies a linear polarizer to the probe @@ -33,6 +38,9 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True # polarizer = t.tensor([polarizer]) if len(polarizer,shape) == 0: polarizer = t.tensor([polarizer]) + + # Abe - Something feels overly complicated about this. Let's take a look + # at it together and do some simplification 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))])) diff --git a/CDTools/tools/polarization/test.py b/CDTools/tools/polarization/test.py deleted file mode 100644 index 071f4ce..0000000 --- a/CDTools/tools/polarization/test.py +++ /dev/null @@ -1,48 +0,0 @@ -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 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 - -dataset = PolarizedPtycho2DDataset.from_cxi('test_ptycho.cxi') -model = PolarizedFancyPtycho.from_dataset(dataset) - -# print('probe', model.probe.shape) -# print('object', model.obj.shape) -# print('transl', dataset.translations.shape) -# print('polarizer', dataset.polarizer.shape) -# print('analyzer', dataset.analyzer.shape) - -exit_waves = model.interaction(*dataset[[0]][0]) -print(exit_waves.shape) -# sim_patterns = model.forward(*dataset[:][0]) - -for loss in model.Adam_optimize(10, dataset, batch_size=1, lr=0.002, schedule=True): - # And we liveplot the updates to the model as they happen - print(model.report()) - - -# probe = t.rand(5, 2, 10, 12) -# obj = t.rand(2, 2, 10, 12) -# polarizer = t.rand(7) -# analyzer = t.rand(7) - -# pol_probes = polarization.apply_linear_polarizer(probe, polarizer) - -# exit_waves = interactions.ptycho_2D_sinc( -# probe, obj, translations, -# shift_probe=True, multiple_modes=True, polarized=True) - -# analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer) - -# print(analyzed_exit_waves.shape) \ No newline at end of file diff --git a/CDTools/tools/polarization/testing_interactions.py b/CDTools/tools/polarization/testing_interactions.py deleted file mode 100644 index aafd16b..0000000 --- a/CDTools/tools/polarization/testing_interactions.py +++ /dev/null @@ -1,246 +0,0 @@ -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 - -dataset = PolarizedPtycho2DDataset.from_cxi('test_ptycho.cxi') -model = PolarizedFancyPtycho.from_dataset(dataset) - -def test_apply_jones_matrix(multiple_modes=False, multiple_patterns=False, diff_jones_across_pixels=False, transpose=True, mode=0): - - 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 - ''' - if not multiple_modes and not multiple_patterns and not diff_jones_across_pixels: - probe = t.rand(2, 4, 4, dtype=t.cfloat) - out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose), quarter_plate, multiple_modes=False, transpose=transpose, mode=0) - - print('expected shape:(2, 3, 4)') - print('actual:', out.shape) - print('simulated:', out) - - '''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 - ''' - if not multiple_modes and not multiple_patterns and diff_jones_across_pixels: - - 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) - - print('expected shape:(2, 4, 4)') - print('simulated:', out_simulated.shape) - print('simulated:', out) - - - ''' - 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 - ''' - if not multiple_modes and multiple_patterns and not diff_jones_across_pixels: - 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) - - ''' - 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 - ''' - if not multiple_modes and multiple_patterns and diff_jones_across_pixels: - 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) - - ''' - probe: multiple modes, 1 diffr pattern - Px2xMxL = 2x2x3x4 - jones_matrix: same jones matrix applied to all the pixels - 2x2 - quarter plate - Nx2xMxL = 3x2x4x4 - ''' - if multiple_modes and not multiple_patterns and not diff_jones_across_pixels: - 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) - - ''' - 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 - ''' - if multiple_modes and not multiple_patterns and diff_jones_across_pixels: - 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) - - - ''' - probe: multiple modes, multiple diffr patterns - NxPx2xMxL = 3x2x2x3x4 - 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 - ''' - if multiple_modes and multiple_patterns and not diff_jones_across_pixels: - 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, probe_mode2), 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, 2, 2, 3, 4)') - print('actual:', out.shape) - print('simulated (patterns in one mode):', out[:, mode, :, :, :]) - - - ''' - probe: multiple modes, multiple diffr pattern - NxPx2xMxL = 3x4x2x4x4 - jones_matrix: jones matrices differ from pixel to pixel - NxPx2x2xMxL = 3x2x2x2x4x4 - (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 - ''' - if multiple_modes and multiple_patterns and diff_jones_across_pixels: - 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[:, mode, :, :, :]) - - - return None - -probe = t.rand(2, 3, 4, dtype=t.cfloat) -angle = 57 -angle_2 = angle - 45 -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) -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]]) - -# pol = t.matmul(polarizer, probe) -# pl = t.matmul(quarter_plate, pol) -# print(pol) - -# print(pl) -# pol = pol.unsqueeze(-1).unsqueeze(-1) -# print(polarization.apply_jones_matrix(polarization.apply_jones_matrix(probe, polarizer), quarter_plate, multiple_modes=False)) -test_apply_jones_matrix(multiple_modes=True, multiple_patterns=True, diff_jones_across_pixels=True, mode=2) \ No newline at end of file diff --git a/tests/test_polarization.py b/tests/tools/test_polarization.py similarity index 91% rename from tests/test_polarization.py rename to tests/tools/test_polarization.py index f4ac70c..04cbbfa 100644 --- a/tests/test_polarization.py +++ b/tests/tools/test_polarization.py @@ -1,23 +1,33 @@ -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 + +# Abe - I removed all the imports that didn't need to be here. + +# Abe - A few issues. First, you could just write "from math import cos, sin" +# Second, we already have numpy imported, so better to use np.cos and np.sin from math import cos as cos, sin as sin +# numpy also has np.pi and np.deg2rad import math +# Further comments: +# +# 1) You should also introduce an assert statement for the output shape +# checking, e.g. assert out.shape == t.Size(). See example in the +# first test. +# +# 2) Instead of using "assert a and b and c and d...", use separate assert +# statements for each condition. This is both more readable, and it allows the +# testing system to pinpoint exactly which of the assert statements failed. +# +# 3) USE SPACES INSTEAD OF TABS!!! 4 spaces per indent is the convention for +# this project. +# +# 4) Docstrings should go after the function definition, rather than before +# +# 5) Still missing cases where Jones matrix has multiple entries (N) but probe +# doesn't. Low priority. angle = 87 @@ -27,6 +37,7 @@ 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)], @@ -54,7 +65,7 @@ jones_matrix: same jones matrix applied to all the pixels ''' transpose = True def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr(): - probe = t.rand(2, 4, 4, dtype=t.cfloat) + probe = t.rand(2, 3, 4, dtype=t.cfloat) print(polarizer) out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose), quarter_plate, multiple_modes=False, transpose=transpose) @@ -62,6 +73,8 @@ def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr(): print('actual:', out.shape) print('simulated:', out) + assert out.shape == t.Size((2,3,4)) + assert np.allclose(np.real(out[0]), np.imag(out[1])) @@ -75,6 +88,7 @@ jones_matrix: jones matrices differ from pixel to pixel 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) @@ -88,12 +102,13 @@ def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr(): 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:])) + # Example of using multiple asserts + assert np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2])) + assert t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) + assert t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) + assert t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) + assert t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) + assert t.allclose(out[0, -2:, -2:], out[0, -2:, -2:]) ''' @@ -320,19 +335,7 @@ def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr(): 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 diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index f45a9b9..a575ff8 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -278,7 +278,7 @@ def test_near_field(): assert z.grad != 0 assert spacing.grad[0] != 0 assert wavelength.grad !=0 - + def test_generalized_near_field():