{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa84188e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib import rc\n",
    "\n",
    "rc(\"font\", family=\"serif\", size=15)\n",
    "rc(\"savefig\", dpi=600)\n",
    "rc(\"figure\", figsize=(8, 6))\n",
    "rc(\"mathtext\", fontset=\"dejavuserif\")\n",
    "rc(\"lines\", linewidth=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "143d36df-a931-4d23-bd77-115d493e7a8b",
   "metadata": {},
   "source": [
    "# Planar Pendulum\n",
    "\n",
    "This example is taken from the `Scikits.odes` documentation. Here approximately the same example is given, but using the STREAM package which kindly depends on Scikits.odes.\n",
    "\n",
    "Imagine a planar (2D) pendulum whose string is of length $l$, at the end of which hangs a weight of mass $m$, given gravitational acceleration $g$. Friction is disregarded, as usual.\n",
    "\n",
    "The Lagrangian of such a problem is simply\n",
    "\n",
    "$$ L = \\frac{1}{2}m\\left(u^2+v^2\\right) - mgy $$\n",
    "\n",
    "Where $\\dot{x} = u, \\dot{y} = v$, and subject to the holonomic constraint\n",
    "\n",
    "$$ x^2 + y^2 = l^2 $$\n",
    "\n",
    "This constraint may be added by a Lagrange multiplier: \n",
    "\n",
    "$$ L = \\frac{1}{2}m\\left(u^2+v^2\\right) - mgy - \\frac{1}{2}\\lambda\\left(l^2 - x^2 - y^2\\right)$$\n",
    "\n",
    "Now using the Euler-Lagrange equation we arrive at 2 additional equation to a total of 4:\n",
    "\n",
    "$$\\begin{align}\\dot{x} &= u \\\\\n",
    "\\dot{y} &= v \\\\\n",
    "\\dot{u} &= \\lambda x / m \\\\\n",
    "\\dot{v} &= \\lambda y / m - g\\end{align}$$\n",
    "\n",
    "A final equation is the constraint itself (which is derived twice to a final form of $u^2 + v^2 + \\lambda l^2/m - gy = 0$, left to the reader as an exercise). Implementation is quite straightforward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "305b18e4-ec3c-4823-bb42-f868be07bac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from stream import Aggregator, Calculation\n",
    "from stream.units import Array1D\n",
    "\n",
    "l, m, g = 1.0, 1.0, 1.0\n",
    "\n",
    "\n",
    "class PendulumCalculation(Calculation):\n",
    "    name = \"Pendulum\"\n",
    "\n",
    "    def calculate(self, variables, **kwargs) -> Array1D:\n",
    "        x, y, u, v, λ = variables\n",
    "        k = λ / m\n",
    "        ẋ = u\n",
    "        ẏ = v\n",
    "        u̇ = k * x\n",
    "        v̇ = k * y - g\n",
    "        constraint = u**2 + v**2 + k * (x**2 + y**2) - y * g\n",
    "        return np.array([ẋ, ẏ, u̇, v̇, constraint])\n",
    "\n",
    "    @property\n",
    "    def mass_vector(self):\n",
    "        return True, True, True, True, False\n",
    "\n",
    "    @property\n",
    "    def variables(self):\n",
    "        return dict(x=0, y=1, u=2, v=3, λ=4)\n",
    "\n",
    "\n",
    "pendulum = PendulumCalculation()\n",
    "agr = Aggregator.from_decoupled(pendulum)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aa6cfec-e90a-4b05-a700-3f8631b74e8e",
   "metadata": {},
   "source": [
    "We'll give the same initial conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb018e0-21eb-432f-8f41-b4fe7ac23fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "λ0 = 0.1\n",
    "θ0 = np.pi / 3\n",
    "x0 = np.sin(θ0)\n",
    "y0 = -((l**2 - x0**2) ** 0.5)\n",
    "z0 = np.array([x0, y0, 0.0, 0.0, λ0])\n",
    "zp0 = np.array([0.0, 0.0, λ0 * x0 / m, (λ0 * y0 / m) - g, -g])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6494f5f7-b4a0-4e45-b044-50850975e47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "time = [0.0, 1.0, 2.0]\n",
    "solution = agr.solve(z0, time, zp0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "092691cf-795e-4715-b14a-16e691b9e10b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(solution.data, time, pendulum.variables.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "955192a4-ff1d-43b6-94a3-bae0412d362f",
   "metadata": {},
   "source": [
    "The 'x' position is given in the example and appears to be in total agreement. Up till here we've confirmed STREAM allows usage of ``Scikits.odes``. Now, let's make some sense of the solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2136208f-cfb0-4966-8f01-e79c33eda99b",
   "metadata": {},
   "outputs": [],
   "source": [
    "time = np.linspace(0, 10, 60)\n",
    "solution = agr.solve(z0, time, zp0)\n",
    "\n",
    "plt.figure(figsize=(6, 6))\n",
    "x, y = agr.at_times(solution, pendulum, \"x\"), agr.at_times(solution, pendulum, \"y\")\n",
    "plt.plot(x, -np.sqrt(l**2 - x**2), color=\"orange\", linestyle=\"dashed\", linewidth=2)\n",
    "plt.scatter(x, y, 30)\n",
    "plt.grid()\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.legend([\"Circular Constraint\", \"Solution\"])\n",
    "plt.show()"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}