{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a4e31467",
   "metadata": {},
   "source": [
    "# Linearise Nonlinear Model Automatically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "431b3477",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beginning notebook setup...\n",
      "\tAdded /home/jhewers/Documents/projects/jdrones/src to path\n",
      "\tImported gymnasium version 0.27.1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "pybullet build time: Feb  2 2023 13:13:41\n",
      "/home/jhewers/.local/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.5\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tImported jdrones version unknown\n",
      "\tImported scipy==1.7.3, numpy==1.23.5, pandas==1.3.5\n",
      "\tImported functools, collections and itertools\n",
      "\tImported tqdm (standard and trange)\n",
      "\tImported seaborn==0.11.2, matplotlib==3.5.1\n",
      "End of notebook setup\n"
     ]
    }
   ],
   "source": [
    "%matplotlib notebook\n",
    "from notebook_quick_setup import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aba89791",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gymnasium.make('NonLinearDynamicModelDroneEnv-v0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "34ed851b",
   "metadata": {},
   "outputs": [],
   "source": [
    "utrim = np.ones(4)*np.sqrt((env.model.mass*env.model.g)/(4*env.model.k_T))\n",
    "xtrim = np.zeros(12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "21f65a03",
   "metadata": {},
   "outputs": [],
   "source": [
    "nstates = len(xtrim)\n",
    "ncontrols = len(utrim)\n",
    "nv = nstates + ncontrols\n",
    "stepSize = 0.01\n",
    "minPertSize = 0.0001\n",
    "cv = np.concatenate([xtrim, utrim])\n",
    "pert = np.clip(np.abs(cv) * stepSize, minPertSize, np.inf)\n",
    "dvar = 2.0 * pert\n",
    "var = cv\n",
    "varpert = np.zeros((2 * nv, nv))\n",
    "for m in range(2 * nv):\n",
    "    for n in range(nv):\n",
    "        varpert[m, n] = var[n]\n",
    "    if m > 1:\n",
    "        ind = np.floor(0.5 * m).astype(int)\n",
    "        if m % 2 == 1:\n",
    "            varpert[m, ind] = var[ind] - pert[ind]\n",
    "        else:\n",
    "            varpert[m, ind] = var[ind] + pert[ind]\n",
    "            \n",
    "# evaluate the Jacobian\n",
    "f = np.empty((nstates, 2 * nv))\n",
    "for j in range(2 * nv):\n",
    "    # Calculate the derivative corresponding to the perturbed state.\n",
    "    xp = varpert[j, 0:nstates].T\n",
    "    up = varpert[j, nstates:nv].T\n",
    "    xd = env.calc_dstate(up, State.from_x(xp), env.model)\n",
    "    xd = State(xd).to_x()\n",
    "    # Now the function\n",
    "    for i in range(nstates):\n",
    "        f[i, j] = xd[i]\n",
    "# calculate the Jacobian using numerical differentiation\n",
    "J = np.empty((nstates, nv))\n",
    "for m in range(nstates):\n",
    "    for n in range(nv):\n",
    "        a = f[m, 2 * n]\n",
    "        b = f[m, 2 * n + 1]\n",
    "        c = dvar[n]\n",
    "        J[m, n]= (a - b) / c\n",
    "\n",
    "sysMatrix = J[0:nstates, 0:nstates]\n",
    "conMatrix = J[0:nstates, nstates:nv]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "520c9f9e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0, 9.81, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, -0.0, 0.0, -9.81, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n"
     ]
    }
   ],
   "source": [
    "for row in sysMatrix:\n",
    "    print(tuple(map(functools.partial(round, ndigits=2), row)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6d7a0025",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.84, 0.84, 0.84, 0.84)\n",
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, 0.0, 0.0, 0.0)\n",
      "(0.0, -1.17, 0.0, 1.17)\n",
      "(-1.17, 0.0, 1.17, 0.0)\n",
      "(5.86, -5.86, 5.86, -5.86)\n"
     ]
    }
   ],
   "source": [
    "for row in conMatrix:\n",
    "    print(tuple(map(functools.partial(round, ndigits=2), row)))"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}