update two link jum lin id

jump_lin_id/two_link.jl

@@ -13,7 +13,7 @@ if !isdefined(:AbstractEnsembleSystem)
     @everywhere include("../flux/nn_prior/linear_sys.jl")
 end
 
-n,m,N,h,σ0 = 4,2,200,0.02,0.01
+n,m,N,h,σ0 = 4,3,200,0.02,0.01
 np         = n*(n+m)
 sys        = TwoLinkSys(N=N, h=h, σ0 = σ0)#; warn("Added noise")
 seed       = 1
@@ -25,30 +25,31 @@ p          = forward_kin(q)
 qd         = centraldiff(q')'/sys.h
 qdd        = centraldiff(qd')'/sys.h
 u          = torque(q,qd,qdd)
-u        .+= 0.2randn(size(u))
+u        .+= 0.1randn(size(u))
 # x,y,u  = generate_data(sys, System, seed; ufun=x->u)
 
 const wall = -0.2
 function affect!(i)
-    q = i.u[1:2]
-    qd = i.u[3:4]
-    p1,p2 = forward_kin(q)[2]
-    J = TwoLink.jacobian(q)
-    v1,v2 = J*qd
-    d = [0, wall-p2]
-    δq = J\d
-    i.u[3:4] = J\[v1, 0]
-    i.u[1:2] .+= δq
+    # q = i.u[1:2]
+    # qd = i.u[3:4]
+    # p1,p2 = forward_kin(q)[2]
+    # J = TwoLink.jacobian(q)
+    # v1,v2 = J*qd
+    # d = [0, wall-p2]
+    # δq = J\d
+    # i.u[3:4] = J\[v1, 0]
+    # i.u[1:2] .+= δq
 end
-condition(t,x,i) = forward_kin(x)[2][2] > wall
+condition(x,t,i) = forward_kin(x)[2][2] > wall
 cb = DiscreteCallback(condition,affect!)
 
 t      = h:h:N*h+h
 q0     = [q[:,1]; qd[:,1]]
-prob   = OrdinaryDiffEq.ODEProblem((t,x)->time_derivative(x, u[:,ceil(Int,t/h)]),q0,(t[[1,end]]...))
+prob   = OrdinaryDiffEq.ODEProblem((x,p,t)->time_derivative(x, u[:,ceil(Int,t/h)]),q0,(t[[1,end]]...))
 sol    = solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8, callback=cb)
 x      = hcat(sol(t)...)
 u      = u[:,1:N]
+u      = [u; ones(1,size(u,2))] # Extend to estimate friction
 x,y    = xy(System,x,N)
 plot(forward_kin(x[1:2,:])')
 plot!([1,N],[wall, wall], l=:dash, c=:black)
@@ -64,29 +65,40 @@ colors = [HSV(h,1,0.8) for h in linspace(0,200,5)]
 
 r2,p0 = 300,7
 model = LTVModels.fit_statespace_admm(x,u,1, extend=false)
-ui = @manipulate for i in 1:24,  λ in linspace(-3,3,50)
-    plot(Jtrue[i,:],lab="True", l=:dash, c=:blue, layout=1, legend=false, size=(1920,1200), show=false)
-    # ui = @manipulate for r2 in linspace(1,1000,20), p0 in linspace(1,7,50), i in 1:24
-    global n,m,np,x,u,N,model
-    function dostuff()
+si = slider(1:np);  sλ = slider(linspace(-3,3,50));
+
+function opt(λ)
+    global n,m,np,x,u,N,model,lastλ
     j = 1
     R1 = eye(n*(n+m))
     R2 = r2*eye(n)
     P0 = 10^p0*R1
-        model = LTVModels.fit_statespace_admm!(model,x,u,10^λ, extend=false, tol=1e-4)
+    model = LTVModels.fit_statespace_admm!(model,x,u,10^λ, extend=false, tol=1e-5)
     # model = LTVModels.fit_statespace_dp(x,u,M,extend=false)
     # model = LTVModels.fit_model(LTVModels.KalmanModel, x,u,R1,R2,P0, extend=true)
-        seed = 1
+end
+lastλ = 0.
+ui = @manipulate for i in si,  λ in sλ
+    global n,m,np,x,u,N,model,lastλ
+    j = 1
+    if λ != lastλ
+        J = opt(λ)
+        lastλ = λ
+    end
     J = reshape(cat(2, model.At, model.Bt), np, N)
+    # plot(Jtrue[i,:],lab="True", l=:dash, c=:blue, layout=1, legend=false, size=(1920,1200), show=false)
+    # ui = @manipulate for r2 in linspace(1,1000,20), p0 in linspace(1,7,50), i in 1:24
+
+    # model = LTVModels.fit_statespace_dp(x,u,M,extend=false)
+    # model = LTVModels.fit_model(LTVModels.KalmanModel, x,u,R1,R2,P0, extend=true)
 
     # yh = LTVModels.predict(model, x,u)
     # plot(y', lab="y", layout=6, show=false)
     # plot!(yh', lab="yh", subplot=1:4, show=false)
     # plot!(u', lab="u", subplot=5:6)
-        plot!(J[i,:], lab="Estimated", c=colors[j], subplot=1, show=false,  show=false)
-        plot!(repmat([1, N],1,N), [0,0], l=:dash, c=:black, linewidth=2)
-    end
-    dostuff()
+    plot(J[i,:], lab="Estimated", c=colors[j], subplot=1, show=false,  show=false)
+    plot!(0.5maximum(abs.(J[i,:]))*sign.(qd)', linewidth=[3 1])
+    hline!([0], l=(:dash, :black), linewidth=2)
 end
 # gui()