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 github.mit.edu:Scattering/CDTools into polarization

Browse Source
This commit is contained in:
Abe Levitan
2021-08-11 14:23:41 -04:00
parent 096ec4ff20 7f1f7cd8f6
commit 45f01264c0
3 changed files with 106 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/polarized_fancy_ptycho.py
+3
View File
@@ -69,6 +69,7 @@ class PolarizedFancyPtycho(FancyPtycho):
# tensor vs tensor.data
return model
polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)]
# WHAT IS INDEX?
def interaction(self, index, translations, polarizer, analyzer, test=False):
@@ -97,6 +98,8 @@ class PolarizedFancyPtycho(FancyPtycho):
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)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
CDTools/tools/polarization/polarization.py
+16 -46
View File
@@ -114,12 +114,10 @@ def apply_phase_retardance(probe, phase_shift):
(...)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)
jones_matrix = t.tensor([[1, 0], [0, phase_shift]]).to(dtype=t.cfloat)
polarized = apply_jones_matrix(probe, jones_matrix)
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
return polarized
def apply_circular_polarizer(probe, left_polarized=True):
"""
@@ -128,25 +126,23 @@ def apply_circular_polarizer(probe, left_polarized=True):
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)
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
return polarized
def apply_quarter_wave_plate(probe, fast_axis_angle):
"""
@@ -165,12 +161,10 @@ def apply_quarter_wave_plate(probe, fast_axis_angle):
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)
return out
def apply_half_wave_plate(probe, fast_axis_angle):
"""
@@ -189,32 +183,8 @@ def apply_half_wave_plate(probe, fast_axis_angle):
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)
# 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)
# print(out.shape)
# print(out2.shape)
# a = t.ones(17, 8, 2, 3, 4)
# b = t.ones(2, 1, 1)
# 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)
return out
temp_tests/simulated_dataset.py
+87
View File
@@ -0,0 +1,87 @@
import numpy as np
import torch as t
from CDTools.models import PolarizedFancyPtycho
#from CDTools.datasets import Polarized2DDataset
import CDTools
from CDTools.tools import polarization
from CDTools import tools
from matplotlib import pyplot as plt
from PIL import Image
# upolad 4 different images representing 4 components of the object
# and 2 gaaussian functionas corresponding to the probe components
a = np.asarray(Image.open('a.jpg'))
b = np.asarray(Image.open('b.jpg'))
c = np.asarray(Image.open('c.jpg'))
d = np.asarray(Image.open('d.jpg'))
#a = np.dot(a[..., :3], [.3, 6., .1])
def simulate_dataset(probe_size, obj_size, num_patt):
translations = []
xs, ys = np.mgrid[:num_patt, :num_patt]
for x, y in zip(xs, ys):
translations.append((x*10e-3, y*10e-3))
translations = t.as_tensor(translations, dtype=t.float32)
a = t.as_tensor(a, dtype=t.cfloat)
a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), 50)
wavefields = tools.interactions.ptycho_2D_sinc(probe, obj, translations)
patterns = tools.propagators.far_field(wavefront)
patterns = np(patterns)
translations = np(t.cat((translations, t.zeros(num_patt)), dim=-1))
# needs to be stored as a cxi file
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi('simulated_dataset.cxi')
dataset.detector_geometry = None
def simulate polarized_datset(probe_size, obj_size, num_patt):
a, b, c, d = t.as_tensor(a, dtype=t.cfloat), t.as_tensor(b, dtype=t.cfloat), t.as_tensor(c, dtype=t.cfloat), t.as_tensor(d, dtype=t.cfloat)
a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
b = t.tensordot(b, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
c = t.tensordot(c, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
d = t.tensordot(d, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
translations = []
xs, ys = np.mgrid[:num_patt, :num_patt]
for x, y in zip(xs, ys):
translations.append((x*10e-3, y*10e-3))
translations = t.as_tensor(translations, dtype=t.float32)
obj = t.stack((t.stack((a, c), dim=0), t.stack((b, d), dim=0)), dim=-3)
probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), 50)
probe = t.stack((probe, probe), dim=-3)
probe = polarization.apply_circular_polarizer(probe)
selections = tools.interactions.ptycho_2D_sinc(t.ones(2, probe_size, probe_size).to(dtype=t.cfloat), obj, translations, polarized=True)
polarizers = [polarization.generate_linear_polarizer(i * 45) for i in range(3)]
pol_probes = [polarization.apply_jones_matrix(probe, polarizers[i]) for i in range(3)]
analyzer = t.stack(([polaryzers[i % 3] for i in range(num_patt)]), dim=0)
# probes = t.stack(([probes[i // 3] for i in num_patt]), dim=0)
wavefields = t.as_tensor([tools.interactions.ptycho_2D_sinc(pol_probes[i], obj, translations, polarized=True) for i in range(3)]).to(dtype=t.cfloat)
wf = t.empty(1, 2, obj_size, obj_size)
for i in range(num_patt):
for j in range(3):
pol_channel = t.stack(([wavefileds[j] for k in range(3)]), dim=0)
wf = t.cat((wf, pol_channel), dim=0)
pol_wavefieds = polarization.apply_jones_matrix(wf, analyzer)
patterns = tools.propagators.far_field(pol_wavefieds)
translations = np(t.cat((translations, t.zeros(num_patt)), dim=-1))
patterns = np(patterns)
dataset.detector_geometry = None
# needs to be stored in a cxi file
dataset = CDTools.datasets.FancyPtycho2DDataset.from_cxi('polarized_simulated_dataset.cxi')
dataset.inspect()
model = tools.models.PolarizedFancyPtycho.from_dataset(dataset)