163 lines
4.9 KiB
Python
163 lines
4.9 KiB
Python
'''
|
|
Vacuum beam propagation
|
|
no plasma, no limiter nor reflections;
|
|
1D beam parameters;
|
|
general astigmatic beam
|
|
'''
|
|
|
|
from .. import TestCase, GrayTest
|
|
from .. import options, load_unit
|
|
|
|
import numpy as np
|
|
|
|
# SI unit converions
|
|
cm = 1e-2
|
|
mm = 1e-3
|
|
deg = np.pi/180
|
|
c = 299792458
|
|
|
|
|
|
def vectorize(*arg, **kwargs):
|
|
'''
|
|
Helper to use np.vectorize with keyword arguments as a decorator
|
|
'''
|
|
return lambda f: np.vectorize(f, *arg, **kwargs)
|
|
|
|
|
|
def rotate(angle: float):
|
|
'''
|
|
Builds a 2D rotation matrix
|
|
'''
|
|
c, s = np.cos(angle), np.sin(angle)
|
|
return np.array([[c, -s], [s, c]])
|
|
|
|
|
|
@vectorize(excluded=['k0', 'W', 'K', 'φ_w', 'φ_k'],
|
|
signature='(),(2,2),(2,2),(),(),(),()->()')
|
|
def eikonal_true(k0: float, W: np.array, K: np.array,
|
|
φ_w: float, φ_k: float, s: float,
|
|
ray: dict[str, float]):
|
|
'''
|
|
Computes the complex eikonal from the initial
|
|
beam parameters and the final ray position.
|
|
'''
|
|
r = np.array([ray['xt'], ray['yt']]) * cm
|
|
# Note: z = s₀ + z̅ where s₀ is the arclength of the
|
|
# central ray and z̅ the local beam coordinate.
|
|
z = (s + ray['zt']) * cm
|
|
|
|
# compute the complex curvature tensor
|
|
K = rotate(φ_k) @ K @ rotate(φ_k).T
|
|
W = rotate(φ_w) @ W @ rotate(φ_w).T
|
|
Q = K - 1j*W
|
|
|
|
# propagate the beam from 0→z
|
|
Q = Q @ np.linalg.inv(np.eye(2) + z*Q)
|
|
|
|
# Compute the eikonal
|
|
# Note: the Gouy phase is not treated by gray
|
|
S = k0 * (z + 1/2 * r @ Q @ r)
|
|
|
|
return S
|
|
|
|
|
|
@vectorize(excluded=['k0', 'W', 'K', 'φ_w', 'φ_k'],
|
|
signature='(),(2,2),(2,2),(),(),()->()')
|
|
def eikonal_gray(k0: float, W: np.array, K: np.array,
|
|
φ_w: float, φ_k: float, ray: dict[str, float]):
|
|
'''
|
|
Extracts the eikonal from the gray results
|
|
|
|
Notes:
|
|
1. the coordinates r,θ in the beam cross section
|
|
are defined by:
|
|
{ ξ_w = w_ξ r cosθ
|
|
{ η_w = w_η r cosθ,
|
|
|
|
2. for the (j,k)-th ray we have:
|
|
{ r = (j-1)/(n_r - 1)
|
|
{ θ = (k-1)/n_θ
|
|
|
|
so S_I = -(ξ_w²/w_ξ² + η_w²/w_η²) = -r²
|
|
and this should be independent of s.
|
|
|
|
3. S_R is given by k₀s, when integrating in
|
|
the phase (idst=2);
|
|
'''
|
|
r_max = 1.0
|
|
n_r = 6
|
|
δr = r_max / (n_r - 1)
|
|
r = (ray['j'] - 1) * δr
|
|
|
|
S_R = k0*ray['sst'] * cm
|
|
S_I = -r**2
|
|
|
|
return S_R + 1j * S_I
|
|
|
|
|
|
class Test(GrayTest, TestCase):
|
|
def test_eikonal(self):
|
|
'''
|
|
Check the correctness of beam propagation
|
|
by comparing the values of eikonal on each ray
|
|
at the final step.
|
|
'''
|
|
# Read the initial beam parameters
|
|
w_ξ, w_η, k_ξ, k_η, φ_w, φ_k = np.loadtxt(
|
|
str(self.inputs / 'beam.txt'),
|
|
usecols=[6, 7, 8, 9, 10, 11], skiprows=2,
|
|
unpack=True)
|
|
|
|
# Convert to SI
|
|
w_ξ, w_η = w_ξ * mm, w_η * mm
|
|
k_ξ, k_η = k_ξ / mm, k_η / mm
|
|
φ_w, φ_k = φ_w * deg, φ_k * deg
|
|
|
|
# Vacuum wavenumber
|
|
k0 = 2*np.pi/c * 170e9
|
|
|
|
# Beam width/curvature
|
|
W = 2/k0 * np.diag([1/w_ξ, 1/w_η])**2
|
|
K = np.diag([k_ξ, k_η])
|
|
|
|
# Load the final rays data
|
|
data = load_unit(self.candidate / 'fort.9')
|
|
|
|
# Complex eikonal: analytical and gray values
|
|
s0 = data['sst'][0]
|
|
S_true = eikonal_true(k0, W, K, φ_w, φ_k, s0, data)
|
|
S_gray = eikonal_gray(k0, W, K, φ_w, φ_k, data)
|
|
|
|
if options.visual:
|
|
import matplotlib.pyplot as plt
|
|
|
|
plt.suptitle(self.__module__ + '.test_eikonal')
|
|
plt.subplot(121, aspect='equal')
|
|
plt.title('$|S_R - S_{R,gray}|$', loc='left')
|
|
plt.tricontour(data['xt'], data['yt'],
|
|
abs(S_true.real - S_gray.real),
|
|
levels=8, linewidths=0.5, colors='k')
|
|
plt.tricontourf(data['xt'], data['yt'],
|
|
abs(S_true.real - S_gray.real),
|
|
levels=8, cmap='summer')
|
|
plt.colorbar()
|
|
plt.scatter(data['xt'], data['yt'], c='k', s=5)
|
|
|
|
plt.subplot(122, aspect='equal')
|
|
plt.title('$|S_I - S_{I,gray}|$', loc='left')
|
|
plt.tricontour(data['xt'], data['yt'],
|
|
abs(S_true.imag - S_gray.imag),
|
|
levels=8, linewidths=0.5, colors='k')
|
|
plt.tricontourf(data['xt'], data['yt'],
|
|
abs(S_true.imag - S_gray.imag),
|
|
levels=8, cmap='summer')
|
|
plt.colorbar()
|
|
plt.scatter(data['xt'], data['yt'], c='k', s=5)
|
|
plt.show()
|
|
|
|
# Compare the eikonal of each ray
|
|
for jk, ray in enumerate(data):
|
|
with self.subTest(ray=(int(ray['j']), int(ray['k']))):
|
|
self.assertAlmostEqual(S_true[jk].imag, S_gray[jk].imag, 2)
|
|
self.assertAlmostEqual(S_true[jk].real, S_gray[jk].real, 1)
|