test_functions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8d700166-c790-4b21-b3b2-3827045e3b85",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-07T00:17:30.437108165Z",
     "start_time": "2024-02-07T00:17:30.389499353Z"
    }
   },
   "outputs": [],
   "source": [
    "from aniso_MLMD.trainer import load_datasets\n",
    "import aniso_MLMD.model.neighbor_ops as neighbor_ops\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7055c055-805f-492a-aace-460869a14075",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-07T00:15:48.220615229Z",
     "start_time": "2024-02-07T00:15:48.217346998Z"
    }
   },
   "outputs": [],
   "source": [
    "data_path =\"/home/marjan/Documents/code-base/ml_datasets/pps_20_new/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dfc9461e-5b9c-4092-a98f-f80af6813f16",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-07T00:15:49.903509790Z",
     "start_time": "2024-02-07T00:15:49.200114251Z"
    }
   },
   "outputs": [],
   "source": [
    "train_df = pd.read_pickle(os.path.join(data_path, 'train.pkl'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1b0ae3d2-afa0-4e93-a1e6-b84ff57411fa",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-07T00:15:51.091564981Z",
     "start_time": "2024-02-07T00:15:50.083418723Z"
    }
   },
   "outputs": [],
   "source": [
    "train_dataloader, valid_dataloader, test_dataloader = load_datasets(data_path, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "bbf74139-571e-45c4-9830-f951f3af5c0c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-07T00:24:28.161678017Z",
     "start_time": "2024-02-07T00:24:27.909319350Z"
    }
   },
   "outputs": [],
   "source": [
    "for i, ((position, q, orientation_R, neighbor_list), target_force, target_torque,\n",
    "                energy) in enumerate(\n",
    "                train_dataloader):\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "box_len = 16.17887"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-07T00:24:30.121250833Z",
     "start_time": "2024-02-07T00:24:30.117450367Z"
    }
   },
   "id": "c8ebb90ea01c7b93",
   "execution_count": 32
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([5, 200, 2])"
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "neighbor_list.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-07T00:24:36.977036142Z",
     "start_time": "2024-02-07T00:24:36.974591357Z"
    }
   },
   "id": "dbf1868685558d19",
   "execution_count": 33
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "84360657-36ed-48ac-90f8-6f0cab714655",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-07T00:24:50.624659989Z",
     "start_time": "2024-02-07T00:24:50.582471163Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "here\n",
      "torch.Size([5, 20, 10, 3, 3, 3])\n"
     ]
    }
   ],
   "source": [
    "batch_size = position.shape[0]\n",
    "N_particles = position.shape[1]\n",
    "\n",
    "# change tuple based neighbor list to (B, N, neighbor_idx)\n",
    "neighbor_list = neighbor_list.reshape(batch_size, N_particles, -1,\n",
    "                                      neighbor_list.shape[-1])[:, :, :,\n",
    "                1]  # (B, N, N_neighbors)\n",
    "N_neighbors = neighbor_list.shape[-1]\n",
    "dr = neighbor_ops.neighbors_distance_vector(position,\n",
    "                                            neighbor_list)  # (B, N, N_neighbors, 3)\n",
    "dr = neighbor_ops.adjust_periodic_boundary(dr, box_len)\n",
    "\n",
    "R = torch.norm(dr, dim=-1, keepdim=True)  # (B, N, N_neighbors, 1)\n",
    "\n",
    "inv_R = 1. / R  # (B, N, N_neighbors, 1)\n",
    "\n",
    "################ orientation related features ################\n",
    "\n",
    "# repeat the neighbors idx to match the shape of orientation_R. This\n",
    "# is necessary to gather each particle's neighbors' orientation\n",
    "NN_expanded = neighbor_list[:, :, :, None, None].expand(\n",
    "    (-1, -1, -1, 3, 3))  # (B, N, N_neighbors, 3, 3)\n",
    "# repeart the orientation_R to match the shape of neighbor_list\n",
    "orientation_R_expanded = orientation_R[:, :, None, :, :].expand(\n",
    "    (-1, -1, N_neighbors, -1, -1))  # (B, N, N_neighbors, 3, 3)\n",
    "# get all neighbors' orientation\n",
    "neighbors_orient_R = torch.gather(orientation_R_expanded, dim=1,\n",
    "                                  index=NN_expanded)  # (B, N, N_neighbors, 3, 3)\n",
    "\n",
    "# dot product: (B, N, N_neighbors, 3, 3)\n",
    "orient_dot_prod = neighbor_ops.orientation_dot_product(\n",
    "    orientation_R_expanded,\n",
    "    neighbors_orient_R)\n",
    "\n",
    "# element product: (B, N, N_neighbors, 3, 3, 3)\n",
    "print('here')\n",
    "orient_element_prod = neighbor_ops.orientation_element_product(\n",
    "    orientation_R_expanded,\n",
    "    neighbors_orient_R)\n",
    "print(orient_element_prod.shape)\n",
    "# element product norm: (B, N, N_neighbors, 3, 3)\n",
    "element_prod_norm = torch.norm(orient_element_prod,\n",
    "                               dim=-1)\n",
    "\n",
    "# principal cross product: (B, N, N_neighbors, 3, 3)\n",
    "orient_cross_prod = neighbor_ops.orientation_principal_cross_product(\n",
    "    orientation_R_expanded,\n",
    "    neighbors_orient_R)\n",
    "# cross product norm: (B, N, N_neighbors, 3)\n",
    "cross_prod_norm = torch.norm(orient_cross_prod,\n",
    "                             dim=-1)\n",
    "\n",
    "# relative orientation: (B, N, N_neighbors, 3, 3)\n",
    "rel_orient = neighbor_ops.relative_orientation(orientation_R_expanded,\n",
    "                                               neighbors_orient_R)\n",
    "\n",
    "################ RBF features ################\n",
    "\n",
    "# RBF for particles:(B, N, N_neighbors, 3)\n",
    "rbf_particle = neighbor_ops.RBF_dr_orientation(dr,\n",
    "                                               orientation_R_expanded)\n",
    "# RBF for neighbors: (B, N, N_neighbors, 3)\n",
    "rbf_neighbors = neighbor_ops.RBF_dr_orientation(dr, neighbors_orient_R)\n",
    "\n",
    "# euler angle (B, N, N_neighbors, 3)\n",
    "angle = neighbor_ops.rot_matrix_to_euler_angle(rel_orient)"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([5, 20, 10, 1])"
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-07T00:18:25.964576743Z",
     "start_time": "2024-02-07T00:18:25.959653706Z"
    }
   },
   "id": "311d1903f0a0a60c",
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "features = torch.cat((R,\n",
    "                      dr / R,\n",
    "                      inv_R,\n",
    "                      orient_dot_prod.flatten(start_dim=-2),\n",
    "                      orient_element_prod.flatten(start_dim=-3),\n",
    "                      element_prod_norm.flatten(start_dim=-2),\n",
    "                      orient_cross_prod.flatten(start_dim=-2),\n",
    "                      cross_prod_norm,\n",
    "                      rel_orient.flatten(start_dim=-2),\n",
    "                      # rbf_particle,\n",
    "                      # rbf_neighbors,\n",
    "                      # angle\n",
    "                      ),\n",
    "                     dim=-1)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-07T00:25:16.481539685Z",
     "start_time": "2024-02-07T00:25:16.475590502Z"
    }
   },
   "id": "f81530e3f398dc6d",
   "execution_count": 43
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([5, 20, 10, 71])"
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-07T00:25:17.655966683Z",
     "start_time": "2024-02-07T00:25:17.652316762Z"
    }
   },
   "id": "65d99407c5aae75",
   "execution_count": 44
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([5, 20, 1])"
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "neighbor_list.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-07T00:24:01.903657166Z",
     "start_time": "2024-02-07T00:24:01.897421009Z"
    }
   },
   "id": "6973f7178bb78e4",
   "execution_count": 30
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[28], line 1\u001B[0m\n\u001B[0;32m----> 1\u001B[0m orient_element_prod \u001B[38;5;241m=\u001B[39m \u001B[43mneighbor_ops\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43morientation_element_product\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m      2\u001B[0m \u001B[43m    \u001B[49m\u001B[43morientation_R_expanded\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m      3\u001B[0m \u001B[43m    \u001B[49m\u001B[43mneighbors_orient_R\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/Documents/code-base/Aniso_MLMD/aniso_MLMD/model/neighbor_ops.py:78\u001B[0m, in \u001B[0;36morientation_element_product\u001B[0;34m(particle_orientations, neighbors_orientations)\u001B[0m\n\u001B[1;32m     61\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21morientation_element_product\u001B[39m(particle_orientations, neighbors_orientations):\n\u001B[1;32m     62\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[1;32m     63\u001B[0m \u001B[38;5;124;03m    Calculate the element wise product of the particle orientation columns with\u001B[39;00m\n\u001B[1;32m     64\u001B[0m \u001B[38;5;124;03m    the neighbors orientations columns.\u001B[39;00m\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m     76\u001B[0m \n\u001B[1;32m     77\u001B[0m \u001B[38;5;124;03m    \"\"\"\u001B[39;00m\n\u001B[0;32m---> 78\u001B[0m     element_prod \u001B[38;5;241m=\u001B[39m \u001B[43mtorch\u001B[49m\u001B[38;5;241m.\u001B[39meinsum(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mijkhl, ijkhm -> ijklmh\u001B[39m\u001B[38;5;124m'\u001B[39m, particle_orientations,\n\u001B[1;32m     79\u001B[0m                                 neighbors_orientations)\n\u001B[1;32m     80\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m element_prod\n",
      "File \u001B[0;32m/app/extra/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_cython.pyx:1187\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython.SafeCallWrapper.__call__\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m/app/extra/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_cython.pyx:627\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython.PyDBFrame.trace_dispatch\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m/app/extra/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_cython.pyx:937\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython.PyDBFrame.trace_dispatch\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m/app/extra/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_cython.pyx:928\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython.PyDBFrame.trace_dispatch\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m/app/extra/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_cython.pyx:585\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython.PyDBFrame.do_wait_suspend\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m/app/extra/plugins/python/helpers/pydev/pydevd.py:1184\u001B[0m, in \u001B[0;36mPyDB.do_wait_suspend\u001B[0;34m(self, thread, frame, event, arg, send_suspend_message, is_unhandled_exception)\u001B[0m\n\u001B[1;32m   1181\u001B[0m         from_this_thread\u001B[38;5;241m.\u001B[39mappend(frame_id)\n\u001B[1;32m   1183\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_threads_suspended_single_notification\u001B[38;5;241m.\u001B[39mnotify_thread_suspended(thread_id, stop_reason):\n\u001B[0;32m-> 1184\u001B[0m     \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_do_wait_suspend\u001B[49m\u001B[43m(\u001B[49m\u001B[43mthread\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mframe\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mevent\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43marg\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43msuspend_type\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mfrom_this_thread\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/app/extra/plugins/python/helpers/pydev/pydevd.py:1199\u001B[0m, in \u001B[0;36mPyDB._do_wait_suspend\u001B[0;34m(self, thread, frame, event, arg, suspend_type, from_this_thread)\u001B[0m\n\u001B[1;32m   1196\u001B[0m             \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_call_mpl_hook()\n\u001B[1;32m   1198\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mprocess_internal_commands()\n\u001B[0;32m-> 1199\u001B[0m         \u001B[43mtime\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43msleep\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m0.01\u001B[39;49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1201\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mcancel_async_evaluation(get_current_thread_id(thread), \u001B[38;5;28mstr\u001B[39m(\u001B[38;5;28mid\u001B[39m(frame)))\n\u001B[1;32m   1203\u001B[0m \u001B[38;5;66;03m# process any stepping instructions\u001B[39;00m\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "orient_element_prod = neighbor_ops.orientation_element_product(\n",
    "    orientation_R_expanded,\n",
    "    neighbors_orient_R)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-07T00:23:41.746121129Z",
     "start_time": "2024-02-07T00:23:25.543564068Z"
    }
   },
   "id": "dcb344d7998c5af8",
   "execution_count": 28
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}