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

Merge branch 'polarization' of https://github.mit.edu/Scattering/CDTools into polarization

Browse Source
This commit is contained in:
Anastasiia Kutakh
2021-08-27 12:35:29 -04:00
parent d48624bf23 24f213d03c
commit 950fe9bdca
4 changed files with 22 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/polarized_fancy_ptycho.py
+12 -8
View File
@@ -26,18 +26,18 @@ class PolarizedFancyPtycho(FancyPtycho):
polarizer_offsets=None, analyzer_offsets=None,
polarizer_scale=1, analyzer_scale=1, mask=None,
weights = None, translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1,
probe_support = None, oversampling=1,
loss='amplitude mse',units='um'):
super(FancyPtycho, self).__init__(wavelength, detector_geometry,
super(PolarizedFancyPtycho, self).__init__(wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
weights = None, translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1,
weights = weights, translation_scale = 1, saturation=None,
probe_support = None, oversampling=1,
loss='amplitude mse',units='um')
if polarizer_offsets is None:
@@ -49,7 +49,7 @@ class PolarizedFancyPtycho(FancyPtycho):
self.analyzer_offsets = None
else:
self.analyzer_offsets = t.nn.Parameter(t.tensor(analyzer_offsets).to(dtype=t.float32)) / analyzer_scale
self.polarizer = polarizer
self.analyzer = analyzer
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
@@ -63,7 +63,8 @@ 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, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True)
# When using this method, remember to pass through the inputs
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
@@ -81,10 +82,13 @@ class PolarizedFancyPtycho(FancyPtycho):
probe = t.stack([probe, ] + probe_stack)
print('probe', type(probe), probe.shape)
model.probe.data = probe
print(model.probe.shape)
# obj = t.stack((model.obj.data, model.obj.data), dim=-3)
# model.obj.data = t.stack((obj.data, obj.data), dim=-4)
# obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
obj = model.obj.detach()
# Abe - Probably something identity matrix-like would be a better
# initialization (e.g. ((obj,0*obj),(0*obj,obj))
obj = t.stack((obj, obj), dim=-3)
obj = t.stack((obj, obj), dim=-4)
print('object', type(obj), obj.shape)
@@ -275,12 +279,12 @@ class PolarizedFancyPtycho(FancyPtycho):
raise NotImplementedError('Unstable Modes not Implemented for polarized light')
pol_probes = polarization.apply_linear_polarizer(prs, polarizer)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj, pix_trans,
pol_probes, self.obj, pix_trans,
shift_probe=True, multiple_modes=True, polarized=True)
analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer)
return analyzed_exit_waves
CDTools/tools/interactions/interactions.py
+4 -4
View File
@@ -423,9 +423,10 @@ 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
@@ -438,8 +439,8 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
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]]
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
@@ -480,7 +481,6 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
else:
raise NotImplementedError('Object shift not yet implemented')
print('ptyvho 2d sinc', output.shape)
if single_translation:
return output[0]
else:
CDTools/tools/plotting/plotting.py
+5 -1
View File
@@ -17,7 +17,11 @@ from matplotlib import ticker, patheffects
__all__ = ['colorize', 'plot_amplitude', 'plot_phase',
'plot_colorized', 'plot_translations', 'get_units_factor',
'plot_nanomap', 'plot_real', 'plot_imag',
'plot_nanomap_with_images']
'plot_nanomap_with_images',
'polarized_plot_component_amplitudes',
'polarized_plot_phase_ret',
'polarized_plot_global_phases',
'polarized_plot_ellipses']
def colorize(z):
CDTools/tools/polarization/polarization.py
+1
View File
@@ -27,6 +27,7 @@ def generate_linear_polarizer(pol_angle):
if pol_angle.dim() == 0:
pol_angle = t.unsqueeze(pol_angle,0)
single_angle = True
pol_angle_rad = t.deg2rad(pol_angle)
a = t.cos(pol_angle_rad) ** 2
b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad)