optimize spring rendering and add options

Author: baggepinnen

ext/Render.jl

@@ -385,16 +385,18 @@ function render!(scene, ::typeof(Spring), sys, sol, t)
     r_0a = get_fun(sol, collect(sys.frame_a.r_0))
     r_0b = get_fun(sol, collect(sys.frame_b.r_0))
     color = get_color(sys, sol, :blue)
+    n_wind = sol(sol.t[1], idxs=sys.num_windings)
+    radius = sol(sol.t[1], idxs=sys.radius) |> Float32
+    N = sol(sol.t[1], idxs=sys.N) |> Int
     thing = @lift begin
         r1 = Point3f(r_0a($t))
         r2 = Point3f(r_0b($t))
-        spring_mesh(r1,r2)
+        spring_mesh(r1,r2; n_wind, radius, N)
     end
     plot!(scene, thing; color)
     true
 end
 
-
 function render!(scene, ::Function, sys, sol, t, args...) # Fallback for systems that have at least two frames
     count(ModelingToolkit.isframe, sys.systems) == 2 || return false
     r_0a = get_fun(sol, collect(sys.frame_a.r_0))
@@ -422,24 +424,19 @@ function render!(scene, ::Function, sys, sol, t, args...) # Fallback for systems
 end
 
 function spring_mesh(p1, p2; n_wind=6, radius=0.1f0, N=200)
-    phi = range(0, n_wind*2π, length=N)
-    
+    phis = range(0, n_wind*2π, length=N)
     d = p2 - p1
-    
-    x = radius*cos.(phi)
-    y = radius*sin.(phi)
     z = range(0, norm(d), length=N) # Correct length
-    points = Point3f.(x, y, z)
+    dn = d ./ norm(d)
+    R = rot_from_line(dn)
 
-    d = d ./ norm(d)
-
-    # Rotate
-    R = rot_from_line(d)
-    points = Ref(R) .* points
-
-    # Translate
-    points = points .+ p1
+    points = map(enumerate(phis)) do (i,phi)
+        x = radius*cos(phi)
+        y = radius*sin(phi)
+        pᵢ = Point3f(x, y, z[i])
 
+        R * pᵢ + p1
+    end
 
     Makie.GeometryBasics.LineString(points)
 end

src/forces.jl

@@ -181,7 +181,7 @@ function LineForceBase(; name, length = 0, s_small = 1e-10, fixed_rotation_at_fr
     compose(ODESystem(eqs, t; name), frame_a, frame_b)
 end
 
-function LineForceWithMass(; name, length = 0, m = 1.0, lengthFraction = 0.5, kwargs...)
+function LineForceWithMass(; name, length = 0, m = 1.0, lengthfraction = 0.5, kwargs...)
     m0 = m
     @named lfb = LineForceBase(; length, kwargs...)
     @unpack length, s, r_rel_0, e_rel_0, frame_a, frame_b = lfb
@@ -195,7 +195,7 @@ function LineForceWithMass(; name, length = 0, m = 1.0, lengthFraction = 0.5, kw
     ]
     @variables v_CM_0(t)[1:3]=zeros(3) [description = "First derivative of r_CM_0"]
     @variables ag_CM_0(t)[1:3]=zeros(3) [description = "D(v_CM_0) - gravityAcceleration"]
-    @parameters lengthFraction=lengthFraction [
+    @parameters lengthfraction=lengthfraction [
         description = "Location of point mass with respect to frame_a as a fraction of the distance from frame_a to frame_b",
         bounds = (0, 1),
     ]
@@ -214,13 +214,13 @@ function LineForceWithMass(; name, length = 0, m = 1.0, lengthFraction = 0.5, kw
     # NOTE, both frames are assumed to be connected, while modelica has special handling if they aren't
     if m0 > 0
         eqs = [eqs
-               collect(r_CM_0 .~ frame_a.r_0 + r_rel_0 * lengthFraction)
+               collect(r_CM_0 .~ frame_a.r_0 + r_rel_0 * lengthfraction)
                v_CM_0 .~ D.(r_CM_0)
                ag_CM_0 .~ D.(v_CM_0) - gravity_acceleration(r_CM_0)
                frame_a.f .~ resolve2(ori(frame_a),
-                                     (m * (1 - lengthFraction)) * ag_CM_0 - e_rel_0 * fa)
+                                     (m * (1 - lengthfraction)) * ag_CM_0 - e_rel_0 * fa)
                frame_b.f .~ resolve2(ori(frame_b),
-                                     (m * lengthFraction) * ag_CM_0 - e_rel_0 * fb)]
+                                     (m * lengthfraction) * ag_CM_0 - e_rel_0 * fb)]
     else
         eqs = [eqs
                r_CM_0 .~ zeros(3)
@@ -264,7 +264,7 @@ function PartialLineForce(; name, kwargs...)
 end
 
 """
-    Spring(; c, name, m = 0, lengthFraction = 0.5, s_unstretched = 0, kwargs)
+    Spring(; c, name, m = 0, lengthfraction = 0.5, s_unstretched = 0, kwargs)
 
 Linear spring acting as line force between `frame_a` and `frame_b`.
 A force `f` is exerted on the origin of `frame_b` and with opposite sign
@@ -284,23 +284,35 @@ additional equations to handle the mass are removed.
 # Arguments:
 - `c`: Spring stiffness
 - `m`: Mass of the spring (can be zero)
-- `lengthFraction`: Location of spring mass with respect to `frame_a` as a fraction of the distance from `frame_a` to `frame_b` (=0: at `frame_a`; =1: at `frame_b`)
+- `lengthfraction`: Location of spring mass with respect to `frame_a` as a fraction of the distance from `frame_a` to `frame_b` (=0: at `frame_a`; =1: at `frame_b`)
 - `s_unstretched`: Length of the spring when it is unstretched
 - `kwargs`: are passed to `LineForceWithMass`
 
+# Rendering
+- `num_windings = 6`
+- `color = [0,0,1,1]`
+- `radius = 0.1`
+- `N = 200`
+
 See also [`SpringDamperParallel`](@ref)
 """
-@component function Spring(; c, name, m = 0, lengthFraction = 0.5, s_unstretched = 0, kwargs...)
+@component function Spring(; c, name, m = 0, lengthfraction = 0.5, s_unstretched = 0, num_windings=6, color=[0,0,1,1], radius=0.1, N=200, kwargs...)
     @named ptf = PartialTwoFrames()
     @unpack frame_a, frame_b = ptf
-    @named lineforce = LineForceWithMass(; length = s_unstretched, m, lengthFraction,
+    pars = @parameters begin
+    #     c=c, [description = "spring constant", bounds = (0, Inf)]
+    #     s_unstretched=s_unstretched, [
+    #         description = "unstretched length of spring",
+    #         bounds = (0, Inf),
+    #     ] # Bug in MTK where parameters only passed to sub components are ignored
+        num_windings = num_windings, [description = "Number of windings of the coil when rendered"]
+        color[1:4] = color
+        radius = radius, [description = "Radius of spring when rendered"]
+        N = N, [description = "Number of points in mesh when rendered"]
+    end
+    @named lineforce = LineForceWithMass(; length = s_unstretched, m, lengthfraction,
                                          kwargs...)
 
-    # @parameters c=c [description = "spring constant", bounds = (0, Inf)]
-    # @parameters s_unstretched=s_unstretched [
-    #     description = "unstretched length of spring",
-    #     bounds = (0, Inf),
-    # ] # Bug in MTK where parameters only passed to sub components are ignored
     @named spring2d = TP.Spring(; c, s_rel0 = s_unstretched)
 
     @variables r_rel_a(t)[1:3]=0 [
@@ -341,7 +353,8 @@ See also [`SpringDamperParallel`](@ref)
            connect(spring2d.flange_b, lineforce.flange_b)
            connect(spring2d.flange_a, lineforce.flange_a)]
 
-    extend(ODESystem(eqs, t; name, systems = [lineforce, spring2d]), ptf)
+    sys = extend(ODESystem(eqs, t; name=:nothing, systems = [lineforce, spring2d]), ptf)
+    add_params(sys, pars; name)
 end
 
 """