{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with Diagrams\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "The [Modeling Dynamical Systems](./dynamical_systems.ipynb) tutorial gave a very basic introduction to Drake's Systems framework, including how to assemble multiple systems into a `Diagram`. In this notebook we'll provide a more advanced/complete overview of working with diagrams.\n",
    "\n",
    "Block diagrams are the standard modeling abstraction in systems theory and control. They provide powerful abstractions for writing modular code describing potentially very complicated dynamical systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pydot\n",
    "from IPython.display import SVG, display\n",
    "from pydrake.examples import PendulumPlant\n",
    "from pydrake.systems.controllers import PidController\n",
    "from pydrake.systems.framework import Diagram, DiagramBuilder, LeafSystem\n",
    "from pydrake.systems.primitives import AffineSystem, LogVectorOutput"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building and visualizing your Diagram\n",
    "\n",
    "Let's start by revisiting the `Diagram` example from the [introductory tutorial](./dynamical_systems.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def MakePidControlledPendulum():\n",
    "    # Use a DiagramBuilder to add and connect the subsystems.\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    # First add the pendulum.\n",
    "    pendulum = builder.AddNamedSystem(\"pendulum\", PendulumPlant())\n",
    "\n",
    "    # Add a PID controller.\n",
    "    controller = builder.AddNamedSystem(\"controller\",\n",
    "                                        PidController(kp=[10.], ki=[1.], kd=[1.]))\n",
    "\n",
    "    # Now \"wire up\" the controller to the plant.\n",
    "    builder.Connect(pendulum.get_state_output_port(),\n",
    "                    controller.get_input_port_estimated_state())\n",
    "    builder.Connect(controller.get_output_port_control(), pendulum.get_input_port())\n",
    "\n",
    "    # Make the desired_state input of the controller an input to the diagram.\n",
    "    builder.ExportInput(controller.get_input_port_desired_state())\n",
    "    # Make the pendulum state an output from the diagram.\n",
    "    builder.ExportOutput(pendulum.get_state_output_port())\n",
    "\n",
    "    # Log the state of the pendulum.\n",
    "    logger = LogVectorOutput(pendulum.get_state_output_port(), builder)\n",
    "    logger.set_name(\"logger\")\n",
    "\n",
    "    pid_controlled_pendulum = builder.Build()\n",
    "    pid_controlled_pendulum.set_name(\"PID-controlled Pendulum\")\n",
    "\n",
    "    return pid_controlled_pendulum, pendulum\n",
    "\n",
    "pid_controlled_pendulum, pendulum = MakePidControlledPendulum()\n",
    "\n",
    "# Visualize the diagram.\n",
    "display(\n",
    "    SVG(\n",
    "        pydot.graph_from_dot_data(\n",
    "            pid_controlled_pendulum.GetGraphvizString(\n",
    "                max_depth=2))[0].create_svg()))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### You can only add a System to one DiagramBuilder.\n",
    "\n",
    "One thing that is important to understand is that when you add a system to a builder, the builder takes ownership of that system.  You cannot add the same system to another builder. Currently, Drake will report a somewhat generic error about ownership of the C++ object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_builder = DiagramBuilder()\n",
    "\n",
    "try:\n",
    "    second_builder.AddSystem(pendulum)\n",
    "except RuntimeError as err:\n",
    "    print(err)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is conceptually straightforward, but it can happen easily in jupyter notebooks. If you define a system in one cell, and add it to a builder in a second cell, then running that second cell twice in a row will cause an error. For this reason, the recommended workflow is to construct the system and add it to the builder immediately (in the same cell), or to use a function to build the system, as we've done above.\n",
    "\n",
    "## Nested diagrams\n",
    "\n",
    "Diagrams are Systems, too, and can be used as subsystems to achieve modularity and encapsulation.  Imagine I want to use the same PID-controlled pendulum in another context, but want to shift the coordinate system by $\\pi$.  I could achieve that by adding more systems around the controlled pendulum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = DiagramBuilder()\n",
    "\n",
    "# We make the PID-controlled pendulum again here, to avoid C++ ownership issues.\n",
    "pid_controlled_pendulum, pendulum = MakePidControlledPendulum()\n",
    "builder.AddSystem(pid_controlled_pendulum)\n",
    "\n",
    "# Shift desired theta by PI with the system y= u + [pi;0]\n",
    "shift_desired = builder.AddNamedSystem(\n",
    "    \"shift_desired\", AffineSystem(D=np.eye(2), y0=[np.pi, 0]))\n",
    "builder.ExportInput(shift_desired.get_input_port(),\n",
    "                    \"controller_desired_state\")\n",
    "\n",
    "# Connect the shift to the *exported input port* from the subdiagram.\n",
    "builder.Connect(shift_desired.get_output_port(),\n",
    "                pid_controlled_pendulum.get_input_port())\n",
    "\n",
    "# Shift actual theta by -PI.\n",
    "shift_actual = builder.AddNamedSystem(\n",
    "    \"shift_actual\", AffineSystem(D=np.eye(2), y0=[-np.pi, 0]))\n",
    "builder.Connect(pid_controlled_pendulum.get_output_port(),\n",
    "                shift_actual.get_input_port())\n",
    "builder.ExportOutput(shift_actual.get_output_port(), \"pendulum_state\")\n",
    "\n",
    "diagram = builder.Build()\n",
    "diagram.set_name(\"PID-controlled Pendulum (with θ shifted by π)\")\n",
    "\n",
    "# Visualize the diagram (max depth=1)\n",
    "display(\n",
    "    SVG(\n",
    "        pydot.graph_from_dot_data(\n",
    "            diagram.GetGraphvizString(max_depth=1))[0].create_svg()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the diagram (max depth=2)\n",
    "display(\n",
    "    SVG(\n",
    "        pydot.graph_from_dot_data(\n",
    "            diagram.GetGraphvizString(max_depth=2))[0].create_svg()))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subsystems and contexts\n",
    "\n",
    "Now that we have a `Diagram`, we can work directly with it's `Context`. But depending on the order of e.g. the state variables in the diagrams' `Context` can be brittle (they may change if you add more systems to the `DiagramBuilder`). Moreover, the individual subsystems might provide methods that help you work with its `Context`.\n",
    "\n",
    "The pointers returned from `DiagramBuilder::AddSystem()`/`AddNamedSystem()` are still valid pointers to the subsystem, and can be used directly. In addition, the `Diagram` provides mechanisms for retrieving pointers to the subsystems that it contains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert(pid_controlled_pendulum.GetSubsystemByName(\"pendulum\") == pendulum)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More interesting is the subsystem `Context`s. You should understand that a diagram's `Context` is just an organized collection of the subsystem's `Context`s. You can access those subsystem `Context`s directly, most often using [`GetMyContextFromRoot()`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1systems_1_1_system.html#ae7fa91d2b2102457ced3361207724e52). Note that since diagrams can be subsystems in other diagrams, this method will recurse through the potentially nested diagram to recover the correct sub-`Context`.\n",
    "\n",
    "In this example, the `Context`s are nested just like the diagram:\n",
    "```\n",
    "  PID-controlled Pendulum (with θ shifted by π) Context (of a Diagram)\n",
    "    ↳ PID-controlled Pendulum Context (of a Diagram)\n",
    "        ↳ pendulum Context\n",
    "        ↳ controller Context\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diagram_context = diagram.CreateDefaultContext()\n",
    "print(diagram_context)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importantly, a subsystem `Context` retrieved from the diagram `Context` in this way is a pointer into the diagram's `Context`. Changing values into the subsystem `Context` changes the values in the root `Context`; in fact it is the recommended way to change those values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pendulum_context = pendulum.GetMyContextFromRoot(diagram_context)\n",
    "pendulum_context.SetContinuousState([1.2, 0.5])\n",
    "# The PendulumPlant class provides some helper methods for working with its\n",
    "# Context.\n",
    "pendulum_state = pendulum.get_state(pendulum_context)\n",
    "print(f\"θ = {pendulum_state.theta()}, θ̇ = {pendulum_state.thetadot()}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Observe that the root system's Context was also updated.\n",
    "print(diagram_context)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In summary, a `Context` for a `LeafSystem` is just a class containing the time, states, parameters, and inputs to the system, which may be accessed in almost any system method.  A `Context` for a `Diagram` provides the same abstraction (a `Diagram` is also a `System`), but additionally provides ways to operate on the sub-`Context`s of the sub-systems directly."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Exporting input and output ports\n",
    "\n",
    "When we [author a leaf system](./authoring_leaf_systems.ipynb), we declare the input and output ports. When we're assembling a `Diagram`, we export the input and output ports of the subsystems to make the inputs and output of the Diagram.  \n",
    "Don't try to access the ports of subsystems directly.\n",
    "\n",
    "One common mistake is to try to connect directly to the input/output ports of a subsystem directly.  Consider our nested `Diagram` example again:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = DiagramBuilder()\n",
    "\n",
    "pid_controlled_pendulum, pendulum = MakePidControlledPendulum()\n",
    "builder.AddSystem(pid_controlled_pendulum)\n",
    "\n",
    "# Shift desired theta by PI with the system y= u + [pi;0]\n",
    "shift_desired = builder.AddNamedSystem(\n",
    "    \"shift_desired\", AffineSystem(D=np.eye(2), y0=[np.pi, 0]))\n",
    "\n",
    "# ** WRONG ** This doesn't work, but gives a helpful message.\n",
    "controller = pid_controlled_pendulum.GetSubsystemByName(\"controller\")\n",
    "try:\n",
    "    builder.Connect(shift_desired.get_output_port(),\n",
    "                    controller.get_input_port_desired_state())\n",
    "except RuntimeError as err:\n",
    "    print(err)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we saw in the example above, the correct approach is to connect the shift to the *exported input port* from the `pid_controlled_pendulum` sub-diagram."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diagrams are Systems, too\n",
    "\n",
    "A `Diagram` implements the `System` class interface. To do this, it dispatches to the subsystem implementations and collects the results. Here is a small example showing how it works for time derivatives and publish events, but all of the `System` methods will perform similarly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyLeafSystem(LeafSystem):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.DeclareContinuousState(1)\n",
    "        self.DeclareForcedPublishEvent(self.Publish)\n",
    "        self.DeclarePeriodicPublishEvent(period_sec=1,\n",
    "                                         offset_sec=0,\n",
    "                                         publish=self.Publish)\n",
    "\n",
    "    def DoCalcTimeDerivatives(self, context, derivatives):\n",
    "        x = context.get_continuous_state_vector().GetAtIndex(0)\n",
    "        print(f\"{self.get_name()}: DoCalcTimeDerivatives()\")\n",
    "        derivatives.get_mutable_vector().SetAtIndex(0, -x)\n",
    "\n",
    "    def Publish(self, context):\n",
    "        print(f\"{self.get_name()}: Publish()\")\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "builder.AddNamedSystem(\"system1\", MyLeafSystem())\n",
    "builder.AddNamedSystem(\"system2\", MyLeafSystem())\n",
    "diagram = builder.Build()\n",
    "\n",
    "context = diagram.CreateDefaultContext()\n",
    "\n",
    "# To evaluate the time derivatives of the diagram, diagram evaluates the time\n",
    "# derivatives of the subsystems.\n",
    "print(\"diagram.EvalTimeDerivatives()\")\n",
    "diagram.EvalTimeDerivatives(context)\n",
    "\n",
    "# A ForcedPublish on the diagram calls ForcedPublish on the subsystems.\n",
    "print(\"diagram.ForcedPublish()\")\n",
    "diagram.ForcedPublish(context)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diagrams and scalar types (double, AutoDiffXd, symbolic::Expression)\n",
    "\n",
    "Diagrams support scalar types, and [scalar-type conversion](https://drake.mit.edu/doxygen_cxx/group__system__scalar__conversion.html) if all of their subsystems support the scalar types.  It is most common to build the `Diagram` with the default scalar first, and then convert with `ToAutoDiffXd()` and/or `ToSymbolic()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = DiagramBuilder()\n",
    "# AffineSystem is a primitive which supports all scalar types, and\n",
    "# scalar-conversion.\n",
    "builder.AddSystem(AffineSystem(y0=[2, 3]))\n",
    "builder.AddSystem(AffineSystem(D=np.eye(2)))\n",
    "diagram = builder.Build()\n",
    "\n",
    "diagram_autodiff = diagram.ToAutoDiffXd()\n",
    "diagram_symbolic = diagram.ToSymbolic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = DiagramBuilder()\n",
    "# MyLeafSystem (defined above) did not implement scalar support, so the\n",
    "# resulting Diagram will not, either.\n",
    "builder.AddSystem(MyLeafSystem())\n",
    "diagram = builder.Build()\n",
    "\n",
    "print(\"calling ToAutoDiffXd()\")\n",
    "try:\n",
    "    diagram_autodiff = diagram.ToAutoDiffXd()\n",
    "except RuntimeError as err:\n",
    "    print(err)\n",
    "\n",
    "print(\"\\ncalling ToSymbolic()\")\n",
    "try:\n",
    "    diagram_symbolic = diagram.ToSymbolic()\n",
    "except RuntimeError as err:\n",
    "    print(err)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See the [authoring leaf systems](./authoring_leaf_systems.ipynb) tutorial for details on how to add that support.\n",
    "\n",
    "# (Advanced) Subclassing from Diagram\n",
    "\n",
    "In most use cases, using `DiagramBuilder()` to construct the diagram will suffice. Sometimes, if you want provide more methods or member variables for the diagram itself, it can be useful to implement your own class which derives from `Diagram`.  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDiagram(Diagram):\n",
    "\n",
    "    def __init__(self):\n",
    "        Diagram.__init__(self)\n",
    "        builder = DiagramBuilder()\n",
    "        builder = DiagramBuilder()\n",
    "        self.system1 = builder.AddNamedSystem(\"system1\", MyLeafSystem())\n",
    "        self.system2 = builder.AddNamedSystem(\"system2\", MyLeafSystem())\n",
    "        # Instead of builder.Build(), we call\n",
    "        builder.BuildInto(self)\n",
    "\n",
    "    def get_system1(self):\n",
    "        return self.system1\n",
    "    \n",
    "    def get_system2(self):\n",
    "        return self.system2\n",
    "    \n",
    "diagram = CustomDiagram()\n",
    "context = diagram.CreateDefaultContext()\n",
    "print(context)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One example in Drake is `RobotDiagram` , which is simply a `Diagram` which is automatically populated with a `MultibodyPlant` and a `SceneGraph`, and provides convenience methods for retrieving them.\n"
   ]
  }
 ],
 "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.11.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}