Merge pull request #9 from SebastianM-C/smc/fix

Fix bug in `GaussLaser` and simplify interface

Author: SebastianM-C

.gitignore

@@ -1,6 +1,6 @@
 *.jl.*.cov
 *.jl.cov
 *.jl.mem
-/Manifest.toml
+/Manifest*.toml
 /docs/Manifest.toml
 /docs/build/

Project.toml

@@ -7,6 +7,7 @@ authors = ["Sebastian Micluța-Câmpeanu <sebastian.mc95@proton.me> and contribu
 HypergeometricFunctions = "34004b35-14d8-5ef3-9330-4cdb6864b03a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -23,6 +24,7 @@ JET = "0.9, 0.10, 0.11"
 LaserTypes = "0.2"
 LinearAlgebra = "1.11"
 ModelingToolkit = "11"
+ModelingToolkitBase = "1.20.0"
 OrdinaryDiffEqNonlinearSolve = "1.10.0"
 OrdinaryDiffEqRosenbrock = "1.11.0"
 OrdinaryDiffEqVerner = "1.2.0"

scripts/benchmark_field_eval.jl

@@ -16,16 +16,19 @@ m_val = 1
 
 # --- FieldEvaluator setup ---
 @named ref_frame = ProperFrame(:atomic)
-@named laser = LaguerreGaussLaser(
-    wavelength=λ_val, a0=a₀_val, beam_waist=w₀_val,
-    radial_index=p_val, azimuthal_index=m_val,
-    ref_frame=ref_frame, temporal_profile=:constant)
+@named laser = LaguerreGaussLaser(;
+    wavelength = λ_val, a0 = a₀_val, beam_waist = w₀_val,
+    radial_index = p_val, azimuthal_index = m_val,
+    ref_frame, temporal_profile = :constant
+)
 
-fe = FieldEvaluator(laser, ref_frame)
+fe = FieldEvaluator(laser)
 
 # --- LaserTypes setup ---
-lt_laser = LaserTypes.LaguerreGaussLaser(:atomic;
-    λ=λ_val, a₀=a₀_val, w₀=w₀_val, p=p_val, m=m_val)
+lt_laser = LaserTypes.LaguerreGaussLaser(
+    :atomic;
+    λ = λ_val, a₀ = a₀_val, w₀ = w₀_val, p = p_val, m = m_val
+)
 
 # Test point
 x, y, z = 0.3w₀_val, 0.1w₀_val, 0.5w₀_val
@@ -60,8 +63,10 @@ display(@benchmark LaserTypes.B($pos, $t_val, $lt_laser))
 println()
 
 println("=== LaserTypes E + B ===")
-display(@benchmark begin
-    LaserTypes.E($pos, $t_val, $lt_laser)
-    LaserTypes.B($pos, $t_val, $lt_laser)
-end)
+display(
+    @benchmark begin
+        LaserTypes.E($pos, $t_val, $lt_laser)
+        LaserTypes.B($pos, $t_val, $lt_laser)
+    end
+)
 println()

scripts/ensemble.jl

@@ -9,7 +9,7 @@ using GLMakie
 using LaTeXStrings
 
 const inch = 96
-const pt = 4/3
+const pt = 4 / 3
 const cm = inch / 2.54
 
 # set_theme!(
@@ -22,21 +22,21 @@ const cm = inch / 2.54
 # Code is using Atomic Units !!!
 # natural constants
 c = 137.03599908330932 # speed of light
-qme = -1. # specific charge
+qme = -1.0 # specific charge
 
 h = 2π
-α = 1/c
-ε₀=qme^2/(2α*h*c)
-μ₀=1/(ε₀*c^2)
+α = 1 / c
+ε₀ = qme^2 / (2α * h * c)
+μ₀ = 1 / (ε₀ * c^2)
 
 # derived
 ω = 0.057
-τ = 150/ω
-λ = 2π*c/ω
+τ = 150 / ω
+λ = 2π * c / ω
 w₀ = 75λ
 Rmax = 3.25w₀
 
-ξx, ξy = (1/√2, im/√2) .|> complex
+ξx, ξy = (1 / √2, im / √2) .|> complex
 
 # Laser parameters in atomic units
 λ_au = λ
@@ -54,20 +54,20 @@ z₀ = 0.0
 # @named ref_frame = LabFrame(:atomic)
 
 @named laser = LaguerreGaussLaser(
-    wavelength=λ_au,
-    a0=a₀,
-    beam_waist=w₀_au,
-    radial_index=p_index,
-    azimuthal_index=m_index,
-    ref_frame=ref_frame,
-    temporal_profile=:gaussian,  # Using Gaussian profile
-    temporal_width=τ_fwhm,
-    focus_position=z₀,
-    polarization=:circular
+    wavelength = λ_au,
+    a0 = a₀,
+    beam_waist = w₀_au,
+    radial_index = p_index,
+    azimuthal_index = m_index,
+    ref_frame,
+    temporal_profile = :gaussian,  # Using Gaussian profile
+    temporal_width = τ_fwhm,
+    focus_position = z₀,
+    polarization = :circular
 )
 
 # Create electron system
-@named lg_elec = ClassicalElectron(; laser, ref_frame)
+@named lg_elec = ClassicalElectron(; laser)
 
 
 # Compile the system
@@ -79,21 +79,21 @@ sys = mtkcompile(lg_elec)
 tspan = (τi, τf)
 
 # Create base problem with placeholder initial position
-x⁰ = [τi*c, 0.0, 0.0, 0.0]
+x⁰ = [τi * c, 0.0, 0.0, 0.0]
 u⁰ = [c, 0.0, 0.0, 0.0]
 
 u0 = [
     (sys.x) => x⁰,
-    (sys.u) => u⁰
+    (sys.u) => u⁰,
 ]
 
-prob = ODEProblem{false, SciMLBase.FullSpecialize}(sys, u0, tspan, u0_constructor=SVector{8}, fully_determined=true)
-sol0 = solve(prob, Vern9(), reltol = 1e-15, abstol = 1e-12)
+prob = ODEProblem{false, SciMLBase.FullSpecialize}(sys, u0, tspan, u0_constructor = SVector{8}, fully_determined = true)
+sol0 = solve(prob, Vern9(), reltol = 1.0e-15, abstol = 1.0e-12)
 
 # Sunflower pattern for initial positions
 N = 900
 
-const ϕ = (1 + √5)/2
+const ϕ = (1 + √5) / 2
 
 function radius(k, n, b)
     if k > n - b
@@ -105,21 +105,21 @@ end
 
 function sunflower(n, α)
     points = []
-    angle_stride = 2π/ϕ^2 # geodesic ? 360 * ϕ :
+    angle_stride = 2π / ϕ^2 # geodesic ? 360 * ϕ :
     b = round(Int, α * sqrt(n))  # number of boundary points
 
     for k in 1:n
         r = radius(k, n, b)
         θ = k * angle_stride
-        append!(points, ([r * cos(θ), r * sin(θ)], ))
+        append!(points, ([r * cos(θ), r * sin(θ)],))
     end
 
     return points
 end
 
 # Generate initial positions in sunflower pattern
-R₀ = Rmax*sunflower(N, 2)
-xμ = [[τi*c, r..., 0.] for r in R₀]
+R₀ = Rmax * sunflower(N, 2)
+xμ = [[τi * c, r..., 0.0] for r in R₀]
 
 # Use SymbolicIndexingInterface to set positions
 set_x = setsym_oop(prob, [Initial(sys.x); Initial(sys.u)]);
@@ -136,59 +136,63 @@ function prob_func(prob, i, repeat)
     # Set new initial conditions
     u0, p = set_x(prob, SVector{8}(x_new..., u_new...))
 
-    remake(prob; u0, p)
+    return remake(prob; u0, p)
 end
 
 # Absolute error tolerance function
 function abserr(a₀)
     amp = log10(a₀)
-    expo = -amp^2/27 + 32amp/27 - 220/27
-    10^expo
+    expo = -amp^2 / 27 + 32amp / 27 - 220 / 27
+    return 10^expo
 end
 
 # Create ensemble problem
-ensemble = EnsembleProblem(prob; prob_func, safetycopy=false)
+ensemble = EnsembleProblem(prob; prob_func, safetycopy = false)
 
 # Solve ensemble
-solution = solve(ensemble, Vern9(), EnsembleThreads();
-                    reltol=1e-12, abstol=abserr(a₀),
-                    trajectories=N)
+solution = solve(
+    ensemble, Vern9(), EnsembleThreads();
+    reltol = 1.0e-12, abstol = abserr(a₀),
+    trajectories = N
+)
 
 # ru = solution.u[1](range(τi, τf, 1001), idxs = [sys.x; sys.u])
 
 # Solve single trajectory for visualization (electron #1)
 x_single = SVector{8}(xμ[1]..., u⁰...)
 u0_single, p_single = set_x(prob, x_single)
-prob_single = remake(prob; u0=u0_single, p=p_single)
+prob_single = remake(prob; u0 = u0_single, p = p_single)
 
-sol = solve(prob_single, Vern9(),
-            reltol=1e-12,
-            abstol=1e-20)
+sol = solve(
+    prob_single, Vern9(),
+    reltol = 1.0e-12,
+    abstol = 1.0e-20
+)
 
 #### eval field
 
 _t = 0
 _x = sol[sys.x, 500]
 
-x_sub = map(x->EvalAt(_t)(x[1])=>x[2], collect(sys.x .=> _x))
-eval_point = [laser.τ=>0; x_sub; sys.t => EvalAt(_t)(sys.x[1]) / c]
+x_sub = map(x -> EvalAt(_t)(x[1]) => x[2], collect(sys.x .=> _x))
+eval_point = [laser.τ => 0; x_sub; sys.t => EvalAt(_t)(sys.x[1]) / c]
 
 all_eqs = Symbolics.fixpoint_sub(equations(laser), merge(initial_conditions(laser), initial_conditions(eval_point)))
-eq_dict = Dict(map(eq->eq.lhs=>eq.rhs, all_eqs[setdiff(1:19, 10:15)]))
+eq_dict = Dict(map(eq -> eq.lhs => eq.rhs, all_eqs[setdiff(1:19, 10:15)]))
 Symbolics.fixpoint_sub(all_eqs, eq_dict)
 
 # using CairoMakie
 # Visualization
-fig = Figure(fontsize=14pt)
+fig = Figure(fontsize = 14pt)
 # ax = Axis3(fig[1, 1], aspect=:data)
-ax = Axis3(fig[1, 1], aspect=(1, 1, 1))
+ax = Axis3(fig[1, 1], aspect = (1, 1, 1))
 
 
 # Extract trajectory
-t_range = range(τi, τf, length=10001)
-x_traj = [sol(t, idxs=sys.x[2]) for t in t_range]
-y_traj = [sol(t, idxs=sys.x[3]) for t in t_range]
-z_traj = [sol(t, idxs=sys.x[4]) for t in t_range]
+t_range = range(τi, τf, length = 10001)
+x_traj = [sol(t, idxs = sys.x[2]) for t in t_range]
+y_traj = [sol(t, idxs = sys.x[3]) for t in t_range]
+z_traj = [sol(t, idxs = sys.x[4]) for t in t_range]
 
 lines!(ax, x_traj, y_traj, z_traj)
 

src/ElectronDynamicsModels.jl

@@ -1,8 +1,8 @@
 module ElectronDynamicsModels
 
 using ModelingToolkit
-using ModelingToolkit: SymbolicT, build_explicit_observed_function
-using SymbolicIndexingInterface: setsym_oop
+using ModelingToolkitBase: AbstractSystem, SymbolicT, build_explicit_observed_function, get_systems
+using SymbolicIndexingInterface: getname, setsym_oop
 using PhysicalConstants, Unitful, UnitfulAtomic
 using PhysicalConstants.CODATA2018: c_0, e, m_e, ε_0
 using LinearAlgebra

src/dynamics.jl

@@ -1,4 +1,5 @@
 @component function ParticleDynamics(; name, mass, ref_frame)
+    @unpack c = ref_frame
     iv = ModelingToolkit.get_iv(ref_frame)
     D = Differential(iv)
 
@@ -12,7 +13,6 @@
 
     if nameof(iv) == :τ
         τ = iv
-        c = ref_frame.c
 
         @variables begin
             t(iv), [description = "Universal time"]
@@ -36,7 +36,6 @@
         System(eqs, iv, [t, γ, x, u, p, F_total], [m]; name, systems = [ref_frame])
     elseif nameof(iv) == :t
         t = iv
-        c = ref_frame.c
 
         @variables begin
             τ(t), [description = "Proper time"]