From 21b618d2e0c723be9cdaf90403a0226d8f1e349a Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 30 Aug 2024 14:13:36 +0200 Subject: [PATCH] Remove the vestigges of the polarization-dependent code. It belongs in it's own repo --- src/cdtools/datasets/__init__.py | 3 +- .../polarization_swept_ptycho_2d_dataset.py | 194 ------ src/cdtools/tools/__init__.py | 1 - .../tools/interactions/interactions.py | 47 +- src/cdtools/tools/polarization/__init__.py | 1 - .../tools/polarization/polarization.py | 217 ------ tests/tools/test_polarization.py | 624 ------------------ 7 files changed, 17 insertions(+), 1070 deletions(-) delete mode 100644 src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py delete mode 100644 src/cdtools/tools/polarization/__init__.py delete mode 100644 src/cdtools/tools/polarization/polarization.py delete mode 100644 tests/tools/test_polarization.py diff --git a/src/cdtools/datasets/__init__.py b/src/cdtools/datasets/__init__.py index afcedf3..84723b1 100644 --- a/src/cdtools/datasets/__init__.py +++ b/src/cdtools/datasets/__init__.py @@ -29,8 +29,7 @@ attempting to do so # I don't believe that __all__ really needed, but it's nice to define it # to be explicit that import * is safe -__all__ = ['CDataset','Ptycho2DDataset','PolarizationSweptPtycho2DDataset'] +__all__ = ['CDataset','Ptycho2DDataset'] from cdtools.datasets.base import CDataset from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset -from cdtools.datasets.polarization_swept_ptycho_2d_dataset import PolarizationSweptPtycho2DDataset diff --git a/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py b/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py deleted file mode 100644 index 32bff52..0000000 --- a/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py +++ /dev/null @@ -1,194 +0,0 @@ -import torch as t -import h5py -import pathlib -from cdtools.datasets import Ptycho2DDataset -from cdtools.tools import data as cdtdata - -__all__ = ['PolarizationSweptPtycho2DDataset'] - - -class PolarizationSweptPtycho2DDataset(Ptycho2DDataset): - """The standard dataset for a 2D ptychography scan - - Subclasses datasets.CDataset - - This class loads and saves 2D ptychography scan data from .cxi files. - It should save and load files compatible with most reconstruction - programs, although it is only tested against SHARP. - """ - def __init__(self, translations, patterns, polarization_indices, - polarization_states, - *args, **kwargs): - """The __init__ function allows construction from python objects. - - The detector_geometry dictionary is defined to have the - entries defined by the outputs of data.get_detector_geometry. - - Note that the created dataset object will not copy the data in the - patterns parameter in order to avoid doubling the memory requiement - for large datasets. - - Parameters - ---------- - translations : array - An nx3 array containing the probe translations at each scan point - patterns : array - An nxmxl array containing the full stack of measured diffraction - patterns - polarization_indices : array(int) - A length-n array containing the index of the polarization state - attached to each pattern - polarization_states : array - An mx2 array, where m is the number of polarization indices, - encoding the polarzation state associated with each index - axes : list(str) - A list of names for the axes of the probe translations - entry_info : dict - A dictionary containing the entry_info metadata - sample_info : dict - A dictionary containing the sample_info metadata - wavelength : float - The wavelength of light used in the experiment - detector_geometry : dict - A dictionary containing the various detector geometry - parameters - mask : array - A mask for the detector, defined as 1 for live pixels, 0 - for dead - background : array - An initial guess for the not-previously-subtracted - detector background - intensities : array - A list of measured shot-to-shot intensities - """ - - super().__init__(translations, patterns, - *args, **kwargs) - - self.polarization_indices = t.tensor(polarization_indices, dtype=t.int32) - self.polarization_states = t.tensor(polarization_states, dtype=t.complex64) - - - def _load(self, index): - """ Internal function to load data - - This function is used internally by the global __getitem__ function - defined in the base class, which handles moving data around when - the dataset is (for example) storing the data on the CPU but - getting data as GPU tensors. - - It loads data in the format (inputs, output) - - The inputs for a 2D ptychogaphy data set are: - - 1) The indices of the patterns to use - 2) The recorded probe positions associated with those points - 3) The angles of the polarizers if polarized=True - - Parameters - ---------- - index (polarized=False): int or slice - The index or indices of the scan points to use - - index (polarized=True): - tuple ((phi1, phi2), ind) - - ind - index or indices of the scan points to use (in a (phi1, phi2) polarization state), int or slice - (phi1, phi2) - angles of the 1st and 2nd polarizers, ints - - Returns - ------- - inputs : tuple - A tuple of the inputs to the related forward models - if polarized: inputs = ((phi1, phi2), ind, transl[(phi1, phi2)][ind]) - outputs : tuple - The output pattern or stack of output patterns - """ - return ((index, self.translations[index], - self.polarization_indices[index]), - self.patterns[index]) - - def to(self, *args, **kwargs): - """Sends the relevant data to the given device and dtype - - This function sends the stored translations, patterns, - mask and background to the specified device and dtype - - Accepts the same parameters as torch.Tensor.to - """ - super().to(*args, **kwargs) - self.polarization_states = self.polarization_states.to(*args, **kwargs) - - - # It sucks that I can't reuse the base factory method here, - # perhaps there is a way but I couldn't figure it out. - @classmethod - def from_cxi(cls, cxi_file, cut_zeros=True): - """Generates a new PolarizationSweptPtycho2DDataset from a .cxi file - - This generates a new PolarizationSweptPtycho2DDataset from a .cxi file - storing the 2D ptychography scan. - - Parameters - ---------- - file : str, pathlib.Path, or h5py.File - The .cxi file to load from - cut_zeros : bool - Default True, whether to set all negative data to zero - - Returns - ------- - dataset : PolarizationSweptPtycho2DDataset - The constructed dataset object - """ - - # If a bare string is passed - if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): - with h5py.File(cxi_file, 'r') as f: - return cls.from_cxi(f, cut_zeros=cut_zeros) - - # Generate a base dataset - dataset = Ptycho2DDataset.from_cxi(cxi_file, cut_zeros=cut_zeros) - - # Mutate the class to this subclass (PolarizedPtycho2DDataset) - dataset.__class__ = cls - - # Now, we save out the polarizer and analyzer states - polarization_indices = cdtdata.get_shot_to_shot_info( - cxi_file, 'polarization_indices') - polarization_states = dataset.entry_info['polarization_states'] - - dataset.polarization_indices = t.tensor(polarization_indices, - dtype=t.int32) - dataset.polarization_states = t.tensor(polarization_states, - dtype=t.complex64) - - return dataset - - def to_cxi(self, cxi_file): - """Saves out a PolarizationSweptPtycho2DDataset as a .cxi file - - This function saves all the compatible information in a - PolarizationSweptPtycho2DDataset object into a .cxi file. This saved - .cxi file should be compatible with any standard .cxi file based - reconstruction tool, such as SHARP. - - Parameters - ---------- - cxi_file : str, pathlib.Path, or h5py.File - The .cxi file to write to - """ - - # If a bare string is passed - if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): - with cdtdata.create_cxi(cxi_file) as f: - return self.to_cxi(f) - - # This saves the translations, patterns, etc. - super().to_cxi(cxi_file) - - # Now, we save out the polarization states - cdtdata.add_shot_to_shot_info(cxi_file, self.polarization_indices, - 'polarization_indices') - cdtdata.add_entry_info(cxi_file, {'polarization_states': - self.polarization_states}) diff --git a/src/cdtools/tools/__init__.py b/src/cdtools/tools/__init__.py index 69b6da1..ade4558 100644 --- a/src/cdtools/tools/__init__.py +++ b/src/cdtools/tools/__init__.py @@ -24,4 +24,3 @@ from cdtools.tools import propagators from cdtools.tools import measurements from cdtools.tools import analysis from cdtools.tools import atoms -from cdtools.tools import polarization diff --git a/src/cdtools/tools/interactions/interactions.py b/src/cdtools/tools/interactions/interactions.py index 5f9b453..2d8642e 100644 --- a/src/cdtools/tools/interactions/interactions.py +++ b/src/cdtools/tools/interactions/interactions.py @@ -7,7 +7,7 @@ for ptychographic reconstruction. import torch as t import numpy as np -from cdtools.tools import propagators, image_processing, polarization +from cdtools.tools import propagators, image_processing __all__ = ['translations_to_pixel', 'pixel_to_translations', 'project_translations_to_sample', @@ -389,7 +389,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True): return t.stack(exit_waves) -def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, probe_support=None, polarized=False, polarizer=None, analyzer=None): +def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, probe_support=None): """Returns a stack of exit waves accounting for subpixel shifts This function returns a collection of exit waves, with the first @@ -438,15 +438,10 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi integer_translations = t.floor(translations) subpixel_translations = translations - integer_translations integer_translations = integer_translations.to(dtype=t.int32) - if not polarized: - selections = t.stack([obj[..., tr[0]:tr[0]+probe.shape[-2], - tr[1]:tr[1]+probe.shape[-1]] - for tr in integer_translations]) - else: - selections = t.stack([obj[:, :, tr[0]:tr[0]+probe.shape[-2], - tr[1]:tr[1]+probe.shape[-1]] - for tr in integer_translations]) - # Nx2x2xMxL tensor + + selections = t.stack([obj[..., tr[0]:tr[0]+probe.shape[-2], + tr[1]:tr[1]+probe.shape[-1]] + for tr in integer_translations]) exit_waves = [] if shift_probe: @@ -459,16 +454,11 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi J = 2 * np.pi * J / probe.shape[-1] phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I -subpixel_translations[:,1,None,None]*J)) - if polarized: - phase_masks = phase_masks[..., None, :, :] - # Nx2x1xMxL tensor - # probe is (N)(P)x2xMxL tensor + + fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2)) if multiple_modes: # Multi-mode probe - if polarized: - shifted_fft_probe = fft_probe * phase_masks[...,None,:,:,:] - else: - shifted_fft_probe = fft_probe * phase_masks[...,None,:,:] + shifted_fft_probe = fft_probe * phase_masks[...,None,:,:] else: shifted_fft_probe = fft_probe * phase_masks shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe, @@ -477,21 +467,16 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi if probe_support is not None: shifted_probe = shifted_probe * probe_support[..., :, :] - if not polarized: - # TODO This is a kludge, I will fix this. I need to handle - # multiple incoherently mixing polarized objects - if multiple_modes and len(selections.shape) == 3: # Multi-mode probe - output = shifted_probe * selections[...,None,:,:] - else: - # This will only work if the - output = shifted_probe * selections - # selections: Nx2x2xMxL - # probe: Nx(P)x2x1xMxL + # TODO This is a kludge, I will fix this. + if multiple_modes and len(selections.shape) == 3: # Multi-mode probe + output = shifted_probe * selections[...,None,:,:] else: - output = polarization.apply_jones_matrix(shifted_probe, selections, multiple_modes=multiple_modes) - + # This will only work if the + output = shifted_probe * selections + else: raise NotImplementedError('Object shift not yet implemented') + if single_translation: return output[0] else: diff --git a/src/cdtools/tools/polarization/__init__.py b/src/cdtools/tools/polarization/__init__.py deleted file mode 100644 index a6f2871..0000000 --- a/src/cdtools/tools/polarization/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from cdtools.tools.polarization.polarization import * diff --git a/src/cdtools/tools/polarization/polarization.py b/src/cdtools/tools/polarization/polarization.py deleted file mode 100644 index 1374895..0000000 --- a/src/cdtools/tools/polarization/polarization.py +++ /dev/null @@ -1,217 +0,0 @@ -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 - -__all__ = ['apply_linear_polarizer', - 'apply_phase_retardance', - 'apply_half_wave_plate', - 'apply_quarter_wave_plate', - 'apply_circular_polarizer', - 'apply_jones_matrix', - 'generate_linear_polarizer', - 'generate_birefringent_obj'] - - -def generate_linear_polarizer(pol_angle): - single_angle = False - - pol_angle = t.as_tensor(pol_angle).to(dtype=t.float32) - if pol_angle.dim() == 0: - pol_angle = t.unsqueeze(pol_angle,0) - single_angle = True - - pol_angle_rad = t.deg2rad(pol_angle) - a = t.cos(pol_angle_rad) ** 2 - b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad) - c = b - d = t.sin(pol_angle_rad) ** 2 - ab = t.stack((a, b), dim=-1) - cd = t.stack((c, d), dim=-1) - jones_matrices = t.stack((ab, cd), dim=-2) - if single_angle: - return jones_matrices[0].to(dtype=t.cfloat) - else: - return jones_matrices.to(dtype=t.cfloat) - - -def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True): - """ - 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)) - - Returns: - -------- - linearly polarized probe: t.Tensor - (N)(P)x2x1xMxL - """ - jones_matrices = generate_linear_polarizer(polarizer) - 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): - # print('probe', probe.shape, 'jones matrix', jones_matrix.shape) - # if jones_matrix.shape == t.Size([5, 2, 2, 2, 2]): - # print('probe', probe.shape, 'jones', jones_matrix.shape) - # print('JONES', jones_matrix) - """ - 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) - - Returns: - -------- - a probe with the jones matrix applied: t.Tensor - (N)(P)x2xMxL - - """ - if transpose: - if jones_matrix.dim() < 4: - jones_matrix = jones_matrix[..., None, None] - if multiple_modes: - jones_matrix = jones_matrix.unsqueeze(-5) - probe = probe[..., None, :, :] - # if jones matrices do not differ from pattern to pattern - if probe.dim() > jones_matrix.dim(): - jones_matrix = jones_matrix.unsqueeze(0) - # vice versa - elif jones_matrix.dim() > probe.dim(): - probe = probe.unsqueeze(0) - # print('apply jonesmatrix: probe', probe.shape, 'matrix:', jones_matrix) - jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4) - probe = probe.transpose(-1, -3).transpose(-2, -4) - output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3) - - else: - raise NotImplementedError - - return output - - -def apply_phase_retardance(probe, phase_shift, multiple_modes=True): - """ - Shifts the y-component of the field wrt the x-component by a given phase shift - - Parameters: - ---------- - probe: t.Tensor - A (...)x2x1xMxL tensor representing the probe - phase_shift: float - phase shift in degrees - - Returns: - -------- - probe: t.Tensor - (...)x2x1xMxL - """ - theta = t.as_tensor(phase_shift, dtype=t.float32) - theta = t.deg2rad(theta) - probe = probe.to(dtype=t.cfloat) - jones_matrix = t.tensor([[1, 0], [0, t.exp(phase_shift)]]).to(dtype=t.cfloat) - polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) - - return polarized - -def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True): - """ - Applies a circular polarizer to the probe - - Parameters: - ---------- - probe: t.Tensor - A (...)x2xMxL tensor representing the probe - left_polarizd: bool - True for the left-polarization, False for the right - - Returns: - -------- - circularly polarized probe: t.Tensor - (...)x2xMxL - """ - probe = probe.to(dtype=t.cfloat) - if left_polarized: - jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]])).to(dtype=t.cfloat) - else: - jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]]).to(dtype=t.cfloat) - polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) - return polarized - -def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True): - """ - 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]]).to(dtype=t.cfloat) - out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) - - return out - -def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True): - """ - 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]]).to(dtype=t.cfloat) - out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) - - return out - -def generate_birefringent_obj(fast_axis=90, phase_ret=10, atten_fast=1, atten_ret=1, global_phase=0): - def to_rad(angle): - angle = t.as_tensor(angle, dtype=t.float32) - angle = t.deg2rad(angle) - return angle - - fast_axis = to_rad(fast_axis) - phase_ret = to_rad(phase_ret) - global_phase = to_rad(global_phase) - - def coord_rot(angle): - a = t.stack((t.cos(angle), t.sin(angle)), dim=-1) - b = t.stack((-t.sin(angle), t.cos(angle)), dim=-1) - return t.stack((a, b), dim=-2).to(dtype=t.cfloat) - - r1 = coord_rot(-fast_axis) - r2 = coord_rot(fast_axis) - p = t.exp(global_phase * 1j) * t.as_tensor([[atten_fast, 0], [0, atten_ret * t.exp(phase_ret*1j)]], dtype=t.cfloat) - return t.matmul(r1, t.matmul(p, r2)) diff --git a/tests/tools/test_polarization.py b/tests/tools/test_polarization.py deleted file mode 100644 index b63696a..0000000 --- a/tests/tools/test_polarization.py +++ /dev/null @@ -1,624 +0,0 @@ -import numpy as np -import torch as t -from cdtools.tools.polarization import apply_linear_polarizer, generate_linear_polarizer -from cdtools.tools.polarization import apply_jones_matrix as jones - - -# 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 - -from numpy import cos, sin, deg2rad - -# 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 -angle_2 = angle - 45 - -def polarizer(angle): - theta = deg2rad(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 * np.pi / 4 * t.ones(2, 2)).to(dtype=t.cfloat) -theta2 = deg2rad(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]]).to(dtype=t.cfloat) - -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) - -transpose = True - -def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr(): - ''' - 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 - ''' - 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) - print('expected shape:(2, 3, 4)') - print('actual:', out.shape) - print('simulated:', out) - - assert np.allclose(np.real(out[0]), np.imag(out[1])) - assert out.shape == t.Size((2, 3, 4)) - - -def test_generate_linear_polarizer(): - pol_angles = [0, 45, 90] - pol_angle1 = 45 - pol_angle2 = t.tensor(90) - pol_angle3 = t.tensor([0]) - pols = generate_linear_polarizer(pol_angles) - pol1 = generate_linear_polarizer(pol_angle1) - pol2 = generate_linear_polarizer(pol_angle2) - pol3 = generate_linear_polarizer(pol_angle3) - print('polarizers 0, 45, 90 (1D tensor) shape:', pols.shape) - print('shape of the polarizer generated from int:', pol1.shape) - print('shape of the polarizer generated from 0D tensor:', pol2.shape) - print('shape of the linear polarizer generated from t.Size(0) tensor:', pol3.shape) - print('90', pol2) - print(jones90) - probe = t.ones(4, 4) - jones_m = [jones0, jones45, jones90] - jones_m = t.stack([matr for matr in jones_m]) - assert pols.shape == t.Size((3, 2, 2)) - assert pol1.shape == t.Size((2, 2)) - assert pol2.shape == t.Size((2, 2)) - assert pol3.shape == t.Size((1, 2, 2)) - assert t.allclose(pol1, jones45) - - -def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr(): - '''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 - ''' - 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) - - # 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:]) - assert out.shape == t.Size((2, 4, 4)) - -def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr(): - ''' - 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 - ''' - 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)) - assert t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat)) - assert t.allclose(out[1, 0, :, :], out[1, 1]) - assert t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat)) - assert out.shape == t.Size((3, 2, 4, 4)) - -def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr(): - ''' - 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 - ''' - 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])) - assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:]) - - assert t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:]) - assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat)) - assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:])) - - assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2]) - assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:])) - assert out.shape == t.Size((3, 2, 4, 4)) - -def test_apply_jones_matrix_mult_modes_1_pattern_one_jones_matr(): - ''' - probe: multiple modes, 1 diffr pattern - Px2xMxL = 2x2x3x4 - jones_matrix: same jones matrix applied to all the pixels - 2x2 - quarter plate - ''' - 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, :, :])) - assert np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :])) - assert out.shape == t.Size((2, 2, 3, 4)) - -def test_apply_jones_matrix_mult_modes_1_pattern_diff_jones_matr(): - ''' - 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 - ''' - 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])) - assert t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:]) - - assert np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2])) - assert t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:]) - - assert np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2])) - assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]) - assert out.shape == t.Size((3, 2, 4, 4)) - -def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr(): - ''' - 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 - ''' - 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)) - assert t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat)) - assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :]) - assert t.allclose(out[2, 6, 0, :, :], 3 * (10**6) * t.zeros(3, 4, dtype=t.cfloat)) - assert t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat)) - assert out.shape == t.Size((3, 7, 2, 3, 4)) - -def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr(): - ''' - 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 - ''' - 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, 4, 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])) - assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:]) - - assert t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:]) - assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat)) - assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:])) - - assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2]) - assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:])) - assert out.shape == t.Size((3, 4, 2, 4, 4)) - -def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr_1(): - ''' - probe: no multiple modes, multiple diffr patterns - Nx2xMxL = 3x2x4x4 - jones_matrix: same jones matrix applied to all the pixels - 2x2 = 2x2 - a quarter waveplate - ''' - 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, :, :])) - assert out.shape == t.Size((3, 2, 4, 4)) - -def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr_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 - - ''' - 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])) - assert t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]) - assert out.shape == t.Size((3, 2, 4, 4)) - -def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr_1(): - ''' - 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 - ''' - 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, :, :])) - assert out.shape == t.Size((3, 7, 2, 3, 4)) - -def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr_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 - ''' - 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, 4, 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])) - assert t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:]) - assert out.shape == t.Size((3, 4, 2, 4, 4)) - - return None - - -def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_1(): - ''' - probe: - 2xMxL = 2x3x4 - jones_matrix: - Nx2x2 = 3x2x2 - jones matrices: - 1: quat plate, 2: 90, 3: 0 - ''' - jones_matr = t.stack((jones_plate, jones90, jones0)) - probe = t.ones(2, 3, 4, dtype=t.cfloat) - out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose) - - print('expected shape: (3, 2, 3, 4)') - print('actual shape:', out.shape) - print('simulated:', out) - - # we'll be checking only one mode (3th) - assert np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :])) - assert t.allclose(out[1, 0, :, :], t.zeros(3, 4, dtype=t.cfloat)) - assert t.allclose(out[1, 1, :, :], t.ones(3, 4, dtype=t.cfloat)) - assert t.allclose(out[2, 0, :, :], t.ones(3, 4, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :, :], t.zeros(3, 4, dtype=t.cfloat)) - assert out.shape == t.Size((3, 2, 3, 4)) - - return None - - -def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_2(): - ''' - probe: - 2xMxL = 2x4x4 - 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 - ''' - 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) - 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])) - assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:]) - - assert t.allclose(out[1, 0, :-2, :-2], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:]) - assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat)) - assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:])) - - assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :-2, :-2], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2]) - assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:])) - assert out.shape == t.Size((3, 2, 4, 4)) - - -def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_1(): - ''' - probe: multiple modes, multiple diffr patterns - Px2xMxL = 7x2x3x4 - jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern) - Nx2x2 = 3x2x2 - 3 different matrices for each pattern: - 1: 0 - 2: 45 - 3: 90 - ''' - probe = t.ones(2, 3, 4, dtype=t.cfloat) - probe = t.stack(([probe * 10 ** i for i in range(7)]), 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=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)) - assert t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat)) - assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :]) - assert t.allclose(out[2, 6, 0, :, :], (10**6) * t.zeros(3, 4, dtype=t.cfloat)) - assert t.allclose(out[2, 6, 1, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat)) - assert out.shape == t.Size((3, 7, 2, 3, 4)) - - -def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_2(): - ''' - probe: multiple modes, multiple diffr pattern - Px2xMxL = 4x2x4x4 - 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 - ''' - 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 * (10 ** i) for i in range(4)]), dim=0) - print('probe:', probe.shape) - print('jones:', jones_matr.shape) - out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose) - - print('expected shape: (3, 4, 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])) - assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:]) - - assert t.allclose(out[1, 3, 0, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:]) - assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[1, 3, 1, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat)) - assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:])) - - assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 1, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 0, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat)) - assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2]) - assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:])) - assert out.shape == t.Size((3, 4, 2, 4, 4))