Add FieldEvaluator

Project.toml

@@ -8,24 +8,31 @@ HypergeometricFunctions = "34004b35-14d8-5ef3-9330-4cdb6864b03a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [compat]
 Aqua = "0.8"
+CSV = "0.10"
 HypergeometricFunctions = "0.3.28"
 JET = "0.9, 0.10, 0.11"
 LaserTypes = "0.2"
-LinearAlgebra = "1.11.0"
+LinearAlgebra = "1.11"
 ModelingToolkit = "11"
 OrdinaryDiffEqNonlinearSolve = "1.10.0"
 OrdinaryDiffEqRosenbrock = "1.11.0"
 OrdinaryDiffEqVerner = "1.2.0"
 PhysicalConstants = "0.2.3"
 Plots = "1"
+Random = "1.11"
 SciMLBase = "2.102.1"
+StaticArrays = "1"
 Statistics = "1"
+SymbolicIndexingInterface = "0.3"
 Symbolics = "7.10"
 Test = "1.11"
 Unitful = "1.21.0"
@@ -34,15 +41,20 @@ julia = "1.11"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LaserTypes = "e07c0bfa-524c-4f35-a151-c3dd916fa2f0"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
 OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "JET", "LaserTypes", "Test", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqVerner", "Plots", "SciMLBase", "Statistics"]
+test = ["Aqua", "CSV", "JET", "LaserTypes", "LinearAlgebra", "Test", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqVerner", "Plots", "Random", "SciMLBase", "StaticArrays", "Statistics", "SymbolicIndexingInterface"]

scripts/benchmark_field_eval.jl

@@ -0,0 +1,67 @@
+using ElectronDynamicsModels
+using ModelingToolkit
+using LaserTypes: LaserTypes
+using BenchmarkTools
+using StaticArrays
+
+# Common parameters (atomic units)
+c = 137.03599908330932
+ω = 0.057
+T₀ = 2π / ω
+λ_val = c * T₀
+w₀_val = 75 * λ_val
+a₀_val = 2.0
+p_val = 1
+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)
+
+fe = FieldEvaluator(laser, ref_frame)
+
+# --- LaserTypes setup ---
+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
+t_val = 0.25T₀
+pos = SVector(x, y, z)
+
+# Verify agreement before benchmarking
+result_fe = fe([t_val, x, y, z])
+lt_E = LaserTypes.E(pos, t_val, lt_laser)
+lt_B = LaserTypes.B(pos, t_val, lt_laser)
+
+println("=== Sanity check ===")
+println("FieldEvaluator E: ", result_fe.E)
+println("LaserTypes     E: ", lt_E)
+println("FieldEvaluator B: ", result_fe.B)
+println("LaserTypes     B: ", lt_B)
+println()
+
+# --- Benchmarks ---
+txyz = SVector(t_val, x, y, z)
+
+println("=== FieldEvaluator (E + B in one call) ===")
+display(@benchmark $fe($txyz))
+println()
+
+println("=== LaserTypes E only ===")
+display(@benchmark LaserTypes.E($pos, $t_val, $lt_laser))
+println()
+
+println("=== LaserTypes B only ===")
+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)
+println()

scripts/ensemble.jl

@@ -55,7 +55,7 @@ z₀ = 0.0
 
 @named laser = LaguerreGaussLaser(
     wavelength=λ_au,
-    amplitude=a₀,
+    a0=a₀,
     beam_waist=w₀_au,
     radial_index=p_index,
     azimuthal_index=m_index,
@@ -173,7 +173,7 @@ _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]
 
-all_eqs = Symbolics.fixpoint_sub(equations(laser), merge(defaults(laser), Dict(eval_point)))
+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)]))
 Symbolics.fixpoint_sub(all_eqs, eq_dict)
 

src/ElectronDynamicsModels.jl

@@ -1,11 +1,14 @@
 module ElectronDynamicsModels
 
 using ModelingToolkit
+using ModelingToolkit: SymbolicT, build_explicit_observed_function
+using SymbolicIndexingInterface: setsym_oop
 using PhysicalConstants, Unitful, UnitfulAtomic
 using PhysicalConstants.CODATA2018: c_0, e, m_e, ε_0
 using LinearAlgebra
 using Symbolics
 using HypergeometricFunctions: HypergeometricFunctions, _₁F₁, pochhammer
+using StaticArrays
 
 m_dot(x, y) = x[1] * y[1] - x[2] * y[2] - x[3] * y[3] - x[4] * y[4]
 
@@ -19,13 +22,15 @@ export ReferenceFrame, ProperFrame, LabFrame,
     ChargedParticle,
     ClassicalElectron,
     RadiatingElectron,
-    LandauLifshitzElectron
+    LandauLifshitzElectron,
+    FieldEvaluator
 
 include("base.jl")
 include("dynamics.jl")
 include("fields.jl")
 include("radiation_reaction.jl")
 include("external_fields.jl")
 include("systems.jl")
+include("field_evaluator.jl")
 
 end

src/external_fields.jl

@@ -6,12 +6,19 @@ The beam propagates along the z-direction with waist w₀ at z=0.
 
 Parameters:
 - λ: wavelength
-- a₀: normalized vector potential amplitude
+- a₀: normalized vector potential (a0 kwarg)
 - w₀: beam waist (defaults to 75λ)
 - T0: pulse duration parameter
 - τ0: temporal width parameter
 """
-@component function GaussLaser(; name, wavelength=1.0, amplitude=10.0, beam_waist=nothing, polarization = :linear, ref_frame)
+@component function GaussLaser(; name, wavelength=nothing, frequency=nothing, a0=10.0, beam_waist=nothing, polarization = :linear, ref_frame)
+    if wavelength === nothing && frequency === nothing
+        wavelength = 1.0
+    end
+    if wavelength !== nothing && frequency !== nothing
+        error("Specify either wavelength or frequency, not both")
+    end
+
     # New interface with spacetime
     @named field_dynamics = EMFieldDynamics(; ref_frame)
 
@@ -25,12 +32,12 @@ Parameters:
     @unpack E, B = field_dynamics
 
     @parameters begin
-        λ = wavelength
-        a₀ = amplitude
-        ω
+        λ, [guess = 1.0]
+        a₀ = a0
+        ω, [guess = 1.0]
         k
         E₀
-        w₀ = beam_waist === nothing ? 75λ : beam_waist
+        w₀
         z_R
         T0 = 100
         τ0
@@ -112,9 +119,14 @@ Parameters:
         w₀ => missing
     ]
 
-    initial_conditions = [
-        τ0 => 10 / ω
-    ]
+    initial_conditions = Pair{SymbolicT,Any}[τ0 => 10 / ω]
+    push!(initial_conditions, w₀ => (beam_waist === nothing ? 75λ : beam_waist))
+    if wavelength !== nothing
+        push!(initial_conditions, λ => wavelength)
+    end
+    if frequency !== nothing
+        push!(initial_conditions, ω => frequency)
+    end
 
     sys = System(eqs, iv, [x, t, wz, z, r], [λ, a₀, ω, k, E₀, w₀, z_R, T0, τ0];
         name,
@@ -136,7 +148,14 @@ electrons can exhibit figure-8 motion when a₀ ~ 1.
 
 Reference: Sarachik & Schappert, Phys. Rev. D 1, 2738 (1970)
 """
-@component function PlaneWave(; name, amplitude=1.0, frequency=1.0, k_direction=[0,0,1], polarization=[1,0,0], ref_frame)
+@component function PlaneWave(; name, amplitude=1.0, wavelength=nothing, frequency=nothing, k_direction=[0,0,1], polarization=[1,0,0], ref_frame)
+    if wavelength === nothing && frequency === nothing
+        frequency = 1.0
+    end
+    if wavelength !== nothing && frequency !== nothing
+        error("Specify either wavelength or frequency, not both")
+    end
+
     # New interface with spacetime
     @named field_dynamics = EMFieldDynamics(; ref_frame)
 
@@ -154,7 +173,13 @@ Reference: Sarachik & Schappert, Phys. Rev. D 1, 2738 (1970)
 
     @unpack E, B = field_dynamics
 
-    @parameters A=amplitude ω=frequency k_dir[1:3]=k_direction pol[1:3]=polarization λ
+    @parameters begin
+        A = amplitude
+        ω, [guess = 1.0]
+        k_dir[1:3] = k_direction
+        pol[1:3] = polarization
+        λ, [guess = 1.0]
+    end
 
     # Normalize k direction to get unit vector k̂
     k_norm = sqrt(k_dir[1]^2 + k_dir[2]^2 + k_dir[3]^2)
@@ -194,9 +219,17 @@ Reference: Sarachik & Schappert, Phys. Rev. D 1, 2738 (1970)
         ω => missing
     ]
 
+    initial_conditions = Pair{SymbolicT,Any}[]
+    if wavelength !== nothing
+        push!(initial_conditions, λ => wavelength)
+    end
+    if frequency !== nothing
+        push!(initial_conditions, ω => frequency)
+    end
+
     vars = nameof(iv) == :τ ? [x, t] : [x]
 
-    sys = System(eqs, iv, vars, [A, ω, k_dir, pol, λ]; name, systems=[ref_frame], initialization_eqs, bindings)
+    sys = System(eqs, iv, vars, [A, ω, k_dir, pol, λ]; name, systems=[ref_frame], initialization_eqs, bindings, initial_conditions)
 
     extend(sys, field_dynamics)
 end
@@ -240,7 +273,7 @@ The beam is characterized by radial index p and azimuthal index m.
 
 Parameters:
 - λ: wavelength
-- a₀: normalized vector potential amplitude
+- a₀: normalized vector potential (a0 kwarg)
 - w₀: beam waist (defaults to 75λ)
 - radial_index (p): radial mode number, p ≥ 0
 - azimuthal_index (m): azimuthal mode number (orbital angular momentum)
@@ -252,8 +285,9 @@ Reference: Allen et al., Phys. Rev. A 45, 8185 (1992)
 """
 @component function LaguerreGaussLaser(;
     name,
-    wavelength=1.0,
-    amplitude=10.0,
+    wavelength=nothing,
+    frequency=nothing,
+    a0=10.0,
     beam_waist=nothing,
     radial_index=0,      # p
     azimuthal_index=1,   # m
@@ -263,6 +297,13 @@ Reference: Allen et al., Phys. Rev. A 45, 8185 (1992)
     focus_position=nothing,       # focal position along z-axis
     polarization = :linear
 )
+    if wavelength === nothing && frequency === nothing
+        wavelength = 1.0
+    end
+    if wavelength !== nothing && frequency !== nothing
+        error("Specify either wavelength or frequency, not both")
+    end
+
     # New interface with spacetime
     @named field_dynamics = EMFieldDynamics(; ref_frame)
 
@@ -289,12 +330,12 @@ Reference: Allen et al., Phys. Rev. A 45, 8185 (1992)
     Npm_val = sqrt(pochhammer(radial_index + 1, mₐ))
 
     params = @parameters begin
-        λ = wavelength
-        a₀ = amplitude
-        ω
+        λ, [guess = 1.0]
+        a₀ = a0
+        ω, [guess = 1.0]
         k
         E₀
-        w₀ = beam_waist === nothing ? 75λ : beam_waist
+        w₀
         z_R
         τ0 = temporal_width === nothing ? 100.0 : temporal_width
 
@@ -439,6 +480,15 @@ Reference: Allen et al., Phys. Rev. A 45, 8185 (1992)
         w₀ => missing
     ]
 
+    initial_conditions = Pair{SymbolicT,Any}[]
+    push!(initial_conditions, w₀ => (beam_waist === nothing ? 75λ : beam_waist))
+    if wavelength !== nothing
+        push!(initial_conditions, λ => wavelength)
+    end
+    if frequency !== nothing
+        push!(initial_conditions, ω => frequency)
+    end
+
     vars = if nameof(iv) == :τ
         [x, t, z, r, θ, wz, σ, rwz, env]
     else
@@ -447,6 +497,6 @@ Reference: Allen et al., Phys. Rev. A 45, 8185 (1992)
     end
 
     sys = System(eqs, iv, vars, params;
-        name, systems=[ref_frame], initialization_eqs, bindings)
+        name, systems=[ref_frame], initialization_eqs, bindings, initial_conditions)
     extend(sys, field_dynamics)
 end