cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-09 13:02:41 +02:00
Code Issues Packages Projects Releases Wiki Activity

'.'

Browse Source
This commit is contained in:
Abe Levitan
2021-08-13 14:17:12 -04:00
parent 45f01264c0
commit 998245ae4a
3 changed files with 33 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/polarized_fancy_ptycho.py
+3 -6
View File
@@ -52,7 +52,7 @@ class PolarizedFancyPtycho(FancyPtycho):
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True):
model = FancyPtycho.from_dataset(dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', polarized=True)
model = FancyPtycho.from_dataset(dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um')
# Mutate the class to its subclass
@@ -97,13 +97,10 @@ class PolarizedFancyPtycho(FancyPtycho):
prs = Ws[...,None,None,None,None] * basis_prs
else:
raise NotImplementedError('Unstable Modes not Implemented for polarized light')
polarizer = tools.polarization.generate_linear_polarizer(polarizer)
analyzer = tools.polarization.generate_linear_polarizer(analyzer)
pol_probes = polarization.apply_linear_polarizer(prs, polarizer)
pol_probes = polarization.apply_linear_polarizer(prs, polarizer)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj_support * self.obj,pix_trans,
prs, self.obj, pix_trans,
shift_probe=True, multiple_modes=True, polarized=True)
analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer)
CDTools/tools/interactions/interactions.py
+8 -16
View File
@@ -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)
@@ -433,16 +433,16 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
integer_translations = t.floor(translations)
subpixel_translations = translations - integer_translations
integer_translations = integer_translations.to(dtype=t.int32)
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:
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 +454,11 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
J = 2 * np.pi * J / probe.shape[-1]
phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
-subpixel_translations[:,1,None,None]*J))
if polarized:
phase_masks = phase_masks[..., None, :, :]
# Nx2x1xMxL tensor
# probe is (N)(P)x2xMxL tensor
fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
if multiple_modes: # Multi-mode probe
if polarized:
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:,:]
@@ -470,11 +466,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,8 +476,7 @@ 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')
CDTools/tools/polarization/polarization.py
+22 -15
View File
@@ -22,15 +22,18 @@ __all__ = ['apply_linear_polarizer',
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 pol_angle_rad])
a = t.cos(pol_angle_rad) ** 2
b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad)
c = b
d = t.sin(pol_angle_rad) ** 2
ab = t.stack((a, b), dim=-1)
cd = t.stack((c, d), dim=-1)
jones_matrices = t.stack((ab, cd), dim=-2)
if single_angle:
return jones_matrices[0].to(dtype=t.cfloat)
else:
@@ -54,11 +57,15 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True
linearly polarized probe: t.Tensor
(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
@@ -87,6 +94,7 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
# 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)
@@ -97,7 +105,7 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
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
@@ -115,11 +123,11 @@ def apply_phase_retardance(probe, phase_shift):
"""
probe = probe.to(dtype=t.cfloat)
jones_matrix = t.tensor([[1, 0], [0, phase_shift]]).to(dtype=t.cfloat)
polarized = apply_jones_matrix(probe, jones_matrix)
polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
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
@@ -140,11 +148,10 @@ def apply_circular_polarizer(probe, left_polarized=True):
jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]])).to(dtype=t.cfloat)
else:
jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]]).to(dtype=t.cfloat)
polarized = apply_jones_matrix(probe, jones_matrix)
polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return polarized
def apply_quarter_wave_plate(probe, fast_axis_angle):
def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
Parameters:
----------
@@ -162,11 +169,11 @@ def apply_quarter_wave_plate(probe, fast_axis_angle):
theta = math.radians(fast_axis_angle)
exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]]).to(dtype=t.cfloat)
out = apply_jones_matrix(probe, jones_matrix)
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:
----------
@@ -184,7 +191,7 @@ def apply_half_wave_plate(probe, fast_axis_angle):
theta = math.radians(fast_axis_angle)
exponent = t.exp(-1j * math.pi / 2 * t.ones(2, 2))
jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]]).to(dtype=t.cfloat)
out = apply_jones_matrix(probe, jones_matrix)
out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return out