03 NeuralODE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchdiffeq import odeint\n",
    "from torchdiffeq import odeint_adjoint as odeint_adj\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "from random import sample\n",
    "import time\n",
    "import os\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ODEFunc(nn.Module):\n",
    "    def __init__(self, twice_dim):\n",
    "        super(ODEFunc, self).__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(twice_dim, 50),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(50,50),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(50, twice_dim),\n",
    "            \n",
    "        )\n",
    "        for m in self.net.modules():\n",
    "            if isinstance(m,nn.Linear):\n",
    "                nn.init.normal_(m.weight,mean=0,std=0.01)\n",
    "                nn.init.constant_(m.bias,val=0)\n",
    "        \n",
    "    def forward(self, t, y):\n",
    "        \n",
    "        #print(y.shape)\n",
    "        out = self.net(y)\n",
    "        #out = torch.ones(y.shape)\n",
    "        #out[:,0] = y[:,2]\n",
    "        #out[:,1] = y[:,3]\n",
    "        #out[:,2] = y[:,0]/(0.6666**2)\n",
    "        #out[:,3] = y[:,1]/(0.3333**2)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ODEFunc_exact_den(nn.Module):\n",
    "    def __init__(self,ODE_model):\n",
    "        super(ODEFunc_exact_den, self).__init__()\n",
    "        self.net = ODE_model\n",
    "        \n",
    "    def forward(self, t, y):\n",
    "#         start = time.time()\n",
    "        #print(\"here\")\n",
    "        zero = torch.tensor([0.0]).view(-1,1)\n",
    "        inp = torch.cat((zero,y[:,:-1]),dim=1)\n",
    "        f = self.net(t,y[:,:-1])\n",
    "        #print(time.time()-start)\n",
    "        Jacobian = torch.autograd.functional.jacobian(self.net, (t,y[:,:-1]))\n",
    "        #print(time.time()-start)\n",
    "        \n",
    "        Tr = torch.trace(Jacobian[1].view(twice_dim,twice_dim)).view(-1,1)\n",
    "        \n",
    "        out = torch.cat((f,-Tr),dim=1)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Evaluate_H(f, q,p,M_inv):\n",
    "    V_o,grad = f(q,p)\n",
    "    H_o = V_o + 0.5*( np.matmul(p.reshape(1,-1),np.matmul(M_inv,p.reshape(-1,1)))   )\n",
    "    H_o = H_o.reshape(-1)\n",
    "    return H_o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Neural_dynamics(f,q_temp,p_temp,M,M_inv,steps,eps,store=True):\n",
    "    dim = q_temp.shape[0]\n",
    "    \n",
    "    batch_t = torch.linspace(0.0,eps*(steps-1),steps)\n",
    "    #print(batch_t.shape)\n",
    "    \n",
    "    q_temp = torch.tensor(q_temp)\n",
    "    p_temp = torch.tensor(p_temp)\n",
    "    Zero = torch.tensor([0.0])\n",
    "    init = torch.cat((q_temp,p_temp,Zero),dim=0).float()\n",
    "    init = init.view(1,-1)\n",
    "#     start = time.time()\n",
    "    pred_y = odeint(func_mod, init, batch_t,method = 'rk4',options = dict(step_size=eps))\n",
    "    #print(time.time()-start)\n",
    "    #print(pred_y.shape)\n",
    "    \n",
    "    \n",
    "    pred_y = pred_y.detach().numpy()\n",
    "    p_y = pred_y[:,0,:-1]\n",
    "    J = pred_y[:,0,-1]\n",
    "    \n",
    "    energies = []\n",
    "    \n",
    "    for i in range(p_y.shape[0]):\n",
    "        q = p_y[i,:dim]\n",
    "        p = p_y[i,dim:]\n",
    "        curr_e = Evaluate_H(f,q,p,M_inv)\n",
    "    \n",
    "        energies.append(curr_e)\n",
    "    \n",
    "    energies = np.asarray(energies)\n",
    "    #print(energies.shape)\n",
    "    #print(p_y.shape)\n",
    "    return (p_y[-1,:dim],p_y[-1,dim:]),J[-1],p_y,energies\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Neural_HMC(f,q0,num_samples,eps,steps,store=False):\n",
    "    \n",
    "    samples = []\n",
    "    accept_rate = 0\n",
    "    q = q0\n",
    "    \n",
    "    if store == True:\n",
    "        stored_vals = np.zeros((num_samples,steps,twice_dim))\n",
    "        energies = np.zeros((num_samples,steps))\n",
    "        \n",
    "    ind = 0\n",
    "    M = np.identity(twice_dim//2)\n",
    "    M_inv = np.linalg.inv(M)\n",
    "    while len(samples) < num_samples:\n",
    "        \n",
    "        ####### Need to fix\n",
    "        mean = np.zeros((twice_dim//2))\n",
    "        cov = M\n",
    "        p = np.random.multivariate_normal(mean, cov)\n",
    "        \n",
    "        q_temp = copy.deepcopy(q)\n",
    "        p_temp = copy.deepcopy(p)\n",
    "        \n",
    "        (q_f,p_f),J,path,energy = Neural_dynamics(f,q_temp,p_temp,M,M_inv,steps,eps,store=True)\n",
    "        stored_vals[ind,:,:] = np.asarray(path)\n",
    "        energies[ind,:] = np.asarray(energy).reshape(-1)\n",
    "        \n",
    "        p_f *= -1\n",
    "        q = q_f\n",
    "        \n",
    "        samples.append(q)\n",
    "        \n",
    "        if len(samples)%500 == 0 and len(samples)< 2000:\n",
    "            if len(samples) == 500:\n",
    "                recent_samples = np.asarray(samples)\n",
    "            else:\n",
    "                recent_samples = np.asarray(samples[-500:])\n",
    "            if (np.abs(np.linalg.det(np.cov(recent_samples.T))) > 0.001):\n",
    "                M_inv = np.cov(recent_samples.T)\n",
    "                M = np.linalg.inv(M_inv)\n",
    "            else:\n",
    "                M = np.identity(twice_dim//2)\n",
    "                M_inv = np.linalg.inv(M)\n",
    "#             print(M_inv)\n",
    "        ind+=1\n",
    "             \n",
    "    acceptance = accept_rate/num_samples\n",
    "        \n",
    "    return samples,stored_vals,energies, acceptance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GaussianXD(q, p):\n",
    "    dim = q.shape[0]\n",
    "    sigma = np.ones(dim)\n",
    "    V = 0.5*np.sum(q**2 / (sigma**2))\n",
    "    grad = q/(sigma**2)\n",
    "    return V, grad\n",
    "\n",
    "def Shell2D(q, p):\n",
    "    r0 = np.sqrt(2)\n",
    "    sigma = 0.5\n",
    "    r = np.sqrt(np.dot(q, q))\n",
    "    V = abs(r-r0)/sigma\n",
    "    \n",
    "    if (r-r0) == 0 or r == 0:\n",
    "        grad = np.array([0, 0])\n",
    "    else:\n",
    "        grad = (q*(r-r0)/(sigma*r*abs(r-r0)))\n",
    "        \n",
    "    return V, grad\n",
    "\n",
    "def Wofe_Quapp(q, p):\n",
    "    x = q[0]\n",
    "    y = q[1]\n",
    "    V = x**4 + y**4 - 2*x**2 - 4*y**2 + x*y + 0.3*x + 0.1*y\n",
    "    grad = np.array([4*x**3 - 4*x + y + 0.3, 4*y**3 - 8*y + x + 0.1])\n",
    "    return V, grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000 1\n",
      "10000 2\n",
      "10000 3\n",
      "10000 4\n",
      "10000 5\n"
     ]
    }
   ],
   "source": [
    "# Wofe-Quapp\n",
    "num_samples = [10000] #[100, 200, 500, 1000, 2000]\n",
    "twice_dim = 4\n",
    "init = np.random.randn(twice_dim//2)*2\n",
    "traj_length = 10\n",
    "traj_step_size = 0.1\n",
    "num_of_runs = 5\n",
    "potential = 'wofe_quapp'\n",
    "potential_function = Wofe_Quapp\n",
    "func = torch.load(f'{potential}/model')\n",
    "func_mod = ODEFunc_exact_den(func)\n",
    "\n",
    "for q in num_samples:\n",
    "    for i in range(1, num_of_runs+1):\n",
    "        start = time.perf_counter()\n",
    "        samps,trajs,energies,acceptance = Neural_HMC(potential_function, init, q, traj_step_size, traj_length, store=True)\n",
    "\n",
    "        if not os.path.exists(f'{potential}/ode_{q}'):\n",
    "            os.makedirs(f'{potential}/ode_{q}')\n",
    "\n",
    "        with open(f'{potential}/ode_{q}/{i}_hmc_samps.npy', 'wb') as f:\n",
    "            np.save(f, samps)\n",
    "\n",
    "        with open(f'{potential}/ode_{q}/{i}_info.npy', 'wb') as f:\n",
    "            np.save(f, np.array([time.perf_counter() - start, acceptance]))\n",
    "        print(q, i)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000 1\n",
      "10000 2\n",
      "10000 3\n",
      "10000 4\n",
      "10000 5\n"
     ]
    }
   ],
   "source": [
    "# 2D Shell\n",
    "num_samples = [10000]#[100, 200, 500, 1000, 2000]\n",
    "twice_dim = 4\n",
    "init = np.random.randn(twice_dim//2)*2\n",
    "traj_length = 10\n",
    "traj_step_size = 0.01\n",
    "num_of_runs = 5\n",
    "potential = '2d_shell'\n",
    "potential_function = Shell2D\n",
    "func = torch.load(f'{potential}/model')\n",
    "func_mod = ODEFunc_exact_den(func)\n",
    "\n",
    "for q in num_samples:\n",
    "    for i in range(1, num_of_runs+1):\n",
    "        start = time.perf_counter()\n",
    "        samps,trajs,energies,acceptance = Neural_HMC(potential_function, init, q, traj_step_size, traj_length, store=True)\n",
    "\n",
    "        if not os.path.exists(f'{potential}/ode_{q}'):\n",
    "            os.makedirs(f'{potential}/ode_{q}')\n",
    "\n",
    "        with open(f'{potential}/ode_{q}/{i}_hmc_samps.npy', 'wb') as f:\n",
    "            np.save(f, samps)\n",
    "\n",
    "        with open(f'{potential}/ode_{q}/{i}_info.npy', 'wb') as f:\n",
    "            np.save(f, np.array([time.perf_counter() - start, acceptance]))\n",
    "        print(q, i)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gaussian 10D\n",
    "num_samples = [10000]#[100, 200, 500, 1000, 2000]\n",
    "twice_dim = 20\n",
    "init = np.random.randn(twice_dim//2)*2\n",
    "traj_length = 10\n",
    "traj_step_size = 0.1\n",
    "num_of_runs = 5\n",
    "potential = '10d_gaussian'\n",
    "potential_function = GaussianXD\n",
    "func = torch.load(f'{potential}/model')\n",
    "func_mod = ODEFunc_exact_den(func)\n",
    "\n",
    "for q in num_samples:\n",
    "    for i in range(1, num_of_runs+1):\n",
    "        start = time.perf_counter()\n",
    "        samps,trajs,energies,acceptance = Neural_HMC(potential_function, init, q, traj_step_size, traj_length, store=True)\n",
    "\n",
    "        if not os.path.exists(f'{potential}/ode_{q}'):\n",
    "            os.makedirs(f'{potential}/ode_{q}')\n",
    "\n",
    "        with open(f'{potential}/ode_{q}/{i}_hmc_samps.npy', 'wb') as f:\n",
    "            np.save(f, samps)\n",
    "\n",
    "        with open(f'{potential}/ode_{q}/{i}_info.npy', 'wb') as f:\n",
    "            np.save(f, np.array([time.perf_counter() - start, acceptance]))\n",
    "        print(q, i)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-3320]",
   "language": "python",
   "name": "conda-env-.conda-3320-py"
  },
  "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.9.7"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}