{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Taking Derivatives of `MultibodyPlant` Computations w.r.t. Mass\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\n",
    "\n",
    "We create a simple `MultibodyPlant` (`MultibodyPlant_[float]`). This means you *could*\n",
    "parse it from a URDF / SDFormat file.\n",
    "We then convert the system to `AutoDiffXd`, and set the parameters and partial derivatives\n",
    "we desire. In this case, we want $\\frac{\\partial{\\boldsymbol{f}}}{\\partial{\\boldsymbol{m}}}$, where $\\boldsymbol{f}$ is just some arbitrary expression.\n",
    "\n",
    "In our case, we choose $\\boldsymbol{f}$ to be generalized forces at the default / home configuration.\n",
    "Also in this case, we choose $\\boldsymbol{m} = \\left[ m_1, m_2 \\right] \\in \\mathbb{R}_+^2$ just to show how to choose gradients for\n",
    "independent values.\n",
    "\n",
    "For related reading, see also:\n",
    "- [Underactuated: System Identification](http://underactuated.csail.mit.edu/sysid.html) - at present, this only presents the symbolic approach for `MultibodyPlant`.\n",
    "- [`Modeling Dynamical Systems`](./dynamical_systems.ipynb)\n",
    "- [`Mathematical Program MultibodyPlant Tutorial`](./mathematical_program_multibody_plant.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from pydrake.multibody.tree import SpatialInertia, UnitInertia\n",
    "from pydrake.multibody.plant import MultibodyPlant_, MultibodyPlant\n",
    "from pydrake.autodiffutils import AutoDiffXd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plant = MultibodyPlant(time_step=0.0)\n",
    "body1 = plant.AddRigidBody(\n",
    "    \"body1\",\n",
    "    M_BBo_B=SpatialInertia(\n",
    "        mass=2.0,\n",
    "        p_PScm_E=[0, 0, 0],\n",
    "        # N.B. Rotational inertia is unimportant for calculations\n",
    "        # in this notebook, and thus is arbitrarily chosen.\n",
    "        G_SP_E=UnitInertia(0.1, 0.1, 0.1),\n",
    "    ),\n",
    ")\n",
    "body2 = plant.AddRigidBody(\n",
    "    \"body2\",\n",
    "    M_BBo_B=SpatialInertia(\n",
    "        mass=0.5,\n",
    "        p_PScm_E=[0, 0, 0],\n",
    "        # N.B. Rotational inertia is unimportant for calculations\n",
    "        # in this notebook, and thus is arbitrarily chosen.\n",
    "        G_SP_E=UnitInertia(0.1, 0.1, 0.1),\n",
    "    ),\n",
    ")\n",
    "plant.Finalize()\n",
    "\n",
    "plant_ad = plant.ToScalarType[AutoDiffXd]()\n",
    "body1_ad = plant_ad.get_body(body1.index())\n",
    "body2_ad = plant_ad.get_body(body2.index())\n",
    "context_ad = plant_ad.CreateDefaultContext()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to take gradients with respect to particular parameters, so we populate those parameters now.\n",
    "\n",
    "For forward-mode automatic differentiation, we must ensure that specify our gradients\n",
    "according to our desired independent variables. In this case, we want $m_1$ and $m_2$\n",
    "to be independent, so we ensure their gradients are distinct unit vectors. (You could\n",
    "use [`InitializeAutoDiff`](https://drake.mit.edu/pydrake/pydrake.autodiffutils.html#pydrake.autodiffutils.InitializeAutoDiff) to do this, but\n",
    "we do it \"by hand\" here for illustration.)\n",
    "\n",
    "If we wanted, we could set more parameters (e.g. center-of-mass), but we'll stick to just mass for simplicity.\n",
    "\n",
    "It is important to note that every other \"autodiff-able\" quantity in this case (state,\n",
    "other parameters) are considered constant w.r.t. our independent values (mass), thus\n",
    "they will have 0-valued gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = AutoDiffXd(2.0, [1.0, 0.0])\n",
    "body1_ad.SetMass(context_ad, m1)\n",
    "\n",
    "m2 = AutoDiffXd(0.5, [0.0, 1.0])\n",
    "body2_ad.SetMass(context_ad, m2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generalized force, in the z-translation direction for each body $i$, should just be $(-m_i \\cdot g)$ with derivative $(-g)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_z_component(plant, body, v):\n",
    "    assert body.is_floating()\n",
    "    # Floating-base velocity dofs are organized as [angular velocity; translation velocity].\n",
    "    v_start = body.floating_velocities_start_in_v()\n",
    "    nv_pose = 6\n",
    "    rxyz_txyz = v[v_start:v_start + nv_pose]\n",
    "    assert len(rxyz_txyz) == nv_pose\n",
    "    txyz = rxyz_txyz[-3:]\n",
    "    z = txyz[2]\n",
    "    return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@np.vectorize\n",
    "def ad_to_string(x):\n",
    "    # Formats an array of AutoDiffXd elements to a string.\n",
    "    # Note that this implementation is for a scalar, but we use `np.vectorize` to\n",
    "    # effectively convert our array to `ndarray` of strings.\n",
    "    return f\"AutoDiffXd({x.value()}, derivatives={x.derivatives()})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tau_g = plant_ad.CalcGravityGeneralizedForces(context_ad)\n",
    "tau_g_z1 = get_z_component(plant_ad, body1_ad, tau_g)\n",
    "tau_g_z2 = get_z_component(plant_ad, body2_ad, tau_g)\n",
    "print(ad_to_string(tau_g_z1))\n",
    "print(ad_to_string(tau_g_z2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}