Model 2 -- AnalogDiffusionFull.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1999104",
   "metadata": {},
   "source": [
    "# GraphGeneration: Modeling and design of hierarchical bio-inspired de novo spider web structures using deep learning and additive manufacturing \n",
    "\n",
    "Spider webs are incredible biological structures, comprising thin but strong silk filament and arranged into highly complex hierarchical architectures with striking mechanical properties (e.g., lightweight but high strength).  While simple 2D orb webs can easily be mimicked, the modeling and synthesis of artificial, bio-inspired 3D-based web structures is challenging, partly due to the rich set of design features. Here we use deep learning as a way to model and synthesize such 3D web structures, where generative models are conditioned based on key geometric parameters (incl.: average edge length, number of nodes, average node degree, and others). To identify construction principles, we use inductive representation sampling of large spider web graphs and develop and train three distinct conditional generative models to accomplish this task: 1) An analog diffusion model with sparse neighbor representation, 2) a discrete diffusion model with full neighbor representation, and 3) an autoregressive transformer architecture with full neighbor representation. We find that all three models can produce complex, de novo bio-inspired spider web mimics and successfully construct samples that meet the design conditioning that reflect key geometric features (including, the number of nodes,   spatial orientation, and edge lengths). We further present an algorithm that assembles inductive samples produced by the generative deep learning models into larger-scale structures based on a series of geometric design targets, including helical forms and parametric curves. \n",
    "\n",
    "[1] W. Lu, N.A. Lee, M.J. Buehler, \"Modeling and design of hierarchical bio-inspired de novo spider web structures using deep learning and additive manufacturing,\" PNAS, 120 (31) e2305273120, 2023, https://www.pnas.org/doi/10.1073/pnas.2305273120 \n",
    "\n",
    "## Model 2: Analog diffusion model with full neighbor representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26a6f27d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os,sys\n",
    "import math\n",
    "\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6420529",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def exists(val):\n",
    "    return val is not None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e62b3a1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b85dbe3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() \n",
    "                                  else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2398028",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#from tqdm import tqdm, trange\n",
    "from tqdm.notebook import trange, tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9183943f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16260708",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "available_gpus = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]\n",
    "available_gpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924ee37e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "num_of_gpus = torch.cuda.device_count()\n",
    "print(num_of_gpus)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd05726",
   "metadata": {},
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ed04b61",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader,Dataset\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c80f92cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    " \n",
    "import torchvision\n",
    " \n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    " \n",
    "from torch import nn\n",
    "from torch import optim, Tensor\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms, models\n",
    "\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import ExponentialLR, StepLR\n",
    "from functools import partial, wraps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1daa07e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Torch version:\", torch.__version__) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b0d9d1d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import QuantileTransformer\n",
    "from sklearn.preprocessing import RobustScaler\n",
    "import ast\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8116f29",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b18d29",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch_geometric.utils import degree\n",
    "from torch_geometric  import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18b66373",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch_geometric.data import Data\n",
    "from torch_geometric.data.datapipes import functional_transform\n",
    "from torch_geometric.transforms import BaseTransform, LinearTransformation\n",
    "import math\n",
    "import numbers\n",
    "import random\n",
    "from typing import Tuple, Union\n",
    "\n",
    " \n",
    "class RandomRotateLoc(BaseTransform):\n",
    "    r\"\"\"Rotates node positions around a specific axis by a randomly sampled\n",
    "    factor within a given interval (functional name: :obj:`random_rotate`).\n",
    "\n",
    "    Args:\n",
    "        degrees (tuple or float): Rotation interval from which the rotation\n",
    "            angle is sampled. If :obj:`degrees` is a number instead of a\n",
    "            tuple, the interval is given by :math:`[-\\mathrm{degrees},\n",
    "            \\mathrm{degrees}]`.\n",
    "        axis (int, optional): The rotation axis. (default: :obj:`0`)\n",
    "    \"\"\"\n",
    "    def __init__(self, degrees: Union[Tuple[float, float], float],\n",
    "                 axis: int = 0):\n",
    "        if isinstance(degrees, numbers.Number):\n",
    "            degrees = (-abs(degrees), abs(degrees))\n",
    "        assert isinstance(degrees, (tuple, list)) and len(degrees) == 2\n",
    "        self.degrees = degrees\n",
    "        self.axis = axis\n",
    "\n",
    "    def __call__(self, data: Data) -> Data:\n",
    "        degree = math.pi * random.uniform(*self.degrees) / 180.0\n",
    "        sin, cos = math.sin(degree), math.cos(degree)\n",
    "        \n",
    "        if data.pos.size(-1) == 2:\n",
    "            matrix = [[cos, sin], [-sin, cos]]\n",
    "        else:\n",
    "            if self.axis == 0:\n",
    "                matrix = [[1, 0, 0], [0, cos, sin], [0, -sin, cos]]\n",
    "            elif self.axis == 1:\n",
    "                matrix = [[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]]\n",
    "            else:\n",
    "                matrix = [[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]]\n",
    "                \n",
    "      \n",
    "        return LinearTransformationLoc(torch.tensor(matrix))(data)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return (f'{self.__class__.__name__}({self.degrees}, '\n",
    "                f'axis={self.axis})')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb73d2d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch_geometric.data import Data, HeteroData\n",
    "class LinearTransformationLoc(BaseTransform):\n",
    "    r\"\"\"Transforms node positions :obj:`data.pos` with a square transformation\n",
    "    matrix computed offline (functional name: :obj:`linear_transformation`)\n",
    "\n",
    "    Args:\n",
    "        matrix (Tensor): Tensor with shape :obj:`[D, D]` where :obj:`D`\n",
    "            corresponds to the dimensionality of node positions.\n",
    "    \"\"\"\n",
    "    def __init__(self, matrix: Tensor):\n",
    "        if not isinstance(matrix, Tensor):\n",
    "            matrix = torch.tensor(matrix)\n",
    "        assert matrix.dim() == 2, (\n",
    "            'Transformation matrix should be two-dimensional.')\n",
    "        assert matrix.size(0) == matrix.size(1), (\n",
    "            f'Transformation matrix should be square (got {matrix.size()})')\n",
    "\n",
    "        # Store the matrix as its transpose.\n",
    "        # We do this to enable post-multiplication in `__call__`.\n",
    "        self.matrix = matrix.t()\n",
    "      #  print (self.matrix.shape)\n",
    "\n",
    "    def __call__(\n",
    "        self,\n",
    "        data: Union[Data, HeteroData],\n",
    "    ) -> Union[Data, HeteroData]:\n",
    "        for store in data.node_stores:\n",
    "            if not hasattr(store, 'pos'):\n",
    "                continue\n",
    "\n",
    "            pos = store.pos.view(-1, 1) if store.pos.dim() == 1 else store.pos\n",
    "            assert pos.size(-1) == self.matrix.size(-2), (\n",
    "                'Node position matrix and transformation matrix have '\n",
    "                'incompatible shape')\n",
    "            # We post-multiply the points by the transformation matrix instead\n",
    "            # of pre-multiplying, because `pos` attribute has shape `[N, D]`,\n",
    "            # and we want to preserve this shape.\n",
    "           # print (self.matrix)\n",
    "            store.pos = pos @ self.matrix.to(pos.device, pos.dtype)\n",
    "      \n",
    "        return data\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f'{self.__class__.__name__}(\\n{self.matrix.cpu().numpy()}\\n)'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6334aee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "input, y_data,node_number_list,labels_y_txt, max_length , posdim_emb, max_neighbors =\\\n",
    "                                        torch.load('dataset_webs_medium.pt')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fb46ca7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Data format for input\n",
    "#0: ordinal numbering of nodes\n",
    "#1,2,3: x,y,z coordinates\n",
    "#4,5,6,7,8,9 - list of neighbors\n",
    "#10,11,12,13,14,1 - list of distances to neighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52aed857",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print (input.shape)\n",
    "print (y_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "117e7b63",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.data.datapipes import functional_transform\n",
    "from torch_geometric.transforms import BaseTransform, LinearTransformation\n",
    "import math\n",
    "import numbers\n",
    "import random\n",
    "from typing import Tuple, Union\n",
    "\n",
    "class RandomRotateDiffusion(BaseTransform):\n",
    "    r\"\"\"Rotates node positions around a specific axis by a randomly sampled\n",
    "    factor within a given interval (functional name: :obj:`random_rotate`).\n",
    "\n",
    "    Args:\n",
    "        degrees (tuple or float): Rotation interval from which the rotation\n",
    "            angle is sampled. If :obj:`degrees` is a number instead of a\n",
    "            tuple, the interval is given by :math:`[-\\mathrm{degrees},\n",
    "            \\mathrm{degrees}]`.\n",
    "        axis (int, optional): The rotation axis. (default: :obj:`0`)\n",
    "    \"\"\"\n",
    "    def __init__(self, degrees: Union[Tuple[float, float], float],\n",
    "                 axis: int = 0):\n",
    "        if isinstance(degrees, numbers.Number):\n",
    "            degrees = (-abs(degrees), abs(degrees))\n",
    "        assert isinstance(degrees, (tuple, list)) and len(degrees) == 2\n",
    "        self.degrees = degrees\n",
    "        self.axis = axis\n",
    "\n",
    "    def __call__(self, pos):\n",
    "        degree = math.pi * random.uniform(*self.degrees) / 180.0\n",
    "        sin, cos = math.sin(degree), math.cos(degree)\n",
    "        \n",
    "        \n",
    "\n",
    "        if data.pos.size(-1) == 2:\n",
    "            matrix = [[cos, sin], [-sin, cos]]\n",
    "        else:\n",
    "            if self.axis == 0:\n",
    "                matrix = [[1, 0, 0], [0, cos, sin], [0, -sin, cos]]\n",
    "            elif self.axis == 1:\n",
    "                matrix = [[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]]\n",
    "            else:\n",
    "                matrix = [[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]]\n",
    "        matrix=torch.Tensor (matrix)        \n",
    "       # print (matrix)\n",
    "        pos=  pos @  matrix#.to(pos.device, pos.dtype)\n",
    "        return pos\n",
    "    \n",
    "    #    return LinearTransformationLoc(torch.tensor(matrix))(data)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return (f'{self.__class__.__name__}({self.degrees}, '\n",
    "                f'axis={self.axis})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e44441c",
   "metadata": {},
   "source": [
    "### Normalize data and prepare dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16e5bfc8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class RegressionDataset(Dataset):\n",
    "    \n",
    "    def __init__(self, X_data, y_data, node_number_list, degrees=0, jitter=0, clamp_neighbors=True,\n",
    "                enforce_symm=True,):\n",
    "        self.X_data = X_data\n",
    "        self.y_data = y_data\n",
    "        self.node_number_list=node_number_list\n",
    "        \n",
    "        self.degrees=degrees\n",
    "        self.jitter=jitter\n",
    "        self.enforce_symm=enforce_symm\n",
    "        \n",
    "        self.randomrotatex= RandomRotateDiffusion (degrees=self.degrees, axis=0)\n",
    "        self.randomrotatey= RandomRotateDiffusion (degrees=self.degrees, axis=1)\n",
    "        self.randomrotatez= RandomRotateDiffusion (degrees=self.degrees, axis=2)\n",
    "        self.clamp_neighbors=clamp_neighbors\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        self.X_data_cl=self.X_data.clone()\n",
    "        resroundL = self.node_number_list [index] #torch.nonzero(self.X_data_cl[index,:,1])[-1]\n",
    "        \n",
    "        if self.degrees>0:\n",
    "           \n",
    "            resroundL = torch.nonzero(self.X_data_cl[index,:,1])[-1]\n",
    "            pos=self.X_data_cl[index,:resroundL,1:4]\n",
    "           \n",
    "            pos =self.randomrotatex(pos)\n",
    "            pos =self.randomrotatey(pos)\n",
    "            pos =self.randomrotatez(pos)\n",
    "            self.X_data_cl[index,:resroundL,1:4]=pos\n",
    "            \n",
    "        if self.jitter >0:\n",
    "            dx=torch.randn(resroundL)*self.jitter \n",
    "            dy=torch.randn(resroundL)*self.jitter \n",
    "            dz=torch.randn(resroundL)*self.jitter \n",
    "            \n",
    "            dx=torch.clamp(dx, min=-self.jitter, max=self.jitter ) \n",
    "            dy=torch.clamp(dy, min=-self.jitter, max=self.jitter ) \n",
    "            dz=torch.clamp(dz, min=-self.jitter, max=self.jitter ) \n",
    " \n",
    "            self.X_data_cl[index,:resroundL,1]=self.X_data_cl[index,:resroundL,1]+dx\n",
    "            self.X_data_cl[index,:resroundL,2]=self.X_data_cl[index,:resroundL,2]+dy\n",
    "            self.X_data_cl[index,:resroundL,3]=self.X_data_cl[index,:resroundL,3]+dz\n",
    "            \n",
    "        self.X_data_cl[index,resroundL:,:]=0\n",
    "            \n",
    "       \n",
    "        dist_matrix=-torch.ones (max_length, max_length)#.to(device)\n",
    "        \n",
    "        for i in range (max_length):\n",
    "             \n",
    "            #neighbors are stored in result [4,5,..., 9]\n",
    "            for j in range (max_neighbors):\n",
    "                #neigh_j=i-int (result[ 3+j, i] )   +1\n",
    "                neigh_j=self.X_data_cl[index, i, 4+j] \n",
    "                \n",
    "                if neigh_j !=0: #zeros are padded values...\n",
    "                     \n",
    "                    neighn=  neigh_j.long()\n",
    "\n",
    "                    if self.clamp_neighbors:\n",
    "                        neighn=max( 1, min (neighn, max_length) )\n",
    "                    \n",
    "                    #print (dist_matrix.shape, j, neighn)\n",
    "                    #dist_matrix[neighn-1, i] = 1\n",
    "                    dist_matrix[i, neighn-1] = 1\n",
    "                    if self.enforce_symm:\n",
    "                        #dist_matrix[i, neighn-1] = 1\n",
    "                        dist_matrix[neighn-1, i] = 1\n",
    "                    \n",
    "                  \n",
    "        output=  torch.cat( (self.X_data_cl[index,:, :4],   dist_matrix), 1)\n",
    "        output = output.permute (1,0)\n",
    "        \n",
    "        \n",
    "        return output, self.y_data[index]\n",
    "        \n",
    "    def __len__ (self):\n",
    "        return len(self.X_data)\n",
    "\n",
    "def scale_data(image2, maxv, minv): \n",
    "    return (image2 -minv)/(maxv-minv) * 2. - 1.0\n",
    " \n",
    "def unscale_data(image2, maxv, minv):\n",
    "   \n",
    "    image2=(image2 +1. )/ 2. * (maxv-minv)+minv \n",
    "    return image2\n",
    "\n",
    "\n",
    "def normalize_data (input, y_data, X_min=None, X_max=None, y_min=None, y_max=None,\n",
    "                   \n",
    "                    X_max_neigh=None, X_min_neigh=None,\n",
    "                    Xscale=0):\n",
    "    if X_min==None:\n",
    "        X_min=input[:,:,1:4].min() \n",
    "    else:\n",
    "        print (\"use provided X_min\", X_min)\n",
    "    if X_max==None:\n",
    "        X_max=input[:,:,1:4].max() \n",
    "    else:\n",
    "        print (\"use provided X_max\", X_max)\n",
    "\n",
    "    input[:,:,1:4]=(input[:,:,1:4]-X_min)/(X_max-X_min)*(2-2*Xscale)-(1-Xscale) #Normalize range -1 to 1\n",
    "\n",
    "    print (\"Check X after norm  \", input[:,:,1:4].min(), input[:,:,1:4].max())\n",
    "\n",
    "    print (\"Check X_neigh before norm  \", input[:,:,4:9].min(), input[:,:,4:9].max())  \n",
    "    \n",
    "\n",
    "    if y_min==None:\n",
    "        y_min=[]\n",
    "        for i in range (y_data.shape[1]):\n",
    "            y_min.append(y_data[:,i].min())\n",
    "        \n",
    "    else:\n",
    "        print (\"use provided y_min\", y_min)\n",
    "    if y_max==None:\n",
    "        y_max=[]\n",
    "        for i in range (y_data.shape[1]):\n",
    "            y_max.append(y_data[:,i].max())\n",
    "        \n",
    "    else:\n",
    "        print (\"use provided y_max\", y_max)\n",
    "    for i in range (y_data.shape[1]):\n",
    "        y_data[:,i]=(y_data[:,i]-y_min[i] )/(y_max[i] -y_min[i])*2-1 #Normalize range -1 to 1\n",
    "    \n",
    "\n",
    "    print (\"Check y_data after norm  \", y_data.min(), y_data.max())\n",
    "    return input, y_data, X_min, X_max, y_min, y_max#,X_min_neigh, X_max_neigh\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "457358b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_scaled, y_data_scaled, X_min, X_max, y_min, y_max = normalize_data (input, y_data,  Xscale=0. )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aba3a4dc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_scaled.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fef0b41",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def pad_sequence_end (output_xyz, max_length_l):         #pad\n",
    "    output=torch.zeros((output_xyz.shape[0] , max_length_l,  output_xyz.shape[2])).to(device)\n",
    "    output[:,:output_xyz.shape[-2],:]=output_xyz  \n",
    "    return output.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79d604c6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "max_neighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75a69459",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_data_loaders (X_scaled,  y_data_scaled, node_number_list, split=0.1, batch_size_=16):\n",
    "\n",
    "    X_train, X_test, y_train, y_test, node_number_list_train, node_number_list_test = train_test_split(X_scaled, \n",
    "                                                                                                       y_data_scaled ,\n",
    "                                                                                                       node_number_list,\n",
    "                                                                                                       test_size=split,random_state=235)\n",
    "\n",
    "\n",
    "    print (f\"Shapes= {X_scaled.shape}, {y_data_scaled.shape}\")\n",
    "    \n",
    "     \n",
    "    print(X_train.shape)\n",
    "    print(y_train.shape)\n",
    "    print(X_test.shape)\n",
    "    print(y_test.shape)\n",
    "    train_dataset = RegressionDataset(X_train, y_train, node_number_list_train, degrees=0, jitter=0.0,\n",
    "                                     enforce_symm=False) #/ynormfac)\n",
    "\n",
    "    test_dataset = RegressionDataset(X_test,y_test,node_number_list_test, degrees=0, jitter=0.0,\n",
    "                                    enforce_symm=False)\n",
    "\n",
    "\n",
    "    train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size_, shuffle=True)\n",
    "    train_loader_noshuffle = DataLoader(dataset=train_dataset, batch_size=batch_size_, shuffle=False)\n",
    "    test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size_)\n",
    "\n",
    "    return train_loader,train_loader_noshuffle, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ea3b2b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_loader,train_loader_noshuffle, test_loader= get_data_loaders (X_scaled,  y_data_scaled,node_number_list, split=0.1, \n",
    "                                                                    batch_size_=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6769a023",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_min, X_max, y_min, y_max"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a00e032",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39b2cc38",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_loop (model,\n",
    "                train_loader,test_loader,\n",
    "                optimizer=None,\n",
    "                print_every=10,\n",
    "                epochs= 300,\n",
    "                start_ep=0,\n",
    "                start_step=0,\n",
    "                train_unet_number=1,\n",
    "                print_loss=1000,\n",
    "                plot_unscaled=False,\n",
    "                save_model=False,\n",
    "                cond_scales=[7.5], #list of cond scales - each sampled...\n",
    "                num_samples=2, #how many samples produced every time tested.....\n",
    "                timesteps=10,\n",
    "                save_loss_images=False,clamp=False,\n",
    "                corplot=False,show_neighbors=False,\n",
    "                xyz_and_graph=False,\n",
    "                 dist_matrix_threshold=0.99,\n",
    "                discretize=True,\n",
    "                get_length_from_result=True,\n",
    "                only_continuous=True,\n",
    "                \n",
    "               ):\n",
    "    \n",
    "    #print_loss=1\n",
    "    #if not exists (optimizer):\n",
    "    #        print (\"ERROR: need to provide optimizer.\")\n",
    "    steps=start_step\n",
    "    start = time.time()\n",
    "    \n",
    "\n",
    "    loss_total=0\n",
    "    for e in range(1, epochs+1):\n",
    "            start = time.time()\n",
    "\n",
    "            torch.cuda.empty_cache()\n",
    "          \n",
    "            # TRAINING\n",
    "            train_epoch_loss = 0\n",
    "            model.train()\n",
    "            \n",
    "\n",
    "            for item  in train_loader:\n",
    "\n",
    "\n",
    "                X_train_batch= item[0].to(device)\n",
    "                y_train_batch=item[1].to(device)\n",
    "                optimizer.zero_grad()\n",
    "                loss=model (  y_train_batch , X_train_batch) #( batch_sentences, output )\n",
    "                loss.backward( )\n",
    "                \n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
    "\n",
    "                optimizer.step()\n",
    "\n",
    "                loss_total=loss_total+loss.item()\n",
    "\n",
    "\n",
    "                if steps % print_every == 0:\n",
    "                    print(\".\", end=\"\")\n",
    "\n",
    "                if steps>0:\n",
    "                    if steps % print_loss == 0:\n",
    "                        norm_loss=loss_total/print_loss\n",
    "                        print (f\"\\nTOTAL LOSS at epoch={e}, step={steps}: {norm_loss}\")\n",
    "\n",
    "                        loss_list.append (norm_loss)\n",
    "                        loss_total=0\n",
    "\n",
    "                        plt.plot (loss_list, label='Loss')\n",
    "                        plt.legend()\n",
    "\n",
    "                        if save_loss_images:\n",
    "                            outname = prefix+ f\"loss_{e}_{steps}.jpg\"\n",
    "                            plt.savefig(outname, dpi=200)\n",
    "                        plt.show()\n",
    "                        \n",
    "                        \n",
    "                        sample_loop (model,device ,\n",
    "                            test_loader,\n",
    "                            cond_scales=cond_scales, #list of cond scales - each sampled...\n",
    "                            num_samples=num_samples, #how many samples produced every time tested.....\n",
    "                            timesteps=timesteps,clamp=clamp,corplot=corplot,\n",
    "                            save_img=save_loss_images,show_neighbors=show_neighbors,\n",
    "                            flag=steps,xyz_and_graph=xyz_and_graph,\n",
    "                            clamp_round_results=True, enforce_symmetry=False,\n",
    "                            dist_matrix_threshold= dist_matrix_threshold , discretize=discretize,\n",
    "                            get_length_from_result=get_length_from_result,\n",
    "                            only_continuous=only_continuous,\n",
    "                                    )\n",
    "                        \n",
    "                        print (f\"\\n\\n-------------------\\nTime passed for {print_loss} epochs at {steps} = {(time.time()-start)/60} mins\\n-------------------\")\n",
    "                        start = time.time()\n",
    "                        if save_model:\n",
    "                         \n",
    "                            fname=f\"{prefix}statedict_save-model-epoch_{e}.pt\"\n",
    "                            torch.save(model.state_dict(), fname)\n",
    "                            print (f\"Model saved: \", fname)\n",
    "                steps=steps+1\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a0fbb2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "y_min, y_max=np.array (y_min), np.array (y_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d5d1072",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "y_min, y_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b891c50d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sns.set_style(\"whitegrid\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8560057a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def delete_isolated (a, clean_symmetrical=True):\n",
    "    a=a.clone ()\n",
    "    delete_list=[]\n",
    "    #b=torch.nonzero(a)[:,-1] #find node number list\n",
    "    b=torch.nonzero(a)[:,1] #find node number list\n",
    "    \n",
    "    c=torch.unique (b)\n",
    "    print (c)\n",
    "    #for i in range (c.shape[0]-1):\n",
    "    for i in range (c.shape[0]-1):\n",
    "        node_id=c[i]\n",
    "        node_id_pone=c[i+1]\n",
    "        \n",
    "        if node_id_pone!=node_id+1:\n",
    "            a [:,node_id_pone]=0\n",
    "            print (\"#########################\", node_id, node_id_pone)\n",
    "            if clean_symmetrical:\n",
    "                a [node_id_pone,:]=0\n",
    "            \n",
    "    return a\n",
    "def select_first_continuous (a, clean_symmetrical=True):\n",
    "    a=a.clone ()\n",
    "    delete_list=[]\n",
    "    \n",
    "    b=torch.nonzero(a)[:,1] #find node number list\n",
    "    \n",
    "    c=torch.unique (b) #get unique entries\n",
    "    print (c)\n",
    "    #for i in range (c.shape[0]-1):\n",
    "    for i in range (c.shape[0]-1):\n",
    "        node_id=c[i]\n",
    "        node_id_pone=c[i+1]\n",
    "        \n",
    "        if node_id_pone!=node_id+1:\n",
    "            a [:,node_id_pone:]=0\n",
    "            print (\"#########################\", node_id, node_id_pone)\n",
    "            if clean_symmetrical:\n",
    "                a [node_id_pone:,:]=0\n",
    "            \n",
    "    return a\n",
    "                        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a446705",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import r2_score\n",
    "\n",
    "def sample_loop (model,device,\n",
    "                train_loader,\n",
    "                cond_scales=[7.5], #list of cond scales - each sampled...\n",
    "                num_samples=2, #how many samples produced every time tested.....\n",
    "                timesteps=100,\n",
    "                 flag=0,clamp=False,\n",
    "                 corplot=False,\n",
    "                 save_img=False,\n",
    "                 show_neighbors=False,\n",
    "                xyz_and_graph=False,GED=False,\n",
    "                 max_length_considered=None,#consider predictions only up to a certain token number \n",
    "                     enforce_symmetry = False, #if True: make distance matrix symmetric\n",
    "                 clamp_round_results=True,dist_matrix_threshold=0.25,\n",
    "                 discretize = True, #whether or not to discretize to 0 or 1\n",
    "                 get_length_from_result=False,\n",
    "                 delete_isolated_nodes=True,\n",
    "                 only_continuous=True,\n",
    "                 show_colorbar=False,\n",
    "               ):\n",
    "    steps=0\n",
    "    e=flag\n",
    "    \n",
    "    if not exists (max_length_considered):\n",
    "        max_length_considered=max_length\n",
    "    for item  in train_loader:\n",
    "            \n",
    "            X_train_batch= item[0]\n",
    "            y_train_batch=item[1].to(device)\n",
    "            \n",
    "            #X_train_batch=torch.permute(X_train_batch, (0,2,1)  )\n",
    "            print (\"###\", X_train_batch.shape)\n",
    "            GT=y_train_batch.cpu().detach().unsqueeze(1) \n",
    "\n",
    "                \n",
    "            ####\n",
    "            num_samples = min (num_samples,y_train_batch.shape[0] )\n",
    "            print (f\"Producing {num_samples} samples...\")\n",
    "            if xyz_and_graph:\n",
    "                y_data_coll_pred=[]\n",
    "                y_data_coll_GT=[]  \n",
    "            for iisample in range (len (cond_scales)):\n",
    "                result=model.sample ( y_train_batch,device,\n",
    "                                         cond_scale=cond_scales[iisample],\n",
    "                                         timesteps=timesteps,clamp=clamp\n",
    "                                          )\n",
    "                \n",
    "                result=result.cpu()#.numpy()\n",
    "                \n",
    "                 \n",
    "                \n",
    "                if enforce_symmetry:\n",
    "                    result[:, 3:3+max_length, :]=0.5*(result[:, 3:3+max_length, :]+\\\n",
    "                                                      torch.transpose (result[:, 3:3+max_length, :], 1, 2 ) )\n",
    "                if clamp_round_results:\n",
    "                    result[:, 3:3+max_length, :]=torch.clamp(result[:, 3:3+max_length, :], 0, 1) \n",
    "                    X_train_batch[:, 4:4+max_length, :]=torch.clamp (X_train_batch[:, 4:4+max_length, :], 0, 1)\n",
    "                \n",
    "                if discretize:\n",
    "                    result[:, 3:3+max_length, :] [result[:, 3:3+max_length, :]< dist_matrix_threshold]=0\n",
    "                    result[:, 3:3+max_length, :] [result[:, 3:3+max_length, :]>= dist_matrix_threshold]=1\n",
    "                    \n",
    "\n",
    "                \n",
    "                result=result.cpu() \n",
    "                \n",
    "                print (f\"sample result for cond_scale={cond_scales[iisample]}....\", result.shape, \"GT shape \", GT.shape)\n",
    "\n",
    "                #print (result.shape, GT.shape)\n",
    "               # X_train_batch= pad_sequence (X_train_batch,  max_length).cpu()\n",
    "                if xyz_and_graph:\n",
    "                    y_data_coll_pred=[]\n",
    "                    y_data_coll_GT=[]                \n",
    "                for samples in range  (num_samples):\n",
    "                    #print (\"############## shape: \", result[samples, 3:3+max_length, :].shape)\n",
    "                    if delete_isolated_nodes:\n",
    "                        result[samples, 3:3+max_length, :]=delete_isolated (result[samples, 3:3+max_length, :])\n",
    "\n",
    "                    if only_continuous:\n",
    "                         result[samples, 3:3+max_length, :]= select_first_continuous(result[samples, 3:3+max_length, :])\n",
    "                \n",
    "                    if get_length_from_result:\n",
    "                        \n",
    "                       \n",
    "                        GTroundL = torch.nonzero(X_train_batch[samples, 4:4+max_length, :])[:,1].flatten().max()+1\n",
    "                        resroundL = torch.nonzero(result[samples, 3:3+max_length, :])[:,1].flatten().max()+1\n",
    "                     \n",
    "                    else:\n",
    "                        GTroundL=max_length\n",
    "                        resroundL=max_length\n",
    "                    \n",
    "                    #resroundL= (resround==-1).nonzero(as_tuple=True)  [0]\n",
    "                      \n",
    "                    fig, ax = plt.subplots(1, 2, figsize=(10, 6) , subplot_kw=dict(projection='3d'))\n",
    "                    \n",
    "                     \n",
    "                    xs=X_train_batch[samples, 1, :GTroundL]\n",
    "                    ys=X_train_batch[samples, 2, :GTroundL]\n",
    "                    zs=X_train_batch[samples, 3, :GTroundL]\n",
    "                    m=6\n",
    "                    ax[0].scatter(xs, ys, zs, c='red', s=m, marker=\"o\")\n",
    "\n",
    "                    ax[0].set_xlabel('X')\n",
    "                    ax[0].set_ylabel('Y')\n",
    "                    ax[0].set_zlabel('Z')\n",
    "                    ax[0].set_title('GT')\n",
    "                    ax[0].set_xlim(-1,1)\n",
    "                    ax[0].set_ylim(-1,1)\n",
    "                    ax[0].set_zlim(-1,1)\n",
    "                    \n",
    "                    \n",
    "                    \n",
    "                    xs=result[samples, 0, :resroundL]\n",
    "                    ys=result[samples, 1, :resroundL]\n",
    "                    zs=result[samples, 2, :resroundL]\n",
    "                    m=6\n",
    "                    ax[1].scatter(xs, ys, zs, c='red', s=m, marker=\"o\")\n",
    "\n",
    "                    ax[1].set_xlabel('X')\n",
    "                    ax[1].set_ylabel('Y')\n",
    "                    ax[1].set_zlabel('Z')\n",
    "                    ax[1].set_xlim(-1,1)\n",
    "                    ax[1].set_ylim(-1,1)\n",
    "                    ax[1].set_zlim(-1,1)                    \n",
    "                    \n",
    "                    \n",
    "                    ax[1].set_title('Prediction')\n",
    "                    if save_img:\n",
    "                            outname = prefix+ f\"img_{flag}.png\"\n",
    "                            plt.savefig(outname, dpi=300)\n",
    "                    \n",
    "                    plt.show()\n",
    "                    \n",
    "                    fig, ax = plt.subplots(1, 2, figsize=(10, 6)  )\n",
    "                    \n",
    "                     \n",
    "                    xs=X_train_batch[samples, 1, :GTroundL]\n",
    "                    ys=X_train_batch[samples, 2, :GTroundL]\n",
    "                    zs=X_train_batch[samples, 3, :GTroundL]\n",
    "                    m=2\n",
    "                    ax[0].plot(xs, ys , 'bo', markersize=m, label='Y over X')\n",
    "                    ax[0].plot(xs, zs,'ro', markersize=m, label='Z over X')\n",
    "\n",
    "                    ax[0].set_xlabel('X')\n",
    "                    ax[0].set_ylabel('Y/Z')\n",
    "                   # ax[0].set_zlabel('Z')\n",
    "                    ax[0].set_title('GT')\n",
    "                    ax[0].axis('square')\n",
    "                    ax[0].set_xlim(-1,1)\n",
    "                    ax[0].set_ylim(-1,1)\n",
    "                    ax[0].legend()\n",
    "                    \n",
    "                   # ax[0].set_zlim(-1,1)\n",
    "                    \n",
    "                    xs=result[samples, 0, :resroundL]\n",
    "                    ys=result[samples, 1, :resroundL]\n",
    "                    zs=result[samples, 2, :resroundL]\n",
    "                     \n",
    "                    ax[1].plot(xs, ys ,'bo', markersize=m, label='Y over X')\n",
    "                    ax[1].plot(xs, zs,'ro', markersize=m, label='Z over X')\n",
    "\n",
    "                    ax[1].set_xlabel('X')\n",
    "                    ax[1].set_ylabel('Y/Z')\n",
    "                    #ax[1].set_zlabel('Z')\n",
    "                    ax[1].axis('square')\n",
    "                    ax[1].set_xlim(-1,1)\n",
    "                    ax[1].set_ylim(-1,1)\n",
    "                    #ax[1].set_zlim(-1,1)                    \n",
    "                    ax[1].legend()\n",
    "                    \n",
    "                    \n",
    "                    ax[1].set_title('Prediction')\n",
    "                    if save_img:\n",
    "                            outname = prefix+ f\"img_proj_{flag}.png\"\n",
    "                            plt.savefig(outname, dpi=300)\n",
    "                    \n",
    "                    plt.show()\n",
    "                    \n",
    "                    \n",
    "                    if show_neighbors:\n",
    "                        fig, ax = plt.subplots(1, 2, figsize=(10, 6)  )\n",
    "                    \n",
    "                        ax[1].imshow (result [samples, 3:3+max(GTroundL,resroundL), :max(GTroundL,resroundL)])\n",
    "                        ax[0].imshow (X_train_batch [samples, 4:4+max(GTroundL,resroundL), :max(GTroundL,resroundL)])\n",
    "                        ax[1].set_title('Prediction')\n",
    "                        ax[0].set_title('GT')\n",
    "                      \n",
    "                        plt.show()\n",
    "                        \n",
    "                        fig, ax = plt.subplots(1, 2, figsize=(10, 6)  )\n",
    "                    \n",
    "                        shw1=ax[1].imshow (result [samples, 3:3+max_length, :])\n",
    "                        shw0=ax[0].imshow (X_train_batch [samples, 4:4+max_length, :]) #X_train_batch[0,4:,:]\n",
    "                        ax[1].set_title('Prediction')\n",
    "                        ax[0].set_title('GT')    \n",
    "                         \n",
    "                        if show_colorbar:\n",
    "                            bar0 = plt.colorbar(shw0)\n",
    "                            bar0.set_label('GT')\n",
    "                            bar1 = plt.colorbar(shw1)\n",
    "                            bar1.set_label('Prediction')                        \n",
    "                        #ax[0].grid(False)\n",
    "                       # ax[1].grid(False)\n",
    "                        plt.show()\n",
    "                        \n",
    "                    if xyz_and_graph:\n",
    "                        G_res, data_res, y_data_pred =construct_xyz_and_graph (result[samples,:3+resroundL,:resroundL].squeeze(),\n",
    "                                                 fname_root=f'{prefix}graph_xyz_{samples}_{flag}_{steps}',\n",
    "                                                                label='Prediction', #limits=[X_min*1.2, X_max*1.2],\n",
    "                                                                dist_matrix=True)\n",
    "                        G_GT, data_GT, y_data_GT=construct_xyz_and_graph (X_train_batch[samples, 1:4+GTroundL, :GTroundL].squeeze(),\n",
    "                                                 fname_root=f'{prefix}graph_GT_xyz_{samples}_{flag}_{steps}',\n",
    "                                                              label='GT',dist_matrix=True)\n",
    "                        \n",
    "                        \n",
    "\n",
    "                        plt.plot ( y_data_GT, y_data_pred, '.',label='Graph properties (GT vs predicted)',markersize=12 )\n",
    "                        plt.legend()\n",
    "                        plt.xlabel ('GT')\n",
    "                        plt.ylabel ('Predicted')\n",
    "                        min_v,max_v=min (min(y_data_GT), min (y_data_pred)),max (max(y_data_GT), max (y_data_pred))\n",
    "                        plt.plot([min_v, max_v], [min_v, max_v], ls=\"--\", c=\".3\")\n",
    "                        plt.axis ('square')\n",
    "                        plt.show()\n",
    "                        \n",
    "                        #print (y_max, y_min,GT.shape[2])\n",
    "                        for i in range (len(y_min)):\n",
    "                            y_data_pred[i]=scale_data(y_data_pred[i], y_max[i],y_min[i]) \n",
    "                            y_data_GT[i]=scale_data(y_data_GT[i], y_max[i],y_min[i]) \n",
    "                            \n",
    "                        plt.plot ( y_data_GT, y_data_pred, '.',label='Graph properties (GT vs predicted)',markersize=12 )\n",
    "                        plt.legend()\n",
    "                        plt.xlabel ('GT')\n",
    "                        plt.ylabel ('Predicted')\n",
    "                        plt.plot([-1, 1], [-1, 1], ls=\"--\", c=\".3\")\n",
    "                        plt.axis ('square')\n",
    "                        plt.show()\n",
    "                        \n",
    "                        y_data_coll_pred.append (y_data_pred)\n",
    "                        y_data_coll_GT.append (y_data_GT)#.cpu().numpy())\n",
    "                        \n",
    "                        if GED:\n",
    "                            GED=nx.graph_edit_distance (G_res, G_GT)\n",
    "                            print (\"Graph edit distance=\", GED)\n",
    "                        print (f\"\")\n",
    "                        \n",
    "                        \n",
    "                        \n",
    "                    print (X_train_batch.shape,result.shape )\n",
    "                    if corplot:\n",
    "                        plt.plot (X_train_batch[samples, 1:4, :].flatten().detach().cpu(), \n",
    "                                  result[samples, 0:3, :].flatten() , '.',label='all',markersize=6 )\n",
    "                        plt.plot (X_train_batch[samples, 1, :].flatten().detach().cpu(), \n",
    "                                  result[samples, 0, :].flatten() , '.',label='x',markersize=2 )\n",
    "                        plt.plot (X_train_batch[samples, 2, :].flatten().detach().cpu(), \n",
    "                                  result[samples, 1, :].flatten() , '.',label='y',markersize=2 )\n",
    "                        plt.plot (X_train_batch[samples, 3, :].flatten().detach().cpu(), \n",
    "                                  result[samples, 2, :].flatten() , '.',label='z',markersize=2 )\n",
    "                        \n",
    "                        plt.legend()\n",
    "                        \n",
    "                        plt.plot([-1, 1], [-1, 1], ls=\"--\", c=\".3\")\n",
    "                        plt.axis ('square')\n",
    "                        plt.show()\n",
    "\n",
    "                        \n",
    "                y_data_coll_pred=np.array(y_data_coll_pred).flatten()\n",
    "                y_data_coll_GT=np.array(y_data_coll_GT).flatten()        \n",
    "                \n",
    "                print (\"collected shape \", y_data_coll_pred.shape, y_data_coll_GT.shape)\n",
    "                \n",
    "                print (\"collected  \", y_data_coll_pred, y_data_coll_GT)\n",
    "                \n",
    "                R2=r2_score(y_data_coll_GT, y_data_coll_pred)\n",
    "                print (\"OVERALL R2: \", R2)\n",
    "                plt.plot ( y_data_coll_GT, y_data_coll_pred, '.', label='Graph properties (GT vs predicted)',markersize=3 )\n",
    "                plt.legend()\n",
    "                plt.xlabel ('GT')\n",
    "                plt.ylabel ('Predicted')\n",
    "                plt.plot([-1, 1], [-1, 1], ls=\"--\", c=\".3\")\n",
    "                plt.axis ('square')\n",
    "                plt.title (\"Correlation prediction vs. GT\")\n",
    "                plt.show()                        \n",
    "                        \n",
    "                        \n",
    "            steps=steps+1\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9b7744b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_sample_cond (model,\n",
    "              \n",
    "                cond_scales=[7.5], #list of cond scales - each sampled...\n",
    "               # num_samples=2, #how many samples produced every time tested.....\n",
    "                timesteps=100,\n",
    "                 flag=0,clamp=False,\n",
    "                 corplot=False,\n",
    "                 save_img=False,\n",
    "                 show_neighbors=False,\n",
    "                xyz_and_graph=False,GED=False,\n",
    "                cond=[1., .5, 1.],\n",
    "                     enforce_symmetry = False, #if True: make distance matrix symmetric\n",
    "                 clamp_round_results=True,dist_matrix_threshold=0.25,\n",
    "                 discretize = True, #whether or not to discretize to 0 or 1\n",
    "                 get_length_from_result=False,\n",
    "                 delete_isolated_nodes=True,\n",
    "                 only_continuous=True,\n",
    "                          show_colorbar=False,\n",
    "                          save_neigh_img=False,\n",
    "                # show_colorbar=False,                 \n",
    "               ):\n",
    "    steps=0\n",
    "    e=flag\n",
    "    #for item  in train_loader:\n",
    "\n",
    "\n",
    "    #X_train_batch= item[0]\n",
    "    y_train_batch=torch.Tensor (cond).to(device)\n",
    "    y_train_batch=y_train_batch.unsqueeze(0) \n",
    "    #X_train_batch=torch.permute(X_train_batch, (0,2,1)  )\n",
    "\n",
    "\n",
    "\n",
    "    for iisample in range (len (cond_scales)):\n",
    "        \n",
    "        samples=0\n",
    "        print (\"y_train_batch \", y_train_batch.shape)\n",
    "       \n",
    "\n",
    "        \n",
    "        result=model.sample ( y_train_batch,device,\n",
    "                                 cond_scale=cond_scales[iisample],\n",
    "                                 timesteps=timesteps,clamp=clamp,\n",
    "                                  )\n",
    "        #print (\"y_train_batch \", y_train_batch.shape)\n",
    "        #result= pad_sequence (result,  max_length)\n",
    "        #print (result.shape)\n",
    "        result=result.cpu()#.numpy()\n",
    "\n",
    "\n",
    "\n",
    "        if enforce_symmetry:\n",
    "            result[:, 3:3+max_length, :]=0.5*(result[:, 3:3+max_length, :]+\\\n",
    "                                              torch.transpose (result[:, 3:3+max_length, :], 1, 2 ) )\n",
    "        if clamp_round_results:\n",
    "            result[:, 3:3+max_length, :]=torch.clamp(result[:, 3:3+max_length, :], 0, 1) \n",
    "           # X_train_batch[:, 4:4+max_length, :]=torch.clamp (X_train_batch[:, 4:4+max_length, :], 0, 1)\n",
    "\n",
    "        if discretize:\n",
    "            result[:, 3:3+max_length, :] [result[:, 3:3+max_length, :]< dist_matrix_threshold]=0\n",
    "            result[:, 3:3+max_length, :] [result[:, 3:3+max_length, :]>= dist_matrix_threshold]=1\n",
    "\n",
    "\n",
    "\n",
    "        result=result.cpu() \n",
    "\n",
    "        print (f\"sample result for cond_scale={cond_scales[iisample]}....\", result.shape )\n",
    "\n",
    "        #print (result.shape, GT.shape)\n",
    "       # X_train_batch= pad_sequence (X_train_batch,  max_length).cpu()\n",
    "        if xyz_and_graph:\n",
    "            y_data_coll_pred=[]\n",
    "            #y_data_coll_GT=[]                \n",
    "\n",
    "        samples=0\n",
    "        #print (\"############## shape: \", result[samples, 3:3+max_length, :].shape)\n",
    "        if delete_isolated_nodes:\n",
    "            result[samples, 3:3+max_length, :]=delete_isolated (result[samples, 3:3+max_length, :])\n",
    "\n",
    "        if only_continuous:\n",
    "             result[samples, 3:3+max_length, :]= select_first_continuous(result[samples, 3:3+max_length, :])\n",
    "\n",
    "        if get_length_from_result:\n",
    "\n",
    "            #print (\"##\", torch.nonzero(result[samples, 3:3+max_length, :]).shape,  \n",
    "            #    result[samples, 3:3+max_length, :])\n",
    "\n",
    "          #  GTroundL = torch.nonzero(X_train_batch[samples, 4:4+max_length, :])[:,1].flatten().max()+1\n",
    "            resroundL = torch.nonzero(result[samples, 3:3+max_length, :])[:,1].flatten().max()+1\n",
    "            #only get length from LENGTH in x direction\n",
    "\n",
    "        else:\n",
    "            #GTroundL=max_length\n",
    "            resroundL=max_length \n",
    "        \n",
    "        \n",
    "        print (f\"sample result for cond_scale={cond_scales[iisample]}....\", result.shape )\n",
    "        print (\"Conditiioning: \", y_train_batch)\n",
    " \n",
    "        fig, ax = plt.subplots(1, 1, figsize=(5, 6) , subplot_kw=dict(projection='3d'))\n",
    "\n",
    "        ax=[ax]\n",
    "\n",
    "        xs=result[samples, 0, :resroundL]\n",
    "        ys=result[samples, 1, :resroundL]\n",
    "        zs=result[samples, 2, :resroundL]\n",
    "        m=6\n",
    "        ax[0].scatter(xs, ys, zs, c='red', s=m, marker=\"o\")\n",
    "\n",
    "        ax[0].set_xlabel('X')\n",
    "        ax[0].set_ylabel('Y')\n",
    "        ax[0].set_zlabel('Z')\n",
    "        ax[0].set_xlim(-1,1)\n",
    "        ax[0].set_ylim(-1,1)\n",
    "        ax[0].set_zlim(-1,1)                    \n",
    "\n",
    "\n",
    "        ax[0].set_title('Prediction')\n",
    "        if save_img:\n",
    "                outname = prefix+ f\"img_{flag}.png\"\n",
    "                plt.savefig(outname, dpi=300)\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "        fig, ax = plt.subplots(1, 1, figsize=(5, 6)  )\n",
    "\n",
    "        ax=[ax]\n",
    "\n",
    "        xs=result[samples, 0, :resroundL]\n",
    "        ys=result[samples, 1, :resroundL]\n",
    "        zs=result[samples, 2, :resroundL]\n",
    "\n",
    "        ax[0].plot(xs, ys ,'bo', markersize=m, label='Y over X')\n",
    "        ax[0].plot(xs, zs,'ro', markersize=m, label='Z over X')\n",
    "\n",
    "        ax[0].set_xlabel('X')\n",
    "        ax[0].set_ylabel('Y/Z')\n",
    "        #ax[1].set_zlabel('Z')\n",
    "        ax[0].axis('square')\n",
    "        ax[0].set_xlim(-1,1)\n",
    "        ax[0].set_ylim(-1,1)\n",
    "        #ax[1].set_zlim(-1,1)                    \n",
    "        ax[0].legend()\n",
    "\n",
    "\n",
    "        ax[0].set_title('Prediction')\n",
    "        if save_img:\n",
    "                outname = prefix+ f\"img_proj_{flag}.png\"\n",
    "                plt.savefig(outname, dpi=300)\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "        \n",
    "        \n",
    "        \n",
    "        if show_neighbors:\n",
    "            fig, ax = plt.subplots(1, 1, figsize=(5, 6)  )\n",
    "            ax=[ax]\n",
    "            ax[0].imshow (result [samples, 3:3+max(resroundL,resroundL), :max(resroundL,resroundL)])\n",
    "            #ax[0].imshow (X_train_batch [samples, 4:4+max(GTroundL,resroundL), :max(GTroundL,resroundL)])\n",
    "            #ax[0].set_title('GT')\n",
    "           # ax[0].grid(False)\n",
    "           # ax[1].grid(False)\n",
    "            if save_neigh_img:\n",
    "                plt.axis('off')\n",
    "                outname = prefix+ f\"neigh_proj_1_{flag}.png\"\n",
    "                plt.savefig(outname, dpi=200)\n",
    "            ax[0].set_title('Prediction')\n",
    "            plt.show()\n",
    "\n",
    "            fig, ax = plt.subplots(1, 1, figsize=(5, 6)  )\n",
    "            ax=[ax]\n",
    "            shw1=ax[0].imshow (result [samples, 3:3+max_length, :])\n",
    "            #shw0=ax[0].imshow (X_train_batch [samples, 4:4+max_length, :]) #X_train_batch[0,4:,:]\n",
    "            #ax[0].set_title('Prediction')\n",
    "            #ax[0].set_title('GT')    \n",
    "\n",
    "            if show_colorbar:\n",
    "                bar0 = plt.colorbar(shw0)\n",
    "                bar0.set_label('Prediction')\n",
    "                #bar1 = plt.colorbar(shw1)\n",
    "                #bar1.set_label('Prediction')                        \n",
    "            #ax[0].grid(False)\n",
    "           # ax[1].grid(False)\n",
    "            if save_neigh_img:\n",
    "                plt.axis('off')\n",
    "                outname = prefix+ f\"neigh_proj_2_{flag}.png\"\n",
    "                plt.savefig(outname, dpi=200)\n",
    "            ax[0].set_title('Prediction')\n",
    "            plt.show()\n",
    "            \n",
    "\n",
    "        if xyz_and_graph:\n",
    "            G_res, data_res, y_data_pred =construct_xyz_and_graph (result[samples,:3+resroundL,:resroundL].squeeze(),\n",
    "                                     fname_root=f'{prefix}graph_xyz_{samples}_{flag}_{steps}',\n",
    "                                                    label='Prediction', #limits=[X_min*1.2, X_max*1.2],\n",
    "                                                    dist_matrix=True)\n",
    "           \n",
    "\n",
    "            if GED:\n",
    "                GED=nx.graph_edit_distance (G_res, G_GT)\n",
    "                print (\"Graph edit distance=\", GED)\n",
    "            print (f\"\")\n",
    "\n",
    "\n",
    " \n",
    " \n",
    "    steps=steps+1\n",
    "    if xyz_and_graph:\n",
    "        return result[samples,:3+resroundL,:resroundL].squeeze().permute (1,0), G_res, data_res,y_data\n",
    "    else:\n",
    "        return result[samples,:3+resroundL,:resroundL].squeeze() \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cee3798a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    " \n",
    "from torch_geometric.utils import to_networkx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eb26a91",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "max_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6be203b4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    " \n",
    "from torch_geometric.utils import to_networkx\n",
    "\n",
    "def add_edge_to_graph(G, e1, e2, w):\n",
    "    G.add_edge(e1, e2, weight=w,\n",
    "              clamp_neighbors=True)\n",
    "def get_properties(item):\n",
    "    #print (item.edge_index, item.num_edges)\n",
    "    #print (item)\n",
    "    length=0\n",
    "    dx,dy,dz=0,0,0\n",
    "    \n",
    "    for jj in range (item.num_edges):\n",
    "        dx_=item.pos[item.edge_index[0,jj],0]-item.pos[item.edge_index[1,jj],0]\n",
    "        dy_=item.pos[item.edge_index[0,jj],1]-item.pos[item.edge_index[1,jj],1]\n",
    "        dz_=item.pos[item.edge_index[0,jj],2]-item.pos[item.edge_index[1,jj],2]\n",
    "        \n",
    "        dx=dx+dx_\n",
    "        dy=dy+dy_\n",
    "        dz=dz+dz_\n",
    "        \n",
    "        length=length + (dx_**2+dy_**2+dz_**2)**0.5\n",
    "        \n",
    "    avg_length = length /item.num_edges     \n",
    "    dx, dy, dz = dx/item.num_edges, dy/item.num_edges, dz/item.num_edges\n",
    "        \n",
    "    num_nodes=item.num_nodes\n",
    "    num_edges=item.num_edges\n",
    "    node_degree=item.num_edges / item.num_nodes\n",
    "    \n",
    "     \n",
    "    return avg_length.numpy(),dx.numpy(), dy.numpy(), dz.numpy(), num_nodes , num_edges, node_degree\n",
    "\n",
    "def construct_xyz_and_graph (result, fname_root='output',\n",
    "                             clamp_neighbors=True,label='Generated',\n",
    "                             limits=None,#axis limits for plot, \n",
    "                             GT_y=None,\n",
    "                             dist_matrix=False,\n",
    "                             \n",
    "                            ):\n",
    "    print (\"##############################################################################\")\n",
    "    print (\"Shape of data provided \", result.shape) \n",
    "    print (f\"Root file: {fname_root}\")\n",
    "    \n",
    "    result[ :3, :]=unscale_data(result[:3,: ], X_max.numpy(),X_min.numpy())\n",
    "\n",
    "    \n",
    "    #result=result.numpy()\n",
    "    xs=result[ 0, :]\n",
    "    ys=result[ 1, :]\n",
    "    zs=result[ 2, :]\n",
    "    \n",
    "    with open(fname_root+'.xyz', 'w') as f:\n",
    "        f.write('#ID x y z \\n')\n",
    "        for i in range (result.shape[1]):\n",
    "            f.write(f'{i+1} {xs[i]} {ys[i]} {zs[i]} \\n')\n",
    "    \n",
    "    \n",
    "    #G=nx.Graph()\n",
    "    node_list=[]\n",
    "    neighbor_list=[]\n",
    "    point_list=[]\n",
    "    \n",
    "    \n",
    "    #now prepare graph\n",
    "    for i in range (result.shape[1]):\n",
    "        node_list.append ( [xs[i], ys[i], zs[i] ])\n",
    "        point_list.append ( (xs[i], ys[i], zs[i]  ))\n",
    "        \n",
    "        if not dist_matrix:\n",
    "            #neighbors are stored in result [4,5,..., 9]\n",
    "            for j in range (max_neighbors):\n",
    "                #neigh_j=i-int (result[ 3+j, i] )   +1\n",
    "                neigh_j=result[ 3+j, i] \n",
    "                #print (neigh_j)\n",
    "                if neigh_j !=0: #zeros are padded values...\n",
    "                    #f neigh_j !=0:\n",
    "                    neighn=  neigh_j-1 #neigh_j is 1+node number (since it encodes 0s....as padding)\n",
    "\n",
    "                    if clamp_neighbors:\n",
    "                        neighn=max( 0, min (neighn, result.shape[1]-1) )\n",
    "\n",
    "\n",
    "                    neighbor_list.append ([i,neighn] )  #val=node1- node2 +1  --> node2= node1- val  +1\n",
    "                    \n",
    "        if dist_matrix:\n",
    "            for j in range (result.shape[1]):\n",
    "                #neigh_j=i-int (result[ 3+j, i] )   +1\n",
    "                neigh_j=result[ 3+j, i] \n",
    "                #print (neigh_j)\n",
    "                if neigh_j >0.5: #zeros are padded values... a value of >0.5 means a neighbor is found\n",
    "                    #f neigh_j !=0:\n",
    "                    neighn=  j\n",
    "\n",
    "                    neighbor_list.append ([i,neighn] )  #val=node1- node2 +1  --> node2= node1- val  +1\n",
    "                \n",
    "                #add_edge_to_graph(G, i, neighn, 1)\n",
    "    #pos = {point: point for point in point_list}\n",
    "    #print (pos)\n",
    "\n",
    "    with open(fname_root+'.dump', 'w') as f:\n",
    "        f.write(f'\\n{result.shape[1]} atoms\\n{len (neighbor_list)} bonds  \\n\\n1 atom types\\n1 bond types\\n\\n')\n",
    "        f.write(f'{min(xs[:])*1.1} {max(xs[:])*1.1} xlo xhi \\n{min(ys[:])*1.1} {max(ys[:])*1.1} ylo yhi \\n{min(zs[:])*1.1} {max(zs[:])*1.1} zlo zhi \\n\\n')\n",
    "        f.write(f'Masses \\n\\n1 100.00  \\n\\n')\n",
    "        f.write(f'Atoms  \\n\\n')\n",
    "\n",
    "\n",
    "        for i in range (result.shape[1]):\n",
    "            f.write(f'{i+1} 1 1 {xs[i]} {ys[i]} {zs[i]} \\n')\n",
    "        f.write(f'\\nBonds  \\n\\n')\n",
    "\n",
    "        for i in range (len (neighbor_list)):\n",
    "            f.write(f'{i+1} 1 {neighbor_list[i][0]+1} { max( 1, min (neighbor_list[i][1], result.shape[1]) )}   \\n')\n",
    "        f.write(f'\\n\\n')\n",
    "    \n",
    "\n",
    "    #rint (node_list)\n",
    "    #print (neighbor_list)\n",
    "    \n",
    "    neighbor_list=torch.Tensor (neighbor_list).long()\n",
    "    print (\"neighborlist shape: \", neighbor_list.shape)\n",
    "    #print (neighbor_list)\n",
    "    \n",
    "    fig = plt.figure(figsize=(8,8))\n",
    "    m=24\n",
    "    ax = fig.add_subplot(projection='3d')\n",
    "    ax.scatter(xs, ys, zs,c='red', s=m, marker=\"o\")\n",
    "    ax.set_proj_type('ortho')\n",
    "    ax.set_title (label)\n",
    "    ax.set_xlabel('X')\n",
    "    ax.set_ylabel('Y')\n",
    "    ax.set_zlabel('Z')\n",
    "    \n",
    "    print (\"STATS of neighbor list: max \", neighbor_list.max(), \"min \",  neighbor_list.min(), \"shape \", neighbor_list.shape)\n",
    "    \n",
    "    \n",
    "    for i in range (len (neighbor_list)):\n",
    "            \n",
    "            N1=neighbor_list[i][0] # because N1 and N2 refers to array indices, not note numbers\n",
    "            \n",
    "            N2=max( 0, min (neighbor_list[i][1], result.shape[1]-1) ) \n",
    "            #print (N1, N2)\n",
    "            ax.plot3D ([xs[N1], xs[N2]], [ys[N1], ys[N2]] , [zs[N1], zs[N2]], 'k-', linewidth=1, )\n",
    "           \n",
    "            \n",
    "    if limits != None:\n",
    "        ax.set_xlim(limits[0],limits[1])\n",
    "        ax.set_ylim(limits[0],limits[1])\n",
    "        ax.set_zlim(limits[0],limits[1])    \n",
    "        \n",
    "    plt.savefig (fname_root+'.png', dpi=400)\n",
    "    plt.savefig (fname_root+'.svg', dpi=400)\n",
    "    plt.show()\n",
    "    \n",
    "    print (\"Neighborlist shape COO^T format: \", neighbor_list.shape)\n",
    "    x = torch.tensor(node_list, dtype=torch.float)  #Node feature matrix with shape [num_nodes, num_node_features]\n",
    "    edge_index =neighbor_list.permute(1,0)  #Graph connectivity in COO format with shape [2, num_edges] and type torch.long\n",
    "    y = None #torch.tensor(graph_label, dtype=torch.float)  #Target to train against (may have arbitrary shape), e.g., node-level targets of shape [num_nodes, *] or graph-level targets of shape [1, *]\n",
    "    data = Data(x=x,pos=x, edge_index=edge_index, y=y)\n",
    "    \n",
    "    torch.save (data, fname_root+'.pt')\n",
    "    \n",
    "    print(f'Number of nodes: {data.num_nodes}')\n",
    "    print(f'Number of edges: {data.num_edges}')\n",
    "    print(f'Average node degree: {data.num_edges / data.num_nodes:.2f}')\n",
    "    G = to_networkx(data, to_undirected=False)\n",
    "    \n",
    "    print (\"##############################################################################\")\n",
    "    \n",
    "    \n",
    "    #Now calculate graph properties\n",
    "    \n",
    "    avg_length,dx, dy, dx, num_nodes , num_edges, node_degree=get_properties(data)\n",
    "    \n",
    "    y_data=np.array([avg_length,dx, dy, dx, num_nodes , num_edges, node_degree])\n",
    "    \n",
    "    print (f\"Graph properties measured: {labels_y_txt}\\n\",y_data)\n",
    "    if exists (GT_y):\n",
    "        print (f\"Graph properties GT: {labels_y_txt}\\n\", GT_y)\n",
    "        \n",
    "        plt.plot ( y_data, GT_y, '.',label='Graph properties (GT vs predicted)',markersize=12 )\n",
    "        plt.legend()\n",
    "        plt.xlabel ('GT')\n",
    "        plt.ylabel ('Predicted')\n",
    "        min_v,max_v=min (min(y_data), min (GT_y)),max (max(y_data), max (GT_y))\n",
    "        plt.plot([min_v, max_v], [min_v, max_v], ls=\"--\", c=\".3\")\n",
    "        plt.axis ('square')\n",
    "        plt.show()\n",
    "    return G,data, y_data\n",
    "    #visualize_graph(G, color=data.y)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3929753a",
   "metadata": {},
   "source": [
    "### Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb118289",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "loss_list=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d2f200",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prefix='./AnalogDiffusionFull/'\n",
    "if not os.path.exists(prefix):\n",
    "        os.mkdir (prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f420059",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    " max_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20751495",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "max_neighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ebada0f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "max_length_diff=64\n",
    "max_length=max_length_diff\n",
    "max_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cc92f84",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sns.set_style(\"whitegrid\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697d19d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from GraphDiffusion import AnalogDiffusionFull, count_parameters,pad_sequence \n",
    "\n",
    "predict_neighbors=True\n",
    "pred_dim=3+ max_length\n",
    "\n",
    "context_embedding_max_length=y_data.shape[1]\n",
    "model =AnalogDiffusionFull( \n",
    "            max_length=max_length,\n",
    "            pred_dim=pred_dim,\n",
    "            channels=256,\n",
    "            unet_type='cfg', #'base', #'cfg',\n",
    "            context_embedding_max_length=context_embedding_max_length,\n",
    "            pos_emb_fourier=True,\n",
    "            pos_emb_fourier_add=False,\n",
    "            text_embed_dim = 256,\n",
    "            embed_dim_position=256,\n",
    "            predict_neighbors=predict_neighbors,\n",
    "                    )  .to(device)  \n",
    "\n",
    "count_parameters (model)\n",
    " \n",
    "optimizer = optim.Adam(model.parameters() , lr=0.0002 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3b7c251",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_loop (model,\n",
    "                train_loader,test_loader,\n",
    "                optimizer=optimizer,\n",
    "                print_every=100,\n",
    "                epochs= 3000000,\n",
    "                start_ep=0,\n",
    "                start_step=0,\n",
    "                train_unet_number=1,\n",
    "                print_loss =  100* (len (train_loader)-1),\n",
    "                plot_unscaled=False,#if unscaled data is plotted\n",
    "                save_model=True,\n",
    "                cond_scales=[1],#[1, 2.5, 3.5, 5., 7.5, 10., 15., 20.],\n",
    "                num_samples=4,\n",
    "                timesteps=150,clamp=True,corplot=False,\n",
    "                save_loss_images=False,show_neighbors=True,\n",
    "                xyz_and_graph=True,\n",
    "                dist_matrix_threshold=0.99,\n",
    "                discretize=True,\n",
    "                get_length_from_result=True,\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ac46c7c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fname=f'{prefix}/statedict_save-model-epoch_2001.pt' #lowest loss mode\n",
    "model.load_state_dict(torch.load(fname))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ace06f7",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "sample_loop (model,device,\n",
    "                test_loader,\n",
    "                cond_scales=[1,  ], #list of cond scales - each sampled...\n",
    "                num_samples=16, #how many samples produced every time tested.....\n",
    "                timesteps=150,clamp=True, corplot=False,show_neighbors=True,\n",
    "           xyz_and_graph=True, clamp_round_results=True, enforce_symmetry=False,discretize=True,\n",
    "            get_length_from_result=True,dist_matrix_threshold=0.25, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec2943aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_sample_cond (model,\n",
    "                cond=[-.5, 0.7, 0.65],\n",
    "                cond_scales=[1.1,  ], #list of cond scales - each sampled...\n",
    "                flag=9992,\n",
    "                timesteps=100,clamp=True, corplot=True,show_neighbors=True,\n",
    "            xyz_and_graph=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff473b9e",
   "metadata": {},
   "source": [
    "### Generate hierarchical structures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bad1c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_length (X1):\n",
    "    return torch.nonzero(X1[:,3:]).flatten().max()+1\n",
    "    \n",
    "\n",
    "def get_length_xyzCOO (result):\n",
    "    return torch.nonzero(result[3 ,: ])[-1]+1 #first neighbor matters...\n",
    "    \n",
    "def get_xyz_and_dist_matrix_fromxyzCOO (X_data_cl, clamp_neighbors=False ,\n",
    "                            visualize=False, max_length=None):\n",
    "\n",
    "    if not exists (max_length):\n",
    "        max_length=get_length_xyzCOO (X_data_cl)\n",
    "    \n",
    "    X_data_cl=X_data_cl.permute (1,0)#bring to [max_length, xyz+max_neighbors]\n",
    "    print (X_data_cl.shape)\n",
    "        \n",
    "    max_neighbors=X_data_cl.shape[1]-3 \n",
    "    print (\"Max neighbors: \", max_neighbors, \"Length: \", max_length)\n",
    "    dist_matrix=torch.zeros (max_length, max_length)#.to(device)\n",
    "    \n",
    "    for i in range (max_length):\n",
    "\n",
    "        #neighbors are stored in result [4,5,..., 9]\n",
    "        for j in range (max_neighbors):\n",
    "           \n",
    "            neigh_j= X_data_cl[ i,3+j] \n",
    "\n",
    "            if neigh_j !=0: #zeros are padded values...\n",
    "\n",
    "                neighn=  neigh_j.long()\n",
    "\n",
    "                if clamp_neighbors:\n",
    "                    neighn=max( 1, min (neighn, max_length) )\n",
    "                    \n",
    "\n",
    "                dist_matrix[i, neighn-1] = 1\n",
    "                dist_matrix[neighn-1, i] = 1\n",
    "\n",
    "    print (dist_matrix.shape,X_data_cl[:max_length, :3].shape )\n",
    "    output=  torch.cat( (X_data_cl[:max_length, :3],   dist_matrix), 1)\n",
    "    \n",
    "    print (\"Shape of new matrix (length, x,y,z+length+): \", output.shape)\n",
    "    \n",
    "        \n",
    "    if visualize:\n",
    "        \n",
    "        plt.imshow (output[:, 3:3+max_length])\n",
    "        plt.show()\n",
    "    return output\n",
    "    \n",
    "def shift (X2, dx, dy, dz):\n",
    "    X2[:, 0] =X2[:, 0] + dx\n",
    "    X2[:, 1] =X2[:, 1] + dy\n",
    "    X2[:, 2] =X2[:, 2] + dz\n",
    "    return X2\n",
    "    \n",
    "def stack_and_extend (X1, X2, stack_node, dx,dy,dz, \n",
    "                     fname_root='file_name',\n",
    "                      visualize=False,\n",
    "                      make_graph = False,\n",
    "                      avg_pos_in_overlap=False,\n",
    "                      \n",
    "                     ): #format: (dist matrix x distmatrix+xyz)\n",
    "    #stacks two  graphs and overlaps them\n",
    "    #stacknode determines where in X1 is second one added\n",
    "    \n",
    "    S1=get_length(X1) #X1.shape[0]\n",
    "    S2=get_length(X2) #X2.shape[0]\n",
    "\n",
    "    \n",
    "    S_new=stack_node+S2\n",
    " \n",
    "    \n",
    "    dist_matrix=torch.zeros (S_new, S_new+3)#.to(device)\n",
    "    \n",
    "  \n",
    "      \n",
    "    X2[:, 0] =X2[:, 0] + dx\n",
    "    X2[:, 1] =X2[:, 1] + dy\n",
    "    X2[:, 2] =X2[:, 2] + dz\n",
    "\n",
    "\n",
    "    \n",
    "    \n",
    "    dist_matrix[:S1, 3:3+S1]= X1 [:S1, 3:3+S1]\n",
    "    dist_matrix[stack_node:stack_node+S2, 3+stack_node:3+stack_node+S2]= X2 [:S2, 3:3+S2]\n",
    "    \n",
    "    #now average positions\n",
    "    \n",
    "    dist_matrix[:S1, :3]= X1[:, :3]\n",
    "    if avg_pos_in_overlap:\n",
    "        dist_matrix[stack_node:stack_node+S2, :3]= dist_matrix[stack_node:stack_node+S2, :3]+X2[:, :3]\n",
    "       # dist_matrix[stack_node:stack_node+S2, :3]= dist_matrix[stack_node:stack_node+S2, :3]+X2[:, :3]\n",
    "        \n",
    "        \n",
    "    else:\n",
    "        dist_matrix[stack_node:stack_node+S2, :3]=  X2[:, :3]\n",
    "    \n",
    "    if avg_pos_in_overlap: \n",
    "        dist_matrix[stack_node:S1, :3]= dist_matrix[stack_node:S1, :3]/2.\n",
    "    \n",
    "\n",
    " \n",
    "    if visualize:\n",
    "        \n",
    "        plt.imshow (dist_matrix[:S_new, 3:3+S_new])\n",
    "        plt.show()\n",
    "        \n",
    "        plt.imshow(X1[:,:3], aspect=.1)\n",
    "        plt.show()\n",
    "        plt.imshow(X2[:,:3], aspect=.1)\n",
    "        plt.show()\n",
    "        plt.imshow(dist_matrix[:,:3], aspect=.1)\n",
    "        plt.show()\n",
    "        \n",
    "    if make_graph:\n",
    "        G_res, data_respr, y_data_pred =construct_xyz_and_graph (dist_matrix.permute(1,0),\n",
    "                                     fname_root=fname_root,\n",
    "                                                    label='Stacked and extended', #limits=[X_min*1.2, X_max*1.2],\n",
    "                                                    dist_matrix=True,\n",
    "                                                  )\n",
    "    \n",
    "    return dist_matrix \n",
    "    \n",
    "    \n",
    "\n",
    "    \n",
    "def make_spiral (X1, radius, \n",
    "                 slope_z, delta_angle, steps,  \n",
    "                stagger_fraction = 0.8,\n",
    "                fname_root='fname',\n",
    "                 visualize=True,visualize_all=False,\n",
    "                 shuffle=False,#shuflle graph\n",
    "                 \n",
    "                 avg_pos_in_overlap=False,\n",
    "                 \n",
    "                 generate_new_every_iteration=False,\n",
    "                 cond_vector_list=None,\n",
    "                 length_cond_vector=7,\n",
    "                 is_COO=False, #whether model predicts COO (sparse) or not\n",
    "                )   :\n",
    "\n",
    "    S1=get_length(X1)\n",
    "    delta_stagg=S1*stagger_fraction\n",
    "    \n",
    "    spiral = torch.clone (X1)\n",
    "    \n",
    "    i=-1\n",
    "    dx= radius * math.cos ((i+1)*delta_angle)\n",
    "    dy= radius * math.sin ((i+1)*delta_angle)\n",
    "    dz= slope_z *(i+1)\n",
    "\n",
    "    spiral = shift (spiral, dx, dy, dz)\n",
    "    for i in tqdm (range (steps)):\n",
    "        \n",
    "        \n",
    "        if generate_new_every_iteration:\n",
    "            if exists (cond_vector_list):\n",
    "                #[-0.7255, -0.2038,  0.5942,  0.3315,  0.286,  0.1852,  0.522]\n",
    "                cond_v=cond_vector_list[i]\n",
    "            else:\n",
    "                cond_v=1.5*torch.rand (length_cond_vector)-0.75\n",
    "            result, G_res, data_res,y_data= generate_sample_cond (GWebT,\n",
    "                cond=cond_v,#[-.5, 0.7, 0.6, 0.2, 0.3, 0.4, -0.9],\n",
    "                cond_scales=[1.,  ], #list of cond scales - each sampled...\n",
    "                flag=9992,\n",
    "              clamp=True, corplot=True,show_neighbors=True,\n",
    "            xyz_and_graph=True)\n",
    "            \n",
    "            \n",
    "            if is_COO:\n",
    "                X1=get_xyz_and_dist_matrix_fromxyzCOO (result, clamp_neighbors=True ,\n",
    "                                            visualize=False)#convert xyz-COO coding to xyz-distance matrix\n",
    "            else:\n",
    "                X1=result\n",
    "            S1=get_length(X1)\n",
    "            delta_stagg=S1*stagger_fraction\n",
    "    \n",
    "\n",
    "\n",
    "        if shuffle:\n",
    "            \n",
    "            if visualize_all:\n",
    "                plt.imshow (X1[:S1, 3:3+S1])\n",
    "                plt.show()               \n",
    "            ar=torch.range (0,S1-1).long()\n",
    "            ar2=torch.range (0,S1-1+3).long()\n",
    "            c=torch.randperm(S1)\n",
    "            c2=torch.cat( (torch.range (0,2), c+3)).long()\n",
    "           \n",
    "            X1=torch.cat( ( X1[:,:3],  X1.clone () [ar][c,3:]), 1)\n",
    "           \n",
    "            X1=X1[:,ar2][:,c2]\n",
    "            \n",
    "        \n",
    "\n",
    "            if visualize_all:\n",
    "                plt.imshow (X1[:S1, 3:3+S1])\n",
    "                plt.show()     \n",
    "    \n",
    "            \n",
    "            #a=a[:,torch.randperm(a.size()[1])]\n",
    "            \n",
    "        stack_node=int (get_length(spiral)-delta_stagg )\n",
    "        \n",
    "        #print (\"stack node: \", stack_node, i)\n",
    "    \n",
    "        dx= radius * math.cos ((i+1)*delta_angle)\n",
    "        dy= radius * math.sin ((i+1)*delta_angle)\n",
    "        dz= slope_z *(i+1)\n",
    "        \n",
    "        #print (dx, dy, dz)\n",
    "        spiral=stack_and_extend (spiral.clone(), X1.clone(), stack_node, dx,dy,dz, \n",
    "                      visualize=False,\n",
    "                      make_graph = False,avg_pos_in_overlap=avg_pos_in_overlap,\n",
    "                     )\n",
    "        #print (spiral.shape)\n",
    "        spiral[:, 3: ]=torch.clamp(spiral[:, 3: ], 0, 1) \n",
    "        \n",
    "        #print (\"#### length spiral \", get_length(spiral), spiral.shape, i)\n",
    "        \n",
    "    S1=get_length(spiral)\n",
    "    \n",
    "    \n",
    "    #print (\"Size of spiral; \", S1)\n",
    "    plt.imshow (spiral[:S1, 3:3+S1],interpolation='none')\n",
    "    plt.show()     \n",
    "    \n",
    "    plt.plot (spiral[:, 0])\n",
    "    plt.plot (spiral[:, 1])\n",
    "    plt.plot (spiral[:, 2])\n",
    "    plt.show()\n",
    "    \n",
    "    plt.imshow (spiral[:S1, :3],interpolation='none')\n",
    "    plt.show()     \n",
    "   \n",
    "        \n",
    "    \n",
    "    G_res, data_res, y_data_pred =construct_xyz_and_graph (spiral.permute(1,0),\n",
    "                                     fname_root=fname_root,\n",
    "                                                    label='Spiral', #limits=[X_min*1.2, X_max*1.2],\n",
    "                                                    dist_matrix=True,\n",
    "                                                  )\n",
    "    \n",
    "    if visualize:\n",
    "        \n",
    "        plt.imshow (spiral[:, 3:],interpolation='none')\n",
    "        plt.show()\n",
    "    #x(t) = rcos(t), y(t) = rsin(t), z(t) = at,\n",
    "    \n",
    "    return spiral\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2b25470",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "output, G_res, data_res,y_data= generate_sample_cond (model,\n",
    "                cond=[-0.7255, -0.2038,  0.5942,  0.3315,  0.1786,  0.1852,  0.0122],#[-.5, 0.7, 0.6, 0.2, 0.3, 0.4, -0.9],\n",
    "                cond_scales=[1.,  ], #list of cond scales - each sampled...\n",
    "                flag=9992,\n",
    "              clamp=True, corplot=True,show_neighbors=True,\n",
    "           xyz_and_graph=True, clamp_round_results=True, enforce_symmetry=True,discretize=True,\n",
    "            get_length_from_result=True,dist_matrix_threshold=0.25,\n",
    "             )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6298222",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(output.shape)\n",
    "plt.imshow (output[:,3:])\n",
    "plt.show()\n",
    "\n",
    "get_length (output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1e9389",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "                      \n",
    "output2=stack_and_extend (output.clone(), output.clone(), 10, 100,60,50,  #\n",
    "                     fname_root='file_name200',\n",
    "                      visualize=True,\n",
    "                      make_graph = True,\n",
    "                                            \n",
    "                     )#{max_neighbors, xyz+max_neighbors}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "447bded3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "spiral=make_spiral (output, radius=20, \n",
    "                 slope_z=5, delta_angle= 10/360*2*math.pi, steps=80,  \n",
    "                stagger_fraction = 0.5,\n",
    "                fname_root='spiral_v900',shuffle=False,\n",
    "                    avg_pos_in_overlap=False,\n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a0bb9fe",
   "metadata": {},
   "source": [
    "### Build attractor structure "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76c472d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_xyz_attractor (t, px=1.5, py=1.5, pz=1.5, dx=2, dy= 2., scale=1):\n",
    "    \n",
    "    x=(dx+math.cos(px*t))*math.cos(t)\n",
    "    y=(dy+math.cos(py*t))*math.sin(t)\n",
    "    z=math.sin(pz*t)\n",
    "    return x*scale, y*scale, z*scale\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaade167",
   "metadata": {},
   "outputs": [],
   "source": [
    "pts_1 = {'x':[],'y':[],'z':[]};\n",
    "for t in np.arange(0,50,0.005):\n",
    " \n",
    " \n",
    "    x=(math.sin(3*t))+math.cos(t)\n",
    "    y=(math.sin(t))*math.sin(5*t)\n",
    "    z=math.sin(2*t)\n",
    "    pts_1['x'].append(x)\n",
    "    pts_1['y'].append(y)\n",
    "    pts_1['z'].append(z)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0c8f395",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.axes(projection='3d')\n",
    "# Data for a three-dimensional line\n",
    "ax.plot3D(pts_1['x'], pts_1['y'], pts_1['z'], 'gray')\n",
    "# Data for three-dimensional scattered points\n",
    "ax.scatter3D(pts_1['x'], pts_1['y'], pts_1['z'], c=pts_1['z'], cmap='Greens');\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d27915c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pts_2 = {'x':[],'y':[],'z':[]};\n",
    "t_r=np.arange(0,10,0.1)\n",
    "print (t_r.shape)\n",
    "for t in t_r:\n",
    "\n",
    "    x,y,z= get_xyz_attractor (t, px=1.5, py=1.5, pz=1.5, dx=2, dy= 2., scale=20)\n",
    "    \n",
    "   \n",
    "    pts_2['x'].append(x)\n",
    "    pts_2['y'].append(y)\n",
    "    pts_2['z'].append(z)\n",
    "    \n",
    "ax = plt.axes(projection='3d')\n",
    "\n",
    "ax.scatter3D(pts_2['x'], pts_2['y'], pts_2['z'], c=pts_2['z'], cmap='Greens');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b303ffd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_attractor (X1,   t,   scale=5,\n",
    "                stagger_fraction = 0.8,\n",
    "                fname_root='fname',\n",
    "                 visualize=True,visualize_all=False,\n",
    "                 shuffle=False,#shuflle graph\n",
    "                 \n",
    "                 avg_pos_in_overlap=False,\n",
    "                 \n",
    "                 generate_new_every_iteration=False,\n",
    "                 cond_vector_list=None,\n",
    "                 length_cond_vector=7,\n",
    "                 is_COO=False, #whether model predicts COO (sparse) or not\n",
    "                )   :\n",
    "\n",
    "    S1=get_length(X1)\n",
    "    delta_stagg=S1*stagger_fraction\n",
    "    \n",
    "    spiral = torch.clone (X1)\n",
    "    \n",
    "\n",
    "    dx, dy, dz =  get_xyz_attractor (t[0], px=1.5, py=1.5, pz=1.5, dx=2, dy= 2., scale=scale)\n",
    "    \n",
    "    spiral = shift (spiral, dx, dy, dz)\n",
    "    for i in tqdm (range (len (t)-1)):\n",
    "        \n",
    "        \n",
    "        if generate_new_every_iteration:\n",
    "            if exists (cond_vector_list):\n",
    "                #[-0.7255, -0.2038,  0.5942,  0.3315,  0.286,  0.1852,  0.522]\n",
    "                cond_v=cond_vector_list[i]\n",
    "            else:\n",
    "                cond_v=1.5*torch.rand (length_cond_vector)-0.75\n",
    "            result, G_res, data_res,y_data= generate_sample_cond (model,\n",
    "                cond=cond_v,#[-.5, 0.7, 0.6, 0.2, 0.3, 0.4, -0.9],\n",
    "                cond_scales=[1.,  ], #list of cond scales - each sampled...\n",
    "                flag=9992,\n",
    "              clamp=True, corplot=True,show_neighbors=True,\n",
    "            xyz_and_graph=True,\n",
    "                                                                 \n",
    "                clamp_round_results=True, enforce_symmetry=True,discretize=True,\n",
    "                get_length_from_result=True,dist_matrix_threshold=0.25,\n",
    "                                                           \n",
    "                                                              )\n",
    "            \n",
    "            \n",
    "            if is_COO:\n",
    "                X1=get_xyz_and_dist_matrix_fromxyzCOO (result, clamp_neighbors=True ,\n",
    "                                            visualize=False)#convert xyz-COO coding to xyz-distance matrix\n",
    "            else:\n",
    "                X1=result\n",
    "            S1=get_length(X1)\n",
    "            delta_stagg=S1*stagger_fraction\n",
    "    \n",
    "\n",
    "\n",
    "        if shuffle:\n",
    "            \n",
    "            if visualize_all:\n",
    "                plt.imshow (X1[:S1, 3:3+S1])\n",
    "                plt.show()               \n",
    "            ar=torch.range (0,S1-1).long()\n",
    "            ar2=torch.range (0,S1-1+3).long()\n",
    "            c=torch.randperm(S1)\n",
    "            c2=torch.cat( (torch.range (0,2), c+3)).long()\n",
    "            #print (X1.shape,\"S1\", S1, \"v \", ar, ar2, c, c2)\n",
    "            X1=torch.cat( ( X1[:,:3],  X1.clone () [ar][c,3:]), 1)\n",
    "            #print (X1.shape)\n",
    "            X1=X1[:,ar2][:,c2]\n",
    "            \n",
    " \n",
    "\n",
    "            #t1[c][:,c2]\n",
    "\n",
    "            if visualize_all:\n",
    "                plt.imshow (X1[:S1, 3:3+S1])\n",
    "                plt.show()     \n",
    "    \n",
    "            \n",
    "            #a=a[:,torch.randperm(a.size()[1])]\n",
    "            \n",
    "        stack_node=int (get_length(spiral)-delta_stagg )\n",
    "        \n",
    "        \n",
    "        dx, dy, dz = get_xyz_attractor(t[i+1],px=1.5, py=1.5, pz=1.5, dx=2, dy= 2., scale=scale)\n",
    "        \n",
    "        #print (dx, dy, dz)\n",
    "        spiral=stack_and_extend (spiral.clone(), X1.clone(), stack_node, dx,dy,dz, \n",
    "                      visualize=False,\n",
    "                      make_graph = False,avg_pos_in_overlap=avg_pos_in_overlap,\n",
    "                     )\n",
    "        #print (spiral.shape)\n",
    "        spiral[:, 3: ]=torch.clamp(spiral[:, 3: ], 0, 1) \n",
    "       \n",
    "    S1=get_length(spiral)\n",
    "    \n",
    "    \n",
    "    #print (\"Size of spiral; \", S1)\n",
    "    plt.imshow (spiral[:S1, 3:3+S1],interpolation='none')\n",
    "    plt.show()     \n",
    "    \n",
    "    plt.plot (spiral[:, 0])\n",
    "    plt.plot (spiral[:, 1])\n",
    "    plt.plot (spiral[:, 2])\n",
    "    plt.show()\n",
    "    \n",
    "    plt.imshow (spiral[:S1, :3],interpolation='none')\n",
    "    plt.show()     \n",
    "   \n",
    "        \n",
    "    \n",
    "    G_res, data_res, y_data_pred =construct_xyz_and_graph (spiral.permute(1,0),\n",
    "                                     fname_root=fname_root,\n",
    "                                                    label='Spiral', #limits=[X_min*1.2, X_max*1.2],\n",
    "                                                    dist_matrix=True,\n",
    "                                                  )\n",
    "    \n",
    "    if visualize:\n",
    "        \n",
    "        plt.imshow (spiral[:, 3:],interpolation='none')\n",
    "        plt.show()\n",
    "    #x(t) = rcos(t), y(t) = rsin(t), z(t) = at,\n",
    "    \n",
    "    return spiral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adad3e15",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_r=np.arange(0,10,0.05)\n",
    "print (t_r.shape)\n",
    "pts_2 = {'x':[],'y':[],'z':[]}\n",
    "for t in t_r:\n",
    " \n",
    "    #x=(6+math.cos(3.5*t))*math.cos(t)\n",
    "    #y=(6+math.cos(3.5*t))*math.sin(t)\n",
    "    #z=math.sin(1.5*t)\n",
    "    \n",
    "    x,y,z= get_xyz_attractor (t, px=1.5, py=1.5, pz=1.5, dx=2, dy= 2., scale=5)\n",
    "    \n",
    "   \n",
    "    pts_2['x'].append(x)\n",
    "    pts_2['y'].append(y)\n",
    "    pts_2['z'].append(z)\n",
    "    \n",
    "ax = plt.axes(projection='3d')\n",
    " \n",
    "ax.scatter3D(pts_2['x'], pts_2['y'], pts_2['z'], c=pts_2['z'], cmap='Greens');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa0de1fb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "make_attractor (output,  t_r,  scale=10,\n",
    "                stagger_fraction = 0.5,\n",
    "                fname_root='attractor_v901',\n",
    "                 visualize=True,visualize_all=False,\n",
    "                 shuffle=True,#shuflle graph\n",
    "                 \n",
    "                 avg_pos_in_overlap=False,\n",
    "                 \n",
    "                 generate_new_every_iteration=False,\n",
    "                 cond_vector_list=None,\n",
    "                 length_cond_vector=7, \n",
    "                 is_COO=False, #whether model predicts COO (sparse) or not\n",
    "                ) "
   ]
  }
 ],
 "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.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}