mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
Clean some stuff that should have been in the polarization channel, get ready for pytorch's move to xy-style meshgrid
This commit is contained in:
Abe Levitan
@@ -39,7 +39,7 @@ def centroid(im, dims=2):
|
||||
"""
|
||||
# For some reason this needs to be a list
|
||||
indices = [t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims)]
|
||||
indices = t.meshgrid(*indices)
|
||||
indices = t.meshgrid(*indices, indexing='ij')
|
||||
|
||||
use_dims = [-dims+i for i in range(dims)]
|
||||
divisor = t.sum(im, dim=use_dims)
|
||||
@@ -102,7 +102,7 @@ def sinc_subpixel_shift(im, shift):
|
||||
|
||||
i = t.arange(im.shape[0]) - im.shape[0]//2
|
||||
j = t.arange(im.shape[1]) - im.shape[1]//2
|
||||
I,J = t.meshgrid(i,j)
|
||||
I,J = t.meshgrid(i,j, indexing='ij')
|
||||
I = 2 * np.pi * I.to(t.float32) / im.shape[0]
|
||||
J = 2 * np.pi * J.to(t.float32) / im.shape[1]
|
||||
I = I.to(dtype=im.dtype,device=im.device)
|
||||
|
||||
@@ -450,7 +450,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
|
||||
- probe.shape[-2]//2
|
||||
j = t.arange(probe.shape[-1],device=probe.device,dtype=t.float32) \
|
||||
- probe.shape[-1]//2
|
||||
I,J = t.meshgrid(i,j)
|
||||
I,J = t.meshgrid(i,j, indexing='ij')
|
||||
I = 2 * np.pi * I / probe.shape[-2]
|
||||
J = 2 * np.pi * J / probe.shape[-1]
|
||||
phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
|
||||
@@ -543,7 +543,7 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
|
||||
if shift_probe:
|
||||
i = t.arange(probe.shape[-2]) - probe.shape[-2]//2
|
||||
j = t.arange(probe.shape[-1]) - probe.shape[-1]//2
|
||||
I,J = t.meshgrid(i,j)
|
||||
I,J = t.meshgrid(i,j, indexing='ij')
|
||||
I = 2 * np.pi * I.to(t.float32) / probe.shape[-2]
|
||||
J = 2 * np.pi * J.to(t.float32) / probe.shape[-1]
|
||||
I = I.to(dtype=probe.dtype,device=probe.device)
|
||||
|
||||
@@ -1,2 +1,2 @@
|
||||
from CDTools.tools.plotting.plotting import *
|
||||
from CDTools.tools.plotting.polarized_plotting import *
|
||||
|
||||
|
||||
@@ -1,342 +0,0 @@
|
||||
import torch as t
|
||||
from matplotlib import pyplot as plt
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from matplotlib.colors import hsv_to_rgb
|
||||
from matplotlib.widgets import Slider
|
||||
from matplotlib import ticker, patheffects
|
||||
from CDTools.tools.polarization import polarization
|
||||
|
||||
all = ['visualize_components_amplitudes',
|
||||
'visualize_attenuations',
|
||||
'visualize_fast_axes',
|
||||
'visualize_global_axes',
|
||||
'visualize_phase_ret',
|
||||
'visualize_probe']
|
||||
|
||||
def iterator(a, b):
|
||||
'''
|
||||
A helper function we can use to facilitate to process of iterating over 2D arrays
|
||||
'''
|
||||
x, y = t.arange(a), t.arange(b)
|
||||
x, y = t.meshgrid(x, y)
|
||||
x, y = t.ravel(x), t.ravel(y)
|
||||
return (x, y)
|
||||
|
||||
def plot_probe_ellipse(a=1, b=1, phase_ret=0, scale=1, x0=4, y0=5):
|
||||
'''
|
||||
Given a probe vector at a point (x0, y0), visualizes ellipticity
|
||||
of its polarization
|
||||
|
||||
Parameters:
|
||||
-----------
|
||||
a : 1D np.array
|
||||
An amplitude of the horizontal component of the probe vector
|
||||
b : 1D np.array
|
||||
An amplitude of the vertical component of the probe vector
|
||||
phase_ret : 1D np.array
|
||||
A phase difference in radians bettween the phases of the y and x components
|
||||
scale : int
|
||||
A scaling factor
|
||||
x0, y0: 1D np.array or float
|
||||
Defines the location of the vector to be plotted
|
||||
|
||||
Returns:
|
||||
--------
|
||||
x, y set of points to plot a single ellipse
|
||||
'''
|
||||
theta = np.linspace(0, 2*np.pi, 20)
|
||||
x = x0 + scale * a * np.real(np.exp(1j * theta))
|
||||
y = y0 + scale * b * np.real(np.exp(1j * (theta + phase_ret)))
|
||||
return x, y
|
||||
|
||||
def plot_attenuations(atten_slow=1, atten_fast=1, fast_ax_angle=0, scale=1, x0=4, y0=4):
|
||||
'''
|
||||
Plots attenuations along the fast and slow axes
|
||||
|
||||
Parameters:
|
||||
-----------
|
||||
atten_slow: 1D np.array
|
||||
Attenuation along the slow axis
|
||||
atten_fast: 1D np.array
|
||||
Attenuation along the fast axis
|
||||
fast_ax_angle: 1D np.array
|
||||
An angle between the fast horizontal and fast axes
|
||||
phase_ret: 1D np.array
|
||||
A difference in phases gained by the slow and the fast components
|
||||
The most clockwise axis is always considered to be the fast one
|
||||
scale: 1D np.array
|
||||
A scaling factor
|
||||
x0, y0: 1D np.array or float
|
||||
Defines the location to be plotted at
|
||||
|
||||
Returns:
|
||||
--------
|
||||
x, y set of points to plot a single Jones matrix of the object
|
||||
'''
|
||||
angle = fast_ax_angle
|
||||
theta = np.linspace(0, 2*np.pi, 20)
|
||||
# collection of points to plot a fast axis
|
||||
print('angle', np.rad2deg(angle))
|
||||
x = x0 + scale * atten_fast * np.real(np.exp(1j * theta))
|
||||
x_f = x0 + (x - x0) * np.cos(angle)
|
||||
y_f = y0 + (x - x0) * np.sin(angle)
|
||||
# collection of point to plot a slow axis
|
||||
y = y0 + scale * atten_slow * np.real(np.exp(1j * theta))
|
||||
x_s = x0 - (y - y0) * np.sin(angle)
|
||||
y_s = y0 + (y - y0) * np.cos(angle)
|
||||
return x_f, y_f, x_s, y_s
|
||||
|
||||
def plot_fast_axis(fast_ax_angle=0, scale=1, x0=4, y0=4):
|
||||
'''
|
||||
Plots directions of the fast axes only
|
||||
'''
|
||||
xf, yf, xs, ys = plot_attenuations(fast_ax_angle=fast_ax_angle, atten_fast=1, atten_slow=0, scale=scale, x0=x0, y0=y0)
|
||||
return xf, yf
|
||||
|
||||
def plot_figures(shape, num_of_el_along_x=20, num_of_el_along_y=20,
|
||||
phases=None, fast_ax_angles=None,
|
||||
atten_fast=None, atten_slow=None, scale=5,
|
||||
probe=False, attenuations=False, fast_axes=False):
|
||||
"""
|
||||
All the parameters - np.arrays of shape (shape)
|
||||
"""
|
||||
x_centers = np.linspace(1, shape[0] - 1, num_of_el_along_x)
|
||||
y_centers = np.linspace(1, shape[1] - 1, num_of_el_along_y)
|
||||
X, Y = np.meshgrid(x_centers, y_centers)
|
||||
X, Y = np.ravel(X), np.ravel(Y)
|
||||
xs, ys = np.array([]), np.array([])
|
||||
for x, y in zip(X, Y):
|
||||
k, m = int(x), int(y)
|
||||
if probe:
|
||||
xx, yy = plot_probe_ellipse(a=atten_fast[k, m], b=atten_slow[k, m],
|
||||
phase_ret=phases[k, m], scale=scale, x0=x, y0=y)
|
||||
title = 'Polarized Probe'
|
||||
elif attenuations:
|
||||
xf, yf, xs, ys = plot_attenuations(atten_slow=atten_slow[k, m], atten_fast=atten_fast[k, m],
|
||||
fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
|
||||
elif fast_axes:
|
||||
xx, yy = plot_fast_axis(fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
|
||||
title = 'Fast Axes'
|
||||
|
||||
if attenuations:
|
||||
plt.plot(xf, yf, c='b')
|
||||
plt.plot(xs, ys, c='b')
|
||||
plt.axis('equal')
|
||||
plt.title('Attenuations')
|
||||
else:
|
||||
plt.plot(xx, yy, c='b')
|
||||
plt.axis('equal')
|
||||
plt.title(title)
|
||||
|
||||
"""
|
||||
VISUALIZATION FUNCTIONS (and helper functions)
|
||||
"""
|
||||
def determine_fast_axis(angle0, angle1):
|
||||
# in radians
|
||||
w0, w1 = t.as_tensor(angle0, dtype=t.float32), t.as_tensor(angle1, dtype=t.float32)
|
||||
diff = (w0 - w1) % (2 * np.pi)
|
||||
if t.tensor(0, dtype=t.float32) <= diff and diff <= t.tensor(np.pi, dtype=t.float32):
|
||||
fast = 1
|
||||
phase_ret = (w0 - w1) % (2 * np.pi)
|
||||
gl_phase = w1
|
||||
else:
|
||||
fast = 0
|
||||
phase_ret = (w1 - w0) % (2 * np.pi)
|
||||
gl_phase = w0
|
||||
|
||||
if t.allclose((w0 - w1) % (2 * np.pi), phase_ret):
|
||||
gl_phase = w1
|
||||
fast = 1
|
||||
|
||||
return fast, phase_ret, gl_phase
|
||||
|
||||
def retrieve_obj_info(obj):
|
||||
obj_t = obj.transpose(-1, -3).transpose(-2, -4)
|
||||
w, v = t.linalg.eig(obj_t)
|
||||
b = obj.shape[-1]
|
||||
a = obj.shape[-2]
|
||||
eigenvectors = t.empty(a, b, 2, dtype=t.cfloat)
|
||||
ret_phases = t.empty(a, b, dtype=t.float32)
|
||||
global_phases = t.empty(a, b, dtype=t.float32)
|
||||
atten_fast = t.empty(a, b, dtype=t.float32)
|
||||
atten_slow = t.empty(a, b, dtype=t.float32)
|
||||
fast_ax_angles = t.empty(a, b, dtype=t.float32)
|
||||
|
||||
for k, m in zip(*iterator(a, b)):
|
||||
angle0, angle1 = t.angle(w[k, m, 0]).to(dtype=t.float32), t.angle(w[k, m, 1]).to(dtype=t.float32)
|
||||
fst, phase_ret, gl_phase = determine_fast_axis(angle0, angle1)
|
||||
atten0, atten1 = t.abs(w[k, m, 0]), t.abs(w[k, m, 1])
|
||||
ret_phases[k, m] = phase_ret
|
||||
def fast(b):
|
||||
if b == 0:
|
||||
eigenvectors[k, m, :] = v[k, m, 0, :]
|
||||
atten_fast[k, m] = atten0
|
||||
atten_slow[k, m] = atten1
|
||||
global_phases[k, m] = angle0
|
||||
elif b == 1:
|
||||
eigenvectors[k, m, :] = v[k, m, 1, :]
|
||||
atten_fast[k, m] = atten1
|
||||
atten_slow[k, m] = atten0
|
||||
global_phases[k, m] = angle1
|
||||
|
||||
if t.allclose(angle0, angle1):
|
||||
# then it's a linear polarizer, fast ax - the one for which attenuation is bigger
|
||||
if atten0 > atten1:
|
||||
fst = 0
|
||||
else:
|
||||
fst = 1
|
||||
fast(fst)
|
||||
|
||||
cos = np.abs(eigenvectors[k, m, 0])
|
||||
sin = np.abs(eigenvectors[k, m, 1])
|
||||
if t.allclose(cos, t.zeros(1, dtype=t.float32)):
|
||||
fast_ax_angles[k, m] = 90
|
||||
else:
|
||||
fast_ax_angles[k, m] = t.atan(sin/cos)
|
||||
|
||||
fast_ax_angles = np.asarray(fast_ax_angles, dtype=np.float32)
|
||||
ret_phases = np.asarray(ret_phases, dtype=np.float32)
|
||||
global_phases = np.asarray(global_phases, dtype=np.float32)
|
||||
atten_fast = np.asarray(atten_fast, dtype=np.float32)
|
||||
atten_slow = np.asarray(atten_slow, dtype=np.float32)
|
||||
|
||||
return fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow
|
||||
|
||||
|
||||
def visualize_components_amplitudes(obj, rot_angle=0, logarithmic=False):
|
||||
# coord_rot angle is the only angle in degrees here
|
||||
def coord_rot(angle):
|
||||
angle = t.as_tensor(angle, dtype=t.float32)
|
||||
angle = t.deg2rad(angle)
|
||||
a = t.stack((t.cos(angle), t.sin(angle)), dim=-1)
|
||||
b = t.stack((-t.sin(angle), t.cos(angle)), dim=-1)
|
||||
return t.stack((a, b), dim=-2).to(dtype=t.cfloat)
|
||||
for k, m in zip(*iterator(obj.shape[-2], obj.shape[-1])):
|
||||
obj[:, :, k, m] = t.matmul(coord_rot(rot_angle), obj[:, :, k, m])
|
||||
components = [obj[i, j, :, :] for i, j in zip(*iterator(2, 2))]
|
||||
if logarithmic:
|
||||
components = [np.log(comp)/np.log(10) for comp in components]
|
||||
titles = ['Amptlitudes of the a components', 'Amplitudes of the b components',
|
||||
'Amplitudes of the c components', 'Amplitudes of the d components']
|
||||
|
||||
for i in range(4):
|
||||
amplitude = np.abs(components[i])
|
||||
plt.imshow(module)
|
||||
plt.colorbar()
|
||||
plt.title(titles[i])
|
||||
|
||||
def visualize_phase_ret(obj, logarithmic=False):
|
||||
print('DHSBCUYLIWGBCGLWIYV')
|
||||
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
|
||||
if logarithmic:
|
||||
ret_phases = np.log(ret_phases)/np.log(10)
|
||||
plt.imshow(ret_phases)
|
||||
plt.colorbar()
|
||||
plt.show()
|
||||
|
||||
def visuallize_global_phases(obj, logarithmic=False):
|
||||
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
|
||||
if logarithmic:
|
||||
global_phases = np.log(global_phases)/np.log(10)
|
||||
plt.imshow(global_phases)
|
||||
plt.colorbar()
|
||||
plt.show()
|
||||
|
||||
def visualize_fast_axes(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
|
||||
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
|
||||
A, B = obj.shape[-2], obj.shape[-1]
|
||||
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
|
||||
fast_ax_angles=fast_ax_angles, scale=scale, fast_axes=True)
|
||||
|
||||
def visualize_attenuations(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
|
||||
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
|
||||
A, B = obj.shape[-2], obj.shape[-1]
|
||||
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
|
||||
phases=ret_phases, atten_slow=atten_slow, atten_fast=atten_fast,
|
||||
fast_ax_angles=fast_ax_angles, scale=scale, attenuations=True)
|
||||
|
||||
def visualize_probe(probe, scale=1, num_of_el_along_x=20, num_of_el_along_y=20):
|
||||
a = np.abs(probe[..., 0, :, :])
|
||||
b = np.abs(probe[..., 1, :, :])
|
||||
phases = np.angle(probe[..., 1, :, :]) - np.angle(probe[..., 0, :, :])
|
||||
A, B = np.asarray(probe.shape[-2]), np.asarray(probe.shape[-1])
|
||||
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
|
||||
phases=phases, atten_fast=a, atten_slow=b, scale=scale,
|
||||
probe=True)
|
||||
|
||||
|
||||
#
|
||||
# def object_from_components(a, b, c, d):
|
||||
# ab = t.stack((a, b), dim=-3)
|
||||
# cd = t.stack((c, d), dim=-3)
|
||||
# return t.stack((ab, cd), dim=-4)
|
||||
#
|
||||
# def object_from_quarters(a, b, c, d):
|
||||
# ab = t.cat((a, b), dim=-1)
|
||||
# cd = t.cat((c, d), dim=-1)
|
||||
# return t.cat((ab, cd), dim=-2)
|
||||
#
|
||||
# def generate_birefringent_obj(shape, func_axes=None, func_ret_phases=None, func_global_phases=None, fast_axes=[0, 0, 90, 90], phases=[0, 18, 40, 18]):
|
||||
# ret = [polarization.generate_birefringent_obj(fast_axis=i, phase_ret=j) for i, j in zip(fast_axes, phases)]
|
||||
# A = shape[0]
|
||||
# B = shape[1]
|
||||
# def to_rad(angle):
|
||||
# angle = t.as_tensor(angle)
|
||||
# return t.deg2rad(angle)
|
||||
#
|
||||
# if func_axes is None:
|
||||
# # 4 sets of components for each quarter
|
||||
# components = [[ret[i][j][k].repeat(A//2, B//2) for j, k in zip(*iterator(2, 2))] for i in range(4)]
|
||||
# # build the quarters from the components
|
||||
# quarters = [object_from_components(*comp) for comp in components]
|
||||
# obj = object_from_quarters(*quarters)
|
||||
# else:
|
||||
# X, Y = t.arange(A), t.arange(B)
|
||||
# X, Y = t.meshgrid(X, Y)
|
||||
# obj = t.empty(2, 2, A, B, dtype=t.cfloat)
|
||||
# for i, j in zip(*iterator(A, B)):
|
||||
# x, y = X[i, j], Y[i, j]
|
||||
# axis = func_axes(x, y)
|
||||
# ret = func_ret_phases(x, y)
|
||||
# gl = func_global_phases(x, y)
|
||||
# obj[:, :, i, j] = polarization.generate_birefringent_obj(fast_axis=axis, phase_ret=ret, global_phase=gl)
|
||||
# w, v = t.linalg.eig(obj[:, :, i, j])
|
||||
#
|
||||
# # t.tensor of shape (2, 2, A, B)
|
||||
# return obj
|
||||
#
|
||||
# def axes(x, y):
|
||||
# # return (x-25)**2 + (y-25)**2
|
||||
# return x * 20
|
||||
#
|
||||
# def phases(x, y):
|
||||
# # return ((x - 5) ** 2 + (y - 5) ** 2) * 20
|
||||
# return (x) * 10
|
||||
#
|
||||
# def glob(x, y):
|
||||
# return x * 10
|
||||
#
|
||||
# def amp(x, y):
|
||||
# return 1 + 0.05 * (x + y)
|
||||
#
|
||||
# def build_probe(shape, phase_ret_func=None, amps_func=None):
|
||||
# probe = t.empty(2, shape[0], shape[1], dtype=t.cfloat)
|
||||
# for x, y in zip(*iterator(probe.shape[-2], probe.shape[-1])):
|
||||
# phase = phase_ret_func(x, y)
|
||||
# phase = t.deg2rad(phase).to(dtype=t.cfloat)
|
||||
# abs = amps_func(x, y).to(dtype=t.cfloat)
|
||||
# probe[0, x, y] = abs
|
||||
# probe[1, x, y] = abs * t.exp(phase * 1j)
|
||||
#
|
||||
# return probe
|
||||
#
|
||||
#
|
||||
# obj = generate_birefringent_obj((10, 10), func_axes=axes, func_ret_phases=phases, func_global_phases=glob)
|
||||
# probe = build_probe((10, 10), phase_ret_func=phases, amps_func=amp)
|
||||
# # visualize_probe(probe, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
|
||||
# visualize_fast_axes(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
|
||||
# # visualize_phase_ret(obj)
|
||||
# # visuallize_global_phases(obj)
|
||||
# # visualize_attenuations(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.3)
|
||||
# plt.show()
|
||||
@@ -385,7 +385,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
|
||||
# No need to multiply by 2pi
|
||||
ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0])
|
||||
kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1])
|
||||
Ki, Kj = t.meshgrid(ki,kj)
|
||||
Ki, Kj = t.meshgrid(ki,kj, indexing='ij')
|
||||
min_radius = min(t.max(ki),t.max(kj))
|
||||
Rs = t.sqrt((Ki/t.max(ki))**2 + (Kj/t.max(kj))**2)
|
||||
propagator = propagator * (Rs < bandlimit)
|
||||
@@ -492,7 +492,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
|
||||
# Then we calculate the frequencies in (i,j) space
|
||||
ki = 2 * np.pi * t.fft.fftfreq(shape[0], dtype=inv_basis.dtype)
|
||||
kj = 2 * np.pi * t.fft.fftfreq(shape[1], dtype=inv_basis.dtype)
|
||||
K_ij = t.stack(t.meshgrid(ki,kj))
|
||||
K_ij = t.stack(t.meshgrid(ki,kj, indexing='ij'))
|
||||
|
||||
# Now we convert these to frequencies in reciprocal space
|
||||
# These frequencies span the 2D plane of the input wavefield,
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
numpy>=1.0
|
||||
scipy>=1.0
|
||||
matplotlib>=2.0
|
||||
python-dateutil
|
||||
pytorch>=1.9.0
|
||||
matplotlib>=2.0 # 2.0 introduces better colormaps which are used by default
|
||||
pytorch>=1.9.0 #1.9.0 implements support for autograd on indexed complex tensors
|
||||
h5py>=2.1
|
||||
python-dateutil
|
||||
pytest
|
||||
sphinx
|
||||
sphinx-argparse
|
||||
sphinx_rtd_theme
|
||||
|
||||
@@ -16,9 +16,9 @@ setuptools.setup(
|
||||
install_requires=[
|
||||
"numpy>=1.0",
|
||||
"scipy>=1.0",
|
||||
"matplotlib>=2.0", # Matplotlib 2.0 introduces better colormaps and no I'm not sorry
|
||||
"matplotlib>=2.0", # 2.0 introduces better colormaps which are used by default
|
||||
"python-dateutil",
|
||||
"torch>=1.9.0", #1.9.0 implements support for autograd on indexed complex tensors, which we need in order to use complex tensors in the forward models
|
||||
"torch>=1.9.0", #1.9.0 implements support for autograd on indexed complex tensors
|
||||
"h5py>=2.1"],
|
||||
extras_require={
|
||||
'tests': ["pytest"],
|
||||
|
||||
@@ -1,132 +0,0 @@
|
||||
import numpy as np
|
||||
import torch as t
|
||||
from scipy import misc
|
||||
from CDTools.models import PolarizedFancyPtycho
|
||||
from CDTools.datasets import PolarizedPtycho2DDataset
|
||||
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
|
||||
|
||||
f = misc.ascent()
|
||||
x , y = np.shape(f)
|
||||
aa = f[:x//2, :y//2]
|
||||
bb = f[:x//2, -y//2:]
|
||||
cc = f[-x//2:, :y//2]
|
||||
dd = f[-x//2:, -y//2:]
|
||||
|
||||
print(1)
|
||||
|
||||
def simulate_polarized_dataset(probe_size, obj_size, num_patt, a, b, c, d):
|
||||
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]
|
||||
xs, ys = np.ravel(xs), np.ravel(ys)
|
||||
for x, y in zip(xs, ys):
|
||||
for j in range(9):
|
||||
translations.append((5*x, 5*y))
|
||||
translations = t.as_tensor(translations, dtype=t.float32)
|
||||
obj = t.stack((t.stack((a, b), dim=-3), t.stack((c, d), dim=-3)), dim=-4)
|
||||
|
||||
# for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1]):
|
||||
# plt.imshow(np.real(obj[i, j, ...]))
|
||||
# plt.show()
|
||||
probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), np.array((2, 2)))
|
||||
probe = t.stack((probe, probe), dim=-3)
|
||||
probe = polarization.apply_circular_polarizer(probe, multiple_modes=False)
|
||||
polarizers = polarization.generate_linear_polarizer(t.tensor([0, 45, 90]))
|
||||
num_transl = len(translations)
|
||||
pol_probes = [polarization.apply_jones_matrix(probe, polarizers[i], multiple_modes=False) for i in range(3)]
|
||||
analyzers = t.stack(([polarizers[i % 3] for i in range(num_transl)]), dim=0)
|
||||
analyzer = t.tensor([(i % 3) * 45 for i in range(num_transl)])
|
||||
polarizer = t.tensor([(i // 3) % 3 * 45 for i in range(num_transl)])
|
||||
polarized_probes = t.stack(([pol_probes[(i // 3) % 3] for i in range(num_transl)]), dim=0)
|
||||
polarized_wavefields = t.stack([tools.interactions.ptycho_2D_sinc(polarized_probes[i], obj, translations[i], polarized=True, multiple_modes=False) for i in range(num_transl)])
|
||||
wavefields = tools.propagators.far_field(polarized_wavefields)
|
||||
wavefields = polarization.apply_jones_matrix(wavefields, analyzers, multiple_modes=False)
|
||||
patterns = t.abs(wavefields[:, 0, :, :])**2 + t.abs(wavefields[:, 1, :, :])**2
|
||||
detector_basis = t.transpose(t.tensor([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]), 0, 1)
|
||||
det_shape = t.Size((obj_size, obj_size))
|
||||
wavelength = 532e-9
|
||||
real_basis = tools.initializers.exit_wave_geometry(detector_basis, det_shape, wavelength, 2.5e-2)[0]
|
||||
# print('f', type(real_basis), translations.shape)
|
||||
real_translations = tools.interactions.pixel_to_translations(real_basis, translations)
|
||||
# print(type(real_translations))
|
||||
patterns = patterns.numpy()
|
||||
detector_geometry = {
|
||||
'corner': np.array([probe_size*4.8e-6/2, probe_size*4.8e-6/2, 2.5e-2]),
|
||||
'basis': np.array([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]).transpose(),
|
||||
'distance': 2.5e-2
|
||||
}
|
||||
|
||||
return PolarizedPtycho2DDataset(real_translations, polarizer, analyzer, patterns,
|
||||
axes=("x", "y"), detector_geometry=detector_geometry, wavelength=wavelength)
|
||||
|
||||
# print(1)
|
||||
dataset = simulate_polarized_dataset(100, 500, 10, aa, bb, cc, dd)
|
||||
# print(2)
|
||||
dataset.inspect()
|
||||
# print(3)
|
||||
plt.show()
|
||||
model = PolarizedFancyPtycho.from_dataset(dataset, propagation_distance=1e-3)
|
||||
model.inspect()
|
||||
# plt.show()
|
||||
model.compare(dataset)
|
||||
plt.show()
|
||||
|
||||
# for loss in model.Adam_optimize(400, dataset, batch_size=5, lr=0.002, schedule=True):
|
||||
# # And we liveplot the updates to the model as they happen
|
||||
# print(model.report())
|
||||
# model.inspect(dataset)
|
||||
# model.save_figures(prefix='simulated', extension='png')
|
||||
# res = model.save_results(dataset)
|
||||
# np.save('simulated_dataset.npy', res)
|
||||
|
||||
dataset = simulate_polarized_dataset(50, 100, 10, aa, bb, cc, dd)
|
||||
|
||||
res = np.load('simulated_dataset.npy', allow_pickle=True)
|
||||
res = res[()]
|
||||
print(type(res))
|
||||
Ws = t.ones(len(dataset))
|
||||
ewg = CDTools.tools.initializers.exit_wave_geometry
|
||||
probe_basis, probe_shape, det_slice = ewg(res['basis'],
|
||||
dataset[0][1].shape,
|
||||
dataset.wavelength,
|
||||
dataset.detector_geometry['distance'],
|
||||
center=None,
|
||||
padding=0)
|
||||
print('object', res['obj'].shape)
|
||||
obj = res['obj'][..., 200:275, 200:275]
|
||||
# print('obj', obj)
|
||||
a = obj[0, 0, :, :]
|
||||
b = obj[0, 1, :, :]
|
||||
c = obj[1, 0, :, :]
|
||||
d = obj[1, 1, :, :]
|
||||
for i in [a, b, c, d]:
|
||||
plt.imshow(np.real(i))
|
||||
# plt.imshow(np.log(np))
|
||||
plt.colorbar()
|
||||
plt.show()
|
||||
print(np.allclose(np.real(a), np.real(b)))
|
||||
print('a', a[..., 10:20, 10:20])
|
||||
print('b', b[..., 10:20, 10:20])
|
||||
models = [CDTools.models.FancyPtycho(dataset.wavelength, dataset.detector_geometry, probe_basis,
|
||||
res['probe'], component, surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
|
||||
min_translation=t.tensor([0, 0], dtype=t.float32),
|
||||
background=t.tensor(res['background']), translation_offsets=None, mask=None,
|
||||
weights=Ws, translation_scale=1, saturation=None,
|
||||
probe_support=None, oversampling=1,
|
||||
loss='amplitude mse', units='um') for component in [a, b, c, d]]
|
||||
# print('model is created')
|
||||
# for model in models:
|
||||
# model.inspect()
|
||||
# plt.show()
|
||||
Reference in New Issue
Block a user