cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-15 15:42:37 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merging a whole buncha work back into master

Browse Source
This commit is contained in:
Abe Levitan
2022-04-18 15:51:45 -07:00
parent 1b5ef2271c 20a1479859
commit f3cee949e0
22 changed files with 3382 additions and 771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/tools/plotting/__init__.py
+1
View File
@@ -1 +1,2 @@
from CDTools.tools.plotting.plotting import *
from CDTools.tools.plotting.polarized_plotting import *
CDTools/tools/plotting/plotting.py
+54 -43
View File
@@ -84,7 +84,7 @@ def get_units_factor(units):
def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, interpolation=None, **kwargs):
"""Plots an image with a colorbar and on an appropriate spatial grid
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
@@ -94,7 +94,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
Finally, if a function is passed to the plot_func argument, this function
will be called on each slice of data before it is plotted. This is used
internally to enable the plot_real, plot_image, plot_phase, etc. functions.
Parameters
----------
@@ -122,7 +122,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
# convert to numpy
if isinstance(im, t.Tensor):
# If final dimension is 2, assume it is a complex array. If not,
@@ -144,7 +144,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
title = plt.gca().get_title()
fig.clear()
# If im only has two dimensions, this reshape will add a leading
# dimension, and update will be called on index 0. If it has 3 or more
# dimensions, then all the leading dimensions will be compressed into
@@ -153,9 +153,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
reshaped_im = im.reshape(-1,s[-2],s[-1])
num_images = reshaped_im.shape[0]
fig.plot_idx = idx % num_images
to_plot = plot_func(reshaped_im[fig.plot_idx])
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
@@ -183,7 +183,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
plt.title(title)
if len(im.shape) >= 3:
@@ -196,14 +196,14 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
result = make_plot(0)
update = make_plot
def on_action(event):
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
update(fig.plot_idx - 1)
elif event.key == 'down' or event.button == 'down':
@@ -219,9 +219,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
fig.my_callbacks = []
fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action))
fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action))
return result
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
"""Plots the real part of a complex array with dimensions NxM
@@ -258,7 +258,7 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
@@ -306,7 +306,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
If a basis is explicitly passed, the image will be plotted in real-space
coordinates.
Parameters
----------
im : array
@@ -418,7 +418,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
units=units, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs):
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
"""Plots a set of probe translations in a nicely formatted way
Parameters
@@ -431,6 +431,8 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
Default is um, units to report in (assuming input in m)
lines : bool
Whether to plot lines indicating the path taken
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -455,6 +457,9 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
translations = translations * factor
plt.plot(translations[:,0], translations[:,1],'k.')
if invert_xaxis:
plt.gca().invert_xaxis()
if lines:
plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
plt.xlabel('X (' + units + ')')
@@ -463,7 +468,7 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
return fig
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe'):
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe', invert_xaxis=True):
"""Plots a set of nanomap data in a flexible way
Parameters
@@ -478,6 +483,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
Default is um, units to report in (assuming input in m)
convention : str
Default is 'probe', alternative is 'obj'. Whether the translations refer to the probe or object.
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
Returns
-------
@@ -511,7 +518,9 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
s /= 4 # A rough value to make the size work out
plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
if invert_xaxis:
plt.gca().invert_xaxis()
plt.gca().set_facecolor('k')
plt.xlabel('Translation x (' + units + ')')
plt.ylabel('Translation y (' + units + ')')
@@ -522,12 +531,12 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs):
"""Plots a nanomap, with an image or stack of images for each point
In many situations, ptychography data or the output of ptychography
reconstructions is formatted as a set of images associated with various
points in real space. This function is designed to allow for browsing
through this kind of data, by making it possible to visualize a
"""
# This should pull heavily from the dataset.inspect function
@@ -560,11 +569,11 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(translations.shape[0]) * s0
s[idx] *= 4
return s
def update_colorbar(im):
#
# This solves the problem of the colorbar being changed
@@ -573,29 +582,29 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
# This is needed to update the colorbar
# only change limits if array contains multiple values
if np.min(im.get_array()) != np.max(im.get_array()):
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
#
# The meatiest part of this program, here we just go through and
# set up the plot how we want it
#
# First we set up the left-hand plot, which shows an overview map
axes[0].set_title('Relative Displacement Map')
translations = translations.detach().cpu().numpy()
if convention.lower() != 'probe':
translations = translations * -1
s = calculate_sizes(0)
nanomap_units_factor = get_units_factor(nanomap_units)
nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0],
nanomap_units_factor * translations[:,1],
@@ -616,7 +625,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# where the colorbar should have been to avoid stretching the
# nanomap plot, while still not showing the (now useless) colorbar.
cb1.remove()
# Now we set up the second plot, which shows the individual
# diffraction patterns
axes[1].set_title(image_title)
@@ -636,7 +645,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This fails if the basis is not rectangular
basis_norm = np.linalg.norm(np_basis, axis = 0)
basis_norm = basis_norm * get_units_factor(image_units)
extent = [0, example_im.shape[-1]*basis_norm[1], 0,
example_im.shape[-2]*basis_norm[0]]
else:
@@ -657,9 +666,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].text_box.set_path_effects(
[patheffects.Stroke(linewidth=2, foreground='black'),
patheffects.Normal()])
meas = axes[1].imshow(im, extent=extent, cmap=cmap)
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
@@ -667,7 +676,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
cb2.ax.tick_params(labelrotation=20)
cb2.ax.set_title(image_colorbar_title, size="medium", pad=5)
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
# This function handles all the updating, except for moving the
# slider value. This is done because the slider widget is
# ultimately responsible for triggering an update, so all other
@@ -677,7 +686,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# We have to explicitly make it an integer because the slider will
# output floats (even if they are still integer-valued)
idx = int(idx)
# Get the new data for this index
im = get_image_func(idx)
if len(im.shape) >= 3:
@@ -688,22 +697,22 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].image_idx = im_idx
axes[1].text_box.set_text(str(im_idx))
im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
# Now we resize the nanomap to show the new selection
axes[0].collections[0].set_sizes(calculate_sizes(idx))
# And we update the data in the image as well
ax_im = axes[1].images[-1]
ax_im.set_data(im)
update_colorbar(ax_im)
#
# Now we define the functions to handle various kinds of events
# that can be thrown our way
#
# We start by creating the slider here, so it can be used
# by the update hooks.
slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d")
@@ -717,7 +726,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# while the mouse is within the image display
im = im.reshape(-1,im.shape[-2],im.shape[-1])
im_idx = axes[1].image_idx
if event.key == 'up' or event.button == 'up' \
or event.key == 'left':
im_idx = (im_idx - 1) % im.shape[0]
@@ -735,7 +744,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up' or event.key == 'left':
idx = slider.val - 1
elif event.key == 'down' or event.button == 'down' or event.key == 'right':
@@ -744,7 +753,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This prevents errors from being thrown on irrelevant key
# or mouse input
return
# Handle the wraparound and trigger the update
idx = int(idx) % translations.shape[0]
slider.set_val(idx)
@@ -755,17 +764,19 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# for example, scroll events that happen over the nanomap
if event.mouseevent.button == 1:
slider.set_val(event.ind[0])
# Here we connect the various update functions
cid1 = fig.canvas.mpl_connect('pick_event',on_pick)
cid2 = fig.canvas.mpl_connect('key_press_event',on_action)
cid3 = fig.canvas.mpl_connect('scroll_event',on_action)
# It's so dumb that matplotlib doesn't automatically track this for you
fig.nanomap_cids = [cid1,cid2,cid3]
fig.nanomap_cids = [cid1,cid2,cid3]
slider.on_changed(update)
# Throw an extra update into the mix just to get rid of any things
# (like the nanomap dot sizes) that otherwise would change on the
# first update
update(0)
return fig
CDTools/tools/plotting/polarized_plotting.py
+342
View File
@@ -0,0 +1,342 @@
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()