cSAXS/cdtools
0
0 Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-11-08 15:33:53 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merging a whole buncha work back into master

Browse Source
This commit is contained in:
Abe Levitan committed 2022-04-18 15:51:45 -07:00
2 parents 1b5ef2271c 20a1479859
commit f3cee949e0
22 files changed
+3382 -771
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CDTools/tools/analysis/analysis.py
+362 -15
View File
@@ -15,7 +15,8 @@ from scipy import special
__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction']
'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction',
'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error']
def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
@@ -456,7 +457,7 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
return cor
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
"""Calculates a Fourier ring correlation between two images
This function requires an input of a basis to allow for FRC calculations
@@ -479,6 +480,8 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
Number of bins to break the FRC up into
snr : float
The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
limit : str
Default is 'side'. What is the highest frequency to calculate the FRC to? If 'side', it chooses the side of the Fourier transform, if 'corner' it goes fully to the corner.
Returns
-------
@@ -507,14 +510,14 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
if nbins is None:
nbins = np.max(im1[im_slice].shape) // 4
nbins = np.max(im1[im_slice].shape) // 8
f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2))
f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2))
cor_fft = f1 * t.conj(f2)
cor_fft = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) * \
t.fft.fftshift(t.conj(t.fft.fft2(im2[im_slice])),dim=(-1,-2))
F1 = t.abs(t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)))**2
F2 = t.abs(t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)))**2
F1 = t.abs(f1)**2
F2 = t.abs(f2)**2
di = np.linalg.norm(basis[:,0])
@@ -526,21 +529,41 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
if limit.lower().strip() == 'side':
max_i = np.max(i_freqs)
max_j = np.max(j_freqs)
frc_range = [0, max(max_i,max_j)]
elif limit.lower().strip() == 'corner':
frc_range = [0, np.max(Rs)]
else:
raise ValueError('Invalid FRC limit: choose "side" or "corner"')
numerator, bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=cor_fft.numpy())
denominator_F1, bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=F1.detach().cpu().numpy())
denominator_F2, bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=F2.detach().cpu().numpy())
n_pix, bins = np.histogram(Rs, bins=nbins, range=frc_range)
numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft.numpy())
denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy())
denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy())
n_pix, bins = np.histogram(Rs,bins=nbins)
frc = np.abs(numerator / np.sqrt(denominator_F1*denominator_F2))
n_pix = n_pix / 4 # This is for an apodized image, apodized with a hann window
frc = np.abs(numerator) / np.sqrt(denominator_F1*denominator_F2)
# This moves from combined-image SNR to single-image SNR
snr /= 2
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
# NOTE: I should update this to produce lots of different threshold curves
# sigma, 2sigma, 3sigma, traditional FRC 1-bit, my better one, n_pix, etc.
threshold = (snr + (2 * np.sqrt(snr) + 1) / np.sqrt(n_pix)) / \
(1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
my_threshold = np.sqrt(snr**2 + (2*snr**2 + 2*snr + 1)/n_pix) / \
np.sqrt(snr**2 + 2 * snr + 1 + 2*snr**2 / n_pix)
twosigma_threshold = 2/ np.sqrt(n_pix)
if not im_np:
bins = t.tensor(bins)
frc = t.tensor(frc)
@@ -581,6 +604,7 @@ def calc_vn_entropy(matrix):
entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig)
return entropy
def calc_top_mode_fraction(matrix):
"""Calculates the fraction of total power in the top mode of a density matrix
@@ -610,3 +634,326 @@ def calc_top_mode_fraction(matrix):
eig = np.linalg.eigh(matrix)[0]
fraction = np.max(eig) / np.sum(eig)
return fraction
def calc_rms_error(field_1, field_2, align_phases=True, normalize=False,
dims=2):
"""Calculates the root-mean-squared error between two complex wavefields
The formal definition of this function is:
output = norm * sqrt(mean(abs(field_1 - gamma * field_2)**2))
Where norm is an optional normalization factor, and gamma is an
optional phase factor which is appropriate when the wavefields suffer
from a global phase degeneracy as is often the case in diffractive
imaging.
The normalization is defined as the square root of the total intensity
contained in field_1, which is appropriate when field_1 represents a
known ground truth:
norm = sqrt(mean(abs(field_1)**2))
The phase offset is an analytic expression for the phase offset which
will minimize the RMS error between the two wavefields:
gamma = exp(1j * angle(sum(field_1 * conj(field_2))))
This implementation is stable even in cases where field_1 and field_2
are completely orthogonal.
In the definitions above, the field_n are n-dimensional wavefields. The
dimensionality of the wavefields can be altered via the dims argument,
but the default is 2 for a 2D wavefield.
Parameters
----------
field_1 : array
The first complex-valued field
field_2 : array
The second complex-valued field
align_phases : bool
Default is True, whether to account for a global phase offset
normalize : bool
Default is False, whether to normalize to the intensity of field_1
dims : (int or tuple of python:ints)
Default is 2, the number of final dimensions to reduce over.
Returns
-------
rms_error : float or t.Tensor
The RMS error, or tensor of RMS errors, depending on the dim argument
"""
sumdims = tuple(d - dims for d in range(dims))
if align_phases:
# Keepdim allows us to broadcast the result correctly when we
# multiply by the fields
gamma = t.exp(1j * t.angle(t.sum(field_1 * t.conj(field_2), dim=sumdims,
keepdim=True)))
else:
gamma = 1
if normalize:
norm = 1 / t.mean(t.abs(field_1)**2, dim=sumdims)
else:
norm = 1
difference = field_1 - gamma * field_2
return t.sqrt(norm * t.mean(t.abs(difference)**2, dim=sumdims))
def calc_fidelity(fields_1, fields_2, dims=2):
"""Calculates the fidelity between two density matrices
The fidelity is a comparison metric between two density matrices
(i.e. mutual coherence functions) that extends the idea of the
overlap to incoherent light. As a reminder, the overlap between two
fields is:
overlap = abs(sum(field_1 * field_2))**2
Whereas the fidelity is defined as:
fidelity = trace(sqrt(sqrt(dm_1) <dot> dm_2 <dot> sqrt(dm_1)))**2
where dm_n refers to the density matrix encoded by fields_n such
that dm_n = fields_n <dot> fields_<n>.conjtranspose(), sqrt
refers to the matrix square root, and <dot> is the matrix product.
This is not a practical implementation, however, as it is not feasible
to explicitly construct the matrices dm_1 and dm_2 in memory. Therefore,
we take advantage of the alternate definition based directly on the
fields_<n> parameter:
fidelity = sum(svdvals(fields_1 <dot> fields_2.conjtranspose()))**2
In the definitions above, the fields_n are regarded as collections of
wavefields, where each wavefield is by default 2-dimensional. The
dimensionality of the wavefields can be altered via the dims argument,
but the fields_n arguments must always have at least one more dimension
than the dims argument. Any additional dimensions are treated as batch
dimensions.
Parameters
----------
fields_1 : array
The first set of complex-valued field modes
fields_2 : array
The second set of complex-valued field modes
dims : int
Default is 2, the number of final dimensions to reduce over.
Returns
-------
fidelity : float or t.Tensor
The fidelity, or tensor of fidelities, depending on the dim argument
"""
fields_1 = t.as_tensor(fields_1)
fields_2 = t.as_tensor(fields_2)
mult = fields_1.unsqueeze(-dims-2) * fields_2.unsqueeze(-dims-1).conj()
sumdims = tuple(d - dims for d in range(dims))
mat = t.sum(mult,dim=sumdims)
# Because I think this is the nuclear norm squared, I would like to swap
# Out the definition for this, but I need to test it before swapping.
# It also probably makes sense to implement sqrt_fidelity separately
# because that's more important
#return t.linalg.matrix_norm(mat, ord='nuc')**2
# I think this is just the nuclear norm.
svdvals = t.linalg.svdvals(mat)
return t.sum(svdvals, dim=-1)**2
def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2):
"""Calculates a generalization of the root-mean-squared error between two complex wavefields
This function calculates an generalization of the RMS error which uses the
concept of fidelity to extend it to capture the error between
incoherent wavefields, defined as a mode decomposition. The extension has
several nice properties, in particular:
1) For coherent wavefields, it precisely matches the RMS error including
a correction for the global phase degeneracy (align_phases=True)
2) All mode decompositions of either field that correspond to the same
density matrix / mutual coherence function will produce the same
output
3) The error will only be zero when comparing mode decompositions that
correspond to the same density matrix.
4) Due to (2), one need not worry about the ordering of the modes,
properly orthogonalizing the modes, and it is even possible to
compare mode decompositions with different numbers of modes.
The formal definition of this function is:
output = norm * sqrt(mean(abs(fields_1)**2)
+ mean(abs(fields_2)**2)
- 2 * sqrt(fidelity(fields_1,fields_2)))
Where norm is an optional normalization factor, and the fidelity is
defined based on the mean, rather than the sum, to match the convention
for the root *mean* squared error.
The normalization is defined as the square root of the total intensity
contained in fields_1, which is appropriate when fields_1 represents a
known ground truth:
norm = sqrt(mean(abs(fields_1)**2))
In the definitions above, the fields_n are regarded as collections of
wavefields, where each wavefield is by default 2-dimensional. The
dimensionality of the wavefields can be altered via the dims argument,
but the fields_n arguments must always have at least one more dimension
than the dims argument. Any additional dimensions are treated as batch
dimensions.
Parameters
----------
fields_1 : array
The first set of complex-valued field modes
fields_2 : array
The second set of complex-valued field modes
normalize : bool
Default is False, whether to normalize to the intensity of fields_1
dims : (int or tuple of python:ints)
Default is 2, the number of final dimensions to reduce over.
Returns
-------
rms_error : float or t.Tensor
The generalized RMS error, or tensor of generalized RMS errors, depending on the dim argument
"""
fields_1 = t.as_tensor(fields_1)
fields_2 = t.as_tensor(fields_2)
npix = t.prod(t.as_tensor(fields_1.shape[-dims:],dtype=t.int32))
sumdims = tuple(d - dims - 1 for d in range(dims+1))
fields_1_intensity = t.sum(t.abs(fields_1)**2,dim=sumdims) / npix
fields_2_intensity = t.sum(t.abs(fields_2)**2,dim=sumdims) / npix
fidelity = calc_fidelity(fields_1, fields_2, dims=dims) / npix**2
result = fields_1_intensity + fields_2_intensity - 2 * t.sqrt(fidelity)
if normalize:
result /= fields_1_intensity
return t.sqrt(result)
def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1.):
"""Calculates a Fourier ring correlation between two images
This function requires an input of a basis to allow for FRC calculations
to be related to physical units.
Like other analysis functions, this can take input in numpy or pytorch,
and will return output in the respective format.
Parameters
----------
im1 : array
The first image, a complex or real valued array
im2 : array
The first image, a complex or real valued array
basis : array
The basis for the images, defined as is standard for datasets
im_slice : slice
Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
nbins : int
Number of bins to break the FRC up into
snr : float
The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
Returns
-------
freqs : array
The frequencies associated with each FRC value
FRC : array
The FRC values
threshold : array
The threshold curve for comparison
"""
im_np = False
if isinstance(fields_1, np.ndarray):
fields_1 = t.as_tensor(fields_)
im_np = True
if isinstance(fields_2, np.ndarray):
fields_2 = t.as_tensor(fields_2)
im_np = True
if isinstance(basis, np.ndarray):
basis = t.tensor(basis)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
if nbins is None:
nbins = np.max(fields_1[...,im_slice].shape[-2:]) // 4
f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2))
f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2))
cor_fft = f1 * t.conj(f2)
F1 = t.abs(f1)**2
F2 = t.abs(f2)**2
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di))
j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
# This line is used to get a set of bins that matches the logic
# used by np.histogram, so that this function will match the choices
# of bin edges that comes from the non-generalized version. This also
# gets us the count on the number of pixels per bin so we can calculate
# the threshold curve
n_pix, bins = np.histogram(Rs,bins=nbins)
frc = []
for i in range(len(bins)-1):
mask = t.logical_and(Rs<bins[i+1], Rs>=bins[i])
masked_f1 = f1 * mask[...,:,:]
masked_f2 = f2 * mask[...,:,:]
numerator = t.sqrt(calc_fidelity(masked_f1, masked_f2))
denominator_f1 = t.sqrt(calc_fidelity(masked_f1, masked_f1))
denominator_f2 = t.sqrt(calc_fidelity(masked_f2, masked_f2))
frc.append(numerator / t.sqrt((denominator_f1 * denominator_f2)))
frc = np.array(frc)
# This moves from combined-image SNR to single-image SNR
snr /= 2
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
(1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
if not im_np:
bins = t.tensor(bins)
frc = t.tensor(frc)
threshold = t.tensor(threshold)
return bins[:-1], frc, threshold
CDTools/tools/data/data.py
+9 -6
View File
@@ -93,7 +93,7 @@ def get_sample_info(cxi_file):
"""
if 'entry_1/sample_1' not in cxi_file:
return None
s1 = cxi_file['entry_1/sample_1']
metadata_attrs = ['name','description','unit_cell_group']
@@ -317,7 +317,7 @@ def get_dark(cxi_file):
return darks
def get_data(cxi_file, cut_zeroes = True):
def get_data(cxi_file, cut_zeros = True):
"""Returns an array with the full stack of detector data defined in the cxi file object
This function will make sure to check all the various places that it's
@@ -333,6 +333,8 @@ def get_data(cxi_file, cut_zeroes = True):
----------
cxi_file : h5py.File
A file object to be read
cut_zeros : bool
Default True, whether to set all negative data to zero
Returns
-------
@@ -351,9 +353,8 @@ def get_data(cxi_file, cut_zeroes = True):
raise KeyError('Data is not defined within cxi file')
data = cxi_file[pull_from][:]
# Use maximum in-place to avoid allocating any more memory than is needed
if cut_zeroes:
if cut_zeros:
np.maximum(data,0,data)
if 'axes' in cxi_file[pull_from].attrs:
@@ -642,7 +643,8 @@ def add_dark(cxi_file, dark):
d1.create_dataset('data_dark',data=dark)
def add_data(cxi_file, data, axes=None):
def add_data(cxi_file, data, axes=None, compression='gzip',
chunks=True):
"""Adds the specified data to the cxi file
It will add the data unchanged to the file, placing it in two spots:
@@ -673,7 +675,8 @@ def add_data(cxi_file, data, axes=None):
if isinstance(data, t.Tensor):
data = data.detach().cpu().numpy()
det1.create_dataset('data', data=data)
det1.create_dataset('data', data=data, compression=compression,
chunks=chunks)
data1['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data')
if axes is not None:
CDTools/tools/initializers/initializers.py
+1 -1
View File
@@ -633,7 +633,7 @@ def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
dominant_fft = far_field(dominant_mode)
shape = dominant_mode.shape
center = ((shape[-2]-1)//2, (shape[-1]-1)//2)
center = ((shape[-2])//2, (shape[-1])//2)
i, j = np.mgrid[:shape[-2], :shape[-1]]
i = t.tensor(i - center[0]).to(dtype=dominant_fft.dtype,
CDTools/tools/interactions/interactions.py
+69 -74
View File
@@ -17,14 +17,14 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations',
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
"""Takes real space translations and outputs them in pixel space
This works for any 2D ptychography geometry. It takes in
A set of translations in (x,y) space and outputs the same translations
in internal pixel units perpendicular to the detector.
in internal pixel units perpendicular to the detector.
It uses information on the wavefield basis and, if defined, the
sample normal, to perform the conversion.
The assumed geometry is incoming radiation with a wavevector parallel
to the +z axis, [0,0,1]. The default sample orientation has a surface
normal parallel to this direction
@@ -33,7 +33,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
----------
basis : torch.Tensor
The real space basis the wavefields are defined in
translations : torch.Tensor
translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
surface_normal : torch.Tensor
Optional, the sample's surface normal
@@ -43,7 +43,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
pixel_translations : torch.Tensor
A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation
"""
projection_1 = t.Tensor([[1,0,0],
[0,1,0],
[0,0,0]]).to(device=translations.device,dtype=translations.dtype)
@@ -68,18 +68,18 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
return pixel_translations[0]
else:
return pixel_translations
def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,0,1])):
"""Takes pixel-space translations and outputs them in real space
This works for any 2D ptychography geometry. It takes in
A set of internal pixel unit translations in (i,j) space and
outputs the same translations real (x,y) space
It uses information on the wavefield basis and, if defined, the
sample normal, to perform the conversion.
The assumed geometry is incoming radiation with a wavevector parallel
to the +z axis, [0,0,1]. The default sample orientation has a surface
normal parallel to this direction. Because of this, the z direction
@@ -96,7 +96,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
Returns
-------
real_translations : torch.Tensor
real_translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
"""
projection_1 = t.Tensor([[1,0,0],
@@ -129,7 +129,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
def project_translations_to_sample(sample_basis, translations):
"""Takes real space translations and outputs them in pixels in a sample basis
This projection function is designed for the Bragg2DPtycho class. More
broadly, it works to take a set of translations in the lab frame and
convert each one into two values. First, an (i,j) value in pixels
@@ -145,7 +145,7 @@ def project_translations_to_sample(sample_basis, translations):
relative amount the probe needs to be propagated to reach any given
location), a positive motion along the z-axis of the probe forming optics
will lead to a negative propagation distance.
The assumed geometry is incoming radiation with a wavevector parallel
to the +z axis, [0,0,1].
@@ -153,7 +153,7 @@ def project_translations_to_sample(sample_basis, translations):
----------
sample_basis : torch.Tensor
The real space basis the wavefields are defined in
translations : torch.Tensor
translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
Returns
@@ -171,7 +171,7 @@ def project_translations_to_sample(sample_basis, translations):
# Then we calculate a matrix which can do the projection
propagation_dir = t.Tensor(np.array([0,0,1])).to(
device=surface_normal.device,
dtype=surface_normal.dtype)
@@ -179,13 +179,13 @@ def project_translations_to_sample(sample_basis, translations):
I = t.eye(3).to(
device=surface_normal.device,
dtype=surface_normal.dtype)
# Here we're setting up a matrix-vector equation mat*answer=input
# At some point ger will need to be replaced by outer, but for now
# outer many places still don't have new enough versions of torch.
mat = t.cat((I - t.ger(propagation_dir,propagation_dir),
surface_normal.unsqueeze(0)))
# And we invert the matrix to do the projection
projector = t.pinverse(mat)[:,:3].to(device=translations.device,
dtype=translations.dtype)
@@ -200,7 +200,7 @@ def project_translations_to_sample(sample_basis, translations):
device=translations.device,
dtype=translations.dtype)
sample_projection = t.mm(basis_vectors_inv, projector).t()
prop_projection = t.mm(propagation_dir_inv, projector).t()
@@ -218,9 +218,9 @@ def project_translations_to_sample(sample_basis, translations):
return pixel_translations[0], propagations[0]
else:
return pixel_translations, propagations
def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj=False):
"""Returns a stack of exit waves without accounting for subpixel shifts
@@ -229,15 +229,15 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
shifting the probe by the rounded value of the translation
If multiple_modes is set to False, any additional dimensions in the
ptycho_2D_round function will be assumed to correspond to the translation
index. If multiple_modes is set to true, the (-4th) dimension of the probe
will always be assumed to be defining a set of (P) incoherently mixing
modes to be broadcast all translation indices. If any additional dimensions
closer to the start exist, they will be assumed to be translation indices
Parameters
----------
probe : torch.Tensor
@@ -251,7 +251,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
Returns
-------
exit_waves : torch.Tensor
exit_waves : torch.Tensor
An (N)x(P)xMxL tensor of the calculated exit waves
"""
@@ -260,9 +260,9 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
translations = translations[None,:]
single_translation = True
integer_translations = t.round(translations).to(dtype=t.int32)
if upsample_obj:
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2]//2,
tr[1]:tr[1]+probe.shape[-1]//2]
@@ -292,7 +292,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
@@ -322,7 +322,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
if translations.dim() == 1:
translations = translations[None,:]
single_translation = True
# Separate the translations into a part that chooses the window
# And a part that defines the windowing function
integer_translations = t.floor(translations)
@@ -342,15 +342,15 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
sel01 = t.cat((probe[:,-1:],probe[:,:-1]),dim=1)
sel10 = t.cat((probe[-1:,:],probe[:-1,:]),dim=0)
sel11 = t.cat((sel01[-1:,:],sel01[:-1,:]),dim=0)
selection = sel00 * (1-sp[0])*(1-sp[1]) + \
sel10 * sp[0]*(1-sp[1]) + \
sel01 * (1-sp[0])*sp[1] + \
sel11 * sp[0]*sp[1]
obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
exit_waves.append(selection * obj_slice)
else:
for tr, sp in zip(integer_translations,
@@ -359,16 +359,16 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
# Here we subpixel shift the object by (-i,-j) after
# slicing out the correct translation of the probe
#
sel00 = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
sel01 = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]+1:tr[1]+1+probe.shape[1]]
sel10 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
sel11 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
tr[1]+1:tr[1]+1+probe.shape[1]]
@@ -387,7 +387,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, polarized=False, polarizer=None, analyzer=None):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
@@ -423,26 +423,27 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
exit_waves : torch.Tensor
An (N)x(P)xMxL tensor of the calculated exit waves
"""
single_translation = False
if translations.dim() == 1:
translations = translations[None,:]
translations = translations[None, :]
single_translation = True
# Separate the translations into a part that chooses the window
# And a part that defines the windowing function
integer_translations = t.floor(translations)
subpixel_translations = translations - integer_translations
integer_translations = integer_translations.to(dtype=t.int32)
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
if polarized:
selections = t.stack([obj[:, :,tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
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
exit_waves = []
if shift_probe:
i = t.arange(probe.shape[-2],device=probe.device,dtype=t.float32) \
@@ -454,15 +455,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,:,:,:]
@@ -470,11 +467,8 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
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,
dim=(-1,-2)))
if not polarized:
if multiple_modes: # Multi-mode probe
output = shifted_probe * selections[...,None,:,:]
@@ -483,12 +477,10 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
# selections: Nx2x2xMxL
# probe: Nx(P)x2x1xMxL
else:
# print('F', shift_probe)
output = polarization.apply_jones_matrix(shifted_probe, selections)
output = polarization.apply_jones_matrix(shifted_probe, selections, multiple_modes=multiple_modes)
else:
raise NotImplementedError('Object shift not yet implemented')
if single_translation:
return output[0]
else:
@@ -497,7 +489,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, padding=10):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
@@ -513,7 +505,7 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
on the input wavefield, and the first two indexes index differences from
that pixel. It is easier to interpret the resulting matrix though if the
latter two indices index locations in the output plane. NOTE: I believe
this change has now been made
this change has now been made
Parameters
----------
@@ -537,17 +529,17 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
if translations.dim() == 1:
translations = translations[None,:]
single_translation = True
# Separate the translations into a part that chooses the window
# And a part that defines the windowing function
integer_translations = t.floor(translations)
subpixel_translations = translations - integer_translations
integer_translations = integer_translations.to(dtype=t.int32)
exit_waves = []
B = s_matrix.shape[0]//2
if shift_probe:
i = t.arange(probe.shape[-2]) - probe.shape[-2]//2
j = t.arange(probe.shape[-1]) - probe.shape[-1]//2
@@ -556,14 +548,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
J = 2 * np.pi * J.to(t.float32) / probe.shape[-1]
I = I.to(dtype=probe.dtype,device=probe.device)
J = J.to(dtype=probe.dtype,device=probe.device)
for tr, sp in zip(integer_translations,
subpixel_translations):
fft_probe = t.fft.fftshift(t.fft.fft2(probe), dim=(-1,-2))
shifted_fft_probe = fft_probe * t.exp(1j*(-sp[0]*I - sp[1]*J))
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[-2]+2*B,
tr[1]:tr[1]+probe.shape[-1]+2*B]
@@ -572,14 +564,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
device=s_matrix_slice.device,
dtype=s_matrix_slice.dtype)
for i in range(s_matrix.shape[0]):
for j in range(s_matrix.shape[1]):
output [i:i+probe.shape[-2],j:j+probe.shape[-1]] += \
shifted_probe * s_matrix_slice[i,j,i:i+probe.shape[-2],j:j+probe.shape[-1]]
exit_waves.append(output)
else:
raise NotImplementedError('Object shift not yet implemented')
@@ -587,11 +579,11 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
return exit_waves[0]
else:
return t.stack(exit_waves)
def RPI_interaction(probe, obj):
"""Returns an exit wave from a high-res probe and a low-res obj
In this interaction, the probe and object arrays are assumed to cover
the same physical region of space, but with the probe array sampling that
region of space more finely. Thus, to do the interaction, the object
@@ -601,7 +593,7 @@ def RPI_interaction(probe, obj):
method and is not commonly used elsewhere.
This also works with object functions that have an extra first dimension
for an incoherently mixing model.
for an incoherently mixing model.
Parameters
@@ -618,17 +610,20 @@ def RPI_interaction(probe, obj):
"""
# TODO: The upsampling only works for arrays of even dimension!
# The far-field propagator is just a 2D FFT but with an fftshift
fftobj = propagators.far_field(obj)
# We calculate the padding that we need to do the upsampling
pad0l = (probe.shape[-2] - obj.shape[-2])//2
pad0r = probe.shape[-2] - obj.shape[-2] - pad0l
pad1l = (probe.shape[-1] - obj.shape[-1])//2
# This is carefully set up to keep the zero-frequency pixel in the correct
# location as the overall shape changes. Don't mess with this without
# having thought about this carefully.
pad2l = probe.shape[-2]//2 - obj.shape[-2]//2
pad2r = probe.shape[-2] - obj.shape[-2] - pad2l
pad1l = probe.shape[-1]//2 - obj.shape[-1]//2
pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad2l, pad2r))
# Again, just an inverse FFT but with an fftshift
upsampled_obj = propagators.inverse_far_field(fftobj)
CDTools/tools/plotting/__init__.py
+1
View File
@@ -1 +1,2 @@
from CDTools.tools.plotting.plotting import *
from CDTools.tools.plotting.polarized_plotting import *
CDTools/tools/plotting/plotting.py
+54 -43
View File
@@ -84,7 +84,7 @@ def get_units_factor(units):
def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, interpolation=None, **kwargs):
"""Plots an image with a colorbar and on an appropriate spatial grid
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
@@ -94,7 +94,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
Finally, if a function is passed to the plot_func argument, this function
will be called on each slice of data before it is plotted. This is used
internally to enable the plot_real, plot_image, plot_phase, etc. functions.
Parameters
----------
@@ -122,7 +122,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
# convert to numpy
if isinstance(im, t.Tensor):
# If final dimension is 2, assume it is a complex array. If not,
@@ -144,7 +144,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
title = plt.gca().get_title()
fig.clear()
# If im only has two dimensions, this reshape will add a leading
# dimension, and update will be called on index 0. If it has 3 or more
# dimensions, then all the leading dimensions will be compressed into
@@ -153,9 +153,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
reshaped_im = im.reshape(-1,s[-2],s[-1])
num_images = reshaped_im.shape[0]
fig.plot_idx = idx % num_images
to_plot = plot_func(reshaped_im[fig.plot_idx])
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
@@ -183,7 +183,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
plt.title(title)
if len(im.shape) >= 3:
@@ -196,14 +196,14 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
result = make_plot(0)
update = make_plot
def on_action(event):
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
update(fig.plot_idx - 1)
elif event.key == 'down' or event.button == 'down':
@@ -219,9 +219,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
fig.my_callbacks = []
fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action))
fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action))
return result
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
"""Plots the real part of a complex array with dimensions NxM
@@ -258,7 +258,7 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
@@ -306,7 +306,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
If a basis is explicitly passed, the image will be plotted in real-space
coordinates.
Parameters
----------
im : array
@@ -418,7 +418,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
units=units, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs):
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
"""Plots a set of probe translations in a nicely formatted way
Parameters
@@ -431,6 +431,8 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
Default is um, units to report in (assuming input in m)
lines : bool
Whether to plot lines indicating the path taken
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -455,6 +457,9 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
translations = translations * factor
plt.plot(translations[:,0], translations[:,1],'k.')
if invert_xaxis:
plt.gca().invert_xaxis()
if lines:
plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
plt.xlabel('X (' + units + ')')
@@ -463,7 +468,7 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
return fig
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe'):
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe', invert_xaxis=True):
"""Plots a set of nanomap data in a flexible way
Parameters
@@ -478,6 +483,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
Default is um, units to report in (assuming input in m)
convention : str
Default is 'probe', alternative is 'obj'. Whether the translations refer to the probe or object.
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
Returns
-------
@@ -511,7 +518,9 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
s /= 4 # A rough value to make the size work out
plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
if invert_xaxis:
plt.gca().invert_xaxis()
plt.gca().set_facecolor('k')
plt.xlabel('Translation x (' + units + ')')
plt.ylabel('Translation y (' + units + ')')
@@ -522,12 +531,12 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs):
"""Plots a nanomap, with an image or stack of images for each point
In many situations, ptychography data or the output of ptychography
reconstructions is formatted as a set of images associated with various
points in real space. This function is designed to allow for browsing
through this kind of data, by making it possible to visualize a
"""
# This should pull heavily from the dataset.inspect function
@@ -560,11 +569,11 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(translations.shape[0]) * s0
s[idx] *= 4
return s
def update_colorbar(im):
#
# This solves the problem of the colorbar being changed
@@ -573,29 +582,29 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
# This is needed to update the colorbar
# only change limits if array contains multiple values
if np.min(im.get_array()) != np.max(im.get_array()):
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
#
# The meatiest part of this program, here we just go through and
# set up the plot how we want it
#
# First we set up the left-hand plot, which shows an overview map
axes[0].set_title('Relative Displacement Map')
translations = translations.detach().cpu().numpy()
if convention.lower() != 'probe':
translations = translations * -1
s = calculate_sizes(0)
nanomap_units_factor = get_units_factor(nanomap_units)
nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0],
nanomap_units_factor * translations[:,1],
@@ -616,7 +625,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# where the colorbar should have been to avoid stretching the
# nanomap plot, while still not showing the (now useless) colorbar.
cb1.remove()
# Now we set up the second plot, which shows the individual
# diffraction patterns
axes[1].set_title(image_title)
@@ -636,7 +645,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This fails if the basis is not rectangular
basis_norm = np.linalg.norm(np_basis, axis = 0)
basis_norm = basis_norm * get_units_factor(image_units)
extent = [0, example_im.shape[-1]*basis_norm[1], 0,
example_im.shape[-2]*basis_norm[0]]
else:
@@ -657,9 +666,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].text_box.set_path_effects(
[patheffects.Stroke(linewidth=2, foreground='black'),
patheffects.Normal()])
meas = axes[1].imshow(im, extent=extent, cmap=cmap)
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
@@ -667,7 +676,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
cb2.ax.tick_params(labelrotation=20)
cb2.ax.set_title(image_colorbar_title, size="medium", pad=5)
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
# This function handles all the updating, except for moving the
# slider value. This is done because the slider widget is
# ultimately responsible for triggering an update, so all other
@@ -677,7 +686,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# We have to explicitly make it an integer because the slider will
# output floats (even if they are still integer-valued)
idx = int(idx)
# Get the new data for this index
im = get_image_func(idx)
if len(im.shape) >= 3:
@@ -688,22 +697,22 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].image_idx = im_idx
axes[1].text_box.set_text(str(im_idx))
im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
# Now we resize the nanomap to show the new selection
axes[0].collections[0].set_sizes(calculate_sizes(idx))
# And we update the data in the image as well
ax_im = axes[1].images[-1]
ax_im.set_data(im)
update_colorbar(ax_im)
#
# Now we define the functions to handle various kinds of events
# that can be thrown our way
#
# We start by creating the slider here, so it can be used
# by the update hooks.
slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d")
@@ -717,7 +726,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# while the mouse is within the image display
im = im.reshape(-1,im.shape[-2],im.shape[-1])
im_idx = axes[1].image_idx
if event.key == 'up' or event.button == 'up' \
or event.key == 'left':
im_idx = (im_idx - 1) % im.shape[0]
@@ -735,7 +744,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up' or event.key == 'left':
idx = slider.val - 1
elif event.key == 'down' or event.button == 'down' or event.key == 'right':
@@ -744,7 +753,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This prevents errors from being thrown on irrelevant key
# or mouse input
return
# Handle the wraparound and trigger the update
idx = int(idx) % translations.shape[0]
slider.set_val(idx)
@@ -755,17 +764,19 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# for example, scroll events that happen over the nanomap
if event.mouseevent.button == 1:
slider.set_val(event.ind[0])
# Here we connect the various update functions
cid1 = fig.canvas.mpl_connect('pick_event',on_pick)
cid2 = fig.canvas.mpl_connect('key_press_event',on_action)
cid3 = fig.canvas.mpl_connect('scroll_event',on_action)
# It's so dumb that matplotlib doesn't automatically track this for you
fig.nanomap_cids = [cid1,cid2,cid3]
fig.nanomap_cids = [cid1,cid2,cid3]
slider.on_changed(update)
# Throw an extra update into the mix just to get rid of any things
# (like the nanomap dot sizes) that otherwise would change on the
# first update
update(0)
return fig
CDTools/tools/plotting/polarized_plotting.py
+342
View File
@@ -0,0 +1,342 @@
import torch as t
from matplotlib import pyplot as plt
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
from matplotlib.widgets import Slider
from matplotlib import ticker, patheffects
from CDTools.tools.polarization import polarization
all = ['visualize_components_amplitudes',
'visualize_attenuations',
'visualize_fast_axes',
'visualize_global_axes',
'visualize_phase_ret',
'visualize_probe']
def iterator(a, b):
'''
A helper function we can use to facilitate to process of iterating over 2D arrays
'''
x, y = t.arange(a), t.arange(b)
x, y = t.meshgrid(x, y)
x, y = t.ravel(x), t.ravel(y)
return (x, y)
def plot_probe_ellipse(a=1, b=1, phase_ret=0, scale=1, x0=4, y0=5):
'''
Given a probe vector at a point (x0, y0), visualizes ellipticity
of its polarization
Parameters:
-----------
a : 1D np.array
An amplitude of the horizontal component of the probe vector
b : 1D np.array
An amplitude of the vertical component of the probe vector
phase_ret : 1D np.array
A phase difference in radians bettween the phases of the y and x components
scale : int
A scaling factor
x0, y0: 1D np.array or float
Defines the location of the vector to be plotted
Returns:
--------
x, y set of points to plot a single ellipse
'''
theta = np.linspace(0, 2*np.pi, 20)
x = x0 + scale * a * np.real(np.exp(1j * theta))
y = y0 + scale * b * np.real(np.exp(1j * (theta + phase_ret)))
return x, y
def plot_attenuations(atten_slow=1, atten_fast=1, fast_ax_angle=0, scale=1, x0=4, y0=4):
'''
Plots attenuations along the fast and slow axes
Parameters:
-----------
atten_slow: 1D np.array
Attenuation along the slow axis
atten_fast: 1D np.array
Attenuation along the fast axis
fast_ax_angle: 1D np.array
An angle between the fast horizontal and fast axes
phase_ret: 1D np.array
A difference in phases gained by the slow and the fast components
The most clockwise axis is always considered to be the fast one
scale: 1D np.array
A scaling factor
x0, y0: 1D np.array or float
Defines the location to be plotted at
Returns:
--------
x, y set of points to plot a single Jones matrix of the object
'''
angle = fast_ax_angle
theta = np.linspace(0, 2*np.pi, 20)
# collection of points to plot a fast axis
print('angle', np.rad2deg(angle))
x = x0 + scale * atten_fast * np.real(np.exp(1j * theta))
x_f = x0 + (x - x0) * np.cos(angle)
y_f = y0 + (x - x0) * np.sin(angle)
# collection of point to plot a slow axis
y = y0 + scale * atten_slow * np.real(np.exp(1j * theta))
x_s = x0 - (y - y0) * np.sin(angle)
y_s = y0 + (y - y0) * np.cos(angle)
return x_f, y_f, x_s, y_s
def plot_fast_axis(fast_ax_angle=0, scale=1, x0=4, y0=4):
'''
Plots directions of the fast axes only
'''
xf, yf, xs, ys = plot_attenuations(fast_ax_angle=fast_ax_angle, atten_fast=1, atten_slow=0, scale=scale, x0=x0, y0=y0)
return xf, yf
def plot_figures(shape, num_of_el_along_x=20, num_of_el_along_y=20,
phases=None, fast_ax_angles=None,
atten_fast=None, atten_slow=None, scale=5,
probe=False, attenuations=False, fast_axes=False):
"""
All the parameters - np.arrays of shape (shape)
"""
x_centers = np.linspace(1, shape[0] - 1, num_of_el_along_x)
y_centers = np.linspace(1, shape[1] - 1, num_of_el_along_y)
X, Y = np.meshgrid(x_centers, y_centers)
X, Y = np.ravel(X), np.ravel(Y)
xs, ys = np.array([]), np.array([])
for x, y in zip(X, Y):
k, m = int(x), int(y)
if probe:
xx, yy = plot_probe_ellipse(a=atten_fast[k, m], b=atten_slow[k, m],
phase_ret=phases[k, m], scale=scale, x0=x, y0=y)
title = 'Polarized Probe'
elif attenuations:
xf, yf, xs, ys = plot_attenuations(atten_slow=atten_slow[k, m], atten_fast=atten_fast[k, m],
fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
elif fast_axes:
xx, yy = plot_fast_axis(fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
title = 'Fast Axes'
if attenuations:
plt.plot(xf, yf, c='b')
plt.plot(xs, ys, c='b')
plt.axis('equal')
plt.title('Attenuations')
else:
plt.plot(xx, yy, c='b')
plt.axis('equal')
plt.title(title)
"""
VISUALIZATION FUNCTIONS (and helper functions)
"""
def determine_fast_axis(angle0, angle1):
# in radians
w0, w1 = t.as_tensor(angle0, dtype=t.float32), t.as_tensor(angle1, dtype=t.float32)
diff = (w0 - w1) % (2 * np.pi)
if t.tensor(0, dtype=t.float32) <= diff and diff <= t.tensor(np.pi, dtype=t.float32):
fast = 1
phase_ret = (w0 - w1) % (2 * np.pi)
gl_phase = w1
else:
fast = 0
phase_ret = (w1 - w0) % (2 * np.pi)
gl_phase = w0
if t.allclose((w0 - w1) % (2 * np.pi), phase_ret):
gl_phase = w1
fast = 1
return fast, phase_ret, gl_phase
def retrieve_obj_info(obj):
obj_t = obj.transpose(-1, -3).transpose(-2, -4)
w, v = t.linalg.eig(obj_t)
b = obj.shape[-1]
a = obj.shape[-2]
eigenvectors = t.empty(a, b, 2, dtype=t.cfloat)
ret_phases = t.empty(a, b, dtype=t.float32)
global_phases = t.empty(a, b, dtype=t.float32)
atten_fast = t.empty(a, b, dtype=t.float32)
atten_slow = t.empty(a, b, dtype=t.float32)
fast_ax_angles = t.empty(a, b, dtype=t.float32)
for k, m in zip(*iterator(a, b)):
angle0, angle1 = t.angle(w[k, m, 0]).to(dtype=t.float32), t.angle(w[k, m, 1]).to(dtype=t.float32)
fst, phase_ret, gl_phase = determine_fast_axis(angle0, angle1)
atten0, atten1 = t.abs(w[k, m, 0]), t.abs(w[k, m, 1])
ret_phases[k, m] = phase_ret
def fast(b):
if b == 0:
eigenvectors[k, m, :] = v[k, m, 0, :]
atten_fast[k, m] = atten0
atten_slow[k, m] = atten1
global_phases[k, m] = angle0
elif b == 1:
eigenvectors[k, m, :] = v[k, m, 1, :]
atten_fast[k, m] = atten1
atten_slow[k, m] = atten0
global_phases[k, m] = angle1
if t.allclose(angle0, angle1):
# then it's a linear polarizer, fast ax - the one for which attenuation is bigger
if atten0 > atten1:
fst = 0
else:
fst = 1
fast(fst)
cos = np.abs(eigenvectors[k, m, 0])
sin = np.abs(eigenvectors[k, m, 1])
if t.allclose(cos, t.zeros(1, dtype=t.float32)):
fast_ax_angles[k, m] = 90
else:
fast_ax_angles[k, m] = t.atan(sin/cos)
fast_ax_angles = np.asarray(fast_ax_angles, dtype=np.float32)
ret_phases = np.asarray(ret_phases, dtype=np.float32)
global_phases = np.asarray(global_phases, dtype=np.float32)
atten_fast = np.asarray(atten_fast, dtype=np.float32)
atten_slow = np.asarray(atten_slow, dtype=np.float32)
return fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow
def visualize_components_amplitudes(obj, rot_angle=0, logarithmic=False):
# coord_rot angle is the only angle in degrees here
def coord_rot(angle):
angle = t.as_tensor(angle, dtype=t.float32)
angle = t.deg2rad(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)
for k, m in zip(*iterator(obj.shape[-2], obj.shape[-1])):
obj[:, :, k, m] = t.matmul(coord_rot(rot_angle), obj[:, :, k, m])
components = [obj[i, j, :, :] for i, j in zip(*iterator(2, 2))]
if logarithmic:
components = [np.log(comp)/np.log(10) for comp in components]
titles = ['Amptlitudes of the a components', 'Amplitudes of the b components',
'Amplitudes of the c components', 'Amplitudes of the d components']
for i in range(4):
amplitude = np.abs(components[i])
plt.imshow(module)
plt.colorbar()
plt.title(titles[i])
def visualize_phase_ret(obj, logarithmic=False):
print('DHSBCUYLIWGBCGLWIYV')
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
if logarithmic:
ret_phases = np.log(ret_phases)/np.log(10)
plt.imshow(ret_phases)
plt.colorbar()
plt.show()
def visuallize_global_phases(obj, logarithmic=False):
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
if logarithmic:
global_phases = np.log(global_phases)/np.log(10)
plt.imshow(global_phases)
plt.colorbar()
plt.show()
def visualize_fast_axes(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
A, B = obj.shape[-2], obj.shape[-1]
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
fast_ax_angles=fast_ax_angles, scale=scale, fast_axes=True)
def visualize_attenuations(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
A, B = obj.shape[-2], obj.shape[-1]
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
phases=ret_phases, atten_slow=atten_slow, atten_fast=atten_fast,
fast_ax_angles=fast_ax_angles, scale=scale, attenuations=True)
def visualize_probe(probe, scale=1, num_of_el_along_x=20, num_of_el_along_y=20):
a = np.abs(probe[..., 0, :, :])
b = np.abs(probe[..., 1, :, :])
phases = np.angle(probe[..., 1, :, :]) - np.angle(probe[..., 0, :, :])
A, B = np.asarray(probe.shape[-2]), np.asarray(probe.shape[-1])
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
phases=phases, atten_fast=a, atten_slow=b, scale=scale,
probe=True)
#
# def object_from_components(a, b, c, d):
# ab = t.stack((a, b), dim=-3)
# cd = t.stack((c, d), dim=-3)
# return t.stack((ab, cd), dim=-4)
#
# def object_from_quarters(a, b, c, d):
# ab = t.cat((a, b), dim=-1)
# cd = t.cat((c, d), dim=-1)
# return t.cat((ab, cd), dim=-2)
#
# def generate_birefringent_obj(shape, func_axes=None, func_ret_phases=None, func_global_phases=None, fast_axes=[0, 0, 90, 90], phases=[0, 18, 40, 18]):
# ret = [polarization.generate_birefringent_obj(fast_axis=i, phase_ret=j) for i, j in zip(fast_axes, phases)]
# A = shape[0]
# B = shape[1]
# def to_rad(angle):
# angle = t.as_tensor(angle)
# return t.deg2rad(angle)
#
# if func_axes is None:
# # 4 sets of components for each quarter
# components = [[ret[i][j][k].repeat(A//2, B//2) for j, k in zip(*iterator(2, 2))] for i in range(4)]
# # build the quarters from the components
# quarters = [object_from_components(*comp) for comp in components]
# obj = object_from_quarters(*quarters)
# else:
# X, Y = t.arange(A), t.arange(B)
# X, Y = t.meshgrid(X, Y)
# obj = t.empty(2, 2, A, B, dtype=t.cfloat)
# for i, j in zip(*iterator(A, B)):
# x, y = X[i, j], Y[i, j]
# axis = func_axes(x, y)
# ret = func_ret_phases(x, y)
# gl = func_global_phases(x, y)
# obj[:, :, i, j] = polarization.generate_birefringent_obj(fast_axis=axis, phase_ret=ret, global_phase=gl)
# w, v = t.linalg.eig(obj[:, :, i, j])
#
# # t.tensor of shape (2, 2, A, B)
# return obj
#
# def axes(x, y):
# # return (x-25)**2 + (y-25)**2
# return x * 20
#
# def phases(x, y):
# # return ((x - 5) ** 2 + (y - 5) ** 2) * 20
# return (x) * 10
#
# def glob(x, y):
# return x * 10
#
# def amp(x, y):
# return 1 + 0.05 * (x + y)
#
# def build_probe(shape, phase_ret_func=None, amps_func=None):
# probe = t.empty(2, shape[0], shape[1], dtype=t.cfloat)
# for x, y in zip(*iterator(probe.shape[-2], probe.shape[-1])):
# phase = phase_ret_func(x, y)
# phase = t.deg2rad(phase).to(dtype=t.cfloat)
# abs = amps_func(x, y).to(dtype=t.cfloat)
# probe[0, x, y] = abs
# probe[1, x, y] = abs * t.exp(phase * 1j)
#
# return probe
#
#
# obj = generate_birefringent_obj((10, 10), func_axes=axes, func_ret_phases=phases, func_global_phases=glob)
# probe = build_probe((10, 10), phase_ret_func=phases, amps_func=amp)
# # visualize_probe(probe, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
# visualize_fast_axes(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
# # visualize_phase_ret(obj)
# # visuallize_global_phases(obj)
# # visualize_attenuations(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.3)
# plt.show()
CDTools/tools/polarization/polarization.py
+81 -108
View File
@@ -12,29 +12,32 @@ __all__ = ['apply_linear_polarizer',
'apply_half_wave_plate',
'apply_quarter_wave_plate',
'apply_circular_polarizer',
'apply_jones_matrix']
'apply_jones_matrix',
'generate_linear_polarizer',
'generate_birefringent_obj']
# Abe - split these into two functions
# Note for the future: this function should
def generate_linear_polarizer(pol_angle):
single_angle = False
pol_angle = t.as_tensor(pol_angle)
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)
jones_matrices = t.stack([t.tensor([[(t.cos(p)) ** 2, t.sin(p) * t.cos(p)],
[t.sin(p) * t.cos(p), (t.sin(p)) ** 2]])
for p in polarizer])
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)
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):
"""
@@ -51,13 +54,17 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True
Returns:
--------
linearly polarized probe: t.Tensor
(N)(P)x2x1xMxL
(N)(P)x2x1xMxL
"""
jones_matrices = generate_linear_polarizer(polarization)
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
@@ -66,64 +73,40 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
probe: t.Tensor
A (N)(P)x2xMxL tensor representing the probe
jones_matrix: t.tensor
(N)x2x2x(M)x(L)
(N)x2x2x(M)x(L)
Returns:
--------
a probe with the jones matrix applied: t.Tensor
(N)(P)x2xMxL
(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, :, :]
# 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 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
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:
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)
raise NotImplementedError
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):
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
Shifts the y-component of the field wrt the x-component by a given phase shift
Parameters:
----------
@@ -135,44 +118,41 @@ def apply_phase_retardance(probe, phase_shift):
Returns:
--------
probe: t.Tensor
(...)x2x1xMxL
(...)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, phase_shift]])
probe = probe.transpose(-1, -3).transpose(-2, -4)
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
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)
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
return polarized
def apply_circular_polarizer(probe, left_polarized=True):
def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True):
"""
Applies a circular polarizer to the probe
Parameters:
----------
probe: t.Tensor
A (...)x2x1xMxL tensor representing the probe
A (...)x2xMxL tensor representing the probe
left_polarizd: bool
True for the left-polarization, False for the right
Returns:
--------
circularly polarized probe: t.Tensor
(...)x2x1xMxL
(...)x2xMxL
"""
probe = probe.to(dtype=t.cfloat)
if left_polarized:
jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]]))
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]])
probe = probe.transpose(-1, -3).transpose(-2, -4)
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
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
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
def apply_quarter_wave_plate(probe, fast_axis_angle):
def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
Parameters:
----------
@@ -184,19 +164,17 @@ def apply_quarter_wave_plate(probe, fast_axis_angle):
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
(...)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)
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):
def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
Parameters:
----------
@@ -208,37 +186,32 @@ def apply_half_wave_plate(probe, fast_axis_angle):
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
(...)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)
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)
# probe = t.rand(17, 7, 2, 6, 4)
# polarizer = t.rand(7)
# out = apply_linear_polarizer(probe, polarizer)
# out2 = apply_linear_polarizer(probe, polarizer, transpose=False)
return out
# print(out.shape)
# print(out2.shape)
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
# a = t.ones(17, 8, 2, 3, 4)
# b = t.ones(2, 1, 1)
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)
# probe = t.ones(5, 2, 3, 3)
# polarizer = t.tensor([45])
# exitw = apply_linear_polarizer(probe, polarizer)
# print(exitw[:, 0, :, :])
# print('y', exitw[:, 1, :, :])
#a = t.ones(2, 4)
#print(t.sum(a, dim=1).shape)
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))
CDTools/tools/propagators/propagators.py
+1
View File
@@ -382,6 +382,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
# Bandlimiting is not implemented in the generalized function, because it
# has a less clear meaning in that setting, so we apply it here instead
if bandlimit is not None:
# No need to multiply by 2pi
ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0])
kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1])
Ki, Kj = t.meshgrid(ki,kj)