Dataset_PyG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset Construction for VEN Data as Heterogeneous Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import os.path as osp\n",
    "from itertools import product\n",
    "from typing import Callable, List, Optional\n",
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import (\n",
    "    HeteroData,\n",
    "    Data,\n",
    "    InMemoryDataset,\n",
    "    download_url,\n",
    "    extract_zip,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/Users/nbai/surfdrive/TUD/Paper/Venice_Graph'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lite Dataset of VEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = HeteroData()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_types = ['vis_only', 'vis_tex']\n",
    "vis = np.load(osp.join(os.getcwd(), 'dataset/Venice/raw/Visual_Features.npy'),allow_pickle=True)[:,2:].astype(float)\n",
    "tex = np.load(osp.join(os.getcwd(), 'dataset/Venice/raw/Textual_Features.npy'),allow_pickle=True)[:,5:].astype(float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((2951, 982), (2951, 771))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vis.shape, tex.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.hstack([vis,np.nan_to_num(tex)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2951, 1753)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_type_idx = np.load('dataset/Venice/raw/node_types.npy')\n",
    "node_type_idx = torch.from_numpy(node_type_idx).to(torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2951])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_type_idx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 1, 1,  ..., 1, 1, 1])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_type_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_only'].num_nodes = int((node_type_idx == 0).sum())\n",
    "data['vis_tex'].num_nodes = int((node_type_idx == 1).sum())\n",
    "data['all'].num_nodes = len(node_type_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  \u001b[1mvis_only\u001b[0m={ num_nodes=1190 },\n",
       "  \u001b[1mvis_tex\u001b[0m={ num_nodes=1761 },\n",
       "  \u001b[1mall\u001b[0m={ num_nodes=2951 }\n",
       ")"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1190, 1753)"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[node_type_idx==0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_only'].x = torch.from_numpy(vis[node_type_idx==0]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_tex'].x = torch.from_numpy(x[node_type_idx==1]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['all'].x = torch.from_numpy(x).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_s = np.load('dataset/Venice/raw/labels.npz',allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_lab = y_s['ATT_LAB'][:,1:10].astype(float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_tag = y_s['ATT_LAB'][:,-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_tag = y_s['VAL_LAB'][:,-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1356, 756, 361)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "att_tag.sum(), val_tag.sum(), (att_tag * val_tag).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_lab = np.nan_to_num(y_s['VAL_LAB'][:,2:13].astype(float))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2951, 20)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ys = np.hstack([att_lab, val_lab])\n",
    "ys.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_only'].y = torch.from_numpy(att_lab[node_type_idx==0]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_tex'].y = torch.from_numpy(ys[node_type_idx==1]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['all'].y = torch.from_numpy(ys).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.node_type = node_type_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[2951],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=1190,\n",
       "    x=[1190, 1753],\n",
       "    y=[1190, 9]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=1761,\n",
       "    x=[1761, 1753],\n",
       "    y=[1761, 20]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=2951,\n",
       "    x=[2951, 1753],\n",
       "    y=[2951, 20]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = np.load('dataset/Venice/train_val_test_idx.npz')\n",
    "for name in ['train', 'val', 'test']:\n",
    "    idx = split[f'{name}_idx']\n",
    "    idx = torch.from_numpy(idx).to(torch.long)\n",
    "    mask = torch.zeros(data['all'].num_nodes, dtype=torch.bool)\n",
    "    mask[idx] = True\n",
    "    data['all'][f'{name}_mask'] = mask\n",
    "    data['vis_only'][f'{name}_mask'] = mask[node_type_idx==0]\n",
    "    data['vis_tex'][f'{name}_mask'] = mask[node_type_idx==1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[2951],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=1190,\n",
       "    x=[1190, 1753],\n",
       "    y=[1190, 9],\n",
       "    train_mask=[1190],\n",
       "    val_mask=[1190],\n",
       "    test_mask=[1190]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=1761,\n",
       "    x=[1761, 1753],\n",
       "    y=[1761, 20],\n",
       "    train_mask=[1761],\n",
       "    val_mask=[1761],\n",
       "    test_mask=[1761]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=2951,\n",
       "    x=[2951, 1753],\n",
       "    y=[2951, 20],\n",
       "    train_mask=[2951],\n",
       "    val_mask=[2951],\n",
       "    test_mask=[2951]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "link_types = ['SOC', 'SPA', 'TEM', 'simp']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(488103,)"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_SOC.data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [],
   "source": [
    "for link in link_types:\n",
    "    A_sub = sp.load_npz(f'dataset/Venice/A_{link}.npz').tocoo()\n",
    "    if A_sub.nnz>0:\n",
    "        row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "        col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "        data['all', f'{link}_link', 'all'].edge_index = torch.stack([row, col], dim=0)\n",
    "        data['all', f'{link}_link', 'all'].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[2951],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=1190,\n",
       "    x=[1190, 1753],\n",
       "    y=[1190, 9],\n",
       "    train_mask=[1190],\n",
       "    val_mask=[1190],\n",
       "    test_mask=[1190]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=1761,\n",
       "    x=[1761, 1753],\n",
       "    y=[1761, 20],\n",
       "    train_mask=[1761],\n",
       "    val_mask=[1761],\n",
       "    test_mask=[1761]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=2951,\n",
       "    x=[2951, 1753],\n",
       "    y=[2951, 20],\n",
       "    train_mask=[2951],\n",
       "    val_mask=[2951],\n",
       "    test_mask=[2951]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 488103],\n",
       "    edge_attr=[488103]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 445779],\n",
       "    edge_attr=[445779]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 501191],\n",
       "    edge_attr=[501191]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 1071977],\n",
       "    edge_attr=[1071977]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 147,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = {}\n",
    "s['vis_only'] = np.arange(len(x))[node_type_idx==0]\n",
    "s['vis_tex'] = np.arange(len(x))[node_type_idx==1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "metadata": {},
   "outputs": [],
   "source": [
    "for src, dst in product(node_types, node_types):\n",
    "    for link in link_types:\n",
    "        A_sub = sp.load_npz(f'dataset/Venice/A_{link}.npz')[s[src]][:,s[dst]].tocoo()\n",
    "        if A_sub.nnz>0:\n",
    "            row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "            col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "            data[src, f'{link}_link', dst].edge_index = torch.stack([row, col], dim=0)\n",
    "            data[src, f'{link}_link', dst].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[2951],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=1190,\n",
       "    x=[1190, 982],\n",
       "    y=[1190, 9],\n",
       "    train_mask=[1190],\n",
       "    val_mask=[1190],\n",
       "    test_mask=[1190]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=1761,\n",
       "    x=[1761, 1753],\n",
       "    y=[1761, 20],\n",
       "    train_mask=[1761],\n",
       "    val_mask=[1761],\n",
       "    test_mask=[1761]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=2951,\n",
       "    x=[2951, 1753],\n",
       "    y=[2951, 20],\n",
       "    train_mask=[2951],\n",
       "    val_mask=[2951],\n",
       "    test_mask=[2951]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 488103],\n",
       "    edge_attr=[488103]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 445779],\n",
       "    edge_attr=[445779]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 501191],\n",
       "    edge_attr=[501191]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 1071977],\n",
       "    edge_attr=[1071977]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 169762],\n",
       "    edge_attr=[169762]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 77626],\n",
       "    edge_attr=[77626]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 175424],\n",
       "    edge_attr=[175424]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 255444],\n",
       "    edge_attr=[255444]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 87752],\n",
       "    edge_attr=[87752]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 104354],\n",
       "    edge_attr=[104354]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 101349],\n",
       "    edge_attr=[101349]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 237794],\n",
       "    edge_attr=[237794]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 87752],\n",
       "    edge_attr=[87752]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 104354],\n",
       "    edge_attr=[104354]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 101349],\n",
       "    edge_attr=[101349]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 237794],\n",
       "    edge_attr=[237794]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 142837],\n",
       "    edge_attr=[142837]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 159445],\n",
       "    edge_attr=[159445]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 123069],\n",
       "    edge_attr=[123069]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 340945],\n",
       "    edge_attr=[340945]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 165,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Large Dataset of VEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = HeteroData()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_types = ['vis_only', 'vis_tex']\n",
    "vis = np.load(osp.join(os.getcwd(), 'dataset/Venice-XL/Visual_Features.npy'),allow_pickle=True)[:,2:].astype(float)\n",
    "tex = np.load(osp.join(os.getcwd(), 'dataset/Venice-XL/Textual_Features.npy'),allow_pickle=True)[:,5:].astype(float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((80963, 982), (80963, 771))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vis.shape, tex.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.hstack([vis,np.nan_to_num(tex)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(80963, 1753)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_type_idx = np.load('dataset/Venice-XL/node_types.npy')\n",
    "node_type_idx = torch.from_numpy(node_type_idx).to(torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([80963])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_type_idx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_only'].num_nodes = int((node_type_idx == 0).sum())\n",
    "data['vis_tex'].num_nodes = int((node_type_idx == 1).sum())\n",
    "data['all'].num_nodes = len(node_type_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  \u001b[1mvis_only\u001b[0m={ num_nodes=31140 },\n",
       "  \u001b[1mvis_tex\u001b[0m={ num_nodes=49823 },\n",
       "  \u001b[1mall\u001b[0m={ num_nodes=80963 }\n",
       ")"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(31140, 1753)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[node_type_idx==0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_only'].x = torch.from_numpy(vis[node_type_idx==0]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_tex'].x = torch.from_numpy(x[node_type_idx==1]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['all'].x = torch.from_numpy(x).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_s = np.load('dataset/Venice-XL/labels.npz',allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_lab = y_s['ATT_LAB'][:,1:10].astype(float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_lab = np.nan_to_num(y_s['VAL_LAB'][:,2:13].astype(float))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(80963, 20)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ys = np.hstack([att_lab, val_lab])\n",
    "ys.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_only'].y = torch.from_numpy(att_lab[node_type_idx==0]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['vis_tex'].y = torch.from_numpy(ys[node_type_idx==1]).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['all'].y = torch.from_numpy(ys).to(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.node_type = node_type_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[80963],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=31140,\n",
       "    x=[31140, 982],\n",
       "    y=[31140, 9]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=49823,\n",
       "    x=[49823, 1753],\n",
       "    y=[49823, 20]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=80963,\n",
       "    x=[80963, 1753],\n",
       "    y=[80963, 20]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = np.load('dataset/Venice-XL/train_val_test_idx.npz')\n",
    "for name in ['train', 'val', 'test']:\n",
    "    idx = split[f'{name}_idx']\n",
    "    idx = torch.from_numpy(idx).to(torch.long)\n",
    "    mask = torch.zeros(data['all'].num_nodes, dtype=torch.bool)\n",
    "    mask[idx] = True\n",
    "    data['all'][f'{name}_mask'] = mask\n",
    "    data['vis_only'][f'{name}_mask'] = mask[node_type_idx==0]\n",
    "    data['vis_tex'][f'{name}_mask'] = mask[node_type_idx==1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[80963],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=31140,\n",
       "    x=[31140, 982],\n",
       "    y=[31140, 9],\n",
       "    train_mask=[31140],\n",
       "    val_mask=[31140],\n",
       "    test_mask=[31140]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=49823,\n",
       "    x=[49823, 1753],\n",
       "    y=[49823, 20],\n",
       "    train_mask=[49823],\n",
       "    val_mask=[49823],\n",
       "    test_mask=[49823]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=80963,\n",
       "    x=[80963, 1753],\n",
       "    y=[80963, 20],\n",
       "    train_mask=[80963],\n",
       "    val_mask=[80963],\n",
       "    test_mask=[80963]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "link_types = ['SOC', 'SPA', 'TEM', 'simp']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "for link in link_types:\n",
    "    A_sub = sp.load_npz(f'dataset/Venice-XL/A_{link}.npz').tocoo()\n",
    "    if A_sub.nnz>0:\n",
    "        row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "        col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "        data['all', f'{link}_link', 'all'].edge_index = torch.stack([row, col], dim=0)\n",
    "        data['all', f'{link}_link', 'all'].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[80963],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=31140,\n",
       "    x=[31140, 982],\n",
       "    y=[31140, 9],\n",
       "    train_mask=[31140],\n",
       "    val_mask=[31140],\n",
       "    test_mask=[31140]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=49823,\n",
       "    x=[49823, 1753],\n",
       "    y=[49823, 20],\n",
       "    train_mask=[49823],\n",
       "    val_mask=[49823],\n",
       "    test_mask=[49823]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=80963,\n",
       "    x=[80963, 1753],\n",
       "    y=[80963, 20],\n",
       "    train_mask=[80963],\n",
       "    val_mask=[80963],\n",
       "    test_mask=[80963]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 76422265],\n",
       "    edge_attr=[76422265]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 202173159],\n",
       "    edge_attr=[202173159]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 71135671],\n",
       "    edge_attr=[71135671]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 290091503],\n",
       "    edge_attr=[290091503]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = {}\n",
    "s['vis_only'] = np.arange(len(x))[node_type_idx==0]\n",
    "s['vis_tex'] = np.arange(len(x))[node_type_idx==1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "for src, dst in product(node_types, node_types):\n",
    "    for link in link_types:\n",
    "        A_sub = sp.load_npz(f'dataset/Venice-XL/A_{link}.npz')[s[src]][:,s[dst]].tocoo()\n",
    "        if A_sub.nnz>0:\n",
    "            row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "            col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "            data[src, f'{link}_link', dst].edge_index = torch.stack([row, col], dim=0)\n",
    "            data[src, f'{link}_link', dst].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[80963],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=31140,\n",
       "    x=[31140, 982],\n",
       "    y=[31140, 9],\n",
       "    train_mask=[31140],\n",
       "    val_mask=[31140],\n",
       "    test_mask=[31140]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=49823,\n",
       "    x=[49823, 1753],\n",
       "    y=[49823, 20],\n",
       "    train_mask=[49823],\n",
       "    val_mask=[49823],\n",
       "    test_mask=[49823]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=80963,\n",
       "    x=[80963, 1753],\n",
       "    y=[80963, 20],\n",
       "    train_mask=[80963],\n",
       "    val_mask=[80963],\n",
       "    test_mask=[80963]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 76422265],\n",
       "    edge_attr=[76422265]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 202173159],\n",
       "    edge_attr=[202173159]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 71135671],\n",
       "    edge_attr=[71135671]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 290091503],\n",
       "    edge_attr=[290091503]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 9813406],\n",
       "    edge_attr=[9813406]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 22404794],\n",
       "    edge_attr=[22404794]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 9618578],\n",
       "    edge_attr=[9618578]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 35574584],\n",
       "    edge_attr=[35574584]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10129888],\n",
       "    edge_attr=[10129888]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 40759807],\n",
       "    edge_attr=[40759807]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10973512],\n",
       "    edge_attr=[10973512]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 57890918],\n",
       "    edge_attr=[57890918]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 10129888],\n",
       "    edge_attr=[10129888]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 40759807],\n",
       "    edge_attr=[40759807]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 10973512],\n",
       "    edge_attr=[10973512]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 57890918],\n",
       "    edge_attr=[57890918]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 46349083],\n",
       "    edge_attr=[46349083]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 98248751],\n",
       "    edge_attr=[98248751]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 39570069],\n",
       "    edge_attr=[39570069]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 138735083],\n",
       "    edge_attr=[138735083]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Formal Class Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VEN(InMemoryDataset):\n",
    "    r\"\"\"A subset of Flickr post collected in Venice annotated with Heritage \n",
    "    Values and Attributes, as collected in the `\"Heri-Graphs: A Workflow of \n",
    "    Creating Datasets for Multi-modal Machine Learning on Graphs of Heritage \n",
    "    Values and Attributes with Social Media\" <https://arxiv.org/abs/2205.07545>`\n",
    "    paper.\n",
    "    VEN is a heterogeneous graph containing two types of nodes - nodes with only \n",
    "    visual features 'vis_only' (1,190 nodes), nodes with both visual and textual\n",
    "    features 'vis_tex' (1,761 nodes) and four types of links - social similarity\n",
    "    'SOC' (488,103 links), spatial similarity (445,779 links), temporal similarity\n",
    "    (501,191 links), and simple composed link (1,071,977 links).\n",
    "    Vis_only nodes are represented with 982-dimensional visual features and are\n",
    "    divided into 9 heritage attribute categories \n",
    "    ('architectural elements', 'form', 'gastronomy', 'interior',\n",
    "    'landscape scenery and natural features', 'monuments', 'people', 'product', \n",
    "    'urban scenery').\n",
    "    Vis_text nodes are represented with 1753-dimensional visual and textual \n",
    "    features and are divided into 9 heritage attribute categories plus 11 \n",
    "    heritage value categories ('criterion i-x', 'other').\n",
    "    Both types of nodes are also merged into a single type of node 'all' with \n",
    "    1753-dimensional features and 20-dimensional label categories.\n",
    "\n",
    "    Args:\n",
    "        root (string): Root directory where the dataset should be saved.\n",
    "        transform (callable, optional): A function/transform that takes in an\n",
    "            :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            every access. (default: :obj:`None`)\n",
    "        pre_transform (callable, optional): A function/transform that takes in\n",
    "            an :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            being saved to disk. (default: :obj:`None`)\n",
    "    \n",
    "    Stats:\n",
    "            * - #nodes\n",
    "              - #edges\n",
    "              - #features\n",
    "              - #classes\n",
    "            * - 2,951\n",
    "              - 1,071,977\n",
    "              - 1753\n",
    "              - 20\n",
    "    \"\"\"\n",
    "\n",
    "    url = 'https://drive.google.com/uc?export=download&id=1sxcKiZr1YGDv06wr03nsk5HVZledgzi9'\n",
    "\n",
    "    def __init__(self, root: str, transform: Optional[Callable] = None,\n",
    "                 pre_transform: Optional[Callable] = None):\n",
    "        super().__init__(root, transform, pre_transform)\n",
    "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
    "        \n",
    "    @property\n",
    "    def raw_file_names(self) -> List[str]:\n",
    "        return [\n",
    "            'A_simp.npz', 'A_SOC.npz', 'A_SPA.npz', 'A_TEM.npz', 'labels.npz',\n",
    "            'node_types.npy', 'Textual_Features.npy', 'train_val_test_idx.npz',\n",
    "            'Visual_Features.npy'\n",
    "        ]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return 'data.pt'\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "        extract_zip(path, self.raw_dir)\n",
    "        os.remove(path)\n",
    "\n",
    "    def process(self):\n",
    "        data = HeteroData()\n",
    "\n",
    "        node_types = ['vis_only', 'vis_tex']\n",
    "        link_types = ['SOC', 'SPA', 'TEM', 'simp']\n",
    "\n",
    "        vis = np.load(osp.join(self.raw_dir, 'Visual_Features.npy'),allow_pickle=True)[:,2:].astype(float)\n",
    "        tex = np.load(osp.join(self.raw_dir, 'Textual_Features.npy'),allow_pickle=True)[:,5:].astype(float)\n",
    "\n",
    "        x = np.hstack([vis,np.nan_to_num(tex)])\n",
    "\n",
    "\n",
    "        node_type_idx = np.load(osp.join(self.raw_dir, 'node_types.npy'))\n",
    "        node_type_idx = torch.from_numpy(node_type_idx).to(torch.long)\n",
    "\n",
    "        data['vis_only'].num_nodes = int((node_type_idx == 0).sum())\n",
    "        data['vis_tex'].num_nodes = int((node_type_idx == 1).sum())\n",
    "        data['all'].num_nodes = len(node_type_idx)\n",
    "\n",
    "        data['vis_only'].x = torch.from_numpy(vis[node_type_idx==0]).to(torch.float)\n",
    "        data['vis_tex'].x = torch.from_numpy(x[node_type_idx==1]).to(torch.float)\n",
    "        data['all'].x = torch.from_numpy(x).to(torch.float)\n",
    "\n",
    "\n",
    "        y_s = np.load(osp.join(self.raw_dir, 'labels.npz'), allow_pickle=True)\n",
    "        att_lab = y_s['ATT_LAB'][:,1:10].astype(float)\n",
    "        val_lab = np.nan_to_num(y_s['VAL_LAB'][:,2:13].astype(float))\n",
    "        ys = np.hstack([att_lab, val_lab])\n",
    "\n",
    "        data['vis_only'].y = torch.from_numpy(att_lab[node_type_idx==0]).to(torch.float)\n",
    "        data['vis_tex'].y = torch.from_numpy(ys[node_type_idx==1]).to(torch.float)\n",
    "        data['all'].y = torch.from_numpy(ys).to(torch.float)\n",
    "\n",
    "        data.node_type = node_type_idx\n",
    "\n",
    "        split = np.load(osp.join(self.raw_dir, 'train_val_test_idx.npz'))\n",
    "        for name in ['train', 'val', 'test']:\n",
    "            idx = split[f'{name}_idx']\n",
    "            idx = torch.from_numpy(idx).to(torch.long)\n",
    "            mask = torch.zeros(data['all'].num_nodes, dtype=torch.bool)\n",
    "            mask[idx] = True\n",
    "            data['all'][f'{name}_mask'] = mask\n",
    "            data['vis_only'][f'{name}_mask'] = mask[node_type_idx==0]\n",
    "            data['vis_tex'][f'{name}_mask'] = mask[node_type_idx==1]\n",
    "\n",
    "        \n",
    "        s = {}\n",
    "        s['vis_only'] = np.arange(len(x))[node_type_idx==0]\n",
    "        s['vis_tex'] = np.arange(len(x))[node_type_idx==1]\n",
    "\n",
    "        for link in link_types:\n",
    "            A_sub = sp.load_npz(osp.join(self.raw_dir, f'A_{link}.npz')).tocoo()\n",
    "            if A_sub.nnz>0:\n",
    "                row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "                col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "                data['all', f'{link}_link', 'all'].edge_index = torch.stack([row, col], dim=0)\n",
    "                data['all', f'{link}_link', 'all'].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)\n",
    "\n",
    "        for src, dst in product(node_types, node_types):\n",
    "            for link in link_types:\n",
    "                A_sub = sp.load_npz(osp.join(self.raw_dir, f'A_{link}.npz'))[s[src]][:,s[dst]].tocoo()\n",
    "                if A_sub.nnz>0:\n",
    "                    row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "                    col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "                    data[src, f'{link}_link', dst].edge_index = torch.stack([row, col], dim=0)\n",
    "                    data[src, f'{link}_link', dst].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)\n",
    "\n",
    "\n",
    "        if self.pre_transform is not None:\n",
    "            data = self.pre_transform(data)\n",
    "\n",
    "        torch.save(self.collate([data]), self.processed_paths[0])\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f'{self.__class__.__name__}()'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VEN_XL(InMemoryDataset):\n",
    "    r\"\"\"A large subset of Flickr post collected in Venice annotated with Heritage \n",
    "    Values and Attributes, as collected in the `\"Heri-Graphs: A Workflow of \n",
    "    Creating Datasets for Multi-modal Machine Learning on Graphs of Heritage \n",
    "    Values and Attributes with Social Media\" <https://arxiv.org/abs/2205.07545>`\n",
    "    paper.\n",
    "    VEN_XL is a heterogeneous graph containing two types of nodes - nodes with only \n",
    "    visual features 'vis_only' (31,140 nodes), nodes with both visual and textual\n",
    "    features 'vis_tex' (49,823 nodes) and four types of links - social similarity\n",
    "    'SOC' (76,422,265 links), spatial similarity (202,173,159 links), temporal similarity\n",
    "    (71,135,671 links), and simple composed link (290,091,503 links).\n",
    "    Vis_only nodes are represented with 982-dimensional visual features and are\n",
    "    divided into 9 heritage attribute categories \n",
    "    ('architectural elements', 'form', 'gastronomy', 'interior',\n",
    "    'landscape scenery and natural features', 'monuments', 'people', 'product', \n",
    "    'urban scenery').\n",
    "    Vis_text nodes are represented with 1753-dimensional visual and textual \n",
    "    features and are divided into 9 heritage attribute categories plus 11 \n",
    "    heritage value categories ('criterion i-x', 'other').\n",
    "    Both types of nodes are also merged into a single type of node 'all' with \n",
    "    1753-dimensional features and 20-dimensional label categories.\n",
    "\n",
    "    Args:\n",
    "        root (string): Root directory where the dataset should be saved.\n",
    "        transform (callable, optional): A function/transform that takes in an\n",
    "            :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            every access. (default: :obj:`None`)\n",
    "        pre_transform (callable, optional): A function/transform that takes in\n",
    "            an :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            being saved to disk. (default: :obj:`None`)\n",
    "    \n",
    "    Stats:\n",
    "            * - #nodes\n",
    "              - #edges\n",
    "              - #features\n",
    "              - #classes\n",
    "            * - 80,963\n",
    "              - 290,091,503\n",
    "              - 1753\n",
    "              - 20\n",
    "    \"\"\"\n",
    "\n",
    "    url = 'https://drive.google.com/uc?export=download&id=1QZ5tyUWs6jYjh7mJrsnpou76iy-vb0CA'\n",
    "\n",
    "    def __init__(self, root: str, transform: Optional[Callable] = None,\n",
    "                 pre_transform: Optional[Callable] = None):\n",
    "        super().__init__(root, transform, pre_transform)\n",
    "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
    "        \n",
    "    @property\n",
    "    def raw_file_names(self) -> List[str]:\n",
    "        return [\n",
    "            'A_simp.npz', 'A_SOC.npz', 'A_SPA.npz', 'A_TEM.npz', 'labels.npz',\n",
    "            'node_types.npy', 'Textual_Features.npy', 'train_val_test_idx.npz',\n",
    "            'Visual_Features.npy'\n",
    "        ]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return 'data.pt'\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "        extract_zip(path, self.raw_dir)\n",
    "        os.remove(path)\n",
    "\n",
    "    def process(self):\n",
    "        data = HeteroData()\n",
    "\n",
    "        node_types = ['vis_only', 'vis_tex']\n",
    "        link_types = ['SOC', 'SPA', 'TEM', 'simp']\n",
    "\n",
    "        vis = np.load(osp.join(self.raw_dir, 'Visual_Features.npy'),allow_pickle=True)[:,2:].astype(float)\n",
    "        tex = np.load(osp.join(self.raw_dir, 'Textual_Features.npy'),allow_pickle=True)[:,5:].astype(float)\n",
    "\n",
    "        x = np.hstack([vis,np.nan_to_num(tex)])\n",
    "\n",
    "\n",
    "        node_type_idx = np.load(osp.join(self.raw_dir, 'node_types.npy'))\n",
    "        node_type_idx = torch.from_numpy(node_type_idx).to(torch.long)\n",
    "\n",
    "        data['vis_only'].num_nodes = int((node_type_idx == 0).sum())\n",
    "        data['vis_tex'].num_nodes = int((node_type_idx == 1).sum())\n",
    "        data['all'].num_nodes = len(node_type_idx)\n",
    "\n",
    "        data['vis_only'].x = torch.from_numpy(vis[node_type_idx==0]).to(torch.float)\n",
    "        data['vis_tex'].x = torch.from_numpy(x[node_type_idx==1]).to(torch.float)\n",
    "        data['all'].x = torch.from_numpy(x).to(torch.float)\n",
    "\n",
    "\n",
    "        y_s = np.load(osp.join(self.raw_dir, 'labels.npz'), allow_pickle=True)\n",
    "        att_lab = y_s['ATT_LAB'][:,1:10].astype(float)\n",
    "        val_lab = np.nan_to_num(y_s['VAL_LAB'][:,2:13].astype(float))\n",
    "        ys = np.hstack([att_lab, val_lab])\n",
    "\n",
    "        data['vis_only'].y = torch.from_numpy(att_lab[node_type_idx==0]).to(torch.float)\n",
    "        data['vis_tex'].y = torch.from_numpy(ys[node_type_idx==1]).to(torch.float)\n",
    "        data['all'].y = torch.from_numpy(ys).to(torch.float)\n",
    "\n",
    "        data.node_type = node_type_idx\n",
    "\n",
    "        split = np.load(osp.join(self.raw_dir, 'train_val_test_idx.npz'))\n",
    "        for name in ['train', 'val', 'test']:\n",
    "            idx = split[f'{name}_idx']\n",
    "            idx = torch.from_numpy(idx).to(torch.long)\n",
    "            mask = torch.zeros(data['all'].num_nodes, dtype=torch.bool)\n",
    "            mask[idx] = True\n",
    "            data['all'][f'{name}_mask'] = mask\n",
    "            data['vis_only'][f'{name}_mask'] = mask[node_type_idx==0]\n",
    "            data['vis_tex'][f'{name}_mask'] = mask[node_type_idx==1]\n",
    "\n",
    "        \n",
    "        s = {}\n",
    "        s['vis_only'] = np.arange(len(x))[node_type_idx==0]\n",
    "        s['vis_tex'] = np.arange(len(x))[node_type_idx==1]\n",
    "\n",
    "        for link in link_types:\n",
    "            A_sub = sp.load_npz(osp.join(self.raw_dir, f'A_{link}.npz')).tocoo()\n",
    "            if A_sub.nnz>0:\n",
    "                row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "                col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "                data['all', f'{link}_link', 'all'].edge_index = torch.stack([row, col], dim=0)\n",
    "                data['all', f'{link}_link', 'all'].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)\n",
    "\n",
    "        for src, dst in product(node_types, node_types):\n",
    "            for link in link_types:\n",
    "                A_sub = sp.load_npz(osp.join(self.raw_dir, f'A_{link}.npz'))[s[src]][:,s[dst]].tocoo()\n",
    "                if A_sub.nnz>0:\n",
    "                    row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "                    col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "                    data[src, f'{link}_link', dst].edge_index = torch.stack([row, col], dim=0)\n",
    "                    data[src, f'{link}_link', dst].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)\n",
    "\n",
    "\n",
    "        if self.pre_transform is not None:\n",
    "            data = self.pre_transform(data)\n",
    "\n",
    "        torch.save(self.collate([data]), self.processed_paths[0])\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f'{self.__class__.__name__}()'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VEN_XL_Homo(InMemoryDataset):\n",
    "    r\"\"\"A large subset of Flickr post collected in Venice annotated with Heritage \n",
    "    Values and Attributes, as collected in the `\"Heri-Graphs: A Workflow of \n",
    "    Creating Datasets for Multi-modal Machine Learning on Graphs of Heritage \n",
    "    Values and Attributes with Social Media\" <https://arxiv.org/abs/2205.07545>`\n",
    "    paper.\n",
    "    VEN_XL is a heterogeneous graph containing two types of nodes - nodes with only \n",
    "    visual features 'vis_only' (31,140 nodes), nodes with both visual and textual\n",
    "    features 'vis_tex' (49,823 nodes) and four types of links - social similarity\n",
    "    'SOC' (76,422,265 links), spatial similarity (202,173,159 links), temporal similarity\n",
    "    (71,135,671 links), and simple composed link (290,091,503 links).\n",
    "    Vis_only nodes are represented with 982-dimensional visual features and are\n",
    "    divided into 9 heritage attribute categories \n",
    "    ('architectural elements', 'form', 'gastronomy', 'interior',\n",
    "    'landscape scenery and natural features', 'monuments', 'people', 'product', \n",
    "    'urban scenery').\n",
    "    Vis_text nodes are represented with 1753-dimensional visual and textual \n",
    "    features and are divided into 9 heritage attribute categories plus 11 \n",
    "    heritage value categories ('criterion i-x', 'other').\n",
    "    Both types of nodes are also merged into a single type of node 'all' with \n",
    "    1753-dimensional features and 20-dimensional label categories.\n",
    "\n",
    "    Args:\n",
    "        root (string): Root directory where the dataset should be saved.\n",
    "        transform (callable, optional): A function/transform that takes in an\n",
    "            :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            every access. (default: :obj:`None`)\n",
    "        pre_transform (callable, optional): A function/transform that takes in\n",
    "            an :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            being saved to disk. (default: :obj:`None`)\n",
    "    \n",
    "    Stats:\n",
    "            * - #nodes\n",
    "              - #edges\n",
    "              - #features\n",
    "              - #classes\n",
    "            * - 80,963\n",
    "              - 290,091,503\n",
    "              - 1753\n",
    "              - 20\n",
    "    \"\"\"\n",
    "\n",
    "    url = 'https://drive.google.com/uc?export=download&id=1sxcKiZr1YGDv06wr03nsk5HVZledgzi9'\n",
    "\n",
    "    def __init__(self, root: str, transform: Optional[Callable] = None,\n",
    "                 pre_transform: Optional[Callable] = None):\n",
    "        super().__init__(root, transform, pre_transform)\n",
    "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
    "        \n",
    "    @property\n",
    "    def raw_file_names(self) -> List[str]:\n",
    "        return [\n",
    "            'A_simp.npz', 'A_SOC.npz', 'A_SPA.npz', 'A_TEM.npz', 'labels.npz',\n",
    "            'node_types.npy', 'Textual_Features.npy', 'train_val_test_idx.npz',\n",
    "            'Visual_Features.npy'\n",
    "        ]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return 'data.pt'\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "        extract_zip(path, self.raw_dir)\n",
    "        os.remove(path)\n",
    "\n",
    "    def process(self):\n",
    "        data = Data()\n",
    "\n",
    "        link_types = ['simp']\n",
    "\n",
    "        vis = np.load(osp.join(self.raw_dir, 'Visual_Features.npy'),allow_pickle=True)[:,2:].astype(float)\n",
    "        tex = np.load(osp.join(self.raw_dir, 'Textual_Features.npy'),allow_pickle=True)[:,5:].astype(float)\n",
    "\n",
    "        x = np.hstack([vis,np.nan_to_num(tex)])\n",
    "\n",
    "        node_type_idx = np.load(osp.join(self.raw_dir, 'node_types.npy'))\n",
    "        node_type_idx = torch.from_numpy(node_type_idx).to(torch.long)\n",
    "\n",
    "        data.num_nodes = len(node_type_idx)\n",
    "\n",
    "        data.x = torch.from_numpy(x).to(torch.float)\n",
    "\n",
    "\n",
    "        y_s = np.load(osp.join(self.raw_dir, 'labels.npz'), allow_pickle=True)\n",
    "        att_lab = y_s['ATT_LAB'][:,1:10].astype(float)\n",
    "        val_lab = np.nan_to_num(y_s['VAL_LAB'][:,2:13].astype(float))\n",
    "        ys = np.hstack([att_lab, val_lab])\n",
    "\n",
    "        data.y = torch.from_numpy(ys).to(torch.float)\n",
    "\n",
    "        data.node_type = node_type_idx\n",
    "\n",
    "        split = np.load(osp.join(self.raw_dir, 'train_val_test_idx.npz'))\n",
    "        for name in ['train', 'val', 'test']:\n",
    "            idx = split[f'{name}_idx']\n",
    "            idx = torch.from_numpy(idx).to(torch.long)\n",
    "            mask = torch.zeros(data.num_nodes, dtype=torch.bool)\n",
    "            mask[idx] = True\n",
    "            data[f'{name}_mask'] = mask\n",
    "                    \n",
    "        s = {}\n",
    "        \n",
    "        for link in link_types:\n",
    "            A_sub = sp.load_npz(osp.join(self.raw_dir, f'A_{link}.npz')).tocoo()\n",
    "            if A_sub.nnz>0:\n",
    "                row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "                col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "                data.edge_index = torch.stack([row, col], dim=0)\n",
    "                data.edge_attr = torch.from_numpy(A_sub.data).to(torch.long)\n",
    "\n",
    "        if self.pre_transform is not None:\n",
    "            data = self.pre_transform(data)\n",
    "\n",
    "        torch.save(self.collate([data]), self.processed_paths[0])\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f'{self.__class__.__name__}()'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VEN_XL_Links(InMemoryDataset):\n",
    "    r\"\"\"A large subset of Flickr post collected in Venice annotated with Heritage \n",
    "    Values and Attributes, as collected in the `\"Heri-Graphs: A Workflow of \n",
    "    Creating Datasets for Multi-modal Machine Learning on Graphs of Heritage \n",
    "    Values and Attributes with Social Media\" <https://arxiv.org/abs/2205.07545>`\n",
    "    paper.\n",
    "    VEN_XL is a heterogeneous graph containing two types of nodes - nodes with only \n",
    "    visual features 'vis_only' (31,140 nodes), nodes with both visual and textual\n",
    "    features 'vis_tex' (49,823 nodes) and four types of links - social similarity\n",
    "    'SOC' (76,422,265 links), spatial similarity (202,173,159 links), temporal similarity\n",
    "    (71,135,671 links), and simple composed link (290,091,503 links).\n",
    "    Vis_only nodes are represented with 982-dimensional visual features and are\n",
    "    divided into 9 heritage attribute categories \n",
    "    ('architectural elements', 'form', 'gastronomy', 'interior',\n",
    "    'landscape scenery and natural features', 'monuments', 'people', 'product', \n",
    "    'urban scenery').\n",
    "    Vis_text nodes are represented with 1753-dimensional visual and textual \n",
    "    features and are divided into 9 heritage attribute categories plus 11 \n",
    "    heritage value categories ('criterion i-x', 'other').\n",
    "    Both types of nodes are also merged into a single type of node 'all' with \n",
    "    1753-dimensional features and 20-dimensional label categories.\n",
    "\n",
    "    Args:\n",
    "        root (string): Root directory where the dataset should be saved.\n",
    "        transform (callable, optional): A function/transform that takes in an\n",
    "            :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            every access. (default: :obj:`None`)\n",
    "        pre_transform (callable, optional): A function/transform that takes in\n",
    "            an :obj:`torch_geometric.data.HeteroData` object and returns a\n",
    "            transformed version. The data object will be transformed before\n",
    "            being saved to disk. (default: :obj:`None`)\n",
    "    \n",
    "    Stats:\n",
    "            * - #nodes\n",
    "              - #edges\n",
    "              - #features\n",
    "              - #classes\n",
    "            * - 80,963\n",
    "              - 290,091,503\n",
    "              - 1753\n",
    "              - 20\n",
    "    \"\"\"\n",
    "\n",
    "    url = 'https://drive.google.com/uc?export=download&id=1QZ5tyUWs6jYjh7mJrsnpou76iy-vb0CA'\n",
    "\n",
    "    def __init__(self, root: str, transform: Optional[Callable] = None,\n",
    "                 pre_transform: Optional[Callable] = None):\n",
    "        super().__init__(root, transform, pre_transform)\n",
    "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
    "        \n",
    "    @property\n",
    "    def raw_file_names(self) -> List[str]:\n",
    "        return [\n",
    "            'A_simp.npz', 'A_SOC.npz', 'A_SPA.npz', 'A_TEM.npz', 'labels.npz',\n",
    "            'node_types.npy', 'Textual_Features.npy', 'train_val_test_idx.npz',\n",
    "            'Visual_Features.npy'\n",
    "        ]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return 'data.pt'\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "        extract_zip(path, self.raw_dir)\n",
    "        os.remove(path)\n",
    "\n",
    "    def process(self):\n",
    "        data = HeteroData()\n",
    "\n",
    "        node_types = ['all']\n",
    "        link_types = ['SOC', 'SPA', 'TEM', 'simp']\n",
    "\n",
    "        vis = np.load(osp.join(self.raw_dir, 'Visual_Features.npy'),allow_pickle=True)[:,2:].astype(float)\n",
    "        tex = np.load(osp.join(self.raw_dir, 'Textual_Features.npy'),allow_pickle=True)[:,5:].astype(float)\n",
    "\n",
    "        x = np.hstack([vis,np.nan_to_num(tex)])\n",
    "\n",
    "\n",
    "        node_type_idx = np.load(osp.join(self.raw_dir, 'node_types.npy'))\n",
    "        node_type_idx = torch.from_numpy(node_type_idx).to(torch.long)\n",
    "\n",
    "        #data['vis_only'].num_nodes = int((node_type_idx == 0).sum())\n",
    "        #data['vis_tex'].num_nodes = int((node_type_idx == 1).sum())\n",
    "        data['all'].num_nodes = len(node_type_idx)\n",
    "\n",
    "        #data['vis_only'].x = torch.from_numpy(vis[node_type_idx==0]).to(torch.float)\n",
    "        #data['vis_tex'].x = torch.from_numpy(x[node_type_idx==1]).to(torch.float)\n",
    "        data['all'].x = torch.from_numpy(x).to(torch.float)\n",
    "\n",
    "\n",
    "        y_s = np.load(osp.join(self.raw_dir, 'labels.npz'), allow_pickle=True)\n",
    "        att_lab = y_s['ATT_LAB'][:,1:10].astype(float)\n",
    "        val_lab = np.nan_to_num(y_s['VAL_LAB'][:,2:13].astype(float))\n",
    "        ys = np.hstack([att_lab, val_lab])\n",
    "\n",
    "        #data['vis_only'].y = torch.from_numpy(att_lab[node_type_idx==0]).to(torch.float)\n",
    "        #data['vis_tex'].y = torch.from_numpy(ys[node_type_idx==1]).to(torch.float)\n",
    "        data['all'].y = torch.from_numpy(ys).to(torch.float)\n",
    "\n",
    "        data.node_type = node_type_idx\n",
    "\n",
    "        split = np.load(osp.join(self.raw_dir, 'train_val_test_idx.npz'))\n",
    "        for name in ['train', 'val', 'test']:\n",
    "            idx = split[f'{name}_idx']\n",
    "            idx = torch.from_numpy(idx).to(torch.long)\n",
    "            mask = torch.zeros(data['all'].num_nodes, dtype=torch.bool)\n",
    "            mask[idx] = True\n",
    "            data['all'][f'{name}_mask'] = mask\n",
    "            #data['vis_only'][f'{name}_mask'] = mask[node_type_idx==0]\n",
    "            #data['vis_tex'][f'{name}_mask'] = mask[node_type_idx==1]\n",
    "\n",
    "        \n",
    "        s = {}\n",
    "        #s['vis_only'] = np.arange(len(x))[node_type_idx==0]\n",
    "        #s['vis_tex'] = np.arange(len(x))[node_type_idx==1]\n",
    "\n",
    "        for link in link_types:\n",
    "            A_sub = sp.load_npz(osp.join(self.raw_dir, f'A_{link}.npz')).tocoo()\n",
    "            if A_sub.nnz>0:\n",
    "                row = torch.from_numpy(A_sub.row).to(torch.long)\n",
    "                col = torch.from_numpy(A_sub.col).to(torch.long)\n",
    "                data['all', f'{link}_link', 'all'].edge_index = torch.stack([row, col], dim=0)\n",
    "                data['all', f'{link}_link', 'all'].edge_attr = torch.from_numpy(A_sub.data).to(torch.long)\n",
    "\n",
    "        if self.pre_transform is not None:\n",
    "            data = self.pre_transform(data)\n",
    "\n",
    "        torch.save(self.collate([data]), self.processed_paths[0])\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f'{self.__class__.__name__}()'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "dataset = VEN_XL_Homo('dataset/Venice_XL_homo')\n",
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "dataset = VEN_XL_Links('dataset/Venice_XL_links')\n",
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#path = osp.join(os.getcwd(), '../../data/VEN')\n",
    "path = 'dataset/Venice'\n",
    "dataset = VEN(path)\n",
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[80963],\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=80963,\n",
       "    x=[80963, 1753],\n",
       "    y=[80963, 20],\n",
       "    train_mask=[80963],\n",
       "    val_mask=[80963],\n",
       "    test_mask=[80963]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 76422265],\n",
       "    edge_attr=[76422265]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 202173159],\n",
       "    edge_attr=[202173159]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 71135671],\n",
       "    edge_attr=[71135671]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 290091503],\n",
       "    edge_attr=[290091503]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(361)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['all']['train_mask'].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#path = osp.join(os.getcwd(), '../../data/VEN')\n",
    "path = 'dataset/Venice-XL'\n",
    "dataset = VEN_XL(path)\n",
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[80963],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=31140,\n",
       "    x=[31140, 982],\n",
       "    y=[31140, 9],\n",
       "    train_mask=[31140],\n",
       "    val_mask=[31140],\n",
       "    test_mask=[31140]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=49823,\n",
       "    x=[49823, 1753],\n",
       "    y=[49823, 20],\n",
       "    train_mask=[49823],\n",
       "    val_mask=[49823],\n",
       "    test_mask=[49823]\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=80963,\n",
       "    x=[80963, 1753],\n",
       "    y=[80963, 20],\n",
       "    train_mask=[80963],\n",
       "    val_mask=[80963],\n",
       "    test_mask=[80963]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 76422265],\n",
       "    edge_attr=[76422265]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 202173159],\n",
       "    edge_attr=[202173159]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 71135671],\n",
       "    edge_attr=[71135671]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 290091503],\n",
       "    edge_attr=[290091503]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 9813406],\n",
       "    edge_attr=[9813406]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 22404794],\n",
       "    edge_attr=[22404794]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 9618578],\n",
       "    edge_attr=[9618578]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 35574584],\n",
       "    edge_attr=[35574584]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10129888],\n",
       "    edge_attr=[10129888]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 40759807],\n",
       "    edge_attr=[40759807]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10973512],\n",
       "    edge_attr=[10973512]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 57890918],\n",
       "    edge_attr=[57890918]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 10129888],\n",
       "    edge_attr=[10129888]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 40759807],\n",
       "    edge_attr=[40759807]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 10973512],\n",
       "    edge_attr=[10973512]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 57890918],\n",
       "    edge_attr=[57890918]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 46349083],\n",
       "    edge_attr=[46349083]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 98248751],\n",
       "    edge_attr=[98248751]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 39570069],\n",
       "    edge_attr=[39570069]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 138735083],\n",
       "    edge_attr=[138735083]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Loader for Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.loader import NeighborLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('all', 'SOC_link', 'all'),\n",
       " ('all', 'SPA_link', 'all'),\n",
       " ('all', 'TEM_link', 'all'),\n",
       " ('all', 'simp_link', 'all'),\n",
       " ('vis_only', 'SOC_link', 'vis_only'),\n",
       " ('vis_only', 'SPA_link', 'vis_only'),\n",
       " ('vis_only', 'TEM_link', 'vis_only'),\n",
       " ('vis_only', 'simp_link', 'vis_only'),\n",
       " ('vis_only', 'SOC_link', 'vis_tex'),\n",
       " ('vis_only', 'SPA_link', 'vis_tex'),\n",
       " ('vis_only', 'TEM_link', 'vis_tex'),\n",
       " ('vis_only', 'simp_link', 'vis_tex'),\n",
       " ('vis_tex', 'SOC_link', 'vis_only'),\n",
       " ('vis_tex', 'SPA_link', 'vis_only'),\n",
       " ('vis_tex', 'TEM_link', 'vis_only'),\n",
       " ('vis_tex', 'simp_link', 'vis_only'),\n",
       " ('vis_tex', 'SOC_link', 'vis_tex'),\n",
       " ('vis_tex', 'SPA_link', 'vis_tex'),\n",
       " ('vis_tex', 'TEM_link', 'vis_tex'),\n",
       " ('vis_tex', 'simp_link', 'vis_tex')]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.edge_types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = NeighborLoader(\n",
    "    data,\n",
    "    # Sample 30 neighbors for each node and edge type for 2 iterations\n",
    "    num_neighbors={key: [5] * 2 for key in data.edge_types if not 'all' in key},\n",
    "    # Use a batch size of 128 for sampling training nodes of type paper\n",
    "    batch_size=8,\n",
    "    input_nodes=('all', data['all'].train_mask),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = NeighborLoader(\n",
    "    data,\n",
    "    # Sample 30 neighbors for each node and edge type for 2 iterations\n",
    "    num_neighbors={key: [5] * 2 for key in data.edge_types},\n",
    "    # Use a batch size of 128 for sampling training nodes of type paper\n",
    "    batch_size=8,\n",
    "    input_nodes=('all', data['all'].train_mask),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = NeighborLoader(\n",
    "    data,\n",
    "    # Sample 30 neighbors for each node and edge type for 2 iterations\n",
    "    num_neighbors={key: [15] * 2 for key in data.edge_types if not 'simp_link' in key},\n",
    "    # Use a batch size of 128 for sampling training nodes of type paper\n",
    "    batch_size=32,\n",
    "    input_nodes=('vis_tex', data['vis_tex'].train_mask),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[2951],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=1161,\n",
       "    x=[1161, 982],\n",
       "    y=[1161, 9],\n",
       "    train_mask=[1161],\n",
       "    val_mask=[1161],\n",
       "    test_mask=[1161]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=1755,\n",
       "    x=[1755, 1753],\n",
       "    y=[1755, 20],\n",
       "    train_mask=[1755],\n",
       "    val_mask=[1755],\n",
       "    test_mask=[1755],\n",
       "    batch_size=32\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=0,\n",
       "    x=[0, 1753],\n",
       "    y=[0, 20],\n",
       "    train_mask=[0],\n",
       "    val_mask=[0],\n",
       "    test_mask=[0]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 6522],\n",
       "    edge_attr=[6522]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 6777],\n",
       "    edge_attr=[6777]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 6913],\n",
       "    edge_attr=[6913]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 6042],\n",
       "    edge_attr=[6042]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10003],\n",
       "    edge_attr=[10003]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 7226],\n",
       "    edge_attr=[7226]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 5775],\n",
       "    edge_attr=[5775]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 7086],\n",
       "    edge_attr=[7086]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 7155],\n",
       "    edge_attr=[7155]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10603],\n",
       "    edge_attr=[10603]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10911],\n",
       "    edge_attr=[10911]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 10454],\n",
       "    edge_attr=[10454]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 191,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sampled_hetero_data = next(iter(train_loader))\n",
    "batch = sampled_hetero_data\n",
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(361)"
      ]
     },
     "execution_count": 192,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch['vis_tex']['train_mask'].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1161, 982])"
      ]
     },
     "execution_count": 195,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.x_dict['vis_only'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = NeighborLoader(\n",
    "    data,\n",
    "    # Sample 30 neighbors for each node and edge type for 2 iterations\n",
    "    num_neighbors={key: [15] * 2 for key in data.edge_types if not 'simp_link' in key},\n",
    "    # Use a batch size of 128 for sampling training nodes of type paper\n",
    "    batch_size=8,\n",
    "    input_nodes=('vis_tex', data['vis_tex'].train_mask),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  node_type=[80963],\n",
       "  \u001b[1mvis_only\u001b[0m={\n",
       "    num_nodes=12351,\n",
       "    x=[12351, 982],\n",
       "    y=[12351, 9],\n",
       "    train_mask=[12351],\n",
       "    val_mask=[12351],\n",
       "    test_mask=[12351]\n",
       "  },\n",
       "  \u001b[1mvis_tex\u001b[0m={\n",
       "    num_nodes=15191,\n",
       "    x=[15191, 1753],\n",
       "    y=[15191, 20],\n",
       "    train_mask=[15191],\n",
       "    val_mask=[15191],\n",
       "    test_mask=[15191],\n",
       "    batch_size=8\n",
       "  },\n",
       "  \u001b[1mall\u001b[0m={\n",
       "    num_nodes=0,\n",
       "    x=[0, 1753],\n",
       "    y=[0, 20],\n",
       "    train_mask=[0],\n",
       "    val_mask=[0],\n",
       "    test_mask=[0]\n",
       "  },\n",
       "  \u001b[1m(all, SOC_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(all, SPA_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(all, TEM_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(all, simp_link, all)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 4361],\n",
       "    edge_attr=[4361]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 4425],\n",
       "    edge_attr=[4425]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 4411],\n",
       "    edge_attr=[4411]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 3988],\n",
       "    edge_attr=[3988]\n",
       "  },\n",
       "  \u001b[1m(vis_only, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 5057],\n",
       "    edge_attr=[5057]\n",
       "  },\n",
       "  \u001b[1m(vis_only, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 4804],\n",
       "    edge_attr=[4804]\n",
       "  },\n",
       "  \u001b[1m(vis_only, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 3774],\n",
       "    edge_attr=[3774]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 4425],\n",
       "    edge_attr=[4425]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 4421],\n",
       "    edge_attr=[4421]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_only)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SOC_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 4947],\n",
       "    edge_attr=[4947]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, SPA_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 5070],\n",
       "    edge_attr=[5070]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, TEM_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 5063],\n",
       "    edge_attr=[5063]\n",
       "  },\n",
       "  \u001b[1m(vis_tex, simp_link, vis_tex)\u001b[0m={\n",
       "    edge_index=[2, 0],\n",
       "    edge_attr=[0]\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sampled_hetero_data = next(iter(train_loader))\n",
    "batch = sampled_hetero_data\n",
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(3490)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch['vis_tex']['train_mask'].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([12351, 982])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.x_dict['vis_only'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}