notebooks/SpecialRelativity.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np # get rid of this eventually\n",
    "import argparse\n",
    "from jax import jit\n",
    "from jax.experimental.ode import odeint\n",
    "from functools import partial # reduces arguments to function by making some subset implicit\n",
    "\n",
    "from jax.experimental import stax\n",
    "from jax.experimental import optimizers\n",
    "\n",
    "import os, sys, time\n",
    "sys.path.append('..')\n",
    "sys.path.append('../hyperopt')\n",
    "from HyperparameterSearch import extended_mlp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import mlp as make_mlp\n",
    "from utils import wrap_coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.ode import odeint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lagrangian_eom(lagrangian, state, conditionals, t=None):\n",
    "    q, q_t = jnp.split(state, 2)\n",
    "    q = q / 10.0 #Normalize\n",
    "    conditionals = conditionals / 10.0\n",
    "    q_t = q_t\n",
    "    q_tt = (jnp.linalg.pinv(jax.hessian(lagrangian, 1)(q, q_t, conditionals))\n",
    "          @ (jax.grad(lagrangian, 0)(q, q_t, conditionals)\n",
    "             - jax.jacobian(jax.jacobian(lagrangian, 1), 0)(q, q_t, conditionals) @ q_t))\n",
    "    return jnp.concatenate([q_t, q_tt])\n",
    "\n",
    "# replace the lagrangian with a parameteric model\n",
    "def learned_dynamics(params, nn_forward_fn):\n",
    "    @jit\n",
    "    def dynamics(q, q_t, conditionals):\n",
    "    #     assert q.shape == (2,)\n",
    "        state = jnp.concatenate([q, q_t, conditionals])\n",
    "        return jnp.squeeze(nn_forward_fn(params, state), axis=-1)\n",
    "    return dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ObjectView(object):\n",
    "    def __init__(self, d): self.__dict__ = d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def qdotdot(q, q_t, conditionals):\n",
    "    g = conditionals\n",
    "    \n",
    "    q_tt = (\n",
    "        g * (1 - q_t**2)**(5./2) / \n",
    "        (1 + 2 * q_t**2)\n",
    "    )\n",
    "    \n",
    "    return q_t, q_tt\n",
    "\n",
    "@jax.jit\n",
    "def ofunc(y, t=None):\n",
    "    q = y[::3]\n",
    "    q_t = y[1::3]\n",
    "    g = y[2::3]\n",
    "    \n",
    "    q_t, q_tt = qdotdot(q, q_t, g)\n",
    "    return jnp.stack([q_t, q_tt, jnp.zeros_like(g)]).T.ravel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.99989885, dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(jnp.tanh(jax.random.uniform(jax.random.PRNGKey(1), (1000,))*10-5)*0.99999).max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([30.,  8.,  5.,  8.,  3.,  2.,  5.,  3.,  5., 31.]),\n",
       " array([-9.9991339e-01, -7.9992497e-01, -5.9993654e-01, -3.9994809e-01,\n",
       "        -1.9995967e-01,  2.8759241e-05,  2.0001718e-01,  4.0000561e-01,\n",
       "         5.9999406e-01,  7.9998249e-01,  9.9997091e-01], dtype=float32),\n",
       " <a list of 10 Patch objects>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXAAAAD4CAYAAAD1jb0+AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8QZhcZAAAPx0lEQVR4nO3dfaxkd13H8feHLhQVtFt6qUth2RYr0sSwJTe1sQkP5amUhJZYdJuAi9YsIBiImLiAiWg0tkZoYiTgYmtXxfJQaLoKiEsfQkigeIul3bIpuy1VS9furaU8xFhp+/WPORfHu3N35t55uPuz71cymTO/8zvnfO9vZj975sw5M6kqJEntecJ6FyBJWhsDXJIaZYBLUqMMcElqlAEuSY3aMMuNnXTSSbVly5ZZblKSmnfLLbc8UFVzy9tnGuBbtmxhYWFhlpuUpOYl+ZdB7R5CkaRGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRs30SkxJWk9bdn563bZ9z6Wvnvg63QOXpEYZ4JLUqKEBnuTJSb6S5GtJ7kjye137qUluTnIgyceSPGn65UqSloyyB/4wcG5VPR/YCpyX5GzgMuDyqjod+DZwyfTKlCQtNzTAq+f73cMndrcCzgWu6dp3AxdOpUJJ0kAjnYWS5DjgFuCngA8AdwEPVdUjXZd7gVNWWHYHsANg8+bNay70/9unx5I0rpE+xKyqR6tqK/BM4CzgeYO6rbDsrqqar6r5ubkjflBCkrRGqzoLpaoeAm4CzgZOSLK0B/9M4L7JliZJOppRzkKZS3JCN/0jwMuA/cCNwEVdt+3AddMqUpJ0pFGOgW8CdnfHwZ8AfLyq/j7J14GPJvkD4J+BK6ZYpyRpmaEBXlW3AWcOaL+b3vFwSdI68EpMSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWrU0ABP8qwkNybZn+SOJG/v2t+b5FtJbu1u50+/XEnSkg0j9HkEeGdVfTXJU4Fbkuzt5l1eVX8yvfIkSSsZGuBVdQg41E1/L8l+4JRpFyZJOrpVHQNPsgU4E7i5a3pbktuSXJlk4wrL7EiykGRhcXFxrGIlSf9r5ABP8hTgk8A7quq7wAeB5wBb6e2hv2/QclW1q6rmq2p+bm5uAiVLkmDEAE/yRHrh/ZGq+hRAVd1fVY9W1WPAh4GzplemJGm5Uc5CCXAFsL+q3t/Xvqmv22uBfZMvT5K0klHOQjkHeANwe5Jbu7Z3Axcn2QoUcA/wpqlUKEkaaJSzUL4IZMCsz0y+HEnSqLwSU5IaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1KihAZ7kWUluTLI/yR1J3t61n5hkb5ID3f3G6ZcrSVoyyh74I8A7q+p5wNnAW5OcAewErq+q04Hru8eSpBkZGuBVdaiqvtpNfw/YD5wCXADs7rrtBi6cVpGSpCOt6hh4ki3AmcDNwMlVdQh6IQ88fYVldiRZSLKwuLg4XrWSpB8aOcCTPAX4JPCOqvruqMtV1a6qmq+q+bm5ubXUKEkaYKQAT/JEeuH9kar6VNd8f5JN3fxNwOHplChJGmSUs1ACXAHsr6r3983aA2zvprcD102+PEnSSjaM0Occ4A3A7Ulu7dreDVwKfDzJJcC/Aq+bTomSpEGGBnhVfRHICrNfOtlyJEmj8kpMSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSo4YGeJIrkxxOsq+v7b1JvpXk1u52/nTLlCQtN8oe+FXAeQPaL6+qrd3tM5MtS5I0zNAAr6ovAA/OoBZJ0iqMcwz8bUlu6w6xbJxYRZKkkaw1wD8IPAfYChwC3rdSxyQ7kiwkWVhcXFzj5iRJy60pwKvq/qp6tKoeAz4MnHWUvruqar6q5ufm5tZapyRpmTUFeJJNfQ9fC+xbqa8kaTo2DOuQ5GrgxcBJSe4Ffhd4cZKtQAH3AG+aYo2SpAGGBnhVXTyg+Yop1CJJWgWvxJSkRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRg0N8CRXJjmcZF9f24lJ9iY50N1vnG6ZkqTlRtkDvwo4b1nbTuD6qjoduL57LEmaoaEBXlVfAB5c1nwBsLub3g1cOOG6JElDrPUY+MlVdQigu3/6Sh2T7EiykGRhcXFxjZuTJC039Q8xq2pXVc1X1fzc3Ny0NydJjxtrDfD7k2wC6O4PT64kSdIo1hrge4Dt3fR24LrJlCNJGtUopxFeDXwJeG6Se5NcAlwKvDzJAeDl3WNJ0gxtGNahqi5eYdZLJ1yLJGkVvBJTkhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGDT0PXLBl56fXZbv3XPrqddnuev29sH5/s9Qi98AlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0a6wcdktwDfA94FHikquYnUZQkabhJ/CLPS6rqgQmsR5K0Ch5CkaRGjbsHXsA/Jingz6tq1/IOSXYAOwA2b9485uYeX9bztyk1O4+331wFX9uTMu4e+DlV9QLgVcBbk7xweYeq2lVV81U1Pzc3N+bmJElLxgrwqrqvuz8MXAucNYmiJEnDrTnAk/xYkqcuTQOvAPZNqjBJ0tGNcwz8ZODaJEvr+duq+oeJVCVJGmrNAV5VdwPPn2AtkqRV8DRCSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMm8X3g0sQ8Hr+Zb734jYDtcw9ckhplgEtSowxwSWqUAS5JjTLAJalRnoUi4RkZapN74JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaNVaAJzkvyZ1JDibZOamiJEnDrTnAkxwHfAB4FXAGcHGSMyZVmCTp6MbZAz8LOFhVd1fVfwMfBS6YTFmSpGHG+TbCU4B/63t8L/Bzyzsl2QHs6B5+P8mda9zeScADa1x2mqxrdaxrdaxrdY7VushlY9X27EGN4wR4BrTVEQ1Vu4BdY2ynt7Fkoarmx13PpFnX6ljX6ljX6hyrdcF0ahvnEMq9wLP6Hj8TuG+8ciRJoxonwP8JOD3JqUmeBGwD9kymLEnSMGs+hFJVjyR5G/A54Djgyqq6Y2KVHWnswzBTYl2rY12rY12rc6zWBVOoLVVHHLaWJDXAKzElqVEGuCQ16pgK8CSvS3JHkseSrHi6zUqX8HcfqN6c5ECSj3Ufrk6irhOT7O3WuzfJxgF9XpLk1r7bfyW5sJt3VZJv9s3bOqu6un6P9m17T1/7eo7X1iRf6p7v25L8Ut+8iY7XsK98SHJ89/cf7MZjS9+8d3XtdyZ55Th1rKGu30zy9W58rk/y7L55A5/TGdX1xiSLfdv/tb5527vn/UCS7TOu6/K+mr6R5KG+edMcryuTHE6yb4X5SfKnXd23JXlB37zxxquqjpkb8DzgucBNwPwKfY4D7gJOA54EfA04o5v3cWBbN/0h4C0TquuPgZ3d9E7gsiH9TwQeBH60e3wVcNEUxmukuoDvr9C+buMF/DRwejf9DOAQcMKkx+tor5e+Pr8OfKib3gZ8rJs+o+t/PHBqt57jZljXS/peQ29Zqutoz+mM6noj8GcDlj0RuLu739hNb5xVXcv6/wa9EyumOl7dul8IvADYt8L884HP0rt25mzg5kmN1zG1B15V+6tq2JWaAy/hTxLgXOCart9u4MIJlXZBt75R13sR8Nmq+s8JbX8lq63rh9Z7vKrqG1V1oJu+DzgMzE1o+/1G+cqH/nqvAV7ajc8FwEer6uGq+iZwsFvfTOqqqhv7XkNfpnetxbSN8xUZrwT2VtWDVfVtYC9w3jrVdTFw9YS2fVRV9QV6O2wruQD4q+r5MnBCkk1MYLyOqQAf0aBL+E8BngY8VFWPLGufhJOr6hBAd//0If23ceSL5w+7t0+XJzl+xnU9OclCki8vHdbhGBqvJGfR26u6q695UuO10utlYJ9uPL5Db3xGWXaadfW7hN5e3JJBz+ks6/qF7vm5JsnSBX3HxHh1h5pOBW7oa57WeI1ipdrHHq9xLqVfkySfB35ywKz3VNV1o6xiQFsdpX3sukZdR7eeTcDP0js/fsm7gH+nF1K7gN8Gfn+GdW2uqvuSnAbckOR24LsD+q3XeP01sL2qHuua1zxegzYxoG353zmV19QQI687yeuBeeBFfc1HPKdVddeg5adQ198BV1fVw0neTO/dy7kjLjvNupZsA66pqkf72qY1XqOY2utr5gFeVS8bcxUrXcL/AL23Jhu6vahVXdp/tLqS3J9kU1Ud6gLn8FFW9YvAtVX1g751H+omH07yl8BvzbKu7hAFVXV3kpuAM4FPss7jleTHgU8Dv9O9tVxa95rHa4BRvvJhqc+9STYAP0HvLfE0vy5ipHUneRm9/xRfVFUPL7Wv8JxOIpCG1lVV/9H38MPAZX3LvnjZsjdNoKaR6uqzDXhrf8MUx2sUK9U+9ni1eAhl4CX81ftU4EZ6x58BtgOj7NGPYk+3vlHWe8Sxty7Elo47XwgM/LR6GnUl2bh0CCLJScA5wNfXe7y65+5aescGP7Fs3iTHa5SvfOiv9yLghm589gDb0jtL5VTgdOArY9SyqrqSnAn8OfCaqjrc1z7wOZ1hXZv6Hr4G2N9Nfw54RVffRuAV/N93olOtq6vtufQ+EPxSX9s0x2sUe4Bf7s5GORv4TreTMv54TeuT2bXcgNfS+1/pYeB+4HNd+zOAz/T1Ox/4Br3/Qd/T134avX9gB4FPAMdPqK6nAdcDB7r7E7v2eeAv+vptAb4FPGHZ8jcAt9MLor8BnjKruoCf77b9te7+kmNhvIDXAz8Abu27bZ3GeA16vdA7JPOabvrJ3d9/sBuP0/qWfU+33J3Aqyb8eh9W1+e7fwdL47Nn2HM6o7r+CLij2/6NwM/0Lfur3TgeBH5llnV1j98LXLpsuWmP19X0zqL6Ab38ugR4M/Dmbn7o/fjNXd325/uWHWu8vJRekhrV4iEUSRIGuCQ1ywCXpEYZ4JLUKANckhplgEtSowxwSWrU/wB/9oZOyVWQ9QAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist((jnp.tanh(jax.random.normal(jax.random.PRNGKey(1), (100,))*2)*0.99999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(jax.jit, static_argnums=(1, 2), backend='cpu')\n",
    "def gen_data(seed, batch, num):\n",
    "    rng = jax.random.PRNGKey(seed)\n",
    "    q0 = jax.random.uniform(rng, (batch,), minval=-10, maxval=10)\n",
    "    qt0 = jax.random.uniform(rng+1, (batch,), minval=-0.99, maxval=0.99)\n",
    "    g = jax.random.normal(rng+2, (batch,))*10\n",
    "\n",
    "    y0 = jnp.stack([q0, qt0, g]).T.ravel()\n",
    "\n",
    "    yt = odeint(ofunc, y0, jnp.linspace(0, 1, num=num), mxsteps=300)\n",
    "\n",
    "    qall = yt[:, ::3]\n",
    "    qtall = yt[:, 1::3]\n",
    "    gall = yt[:, 2::3]\n",
    "    \n",
    "    return jnp.stack([qall, qtall]).reshape(2, -1).T, gall.reshape(1, -1).T, qdotdot(qall, qtall, gall)[1].reshape(1, -1).T\n",
    "\n",
    "@partial(jax.jit, static_argnums=(1,))\n",
    "def gen_data_batch(seed, batch):\n",
    "    rng = jax.random.PRNGKey(seed)\n",
    "    q0 = jax.random.uniform(rng, (batch,), minval=-10, maxval=10)\n",
    "    qt0 = (jnp.tanh(jax.random.normal(jax.random.PRNGKey(1), (batch,))*2)*0.99999)#jax.random.uniform(rng+1, (batch,), minval=-1, maxval=1)\n",
    "    g = jax.random.normal(rng+2, (batch,))*10\n",
    "    \n",
    "    return jnp.stack([q0, qt0]).reshape(2, -1).T, g.reshape(1, -1).T, qdotdot(q0, qt0, g)[1].reshape(1, -1).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "qt [ 0.13323684 -0.12181558  0.9635672   0.9999864   0.98424286] g [ -7.3825674    0.48637673   5.567146     3.391427   -11.220732  ] qtt [-6.8172374e+00  4.5502922e-01  2.6673502e-03  4.3539304e-12\n",
      " -6.6027965e-04]\n"
     ]
    }
   ],
   "source": [
    "print('qt', gen_data_batch(0, 128)[0][:5, 1], 'g', gen_data_batch(0, 128)[1][:5, 0], 'qtt', gen_data_batch(0, 128)[2][:5, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# qdotdot(jnp.array([0]), jnp.array([0.9]), jnp.array([10]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 0.29830917716026306 {'act': [4],\n",
    "# 'batch_size': [27.0], 'dt': [0.09609870774790222],\n",
    "# 'hidden_dim': [596.0], 'l2reg': [0.24927677946969878],\n",
    "# 'layers': [4.0], 'lr': [0.005516656601005163],\n",
    "# 'lr2': [1.897157209816416e-05], 'n_updates': [4.0]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle as pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loaded = pkl.load(open('./params_for_loss_0.29429444670677185_nupdates=1.pkl', 'rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = ObjectView({'dataset_size': 200,\n",
    " 'fps': 10,\n",
    " 'samples': 100,\n",
    " 'num_epochs': 80000,\n",
    " 'seed': 0,\n",
    " 'loss': 'l1',\n",
    " 'act': 'softplus',\n",
    " 'hidden_dim': 500,\n",
    " 'output_dim': 1,\n",
    " 'layers': 4,\n",
    " 'n_updates': 1,\n",
    " 'lr': 0.001,\n",
    " 'lr2': 2e-05,\n",
    " 'dt': 0.1,\n",
    " 'model': 'gln',\n",
    " 'batch_size': 68,\n",
    " 'l2reg': 5.7e-07,\n",
    "})\n",
    "# args = loaded['args']\n",
    "rng = jax.random.PRNGKey(args.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental.ode import odeint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_params = None\n",
    "best_loss = np.inf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_random_params, nn_forward_fn = extended_mlp(args)\n",
    "rng = jax.random.PRNGKey(0)\n",
    "_, init_params = init_random_params(rng, (-1, 3))\n",
    "rng += 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the output. Now, let's train it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Idea: add identity before inverse:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's train it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_small_loss = np.inf\n",
    "iteration = 0\n",
    "total_epochs = 100\n",
    "minibatch_per = 3000\n",
    "train_losses, test_losses = [], []\n",
    "\n",
    "lr = 1e-3 #1e-3\n",
    "\n",
    "final_div_factor=1e4\n",
    "\n",
    "#OneCycleLR:\n",
    "@jax.jit\n",
    "def OneCycleLR(pct):\n",
    "    #Rush it:\n",
    "    start = 0.3 #0.2\n",
    "    pct = pct * (1-start) + start\n",
    "    high, low = lr, lr/final_div_factor\n",
    "    \n",
    "    scale = 1.0 - (jnp.cos(2 * jnp.pi * pct) + 1)/2\n",
    "    \n",
    "    return low + (high - low)*scale\n",
    "    \n",
    "\n",
    "opt_init, opt_update, get_params = optimizers.adam(\n",
    "    OneCycleLR\n",
    ")\n",
    "from lnn import custom_init\n",
    "init_params = custom_init(init_params, seed=0)\n",
    "opt_state = opt_init(init_params)\n",
    "# opt_state = opt_init(best_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'lr schedule')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXwAAAEICAYAAABcVE8dAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8QZhcZAAAgAElEQVR4nO3deXxc5X3v8c9vZrSv1mpLtiwvso0BY0BgdkiAxlAcJ21DgWwFGkLuJYWmacO9aZP0Nnlx29w2996GhppA2BIIhDQxhCahScEEMNjG2HgBb9iWLNtarMXat6d/zJGRhWRkazTnaOb7fr300syjmXN+Ohp955nnLI855xARkcQX8rsAERGJDwW+iEiSUOCLiCQJBb6ISJJQ4IuIJAkFvohIklDgy5RkZnvN7KpJXscLZvanMVrWN8zssVg/VuRkKPBFRJKEAl8SjplF/K5BJIgU+DLleUMgPzGzx8ysDfiTUR5zrZltM7OjZnbAzL487GcrzexNM2szs91mtnzYU2eb2cve835tZkXDnneBmb1iZi1mtsnMrhj2szlm9qL3vOeB4c+7wsxqR9Q35hDVidYjcjIU+JIoVgI/AfKBH47y8weAzzvncoAzgN8CmNn5wCPAX3rPvQzYO+x5NwE3AyVAKvBl73nlwC+AbwIFXvvTZlbsPe9HwAaiQf93wGdP5Zcax3pExk2BL4niVefcz5xzg865rlF+3gcsNrNc51yzc+4Nr/1W4EHn3PPecw84594e9rwfOOd2eMt8EljqtX8KeM4595z3vOeB9cC1ZlYBnAf8jXOuxzm3BnjmFH+vMddzisuTJKbAl0RR8wE//0OiIbnPG2q50GufBew+wfMODbvdCWR7t2cDn/CGWVrMrAW4BJgBlAHNzrmOYc/dN87fY6QTrUfkpGjnliSKE1721Tm3DlhpZinAHUR767OIvlHMO4X11QCPOuc+N/IHZjYbmGZmWcNCv2JYjR1A5rDHh4GxhmjGXI/IyVIPXxKemaWa2SfNLM851we0AQPejx8AbjazK80sZGblZrZoHIt9DFhhZh8xs7CZpXs7Y2c65/YRHXb5W2/dlwArhj13B5BuZr/vvQH9NZB2sus5hU0hSU6BL8ni08Be7yie24mOjeOce53oTtnvAK3Ai0SHUU7IOVdDdEfx/wQaiPbE/5L3/qduApYBR4CvE90xPPTcVuC/Ad8HDhDt8R931M5JrEdk3EwToIiIJAf1EkREkoQCX0QkSSjwRUSShAJfRCRJBPo4/KKiIldZWel3GSIiU8qGDRsanXPvO7cj0IFfWVnJ+vXr/S5DRGRKMbNRz+zWkI6ISJJQ4IuIJAkFvohIklDgi4gkCQW+iEiSiFvgm9lpZnafNxXdF+K1XhERiRpX4JvZg2ZWb2ZbRrQvN7N3zGyXmd19omU457Y7524HrgeqT71kERE5FeM9Dv8h4LsMu8SrN2nDvcDVRC/tus7MVgNh4J4Rz7/FOVdvZh8F7vaWJXEwMOjo7O2nq3eAjt6BY7e7+wbpGxykr3+Q/kFH38AgvcNu9w1Ev/cPDDJ0QVWzoe92bPnH2jDMICUcIjUSIi0cIiVipIbDpEaibSlhIy0SIjUcJj0lRGZahOzUCJlpYVLCGl0UmWzjCnzn3BozqxzRfD6wyzm3B8DMngBWOufuAa4bYzmrgdVm9guikzy/j5ndBtwGUFFRMZ7yEppzjqM9/TS199LU3kNLZx+tXe99tXV734e1dfREg72zd4Ce/kG/f4VxSY2EyEoNk5UWISs1QlZa9HZ+Zir5GSnkZ6aQl5Fy3P1oWyoFWamEQ/bBKxFJchM507ac4+cRrSU64cOozOwK4A+Izuzz3FiPc86tAlYBVFdXJ+zF+vsGBjnc1s2h1m4OtnZzuK2bhvYemtp7afS+N7X30NjRS+8JQjsnLUJuRgq5GSnkZUSoLMwiOz0ampmpYTJSw2SlRshIDZOZGibTa0871uuOfkXCRuqw29F2IxIKEbL35uYb6u07HCOnUhh0jr5+R8/AAL390U8Jvf3RTw69AwP09jt6vU8S3X3RN6X2ngE6e/rp6B2go6efjt5+Onqib1Zt3f3UNnfR0tlLa1cfg2O8GkIGBVmpFGWnUZyTRrH3/dj9nDSm56VTlpdBRmp44n88kSlqIoE/WpdqzIB2zr0AvDCB9U0prZ197D/Syf4jndQ0d3rB3nUs4Bvae94XmClhoyg7jcLsaHgtKM2hKCeVoqw0inJSKchKY5rX083LSCE7LUIkaEMhqQApMV/s4GD0k05rZx8tXb20dPbR0tVHc0f0jbGhvYeGo700tPewp6GDhvaeUd8op2WmMCMvg7L8DMry0ynLz2BGXvR7RUEmJTlpxw1ZiSSSiQR+LdFJoIfMBOomVk6Uma0AVsyfPz8Wi5sUzjnqWrvZ29jB/iOd7GvqpMYL+P1HOmnt6jvu8dlpEWbkpTM9L52F03OYnpdx7H5ZXgbTc9PJzYgobMYQCtmxN7qK9+b/HpNzjrbufhrbe6hv6+FQWxd1Ld3UtXRxsLWb2uZOXn+3ibbu/uOel54SYnZBFrMLM72vLCoLo/fL8jM0dCRT2rinOPTG8J91zp3h3Y8QnYz5SqLzcq4DbnLObY1VcdXV1c7vi6f1DQyyr6mDXfXt7KpvZ3dDh/e9nc7egWOPSwkbM6dlMqsgk4qCaG+xoiCLioJMZhZkkJse+16vTFx7Tz8HW7o40NJFjffGvbepk/1HOtjX1HncPpCUsFFRkElVSQ5VpdnML8mmqiSHucVZpKdoqEiCw8w2OOfedzTkuHr4ZvY4cAVQZGa1wNedcw+Y2R3Ar4gemfNgrMLejx7+UI99e10b2w62sf1gGzsOH2VfUyf9wwaPy/LSmVeSzfXVs5hfks3c4ixmF2YxPTddvb8pKDstQlVpDlWlOe/72eCg4/DRbvY1dbKvqYO9TZ3saWhnR/1Rnt9+mAHvdREymF2Y5b0BZLOgNIfTy3KZW5yt14QESqAnMZ+sHn5v/yC76tvZdrCNbXXRcN92sO24YZg5RVks8Hpx80uymVcc/cpKC/QVpSVOevoH2NvYyc76o+w83H7s+7uNHcc6COkpIRZNz+X0slzOKM/j9LJcFpTm6NOATLqxevgJH/jOOfY2dfJmTTObalp5s6aFbXVt9A5EP6oP/VOeNiOXxWW5LJ6Rw8LpuWQr2OUU9A0Msruhna0H2tha18bWula2HWzjqLevIBIy5pdkc3pZHksr8jmnIp+FpTnB2/kuU9qUCvxhQzqf27lz50k//7U9Tby8q5GNNS1srm091nPPTA1zZnkeS2flc0Z5HqfNyGVOUZY+dsukcs5Rc6SLrXWtbKlrZWtdG1sOtNLY3gtARkqYJTPzOLtiGmdX5HN2RT4lOek+Vy1T2ZQK/CGn2sP/yk8289SGGhZOz2XprDzOmpnP0op8qkpyFO4SCM45apu7eGN/Mxv3t7CxpoVtda30DUT/H2cVZLBsTiHL5hRwwdxCZk7L0BFcMm5JFfgNR3vISoueZCQyVXT3DbC1ro2N+5tZt/cIr797hObO6KfTsrx0ls2NvgEsm1tIZWGm3gBkTFMq8Cc6pCOSCAYHHTvr23nt3SZe23OE195tOjYMVJ6fwaVVRVy2oJiL5xWRl6nDfuU9UyrwhwThOHyRoHDOsbuhg7V7mvjdzkZe3t3I0e5+QgZLZuZz2YJiLqsqYumsfO0ETnIKfJEE0z8wyKbaFtbsaGTNzgY21bQw6CAnPcKHFpZw9eJSLl9YrJP+kpACXyTBtXb28cruRn77dj2/fbuepo5eUsLGBXMLuXpxKVeeVkp5fobfZUocTKnA1xi+yMQMDDo27m/m+W2HeX77YfY0dABwelkuy0+fznVnlTGnKMvnKmWyTKnAH6Ievkhs7G5o5z+2HebX2w6zYV8zAGeU57JiSRnXnVWmnn+CUeCLCAAHW7v4xeaDPLP5IJtqWgA4d/Y0ViyZwbVLZuikrwSgwBeR99nf1Mkzm+t4ZlMdbx86Ssjgkqpirq+eyVWnleq6P1OUAl9ETmjn4aOs3lTH0xtqqWvtJi8jhZVLy/jEubM4ozxXJ3pNIVMq8LXTVsQ/A4OOV3Y38tT6Wn659RC9/YMsmp7DTcsq+PjZ5eToMM/Am1KBP0Q9fBF/tXb18cymOp5cX8Pm2lYyU8N8/OxyPn3hbBZNz/W7PBmDAl9EJmRTTQuPrd3H6k119PQPcn5lAZ+6cDbLT59OakRn9gaJAl9EYqK5o5enNtTw2Nr97D/SyfTcdG6+uJIbl1XorN6AUOCLSEwNDjpe3NHA/S/t4ZXdTWSnRbhpWQU3X1zJjDwd1+8nBb6ITJotB1pZtWYPv3jrIAZ89Kwybr9iHgtGmStYJt+UCnwdpSMyNdUc6eTBl9/lx+tq6Oob4NozZ3DXlVWjThIvk2dKBf4Q9fBFpqaWzl6+/9K7/ODld+nsG+C6JWXceeV85pco+ONBgS8icdfc0cv9L+3hoVf20tU3wIolZdx5VRXzirP9Li2hKfBFxDdHOnpZtWYPj7y6l57+QW44bxZ3XlWl6/ZMEgW+iPiuqb2Hf/7tLh5bu4/USIjPXzaPz102R/NPx9hYga+zJUQkbgqz0/jGR0/n+S9dzuULivnOf+zg8m+/wE821DI4GNzOZ6JQ4ItI3M0pyuJ7nzqXp79wETOnZfDlpzbxR/e9wpYDrX6XltAU+CLim3NnT+Pp2y/i23+0hP1HOlnx3d/x1z97i5bOXr9LS0gKfBHxVShkfKJ6Fr/5iyv4k4sqefz1Gj78jy/ys40HCPI+xqkokIFvZivMbFVrqz7eiSSLvIwUvr7idH7xZ5cwuzCTu378Jrc8tI66li6/S0sYOkpHRAJnYNDxyKt7+YdfvkPI4CvXLOJTy2YTCmkSlvHQUToiMmWEQ8bNF8/h139+GefMnsbXfr6Vm76/lgPq7U+IAl9EAmtWQSaP3HI+//CHS3irtpXl/3cNqzfV+V3WlKXAF5FAMzOuP28Wz915KVUl2fzZ4xu564mNtHb1+V3alKPAF5EpYXZhFk9+/kL+/KoFPLP5INf+v5fYuL/Z77KmFAW+iEwZkXCIO6+q4qnbL8QMrv/XV3n01b06fHOcFPgiMuWcUzGNZ794CZfML+Jvfr6VP//xm3T29vtdVuAp8EVkSsrPTOWBz57HX1y9gJ9vquNj977MnoZ2v8sKNAW+iExZoZDxxSureOSW82k42sPKe1/mlV2NfpcVWAp8EZnyLq0q5pkvXsKMvHQ+8+DrPLmuxu+SAimugW9mWWa2wcyui+d6RSTxzZyWyU++cBEXzivkr57ezP/+97d1yeURxhX4ZvagmdWb2ZYR7cvN7B0z22Vmd49jUV8BnjyVQkVEPkhuego/+JPzuGlZBfe9uJs7Hn+D7r4Bv8sKjPFOM/MQ8F3gkaEGMwsD9wJXA7XAOjNbDYSBe0Y8/xZgCbAN0JxmIjJpIuEQ3/rYGcwtyuJbz22nsf11vv/ZanLTU/wuzXfjCnzn3BozqxzRfD6wyzm3B8DMngBWOufuAd43ZGNmHwKygMVAl5k955wbnEDtIiKjMjP+9NK5lOSm86Ufv8lN96/l4ZvPpzA7ze/SfDWRMfxyYPiekVqvbVTOua865+4CfgTcP1bYm9ltZrbezNY3NDRMoDwRSXYfPauM+z9Tzc7D7XziX1/lYGtyX3xtIoE/2nVKP3APiXPuIefcsyf4+SrnXLVzrrq4uHgC5YmIwIcWlfDorcuob+vhhlVrOdTa7XdJvplI4NcCs4bdnwnE5DJ2mgBFRGLp/DkFPHzL+TS193Lj/Ws53JacoT+RwF8HVJnZHDNLBW4AVseiKOfcM8652/Ly8mKxOBERzp09jYdvOY/6tm5uXLWW+iQM/fEelvk48Cqw0MxqzexW51w/cAfwK2A78KRzbmssilIPX0Qmw7mzoz39w23dfPqB12ntTK5LLGuKQxFJOi/vauTmH6xjycw8Hr11GRmpYb9LiilNcSgi4rl4fhHf+eOlbNjfzB0/eoO+geQ4QjyQga8hHRGZbL+/ZAb/a+UZ/Obtev7mZ1uS4pr6gQx87bQVkXj49AWzueND83liXQ0P/O5dv8uZdOO9tIKISEL60tUL2N3Qzree286coiyuPK3U75ImTSB7+CIi8RIKGf94/VmcXpbLnz2+kbcPtfld0qQJZOBrDF9E4ikzNcL3P3MeWWkRbn90A23diXm4ZiADX2P4IhJv0/PSufeT51DT3MVfPbU5IXfiBjLwRUT8cF5lAXcvX8Qvtx7iwZf3+l1OzCnwRUSG+dNL5/B7i0u557ntbNjX7Hc5MRXIwNcYvoj4xcz49ifOoiw/gzuf2Eh7T7/fJcVMIANfY/gi4qe8jBT+6fqzqGvp4pvPbvO7nJgJZOCLiPiturKAz18+jyfW1fCb7Yf9LicmFPgiImO466oqFk3P4StPv8WRjl6/y5mwQAa+xvBFJAjSImG+88dLae3q5Ws/3+J3ORMWyMDXGL6IBMVpM3L54oereHbzQV7cMbXn2Q5k4IuIBMnnL5/L3KIsvvbzLXT3DfhdzilT4IuIfIC0SJhvfuwM9jV18i8v7Pa7nFOmwBcRGYeL5hexcmkZ972wmz0N7X6Xc0oU+CIi4/TV3z+NtJQQX18dk+m7406BLyIyTiU56dx5ZRUv7WzkpZ1TbwduIANfh2WKSFB9+sLZlOdn8Pe/fJvBwal1Rc1ABr4OyxSRoEqLhPnyRxaw5UAbz2yu87uckxLIwBcRCbKVZ5Vz2oxc/s+v36G3f9DvcsZNgS8icpJCIePuaxZRc6SLH762z+9yxk2BLyJyCi6rKuKieYXc+5+7pszJWAp8EZFTYGZ88cNVNLb38pMNtX6XMy4KfBGRU3TB3AKWzspn1Zo99A8EfyxfgS8icorMjNsvn8f+I538+5ZDfpfzgQIZ+DoOX0Smit9bXMrc4iy+98JunAv2cfmBDHwdhy8iU0UoFO3lbzvYxpqdjX6Xc0KBDHwRkankY0vLmZ6bzvde2OV3KSekwBcRmaDUSIibL65k7Z4jvHPoqN/ljEmBLyISA9dXzyI1EuKxtcE9EUuBLyISA9OyUlmxpIyfvlFLe0+/3+WMSoEvIhIjn75wNh29A/xs4wG/SxmVAl9EJEbOmpnHouk5PLW+xu9SRqXAFxGJETPjj86dyabaVnYcDt7OWwW+iEgMffzsciIhC+T1dRT4IiIxVJidxocXlfDTNw4E7vo6CnwRkRj7+NnlNLb38Nq7R/wu5ThxC3wzu8LMXjKz+8zsinitV0Qk3j60qISs1DDPbArWFIjjCnwze9DM6s1sy4j25Wb2jpntMrO7P2AxDmgH0oHgDW6JiMRIekqYqxeX8suthwI1BeJ4e/gPAcuHN5hZGLgXuAZYDNxoZovN7Ewze3bEVwnwknPuGuArwN/G7lcQEQme65aU0dLZx+92NfhdyjGR8TzIObfGzCpHNJ8P7HLO7QEwsyeAlc65e4DrTrC4ZiBtrB+a2W3AbQAVFRXjKU9EJHAuXVBETlqEX245xIcXlfpdDjCxMfxyYPjZBbVe26jM7A/M7F+BR4HvjvU459wq51y1c666uLh4AuWJiPgnLRLm8oXF/GZ7PQODwbhO/kQC30ZpG/O3cs791Dn3eefcHzvnXjjhgjUBiogkgKsXl9LU0cubNc1+lwJMLPBrgVnD7s8EYrJLWhOgiEgiuGJhCZGQ8fy2er9LASYW+OuAKjObY2apwA3A6tiUJSIy9eVlpLBsbgHPbwvGfLfjPSzzceBVYKGZ1ZrZrc65fuAO4FfAduBJ59zWWBSlIR0RSRRXLipld0MHNUc6/S5lfIHvnLvROTfDOZfinJvpnHvAa3/OObfAOTfPOfetWBWlIR0RSRSXLSgC4KUAzHerSyuIiEyiecXZzMhLD8Tx+IEMfA3piEiiMDMumV/Ey7uafD88M5CBryEdEUkkly4oprWrj7cO+NuJDWTgi4gkkovnFWIGv9vp77BOIANfQzoikkgKs9NYWJrj++WSAxn4GtIRkURzXmUBb+xr9nVSlEAGvohIoqmunEZH7wBvH/JvrlsFvohIHJxXWQDAur3+DesEMvA1hi8iiaYsP4Py/AzW7/XvQmqBDHyN4YtIIjqvchqv7z2Cc/4cjx/IwBcRSUTnVhbQcLSH2uYuX9avwBcRiZMl5dFRiy0+nYClwBcRiZNFM3JICRubFfjv0U5bEUlEaZEwC0pz1MMfTjttRSRRLZmZx+baVl923AYy8EVEEtUZ5Xm0dvX5suNWgS8iEkdnejtuN9fGf1hHgS8iEkcLpw/tuG2J+7oV+CIicZQWCTOvOJsdPlxTJ5CBr6N0RCSRLSjNYcfh9rivN5CBr6N0RCSRVZVkc6Cli46e/riuN5CBLyKSyKpKcwDYVR/fXr4CX0QkzhaUZgOw43B8x/EV+CIicVZRkElqOKQevohIoouEQ8wtzlIPX0QkGfhxpI4CX0TEB34cqaPAFxHxwbyS6I7bvU0dcVtnIANfJ16JSKKrKMgEoOZIZ9zWGcjA14lXIpLoZnmBvz/ZA19EJNHlZaSQl5GiwBcRSQYVBZnsPxK/6+Ir8EVEfFJRkEmtevgiIolvVkEmtc1dDAzGZ7pDBb6IiE8qCjLpHRjkcFt3XNanwBcR8UlFnI/UUeCLiPhEgS8ikiRm5KcTDlncTr5S4IuI+CQlHGJGXnrceviRuKwFMLMQ8HdALrDeOfdwvNYtIhJUM/LSg7XT1sweNLN6M9syon25mb1jZrvM7O4PWMxKoBzoA2pPrVwRkcRSkptO/dGeuKxrvEM6DwHLhzeYWRi4F7gGWAzcaGaLzexMM3t2xFcJsBB41Tn3JeALsfsVRESmrpKcNOrb4hP44xrScc6tMbPKEc3nA7ucc3sAzOwJYKVz7h7gupHLMLNaoNe7OzDWuszsNuA2gIqKivGUJyIyZZXmptPe009HTz9ZaZM7yj6RnbblQM2w+7Ve21h+CnzEzP4ZWDPWg5xzq5xz1c656uLi4gmUJyISfKW5aQBxGdaZyNuJjdI25vnBzrlO4NYJrE9EJOGU5KQDcLitmzlFWZO6ron08GuBWcPuzwTqJlZOlCZAEZFkMdTDj8eROhMJ/HVAlZnNMbNU4AZgdSyK0gQoIpIsir0efkMchnTGe1jm48CrwEIzqzWzW51z/cAdwK+A7cCTzrmtsShKPXwRSRa56RHSU0Jx6eGP9yidG8dofw54LqYVRZf7DPBMdXX152K9bBGRIDEzSuN0LL4urSAi4rOSnLTAj+FPGg3piEgyKclNj8vJV4EMfO20FZFkUpqjIR0RkaRQkptGe08/7T39k7qeQAa+hnREJJkcO9t2ksfxAxn4GtIRkWRSmBUN/CMdvR/wyIkJZOCLiCST/MwUAFo6+yZ1PQp8ERGfTctMBaC5Mwl7+BrDF5FkktQ9fI3hi0gyyU6LEAlZcvbwRUSSiZmRn5lCS1cS9vBFRJJNfmYqLerhi4gkvvyMFJo7krCHr522IpJs8jNTk3MMXzttRSTZTMtMoVVj+CIiiW9aVpL28EVEkk1eRgrdfYN09w1M2joU+CIiARCPs20V+CIiATAtDmfbBjLwdZSOiCSbPC/wk66Hr6N0RCTZDA3pJF0PX0Qk2WgMX0QkScTjipkKfBGRAEhPCZOeEprU6+ko8EVEAmJaZirN6uGLiCS+yb5ipgJfRCQgCrNSaWxX4IuIJLyi7FSaOnombfmBDHydeCUiyagoO43Go0nWw9eJVyKSjIpy0ujqG6Cjp39Slh/IwBcRSUaFWdGTrxrbJ2dYR4EvIhIQRTlpgAJfRCThFWdHA79hksbxFfgiIgFR5AX+ZB2po8AXEQmIwmxvDF89fBGRxJYSDpGfmaIxfBGRZFCUnabAFxFJBkXZqQp8EZFkUJidNmnX01Hgi4gESPEkDulEJmWpozCzS4FPeutc7Jy7KF7rFhGZKoqyUzna3U933wDpKeGYLntcPXwze9DM6s1sy4j25Wb2jpntMrO7T7QM59xLzrnbgWeBh0+9ZBGRxPXesfixH9YZ75DOQ8Dy4Q1mFgbuBa4BFgM3mtliMzvTzJ4d8VUy7Kk3AY/HoHYRkYQzFPiNR2M/rDOuIR3n3BozqxzRfD6wyzm3B8DMngBWOufuAa4bbTlmVgG0OufaxlqXmd0G3AZQUVExnvJERBJGZVEm1545PebDOTCxnbblQM2w+7Ve24ncCvzgRA9wzq1yzlU756qLi4snUJ6IyNQzvySHf/nkuSycnhPzZU9kp62N0uZO9ATn3NfHtWCzFcCK+fPnn0pdIiIyion08GuBWcPuzwTqJlZOlCZAERGJvYkE/jqgyszmmFkqcAOwOjZliYhIrI33sMzHgVeBhWZWa2a3Ouf6gTuAXwHbgSedc1tjUZTmtBURiT1z7oTD7r6qrq5269ev97sMEZEpxcw2OOeqR7br0goiIkkikIGvIR0RkdgLZODrKB0RkdgL9Bi+mTUA+07x6UVAYwzLiZWg1gXBrU11nZyg1gXBrS3R6prtnHvfmauBDvyJMLP1o+208FtQ64Lg1qa6Tk5Q64Lg1pYsdQVySEdERGJPgS8ikiQSOfBX+V3AGIJaFwS3NtV1coJaFwS3tqSoK2HH8EVE5HiJ3MMXEZFhFPgiIkkiIQP/ZObaneQ6ZpnZf5rZdjPbamZ3eu3fMLMDZvam93WtD7XtNbO3vPWv99oKzOx5M9vpfZ8W55oWDtsmb5pZm5nd5df2Gm0u57G2kUX9f+81t9nMzolzXd82s7e9df+bmeV77ZVm1jVs290X57rG/NuZ2f/wttc7ZvaRONf142E17TWzN732eG6vsfJh8l5jzrmE+gLCwG5gLpAKbAIW+1TLDOAc73YOsIPo/L/fAL7s83baCxSNaPsH4G7v9t3A3/v8dzwEzPZrewGXAecAWz5oGwHXAv9OdGKgC4DX4lzX7wER7/bfD6urcvjjfNheo/7tvP+DTUAaMMf7nw3Hq64RP/9H4Gs+bK+x8mHSXmOJ2MM/Nteuc64XeAJY6UchzrmDzrk3vNtHiV5G+oOmgfTTSqokAsMAAAMDSURBVOBh7/bDwMd8rOVKYLdz7lTPtJ4w59wa4MiI5rG20UrgERe1Fsg3sxnxqss592sXvWQ5wFqiExLF1RjbaywrgSeccz3OuXeBXUT/d+Nal5kZcD3w+GSs+0ROkA+T9hpLxMA/lbl2J51FJ4E/G3jNa7rD+1j2YLyHTjwO+LWZbbDoxPEApc65gxB9MQIlPtQ15AaO/yf0e3sNGWsbBel1dwvRnuCQOWa20cxeNLNLfahntL9dULbXpcBh59zOYW1x314j8mHSXmOJGPgnPdfuZDOzbOBp4C7nXBvwPWAesBQ4SPQjZbxd7Jw7B7gG+O9mdpkPNYzKojOofRR4ymsKwvb6IIF43ZnZV4F+4Ide00Ggwjl3NvAl4EdmlhvHksb62wViewE3cnzHIu7ba5R8GPOho7Sd1DZLxMCftLl2T4WZpRD9Y/7QOfdTAOfcYefcgHNuELifSfooeyLOuTrvez3wb14Nh4c+Inrf6+Ndl+ca4A3n3GGvRt+31zBjbSPfX3dm9lngOuCTzhv09YZMmrzbG4iOlS+IV00n+NsFYXtFgD8AfjzUFu/tNVo+MImvsUQM/MDMteuNDz4AbHfO/dOw9uHjbh8Htox87iTXlWVmOUO3ie7w20J0O33We9hngZ/Hs65hjut1+b29RhhrG60GPuMdSXEB0Dr0sTwezGw58BXgo865zmHtxWYW9m7PBaqAPXGsa6y/3WrgBjNLM7M5Xl2vx6suz1XA28652qGGeG6vsfKByXyNxWNvdLy/iO7N3kH03fmrPtZxCdGPXJuBN72va4FHgbe89tXAjDjXNZfoERKbgK1D2wgoBH4D7PS+F/iwzTKBJiBvWJsv24vom85BoI9o7+rWsbYR0Y/b93qvubeA6jjXtYvo+O7Q6+w+77F/6P2NNwFvACviXNeYfzvgq972ege4Jp51ee0PAbePeGw8t9dY+TBprzFdWkFEJEkk4pCOiIiMQoEvIpIkFPgiIklCgS8ikiQU+CIiSUKBLyKSJBT4IiJJ4r8ArUnT0Hilux0AAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(OneCycleLR(jnp.linspace(0, 1, num=200)))\n",
    "plt.yscale('log')\n",
    "plt.title('lr schedule')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.tree_util import tree_flatten\n",
    "\n",
    "@jax.jit\n",
    "def loss(params, cstate, cconditionals, ctarget):\n",
    "    runner = jax.vmap(\n",
    "        partial(\n",
    "            lagrangian_eom,\n",
    "            learned_dynamics(params, nn_forward_fn)), (0, 0), 0)\n",
    "    preds = runner(cstate, cconditionals)[:, [1]]\n",
    "    \n",
    "    error = jnp.abs(preds - ctarget)\n",
    "    #Weight additionally by proximity to c!\n",
    "    error_weights = (1 + 1/jnp.sqrt(1.0-cstate[:, [1]]**2))\n",
    "    \n",
    "    return jnp.sum(error * error_weights)*len(preds)/jnp.sum(error_weights)\n",
    "\n",
    "@jax.jit\n",
    "def update_derivative(i, opt_state, cstate, cconditionals, ctarget):\n",
    "    params = get_params(opt_state)\n",
    "    param_update = jax.grad(\n",
    "            lambda *args: loss(*args)/len(cstate),\n",
    "            0\n",
    "        )(params, cstate, cconditionals, ctarget)\n",
    "    params = get_params(opt_state)\n",
    "    return opt_update(i, param_update, opt_state), params\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "cstate, cconditionals, ctarget = gen_data_batch(epoch, 128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(5611.877, dtype=float32)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss(get_params(opt_state), cstate, cconditionals, ctarget)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "update_derivative(0, opt_state, cstate, cconditionals, ctarget);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = jax.random.PRNGKey(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(128, 2)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gen_data_batch(0, 128)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ -7.3825674 ],\n",
       "             [  0.48637673],\n",
       "             [  5.567146  ],\n",
       "             [  3.391427  ],\n",
       "             [-11.220732  ]], dtype=float32)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cconditionals[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[-8.077726  ,  0.13323684],\n",
       "             [-2.589147  , -0.12181558],\n",
       "             [-8.912823  ,  0.9635672 ],\n",
       "             [ 4.5006466 ,  0.9999864 ],\n",
       "             [ 6.4032707 ,  0.98424286]], dtype=float32)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cstate[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[-6.8172374e+00],\n",
       "             [ 4.5502922e-01],\n",
       "             [ 2.6673502e-03],\n",
       "             [ 4.3539304e-12],\n",
       "             [-6.6027965e-04]], dtype=float32)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ctarget[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_loss = np.inf\n",
    "best_params = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy as copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fefe77ee6da147fba7b3c27346801dc7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=0 lr=0.0006752928602509201 loss=0.45701903104782104\n",
      "epoch=1 lr=0.0006957106525078416 loss=0.4204932749271393\n",
      "epoch=2 lr=0.0007157498621381819 loss=0.41228559613227844\n",
      "epoch=3 lr=0.0007353719556704164 loss=0.410251259803772\n",
      "epoch=4 lr=0.0007545390981249511 loss=0.4115859270095825\n",
      "epoch=5 lr=0.000773213803768158 loss=0.409858763217926\n",
      "epoch=6 lr=0.0007913602748885751 loss=0.4101628363132477\n",
      "epoch=7 lr=0.0008089432376436889 loss=0.410196989774704\n",
      "epoch=8 lr=0.0008259288733825088 loss=0.4147755205631256\n",
      "epoch=9 lr=0.0008422840619459748 loss=0.40926679968833923\n",
      "epoch=10 lr=0.0008579773711971939 loss=0.4101622402667999\n",
      "epoch=11 lr=0.000872978416737169 loss=0.4091649651527405\n",
      "epoch=12 lr=0.0008872582111507654 loss=0.41075968742370605\n",
      "epoch=13 lr=0.0009007891057990491 loss=0.4095885455608368\n",
      "epoch=14 lr=0.0009135449072346091 loss=0.41020146012306213\n",
      "epoch=15 lr=0.0009255009354092181 loss=0.42736509442329407\n",
      "epoch=16 lr=0.0009366340818814933 loss=0.4159628450870514\n",
      "epoch=17 lr=0.000946922751609236 loss=0.41024911403656006\n",
      "epoch=18 lr=0.0009563472121953964 loss=0.4100976586341858\n",
      "epoch=19 lr=0.0009648891282267869 loss=0.44837766885757446\n",
      "epoch=20 lr=0.0009725319687277079 loss=0.4975762963294983\n",
      "epoch=21 lr=0.0009792608907446265 loss=0.41222891211509705\n",
      "epoch=22 lr=0.000985063030384481 loss=0.408489465713501\n",
      "epoch=23 lr=0.0009899272117763758 loss=0.4279766082763672\n",
      "epoch=24 lr=0.0009938437724485993 loss=0.4090638756752014\n",
      "epoch=25 lr=0.000996805145405233 loss=0.40898397564888\n",
      "epoch=26 lr=0.0009988059755414724 loss=0.41019389033317566\n",
      "epoch=27 lr=0.0009998419554904103 loss=0.43750011920928955\n",
      "epoch=28 lr=0.0009999113390222192 loss=0.4190872311592102\n",
      "epoch=29 lr=0.0009990138933062553 loss=0.5750089287757874\n",
      "epoch=30 lr=0.000997151480987668 loss=0.44808685779571533\n",
      "epoch=31 lr=0.000994327594526112 loss=0.50361567735672\n",
      "epoch=32 lr=0.0009905475890263915 loss=0.42528727650642395\n",
      "epoch=33 lr=0.000985819031484425 loss=0.40401551127433777\n",
      "epoch=34 lr=0.0009801508858799934 loss=0.40294763445854187\n",
      "epoch=35 lr=0.0009735542116686702 loss=0.4842650592327118\n",
      "epoch=36 lr=0.0009660416981205344 loss=0.4176408648490906\n",
      "epoch=37 lr=0.0009576278389431536 loss=0.408379465341568\n",
      "epoch=38 lr=0.0009483289322815835 loss=0.4007442891597748\n",
      "epoch=39 lr=0.0009381631389260292 loss=0.2769946753978729\n",
      "epoch=40 lr=0.0009271499002352357 loss=0.03971005231142044\n",
      "epoch=41 lr=0.0009153104037977755 loss=0.020392214879393578\n",
      "epoch=42 lr=0.000902667990885675 loss=0.028213180601596832\n",
      "epoch=43 lr=0.000889246875885874 loss=0.01611640863120556\n",
      "epoch=44 lr=0.0008750728447921574 loss=0.015939785167574883\n",
      "epoch=45 lr=0.0008601734298281372 loss=0.013999160379171371\n",
      "epoch=46 lr=0.00084457773482427 loss=0.012936053797602654\n",
      "epoch=47 lr=0.0008283156785182655 loss=0.012137713842093945\n",
      "epoch=48 lr=0.0008114183438010514 loss=0.011220214888453484\n",
      "epoch=49 lr=0.0007939189672470093 loss=0.010879701934754848\n",
      "epoch=50 lr=0.0007758514257147908 loss=0.01027984544634819\n",
      "epoch=51 lr=0.0007572502945549786 loss=0.00977717898786068\n",
      "epoch=52 lr=0.000738151662517339 loss=0.009415862150490284\n",
      "epoch=53 lr=0.0007185924332588911 loss=0.008636108599603176\n",
      "epoch=54 lr=0.0006986107910051942 loss=0.008344546891748905\n",
      "epoch=55 lr=0.0006782450363971293 loss=0.007729768753051758\n",
      "epoch=56 lr=0.0006575344013981521 loss=0.016900360584259033\n",
      "epoch=57 lr=0.0006365194567479193 loss=0.006467245519161224\n",
      "epoch=58 lr=0.0006152403657324612 loss=0.006712664850056171\n",
      "epoch=59 lr=0.0005937386304140091 loss=0.006442787125706673\n",
      "epoch=60 lr=0.0005720554618164897 loss=0.006018386222422123\n",
      "epoch=61 lr=0.0005502330604940653 loss=0.005452815908938646\n",
      "epoch=62 lr=0.0005283138016238809 loss=0.005117522086948156\n",
      "epoch=63 lr=0.000506339652929455 loss=0.004767376463860273\n",
      "epoch=64 lr=0.00048435357166454196 loss=0.00455406354740262\n",
      "epoch=65 lr=0.0004623976710718125 loss=0.004279023502022028\n",
      "epoch=66 lr=0.00044051450095139444 loss=0.004121516365557909\n",
      "epoch=67 lr=0.0004187468148302287 loss=0.0037611983716487885\n",
      "epoch=68 lr=0.0003971360856667161 loss=0.0035603572614490986\n",
      "epoch=69 lr=0.00037572436849586666 loss=0.003439656225964427\n",
      "epoch=70 lr=0.00035455342731438577 loss=0.0032285756897181273\n",
      "epoch=71 lr=0.0003336636000312865 loss=0.0029799824114888906\n",
      "epoch=72 lr=0.00031309586483985186 loss=0.0027777303475886583\n",
      "epoch=73 lr=0.0002928894537035376 loss=0.0026550856418907642\n",
      "epoch=74 lr=0.0002730839769355953 loss=0.002424228237941861\n",
      "epoch=75 lr=0.0002537174441386014 loss=0.0022690424229949713\n",
      "epoch=76 lr=0.00023482716642320156 loss=0.0021453076042234898\n",
      "epoch=77 lr=0.00021644984371960163 loss=0.0019376821583136916\n",
      "epoch=78 lr=0.0001986212591873482 loss=0.0018619343172758818\n",
      "epoch=79 lr=0.00018137549341190606 loss=0.0017642470775172114\n",
      "epoch=80 lr=0.0001647462631808594 loss=0.0016290233470499516\n",
      "epoch=81 lr=0.0001487653498770669 loss=0.0014907082077115774\n",
      "epoch=82 lr=0.00013346386549528688 loss=0.0014077556552365422\n",
      "epoch=83 lr=0.00011887160508194938 loss=0.0013333468232303858\n",
      "epoch=84 lr=0.00010501645738258958 loss=0.0012401037383824587\n",
      "epoch=85 lr=9.192524885293096e-05 loss=0.001162555068731308\n",
      "epoch=86 lr=7.96236636233516e-05 loss=0.0011141973081976175\n",
      "epoch=87 lr=6.813512300141156e-05 loss=0.0010512381559237838\n",
      "epoch=88 lr=5.748197509092279e-05 loss=0.0010054815793409944\n",
      "epoch=89 lr=4.768490543938242e-05 loss=0.0009749584714882076\n",
      "epoch=90 lr=3.876250411849469e-05 loss=0.0009298786171711981\n",
      "epoch=91 lr=3.073247353313491e-05 loss=0.0009030640940181911\n",
      "epoch=92 lr=2.3610193238710053e-05 loss=0.000874264573212713\n",
      "epoch=93 lr=1.740936750138644e-05 loss=0.0008505667210556567\n",
      "epoch=94 lr=1.2142033483542036e-05 loss=0.0008347271359525621\n",
      "epoch=95 lr=7.818327503628097e-06 loss=0.0008203327306546271\n",
      "epoch=96 lr=4.446651473699603e-06 loss=0.0008076290250755847\n",
      "epoch=97 lr=2.0336199213488726e-06 loss=0.0007957015186548233\n",
      "epoch=98 lr=5.838221568410518e-07 loss=0.0007923548109829426\n",
      "epoch=99 lr=1.0000000116860974e-07 loss=0.0007864094222895801\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for epoch in tqdm(range(epoch, total_epochs)):\n",
    "    epoch_loss = 0.0\n",
    "    num_samples = 0\n",
    "    batch = 512\n",
    "    ocstate, occonditionals, octarget = gen_data_batch(epoch, minibatch_per*batch)\n",
    "    for minibatch in range(minibatch_per):\n",
    "        fraction = (epoch + minibatch/minibatch_per)/total_epochs\n",
    "        s = np.s_[minibatch*batch:(minibatch+1)*batch]\n",
    "        \n",
    "        cstate, cconditionals, ctarget = ocstate[s], occonditionals[s], octarget[s]\n",
    "        opt_state, params = update_derivative(fraction, opt_state, cstate, cconditionals, ctarget);\n",
    "        rng += 10\n",
    "        \n",
    "        cur_loss = loss(params, cstate, cconditionals, ctarget)\n",
    "        \n",
    "        epoch_loss += cur_loss\n",
    "        num_samples += len(cstate)\n",
    "    closs = epoch_loss/num_samples\n",
    "    print('epoch={} lr={} loss={}'.format(\n",
    "        epoch, OneCycleLR(fraction), closs)\n",
    "         )\n",
    "    if closs < best_loss:\n",
    "        best_loss = closs\n",
    "        best_params = [[copy(jax.device_get(l2)) for l2 in l1] if len(l1) > 0 else () for l1 in params]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle as pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pkl.dump({'params': best_params, 'description': 'q and g are divided by 10. hidden=500. act=Softplus'},\n",
    "#          open('best_sr_params_v2.pkl', 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_params = pkl.load(open('best_sr_params_v2.pkl', 'rb'))['params']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_state = opt_init(best_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "cstate, cconditionals, ctarget = gen_data(0, 1, 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50, 2)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cstate.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x2aab5b28d950>]"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAD4CAYAAAD8Zh1EAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8QZhcZAAAeAElEQVR4nO3deXRc5Znn8e+jfV8syZYteQUDXgADGrOFhCUQQ9KQhUwgdBISgtMnIUtP0n0gJ02nSeck3WdOkzkZMjME0qEzaQhNEuJm6GYJ0Ek6BixjMFhe8IJtLZZkLaWtqrQ980eVTSFkq2xLLtet3+ecOlX31nXpuXbp5/e8933fa+6OiIikv6xUFyAiItNDgS4iEhAKdBGRgFCgi4gEhAJdRCQgclL1g6urq33RokWp+vEiImlp48aNB929ZrL3UhboixYtorGxMVU/XkQkLZnZ3iO9py4XEZGAUKCLiASEAl1EJCAU6CIiAaFAFxEJCAW6iEhAKNBFRAIiZePQRUSCxN2Jjo4zEB1lIDLKQHSUwWjsOfZ67PD2lWfN5tz5FdNegwJdRDJedHSM/sho/DHCQGSUvkOvo7H9ic8DkZF37BscjoX46Hhy95eoKc1XoIuITOTuREbG6YuM0BceoS8yQig8Ql94NGFfLJwT9/XHQ7svMsLw6PiUPyc/J4vSglxKC3IoyY895s8qojQ/h5L4vuL8nMPvF+fnUBp/Ls4/tC+b4rwcsrJsRv4uFOgickoYGRund2iEUHg4/jzy9nM4FsK9Q8OEwiMJj1H6wiMMjx09kAtysygryKWsMBbI5UV5sTAuyKWsIBbCh94rzc+lJL6vND+2rzg/h7ycU/+SowJdRKaVuzM4PEbP4DA9Q8N0D8YCOvY8TM/QCD1DsX294WF6BmPhPBAdPeJnmkFpfg7lRbmUF8YeteUFlBfmUV6YS1lhTuy5IDe+HQvqQyGdn5N9Ev8GUkeBLiJHNT7u9IZH6B6McnBgmK6BYboHo3QNDtMzOEzXYCy0u+MB3jN45BazGZQX5lJZlEdFUS6zSws4Y3YpFfHtinhgVxTlUVH49nZpQS7ZM9RNESQKdJEMNDw6zsGBKJ39UboGoxzsH6ZzIErXwDAHB6IcHIjSPTjMwYFYSI8d4WJfWUEOVSX5zCrOo76yiHPrK6gszmNWcSy0ZxXnUVGUR2VRbLusUME8kxToIgHh7vSFR+noj9DRH6WjP0Jnf5SOvigd/dHDAd45EKV3aGTSzyjJz6G6JI+qknwWzCrivAWVse3iPGaV5FNVnEdVSSyoK4vyyM0+9fuVM4kCXSQNREfHaA9FaQuFOdAXob0vQntflPa+CB190cP7opOM1ijMzaamNJ/ZpfmcPruEi0+roqYkn+rSfKpL8qkpjQV1TWk+BbmZ0dccVAp0kRQbHh3nQChCayhMWyhMa2+EtlCYtt4IbaFYUHcNDr/rzxXmZlNbXsDs0nxWza84/DoW3gXMLouFeEl+Dmbq5sgECnSRGTYYHaW5J0xzzxAtvWFaesI094Zpjb/uHIjiE7qoywtzmVtewLyKQlYtqKC2rIDa8gLmlhdQW1bAnPICShXUMoECXeQEjYyN09ITZl/3EHu7h2juHmJ/zxDNPWH2dw/RM6G/Oi87i3kVBdRVFvK+M2qYV1FIXUUhcysKmFteyLyKAory9Kspxy6pb42ZrQH+B5ANPODu35/w/kLgJ0AN0A38qbs3T3OtIikTGRljX/cQew4OsrdrkD0Hh9jXPci+7iFaeyPvGAWSl51FfWUhdZWFrDx7LvWVhcyvLDq8r7o4f8ZmCkpmmzLQzSwbuA+4GmgGNpjZOndvSjjsvwP/5O4PmdmVwPeAT81EwSIzZWzcaekJs+vgALs7B9ndGXve2zVIayjyjmMri3JZWFXM+Qsq+fCqIhbMij0WVhUzu1SBLamRTAt9NbDT3XcDmNkjwA1AYqAvB/48/vp54PHpLFJkOkVGxtjdOcibHf3s6hjgzY4BdnUO8NbBoXdMiCkvzGVxdTEXLaliYVUxi6qLWFxdzMKqYsoLc1N4BiKTSybQ64D9CdvNwIUTjnkN+BixbpmPAKVmVuXuXYkHmdlaYC3AggULjrdmkaSMjI2z5+Ag2w70s+NAP9sO9PNmRz/7uocOX4TMMlhYVcxpNSVccdZsllQXs6SmhCXVxcwqztNFR0kryQT6ZN/oidPGvgH8TzO7Ffgd0AK8a2EGd78fuB+goaEhuXUmRZJwcCBKU2sfW9v6aGrrY/uBfnZ1DjAyFvuaZWcZS6qLWTmvnA+vqmPpnBJOn13C4urijFnnQ4IvmUBvBuYnbNcDrYkHuHsr8FEAMysBPubuoekqUuQQd6e5J8zrLSFebwkdDvGO/ujhY+aVF3DW3DKuOGs2Z84p5czaUpbUKLgl+JIJ9A3AUjNbTKzlfRPwycQDzKwa6Hb3ceAuYiNeRE5YWyjMq/t62dwS4o14iB+atp6TZSydU8p7llazfG4Zy+eVsay2jMrivBRXLZIaUwa6u4+a2R3AU8SGLf7E3beY2T1Ao7uvAy4HvmdmTqzL5UszWLME1NDwKJubQ7y6v5dN+3p4dX8v7X2xlndOlnFmbSlrVtSysq6cs+vKObO2VFPVRRKYT5yidpI0NDR4Y2NjSn62nBo6+iI07u3h5T3dNO7tZmtb/+Hx3AurijhvfgWr5lewakElZym8RQAws43u3jDZe5qOJifN/u4h1u/q4qV4gO/tGgJid5M5b34lX7z8NM5bUMG59RVUleSnuFqR9KNAlxnT2htm/a4u1u/uYv2uLlp6wwBUFefRsKiST120kIZFs1gxr0zLsIpMAwW6TJuh4VFe3N3F73Yc5D92dLLn4CAAFUW5XLS4irXvXcIlp1Vx+uwSje8WmQEKdDlu7s729n7+Y3snv3uzkw17ehgeG6cgN4uLl1Rxy4ULuOS0as6qLdVUeJGTQIEux2RkbJwNe7p5uqmdZ7e209wT60Y5c04pt166iPcuraFhUaUuYIqkgAJdpjQYHeW5bR08u7Wd57d10BcZJT8ni/ecXs2XrjidK86cTW15QarLFMl4CnSZVHh4jOe2dfDE5lae29ZBdHScquI8PrCilquXz+E9S6u1ZrfIKUa/kXJYZGSMF7Z38sTmVn67tYPwyBjVJfnc9F/mc93Zc2lYNEt3bBc5hSnQM5y7s7k5xGMbm1n3Wiuh8AhVxXl89Pw6PnTOPFYvVoiLpAsFeobq6I/w+KYWHtvYzI72AfJzslizspaPnV/PJadVkaNx4SJpR4GeQdyd/9zZxUPr3+K5bR2MjTsXLKzkex89mw+eM5eyAt20QSSdKdAzQH9khF9ubOafXtzL7s5BqorzuP2yJXy8oZ7TakpSXZ6ITBMFeoDt7OjnoT/u5VevNDM4PMaq+RXc+4lzue7suVobXCSAFOgB9Nr+Xu57fidPN7WTl5PF9efO49MXL+Sc+opUlyYiM0iBHhDuzvrdXfzo+V38YedBygtz+epVS/n0xQu1cqFIhlCgpzl357dbO7jvhZ1s2tdLTWk+37zuLD554UJK8vXPK5JJ9Bufxja81c33ntzKK/t6qa8s5G8/vJIbL6jXOioiGUqBnoZ2dvTz/X/bzrNb25lTls/3P3o2N15Qr7HjIhlOgZ5G2vsi3PvMDh5t3E9xXg5/8YEz+dyliynMU4tcRBToaWF4dJwf/343P3zuTcbGnVsvWcwdV57OLN3dXkQSKNBPcRve6uabv3qdNzsGuHZlLXddu4wFVUWpLktETkEK9FNUaGiE7//7Nh5+eR91FYU8+JkGrlo2J9VlicgpTIF+inF31r3WyneeaKJnaIS1713CV69aSrGGIIrIFJQSp5CewWH+8pebeaapnXPry3noc6tZMa881WWJSJpQoJ8iXtzdxdceeZWuwSjf+uAyPnvpYq1DLiLHRIGeYqNj4/zwuZ388Lk3WVhVzK8/cykr69QqF5Fjp0BPodbeMF975FVefqubj55fxz03rNR0fRE5bkqPFHm2qZ2v/8trjI6Nc+8nzuUj59WnuiQRSXMK9JPM3XnwD3v47pNbWTGvjB/efD6Lq4tTXZaIBEBSi3+Y2Roz225mO83szkneX2Bmz5vZJjPbbGbXTX+p6W9s3Pmbf23ib//fVj6wvJbH/uwShbmITJspW+hmlg3cB1wNNAMbzGyduzclHPYt4FF3/19mthx4Elg0A/WmraHhUb7y8Cae3drB7Zct5q5rl5GlUSwiMo2S6XJZDex0990AZvYIcAOQGOgOlMVflwOt01lkuuvoj3DbTxvZ0hrinhtW8OmLF6W6JBEJoGQCvQ7Yn7DdDFw44ZhvA0+b2ZeBYuD9k32Qma0F1gIsWLDgWGtNSzva+/nsP26ge3CYH39a0/dFZOYk04c+Wb+AT9i+Gfipu9cD1wE/M7N3fba73+/uDe7eUFNTc+zVppmm1j4+/r/XMzw2zqNfuFhhLiIzKpkWejMwP2G7nnd3qdwGrAFw9/VmVgBUAx3TUWQ62tU5wKcefImivGwe/cLFzJ+lFRJFZGYl00LfACw1s8VmlgfcBKybcMw+4CoAM1sGFACd01loOtnfPcQtP34JM/j55y9UmIvISTFloLv7KHAH8BSwldholi1mdo+ZXR8/7OvA7Wb2GvAwcKu7T+yWyQgHQhFueeAlwiNj/Oy2C1lSU5LqkkQkQyQ1scjdnyQ2FDFx390Jr5uAS6e3tPTTNRDllgdepGsgys9vv4hlc8um/kMiItNEM0WnSSg8wqcefJmW3jAPfXY1q+ZXpLokEckwuk38NAgPj/HZf3yZNzv6+T+fauDCJVWpLklEMpBa6CfI3fnW42+waX8vP/rk+bzvjOAPxxSRU5Na6Cfo0cb9/PKVZr585VKuPXtuqssRkQymQD8Bb7SE+KvfbOGypdV89aqlqS5HRDKcAv04hcIjfPHnrzCrKI8ffGKVbhcnIimnPvTj4O58419eo7U3zC++cBFVJfmpLklERC304/Hj3+/mmaZ27rpuGRcsnJXqckREAAX6MXtpdxd/9+/buXZlLZ+7dFGqyxEROUyBfgw6+6N8+eFNLJhVxN/feA5m6jcXkVOHAv0YfOeJJnrDI/zolvMpLchNdTkiIu+gQE/S+l1drHutlT9732lao0VETkkK9CSMjI3z1+veoL6ykC9eflqqyxERmZQCPQkP/fEtdrQPcPeHllOQm53qckREJqVAn0JHX4QfPPsml59Zw9XLdQs5ETl1KdCn8L1/28bw6Djf/pMVGtUiIqc0BfpRvLynm19vamHte5ewqLo41eWIiByVAv0IRsfGufs3b1BXUciXrjg91eWIiExJgX4EP3txL9sO9PNXH1pGYZ4uhIrIqU+BPonO/ij/8PQO3ntGDR9YUZvqckREkqJAn8S9z+4gMjrGt/9kuS6EikjaUKBPcHAgymMbm7nxgvksqSlJdTkiIklToE/wf1/cy/DoOLe9Z3GqSxEROSYK9ASRkTF+tn4vV501m9Nnq3UuIulFgZ7g8U0tdA0Oc9tlap2LSPpRoMe5Ow/8YQ8r5pVx8ZKqVJcjInLMFOhxL+zoZGfHAJ+/bLFGtohIWlKgxz34+z3MKcvng2fPS3UpIiLHRYEObG3r4w87D3LrJYvJy9FfiYikp6TSy8zWmNl2M9tpZndO8v69ZvZq/LHDzHqnv9SZ88Dv91CYm80nVy9IdSkiIsctZ6oDzCwbuA+4GmgGNpjZOndvOnSMu/95wvFfBs6bgVpnREdfhHWvtfDJ1QsoL9J9QkUkfSXTQl8N7HT33e4+DDwC3HCU428GHp6O4k6Gh9a/xei48zlNJBKRNJdMoNcB+xO2m+P73sXMFgKLgeeO8P5aM2s0s8bOzs5jrXXaDQ2P8vOX9nHN8jksrNJ65yKS3pIJ9MnG8PkRjr0JeMzdxyZ7093vd/cGd2+oqalJtsYZ88tXWugdGuH2y5akuhQRkROWTKA3A/MTtuuB1iMcexNp1N3ys/VvcW59ORcsrEx1KSIiJyyZQN8ALDWzxWaWRyy01008yMzOBCqB9dNb4szY1TnAjvYBPnJenSYSiUggTBno7j4K3AE8BWwFHnX3LWZ2j5ldn3DozcAj7n6k7phTytNb2gG4RjewEJGAmHLYIoC7Pwk8OWHf3RO2vz19Zc28p5sOcHZdOfMqClNdiojItMjIaZHtfRE27evlAyvmpLoUEZFpk5GB/kyTultEJHgyMtCf2nKAxdXFLNVNLEQkQDIu0EPhEdbv6uKa5XM0ukVEAiXjAv2F7R2MjjvXqP9cRAIm4wL96S3tVJfkc958TSYSkWDJqECPjIzxwvYOrl4+h6wsdbeISLBkVKD/cddBBofHNFxRRAIpowL96S3tlOTncPFpugm0iARPxgT62LjzTFM7V5w1m/yc7FSXIyIy7TIm0F/Z10PX4DDXLFd3i4gEU8YE+lNvHCAvO4vLz0z9OuwiIjMhIwLd3Xm6qZ1LTq+itED3DRWRYMqIQN92oJ993UNcs1xrt4hIcGVEoD+9pR0zeP/y2akuRURkxmREoD+15QDnL6hkdmlBqksREZkxgQ/0gwNRmtr6uGqZWuciEmyBD/TXW0IAWrtFRAIv+IHeHAv0lXVlKa5ERGRmBT/QW0IsqSnWcEURCbzgB3pziHPqylNdhojIjAt0oHf0RTjQF2GlAl1EMkCgA/3QBdFz6itSXImIyMwLdKBvbg5hBivm6YKoiARfoAP9jZYQp9eUUJyfk+pSRERmXGAD3d3Z3BLi7Hr1n4tIZghsoLf3Rensj3K2LoiKSIYIbKBvbu4F4By10EUkQwQ20N9oCZFlsHyuAl1EMkNSgW5ma8xsu5ntNLM7j3DMfzWzJjPbYmb/PL1lHrvNLSHOmFNKYZ7uHyoimWHK4R9mlg3cB1wNNAMbzGyduzclHLMUuAu41N17zCylSxu6O683h7jyLK2wKCKZI5kW+mpgp7vvdvdh4BHghgnH3A7c5+49AO7eMb1lHpvWUISuwWGNcBGRjJJMoNcB+xO2m+P7Ep0BnGFm/2lmL5rZmsk+yMzWmlmjmTV2dnYeX8VJOLTCoka4iEgmSSbQbZJ9PmE7B1gKXA7cDDxgZu+ab+/u97t7g7s31NTUHGutSXu9pZecLGPZXM0QFZHMkUygNwPzE7brgdZJjvmNu4+4+x5gO7GAT4nNzbELogW5uiAqIpkjmUDfACw1s8VmlgfcBKybcMzjwBUAZlZNrAtm93QWmix35/WWkLpbRCTjTBno7j4K3AE8BWwFHnX3LWZ2j5ldHz/sKaDLzJqA54G/cPeumSr6aJp7wvQOjeiCqIhknKRWrXL3J4EnJ+y7O+G1A/8t/kipt5fMVaCLSGYJ3EzRzc0hcrONM2tLU12KiMhJFbhAf72ll7Nqy8jP0QVREcksgQr0QzNEdcs5EclEgQr0fd1D9EVG1X8uIhkpUIG+WTNERSSDBSrQX28JkZeTxRlzdEFURDJPoAJ9c3Mvy2pLycsJ1GmJiCQlMMk3Pu5saenThCIRyViBCfS3ugbpj45yTt271gQTEckIgQn0XZ2DACydU5LiSkREUiMwgd4WCgNQV1GY4kpERFIjMIHe2hshN9uoLslPdSkiIikRmEBvC4WZU1ZAVtZk9+MQEQm+4AR6b4R55epuEZHMFZhAbw2FmVtRkOoyRERSJhCBPj7utPdFmKsWuohksEAE+sGBKCNjzjy10EUkgwUi0FtDEQC10EUkowUi0Nt6Y2PQ55arhS4imSsQgX6ohT5Pk4pEJIMFItDbesPk52RRWZSb6lJERFImGIEeijCvohAzTSoSkcwViEBvDYXVfy4iGS8Qgd7WqzHoIiJpH+ijY+N09Ec0Bl1EMl7aB3p7f5Rx1xh0EZG0D/QD8XXQtY6LiGS6tA/01t74GHS10EUkw6V9oLephS4iAiQZ6Ga2xsy2m9lOM7tzkvdvNbNOM3s1/vj89Jc6udbeCCX5OZQVaFKRiGS2nKkOMLNs4D7gaqAZ2GBm69y9acKhv3D3O2agxqNq0xh0EREguRb6amCnu+9292HgEeCGmS0reW2hCHO1houISFKBXgfsT9huju+b6GNmttnMHjOz+ZN9kJmtNbNGM2vs7Ow8jnLfrbU3wjy10EVEkgr0yRZI8Qnb/woscvdzgGeBhyb7IHe/390b3L2hpqbm2CqdRHR0jIMDUY1BFxEhuUBvBhJb3PVAa+IB7t7l7tH45o+BC6anvKNrD8V+pEa4iIgkF+gbgKVmttjM8oCbgHWJB5jZ3ITN64Gt01fikbXGhyxqDLqISBKjXNx91MzuAJ4CsoGfuPsWM7sHaHT3dcBXzOx6YBToBm6dwZoP0xh0EZG3TRnoAO7+JPDkhH13J7y+C7hrekubmmaJioi8La1niraFwlQU5VKYl53qUkREUi69A13roIuIHJbWgd4a0hh0EZFD0jrQ20JhXRAVEYlL20APD4/ROzSiLhcRkbi0DfTDY9DVQhcRAdI40NviQxbVQhcRiUnbQNcsURGRd0rbQD/UQp9Tnp/iSkRETg3pG+ihMNUl+eTnaFKRiAikcaC3hiK6ICoikiBtA72tV7eeExFJlL6BHtK0fxGRRGkZ6H2REQaio+pyERFJkJaBrjHoIiLvlpaBrlmiIiLvlpaBrha6iMi7pWegh8JkGcwu1aQiEZFD0jLQW3sjzCkrICc7LcsXEZkRaZmIbSGNQRcRmShNAz3C3Ar1n4uIJEq7QHd3WnvDuvWciMgEaRfoPUMjREfHNcJFRGSCtAv01l6NQRcRmUzaBXpbSGPQRUQmk4aBHmuhz1ULXUTkHdIu0GvLCrh6+RyqizWpSEQkUU6qCzhW16yo5ZoVtakuQ0TklJN2LXQREZlcUoFuZmvMbLuZ7TSzO49y3I1m5mbWMH0liohIMqYMdDPLBu4DrgWWAzeb2fJJjisFvgK8NN1FiojI1JJpoa8Gdrr7bncfBh4BbpjkuO8Afw9EprE+ERFJUjKBXgfsT9huju87zMzOA+a7+xNH+yAzW2tmjWbW2NnZeczFiojIkSUT6DbJPj/8plkWcC/w9ak+yN3vd/cGd2+oqalJvkoREZlSMoHeDMxP2K4HWhO2S4GVwAtm9hZwEbBOF0ZFRE6uZAJ9A7DUzBabWR5wE7Du0JvuHnL3andf5O6LgBeB6929cUYqFhGRSU05scjdR83sDuApIBv4ibtvMbN7gEZ3X3f0T5jcxo0bD5rZ3uP5s0A1cPA4/2w6y9Tzhsw9d513ZknmvBce6Q1z9yO9d8oys0Z3z7gunUw9b8jcc9d5Z5YTPW/NFBURCQgFuohIQKRroN+f6gJSJFPPGzL33HXemeWEzjst+9BFROTd0rWFLiIiEyjQRUQCIu0CPdmlfNOdmf3EzDrM7I2EfbPM7BkzezP+XJnKGmeCmc03s+fNbKuZbTGzr8b3B/rczazAzF42s9fi5/038f2Lzeyl+Hn/Ij65L3DMLNvMNpnZE/HtwJ+3mb1lZq+b2atm1hjfd0Lf87QK9GSX8g2InwJrJuy7E/ituy8FfhvfDppR4OvuvozYMhJfiv8bB/3co8CV7n4usApYY2YXAX8H3Bs/7x7gthTWOJO+CmxN2M6U877C3VcljD0/oe95WgU6yS/lm/bc/XdA94TdNwAPxV8/BHz4pBZ1Erh7m7u/En/dT+yXvI6An7vHDMQ3c+MPB64EHovvD9x5A5hZPfBB4IH4tpEB530EJ/Q9T7dAn3Ip34Cb4+5tEAs+YHaK65lRZrYIOI/YTVMCf+7xbodXgQ7gGWAX0Ovuo/FDgvp9/wHwl8B4fLuKzDhvB542s41mtja+74S+5+l2k+ijLuUrwWFmJcAvga+5e1+s0RZs7j4GrDKzCuDXwLLJDju5Vc0sM/sQ0OHuG83s8kO7Jzk0UOcdd6m7t5rZbOAZM9t2oh+Ybi30qZbyDbp2M5sLEH/uSHE9M8LMcomF+c/d/Vfx3Rlx7gDu3gu8QOwaQoWZHWp4BfH7filwfXzp7UeIdbX8gOCfN+7eGn/uIPYf+GpO8HueboF+1KV8M8A64DPx158BfpPCWmZEvP/0QWCru/9DwluBPnczq4m3zDGzQuD9xK4fPA/cGD8scOft7ne5e3186e2bgOfc/RYCft5mVhy/DzNmVgxcA7zBCX7P026mqJldR+x/8ENL+X43xSXNCDN7GLic2HKa7cBfA48DjwILgH3Ax9194oXTtGZm7wF+D7zO232q3yTWjx7Yczezc4hdBMsm1tB61N3vMbMlxFqus4BNwJ+6ezR1lc6ceJfLN9z9Q0E/7/j5/Tq+mQP8s7t/18yqOIHvedoFuoiITC7dulxEROQIFOgiIgGhQBcRCQgFuohIQCjQRUQCQoEuIhIQCnQRkYD4/4uPZ4TkF57aAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(cstate[:, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = get_params(opt_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rc('font', family='serif')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "251ef73ac5f247918ab75f26822bce73",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAARkAAAEYCAYAAABoTIKyAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8QZhcZAAAgAElEQVR4nO2deXgV5fXHPychISxhDSFACAn7voMoilDUalvbqrXuFlTUVmutP7RqW+suWutSW7W40eLSDRWtRbEqriCyh30JARIICQHCEhKynN8fM4FrTG4m5N7MTXI+z5PnzrzzzrzfmTv35F3PEVXFMAwjXET5LcAwjMaNGRnDMMKKGRnDMMKKGRnDMMKKGRnDMMKKGRnDMMJKkzIyInKmiKwQERWRj0Wkg9+a6oKIPCsiV4bhun8RkRwRWSQiUQHp14rIevfvpyJysptHReS7AfliRWSBiOx3P1ucgAYRkbtFZLGIfCQiX4jI0yLSKVT3WU25np6piNzlPqO7qzl+q4hkikiB+wy+EJEtInKHRx1Pu89vSu3u4Nj5qSKy1mPesLxHx1DVJvUHTAQUaOa3lhDcS+tw3QcwCygF/q9S+hRgSsB+qpsvC2hbKe+COpT/E+AzINbdjwfSgdGR8kzdZ3R3kON3A58F7J/ivnsTPV5/QeCzriFvZuXrAu1O5J6rulZd/ppUTaaxoaqHVLU0jEX8HrhPRHrXkO8NoAh4NIRljwW+UNWjAKp6ELgP2BvCMr5BOJ+pqn4B5AOjw3H9Ksrb7zFfWN8jMzKVcKupX7hV3NdEpE3AsdEiskREPheRZ0TkM7fp8P2A6u19IvK6iOS61f1BIjLPbZ4tFJFrA673uogUicivRORNEdksIj8LON5KRF4RkeUi8o6I/MEt42kRucKtjs+qSbvbzMkUkb+7TaFlIvJfEYmr4XH8Bac28YKISJB8h4GpwFUickatHnj1bAN+LCKDKhJU9Z+qmiEiKQHNtF+IyAcislpEflKRV0Rai8iL7nf0hYhcH3CsmYjMcNM/FpF/ikjP2jzTOtAMp9ZXcf1vu+/FAhF5W0S6VnWSez+z3fv51H3/mrnHXgKSgCfc64xyn4m6zaa7ROSo+x71E5GBIrLGfXe/ds9VXOv/RCRPRLaKyI/cPJ+KyC4ROdPTHYez6hmJf9TQXAJ+EbB9N3Cfux0L7AAucfeH4zQTpgTkXwDMB6KB/sA1wEnASe7xGGAd0KdS1fRpd3sscKhCG/AIMA/nn0EssJCAJoirb1ZN2gP2s4H27vVWV9xLNc9hFk5TqDtQANzgpk/hm82lWe7248BWoFXF86jD99QW+BgoBz4CrgVaVipXgZ+6+2lAITDA3X8O+Ju7HQ9kAKe6+3cC7wPR7v6fKu6pls90FrVrLp0DvBDw/aYBB4F+7v4NwP8qvU8VujoAl1cq++pK79HESuUrkOpuvw7cGXDsMdz3sIp7/tq1gFuA+QH7FwBTvX6XVpP5JkdcS/0xcDEwyk0/GUgE/gmgqiuAqjrW3lHVMlVdr6rPA5uAq0Xkc5wXuwswotI577qfq4BWbjkAFwKvqmq5Os2Gf52g9gq+VNV9qlqOY2TSargeqroD5yWbISI9ash+J3AUmFHTdT2UW6CqpwOTgM04TbHVItKtUta/u/m34hjhC8XprL4C5weNOk2tt900cGpds1W1zN1/EMegVUVNz7Qmhrg1ggwcw/CiHm+aXAosUdUN7v6rwGQR6VLFdfYBPdyazAKcf5a10fI34EoAEYkG0lR1k8dzXwFOD6hlXQjM8VqwGZkARGQi8AfgCvcFnwG0dA93AfYHvJhQdf9AQaX9x3CMxgRVnQisCLhmBQcAVLXI3Y8NKHNPDeV50f61clyKAsoJiqq+AHyKUzsIlu8ITk3nehE5NVhe94e3wP3BBLvmx6o6DUjB+aFdV+n4voDdfJxn1gloDjwSUMbpbhpAMpAXcI2drpGqrHEiNT/Tmkh3v/fBwAbgjwHHkoGBARrfwGkmdq7iOj/Buffvu9ebVUst7wAdReQk4Ns4NW5PqOpu4APgchFJAI6q6oEaTjtGs1qIbLSISFvgVGAQsEFVM91DMQHZdgHtRKRZwH+ijh4uPxanOVRhnGKCZa7ELpwfTAXByhtL9dpDwTSc2k9z4KXqMqnqQhF5AngR50dfXb6JwQoTkZuA91V1nZt/v4h8gNOMCszXQVUrjG+CqzEPKAZuVNWv3HwxHP9R7iDguYpIRyA+4NlVELJnqqqFInIL8JWInKaqn7o6lqhq4PB/e77+zyBQy+KAe62VFlUtEZF/4BirVsAva3kLfwN+AxwB/lGbE60m49Aepwq4GejtvnTgWPwKFgK5wEUAIjIc6OPh2ptx+mVwq8FDa6Hrn8BlIhIlIrHAD2sopzrtdUZVs4GbgQkesv8Wp7+qqmq/VzoA11V0OItIPM49fVIpX0VnZE9gHPBPtzn4N443j8D5gVTMBZkFXOE2G8CpoQyrQkNIn6mqLgH+B/zCTXoNOKmiGSoiiTjNtqp+l5uBYSLS3O3wnVzp+EGgpYhMEpFffPN0wHkmFwNRAcaqKqq61ps4Na9rgPeC3ec3ONGOuYb4B5yG8yUqTpvy3+7ff3FevCjgeZwv9HX3S9kPPOKePwZYijPi8jjOC/8T99gjbt71wC0BZfYHluAYqRdx+l3WA99yr1+E04RKc8tUYBHOj6wVTnt4BfAW8ADwkXvdK3A66HKAXwfTjtP2r8j7U5xO1Bw37dIqntNf3OOLgLGVjr3F8c7Ik908Obid1wH5xgKb6/BdjcD5EX6B0wG6BLg14Hiq+6ym4fxw11R8F+7x1u7z+ML9zp/geEdvDI5hWeh+lw+cwDO9K+AZXl2F/lvdYwV8vdN0Eo4BXuBqPAv43N3/EBjn5nua4+/Td9134Q2cgYN/uds5uO8acCNOH+EinBr5Bxx/l7oFlL8W+EHA/tfuuaprBeR9Hniytt+luCcbHqhUNUdE1gDTVXVemMprjdP+Peru34ozGe2icJTXkBCRVGCrqgYbWjdCiIg8DPxb3SaoV6y5VDtedju+EJFROM2BL8NY3sU4IyGISHOcJl3tqqqGUUdE5Eq3T2tEbQ0MWMdvbXkfeFdEDuN0gP5Ig7dt68oynElRl+BUrefjNOuaNCKSgjt07Y7KTNavj/oZoeU+nH6ku0/kZN+aSyKSBNwPDFPVMVUcj8KZv3AI6AG8oKqL3GNnAOfjdMSqqt5Tb8INw6gVftZkTgXm4sycrYofA21U9XZxVksvEpEBODWIZ3E6pIpFZI6ITFbVD+pHtmEYtcE3I6Oq/3YnO1XHd3EnDKnqXhEpwuk17wRsU9ViN9/nbt5vGBlx1gldC9CqVatR/fv3D90NGIYBwNKlS/eoarUuOCK5TyYRZ7y+ggNuWqdq0r+Bqs4EZgKMHj1alyxZEh6lhtGEEZFtwY5HspHJxVnYVkEbN02rSTcMIwKJqCFscVwbVFS73sGZ7IXbJxOHM+FqIc5CsYp1KOPdvIZhRCC+1WRE5HSc2YZdROQ3OAvRpgBDgOtxptSPEJHf4SyOu9IdpiwUkZ8CfxSRPGCVdfoaRuTSZGb8VtUnU1JSQlZWFkVFRdWcZVQmLi6O5ORkYmJCvf7SaKiIyFJVrdbbXyT3yYSdrKws4uPjSU1NRYI6fjPAWeeWn59PVlYWaWk1uqIxDCDC+mTqm6KiIjp27GgGxiMiQseOHa3mZ9SKJm1kADMwtcSel1FbmryRMQwjvJiR8ZHTTjuN6dOnM3XqVNq2bcv06dOZPn06U6ZMqdV1Zs2axf79TvSLzz77jJEjR7JgwYLQCzaME6BJd/z6zVVXXcXUqVNZvXo1H330EY8+6oQteumlar1bVsmsWbOYOHEi7dq149RTT2Xo0No43zOM8GJGxuWet9ewdqdn38ieGNi1Db87d1C1x6dOnVpl+vr160lJSWHatGl8+eWXJCcnH+tsnTVrFjNnzuTBBx8kMzOT+fPnk5mZyRNPPEH//v25/nonvND8+fOZM2cOy5Yt45VXXiE1NTWk92YYXrHmUgTy8MMPk5uby4033sjcuXO5/vrrv9aEuvbaY/HhOOuss0hNTeXmm28+ZmAA0tLSeOqppzjvvPOYM8dz9ArDCDlWk3EJVuPwg86dO9O+fXsAhg8fXus+lt69nciyCQkJZGZmhlid0WgpLebIvp3sz9nOofwskkd9lxbx7ep0STMyEUrloeL4+HgOHHCac9u3b//asejoaFSV9PR0Bg4cWOX5hqFHCynI3ca+nRkczttO6b7tcGAnsYU5tCzOpX1pHm31IC2AFu45m9u9Te/hXgJUVI8ZGZ85cuQIM2fOpKCggBdffJGrrrqK559/noKCAh577DFuueUWwKnNlJeXc99999GjRw8KCgp47bXXuOSSSzj77LOZMWMGRUVF3HDDDaxatYrZs2fTr18/3n77bfbt28fmzZuP1W6MxomWlVCQs5X8HRs5nLOJ0vytNDuYRcvCbDqU5NBB99MOCKyX7NE25EclkBfTie1tBlHaMglpk0RM2260TOhGn17V+ZTzTpNeu7Ru3ToGDBjgk6KGiz03H1HlQO52dmes4mDWWjR/C3EHMmlXtJ3Eslxi5Lir46MaTY50Ij8micNxXSmJTyaqXXfiElJo07knHbv0IKFdG6Ki6lbrtbVLhtEQUWXf7m3kbFrG4R3pROWto82hzXQp2UEbimjjZjuszdkZ3ZXsuD5siT8TOvSkZedetEvuS5duaaS0aE6KrzdiRsYwfEfLy9i1JZ3cDYsoyV5Jq31r6Vq0hfYcpL2bJ5f27IpNZUWH71Ke0IeWXQeSkDqIrslp9GkWHfT6fmNGxjDqmYK8LLatXEBRxkJa56eTUryRrhyhK1CkMWxrlsq6dqdTnjiYVilD6dJ7BImdu5DYQDvzzcgYRpjJ3b6eHcvmQ+bnJBWsoJvmMBQo1mZkNksjvcPZ0G0kHfuMo0f/YfSLbV7jNRsSZmQMI8Qc3J/HloVvU7rxfZL3LyFJc0kE9tKGrS2Hkpl0EW37nUbPIafQr2Urv+WGHTMyhlFXVMnatIKdC/9Fm+yP6FO8juGiFGgrNrcaQUbyVDoNOYNeA0czKrrpTbI3I+Mzixcv5rbbbuPo0aM89thjjBs3zm9JgLOa+6abbuKxxx5j4sSJfsuJPFTZvmYROYv+QZed8+lenk0ysDG6N4uSp9J26HfoN3Iio8xNqa+OxIOGmhWRF4BeAUlDgZGqmikimUCmm56tqpeFX3F4GDt2LBMnTuTQoUMRY2AAW81dDft2ZbL5gxfpnPE6KeU76KpRrGs+hB1pV5I24WL6dkulr98iIwxfjIyItKTmULPzVfUfbv42wCxVzXSPzVLVu0Mqat7tkJMe0kuSNATOmXFCp2ZnZ3PHHXcwePBgNm/ezHXXXceoUaN48803mTt3Lv369SM9PZ1nnnmGNm3acNFFF5GRkcGZZ57Jp59+yvnnn897771HdHQ0Q4cOZdGiRVx66aVMmzYNgLvuuovS0lKio6OJj4/ntttuA+Cmm26ipKSEnj17kpWVFbJH0ZDRslLWfvIvdPELDChcwhhR1jYbwBf9fkOfiZcypHM3vyVGNDUaGRGJVdWj7nYCEKOqu+pY7snUEGq2wsC4XA28GLA/QURuwwnyNk9Vv6ijnohj+vTpnHvuuVx66aVkZmZy3nnnsXz5ctq3b88TTzxB27Zteeyxx5g9ezY33HADDz/8MOPHj+eee+6hqKiIXbt2MWLECO644w4eeugh8vLy+Na3vsW0adN47733WLRoEfPnzwdg4sSJnHXWWWRnZ7Np0ybmzZsHwNy5c/18BL5TeCCftf99hq4bZjNIc8ihIwu7TaHrhKkM7D/Mb3kNBi81mduBe93tWOAR4PI6lltdCNpvICJRwLeBJwI1qepit0a0TES+p6qbqzj3WCzslJQa5j2eYI0jXKxatYrExES2b9+OqpKYmEh5eTmtW7fm3nvvJSEhgWXLljFo0PHV47179yYmJoaYmBji4+PZuXMnffs6lfdOnTpx8ODBY9cuLCxkxgznnrt3705eXh5r1qyhT58+x67Xs2fPerzjyOHQvlzWv/4gA3a8xmiKWNNsINuH3caIb1/OqY1seLk+qNbIiMhQYDgwXESudJOjgFCMuVUXgrYqfgD8RwMWWanqYvezUERW4ESR/IaRqRwLOwS664X09HT27t3L5MmT+f73v4+q0q1bN6Kiorjmmmt48sknmTBhAjNnzmTnzp3Hzqtq5XVVacOGDWPhwoXcfvvtAHz44Yf07t2b4uJiPvzww2P5MjIywnB3kUvhgXzWznmI/tteZqQWsaT16bSadAsDR02wVe11IFhNpj2QFvAJUAY8HoJyj4WadZtM44Gn3XC0paoa6KJuCnCsY1dEJuM02d51k3oDW0KgyReWLFnCJ598wtGjR7n//vsB2Lp1Kz//+c956623SE9PJycnh0mTJgFw9dVXc9999zFp0iSWLl16bIX1rFmz2LZt27GV3MXFxcyePZtVq1axZMkSVq9eTUFBAXPmzOGCCy5g8eLF3HHHHTRr1oyioiJmzJhB9+7dmTdvHtdccw3du3dHVZk9ezajRo0iPj4+2G00aLSslOVvPk7v9McZzWEWtziNtmf/lrHDTvJbWqOgxlXYItJHVTcF7Ee74WLrVrDImcCPgDygRFXvEZFHgL2qOsPNMxy4TFVvDThvCHA3sBToijO69FBN5dkq7NDRmJ7btvTPKJ17M71KN7EyZjgx5zzAwJGn+i2rQRGKVdibRWQEx5s3VwDT6ipMVd8H3q+Udlul/RXAikpp6cAFdS3faNoUHS5gzV9/yYjdr7NH2vH58N9z8rnXENUEJ8uFGy9G5m1AgT3u/pDwyTGM8JO9YSnl/7iSEWXZfJFwPoMue4TxHRL8ltVo8WJkclX1qoodt1bTaFBV69SrBQ3ayZkqK9/+M32X3sNhacGKSbM4deIP/VbV6PFSN1woIoFjmY1mgkBcXBz5+fkN+4dTj6gq+fn5xMXF+S2l1pQUF7L8qYsZtuzXbIodQMm0TxlpBqZeCDaEvRfYDwhwp4iou90GmFUv6sJMcnIyWVlZ5OXl+S2lwRAXF0dycrLfMmpF0eECtj71A0YULeeTrldz8tSHibE1RfVGsObSjar6auVEEflRGPXUKzExMaSlpdWc0WiwHN6/h51//h59jm7g86EPMOGCG/2W1OSotrlUYWBE5HuV0v8dblGGEQoO5GWT+6czSDm6iSVjH2e8GRhf8NLx+5SI3BKwrziT3+5R1ezwyDKMunEgbycHnjmTzmV7SD99JuO+ZbMe/MKLkXkNZ+FiBtATGAt8BtwFXBc+aYZxYpSVFLPr+QtJKctjw1mzGT3+bL8lNWm8jC4dUdUPVHWr64ohVlU/BTbVdKJh+MHKmdfRr3g1y0bcz3AzML7jxciMEZHRItJORMYCY0UkFugXZm2GUWtWvvEHRua9wceJlzP+h1bRjgS8NJd+B/wFGASsAX4KDAQ+DqMuw6g1GUveY+CKB1gaN5aTr3mi5hOMeqFGI6Oqy3H6YQAQkS6u06oV1Z9lGPXLgdws2v/nGrKjkki99lViY20eTKQQbDLeMFVdGeBLpoJzgQvDK8swakfGa7cwQAvZfcEbpHbs5LccI4BgfTI3uZ9TcfzJVPx1CLcow6gNmUvfZ/i+91iUdBn9h1TrccDwiWprMqp6tbt5k+teAQARGVTNKYZR72hZCTJvOrtIYNil99Z8glHvVFuTEZEUEUkBCiq23f0GG37EaHysfvMP9CjNZPPIX9OubTu/5RhVEKzjdwFObKPKfhBSgDvDpMcwPHN4707S0p9kWcxITvnuFL/lGNVQ0wLJ/1ZOFJHvhFGPYXhmy6vT6a/FNP/+74k2j3YRS7AFkt8wMMHSDaM+ydm4lKF73uGzThcxyDp7Ixoz/0aDZNf//kSRxjDgR7/1W4pRA8HmyfwKeDQUkQmquX5NsbCnANcDRW7SC6o62z12OTACJ0TLFlX9Szg0GpHJ0cMF9M39L0vjv8X4pK5+yzFqIFifTIyqlonI9ar6bEWiiJyuqnVaUuAxFjbAxQHxryvOTQamAyNUVUXkKxH5MDBsi9G4WT//eYZSRNwpdQ6aYdQDwYxMP7c2cZaIFAakn0vd1y3VGAvb5UYRyQFaAn9S1b04IWuXBkSUXAicg60Kbxqo0nbN39ggPRlx0mS/1RgeCNYnMwMneFpFBMlQzvj1Egv7Y+BhVX0UWAL8qxbnAk4sbBFZIiJLzI9v4yBr1QJ6lGays8+lFiOpgRBsxu8aYI2IvB2GGb81xsJW1a0Bux8Cb4lItJuvd6VzvxEH271Gg4yFbVTP3o+foY22ZMjZV9ec2YgIvPwryBCR+0XkbRG5D2eCXl05Fgvb3R8PvCMiHUSkDYCIPCQiFUawD7DV7YR+Dxglx4MlnQzMC4EmI8Ip2r+b/ns/YFn7s0noYEvoGgpe/Mk8huPT9yWcH/vjwLV1KVRVC0Xkp8AfRSQPWKWqH1TEwsZpquUAz4jIVpyolVe452aJyKPA4yJSBjxvnb5Ng43vPsNQSmk/wZxRNSS8GJktqvpIxY6I/DoUBdcUC1tVnwxy7svAy6HQYTQQystJ3PgaK6MHM3TESX6rMWqBl+ZSd7cvBLf50i28kgzjm2Slf0xSeQ77BlxqYYUbGF6MzPtApoiswIlYYP0fRr2Tu/I9ylXoM95CyzY0vLjffEtEPsEZ0dmsqvvDL8swvk7r7M/YGN2T/l2sIt3Q8DTRQFX3q+oSMzCGH5QeOUBa0VpyOo7zW4pxAthsJiPiyVz2P2KkjLi+k/yWYpwAZmSMiOfgug8o1hj6jD7TbynGCVBrIyMi14dDiGFUR/ucL1gXM4CO7c29ZkMkmI/fvSKSISI7ROSwiGS6CyV/V4/6jCZO0f7dpJZmsK/zyX5LMU6QYDWZG1W1J/AIkKiqqTgLEf9cH8IMAyDzK2fGRJuBZ/isxDhRgrnffNXd7Kyqh920QzgLEg2jXija9CEHtCX9R07wW4pxgnhZVjBARG7F8dfSF2f9kmHUC0l7FrEubhgntYjzW4pxgnjp+J0GdHI/E4BrwqrIMFwO7NpEUvluDnc71W8pRh3wMuN3r4jcDnQE8lW1PPyyDAN2LJnHIKDjkLP8lmLUgRprMiJyFs6apReAS0TE1tkb9UL5lgXs1vYMGDzKbylGHfDSXDoX6A98rqqvAL3CK8kwgPJyuhd8xabWo4mNifZbjVEHvBiZLFUtAircV9r6JSPs5G9bRTs9QEn38X5LMeqIFyPT1+2TGSgiN2L+ZIx6YPemZQAk9B7jsxKjrngxMjfjzI1JAJKAX4VVkWEARTvXUKZCcp+hfksx6oiX0aWDIvIbnFAoe210yagPYvI3kCVd6NHW5n42dGozuvQiNrpk1BMdCjPYHZfmtwwjBHiZ8VsxuvQLVX3FjShQZzzEwv4VTvMsBxgF3KWq691jmRwPzZKtqpeFQpMRGWjJEZJKd7Kpk7l2aAx4MTJZqlokIiEbXfIYC7s1cIsb7/oi4Pc4Bg9glqreXVcdRmSSm7mGzqJEdx7otxQjBPg1ulRdLOxjqOpvA+JdRwGHAg5PEJHbROQ+ETmlukIsTG3DZE/GCgDa9bBO38aAX6NLtYlnHQv8BPhNQPLtbiyoh4AXRaR3Veeq6kxVHa2qozt16hQC2UZ9cHTnWko0muTeQ/yWYoQAr6NLv8eZ6bvFdfdQV2qMhQ3HDMwzwK9VdUuApsXuZ6EbqmU81cTDNhoeMXs3sEO60LNNa7+lGCHAy+jSFcAanNGltSJyZQjK9RILuwXwF+AxVV0qIhe46ZNF5OyAa/XGCaNrNBI6FGaQZyNLjQYvHb/fB1JV9aiIxAGvAn+rS6EeY2G/AgwG0tyIga2AOTg1nrtFZCTQFZijqp/VRY8ROejRQpJKd7Gx09k1ZzYaBF6MzHJVPQrgjjJ9BSAi3VQ1+0QL9hAL+/xqzksHLjjRco3IJi8znURRopJsZKmx4MXIDBaRe3Em5PUEEt0m07nAheEUZzQ99mSsJBEbWWpMeBld6gKUAT3cz11AGs4yA8MIKUd3ruGoRpPSe7DfUowQ4aUmc5PbRPkaIjIoDHqMJk7s3o3skK70im/ltxQjRASLuzRIRE5R1XQRaSUij4jIEyKSCKCqa+pPptFU6FCYQW4LG1lqTARrLj2KMzwM8AAwFNgOPB5uUUbTRIsPkVSeQ1G7vn5LMUJIsObSUlWtGKq+CBilqjtF5KF60GU0QXK3rqIzEJ00wG8pRggJVpMpBRCRcTgrnXe66UVhV2U0SfZuXQVA+x7DfFZihJJgNZle7oLIy4GXAEQkGWeCnGGEnOKdayjWZnTvbWMKjYlgNZlbcYat/wU87abdB/w33KKMpknzfRvZJt1o17ql31KMEBKsJnMyzsLEoxUJqjo1/JKMpkqHwgy2tBiIdfs2LoLVZHoB80TkJRH5toh4mbhnGCdE+ZEDdC7P5YiNLDU6qjUcqvqoqk7GGb4eCywQkadFZEK9qTOaDHnbVgMQk9TfZyVGqKmxdqKqm1X1PlWdgOMy8xwR+Un4pRlNiX07NgLQtpsZmcaGl2UFiMi5OI6lVgD3qWphWFUZTY6i3E0AdE7t57MSI9R4cVr1CE5UgQlALI6vF8MIKbJvK7u1PYkdOvotxQgxXjpz97ujShmquhzHqZRhhJSWh7axu1lXoqLEbylGiPFiZBLcz4rIAfHVZTSME6VDcTYHWnb3W4YRBrz0yWwUkbVAuYj8CMfvrmGEjPKig3TUfZS0SfVbihEGvEQreFZEPgYGAemquiH8soymRH7WBjoB0Qm9/JZihAFPo0uqug5YB07ANFWdWdeCPYSpjcNxN5EN9AFmqOpG99jlwAgcT31bVNVqVw2YvdvX0wmI72oT8Roj1RoZEdmLE5JWON4fIzhD2XUyMh7D1N4MbFfVR0RkCPACcJq7SHM6MMINYfuViHyoqpvqosnwjyO7nZBZnVJsjkxjJFjH742q2lNV09zPnqqaBvw8BOXWGKbW3V8IxyIUDHNjMsR0mlUAABaASURBVH0bx9dNheFbCJwTAk2GX+zNIF/b0KVzZ7+VGGEg2LKCVwFEZLSIjHG3xwFvhaBcL2Fqq8tTmxC3Fgu7AdDi0DZyorsQbcPXjRIvQ9i3Aofd7cPA/SEo10uY2uryeApxCxYLu6HQviib/S1s+Lqx4sXIfKWqa+FYs2VXCMqtMUwt8A5Oswq3T2alqh4A3gNGiRtW0s0zLwSaDB/Qo4Ukah5Hbfi60eJldKmviCSo6h4RScAZ6akTHsPUPgk8KiK/wXFofrV7bpaIPAo8LiJlwPPW6dtw2bdzMx2AqI49/ZZihAkvRmYWsFxE4oEC4OJQFOwhTO0R4IZqzn0ZeDkUOgx/yd++jg5AKxu+brR4mYz3BdA9oDbToh50GU2Ewhx3+Lq7DV83VoLNk6kwKhMC0gCuAKbVgzajCaB7MyjQVnRJ6uq3FCNMBKvJzMaZf/Ikjh+ZCoaEVZHRpGh+YBs7o7swICbabylGmKjWyKhqxQS3m1T1UwDXz+/J9SHMaBq0K9pBZpw1lRozXoawA8P5DQSuCZMWo4mhpcUklu+mOL6H31KMMBKsT6YN0A7oLyIpbvIhoLi6cwyjNhzIyaAtitjwdaMmWJ/MecAUIBUYjrM4shR4N+yqjCbBnm3raAu0TLLh68ZMsD6ZvwJ/FZELgTdUtbT+ZBlNgUPu8HWCrb5u1Hjpk3kKsLfACDnl+Vs4pHF06ZpSc2ajweLFyLysqqsrdkRkZBj1GE2I2APb2BWVRFysJ99pRgPFy7fbTkRm4HjGU+Bc4MKwqjKaBG2P7CC7earfMoww46UmMwo4gtMBnAZ0CKcgo4lQXkbnshyO2PB1o8dLTeanqrqoYkdE6rwK2zAO5m4lnlKkQ5rfUoww4yUW9iIRaSUiKe58mcvqQZfRyMndvAyA5t0G+6zECDc11mRE5BacRZHxwG6gK3B3eGUZjZ1DmcspVyG572i/pRhhxkufTJKqjgCeU9XxwJ/CrMloAjTLW802SaJb54SaMxsNGi9G5pD7WeFXt1+YtBhNiI4HN7Irrg/HvagajRUvRiZZRM4FdojIFqBLmDUZjZyywv0kledwpMOAmjMbDR4vnvGurdgWkYXAxrAqMho9uzYuJRlonjzcbylGPVBjTUZE2onIwyLyH+ByoGX4ZRmNmX0ZSwFI7GOdvk0BL82lF4B84EWcSAIv1KVAN+zJTBG5XUReEJFvhA0UkTEi8oqITBeR50RkWsCxZ0VkQcCfeeprYJTvSmevxpOa2stvKUY94GUy3gZVfaRiR0QermOZDwL/U9V/un09j+IMkQfSBXhSVReLSAyQKyJvqOoeIEdVr6+jBsNH2hSsY1tML0aYy80mgZeazMGKCAUi0hLY6W5fcoJlHotxTdUxsFHVt1R1cUBSKVDibseLyK9F5FcicqOI2Oq6hkRZKV2PZlLQ1hb2NxW8/EB/AfzGDcKWCOwVkZtxwsO+VtUJIvIeUFX09Lv4eizrA0B7EWkWxF/NjcCDqlrg7r+CEwyu1A0GdwdwXzU6rgWuBUhJMXcCkUD+9tV0pATpYq3cpoIXI/OAqj5VOdGNAFklqvrt6o6JSEUs6/04hmpfdQZGRC4FWqnqsfjbqrosIMuHwK+oxsio6kxgJsDo0aO1Ok1G/bF741d0BNqnmceQpoKXtUvfMDBu+jMnWOaxGNe4MbDBiYQQ4EsYEbkGSFTV+0VkiIj0ddN/H3CtPsDmE9Rh+EDxjpUUawyp/W34uqngR3/GncDDrtHoBUx304fixHoaIiI/AP6AEx73h0BH4Oc4c3Q6uf5tCnFmH99Sz/qNOhC3dy1bo1Po38pmQjQVgkUrOB34RFVD2sxQ1b1UEYFSVVfgBo5T1blA22rOnxJKPUY9okrSkc2saW2hu5oSwZpLF6uqisjZgYkikhpWRUajpTA/i/ZaQGknc+/QlAhmZEpcg3JWhS8Zt8/kZ/WizGh0ZK//CoDWPUb4rMSoT4L1ySzEGZnpAwS+FSnAbeEUZTRODm1zBga79R/jsxKjPgkWd+k14DUROVdV365IF5Hv1Isyo9ERnbuGLBLplpjotxSjHvGyCvttETkDGAasUNX/hl+W0RjpeGgD2XF9SDYfMk0KL6uwf4szTNwDmO7uG0atKD1ygC6lO82HTBPEyzyZWFU91kQSkYfCqMdopGStXECqKDEp5t6hqeFlgWRZpf3ycAgxGjcFK9/miMbSZ+w5fksx6hkvNZlSEXkLyMCZoftleCUZjQ5VknZ/THrzEYzt0M5vNUY946Xj934ROQtn2v87qvp++GUZjYm9mavoXL6bNSlX+y3F8AFPa5dUdT4wP8xajEZK1pdv0AHodtIP/ZZi+ICXPhnDqBMtMt9ng6TRt3dfv6UYPmBGxggrRQV59DyyhuzE0y3GUhPFyzyZB+tDiNE42bpoLtGitBn6Pb+lGD7hpSYzRET+LCI/E5FWYVdkNCpK1r/LHm3L4DET/ZZi+ISXjt8fq+oRERkI/FlE8oE/qerWMGszGjhaVkLqvi9YFX8ap8bG+C3H8AkvNZnJIjIKx2H3ycAu4AIRuTOsyowGz/aVH9GGw9C3WpfPRhPAS03mZWA58BTwE1UtBxCRJ8MpzGj45C97my4aTb9Tvu+3FMNHvBiZ31Y4ExeRliJSjrPUYEVYlRkNnk67PmJN7FBGJCT4LcXwES/NpRYB212AF1S1RFVfCpMmoxGQv3Ul3ct2UND9W35LMXwmmCPxFCAV6C8iE9zkKKBOjsVFpAMwA2ctVB/gTlXdXUW+TCDT3c1W1cvc9FTgtzihUFKB/1PVQ3XRZISerHl/oJXGkDZpit9SDJ8J1lwaAfwQGA5UzKIqA/5TxzK9xMIGmKWqd1eR/ixwlxsn++c4wd3Mx00EUZifzYDcd1jY9juc3t0idzZ1grnfnAvMFZExqvpVCMv8LvCAu/058Ndq8k0Qkdtwok3OU9UvRCQGmAR8FXD+81RjZCxMrT9s/M/jDNUyEs78pd9SjAggWHNJ3JhLuwMjOwLXq2rQ4esQxcK+3a2ttASWicj3gMPAkYBYUAfc61WJhamtf0qPHCRt62t8FXcyJw2xULRG8ObSl8BY4GNgK8ebTCk4USCrJRSxsFV1sftZKCIrcELavgq0CDCAbYDcYFqM+mXtvGcZyiE45ed+SzEihGDNpbHu5k0hjlZQEQt7B5ViYQPJqrpdRCYDMar6rntOb2CLqpaIyEfAGGBx4PmG/2hZKYmrn2NNVD/GnGYe8AwHL/NkdlX0y4jIOOCTOpZZYyxsnNrJ3SIyEugKzFHVz9x81wN3uY60UrBY2BHDxgWv0a98N5tG3k5UlK24NhykplDXIvJ34F5VXSsiQ4CrVfXmelEXQkaPHq1LlizxW0bjRZUtD40j9uh+Ot2RTlzzWL8VGfWEiCxV1Wo9xHuZjLdEVdcCqGo6ztolw/gamz/5O72OrmdL75+YgTG+hhcj01dEEgDcz97hlWQ0NIoO7qXdgjvYJKmMOd+GrY2v46VPZhawXETa4IwIXRxWRUaDY8Psmxlcvp/t57xEq5Ytaj7BaFJ4iVbwBdBdRBJUdU89aDIaEFsX/5dhuXNZkHAJE8dN8luOEYF4cb+ZJCKvAB+JyGwRqWqSndEEKSk6RNy7v2QHSYy48mG/5RgRipc+mQeAucBPcNYtzQirIqPBsPrl2+lSnsOuCQ/Ttm1bv+UYEYqXPpn1qvpPd3uZiPQMpyCjYbD5izcZuuNlPmv7HU79lsVTMqrHS02mt+ueoWJ0yYxME2fn+sV0mX8dGVE9GDDlT37LMSIcLzWZvwIrRSQeKMBGl5o0+3IyifnHxRyiJbFX/puOHTr6LcmIcGo9uiQiNkbZRCk6tI+C539IQnkh2384h4FpffyWZDQAgrl6mFBFGjgOpqaFUZMRgZQWF5Lx5x/Rt2Qby0+byZgR4/2WZDQQgtVkHgdWctzFQwVDwifHiEQO79/Djmd+yMDidD4ZeDcTzrjQb0lGAyKYkblJVT+vnCgi9i+sCZGfncHhF39Az9JsPh/+MBPOu95vSUYDI5g/mc8BRCQWuA6IwXFktal+pBl+k7VuMc3/eRHty4+wetJLjJ/4A78lGQ0QL0PYjwMdcHy37ALuDasiIyJYNe852v3j+5Qr7Dz/TUaagTFOEC9GJlNV7wF2qWoGkB1mTYaPFB7IZ8UTFzD0y+lsb5ZKyZT36DdsnN+yjAaMl3kyPUWkOaCui0xbu9RI2bz4PVrPu4HB5fl80v06xl15P7Gx5hvGqBvBhrCHquoqYD6OI3HFCS9iDkMaGQW5O9jy918xcu877JAurDvn30wYN9lvWUYjIVhN5hkRuVNV33Cdd/cGNqvq/nrSZoSZ0uIjrJwzg34b/8JgPcqnnS9l6KUP0r1de7+lGY2IYEZmNtBVRP4CbAD+GgoD4yVMrYhMBP4M5LlJicA/VfVuEXkW6B+Q/eeuW1DDI+UlR0l/9zk6LX+KUeW7WNp8HO3Oe4TT+g/zW5rRCAk2hP2su/maiPQBful6x3tdVRfUoUwvYWp3Aper6nIAEXkBeMk9lqOqNlnjBCgtLmT1O8+QlP4swzSXTVE9+Wr8c4w+48KK2dyGEXK8dPyCE/h+DXAjcDnOkPaJUmOYWlXdWLHtOslqrqrb3KR4Efk1UIoTUfLZqoLDGcc5sHsbm9/9Ez0y/8Vw3cfa6H5sG3svo8+4iOhoLwOMhnHiBOv4PQunw/c64EpgM/Ac8PeaLhqiMLUV/Ax4NmD/FWCVqpaKyCPAHcB91ehosrGwtayUjK/mUfj5Xxhw4HOGo6xoPoot425k9Ok/IMqMi1FPVBt3SUTycObRvAI8F6p+DxHZAZyiqjvc/pnNqlplzcgdOn9dVb9bzfGzgV+pao3OZZtE3CVVcjZ8ya5P/0b3nfNI0L3s1XjSE8+l6+Sf0ae/LTszQk9NcZeCNZfmAdeqalGINdUYpjYg76XAa4Eni8jvVfVWd7cPTg2r6VJezq51X7B78RwSst4nuWwHHTSalXFjWDfgAoadcQmnt473W6XRhAlmZH4aBgMD3sLUVnAhUHk+eycRmQEUAv1ogmFqSwv3k7nkXQ6teY/k3AV00b100ihWxwxmc58r6DvpMsZ0TfZbpmEAHsLUNhYacnNJS4rIXvMZe1Z/SKusT0g9spYYKeOwNie9xRiKe51N7/Hn061rN7+lGk2QujSXDJ8oObSXrNWfUbDxM1ru+pKUI2tJ5ijJwHrpyWeJl9BiwFkMGHMG4+Jb+S3XMIJiRsZn9GghuzYuZc+mL9Hs5XTct4rksu2kAWUqbIpKY1GHHxCddiopwyfTLzmZ/janxWhAmJGpR4r272bnhq84kLkc2b2adgc20K10G10ppyuwV+PZ0rw/GzufQ1zPcfQYcir9EjuZUTEaNGZkQo0qh/bsYHdGOgez11G+ex0tCzbRuTiT9lpwLJ5MrrZnR2xPMhJPpVnyKBL7nURaz36MiYn2Vb5hhBozMieCKgfyssnL2sDB7I2U5G0ipiCT+MLtJJXsoDVFtHazHtIWbG+Wwtr48ZR06EeL7kNJ6jua5G7dSYyyGorR+DEjUwXlZWXs25PNvuytHMzN4Gj+NmT/duIOZ9G2aBeJ5Tm04Sht3PxlKuySRPbEdmNXh+9S1rE3LZIG0Cl1IN1SejPQaidGE8aMjEvh4QNsffI7tC/JJaF8Dx2ljMCwZQXairzoRPY0705W61OgfSrNO6XRrlt/klL7ktyyJTYzxTC+iRkZlxYtWqMSzc7Wg9nWuivSthuxHVOI79yLhOTetGvfkbbWAWsYtcaMjItERTH4jo/9lmEYjQ5bimsYRlgxI2MYRlgxI2MYRlgxI2MYRlgxI2MYRlgxI2MYRlgxI2MYRlgxI2MYRlgxI2MYRlgxI2MYRlgxI2MYRlipdyMjIlEicp2I5IrI4CD5zhCRp0XkbhH5XUB6BxGZKSK3i8gLboRJwzAiFD8WSA4DvsQJaVIlItISJ2rkIFUtFpE5IjJZVT/AWyxtwzAihHqvyajqclVdUUO2k4Ftqlrs7n+OE0Mb93NhFemGYUQgYanJBIuFrapvebhEYLxscGJmJ1ZxLGgs7cBY2MAhEdngoewEYI+HfH5h+upGJOuLZG1Qvb4ewU4Ki5FR1W/X8RK5QGBs1TZuWuCx/W76vqoMjKtjJjCzNgWLyJJggar8xvTVjUjWF8na4MT1RdTokoikuZsLgR4i0tzdPxYzm+OxtCunG4YRgdR7x6+ItAduANoC14rIq6q6SEQ6AZ+JSC9VLRSRnwJ/FJE8YJXb6QvVx9I2DCMCqXcjo6r7gPvdv8D0PKBbwP77wPtVnL8XmBZGibVqXvmA6asbkawvkrXBCeoTVQ21EMMwjGNEVJ+MYRiNDzMyhmGEFQuJ4iIiZwDn4wyRq6re47MkRCQJp+9qmKqOcdPicGY5ZwN9gBmqutEHbb1cbcuAZCBfVe8VkQ7ADCDD1Xenqu72QV8U8DbO7PJYnEGCq4AWkaDP1djC1TdfVadHyncboG8RUOTulqnq5BP6flW1yf8BLYHNQHN3fw4wOQJ0/Qg4F1gSkHY7cJu7PQT41CdtY4AfBOyvBUbhLAf5sZt2LjDbJ31RwG8C9ucCl0WKPrf8PwB/BR6NpO82QN/dVaTV+vlZc8kh2DIG31DVf/P1mc8QsKxCVdOBYSLSpvK59aDtK1WdG5AUBRwmQpZ9qGq5qt4PICLNcGpbGyJFn4hc4Za/NSA5Ir7bAIaIyK/cRconvKzHmksOwZYxRBrVaT3gjxwQkfOA91R1vYh4XvZRT9q+DfwS+I+qLokEfSIyEBigqneKyNCAQ5H23T6sqotFJBr4REQOUotlPRVYTcYh2DKGSCOitIrIJGASzg8Zvq4v6LKP+kBV31PVs4E0EflZhOg7DygSkduBU4GxInIzEfbdqupi97MM+BTne67187OajMOxZQxuk2k88LTPmqqjYlnFpyIyBFipqr78p3Or0KcBvwC6iEiPAH078HHZh1tbSFPVivK3Aj0jQZ+qPhCgMw5orapPuNuR8t32B8ar6gtuUh/gdU7g+dlkPBcROROnozUPKNHIGF06HbgSOBt4BqejEJwRiF1Ab+BB9Wd0aRTwMbDETWoF/Bl4C3gY2IYzonO7+jO61Av4Pc7oVwwwALgJOBoJ+lyNF+AssYnFeXZvEgHfrautq6tpGU6NJQa4BWhHLZ+fGRnDMMKK9ckYhhFWzMgYhhFWzMgYhhFWzMgYhhFWzMgYhhFWbJ6MEXZE5FOchYAdcRahPuce6oYzwnmxX9qM8GND2EbYEZGpqvqSG8zvP6qaWpEOzFJ7CRs11lwywo6qvlTNoXjcBYIiMlVEckTkVhGZLSLzROTHbpTQTyoWCorIIBH5m5vvBRHpWV/3YZwYZmQM31DVPwZsvwSsB5ap6hVAMRCvqlcDy4Ez3azPA8+q6u+B2RyfBW1EKNYnY0QaW9zP/QHb+zi+KG8ocJaITMBxQHWofuUZtcWMjNHQWAm8rqqr3Lhc5/ktyAiOGRmjXnBdTV4LtBWRq1T1Rdf1QlsRuQQn/GkPYIqIvIVTY7lCRHYCE3AcKM0Drgb+T0S2At2Bl/24H8M7NrpkGEZYsY5fwzDCihkZwzDCihkZwzDCihkZwzDCihkZwzDCihkZwzDCihkZwzDCyv8Dr0ezvQSVMjcAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4*1, 4*1), sharex=True, sharey=True)\n",
    "ax_idx = [(i, j) for i in range(1) for j in range(1)]\n",
    "\n",
    "for i in tqdm(range(1)):\n",
    "    \n",
    "    ci = ax_idx[i]\n",
    "    \n",
    "    cstate, cconditionals, ctarget = gen_data((i+4)*(i+1), 1, 50)\n",
    "    \n",
    "    runner = jax.jit(jax.vmap(\n",
    "        partial(\n",
    "            lagrangian_eom,\n",
    "            learned_dynamics(params, nn_forward_fn)), (0, 0), 0))\n",
    "\n",
    "    @jax.jit\n",
    "    def odefunc_learned(y, t):\n",
    "        return jnp.concatenate((runner(y[None, :2], y[None, [2]])[0], jnp.zeros(1)))\n",
    "\n",
    "    yt_learned = odeint(\n",
    "        odefunc_learned,\n",
    "        jnp.concatenate([cstate[0], cconditionals[0]]),\n",
    "        np.linspace(0, 1, 50),\n",
    "        mxsteps=100)\n",
    "    \n",
    "    cax = ax#[ci[0], ci[1]]\n",
    "    cax.plot(cstate[:, 1], label='Truth')\n",
    "    cax.plot(yt_learned[:, 1], label='Learned')\n",
    "    cax.legend()\n",
    "    if ci[1] == 0:\n",
    "        cax.set_ylabel('Velocity of particle/Speed of light')\n",
    "    if ci[0] == 0:\n",
    "        cax.set_xlabel('Time')\n",
    "        \n",
    "    cax.set_ylim(-1, 1)\n",
    "    \n",
    "plt.title(\"Lagrangian NN - Special Relativity\")\n",
    "plt.tight_layout()\n",
    "plt.savefig('sr_lnn.png', dpi=150)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "main2",
   "language": "python",
   "name": "main2"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}