mirror of
https://github.com/yosata/fourierOptics.git
synced 2026-10-29 20:51:10 +00:00
340 lines
11 KiB
Python
340 lines
11 KiB
Python
import numpy as np
|
|
import pickle
|
|
import matplotlib.pyplot as plt
|
|
|
|
'''
|
|
Miscellaneous functions that do not belong specifically to any of the
|
|
other files.
|
|
|
|
Author: Joost van den Born
|
|
Contact: born@astron.nl
|
|
'''
|
|
|
|
|
|
def perform_zero_padding(A, xs, ys):
|
|
'''
|
|
Add zero padding around the array
|
|
|
|
Input:
|
|
A - 2D-Input array
|
|
xs,ys - The new shape after padding.
|
|
|
|
Output:
|
|
- Array with the new padded zeros
|
|
|
|
'''
|
|
assert A.shape[0] <= xs and A.shape[1] <= ys, 'New size is smaller than array size.'
|
|
x_margin = int((xs - A.shape[0])/2)
|
|
y_margin = int((ys - A.shape[1])/2)
|
|
return np.pad(A, (x_margin, y_margin))
|
|
|
|
def find_nearest(array, value):
|
|
'''
|
|
For a value it finds the nearest element in an array
|
|
Returns the index of that element in the original
|
|
array.
|
|
'''
|
|
array = np.asarray(array)
|
|
diff = (np.abs(array - value)).argmin()
|
|
idx = np.unravel_index(diff, array.shape)
|
|
return idx
|
|
|
|
def find_maximum(array):
|
|
'''
|
|
Finds the index of the maximum of the given array
|
|
'''
|
|
idx = np.unravel_index(array.argmax(), array.shape)
|
|
return idx
|
|
|
|
def select_subsection(x, y , outputCoordinates):
|
|
'''
|
|
Returns the indices of a subsection of an x,y meshgrid
|
|
'''
|
|
select_1 = find_nearest(x, outputCoordinates[0])[1]
|
|
select_2 = find_nearest(x, outputCoordinates[1])[1]
|
|
select_3 = find_nearest(y, outputCoordinates[2])[0]
|
|
select_4 = find_nearest(y, outputCoordinates[3])[0]
|
|
return [select_1, select_2, select_3, select_4]
|
|
|
|
def perform_extrapolation(x, y, n, m):
|
|
'''
|
|
Performs linear extrapolation of a meshgrid to a new
|
|
shape n,m
|
|
|
|
Returns the new extrapolated x,y meshgrid
|
|
'''
|
|
|
|
# Determine properties of x,y meshgrid
|
|
xshape = x.shape[0]
|
|
yshape = y.shape[0]
|
|
dpix = (np.abs(np.min(x)) + np.abs(np.max(x))) / (xshape-1)
|
|
|
|
# Extrapolate from, and include, the old x,y meshgrid
|
|
xnew = np.linspace(-(n-1)/2*dpix, (n-1)/2*dpix, n)
|
|
ynew = np.linspace(-(m-1)/2*dpix, (m-1)/2*dpix, m)
|
|
xnew, ynew = np.meshgrid(xnew, ynew)
|
|
return xnew, ynew
|
|
|
|
def plot_results(E, x, y, logscale=False, fig_size=(8,3),grid=False, show=True, colormap='hot'):
|
|
x = x.real
|
|
y = y.real
|
|
_fig = plt.figure(figsize=fig_size)
|
|
_ax = _fig.add_subplot(121)
|
|
if logscale:
|
|
from matplotlib.colors import LogNorm
|
|
_im = _ax.imshow(np.abs(E), cmap=colormap, extent=[x.min(), x.max(), y.min(), y.max()], origin='lower', norm=LogNorm())
|
|
else:
|
|
_im = _ax.imshow(np.abs(E), cmap=colormap, extent=[x.min(), x.max(), y.min(), y.max()], origin='lower')
|
|
if grid:
|
|
_ax.grid(True)
|
|
_ax.set_xlim([x.min(), x.max()])
|
|
_ax.set_ylim([y.min(), y.max()])
|
|
_ax.set_xlabel('x [mm]')
|
|
_ax.set_ylabel('y [mm]')
|
|
_fig.colorbar(_im, ax=_ax, label='Amplitude')
|
|
_ax2 = _fig.add_subplot(122)
|
|
_im2 = _ax2.imshow(np.angle(E), cmap=colormap, extent=[x.min(), x.max(), y.min(), y.max()], origin='lower')
|
|
_ax2.set_xlim([x.min(), x.max()])
|
|
_ax2.set_ylim([y.min(), y.max()])
|
|
_ax2.set_xlabel('x [mm]')
|
|
_ax2.set_ylabel('y [mm]')
|
|
_fig.colorbar(_im2, ax=_ax2, label='Phase')
|
|
plt.tight_layout()
|
|
if show:
|
|
plt.show()
|
|
return _fig, _ax, _ax2
|
|
|
|
|
|
def plot_real_valued_results(E, x, y, logscale=False, fig_size=(8,8), grid=False, show=True, colormap='hot', labeltext=None):
|
|
x = x.real
|
|
y = y.real
|
|
_fig = plt.figure(figsize=fig_size)
|
|
_ax = _fig.add_subplot(111)
|
|
if logscale:
|
|
from matplotlib.colors import LogNorm
|
|
_im = _ax.imshow(E, cmap=colormap, extent=[x.min(), x.max(), y.min(), y.max()], origin='lower', norm=LogNorm())
|
|
else:
|
|
_im = _ax.imshow(E, cmap=colormap, extent=[x.min(), x.max(), y.min(), y.max()], origin='lower')
|
|
if grid:
|
|
_ax.grid(True)
|
|
_ax.set_xlim([x.min(), x.max()])
|
|
_ax.set_ylim([y.min(), y.max()])
|
|
_ax.set_xlabel('x [mm]')
|
|
_ax.set_ylabel('y [mm]')
|
|
_fig.colorbar(_im, ax=_ax, label=labeltext)
|
|
plt.tight_layout()
|
|
if show:
|
|
plt.show()
|
|
return _fig, _ax
|
|
|
|
|
|
def save_variable(filename, variable):
|
|
'''
|
|
Dumps the given variable to a binary pickle file
|
|
for later use.
|
|
'''
|
|
with open(filename, 'wb') as f:
|
|
pickle.dump(variable, f, protocol=-1 )
|
|
|
|
|
|
def load_variable(filename):
|
|
'''
|
|
Loads a binary pickle file for use.
|
|
'''
|
|
with open(filename, 'rb') as f:
|
|
variable = pickle.load(f)
|
|
return variable
|
|
|
|
|
|
def geometric_distortions(xp, yp, x, y, magnification=1):
|
|
'''
|
|
Plots a grid showing the geometric distortions. Effects can be magnified.
|
|
|
|
Parameters
|
|
----------
|
|
xp : array
|
|
Meshgrid array denoting the perturbed x-coordinates
|
|
yp : array
|
|
Meshgrid array denoting the perturbed y-coordinates
|
|
x : array
|
|
Meshgrid array denoting the unperturbed x-coordinates
|
|
y : array
|
|
Meshgrid array denoting the unperturbed y-coordinates
|
|
magnification : float (optional)
|
|
If given, amplifies the perturbation, so that it is more visible.
|
|
|
|
Returns
|
|
-------
|
|
fig, ax : matplotlib.pyplot object
|
|
Figure that can be shown using plt.show(), or modified by the user.
|
|
|
|
'''
|
|
dx = (xp - x) * magnification
|
|
dy = (yp - y) * magnification
|
|
xpn = x + dx
|
|
ypn = y + dy
|
|
|
|
fig = plt.figure(figsize=(10,10))
|
|
ax = fig.add_subplot(111, aspect=1)
|
|
for i in range(len(xp)):
|
|
ax.plot(x[:,i], y[:,i], 'k-', alpha=0.2)
|
|
ax.plot(x[i,:], y[i,:], 'k-', alpha=0.2)
|
|
ax.plot(xpn[:,i], ypn[:,i], 'k-', alpha=1)
|
|
ax.plot(xpn[i,:], ypn[i,:], 'k-', alpha=1)
|
|
ax.plot(xpn, ypn, 'rx')
|
|
ax.set_xlabel('x [mm]')
|
|
ax.set_ylabel('y [mm]')
|
|
return fig, ax
|
|
|
|
|
|
def get_prism_rtm(x_in, y_in, theta_x_in, theta_y_in, thickness, a1x, a2x, a1y, a2y, n_glass):
|
|
'''
|
|
Finds the ray transfer matrix for a prism with the given parameters.
|
|
It does not assume small angles. It does not take into account any
|
|
(internal) reflections. The results were confirmed using CODEV.
|
|
It is assumed that the prism is located inside a vacuum (i.e. n=1
|
|
outside the prism).
|
|
|
|
Parameters
|
|
----------
|
|
x_in : float
|
|
Input x-coordinate at the point in the plane where the optical
|
|
axis reaches the front of the prism
|
|
y_in : float
|
|
Input y-coordinate at the point in the plane where the optical
|
|
axis reaches the front of the prism
|
|
theta_x_in : float
|
|
Input angle about the y-axis (so it affects the x-coordinates)
|
|
theta_y_in : float
|
|
Input angle about the x-axis (so it affects the y-coordinates)
|
|
thickness : float
|
|
Thickness of the prism along the optical axis
|
|
a1x : float
|
|
Rotation angle of the first prism surface about the y-axis,
|
|
so that the x-coordinates are affected.
|
|
a2x : float
|
|
Rotation angle of the second prism surface about the y-axis,
|
|
so that the x-coordinates are affected.
|
|
a1y : float
|
|
Rotation angle of the first prism surface about the x-axis,
|
|
so that the y-coordinates are affected.
|
|
a2y : float
|
|
Rotation angle of the second prism surface about the x-axis,
|
|
so that the y-coordinates are affected.
|
|
n_glass : float
|
|
Refractive index of the prism glass
|
|
|
|
Returns
|
|
-------
|
|
x_out : float
|
|
Output x-coordinate at the point in the plane where the optical
|
|
axis reaches the back of the prism
|
|
y_out : float
|
|
Output y-coordinate at the point in the plane where the optical
|
|
axis reaches the back of the prism
|
|
theta_x_out : float
|
|
Output angle about the y-axis (so it affects the x-coordinates)
|
|
theta_y_out : float
|
|
Output angle about the x-axis (so it affects the y-coordinates)
|
|
'''
|
|
theta_x_out = -1 * np.arcsin( -n_glass * np.sin( np.arcsin( 1/n_glass * np.sin(theta_x_in - a1x) ) + a1x - a2x ) ) + a2x
|
|
theta_y_out = -1 * np.arcsin( -n_glass * np.sin( np.arcsin( 1/n_glass * np.sin(theta_y_in - a1y) ) + a1y - a2y ) ) + a2y
|
|
|
|
xp1 = x_in / (1 + np.tan(a1x) * np.tan(theta_x_in))
|
|
yp1 = y_in / (1 + np.tan(a1y) * np.tan(theta_y_in))
|
|
|
|
tp1x = -1 * np.arcsin( 1 / n_glass * np.sin(theta_x_in - a1x) ) + a1x
|
|
xp2 = xp1 - (thickness + xp1 * (np.tan(a1x) - np.tan(a2x))) * (np.cos(a2x) * np.sin(tp1x)) / np.cos(tp1x - a2x)
|
|
|
|
tp1y = -1 * np.arcsin( 1 / n_glass * np.sin(theta_y_in - a1y) ) + a1y
|
|
yp2 = yp1 - (thickness + yp1 * (np.tan(a1y) - np.tan(a2y))) * (np.cos(a2y) * np.sin(tp1y)) / np.cos(tp1y - a2y)
|
|
|
|
x_out = xp2 * (1 + np.tan(a2x) * np.tan(theta_x_out))
|
|
y_out = yp2 * (1 + np.tan(a2y) * np.tan(theta_y_out))
|
|
|
|
return x_out, y_out, theta_x_out, theta_y_out
|
|
|
|
|
|
def get_paraxial_lens_rtm(x_in, y_in, theta_x_in, theta_y_in, f_lens_x, f_lens_y):
|
|
'''
|
|
Finds the ray transfer matrix of a paraxial lens, without small angle
|
|
approximations.
|
|
|
|
Parameters
|
|
----------
|
|
x_in : float
|
|
Input x-coordinate at the point in the plane where the optical
|
|
axis reaches the front of the prism
|
|
y_in : float
|
|
Input y-coordinate at the point in the plane where the optical
|
|
axis reaches the front of the prism
|
|
theta_x_in : float
|
|
Input angle about the y-axis (so it affects the x-coordinates)
|
|
theta_y_in : float
|
|
Input angle about the x-axis (so it affects the y-coordinates)
|
|
f_lens_x : float
|
|
Focal length of the lens for the points with a nonzero x-coordinate
|
|
value.
|
|
f_lens_y : float
|
|
Focal length of the lens for the points with a nonzero y-coordinate
|
|
value.
|
|
|
|
Returns
|
|
-------
|
|
x_out : float
|
|
Output x-coordinate at the point in the plane where the optical
|
|
axis reaches the back of the prism
|
|
y_out : float
|
|
Output y-coordinate at the point in the plane where the optical
|
|
axis reaches the back of the prism
|
|
theta_x_out : float
|
|
Output angle about the y-axis (so it affects the x-coordinates)
|
|
theta_y_out : float
|
|
Output angle about the x-axis (so it affects the y-coordinates)
|
|
'''
|
|
x_out = x_in
|
|
y_out = y_in
|
|
theta_x_out = np.arctan( (x_in + -f_lens_x * np.tan(theta_x_in)) / -f_lens_x )
|
|
theta_y_out = np.arctan( (y_in + -f_lens_y * np.tan(theta_y_in)) / -f_lens_y )
|
|
return x_out, y_out, theta_x_out, theta_y_out
|
|
|
|
|
|
def get_free_space_propagation_rtm(x_in, y_in, theta_x_in, theta_y_in, distance):
|
|
'''
|
|
Finds the ray transfer matrix of a free space propagation, without small
|
|
angle approximations.
|
|
|
|
Parameters
|
|
----------
|
|
x_in : float
|
|
Input x-coordinate at the point in the plane where the optical
|
|
axis reaches the front of the prism
|
|
y_in : float
|
|
Input y-coordinate at the point in the plane where the optical
|
|
axis reaches the front of the prism
|
|
theta_x_in : float
|
|
Input angle about the y-axis (so it affects the x-coordinates)
|
|
theta_y_in : float
|
|
Input angle about the x-axis (so it affects the y-coordinates)
|
|
distance : float
|
|
Distance to travel through free space.
|
|
|
|
Returns
|
|
-------
|
|
x_out : float
|
|
Output x-coordinate at the point in the plane where the optical
|
|
axis reaches the back of the prism
|
|
y_out : float
|
|
Output y-coordinate at the point in the plane where the optical
|
|
axis reaches the back of the prism
|
|
theta_x_out : float
|
|
Output angle about the y-axis (so it affects the x-coordinates)
|
|
theta_y_out : float
|
|
Output angle about the x-axis (so it affects the y-coordinates)
|
|
'''
|
|
x_out = x_in + distance * np.tan(theta_x_in)
|
|
y_out = y_in + distance * np.tan(theta_y_in)
|
|
theta_x_out = theta_x_in
|
|
theta_y_out = theta_y_in
|
|
return x_out, y_out, theta_x_out, theta_y_out
|