Skip to content

Commit

Permalink
Merge pull request #2378 from stack-of-tasks/pre-commit-ci-update-config
Browse files Browse the repository at this point in the history 
[pre-commit.ci] pre-commit autoupdate
  • Loading branch information
@jcarpent
jcarpent authored Aug 20, 2024
2 parents 7481a1c + 89f326c commit 178d996
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 43 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ repos:
doc/doxygen-awesome.*
)$
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.5.7
rev: v0.6.1
hooks:
- id: ruff
- id: ruff-format
Expand Down
90 changes: 48 additions & 42 deletions examples/simulation-pendulum-variational-casadi.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,11 +111,11 @@
" model = pin.Model()\n",
" geom_model = pin.GeometryModel()\n",
"\n",
" RED_COLOR = np.array([1., 0.1, 0.1, 1.])\n",
" RED_COLOR = np.array([1.0, 0.1, 0.1, 1.0])\n",
" WHITE_COLOR = np.ones(4)\n",
" \n",
"\n",
" joint_pl = pin.SE3.Identity()\n",
" mass = 1.\n",
" mass = 1.0\n",
" radius_base = 0.08\n",
" shape0 = fcl.Sphere(radius_base)\n",
" geom_obj = pin.GeometryObject(\"base\", 0, shape0, pin.SE3.Identity())\n",
Expand All @@ -131,7 +131,7 @@
"\n",
" radius = 0.06\n",
" mass = 0.5\n",
" length = 1.\n",
" length = 1.0\n",
" w_ = 0.02\n",
" sizes_ = (w_, length, w_)\n",
"\n",
Expand All @@ -149,15 +149,15 @@
"\n",
" geom_name2 = f\"capsule{k}\"\n",
" geom_pl2 = body_pl.copy()\n",
" geom_pl2.translation = np.array([0., 0., length / 2])\n",
" shape2 = fcl.Cylinder(radius/3, length)\n",
" geom_pl2.translation = np.array([0.0, 0.0, length / 2])\n",
" shape2 = fcl.Cylinder(radius / 3, length)\n",
" geom_obj2 = pin.GeometryObject(geom_name2, j_id, shape2, geom_pl2)\n",
" geom_obj2.meshColor = WHITE_COLOR\n",
" geom_model.addGeometryObject(geom_obj2)\n",
"\n",
" parent_id = j_id\n",
" joint_pl = body_pl.copy()\n",
" \n",
"\n",
" return model, geom_model"
]
},
Expand Down Expand Up @@ -223,7 +223,7 @@
],
"source": [
"vizer.initViewer()\n",
"vizer.loadViewerModel('pinocchio')"
"vizer.loadViewerModel(\"pinocchio\")"
]
},
{
Expand Down Expand Up @@ -285,15 +285,18 @@
"source": [
"# define the system Lagrangian's graph\n",
"\n",
"\n",
"def continuousLagrangian(model, data, q, v, ns=cpin):\n",
" return ns.computeKineticEnergy(model, data, q, v) - ns.computePotentialEnergy(model, data, q)\n",
" return ns.computeKineticEnergy(model, data, q, v) - ns.computePotentialEnergy(\n",
" model, data, q\n",
" )\n",
"\n",
"\n",
"def discreteLagrangian(model, data, q0, q1, alpha=0.5, ns=cpin):\n",
" q = ns.interpolate(model, q0, q1, alpha)\n",
" v = ns.difference(model, q0, q1) / timestep\n",
" return continuousLagrangian(model, data, q, v, ns=ns) * timestep\n",
" \n",
"\n",
"\n",
"## define the symbols\n",
"cmodel = cpin.Model(model)\n",
Expand Down Expand Up @@ -348,11 +351,15 @@
"nv = model.nv\n",
"cargs = [cq0, cq1, cq2, cu, cdq0_, cdq1_, cdq2_]\n",
"fused_dq = casadi.vertcat(cdq0_, cdq1_, cdq2_)\n",
"carg_names = ['q0', 'q1', 'q2', 'u', 'dq0_', 'dq1_', 'dq2_']\n",
"\n",
"residual_fun = casadi.Function(\"f\", cargs[:4], [\n",
" casadi.substitute(residual, fused_dq, SX.zeros(nv+nv+nv))\n",
"], carg_names[:4], [\"f\"])\n",
"carg_names = [\"q0\", \"q1\", \"q2\", \"u\", \"dq0_\", \"dq1_\", \"dq2_\"]\n",
"\n",
"residual_fun = casadi.Function(\n",
" \"f\",\n",
" cargs[:4],\n",
" [casadi.substitute(residual, fused_dq, SX.zeros(nv + nv + nv))],\n",
" carg_names[:4],\n",
" [\"f\"],\n",
")\n",
"residual_fun"
]
},
Expand All @@ -374,9 +381,7 @@
}
],
"source": [
"residual_fun(\n",
" pin.neutral(model), pin.neutral(model), pin.neutral(model), np.ones(nu)\n",
")"
"residual_fun(pin.neutral(model), pin.neutral(model), pin.neutral(model), np.ones(nu))"
]
},
{
Expand All @@ -392,21 +397,23 @@
" # define a function to compute the Jacobians of the residual\n",
" # wrt variables q0, q1, q2 and u\n",
" fused_args = casadi.vertcat(fused_dq, cu)\n",
" \n",
"\n",
" jac_residual_expr_full = casadi.jacobian(residual, fused_args)\n",
" nv = model.nv\n",
" jac_residual_blocks = [\n",
" jac_residual_expr_full[:, :nv],\n",
" jac_residual_expr_full[:, nv:2*nv],\n",
" jac_residual_expr_full[:, 2*nv:3*nv],\n",
" jac_residual_expr_full[:, 3*nv:]\n",
" jac_residual_expr_full[:, nv : 2 * nv],\n",
" jac_residual_expr_full[:, 2 * nv : 3 * nv],\n",
" jac_residual_expr_full[:, 3 * nv :],\n",
" ]\n",
" jac_residual_blocks = [\n",
" casadi.substitute(J, fused_dq, SX.zeros(nv+nv+nv))\n",
" casadi.substitute(J, fused_dq, SX.zeros(nv + nv + nv))\n",
" for J in jac_residual_blocks\n",
" ]\n",
" assert jac_residual_blocks[-1].shape == (nu, nv)\n",
" jac_residual_fun = casadi.Function(\"Df\", cargs[:4], jac_residual_blocks, carg_names[:4], [\"f0\", \"f1\", \"f2\", \"fu\"])\n",
" jac_residual_fun = casadi.Function(\n",
" \"Df\", cargs[:4], jac_residual_blocks, carg_names[:4], [\"f0\", \"f1\", \"f2\", \"fu\"]\n",
" )\n",
" return jac_residual_fun\n",
"\n",
"\n",
Expand Down Expand Up @@ -460,19 +467,17 @@
"def variational_integrator(q0, q1, u, Tf):\n",
" states_ = [q0, q1]\n",
" nsteps = int(Tf / timestep)\n",
" \n",
" Em_vals = [\n",
" mechanicalEnergyFromQs(model, pdata, q0, q1)\n",
" ]\n",
"\n",
" Em_vals = [mechanicalEnergyFromQs(model, pdata, q0, q1)]\n",
" errs = []\n",
" u0 = np.zeros(nu)\n",
" TOL = 1e-8\n",
" \n",
"\n",
" tnr = tqdm.trange(1, nsteps)\n",
" for t in tnr:\n",
" qp = states_[t-1].copy()\n",
" qp = states_[t - 1].copy()\n",
" qc = states_[t].copy()\n",
" \n",
"\n",
" # newton procedure\n",
" # semi-explicit warm start\n",
" vc = pin.difference(model, qp, qc) / timestep\n",
Expand All @@ -488,16 +493,16 @@
" break\n",
" Jf2 = jac_residual_fun(qp, qc, qn_guess, u)[2].full() # derivative in qnext\n",
" z = -np.linalg.solve(Jf2, fval)\n",
" \n",
"\n",
" # adaptive Newton\n",
" atau = 0.4\n",
" alpha = min(np.sqrt(2*atau/res_norm), 1.)\n",
" alpha = min(np.sqrt(2 * atau / res_norm), 1.0)\n",
" qn_guess[:] = pin.integrate(model, qn_guess, z)\n",
" \n",
"\n",
" conv = abs(max(fval)) < TOL\n",
" if not conv:\n",
" print(\"Did not converge.\")\n",
" \n",
"\n",
" qn = qn_guess\n",
" states_.append(qn)\n",
" e = mechanicalEnergyFromQs(model, pdata, qc, qn)\n",
Expand Down Expand Up @@ -649,10 +654,11 @@
"source": [
"import matplotlib.pyplot as plt\n",
"import seaborn\n",
"\n",
"seaborn.set_style()\n",
"\n",
"times = np.linspace(0., Tf, int(Tf / timestep))\n",
"plt.plot(times, Em_vals, lw=1., ls='--')\n",
"times = np.linspace(0.0, Tf, int(Tf / timestep))\n",
"plt.plot(times, Em_vals, lw=1.0, ls=\"--\")\n",
"plt.title(\"Mechanical energy\")\n",
"plt.xlabel(\"Time $t$\")\n",
"plt.tight_layout()"
Expand Down Expand Up @@ -680,13 +686,13 @@
" Em_vals = []\n",
" for t in range(nsteps):\n",
" xc = states_[t]\n",
" qc, vc = xc[:model.nq], xc[model.nq:]\n",
" qc, vc = xc[: model.nq], xc[model.nq :]\n",
"\n",
" ac = pin.aba(model, pdata, qc, vc, tau=u)\n",
" vn = vc + ac * timestep\n",
" qn = pin.integrate(model, qc, vn * timestep)\n",
" xn = np.concatenate([qn, vn])\n",
" \n",
"\n",
" states_.append(xn)\n",
" e = pin.computeMechanicalEnergy(model, pdata, qn, vn)\n",
" Em_vals.append(e)\n",
Expand Down Expand Up @@ -752,7 +758,7 @@
"metadata": {},
"outputs": [],
"source": [
"q_traj_si = states_si[:model.nq]\n",
"q_traj_si = states_si[: model.nq]\n",
"\n",
"vizer.play(q_traj, timestep)"
]
Expand Down Expand Up @@ -836,8 +842,8 @@
"source": [
"# other comparison: angle theta0\n",
"plt.figure(dpi=100)\n",
"plt.plot(times, q_traj_si[0,:-1], label='semi-implicit')\n",
"plt.plot(times, q_traj[0,:-1], label='variational')\n",
"plt.plot(times, q_traj_si[0, :-1], label=\"semi-implicit\")\n",
"plt.plot(times, q_traj[0, :-1], label=\"variational\")\n",
"plt.legend()\n",
"plt.title(r\"Comparison: trajectory of angle $\\theta_0$\")\n",
"plt.xlabel(\"Time $t$\")"
Expand Down

0 comments on commit 178d996

Please sign in to comment.