mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
439 lines
17 KiB
Python
439 lines
17 KiB
Python
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
|
|
|