|
| 1 | +using ElectronDynamicsModels |
| 2 | +using ModelingToolkit |
| 3 | +using OrdinaryDiffEqVerner, OrdinaryDiffEqTsit5 |
| 4 | +using OrdinaryDiffEqNonlinearSolve |
| 5 | +using SciMLBase |
| 6 | +using StaticArrays |
| 7 | +using SymbolicIndexingInterface |
| 8 | +using LinearAlgebra |
| 9 | +using FFTW |
| 10 | +using CairoMakie |
| 11 | + |
| 12 | +# Atomic units |
| 13 | +const c = 137.03599908330932 |
| 14 | + |
| 15 | +# Laser parameters |
| 16 | +ω = 0.057 |
| 17 | +τ = 150 / ω |
| 18 | +λ = 2π * c / ω |
| 19 | +w₀ = 75λ |
| 20 | +Rmax = 3.25w₀ |
| 21 | + |
| 22 | +a₀ = 10.0 |
| 23 | + |
| 24 | +@named ref_frame = ProperFrame(:atomic) |
| 25 | + |
| 26 | +@named laser = LaguerreGaussLaser(; |
| 27 | + wavelength = λ, |
| 28 | + a0 = a₀, |
| 29 | + beam_waist = w₀, |
| 30 | + radial_index = 2, |
| 31 | + azimuthal_index = -2, |
| 32 | + ref_frame, |
| 33 | + temporal_profile = :gaussian, |
| 34 | + temporal_width = τ, |
| 35 | + focus_position = 0.0, |
| 36 | + polarization = :circular |
| 37 | +) |
| 38 | +@named elec = ClassicalElectron(; laser) |
| 39 | +sys = mtkcompile(elec) |
| 40 | + |
| 41 | +# Time span |
| 42 | +τi = -8τ |
| 43 | +τf = 8τ |
| 44 | +tspan = (τi, τf) |
| 45 | + |
| 46 | +# Single electron solve (for parameter access) |
| 47 | +x⁰ = [τi * c, 0.0, 0.0, 0.0] |
| 48 | +u⁰ = [c, 0.0, 0.0, 0.0] |
| 49 | + |
| 50 | +u0 = [ |
| 51 | + sys.x => x⁰, |
| 52 | + sys.u => u⁰, |
| 53 | +] |
| 54 | + |
| 55 | +prob = ODEProblem{false, SciMLBase.FullSpecialize}( |
| 56 | + sys, u0, tspan, u0_constructor = SVector{8}, fully_determined = true |
| 57 | +) |
| 58 | +sol0 = solve(prob, Vern9(), reltol = 1.0e-15, abstol = 1.0e-12) |
| 59 | + |
| 60 | +# Sunflower distribution for electron positions |
| 61 | +const ϕ = (1 + √5) / 2 |
| 62 | + |
| 63 | +function radius(k, n, b) |
| 64 | + if k > n - b |
| 65 | + return 1.0 |
| 66 | + else |
| 67 | + return sqrt(k - 0.5) / sqrt(n - (b + 1) / 2) |
| 68 | + end |
| 69 | +end |
| 70 | + |
| 71 | +function sunflower(n, α) |
| 72 | + points = Vector{Vector{Float64}}() |
| 73 | + angle_stride = 2π / ϕ^2 |
| 74 | + b = round(Int, α * sqrt(n)) |
| 75 | + for k in 1:n |
| 76 | + r = radius(k, n, b) |
| 77 | + θ = k * angle_stride |
| 78 | + push!(points, [r * cos(θ), r * sin(θ)]) |
| 79 | + end |
| 80 | + return points |
| 81 | +end |
| 82 | + |
| 83 | +# Ensemble solve |
| 84 | +N = 300 |
| 85 | +R₀ = Rmax * sunflower(N, 2) |
| 86 | +xμ = [[τi * c, r..., 0.0] for r in R₀] |
| 87 | + |
| 88 | +set_x = setsym_oop(prob, [Initial(sys.x); Initial(sys.u)]) |
| 89 | + |
| 90 | +function prob_func(prob, i, repeat) |
| 91 | + x_new = SVector{4}(xμ[i]...) |
| 92 | + u_new = SVector{4}(c, 0.0, 0.0, 0.0) |
| 93 | + u0, p = set_x(prob, SVector{8}(x_new..., u_new...)) |
| 94 | + return remake(prob; u0, p) |
| 95 | +end |
| 96 | + |
| 97 | +function abserr(a₀) |
| 98 | + amp = log10(a₀) |
| 99 | + expo = -amp^2 / 27 + 32amp / 27 - 220 / 27 |
| 100 | + return 10^expo |
| 101 | +end |
| 102 | + |
| 103 | +ensemble = EnsembleProblem(prob; prob_func, safetycopy = false) |
| 104 | +solution = solve( |
| 105 | + ensemble, Vern9(), EnsembleThreads(); |
| 106 | + reltol = 1.0e-12, abstol = abserr(a₀), trajectories = N |
| 107 | +) |
| 108 | + |
| 109 | +# Radiation computation |
| 110 | + |
| 111 | +trajs = trajectory_interpolants(solution) |
| 112 | + |
| 113 | +# Screen parameters |
| 114 | +const Z = 2.0e5λ |
| 115 | +const δt = 2π / ω / 4 |
| 116 | +const N_samples = floor(Int, (τf - τi) / δt) |
| 117 | +const x⁰_start = c * τi + hypot(Z, 25w₀ + Rmax) |
| 118 | + |
| 119 | +Nx = 50 |
| 120 | +Ny = 50 |
| 121 | + |
| 122 | +x⁰_samples = range(start = x⁰_start, step = c * δt, length = N_samples) |
| 123 | + |
| 124 | +screen = ObserverScreen( |
| 125 | + LinRange(-25w₀, 25w₀, Nx), |
| 126 | + LinRange(-25w₀, 25w₀, Ny), |
| 127 | + Z, |
| 128 | + x⁰_samples |
| 129 | +) |
| 130 | + |
| 131 | +A_s = accumulate_potential(trajs, screen, Tsit5()) |
| 132 | + |
| 133 | +# GPU |
| 134 | +# accumulate_potential(trajs, screen, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend())) |
| 135 | +A_ω = rfft(A_s, 1) |
| 136 | + |
| 137 | +# Find fundamental frequency bin |
| 138 | +freqs = rfftfreq(N_samples, 1 / δt) |
| 139 | +idx_f1 = findmin(x -> abs(x - ω / 2π), freqs)[2] |
| 140 | + |
| 141 | +# Plot the y-component of the potential at the fundamental frequency |
| 142 | +fig = Figure() |
| 143 | +ax = Axis(fig[1, 1], aspect = 1, xlabel = "x", ylabel = "y", title = "Thomson scattering (ω₁)") |
| 144 | +field = real.(A_ω[idx_f1, 3, :, :]) |
| 145 | +heatmap!( |
| 146 | + ax, collect(screen.x_grid), collect(screen.y_grid), field, |
| 147 | + colorrange = maximum(abs, field) .* (-1, 1), colormap = :seismic |
| 148 | +) |
| 149 | +fig |
0 commit comments