dgl_graphs_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7571cfdc-e098-465b-87aa-0a4d62e42f41",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import h5py\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d451d09-0f8a-4b60-8a65-d49e07b3a4d3",
   "metadata": {},
   "source": [
    "Importing the data file (from Hongyang's file), of a single simulation. The name of the file has the following structure: simState_typeExperiment_SampleID_numParticles_graphs.hdf5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f8281498-6277-410a-9945-2ab14d5877f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#File containing information of sample 200\n",
    "#f_200 = h5py.File('/Users/luisaorozco/Documents/Projects/GrainLearning/data/gnn_data/simState_drained_200_10000_graphs.hdf5', 'r')\n",
    "#File containining all samples (200)\n",
    "f_all = h5py.File('/Users/luisaorozco/Documents/Projects/GrainLearning/data/gnn_data/simState_path_sampling_5000_graphs.hdf5', 'r')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09605588-f713-4f5b-bcb8-23116943404f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Initial testing of the data (exploratory)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "eefad184-f77b-4dfc-907a-5018313239a9",
   "metadata": {},
   "source": [
    "Getting the most basic node and edge features\n",
    "---- Node features:\n",
    "[0] radius\n",
    "[1] mass\n",
    "[2] position x\n",
    "[3] position y\n",
    "[4] position z\n",
    "[5] vel x \n",
    "[6] vel y\n",
    "[7] vel z\n",
    "[8] angVel 0 #angular velocities in the 3 angular directions\n",
    "[9] angVel 1\n",
    "[10] angVel 3\n",
    "[11] angMom[0] # redundant\n",
    "[12] angMom[1]\n",
    "[13]angMom[2]\n",
    "[14]inertia[0] # redundant (spheres)\n",
    "[15]inertia[1]\n",
    "[16]inertia[2]\n",
    "[17]refPos x # position at t=0, with respect to the origin\n",
    "[18]refPos y\n",
    "[19]refPos z\n",
    "[20]refOri[0] # orientation, redundant for spheres\n",
    "[21]refOri[1]\n",
    "[22]refOri[2]\n",
    "[23]refOri[3]\n",
    "\n",
    "---- Edge features:\n",
    "[0] radius particle 1 # radius particle 1, redundant because already in node\n",
    "[1] radius particle 2 # as above\n",
    "[2] ks  # tangential stiffness\n",
    "[3] kn  # normal stiffness\n",
    "[4] penetrationDepth # overlap between spheres\n",
    "[5] shear increment x between particles\n",
    "[6] shear increment y between particles\n",
    "[7] shear increment z between particles\n",
    "[8] contact normal vector x component, can get rid of and just keep normal force\n",
    "[9] contact normal vector y component\n",
    "[10] contact normal vectorz component\n",
    "[11] contactPoint x in the cross section of the overlap\n",
    "[12] contactPoint y \n",
    "[13] contactPoint z \n",
    "[14] elastic component of the shear (tangential) force x, redundant equal to shearForce\n",
    "[15] elastic component of the shear (tangential) force y\n",
    "[16] elastic component of the shear (tangential) force z\n",
    "[17] elastic component of the shear displacement x\n",
    "[18] elastic component of the shear displacement y\n",
    "[19] elastic component of the shear displacement z\n",
    "[20] total shear displacement x\n",
    "[21] total shear displacement y\n",
    "[22] total shear displacement z\n",
    "[23] shear force x\n",
    "[24] shear force y\n",
    "[25] shear force z\n",
    "[26] normal force x\n",
    "[27] normal force y\n",
    "[28] normal force z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "id": "7a8a286c-5b91-4e0a-b7e9-be6424c2dc74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([6.96000000e+10, 2.00000000e-01, 2.91456794e-01])"
      ]
     },
     "execution_count": 245,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#contact parameters of the simulation of sample 200\n",
    "contact_params = f_200['contact_params']\n",
    "contact_params[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 246,
   "id": "eccd52ba-a736-4b59-b10c-ed6a1a4e873a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([6.96000000e+10, 2.00000000e-01, 2.91456794e-01])"
      ]
     },
     "execution_count": 246,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#contact parameters of the simulation of sample 200 retrieved from the big dataset\n",
    "contact_params = f_all['200/contact_params']\n",
    "contact_params[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 247,
   "id": "4ce365ad-d44c-4cf3-ab1b-1cd49a2520cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000.0"
      ]
     },
     "execution_count": 247,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#example for step 0 of sample 200\n",
    "some_n_feats = f_200['0']['node_features'][:,5:8] #particle velocity\n",
    "some_e_feats = f_200['0']['edge_features'][:,-3:] #the last 3 features for all nodes\n",
    "pos_nodes = f_200['0']['node_features'][:,2:5] #position features\n",
    "num_nodes = f_200['0']['input_features'][5]\n",
    "num_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77d22f76-b902-4149-8a0d-e7842f2667a8",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Encoder: Definition of features embeddings\n",
    "To be able to do operations between node and edges features, the dimensions of such tensors musht be the same. Thus, we can create an [*embedding*](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) that will encode *node_features* into a tensor of dimension *emb_dim* and another one for the *edge_features*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 248,
   "id": "96741492-45d5-4b53-9274-497d52a5d8a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class featuresEncoder(torch.nn.Module):\n",
    "    def __init__(self,emb_dim,num_features):\n",
    "        super(featuresEncoder, self).__init__()\n",
    "\n",
    "        self.mlp = torch.nn.Linear(num_features,emb_dim,bias=True)\n",
    "        self.mlp.reset_parameters()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        h_x = self.mlp(x)\n",
    "        return h_x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c417c72-1cf7-443e-859a-b8daac35d2a7",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Creation of Graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edfdc897-6233-4180-b6d0-b73939d39135",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Single sample\n",
    "Definition of a function that creates a dgl graph for a given step in the simulation.\n",
    "It assigns the node and edge features to the nodes and edges of the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a4f55635-0bf9-4c02-8c81-56c3c55ed8f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_graph_dgl(f_sample, step):\n",
    "    \"\"\"\n",
    "    Function that creates a graph given the dictionary of a sample and a step.\n",
    "    We are creating one graph per time step.\n",
    "    Returns a dglGraph homogeneous (single node type representing the particles)\n",
    "\n",
    "    Arguments:\n",
    "    ----------------\n",
    "    f_sample : HDF5 file \n",
    "        Contains the information of a simulation for a given sample.\n",
    "        The keys are numbers representing each step [0,200]\n",
    "    step : int\n",
    "        Index of the step [0,200]\n",
    "    \"\"\"\n",
    "    assert isinstance(step, int), 'Argument of wrong type!'\n",
    "    step_data = f_sample[f'time_sequence/{step}']\n",
    "    sources = step_data['sources'][:]\n",
    "    destinations = step_data['destinations'][:]\n",
    "    num_nodes = len(step_data['node_features'])\n",
    "    #input_params = step_data['macro_input_features'][:]\n",
    "    \n",
    "    graph = dgl.graph((sources, destinations), num_nodes=num_nodes)\n",
    "    some_n_feats=step_data['node_features'][:]\n",
    "    #WARNING! node features and edge features should have the same dimension to be able to make operations \n",
    "\n",
    "    graph.ndata['n_feats'] = torch.tensor(some_n_feats)\n",
    "    #graph.theta = torch.tensor(input_params)\n",
    "\n",
    "    return dgl.to_bidirected(graph,copy_ndata=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2ee7a926-63e6-4c8b-a313-136046bfa4a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'dgl.heterograph.DGLHeteroGraph'>\n"
     ]
    }
   ],
   "source": [
    "graph_0 = create_graph_dgl(f_all['0_7'],100)\n",
    "print(type(graph_0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "80b290c0-ae8f-4bef-9118-3baf815cb307",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['_N']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#test out WARNING! node features and edge features should have \n",
    "#the same dimension to be able to make operations with edges and nodes features\n",
    "graph_0.ndata['n_feats']\n",
    "#graph_0.apply_edges(dgl.function.u_add_e('n_feats', 'e_feats', 'm'))\n",
    "graph_0.theta_2 = [1,2,3] #adding a graph attribute\n",
    "graph_0.theta_2\n",
    "graph_0.ntypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9aaee2b7-b2b8-450b-94c2-46ac1f419b24",
   "metadata": {
    "tags": []
   },
   "source": [
    "About **graph level features** see this in [dgl discuss forum](https://discuss.dgl.ai/t/graph-level-features/1974):\n",
    "\n",
    "*Option 1:* Store the graph level features in a special type of node. -> heterogeneous graph with different node types\n",
    "\n",
    "*Option 2:* Store the graph level features in a separate tensor. (This seems complicated to handle when shuffling the data.)\n",
    "\n",
    "*Option 3:* Store the graph level features as an attribute of the object DGLgraph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19e56b29-c5ee-4592-81b3-f1a0e1332a13",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### *Option 1*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6df7ea0d-39eb-48c1-832f-e63bf1a4be19",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_graph_dgl_heterogeneous(f_sample, step, num_nodes=10_000):\n",
    "    \"\"\"\n",
    "    Function that creates a graph given the dictionary of a sample and an step.\n",
    "    We are creating one graph per time step.\n",
    "    In this graph we have 2 node types 'particle' and 'contact_params'.\n",
    "    the second is a single node that carries the tensor with the \n",
    "    contact parameters.\n",
    "\n",
    "    Arguments:\n",
    "    ----------------\n",
    "    f_sample : HDF5 file or HDF5 group\n",
    "        Contains the information of a simulation for a given sample.\n",
    "        The keys are numbers representing each step [0,200]\n",
    "    step : int\n",
    "        Index of the step [0,200]\n",
    "    num_nodes : int\n",
    "        number of particles, this could also be retrieved from 'input_features' of each step.\n",
    "    \"\"\"\n",
    "    assert isinstance(step, int), 'Argument of wrong type!'\n",
    "    step_data = f_sample[f'time_sequence/{step}']\n",
    "    sources = torch.tensor(step_data['sources'][:])\n",
    "    destinations = torch.tensor(step_data['destinations'][:])\n",
    "    num_nodes = len(step_data['node_features'])\n",
    "    #input_params = np.concatenate((step_data['other_features'][:],step_data['macro_input_features'][:]))\n",
    "    input_params = step_data['macro_input_features'][:]\n",
    "    #creating the data of the graph with nodes => particles\n",
    "    data_dict = {\n",
    "        #n_type source, e_type, n_stype dest:(1D tensor node source ID, 1D tensor node dest ID)\n",
    "        ('particle','contact','particle'):(sources,destinations)\n",
    "    }\n",
    "    #Graph creation. adding a special node type for input_features\n",
    "    graph = dgl.heterograph(data_dict,\n",
    "                            num_nodes_dict={'particle':num_nodes,'theta':1,'macro':1})\n",
    "    \n",
    "    #Retrieve nodes and edge features\n",
    "    some_n_feats=step_data['node_features'][:] #all node features\n",
    "    #Assign features to nodes\n",
    "    graph.nodes['particle'].data['n_feats'] = torch.tensor(some_n_feats)\n",
    "    #graph.nodes['theta'].data['theta_params'] = torch.tensor([input_params]) \n",
    "    #tensor of size [1,len(input_params)]\n",
    "    \n",
    "    return dgl.to_bidirected(graph,copy_ndata=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "22ea0654-f0e6-4e70-8d69-59c1965ac8bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_hetero = create_graph_dgl_heterogeneous(f_all['5_2'],1)\n",
    "graph_hetero.ntypes\n",
    "graph_hetero.etypes\n",
    "#To get the data:\n",
    "#graph_hetero.nodes['theta'].data['theta_params'] #This returns a tensor\n",
    "#graph_hetero.ndata['theta_params'] #This returns a dict. Only possible if the feature is of a single type\n",
    "#graph_hetero.srcdata['position'] #This returns a dict. Only possible if the feature is of a single type\n",
    "#Graph operations\n",
    "graph_hetero.update_all(dgl.function.copy_u('n_feats','m'),dgl.function.mean('m','h_avg'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1baacb0a-887a-4a80-9081-16e98db9e19e",
   "metadata": {},
   "source": [
    "Loop to create `graphs_list`, containing one graph per simulation step for a given sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bc38eff8-3aa1-4a73-9be0-3e56ea59a998",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_graphs_dgl_of_sample(f_sample,homogeneous=True):\n",
    "    \"\"\"\n",
    "    Function that creates a graph given the dictionary of a sample and an step.\n",
    "    We are creating one graph per time step.\n",
    "\n",
    "    Arguments:\n",
    "    ----------------\n",
    "    f_sample : HDF5 group \n",
    "        Contains the information of a simulation for a given sample.\n",
    "    homogeneous : Bool\n",
    "        True (default) if generating homogeneous graphs single node type (particles)\n",
    "        False for having a graph with more node types in which the graph-level features would be stored\n",
    "    \"\"\"\n",
    "    num_steps = f_sample['metadata/num_steps'][()] #[()] is needed because is a scalar dataspace\n",
    "    strain_rate = torch.tensor([f_sample['metadata/compressive_strain_rate'][()]])\n",
    "    initial_friction = torch.tensor([f_sample['metadata/initial_friction'][()]])\n",
    "    pressure = torch.tensor([f_sample['metadata/pressure'][()]])\n",
    "    shear_strain_rate = torch.tensor([f_sample['metadata/shear_strain_rate'][()]])\n",
    "    graphs_sample = [] #list of graphs of sample f\n",
    "    \n",
    "    for i in range(num_steps - 2): #We don't create a graph for the last step\n",
    "        #position_next = f_sample[str(i+1)]['node_features'][:,2:5]\n",
    "        velocity_next = f_sample[f'time_sequence/{i+1}/node_features'][:,3:6]\n",
    "        macro = f_sample[f'time_sequence/{i+1}/macro_output_features'] #stress\n",
    "        inputs_this = torch.tensor(f_sample[f'time_sequence/{i}/macro_input_features']) #strain this t\n",
    "        inputs_next = torch.tensor(f_sample[f'time_sequence/{i+1}/macro_input_features']) #strain next t\n",
    "        delta_T = torch.tensor([f_sample[f'time_sequence/{i+1}/time'][0]-f_sample[f'time_sequence/{i}/time'][0]])\n",
    "        theta = np.concatenate((inputs_this,inputs_next,delta_T,strain_rate,initial_friction,pressure,shear_strain_rate))\n",
    "        if homogeneous:\n",
    "            graph_i = create_graph_dgl(f_sample,i)\n",
    "            #Labels -> Node features\n",
    "            graph_i.ndata['label'] = torch.tensor(velocity_next)\n",
    "            graph_i.label_macro = torch.tensor(macro)\n",
    "            graph_i.theta = torch.tensor(theta)\n",
    "        else:\n",
    "            graph_i = create_graph_dgl_heterogeneous(f_sample,i)\n",
    "            graph_i.nodes['particle'].data['label'] = torch.tensor(velocity_next)\n",
    "            graph_i.nodes['macro'].data['label_macro'] = torch.tensor([macro])\n",
    "            graph_i.nodes['theta'].data['theta_params'] = torch.tensor([theta])\n",
    "            \n",
    "        #graphs_sample.append(dgl.to_bidirected(graph_i,copy_ndata=True))\n",
    "        graphs_sample.append(graph_i)\n",
    "    \n",
    "    return graphs_sample"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "467ac7f4-9471-4879-bf97-c75ef30663bf",
   "metadata": {},
   "source": [
    "## Multiple samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a8af062b-a7c7-431e-a261-adfe0ee99449",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/ls/dj0bf91x7715kphjd6dghv0w0000gn/T/ipykernel_45638/3031097362.py:39: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  /Users/runner/miniforge3/conda-bld/pytorch-recipe_1647804326539/work/torch/csrc/utils/tensor_new.cpp:210.)\n",
      "  graph_i.nodes['theta'].data['theta_params'] = torch.tensor([theta])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "dgl.heterograph.DGLHeteroGraph"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list_g_i = create_graphs_dgl_of_sample(f_all['5_2'],homogeneous =False)\n",
    "type(list_g_i[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "12be92f0-222d-428e-9ebc-3f1d8de355dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list_g_i[80].nodes['theta'].data['theta_params'].size(dim=1)\n",
    "len(list_g_i[80].nodes['theta'].data['theta_params'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b49117d5-1aa5-47a3-9b4d-9c6dae4f1d53",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### (A) Data stored structurally in a tensor of fixed dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "5651bc97-5e03-4bf2-b1e8-c09b30ca7d3d",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "\"Unable to open object (object 'num_steps' doesn't exist)\"",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "Input \u001b[0;32mIn [67]\u001b[0m, in \u001b[0;36m<cell line: 3>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m tensor_data \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray([])\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m sample_id \u001b[38;5;129;01min\u001b[39;00m f_all:\n\u001b[0;32m----> 4\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[43mf_all\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43msample_id\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m/num_steps\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m[()] \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m201\u001b[39m: \u001b[38;5;66;03m#only considering samples that have 200 steps\u001b[39;00m\n\u001b[1;32m      5\u001b[0m         graphs_sample_i \u001b[38;5;241m=\u001b[39m create_graphs_dgl_of_sample(f_all[\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00msample_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m      6\u001b[0m         \u001b[38;5;66;03m#Staking the lists 2 axis: samples, steps\u001b[39;00m\n",
      "File \u001b[0;32mh5py/_objects.pyx:54\u001b[0m, in \u001b[0;36mh5py._objects.with_phil.wrapper\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mh5py/_objects.pyx:55\u001b[0m, in \u001b[0;36mh5py._objects.with_phil.wrapper\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/h5py/_hl/group.py:305\u001b[0m, in \u001b[0;36mGroup.__getitem__\u001b[0;34m(self, name)\u001b[0m\n\u001b[1;32m    303\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInvalid HDF5 object reference\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    304\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(name, (\u001b[38;5;28mbytes\u001b[39m, \u001b[38;5;28mstr\u001b[39m)):\n\u001b[0;32m--> 305\u001b[0m     oid \u001b[38;5;241m=\u001b[39m \u001b[43mh5o\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mopen\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mid\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_e\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlapl\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_lapl\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    306\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    307\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAccessing a group is done with bytes or str, \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    308\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m not \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\u001b[38;5;28mtype\u001b[39m(name)))\n",
      "File \u001b[0;32mh5py/_objects.pyx:54\u001b[0m, in \u001b[0;36mh5py._objects.with_phil.wrapper\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mh5py/_objects.pyx:55\u001b[0m, in \u001b[0;36mh5py._objects.with_phil.wrapper\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mh5py/h5o.pyx:190\u001b[0m, in \u001b[0;36mh5py.h5o.open\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: \"Unable to open object (object 'num_steps' doesn't exist)\""
     ]
    }
   ],
   "source": [
    "#creating a tensor (n_sample x n_step) containing a graph per step of each sample\n",
    "tensor_data = np.array([])\n",
    "for sample_id in f_all:\n",
    "    if f_all[f'{sample_id}/num_steps'][()] == 201: #only considering samples that have 200 steps\n",
    "        graphs_sample_i = create_graphs_dgl_of_sample(f_all[f'{sample_id}'])\n",
    "        #Staking the lists 2 axis: samples, steps\n",
    "        if len(tensor_data)==0: tensor_data = graphs_sample_i\n",
    "        else: tensor_data = np.vstack((tensor_data,graphs_sample_i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "5ef95c53-93b1-4d24-8524-90ed5c8ffea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#torch.tensor(tensor_data) #this is not supported because the elements are graphs (objects), not a common python type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "9f8ae3fa-b697-47a3-a8f8-a5a66d07bd1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tensor_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "174cd4d6-03e8-42cd-8573-ec757053b055",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### (B) Data stored unstructurally in a list (samples) of lists (steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 301,
   "id": "23343354-332e-4355-9211-26cc66ee465f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#creating a list of samples with variable step number\n",
    "unstructured_data = []\n",
    "for sample_id in f_all:\n",
    "    graphs_sample_i = create_graphs_dgl_of_sample(f_all[f'{sample_id}'])\n",
    "    if len(graphs_sample_i)>0: unstructured_data.append(graphs_sample_i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f72b7d7-f439-4539-9e04-1d93fcf5b323",
   "metadata": {},
   "source": [
    "### (C) Data stored as a list of graphs\n",
    "Samples and steps might be combined thus, the labels, contact params and input parameters should be stored at the graph ```homogeneous=False```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a86d7137-8726-4b99-8ce6-f6b215ec7d92",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/ls/dj0bf91x7715kphjd6dghv0w0000gn/T/ipykernel_48452/3031097362.py:39: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  /Users/runner/miniforge3/conda-bld/pytorch-recipe_1647804326539/work/torch/csrc/utils/tensor_new.cpp:210.)\n",
      "  graph_i.nodes['theta'].data['theta_params'] = torch.tensor([theta])\n"
     ]
    }
   ],
   "source": [
    "list_graphs = []\n",
    "for i,sample_id in enumerate(f_all):\n",
    "    if i < len(f_all)-1:\n",
    "        graphs_sample_i = create_graphs_dgl_of_sample(f_all[sample_id],homogeneous=False)\n",
    "        list_graphs.extend(graphs_sample_i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cc66b82f-3d16-4052-88db-6ee124c4e665",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'DGLHeteroGraph' object has no attribute 'theta'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Input \u001b[0;32mIn [14]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#homogeneous testing\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[43mlist_graphs\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtheta\u001b[49m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'DGLHeteroGraph' object has no attribute 'theta'"
     ]
    }
   ],
   "source": [
    "#homogeneous testing\n",
    "list_graphs[3].theta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 242,
   "id": "3b6f59ea-145c-45fc-b80a-45e6e92d76e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#heterogeneous testing\n",
    "t = list_graphs[3].nodes['theta'].data['theta_params']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "id": "189dc0d5-6083-4761-a8b5-19cc428e4888",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Sizes of tensors must match except in dimension 0. Expected size 11 but got size 9 for tensor number 1 in the list.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Input \u001b[0;32mIn [245]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m n\u001b[38;5;241m=\u001b[39mlist_graphs[\u001b[38;5;241m3\u001b[39m]\u001b[38;5;241m.\u001b[39mnodes[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mparticle\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mdata[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mn_feats\u001b[39m\u001b[38;5;124m'\u001b[39m]\n\u001b[0;32m----> 2\u001b[0m r \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcat\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43mn\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Sizes of tensors must match except in dimension 0. Expected size 11 but got size 9 for tensor number 1 in the list."
     ]
    }
   ],
   "source": [
    "n=list_graphs[3].nodes['particle'].data['n_feats']\n",
    "r = torch.cat((t,n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f12c0a37-3672-4a77-8dfc-1c62229861d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set which type of data is going to be take into account: \n",
    "#unstructured (steps of variable size) or tensor (only considering samples with 200 steps) \n",
    "#or list_graphs (all steps and samples combined)\n",
    "graph_data = list_graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85082e56-0e04-435f-828b-1100a5f950ab",
   "metadata": {},
   "source": [
    "# Data Split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d7024f-f5c4-418e-af44-9c60b2474418",
   "metadata": {},
   "source": [
    "### DGL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "6d0dca97-882a-4346-a3b0-00b3a68f81b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Option 1: very basic in DGL\n",
    "data_train, data_val, data_test = dgl.data.utils.split_dataset(\n",
    "    graph_data,\n",
    "    frac_list=[0.8, 0.1, 0.1],\n",
    "    shuffle=True,random_state=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2116c302-f584-456c-bd53-2e17a1380a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Option 2 (Preferred): from DGL docs\n",
    "from dgl.dataloading import GraphDataLoader\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "\n",
    "num_examples = len(graph_data)\n",
    "#number of samples in train, validation and test dataset\n",
    "num_train = int(num_examples * 0.8)\n",
    "num_val = int(num_examples * 0.1)\n",
    "num_test = num_examples - num_train - num_val\n",
    "\n",
    "bs = 1 #batch size\n",
    "\n",
    "train_sampler = SubsetRandomSampler(torch.arange(num_train))\n",
    "val_sampler = SubsetRandomSampler(torch.arange(num_train, num_train+num_val))\n",
    "test_sampler = SubsetRandomSampler(torch.arange(num_train+num_val, num_examples))\n",
    "\n",
    "train_dataloader = GraphDataLoader(\n",
    "    graph_data, sampler=train_sampler, batch_size=bs, drop_last=False)\n",
    "val_dataloader = GraphDataLoader(\n",
    "    graph_data, sampler=val_sampler, batch_size=bs, drop_last=False)\n",
    "test_dataloader = GraphDataLoader(\n",
    "    graph_data, sampler=test_sampler, batch_size=bs, drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecdb2631-9a9f-4c52-b0b7-d14b297a4376",
   "metadata": {},
   "source": [
    "If batch_size = 1 this technique is equivalent to the next one (pytorch). `GraphDataLoader` may combine the graphs is a batch in a single structure, but you can only take the advantages of having this feature if batch_size >1 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c558504-c982-4a5f-b1d6-5fb9bf49ca64",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### pytorch\n",
    "For a specific pytorch geometric function see [torch_geometric.loader.DataLoader](https://pytorch-geometric.readthedocs.io/en/latest/notes/batching.html?highlight=DataLoader) But this one doen't work with 'object' data types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd2d49be-2be5-4a79-aa97-20a6f84f2ad7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO\n",
    "def get_paths_to_data(file_hdf5):\n",
    "    \"\"\"\n",
    "    Function that returns a list with tuples:\n",
    "    (sample_id (str), step_id (int))\n",
    "    Arguments:\n",
    "    file_hdf5: and hdf5 file havign our structure {sample_id}/'time_sequence'/{time_step}\n",
    "    \"\"\"\n",
    "    list_keys =[]\n",
    "    for i,sample_id in enumerate(file_hdf5.keys()):\n",
    "        if i < len(file_hdf5)-1: #getting rid of metadata (last one)\n",
    "            str_acces_step = sample_id+'/time_sequence/'\n",
    "            num_steps = file_hdf5[sample_id+'metadata/num_steps'][()]\n",
    "            for j in range(num_steps):\n",
    "                list_keys.append((sample_id,j))\n",
    "    return list_keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "c99ce77f-d673-447f-998f-69998c4ef964",
   "metadata": {},
   "outputs": [],
   "source": [
    "#option 1: very general Pytorch TODO\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "num_examples = len(list_keys)\n",
    "#number of samples in train, validation and test dataset\n",
    "num_train = int(num_examples * 0.8)\n",
    "num_val = int(num_examples * 0.1)\n",
    "num_test = num_examples - num_train - num_val\n",
    "\n",
    "bs = 5 #batch size\n",
    "\n",
    "train_sampler = SubsetRandomSampler(torch.arange(num_train))\n",
    "val_sampler = SubsetRandomSampler(torch.arange(num_train, num_train+num_val))\n",
    "test_sampler = SubsetRandomSampler(torch.arange(num_train+num_val, num_examples))\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    graph_data, sampler=train_sampler, batch_size=bs, drop_last=False)\n",
    "val_dataloader = DataLoader(\n",
    "    graph_data, sampler=val_sampler, batch_size=bs, drop_last=False)\n",
    "test_dataloader = DataLoader(\n",
    "    graph_data, sampler=test_sampler, batch_size=bs, drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84908cb1-e0a0-4066-9a0a-fd7d632117af",
   "metadata": {},
   "source": [
    "# Creation of GNN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a394df17-d39e-43a8-9270-b6d1a895f527",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 1. A simple GNN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf6b8044-693c-4735-95ed-c95de987da6a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### GNN at the node level\n",
    "We create this layer that is going to do the message passing and then prediction of each node feature as follows:\n",
    "$$\\mathbf{h}_i^{(l+1)} = \\text{MLP_1}[ (1+\\varepsilon) \\cdot \\mathbf{h}_i^{(l)} + \\text{ReLU}\\{W (\\text{concat}(\\sum_{j \\in \\mathcal{N}_i} \\mathbf{h}_j^{l}, \\theta_{DEM})+bias)\\}]$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "7f35fd26-1582-4be1-8f45-a261f8d3b92e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class myGnn(nn.Module):\n",
    "    def __init__(self,emb_dim,theta_dim):\n",
    "        \"\"\"\n",
    "        This class implements a single GNN layer with node and edge features.\n",
    "    \n",
    "        Arguments\n",
    "        ----------\n",
    "        emb_dim : int\n",
    "            Number of dimensions of the node embeding (features)\n",
    "        theta_dim : int\n",
    "            Number of dimensions of the theta parameters\n",
    "        \"\"\"\n",
    "        super(myGnn, self).__init__()\n",
    "        \n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(emb_dim, 2*emb_dim), \n",
    "            #nn.BatchNorm1d(2*emb_dim), \n",
    "            nn.ReLU(), \n",
    "            nn.Linear(2*emb_dim, emb_dim)\n",
    "        )\n",
    "        self.linear = nn.Linear(emb_dim+theta_dim,emb_dim,bias=True)\n",
    "        self.epsilon = nn.Parameter(torch.FloatTensor([0]))\n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        \"\"\"\n",
    "        A method that defines how to re-initialize all the parameters of\n",
    "        the model.\n",
    "        \"\"\"\n",
    "        # MLP\n",
    "        for layer in self.mlp:\n",
    "            if hasattr(layer, 'reset_parameters'):\n",
    "                layer.reset_parameters()\n",
    "        self.linear.reset_parameters()\n",
    "        # epsilon\n",
    "        nn.init.constant_(self.epsilon, 0.)\n",
    "    \n",
    "    def forward(self, g, x_node):\n",
    "        \"\"\"\n",
    "        Applies the update equation, which includes edge features\n",
    "        Arguments\n",
    "        ----------\n",
    "        g : DGLGraph\n",
    "            The graph used for message passing\n",
    "        x_node : Tensor\n",
    "            The node features (i.e., velocity)\n",
    "        \"\"\"\n",
    "        with g.local_scope():\n",
    "            g.ndata['h'] = x_node\n",
    "            g.update_all(dgl.function.copy_u('h', 'm'), dgl.function.sum('m', 'h_mp'))\n",
    "            h_mp = g.ndata['h_mp']\n",
    "            # Concatenation of node features and theta and passing it to the linear layer\n",
    "            h_n = self.linear(torch.cat(h_mp, g.theta))\n",
    "            h_n = F.relu(h_n)\n",
    "            return self.mlp((1+self.epsilon) * x_node + h_n)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cad58e9-b3df-4aeb-a552-b80a82a18bfa",
   "metadata": {
    "tags": []
   },
   "source": [
    "### GNN at the graph level\n",
    "We create this model in which we will have several myGnn layers stacked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "0a663a90-c217-4dd0-b627-ac7ec33c0db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class model_gnn(nn.Module):\n",
    "    def __init__(self, emb_dim, theta_dim, num_layers, add_bn=False, dropout=0):\n",
    "        \"\"\"\n",
    "        Defines the architecture of the model.\n",
    "        The `forward` method will define how to use the \n",
    "        layers created here.\n",
    "        \"\"\"\n",
    "        super(model_gnn, self).__init__()\n",
    "        \n",
    "        self.num_layers = num_layers\n",
    "        #batch normalization and droput\n",
    "        self.add_bn = add_bn\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        #Stacked layers\n",
    "        self.layers = nn.ModuleList()\n",
    "        for _ in range(num_layers):\n",
    "            self.layers.append(myGnn(emb_dim,theta_dim))\n",
    "        \n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        \"\"\"\n",
    "        A method that defines how to re-initialize all the parameters of\n",
    "        the model.  This method will be invoked to start a new experiment,\n",
    "        so it is important that all parameters are reset, or else there will \n",
    "        be an information leak between experiments.\n",
    "        \"\"\"\n",
    "        for ly in self.layers:\n",
    "            ly.reset_parameters()\n",
    "        \n",
    "    def forward(self, g):\n",
    "        \"\"\"\n",
    "        Arguments\n",
    "        ----------\n",
    "        g : DGLGraph\n",
    "            The graph used for message passing\n",
    "        \"\"\"\n",
    "        h_node = g.ndata['n_feats']\n",
    "        for i,ly in enumerate(self.layers[:-1]):\n",
    "            with g.local_scope(): #So that we will not modify the graph\n",
    "                h_node = ly(g, h_node)\n",
    "                h_node = F.relu(h_node)\n",
    "        return h_node"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01201c5a-3862-43d6-a973-b515b0755fa2",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 2. A GNN that handles heterogenous graph\n",
    "### GNN at the node level\n",
    "\n",
    "We create this layer that is going to do the message passing and then prediction of each node feature as follows:\n",
    "$$\\mathbf{h}_i^{(l+1)} = MLP_1 \\left[ (1+\\varepsilon) \\cdot \\mathbf{h}_i^{(l)} + MLP_2 \\left( CONCAT\\left\\{ \\sum_{j \\in \\mathcal{N}_i} \\frac{\\mathbf{h}_j^{(l)}}{\\mathcal{N}_i}, \\theta_{DEM} \\right\\} \\right) \\right]$$\n",
    "Where $\\theta_{DEM}$ is `input_params` like $\\dot{\\varepsilon},\\sigma , \\mu_0$. Note that in this model we do ***not consider*** edge features for the Message passing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3aef1a44-f851-40a4-9c92-2ce0d1946c7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class myGnn_heterogeneous(nn.Module):\n",
    "    def __init__(self,emb_dim,theta_dim,n_type_part,n_type_th):\n",
    "        \"\"\"\n",
    "        This class implements a single GNN layer with node and edge features.\n",
    "        Arguments\n",
    "        ----------\n",
    "        emb_dim : int\n",
    "            Dimension of the node embedding (feature)\n",
    "        theta_dim : int\n",
    "            Dimension of the control parameters \\theta\n",
    "        n_type_part : str\n",
    "            String identifier of the nodes representing particles\n",
    "        n_type_th : str\n",
    "            String identifier of the nodes representing theta \n",
    "            (control parameters of the experiment)\n",
    "        \"\"\"\n",
    "        super(myGnn_heterogeneous, self).__init__()\n",
    "        \n",
    "        self.n_type_particles = n_type_part\n",
    "        self.n_type_theta = n_type_th\n",
    "        self.theta_size = theta_dim\n",
    "        self.mlp_1 = nn.Linear(emb_dim,emb_dim,bias=True)\n",
    "        self.mlp_2 = nn.Linear(emb_dim+theta_dim,emb_dim,bias=True)\n",
    "        self.epsilon = nn.Parameter(torch.FloatTensor([0]))\n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        \"\"\"\n",
    "        A method that defines how to re-initialize all the parameters of\n",
    "        the model.\n",
    "        \"\"\"\n",
    "        # MLP\n",
    "        self.mlp_1.reset_parameters()\n",
    "        self.mlp_2.reset_parameters()\n",
    "        \n",
    "        # epsilon\n",
    "        nn.init.constant_(self.epsilon, 0.)\n",
    "    \n",
    "    def forward(self, g, x_node):\n",
    "        \"\"\"\n",
    "        Applies the update equation, which includes edge features\n",
    "        Arguments\n",
    "        ----------\n",
    "        g : DGLGraph heterogeneous\n",
    "            The graph used for message passing\n",
    "        x_node : Tensor\n",
    "            The node features or embeddings (i.e., position)\n",
    "        \"\"\"\n",
    "        theta = g.nodes[self.n_type_theta].data['theta_params']\n",
    "        num_nodes_in_batch = g.nodes[self.n_type_particles].data['n_feats'].size(dim=0)\n",
    "        bs = theta.size(dim=0)\n",
    "        repeats = torch.ones(bs,dtype=int)*int(num_nodes_in_batch/bs)\n",
    "        theta_repeated = torch.repeat_interleave(theta,repeats,dim=0)\n",
    "        #we obtain a tensor of size (num_nodes,num_theta_params)\n",
    "        with g.local_scope():\n",
    "            # Store node features to 'h_n' key in g.ndata\n",
    "            g.nodes[self.n_type_particles].data['h_n'] = x_node\n",
    "            # m-> message , h_avg -> aggregated message\n",
    "            g.update_all(dgl.function.copy_u('h_n', 'm'), dgl.function.mean('m', 'h_avg'))\n",
    "            h_avg = g.nodes[self.n_type_particles].data['h_avg']\n",
    "            cat_vector = torch.cat((h_avg,theta_repeated),dim=-1)\n",
    "            h_n = self.mlp_2(cat_vector.float())\n",
    "            h_self = ((1+ self.epsilon) * x_node)\n",
    "            return self.mlp_1(h_self + h_n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "11926559-d15e-4742-ac00-c197bdc1db01",
   "metadata": {},
   "outputs": [],
   "source": [
    "class model_gnn_heterogeneous(nn.Module):\n",
    "    def __init__(self,emb_dim,theta_dim,out_dim,\n",
    "                  num_layers,n_type_part,n_type_th,add_bn=False,dropout=0):\n",
    "        \"\"\"\n",
    "        Defines the architecture of the model.\n",
    "        Arguments\n",
    "        ----------\n",
    "        emb_dim : int\n",
    "            Dimension of the node embedding (feature)\n",
    "        theta_dim : int\n",
    "            Dimension of the control parameters \\theta\n",
    "        out_dim : int\n",
    "            Dimension of the output\n",
    "        n_type_part : str\n",
    "            String identifier of the nodes representing particles\n",
    "        n_type_th : str\n",
    "            String identifier of the nodes representing theta \n",
    "            (control parameters of the experiment)\n",
    "        \"\"\"\n",
    "        super(model_gnn_heterogeneous, self).__init__()\n",
    "        \n",
    "        self.num_layers = num_layers\n",
    "        self.n_type_particles = n_type_part\n",
    "        self.n_type_theta = n_type_th\n",
    "        \n",
    "        #batch normalization and droput\n",
    "        self.add_bn = add_bn\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        #Stacked layers\n",
    "        self.layers = nn.ModuleList()\n",
    "        for _ in range(num_layers):\n",
    "            self.layers.append(myGnn_heterogeneous(emb_dim,theta_dim,n_type_part,n_type_th))\n",
    "            \n",
    "        self.layers.append(nn.Linear(emb_dim,out_dim))\n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        \"\"\"\n",
    "        Defines how to re-initialize all parameters of the model.\n",
    "        This method will be invoked to start a new experiment,\n",
    "        so it is important that all parameters are reset, or else there will \n",
    "        be an information leak between experiments.\n",
    "        \"\"\"\n",
    "        for ly in self.layers:\n",
    "            ly.reset_parameters()\n",
    "        \n",
    "    def forward(self, g):\n",
    "        \"\"\"\n",
    "        Arguments\n",
    "        ----------\n",
    "        g : DGLGraph heterogeneous\n",
    "            The graph used for message passing\n",
    "        \"\"\"\n",
    "        h_node = g.nodes[self.n_type_particles].data['n_feats']\n",
    "        for ly in self.layers[:-1]:\n",
    "            with g.local_scope(): #So that we will not modify the graph\n",
    "                h_node = ly(g, h_node)\n",
    "                h_node = F.relu(h_node)\n",
    "        return self.layers[-1](h_node)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13b0854d-c9d1-4d5e-a0be-c2843b1c2834",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Creation of the model\n",
    "<a id='creation_model'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9cbfd58f-9b2c-49a2-a79c-688315f217e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model parameters definition\n",
    "num_layers = 2\n",
    "node_types = graph_data[0].ntypes\n",
    "if len(node_types)>1 : #heterogeneous graph\n",
    "    node_type_macro, node_type_particles, node_type_theta = node_types\n",
    "    emb_dim = graph_data[0].nodes[node_type_particles].data['n_feats'].size(dim=1)\n",
    "    theta_dim = graph_data[0].nodes[node_type_theta].data['theta_params'].size(dim=1)\n",
    "    out_feats = graph_data[0].nodes[node_type_particles].data['label'].size(dim=1)\n",
    "    model = model_gnn_heterogeneous(emb_dim,theta_dim,out_feats,num_layers,node_type_particles,node_type_theta)\n",
    "else: #homogeneous graph\n",
    "    emb_dim = graph_data[0].ndata['n_feats'].size(dim=1)\n",
    "    theta_dim = graph_data[0].theta.size(dim=0)\n",
    "    out_feats = graph_data[0].ndata['label'].size(dim=1)+graph_data[0].label_macro.size(dim=0)\n",
    "    model = model_gnn(emb_dim,theta_dim,out_feats,num_layers)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72d4152b-1000-49c7-966a-e61d9a5c2c03",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 3. Model of Brandstetter et al. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04cb05be-f685-4881-85ed-872d20f5d51d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20893571-dc1f-4e46-b6a6-5a4f629dd603",
   "metadata": {},
   "source": [
    "# Model training\n",
    "The following steps will be carried up:\n",
    "1. Definition of the optimizer\n",
    "2. Definition of the loss function\n",
    "3. Definition of the trainning loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bec49bb-07ae-4b0a-ab28-a057b711768d",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Optimizer definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44369de-9894-443c-8a46-19d07ffba660",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8ca8b67-3157-4ae3-9862-ce17569c83b7",
   "metadata": {},
   "source": [
    "### Loss Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78bd92fc-7937-4608-82a3-19f3acaeb416",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = torch.nn.L1Loss()\n",
    "#loss_fn = torch.nn.mse_loss()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97bb77a1-fdc2-4e5a-8bc3-4c6fcb8b1b8b",
   "metadata": {
    "tags": []
   },
   "source": [
    "### The training Loop proposed by [Pytorch](https://pytorch.org/tutorials/beginner/introyt/trainingyt.html)\n",
    "This is a very basic template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "62834559-7d5c-4c18-9c93-d506c5f28d93",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_labels_from_graphList(list_graphs):\n",
    "    \"\"\"\n",
    "    Returns the labels of a graph\n",
    "    Arguments:\n",
    "    --------------\n",
    "    list_graphs : graph\n",
    "        Super graph containing batch_size DGLGraphs\n",
    "    \"\"\"\n",
    "    node_types = list_graphs.ntypes\n",
    "    if len(node_types)>1: #heterogeneous\n",
    "        labels = list_graphs.nodes[node_types[1]].data['label']\n",
    "        labels_macro = list_graphs.nodes[node_types[0]].data['label_macro'].squeeze(dim=0)\n",
    "    else: #homogeneous\n",
    "        labels = list_graphs.ndata['label']\n",
    "        labels_macro = list_graphs.label_macro\n",
    "    return labels, labels_macro"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91187faf-8dd0-40d2-b18e-80f016e623cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_one_epoch(model,epoch_index,loss_fn,training_loader,tb_writer=None):\n",
    "    \"\"\"\n",
    "    Function to train an epoch, default pytorch using TensorBoard.\n",
    "    Arguments:\n",
    "    -----------------\n",
    "        model : GNN model\n",
    "        epoch_index: int\n",
    "        loss_fn: \n",
    "            Loss function\n",
    "        tb_writer: SummaryWriter (TensorBoard)\n",
    "        training_loader: Data loader\n",
    "    \"\"\"\n",
    "    running_loss = 0.\n",
    "    last_loss = 0.\n",
    "\n",
    "    # Here, we use enumerate(training_loader) instead of\n",
    "    # iter(training_loader) so that we can track the batch\n",
    "    # index and do some intra-epoch reporting\n",
    "    for i, batch_graphs in enumerate(training_loader):\n",
    "        # Zero your gradients for every batch!\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Make predictions for this batch\n",
    "        outputs = model(batch_graphs)\n",
    "        \n",
    "        #get the labels\n",
    "        labels, labels_macro = get_labels_from_graphList(batch_graphs)\n",
    "        \n",
    "        # Compute features loss and its gradients\n",
    "        loss = loss_fn(outputs, labels) #For the moment labels only node features\n",
    "        loss.backward()\n",
    "        #Compute macro loss and its gradients\n",
    "        #loss_macro = loss_fn(outputs[3:], labels_macro)\n",
    "        #loss_macro.backward()\n",
    "        \n",
    "        # Adjust learning weights\n",
    "        optimizer.step()\n",
    "\n",
    "        # Gather data and report\n",
    "        running_loss += loss.item()\n",
    "        if i % 1000 == 999:\n",
    "            last_loss = running_loss / 1000 # loss per batch\n",
    "            print('  batch {} loss: {}'.format(i + 1, last_loss))\n",
    "            tb_x = epoch_index * len(training_loader) + i + 1\n",
    "            if (tb_writer==None): wandb.log({'Loss/train': last_loss,\n",
    "                             'trainning step': tb_x})\n",
    "            else: tb_writer.add_scalar('Loss/train', last_loss, tb_x)\n",
    "            running_loss = 0.\n",
    "\n",
    "    return last_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a38333a-2d85-41c1-9ef1-9897f7c68355",
   "metadata": {},
   "source": [
    "**Per-Epoch Activity**\n",
    "There are a couple of things we’ll want to do once per epoch:\n",
    "\n",
    "* Perform validation by checking our relative loss on a set of data that was not used for training, and report this\n",
    "* Save a copy of the model\n",
    "\n",
    "Here, we’ll do our reporting in TensorBoard. During the training a folder will be creatied inside runs. To visualize them open the folder containing runs and run `tensorboard --logdir=runs`, then copy `http://localhost:6006/` to your browser to open the tensorboard visualizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "517a2dc4-2546-4284-9129-7a758a933fb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_tensorboard(model,num_epochs,loss_fn,training_loader,validation_loader,tb_writer):\n",
    "    best_vloss = 1_000_000.\n",
    "    for epoch in range(num_epochs):\n",
    "        print('EPOCH {}:'.format(epoch + 1))\n",
    "\n",
    "        # Make sure gradient tracking is on, and do a pass over the data\n",
    "        model.train(True)\n",
    "        avg_loss = train_one_epoch(model,epoch,loss_fn,training_loader,tb_writer=tb_writer)\n",
    "\n",
    "        # We don't need gradients on to do reporting\n",
    "        model.train(False)\n",
    "\n",
    "        running_vloss = 0.0\n",
    "        for i, vdata in enumerate(validation_loader):\n",
    "            voutputs = model(vdata)\n",
    "            vlabels, vlabels_macro = get_labels_from_graphList(vdata)\n",
    "            vloss = loss_fn(voutputs, vlabels)\n",
    "            running_vloss += vloss\n",
    "\n",
    "        avg_vloss = running_vloss / (i + 1)\n",
    "        print('LOSS train {} valid {}'.format(avg_loss, avg_vloss))\n",
    "\n",
    "        # Log the running loss averaged per batch\n",
    "        # for both training and validation\n",
    "        tb_writer.add_scalars('Training vs. Validation Loss',\n",
    "                        { 'Training' : avg_loss, 'Validation' : avg_vloss },\n",
    "                        epoch)\n",
    "        tb_writer.flush()\n",
    "\n",
    "        # Track best performance, and save the model's state\n",
    "        if avg_vloss < best_vloss:\n",
    "            best_vloss = avg_vloss\n",
    "            model_path = 'model_{}_{}'.format(timestamp, epoch)\n",
    "            torch.save(model.state_dict(), model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1030f32-15e4-42b5-9a15-fe0def407274",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from datetime import datetime\n",
    "# Initializing in a separate cell so we can easily add more epochs to the same run\n",
    "timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')\n",
    "writer = SummaryWriter('runs/fashion_trainer_{}'.format(timestamp))\n",
    "EPOCHS = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 345,
   "id": "d1f22840-0bb8-4386-bde5-ea71ae465235",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH 1:\n",
      "  batch 1000 loss: 0.6874793180823326\n",
      "  batch 2000 loss: 0.6861921851038932\n",
      "  batch 3000 loss: 0.6833268201947212\n",
      "  batch 4000 loss: 0.6827743496894837\n",
      "  batch 5000 loss: 0.6816953321099282\n",
      "  batch 6000 loss: 0.6831414459347724\n",
      "  batch 7000 loss: 0.6823037651181221\n",
      "  batch 8000 loss: 0.6821828244924545\n",
      "  batch 9000 loss: 0.6829529471993446\n",
      "  batch 10000 loss: 0.6841480988264084\n",
      "  batch 11000 loss: 0.6853966968059539\n",
      "  batch 12000 loss: 0.6864752916097641\n",
      "LOSS train 0.6864752916097641 valid 0.6797007322311401\n",
      "EPOCH 2:\n",
      "  batch 1000 loss: 0.68649784553051\n",
      "  batch 2000 loss: 0.6822987530827522\n",
      "  batch 3000 loss: 0.6818726670145988\n",
      "  batch 4000 loss: 0.6866000335216522\n",
      "  batch 5000 loss: 0.6844046101570129\n",
      "  batch 6000 loss: 0.6814422015547752\n",
      "  batch 7000 loss: 0.6851514179706574\n",
      "  batch 8000 loss: 0.6843487210273743\n",
      "  batch 9000 loss: 0.6830541912317276\n",
      "  batch 10000 loss: 0.6842596131563187\n",
      "  batch 11000 loss: 0.6832783486247063\n",
      "  batch 12000 loss: 0.6853938121795654\n",
      "LOSS train 0.6853938121795654 valid 0.6797045469284058\n",
      "EPOCH 3:\n",
      "  batch 1000 loss: 0.6823789322972298\n",
      "  batch 2000 loss: 0.6848776631355286\n",
      "  batch 3000 loss: 0.6829944735169411\n",
      "  batch 4000 loss: 0.6837641568779945\n",
      "  batch 5000 loss: 0.6874307312369347\n",
      "  batch 6000 loss: 0.6883712704777718\n",
      "  batch 7000 loss: 0.6823727006316185\n",
      "  batch 8000 loss: 0.6834711146950722\n",
      "  batch 9000 loss: 0.6820664675235748\n",
      "  batch 10000 loss: 0.6849515787959098\n",
      "  batch 11000 loss: 0.6822153873443604\n",
      "  batch 12000 loss: 0.6830203679203987\n",
      "LOSS train 0.6830203679203987 valid 0.6797103881835938\n",
      "EPOCH 4:\n",
      "  batch 1000 loss: 0.6823081954717636\n",
      "  batch 2000 loss: 0.6829553837776184\n",
      "  batch 3000 loss: 0.6839974483847618\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [345]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtrain_tensorboard\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43mEPOCHS\u001b[49m\u001b[43m,\u001b[49m\u001b[43mloss_fn\u001b[49m\u001b[43m,\u001b[49m\u001b[43mtrain_dataloader\u001b[49m\u001b[43m,\u001b[49m\u001b[43mval_dataloader\u001b[49m\u001b[43m,\u001b[49m\u001b[43mwriter\u001b[49m\u001b[43m)\u001b[49m\n",
      "Input \u001b[0;32mIn [344]\u001b[0m, in \u001b[0;36mtrain_tensorboard\u001b[0;34m(model, num_epochs, loss_fn, training_loader, validation_loader, tb_writer)\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;66;03m# Make sure gradient tracking is on, and do a pass over the data\u001b[39;00m\n\u001b[1;32m      7\u001b[0m model\u001b[38;5;241m.\u001b[39mtrain(\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m----> 8\u001b[0m avg_loss \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_one_epoch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43mepoch\u001b[49m\u001b[43m,\u001b[49m\u001b[43mloss_fn\u001b[49m\u001b[43m,\u001b[49m\u001b[43mtraining_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43mtb_writer\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtb_writer\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;66;03m# We don't need gradients on to do reporting\u001b[39;00m\n\u001b[1;32m     11\u001b[0m model\u001b[38;5;241m.\u001b[39mtrain(\u001b[38;5;28;01mFalse\u001b[39;00m)\n",
      "Input \u001b[0;32mIn [338]\u001b[0m, in \u001b[0;36mtrain_one_epoch\u001b[0;34m(model, epoch_index, loss_fn, training_loader, tb_writer)\u001b[0m\n\u001b[1;32m     14\u001b[0m last_loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.\u001b[39m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;66;03m# Here, we use enumerate(training_loader) instead of\u001b[39;00m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# iter(training_loader) so that we can track the batch\u001b[39;00m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# index and do some intra-epoch reporting\u001b[39;00m\n\u001b[0;32m---> 19\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, batch_graphs \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(training_loader):\n\u001b[1;32m     20\u001b[0m     \u001b[38;5;66;03m# Zero your gradients for every batch!\u001b[39;00m\n\u001b[1;32m     21\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[1;32m     23\u001b[0m     \u001b[38;5;66;03m# Make predictions for this batch\u001b[39;00m\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/torch/utils/data/dataloader.py:530\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    528\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    529\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m--> 530\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    531\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    532\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    533\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    534\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/torch/utils/data/dataloader.py:570\u001b[0m, in \u001b[0;36m_SingleProcessDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    568\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_next_data\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    569\u001b[0m     index \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_next_index()  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[0;32m--> 570\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_dataset_fetcher\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfetch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindex\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[1;32m    571\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory:\n\u001b[1;32m    572\u001b[0m         data \u001b[38;5;241m=\u001b[39m _utils\u001b[38;5;241m.\u001b[39mpin_memory\u001b[38;5;241m.\u001b[39mpin_memory(data)\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py:52\u001b[0m, in \u001b[0;36m_MapDatasetFetcher.fetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     50\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     51\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n\u001b[0;32m---> 52\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollate_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/dgl-0.9-py3.9-macosx-11.0-arm64.egg/dgl/dataloading/dataloader.py:1035\u001b[0m, in \u001b[0;36mGraphCollator.collate\u001b[0;34m(self, items)\u001b[0m\n\u001b[1;32m   1033\u001b[0m elem_type \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtype\u001b[39m(elem)\n\u001b[1;32m   1034\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(elem, DGLHeteroGraph):\n\u001b[0;32m-> 1035\u001b[0m     batched_graphs \u001b[38;5;241m=\u001b[39m \u001b[43mbatch_graphs\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitems\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1036\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m batched_graphs\n\u001b[1;32m   1037\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m F\u001b[38;5;241m.\u001b[39mis_tensor(elem):\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/dgl-0.9-py3.9-macosx-11.0-arm64.egg/dgl/batch.py:177\u001b[0m, in \u001b[0;36mbatch\u001b[0;34m(graphs, ndata, edata, node_attrs, edge_attrs)\u001b[0m\n\u001b[1;32m    174\u001b[0m ntype_ids \u001b[38;5;241m=\u001b[39m [graphs[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mget_ntype_id(n) \u001b[38;5;28;01mfor\u001b[39;00m n \u001b[38;5;129;01min\u001b[39;00m ntypes]\n\u001b[1;32m    175\u001b[0m etypes \u001b[38;5;241m=\u001b[39m [etype \u001b[38;5;28;01mfor\u001b[39;00m _, etype, _ \u001b[38;5;129;01min\u001b[39;00m relations]\n\u001b[0;32m--> 177\u001b[0m gidx \u001b[38;5;241m=\u001b[39m \u001b[43mdisjoint_union\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraphs\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_graph\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmetagraph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[43mg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_graph\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mg\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mgraphs\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    178\u001b[0m retg \u001b[38;5;241m=\u001b[39m DGLHeteroGraph(gidx, ntypes, etypes)\n\u001b[1;32m    180\u001b[0m \u001b[38;5;66;03m# Compute batch num nodes\u001b[39;00m\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/dgl-0.9-py3.9-macosx-11.0-arm64.egg/dgl/heterograph_index.py:1254\u001b[0m, in \u001b[0;36mdisjoint_union\u001b[0;34m(metagraph, graphs)\u001b[0m\n\u001b[1;32m   1239\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdisjoint_union\u001b[39m(metagraph, graphs):\n\u001b[1;32m   1240\u001b[0m     \u001b[38;5;124;03m\"\"\"Return a disjoint union of the input heterographs.\u001b[39;00m\n\u001b[1;32m   1241\u001b[0m \n\u001b[1;32m   1242\u001b[0m \u001b[38;5;124;03m    Parameters\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1252\u001b[0m \u001b[38;5;124;03m        Batched Heterograph.\u001b[39;00m\n\u001b[1;32m   1253\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m-> 1254\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_CAPI_DGLHeteroDisjointUnion_v2\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmetagraph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgraphs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/dgl-0.9-py3.9-macosx-11.0-arm64.egg/dgl/_ffi/_ctypes/function.py:188\u001b[0m, in \u001b[0;36mFunctionBase.__call__\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m    186\u001b[0m ret_val \u001b[38;5;241m=\u001b[39m DGLValue()\n\u001b[1;32m    187\u001b[0m ret_tcode \u001b[38;5;241m=\u001b[39m ctypes\u001b[38;5;241m.\u001b[39mc_int()\n\u001b[0;32m--> 188\u001b[0m check_call(\u001b[43m_LIB\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mDGLFuncCall\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    189\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mhandle\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalues\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtcodes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mctypes\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mc_int\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnum_args\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    190\u001b[0m \u001b[43m    \u001b[49m\u001b[43mctypes\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbyref\u001b[49m\u001b[43m(\u001b[49m\u001b[43mret_val\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mctypes\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbyref\u001b[49m\u001b[43m(\u001b[49m\u001b[43mret_tcode\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    191\u001b[0m _ \u001b[38;5;241m=\u001b[39m temp_args\n\u001b[1;32m    192\u001b[0m _ \u001b[38;5;241m=\u001b[39m args\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train_tensorboard(model,EPOCHS,loss_fn,train_dataloader,val_dataloader,writer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40c9b593-bcc9-4f7b-b4a5-c40c62d38afd",
   "metadata": {},
   "source": [
    "# Training using wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52b1092a-60ec-4c91-978d-20d16932ae1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "wandb.init(project='GNN_initial_test') #start a run\n",
    "config = wandb.config\n",
    "wandb.watch(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4deb3c29-fb25-4ff0-931a-967496f5f7a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_wandb(model,num_epochs,loss_fn,training_loader,validation_loader):\n",
    "    \"\"\"\n",
    "    Function to train a model, usign experiment tracking wandb.\n",
    "    Arguments:\n",
    "    -----------------\n",
    "        model : GNN model\n",
    "        num_epochs: int\n",
    "        loss_fn: \n",
    "            Loss function\n",
    "        training_loader: Data loader\n",
    "        validation_loader: Data loader\n",
    "    \"\"\"\n",
    "    best_vloss = 1_000_000.\n",
    "    for epoch in range(num_epochs):\n",
    "        print('EPOCH {}:'.format(epoch + 1))\n",
    "\n",
    "        # Make sure gradient tracking is on, and do a pass over the data\n",
    "        model.train(True)\n",
    "        avg_loss = train_one_epoch(model, epoch, loss_fn, training_loader)\n",
    "\n",
    "        # We don't need gradients on to do reporting\n",
    "        model.train(False)\n",
    "\n",
    "        running_vloss = 0.0\n",
    "        for i, vdata in enumerate(validation_loader):\n",
    "            voutputs = model(vdata)\n",
    "            vlabels, vlabels_macro = get_labels_from_graphList(vdata)\n",
    "            vloss = loss_fn(voutputs, vlabels)\n",
    "            running_vloss += vloss\n",
    "\n",
    "        avg_vloss = running_vloss / (i + 1)\n",
    "        print('LOSS train {} valid {}'.format(avg_loss, avg_vloss))\n",
    "\n",
    "        # Log the running loss averaged per batch\n",
    "        # for both training and validation\n",
    "        wandb.log({'Training Loss' : avg_loss, \n",
    "                   'Validation Loss' : avg_vloss, \n",
    "                   'epoch': epoch},\n",
    "                )\n",
    "\n",
    "        # Track best performance, and save the best model's state\n",
    "        if avg_vloss < best_vloss:\n",
    "            best_vloss = avg_vloss\n",
    "            model_path = 'model_best{}'.format(timestamp)\n",
    "            torch.save(model.state_dict(), model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ece20324-a356-45d8-9b37-3628f2a703e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH 1:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/Caskroom/miniforge/base/envs/GrainLearning/lib/python3.9/site-packages/torch/autocast_mode.py:162: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
      "  warnings.warn('User provided device_type of \\'cuda\\', but CUDA is not available. Disabling')\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  batch 1000 loss: 0.15496826105564832\n",
      "  batch 2000 loss: 0.12373793717846274\n",
      "  batch 3000 loss: 0.12284514275565743\n",
      "  batch 4000 loss: 0.12386507965624333\n",
      "  batch 5000 loss: 0.12447568797692657\n",
      "  batch 6000 loss: 0.12517126007378102\n",
      "  batch 7000 loss: 0.1271480363085866\n",
      "  batch 8000 loss: 0.12534524934738875\n",
      "  batch 9000 loss: 0.11943239637091756\n",
      "  batch 10000 loss: 0.12406054601073264\n",
      "  batch 11000 loss: 0.1254925480633974\n",
      "  batch 12000 loss: 0.12396180425211788\n",
      "  batch 13000 loss: 0.1261508788689971\n",
      "  batch 14000 loss: 0.12399742457270622\n",
      "  batch 15000 loss: 0.12615143145248295\n",
      "  batch 16000 loss: 0.12694257579743862\n",
      "  batch 17000 loss: 0.1256946885250509\n",
      "  batch 18000 loss: 0.1232190890610218\n",
      "  batch 19000 loss: 0.1219453156478703\n",
      "  batch 20000 loss: 0.12264493888244032\n",
      "  batch 21000 loss: 0.12337215861305595\n",
      "  batch 22000 loss: 0.1265384682379663\n",
      "  batch 23000 loss: 0.12764782214909792\n",
      "  batch 24000 loss: 0.12391883159801365\n",
      "  batch 25000 loss: 0.12176994782686233\n",
      "  batch 26000 loss: 0.12280244205519557\n",
      "  batch 27000 loss: 0.12561727103218437\n",
      "  batch 28000 loss: 0.12351453598216176\n",
      "  batch 29000 loss: 0.12501726735383273\n",
      "  batch 30000 loss: 0.1250840394757688\n",
      "  batch 31000 loss: 0.11988069978356361\n",
      "  batch 32000 loss: 0.12676749895885586\n",
      "  batch 33000 loss: 0.1212149008475244\n",
      "  batch 34000 loss: 0.125427221596241\n",
      "  batch 35000 loss: 0.12498089714348316\n",
      "  batch 36000 loss: 0.12577721587941051\n",
      "  batch 37000 loss: 0.12228921563923359\n",
      "  batch 38000 loss: 0.12297845617681742\n",
      "  batch 39000 loss: 0.1235311576090753\n",
      "  batch 40000 loss: 0.1232716122455895\n",
      "  batch 41000 loss: 0.12589956513419748\n",
      "  batch 42000 loss: 0.12325141790136696\n",
      "  batch 43000 loss: 0.12469953605532647\n",
      "  batch 44000 loss: 0.12392593870684504\n",
      "  batch 45000 loss: 0.12375503939390183\n",
      "  batch 46000 loss: 0.1261146176457405\n",
      "  batch 47000 loss: 0.12786835538595914\n",
      "  batch 48000 loss: 0.12263902110606432\n",
      "  batch 49000 loss: 0.12115604761242867\n",
      "  batch 50000 loss: 0.12294317320734262\n",
      "  batch 51000 loss: 0.12368847335502506\n",
      "  batch 52000 loss: 0.12473918367549777\n",
      "  batch 53000 loss: 0.12689538627117872\n",
      "  batch 54000 loss: 0.12247466841712594\n",
      "  batch 55000 loss: 0.12497652699053287\n",
      "  batch 56000 loss: 0.12140736382082104\n",
      "  batch 57000 loss: 0.12290642410889267\n",
      "  batch 58000 loss: 0.12521744979172944\n",
      "  batch 59000 loss: 0.12481846584752201\n",
      "  batch 60000 loss: 0.12477297661453486\n",
      "LOSS train 0.12477297661453486 valid 0.12541118264198303\n",
      "EPOCH 2:\n",
      "  batch 1000 loss: 0.12337850755825638\n",
      "  batch 2000 loss: 0.12591436993330718\n",
      "  batch 3000 loss: 0.12479320151731371\n",
      "  batch 4000 loss: 0.12488980619236827\n",
      "  batch 5000 loss: 0.12426028973609209\n",
      "  batch 6000 loss: 0.1242395628541708\n",
      "  batch 7000 loss: 0.12582759673148394\n",
      "  batch 8000 loss: 0.12249622917920351\n",
      "  batch 9000 loss: 0.12600861676782368\n",
      "  batch 10000 loss: 0.12477306272834539\n",
      "  batch 11000 loss: 0.12577148696035148\n",
      "  batch 12000 loss: 0.1264539649039507\n",
      "  batch 13000 loss: 0.12274048418551683\n",
      "  batch 14000 loss: 0.12270843057334423\n",
      "  batch 15000 loss: 0.12403509368747473\n",
      "  batch 16000 loss: 0.12359536886960268\n",
      "  batch 17000 loss: 0.12380667524784804\n",
      "  batch 18000 loss: 0.11976245029643177\n",
      "  batch 19000 loss: 0.12422943462431431\n",
      "  batch 20000 loss: 0.12374801117926836\n",
      "  batch 21000 loss: 0.12543812850862743\n",
      "  batch 22000 loss: 0.12699768985807897\n",
      "  batch 23000 loss: 0.1245639973655343\n",
      "  batch 24000 loss: 0.1250422996096313\n",
      "  batch 25000 loss: 0.12317419398948551\n",
      "  batch 26000 loss: 0.1242206019833684\n",
      "  batch 27000 loss: 0.1227801818549633\n",
      "  batch 28000 loss: 0.12104679480567575\n",
      "  batch 29000 loss: 0.12622661795094609\n",
      "  batch 30000 loss: 0.12683756674826147\n",
      "  batch 31000 loss: 0.12255142204090953\n",
      "  batch 32000 loss: 0.1235113055370748\n",
      "  batch 33000 loss: 0.1274374019354582\n",
      "  batch 34000 loss: 0.12393814677745103\n",
      "  batch 35000 loss: 0.12271419787406922\n",
      "  batch 36000 loss: 0.12298332161828876\n",
      "  batch 37000 loss: 0.1250500760525465\n",
      "  batch 38000 loss: 0.12261295894160867\n",
      "  batch 39000 loss: 0.12249524988606572\n",
      "  batch 40000 loss: 0.12094663958624005\n",
      "  batch 41000 loss: 0.12231772582605481\n",
      "  batch 42000 loss: 0.12593367437645792\n",
      "  batch 43000 loss: 0.12243173095583916\n",
      "  batch 44000 loss: 0.12506796926632524\n",
      "  batch 45000 loss: 0.12317819774523378\n",
      "  batch 46000 loss: 0.12701101091131567\n",
      "  batch 47000 loss: 0.1194572785757482\n",
      "  batch 48000 loss: 0.12538708506524562\n",
      "  batch 49000 loss: 0.12337341235950589\n",
      "  batch 50000 loss: 0.1256310261785984\n",
      "  batch 51000 loss: 0.12496960350126028\n",
      "  batch 52000 loss: 0.12395965887978673\n",
      "  batch 53000 loss: 0.12234001256525516\n",
      "  batch 54000 loss: 0.1288700444623828\n",
      "  batch 55000 loss: 0.12783182503283025\n",
      "  batch 56000 loss: 0.12467896924167872\n",
      "  batch 57000 loss: 0.12382294619455933\n",
      "  batch 58000 loss: 0.12625313194841145\n",
      "  batch 59000 loss: 0.12262469804286957\n",
      "  batch 60000 loss: 0.12554356118664145\n",
      "LOSS train 0.12554356118664145 valid 0.12486094981431961\n",
      "EPOCH 3:\n",
      "  batch 1000 loss: 0.1267492844760418\n",
      "  batch 2000 loss: 0.12124988866224885\n",
      "  batch 3000 loss: 0.12378492916747928\n",
      "  batch 4000 loss: 0.1218448914475739\n",
      "  batch 5000 loss: 0.12313714226707816\n",
      "  batch 6000 loss: 0.12373926699906587\n",
      "  batch 7000 loss: 0.1259697214998305\n",
      "  batch 8000 loss: 0.1255056512504816\n",
      "  batch 9000 loss: 0.12480390736833215\n",
      "  batch 10000 loss: 0.12386708515509963\n",
      "  batch 11000 loss: 0.12307389721274375\n",
      "  batch 12000 loss: 0.12402595705911518\n",
      "  batch 13000 loss: 0.12468054101616144\n",
      "  batch 14000 loss: 0.12355434768274426\n",
      "  batch 15000 loss: 0.12501511226594447\n",
      "  batch 16000 loss: 0.1256584247164428\n",
      "  batch 17000 loss: 0.12468628700077533\n",
      "  batch 18000 loss: 0.12582914667204023\n",
      "  batch 19000 loss: 0.1251026953496039\n",
      "  batch 20000 loss: 0.12225597032532096\n",
      "  batch 21000 loss: 0.12303161270171403\n",
      "  batch 22000 loss: 0.12311408957839012\n",
      "  batch 23000 loss: 0.1254812586121261\n",
      "  batch 24000 loss: 0.12452740606665612\n",
      "  batch 25000 loss: 0.12084807667136192\n",
      "  batch 26000 loss: 0.127914171744138\n",
      "  batch 27000 loss: 0.1251149436943233\n",
      "  batch 28000 loss: 0.12325485156476498\n",
      "  batch 29000 loss: 0.1241563875451684\n",
      "  batch 30000 loss: 0.12467803585156799\n",
      "  batch 31000 loss: 0.12363806174322962\n",
      "  batch 32000 loss: 0.12336360821127891\n",
      "  batch 33000 loss: 0.12735376516357064\n",
      "  batch 34000 loss: 0.12481552628800273\n",
      "  batch 35000 loss: 0.12236618760600686\n",
      "  batch 36000 loss: 0.12262018664181232\n",
      "  batch 37000 loss: 0.12297423890605569\n",
      "  batch 38000 loss: 0.1226554827094078\n",
      "  batch 39000 loss: 0.11968631539866328\n",
      "  batch 40000 loss: 0.124038314063102\n",
      "  batch 41000 loss: 0.12628419122472406\n",
      "  batch 42000 loss: 0.12747920710220934\n",
      "  batch 43000 loss: 0.12671244589984418\n",
      "  batch 44000 loss: 0.12550855991244317\n",
      "  batch 45000 loss: 0.12225767609849572\n",
      "  batch 46000 loss: 0.12269987985119224\n",
      "  batch 47000 loss: 0.12456299379095435\n",
      "  batch 48000 loss: 0.12446066634729505\n",
      "  batch 49000 loss: 0.12395162603259087\n",
      "  batch 50000 loss: 0.12326756182685494\n",
      "  batch 51000 loss: 0.12323537880927324\n",
      "  batch 52000 loss: 0.12536478251963853\n",
      "  batch 53000 loss: 0.12701897514238952\n",
      "  batch 54000 loss: 0.123670763887465\n",
      "  batch 55000 loss: 0.12740181799978018\n",
      "  batch 56000 loss: 0.12124287080019712\n",
      "  batch 57000 loss: 0.12556054605543612\n",
      "  batch 58000 loss: 0.1229668642282486\n",
      "  batch 59000 loss: 0.12410794118791818\n",
      "  batch 60000 loss: 0.12644923799857496\n",
      "LOSS train 0.12644923799857496 valid 0.12489263713359833\n",
      "EPOCH 4:\n",
      "  batch 1000 loss: 0.1258825623765588\n",
      "  batch 2000 loss: 0.125282048933208\n",
      "  batch 3000 loss: 0.12376842402666807\n",
      "  batch 4000 loss: 0.12592085806280376\n",
      "  batch 5000 loss: 0.12149813099205493\n",
      "  batch 6000 loss: 0.12349951053783298\n",
      "  batch 7000 loss: 0.12353512549772859\n",
      "  batch 8000 loss: 0.1262033420279622\n",
      "  batch 9000 loss: 0.12608837826550007\n",
      "  batch 10000 loss: 0.12417587691172957\n",
      "  batch 11000 loss: 0.12497485759109259\n",
      "  batch 12000 loss: 0.12252860714122653\n",
      "  batch 13000 loss: 0.12568679336085917\n",
      "  batch 14000 loss: 0.1243427089639008\n",
      "  batch 15000 loss: 0.12453086640313268\n",
      "  batch 16000 loss: 0.125678003244102\n",
      "  batch 17000 loss: 0.12603500666841866\n",
      "  batch 18000 loss: 0.1270558275319636\n",
      "  batch 19000 loss: 0.12376801667362451\n",
      "  batch 20000 loss: 0.12320682920143008\n",
      "  batch 21000 loss: 0.12210147245228291\n",
      "  batch 22000 loss: 0.12265797811746597\n",
      "  batch 23000 loss: 0.12311532448977232\n",
      "  batch 24000 loss: 0.12154117906093598\n",
      "  batch 25000 loss: 0.12193005976080895\n",
      "  batch 26000 loss: 0.12466374470666051\n",
      "  batch 27000 loss: 0.12570986296981573\n",
      "  batch 28000 loss: 0.1254268309958279\n",
      "  batch 29000 loss: 0.1268746679276228\n",
      "  batch 30000 loss: 0.12453907912224531\n",
      "  batch 31000 loss: 0.12213657311350107\n",
      "  batch 32000 loss: 0.12273425878956915\n",
      "  batch 33000 loss: 0.12583002249151468\n",
      "  batch 34000 loss: 0.12547096455097198\n",
      "  batch 35000 loss: 0.11876430472359062\n",
      "  batch 36000 loss: 0.12352541594952345\n",
      "  batch 37000 loss: 0.12725493463873863\n",
      "  batch 38000 loss: 0.12424309403821826\n",
      "  batch 39000 loss: 0.12562987093627453\n",
      "  batch 40000 loss: 0.1275841730758548\n",
      "  batch 41000 loss: 0.12401417327672243\n",
      "  batch 42000 loss: 0.12543116005137564\n",
      "  batch 43000 loss: 0.12364816131815315\n",
      "  batch 44000 loss: 0.12053742807731033\n",
      "  batch 45000 loss: 0.12426743441075086\n",
      "  batch 46000 loss: 0.12345856242254376\n",
      "  batch 47000 loss: 0.12149448539316654\n",
      "  batch 48000 loss: 0.12268452524021269\n",
      "  batch 49000 loss: 0.12538158874586225\n",
      "  batch 50000 loss: 0.12342513085901738\n",
      "  batch 51000 loss: 0.1266300481520593\n",
      "  batch 52000 loss: 0.12416659922525286\n",
      "  batch 53000 loss: 0.1228267096951604\n",
      "  batch 54000 loss: 0.12286736056208611\n",
      "  batch 55000 loss: 0.1222490508519113\n",
      "  batch 56000 loss: 0.12528828114643692\n",
      "  batch 57000 loss: 0.1269216073229909\n",
      "  batch 58000 loss: 0.1253466901741922\n",
      "  batch 59000 loss: 0.12283193504810333\n",
      "  batch 60000 loss: 0.12333381924778224\n",
      "LOSS train 0.12333381924778224 valid 0.12675829231739044\n",
      "EPOCH 5:\n",
      "  batch 1000 loss: 0.12209192098304629\n",
      "  batch 2000 loss: 0.12265527293086052\n",
      "  batch 3000 loss: 0.12392087165266276\n",
      "  batch 4000 loss: 0.12651339977234602\n",
      "  batch 5000 loss: 0.12718593437969686\n",
      "  batch 6000 loss: 0.12334462484717369\n",
      "  batch 7000 loss: 0.12613103463873268\n",
      "  batch 8000 loss: 0.12291083782538771\n",
      "  batch 9000 loss: 0.12330233150348067\n",
      "  batch 10000 loss: 0.12355072239413857\n",
      "  batch 11000 loss: 0.1253525351025164\n",
      "  batch 12000 loss: 0.12582643768191337\n",
      "  batch 13000 loss: 0.12489730352163315\n",
      "  batch 14000 loss: 0.12482250244542957\n",
      "  batch 15000 loss: 0.12416515523940325\n",
      "  batch 16000 loss: 0.12541240277141333\n",
      "  batch 17000 loss: 0.12281225112453104\n",
      "  batch 18000 loss: 0.12364689292386175\n",
      "  batch 19000 loss: 0.12490173096582294\n",
      "  batch 20000 loss: 0.12319588858261704\n",
      "  batch 21000 loss: 0.1223425467275083\n",
      "  batch 22000 loss: 0.12201417791470885\n",
      "  batch 23000 loss: 0.12476051906123757\n",
      "  batch 24000 loss: 0.12704875804111362\n",
      "  batch 25000 loss: 0.12440832203999162\n",
      "  batch 26000 loss: 0.12250706413760781\n",
      "  batch 27000 loss: 0.1255154955163598\n",
      "  batch 28000 loss: 0.12338698727637529\n",
      "  batch 29000 loss: 0.12378503259271384\n",
      "  batch 30000 loss: 0.12398687796667218\n",
      "  batch 31000 loss: 0.12593983556702734\n",
      "  batch 32000 loss: 0.12443839297816157\n",
      "  batch 33000 loss: 0.12350292140990496\n",
      "  batch 34000 loss: 0.12547913343831896\n",
      "  batch 35000 loss: 0.1264746130183339\n",
      "  batch 36000 loss: 0.12252803146466613\n",
      "  batch 37000 loss: 0.12392948725819587\n",
      "  batch 38000 loss: 0.12316719190403819\n",
      "  batch 39000 loss: 0.12340258066356181\n",
      "  batch 40000 loss: 0.12308789519593119\n",
      "  batch 41000 loss: 0.12313237976655364\n",
      "  batch 42000 loss: 0.12683557751402258\n",
      "  batch 43000 loss: 0.12410974799841643\n",
      "  batch 44000 loss: 0.12456187699735165\n",
      "  batch 45000 loss: 0.12041464603319764\n",
      "  batch 46000 loss: 0.12299171421304345\n",
      "  batch 47000 loss: 0.12506777531653643\n",
      "  batch 48000 loss: 0.12221631077304483\n",
      "  batch 49000 loss: 0.12602249264344573\n",
      "  batch 50000 loss: 0.12225242717936635\n",
      "  batch 51000 loss: 0.12612645713984966\n",
      "  batch 52000 loss: 0.12326804797351361\n",
      "  batch 53000 loss: 0.12759554669633508\n",
      "  batch 54000 loss: 0.12784206495806574\n",
      "  batch 55000 loss: 0.12516228882223368\n",
      "  batch 56000 loss: 0.12187858945131302\n",
      "  batch 57000 loss: 0.12236075780168176\n",
      "  batch 58000 loss: 0.12156338589265943\n",
      "  batch 59000 loss: 0.12656260067224503\n",
      "  batch 60000 loss: 0.12533757996559142\n",
      "LOSS train 0.12533757996559142 valid 0.12553371489048004\n"
     ]
    }
   ],
   "source": [
    "train_wandb(model,5,loss_fn,train_dataloader,val_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ea59ef9-a394-4a30-baa1-9f65c09a3bfd",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Trainning loop of Brandstetter et al."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "665123dc-99b7-49dd-bb89-a2dc46b26faa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#TODO\n",
    "#def training_loop(model,batch_size,optimizer,loss_fn):\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d28af5a-f5b2-49eb-a3d7-d465399146a6",
   "metadata": {},
   "source": [
    "# Loading a model\n",
    "First define a model in: [Creation of the model](#creation_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "edf2bb21-c75b-4bb9-8cef-a3853708493f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "model_gnn_heterogeneous(\n",
       "  (dropout): Dropout(p=0, inplace=False)\n",
       "  (layers): ModuleList(\n",
       "    (0): myGnn_heterogeneous(\n",
       "      (mlp_1): Linear(in_features=9, out_features=9, bias=True)\n",
       "      (mlp_2): Linear(in_features=20, out_features=9, bias=True)\n",
       "    )\n",
       "    (1): myGnn_heterogeneous(\n",
       "      (mlp_1): Linear(in_features=9, out_features=9, bias=True)\n",
       "      (mlp_2): Linear(in_features=20, out_features=9, bias=True)\n",
       "    )\n",
       "    (2): Linear(in_features=9, out_features=3, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Run first cell defining the model\n",
    "PATH = \"./model_20220815_113137_1\"\n",
    "model.load_state_dict(torch.load(PATH))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74038ec5-3797-4960-a32c-d597fdd2f42f",
   "metadata": {},
   "source": [
    "## Plot model predictions on the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "2d0ee586-88ab-46a0-b960-bcc4467c1689",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "def plot_predictions(test_loader):\n",
    "    %matplotlib inline\n",
    "    fig, ax = plt.subplots(1, 3)\n",
    "    prediciton,label= ([] for _ in range(2))\n",
    "    for i, test_data in enumerate(test_loader):\n",
    "        if i>=1: break\n",
    "        prediciton_i = model(test_data).detach().numpy()\n",
    "        label_i, labels_macro = get_labels_from_graphList(test_data)\n",
    "        label_i = label_i.detach().numpy()\n",
    "        #print(prediciton_i[:,0].shape)\n",
    "        #print(type(label_i[:,0]))\n",
    "        for j in range(3):\n",
    "            ax[j].plot(label_i[:,j],prediciton_i[:,j]) \n",
    "            #TODO check structure of labels and predictions\n",
    "    \n",
    "    #labels of the plot\n",
    "    for j in range(3): \n",
    "        ax[j].set_xlabel('truth')\n",
    "        ax[j].set_ylabel('prediction')\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8d2ab1cb-a598-4cd2-b4f6-b8a0cdd9ea0d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAagAAAEYCAYAAAAJeGK1AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAApd0lEQVR4nO3df7hcVX3v8fdHAsEAtyQm/JAQUhGuF4gPwhHsrTxyNQSrlFCxVM3V8BSaUqFwtRjCRUsKxMZci/YWam/gQmMpIqXtJSlqmqSmVmsoJxgSiGL4aRL5kRBEgoIYvvePvQ7ZGWbOmTOzZ2bPzOf1PPuZ2WvW3rN28n3Od/bea62tiMDMzKxsXtfpBpiZmVXjBGVmZqXkBGVmZqXkBGVmZqXkBGVmZqU0ptMNKLOJEyfG1KlTO92Mtlu7du32iJjU6XZ0s36MHcdN8/oxbqB27DhBDWPq1KkMDg52uhltJ+nxTreh2/Vj7DhumtePcQO1Y8eX+MzMrJScoMysb0maIGmFpE3pdXyNerNTnU2SZufKV0t6UNK6tBxUsd3ZkkLSQKuPpRc5QZlZP5sHrIqIo4BVaX0PkiYAVwInAycBV1YkslkRcXxans5tdwBwCXB3Kw+glzlBmVk/mwksSe+XAGdVqXM6sCIidkTEs8AK4L117Ptq4HPAiwW0sy85QZlZPzs4Ip5I758EDq5S5zBgc259SyobcnO6vPcZSQKQdAJweETcNVIDJM2RNChpcNu2bY0dRY9yLz4z62mSVgKHVPnoivxKRISk0c6ePSsitqbLeX8PfFTSLcC1wLn17CAiFgOLAQYGBjx7d44TlJn1tIiYXuszSU9JOjQinpB0KPB0lWpbgVNz65OB1WnfW9Pr85JuJbtHdSdwHLA6nVAdAiyVdGZE9F8f8ib4Ep+Z9bOlwFCvvNlkyaXScmCGpPGpc8QMYLmkMZImAkjaGzgDuD8inouIiRExNSKmAmsAJ6cGOEGZWT9bCJwmaRMwPa0jaUDSjQARsYOsw8M9abkqlY0lS1TrgXVkZ1o3tP0Iepgv8ZlZ34qIZ4D3VCkfBM7Prd8E3FRR5wXgxDq+49SmG9qnfAZlZmal1JEEVcDo7QWSNkvaWVH/AkkbUpfPb0s6JpWfJmlt+mytpHe39gjNzKxZnTqDanb09rJUVunWiJgWEccDi8i6egJsB34zIqaR3Qj9mwKPxczMWqBTCaqp0dsRsSY3uO5VEfHT3Op+QKTy70XEj1P5A8DrJY0t4kDMzKw1OtVJoojR21VJuhD4JLAPUO1S3tnAvRHxUo3t5wBzAKZMmTLS15mZWYu07AxK0kpJ91dZZubrRUSQznSKEBHXR8SRwGXApyvadCzZ3Fi/P8z2iyNiICIGJk3ys9fMzDqlZWdQrRy9XafbgC/lvnMy8I/AxyLi4VHsx8zMOqBT96AaHr093E4lHZVbfT+wKZUfCNwFzIuI7zTXdDMza4dOJahmRm8jaZGkLcA4SVskzU/7vUjSA5LWkd2HGkqCFwFvBv641oPFzMysXDrSSaKZ0dupfC4wt0r5JTW+7xrgmiaabCWRhh98FZgKPAack3p5Vtabze57kNdExJJUvgD4GDA+IvbP1b8AuBDYBewE5kTExtznU4CNwPyI+HzxR2ZmlTyThHWbdo+hG3It8PVCjsDM6uIEZd2mrWPoACSdBTxKNobOzNrEk8Vat2nrGDpJ+5MNWTgNuHSE7T2GzqxAPoOy0inZGLr5wBciYmfNDXdv7zF0ZgXyGZSVTsnG0J0MfFDSIuBA4BVJL0bEdaPYn5k1wGdQ1m3aOoYuIk7JPRn1i8BnnZzM2sMJyrpNu8fQmVmH+BKfdZV2j6GrqDN/lM01syb4DMrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrM+pakCZJWSNqUXsfXqDc71dkkaXaufLWkByWtS8tBqfxcSdty5edX268NzwnKzPrZPGBVRBwFrErre5A0AbiS7NlgJwFXViSyWRFxfFryzyf7aq78xhYeQ89ygjKzfjYTWJLeLwHOqlLndGBFROyIiGeBFcB729O8/taRBFXAafUCSZsl7ayof4GkDemU+tuSjqn4fIqknZIubc2RmVmXOTginkjvnwQOrlLnMGBzbn1LKhtyc/qb8xlJypWfLWm9pDskHV6rAZLmSBqUNLht27ZGj6MndeoMqtnT6mWprNKtETEtIo4HFgHXVnx+LfD1Qo7AzLqCpJWS7q+yzMzXi4gAYpS7nxUR04BT0vLRVL4MmBoRbyU741pSY3siYnFEDETEwKRJk0b59b2tUwmqqdPqiFiT+9Xzqoj4aW51P3LBJuks4FHggQLab2ZdIiKmR8RxVZY7gackHQqQXp+usoutQP4MaHIqIyKGXp8HbiX9cI6IZyLipVT/RuDEVhxbr+tUgiritLoqSRdKepjsDOriVLY/cBnwJ3Vs79Nts/6xFBi6fTAbuLNKneXADEnj01WcGcBySWMkTQSQtDdwBnB/Wj80t/2ZwPdb1P6e1rJHvktaCRxS5aMr8isREZJGe1pdU0RcD1wv6SPAp8mCbj7whYjYuecl4qrbLwYWAwwMDBTWLjMrpYXA7ZLOAx4HzgGQNABcEBHnR8QOSVcD96Rtrkpl+5Elqr2BvYCVwA2pzsWSzgR+CewAzm3bEfWQliWoiJhe6zNJT0k6NCKeGOG0+tTc+mRg9SiacBvwpfT+ZOCDkhYBBwKvSHoxIq4bxf7MrMdExDPAe6qUDwLn59ZvAm6qqPMCNS7dRcTlwOWFNrYPdeoSX8On1cPtVNJRudX3A5sAIuKUiJgaEVOBLwKfdXIyMyu3TiWohcBpkjYB09M6kgYk3QgQETuAodPqe0in1aneIklbgHGStkian/Z7kaQHJK0DPsnuJGhmZl2mZZf4htPMaXUqnwvMrVJ+SR3fPX+UzbUSScMPvgpMBR4Dzkm9PCvrzSa7BwlwTUQsSeULgI8B4yNi/1z9C4ALgV3ATmBORGxMn70V+D/AfwJeAd4eES+24vjMbDfPJGHdpq1j6CSNAW4hu2F+LNl90ZeLPCAzq84JyrpNu8fQzQDWR8R9qd4zEbGriAMxs+F15BKfWRNaOoaO7N7lPsC7U/HRQEhaDkwCbouIRTW2nwPMAZgyZcrIR2Jmw/IZlJVOi6emqSkiro+II8kGdQ/dvxoDvBOYlV5/S9Jr7p+m7T1ljVmBfAZlpVOyMXRbgG9FxPb0/V8DTiC7/2VmLeQzKOs2bR1Dl7abJmlc6jDxLmBjE+03szo5QVm3aesYutTJ4tq0n3XAvRFxVzsO1Kzf+RKfdZVOjKGLiFvIupqbWRv5DMrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzEqpIwlK0gRJKyRtSq/ja9SbnepskjQ7V75A0mZJOyvqXyBpg6R1kr4t6ZjcZ2+V9N30zJ8NkvZt3RGamVmzOnUGNQ9YFRFHkT06e15lBUkTgCuBk4GTgCtziWxZKqt0a0RMi4jjgUVkD5ojPQn1FuCCiDiW7HHgLxd5QGZmVqxOJaiZwJL0fglwVpU6pwMrImJHeqrpCuC9ABGxJiKeqNwgIn6aW90PiPR+BrA+Iu5L9Z6JiF1FHIiZmbVGpxLUwbkE8yRwcJU6hwGbc+tbUtmwJF0o6WGyM6iLU/HRQEhaLuleSa95ompu+zmSBiUNbtu2rZ5jMbMuVcDthtWSHky3FdZJOij32TmSNqbbCre243h6TcsSlKSVku6vsszM14uIYPeZTtMi4vqIOBK4DPh0Kh4DvBOYlV5/S9JrHhuetl8cEQMRMTBp0qSimmVm5dTs7QaAWRFxfFqeTtscBVwO/Hq6rfA/WnsYvallCSoipkfEcVWWO4GnJB0KkF6frrKLrcDhufXJqaxet7H70uEW4FsRsT0ifgZ8DThhlIdkZr2nqdsNw/g94PpUn6HEZaPTqUt8S4Gh0+TZwJ1V6iwHZkgan36tzEhlNaVfLUPeD2zK7WuapHGpw8S7gI1NtN/MekMRtxtuTpf3PiNJqexo4GhJ35G0RlLNhObbCrV1KkEtBE6TtAmYntaRNCDpRoCI2AFcDdyTlqtSGZIWSdoCjJO0RdL8tN+L0vXedcAnSUkw/Yq5Nu1nHXBvRNzVjgO1YrV7iIKkvSUtSZ99X9LlrT1CK1qLbzfMiohpwClp+WgqHwMcRdZj+MPADZIOrLYD31aobUwnvjQingFecw8oIgaB83PrNwE3Vak3F3hNR4eIuGSY77yFrKu5dbehewYLJc1L65flK+TuGQyQ/cFZK2lp+qGyDLiO3WfXQ26NiL9K259J9oPmvcBvA2MjYpqkccBGSV+JiMdadoRWqIiYXuszSU9JOjQinhjhdsOpufXJwOq0763p9fnUEeIk4MtkZ1l3R8TLwKOSfkiWsO5p/oj6h2eSsG7T7iEKAeyXLg2/HvgF8FOsVzR8u0HSGEkTITvTBs4A7k/b/D9SUkt1jgYeadEx9KyOnEGZNaGlQxTILg3vA7w7Fd9BlhSfAMYBnxi61Fxl+znAHIApU6aMeCBWCguB2yWdBzwOnAPZ7Qaygf3nR8QOSUO3GyDdbpC0H1mi2hvYC1gJ3JDqDCW1jcAu4FPpypGNghOUlY6klcAhVT66Ir8SESGp0CEKwPWSPkI2RGE22SWbXcAbgfHAv0laGRGv+TUcEYuBxQADAwOFtctap5nbDRHxAnBijf0G2Y+dTxbZ3n7jBGWl08p7BnW6DfhSev8R4BvpXsLTkr5Ddm/Ll2vMWsz3oKzbtHuIwo9Il/vSJZ13AD9ouPVmVjcnKOs2bR2iAFwP7C/pgbSvmyNifRuO06zv+RKfdZV2D1GIiJ1kXc3NrM18BmVmZqXkBGVmZqXkBGVmZqXkBGVmZqXkBGVmZqVUdy8+SYcBR+S3iYhvtaJR1jscN1Y0x1T/qCtBSfoc8Dtkz1DalYoDcFBYTY4bK5pjqr/UewZ1FvCfI+KlFrbFes9ZOG6sWGfhmOob9d6DegTYu5UNsZ7kuLGiOab6SL1nUD8D1klaBbz6yyUiLm5Jq6xXOG6saI6pPlJvglqaFrPRcNxY0RxTfaSuBBURSyTtQ/ZUSIAH0+MHzGpy3FjRHFP9pd5efKeSPV77MUDA4ZJmu2unDcdxY0VzTPWXei/x/RkwIyIeBJB0NPAVajxN0ixx3FjRHFN9pN5efHsPBQRARPyQJnrSSJogaYWkTel1fI16s1OdTZJm58oXSNosaWdF/QskbZC0TtK3JR2TyveWtCR99n1JlzfadhuVQuPGDMdUX6k3QQ1KulHSqWm5ARhs4nvnAasi4ihgVVrfg6QJwJXAycBJwJW5RLYslVW6NSKmRcTxwCLg2lT+28DYiJhG9kvr9yVNbaL9Vp+i48bMMdVH6k1Qf0A2cvvitGxMZY2aSXYdmfR6VpU6pwMrImJHRDwLrADeCxARayLiicoNIuKnudX9yEaYk173kzQGeD3wC+CnWKsVHTdmjqk+Um8vvpfIzkauHalunQ7OJZgngYOr1DkM2Jxb35LKhiXpQrJHdu8DvDsV30GWFJ8AxgGfGHoEuLVOC+LG+pxjqr8Mm6Ak3R4R50jawO6zkVdFxFuH2XYlcEiVj66o2EdIes2+GxUR1wPXS/oI8GlgNtnlwF3AG4HxwL9JWhkRj1Rp9xxgDsCUKVOKalZfaSZuzKpxTPWnkc6gLkmvZ4x2xxExvdZnkp6SdGhEPCHpUODpKtW2Aqfm1icDq0fRhNuAL6X3HwG+kcZLPC3pO8AA2bQple1eDCwGGBgYKCxx9pmG48asBsdUHxr2HlTuMtzHI+Lx/AJ8vInvXUp2ZkN6vbNKneXADEnjU+eIGamsJklH5VbfD2xK739EutwnaT/gHcAPGm69DauFcWN9yjHVn+rtJHFalbLfaOJ7FwKnSdoETE/rSBqQdCNAukd0NXBPWq4aum8kaZGkLcA4SVskzU/7vUjSA5LWkd2HGkqC1wP7S3og7evmiFjfRPutPkXHjZljqo+MdA/qD8h+nRwpKf8H/QDg3xv90oh4BnhPlfJB4Pzc+k3ATVXqzQXmVim/pLIsle8k62pubdCquEn7ngB8FZhKNpvAOamXZ2W92WT3IAGuiYglqXwB8DFgfETsX2W7s8k61bw9xSNp3Nx5ZPcxL46IYc/krXitjCkrr5HuQd0KfB34U/Ycq/S8e8HZMFoZN0Nj6BZKmpfWL8tXyI2hGyC7ob5W0tKUyJYB17H78m9+uwPI7nXcnSs7BvgQcCxZJ5uVko6OiF2V21tL+W9RHxrpHtRzEfEY8OfAjtw1319KOrkdDbTu0+K4ackYuuRq4HPAixXfd1tEvBQRjwIPUX2QuLWQ/xb1p3rvQX0JyE8rtJPdPeTMamlF3LRkDJ2kE4DDI+KuRvclaY6kQUmD27ZtG+7rrHGFxlQB066tlvRgml5tnaSDUvkXcmU/lPSTRtvYz+qdLFYR8WqX64h4Jc3K0Bdu+vajXPVPG0est//YMbz7LQe9ur7jhV/w7Ye2t7JphXrks+/jda9TkbtsKG7aPYZO0uvIBn6e28x+Koco3PujZ/nAX/b+7ZFVf/Qujpz0mtt5rVL036JmLxkDzBq6X5lr1ydy2/8h8LZ6G/T8iy8zbf4/N3Qw1Zx2zMG8fu+9an7+k5+/zLd+2P4fVH/+oeOZefzwcy/U+x/7iKSL2f1L5eNUGUPUq9Zv+Uld9Xa+9Es2bH3u1fUnn3txmNrl04JBXw3FTQfG0B0AHAeslgRZclwq6cy0r8Mr9rV1pGMA2Prsz+up1vW2P/9SOxNU0X+LZrI7VpaQxcllFXVevWQMIGnokvFX6vyOD5MluLq89MtX6q1al/s2/4T9xtb+U7/9+ZdqftZKP3jyeWaOVCkiRlyAg8gGvj4NPEV2w/Kgerbt5uXEE0+MfgQMRgH/fq2IG+B/AfPS+3nAoip1JgCPks0aMj69n1BRZ+cw37EaGEjvjwXuA8YCv0r2x3CvkdrZj7FTVNwMtxQdU8BPcu+VX8+VXwp8Orf+GeDS2B0rG4B1qVwV2x5BNsVazZghm7lmEBicMmVKcf8hXaRW7NQ7F9/TZD2ZzOrWorhZCNwu6TzgceAcyMbQARdExPkRsUPS0Bg6qBhDRzazyLg0lu7GiJg/zDE8IOl2sklJfwlcGO7B1zGNxFSLLxnPioitqQfo3wMfBb6c+/xDwB3DxUx49pqaRhoHNTciFkn6C6rPf3Vxy1pmXauVcRMtGkNXUefUivUFwILGWmxFaCamooWXjCNia3p9XtKtZD08KxPUhbW+34Y30hnU99Orn7dio+G4saK1KqaGpl1byPDTrn0218NvBnB56pxxYERsl7Q32TyBK4c2kvQWskvM3y24zX1j2AQVEcvS65Lh6pnlOW6saC2MqYYvGad5PZen5LQXWXK6IbfvD5GNofNluwaNdIlvGcN07oqIMwtvkXU9x40VrVUx1cwl44h4gewJ3bX2Pb+RNtluI13i+3x6/QDZTcZb0vqHyXrQmFXjuLGiOab60EiX+P4VQNKfRcRA7qNlknx/wapy3FjRHFP9qd6pjvaT9KahFUm/CuzXmiZZD3HcWNEcU32k3pkkPkE2yv4RssFsRwC/37JWWa9w3FjRHFN9pN6But9IT6t9Syr6QUR0Zn4M6xqOGyuaY6q/1HWJT9I44FPARRFxHzBF0hktbZl1PceNFc0x1V/qvQd1M/AL4NfS+lbgmpa0yHqJ48aK5pjqI/UmqCMjYhHwMkBE/Izs+q/ZcBw3VjTHVB+pN0H9QtLrSQPlJB0J+LqvjcRxY0VzTPWRenvxXQl8Azhc0t8Cv06TD3ezvuC4saI5pvpIPU83fR3ZhIcfAN5Bdjp9SUR0z6Nire0cN1Y0x1T/GfESX0S8AsyNiGci4q6I+KdmA0LSBEkrJG1Kr+Nr1Jud6mySNDtXvkDSZkk7a2x3tqRIEz4OlV0u6SFJD0o6vZn228haETfW3xxT/afee1ArJV0q6fCUXCZImtDE984DVkXEUcCqtL6HtP8rgZPJnrFyZS6RLUtlr5EeHHYJcHeu7BiymYWPJXtU819K2quJ9lt9io4bM8dUH6n3HtTvkN2U/HhF+Zuq1K3HTHY/AGwJ2cO/LquoczqwIvck1BVkyeUrEbEmlVXb99XA58jGSuS/77Y0oO9RSQ+RJTg/p6W1io4bM8dUH6n3DOoY4HrgPmAd8BdkZyONOjginkjvnwQOrlLnMGBzbn1LKqtJ0gnA4RFxV6P7kjRH0qCkwW3btg33dTayouPGzDHVR+o9g1oC/BT432n9I6nsnFobSFpJNi1+pSvyKxERkpp+oFe6gXotTfboiYjFwGKAgYEBP2isOaOOG7MROKb6SL0J6riIOCa3/k1JG4fbICKm1/pM0lOSDo2IJyQdCjxdpdpWdl8GBJhMdimwlgOA48gmkoQsOS6VdGba1+EV+9o6XPutEKOOG7MROKb6SL2X+O6V9I6hFUknA808g2UpMNQrbzZwZ5U6y4EZksanzhEzUllVEfFcREyMiKkRMRVYA5yZnoy5FPiQpLFpev6jgP9oov1Wn6Ljxswx1UfqTVAnAv8u6TFJj5F1Lni7pA2S1jfwvQuB0yRtAqandSQNSLoRIHWOuBq4Jy1X5TpMLJK0BRgnaYuk+cN9WUQ8ANwObCQb5HdhROxqoN02OkXHTduHKEg6TdLa1Oa1kt7dSLutMIXHlJWXIka+zSLpiOE+j4jHC2tRiQwMDMTgYP/9OJO0tuKppY3up/C4kbQI2BERCyXNA8ZHxGUVdSaQ/aoeIOvxtRY4MSKeTb++Hwc2RcT+FdsdANwF7EM2W/agpLcBT0XEjyUdByyPiGE760B/xk5RcTPCd/T036J+jBuoHTv1Pg+qq//TrTNaFDdtHaIQEd/Lff4A8HpJY/0Mos7w36L+Uu8lPrOyaPcQhbyzgXtrJScPUTArVr29+MzapoxDFCQdS3Z2NaNWHQ9RMCuWE5SVTpmGKKT7UJOBfwQ+FhEPj+pgzKxhvsRn3aatQxQkHUjWcWJeRHynwOMwsxE4QVm3aesQBeAi4M3AH0tal5aDWnFgZrYnX+KzrhIRzwDvqVI+CJyfW78JuKlKvbnA3BG+49Tc+2uAaxpvsZk1ymdQZmZWSk5QZta3CpiZZLWyh6DucflX0hRJ35T0PUnrJb2vXcfUS5ygzKyfNfvwVIBZEXF8WoZ6lX4auD0i3kb2sNS/bOVB9ConKDPrZzPJZiQhvZ5Vpc6rM5NExLPA0MwkwwngP6X3vwL8uPmm9h93kjCzflbEzCQ3S9oF/D1wTWQTnM4H/lnSHwL7kfU4tVHyGZSZ9TRJKyXdX2WZma+XEstoZwCZFRHTgFPS8tFU/mHgryNiMvA+4G/SjCXV2ucpsmrwGZSZ9bRWzkwSEVvT6/OSbiW7R/Vl4DzSZcCI+K6kfYGJ1fbvKbJq8xmUmfWzhmcmkTRG0kQASXsDZwD3p21+RBqvJ+m/APsCPj0aJScoM+tnzcxMMpYsUa0H1pGdad2Q9vtHwO9Jug/4CnBu1PPwPduDL/GZWd9qZmaSiHiB7Am/1fa7Efj1Qhvbh3wGZWZmpeQEZWZmpeQEZWZmpeQEZWZmpeQEZWZmpdSRBFXADMILJG2WtLPGdmdLCkkDaf00SWslbUiv727NkZmZWVE6dQbV7AzCy1LZa0g6ALgEuDtXvB34zTQlyWzgbwo6DjMza5FOJaimZhCOiDW5CR4rXQ18DnhxqCAivhcRQ7MJPwC8XtLYpo/CzMxaplMJqogZhF9D0gnA4RFx1zDVzgbujYiXauzDEzeamZVAy2aSkLQSOKTKR1fkVyIiJDU9BUiaKfha4Nxh6hxLdnY1o1YdT9xoZlYOLUtQrZxBuIYDgOOA1ZIgS45LJZ0ZEYOSJgP/CHwsIh4e1cGYmVnbdeoSX8MzCNfaYUQ8FxETI2JqREwF1gBDyelA4C5gXkR8p8DjMDOzFulUgmpmBmEkLZK0BRgnaYuk+SN830XAm4E/lrQuLQe14sCstdo9RCFXPkXSTkmXFntEZlZLR2Yzb2YG4VQ+F5g7wnecmnt/DXBN4y22EhkaorBQ0ry0flm+Qm6IwgDZE1LXSlqaeoMuA64DNlXuuMYQhSHXAl8v8kDMbHieScK6TVuHKABIOgt4lGyIgpm1iROUdZu2DlGQtD/ZGdqfjNQwD1EwK5YfWGilU7IhCvOBL0TEztQ7tCYPUTArlhOUlU6ZhiiQTbX1QUmLgAOBVyS9GBHX1X9EZtYIX+KzbtPWIQoRcUqu/IvAZ52czNrDCcq6TbuHKJhZh/gSn3WVdg9RqCifP7rWmlkzfAZlZmal5ARlZmal5ARlZmal5ARlZmal5ARlZmal5ARlZmal5ARlZmal5ARlZmal5ARlZn2rgAdgrpb0YOWDUCUdIWmVpPWpzuR2HVMvcYIys3429ADMo4BVaX0PuQdgngycBFxZkchmRcTxaRmavPjzwJcj4q3AVcCftvIgepUTlJn1s6YegDmMY4B/Se+/mb7HRskJysz6WREPwLw5Xd77jHY/NOw+4APp/W8BB0h6Q7UG+EGXtTlBmVlPk7RS0v1Vlj3OaiIigNE+aHJWREwDTknLR1P5pcC7JH0PeBfZM8p2VdtBRCyOiIGIGJg0adIov763eTZzM+tprXwAZkRsTa/PS7qV7B7VlyPix6QzKEn7A2dHxE+aP5r+4jMoM+tnDT8AU9IYSRMBJO0NnAHcn9YnShr6+3o5VR79YiPrSIIqoGvnAkmbJe2ssd3ZkkLSQEX5FEk7JV1a7BGZWZdq5gGYY8kS1XpgHdmZ1g1pv6cCD0r6Idl9rQXtOqBe0qlLfENdOxdKmpfWL8tXyHXtHCC7LrxW0tLUi2YZcB2wqXLHkg4ALgHurvK91wJfL/JAzKx7NfMAzIh4ATixxn7vAO4otLF9qFOX+Jrq2hkRa3I9bypdDXwOeDFfKOks4FHggWYbb2ZmrdepBFVE187XkHQCcHhE3FVRvj/ZGdqfjNQwd/k0MyuHll3ik7QSOKTKR1fkVyIiJI22a2e173sd2SW8c6t8PB/4QkTs3D1MobqIWAwsBhgYGGi6XWZm1piWJahWdu2s4QDgOGB1SkKHAEslnUk2RckHJS0CDgRekfRiRFxX/xGZmVk7deoSX8NdO2vtMCKei4iJETE1IqYCa4AzI2IwIk7JlX8R+KyTU3fqRA9QSW+V9F1JD0jaIGnf4o/MzCp1KkE107UTSYskbQHGSdoiaX4HjsE6o9nJPZelsteo1gNU0hjgFuCCiDiW7Kz+5aIOxsxq60g382a6dqbyucDcEb7j1Brl80fXWiuZmey+9LuE7LLvZRV1Xu0BCiBpqAfoVyJiTSqrtu+hHqCfypXNANZHxH3wauyaWRt4JgnrNm3tAQocDYSk5ZLulVTzh5F7gJoVy3PxWemUrAfoGOCdwNuBnwGrJK2NiFWVFd0D1KxYTlBWOiXrAboF+FZEbE/f/zXgBLL7X2bWQr7EZ92mrT1A03bTJI1LHSbeBWws7nDMrBYnKOs2be0BmqbZujbtZx1wb5X7VGbWAr7EZ12lEz1AI+IWsq7mZtZGPoMyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIyM7NScoIys74laYKkFZI2pdfxNerNTnU2SZqdK99H0mJJP5T0A0lnp/Kxkr4q6SFJd0ua2qZD6ikdSVAFBMUCSZsl7ayx3dmSQtJAruytkr4r6QFJGyTtW/yRmVmXmQesioijgFVpfQ+SJgBXAicDJwFX5v5mXQE8HRFHA8cA/5rKzwOejYg3A18APtfSo+hRnTqDajYolqWy15B0AHAJcHeubAzZE1EviIhjgVOBl4s6GDPrWjOBJen9EuCsKnVOB1ZExI6IeBZYAbw3ffa7wJ8CRMQrEbG9yn7vAN4jScU3v7d1KkE1FRQRsSYinqix76vJfq28mCubAayPiPvS9s9ExK6mj8LMut3Bub8lTwIHV6lzGLA5t74FOEzSgWn9akn3Svo7SQdXbhMRvwSeA95QdON7XacSVMNBMdxOJZ0AHB4Rd1V8dDQQkpanQJo7zD7mSBqUNLht27YRD8TMyk3SSkn3V1lm5utFRAAxil2PASYD/x4RJwDfBT7fQPv8N6eGMa3asaSVwCFVProivxIRIWk0QVHr+14HXAucW+XjMcA7gbcDPwNWSVobEasqK0bEYmAxwMDAQNPtMrPOiojptT6T9JSkQyPiCUmHAk9XqbaV7LbAkMnAauAZsr8n/5DK/47s3tPQNocDW9Ithl9J9au1z39zamjZGVRETI+I46osdwJPpWBghKA4PLc+OZXVcgBwHLBa0mPAO4ClqaPEFuBbEbE9In4GfA04odljNLOutxQY6oA1G7izSp3lwAxJ49N98BnA8nTGtYzdyes9wMYq+/0g8C+pvo1Cpy7xNRwUtXYYEc9FxMSImBoRU4E1wJkRMZi2myZpXPo18y52B5KZ9a+FwGmSNgHT0zqSBiTdCBARO8jubd+TlqtSGcBlwHxJ64GPAn+Uyv8v8AZJDwGfpEpHMBuZOpHUJb0BuB2YAjwOnBMRO9LZzgURcX6q97vA/0ybLYiIm1P5IuAjwBuBHwM3RsT8iu9YDVyaEhSS/jtwOdk15q9FRM37ULl9bEvtG85EYPsIdbrJRGC/iJjU6YZ0sxFipxdjZjtwhOOmOcPETZljpoi2VY2djiSoXiJpMCIGRq7ZHXrteMqo1/6Ne+14yqjM/8atbJtnkjAzs1JygjIzs1Jygmre4k43oGC9djxl1Gv/xr12PGVU5n/jlrXN96DMzKyUfAZlZmal5ARlZmal5ARVJ0nvlfRger5LtdnXu+b5L3UcyyclbZS0XtIqSUd0op3drpdiBhw3rdJMnEi6PJU/KOn0Nrer5v+3pF2S1qVlacONiAgvIyzAXsDDwJuAfYD7gGMq6nwc+Kv0/kPAVzvd7iaO5b8B49L7PyjrsZR56aWYcdx0/N+1apyQPX/qPmAs8KtpP3uV4f8b2FlEO3wGVZ+TgIci4pGI+AVwG9kjQ/K65fkvIx5LRHwzsjkLIZsyanKb29gLeilmwHHTKs3EyUzgtoh4KSIeBR6ixnPyWtGudvx/O0HVp55Hf3TL819G+xiT84Cvt7RFvamXYgYcN63STJyM+pFEBbcrr/L/e9/0CJE1ks5qtBEte9yGdb80f+EA2eS6ZnVx3PSXGv/fR0TEVklvAv5F0oaIeHi0+/YZVH3qefTHq3VGev5Lh9X1GBNJ08me3XVmRLzUprb1kl6KGXDctEozcTLaRxIV3a6a/98RsTW9PkL27Ky3NdSKTt8k7IaF7EzzEbIbkUM3DI+tqHMhe97IvL3T7W7iWN5GdoP0qE63t1uXXooZx03H/12rxglwLHt2kniE4jpJNPz/DYwHxqb3E4FNVHSwqLsdnf4P6pYFeB/ww/QfckUqu4rslwPAvmRP1HwI+A/gTZ1ucxPHshJ4CliXlqWdbnM3Lr0UM46bcsYJ2dnLw8CDwG+U4f8b+K/AhpTUNgDnNdoGT3VkZmal5HtQZmZWSk5QZmZWSk5QZmZWSk5QZmZWSk5QZmZWSk5QbSLpQEkfb2C7cyW9Mbf+mKSJxbbOyspxY43qhdhxgmqfA8lmJd5DGhk+nHOBN45Qx3rXgThurDEH0uWx47n42mchcKSkdcDLwIvAs8BbJM0A/ikijgOQdCmwP3A/2RxXfyvp58CvpX39oaTfBPYGfjsiftDWI7F2ctxYo7o+dnwG1T7zgIcj4njgU8AJwCURcXStDSLiDmAQmBURx0fEz9NH2yPiBOBLwKWtbbZ1mOPGGtX1seME1Tn/EdkzXBrxD+l1LTC1mOZYl3DcWKO6LnacoDrnhdz7X7Ln/8W+I2w7NGvwLnyZtt84bqxRXRc7TlDt8zxwQI3PngIOkvQGSWOBM+rcznqf48Ya1fWx419RbRIRz0j6jqT7gZ+TBcjQZy9LuopspuKtQP4G5F8Df1Vxw9L6hOPGGtULsePZzM3MrJR8ic/MzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErJCcrMzErp/wP/4wOhbjObtwAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_predictions(test_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "acc341d1-7327-40c2-8ccc-2707173fe905",
   "metadata": {},
   "outputs": [],
   "source": [
    "next(iter(test_dataloader))\n",
    "prediciton_i = model(next(iter(test_dataloader)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c312c309-ef9b-4bef-9ed0-04f8b6ca3b44",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0139, -0.0140, -0.0570],\n",
       "        [-0.0139, -0.0140, -0.0570],\n",
       "        [-0.0139, -0.0140, -0.0570],\n",
       "        ...,\n",
       "        [-0.0139, -0.0140, -0.0570],\n",
       "        [-0.0139, -0.0140, -0.0570],\n",
       "        [-0.0139, -0.0140, -0.0570]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prediciton_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff2ec930-79d9-4a13-9341-cab9f2aeec30",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.12"
  },
  "toc-autonumbering": true,
  "toc-showcode": false,
  "toc-showmarkdowntxt": true,
  "toc-showtags": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}