From 58bc9624e234e086d567acc219e2a26839b81199 Mon Sep 17 00:00:00 2001
From: Fredrik Bagge Carlson <baggepinnen@gmail.com>
Date: Thu, 20 Jun 2024 09:16:23 +0200
Subject: [PATCH] switch quaternion representation and add tests (#81)

* switch quaternion representation and add tests

* move testset

* docs improvement and clean up
---
 docs/make.jl                       |   2 +-
 docs/src/examples/free_motion.md   |  10 +-
 docs/src/examples/space.md         |   2 +
 docs/src/examples/three_springs.md |  30 ++-
 src/components.jl                  |  16 +-
 src/joints.jl                      |  13 +-
 src/orientation.jl                 |  63 +++---
 test/runtests.jl                   | 299 +-------------------------
 test/test_quaternions.jl           | 323 +++++++++++++++++++++++++++++
 9 files changed, 421 insertions(+), 337 deletions(-)
 create mode 100644 test/test_quaternions.jl

diff --git a/docs/make.jl b/docs/make.jl
index 7af9ffd4..5f3b6491 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -10,7 +10,7 @@ makedocs(;
          authors = "JuliaHub Inc.",
          #  strict = [:example_block, :setup_block, :eval_block],
          sitename = "Multibody.jl",
-         warnonly = [:missing_docs, :cross_references, :example_block, :docs_block],
+         warnonly = [:missing_docs, :cross_references, :docs_block],
          pagesonly = true,
          format = Documenter.HTML(;
                                   prettyurls = get(ENV, "CI", nothing) == "true",
diff --git a/docs/src/examples/free_motion.md b/docs/src/examples/free_motion.md
index ec160c23..a743524c 100644
--- a/docs/src/examples/free_motion.md
+++ b/docs/src/examples/free_motion.md
@@ -49,9 +49,11 @@ nothing # hide
 ```
 
 
-If we instead model a body suspended in springs without the presence of any joint, we need to give the body state variables. We do this by saying `isroot = true` when we create the body, we also use quaternions to represent angles using `quat = true`.
+## Body suspended in springs
+If we instead model a body suspended in springs without the presence of any joint, we need to _give the body state variables_. We do this by saying `isroot = true` when we create the body, we also use quaternions to represent angular state using `quat = true`.
 
 ```@example FREE_MOTION
+using Multibody.Rotations: QuatRotation, RotXYZ
 @named begin
     body = BodyShape(m = 1, I_11 = 1, I_22 = 1, I_33 = 1, r = [0.4, 0, 0],
                      r_0 = [0.2, -0.5, 0.1], isroot = true, quat=true)
@@ -78,14 +80,16 @@ ssys = structural_simplify(IRSystem(model))
 prob = ODEProblem(ssys, [
     collect(body.body.v_0 .=> 0);
     collect(body.body.w_a .=> 0);
-    collect(body.body.Q .=> Multibody.to_q([0.0940609, 0.0789265, 0.0940609, 0.987965]));
-    collect(body.body.Q̂ .=> Multibody.to_q([0.0940609, 0.0789265, 0.0940609, 0.987965]));
+    collect(body.body.Q) .=> vec(QuatRotation(RotXYZ(deg2rad.((10,10,10))...)));
+    collect(body.body.Q̂) .=> vec(QuatRotation(RotXYZ(deg2rad.((10,10,10))...)));
 ], (0, 4))
 
 sol = solve(prob, Rodas5P())
 @assert SciMLBase.successful_retcode(sol)
 
 plot(sol, idxs = [body.r_0...])
+# plot(sol, idxs = [body.body.Q...])
+
 ```
 
 ```@example FREE_MOTION
diff --git a/docs/src/examples/space.md b/docs/src/examples/space.md
index 25d942b7..8cb54ca3 100644
--- a/docs/src/examples/space.md
+++ b/docs/src/examples/space.md
@@ -35,6 +35,7 @@ W(;kwargs...) = Multibody.world
             I_33=0.1,
             r_0=[0,0.6,0],
             isroot=true,
+            # quat=true, # Activate to use quaternions as state instead of Euler angles
             v_0=[1,0,0])
         body2 = Body(
             m=1,
@@ -43,6 +44,7 @@ W(;kwargs...) = Multibody.world
             I_33=0.1,
             r_0=[0.6,0.6,0],
             isroot=true,
+            # quat=true, # Activate to use quaternions as state instead of Euler angles
             v_0=[0.6,0,0])
     end
 end
diff --git a/docs/src/examples/three_springs.md b/docs/src/examples/three_springs.md
index 664e7c6d..fcada9ec 100644
--- a/docs/src/examples/three_springs.md
+++ b/docs/src/examples/three_springs.md
@@ -20,13 +20,13 @@ world = Multibody.world
 
 systems = @named begin
     body1 = Body(m = 0.8, I_11 = 0.1, I_22 = 0.1, I_33 = 0.1, r_0 = [0.5, -0.3, 0],
-                 r_cm = [0, 0, 0], isroot=false, quat=false)
+                 r_cm = [0, -0.2, 0], isroot=true, quat=true)
     bar1 = FixedTranslation(r = [0.3, 0, 0])
     bar2 = FixedTranslation(r = [0, 0, 0.3])
-    spring1 = Multibody.Spring(c = 20, m = 0, s_unstretched = 0.1,
+    spring1 = Multibody.Spring(c = 20, m = 0, s_unstretched = 0.1, fixed_rotation_at_frame_b=true,
                                r_rel_0 = [-0.2, -0.2, 0.2])
-    spring2 = Multibody.Spring(c = 40, m = 0, s_unstretched = 0.1, fixed_rotation_at_frame_b=true, fixed_rotation_at_frame_a=true)
-    spring3 = Multibody.Spring(c = 20, m = 0, s_unstretched = 0.1)
+    spring2 = Multibody.Spring(c = 40, m = 0, s_unstretched = 0.1)
+    spring3 = Multibody.Spring(c = 20, m = 0, s_unstretched = 0.1, fixed_rotation_at_frame_a=false)
 end
 eqs = [connect(world.frame_b, bar1.frame_a)
        connect(world.frame_b, bar2.frame_a)
@@ -40,11 +40,29 @@ eqs = [connect(world.frame_b, bar1.frame_a)
 ssys = structural_simplify(IRSystem(model))
 prob = ODEProblem(ssys, [], (0, 10))
 
-sol = solve(prob, FBDF(), u0=prob.u0 .+ 1e-9*randn(length(prob.u0)))
+sol = solve(prob, FBDF(), u0=prob.u0 .+ 1e-12*randn(length(prob.u0)))
 @assert SciMLBase.successful_retcode(sol)
 
 Plots.plot(sol, idxs = [body1.r_0...])
-# Plots.plot(sol, idxs = ori(body1.frame_a).R[1])
+```
+
+```@setup
+# States selected by Dymola
+# Statically selected continuous time states
+# body1.frame_a.r_0[1]
+# body1.frame_a.r_0[2]
+# body1.frame_a.r_0[3]
+# body1.v_0[1]
+# body1.v_0[2]
+# body1.v_0[3]
+# body1.w_a[1]
+# body1.w_a[2]
+# body1.w_a[3]
+
+# Dynamically selected continuous time states
+# There is one set of dynamic state selection.
+# There are 3 states to be selected from:
+# body1.Q[]
 ```
 
 ## 3D animation
diff --git a/src/components.jl b/src/components.jl
index 8324e880..d85436fb 100644
--- a/src/components.jl
+++ b/src/components.jl
@@ -230,6 +230,8 @@ Representing a body with 3 translational and 3 rotational degrees-of-freedom.
               I_31 = 0,
               I_32 = 0,
               isroot = false,
+              state = false,
+              vel_from_R = false,
               phi0 = zeros(3),
               phid0 = zeros(3),
               r_0 = 0,
@@ -283,13 +285,18 @@ Representing a body with 3 translational and 3 rotational degrees-of-freedom.
 
     # DRa = D(Ra)
 
+    if state
+        # @warn "Make the body have state variables by using isroot=true rather than state=true"
+        isroot = true
+    end
+
     dvs = [r_0;v_0;a_0;g_0;w_a;z_a;]
     eqs = if isroot # isRoot
         
         if quat
             @named frame_a = Frame(varw = false)
             Ra = ori(frame_a, false)
-            nonunit_quaternion_equations(Ra, w_a);
+            qeeqs = nonunit_quaternion_equations(Ra, w_a)
         else
             @named frame_a = Frame(varw = true)
             Ra = ori(frame_a, true)
@@ -303,8 +310,11 @@ Representing a body with 3 translational and 3 rotational degrees-of-freedom.
                     phidd .~ D.(phid)
                     Ra ~ ar
                     Ra.w .~ w_a
-                    # w_a .~ angular_velocity2(ori(frame_a, false)) # This is required for FreeBody tests to pass, but the other one required for harmonic osciallator without joint to pass. FreeBody passes with quat=true so we use that instead
-                    w_a .~ ar.w # This one for most systems
+                    if vel_from_R
+                        w_a .~ angular_velocity2(ori(frame_a, false)) # This is required for FreeBody and ThreeSprings tests to pass, but the other one required for harmonic osciallator without joint to pass. FreeBody passes with quat=true so we use that instead
+                    else
+                        w_a .~ ar.w # This one for most systems
+                    end
                     ]
         end
     else
diff --git a/src/joints.jl b/src/joints.jl
index 670e39e1..c19bc58c 100644
--- a/src/joints.jl
+++ b/src/joints.jl
@@ -724,8 +724,10 @@ The relative position vector `r_rel_a` from the origin of `frame_a` to the origi
         (a_rel_a(t)[1:3] = a_rel_a), [description = "= der(v_rel_a)"]
     end
 
-    @named R_rel = NumRotationMatrix()
-    @named R_rel_inv = NumRotationMatrix()
+    @named R_rel_f = Frame()
+    @named R_rel_inv_f = Frame()
+    R_rel = ori(R_rel_f)
+    R_rel_inv = ori(R_rel_inv_f)
 
     eqs = [
            # Cut-forces and cut-torques are zero
@@ -750,7 +752,6 @@ The relative position vector `r_rel_a` from the origin of `frame_a` to the origi
         end
 
         if quat
-            # @named Q = NumQuaternion() 
             append!(eqs, nonunit_quaternion_equations(R_rel, w_rel_b))
 
         else
@@ -769,7 +770,11 @@ The relative position vector `r_rel_a` from the origin of `frame_a` to the origi
                                         R, angular_velocity1(frame_a)))
         end
     end
-    compose(ODESystem(eqs, t; name), frame_a, frame_b)
+    if state && !isroot
+        compose(ODESystem(eqs, t; name), frame_a, frame_b, R_rel_f, R_rel_inv_f)
+    else
+        compose(ODESystem(eqs, t; name), frame_a, frame_b, R_rel_f, )
+    end
 end
 
 """
diff --git a/src/orientation.jl b/src/orientation.jl
index 3a773cd1..8c842f2f 100644
--- a/src/orientation.jl
+++ b/src/orientation.jl
@@ -36,8 +36,11 @@ Create a new [`RotationMatrix`](@ref) struct with symbolic elements. `R,w` deter
 The primary difference between `NumRotationMatrix` and `RotationMatrix` is that the `NumRotationMatrix` constructor is used in the constructor of a [`Frame`](@ref) in order to introduce the frame variables, whereas `RorationMatrix` (the struct) only wraps existing variables.
 
 - `varw`: If true, `w` is a variable, otherwise it is derived from the derivative of `R` as `w = get_w(R)`.
+
+Never call this function directly from a component constructor, instead call `f = Frame(); R = ori(f)` and add `f` to the subsystems.
 """
-function NumRotationMatrix(; R = collect(1.0I(3)), w = zeros(3), name, varw = false, state_priority=nothing)
+function NumRotationMatrix(; R = collect(1.0I(3)), w = zeros(3), name=:R, varw = false, state_priority=nothing)
+    # The reason for not calling this directly is that all R vaiables have to have the same name since they are treated as connector variables (otherwise a connection error is thrown). A component with more than one rotation matrix will thus have two different R variables that overwrite each other
     R = at_variables_t(:R, 1:3, 1:3; default = R, state_priority) #[description="Orientation rotation matrix ∈ SO(3)"]
     # @variables w(t)[1:3]=w [description="angular velocity"]
     # R = collect(R)
@@ -251,9 +254,13 @@ end
 
 Base.getindex(Q::Quaternion, i) = Q.Q[i]
 
-function NumQuaternion(; Q = [0.0, 0, 0, 1.0], w = zeros(3), name, varw = false)
+"""
+Never call this function directly from a component constructor, instead call `f = Frame(); R = ori(f)` and add `f` to the subsystems.
+"""
+function NumQuaternion(; Q = [1.0, 0, 0, 0.0], w = zeros(3), name, varw = false)
+    # The reason for not calling this directly is that all R vaiables have to have the same name since they are treated as connector variables (otherwise a connection error is thrown). A component with more than one rotation matrix will thus have two different R variables that overwrite each other
     # Q = at_variables_t(:Q, 1:4, default = Q) #[description="Orientation rotation matrix ∈ SO(3)"]
-    @variables Q(t)[1:4] = [0.0,0,0,1.0]
+    @variables Q(t)[1:4] = [1.0,0,0,0]
     if varw
         @variables w(t)[1:3]=w [description="angular velocity"]
         # w = at_variables_t(:w, 1:3, default = w)
@@ -274,23 +281,26 @@ orientation_constraint(q::Quaternion) = orientation_constraint(q.Q)
 # end
 
 Base.:/(q::Rotations.Quaternions.Quaternion, x::Num) = Rotations.Quaternions.Quaternion(q.s / x, q.v1 / x, q.v2 / x, q.v3 / x)
-function from_Q(Q, w)
-    Q2 = to_q(Q) # Due to different conventions
+function from_Q(Q2, w)
+    # Q2 = to_q(Q) # Due to different conventions
     q = Rotations.QuatRotation(Q2)
     R = RotMatrix(q)
     RotationMatrix(R, w)
 end
 
-to_q(Q) = SA[Q[4], Q[1], Q[2], Q[3]]
-to_mb(Q) = SA[Q[2], Q[3], Q[4], Q[1]]
+to_q(Q::AbstractVector) = SA[Q[4], Q[1], Q[2], Q[3]]
+to_q(Q::Rotations.QuatRotation) = to_q(vec(Q))
+to_mb(Q::AbstractVector) = SA[Q[2], Q[3], Q[4], Q[1]]
+to_mb(Q::Rotations.QuatRotation) = to_mb(vec(Q))
+Base.vec(Q::Rotations.QuatRotation) = SA[Q.q.s, Q.q.v1, Q.q.v2, Q.q.v3]
 
-function angular_velocity1(Q, der_Q)
-    2*([Q[4] -Q[3] Q[2] -Q[1]; Q[3] Q[4] -Q[1] -Q[2]; -Q[2] Q[1] Q[4] -Q[3]]*der_Q)
-end
+# function angular_velocity1(Q, der_Q)
+#     2*([Q[4] -Q[3] Q[2] -Q[1]; Q[3] Q[4] -Q[1] -Q[2]; -Q[2] Q[1] Q[4] -Q[3]]*der_Q)
+# end
 
-function angular_velocity2(Q, der_Q)
-    2*([Q[4]  Q[3] -Q[2] -Q[1]; -Q[3] Q[4] Q[1] -Q[2]; Q[2] -Q[1] Q[4] -Q[3]]*der_Q)
-end
+# function angular_velocity2(Q, der_Q)
+#     2*([Q[4]  Q[3] -Q[2] -Q[1]; -Q[3] Q[4] Q[1] -Q[2]; Q[2] -Q[1] Q[4] -Q[3]]*der_Q)
+# end
 
 
 ## Euler
@@ -446,6 +456,8 @@ The rotation matrix returned, ``R_W^F``, is such that when a vector ``p_F`` expr
 p_W = R_W^F  p_F
 ```
 
+The columns of ``R_W_F`` indicate are the basis vectors of the frame ``F`` expressed in the world coordinate frame.
+
 See also [`get_trans`](@ref), [`get_frame`](@ref), [Orientations and directions](@ref) (docs section).
 """
 function get_rot(sol, frame, t)
@@ -481,26 +493,25 @@ function get_frame(sol, frame, t)
 end
 
 function nonunit_quaternion_equations(R, w)
-    @variables Q(t)[1:4]=[0,0,0,1], [description="Unit quaternion with [i,j,k,w]"] # normalized
-    @variables Q̂(t)[1:4]=[0,0,0,1], [description="Non-unit quaternion with [i,j,k,w]"] # Non-normalized
+    @variables Q(t)[1:4]=[1,0,0,0], [description="Unit quaternion with [w,i,j,k]"] # normalized
+    @variables Q̂(t)[1:4]=[1,0,0,0], [description="Non-unit quaternion with [w,i,j,k]"] # Non-normalized
     @variables n(t)=1 c(t)=0
     @parameters k = 0.1
     Q̂ = collect(Q̂)
+    Q = collect(Q)
+    # w is used in Ω, and Ω determines D(Q̂)
+    # This corresponds to modelica's 
+    # frame_a.R = from_Q(Q, angularVelocity2(Q, der(Q)));
+    # where angularVelocity2(Q, der(Q)) = 2*([Q[4]  Q[3] -Q[2] -Q[1]; -Q[3] Q[4] Q[1] -Q[2]; Q[2] -Q[1] Q[4] -Q[3]]*der_Q)
+    # They also have w_a = angularVelocity2(frame_a.R) even for quaternions, so w_a = angularVelocity2(Q, der(Q)), this is their link between w_a and D(Q), while ours is D(Q̂) .~ (Ω * Q̂)
     Ω = [0 -w[1] -w[2] -w[3]; w[1] 0 w[3] -w[2]; w[2] -w[3] 0 w[1]; w[3] w[2] -w[1] 0]
-    P = [
-        0 0 0 1
-        1 0 0 0
-        0 1 0 0
-        0 0 1 0
-    ]
-    Ω = P'Ω*P
+    # QR = from_Q(Q, angular_velocity2(Q, D.(Q)))
+    QR = from_Q(Q, w)
     [
         n ~ Q̂'Q̂
-        # n ~ _norm(Q̂)^2
         c ~ k * (1 - n)
-        D.(Q̂) .~ (Ω * Q̂) ./ 2 + c * Q̂
+        D.(Q̂) .~ (Ω' * Q̂) ./ 2 + c * Q̂ # We use Ω' which is the same as using -w to handle the fact that w coming in here is typically described frame_a rather than in frame_b, the paper is formulated with w being expressed in the rotating body frame (frame_b)
         Q .~ Q̂ ./ sqrt(n)
-        R ~ from_Q(Q, w)
-        # R.w ~ w
+        R ~ QR
     ]
 end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 1add1630..ac7ef177 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -40,6 +40,11 @@ end
     include("test_orientation_getters.jl")
 end
 
+@testset "quaternions" begin
+    @info "Testing quaternions"
+    include("test_quaternions.jl")
+end
+
 
 # ==============================================================================
 ## Add spring to make a harmonic oscillator ====================================
@@ -508,71 +513,7 @@ doplot() && plot(sol, idxs = [body1.r_0...]) |> display
 # end
 # fixed_rotation_at_frame_a and b = true required
 end
-# ==============================================================================
-## FreeBody ====================================================================
-# ==============================================================================
-# https://doc.modelica.org/om/Modelica.Mechanics.MultiBody.Examples.Elementary.FreeBody.html
-using Multibody
-using ModelingToolkit
-# using Plots
-using JuliaSimCompiler
-using OrdinaryDiffEq
 
-@testset "FreeBody" begin
-t = Multibody.t
-D = Differential(t)
-world = Multibody.world
-
-@named begin
-    body = BodyShape(m = 1, I_11 = 1, I_22 = 1, I_33 = 1, r = [0.4, 0, 0],
-                     r_0 = [0.2, -0.5, 0.1], r_cm = [0.2, 0, 0], isroot = true, quat=true)
-    bar2 = FixedTranslation(r = [0.8, 0, 0])
-    spring1 = Multibody.Spring(c = 20, s_unstretched = 0)
-    spring2 = Multibody.Spring(c = 20, s_unstretched = 0)
-end
-
-eqs = [connect(bar2.frame_a, world.frame_b)
-       connect(spring1.frame_b, body.frame_a)
-       connect(bar2.frame_b, spring2.frame_a)
-       connect(spring1.frame_a, world.frame_b)
-       connect(body.frame_b, spring2.frame_b)]
-
-@named model = ODESystem(eqs, t,
-                         systems = [
-                             world,
-                             body,
-                             bar2,
-                             spring1,
-                             spring2,
-                         ])
-ssys = structural_simplify(IRSystem(model))#, alias_eliminate = true)
-# ssys = structural_simplify(model, allow_parameters = false)
-prob = ODEProblem(ssys,
-                  [world.g => 9.80665;
-                    collect(body.body.w_a .=> 0);
-                  collect(body.body.v_0 .=> 0);
-                #   collect(body.body.phi .=> deg2rad(10));
-                    collect(body.body.Q) .=> Multibody.to_mb(Rotations.QuatRotation(Rotations.RotXYZ(deg2rad.((10,10,10))...)));
-                    collect(body.body.Q̂) .=> Multibody.to_mb(Rotations.QuatRotation(Rotations.RotXYZ(deg2rad.((10,10,10))...)));
-                  ], (0, 10))
-
-# @test_skip begin # The modelica example uses angles_fixed = true, which causes the body component to run special code for variable initialization. This is not yet supported by MTK
-sol = solve(prob, Rodas5P(), abstol=1e-6, reltol=1e-6)
-@test SciMLBase.successful_retcode(sol)
-
-@test sol(10, idxs=[
-    body.r_0;
-]) ≈ [
-    0.196097
-    -0.46243
-    0.080562
-] atol=1e-1
-
-@test_broken get_rot(sol, body.frame_b, 10)[1,2] ≈ 0.104409 atol=0.01
-
-doplot() && plot(sol, idxs = [body.r_0...; body.body.w_a; body.body.v_0], layout=(3,3), size=(1000, 1000)) |> display
-
-end
 # ==============================================================================
 ## Sperical-joint pendulum ===================================================
 # ==============================================================================
@@ -860,44 +801,6 @@ doplot() && plot(sol, idxs = collect(body.body.phi), title = "Dzhanibekov effect
 end
 
 
-# ==============================================================================
-## Harmonic oscillator with Body as root and quaternions as state variables
-# ==============================================================================
-using LinearAlgebra
-@testset "Harmonic oscillator with Body as root and quaternions as state variables" begin
-
-@named body = Body(; m = 1, isroot = true, r_cm = [0.0, 0, 0], phi0 = [0, 0.9, 0], quat=true) # This time the body isroot since there is no joint containing state
-@named spring = Multibody.Spring(c = 1)
-
-connections = [connect(world.frame_b, spring.frame_a)
-               connect(spring.frame_b, body.frame_a)]
-
-@named model = ODESystem(connections, t, systems = [world, spring, body])
-model = complete(model)
-# ssys = structural_simplify(model, allow_parameter = false)
-
-irsys = IRSystem(model)
-ssys = structural_simplify(irsys)
-@test length(unknowns(ssys)) == 13 # One extra due to quaternions
-D = Differential(t)
-
-# du = prob.f.f.f_oop(prob.u0, prob.p, 0)
-# @test all(isfinite, du)
-
-# prob = ODEProblem(ssys, ModelingToolkit.missing_variable_defaults(ssys), (0, 10))
-prob = ODEProblem(ssys, [collect(body.v_0 .=> [0, 0, 0]); collect(body.w_a .=> [0, 0, 0]); ], (0, 10))
-sol = solve(prob, Rodas5P(), u0 = prob.u0 .+ 1e-12 .* randn.())
-
-doplot() &&
-    plot(sol, idxs = [collect(body.r_0); collect(body.v_0)], layout = 6) |> display
-
-@test sol(2pi, idxs = body.r_0[1])≈0 atol=1e-3
-@test sol(2pi, idxs = body.r_0[2])≈0 atol=1e-3
-@test sol(2pi, idxs = body.r_0[3])≈0 atol=1e-3
-@test sol(pi, idxs = body.r_0[2]) < -2
-
-end
-
 ## Actuated joint
 using Multibody
 using ModelingToolkit
@@ -1085,198 +988,6 @@ tt = 0:0.1:10
     # Multibody.render(mounted_chain, sol, x=3, filename = "mounted_chain.gif") # May take long time for n>=10
 end
 
-# ==============================================================================
-## Simple motion with quaternions=============================================================
-# ==============================================================================
-using LinearAlgebra, ModelingToolkit, Multibody, JuliaSimCompiler
-t = Multibody.t
-@named joint = Multibody.FreeMotion(isroot = true, state=false, quat=true)
-@named body = Body(; m = 1, r_cm = [0.0, 0, 0], isroot=true, quat=true, w_a=[1,0.5,0.2])
-
-# @named joint = Multibody.FreeMotion(isroot = true, state=true, quat=true)
-# @named body = Body(; m = 1, r_cm = [0.0, 0, 0], isroot=false, w_a=[1,1,1])
-
-world = Multibody.world
-
-
-connections = [connect(world.frame_b, joint.frame_a)
-               connect(joint.frame_b, body.frame_a)]
-
-
-@named model = ODESystem(connections, t,
-                         systems = [world, joint, body])
-irsys = IRSystem(model)
-ssys = structural_simplify(irsys)
-
-
-
-##
-D = Differential(t)
-prob = ODEProblem(ssys, [collect(body.w_a) .=> [1,0,0];], (0, 2pi))
-
-using OrdinaryDiffEq, Test
-sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 0 .* randn.())#, initializealg=ShampineCollocationInit(0.01))
-# sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 1e-12 .* randn.(), dtmin=1e-8, force_dtmin=true)
-@test SciMLBase.successful_retcode(sol)
-doplot() && plot(sol, layout=13)
-# end
-
-ts = 0:0.1:2pi
-Q = Matrix(sol(ts, idxs = [body.Q...]))
-Qh = Matrix(sol(ts, idxs = [body.Q̂...]))
-n = Matrix(sol(ts, idxs = [body.n...]))
-@test mapslices(norm, Qh, dims=1).^2 ≈ n
-@test Q ≈ Qh ./ sqrt.(n) atol=1e-2
-@test norm(mapslices(norm, Q, dims=1) .- 1) < 1e-2
-
-
-
-function get_R(frame, t)
-    reshape(sol(t, idxs=vec(ori(frame).R.mat)), 3, 3)
-end
-function get_r(frame, t)
-    sol(t, idxs=collect(frame.r_0))
-end
-
-import Multibody.Rotations.QuatRotation as Quat
-import Multibody.Rotations
-ti = 5
-for (i, ti) in enumerate(ts)
-    R = get_R(body.frame_a, ti)
-    @test norm(R'R - I) < 1e-10
-    @test Multibody.from_Q(Q[:, i], 0).R ≈ R atol=1e-3
-end
-
-
-# Test that rotational velocity of 1 results in one full rotation in 2π seconds. Test rotation around all three major axes
-@test get_R(body.frame_a, 0pi) ≈ I
-@test get_R(body.frame_a, pi/2) ≈ [1 0 0; 0 0 -1; 0 1 0] atol=1e-3
-@test get_R(body.frame_a, 1pi) ≈ diagm([1, -1, -1]) atol=1e-3
-@test get_R(body.frame_a, 2pi) ≈ I atol=1e-3
-
-prob = ODEProblem(ssys, [collect(body.w_a) .=> [0,1,0];], (0, 2pi))
-sol = solve(prob, Rodas4())#, 
-@test get_R(body.frame_a, 0pi) ≈ I
-@test get_R(body.frame_a, 1pi) ≈ diagm([-1, 1, -1]) atol=1e-3
-@test get_R(body.frame_a, 2pi) ≈ I atol=1e-3
-
-
-prob = ODEProblem(ssys, [collect(body.w_a) .=> [0,0,1];], (0, 2pi))
-sol = solve(prob, Rodas4())#, 
-@test get_R(body.frame_a, 0pi) ≈ I
-@test get_R(body.frame_a, 1pi) ≈ diagm([-1, -1, 1]) atol=1e-3
-@test get_R(body.frame_a, 2pi) ≈ I atol=1e-3
-
-
-
-## states in joint instead of body
-
-using LinearAlgebra, ModelingToolkit, Multibody, JuliaSimCompiler
-t = Multibody.t
-world = Multibody.world
-
-@named joint = Multibody.FreeMotion(isroot = true, state=true, quat=true)
-@named body = Body(; m = 1, r_cm = [0.0, 0, 0])
-
-connections = [connect(world.frame_b, joint.frame_a)
-               connect(joint.frame_b, body.frame_a)]
-
-@named model = ODESystem(connections, t,
-                         systems = [world, joint, body])
-irsys = IRSystem(model)
-ssys = structural_simplify(irsys)
-
-##
-D = Differential(t)
-# q0 = randn(4); q0 ./= norm(q0)
-q0 = [0,0,0,1]
-prob = ODEProblem(ssys, [
-    collect(body.w_a) .=> [1,1,1];
-    collect(joint.Q) .=> q0;
-    collect(joint.Q̂) .=> q0;
-    ], (0, 2pi))
-
-using OrdinaryDiffEq, Test
-sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 1e-6 .* randn.())
-@test SciMLBase.successful_retcode(sol)
-# doplot() && plot(sol, layout=21)
-
-
-ts = 0:0.1:2pi
-Q = Matrix(sol(ts, idxs = [joint.Q...]))
-Qh = Matrix(sol(ts, idxs = [joint.Q̂...]))
-n = Matrix(sol(ts, idxs = [joint.n...]))
-@test mapslices(norm, Qh, dims=1).^2 ≈ n
-@test Q ≈ Qh ./ sqrt.(n) atol=1e-2
-@test norm(mapslices(norm, Q, dims=1) .- 1) < 1e-2
-
-@test get_R(joint.frame_b, 0pi) ≈ I
-@test_broken get_R(joint.frame_b, 1pi) ≈ diagm([1, -1, -1]) atol=1e-3
-@test get_R(joint.frame_b, 2pi) ≈ I atol=1e-3
-
-Matrix(sol(ts, idxs = [joint.w_rel_b...]))
-
-## Spherical joint pendulum with quaternions
-
-using LinearAlgebra, ModelingToolkit, Multibody, JuliaSimCompiler
-t = Multibody.t
-world = Multibody.world
-
-
-@named joint = Multibody.Spherical(isroot=false, state=false, quat=false)
-@named rod = FixedTranslation(; r = [1, 0, 0])
-@named body = Body(; m = 1, isroot=true, quat=true, air_resistance=0.0)
-
-# @named joint = Multibody.Spherical(isroot=true, state=true, quat=true)
-# @named body = Body(; m = 1, r_cm = [1.0, 0, 0], isroot=false)
-
-connections = [connect(world.frame_b, joint.frame_a)
-               connect(joint.frame_b, rod.frame_a)
-               connect(rod.frame_b, body.frame_a)]
-
-
-@named model = ODESystem(connections, t,
-                         systems = [world, joint, body, rod])
-irsys = IRSystem(model)
-@test_skip begin # https://github.com/JuliaComputing/JuliaSimCompiler.jl/issues/298
-    ssys = structural_simplify(irsys)
-
-
-    D = Differential(t)
-    q0 = randn(4); q0 ./= norm(q0)
-    # q0 = [0,0,0,1]
-    prob = ODEProblem(ssys, [
-        # collect(body.w_a) .=> [1,0,0];
-        # collect(body.Q) .=> q0;
-        # collect(body.Q̂) .=> q0;
-        ], (0, 30))
-
-    using OrdinaryDiffEq, Test
-    sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 0 .* randn.())
-    @test SciMLBase.successful_retcode(sol)
-    # doplot() && plot(sol, layout=21)
-
-
-    ts = 0:0.1:2pi
-    Q = Matrix(sol(ts, idxs = [body.Q...]))
-    Qh = Matrix(sol(ts, idxs = [body.Q̂...]))
-    n = Matrix(sol(ts, idxs = [body.n...]))
-    @test mapslices(norm, Qh, dims=1).^2 ≈ n
-    @test Q ≈ Qh ./ sqrt.(n) atol=1e-2
-    @test norm(mapslices(norm, Q, dims=1) .- 1) < 1e-2
-
-    Matrix(sol(ts, idxs = [body.w_a...]))
-
-    @test get_R(joint.frame_b, 0pi) ≈ I
-    @test get_R(joint.frame_b, sqrt(9.81/1)) ≈ diagm([1, -1, -1]) atol=1e-3
-    @test get_R(joint.frame_b, 2pi) ≈ I atol=1e-3
-
-
-    Matrix(sol(ts, idxs = [joint.w_rel_b...]))
-
-    # render(model, sol)
-end
-
 
 # ==============================================================================
 ## BodyCylinder ================================================================
diff --git a/test/test_quaternions.jl b/test/test_quaternions.jl
new file mode 100644
index 00000000..7e71188a
--- /dev/null
+++ b/test/test_quaternions.jl
@@ -0,0 +1,323 @@
+# test utils
+import Multibody.Rotations.QuatRotation as Quat
+import Multibody.Rotations
+function get_R(sol, frame, t)
+    reshape(sol(t, idxs=vec(ori(frame).R.mat)), 3, 3)
+end
+function get_r(sol, frame, t)
+    sol(t, idxs=collect(frame.r_0))
+end
+
+
+
+
+# ==============================================================================
+## Harmonic oscillator with Body as root and quaternions as state variables
+# ==============================================================================
+using LinearAlgebra
+@testset "Harmonic oscillator with Body as root and quaternions as state variables" begin
+
+@named body = Body(; m = 1, isroot = true, r_cm = [0.0, 0, 0], phi0 = [0, 0.9, 0], quat=true) # This time the body isroot since there is no joint containing state
+@named spring = Multibody.Spring(c = 1)
+
+connections = [connect(world.frame_b, spring.frame_a)
+               connect(spring.frame_b, body.frame_a)]
+
+@named model = ODESystem(connections, t, systems = [world, spring, body])
+model = complete(model)
+# ssys = structural_simplify(model, allow_parameter = false)
+
+irsys = IRSystem(model)
+ssys = structural_simplify(irsys)
+@test length(unknowns(ssys)) == 13 # One extra due to quaternions
+D = Differential(t)
+
+# du = prob.f.f.f_oop(prob.u0, prob.p, 0)
+# @test all(isfinite, du)
+
+# prob = ODEProblem(ssys, ModelingToolkit.missing_variable_defaults(ssys), (0, 10))
+prob = ODEProblem(ssys, [collect(body.v_0 .=> [0, 0, 0]); collect(body.w_a .=> [0, 1, 0]); ], (0, 10))
+sol = solve(prob, Rodas5P(), u0 = prob.u0 .+ 1e-12 .* randn.())
+
+doplot() &&
+    plot(sol, idxs = [collect(body.r_0); collect(body.v_0)], layout = 6) |> display
+
+@test sol(2pi, idxs = body.r_0[1])≈0 atol=1e-3
+@test sol(2pi, idxs = body.r_0[2])≈0 atol=1e-3
+@test sol(2pi, idxs = body.r_0[3])≈0 atol=1e-3
+@test sol(pi, idxs = body.r_0[2]) < -2
+
+# TODO: add more tests
+ts = 0:0.1:2pi
+
+ti = 5
+for (i, ti) in enumerate(ts)
+    R = get_R(sol, body.frame_a, ti)
+    Q = sol(ti, idxs=body.Q)
+    @test Multibody.from_Q(Q, 0).R ≈ R atol=1e-6 # Test that from_Q yields a rotation matrix consistent with R
+    R2 = get_R(sol, spring.frame_b, ti)
+    @test norm(R'R2 - I) < 1e-10
+end
+
+end
+
+
+
+
+# ==============================================================================
+## Simple motion with quaternions===============================================
+# ==============================================================================
+using LinearAlgebra, ModelingToolkit, Multibody, JuliaSimCompiler
+using OrdinaryDiffEq, Test
+
+@testset "Simple motion with quaternions and state in Body" begin
+
+t = Multibody.t
+@named joint = Multibody.FreeMotion(isroot = true, state=false)
+@named body = Body(; m = 1, r_cm = [0.0, 0, 0], isroot=true, quat=true, w_a=[1,0.5,0.2])
+
+# @named joint = Multibody.FreeMotion(isroot = true, state=true, quat=true)
+# @named body = Body(; m = 1, r_cm = [0.0, 0, 0], isroot=false, w_a=[1,1,1])
+
+world = Multibody.world
+
+
+connections = [connect(world.frame_b, joint.frame_a)
+               connect(joint.frame_b, body.frame_a)]
+
+
+@named model = ODESystem(connections, t,
+                         systems = [world, joint, body])
+irsys = IRSystem(model)
+ssys = structural_simplify(irsys)
+
+D = Differential(t)
+prob = ODEProblem(ssys, [collect(body.w_a) .=> [1,0,0];], (0, 2pi))
+
+sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 0 .* randn.())#, initializealg=ShampineCollocationInit(0.01))
+# sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 1e-12 .* randn.(), dtmin=1e-8, force_dtmin=true)
+@test SciMLBase.successful_retcode(sol)
+doplot() && plot(sol, layout=13)
+# end
+
+ts = 0:0.1:2pi
+Q = Matrix(sol(ts, idxs = [body.Q...]))
+Qh = Matrix(sol(ts, idxs = [body.Q̂...]))
+n = Matrix(sol(ts, idxs = [body.n...]))
+@test mapslices(norm, Qh, dims=1).^2 ≈ n
+@test Q ≈ Qh ./ sqrt.(n) atol=1e-2
+@test norm(mapslices(norm, Q, dims=1) .- 1) < 1e-2
+
+
+ti = 5
+for (i, ti) in enumerate(ts)
+    R = get_R(sol, body.frame_a, ti)
+    @test norm(R'R - I) < 1e-10
+    @test Multibody.from_Q(Q[:, i], 0).R ≈ R atol=1e-3
+end
+
+# Test that rotational velocity of 1 results in one full rotation in 2π seconds. Test rotation around all three major axes
+@test get_R(sol, body.frame_a, 0pi) ≈ I
+@test get_R(sol, body.frame_a, pi/2) ≈ [1 0 0; 0 0 -1; 0 1 0]' atol=1e-3
+@test get_rot(sol, body.frame_a, pi/2)*[0,1,0] ≈ [0,0,1] atol=1e-3 # Same test as above, but in one other way for sanity. They y-axis basis vector is rotated around x axis which aligns it with the z-axis basis vector. get_rot is used for this, which is the transpose of get_R
+@test get_R(sol, body.frame_a, 1pi) ≈ diagm([1, -1, -1]) atol=1e-3
+@test get_R(sol, body.frame_a, 2pi) ≈ I atol=1e-3
+
+prob = ODEProblem(ssys, [collect(body.w_a) .=> [0,1,0];], (0, 2pi))
+sol = solve(prob, Rodas4())#, 
+@test get_R(sol, body.frame_a, 0pi) ≈ I
+@test get_R(sol, body.frame_a, 1pi) ≈ diagm([-1, 1, -1]) atol=1e-3
+@test get_R(sol, body.frame_a, 2pi) ≈ I atol=1e-3
+
+
+prob = ODEProblem(ssys, [collect(body.w_a) .=> [0,0,1];], (0, 2pi))
+sol = solve(prob, Rodas4())#, 
+@test get_R(sol, body.frame_a, 0pi) ≈ I
+@test get_R(sol, body.frame_a, 1pi) ≈ diagm([-1, -1, 1]) atol=1e-3
+@test get_R(sol, body.frame_a, 2pi) ≈ I atol=1e-3
+
+end
+
+
+# ============================================================
+## Quaternion states in joint instead of body
+# ============================================================
+
+@testset "Quaternions and state in free motion" begin
+    using LinearAlgebra, ModelingToolkit, Multibody, JuliaSimCompiler
+    t = Multibody.t
+    world = Multibody.world
+
+    @named joint = Multibody.FreeMotion(isroot = true, state=true, quat=true)
+    @named body = Body(; m = 1, r_cm = [0.0, 0, 0])
+
+    connections = [connect(world.frame_b, joint.frame_a)
+                connect(joint.frame_b, body.frame_a)]
+
+    @named model = ODESystem(connections, t,
+                            systems = [world, joint, body])
+    irsys = IRSystem(model)
+    ssys = structural_simplify(irsys)
+
+    D = Differential(t)
+    # q0 = randn(4); q0 ./= norm(q0)
+    q0 = [1,0,0,0]
+    prob = ODEProblem(ssys, [
+        collect(body.w_a) .=> [1,0,0];
+        collect(joint.Q) .=> q0;
+        collect(joint.Q̂) .=> q0;
+        ], (0, 2pi))
+
+    using OrdinaryDiffEq, Test
+    sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 1e-6 .* randn.())
+    @test SciMLBase.successful_retcode(sol)
+    # doplot() && plot(sol, layout=21)
+
+
+    # These tests all pass with Euler angles as state, but with quaternions as state there are rotation-matrix elements in the state and even though the system simplifies and solves, it does not solve correctly
+    ts = 0:0.1:2pi
+    Q = Matrix(sol(ts, idxs = [joint.Q...]))
+    Qh = Matrix(sol(ts, idxs = [joint.Q̂...]))
+    n = Matrix(sol(ts, idxs = [joint.n...]))
+    @test mapslices(norm, Qh, dims=1).^2 ≈ n
+    @test Q ≈ Qh ./ sqrt.(n) atol=1e-2
+    @test norm(mapslices(norm, Q, dims=1) .- 1) < 1e-2
+
+    @test get_R(sol, joint.frame_b, 0pi) ≈ I
+    @test_broken get_R(sol, joint.frame_b, 1pi) ≈ diagm([1, -1, -1]) atol=1e-3
+    @test get_R(sol, joint.frame_b, 2pi) ≈ I atol=1e-3
+
+    Matrix(sol(ts, idxs = [joint.w_rel_b...]))
+end
+
+# ============================================================
+## Spherical joint pendulum with quaternions
+# ============================================================
+
+@testset "Spherical joint with quaternion state" begin
+    using LinearAlgebra, ModelingToolkit, Multibody, JuliaSimCompiler
+    t = Multibody.t
+    world = Multibody.world
+
+
+    @named joint = Multibody.Spherical(isroot=false, state=false, quat=false)
+    @named rod = FixedTranslation(; r = [1, 0, 0])
+    @named body = Body(; m = 1, isroot=true, quat=true, air_resistance=0.0)
+
+    # @named joint = Multibody.Spherical(isroot=true, state=true, quat=true)
+    # @named body = Body(; m = 1, r_cm = [1.0, 0, 0], isroot=false)
+
+    connections = [connect(world.frame_b, joint.frame_a)
+                connect(joint.frame_b, rod.frame_a)
+                connect(rod.frame_b, body.frame_a)]
+
+
+    @named model = ODESystem(connections, t,
+                            systems = [world, joint, body, rod])
+    irsys = IRSystem(model)
+    @test_skip begin # https://github.com/JuliaComputing/JuliaSimCompiler.jl/issues/298
+        ssys = structural_simplify(irsys)
+
+
+        D = Differential(t)
+        q0 = randn(4); q0 ./= norm(q0)
+        # q0 = [1,0,0,0]
+        prob = ODEProblem(ssys, [
+            # collect(body.w_a) .=> [1,0,0];
+            # collect(body.Q) .=> q0;
+            # collect(body.Q̂) .=> q0;
+            ], (0, 30))
+
+        using OrdinaryDiffEq, Test
+        sol = solve(prob, Rodas4(); u0 = prob.u0 .+ 0 .* randn.())
+        @test SciMLBase.successful_retcode(sol)
+        # doplot() && plot(sol, layout=21)
+
+
+        ts = 0:0.1:2pi
+        Q = Matrix(sol(ts, idxs = [body.Q...]))
+        Qh = Matrix(sol(ts, idxs = [body.Q̂...]))
+        n = Matrix(sol(ts, idxs = [body.n...]))
+        @test mapslices(norm, Qh, dims=1).^2 ≈ n
+        @test Q ≈ Qh ./ sqrt.(n) atol=1e-2
+        @test norm(mapslices(norm, Q, dims=1) .- 1) < 1e-2
+
+        Matrix(sol(ts, idxs = [body.w_a...]))
+
+        @test get_R(joint.frame_b, 0pi) ≈ I
+        @test get_R(joint.frame_b, sqrt(9.81/1)) ≈ diagm([1, -1, -1]) atol=1e-3
+        @test get_R(joint.frame_b, 2pi) ≈ I atol=1e-3
+
+
+        Matrix(sol(ts, idxs = [joint.w_rel_b...]))
+
+        # render(model, sol)
+    end
+end
+
+
+# ==============================================================================
+## FreeBody ====================================================================
+# ==============================================================================
+# https://doc.modelica.org/om/Modelica.Mechanics.MultiBody.Examples.Elementary.FreeBody.html
+using Multibody
+using ModelingToolkit
+# using Plots
+using JuliaSimCompiler
+using OrdinaryDiffEq
+
+@testset "FreeBody" begin
+    t = Multibody.t
+    D = Differential(t)
+    world = Multibody.world
+
+    @named begin
+        body = BodyShape(m = 1, I_11 = 1, I_22 = 1, I_33 = 1, r = [0.4, 0, 0],
+                        r_0 = [0.2, -0.5, 0.1], r_cm = [0.2, 0, 0], isroot = true, quat=true)
+        bar2 = FixedTranslation(r = [0.8, 0, 0])
+        spring1 = Multibody.Spring(c = 20, s_unstretched = 0)
+        spring2 = Multibody.Spring(c = 20, s_unstretched = 0)
+    end
+
+    eqs = [connect(bar2.frame_a, world.frame_b)
+        connect(spring1.frame_b, body.frame_a)
+        connect(bar2.frame_b, spring2.frame_a)
+        connect(spring1.frame_a, world.frame_b)
+        connect(body.frame_b, spring2.frame_b)]
+
+    @named model = ODESystem(eqs, t,
+                            systems = [
+                                world,
+                                body,
+                                bar2,
+                                spring1,
+                                spring2,
+                            ])
+    ssys = structural_simplify(IRSystem(model))#, alias_eliminate = true)
+    # ssys = structural_simplify(model, allow_parameters = false)
+    prob = ODEProblem(ssys,
+                    [world.g => 9.80665;
+                        collect(body.body.w_a .=> 0);
+                    collect(body.body.v_0 .=> 0);
+                    #   collect(body.body.phi .=> deg2rad(10));
+                        collect(body.body.Q) .=> vec(QuatRotation(RotXYZ(deg2rad.((10,10,10))...)));
+                        collect(body.body.Q̂) .=> vec(QuatRotation(RotXYZ(deg2rad.((10,10,10))...)));
+                    ], (0, 10))
+
+    # @test_skip begin # The modelica example uses angles_fixed = true, which causes the body component to run special code for variable initialization. This is not yet supported by MTK
+    sol = solve(prob, Rodas5P(), abstol=1e-6, reltol=1e-6)
+    @test SciMLBase.successful_retcode(sol)
+
+    @test sol(10, idxs=[
+        body.r_0;
+    ]) ≈ [
+        0.196097
+        -0.46243
+        0.080562
+    ] atol=1e-1
+
+    @test_broken get_rot(sol, body.frame_b, 10)[1,2] ≈ 0.104409 atol=0.01
+
+    doplot() && plot(sol, idxs = [body.r_0...; body.body.w_a; body.body.v_0], layout=(3,3), size=(1000, 1000)) |> display
+
+end
\ No newline at end of file