{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f4b7c42",
   "metadata": {},
   "source": [
    "# Parallel vehicle propagation\n",
    "\n",
    "This notebook runs two independent vehicle {py:class}`SubTrees <Karana.Dynamics.SubTree>` in parallel via {py:class}`Karana.Models.ParallelVehicles`. {py:class}`ParallelVehicles <Karana.Models.ParallelVehicles>` registers on a global {{ spcls }} and advances a list of vehicle {{ spcls }} instances during the global propagator's hop.\n",
    "\n",
    "This example uses small serial-chain vehicles so the setup is easy to see. Each vehicle has its own {py:class}`SubTree <Karana.Dynamics.SubTree>`, its own {{ spcls }}, its own integrator, and a recurring pre-hop event that sleeps for 0.1 seconds. The sleep events make the wall-clock difference between parallel and serial execution visible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4c3a3b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import atexit\n",
    "\n",
    "import numpy as np\n",
    "import Karana.Core as kc\n",
    "import Karana.Dynamics as kd\n",
    "import Karana.Math as km\n",
    "import Karana.Integrators as ki\n",
    "import Karana.Models as kmdl\n",
    "from Karana.KUtils.Sim import Sim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af35c92d",
   "metadata": {},
   "source": [
    "## Build one simulation with two vehicle propagators\n",
    "\n",
    "The helper below creates one simple serial-chain vehicle, creates a vehicle {{ spcls }} for that vehicle's {py:class}`SubTree <Karana.Dynamics.SubTree>`, and registers a recurring pre-hop sleep event. It returns the vehicle {{ spcls }}.\n",
    "\n",
    "The notebook uses one {py:class}`Karana.KUtils.Sim.Sim` for the entire example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d8a6af35",
   "metadata": {},
   "outputs": [],
   "source": [
    "bc = kc.BaseContainer.singleton()\n",
    "\n",
    "\n",
    "def createVehicleStatePropagator(\n",
    "    sim: Sim,\n",
    "    prefix: str,\n",
    "    integrator: ki.IntegratorType,\n",
    "    sleep_time: float = 0.1,\n",
    ") -> kd.StatePropagator:\n",
    "    \"\"\"Create one vehicle and its state propagator.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    sim : Sim\n",
    "        Simulation that owns the global multibody and global state propagator.\n",
    "    prefix : str\n",
    "        Prefix used to name the vehicle bodies, subtree, and timed event.\n",
    "    integrator : ki.IntegratorType\n",
    "        Integrator type to use for the vehicle state propagator.\n",
    "    sleepTime : float\n",
    "        Wall-clock sleep duration, in seconds, for each recurring pre-hop event.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    kd.StatePropagator\n",
    "        State propagator for the vehicle subtree.\n",
    "    \"\"\"\n",
    "    # Create the bodies for the vehicle\n",
    "    kd.PhysicalBody.addSerialChain(\n",
    "        f\"{prefix}_bd\",\n",
    "        3,\n",
    "        sim.mb.virtualRoot(),\n",
    "        kd.PhysicalBodyParams(\n",
    "            spatial_inertia=km.SpatialInertia(\n",
    "                2.0, np.array([0.3, 0.2, 0.1]), np.diag([3.0, 2.0, 1.0])\n",
    "            ),\n",
    "            body_to_joint_transform=km.HomTran(\n",
    "                km.UnitQuaternion(0.8, 0.6, 0.0, 0.0), np.array([0.2, 0.3, 0.4])\n",
    "            ),\n",
    "            hinge_params=kd.PhysicalHingeParams(\n",
    "                hinge_type=kd.HingeType.REVOLUTE,\n",
    "                subhinge_params=[kd.PinSubhingeParams(unit_axis=np.array([1.0, 0.0, 0.0]))],\n",
    "            ),\n",
    "            parent_body_attachment_params=kd.ParentPhysicalBodyAttachmentParams(\n",
    "                inb_to_joint_transform=km.HomTran(\n",
    "                    km.UnitQuaternion(0.5, 0.5, 0.5, 0.5),\n",
    "                    np.array([0.4, 0.2, 0.3]),\n",
    "                )\n",
    "            ),\n",
    "            scene_part_specs=[],\n",
    "            scene_file_object_specs=[],\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    sim.mb.ensureHealthy()\n",
    "    sim.mb.resetData()\n",
    "\n",
    "    # Create the SubTree for the vehicle\n",
    "    subtree = kd.SubTree(\n",
    "        prefix,\n",
    "        sim.mb,\n",
    "        sim.mb.virtualRoot(),\n",
    "        [sim.mb.getBody(f\"{prefix}_bd_0\")],\n",
    "    )\n",
    "\n",
    "    # Create the StatePropagator for the vehicle\n",
    "    sp = kd.StatePropagator(subtree)\n",
    "    bc.at_exit_fns.insertBefore(\n",
    "        f\"_discard_sp_{sim.sp.id()}\",\n",
    "        f\"_discard_vsp_{sp.id()}\",\n",
    "        lambda: kc.discard(sp),\n",
    "    )\n",
    "    sp.setIntegrator(integrator)\n",
    "    sp.ensureHealthy()\n",
    "    sp.setTime(0.0)\n",
    "    sp.hardReset()\n",
    "    sp.setState(sp.assembleState())\n",
    "    sp.setMaxHopSize(1.0)\n",
    "\n",
    "    # Register the preHop sleep event\n",
    "    def sleepPreHop(t, _) -> None:\n",
    "        \"\"\"Sleep before a vehicle hop.\"\"\"\n",
    "        time.sleep(sleep_time)\n",
    "\n",
    "    sp.fns.pre_hop_fns[\"sleep\"] = sleepPreHop\n",
    "\n",
    "    return sp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9696e3d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = Sim()\n",
    "\n",
    "sp1 = createVehicleStatePropagator(sim, \"v1\", ki.IntegratorType.RK4)\n",
    "sp2 = createVehicleStatePropagator(sim, \"v2\", ki.IntegratorType.CVODE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac75630d",
   "metadata": {},
   "source": [
    "## Register `ParallelVehicles`\n",
    "\n",
    "Size the shared {py:class}`Karana.Core.ThreadPool` with {py:meth}`ThreadPool.setThreadPoolSize <Karana.Core.ThreadPool.setThreadPoolSize>`, register `ParallelVehicles` on the global {{ spcls }}, and set the global propagator to {py:attr}`Karana.Integrators.IntegratorType.NO_OP`. The global {{ spcls }} coordinates the vehicle {py:class}`state propagators <Karana.Dynamics.StatePropagator>`; the vehicle {py:class}`state propagators <Karana.Dynamics.StatePropagator>` do the integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d467b3b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Size the thread pool to be at least the number of vehicles + 1\n",
    "kc.ThreadPool.singleton().setThreadPoolSize(3)\n",
    "\n",
    "# Create the ParallelVehicles KModel on the global StatePropagator (sim.sp)\n",
    "parallel_vehicles = kmdl.ParallelVehicles(\"parallel_vehicles\", sim.sp, [sp1, sp2])\n",
    "del sp1, sp2\n",
    "\n",
    "# Set the global state propagator to be NO_OP\n",
    "sim.sp.setIntegrator(ki.IntegratorType.NO_OP)\n",
    "sim.sp.ensureHealthy()\n",
    "sim.sp.hardReset()\n",
    "sim.sp.setState(sim.sp.assembleState())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9645bc93",
   "metadata": {},
   "source": [
    "## Run in parallel mode\n",
    "\n",
    "First run the simulation for 5 seconds with. This lets the vehicles run in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d61cba95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'mode': 'parallel', 'elapsed_s': 0.5023632100783288}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tic = time.perf_counter()\n",
    "sim.sp.advanceBy(5.0)\n",
    "parallel_elapsed = time.perf_counter() - tic\n",
    "\n",
    "parallel_result = {\"mode\": \"parallel\", \"elapsed_s\": parallel_elapsed}\n",
    "parallel_result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fba97f6",
   "metadata": {},
   "source": [
    "## Run the same simulation in serial mode\n",
    "\n",
    "Now run the same `Sim` for another 5 seconds with {py:meth}`ParallelVehicles.setRunInSerial <Karana.Models.ParallelVehicles.setRunInSerial>` set to `True`. This runs the same simulation, but executes the vehicles in serial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "25537528",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'mode': 'serial', 'elapsed_s': 1.0018136049620807}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parallel_vehicles.setRunInSerial(True)\n",
    "\n",
    "tic = time.perf_counter()\n",
    "sim.sp.advanceBy(5.0)\n",
    "serial_elapsed = time.perf_counter() - tic\n",
    "\n",
    "serial_result = {\"mode\": \"serial\", \"elapsed_s\": serial_elapsed}\n",
    "serial_result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45bf165c",
   "metadata": {},
   "source": [
    "## Compare results\n",
    "\n",
    "Both timing measurements cover a 5 second simulation interval. The first interval advances from 0 to 5 seconds in parallel mode, and the second advances from 5 to 10 seconds in serial mode. Because each vehicle has a recurring 0.1 second pre-hop sleep event, the parallel run should take less wall-clock time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9df0bfc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "parallel: elapsed=0.502 s, \n",
      "serial: elapsed=1.002 s, \n",
      "parallel speedup over serial: 1.99x\n"
     ]
    }
   ],
   "source": [
    "speedup = serial_result[\"elapsed_s\"] / parallel_result[\"elapsed_s\"]\n",
    "\n",
    "for result in (parallel_result, serial_result):\n",
    "    print(f\"{result['mode']}: elapsed={result['elapsed_s']:.3f} s, \")\n",
    "\n",
    "print(f\"parallel speedup over serial: {speedup:.2f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1080b61",
   "metadata": {},
   "source": [
    "## Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "083ddd72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.cleanup()>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def cleanup():\n",
    "    \"\"\"Clean up the simulation.\"\"\"\n",
    "    global sim, parallel_vehicles\n",
    "    del sim, parallel_vehicles\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  }
 ],
 "metadata": {
  "docs": {
   "tags": [
    "time domain",
    "parallelization"
   ]
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
