mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-10-11 18:40:25 +02:00
Remove the vestigges of the polarization-dependent code. It belongs in it's own repo
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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})
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
from cdtools.tools.polarization.polarization import *
|
||||
@@ -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))
|
||||
@@ -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(<shape>). 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))
|
||||
Reference in New Issue
Block a user