ContinuumSimulation3DFast.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sortedcontainers import SortedSet\n",
    "from utils import *\n",
    "import cProfile\n",
    "import pstats\n",
    "import time\n",
    "import warnings\n",
    "import json\n",
    "import copy\n",
    "\n",
    "# taken from https://github.com/grantjenks/python-sortedcontainers/issues/53\n",
    "class FasterSortedSet(SortedSet):\n",
    "    def union(self, *iterables):\n",
    "        if (self and len(iterables) == 1\n",
    "                and isinstance(iterables[0], SortedSet)\n",
    "                and iterables[0]\n",
    "                and self._list._lists[-1][-1] < iterables[0]._list._lists[0][0]):\n",
    "            that = iterables[0]\n",
    "            result = FasterSortedSet()\n",
    "            _set = result._set\n",
    "            _set.update(self)\n",
    "            _set.update(that)\n",
    "            _list = result._list\n",
    "            _list._lists.extend(list(sublist) for sublist in self._list._lists)\n",
    "            _list._maxes.extend(self._list._maxes)\n",
    "            _list._lists.extend(list(sublist) for sublist in that._list._lists)\n",
    "            _list._maxes.extend(that._list._maxes)\n",
    "            _list._len = len(_set)\n",
    "            return result\n",
    "        else:\n",
    "            return self.__class__(\n",
    "                chain(iter(self), *iterables),\n",
    "                key=self._key,\n",
    "                load=self._load\n",
    "            )\n",
    "\n",
    "    __or__ = union\n",
    "    __ror__ = __or__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Cython\n",
    "%load_ext Cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%cython\n",
    "cimport cython\n",
    "import numpy as np\n",
    "cimport numpy as np\n",
    "\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def superimpose_potential(np.float32_t[:,:,:,:] pot_field, np.float64_t[:] center, unsigned long scale, double radius, np.float32_t[:,:,:,:] template, unsigned long diameter):\n",
    "    cdef int[:] scaled_center = np.array([0,0,0], dtype=np.int32)\n",
    "    scaled_center[0] = np.round(center[0] * scale).astype(np.int32) % pot_field.shape[0]\n",
    "    scaled_center[1] = np.round(center[1] * scale).astype(np.int32) % pot_field.shape[1]\n",
    "    scaled_center[2] = np.round(center[2] * scale).astype(np.int32) % pot_field.shape[2]\n",
    "    #print(\"Scaled center {}\".format(scaled_center))\n",
    "    #Vmin is at r=2^(1/6)*zero_pot\n",
    "    cdef np.float32_t zero_pot = radius / 2**(1/5) * scale # in indices\n",
    "    cdef int low = -((diameter-1)//2)\n",
    "    cdef int high = (diameter+1)//2\n",
    "    cdef Py_ssize_t x = (pot_field.shape[0])\n",
    "    cdef Py_ssize_t y = (pot_field.shape[1])\n",
    "    cdef Py_ssize_t z = (pot_field.shape[2])\n",
    "    \n",
    "    cdef np.float32_t[:,:,:,:] pot_view = pot_field\n",
    "    cdef np.float32_t[:,:,:,:] tmp_view = template\n",
    "    \n",
    "    cdef Py_ssize_t i, ii, j, jj, k, kk\n",
    "    for i in range(low, high):\n",
    "        ii = (scaled_center[0] + i) % x\n",
    "        idist = (i)**2\n",
    "        for j in range(low, high):\n",
    "            jj = (scaled_center[1] + j) % y\n",
    "            jdist = (j)**2\n",
    "            for k in range(low, high):\n",
    "                kk = scaled_center[2] + k\n",
    "                if kk < 0 or kk > z-1:\n",
    "                    continue\n",
    "                if i == 0 and j == 0 and k == 0:\n",
    "                    pot_view[ii,jj,kk,0] = 0\n",
    "                    pot_view[ii,jj,kk,1] = 0\n",
    "                    pot_view[ii,jj,kk,2] = 0\n",
    "                    pot_view[ii,jj,kk,3] = np.inf\n",
    "                \n",
    "                #r = np.sqrt(idist + jdist + (k)**2)\n",
    "                #magnitude = tmp_view[i-low,j-low,k-low,3]#(zero_pot/r)**12 - (zero_pot/r)**6\n",
    "                #direction = tmp_view[i-low,j-low,k-low,:3]#np.array([i,j,k]) / distance3D([i,j,k], [0,0,0]) * magnitude/np.abs(magnitude) # get direction\n",
    "                #magnitude = np.abs(magnitude) # set magnitude to the absolute value\n",
    "                \n",
    "                pot_view[ii,jj,kk,0] = pot_view[ii,jj,kk,3]*pot_view[ii,jj,kk,0] + tmp_view[i-low,j-low,k-low,3]*tmp_view[i-low,j-low,k-low,0]\n",
    "                pot_view[ii,jj,kk,1] = pot_view[ii,jj,kk,3]*pot_view[ii,jj,kk,1] + tmp_view[i-low,j-low,k-low,3]*tmp_view[i-low,j-low,k-low,1]\n",
    "                pot_view[ii,jj,kk,2] = pot_view[ii,jj,kk,3]*pot_view[ii,jj,kk,2] + tmp_view[i-low,j-low,k-low,3]*tmp_view[i-low,j-low,k-low,2]\n",
    "                #pot_view[ii,jj,kk,3] = np.sqrt(pot_view[ii,jj,kk,0]**2 + pot_view[ii,jj,kk,1]**2 + pot_view[ii,jj,kk,2]**2)\n",
    "                pot_view[ii,jj,kk,3] = pot_view[ii,jj,kk,3] + tmp_view[i-low, j-low, k-low, 3]\n",
    "                if pot_view[ii,jj,kk,3] == 0.0:\n",
    "                    pot_view[ii,jj,kk,0] = 0.0\n",
    "                    pot_view[ii,jj,kk,1] = 0.0\n",
    "                    pot_view[ii,jj,kk,2] = 0.0\n",
    "                else:\n",
    "                    pot_view[ii,jj,kk,0] = pot_view[ii,jj,kk,0]/pot_view[ii,jj,kk,3]\n",
    "                    pot_view[ii,jj,kk,1] = pot_view[ii,jj,kk,1]/pot_view[ii,jj,kk,3]\n",
    "                    pot_view[ii,jj,kk,2] = pot_view[ii,jj,kk,2]/pot_view[ii,jj,kk,3]\n",
    "#     ii = scaled_center[0] + low\n",
    "#     jj = scaled_center[1] + low\n",
    "#     kk = scaled_center[2] + low\n",
    "#     vectors = [0, 1, 2]\n",
    "#     if scaled_center[0] > -low and scaled_center[0] + high < x-1:\n",
    "#         if scaled_center[1] > -low and scaled_center[1] + high < y-1:\n",
    "#             if scaled_center[2] > -low and scaled_center[2] + high < z-1:\n",
    "#                 # easy case: entire added field doesn't cross boundaries\n",
    "#                 region = pot_field[ii:ii+diameter,jj:jj+diameter,kk:kk+diameter,:]\n",
    "#                 print(region.shape, template.shape)\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:,:,3]*template[:,:,:,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "#             elif scaled_center[2] <= -low:\n",
    "#                 # 0<x<L, 0<y<L, 0=z<h\n",
    "#                 region = pot_field[ii:ii+diameter,jj:jj+diameter,0:kk+diameter,:]\n",
    "#                 t = template[:,:,-kk-diameter:,:]\n",
    "#                 print(scaled_center, region.shape, template[:,:,-kk-diameter:,:].shape, low, kk)\n",
    "#                 cs = [region[:,:,:,3]*region[:,:,:,i] + t[:,:,:,3]*t[:,:,:,i] for i in vectors]\n",
    "#                 region[:,:,:,:3] = np.stack(cs, axis=3)\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.stack([region[:,:,:,0]/region[:,:,:,3], region[:,:,:,1]/region[:,:,:,3], region[:,:,:,2]/region[:,:,:,3], region[:,:,:,3]], axis=3)\n",
    "#                 region = np.where(region==np.nan, 0, region)\n",
    "#             else:\n",
    "#                 # 0<x<L, 0<y<L, 0<z=h\n",
    "#                 region = pot_field[ii:ii+diameter,jj:jj+diameter,kk:-1,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:,:low,3]*template[:,:,:low,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "#         elif scaled_center[0] <= -low:\n",
    "#             if scaled_center[2] > -low and scaled_center[2] + high < z-1:\n",
    "#                 # 0<x<L, 0=y<L, 0<z<h\n",
    "#                 region = pot_field[ii:ii+diameter,0:jj+diameter,kk:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,-low:,:,3]*template[:,-low:,:,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "                \n",
    "#                 #0<x<L, y<0, 0<z<h\n",
    "#                 region = pot_field[ii:ii+diameter,low:,kk:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:low,:,3]*template[:,:low,:,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "#             elif scaled_center[2] <= -low:\n",
    "#                 # 0<x<L, 0=y<L, 0=z<h\n",
    "#                 region = pot_field[ii:ii+diameter,0:jj+diameter,0:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,-low:,-low:,3]*template[:,-low:,-low:,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "                \n",
    "#                 #0<x<L, y<0, 0=z<h\n",
    "#                 region = pot_field[ii:ii+diameter,low:,0:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:low,:low,3]*template[:,:low,:low,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "#             else:\n",
    "#                 # 0<x<L, 0=y<L, 0<z=h\n",
    "#                 region = pot_field[ii:ii+diameter,0:jj+diameter,kk:-1,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,-low:,:low,3]*template[:,-low:,:low,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "                \n",
    "#                 # 0<x<L, y<0, 0<z=h\n",
    "#                 region = pot_field[ii:ii+diameter,low:,kk:-1,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:low,:low,3]*template[:,:low,:low,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "#         else:\n",
    "#             if scaled_center[2] > -low and scaled_center[2] + high < z-1:\n",
    "#                 # 0<x<L, 0<y=L, 0<z<h\n",
    "#                 region = pot_field[ii:ii+diameter,jj:,kk:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,low:,:,3]*template[:,low:,:,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "                \n",
    "#                 #0<x<L, y>L, 0<z<h\n",
    "#                 region = pot_field[ii:ii+diameter,low:,kk:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:low,:,3]*template[:,:low,:,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "#             elif scaled_center[2] <= -low:\n",
    "#                 # 0<x<L, 0=y<L, 0=z<h\n",
    "#                 region = pot_field[ii:ii+diameter,0:jj+diameter,0:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,-low:,-low:,3]*template[:,-low:,-low:,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "                \n",
    "#                 #0<x<L, y<0, 0=z<h\n",
    "#                 region = pot_field[ii:ii+diameter,low:,0:kk+diameter,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:low,:low,3]*template[:,:low,:low,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "#             else:\n",
    "#                 # 0<x<L, 0=y<L, 0<z=h\n",
    "#                 region = pot_field[ii:ii+diameter,0:jj+diameter,kk:-1,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,-low:,:low,3]*template[:,-low:,:low,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "                \n",
    "#                 # 0<x<L, y<0, 0<z=h\n",
    "#                 region = pot_field[ii:ii+diameter,low:,kk:-1,:]\n",
    "#                 region[:,:,:,:3] = region[:,:,:,3]*region[:,:,:,:3] + template[:,:low,:low,3]*template[:,:low,:low,:3]\n",
    "#                 region[:,:,:,3] = np.sqrt(region[:,:,:,0]**2 + region[:,:,:,1]**2 + region[:,:,:,2]**2)\n",
    "#                 region = np.where(region[:,:,:,3]==0.0, [0,0,0,0], [region[:,:,:,:3]/region[:,:,:,3], region[:,:,:,3]])\n",
    "    return\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def shortest_distance_line_to_point(double[:] line_1, double[:] line_2, np.float32_t[:] point):\n",
    "    x = np.array([line_1[0] - line_2[0], line_1[1] - line_2[1], line_1[2] - line_2[2]])\n",
    "    pmq = np.array([point[0] - line_2[0], point[1] - line_2[1], point[2] - line_2[2]])\n",
    "    t = np.dot(pmq, x)/np.dot(x, x)\n",
    "    r = np.linalg.norm(t*x - pmq)\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def distance3D(p1, p2):\n",
    "    return np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2 + (p1[2] - p2[2])**2)\n",
    "\n",
    "def distance3DPBC(p1, p2, L):\n",
    "    # distance in x considering PBC\n",
    "    if (p1[0] - p2[0]) > L/2:\n",
    "        dx = p1[0] - p2[0] - L\n",
    "    elif p1[0] - p2[0] < -L/2:\n",
    "        dx = p1[0] - p2[0] + L\n",
    "    else:\n",
    "        dx = p1[0] - p2[0]\n",
    "    # distance in y considering PBC\n",
    "    if (p1[1] - p2[1]) > L/2:\n",
    "        dy = p1[1] - p2[1] - L\n",
    "    elif p1[1] - p2[1] < -L/2:\n",
    "        dy = p1[1] - p2[1] + L\n",
    "    else:\n",
    "        dy = p1[1] - p2[1]\n",
    "    return np.sqrt(dx**2 + dy**2 + (p1[2] - p2[2])**2)\n",
    "\n",
    "# returns distance^2 (for faster comparisons)\n",
    "def distance3D2PBC(p1, p2, L):\n",
    "    # distance in x considering PBC\n",
    "    if (p1[0] - p2[0]) > L/2:\n",
    "        dx = p1[0] - p2[0] - L\n",
    "    elif p1[0] - p2[0] < -L/2:\n",
    "        dx = p1[0] - p2[0] + L\n",
    "    else:\n",
    "        dx = p1[0] - p2[0]\n",
    "    # distance in y considering PBC\n",
    "    if (p1[1] - p2[1]) > L/2:\n",
    "        dy = p1[1] - p2[1] - L\n",
    "    elif p1[1] - p2[1] < -L/2:\n",
    "        dy = p1[1] - p2[1] + L\n",
    "    else:\n",
    "        dy = p1[1] - p2[1]\n",
    "    return dx**2 + dy**2 + (p1[2] - p2[2])**2\n",
    "\n",
    "def distance2D2PBC(p1, p2, L):\n",
    "    # distance in x considering PBC\n",
    "    if (p1[0] - p2[0]) > L/2:\n",
    "        dx = p1[0] - p2[0] - L\n",
    "    elif p1[0] - p2[0] < -L/2:\n",
    "        dx = p1[0] - p2[0] + L\n",
    "    else:\n",
    "        dx = p1[0] - p2[0]\n",
    "    # distance in y considering PBC\n",
    "    if (p1[1] - p2[1]) > L/2:\n",
    "        dy = p1[1] - p2[1] - L\n",
    "    elif p1[1] - p2[1] < -L/2:\n",
    "        dy = p1[1] - p2[1] + L\n",
    "    else:\n",
    "        dy = p1[1] - p2[1]\n",
    "    return dx**2 + dy**2\n",
    "\n",
    "# p is the current point position, n the unit vector, and step_size the magnitude\n",
    "def step_along_line(p, n, step_size):\n",
    "    return p + n * step_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SlantedCorridors:\n",
    "    def __init__(self, length, bin_size, theta, height, verbose=False, *args, **kwargs):\n",
    "        self.length = length\n",
    "        self.theta = theta\n",
    "        self.tantheta = tand(theta)\n",
    "        self.cos1theta = cosd(theta)\n",
    "        self.sin1theta = sind(theta)\n",
    "        self.height = height\n",
    "        self.bin_size = bin_size\n",
    "        self.bins_side = int(length / bin_size)\n",
    "        \n",
    "        self.bin_nums = int(self.bins_side**2)\n",
    "        self.bins = []\n",
    "        #self.bins = [[] for _ in range(self.bin_nums)]\n",
    "        #for b in self.bins:\n",
    "        #    b = SortedSet(key=lambda position: position[1])\n",
    "        #self.bins = [FasterSortedSet(key=lambda position: position[1])] * self.bin_nums\n",
    "        for i in range(self.bin_nums):\n",
    "            self.bins.append(FasterSortedSet(key=lambda position: position[1]))\n",
    "        print(\"Total number of bins: {}\".format(len(self.bins)))\n",
    "        \n",
    "        self.verbose = verbose\n",
    "        \n",
    "    def find_bin(self, position):\n",
    "        try:\n",
    "            adjusted_x = position[2] * self.tantheta + position[0]\n",
    "            spacex = int(np.floor((adjusted_x % self.length)//self.bin_size))\n",
    "            spacey = int(np.floor((position[1] % self.length)//self.bin_size))\n",
    "        except FloatingPointError:\n",
    "            raise FloatingPointError(\"\\n*\\tposition: {}\\n*\\tself.tantheta: {}\\n*\\tself.length: {}\\n*\\tself.bin_size: {}\".format(position, self.tantheta, self.length, self.bin_size))\n",
    "        \n",
    "        idx = spacex + spacey*self.bins_side\n",
    "        #print(position % (self.length, self.length, np.inf), idx, spacex, spacey, spacez, bin_size)\n",
    "        \n",
    "        return self.bins[idx]\n",
    "    \n",
    "    # over-inclusive as it assumes a square with side length 2x radius\n",
    "    def add_to_bins(self, center, radius, idx):\n",
    "        bins = self.find_bins(center, radius)\n",
    "        priority = self.calc_priority(center)\n",
    "        for b in bins:\n",
    "            b.add((idx, priority))\n",
    "        return bins\n",
    "    \n",
    "    # over-inclusive as it assumes a cube with side length 2x radius\n",
    "    def find_bins(self, center, radius):\n",
    "        center_bin = self.find_bin(center)\n",
    "        temp = [center_bin]\n",
    "        \n",
    "        left_back_lower_bin = self.find_bin((center + np.array(([-radius, -radius, -radius]))) % [self.length, self.length, np.inf])\n",
    "        if left_back_lower_bin not in temp:\n",
    "            temp.append(left_back_lower_bin)\n",
    "            \n",
    "        left_front_lower_bin = self.find_bin((center + np.array([-radius, radius, -radius])) % [self.length, self.length, np.inf])\n",
    "        if left_front_lower_bin not in temp:\n",
    "            temp.append(left_front_lower_bin)\n",
    "        \n",
    "        left_back_upper_bin = self.find_bin((center + np.array([-radius, -radius, radius])) % [self.length, self.length, np.inf])\n",
    "        if left_back_upper_bin not in temp:\n",
    "            temp.append(left_back_upper_bin)\n",
    "            \n",
    "        left_front_upper_bin = self.find_bin((center + np.array([-radius, radius, radius])) % [self.length, self.length, np.inf])\n",
    "        if left_front_upper_bin not in temp:\n",
    "            temp.append(left_front_upper_bin)\n",
    "            \n",
    "        right_back_lower_bin = self.find_bin((center + np.array([radius, -radius, -radius])) % [self.length, self.length, np.inf])\n",
    "        if right_back_lower_bin not in temp:\n",
    "            temp.append(right_back_lower_bin)\n",
    "            \n",
    "        right_front_lower_bin = self.find_bin((center + np.array([radius, radius, -radius])) % [self.length, self.length, np.inf])\n",
    "        if right_front_lower_bin not in temp:\n",
    "            temp.append(right_front_lower_bin)\n",
    "            \n",
    "        right_back_upper_bin = self.find_bin((center + np.array([radius, -radius, radius])) % [self.length, self.length, np.inf])\n",
    "        if right_back_upper_bin not in temp:\n",
    "            temp.append(right_back_upper_bin)\n",
    "            \n",
    "        right_front_upper_bin = self.find_bin((center + np.array([radius, radius, radius])) % [self.length, self.length, np.inf])\n",
    "        if right_front_upper_bin not in temp:\n",
    "            temp.append(right_front_upper_bin)\n",
    "        \n",
    "        return temp\n",
    "    \n",
    "    def find_bin_idx(self, position):\n",
    "        adjusted_x = position[2] * self.tantheta + position[0]\n",
    "        spacex = int(np.floor((adjusted_x % self.length)//self.bin_size))\n",
    "        spacey = int(np.floor((position[1] % self.length)//self.bin_size))\n",
    "        \n",
    "        idx = spacex + spacey*self.bins_side\n",
    "        if self.verbose:\n",
    "            print(\"Position: {}\".format(position))\n",
    "            #print(\"self.tandtheta: {}, {}\".format(self.tantheta, tand(self.theta)))\n",
    "            print(\"adjusted_x: {}\".format(adjusted_x))\n",
    "            print(\"Indices: {}, {}, {}\".format(spacex, spacey, idx))\n",
    "        return idx\n",
    "    \n",
    "    # over-inclusive as it assumes a cube with side length 2x radius\n",
    "    def find_bins_idx(self, center, radius):\n",
    "        center_bin = self.find_bin_idx(center)\n",
    "        temp = [center_bin]\n",
    "        \n",
    "        left_back_lower_bin = self.find_bin_idx((center + np.array(([-radius, -radius, -radius]))) % [self.length, self.length, self.height])\n",
    "        if left_back_lower_bin not in temp:\n",
    "            temp.append(left_back_lower_bin)\n",
    "            \n",
    "        left_front_lower_bin = self.find_bin_idx((center + np.array([-radius, radius, -radius])) % [self.length, self.length, self.height])\n",
    "        if left_front_lower_bin not in temp:\n",
    "            temp.append(left_front_lower_bin)\n",
    "        \n",
    "        left_back_upper_bin = self.find_bin_idx((center + np.array([-radius, -radius, radius])) % [self.length, self.length, self.height])\n",
    "        if left_back_upper_bin not in temp:\n",
    "            temp.append(left_back_upper_bin)\n",
    "            \n",
    "        left_front_upper_bin = self.find_bin_idx((center + np.array([-radius, radius, radius])) % [self.length, self.length, self.height])\n",
    "        if left_front_upper_bin not in temp:\n",
    "            temp.append(left_front_upper_bin)\n",
    "            \n",
    "        right_back_lower_bin = self.find_bin_idx((center + np.array([radius, -radius, -radius])) % [self.length, self.length, self.height])\n",
    "        if right_back_lower_bin not in temp:\n",
    "            temp.append(right_back_lower_bin)\n",
    "            \n",
    "        right_front_lower_bin = self.find_bin_idx((center + np.array([radius, radius, -radius])) % [self.length, self.length, self.height])\n",
    "        if right_front_lower_bin not in temp:\n",
    "            temp.append(right_front_lower_bin)\n",
    "            \n",
    "        right_back_upper_bin = self.find_bin_idx((center + np.array([radius, -radius, radius])) % [self.length, self.length, self.height])\n",
    "        if right_back_upper_bin not in temp:\n",
    "            temp.append(right_back_upper_bin)\n",
    "            \n",
    "        right_front_upper_bin = self.find_bin_idx(center + np.array([radius, radius, radius]))\n",
    "        if right_front_upper_bin not in temp:\n",
    "            temp.append(right_front_upper_bin)\n",
    "            \n",
    "        return temp\n",
    "    \n",
    "    # return particles that share bins\n",
    "    def get_nearby_particles(self, center, radius):\n",
    "        bins = self.find_bins(center, radius)\n",
    "        \n",
    "        if bins is []:\n",
    "            return [] # should never happen\n",
    "        \n",
    "        temp = []\n",
    "        for b in bins:\n",
    "            if len(b) > 0:\n",
    "                for i in b:\n",
    "                    if i not in temp:\n",
    "                        temp.append(i)\n",
    "        return temp\n",
    "    \n",
    "    def drop_particle(self, center, radius, atoms, profile=False):\n",
    "        if profile:\n",
    "            pr = cProfile.Profile()\n",
    "            pr.enable()\n",
    "        bins = FasterSortedSet(key=lambda position: position[1])\n",
    "        # combine the bins into one larger bin\n",
    "        fbins = self.find_bins(center, radius)\n",
    "        for b in fbins:\n",
    "            bins |= b\n",
    "        adjusted_x = center[2] * self.tantheta + center[0]\n",
    "        factor = adjusted_x // self.length\n",
    "        adjusted_x %= self.length\n",
    "        center_float = center.astype(float)\n",
    "        \n",
    "        \n",
    "        for i in range(0, len(bins)):\n",
    "            # check collision with bins[i][0]-th particle\n",
    "            particle = atoms[bins[len(bins) - i-1][0],:]\n",
    "            rs = radius + particle[4]\n",
    "                \n",
    "            x_on_particle_z = (center[2] - particle[2]) * self.tantheta + center[0] # x coordinate on z=particle[2] plane\n",
    "            factor = x_on_particle_z // self.length                   # get multiple of self.length to shift by\n",
    "            x_on_particle_z %= self.length                            # restrict x coordinate within PBC\n",
    "            \n",
    "            #center_proj_sub = np.array([x_on_particle_z, center[1], particle[2]])\n",
    "            #center_proj_height = center_float - [factor*self.length, 0, 0]\n",
    "            \n",
    "            offset = [0,0,0]\n",
    "            # offset_x, shift center nearer to account for PBC\n",
    "            #if particle[0] < rs and center_proj_sub[0] > self.length - rs:\n",
    "            if particle[0] - x_on_particle_z < -self.length/2:\n",
    "                offset[0] = -self.length\n",
    "            #elif particle[0] > self.length - rs and center_proj_sub[0] < rs:\n",
    "            elif particle[0] - x_on_particle_z > self.length/2:\n",
    "                offset[0] = self.length\n",
    "            # offset_y, shift center nearer to account for PBC\n",
    "            #if particle[1] < rs and center_proj_sub[1] > self.length - rs:\n",
    "            if particle[1] - center[1] < -self.length/2:\n",
    "                offset[1] = -self.length\n",
    "            #elif particle[1] > self.length - rs and center_proj_sub[1] < rs:\n",
    "            elif particle[1] - center[1] > self.length/2:\n",
    "                offset[1] = self.length\n",
    "            \n",
    "            #proj_x = particle[2] * self.tantheta + center[0]\n",
    "            #print('\\t{}'.format(particle[:3]))\n",
    "            \n",
    "            #d = shortest_distance_line_to_point(center_proj_height + offset, center_proj_sub+offset, particle)\n",
    "            radii2 = rs**2\n",
    "            \n",
    "            \n",
    "                \n",
    "            # rotate axes by theta degrees about y-axis into new coordinate system so that z' axis is unimportant\n",
    "            #rotx = (center_proj_height[0] - particle[0] + offset[0])*self.cos1theta + (center_proj_height[2] - particle[2])*self.sin1theta\n",
    "            #d2 = rotx**2 + (center_proj_height[1] - particle[1] + offset[1])**2\n",
    "            rotx = (center_float[0] - factor*self.length - particle[0] + offset[0])*self.cos1theta + (center_float[2] - particle[2])*self.sin1theta\n",
    "            d2 = rotx**2 + (center_float[1] - particle[1] + offset[1])**2\n",
    "            \n",
    "            if self.verbose:\n",
    "                print(\"Generated position: {}\".format(center))\n",
    "                print(\"x-coordinate in new coordinate system: {}\".format(rotx))\n",
    "                print(\"Offset: {}\".format(offset))\n",
    "                #print(\"center_proj_height + offset: {}\".format(center_proj_height+offset))\n",
    "                #print(\"center_proj_sub + offset: {}\".format(center_proj_sub+offset))\n",
    "                #print(\"Adjusted generated position: {}\".format(center_proj_height))\n",
    "                print(\"Existing particle position being considered: {}\".format(particle[:3]))\n",
    "                #print(\"Offset (if any): {}\".format(offset))\n",
    "                print(\"Minimum distance^2 between: {}\".format(d2))\n",
    "                print(\"Radii added together: {}\".format(rs))\n",
    "                \n",
    "            if radii2 > d2:\n",
    "                # it collides, calculate the new position\n",
    "                rotz = np.sqrt(radii2 - d2)\n",
    "                z = rotx*self.sin1theta + rotz*self.cos1theta + particle[2]\n",
    "                x = rotx*self.cos1theta - rotz*self.sin1theta + particle[0]\n",
    "\n",
    "                if z < radius:\n",
    "                    # if it would collide with the surface below the substrate for this particle then we can skip this particle\n",
    "                    continue\n",
    "                if self.verbose:\n",
    "                    print(\"Particles collide.\")\n",
    "                    print(\"Coordinates in new coordinate system: {}\".format([rotx, center[1], rotz]))\n",
    "                    print(\"Distance: {}\".format(np.sqrt((x - particle[0])**2 + (center[1] - particle[1])**2 + (z - particle[2])**2)))\n",
    "#                     print(\"Start: {}\".format(center_proj_height))\n",
    "#                     print(\"Start (reprojected): {}\".format(center_proj_height + offset))\n",
    "#                     print(\"Target: {}\".format(center_proj_sub))\n",
    "#                     print(\"Target (reprojected): {}\".format(center_proj_sub+offset))\n",
    "                    print(\"Collision position: {}\".format([x, center[1], z]))\n",
    "                    print(\"Colliding atom position: {}\\n\".format(particle[:3]))\n",
    "                #print(\"Hit particle: {} -> {}\".format(center, [x, center[1], z]))\n",
    "                if profile:\n",
    "                    pr.disable()\n",
    "                    p = pstats.Stats(pr)\n",
    "                    p.sort_stats('time').print_stats()\n",
    "                    #pr.print_stats()\n",
    "                return np.array([x, center[1], z]), bins[len(bins) - i-1][0]\n",
    "        x = adjusted_x - radius*self.tantheta\n",
    "        #print(\"Didn't hit particle: {} -> {}\".format(center, [x, center[1], radius]))\n",
    "        if self.verbose:\n",
    "            print(\"Particle did not collide with existing surface; landing at {}.\\n\".format([x, center[1], radius]))\n",
    "        if profile:\n",
    "            pr.disable()\n",
    "            pr.print_stats()\n",
    "        return np.array([x, center[1], radius]), -1\n",
    "    '''\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    z\n",
    "    '''\n",
    "    def drop_particle_waffle(self, center, radius, atoms, profile=False):\n",
    "        if profile:\n",
    "            pr = cProfile.Profile()\n",
    "            pr.enable()\n",
    "        #bins = FasterSortedSet(key=lambda position: position[1])\n",
    "        # combine the bins into one larger bin\n",
    "        fbins = self.find_bins(center, radius)\n",
    "        #for b in fbins:\n",
    "        #    bins |= b\n",
    "        bin_len = 0\n",
    "        val_array = np.zeros(len(fbins))\n",
    "        particle_array = [None]*len(fbins)\n",
    "        used_array = [1]*len(fbins)\n",
    "        for i in range(0, len(fbins)):\n",
    "            bin_len += len(fbins[i])\n",
    "            j = 1\n",
    "            if len(fbins[i]) > 0:\n",
    "                temp_particle = fbins[i][-j]\n",
    "                while fbins[i][-j][0] in particle_array and j < len(fbins[i]):\n",
    "                    j += 1\n",
    "                    if j > len(fbins[i]):\n",
    "                        temp_particle = [-np.inf, 0]\n",
    "                    else:\n",
    "                        temp_particle = fbins[i][-j]\n",
    "                particle_array[i] = temp_particle[0]\n",
    "                val_array[i] = temp_particle[1]\n",
    "            else:\n",
    "                particle_array[i] = -np.inf\n",
    "                val_array[i] = 0\n",
    "        print(\"fbins: {}\\nparticles: {}\\nvalues: {}\\nused: {}\".format(fbins, particle_array, val_array, used_array))\n",
    "        used_array[i] = j\n",
    "            \n",
    "        adjusted_x = center[2] * self.tantheta + center[0]\n",
    "        factor = adjusted_x // self.length\n",
    "        adjusted_x %= self.length\n",
    "        center_float = center.astype(float)\n",
    "        \n",
    "        \n",
    "        for i in range(0, bin_len):\n",
    "            print(\"iteration {}\".format(i))\n",
    "            # figure out which bin has the highest priority particle and replace\n",
    "            highest_priority_bin = np.argmax(val_array)\n",
    "            print(\"\\tfbins: {}\\n\\tparticles: {}\\n\\tvalues: {}\\n\\tused: {}\".format(fbins, particle_array, val_array, used_array))\n",
    "            particle = atoms[particle_array[highest_priority_bin], :]\n",
    "            j = used_array[highest_priority_bin]\n",
    "            if j < len(fbins[highest_priority_bin]):\n",
    "                while fbins[highest_priority_bin][-j][0] in particle_array and j < len(fbins[highest_priority_bin]):\n",
    "                    print(\"\\t\\tj: {}\".format(j))\n",
    "                    j += 1\n",
    "                    if j == len(fbins[highest_priority_bin]):\n",
    "                        temp_particle = [-np.inf, 0]\n",
    "                    else:\n",
    "                        temp_particle = fbins[highest_priority_bin][-j]\n",
    "                \n",
    "                particle_array[highest_priority_bin] = temp_particle[0]\n",
    "                val_array[highest_priority_bin] = temp_particle[1]\n",
    "            else:\n",
    "                particle_array[highest_priority_bin] = -np.inf\n",
    "                val_array[highest_priority_bin] = 0\n",
    "            used_array[highest_priority_bin] = j\n",
    "                \n",
    "            \n",
    "            # check collision with bins[i][0]-th particle\n",
    "            #particle = atoms[bins[len(bins) - i-1][0],:]\n",
    "            rs = radius + particle[4]\n",
    "            \n",
    "            x_on_particle_z = (center[2] - particle[2]) * self.tantheta + center[0] # x coordinate on z=particle[2] plane\n",
    "            factor = x_on_particle_z // self.length                   # get multiple of self.length to shift by\n",
    "            x_on_particle_z %= self.length                            # restrict x coordinate within PBC\n",
    "            \n",
    "            #center_proj_sub = np.array([x_on_particle_z, center[1], particle[2]])\n",
    "            #center_proj_height = center_float - [factor*self.length, 0, 0]\n",
    "            \n",
    "            offset = [0,0,0]\n",
    "            # offset_x, shift center nearer to account for PBC\n",
    "            #if particle[0] < rs and center_proj_sub[0] > self.length - rs:\n",
    "            if particle[0] - x_on_particle_z < -self.length/2:\n",
    "                offset[0] = -self.length\n",
    "            #elif particle[0] > self.length - rs and center_proj_sub[0] < rs:\n",
    "            elif particle[0] - x_on_particle_z > self.length/2:\n",
    "                offset[0] = self.length\n",
    "            # offset_y, shift center nearer to account for PBC\n",
    "            #if particle[1] < rs and center_proj_sub[1] > self.length - rs:\n",
    "            if particle[1] - center[1] < -self.length/2:\n",
    "                offset[1] = -self.length\n",
    "            #elif particle[1] > self.length - rs and center_proj_sub[1] < rs:\n",
    "            elif particle[1] - center[1] > self.length/2:\n",
    "                offset[1] = self.length\n",
    "            \n",
    "            #proj_x = particle[2] * self.tantheta + center[0]\n",
    "            #print('\\t{}'.format(particle[:3]))\n",
    "            \n",
    "            #d = shortest_distance_line_to_point(center_proj_height + offset, center_proj_sub+offset, particle)\n",
    "            radii2 = rs**2\n",
    "            \n",
    "            \n",
    "                \n",
    "            # rotate axes by theta degrees about y-axis into new coordinate system so that z' axis is unimportant\n",
    "            #rotx = (center_proj_height[0] - particle[0] + offset[0])*self.cos1theta + (center_proj_height[2] - particle[2])*self.sin1theta\n",
    "            #d2 = rotx**2 + (center_proj_height[1] - particle[1] + offset[1])**2\n",
    "            rotx = (center_float[0] - factor*self.length - particle[0] + offset[0])*self.cos1theta + (center_float[2] - particle[2])*self.sin1theta\n",
    "            d2 = rotx**2 + (center_float[1] - particle[1] + offset[1])**2\n",
    "            \n",
    "            if self.verbose:\n",
    "                print(\"Generated position: {}\".format(center))\n",
    "                print(\"x-coordinate in new coordinate system: {}\".format(rotx))\n",
    "                print(\"Offset: {}\".format(offset))\n",
    "                #print(\"center_proj_height + offset: {}\".format(center_proj_height+offset))\n",
    "                #print(\"center_proj_sub + offset: {}\".format(center_proj_sub+offset))\n",
    "                #print(\"Adjusted generated position: {}\".format(center_proj_height))\n",
    "                print(\"Existing particle position being considered: {}\".format(particle[:3]))\n",
    "                #print(\"Offset (if any): {}\".format(offset))\n",
    "                print(\"Minimum distance^2 between: {}\".format(d2))\n",
    "                print(\"Radii added together: {}\".format(rs))\n",
    "                \n",
    "            if radii2 > d2:\n",
    "                # it collides, calculate the new position\n",
    "                rotz = np.sqrt(radii2 - d2)\n",
    "                z = rotx*self.sin1theta + rotz*self.cos1theta + particle[2]\n",
    "                x = rotx*self.cos1theta - rotz*self.sin1theta + particle[0]\n",
    "\n",
    "                if z < radius:\n",
    "                    # if it would collide with the surface below the substrate for this particle then we can skip this particle\n",
    "                    continue\n",
    "                if self.verbose:\n",
    "                    print(\"Particles collide.\")\n",
    "                    print(\"Coordinates in new coordinate system: {}\".format([rotx, center[1], rotz]))\n",
    "                    print(\"Distance: {}\".format(np.sqrt((x - particle[0])**2 + (center[1] - particle[1])**2 + (z - particle[2])**2)))\n",
    "#                     print(\"Start: {}\".format(center_proj_height))\n",
    "#                     print(\"Start (reprojected): {}\".format(center_proj_height + offset))\n",
    "#                     print(\"Target: {}\".format(center_proj_sub))\n",
    "#                     print(\"Target (reprojected): {}\".format(center_proj_sub+offset))\n",
    "                    print(\"Collision position: {}\".format([x, center[1], z]))\n",
    "                    print(\"Colliding atom position: {}\\n\".format(particle[:3]))\n",
    "                #print(\"Hit particle: {} -> {}\".format(center, [x, center[1], z]))\n",
    "                if profile:\n",
    "                    pr.disable()\n",
    "                    p = pstats.Stats(pr)\n",
    "                    p.sort_stats('time').print_stats()\n",
    "                    #pr.print_stats()\n",
    "                return np.array([x, center[1], z]), particle\n",
    "        x = adjusted_x - radius*self.tantheta\n",
    "        #print(\"Didn't hit particle: {} -> {}\".format(center, [x, center[1], radius]))\n",
    "        if self.verbose:\n",
    "            print(\"Particle did not collide with existing surface; landing at {}.\\n\".format([x, center[1], radius]))\n",
    "        if profile:\n",
    "            pr.disable()\n",
    "            pr.print_stats()\n",
    "        return np.array([x, center[1], radius]), -1\n",
    "        \n",
    "    def check_collision(self, center, radius, atoms):\n",
    "        if(center[2] - radius < 0):\n",
    "            return 2 # collision with substrate\n",
    "        potential_collisions = self.get_nearby_particles(center, radius)\n",
    "        if potential_collisions == []:\n",
    "            return 0 # no nearby atoms\n",
    "        for c in potential_collisions:\n",
    "            if distance3D2PBC(atoms[c,:3], center, self.length) < (atoms[c,4] + radius)**2:\n",
    "                return 3 + c # collision\n",
    "        return 1 # no collision but atoms nearby\n",
    "    \n",
    "    def calc_priority(self, position):\n",
    "        #dx = position[0] % self.bin_size\n",
    "        return position[2]#-(dx * self.sin1theta - position[2] * self.cos1theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def follow_potential(pot_field, center, scale, L, h, distance, current_atom, lennardjones=True, deterministic=False, ratio=1.0):\n",
    "    if distance == 0:\n",
    "        return center\n",
    "    if lennardjones:\n",
    "        steps = int(round(scale*distance))\n",
    "        position = center\n",
    "        scaled_position = np.round(position * scale).astype(np.int32) % pot_field.shape[:3]\n",
    "        \n",
    "        if deterministic:\n",
    "            remaining_distance = steps\n",
    "            rads = int(round(scale*2*radius*ratio))\n",
    "            while remaining_distance > 0:\n",
    "                it = np.min([remaining_distance, rads])\n",
    "                remaining_distance -= it\n",
    "            \n",
    "                minx = scaled_position[0]-it#steps\n",
    "                if minx < 0:\n",
    "                    minx = 0\n",
    "                miny = scaled_position[1]-it#steps\n",
    "                if miny < 0:\n",
    "                    miny = 0\n",
    "                minz = scaled_position[2]-it#steps\n",
    "                if minz < int(np.ceil(scale*radius)):\n",
    "                    minz = int(np.ceil(scale*radius))\n",
    "                maxx = scaled_position[0]+it+1#steps+1\n",
    "                if maxx > pot_field.shape[0]-1:\n",
    "                    maxx = pot_field.shape[0]-1\n",
    "                maxy = scaled_position[1]+it+1#steps+1\n",
    "                if maxy > pot_field.shape[1]-1:\n",
    "                    maxy = pot_field.shape[1]-1\n",
    "                maxz = scaled_position[2]+it+1#steps+1\n",
    "\n",
    "                region = pot_field[minx:maxx, miny:maxy, minz:maxz, 3]\n",
    "                if 0 in region.shape:\n",
    "                    print(\"steps: {}\".format(steps))\n",
    "                    print(\"real position: {}\".format(position))\n",
    "                    print(\"scaled_position: {}\".format(scaled_position))\n",
    "                    print(\"pot_field.shape: {}\".format(pot_field.shape))\n",
    "                    print(\"region.shape: {}\".format(region.shape))\n",
    "                    print(\"{}->{}, {}->{}, {}->{}\".format(minx, maxx, miny, maxy, minz, maxz))\n",
    "                minimum = np.argmin(region)\n",
    "                minimum_i = np.unravel_index(minimum, region.shape)\n",
    "\n",
    "                #print(pot_field[scaled_position[0], scaled_position[1], scaled_position[2], 3], minimum_i, region[minimum_i])\n",
    "                if pot_field[scaled_position[0], scaled_position[1], scaled_position[2], 3] == region[minimum_i]:\n",
    "                    return position\n",
    "                minimum_i = np.array(minimum_i) + (minx, miny, minz)\n",
    "            return np.array(minimum_i)/np.array((scale,scale,scale)) % (L,L,np.inf)\n",
    "        \n",
    "        for i in range(steps):\n",
    "            try:\n",
    "                direction = pot_field[scaled_position[0], scaled_position[1], scaled_position[2], :3]\n",
    "            except IndexError:\n",
    "                print(position, scaled_position)\n",
    "                return position\n",
    "            # check if all zero because we stop there\n",
    "            if not np.any(direction):\n",
    "                return position % (L,L,h)\n",
    "            position = position + direction / scale\n",
    "            if position[2] < 0:\n",
    "                position[2] = 0\n",
    "            elif position[2] > h:\n",
    "                position[2] = h\n",
    "            #print(direction, position, position % (L,L,h))\n",
    "            position = position % (L,L,h)\n",
    "            scaled_position = np.round(position * scale).astype(np.int32) % pot_field.shape[:3]\n",
    "        return position % (L,L,h)\n",
    "    else:\n",
    "        k = 1\n",
    "        diameter = 1\n",
    "        position = center\n",
    "        L2 = L**2\n",
    "        step_size = 0.01\n",
    "        for d in np.arange(0, distance, step_size):\n",
    "            # get nearest neighbors within diameter radius\n",
    "            displacementx = position[0] - atom_pos[:current_atom,0]\n",
    "            displacementx = np.where(displacementx > L/2, -displacementx + L/2, displacementx)\n",
    "            displacementx = np.where(displacementx < -L/2, -displacementx - L/2, displacementx)\n",
    "            displacementy = position[1] - atom_pos[:current_atom,1]\n",
    "            displacementy = np.where(displacementy > L/2, -displacementy + L/2, displacementy)\n",
    "            displacementy = np.where(displacementy < -L/2, -displacementy - L/2, displacementy)\n",
    "            displacementz = position[2] - atom_pos[:current_atom,2]\n",
    "            distances = np.sqrt((displacementx)**2 + (displacementy)**2 + displacementz**2)\n",
    "            idx = np.argwhere(np.where(distances < diameter, distances, 0)) # indices of nearest neighbor\n",
    "            #print(d, distances.shape, distances[idx].shape)\n",
    "            distances = distances[idx]\n",
    "            displacement = np.hstack([displacementx[idx], displacementy[idx], displacementz[idx]])\n",
    "            #print(displacement.shape)\n",
    "            f = -np.sum(k * (distances - radius*2) * displacement/np.hstack([distances, distances, distances]), axis=0)\n",
    "            if not ((f[0] == -0.0 and f[1] == -0.0 and f[2] == -0.0) or np.any(np.isnan(f))):\n",
    "                position += step_size * f/np.linalg.norm(f)\n",
    "                position = position % (L,L,np.inf)\n",
    "            #print(position)\n",
    "        #print(center)\n",
    "        #print(position)\n",
    "        #print(np.hstack([displacement, atom_pos[:,0][idx], atom_pos[:,1][idx], atom_pos[:,2][idx]]))\n",
    "        return position % (L,L,np.inf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "L: \t48\n",
      "theta: \t85\n",
      "phi: \t0\n",
      "H: \t0.0\n",
      "D: \t0.45\n",
      "R: \t1\n",
      "A: \t1\n",
      "turns: \t0\n",
      "stepper resolution: \t0\n",
      "species: \t1\n",
      "radius: \t0.15\n",
      "weights: \tNone\n",
      "repetition: \t131072\n",
      "spread: \t0\n",
      "time: \t25748.940731048584\n",
      "Total number of bins: 36\n",
      "         4063242 function calls in 23.723 seconds\n",
      "\n",
      "   Ordered by: internal time\n",
      "   List reduced from 69 to 10 due to restriction <10>\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "  1179648   11.765    0.000   11.765    0.000 <ipython-input-134-729a6894b696>:21(find_bin)\n",
      "   131072    8.061    0.000   22.104    0.000 <ipython-input-134-729a6894b696>:43(find_bins)\n",
      "  1048576    2.277    0.000    2.277    0.000 {built-in method numpy.array}\n",
      "   131072    0.385    0.000    0.893    0.000 F:\\Programs\\Anaconda3_2020.11\\lib\\site-packages\\sortedcontainers\\sortedlist.py:1777(add)\n",
      "        1    0.239    0.239   23.723   23.723 <ipython-input-135-83996588ce4f>:8(<module>)\n",
      "   131072    0.216    0.000   23.484    0.000 <ipython-input-134-729a6894b696>:35(add_to_bins)\n",
      "   261768    0.207    0.000    0.207    0.000 {built-in method _bisect.bisect_right}\n",
      "   131072    0.193    0.000    1.118    0.000 F:\\Programs\\Anaconda3_2020.11\\lib\\site-packages\\sortedcontainers\\sortedset.py:341(add)\n",
      "   261664    0.123    0.000    0.123    0.000 {method 'insert' of 'list' objects}\n",
      "   131071    0.121    0.000    0.136    0.000 F:\\Programs\\Anaconda3_2020.11\\lib\\site-packages\\sortedcontainers\\sortedlist.py:1821(_expand)\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<pstats.Stats at 0x1ac545b8850>"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load an old simulation to get a BSP object ready\n",
    "atoms, deposited, p = loadContinuum('STF_Si_L48_Th85_D0.45_N131072_1672876042' + '.npz')\n",
    "for k in p[0]:\n",
    "    print(\"{}: \\t{}\".format(k, p[0][k]))\n",
    "pr = cProfile.Profile()\n",
    "pr.enable()\n",
    "bins = SlantedCorridors(p[0]['L'], 8, p[0]['theta'], 32, verbose=False)\n",
    "for i in range(atoms.shape[0]):\n",
    "    bins.add_to_bins(atoms[i,:3], p[0]['radius'], i)\n",
    "pr.disable()\n",
    "ps = pstats.Stats(pr)\n",
    "ps.sort_stats('time').print_stats(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         29305818 function calls (29029175 primitive calls) in 33.224 seconds\n",
      "\n",
      "   Ordered by: internal time\n",
      "   List reduced from 55 to 10 due to restriction <10>\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "      200    9.370    0.047   11.764    0.059 {built-in method builtins.sorted}\n",
      "      100    9.285    0.093   32.461    0.325 <ipython-input-134-729a6894b696>:148(drop_particle)\n",
      " 26214400    4.545    0.000    4.545    0.000 <ipython-input-134-729a6894b696>:152(<lambda>)\n",
      "    13400    3.256    0.000    5.407    0.000 F:\\Programs\\Anaconda3_2020.11\\lib\\site-packages\\sortedcontainers\\sortedlist.py:1893(<genexpr>)\n",
      "      100    2.915    0.029   20.966    0.210 F:\\Programs\\Anaconda3_2020.11\\lib\\site-packages\\sortedcontainers\\sortedset.py:664(update)\n",
      "   276543    1.043    0.000    1.096    0.000 F:\\Programs\\Anaconda3_2020.11\\lib\\site-packages\\sortedcontainers\\sortedlist.py:600(_pos)\n",
      "        1    0.763    0.763   33.224   33.224 <ipython-input-136-23b370a901b4>:15(<module>)\n",
      "      100    0.673    0.007    0.673    0.007 {method 'update' of 'set' objects}\n",
      "   276543    0.563    0.000    1.753    0.000 F:\\Programs\\Anaconda3_2020.11\\lib\\site-packages\\sortedcontainers\\sortedlist.py:815(__getitem__)\n",
      "1661358/1384715    0.268    0.000    0.365    0.000 {built-in method builtins.len}\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<pstats.Stats at 0x1ac4ecbe190>"
      ]
     },
     "execution_count": 136,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# generate random particle\n",
    "rng = np.random.default_rng()\n",
    "dests = np.hstack([rng.random(size=(100, 2), dtype=np.float32)*48, np.zeros((100,1), dtype=np.float32)+32]).squeeze()\n",
    "#print(\"Generated particle position: {}\".format(dest))\n",
    "\n",
    "# drop test particle\n",
    "#bins.verbose = False\n",
    "#bins.find_bins(dest, p[0]['radius'])\n",
    "#bins.find_bins(dest, radius)\n",
    "radius = p[0]['radius']\n",
    "#print(radius)\n",
    "L = p[0]['L']\n",
    "pr = cProfile.Profile()\n",
    "pr.enable()\n",
    "for i in range(dests.shape[0]):\n",
    "    bins.drop_particle(dests[i], radius, atoms, False)\n",
    "pr.disable()\n",
    "ps = pstats.Stats(pr)\n",
    "ps.sort_stats('time').print_stats(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of bins: 16\n",
      "Scaled:\tZero-potential point: 2.672696154421018\n",
      "\tMinimum-potential point: 3.0\n",
      "Real:\tZero-potential point: 0.2672696154421018\n",
      "\tMinimum-potential point: 0.3\n",
      "256/32768 complete; 0.285 seconds taken; 898.044 reps per second.\n",
      "512/32768 complete; 0.731 seconds taken; 700.243 reps per second.\n",
      "768/32768 complete; 1.272 seconds taken; 603.633 reps per second.\n",
      "1024/32768 complete; 1.844 seconds taken; 555.189 reps per second.\n",
      "1280/32768 complete; 2.494 seconds taken; 513.267 reps per second.\n",
      "1536/32768 complete; 3.144 seconds taken; 488.551 reps per second.\n",
      "1792/32768 complete; 3.780 seconds taken; 474.053 reps per second.\n",
      "2048/32768 complete; 4.462 seconds taken; 459.001 reps per second.\n",
      "2304/32768 complete; 5.153 seconds taken; 447.082 reps per second.\n",
      "2560/32768 complete; 5.871 seconds taken; 436.061 reps per second.\n",
      "2816/32768 complete; 6.632 seconds taken; 424.619 reps per second.\n",
      "3072/32768 complete; 7.414 seconds taken; 414.364 reps per second.\n",
      "3328/32768 complete; 8.258 seconds taken; 403.024 reps per second.\n",
      "3584/32768 complete; 9.049 seconds taken; 396.050 reps per second.\n",
      "3840/32768 complete; 9.880 seconds taken; 388.676 reps per second.\n",
      "4096/32768 complete; 10.819 seconds taken; 378.584 reps per second.\n",
      "4352/32768 complete; 11.696 seconds taken; 372.078 reps per second.\n",
      "4608/32768 complete; 12.636 seconds taken; 364.685 reps per second.\n",
      "4864/32768 complete; 13.581 seconds taken; 358.137 reps per second.\n",
      "5120/32768 complete; 14.574 seconds taken; 351.299 reps per second.\n",
      "5376/32768 complete; 15.567 seconds taken; 345.346 reps per second.\n",
      "5632/32768 complete; 16.572 seconds taken; 339.844 reps per second.\n",
      "5888/32768 complete; 17.599 seconds taken; 334.559 reps per second.\n",
      "6144/32768 complete; 18.586 seconds taken; 330.576 reps per second.\n",
      "6400/32768 complete; 19.659 seconds taken; 325.543 reps per second.\n",
      "6656/32768 complete; 20.833 seconds taken; 319.486 reps per second.\n",
      "6912/32768 complete; 21.855 seconds taken; 316.267 reps per second.\n",
      "7168/32768 complete; 22.969 seconds taken; 312.068 reps per second.\n",
      "7424/32768 complete; 24.151 seconds taken; 307.403 reps per second.\n",
      "7680/32768 complete; 25.344 seconds taken; 303.026 reps per second.\n",
      "7936/32768 complete; 26.534 seconds taken; 299.085 reps per second.\n",
      "8192/32768 complete; 27.781 seconds taken; 294.875 reps per second.\n",
      "8448/32768 complete; 29.063 seconds taken; 290.676 reps per second.\n",
      "8704/32768 complete; 30.281 seconds taken; 287.440 reps per second.\n",
      "8960/32768 complete; 31.523 seconds taken; 284.241 reps per second.\n",
      "9216/32768 complete; 32.715 seconds taken; 281.702 reps per second.\n",
      "9472/32768 complete; 33.981 seconds taken; 278.741 reps per second.\n",
      "9728/32768 complete; 35.284 seconds taken; 275.708 reps per second.\n",
      "9984/32768 complete; 36.699 seconds taken; 272.050 reps per second.\n",
      "10240/32768 complete; 38.107 seconds taken; 268.718 reps per second.\n",
      "10496/32768 complete; 39.436 seconds taken; 266.154 reps per second.\n",
      "10752/32768 complete; 40.784 seconds taken; 263.634 reps per second.\n",
      "11008/32768 complete; 42.214 seconds taken; 260.765 reps per second.\n",
      "11264/32768 complete; 43.724 seconds taken; 257.618 reps per second.\n",
      "11520/32768 complete; 45.236 seconds taken; 254.666 reps per second.\n",
      "11776/32768 complete; 46.769 seconds taken; 251.791 reps per second.\n",
      "12032/32768 complete; 48.191 seconds taken; 249.671 reps per second.\n",
      "12288/32768 complete; 49.712 seconds taken; 247.186 reps per second.\n",
      "12544/32768 complete; 51.183 seconds taken; 245.079 reps per second.\n",
      "12800/32768 complete; 52.817 seconds taken; 242.347 reps per second.\n",
      "13056/32768 complete; 54.376 seconds taken; 240.105 reps per second.\n",
      "13312/32768 complete; 55.922 seconds taken; 238.047 reps per second.\n",
      "13568/32768 complete; 57.451 seconds taken; 236.166 reps per second.\n",
      "13824/32768 complete; 59.062 seconds taken; 234.061 reps per second.\n",
      "14080/32768 complete; 60.782 seconds taken; 231.649 reps per second.\n",
      "14336/32768 complete; 62.429 seconds taken; 229.639 reps per second.\n",
      "14592/32768 complete; 64.029 seconds taken; 227.896 reps per second.\n",
      "14848/32768 complete; 65.614 seconds taken; 226.295 reps per second.\n",
      "15104/32768 complete; 67.196 seconds taken; 224.777 reps per second.\n",
      "15360/32768 complete; 68.781 seconds taken; 223.317 reps per second.\n",
      "15616/32768 complete; 70.493 seconds taken; 221.525 reps per second.\n",
      "15872/32768 complete; 72.167 seconds taken; 219.933 reps per second.\n",
      "16128/32768 complete; 73.841 seconds taken; 218.415 reps per second.\n",
      "16384/32768 complete; 75.492 seconds taken; 217.029 reps per second.\n",
      "16640/32768 complete; 77.200 seconds taken; 215.545 reps per second.\n",
      "16896/32768 complete; 78.959 seconds taken; 213.985 reps per second.\n",
      "17152/32768 complete; 80.628 seconds taken; 212.729 reps per second.\n",
      "17408/32768 complete; 82.432 seconds taken; 211.180 reps per second.\n",
      "17664/32768 complete; 84.152 seconds taken; 209.906 reps per second.\n",
      "17920/32768 complete; 86.004 seconds taken; 208.362 reps per second.\n",
      "18176/32768 complete; 87.916 seconds taken; 206.743 reps per second.\n",
      "18432/32768 complete; 89.634 seconds taken; 205.637 reps per second.\n",
      "18688/32768 complete; 91.221 seconds taken; 204.865 reps per second.\n",
      "18944/32768 complete; 93.103 seconds taken; 203.474 reps per second.\n",
      "19200/32768 complete; 95.013 seconds taken; 202.078 reps per second.\n",
      "19456/32768 complete; 97.046 seconds taken; 200.482 reps per second.\n",
      "19712/32768 complete; 98.867 seconds taken; 199.380 reps per second.\n",
      "19968/32768 complete; 100.750 seconds taken; 198.193 reps per second.\n",
      "20224/32768 complete; 102.686 seconds taken; 196.951 reps per second.\n",
      "20480/32768 complete; 104.525 seconds taken; 195.934 reps per second.\n",
      "20736/32768 complete; 106.531 seconds taken; 194.648 reps per second.\n",
      "20992/32768 complete; 108.615 seconds taken; 193.269 reps per second.\n",
      "21248/32768 complete; 110.644 seconds taken; 192.039 reps per second.\n",
      "21504/32768 complete; 112.861 seconds taken; 190.536 reps per second.\n",
      "21760/32768 complete; 114.777 seconds taken; 189.585 reps per second.\n",
      "22016/32768 complete; 116.727 seconds taken; 188.612 reps per second.\n",
      "22272/32768 complete; 118.787 seconds taken; 187.496 reps per second.\n",
      "22528/32768 complete; 120.718 seconds taken; 186.617 reps per second.\n",
      "22784/32768 complete; 122.608 seconds taken; 185.828 reps per second.\n",
      "23040/32768 complete; 124.634 seconds taken; 184.861 reps per second.\n",
      "23296/32768 complete; 126.792 seconds taken; 183.734 reps per second.\n",
      "23552/32768 complete; 128.821 seconds taken; 182.828 reps per second.\n",
      "23808/32768 complete; 130.839 seconds taken; 181.964 reps per second.\n",
      "24064/32768 complete; 133.022 seconds taken; 180.902 reps per second.\n",
      "24320/32768 complete; 135.218 seconds taken; 179.858 reps per second.\n",
      "24576/32768 complete; 137.453 seconds taken; 178.795 reps per second.\n",
      "24832/32768 complete; 139.646 seconds taken; 177.821 reps per second.\n",
      "25088/32768 complete; 141.758 seconds taken; 176.978 reps per second.\n",
      "25344/32768 complete; 143.994 seconds taken; 176.007 reps per second.\n",
      "25600/32768 complete; 146.357 seconds taken; 174.915 reps per second.\n",
      "25856/32768 complete; 148.722 seconds taken; 173.855 reps per second.\n",
      "26112/32768 complete; 150.815 seconds taken; 173.140 reps per second.\n",
      "26368/32768 complete; 153.223 seconds taken; 172.089 reps per second.\n",
      "26624/32768 complete; 155.535 seconds taken; 171.177 reps per second.\n",
      "26880/32768 complete; 157.989 seconds taken; 170.138 reps per second.\n",
      "27136/32768 complete; 160.313 seconds taken; 169.269 reps per second.\n",
      "27392/32768 complete; 162.695 seconds taken; 168.364 reps per second.\n",
      "27648/32768 complete; 165.093 seconds taken; 167.469 reps per second.\n",
      "27904/32768 complete; 167.424 seconds taken; 166.666 reps per second.\n",
      "28160/32768 complete; 169.651 seconds taken; 165.987 reps per second.\n",
      "28416/32768 complete; 171.966 seconds taken; 165.242 reps per second.\n",
      "28672/32768 complete; 174.245 seconds taken; 164.549 reps per second.\n",
      "28928/32768 complete; 176.447 seconds taken; 163.947 reps per second.\n",
      "29184/32768 complete; 178.909 seconds taken; 163.122 reps per second.\n",
      "29440/32768 complete; 181.380 seconds taken; 162.311 reps per second.\n",
      "29696/32768 complete; 183.883 seconds taken; 161.494 reps per second.\n",
      "29952/32768 complete; 186.143 seconds taken; 160.908 reps per second.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30208/32768 complete; 188.511 seconds taken; 160.245 reps per second.\n",
      "30464/32768 complete; 190.793 seconds taken; 159.671 reps per second.\n",
      "30720/32768 complete; 193.074 seconds taken; 159.110 reps per second.\n",
      "30976/32768 complete; 195.464 seconds taken; 158.474 reps per second.\n",
      "31232/32768 complete; 197.871 seconds taken; 157.840 reps per second.\n",
      "31488/32768 complete; 200.276 seconds taken; 157.223 reps per second.\n",
      "31744/32768 complete; 202.812 seconds taken; 156.519 reps per second.\n",
      "32000/32768 complete; 205.301 seconds taken; 155.868 reps per second.\n",
      "32256/32768 complete; 207.694 seconds taken; 155.306 reps per second.\n",
      "32512/32768 complete; 210.260 seconds taken; 154.628 reps per second.\n",
      "32768/32768 complete; 212.803 seconds taken; 153.983 reps per second.\n",
      "32768 reps completed in 212.80290007591248 seconds;\n",
      "153.98286390040164 reps per second.\n",
      "196.97659182548523 seconds were traversal and collision.\n",
      "3.915703296661377 seconds were potential following.\n",
      "6.584975719451904 seconds were potential addition.\n",
      "5.152017593383789 seconds were adding to collections.\n",
      "Maximum diffusion distance: 0\n",
      "structures/continuum/STF_Si_L32_Th85_D0.45_N32768_1673398759\n"
     ]
    }
   ],
   "source": [
    "''' Slanted columns (no flux distribution) '''\n",
    "# unit length is 1 nm\n",
    "# diffusion remains to be solved\n",
    "# need both a tree (to store roots/fiber IDs) and a spacial container (to improve collision detection)\n",
    "# PBC\n",
    "np.set_printoptions(precision=4)\n",
    "np.seterr(all='raise')\n",
    "reps = 32768\n",
    "L = 32\n",
    "theta = 85\n",
    "phi = 0\n",
    "turns = 0\n",
    "species = 1\n",
    "bin_size = 8\n",
    "radius = 0.15\n",
    "height = 32\n",
    "D = 0.45\n",
    "R = 1\n",
    "A = 1\n",
    "\n",
    "batch_size = 128\n",
    "poses = [np.array([0,0,0]) for i in range(batch_size)]\n",
    "cs = [-1 for i in range(batch_size)]\n",
    "\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "    PROCESS:\n",
    "        1) generate particle with transport angle\n",
    "        2) send towards substrate\n",
    "        3) during travel check that all particle centers in area are R1+R2 away at least\n",
    "        4) once that isn't true a collision occurred: move the particle back until there isn't a collision\n",
    "        5) diffuse the particle with a pseudo-Lennard Jones potential\n",
    "        6) add particle to particle list, tree, and spacial container\n",
    "        7) repeat a large number of times\n",
    "        \n",
    "    VISUALIZATION:\n",
    "        1) It's a point list with float locations; use any point cloud viewer to view it\n",
    "\"\"\"\n",
    "\n",
    "# setup\n",
    "atom_pos = np.zeros((reps, 7), dtype=np.float32) # x,y,z,sp,radius,fiber\n",
    "bins = SlantedCorridors(L, bin_size, theta, height, verbose=False)\n",
    "#print(bins.bins)\n",
    "scale = 10\n",
    "potentials = np.zeros((L*scale, L*scale, height*scale, 4), dtype=np.float32)\n",
    "p = [{'L': L, \n",
    "    'theta': theta, \n",
    "    'phi': phi, \n",
    "    'H': height, \n",
    "    'D': D, \n",
    "    'R': R,\n",
    "    'A': A,\n",
    "    'turns': turns, \n",
    "    'stepper resolution': 0,\n",
    "    'species': species,\n",
    "    'radius': radius,\n",
    "    'weights': None, \n",
    "    'repetition': reps, \n",
    "    'spread': 0}]\n",
    "\n",
    "Vm = scale*(2*radius)\n",
    "zero_pot = Vm / 2**(1/6)\n",
    "\n",
    "if R != 1.0 and A != 1.0:\n",
    "    print(\"Scaled:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format((R/A)**(1/6)*zero_pot, (2*R/A)**(1/6)*Vm))\n",
    "    print(\"Real:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format((R/A)**(1/6)*zero_pot/scale, (2*R/A)**(1/6)*Vm/scale))\n",
    "else:\n",
    "    print(\"Scaled:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format(zero_pot, Vm))\n",
    "    print(\"Real:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format(zero_pot/scale, Vm/scale))\n",
    "diameter = 2 * scale\n",
    "if diameter % 2 == 0:\n",
    "    diameter += 1\n",
    "template = np.zeros((diameter, diameter, diameter, 4), dtype=np.float32)\n",
    "for i in range(diameter):\n",
    "    ii = i - (diameter-1) // 2\n",
    "    for j in range(diameter):\n",
    "        jj = j - (diameter-1) // 2\n",
    "        for k in range(diameter):\n",
    "            kk = k - (diameter-1) // 2\n",
    "            if ii == 0 and jj == 0 and kk == 0:\n",
    "                template[i,j,k,:] = [0,0,0,np.inf]\n",
    "            else:\n",
    "                r = np.sqrt((ii)**2 + (jj)**2 + (kk)**2)\n",
    "                magnitude = A*(zero_pot/r)**12 - R*(zero_pot/r)**6\n",
    "                direction = np.array([ii,jj,kk])/r# * magnitude/np.abs(magnitude)\n",
    "                template[i,j,k, 3] = magnitude\n",
    "                template[i,j,k,:3] = direction\n",
    "                if r > Vm:\n",
    "                    template[i,j,k,:3] *= -1\n",
    "                \n",
    "rng = np.random.default_rng()\n",
    "dests = np.hstack([rng.random(size=(reps, 2), dtype=np.float32)*L, np.zeros((reps,1), dtype=np.float32)+height])\n",
    "#src = np.array([genh/tand(theta-90)*cosd(phi), genh/tand(theta-90)*sind(phi), 0])\n",
    "#distance = distance3D(src, [0,0,0]) # distance from src to destination\n",
    "#unit_vector = -src / distance # unit vector in travel direction\n",
    "#print(src, distance, unit_vector)\n",
    "\n",
    "update = max((int(reps // batch_size / 128), 1)) # frequency of printed updates\n",
    "current_fiber_id = 0\n",
    "\n",
    "# pr = cProfile.Profile()\n",
    "# pr.enable()\n",
    "# each particle must be treated in sequence\n",
    "#print(\"Beginning simulation\")\n",
    "endsurf = 0\n",
    "endpot = 0\n",
    "endretreat = 0\n",
    "endfly = 0\n",
    "endsettle = 0\n",
    "start = time.time()\n",
    "dmax = 0\n",
    "for i in range(reps//batch_size):\n",
    "    #print(\"Particle {}\".format(i+1))\n",
    "    # get destination and source point\n",
    "    for j in range(batch_size):\n",
    "        startfly = time.time()\n",
    "        dest = dests[i*batch_size + j]\n",
    "        #shift = src + [dest[0], dest[1], dest[2]]\n",
    "\n",
    "        #print(\"Traversing\")\n",
    "        startfly = time.time()\n",
    "        if i < reps - 1:\n",
    "            pos, cs[j] = bins.drop_particle(dest, radius, atom_pos)\n",
    "        else:\n",
    "            #pr = cProfile.Profile()\n",
    "            #pr.enable()\n",
    "            pos, cs[j] = bins.drop_particle(dest, radius, atom_pos, profile=False)\n",
    "            #pr.disable()\n",
    "\n",
    "        poses[j] = pos % [L,L,np.inf]\n",
    "\n",
    "        if np.any(np.isnan(position)):\n",
    "            print(position)\n",
    "        endfly += time.time() - startfly\n",
    "        #print(\"{} -> {}\".format(dest, pos))\n",
    "\n",
    "        # time to try diffusing\n",
    "        #print(\"Following potential\")\n",
    "    if D > 0.0:\n",
    "        for j in range(batch_size):\n",
    "            startsurf = time.time()\n",
    "            poses[j] = follow_potential(potentials, poses[j], scale, L, height, D, current_atom=i*batch_size + j, lennardjones=True, deterministic=True, ratio=1.0)\n",
    "            endsurf += time.time() - startsurf\n",
    "            #print(\"Adding potential\")\n",
    "            startpot = time.time()\n",
    "            superimpose_potential(potentials, poses[j], scale, radius, template, diameter)\n",
    "            endpot += time.time() - startpot\n",
    "\n",
    "        #print(\"Settling atom\")\n",
    "    for j in range(batch_size):\n",
    "        startsettle = time.time()\n",
    "        if not np.any(np.isnan(poses[j])):\n",
    "            atom_pos[i*batch_size+j,:3] = poses[j]\n",
    "            #print(\"\\tSettled at {}\".format(position))\n",
    "            atom_pos[i*batch_size+j,3] = species\n",
    "            atom_pos[i*batch_size+j,4] = radius\n",
    "            #if D > 0.0:\n",
    "            #    superimpose_potential(potentials, poses[j], scale, radius, template, diameter)\n",
    "\n",
    "    #         atom_pos[i+reps//2,:3] = dest / np.array([1,1,2])\n",
    "    #         atom_pos[i+reps//2,3] = species\n",
    "    #         atom_pos[i,4] = radius\n",
    "    #         atom_pos[i+reps,:3] = dest / np.array([1,1,4])\n",
    "    #         atom_pos[i+reps,3] = species\n",
    "    #         atom_pos[i,4] = radius\n",
    "\n",
    "            if cs[j] == -1:\n",
    "                atom_pos[i*batch_size+j,5] = current_fiber_id\n",
    "                current_fiber_id += 1\n",
    "            else:\n",
    "                atom_pos[i*batch_size+j,5] = atom_pos[cs[j], 5]\n",
    "            #print(position, radius, i)\n",
    "            bins.add_to_bins(poses[j], radius, i*batch_size+j)\n",
    "        else:\n",
    "            print(position)\n",
    "        endsettle += time.time() - startsettle\n",
    "    \n",
    "        \n",
    "    if i % update == update - 1:\n",
    "        up_time = time.time() - start\n",
    "        if i*batch_size+j+1 > up_time:\n",
    "            print(\"{}/{} complete; {:.3f} seconds taken; {:.3f} reps per second.\".format(i*batch_size+j+1, reps, up_time, (i*batch_size+j+1)/up_time))\n",
    "        else:\n",
    "            print(\"{}/{} complete; {:.3f} seconds taken; {:.3f} seconds per rep.\".format(i*batch_size+j+1, reps, up_time, up_time/(i*batch_size+j+1)))\n",
    "end = time.time() - start\n",
    "if reps > end:\n",
    "    print(\"{} reps completed in {} seconds;\\n{} reps per second.\".format(reps, end, reps/end))\n",
    "else:\n",
    "    print(\"{} reps completed in {} seconds;\\n{} seconds per rep.\".format(reps, end, end/reps))\n",
    "print(\"{} seconds were traversal and collision.\".format(endfly))\n",
    "print(\"{} seconds were potential following.\".format(endsurf))\n",
    "print(\"{} seconds were potential addition.\".format(endpot))\n",
    "print(\"{} seconds were adding to collections.\".format(endsettle))\n",
    "for i in range(atom_pos.shape[0]):\n",
    "    atom_pos[i, 6] = bins.find_bin_idx(atom_pos[i, :3])\n",
    "#warnings.filterwarnings('default', category=RuntimeWarning)\n",
    "print(\"Maximum diffusion distance: {}\".format(dmax))\n",
    "\n",
    "p[0]['time'] = end\n",
    "#pr.disable()\n",
    "filename = saveContinuum(atom_pos, p, \"Si\", zipped=True)\n",
    "print(filename)\n",
    "# pr.disable()\n",
    "# ps = pstats.Stats(pr)\n",
    "# ps.sort_stats('time').print_stats()\n",
    "\n",
    "# snapshot = tracemalloc.take_snapshot()\n",
    "# top_stats = snapshot.statistics('lineno')\n",
    "# for stat in top_stats[:10]:\n",
    "#     print(stat)\n",
    "# pr.print_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(73728, 7)\n",
      "Scaled:\tZero-potential point: 8.908987181403393\n",
      "\tMinimum-potential point: 10.0\n",
      "Real:\tZero-potential point: 0.8908987181403394\n",
      "\tMinimum-potential point: 1.0\n",
      "Scaled:\tZero-potential point: 2.672696154421018\n",
      "\tMinimum-potential point: 3.0\n",
      "Real:\tZero-potential point: 0.2672696154421018\n",
      "\tMinimum-potential point: 0.3\n",
      "64/8192 complete; 1.126 seconds taken; 56.825 reps per second.\n",
      "128/8192 complete; 2.157 seconds taken; 59.328 reps per second.\n",
      "192/8192 complete; 2.988 seconds taken; 64.264 reps per second.\n",
      "256/8192 complete; 3.907 seconds taken; 65.525 reps per second.\n",
      "320/8192 complete; 4.689 seconds taken; 68.244 reps per second.\n",
      "384/8192 complete; 5.733 seconds taken; 66.977 reps per second.\n",
      "448/8192 complete; 6.536 seconds taken; 68.538 reps per second.\n",
      "512/8192 complete; 7.440 seconds taken; 68.820 reps per second.\n",
      "576/8192 complete; 8.279 seconds taken; 69.578 reps per second.\n",
      "640/8192 complete; 9.023 seconds taken; 70.932 reps per second.\n",
      "704/8192 complete; 9.909 seconds taken; 71.047 reps per second.\n",
      "768/8192 complete; 10.726 seconds taken; 71.601 reps per second.\n",
      "832/8192 complete; 11.435 seconds taken; 72.757 reps per second.\n",
      "896/8192 complete; 12.129 seconds taken; 73.870 reps per second.\n",
      "960/8192 complete; 12.953 seconds taken; 74.116 reps per second.\n",
      "1024/8192 complete; 13.789 seconds taken; 74.263 reps per second.\n",
      "1088/8192 complete; 14.483 seconds taken; 75.123 reps per second.\n",
      "1152/8192 complete; 15.195 seconds taken; 75.814 reps per second.\n",
      "1216/8192 complete; 15.854 seconds taken; 76.699 reps per second.\n",
      "1280/8192 complete; 16.635 seconds taken; 76.944 reps per second.\n",
      "1344/8192 complete; 17.557 seconds taken; 76.552 reps per second.\n",
      "1408/8192 complete; 18.295 seconds taken; 76.962 reps per second.\n",
      "1472/8192 complete; 18.950 seconds taken; 77.678 reps per second.\n",
      "1536/8192 complete; 19.580 seconds taken; 78.447 reps per second.\n",
      "1600/8192 complete; 20.289 seconds taken; 78.860 reps per second.\n",
      "1664/8192 complete; 21.021 seconds taken; 79.157 reps per second.\n",
      "1728/8192 complete; 21.706 seconds taken; 79.611 reps per second.\n",
      "1792/8192 complete; 22.381 seconds taken; 80.069 reps per second.\n",
      "1856/8192 complete; 23.048 seconds taken; 80.528 reps per second.\n",
      "1920/8192 complete; 23.754 seconds taken; 80.828 reps per second.\n",
      "1984/8192 complete; 24.408 seconds taken; 81.284 reps per second.\n",
      "2048/8192 complete; 25.105 seconds taken; 81.576 reps per second.\n",
      "2112/8192 complete; 25.836 seconds taken; 81.745 reps per second.\n",
      "2176/8192 complete; 26.554 seconds taken; 81.947 reps per second.\n",
      "2240/8192 complete; 27.282 seconds taken; 82.106 reps per second.\n",
      "2304/8192 complete; 27.962 seconds taken; 82.398 reps per second.\n",
      "2368/8192 complete; 28.584 seconds taken; 82.843 reps per second.\n",
      "2432/8192 complete; 29.307 seconds taken; 82.983 reps per second.\n",
      "2496/8192 complete; 30.014 seconds taken; 83.160 reps per second.\n",
      "2560/8192 complete; 30.645 seconds taken; 83.538 reps per second.\n",
      "2624/8192 complete; 31.392 seconds taken; 83.589 reps per second.\n",
      "2688/8192 complete; 32.007 seconds taken; 83.982 reps per second.\n",
      "2752/8192 complete; 32.669 seconds taken; 84.239 reps per second.\n",
      "2816/8192 complete; 33.342 seconds taken; 84.458 reps per second.\n",
      "2880/8192 complete; 34.084 seconds taken; 84.496 reps per second.\n",
      "2944/8192 complete; 34.800 seconds taken; 84.599 reps per second.\n",
      "3008/8192 complete; 35.465 seconds taken; 84.817 reps per second.\n",
      "3072/8192 complete; 36.075 seconds taken; 85.156 reps per second.\n",
      "3136/8192 complete; 36.813 seconds taken; 85.187 reps per second.\n",
      "3200/8192 complete; 37.451 seconds taken; 85.445 reps per second.\n",
      "3264/8192 complete; 38.104 seconds taken; 85.660 reps per second.\n",
      "3328/8192 complete; 38.829 seconds taken; 85.708 reps per second.\n",
      "3392/8192 complete; 39.513 seconds taken; 85.846 reps per second.\n",
      "3456/8192 complete; 40.271 seconds taken; 85.819 reps per second.\n",
      "3520/8192 complete; 40.989 seconds taken; 85.877 reps per second.\n",
      "3584/8192 complete; 41.687 seconds taken; 85.974 reps per second.\n",
      "3648/8192 complete; 42.411 seconds taken; 86.015 reps per second.\n",
      "3712/8192 complete; 43.170 seconds taken; 85.985 reps per second.\n",
      "3776/8192 complete; 43.929 seconds taken; 85.958 reps per second.\n",
      "3840/8192 complete; 44.658 seconds taken; 85.987 reps per second.\n",
      "3904/8192 complete; 45.414 seconds taken; 85.965 reps per second.\n",
      "3968/8192 complete; 46.202 seconds taken; 85.883 reps per second.\n",
      "4032/8192 complete; 47.000 seconds taken; 85.787 reps per second.\n",
      "4096/8192 complete; 47.703 seconds taken; 85.864 reps per second.\n",
      "4160/8192 complete; 48.425 seconds taken; 85.907 reps per second.\n",
      "4224/8192 complete; 49.237 seconds taken; 85.789 reps per second.\n",
      "4288/8192 complete; 50.047 seconds taken; 85.679 reps per second.\n",
      "4352/8192 complete; 50.893 seconds taken; 85.512 reps per second.\n",
      "4416/8192 complete; 51.572 seconds taken; 85.627 reps per second.\n",
      "4480/8192 complete; 52.312 seconds taken; 85.641 reps per second.\n",
      "4544/8192 complete; 53.038 seconds taken; 85.675 reps per second.\n",
      "4608/8192 complete; 53.770 seconds taken; 85.698 reps per second.\n",
      "4672/8192 complete; 54.577 seconds taken; 85.604 reps per second.\n",
      "4736/8192 complete; 55.281 seconds taken; 85.671 reps per second.\n",
      "4800/8192 complete; 56.039 seconds taken; 85.655 reps per second.\n",
      "4864/8192 complete; 56.756 seconds taken; 85.701 reps per second.\n",
      "4928/8192 complete; 57.477 seconds taken; 85.739 reps per second.\n",
      "4992/8192 complete; 58.209 seconds taken; 85.760 reps per second.\n",
      "5056/8192 complete; 59.101 seconds taken; 85.548 reps per second.\n",
      "5120/8192 complete; 59.905 seconds taken; 85.469 reps per second.\n",
      "5184/8192 complete; 60.727 seconds taken; 85.366 reps per second.\n",
      "5248/8192 complete; 61.528 seconds taken; 85.295 reps per second.\n",
      "5312/8192 complete; 62.302 seconds taken; 85.262 reps per second.\n",
      "5376/8192 complete; 63.069 seconds taken; 85.240 reps per second.\n",
      "5440/8192 complete; 63.814 seconds taken; 85.247 reps per second.\n",
      "5504/8192 complete; 64.638 seconds taken; 85.152 reps per second.\n",
      "5568/8192 complete; 65.434 seconds taken; 85.094 reps per second.\n",
      "5632/8192 complete; 66.290 seconds taken; 84.960 reps per second.\n",
      "5696/8192 complete; 67.056 seconds taken; 84.944 reps per second.\n",
      "5760/8192 complete; 67.840 seconds taken; 84.905 reps per second.\n",
      "5824/8192 complete; 68.608 seconds taken; 84.887 reps per second.\n",
      "5888/8192 complete; 69.356 seconds taken; 84.896 reps per second.\n",
      "5952/8192 complete; 70.105 seconds taken; 84.901 reps per second.\n",
      "6016/8192 complete; 70.931 seconds taken; 84.815 reps per second.\n",
      "6080/8192 complete; 71.783 seconds taken; 84.700 reps per second.\n",
      "6144/8192 complete; 72.557 seconds taken; 84.678 reps per second.\n",
      "6208/8192 complete; 73.353 seconds taken; 84.632 reps per second.\n",
      "6272/8192 complete; 74.104 seconds taken; 84.638 reps per second.\n",
      "6336/8192 complete; 74.928 seconds taken; 84.561 reps per second.\n",
      "6400/8192 complete; 75.777 seconds taken; 84.458 reps per second.\n",
      "6464/8192 complete; 76.552 seconds taken; 84.439 reps per second.\n",
      "6528/8192 complete; 77.310 seconds taken; 84.439 reps per second.\n",
      "6592/8192 complete; 78.185 seconds taken; 84.313 reps per second.\n",
      "6656/8192 complete; 79.029 seconds taken; 84.222 reps per second.\n",
      "6720/8192 complete; 79.845 seconds taken; 84.163 reps per second.\n",
      "6784/8192 complete; 80.705 seconds taken; 84.059 reps per second.\n",
      "6848/8192 complete; 81.537 seconds taken; 83.986 reps per second.\n",
      "6912/8192 complete; 82.429 seconds taken; 83.854 reps per second.\n",
      "6976/8192 complete; 83.241 seconds taken; 83.805 reps per second.\n",
      "7040/8192 complete; 84.072 seconds taken; 83.738 reps per second.\n",
      "7104/8192 complete; 84.837 seconds taken; 83.737 reps per second.\n",
      "7168/8192 complete; 85.662 seconds taken; 83.677 reps per second.\n",
      "7232/8192 complete; 86.441 seconds taken; 83.664 reps per second.\n",
      "7296/8192 complete; 87.245 seconds taken; 83.627 reps per second.\n",
      "7360/8192 complete; 88.081 seconds taken; 83.560 reps per second.\n",
      "7424/8192 complete; 88.935 seconds taken; 83.477 reps per second.\n",
      "7488/8192 complete; 89.773 seconds taken; 83.410 reps per second.\n",
      "7552/8192 complete; 90.607 seconds taken; 83.349 reps per second.\n",
      "7616/8192 complete; 91.354 seconds taken; 83.368 reps per second.\n",
      "7680/8192 complete; 92.103 seconds taken; 83.385 reps per second.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7744/8192 complete; 92.936 seconds taken; 83.326 reps per second.\n",
      "7808/8192 complete; 93.932 seconds taken; 83.124 reps per second.\n",
      "7872/8192 complete; 94.698 seconds taken; 83.127 reps per second.\n",
      "7936/8192 complete; 95.645 seconds taken; 82.974 reps per second.\n",
      "8000/8192 complete; 96.468 seconds taken; 82.929 reps per second.\n",
      "8064/8192 complete; 97.271 seconds taken; 82.902 reps per second.\n",
      "8128/8192 complete; 98.054 seconds taken; 82.893 reps per second.\n",
      "8192/8192 complete; 98.782 seconds taken; 82.930 reps per second.\n",
      "8192 reps completed in 98.78324961662292 seconds;\n",
      "82.92903940488992 reps per second.\n",
      "92.09034204483032 seconds were traversal before collision.\n",
      "0 seconds were backing up from collision.\n",
      "3.3561787605285645 seconds were potential following.\n",
      "1.8731229305267334 seconds were potential addition.\n",
      "Maximum diffusion distance: 0.2529737420605582\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'structures/continuum/STF_Si_Au_L32_Th85_D3_N73728_1673375548'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#### tracemalloc.start()\n",
    "\n",
    "# pr = cProfile.Profile()\n",
    "# pr.enable()\n",
    "#filename = 'structures\\\\continuum\\\\STF_Si_L32_Th85_D0.45_N65536_1673328957'\n",
    "atom_pos, deposited, p = loadContinuum(filename + '.npz', folder = '')\n",
    "reps2 = reps//8\n",
    "atom_pos = np.vstack([atom_pos[:reps,:], np.zeros((reps2, atom_pos.shape[1]), dtype=np.float32)])\n",
    "print(atom_pos.shape)\n",
    "species = 2\n",
    "radius = 0.15\n",
    "height = 32\n",
    "\n",
    "D = 3\n",
    "p.append(p[0])\n",
    "p[1]['repetition'] = reps2\n",
    "p[1]['D'] = D\n",
    "p[1]['species'] = species\n",
    "p[1]['radius'] = radius\n",
    "\n",
    "''' Do potentials twice, once for each species '''\n",
    "potentials = np.zeros((L*scale, L*scale, height*scale, 4), dtype=np.float32)\n",
    "\n",
    "for Vm in [scale*(2*.5), scale*(2*radius)]:\n",
    "    R = 1.0\n",
    "    A = 1.0\n",
    "    zero_pot = Vm / 2**(1/6)\n",
    "    if R != 1.0 or A != 1.0:\n",
    "        print(\"Scaled:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format((R/A)**(1/6)*zero_pot, (2*R/A)**(1/6)*Vm))\n",
    "        print(\"Real:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format((R/A)**(1/6)*zero_pot/scale, (2*R/A)**(1/6)*Vm/scale))\n",
    "    else:\n",
    "        print(\"Scaled:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format(zero_pot, Vm))\n",
    "        print(\"Real:\\tZero-potential point: {}\\n\\tMinimum-potential point: {}\".format(zero_pot/scale, Vm/scale))\n",
    "    diameter = 2 * scale\n",
    "    if diameter % 2 == 0:\n",
    "        diameter += 1\n",
    "    template = np.zeros((diameter, diameter, diameter, 4), dtype=np.float32)\n",
    "    for i in range(diameter):\n",
    "        ii = i - (diameter-1) // 2\n",
    "        for j in range(diameter):\n",
    "            jj = j - (diameter-1) // 2\n",
    "            for k in range(diameter):\n",
    "                kk = k - (diameter-1) // 2\n",
    "                if ii == 0 and jj == 0 and kk == 0:\n",
    "                    template[i,j,k,:] = [0,0,0,np.inf]\n",
    "                else:\n",
    "                    r = np.sqrt((ii)**2 + (jj)**2 + (kk)**2)\n",
    "                    magnitude = A*(zero_pot/r)**12 - R*(zero_pot/r)**6\n",
    "                    direction = np.array([ii,jj,kk])/r# * magnitude/np.abs(magnitude)\n",
    "                    template[i,j,k, 3] = magnitude\n",
    "                    template[i,j,k,:3] = direction\n",
    "                    if r > Vm:\n",
    "                        template[i,j,k,:3] *= -1\n",
    "                        \n",
    "rng = np.random.default_rng()\n",
    "dests = np.hstack([rng.random(size=(reps2, 2), dtype=np.float32)*L, np.zeros((reps2,1), dtype=np.float32)+height])\n",
    "\n",
    "update = max((int(reps2 / 128), 1)) # frequency of printed updates\n",
    "current_fiber_id = np.max(atom_pos[:,5] + 1)\n",
    "\n",
    "# each particle must be treated in sequence\n",
    "#print(\"Beginning simulation\")\n",
    "endsurf = 0\n",
    "endpot = 0\n",
    "endretreat = 0\n",
    "endfly = 0\n",
    "start = time.time()\n",
    "dmax = 0\n",
    "for i in range(reps, reps + reps2):\n",
    "    #print(\"Particle {}\".format(i+1))\n",
    "    # get destination and source point\n",
    "    startfly = time.time()\n",
    "    dest = dests[i-reps]\n",
    "    #shift = src + [dest[0], dest[1], dest[2]]\n",
    "    \n",
    "    #print(\"Traversing\")\n",
    "    \n",
    "    pos, c = bins.drop_particle(dest, radius, atom_pos)\n",
    "    \n",
    "    position = pos % [L,L,np.inf]\n",
    "    if np.any(np.isnan(position)):\n",
    "        print(position)\n",
    "    #print(\"{} -> {}\".format(dest, pos))\n",
    "    endfly += time.time() - startfly\n",
    "    \n",
    "    # time to try diffusing\n",
    "    #print(\"Following potential\")\n",
    "    if D > 0.0:\n",
    "        startsurf = time.time()\n",
    "        position_new = follow_potential(potentials, position, scale, L, height, D/2, current_atom=i, lennardjones=True, deterministic=True, ratio=0.5)\n",
    "        d = distance3DPBC(position, position_new, L)\n",
    "        if d > dmax:\n",
    "            dmax = d\n",
    "        if d > D:\n",
    "            print(\"Former position: {}\".format(position))\n",
    "            print(\"New position: {}\".format(position_new))\n",
    "            print(\"Distance: {}\".format(d))\n",
    "        position = position_new\n",
    "        endsurf += time.time() - startsurf\n",
    "        #print(\"Adding potential\")\n",
    "        startpot = time.time()\n",
    "        superimpose_potential(potentials, position, scale, radius, template, diameter)\n",
    "        endpot += time.time() - startpot\n",
    "    \n",
    "    #print(\"Settling atom\")\n",
    "    if not np.any(np.isnan(position)):\n",
    "        atom_pos[i,:3] = position\n",
    "        #print(\"\\tSettled at {}\".format(position))\n",
    "        atom_pos[i,3] = species\n",
    "        atom_pos[i,4] = radius\n",
    "        \n",
    "#         atom_pos[i+reps//2,:3] = dest / np.array([1,1,2])\n",
    "#         atom_pos[i+reps//2,3] = species\n",
    "#         atom_pos[i,4] = radius\n",
    "#         atom_pos[i+reps,:3] = dest / np.array([1,1,4])\n",
    "#         atom_pos[i+reps,3] = species\n",
    "#         atom_pos[i,4] = radius\n",
    "        \n",
    "        if c == -1:\n",
    "            atom_pos[i,5] = current_fiber_id\n",
    "            current_fiber_id += 1\n",
    "        else:\n",
    "            atom_pos[i,5] = atom_pos[c, 5]\n",
    "        #print(position, radius, i)\n",
    "        bins.add_to_bins(position, radius, i)\n",
    "    else:\n",
    "        print(position)\n",
    "    \n",
    "        \n",
    "    if i % update == update - 1:\n",
    "        up_time = time.time() - start\n",
    "        if i+1-reps > up_time:\n",
    "            print(\"{}/{} complete; {:.3f} seconds taken; {:.3f} reps per second.\".format(i+1-reps, reps2, up_time, (i+1-reps)/up_time))\n",
    "        else:\n",
    "            print(\"{}/{} complete; {:.3f} seconds taken; {:.3f} seconds per rep.\".format(i+1-reps, reps, up_time, up_time/(i+1-reps)))\n",
    "end = time.time() - start\n",
    "if reps > end:\n",
    "    print(\"{} reps completed in {} seconds;\\n{} reps per second.\".format(reps2, end, reps2/end))\n",
    "else:\n",
    "    print(\"{} reps completed in {} seconds;\\n{} seconds per rep.\".format(reps2, end, end/reps2))\n",
    "print(\"{} seconds were traversal before collision.\".format(endfly))\n",
    "print(\"{} seconds were backing up from collision.\".format(endretreat))\n",
    "print(\"{} seconds were potential following.\".format(endsurf))\n",
    "print(\"{} seconds were potential addition.\".format(endpot))\n",
    "for i in range(atom_pos.shape[0]):\n",
    "    atom_pos[i, 6] = bins.find_bin_idx(atom_pos[i, :3])\n",
    "#warnings.filterwarnings('default', category=RuntimeWarning)\n",
    "print(\"Maximum diffusion distance: {}\".format(dmax))\n",
    "\n",
    "p[1]['time'] = end\n",
    "saveContinuum(atom_pos, p, \"Si_Au\", zipped=True)\n",
    "\n",
    "# snapshot = tracemalloc.take_snapshot()\n",
    "# top_stats = snapshot.statistics('lineno')\n",
    "# for stat in top_stats[:10]:\n",
    "#     print(stat)\n",
    "\n",
    "# pr.print_stats()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}