diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index dada686..d2483f3 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -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( diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index f6352b9..81d2edf 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -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 + \ No newline at end of file diff --git a/temp_tests/simulated_dataset.py b/temp_tests/simulated_dataset.py new file mode 100644 index 0000000..8972e57 --- /dev/null +++ b/temp_tests/simulated_dataset.py @@ -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)