mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
88 lines
3.8 KiB
Python
88 lines
3.8 KiB
Python
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)
|