mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
Merge branch 'polarization' of github.mit.edu:Scattering/CDTools into polarization
This commit is contained in:
Abe Levitan
@@ -68,9 +68,9 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
|
||||
|
||||
polarizer = []
|
||||
analyzer = []
|
||||
for k in t.tensor(translations).shape[0]:
|
||||
polarizer.append((k//3)%3)
|
||||
analyzer.append((k%3))
|
||||
for k in range(t.tensor(translations).shape[0]):
|
||||
polarizer.append((k//3)%3 * 45)
|
||||
analyzer.append((k%3 * 45))
|
||||
self.polarizer = t.tensor(polarizer)
|
||||
self.analyzer = t.tensor(analyzer)
|
||||
|
||||
|
||||
@@ -42,6 +42,8 @@ import time
|
||||
#import pytorch_warmup
|
||||
from .complex_adam import MyAdam
|
||||
from .complex_lbfgs import MyLBFGS
|
||||
from matplotlib.backends.backend_pdf import PdfPages
|
||||
|
||||
|
||||
__all__ = ['CDIModel']
|
||||
|
||||
@@ -497,12 +499,16 @@ class CDIModel(t.nn.Module):
|
||||
try:
|
||||
plotter(self,fig)
|
||||
plt.title(name)
|
||||
|
||||
print('f', idx)
|
||||
plt.savefig('img-{0}.pdf'.format(idx), bbox_inches='tight')
|
||||
|
||||
except TypeError as e:
|
||||
if dataset is not None:
|
||||
try:
|
||||
plotter(self, fig, dataset)
|
||||
plt.title(name)
|
||||
print('f', idx)
|
||||
plt.savefig('img-{0}.pdf'.format(idx), bbox_inches='tight')
|
||||
except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
|
||||
pass
|
||||
|
||||
|
||||
@@ -5,196 +5,201 @@ from math import sin
|
||||
from math import cos
|
||||
|
||||
__all__ = ['apply_linear_polarizer',
|
||||
'apply_phase_retardance',
|
||||
'apply_half_wave_plate',
|
||||
'apply_quarter_wave_plate',
|
||||
'apply_circular_polarizer',
|
||||
'apply_jones_matrix']
|
||||
|
||||
'apply_phase_retardance',
|
||||
'apply_half_wave_plate',
|
||||
'apply_quarter_wave_plate',
|
||||
'apply_circular_polarizer',
|
||||
'apply_jones_matrix']
|
||||
print()
|
||||
|
||||
def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True):
|
||||
"""
|
||||
Applies a linear polarizer to the probe
|
||||
"""
|
||||
Applies a linear polarizer to the probe
|
||||
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (N)(P)x2xMxL tensor representing the probe, MxL - the size of the probe
|
||||
The angle between the fast-axis of the linear polarizer and the horizontal axis
|
||||
polarizer: t.Tensor
|
||||
A 1D tensor (N) representing the polarizer angles for each of the patterns (or a single tensor of shape (1))
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (N)(P)x2xMxL tensor representing the probe, MxL - the size of the probe
|
||||
The angle between the fast-axis of the linear polarizer and the horizontal axis
|
||||
polarizer: t.Tensor
|
||||
A 1D tensor (N) representing the polarizer angles for each of the patterns (or a single tensor of shape (1))
|
||||
|
||||
Returns:
|
||||
--------
|
||||
linearly polarized probe: t.Tensor
|
||||
(N)(P)x2x1xMxL
|
||||
"""
|
||||
Returns:
|
||||
--------
|
||||
linearly polarized probe: t.Tensor
|
||||
(N)(P)x2x1xMxL
|
||||
"""
|
||||
# if len(polarizer.shape) == 0:
|
||||
# polarizer = t.tensor([polarizer])
|
||||
if len(polarizer,shape) == 0:
|
||||
polarizer = t.tensor([polarizer])
|
||||
pol_cos = lambda idx: cos(math.radians(polarizer[idx]))
|
||||
pol_sin = lambda idx: sin(math.radians(polarizer[idx]))
|
||||
jones_matrices = t.stack(([t.tensor([[(pol_cos(idx)) ** 2, pol_sin(idx) * pol_cos(idx)], [pol_sin(idx) * pol_cos(idx), (pol_sin(idx)) ** 2]]).to(dtype=t.cfloat) for idx in range(len(polarizer))]))
|
||||
|
||||
if len(polarizer) == 1:
|
||||
theta = math.radians(polarizer)
|
||||
jones_matrices = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat)
|
||||
else:
|
||||
pol_cos = lambda idx: cos(math.radians(polarizer[idx]))
|
||||
pol_sin = lambda idx: sin(math.radians(polarizer[idx]))
|
||||
jones_matrices = t.stack(([t.tensor([[(pol_cos(idx)) ** 2, pol_sin(idx) * pol_cos(idx)], [pol_sin(idx) * pol_cos(idx), (pol_sin(idx)) ** 2]]).to(dtype=t.cfloat) for idx in range(len(polarizer))]))
|
||||
|
||||
return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes)
|
||||
return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes)
|
||||
|
||||
def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True):
|
||||
"""
|
||||
Applies a given Jones matrix to the probe
|
||||
"""
|
||||
Applies a given Jones matrix to the probe
|
||||
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (N)(P)x2xMxL tensor representing the probe
|
||||
jones_matrix: t.tensor
|
||||
(N)x2x2x(M)x(L)
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (N)(P)x2xMxL tensor representing the probe
|
||||
jones_matrix: t.tensor
|
||||
(N)x2x2x(M)x(L)
|
||||
|
||||
Returns:
|
||||
--------
|
||||
a probe with the jones matrix applied: t.Tensor
|
||||
(N)(P)x2xMxL
|
||||
Returns:
|
||||
--------
|
||||
a probe with the jones matrix applied: t.Tensor
|
||||
(N)(P)x2xMxL
|
||||
|
||||
Assume that if the probe has a dimension (N), so does the jones matrix
|
||||
"""
|
||||
if multiple_modes:
|
||||
if transpose:
|
||||
if len(jones_matrix.shape) >= 4:
|
||||
jones_matrix = jones_matrix[..., None, :, :, :, :]
|
||||
else:
|
||||
jones_matrix = jones_matrix[..., None, :, :, None, None]
|
||||
probe = probe[..., None, :, :]
|
||||
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
|
||||
# (N)1xMxLx2x2 or (N)1x1x1x2x2
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
|
||||
# (N)Px2xMxL
|
||||
Assume that if the probe has a dimension (N), so does the jones matrix
|
||||
"""
|
||||
if multiple_modes:
|
||||
if transpose:
|
||||
if len(jones_matrix.shape) >= 4:
|
||||
jones_matrix = jones_matrix[..., None, :, :, :, :]
|
||||
|
||||
else:
|
||||
if len(jones_matrix.shape) < 4:
|
||||
jones_matrix = jones_matrix[..., None, :, :, :, :]
|
||||
probe = t.stack((probe, probe), dim=-4)
|
||||
output = t.sum(jones_matrix * probe, dim=-3)
|
||||
#(N)x2xMxL
|
||||
else:
|
||||
jones_matrix = jones_matrix[..., None, :, :, None, None]
|
||||
probe = probe[..., None, :, :]
|
||||
# if jones matrices do not differ from pattern to pattern
|
||||
if len(probe.shape) > len(jones_matrix.shape):
|
||||
jones_matrix = jones_matrix[None, ...]
|
||||
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
|
||||
# (N)1xMxLx2x2 or (N)1x1x1x2x2
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
|
||||
# (N)Px2xMxL
|
||||
|
||||
else:
|
||||
if transpose:
|
||||
if len(jones_matrix.shape) < 4:
|
||||
jones_matrix = jones_matrix[..., None, None]
|
||||
probe = probe[..., None, :, :]
|
||||
print('probs', probe.shape)
|
||||
print('joness', jones_matrix.shape)
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
|
||||
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
|
||||
else:
|
||||
if len(jones_matrix.shape) < 4:
|
||||
jones_matrix = jones_matrix[..., None, :, :, :, :]
|
||||
probe = t.stack((probe, probe), dim=-4)
|
||||
output = t.sum(jones_matrix * probe, dim=-3)
|
||||
#(N)x2xMxL
|
||||
|
||||
else:
|
||||
if transpose:
|
||||
if len(jones_matrix.shape) < 4:
|
||||
jones_matrix = jones_matrix[..., None, None]
|
||||
probe = probe[..., None, :, :]
|
||||
# if jones matrices do not differ from pattern to pattern
|
||||
if len(probe.shape) > len(jones_matrix.shape):
|
||||
jones_matrix = jones_matrix[None, ...]
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
|
||||
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
|
||||
|
||||
|
||||
else:
|
||||
if len(jones_matrix.shape) < 4:
|
||||
jones_matrix = jones_matrix[..., None, None]
|
||||
probe = t.stack((probe, probe), dim=-4)
|
||||
output = t.sum(jones_matrix * probe, dim=-3)
|
||||
|
||||
return output
|
||||
else:
|
||||
if len(jones_matrix.shape) < 4:
|
||||
jones_matrix = jones_matrix[..., None, None]
|
||||
probe = t.stack((probe, probe), dim=-4)
|
||||
output = t.sum(jones_matrix * probe, dim=-3)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def apply_phase_retardance(probe, phase_shift):
|
||||
"""
|
||||
Shifts the y-component of the field wrt the x-component by a given phase shift
|
||||
"""
|
||||
Shifts the y-component of the field wrt the x-component by a given phase shift
|
||||
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe
|
||||
phase_shift: float
|
||||
phase shift in degrees
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe
|
||||
phase_shift: float
|
||||
phase shift in degrees
|
||||
|
||||
Returns:
|
||||
--------
|
||||
probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
jones_matrix = t.tensor([[1, 0], [0, phase_shift]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
Returns:
|
||||
--------
|
||||
probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
jones_matrix = t.tensor([[1, 0], [0, phase_shift]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
|
||||
def apply_circular_polarizer(probe, left_polarized=True):
|
||||
"""
|
||||
Applies a circular polarizer to the probe
|
||||
"""
|
||||
Applies a circular polarizer to the probe
|
||||
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe
|
||||
left_polarizd: bool
|
||||
True for the left-polarization, False for the right
|
||||
|
||||
Returns:
|
||||
--------
|
||||
circularly polarized probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
if left_polarized:
|
||||
jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]]))
|
||||
else:
|
||||
jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe
|
||||
left_polarizd: bool
|
||||
True for the left-polarization, False for the right
|
||||
|
||||
Returns:
|
||||
--------
|
||||
circularly polarized probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
if left_polarized:
|
||||
jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]]))
|
||||
else:
|
||||
jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
|
||||
def apply_quarter_wave_plate(probe, fast_axis_angle):
|
||||
"""
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
|
||||
fast_axis_angle: float
|
||||
The angle between the fast-axis of the polarizer and the horizontal axis
|
||||
"""
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
|
||||
fast_axis_angle: float
|
||||
The angle between the fast-axis of the polarizer and the horizontal axis
|
||||
|
||||
Returns:
|
||||
--------
|
||||
polarized probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
theta = math.radians(fast_axis_angle)
|
||||
exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
|
||||
jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
Returns:
|
||||
--------
|
||||
polarized probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
theta = math.radians(fast_axis_angle)
|
||||
exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
|
||||
jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
|
||||
|
||||
def apply_half_wave_plate(probe, fast_axis_angle):
|
||||
"""
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
|
||||
fast_axis_angle: float
|
||||
The angle between the fast-axis of the polarizer and the horizontal axis
|
||||
"""
|
||||
Parameters:
|
||||
----------
|
||||
probe: t.Tensor
|
||||
A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
|
||||
fast_axis_angle: float
|
||||
The angle between the fast-axis of the polarizer and the horizontal axis
|
||||
|
||||
Returns:
|
||||
--------
|
||||
polarized probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
theta = math.radians(fast_axis_angle)
|
||||
exponent = t.exp(-1j * math.pi / 2 * t.ones(2, 2))
|
||||
jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
Returns:
|
||||
--------
|
||||
polarized probe: t.Tensor
|
||||
(...)x2x1xMxL
|
||||
"""
|
||||
probe = probe.to(dtype=t.cfloat)
|
||||
theta = math.radians(fast_axis_angle)
|
||||
exponent = t.exp(-1j * math.pi / 2 * t.ones(2, 2))
|
||||
jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]])
|
||||
probe = probe.transpose(-1, -3).transpose(-2, -4)
|
||||
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
|
||||
# Transpose it back
|
||||
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
|
||||
|
||||
# probe = t.rand(17, 7, 2, 6, 4)
|
||||
# polarizer = t.rand(7)
|
||||
|
||||
@@ -0,0 +1,438 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import numpy as np
|
||||
import torch as t
|
||||
from copy import copy
|
||||
import h5py
|
||||
import pathlib
|
||||
from CDTools.datasets import CDataset, Ptycho2DDataset, PolarizedPtycho2DDataset
|
||||
from CDTools.tools import data as cdtdata, initializers, polarization, interactions
|
||||
from CDTools.tools.polarization import apply_jones_matrix as jones
|
||||
from CDTools.tools import plotting
|
||||
from torch.utils import data as torchdata
|
||||
from matplotlib import pyplot as plt
|
||||
from matplotlib.widgets import Slider
|
||||
from matplotlib import ticker
|
||||
from CDTools.models import FancyPtycho, PolarizedFancyPtycho
|
||||
from torch.utils import data as torchdata
|
||||
from math import cos as cos, sin as sin
|
||||
import math
|
||||
|
||||
|
||||
|
||||
|
||||
angle = 87
|
||||
angle_2 = angle - 45
|
||||
|
||||
def polarizer(angle):
|
||||
theta = math.radians(angle)
|
||||
polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat)
|
||||
return polarizer
|
||||
exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
|
||||
theta2 = math.radians(angle_2)
|
||||
quarter_plate = t.tensor([[(cos(theta2))**2 + 1j * (sin(theta2))**2, (1 - 1j) * sin(theta2) * cos(theta2)],
|
||||
[(1 - 1j) * sin(theta2) * cos(theta2), (sin(theta2))**2 + 1j * (cos(theta2))**2]])
|
||||
|
||||
def build_from_quarters(jones1, jones2, jones3, jones4):
|
||||
x = t.cat((t.stack((jones1, jones1), dim=-1), t.stack((jones2, jones2), dim=-1)), dim=-1)
|
||||
y = t.cat((t.stack((jones3, jones3), dim=-1), t.stack((jones4, jones4), dim=-1)), dim=-1)
|
||||
x = t.stack((x, x), dim=-2)
|
||||
y = t.stack((y, y), dim=-2)
|
||||
return t.cat((x, y), dim=-2).to(dtype=t.cfloat)
|
||||
|
||||
jones_plate = t.matmul(quarter_plate, polarizer(angle))
|
||||
jones0 = polarizer(0)
|
||||
jones90 = polarizer(90)
|
||||
jones45 = polarizer(45)
|
||||
|
||||
|
||||
'''
|
||||
after applying the polarizer and the quarter_plate, the probe should get circularly polarized
|
||||
probe: no multiple modes, 1 diffr pattern
|
||||
2xMxL
|
||||
jones_matrix: same jones matrix applied to all the pixels
|
||||
2x2
|
||||
'''
|
||||
transpose = True
|
||||
def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr():
|
||||
probe = t.rand(2, 4, 4, dtype=t.cfloat)
|
||||
print(polarizer)
|
||||
out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose),
|
||||
quarter_plate, multiple_modes=False, transpose=transpose)
|
||||
print('expected shape:(2, 3, 4)')
|
||||
print('actual:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
assert np.allclose(np.real(out[0]), np.imag(out[1]))
|
||||
|
||||
|
||||
'''probe: no multiple modes, 1 diffr pattern
|
||||
2xMxL = 2x4x4
|
||||
jones_matrix: jones matrices differ from pixel to pixel
|
||||
2x2xMxL = 2x2x4x4
|
||||
4 quarters:
|
||||
1: [:, :, :-2, :-2] - circular_polarizer,
|
||||
2: [:, :, :-2, -2:] - 0
|
||||
3: [:, :, -2:, :-2] - 90
|
||||
4: [:, :, -2:, -2:] - 45
|
||||
'''
|
||||
def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr():
|
||||
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
|
||||
probe = t.ones(2, 4, 4).to(dtype=t.cfloat)
|
||||
print('jones:', jones_matr)
|
||||
# print('jones:', jones_matr)
|
||||
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
|
||||
|
||||
# jones -> (2,2,x,y), probe, output probe
|
||||
# interaction ->
|
||||
print('expected shape:(2, 4, 4)')
|
||||
print('simulated:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
assert (np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2]))
|
||||
and t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, -2:, -2:], out[0, -2:, -2:]))
|
||||
|
||||
|
||||
'''
|
||||
probe: no multiple modes, multiple diffr patterns
|
||||
Nx2xMxL = 3x2x4x4
|
||||
jones_matrix: same jones matrix applied to all the pixels
|
||||
Nx2x2 = 3x2x2
|
||||
3 different matrices for each probe:
|
||||
1: 0
|
||||
2: 45
|
||||
3: 90
|
||||
'''
|
||||
def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr():
|
||||
probe = t.ones(2, 4, 4, dtype=t.cfloat)
|
||||
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
|
||||
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
|
||||
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 2, 4, 4)')
|
||||
print('actual:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
assert (t.allclose(out[0, 0, :, :], t.ones(4, 4, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 0, :, :], out[1, 1])
|
||||
and t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat)))
|
||||
|
||||
|
||||
'''
|
||||
probe: no multiple modes, multiple diffr pattern
|
||||
Nx2xMxL = 3x2x4x4
|
||||
jones_matrix: jones matrices differ from pixel to pixel
|
||||
Nx2x2xMxL = 3x2x2x4x4
|
||||
jones matrix for the 1st pattern:
|
||||
1: quat plate, 2: 90, 3: 0, 4: 45
|
||||
jones matrix for the 2nd pattern:
|
||||
1: 0, 2: 45, 3: 90, 4: quat plate
|
||||
jones matrix for the 3rd pattern:
|
||||
1: 90, 2: 0, 3: 45, 4: quat plate
|
||||
'''
|
||||
def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr():
|
||||
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
|
||||
build_from_quarters(jones0, jones45, jones90, jones_plate),
|
||||
build_from_quarters(jones90, jones0, jones45, jones_plate)]
|
||||
jones_matr = t.stack(([i for i in jones_m]), dim=0)
|
||||
probe = t.ones(2, 4, 4, dtype=t.cfloat)
|
||||
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 2, 4, 4)')
|
||||
print('actual shape:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
|
||||
assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
|
||||
and t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
|
||||
|
||||
and t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
|
||||
and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
|
||||
|
||||
and t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
|
||||
and np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:])))
|
||||
|
||||
|
||||
'''
|
||||
probe: multiple modes, 1 diffr pattern
|
||||
Px2xMxL = 2x2x3x4
|
||||
jones_matrix: same jones matrix applied to all the pixels
|
||||
2x2 - quarter plate
|
||||
'''
|
||||
def test_apply_jones_matrix_mult_modes_1_pattern_one_jones_matr():
|
||||
probe = t.rand(2, 2, 3, 4, dtype=t.cfloat)
|
||||
out = jones(jones(probe, polarizer(angle), multiple_modes=True, transpose=transpose), quarter_plate, multiple_modes=True, transpose=transpose)
|
||||
|
||||
print('expected shape: (2, 2, 3, 4)')
|
||||
print('actual:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
assert (np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
|
||||
and np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :])))
|
||||
|
||||
|
||||
'''
|
||||
probe: multiple modes, 1 diffr pattern
|
||||
Px2xMxL = 3x2x4x4
|
||||
jones_matrix: jones matrices differ from pixel to pixel
|
||||
2xMxL = 2x4x4
|
||||
4 quarters:
|
||||
1: [:, :, :-2, :-2] - circular_polarizer,
|
||||
2: [:, :, :-2, -2:] - 0
|
||||
3: [:, :, -2:, :-2] - 90
|
||||
4: [:, :, -2:, -2:] - 45
|
||||
'''
|
||||
def test_apply_jones_matrix_mult_modes_1_pattern_diff_jones_matr():
|
||||
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
|
||||
probe = t.ones(2, 4, 4, dtype=t.cfloat)
|
||||
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 2, 4, 4)')
|
||||
print('actual shape:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
|
||||
and t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
|
||||
|
||||
and np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2]))
|
||||
and t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:])
|
||||
|
||||
and np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
|
||||
and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]))
|
||||
|
||||
|
||||
'''
|
||||
probe: multiple modes, multiple diffr patterns
|
||||
NxPx2xMxL = 3x7x2x3x4
|
||||
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
|
||||
Nx2x2 = 3x2x2
|
||||
3 different matrices for each probe in one mode:
|
||||
1: 0
|
||||
2: 45
|
||||
3: 90
|
||||
'''
|
||||
def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr():
|
||||
probe = t.ones(2, 3, 4, dtype=t.cfloat)
|
||||
# 1st mode
|
||||
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
# 2nd mode
|
||||
probe_mode2 = 10 * t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
|
||||
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
|
||||
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
|
||||
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 7, 2, 3, 4)')
|
||||
print('actual:', out.shape)
|
||||
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
|
||||
|
||||
# we'll be checking only one mode
|
||||
assert (t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
|
||||
and t.allclose(out[2, 6, 0, :, :], 3 * (10**6) * t.zeros(3, 4, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat)))
|
||||
|
||||
|
||||
'''
|
||||
probe: multiple modes, multiple diffr pattern
|
||||
NxPx2xMxL = 3x4x2x4x4
|
||||
jones_matrix: jones matrices differ from pixel to pixel
|
||||
Nx2x2xMxL = 3x2x2x4x4
|
||||
(differs across the patterns in each mode)
|
||||
jones matrix for the 1st pattern:
|
||||
1: quat plate, 2: 90, 3: 0, 4: 45
|
||||
jones matrix for the 2nd pattern:
|
||||
1: 0, 2: 45, 3: 90, 4: quat plate
|
||||
jones matrix for the 3rd pattern:
|
||||
1: 90, 2: 0, 3: 45, 4: quat plate
|
||||
'''
|
||||
def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr():
|
||||
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
|
||||
build_from_quarters(jones0, jones45, jones90, jones_plate),
|
||||
build_from_quarters(jones90, jones0, jones45, jones_plate)]
|
||||
jones_matr = t.stack(([i for i in jones_m]), dim=0)
|
||||
probe = t.ones(2, 4, 4, dtype=t.cfloat)
|
||||
# one mode
|
||||
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
|
||||
print('probe:', probe.shape)
|
||||
print('jones:', jones_matr.shape)
|
||||
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 2, 2, 4, 4)')
|
||||
print('actual shape:', out.shape)
|
||||
print('simulated patterns in one mode:', out[:, 3, :, :, :])
|
||||
|
||||
o = 10 ** 3
|
||||
# we'll be checking only one mode (4th)
|
||||
assert (np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
|
||||
and t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
|
||||
|
||||
and t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
|
||||
and t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
|
||||
|
||||
and t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
|
||||
and np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:])))
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
'''
|
||||
I REALIZED THAT ANOTHER POSSIBLE SISTUATION WE HAVEN'T CONSIDERED BEFORE IS WHEN THE JONES MATRIX DOESN'T DIFFER ACROSS THE PATTERNS
|
||||
IN OTHER WORDS, PROBE.SHAPE CONTAINS N (NUM OF PATTERNS) BUT JONES_MATRIX.SHAPE DOESN'T
|
||||
|
||||
|
||||
NOW, WE'LL TEST THE CASES WHEN THE PROBE DOESN'T DIFFER FROM PATTERN TO PATTERN
|
||||
'''
|
||||
|
||||
|
||||
'''
|
||||
probe: no multiple modes, multiple diffr patterns
|
||||
Nx2xMxL = 3x2x4x4
|
||||
jones_matrix: same jones matrix applied to all the pixels
|
||||
2x2 = 2x2 - a quarter waveplate
|
||||
'''
|
||||
def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr_1():
|
||||
probe = t.ones(2, 4, 4, dtype=t.cfloat)
|
||||
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
jones_matr = jones_plate
|
||||
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
|
||||
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 2, 4, 4)')
|
||||
print('actual:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
assert np.allclose(np.real(out[:, 0, :, :]), np.imag(out[:, 1, :, :]))
|
||||
|
||||
|
||||
'''
|
||||
probe: no multiple modes, multiple diffr pattern
|
||||
Nx2xMxL = 3x2x4x4
|
||||
jones_matrix: jones matrices differ from pixel to pixel
|
||||
2x2xMxL = 2x2x4x4
|
||||
1: quat plate, 2: 90, 3: 0, 4: 45
|
||||
|
||||
'''
|
||||
def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr_1():
|
||||
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
|
||||
probe = t.ones(2, 4, 4, dtype=t.cfloat)
|
||||
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 2, 4, 4)')
|
||||
print('actual shape:', out.shape)
|
||||
print('simulated:', out)
|
||||
|
||||
# check the 3rd pattern
|
||||
assert (np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
|
||||
and t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]))
|
||||
|
||||
|
||||
'''
|
||||
probe: multiple modes, multiple diffr patterns
|
||||
NxPx2xMxL = 3x7x2x3x4
|
||||
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
|
||||
2x2 = 2x2
|
||||
quarter wave plate
|
||||
'''
|
||||
def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr_1():
|
||||
probe = t.ones(2, 3, 4, dtype=t.cfloat)
|
||||
# 1st mode
|
||||
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
|
||||
jones_matr = jones_plate
|
||||
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
|
||||
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 7, 2, 3, 4)')
|
||||
print('actual:', out.shape)
|
||||
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
|
||||
|
||||
# we'll be checking only one mode
|
||||
assert np.allclose(np.real(out[:, :, 0, :, :]), np.imag(out[:, :, 1, :, :]))
|
||||
|
||||
|
||||
'''
|
||||
probe: multiple modes, multiple diffr pattern
|
||||
NxPx2xMxL = 3x4x2x4x4
|
||||
jones_matrix: jones matrices differ from pixel to pixel
|
||||
2x2xMxL = 2x2x4x4
|
||||
jones matrix:
|
||||
1: quat plate, 2: 90, 3: 0, 4: 45
|
||||
'''
|
||||
def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr_1():
|
||||
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
|
||||
probe = t.ones(2, 4, 4, dtype=t.cfloat)
|
||||
# one mode
|
||||
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
|
||||
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
|
||||
print('probe:', probe.shape)
|
||||
print('jones:', jones_matr.shape)
|
||||
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
|
||||
|
||||
print('expected shape: (3, 2, 2, 4, 4)')
|
||||
print('actual shape:', out.shape)
|
||||
print('simulated patterns in one mode:', out[:, 3, :, :, :])
|
||||
|
||||
o = 10 ** 2
|
||||
# we'll be checking only one mode (3th)
|
||||
assert (np.allclose(np.real(out[0, 2, 0, :-2, :-2]), np.imag(out[0, 2, 1, :-2, :-2]))
|
||||
and t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
|
||||
and t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:]))
|
||||
|
||||
return None
|
||||
|
||||
Reference in New Issue
Block a user