{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:34:30.787475200Z",
     "start_time": "2024-06-03T03:34:28.159120700Z"
    }
   },
   "outputs": [],
   "source": [
    "from algorithm.neat.gene.conn.bspline import BSplineConn\n",
    "from algorithm.neat.gene.conn.cache_bspline import CacheBSplineConn\n",
    "import jax, jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "outputs": [],
   "source": [
    "normal_gene = BSplineConn(grid_cnt=6, spline_order=3, grid_init_range=[-1, 1])\n",
    "cache_gene = CacheBSplineConn(grid_cnt=6, spline_order=3, grid_range=[-1, 1], cache_num=1000000)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:41:35.807924900Z",
     "start_time": "2024-06-03T03:41:35.793415500Z"
    }
   },
   "id": "a09d6ccf956606b3"
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [
    {
     "data": {
      "text/plain": "State ({'kan_initial_grids': Array([-1.        , -0.6       , -0.20000002,  0.20000005,  0.6       ,\n        1.        ], dtype=float32), 'bspline_cache': Array([[0.16666666, 0.6666666 , 0.16666667, ..., 0.        , 0.        ,\n        0.        ],\n       [0.16666412, 0.6666665 , 0.1666692 , ..., 0.        , 0.        ,\n        0.        ],\n       [0.16666159, 0.66666657, 0.16667175, ..., 0.        , 0.        ,\n        0.        ],\n       ...,\n       [0.        , 0.        , 0.        , ..., 0.1666717 , 0.6666666 ,\n        0.16666172],\n       [0.        , 0.        , 0.        , ..., 0.16666915, 0.6666666 ,\n        0.16666426],\n       [0.        , 0.        , 0.        , ..., 0.16666675, 0.6666665 ,\n        0.16666663]], dtype=float32)})"
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state = normal_gene.setup()\n",
    "state = cache_gene.setup(state)\n",
    "state"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:41:37.462678200Z",
     "start_time": "2024-06-03T03:41:36.459771900Z"
    }
   },
   "id": "57fbeab2e4d4c511"
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "outputs": [
    {
     "data": {
      "text/plain": "Array([ 0.08086783, -0.38624713, -0.37565565,  1.6689739 , -1.2758198 ,\n        2.1192005 , -0.85821223,  1.1305932 ], dtype=float32)"
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key = jax.random.PRNGKey(0)\n",
    "normal_attrs = normal_gene.new_random_attrs(state, key)\n",
    "normal_attrs\n",
    "weights = normal_attrs[normal_gene.grid_cnt:]\n",
    "weights"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:41:38.867812Z",
     "start_time": "2024-06-03T03:41:38.789154200Z"
    }
   },
   "id": "9d9cacf5af5f38c3"
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "outputs": [
    {
     "data": {
      "text/plain": "(Array(-0.0304966, dtype=float32), Array(-0.03049916, dtype=float32))"
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = 0.99999\n",
    "normal_res = normal_gene.forward(state, normal_attrs, t)\n",
    "cache_res = cache_gene.forward(state, weights, t)\n",
    "normal_res, cache_res"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:43:23.083384100Z",
     "start_time": "2024-06-03T03:43:23.002384100Z"
    }
   },
   "id": "9177b012b7ab25cd"
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "outputs": [],
   "source": [
    "batch = 100000\n",
    "t = jnp.linspace(-1, 1, batch)\n",
    "batch_normal_forward = jax.jit(jax.vmap(normal_gene.forward, in_axes=(None, None, 0)))\n",
    "batch_cache_forward = jax.jit(jax.vmap(cache_gene.forward, in_axes=(None, None, 0)))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:50:05.806318700Z",
     "start_time": "2024-06-03T03:50:05.785312900Z"
    }
   },
   "id": "878f5eda35df17f7"
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.34 ms ± 99.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit batch_normal_forward(state, normal_attrs, t).block_until_ready()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:52:23.070679300Z",
     "start_time": "2024-06-03T03:52:21.768451Z"
    }
   },
   "id": "2f649d81e1945757"
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "134 µs ± 28.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit batch_cache_forward(state, weights, t).block_until_ready()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:52:43.066619300Z",
     "start_time": "2024-06-03T03:52:42.599706100Z"
    }
   },
   "id": "cac106d6792ea53b"
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "outputs": [],
   "source": [
    "normal_res = batch_normal_forward(state, normal_attrs, t).block_until_ready()\n",
    "cache_res = batch_cache_forward(state, weights, t).block_until_ready()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:53:15.153532800Z",
     "start_time": "2024-06-03T03:53:15.145531800Z"
    }
   },
   "id": "27bc7e5abf9b65d4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "1b25ee5a2fe153c4"
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "outputs": [
    {
     "data": {
      "text/plain": "[<matplotlib.lines.Line2D at 0x7f5b32908a30>]"
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 640x480 with 1 Axes>",
      "image/png": "[encoded data, 37028 characters]"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "plt.plot(t, normal_res, color='blue')\n",
    "plt.plot(t, cache_res, color='red')"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-03T03:53:58.803332700Z",
     "start_time": "2024-06-03T03:53:57.867105200Z"
    }
   },
   "id": "97ace2f9183fef16"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "4edc26eea7760479"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
