{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e8bf9c66-d3db-483d-bd12-1f5ff2c6d0ed",
   "metadata": {},
   "source": [
    "# Coupled Masses and Springs\n",
    "\n",
    "Consider a 1D system of point bodies to which two springs are attached in series. The force acting on each body is simply:\n",
    "\n",
    "$$ F_i = k_{i, l} \\left[L_{i, l} - (x_i - x_{i - 1})\\right] - k_{i, r} \\left[L_{i, r} - (x_{i + 1} - x_i)\\right] $$\n",
    "\n",
    "Therefore, a reasonable approach in building a calculation is implementing the following equations:\n",
    "\n",
    "$$ \\begin{align} \\frac{dx}{dt} &= v \\\\ \n",
    "\\frac{dv}{dt} &= a \\\\ \n",
    "F &= ma \\end{align}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "speaking-mouth",
   "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)\n",
    "\n",
    "from networkx import DiGraph\n",
    "\n",
    "from stream import Aggregator, Calculation, unpacked\n",
    "from stream.aggregator import vars_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "outdoor-bosnia",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mass(Calculation):\n",
    "    def __init__(self, m, k_right, k_left, L_right, L_left, name):\n",
    "        self.name = name\n",
    "        self.m = m\n",
    "        self.k_left, self.k_right = k_left, k_right\n",
    "        self.L_left, self.L_right = L_left, L_right\n",
    "\n",
    "    @property\n",
    "    def mass_vector(self):\n",
    "        return True, True, False\n",
    "\n",
    "    @unpacked\n",
    "    def calculate(self, variables, *, x_left, x_right):\n",
    "        x, v, a = variables\n",
    "        left_spring = -self.k_left * (self.L_left - (x - x_left))\n",
    "        right_spring = -self.k_right * (self.L_right - (x_right - x))\n",
    "        total_force = -left_spring + right_spring\n",
    "        return self.load(dict(x=v, v=a, a=total_force - a * self.m))\n",
    "\n",
    "    def indices(self, variable, asking=None):\n",
    "        try:\n",
    "            return super().indices(variable)\n",
    "        except KeyError:\n",
    "            pass\n",
    "        return dict(x_left=0, x_right=1)[variable]\n",
    "\n",
    "    @property\n",
    "    def variables(self):\n",
    "        return dict(x=0, v=1, a=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824e74e9-c101-4f1c-8284-15bdcdd71493",
   "metadata": {},
   "outputs": [],
   "source": [
    "M1 = Mass(name=\"M1\", m=1.5, k_left=1.0, k_right=1.0, L_left=1.0, L_right=1.0)\n",
    "M2 = Mass(name=\"M2\", m=1.0, k_left=M1.k_right, k_right=1.0, L_left=M1.L_right, L_right=1.0)\n",
    "\n",
    "agr = Aggregator(\n",
    "    graph=DiGraph([(M1, M2, vars_(\"x_left\")), (M2, M1, vars_(\"x_right\"))]),\n",
    "    funcs={M1: dict(x_left=0.0), M2: dict(x_right=3.0)},\n",
    ")\n",
    "\n",
    "from stream.analysis.report import report\n",
    "\n",
    "report(agr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5a7cf94-95e2-4e41-9c45-bb7ce8f32f7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 2))\n",
    "agr.draw(\n",
    "    node_options=dict(\n",
    "        pos={M1: (0, 0), M2: (1, 0)},\n",
    "        labels={M1: \"M1\", M2: \"M2\"},\n",
    "        node_size=2000,\n",
    "        font_size=20,\n",
    "    ),\n",
    "    edge_options=dict(label_pos=0.3, font_size=18),\n",
    ")\n",
    "plt.box()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "honey-allowance",
   "metadata": {},
   "outputs": [],
   "source": [
    "y0 = np.array([1.2, 0.0, 0.0, 1.9, 0.0, 0.0])\n",
    "time = np.linspace(0, 20.0, 400)\n",
    "\n",
    "sol = agr.solve(y0=y0, time=time, yp0=agr.compute(y0, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "laden-batch",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(time, sol[:, 0] - 1, label=\"M1 position\")\n",
    "plt.plot(time, sol[:, 3] - 2, label=\"M2 position\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "enhanced-format",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(time, sol[:, 1], label=\"M1 velocity\")\n",
    "plt.plot(time, sol[:, 4], label=\"M2 velocity\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "assured-subsection",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(time, sol[:, 2] / M1.m, label=\"M1 force\")\n",
    "plt.plot(time, sol[:, 5] / M2.m, label=\"M2 force\")\n",
    "plt.grid()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "stream-env",
   "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.14.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}