Scalarize Lagrangians and Hamiltonians. Rename keyword argument dosim…

…plify to simplify.
JuliaGNI · Dec 5, 2024 · e43f59b · e43f59b
1 parent dba48c7
commit e43f59b
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 33 deletions.
diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl
@@ -42,7 +42,7 @@ struct HamiltonianSystem
     equations
     functions
 
-    function HamiltonianSystem(H, t, q, p, params = NamedTuple(); dosimplify = true)
+    function HamiltonianSystem(H, t, q, p, params = NamedTuple(); simplify = true, scalarize = true)
 
         @assert eachindex(q) == eachindex(p)
 
@@ -51,7 +51,8 @@ struct HamiltonianSystem
 
         Dt, Dq, Dp = hamiltonian_derivatives(t, q, p)
 
-        Hs = dosimplify ? simplify(H) : H
+        Hs = scalarize ? Symbolics.scalarize(H) : H
+        Hs = simplify ? Symbolics.simplify(Hs) : Hs
 
         EHq = [expand_derivatives(Dt(q[i]) - Dp[i](Hs)) for i in eachindex(Dp,q)]
         EHp = [expand_derivatives(Dt(p[i]) + Dq[i](Hs)) for i in eachindex(Dq,p)]
@@ -84,7 +85,7 @@ struct HamiltonianSystem
 
         funcs = generate_code(code)
 
-        new(H, t, q, p, params, equs, funcs)
+        new(Hs, t, q, p, params, equs, funcs)
     end
 end
 

diff --git a/src/lagrangian.jl b/src/lagrangian.jl
@@ -10,7 +10,7 @@ struct LagrangianSystem
     equations
     functions
 
-    function LagrangianSystem(L, t, x, v, params = NamedTuple(); dosimplify = true)
+    function LagrangianSystem(L, t, x, v, params = NamedTuple(); simplify = true, scalarize = true)
 
         @assert eachindex(x) == eachindex(v)
 
@@ -29,22 +29,25 @@ struct LagrangianSystem
         ẋ  = collect(Dt.(x))
         ṗ  = collect(Dt.(p))
 
-        EL = [expand_derivatives(Dx[i](L) - Dt(Dv[i](L))) for i in eachindex(Dx,Dv)]
-        f  = [expand_derivatives(dx(L)) for dx in Dx]
-        g  = [expand_derivatives(Dt(dv(L))) for dv in Dv]
-        ϑ  = [expand_derivatives(dv(L)) for dv in Dv]
-        θ  = [expand_derivatives(dz(L)) for dz in Dz]
+        Ls = scalarize ? Symbolics.scalarize(L) : L
+        Ls = simplify ? Symbolics.simplify(Ls) : Ls
+
+        f  = [expand_derivatives(dx(Ls)) for dx in Dx]
+        g  = [expand_derivatives(Dt(dv(Ls))) for dv in Dv]
+        ϑ  = [expand_derivatives(dv(Ls)) for dv in Dv]
+        θ  = [expand_derivatives(dz(Ls)) for dz in Dz]
+        EL = [f[i] - g[i] for i in eachindex(f,g)]
 
         Ω  = [Dx[i](ϑ[j]) - Dx[j](ϑ[i]) for i in eachindex(Dx,ϑ), j in eachindex(Dx,ϑ)]
         ω  = [Dz[i](θ[j]) - Dz[j](θ[i]) for i in eachindex(Dz,θ), j in eachindex(Dz,θ)]
         M  = [Dv[i](ϑ[j]) for i in eachindex(Dv), j in eachindex(ϑ)]
         N  = [Dx[i](ϑ[j]) for i in eachindex(Dv), j in eachindex(ϑ)]
 
-        if dosimplify
-            Ω = simplify.(Ω)
-            ω = simplify.(ω)
-            M = simplify.(M)
-            N = simplify.(N)
+        if simplify
+            Ω = Symbolics.simplify.(Ω)
+            ω = Symbolics.simplify.(ω)
+            M = Symbolics.simplify.(M)
+            N = Symbolics.simplify.(N)
         end
 
         Ω = expand_derivatives.(Ω)
@@ -53,7 +56,7 @@ struct LagrangianSystem
         N = expand_derivatives.(N)
 
         equs = (
-            L = L,
+            L = Ls,
             EL = EL,
             f = f,
             g = g,
@@ -99,7 +102,7 @@ struct LagrangianSystem
             # P  = substitute_parameters(build_function(equs_subs.Σ,  t, X, V, params...; nanmath = false)[2], params),
         )
 
-        new(L, t, x, v, params, equs, generate_code(code))
+        new(Ls, t, x, v, params, equs, generate_code(code))
     end
 end
 

diff --git a/src/lagrangian_degenerate.jl b/src/lagrangian_degenerate.jl
@@ -10,7 +10,7 @@ struct DegenerateLagrangianSystem
     equations
     functions
 
-    function DegenerateLagrangianSystem(K, H, t, x, v, params = NamedTuple())
+    function DegenerateLagrangianSystem(K, H, t, x, v, params = NamedTuple(); simplify = true, scalarize = true)
 
         @assert eachindex(x) == eachindex(v)
 
@@ -27,21 +27,27 @@ struct DegenerateLagrangianSystem
 
         ẋ  = collect(Dt.(x))
 
-        L  = K - H
-        EL = [expand_derivatives(Dx[i](L) - Dt(Dv[i](L))) for i in eachindex(Dx,Dv)]
-        ∇H = [expand_derivatives(dx(H)) for dx in Dx]
-        ϑ  = [expand_derivatives(dv(K)) for dv in Dv]
-        f  = [expand_derivatives(dx(L)) for dx in Dx]
+        Ks = scalarize ? Symbolics.scalarize(K) : K
+        Hs = scalarize ? Symbolics.scalarize(H) : H
+
+        Ks = simplify ? Symbolics.simplify(Ks) : Ks
+        Hs = simplify ? Symbolics.simplify(Hs) : Hs
+
+        Ls = Ks - Hs
+        ∇H = [expand_derivatives(dx(Hs)) for dx in Dx]
+        ϑ  = [expand_derivatives(dv(Ls)) for dv in Dv]
+        f  = [expand_derivatives(dx(Ls)) for dx in Dx]
+        g  = [expand_derivatives(Dt(θ)) for θ in ϑ]
+        ḡ  = [expand_derivatives(dx(Ks)) for dx in Dx]
+        ω  = [expand_derivatives(Symbolics.simplify(Dx[i](ϑ[j]) - Dx[j](ϑ[i]))) for i in eachindex(Dx,ϑ), j in eachindex(Dx,ϑ)]
+        N  = [expand_derivatives(Symbolics.simplify(Dx[i](ϑ[j]))) for i in eachindex(Dv), j in eachindex(ϑ)]
         u  = [u for u in ẋ]
-        g  = [expand_derivatives(Dt.(ϑ))]
         ū  = [u for u in ẋ]
-        ḡ  = [expand_derivatives(dx(K)) for dx in Dx]
-        ω  = [expand_derivatives(simplify(Dx[i](ϑ[j]) - Dx[j](ϑ[i]))) for i in eachindex(Dx,ϑ), j in eachindex(Dx,ϑ)]
-        N  = [expand_derivatives(simplify(Dx[i](ϑ[j]))) for i in eachindex(Dv), j in eachindex(ϑ)]
+        EL = [f[i] - g[i] for i in eachindex(f,g)]
 
         equs = (
-            L = L,
-            H = H,
+            L = Ls,
+            H = Hs,
             EL = EL,
             ∇H = ∇H,
             f = f,
@@ -55,14 +61,19 @@ struct DegenerateLagrangianSystem
         )
 
         equs_subs = substitute_lagrangian_variables(equs, x, ẋ, v)
+
+        σ = inv(equs_subs.ω)
+
         equs_subs = merge(equs_subs, (
             ϕ = P .- equs_subs.ϑ,
             ψ = F .- equs_subs.g,
-            σ = simplify.(inv(equs_subs.ω)),
+            σ = simplify ? Symbolics.simplify.(σ) : σ,
         ))
 
+        ẋeq = equs_subs.σ * equs_subs.∇H
+
         equs_subs = merge(equs_subs, (
-            ẋ = equs_subs.σ * equs_subs.∇H,
+            ẋ = simplify ? Symbolics.simplify.(ẋeq) : ẋeq,
         ))
 
         # equs = substitute_v_with_ẋ(equs, v, ẋ)
@@ -91,7 +102,7 @@ struct DegenerateLagrangianSystem
             P  = substitute_parameters(build_function(equs_subs.σ,  t, X, V, params...; nanmath = false)[2], params),
         )
 
-        new(L, t, x, v, params, equs, generate_code(code))
+        new(Ls, t, x, v, params, equs, generate_code(code))
     end
 end
 

diff --git a/test/hamiltonian_particle.jl b/test/hamiltonian_particle.jl
@@ -1,6 +1,7 @@
 using EulerLagrange
 using GeometricEquations
 using LinearAlgebra
+using Symbolics
 using Test
 
 
@@ -14,7 +15,7 @@ t, q, p = hamiltonian_variables(2)
 sym_ham = H(t, q, p, params)
 ham_sys = HamiltonianSystem(sym_ham, t, q, p)
 
-@test isequal(hamiltonian(ham_sys), sym_ham)
+@test isequal(hamiltonian(ham_sys), Symbolics.simplify(sym_ham))
 @test isequal(variables(ham_sys), (t, q, p))
 @test isequal(EulerLagrange.parameters(ham_sys), NamedTuple())
 

diff --git a/test/lagrangian_solar_system.jl b/test/lagrangian_solar_system.jl
@@ -139,7 +139,7 @@ function test_solar_system(ss)
     sparams = symbolize(params)
 
     # LagrangianSystem
-    lag_sys = LagrangianSystem(lagrangian(t, x, v, sparams, d, n), t, x, v, sparams; dosimplify = false)
+    lag_sys = LagrangianSystem(lagrangian(t, x, v, sparams, d, n), t, x, v, sparams; simplify = false)
 
     p₁, p₂ = zero(p₀), zero(p₀)
     ṗ₁, ṗ₂ = zero(p₀), zero(p₀)