From 7f1f7cd8f62d53198e312e15f466879f7be16f13 Mon Sep 17 00:00:00 2001
From: Anastasiia Kutakh <anastkut@mit.edu>
Date: Tue, 10 Aug 2021 15:42:12 -0400
Subject: [PATCH] .

---
 temp_tests/simulated_dataset.py | 87 +++++++++++++++++++++++++++++++++
 1 file changed, 87 insertions(+)
 create mode 100644 temp_tests/simulated_dataset.py

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)