Merge pull request #336 from moorepants/add-grf

Add ground reaction force plot to human gait example

examples-gallery/advanced/plot_human_gait.py

@@ -2,38 +2,46 @@
 Human Gait
 ==========
 
-This example replicates a similar solution as shown in [Ackermann2010]_ using
-joint torques as inputs instead of muscle activations [1]_.
-
-pygait2d and symmeplot and their dependencies must be installed first to run
-this example. Note that pygait2d has not been released to PyPi or Conda Forge::
+.. note::
 
-    conda install cython pip pydy pyyaml setuptools symmeplot sympy
-    python -m pip install --no-deps --no-build-isolation git+https://github.com/csu-hmc/gait2d
+    pygait2d and symmeplot and their dependencies must be installed first to
+    run this example. Note that pygait2d has not been released to PyPi or Conda
+    Forge::
 
-gait2d provides a joint torque driven 2D bipedal human dynamical model with
-seven body segments (trunk, thighs, shanks, feet) and foot-ground contact
-forces based on the description in [Ackermann2010]_.
+        conda install cython pip pydy pyyaml setuptools symmeplot sympy
+        python -m pip install --no-deps --no-build-isolation git+https://github.com/csu-hmc/gait2d
 
-The optimal control goal is to find the joint torques (hip, knee, ankle) that
-generate a minimal mean-torque periodic motion to ambulate at a specified
-average speed over half a period.
+Objectives
+----------
 
-This example highlights two points of interest that the other examples may not
-have:
+This example highlights three points of interest compared to the others:
 
 - Instance constraints are used to make the start state the same as the end
-  state, for example :math:`q_b(t_0) = q_e(t_f)`, except for forward
-  translation.
+  state except for forward translation to solve for a cyclic trajectory, for
+  example :math:`q_b(t_0) = q_e(t_f)`.
 - The average speed is constrained in this variable time step solution by
   introducing an additional differential equation that, when integrated, gives
   the duration at :math:`\Delta_t(t)`, which can be used to calculate distance
   traveled with :math:`q_{ax}(t_f) = v_\textrm{avg} (t_f - t_0)` and used as a
   constraint.
-- The parallel option is enabled because the equations of motion are on the
-  large side. This speeds up the evaluation of the constraints and its Jacobian
+- The parallel option is enabled because the equations of motion are relatively
+  large. This speeds up the evaluation of the constraints and its Jacobian
   about 1.3X.
 
+Introduction
+------------
+
+This example replicates a similar solution as shown in [Ackermann2010]_ using
+joint torques as inputs instead of muscle activations [1]_.
+
+gait2d provides a joint torque driven 2D bipedal human dynamical model with
+seven body segments (trunk, thighs, shanks, feet) and foot-ground contact
+forces based on the description in [Ackermann2010]_.
+
+The optimal control goal is to find the joint torques (hip, knee, ankle) that
+generate a minimal mean-torque periodic motion to ambulate at a specified
+average speed over half a period.
+
 Import all necessary modules, functions, and classes:
 """
 import os
@@ -46,6 +54,7 @@
 from symmeplot.matplotlib import Scene3D
 import matplotlib
 import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
 import numpy as np
 import sympy as sm
 
@@ -245,10 +254,13 @@ def animate():
     ground, origin, segments = symbolics[8], symbolics[9], symbolics[10]
     trunk, rthigh, rshank, rfoot, lthigh, lshank, lfoot = segments
 
-    fig, ax = plt.subplots(subplot_kw={'projection': '3d'})
+    fig = plt.figure(figsize=(10.0, 4.0))
+
+    ax3d = fig.add_subplot(1, 2, 1, projection='3d')
+    ax2d = fig.add_subplot(1, 2, 2)
 
     hip_proj = origin.locatenew('m', qax*ground.x)
-    scene = Scene3D(ground, hip_proj, ax=ax)
+    scene = Scene3D(ground, hip_proj, ax=ax3d)
 
     # creates the stick person
     scene.add_line([
@@ -304,14 +316,51 @@ def animate():
 
     scene.axes.set_proj_type("ortho")
     scene.axes.view_init(90, -90, 0)
-    scene.plot()
-
-    ax.set_xlim((-0.8, 0.8))
-    ax.set_ylim((-0.2, 1.4))
-    ax.set_aspect('equal')
-
-    ani = scene.animate(lambda i: gait_cycle[:, i], frames=len(times),
-                        interval=h_val*1000)
+    scene.plot(prettify=False)
+
+    ax3d.set_xlim((-0.8, 0.8))
+    ax3d.set_ylim((-0.2, 1.4))
+    ax3d.set_aspect('equal')
+    for axis in (ax3d.xaxis, ax3d.yaxis, ax3d.zaxis):
+        axis.set_ticklabels([])
+        axis.set_ticks_position("none")
+
+    eval_rforce = sm.lambdify(
+        states + specified + constants,
+        (contact_force(rfoot.toe, ground, origin) +
+         contact_force(rfoot.heel, ground, origin)).to_matrix(ground),
+        cse=True)
+
+    eval_lforce = sm.lambdify(
+        states + specified + constants,
+        (contact_force(lfoot.toe, ground, origin) +
+         contact_force(lfoot.heel, ground, origin)).to_matrix(ground),
+        cse=True)
+
+    rforces = np.array([eval_rforce(*gci).squeeze() for gci in gait_cycle.T])
+    lforces = np.array([eval_lforce(*gci).squeeze() for gci in gait_cycle.T])
+
+    ax2d.plot(times, rforces[:, :2], times, lforces[:, :2])
+    ax2d.grid()
+    ax2d.set_ylabel('Force [N]')
+    ax2d.set_xlabel('Time [s]')
+    ax2d.legend(['Horizontal GRF (r)', 'Vertical GRF (r)',
+                 'Horizontal GRF (l)', 'Vertical GRF (l)'], loc='upper right')
+    ax2d.set_title('Foot Ground Reaction Force Components')
+    vline = ax2d.axvline(times[0], color='black')
+
+    def update(i):
+        scene.evaluate_system(*gait_cycle[:, i])
+        scene.update()
+        vline.set_xdata([times[i], times[i]])
+        return scene.artists + (vline,)
+
+    ani = FuncAnimation(
+        fig,
+        update,
+        frames=range(len(times)),
+        interval=h_val*1000,
+    )
 
     return ani