foot-hits-ground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python\n",
    "# -*- coding: utf-8 -*-\n",
    "\n",
    "# external libraries\n",
    "from sympy import symbols, exp, acos, pi, Function\n",
    "import sympy.physics.mechanics as me\n",
    "from pydy.viz import VisualizationFrame, Cylinder, Sphere\n",
    "time_symbol = symbols('t', real=True)\n",
    "me.dynamicsymbols._t = time_symbol\n",
    "\n",
    "# Most of the Symbols are positive and real.\n",
    "sym_kwargs = {'positive': True,\n",
    "              'real': True}\n",
    "\n",
    "\n",
    "def time_varying(sym_string):\n",
    "    \"\"\"Retuns functions of time given a standard string for\n",
    "    sympy.symbols().\"\"\"\n",
    "\n",
    "    # TODO : Would it be better if these functions had the real=True\n",
    "    # assumption?\n",
    "\n",
    "    funcs = symbols(sym_string, cls=Function)\n",
    "\n",
    "    try:\n",
    "        return tuple([f(time_symbol) for f in funcs])\n",
    "    except TypeError:\n",
    "        return funcs(time_symbol)\n",
    "\n",
    "\n",
    "class BodySegment(object):\n",
    "\n",
    "    viz_sphere_radius = 0.07  # meters\n",
    "    viz_cylinder_radius = 0.035  # meters\n",
    "\n",
    "    def __init__(self, label, description, parent_reference_frame,\n",
    "                 origin_joint, joint_description, inertial_frame):\n",
    "        \"\"\"Initializes a body segment.\n",
    "        Parameters\n",
    "        ==========\n",
    "        label : string\n",
    "            A short label for the segment, like 'A'.\n",
    "        description : string\n",
    "            A short description of the segment, like 'Trunk'.\n",
    "        parent_reference_frame : sympy.physics.vector.ReferenceFrame\n",
    "            The parent reference frame for this segment.\n",
    "        origin_joint : sympy.physics.vector.Point\n",
    "            The point where this segment is connected to its parent.\n",
    "        joint_description : string\n",
    "            A short description of a new joint point, e.g. 'knee', for this\n",
    "            segemt to attach to its child segment.\n",
    "        inertial_frame : sympy.physics.mechanics.ReferenceFrame\n",
    "            The global inertial reference frame of the system. This is used\n",
    "            to apply gravity to the segment (in the negative y direction of\n",
    "            this frame).\n",
    "        \"\"\"\n",
    "\n",
    "        self.label = label\n",
    "        self.description = description\n",
    "        self.parent_reference_frame = parent_reference_frame\n",
    "        self.origin_joint = origin_joint\n",
    "        self.joint_description = joint_description\n",
    "        self.inertial_frame = inertial_frame\n",
    "\n",
    "        self._create_symbols()\n",
    "        self._kinematic_differential_equations()\n",
    "        self._orient()\n",
    "        self._set_angular_velocity()\n",
    "        self._locate_joint()\n",
    "        self._locate_mass_center()\n",
    "        self._set_linear_velocities()\n",
    "        self._inertia_dyadic()\n",
    "        self._create_rigid_body()\n",
    "        self._joint_torque()\n",
    "        self._gravity()\n",
    "\n",
    "    def _create_symbols(self):\n",
    "        \"\"\"Generates all of the SymPy symbols and functions of time\n",
    "        associated with this segment.\"\"\"\n",
    "\n",
    "        subscript = self.label.lower()\n",
    "\n",
    "        self.t = time_symbol\n",
    "\n",
    "        # constants\n",
    "        self.g = symbols('g', **sym_kwargs)\n",
    "        self.mass_symbol = symbols('m{}'.format(subscript), **sym_kwargs)\n",
    "        self.inertia_symbol = \\\n",
    "            symbols('i{}'.format(subscript), **sym_kwargs)\n",
    "        self.length_symbol = \\\n",
    "            symbols('l{}'.format(subscript), **sym_kwargs)\n",
    "        self.mass_center_x_symbol = symbols('x{}'.format(subscript), real=True)\n",
    "        self.mass_center_y_symbol = symbols('y{}'.format(subscript), real=True)\n",
    "\n",
    "        self.constants = [self.g,\n",
    "                          self.mass_symbol,\n",
    "                          self.inertia_symbol,\n",
    "                          self.length_symbol,\n",
    "                          self.mass_center_x_symbol,\n",
    "                          self.mass_center_y_symbol]\n",
    "\n",
    "        # functions of time\n",
    "        self.generalized_coordinate_symbol = \\\n",
    "            time_varying('q{}'.format(subscript))\n",
    "        self.generalized_coordinate_derivative_symbol = \\\n",
    "            self.generalized_coordinate_symbol.diff(self.t)\n",
    "        self.generalized_speed_symbol = time_varying('u{}'.format(subscript))\n",
    "        self.joint_torque_symbol = time_varying('T{}'.format(subscript))\n",
    "\n",
    "    def _kinematic_differential_equations(self):\n",
    "        \"\"\"Creates a list of the kinematic differential equations. This is\n",
    "        the simple definition:\n",
    "        0 = \\dot{q}_i - u_i\n",
    "        \"\"\"\n",
    "        self.kinematic_equations = \\\n",
    "            [self.generalized_coordinate_derivative_symbol -\n",
    "             self.generalized_speed_symbol]\n",
    "\n",
    "    def _orient(self):\n",
    "        \"\"\"Generates and orients the segment's reference frame relative to\n",
    "        the parent reference frame by body fixed simple rotation about the\n",
    "        generalized coordinate.\"\"\"\n",
    "        self.reference_frame = \\\n",
    "            self.parent_reference_frame.orientnew(\n",
    "                self.label, 'Axis', (self.generalized_coordinate_symbol,\n",
    "                                     self.parent_reference_frame.z))\n",
    "\n",
    "    def _set_angular_velocity(self):\n",
    "        \"\"\"Sets the angular velocity with the generalized speed.\"\"\"\n",
    "        self.reference_frame.set_ang_vel(self.parent_reference_frame,\n",
    "                                         self.generalized_speed_symbol *\n",
    "                                         self.parent_reference_frame.z)\n",
    "\n",
    "    def _locate_joint(self):\n",
    "        \"\"\"Creates a point with respect to the origin joint for the next\n",
    "        joint in the segment. This assumes that new joint is in the negative\n",
    "        y direction with repect to the origin joint.\"\"\"\n",
    "        self.joint = self.origin_joint.locatenew(self.joint_description,\n",
    "                                                 -self.length_symbol *\n",
    "                                                 self.reference_frame.y)\n",
    "\n",
    "    def _locate_mass_center(self):\n",
    "        \"\"\"Creates a point with respect the origin joint for the mass center\n",
    "        of the segment.\"\"\"\n",
    "        self.mass_center = self.origin_joint.locatenew(\n",
    "            '{} mass center'.format(self.description),\n",
    "            self.mass_center_x_symbol * self.reference_frame.x +\n",
    "            self.mass_center_y_symbol * self.reference_frame.y)\n",
    "\n",
    "    def _set_linear_velocities(self):\n",
    "        \"\"\"Sets the linear velocities of the mass center and new joint.\"\"\"\n",
    "        self.mass_center.v2pt_theory(self.origin_joint, self.inertial_frame,\n",
    "                                     self.reference_frame)\n",
    "        self.joint.v2pt_theory(self.origin_joint, self.inertial_frame,\n",
    "                               self.reference_frame)\n",
    "\n",
    "    def _inertia_dyadic(self):\n",
    "        \"\"\"Creates an inertia dyadic for the segment.\"\"\"\n",
    "        self.inertia_dyadic = me.inertia(self.reference_frame, 0, 0,\n",
    "                                         self.inertia_symbol)\n",
    "\n",
    "    def _create_rigid_body(self):\n",
    "        \"\"\"Creates a rigid body for the segment.\"\"\"\n",
    "        self.rigid_body = me.RigidBody(self.description, self.mass_center,\n",
    "                                       self.reference_frame,\n",
    "                                       self.mass_symbol,\n",
    "                                       (self.inertia_dyadic, self.mass_center))\n",
    "\n",
    "    def _joint_torque(self):\n",
    "        \"\"\"Creates the joint torque vector acting on the segment.\"\"\"\n",
    "        self.torque = self.joint_torque_symbol * self.reference_frame.z\n",
    "        # TODO : add in passive joint stiffness and damping\n",
    "\n",
    "    def _gravity(self):\n",
    "        \"\"\"Creates the gravitational force vector acting on the segment.\"\"\"\n",
    "        self.gravity = -self.mass_symbol * self.g * self.inertial_frame.y\n",
    "\n",
    "    def visualization_frames(self):\n",
    "        \"\"\"Returns visualization frames for the animation of the system.\n",
    "        The requires numerical values of the cylinders and spheres.\"\"\"\n",
    "\n",
    "        viz_frames = []\n",
    "\n",
    "        cylinder = Cylinder(color='red',\n",
    "                            length=self.length_symbol,\n",
    "                            radius=self.viz_cylinder_radius)\n",
    "\n",
    "        center_point = self.origin_joint.locatenew('Cylinder Center',\n",
    "                                                   -self.length_symbol / 2 *\n",
    "                                                   self.reference_frame.y)\n",
    "\n",
    "        viz_frames.append(VisualizationFrame('VizFrame',\n",
    "                                             self.reference_frame,\n",
    "                                             center_point, cylinder))\n",
    "\n",
    "        viz_frames.append(VisualizationFrame('OriginJointFrame',\n",
    "                                             self.reference_frame,\n",
    "                                             self.origin_joint,\n",
    "                                             Sphere(color='blue',\n",
    "                                                    radius=self.viz_sphere_radius)))\n",
    "\n",
    "        return viz_frames\n",
    "\n",
    "\n",
    "class TrunkSegment(BodySegment):\n",
    "    def __init__(self, *args):\n",
    "        super(TrunkSegment, self).__init__(*args)\n",
    "        self._trunk_extra_kinematic_equations()\n",
    "\n",
    "    def _create_symbols(self):\n",
    "        super(TrunkSegment, self)._create_symbols()\n",
    "        # TODO : Format these with the subscript instead of a directly.\n",
    "        self.qa = time_varying('qax, qay')\n",
    "        self.ua = time_varying('uax, uay')\n",
    "        self.constants.remove(self.length_symbol)\n",
    "        del self.length_symbol\n",
    "\n",
    "    def _trunk_extra_kinematic_equations(self):\n",
    "        qaxd, qayd = [f.diff(self.t) for f in self.qa]\n",
    "        self.kinematic_equations += [self.ua[0] - qaxd, self.ua[1] - qayd]\n",
    "\n",
    "    def _locate_joint(self):\n",
    "        \"\"\"The trunk only has one joint, the hip, there is no other point.\"\"\"\n",
    "        # This locates the hip joint relative to the ground origin point.\n",
    "        self.joint = self.origin_joint.locatenew(self.joint_description,\n",
    "                                                 self.qa[0] *\n",
    "                                                 self.inertial_frame.x +\n",
    "                                                 self.qa[1] *\n",
    "                                                 self.inertial_frame.y)\n",
    "\n",
    "    def _locate_mass_center(self):\n",
    "        \"\"\"Creates a point with respect the hip joint for the mass center\n",
    "        of the segment.\"\"\"\n",
    "        self.mass_center = self.joint.locatenew(\n",
    "            '{} mass center'.format(self.description),\n",
    "            self.mass_center_x_symbol * self.reference_frame.x +\n",
    "            self.mass_center_y_symbol * self.reference_frame.y)\n",
    "\n",
    "    def _set_linear_velocities(self):\n",
    "        \"\"\"Sets the linear velocities of the mass center and new joint.\"\"\"\n",
    "        # The joint is the hip. The origin joint is the ground's origin.\n",
    "        self.joint.set_vel(self.inertial_frame, self.ua[0] *\n",
    "                           self.inertial_frame.x + self.ua[1] *\n",
    "                           self.inertial_frame.y)\n",
    "        self.mass_center.v2pt_theory(self.joint, self.inertial_frame,\n",
    "                                     self.reference_frame)\n",
    "\n",
    "    def visualization_frames(self):\n",
    "        \"\"\"This should go from the hip to the mass center.\"\"\"\n",
    "\n",
    "        viz_frames = []\n",
    "\n",
    "        hip_to_mc_vector = self.mass_center.pos_from(self.joint)\n",
    "\n",
    "        cylinder = Cylinder(color='red', length=hip_to_mc_vector.magnitude(),\n",
    "                            radius=self.viz_cylinder_radius)\n",
    "\n",
    "        center_point = \\\n",
    "            self.joint.locatenew('Cylinder Center',\n",
    "                                 hip_to_mc_vector.magnitude() / 2 *\n",
    "                                 hip_to_mc_vector.normalize())\n",
    "\n",
    "        viz_frames.append(VisualizationFrame('VizFrame',\n",
    "                                             self.reference_frame,\n",
    "                                             center_point, cylinder))\n",
    "\n",
    "        viz_frames.append(VisualizationFrame('MassCenterFrame',\n",
    "                                             self.reference_frame,\n",
    "                                             self.mass_center,\n",
    "                                             Sphere(color='blue',\n",
    "                                                    radius=self.viz_sphere_radius)))\n",
    "\n",
    "        return viz_frames\n",
    "\n",
    "\n",
    "class FootSegment(BodySegment):\n",
    "\n",
    "    viz_sphere_radius = 0.03\n",
    "    viz_cylinder_radius = 0.01\n",
    "\n",
    "    def __init__(self, *args):\n",
    "        super(FootSegment, self).__init__(*args)\n",
    "        self._locate_foot_points()\n",
    "        self._set_foot_linear_velocities()\n",
    "\n",
    "    def _create_symbols(self):\n",
    "        super(FootSegment, self)._create_symbols()\n",
    "        self.heel_distance = symbols('hx{}'.format(self.label.lower()),\n",
    "                                     real=True)\n",
    "        self.toe_distance = symbols('tx{}'.format(self.label.lower()),\n",
    "                                    real=True)\n",
    "        self.foot_depth = symbols('fy{}'.format(self.label.lower()),\n",
    "                                  real=True)\n",
    "\n",
    "        self.constants.remove(self.length_symbol)\n",
    "        del self.length_symbol\n",
    "\n",
    "        self.constants += [self.heel_distance, self.toe_distance,\n",
    "                           self.foot_depth]\n",
    "\n",
    "    def _locate_joint(self):\n",
    "        \"\"\"The foot has no joint.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def _locate_foot_points(self):\n",
    "\n",
    "        self.heel = self.origin_joint.locatenew(\n",
    "            '{} heel'.format(self.description), self.heel_distance *\n",
    "            self.reference_frame.x + self.foot_depth *\n",
    "            self.reference_frame.y)\n",
    "\n",
    "        self.toe = self.origin_joint.locatenew(\n",
    "            '{} toe'.format(self.description), self.toe_distance *\n",
    "            self.reference_frame.x + self.foot_depth *\n",
    "            self.reference_frame.y)\n",
    "\n",
    "    def _set_linear_velocities(self):\n",
    "        \"\"\"There is no joint so pass this.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def _set_foot_linear_velocities(self):\n",
    "        \"\"\"Sets the linear velocities of the mass center and new joint.\"\"\"\n",
    "        self.mass_center.v2pt_theory(self.origin_joint, self.inertial_frame,\n",
    "                                     self.reference_frame)\n",
    "        self.heel.v2pt_theory(self.origin_joint, self.inertial_frame,\n",
    "                              self.reference_frame)\n",
    "        self.toe.v2pt_theory(self.origin_joint, self.inertial_frame,\n",
    "                             self.reference_frame)\n",
    "\n",
    "    def visualization_frames(self):\n",
    "        \"\"\"Returns a list of visualization frames needed to visualize the\n",
    "        foot.\"\"\"\n",
    "\n",
    "        viz_frames = []\n",
    "\n",
    "        heel_to_toe_length = self.toe.pos_from(self.heel).magnitude()\n",
    "        bottom_cylinder = Cylinder(color='red',\n",
    "                                   length=heel_to_toe_length,\n",
    "                                   radius=self.viz_cylinder_radius)\n",
    "        bottom_center_point = self.heel.locatenew('BottomCenter',\n",
    "                                                  heel_to_toe_length / 2 *\n",
    "                                                  self.reference_frame.x)\n",
    "        # Creates a reference frame with the Y axis pointing from the heel\n",
    "        # to the toe.\n",
    "        bottom_rf = self.reference_frame.orientnew('Bottom', 'Axis',\n",
    "                                                   (-pi / 2,\n",
    "                                                    self.reference_frame.z))\n",
    "        viz_frames.append(VisualizationFrame('BottomVizFrame',\n",
    "                                             bottom_rf,\n",
    "                                             bottom_center_point,\n",
    "                                             bottom_cylinder))\n",
    "\n",
    "        # top of foot\n",
    "        ankle_to_toe_vector = self.toe.pos_from(self.origin_joint)\n",
    "        top_cylinder = Cylinder(color='red',\n",
    "                                length=ankle_to_toe_vector.magnitude(),\n",
    "                                radius=self.viz_cylinder_radius)\n",
    "        angle = -acos(ankle_to_toe_vector.normalize().dot(bottom_rf.y))\n",
    "        top_foot_rf = bottom_rf.orientnew('Top', 'Axis', (angle,\n",
    "                                                          bottom_rf.z))\n",
    "        top_foot_center_point = \\\n",
    "            self.origin_joint.locatenew('TopCenter',\n",
    "                                        ankle_to_toe_vector.magnitude() / 2\n",
    "                                        * top_foot_rf.y)\n",
    "        viz_frames.append(VisualizationFrame('TopVizFrame', top_foot_rf,\n",
    "                                             top_foot_center_point,\n",
    "                                             top_cylinder))\n",
    "\n",
    "        # back of foot\n",
    "        heel_to_ankle_vector = self.origin_joint.pos_from(self.heel)\n",
    "        back_cylinder = Cylinder(color='red',\n",
    "                                 length=heel_to_ankle_vector.magnitude(),\n",
    "                                 radius=self.viz_cylinder_radius)\n",
    "        angle = acos(heel_to_ankle_vector.normalize().dot(bottom_rf.y))\n",
    "        back_foot_rf = bottom_rf.orientnew('Back', 'Axis', (angle,\n",
    "                                                            bottom_rf.z))\n",
    "        back_foot_center_point = \\\n",
    "            self.heel.locatenew('BackCenter',\n",
    "                                heel_to_ankle_vector.magnitude() / 2 *\n",
    "                                back_foot_rf.y)\n",
    "        viz_frames.append(VisualizationFrame('BackVizFrame', back_foot_rf,\n",
    "                                             back_foot_center_point,\n",
    "                                             back_cylinder))\n",
    "\n",
    "        # spheres for the ankle, toe, and heel\n",
    "        viz_frames.append(VisualizationFrame('AnkleVizFrame',\n",
    "                                             self.reference_frame,\n",
    "                                             self.origin_joint,\n",
    "                                             Sphere(color='blue',\n",
    "                                                    radius=self.viz_sphere_radius)))\n",
    "        viz_frames.append(VisualizationFrame('ToeVizFrame',\n",
    "                                             self.reference_frame, self.toe,\n",
    "                                             Sphere(color='blue',\n",
    "                                                    radius=self.viz_sphere_radius)))\n",
    "        viz_frames.append(VisualizationFrame('HeelVizFrame',\n",
    "                                             self.reference_frame, self.heel,\n",
    "                                             Sphere(color='blue',\n",
    "                                                    radius=self.viz_sphere_radius)))\n",
    "\n",
    "        return viz_frames\n",
    "\n",
    "\n",
    "def contact_force(point, ground, origin):\n",
    "    \"\"\"Returns a contact force vector acting on the given point made of\n",
    "    friction along the contact surface and elastic force in the vertical\n",
    "    direction.\n",
    "    Parameters\n",
    "    ==========\n",
    "    point : sympy.physics.mechanics.Point\n",
    "        The point which the contact force should be computed for.\n",
    "    ground : sympy.physics.mechanics.ReferenceFrame\n",
    "        A reference frame which represents the inerital ground in 2D space.\n",
    "        The x axis defines the ground line and positive y is up.\n",
    "    origin : sympy.physics.mechanics.Point\n",
    "        An origin point located on the ground line.\n",
    "    Returns\n",
    "    =======\n",
    "    force : sympy.physics.mechanics.Vector\n",
    "        The contact force between the point and the ground.\n",
    "    \"\"\"\n",
    "    # This is the \"height\" of the point above the ground, where a negative\n",
    "    # value means that the point is below the ground.\n",
    "    y_location = point.pos_from(origin).dot(ground.y)\n",
    "\n",
    "    # The penetration into the ground is mathematically defined as:\n",
    "    #\n",
    "    #               { 0 if y_location > 0\n",
    "    # deformation = {\n",
    "    #               { abs(y_location) if y_location < 0\n",
    "    #\n",
    "\n",
    "    penetration = (abs(y_location) - y_location) / 2\n",
    "\n",
    "    velocity = point.vel(ground)\n",
    "\n",
    "    # The addition of \"- y_location\" here adds a small linear term to the\n",
    "    # cubic stiffness and creates a light attractive force torwards the\n",
    "    # ground. This is in place to ensure that gradients can be computed for\n",
    "    # the optimization used in Ackermann and van den Bogert 2010.\n",
    "    contact_stiffness, contact_damping = symbols('kc, cc', **sym_kwargs)\n",
    "    contact_friction_coefficient, friction_scaling_factor = \\\n",
    "        symbols('mu, vs', **sym_kwargs)\n",
    "\n",
    "    vertical_force = (contact_stiffness * penetration ** 3 - y_location) * \\\n",
    "        (1 - contact_damping * velocity.dot(ground.y))\n",
    "\n",
    "    friction = -contact_friction_coefficient * vertical_force * \\\n",
    "        ((2 / (1 + exp(-velocity.dot(ground.x) /\n",
    "                       friction_scaling_factor))) - 1)\n",
    "\n",
    "    return friction * ground.x + vertical_force * ground.y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}