creating_graphs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "06c9bbc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "SAVE_DIR = '/Users/aronjansen/Documents/grainsData/graph_test_data/graphs.hdf5'\n",
    "file = h5py.File(SAVE_DIR, 'r')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f5d478",
   "metadata": {},
   "source": [
    "# Structure of data:\n",
    "\n",
    "- Each timestep has its own entry, i.e. is its own group.\n",
    "- There is a separate dataset 'contact_params' for the contact parameters\n",
    "- Also a 'num_steps' just indicating the number of steps\n",
    "- For each timestep there are the datasets:\n",
    "    - 'sources'\n",
    "    - 'destinations'\n",
    "    - 'node_features'\n",
    "    - 'edge_features'\n",
    "    - 'input_features': experimental setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "a4c941ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<KeysViewHDF5 ['0', '1', '10', '100', '101', '102', '103', '104', '105', '106', '107', '108', '109', '11', '110', '111', '112', '113', '114', '115', '116', '117', '118', '119', '12', '120', '121', '122', '123', '124', '125', '126', '127', '128', '129', '13', '130', '131', '132', '133', '134', '135', '136', '137', '138', '139', '14', '140', '141', '142', '143', '144', '145', '146', '147', '148', '149', '15', '150', '151', '152', '153', '154', '155', '156', '157', '158', '159', '16', '160', '161', '162', '163', '164', '165', '166', '167', '168', '169', '17', '170', '171', '172', '173', '174', '175', '176', '177', '178', '179', '18', '180', '181', '182', '183', '184', '185', '186', '187', '188', '189', '19', '190', '191', '192', '193', '194', '195', '196', '197', '198', '199', '2', '20', '200', '21', '22', '23', '24', '25', '26', '27', '28', '29', '3', '30', '31', '32', '33', '34', '35', '36', '37', '38', '39', '4', '40', '41', '42', '43', '44', '45', '46', '47', '48', '49', '5', '50', '51', '52', '53', '54', '55', '56', '57', '58', '59', '6', '60', '61', '62', '63', '64', '65', '66', '67', '68', '69', '7', '70', '71', '72', '73', '74', '75', '76', '77', '78', '79', '8', '80', '81', '82', '83', '84', '85', '86', '87', '88', '89', '9', '90', '91', '92', '93', '94', '95', '96', '97', '98', '99', 'contact_params', 'num_steps']>\n",
      "contact_params:  [[6.96000000e+10 2.00000000e-01 2.91456794e-01]]\n",
      "num_steps:  201\n"
     ]
    }
   ],
   "source": [
    "print(file.keys())\n",
    "print(\"contact_params: \", file['contact_params'][:])\n",
    "print(\"num_steps: \", np.array(file['num_steps']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2662637",
   "metadata": {},
   "source": [
    "# Turn into pytorch geometric graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f80d7bb7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/Caskroom/miniforge/base/envs/grain-graphs/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch_geometric as tg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "367e04fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_edge_index(data):\n",
    "    s = torch.tensor(data['sources'][:], dtype=torch.int)\n",
    "    d = torch.tensor(data['destinations'][:], dtype=torch.int)\n",
    "    edge_index = torch.stack([s, d])\n",
    "    \n",
    "    return edge_index\n",
    "\n",
    "def make_tg_graph(data):\n",
    "    edge_index = make_edge_index(data)\n",
    "    node_features = torch.tensor(data['node_features'][:])\n",
    "    edge_features = torch.tensor(data['edge_features'][:])\n",
    "    \n",
    "    # the node features include position, which we can actually provide separately\n",
    "    position_idx = 15\n",
    "    positions = node_features[:, position_idx:position_idx + 3]\n",
    "    node_features = torch.concat(\n",
    "        [node_features[:, :position_idx], node_features[:, position_idx + 3:]], axis=1)\n",
    "    \n",
    "    graph = tg.data.Data(\n",
    "        x=node_features,\n",
    "        edge_index=edge_index,\n",
    "        edge_attr=edge_features,\n",
    "        pos=positions,\n",
    "    )\n",
    "    return graph\n",
    "\n",
    "\n",
    "def make_graphs(hdf5_file):\n",
    "    contact_params = torch.tensor(hdf5_file['contact_params'][:], dtype=torch.float)\n",
    "    num_steps = int(np.array(hdf5_file['num_steps']))\n",
    "    \n",
    "    graphs = [make_tg_graph(hdf5_file[f'{step}']) for step in range(num_steps)]\n",
    "    for step in range(num_steps - 1):\n",
    "        graphs[step].y = graphs[step + 1].pos\n",
    "        \n",
    "    return graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "597169a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data(x=[10000, 18], edge_index=[2, 27353], edge_attr=[27353, 29], pos=[10000, 3])\n",
      "tensor([[ 6.9998e-01,  1.3319e-07,  1.4410e-08,  ...,  1.7645e-01,\n",
      "          2.1026e-02,  9.4487e-01],\n",
      "        [ 6.9994e-01,  6.9876e-08,  3.8542e-07,  ..., -1.0430e-01,\n",
      "         -3.1843e-02,  9.8436e-01],\n",
      "        [ 6.9990e-01,  4.9942e-08,  1.8485e-07,  ..., -1.0288e-01,\n",
      "          6.1068e-02,  9.9264e-01],\n",
      "        ...,\n",
      "        [ 3.0014e-01,  2.8647e-07,  9.7508e-08,  ..., -3.4116e-01,\n",
      "         -3.6463e-01,  7.9001e-01],\n",
      "        [ 3.0010e-01,  3.0145e-04, -4.2226e-04,  ..., -5.7848e-01,\n",
      "          5.5913e-02, -7.1321e-01],\n",
      "        [ 3.0006e-01,  3.6419e-07,  4.4413e-07,  ..., -7.2468e-01,\n",
      "          1.5819e-01,  6.7068e-01]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "g = make_tg_graph(file['1'])\n",
    "print(g)\n",
    "print(g.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "60b81fd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "tst = make_edge_index(file['1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "f2784570",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "print(tg.utils.is_undirected(tst))\n",
    "bla = torch.tensor([[3, 7], [7, 3]])\n",
    "print(tg.utils.is_undirected(bla))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "c440bd05",
   "metadata": {},
   "outputs": [],
   "source": [
    "graphs = make_graphs(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "5066f36b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Data(x=[10000, 18], edge_index=[2, 27355], edge_attr=[27355, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 27353], edge_attr=[27353, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 23165], edge_attr=[23165, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 22230], edge_attr=[22230, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21769], edge_attr=[21769, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21535], edge_attr=[21535, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21360], edge_attr=[21360, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21291], edge_attr=[21291, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21146], edge_attr=[21146, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21095], edge_attr=[21095, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21089], edge_attr=[21089, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 21035], edge_attr=[21035, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20908], edge_attr=[20908, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20966], edge_attr=[20966, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20940], edge_attr=[20940, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20894], edge_attr=[20894, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20714], edge_attr=[20714, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20697], edge_attr=[20697, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20770], edge_attr=[20770, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20851], edge_attr=[20851, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20784], edge_attr=[20784, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20876], edge_attr=[20876, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20764], edge_attr=[20764, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20844], edge_attr=[20844, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20830], edge_attr=[20830, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20826], edge_attr=[20826, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20881], edge_attr=[20881, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20809], edge_attr=[20809, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20899], edge_attr=[20899, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20867], edge_attr=[20867, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20828], edge_attr=[20828, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20847], edge_attr=[20847, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20675], edge_attr=[20675, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20687], edge_attr=[20687, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20636], edge_attr=[20636, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20683], edge_attr=[20683, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20757], edge_attr=[20757, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20797], edge_attr=[20797, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20808], edge_attr=[20808, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20870], edge_attr=[20870, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20807], edge_attr=[20807, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20858], edge_attr=[20858, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20842], edge_attr=[20842, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20963], edge_attr=[20963, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20939], edge_attr=[20939, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20899], edge_attr=[20899, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20820], edge_attr=[20820, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20867], edge_attr=[20867, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20848], edge_attr=[20848, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20849], edge_attr=[20849, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20826], edge_attr=[20826, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20810], edge_attr=[20810, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20819], edge_attr=[20819, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20820], edge_attr=[20820, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20773], edge_attr=[20773, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20804], edge_attr=[20804, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20746], edge_attr=[20746, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20728], edge_attr=[20728, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20642], edge_attr=[20642, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20720], edge_attr=[20720, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20793], edge_attr=[20793, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20704], edge_attr=[20704, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20761], edge_attr=[20761, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20798], edge_attr=[20798, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20738], edge_attr=[20738, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20698], edge_attr=[20698, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20651], edge_attr=[20651, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20626], edge_attr=[20626, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20719], edge_attr=[20719, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20622], edge_attr=[20622, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20688], edge_attr=[20688, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20684], edge_attr=[20684, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20722], edge_attr=[20722, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20684], edge_attr=[20684, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20618], edge_attr=[20618, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20614], edge_attr=[20614, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20565], edge_attr=[20565, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20632], edge_attr=[20632, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20578], edge_attr=[20578, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20577], edge_attr=[20577, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20600], edge_attr=[20600, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20704], edge_attr=[20704, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20727], edge_attr=[20727, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20682], edge_attr=[20682, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20735], edge_attr=[20735, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20743], edge_attr=[20743, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20764], edge_attr=[20764, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20745], edge_attr=[20745, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20759], edge_attr=[20759, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20834], edge_attr=[20834, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20698], edge_attr=[20698, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20650], edge_attr=[20650, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20615], edge_attr=[20615, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20646], edge_attr=[20646, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20694], edge_attr=[20694, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20715], edge_attr=[20715, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20705], edge_attr=[20705, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20612], edge_attr=[20612, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20606], edge_attr=[20606, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20570], edge_attr=[20570, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20670], edge_attr=[20670, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20685], edge_attr=[20685, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20662], edge_attr=[20662, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20648], edge_attr=[20648, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20629], edge_attr=[20629, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20587], edge_attr=[20587, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20499], edge_attr=[20499, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20710], edge_attr=[20710, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20681], edge_attr=[20681, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20624], edge_attr=[20624, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20649], edge_attr=[20649, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20614], edge_attr=[20614, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20666], edge_attr=[20666, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20649], edge_attr=[20649, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20692], edge_attr=[20692, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20675], edge_attr=[20675, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20693], edge_attr=[20693, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20524], edge_attr=[20524, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20558], edge_attr=[20558, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20662], edge_attr=[20662, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20722], edge_attr=[20722, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20615], edge_attr=[20615, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20621], edge_attr=[20621, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20681], edge_attr=[20681, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20693], edge_attr=[20693, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20551], edge_attr=[20551, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20644], edge_attr=[20644, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20586], edge_attr=[20586, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20504], edge_attr=[20504, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20574], edge_attr=[20574, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20596], edge_attr=[20596, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20757], edge_attr=[20757, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20835], edge_attr=[20835, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20765], edge_attr=[20765, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20702], edge_attr=[20702, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20635], edge_attr=[20635, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20667], edge_attr=[20667, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20719], edge_attr=[20719, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20721], edge_attr=[20721, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20575], edge_attr=[20575, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20589], edge_attr=[20589, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20480], edge_attr=[20480, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20440], edge_attr=[20440, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20478], edge_attr=[20478, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20406], edge_attr=[20406, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20415], edge_attr=[20415, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20336], edge_attr=[20336, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20400], edge_attr=[20400, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20396], edge_attr=[20396, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20374], edge_attr=[20374, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20350], edge_attr=[20350, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20349], edge_attr=[20349, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20331], edge_attr=[20331, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20491], edge_attr=[20491, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20362], edge_attr=[20362, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20389], edge_attr=[20389, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20564], edge_attr=[20564, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20446], edge_attr=[20446, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20516], edge_attr=[20516, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20398], edge_attr=[20398, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20406], edge_attr=[20406, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20459], edge_attr=[20459, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20403], edge_attr=[20403, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20422], edge_attr=[20422, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20511], edge_attr=[20511, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20461], edge_attr=[20461, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20471], edge_attr=[20471, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20433], edge_attr=[20433, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20457], edge_attr=[20457, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20477], edge_attr=[20477, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20435], edge_attr=[20435, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20425], edge_attr=[20425, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20324], edge_attr=[20324, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20541], edge_attr=[20541, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20544], edge_attr=[20544, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20596], edge_attr=[20596, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20442], edge_attr=[20442, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20556], edge_attr=[20556, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20581], edge_attr=[20581, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20595], edge_attr=[20595, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20574], edge_attr=[20574, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20507], edge_attr=[20507, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20457], edge_attr=[20457, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20617], edge_attr=[20617, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20560], edge_attr=[20560, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20576], edge_attr=[20576, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20528], edge_attr=[20528, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20467], edge_attr=[20467, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20492], edge_attr=[20492, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20466], edge_attr=[20466, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20403], edge_attr=[20403, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20523], edge_attr=[20523, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20424], edge_attr=[20424, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20447], edge_attr=[20447, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20439], edge_attr=[20439, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20482], edge_attr=[20482, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20500], edge_attr=[20500, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20574], edge_attr=[20574, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20513], edge_attr=[20513, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20598], edge_attr=[20598, 29], pos=[10000, 3], y=[10000, 3]),\n",
       " Data(x=[10000, 18], edge_index=[2, 20511], edge_attr=[20511, 29], pos=[10000, 3])]"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6c444e76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.return_types.sort(\n",
       "values=tensor([    0,     5,    17,  ..., 11735, 11833, 11882]),\n",
       "indices=tensor([    0, 16086, 24705,  ..., 17032,  9902,  1516]))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.abs(tst2 - tst2[0]).sum(axis=1).sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "377a8348",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = '/Users/aronjansen/Documents/grainsData/graph_test_data/'\n",
    "step = 0\n",
    "fname = f'simState_drained_200_10000_{step}_gnn_data.hdf5'\n",
    "f_input = h5py.File(DATA_DIR + fname, 'r')\n",
    "\n",
    "PRESSURE = '0.1e6'\n",
    "EXPERIMENT_TYPE = 'drained'\n",
    "NUM_PARTICLES = '10000'\n",
    "\n",
    "f_input = f_input[PRESSURE][EXPERIMENT_TYPE][NUM_PARTICLES]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "958264c9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<KeysViewHDF5 ['contact_params', 'inputs', 'outputs_bodies', 'outputs_inters']>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_input.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "926f3150",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[6.96000000e+10, 2.00000000e-01, 2.91456794e-01]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_input['contact_params'][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ecab3e60",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<HDF5 dataset \"contact_params\": shape (1, 3), type \"<f8\">\n",
      "<HDF5 dataset \"inputs\": shape (1, 3), type \"<f8\">\n",
      "<HDF5 dataset \"outputs_bodies\": shape (1, 21, 10000), type \"<f8\">\n",
      "<HDF5 dataset \"outputs_inters\": shape (1, 31, 27355), type \"<f8\">\n"
     ]
    }
   ],
   "source": [
    "for key in f_input.keys():\n",
    "    print(f_input[key])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c50f42f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 31, 27355)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tst = f_input['outputs_inters']\n",
    "tst.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "41bfcf71",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(31, 27355)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tst[:][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d62f199f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(31, 27355)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tst[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f5e72c31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 5.47688730e-01, -1.00065869e+05,  1.00000000e+04])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_input['inputs'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "664d1ab2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "grain-graphs",
   "language": "python",
   "name": "grain-graphs"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}