From c3fef2831492be72caebe0ebdf2e8b14e3daa28f Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 8 Jun 2022 11:11:46 -0700 Subject: [PATCH] Clean some stuff that should have been in the polarization channel, get ready for pytorch's move to xy-style meshgrid --- .../image_processing/image_processing.py | 4 +- CDTools/tools/interactions/interactions.py | 4 +- CDTools/tools/plotting/__init__.py | 2 +- CDTools/tools/plotting/polarized_plotting.py | 342 ------------------ CDTools/tools/propagators/propagators.py | 4 +- conda_requirements.txt | 7 +- setup.py | 4 +- temp_tests/simulated_dataset.py | 132 ------- 8 files changed, 13 insertions(+), 486 deletions(-) delete mode 100644 CDTools/tools/plotting/polarized_plotting.py delete mode 100644 temp_tests/simulated_dataset.py diff --git a/CDTools/tools/image_processing/image_processing.py b/CDTools/tools/image_processing/image_processing.py index 8e57afc..3326dc1 100644 --- a/CDTools/tools/image_processing/image_processing.py +++ b/CDTools/tools/image_processing/image_processing.py @@ -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) diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index 42c2544..2bcc2ee 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/CDTools/tools/interactions/interactions.py @@ -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) diff --git a/CDTools/tools/plotting/__init__.py b/CDTools/tools/plotting/__init__.py index 3ed18c3..0ff7929 100644 --- a/CDTools/tools/plotting/__init__.py +++ b/CDTools/tools/plotting/__init__.py @@ -1,2 +1,2 @@ from CDTools.tools.plotting.plotting import * -from CDTools.tools.plotting.polarized_plotting import * + diff --git a/CDTools/tools/plotting/polarized_plotting.py b/CDTools/tools/plotting/polarized_plotting.py deleted file mode 100644 index e17cc16..0000000 --- a/CDTools/tools/plotting/polarized_plotting.py +++ /dev/null @@ -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() diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py index bca164e..25c6769 100644 --- a/CDTools/tools/propagators/propagators.py +++ b/CDTools/tools/propagators/propagators.py @@ -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, diff --git a/conda_requirements.txt b/conda_requirements.txt index 4726895..d3242fa 100644 --- a/conda_requirements.txt +++ b/conda_requirements.txt @@ -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 diff --git a/setup.py b/setup.py index 10cca68..9c8cd79 100644 --- a/setup.py +++ b/setup.py @@ -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"], diff --git a/temp_tests/simulated_dataset.py b/temp_tests/simulated_dataset.py deleted file mode 100644 index 817996b..0000000 --- a/temp_tests/simulated_dataset.py +++ /dev/null @@ -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()