Model 1 -- AnalogDiffusionSparse.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 1: Analog diffusion model with sparse neighbor representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26a6f27d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os,sys\n",
    " \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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a281c29d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os,sys\n",
    "import math\n",
    "import torch\n",
    "import time\n",
    "\n",
    "from tqdm.notebook import trange, tqdm\n",
    "from sklearn.model_selection import train_test_split "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e62b3a1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() \n",
    "                                  else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d862f28-4735-4c53-8940-4f0be3d8276d",
   "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\n",
    "from torch.utils.data import Dataset\n",
    "import pandas as pd\n",
    "import seaborn as sns"
   ]
  },
  {
   "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": [
    "\n",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fc9453b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "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",
    "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",
    "        #print (\"Rotation: \", degree*180)\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",
    "       # print (matrix)\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"
   ]
  },
  {
   "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": "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):\n",
    "        self.X_data = X_data\n",
    "        self.y_data = y_data\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        return self.X_data[index], self.y_data[index]\n",
    "        \n",
    "    def __len__ (self):\n",
    "        return len(self.X_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "457358b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def scale_data(image2, maxv, minv): \n",
    "    return (image2 -minv)/(maxv-minv) * 2. - 1.0\n",
    " \n",
    "def unscale_data(image2, maxv, minv):\n",
    "    image2=(image2 +1. )/ 2. * (maxv-minv)+minv \n",
    "    return image2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fef0b41",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def pad_sequence_end (output_xyz, max_length_l):        \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)"
   ]
  },
  {
   "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": [
    "max_length_diff=128 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ea3b2b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def normalize_data (input, y_data, X_min=None, X_max=None, y_min=None, y_max=None,\n",
    "                     y_min2=None, y_max2=None,\n",
    "                    X_max_neigh=None, X_min_neigh=None,\n",
    "                    Xscale=0, max_length_pad=256):\n",
    "    \n",
    "    if exists (max_length_pad):\n",
    "       \n",
    "        input=pad_sequence_end (input, max_length_pad)\n",
    "        \n",
    "        plt.plot (input[0, :, 4].cpu())\n",
    "        plt.show()\n",
    "\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:4+max_neighbors].min(), input[:,:,4:4+max_neighbors].max())  \n",
    "    if X_min_neigh==None:\n",
    "        X_min_neigh=input[:,:,4:4+max_neighbors].min() \n",
    "    else:\n",
    "        print (\"use provided X_min_neigh\", X_min_neigh)\n",
    "    if X_max_neigh==None:\n",
    "        X_max_neigh=input[:,:,4:4+max_neighbors].max() \n",
    "    else:\n",
    "        print (\"use provided X_max_neigh\", X_max_neigh)\n",
    "\n",
    "    input[:,:,4:4+max_neighbors]=(input[:,:,4:4+max_neighbors]-X_min_neigh)/(X_max_neigh-X_min_neigh)*2-1 #Normalize range -1 to 1\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",
    "    print (\"Check y_data after norm  \", y_data.min(), y_data.max())\n",
    "    return input, y_data, X_min.cpu(), X_max.cpu(), y_min , y_max , X_min_neigh.cpu(), X_max_neigh.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40ab97db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "y_data[:,0:2].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc805dcf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_scaled, y_data_scaled, X_min, X_max, y_min, y_max,  X_min_neigh, X_max_neigh= normalize_data (input, y_data,Xscale=0.,\n",
    "                                                                                               max_length_pad=max_length_diff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f952d108",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_min_neigh, X_max_neigh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6769a023",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_min, X_max, y_min, y_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7b55a84",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_data_loaders (X_scaled,  y_data_scaled, split=0.1, batch_size_=16):\n",
    "\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X_scaled,  \n",
    "                                                        y_data_scaled , test_size=split,random_state=235)\n",
    "    \n",
    "    print (f\"Shapes= {X_scaled.shape}, {y_data_scaled.shape}\")\n",
    "     \n",
    "    print(X_train.shape)\n",
    "    print(y_train.shape)\n",
    "    print(X_test.shape)\n",
    "    print(y_test.shape)\n",
    "    print (len (X_train))\n",
    "    train_dataset = RegressionDataset(X_train, y_train) \n",
    "\n",
    "    test_dataset = RegressionDataset(X_test,y_test)\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": "cd0a9f40",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_loader,train_loader_noshuffle, test_loader= get_data_loaders (X_scaled,  y_data_scaled, split=0.1, batch_size_=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20c2e085",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_scaled[:,:,1:4].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c343d778",
   "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\n",
    "\n",
    "sns.set_style(\"whitegrid\")\n",
    "fig=sns.histplot(X_scaled[:,:,4].flatten().cpu() ,bins=100, binrange=(-1,1 ) )\n",
    "fig.set_xlabel( \"positions\", fontsize = 10 )\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b33ccc1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in range (len (labels_y_txt)):\n",
    "    sns.set_style(\"whitegrid\")\n",
    "    fig=sns.histplot(y_data_scaled[:,i].flatten() ,bins=20, binrange=(-1, 1) )\n",
    "    fig.set_xlabel( labels_y_txt[i], fontsize = 10 )\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea9058aa",
   "metadata": {},
   "source": [
    "## Define auxiliary functions, e.g. training loop, sampling etc. "
   ]
  },
  {
   "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",
    "               ):\n",
    "    \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",
    "            train_epoch_loss = 0\n",
    "            model.train()\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",
    "\n",
    "                X_train_batch=torch.permute(X_train_batch, (0,2,1)  )\n",
    "                optimizer.zero_grad()\n",
    "                \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",
    "                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",
    "                        sample_loop (model,\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",
    "                        \n",
    "                        print (f\"\\n\\n-------------------\\nTime passed for {print_loss} epochs at {steps} = {(time.time()-start)/60} mins\\n-------------------\")\n",
    "                        start = time.time()\n",
    "                        \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"
   ]
  },
  {
   "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": "6a446705",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import r2_score\n",
    "\n",
    "def sample_loop (model,\n",
    "                train_loader,\n",
    "                 device,\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",
    "                 \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",
    "\n",
    "            GT=y_train_batch.cpu().detach().unsqueeze(1) \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,\n",
    "                                         cond_scale=cond_scales[iisample],\n",
    "                                         timesteps=timesteps,clamp=clamp,\n",
    "                                     device=device,\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",
    "                X_train_batch= pad_sequence (X_train_batch,  max_length).cpu()\n",
    "                \n",
    "                for samples in range  (num_samples):\n",
    "                    \n",
    "                    #unscale neigh bors....\n",
    "                    result[samples,3:3+max_neighbors,: ]=unscale_data(result[samples,3:3+max_neighbors,: ], \n",
    "                                                                      X_max_neigh.numpy(),\n",
    "                                                                      X_min_neigh.numpy())\n",
    "                    result[samples,3:3+max_neighbors,: ]=(result[samples,3:3+max_neighbors,: ]).round ()\n",
    "                    \n",
    "                    X_train_batch[samples,4:4+max_neighbors,: ]=unscale_data(X_train_batch[samples,4:4+max_neighbors,: ], \n",
    "                                                                             X_max_neigh.numpy(),\n",
    "                                                                             X_min_neigh.numpy())\n",
    "                    X_train_batch[samples,4:4+max_neighbors,: ]=(X_train_batch[samples,4:4+max_neighbors,: ]).round ()\n",
    "\n",
    "                    #get length\n",
    "                    GTroundL=torch.nonzero(X_train_batch[samples,4 ,:max_length_considered ])[-1] \n",
    "                    resroundL = torch.nonzero(result[samples,3 ,:max_length_considered ])[-1] \n",
    "                    \n",
    "                    print (\"lengths GT vs pred \", GTroundL,resroundL)\n",
    "                    \n",
    "                    print (f\"### Ground truth {labels_y_txt}  = \", GT[samples,:])\n",
    "                    \n",
    "                    fig, ax = plt.subplots(1, 2, figsize=(10, 6) , subplot_kw=dict(projection='3d'))\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",
    "                    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",
    "                    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",
    "                    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",
    "                 \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",
    "                    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",
    "                    \n",
    "                    ax[1].axis('square')\n",
    "                    ax[1].set_xlim(-1,1)\n",
    "                    ax[1].set_ylim(-1,1)\n",
    "                                  \n",
    "                    ax[1].legend()\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",
    "                    if show_neighbors:\n",
    "                         \n",
    "                        fig, (ax1, ax2) = plt.subplots(2,1)\n",
    "\n",
    "                        ax2.imshow (result[samples, 3:3+max_neighbors, :resroundL],aspect=2)\n",
    "                        ax2.set_title('Prediction')\n",
    "                        ax1.imshow (X_train_batch[samples, 4:4+max_neighbors, :GTroundL],aspect=2)\n",
    "                        ax1.set_title('GT')\n",
    "                        plt.show()\n",
    "                        \n",
    "                        fig, (ax1, ax2) = plt.subplots(2,1)\n",
    "\n",
    "                        ax2.imshow (result[samples, 3:3+max_neighbors, :max_length],aspect=2)\n",
    "                        ax2.set_title('Prediction')\n",
    "                        ax1.imshow (X_train_batch[samples, 4:4+max_neighbors, :max_length],aspect=2)\n",
    "                        ax1.set_title('GT')\n",
    "                        plt.show()\n",
    "                        \n",
    "                        \n",
    "                    if xyz_and_graph:\n",
    "                        G_res, data_res,y_data_pred=construct_xyz_and_graph (result[samples, :, :resroundL].squeeze(),\n",
    "                                                 fname_root=f'{prefix}graph_xyz_{samples}_{flag}_{steps}', \n",
    "                                                                 #limits=[X_min*1., X_max*1.] ,\n",
    "                                                              label =f'Prediction {flag} {steps}')\n",
    "                        G_GT, data_GT,y_data_GT=construct_xyz_and_graph (X_train_batch[samples, 1:4+max_neighbors, :GTroundL].squeeze(),\n",
    "                                                 fname_root=f'{prefix}graph_GT_xyz_{samples}_{flag}_{steps}',\n",
    "                                                              label =f'GT {flag} {steps}',\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",
    "                        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",
    "                        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",
    "                    print (X_train_batch.shape,result.shape )\n",
    "                    if corplot:\n",
    "                        plt.plot (X_train_batch[samples, 1:4, :resroundL].flatten().detach().cpu(), \n",
    "                                  result[samples, 0:3, :resroundL].flatten() , '.',label='all',markersize=6 )\n",
    "                        plt.plot (X_train_batch[samples, 1, :resroundL].flatten().detach().cpu(), \n",
    "                                  result[samples, 0, :resroundL].flatten() , '.',label='x',markersize=2 )\n",
    "                        plt.plot (X_train_batch[samples, 2, :resroundL].flatten().detach().cpu(), \n",
    "                                  result[samples, 1, :resroundL].flatten() , '.',label='y',markersize=2 )\n",
    "                        plt.plot (X_train_batch[samples, 3, :resroundL].flatten().detach().cpu(), \n",
    "                                  result[samples, 2, :resroundL].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",
    "                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",
    "            steps=steps+1\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9b7744b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_sample_cond (model,\n",
    "              device,\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",
    "                        \n",
    "                 \n",
    "               ):\n",
    "    steps=0\n",
    "    e=flag\n",
    "\n",
    "    y_train_batch=torch.Tensor (cond).to(device)\n",
    "    y_train_batch=y_train_batch.unsqueeze(0) \n",
    "\n",
    "    for iisample in range (len (cond_scales)):\n",
    "        \n",
    "        samples=0\n",
    "        result=model.sample ( y_train_batch,\n",
    "                                 cond_scale=cond_scales[iisample],\n",
    "                                 timesteps=timesteps,clamp=clamp,device=device,\n",
    "                                  ).cpu()\n",
    "        result[samples,3:3+max_neighbors,: ]=unscale_data(result[samples,3:3+max_neighbors,: ], \n",
    "                                                          X_max_neigh.cpu().numpy(),\n",
    "                                                          X_min_neigh.cpu().numpy())\n",
    "        result[samples,3:3+max_neighbors,: ]=(result[samples,3:3+max_neighbors,: ]).round ()\n",
    " \n",
    "\n",
    "         \n",
    "        resroundL = torch.nonzero(result[samples,3 ,: ])[-1] \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",
    "        if show_neighbors:\n",
    "             \n",
    "            fig, ax1 = plt.subplots(1,1)\n",
    "\n",
    "            ax1.imshow (result[samples, 3:3+max_neighbors, :resroundL],aspect=2)\n",
    "            ax1.set_title('Prediction')\n",
    "            plt.show()\n",
    "\n",
    "            fig, ax1 = plt.subplots(1,1)\n",
    "\n",
    "            ax1.imshow (result[samples, 3:3+max_neighbors, :max_length],aspect=2)\n",
    "            ax1.set_title('Prediction')\n",
    "            plt.show()\n",
    "        if xyz_and_graph:\n",
    "            \n",
    "            for i in range (len(y_min)):\n",
    "                y_train_batch[0,i]=unscale_data(y_train_batch[0,i], y_max[i],y_min[i]) \n",
    "                \n",
    "\n",
    "            G_res, data_res,y_data=construct_xyz_and_graph (result[samples, :, :resroundL].squeeze(),\n",
    "                                     fname_root=f'{prefix}graph_xyz_{samples}_{flag}_{steps}',\n",
    "                                                              label =f'Prediction {flag} {steps}',\n",
    "                                                    GT_y=y_train_batch[0,:].cpu().numpy())\n",
    "            \n",
    "\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",
    "    steps=steps+1\n",
    "    if xyz_and_graph:\n",
    "        return result[samples, :resroundL+3, :resroundL].squeeze(), G_res, data_res,y_data\n",
    "    else:\n",
    "        return result[samples, :resroundL+3, :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": [
    "X_max_neigh,X_min_neigh"
   ]
  },
  {
   "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",
    " \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",
    "    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",
    "    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 \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",
    "\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'0 500 xlo xhi \\n0 500 ylo yhi \\n0 500 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]} { max( 1, min (neighbor_list[i][1], result.shape[1]) )}   \\n')\n",
    "        f.write(f'\\n\\n')\n",
    "    \n",
    "    \n",
    "    neighbor_list=torch.Tensor (neighbor_list).long()\n",
    "   \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",
    "    for i in range (len (neighbor_list)):\n",
    "            \n",
    "            N1=neighbor_list[i][0] #  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(f\"radius: {nx.radius(G)} center: {nx.center(G)} density: {nx.density(G)}\")\n",
    "    \n",
    "    #nx.draw(G, with_labels=False,  )\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",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3929753a",
   "metadata": {},
   "source": [
    "## Define diffusion model "
   ]
  },
  {
   "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='./AnalogDiffusionSparse/'\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=max_length_diff\n",
    "max_length, max_neighbors, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697d19d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from GraphDiffusion import AnalogDiffusionSparse, count_parameters,pad_sequence \n",
    "\n",
    "predict_neighbors=True\n",
    "pred_dim=3+max_neighbors*predict_neighbors\n",
    "\n",
    "context_embedding_max_length=y_data.shape[1]\n",
    "model =AnalogDiffusionSparse( \n",
    "            max_length=max_length,\n",
    "            pred_dim=pred_dim,\n",
    "            channels=128,\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
   },
   "outputs": [],
   "source": [
    "train=False\n",
    "if train:\n",
    "    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 =  500* (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",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d9f373d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname=f'{prefix}/statedict_save-model-epoch_4327.pt' #lowest loss model\n",
    "model.load_state_dict(torch.load(fname))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ace06f7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sample_loop (model,\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, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a2ff3dd",
   "metadata": {},
   "source": [
    "### Generate hierarchical structures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65049189",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_length (X1):\n",
    "    return torch.nonzero(X1[:,3:]).flatten().max()+1\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, enforce_symm=True):\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",
    "            #neigh_j=i-int (result[ 3+j, i] )   +1\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",
    "\n",
    "                #print (dist_matrix.shape, j, neighn)\n",
    "                #dist_matrix[neighn-1, j] = 1\n",
    "                \n",
    "                \n",
    "                dist_matrix[i, neighn-1] = 1\n",
    "                if enforce_symm:\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",
    "    #print (S1, S2)\n",
    "\n",
    "    #plt.imshow(X1[:,:3], aspect=.1)\n",
    "    #plt.show()\n",
    "    #plt.imshow(X2[:,:3], aspect=.1)\n",
    "    #plt.show()\n",
    "    \n",
    "    S_new=stack_node+S2\n",
    "    \n",
    "    #print (\"sizes: \", S1, S2, S_new)\n",
    "    \n",
    "    dist_matrix=torch.zeros (S_new, S_new+3)#.to(device)\n",
    "    \n",
    "    #shift second web\n",
    "    #for i in range (X2.shape[0]):\n",
    "        \n",
    "        #print (X2[i, 0])\n",
    "    #    X2[i, 0] =X2[i, 0] + dx\n",
    "    #    X2[i, 1] =X2[i, 1] + dy\n",
    "    #    X2[i, 2] =X2[i, 2] + dz\n",
    "      \n",
    "    X2[:, 0] =X2[:, 0] + dx\n",
    "    X2[:, 1] =X2[:, 1] + dy\n",
    "    X2[:, 2] =X2[:, 2] + dz\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",
    "    else:\n",
    "        dist_matrix[stack_node:stack_node+S2, :3]=  X2[:, :3]\n",
    "   \n",
    "    #dist_matrix[stack_node:stack_node+S2, :3]= dist_matrix[stack_node:stack_node+S2, :3]/2.\n",
    "    if avg_pos_in_overlap: \n",
    "        dist_matrix[stack_node:S1, :3]= dist_matrix[stack_node:S1, :3]/2.\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",
    "                )   :\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",
    "        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",
    "            X1=get_xyz_and_dist_matrix_fromxyzCOO (result, clamp_neighbors=True ,\n",
    "                                        visualize=False)#convert xyz-COO coding to xyz-distance matrix\n",
    "            \n",
    "            S1=get_length(X1)\n",
    "            delta_stagg=S1*stagger_fraction\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",
    "            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",
    "        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",
    "      \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": "e242df28",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result, G_res, data_res, y_data_pred=generate_sample_cond (model,\n",
    "                cond=[-.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",
    "                timesteps=150,clamp=True, corplot=True,show_neighbors=True,\n",
    "            xyz_and_graph=True, device=device)\n",
    "\n",
    "output=get_xyz_and_dist_matrix_fromxyzCOO (result, clamp_neighbors=False ,\n",
    "                            visualize=True)#convert xyz-COO coding to xyz-distance matrix\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a45bb74c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(output.shape)\n",
    "\n",
    "plt.imshow (output[:,3:])\n",
    "plt.show()\n",
    "\n",
    "get_length (output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc2c2d37",
   "metadata": {},
   "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": "e1cf4e28",
   "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_v400',shuffle=False,\n",
    "                    avg_pos_in_overlap=False,\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4408e596",
   "metadata": {},
   "outputs": [],
   "source": [
    "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)"
   ]
  }
 ],
 "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
}