add Jacobian wrt data to Zhong06_CCP_constant_mass_mono.jl;

update `Moreau_CCP_constant_mass.jl`;
update `superball.jl` and notebook;
fix `record`
workaround `slack_linestyle`, need to find another way;

examples/robots/superball.jl

@@ -11,6 +11,7 @@ function superball(c=0.0;
         d = l/2,
         k = 4000.0,
         visible = true,
+        constrained = false,
         loadmesh = true,
     )
 
@@ -40,9 +41,9 @@ function superball(c=0.0;
             m = 5.0,
             μ,
             e,
-            visible = ifelse(i==1 && visible,true,false),
-            ci = ifelse(i==1 && visible,collect(1:6),Int[]),
-            cstr_idx = ifelse(i==1 && visible,Int[],[1]),
+            visible,
+            ci = ifelse(i==1 && constrained,collect(1:6),Int[]),
+            cstr_idx = ifelse(i==1 && constrained,Int[],[1]),
             loadmesh,
             isbody =false,
         )

src/control/control.jl

@@ -322,4 +322,8 @@ end
 function get_num_of_params(policy::AbstractPolicy)
     (;ps) = policy.nt
     length(ps)
+end
+
+function get_num_of_params(policy::NoPolicy)
+    0
 end

src/mechanics/dynamics_solvers/Moreau_family/Moreau_CCP_constant_mass.jl

@@ -6,6 +6,7 @@ end
 function generate_cache(
         simulator::Simulator{<:DynamicsProblem{
             RobotType,
+            policyType,
             EnvType,
             RestitutionFrictionCombined{NewtonRestitution,CoulombFriction}
         }},
@@ -15,12 +16,17 @@ function generate_cache(
         },
         ::Val{true};
         dt,kargs...
-    )   where {RobotType,EnvType}
+    )   where {RobotType,policyType,EnvType}
     (;prob) = simulator
-    (;bot,env) = prob
+    (;bot,policy,env) = prob
     (;structure) = bot
-    F!(F,q,q̇,t) = generalized_force!(F,bot,q,q̇,t;gravity=true)
-    Jac_F!(∂F∂q̌,∂F∂q̌̇,q,q̇,t) = generalized_force_jacobian!(∂F∂q̌,∂F∂q̌̇,bot,q,q̇,t)
+    options = merge(
+        (gravity=true,factor=1,checkpersist=true), #default
+        prob.options,
+        solver.options,
+    )
+    F!(F,q,q̇,t) = generalized_force!(F,bot,policy,q,q̇,t;gravity=options.gravity)
+    Jac_F!(∂F∂q̌,∂F∂q̌̇,q,q̇,t) = generalized_force_jacobian!(∂F∂q̌,∂F∂q̌̇,bot,policy,q,q̇,t)
 
     M = Matrix(assemble_M(structure))
     Φ = make_cstr_function(bot)

src/mechanics/dynamics_solvers/Zhong06_family/Adjoint_Zhong06_constant_mass.jl

@@ -199,7 +199,7 @@ function generate_cache(
     ∂Fₘ∂θ = zeros(T,nq,nθ)
     Jacᵏ⁺¹ₘ = zeros(T,nx,nθ)
     function ∂f∂θ!(Jacᵏ⁺¹ₘ,∂Fₘ∂θ,q,q̇,t,h)
-        generalized_force_jacobian!(∂Fₘ∂θ,bot,policy,q,q̇,t;gravity=options.gravity)
+        generalized_force_jacobian!(∂Fₘ∂θ,structure,policy,q,q̇,t;gravity=options.gravity)
         Jacᵏ⁺¹ₘ[    1: nq,1:nθ] .= -h^2/2*∂Fₘ∂θ
         Jacᵏ⁺¹ₘ[ nq+1:2nq,1:nθ] .= -h^2/2*∂Fₘ∂θ
         Jacᵏ⁺¹ₘ[2nq+1:end,1:nθ] .= 0.0

src/mechanics/dynamics_solvers/Zhong06_family/Zhong06_CCP_constant_mass_mono.jl

@@ -25,8 +25,9 @@ function generate_cache(
         prob.options,
         solver.options,
     )
-    F!(F,q,q̇,t) = generalized_force!(F,bot,policy,q,q̇,t;gravity=true)
+    F!(F,q,q̇,t) = generalized_force!(F,bot,policy,q,q̇,t;gravity=options.gravity)
     Jac_F!(∂F∂q̌,∂F∂q̌̇,q,q̇,t) = generalized_force_jacobian!(∂F∂q̌,∂F∂q̌̇,bot,policy,q,q̇,t)
+    ## ∂F∂θ!(∂Fₘ∂θ,q,q̇,t) = generalized_force_jacobian!(∂Fₘ∂θ,structure,policy,q,q̇,t;)
 
     M = Matrix(assemble_M(structure))
     Φ = make_cstr_function(bot)
@@ -51,7 +52,7 @@ function generate_cache(
     ) = prepare_contacts(bot,env)
 
     cache = @eponymtuple(
-        F!,Jac_F!,
+        F!,Jac_F!,#∂F∂θ!,
         M,Φ,A,Ψ,B,∂Ψ∂q,∂Aᵀλ∂q,∂Bᵀμ∂q,
         contacts_bits,
         persistent_bits,
@@ -66,19 +67,19 @@ end
 function make_step_k(
         solver_cache::Zhong06_CCP_Constant_Mass_Mono_Cache,
         nq,nλ,na,
-        qₖ₋₁,vₖ₋₁,pₖ₋₁,tₖ₋₁,
+        step_data,tₖ₋₁,
         pₖ,vₖ,
         invM,
         h,mass_norm)
     (;F!,Jac_F!,M,Φ,A,∂Aᵀλ∂q) = solver_cache.cache
-
+    (;qₖ₋₁,vₖ₋₁,pₖ₋₁) = step_data
     n1 = nq
     n2 = nq+nλ
     nΛ = 3na
     nx = n2
     function ns_stepk!(
-            𝐫𝐞𝐬,𝐉,
-            Fₘ,∂F∂q,∂F∂q̇,
+            𝐫𝐞𝐬,𝐉,𝐉_data,
+            Fₘ,∂F∂q,∂F∂q̇,#∂Fₘ∂θ,
             𝐰,x,Λₖ,y,∂y∂x,
             Λ_split,y_split,
             structure,
@@ -98,6 +99,7 @@ function make_step_k(
         Aₖ₋₁ = A(qₖ₋₁)
         Aₖ   = A(qₖ)
 
+        # pₖ₋₁ = M*vₖ₋₁
         𝐫𝐞𝐬[   1:n1] .= -h.*pₖ₋₁ .+ M*(qₖ.-qₖ₋₁) .-
                         mass_norm.*transpose(Aₖ₋₁)*λₘ .-
                         (h^2)/2 .*Fₘ
@@ -108,6 +110,18 @@ function make_step_k(
         𝐉[   1:n1,n1+1:n2] .= -mass_norm.*transpose(Aₖ₋₁)
 
         𝐉[n1+1:n2,   1:n1] .=  -mass_norm.*Aₖ
+        𝐉_data .= 0.0
+        #position
+        𝐉_data[   1:nq,   1:nq] .= -M .-h^2/2 .*(1/2 .*∂F∂q .- 1/h.*∂F∂q̇) .- mass_norm.*∂Aᵀλ∂q(qₖ₋₁,λₘ)
+        #velocity
+        𝐉_data[   1:nq,nq+1:2nq] .= -h*M
+        #control
+        ## ∂F∂θ!(∂Fₘ∂θ,qₘ,q̇ₘ,tₘ)
+        ## 𝐉_data[   1:nq,2nq+1:nθ] .= -h^2/2*∂Fₘ∂θ
+        #structural
+        ## 𝐉_data[   1:n1,n1+1:n2]
+        #contact
+        ## 𝐉_data[   1:n1,n1+1:n2]
 
         if na != 0
             get_distribution_law!(structure,contact_cache,qₖ)
@@ -178,7 +192,7 @@ function solve!(sim::Simulator,solver_cache::Zhong06_CCP_Constant_Mass_Mono_Cach
         progress=true,exception=true
     )
     (;prob,controller,tspan,restart,totalstep) = sim
-    (;bot,env) = prob
+    (;bot,policy,env) = prob
     (;structure,traj,contacts_traj) = bot
     (;M,A,contacts_bits) = solver_cache.cache
     q0 = traj.q[begin]
@@ -194,11 +208,9 @@ function solve!(sim::Simulator,solver_cache::Zhong06_CCP_Constant_Mass_Mono_Cach
     F = zeros(T,nq)
     ∂F∂q = zeros(T,nq,nq)
     ∂F∂q̇ = zeros(T,nq,nq)
+    nθ = get_num_of_params(policy)
+    ∂Fₘ∂θ = zeros(T,nq,nθ)
     nx = nq + nλ
-    Δx = zeros(T,nx)
-    x = zero(Δx)
-    Res = zero(Δx)
-    Jac = zeros(T,nx,nx)
     mr = norm(M,Inf)
     mass_norm = mr
 
@@ -238,6 +250,7 @@ function solve!(sim::Simulator,solver_cache::Zhong06_CCP_Constant_Mass_Mono_Cach
         x = zero(Δx)
         Res = zero(Δx)
         Jac = zeros(T,nx,nx)
+        Jac_data = zeros(T,nx,nq+nq+nθ)
         Λₖ = @view x[(n2+1):n2+nΛ]
         y  = @view x[n2+nΛ+1:n2+2nΛ]
         𝐰 = zeros(T,nΛ)
@@ -259,13 +272,17 @@ function solve!(sim::Simulator,solver_cache::Zhong06_CCP_Constant_Mass_Mono_Cach
         ΔΛc_split = split_by_lengths(ΔΛc,3)
         Δyc_split = split_by_lengths(Δyc,3)
         get_frictional_directions_and_positions!(structure, contact_cache, qₖ₋₁, q̇ₖ₋₁, Λₖ)
+        step_data = ComponentArray(
+            @eponymtuple(
+                qₖ₋₁,vₖ₋₁=q̇ₖ₋₁,pₖ₋₁,
+            )
+        )
         ns_stepk! = make_step_k(
             solver_cache,
             nq,nλ,na,
-            qₖ₋₁,q̇ₖ₋₁,pₖ₋₁,tₖ₋₁,
+            step_data,tₖ₋₁,
             pₖ,q̇ₖ,
-            invM,
-            dt,mass_norm
+            invM,dt,mass_norm
         )
         restart_count = 0
         Λ_guess = 0.1
@@ -276,8 +293,8 @@ function solve!(sim::Simulator,solver_cache::Zhong06_CCP_Constant_Mass_Mono_Cach
             x[   nq+1:nq+nλ]       .= 0.0
             for iteration = 1:maxiters
                 ns_stepk!(
-                    Res,Jac,
-                    F,∂F∂q,∂F∂q̇,
+                    Res,Jac,Jac_data,
+                    F,∂F∂q,∂F∂q̇,#∂Fₘ∂θ,
                     𝐰,x,Λₖ,y,∂y∂x,
                     Λ_split,y_split,
                     structure,
@@ -351,6 +368,7 @@ function solve!(sim::Simulator,solver_cache::Zhong06_CCP_Constant_Mass_Mono_Cach
                         @show (y_split)⊙(Λ_split)
                         @show (Δyp_split)⊙(ΔΛp_split)
                     end
+                    @show lu𝐉\(-Jac_data)
                     ## @show Λₖ, μ
                 end
             end

src/postprocess/vis.jl

@@ -132,7 +132,7 @@ function MakieCore.plot!(viz::Viz{Tuple{S}};
         if cab.force.state.length > cab.force.state.restlen
             return :solid
         else
-            return viz.slack_linestyle[]
+            return :solid
         end
     end
 

test/vis.jl

@@ -356,7 +356,7 @@ function plot_traj!(bot::RB.Robot;
                 framerate = 30 
                 skipstep = round(Int,1/framerate/dt*speedup)
                 recordsteps = 1:skipstep:length(traj.t)
-                record(fig, figname, recordsteps;
+                record(fig, figname, recordsteps; px_per_unit = 2,
                     framerate) do this_step
                     if actuate
                         RB.actuate!(bot,[traj.t[this_step]])
@@ -367,7 +367,7 @@ function plot_traj!(bot::RB.Robot;
                     # @show RB.mechanical_energy(structure)
                     #
                     this_time[] = traj.t[this_step]
-                    RB.analyse_slack(structure,true)
+                    ## RB.analyse_slack(structure,true)
                     tgob[] = structure
                 end
             else
@@ -401,7 +401,7 @@ function plot_traj!(bot::RB.Robot;
         colgap!(grid1,colgap)
         rowgap!(grid1,rowgap)
     end
-    if fig isa Figure
+    if (fig isa Figure) && !dorecord
         savefig(fig,figname)
         DataInspector(fig)
     end

yard/nonsmooth/make_notebooks.jl

@@ -28,20 +28,20 @@ using GLMakie
 # )
 
 
-Literate.notebook(
-    joinpath(@__DIR__,"meteor_hammer.jl"), 
-    joinpath(@__DIR__,"../../../notebooks/");
-    name="meteor_hammer",
-    execute = true,
-    # documenter = false,
-)
-
-
 ## Literate.notebook(
-##     joinpath(@__DIR__,"superball.jl"), 
+##     joinpath(@__DIR__,"meteor_hammer.jl"), 
 ##     joinpath(@__DIR__,"../../../notebooks/");
-##     name="superball",
+##     name="meteor_hammer",
 ##     execute = true,
 ##     # documenter = false,
 ## )
+
+
+Literate.notebook(
+    joinpath(@__DIR__,"superball.jl"), 
+    joinpath(@__DIR__,"../../../notebooks/");
+    name="superball",
+    execute = true,
+    # documenter = false,
+)
 GLMakie.closeall()

yard/nonsmooth/superball.jl

@@ -48,9 +48,11 @@ ballbot = superball(
     e = 0.8,
     l,d,
     z0 = l^2/(sqrt(5)*d) + 2.0,
+    visible = true,
     constrained = false,
     loadmesh = false,
 )
+
 # time 
 tspan = (0.0,5.0)
 h = 1e-2
@@ -125,7 +127,7 @@ GM.activate!(;scalefactor); with_theme(theme_pub;
                 meshcolor = Makie.RGBAf(r,g,b,alphas[i]))
             end
             RB.hidey(ax)
-            lines!(ax,r2p1)
+            lines!(ax,Matrix(r2p1))
         end
     )
     ax31 = Axis(gd2[1,1],
@@ -189,6 +191,54 @@ GM.activate!(;scalefactor); with_theme(theme_pub;
     fig
 end
 
+GM.activate!(;scalefactor=4,px_per_unit=4); with_theme(theme_pub;
+        figure_padding = (0,0,0,-fontsize),
+        size = (1.0tw,0.30tw),
+        Axis3 = (
+            azimuth = 4.7855306333269805,
+            elevation = 0.03269908169872391,
+            xlabeloffset = fontsize,
+            zlabeloffset = 0.6fontsize,
+        )
+    ) do
+    bot = ballbot
+    (;t) = bot.traj
+    imptimes = [0.25,0.29,0.30,0.34]
+    impstep = RB.time2step(imptimes[1],bot.traj.t)
+    steps = vcat(1,impstep,collect(impstep:50:length(t)))
+    nstep = length(steps)
+    alphas = fill(0.15,nstep)
+    alphas[1:3] = [1,0.2,0.2]
+    alphas[end] = 1
+    cg = cgrad(:winter, nstep, categorical = true, rev = true)
+    step_start = RB.time2step(0.1,bot.traj.t)
+    step_stop = RB.time2step(0.35,bot.traj.t)
+    v2p1 = RB.get_velocity!(bot,2,1,step_start:step_stop)
+    v1p2 = RB.get_velocity!(bot,1,2,step_start:step_stop)
+    v6p2 = RB.get_velocity!(bot,6,2,step_start:step_stop)
+    r2p1 = RB.get_trajectory!(bot,2,1)
+    me = RB.mechanical_energy!(bot,)
+    plot_traj!(bot;
+        AxisType = Axis3,
+        figname = "ballbot_rolling.gif",
+        xlims = [-1,20],
+        ylims = [-1,8],
+        zlims = [-1e-3,3.0],
+        doslide = false,
+        dorecord = true,
+        showinfo = false,
+        showpoints = false,
+        showlabels = false,
+        showmesh = true,
+        showwire = true,
+        showtitle = false,
+        showcables  = true,
+        sup! = (ax,_,_) -> begin
+            RB.hidey(ax)
+        end
+    )
+end
+
 # ## Second Scenario
 ballbot = superball(
     0.0;
@@ -272,7 +322,7 @@ GM.activate!(;scalefactor); with_theme(theme_pub;
                 cablecolor=Makie.RGBAf(db.r,db.g,db.b,Makie.N0f8(alphas[i])),
                 meshcolor = Makie.RGBAf(r,g,b,alphas[i]))
             end
-            lines!(ax,r2p1)
+            lines!(ax,Matrix(r2p1))
         end
     )
     ax2 = Axis(gd2[1,1];tellheight=false,xlabel=tlabel,ylabel = "Energy (J)")
@@ -325,7 +375,7 @@ stats_superballs_dt = [
             );
             tspan=(0.0,0.1),dt,ftol=ifelse(dt==5e-6,1e-14,1e-12),
             maxiters=500,exception=false
-        ).prob.bot
+        )[1].prob.bot
     end
     for dt in dts, solver in (RB.Zhong06(),RB.Moreau(0.5))
 ];