{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f1ce1d-1a09-46d1-9304-3d02c715903f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Benchmarking...\n",
      "Single-process execution time: 37.67 seconds\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e29a55948eae404da54cadb0ff8396ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Running Simulations:   0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parallel execution time: 5.86 seconds\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from scipy.integrate import odeint\n",
    "from joblib import Parallel, delayed\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "from tqdm.notebook import tqdm\n",
    "from numba import jit\n",
    "\n",
    "# Accelerate the ODE system using numba\n",
    "@jit(nopython=True)\n",
    "def simple_oscillator(y, t):\n",
    "    x, v = y  # y[0] = position (x), y[1] = velocity (v)\n",
    "    dxdt = v  # dx/dt = v\n",
    "    dvdt = -x  # dv/dt = -x\n",
    "    return [dxdt, dvdt]\n",
    "\n",
    "# Function to simulate a batch of initial conditions\n",
    "def simulate_batch(batch):\n",
    "    t = np.linspace(0, 1000, 100000)  # Larger time steps for benchmarking\n",
    "    results = [(t, odeint(simple_oscillator, ic, t, atol=1e-9)) for ic in batch]\n",
    "    return results\n",
    "\n",
    "# Function to run simulations using joblib for parallelism\n",
    "def parallel_simulations(initial_conditions_list, batch_size=2):\n",
    "    # Split the initial conditions into batches to reduce overhead\n",
    "    batches = [\n",
    "        initial_conditions_list[i:i + batch_size]\n",
    "        for i in range(0, len(initial_conditions_list), batch_size)\n",
    "    ]\n",
    "    \n",
    "    # Run parallel simulations with joblib\n",
    "    results = Parallel(n_jobs=-1)(\n",
    "        delayed(simulate_batch)(batch) for batch in tqdm(batches, desc=\"Running Simulations\")\n",
    "    )\n",
    "    \n",
    "    # Flatten the list of batches into a single list of results\n",
    "    return [item for batch_result in results for item in batch_result]\n",
    "\n",
    "# Benchmarking function\n",
    "def benchmark(initial_conditions_list):\n",
    "    print(\"Benchmarking...\")\n",
    "\n",
    "    # Single-process execution\n",
    "    start_time = time.time()\n",
    "    single_process_results = [simulate_batch([ic])[0] for ic in initial_conditions_list]\n",
    "    single_process_time = time.time() - start_time\n",
    "    print(f\"Single-process execution time: {single_process_time:.2f} seconds\")\n",
    "\n",
    "    # Parallel execution with joblib\n",
    "    start_time = time.time()\n",
    "    parallel_results = parallel_simulations(initial_conditions_list, batch_size=2)\n",
    "    parallel_time = time.time() - start_time\n",
    "    print(f\"Parallel execution time: {parallel_time:.2f} seconds\")\n",
    "\n",
    "    return single_process_results, parallel_results\n",
    "\n",
    "# List of initial conditions for benchmarking\n",
    "initial_conditions_list = [\n",
    "    [1.0, 0.0],  # Initial position 1.0, velocity 0.0\n",
    "    [0.5, 0.5],  # Initial position 0.5, velocity 0.5\n",
    "    [-1.0, 0.0], # Initial position -1.0, velocity 0.0\n",
    "    [0.0, 1.0],  # Initial position 0.0, velocity 1.0\n",
    "    [2.0, -1.0]  # Initial position 2.0, velocity -1.0\n",
    "] * 400  # Increase the workload for meaningful benchmarking\n",
    "\n",
    "# Run the benchmark\n",
    "single_process_results, parallel_results = benchmark(initial_conditions_list)\n",
    "\n",
    "# Plotting all results on the same plot (using the parallel results)\n",
    "plt.figure(figsize=(10, 6))\n",
    "for i, (t, solution) in enumerate(parallel_results):\n",
    "    plt.plot(t, solution[:, 0], label=f'Initial Condition {i+1}')\n",
    "\n",
    "# Add labels, title, legend, and grid\n",
    "plt.xlabel('Time (t)')\n",
    "plt.ylabel('Position (x)')\n",
    "plt.title('ODE Simulations: Simple Harmonic Oscillator')\n",
    "plt.legend(loc='upper right', bbox_to_anchor=(1.15, 1))\n",
    "plt.grid(True)\n",
    "\n",
    "# Display the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "367b156c-a843-4fec-9d95-5aeeb2de38c6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03f1d22f-7efa-41de-9c8e-d165213172e2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}