{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1409d50c-9251-4f36-af8d-8f8481182601",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ebf9cbc-a151-4178-aba5-9f5f66c1209b",
   "metadata": {},
   "source": [
    "Manipulating the norm into an usable equation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "33137bb3-e298-455b-9174-27c005a84e45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}b \\frac{d}{d t} x{\\left(t \\right)} + x{\\left(t \\right)} + \\frac{d^{2}}{d t^{2}} x{\\left(t \\right)}\\\\b \\frac{d}{d t} y{\\left(t \\right)} + y{\\left(t \\right)} + \\frac{d^{2}}{d t^{2}} y{\\left(t \\right)}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[b*Derivative(x(t), t) + x(t) + Derivative(x(t), (t, 2))],\n",
       "[b*Derivative(y(t), t) + y(t) + Derivative(y(t), (t, 2))]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Working with their LHS of Eq 2:\n",
    "t = sm.symbols('t')\n",
    "x = sm.Function('x')(t)\n",
    "y = sm.Function('y')(t)\n",
    "\n",
    "#Create State Vector\n",
    "X = sm.Matrix([x, y])\n",
    "\n",
    "#Create EQ 2 LHS:\n",
    "b = sm.symbols('b')\n",
    "LHS = X.diff('t','t') + b*X.diff('t') + X\n",
    "LHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3edae373-94fd-4458-887a-8a01c220a1f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}\\frac{p_{n} \\left(x_{n} - x{\\left(t \\right)}\\right)}{\\left(h^{2} + \\left|{x_{n} - x{\\left(t \\right)}}\\right|^{2} + \\left|{y_{n} - y{\\left(t \\right)}}\\right|^{2}\\right)^{2.5}}\\\\\\frac{p_{n} \\left(y_{n} - y{\\left(t \\right)}\\right)}{\\left(h^{2} + \\left|{x_{n} - x{\\left(t \\right)}}\\right|^{2} + \\left|{y_{n} - y{\\left(t \\right)}}\\right|^{2}\\right)^{2.5}}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[p_n*(x_n - x(t))/(h**2 + Abs(x_n - x(t))**2 + Abs(y_n - y(t))**2)**2.5],\n",
       "[p_n*(y_n - y(t))/(h**2 + Abs(x_n - x(t))**2 + Abs(y_n - y(t))**2)**2.5]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Now, create the RHS for a given x_n and y_n of a magnet\n",
    "p_n, x_n, y_n, h = sm.symbols('p_n, x_n, y_n, h')\n",
    "X_n = sm.Matrix([x_n, y_n])\n",
    "\n",
    "RHS_gen = p_n * (X_n - X)/((X_n-X).norm()**2 + h**2)**(5/2)\n",
    "RHS_gen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a0920c7a-4abd-4f63-9804-a5207a6c2693",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}b \\frac{d}{d t} x{\\left(t \\right)} + x{\\left(t \\right)} + \\frac{d^{2}}{d t^{2}} x{\\left(t \\right)}\\\\b \\frac{d}{d t} y{\\left(t \\right)} + y{\\left(t \\right)} + \\frac{d^{2}}{d t^{2}} y{\\left(t \\right)}\\end{matrix}\\right] = \\left[\\begin{matrix}- \\frac{x{\\left(t \\right)} - 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} - 1}\\right|^{2} + \\left|{y{\\left(t \\right)} - 1}\\right|^{2}\\right)^{2.5}} - \\frac{x{\\left(t \\right)} - 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} - 1}\\right|^{2} + \\left|{y{\\left(t \\right)} + 1}\\right|^{2}\\right)^{2.5}} - \\frac{x{\\left(t \\right)} + 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} + 1}\\right|^{2} + \\left|{y{\\left(t \\right)} - 1}\\right|^{2}\\right)^{2.5}} - \\frac{x{\\left(t \\right)} + 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} + 1}\\right|^{2} + \\left|{y{\\left(t \\right)} + 1}\\right|^{2}\\right)^{2.5}}\\\\- \\frac{y{\\left(t \\right)} - 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} - 1}\\right|^{2} + \\left|{y{\\left(t \\right)} - 1}\\right|^{2}\\right)^{2.5}} - \\frac{y{\\left(t \\right)} - 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} + 1}\\right|^{2} + \\left|{y{\\left(t \\right)} - 1}\\right|^{2}\\right)^{2.5}} - \\frac{y{\\left(t \\right)} + 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} - 1}\\right|^{2} + \\left|{y{\\left(t \\right)} + 1}\\right|^{2}\\right)^{2.5}} - \\frac{y{\\left(t \\right)} + 1}{\\left(h^{2} + \\left|{x{\\left(t \\right)} + 1}\\right|^{2} + \\left|{y{\\left(t \\right)} + 1}\\right|^{2}\\right)^{2.5}}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Eq(Matrix([\n",
       "[b*Derivative(x(t), t) + x(t) + Derivative(x(t), (t, 2))],\n",
       "[b*Derivative(y(t), t) + y(t) + Derivative(y(t), (t, 2))]]), Matrix([\n",
       "[-(x(t) - 1)/(h**2 + Abs(x(t) - 1)**2 + Abs(y(t) - 1)**2)**2.5 - (x(t) - 1)/(h**2 + Abs(x(t) - 1)**2 + Abs(y(t) + 1)**2)**2.5 - (x(t) + 1)/(h**2 + Abs(x(t) + 1)**2 + Abs(y(t) - 1)**2)**2.5 - (x(t) + 1)/(h**2 + Abs(x(t) + 1)**2 + Abs(y(t) + 1)**2)**2.5],\n",
       "[-(y(t) - 1)/(h**2 + Abs(x(t) - 1)**2 + Abs(y(t) - 1)**2)**2.5 - (y(t) - 1)/(h**2 + Abs(x(t) + 1)**2 + Abs(y(t) - 1)**2)**2.5 - (y(t) + 1)/(h**2 + Abs(x(t) - 1)**2 + Abs(y(t) + 1)**2)**2.5 - (y(t) + 1)/(h**2 + Abs(x(t) + 1)**2 + Abs(y(t) + 1)**2)**2.5]]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Now we combine these two halves, and substitue the RHS for magnet locations on the unit circle.\n",
    "X_1 = [1, 1]; X_2 = [1, -1]; X_3 = [-1, -1]; X_4 = [-1, 1]\n",
    "p = 1\n",
    "EQ_2 = sm.Eq(LHS,\n",
    "             RHS_gen.subs({x_n: X_1[0], y_n: X_1[1], p_n: p}) +\n",
    "             RHS_gen.subs({x_n: X_2[0], y_n: X_2[1], p_n: p}) +\n",
    "             RHS_gen.subs({x_n: X_3[0], y_n: X_3[1], p_n: p}) +\n",
    "             RHS_gen.subs({x_n: X_4[0], y_n: X_4[1], p_n: p})\n",
    "            )\n",
    "EQ_2.simplify()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9ef3e10-851b-4fbb-bc20-c758dd7f66b5",
   "metadata": {},
   "source": [
    "So with this we have four states:\n",
    "\n",
    "$\\left[\\matrix{x \\\\ \\dot x \\\\ y \\\\ \\dot y}\\right] = \\left[\\matrix{z_1 \\\\ z_2 \\\\ z_3 \\\\ z_4}\\right]$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ce02f2-1a39-44b3-90a2-bffdac9fe9bc",
   "metadata": {},
   "source": [
    "We can rearrange our equation that we got from SymPy into two first order differential equations:\n",
    "\n",
    "$$\\left[\\matrix{\\dot z_1 \\\\ \\dot z_2 \\\\ \\dot z_3 \\\\ \\dot z_4} \\right] = \n",
    "\\left[\\matrix{\n",
    "z_2 \\\\\n",
    "-\\frac{z_1 - 1}{\\left(h^2+(z_1-1)^2 + (z_3 - 1)^2 \\right)^{2.5}} \n",
    "    -\\frac{z_1 - 1}{\\left(h^2+(z_1-1)^2 + (z_3 + 1)^2 \\right)^{2.5}} \n",
    "    -\\frac{z_1 + 1}{\\left(h^2+(z_1+1)^2 + (z_3 - 1)^2 \\right)^{2.5}} \n",
    "    -\\frac{z_1 + 1}{\\left(h^2+(z_1+1)^2 + (z_3 + 1)^2 \\right)^{2.5}} \n",
    "    - b z_2 - z_1 \\\\\n",
    "z_4 \\\\\n",
    "-\\frac{z_3 - 1}{\\left(h^2+(z_1-1)^2 + (z_3 - 1)^2 \\right)^{2.5}} \n",
    "    -\\frac{z_3 - 1}{\\left(h^2+(z_1+1)^2 + (z_3 - 1)^2 \\right)^{2.5}} \n",
    "    -\\frac{z_3 + 1}{\\left(h^2+(z_1-1)^2 + (z_3 + 1)^2 \\right)^{2.5}} \n",
    "    -\\frac{z_3 + 1}{\\left(h^2+(z_1+1)^2 + (z_3 + 1)^2 \\right)^{2.5}} \n",
    "    - b z_4 - z_3\\\\\n",
    "}\\right]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fddb7795-21f3-409a-bbc9-e2bba02eb8a2",
   "metadata": {},
   "source": [
    "and this we can use with SciPy's odeint.\n",
    "\n",
    "Let's do it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9cd6761f-ab45-4e05-a658-d7ca4f715536",
   "metadata": {},
   "outputs": [],
   "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",
    "''' \n",
    "Code below is heavily borrowed from Nikhil Bajaj, University of Pittsburgh from ME2016 Nonlinear Dynamics 1.\n",
    "\n",
    "Major Changes:\n",
    "Oscilator was changed into the magnetic pendulum. 4 states instead of 2\n",
    "\n",
    "'''\n",
    "# Declare Global Variables\n",
    "h = 0.5\n",
    "b_global = 0.175\n",
    "\n",
    "# Accelerate the ODE system using numba\n",
    "@jit(nopython=True)\n",
    "def magnetic_pendulum(Z, t):\n",
    "    z_1, z_2, z_3, z_4, b = Z #z_1 = x, z_2 = x_dot, z_3 = y, z_4 = y_dot, z_5 = b (for convenience to change damping)\n",
    "    dz_1dt = z_2  \n",
    "    dz_2dt = -(z_1 - 1)/(h**2 + (z_1 - 1)**2 + (z_3 - 1)**2)**2.5 \\\n",
    "                - (z_1 - 1)/(h**2 + (z_1 - 1)**2 + (z_3 + 1)**2)**2.5 \\\n",
    "                - (z_1 + 1)/(h**2 + (z_1 + 1)**2 + (z_3 - 1)**2)**2.5 \\\n",
    "                - (z_1 + 1)/(h**2 + (z_1 + 1)**2 + (z_3 + 1)**2)**2.5 \\\n",
    "                - b*z_2 - z_1\n",
    "    dz_3dt = z_4  \n",
    "    dz_4dt = -(z_3 - 1)/(h**2 + (z_1 - 1)**2 + (z_3 - 1)**2)**2.5 \\\n",
    "                - (z_3 - 1)/(h**2 + (z_1 + 1)**2 + (z_3 - 1)**2)**2.5 \\\n",
    "                - (z_3 + 1)/(h**2 + (z_1 - 1)**2 + (z_3 + 1)**2)**2.5 \\\n",
    "                - (z_3 + 1)/(h**2 + (z_1 + 1)**2 + (z_3 + 1)**2)**2.5 \\\n",
    "                - b*z_4 - z_3\n",
    "    db_dt = 0\n",
    "    return [dz_1dt, dz_2dt, dz_3dt, dz_4dt, db_dt]\n",
    "\n",
    "# Function to simulate a batch of initial conditions\n",
    "def simulate_batch(batch):\n",
    "    t = np.linspace(0, 1000, 500)  # Larger time steps for benchmarking\n",
    "    results = [(t, odeint(magnetic_pendulum, ic, t)) 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 None, parallel_results\n",
    "\n",
    "# List of initial conditions for benchmarking\n",
    "num_pixels = 200\n",
    "nudge = 1e-1\n",
    "value_range = np.linspace(-2+nudge, 2+nudge, num_pixels)\n",
    "zeros = np.zeros([num_pixels**2])\n",
    "b_vec = np.ones([num_pixels**2])*b_global\n",
    "x_values, y_values = np.meshgrid(value_range, value_range)\n",
    "\n",
    "initial_conditions_list = np.column_stack((x_values.flatten(), zeros, y_values.flatten(), zeros, b_vec))\n",
    "#initial_conditions_list\n",
    "#plt.plot(x_values,y_values, 'x')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed23a81b-9596-43ef-8a4d-c634f3ee1f4a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Benchmarking...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee773cb239244670aca79c0493fc6d24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Running Simulations:   0%|          | 0/20000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run the benchmark\n",
    "single_process_results, parallel_results = benchmark(initial_conditions_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a23ca9a5-d423-4f26-9810-1a988ca5b0f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Figure Out Final Magnets\n",
    "final_magnet = np.zeros([num_pixels**2,3])\n",
    "smallest_distance = np.zeros(num_pixels**2)\n",
    "\n",
    "for i, (t, solution) in enumerate(parallel_results):\n",
    "    last_location = np.asarray([solution[-1,0], solution[-1,2]])\n",
    "    \n",
    "    dist = np.asarray([np.linalg.norm(X_1-last_location), \n",
    "            np.linalg.norm(X_2-last_location),\n",
    "            np.linalg.norm(X_3-last_location),\n",
    "            np.linalg.norm(X_4-last_location)])\n",
    "    final_magnet[i,:] = np.asarray([solution[0,0], solution[0,2], np.argmin(dist)])\n",
    "    smallest_distance[i] = np.min(dist)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3712c30d-b17b-42f5-be73-c711be7b0f2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "magnet_map = np.zeros([num_pixels**2, 3])\n",
    "for i, magnet in enumerate(final_magnet[:,2]):\n",
    "    if magnet == 0:\n",
    "        magnet_map[i,:] = [1.0, 0, 0]\n",
    "    elif magnet == 1:\n",
    "        magnet_map[i,:] = [0, 1.0, 0]\n",
    "    elif magnet == 2:\n",
    "        magnet_map[i,:] = [0, 0, 1.0]\n",
    "    elif magnet == 3:\n",
    "        magnet_map[i,:] = [1.0, 1.0, 0]\n",
    "    else:\n",
    "        print('dear god')\n",
    "\n",
    "magnet_map = magnet_map.reshape([num_pixels, num_pixels, 3])        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda0bc4f-a42b-4732-a2f0-7e41a0737023",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize = (9, 10/4))\n",
    "\n",
    "axs[0].imshow(magnet_map, extent = [-2, 2, -2, 2])\n",
    "axs[1].imshow(magnet_map[int(num_pixels*0.5/2):int(num_pixels*1.5/2), int(num_pixels*1.1/2):int(num_pixels*1.9/2), :], extent = [0.9, 1.9, -0.5, 0.5])\n",
    "axs[2].imshow(magnet_map[int(num_pixels*1.83/2):int(num_pixels*1.93/2), int(num_pixels*1.3/2):int(num_pixels*1.4/2), :], extent = [1.3, 1.4, 0.33, 0.44])\n",
    "fig.suptitle(f'Basin of Attraction with b = {b_global:.3f} and {num_pixels**2} initial conditions.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f18d7de-b8fc-4ee9-9879-f406cc263874",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting all results on the same plot (using the parallel results)\n",
    "plt.figure(figsize=(10, 10))\n",
    "for i, (t, solution) in enumerate(parallel_results):\n",
    "    #print(solution[:10,:])\n",
    "    plt.plot(solution[:,0], solution[:, 2], label=f'Initial Condition {i+1}')"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
