diff --git a/README.md b/README.md
index f27ca50..04c7ac4 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,9 @@
 # PointCNN
 PyTorch implementation of PointCNN model specified in the white paper located here: https://arxiv.org/pdf/1801.07791.pdf
+
+Current MNIST accuracy: ~96%
+
+My coding style is somewhat unique, but ultimately geared towards maximal
+readability. Along with extensive documentation in the code, I use type , and
+code comments indicating input/outputs shapes.  (x,y,z) just indicate that any
+value is accepted at runtime.
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..31221ea
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1 @@
+from PointCNN.core import RandPointCNN, PointCNN, XConv, knn_indices_func_cpu, knn_indices_func_gpu
diff --git a/core/__init__.py b/core/__init__.py
new file mode 100644
index 0000000..fdef393
--- /dev/null
+++ b/core/__init__.py
@@ -0,0 +1,2 @@
+from PointCNN.core.model import XConv, PointCNN, RandPointCNN
+from PointCNN.core.util_funcs import knn_indices_func_cpu, knn_indices_func_gpu, UFloatTensor, ULongTensor
diff --git a/pointcnn/context.py b/core/context.py
similarity index 100%
rename from pointcnn/context.py
rename to core/context.py
diff --git a/pointcnn/core.py b/core/model.py
similarity index 53%
rename from pointcnn/core.py
rename to core/model.py
index 0bcd63d..ebe3e03 100644
--- a/pointcnn/core.py
+++ b/core/model.py
@@ -1,40 +1,24 @@
 """
-PyTorch implementation of the PointCNN paper, as specified in:
-  https://arxiv.org/pdf/1801.07791.pdf
-
 Author: Austin J. Garrett
 
-I make liberal use of the mypy static type checker for Python.
-It should be mostly intuitive, but further documentation can be found at:
-  http://mypy-lang.org/
+PyTorch implementation of the PointCNN paper, as specified in:
+  https://arxiv.org/pdf/1801.07791.pdf
+Original paper by: Yangyan Li, Rui Bu, Mingchao Sun, Baoquan Chen
 """
 
-# Standard Modules
-import time
-
 # External Modules
 import torch
 import torch.nn as nn
-from torch import Tensor, LongTensor
-from torch.autograd import Variable
+from torch import FloatTensor
 import numpy as np
-import matplotlib.pyplot as plt
-from typing import Tuple, Callable
+from typing import Tuple, Callable, Optional
 
 # Internal Modules
-try:
-    from .util import knn_indices_func, knn_indices_func_gpu, plot
-    from .layers import MLP, LayerNorm, Conv, SepConv, Dense, end_channels
-    # from .context import timed
-except SystemError:
-    from util import knn_indices_func, knn_indices_func_gpu, plot
-    from layers import MLP, LayerNorm, Conv, SepConv, Dense, end_channels
-    # from context import timed
+from PointCNN.core.util_funcs import UFloatTensor, ULongTensor
+from PointCNN.core.util_layers import Conv, SepConv, Dense, EndChannels
 
 class XConv(nn.Module):
-    """
-    Vectorized pointwise convolution.
-    """
+    """ Convolution over a single point and its neighbors.  """
 
     def __init__(self, C_in : int, C_out : int, dims : int, K : int,
                  P : int, C_mid : int, depth_multiplier : int) -> None:
@@ -43,7 +27,7 @@ class XConv(nn.Module):
         :param C_out: Output dimension of the representative point features.
         :param dims: Spatial dimensionality of points.
         :param K: Number of neighbors to convolve over.
-        :param P: Number of representative points
+        :param P: Number of representative points.
         :param C_mid: Dimensionality of lifted point features.
         :param depth_multiplier: Depth multiplier for internal depthwise separable convolution.
         """
@@ -67,7 +51,7 @@ class XConv(nn.Module):
 
         # Layers to generate X
         self.x_trans = nn.Sequential(
-            end_channels(Conv(
+            EndChannels(Conv(
                 in_channels = dims,
                 out_channels = K*K,
                 kernel_size = (1, K),
@@ -77,30 +61,29 @@ class XConv(nn.Module):
             Dense(K*K, K*K, with_bn = False, activation = None)
         )
 
-        self.end_conv = end_channels(SepConv(
+        self.end_conv = EndChannels(SepConv(
             in_channels = C_mid + C_in,
             out_channels = C_out,
             kernel_size = (1, K),
             depth_multiplier = depth_multiplier
         )).cuda()
 
-    def forward(self, x : Tuple[Tensor, Tensor, Tensor]) -> Tensor:
+    def forward(self, x : Tuple[UFloatTensor,            # (N, P, dims)
+                                UFloatTensor,            # (N, P, K, dims)
+                                Optional[UFloatTensor]]  # (N, P, K, C_in)
+               ) -> UFloatTensor:                        # (N, K, C_out)
         """
         Applies XConv to the input data.
-        :type rep_pt: (N, P, dims)
-        :type pts: (N, P, K, dims)
-        :type fts: (N, P, K, C_in)
-        :rtype: (TODO: shape)
-        :param x: (rep_pt, pts, fts)
-        :param rep_pt: Representative point
-        :param pts: Regional point cloud such that fts[:,p_idx,:] is the feature associated with pts[:,p_idx,:]
-        :param fts: Regional features such that pts[:,p_idx,:] is the feature associated with fts[:,p_idx,:]
+        :param x: (rep_pt, pts, fts) where
+          - rep_pt: Representative point.
+          - pts: Regional point cloud such that fts[:,p_idx,:] is the feature
+          associated with pts[:,p_idx,:].
+          - fts: Regional features such that pts[:,p_idx,:] is the feature
+          associated with fts[:,p_idx,:].
         :return: Features aggregated into point rep_pt.
         """
         rep_pt, pts, fts = x
-        N = len(pts)
 
-        #== RUNTIME ASSERTIONS ==#
         if fts is not None:
             assert(rep_pt.size()[0] == pts.size()[0] == fts.size()[0])  # Check N is equal.
             assert(rep_pt.size()[1] == pts.size()[1] == fts.size()[1])  # Check P is equal.
@@ -111,15 +94,15 @@ class XConv(nn.Module):
             assert(rep_pt.size()[1] == pts.size()[1])                   # Check P is equal.
             assert(pts.size()[2] == self.K)                             # Check K is equal.
         assert(rep_pt.size()[2] == pts.size()[3] == self.dims)          # Check dims is equal.
-        #========================#
 
+        N = len(pts)
         P = rep_pt.size()[1]  # (N, P, K, dims)
         p_center = torch.unsqueeze(rep_pt, dim = 2)  # (N, P, 1, dims)
 
         # Move pts to local coordinate system of rep_pt.
         pts_local = pts - p_center  # (N, P, K, dims)
         # pts_local = self.pts_layernorm(pts - p_center)
-    
+
         # Individually lift each point into C_mid space.
         fts_lifted0 = self.dense1(pts_local)
         fts_lifted  = self.dense2(fts_lifted0)  # (N, P, K, C_mid)
@@ -134,125 +117,120 @@ class XConv(nn.Module):
         X = self.x_trans(pts_local)
         X = X.view(*X_shape)
 
-        """
-        X = self.mid_conv(pts_local)
-        X = X.contiguous().view(*X_shape)
-        X = self.mid_dwconv1(X)
-        X = X.contiguous().view(*X_shape)
-        X = self.mid_dwconv2(X)
-        X = X.contiguous().view(*X_shape)
-        """
-
         # Weight and permute fts_cat with the learned X.
         fts_X = torch.matmul(X, fts_cat)
         fts_p = self.end_conv(fts_X).squeeze(dim = 2)
         return fts_p
 
 class PointCNN(nn.Module):
-    """
-    TODO: Insert documentation
-    """
+    """ Pointwise convolutional model. """
 
     def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int,
-                 r_indices_func : Callable[[Tensor, Tensor, int, int], LongTensor]) -> None:
+                 r_indices_func : Callable[[UFloatTensor,  # (N, P, dims)
+                                            UFloatTensor,  # (N, x, dims)
+                                            int, int],
+                                           ULongTensor]    # (N, P, K)
+                ) -> None:
         """
         :param C_in: Input dimension of the points' features.
         :param C_out: Output dimension of the representative point features.
         :param dims: Spatial dimensionality of points.
         :param K: Number of neighbors to convolve over.
-        :param P: Number of representative points.
         :param D: "Spread" of neighboring points.
+        :param P: Number of representative points.
         :param r_indices_func: Selector function of the type,
-            INP
-            ======
-            rep_pts : (N, P, dims) Representative points
-            pts  : (N, *, dims) Point cloud
-            K : Number of points for each region.
-            D : "Spread" of neighboring points (analogous to stride).
+          INPUTS
+          rep_pts : Representative points.
+          pts  : Point cloud.
+          K : Number of points for each region.
+          D : "Spread" of neighboring points.
 
-            OUT
-            ======
-            pts_idx : (N, P, K) Array of indices into pts such that
-            pts[pts_idx] is the set of points in the "region" around rep_pt.
-
-        a representative point rep_pt and a point cloud pts. From these it returns an
-        array of K
-        :param C_mid: Dimensionality of lifted point features.
-        :param mlp_width: Number of hidden layers in MLPs.
+          OUTPUT
+          pts_idx : Array of indices into pts such that pts[pts_idx] is the set
+          of points in the "region" around rep_pt.
         """
         super(PointCNN, self).__init__()
 
         C_mid = C_out // 2 if C_in == 0 else C_out // 4
         depth_multiplier = min(int(np.ceil(C_out / C_in)), 4)
 
-        self.r_indices_func = r_indices_func
+        self.r_indices_func = lambda rep_pts, pts: r_indices_func(rep_pts, pts, K, D)
         self.dense = Dense(C_in, C_out // 2) if C_in != 0 else None
         self.x_conv = XConv(C_out // 2 if C_in != 0 else C_in, C_out, dims, K, P, C_mid, depth_multiplier)
         self.D = D
 
-    def select_region(self, pts : Tensor, pts_idx : LongTensor) -> Tensor:
+    def select_region(self, pts : UFloatTensor,  # (N, x, dims)
+                      pts_idx : ULongTensor      # (N, P, K)
+                     ) -> UFloatTensor:          # (P, K, dims)
         """
-        Selects
-        :type pts: (N, *, dims)
-        :type pts_idx: (N, P, K)
-        :rtype pts_region: (P, K, dims)
-        :param pts: Point cloud to select regional points from
-        :param pts_idx: Indices of points in region to be selected
-        :return:
+        Selects neighborhood points based on output of r_indices_func.
+        :param pts: Point cloud to select regional points from.
+        :param pts_idx: Indices of points in region to be selected.
+        :return: Local neighborhoods around each representative point.
         """
         regions = torch.stack([
             pts[n][idx,:] for n, idx in enumerate(torch.unbind(pts_idx, dim = 0))
         ], dim = 0)
         return regions
 
-    def forward(self, x : Tuple[Tensor, Tensor, Tensor]) -> Tensor:
+    def forward(self, x : Tuple[FloatTensor,  # (N, P, dims)
+                                FloatTensor,  # (N, x, dims)
+                                FloatTensor]  # (N, x, C_in)
+               ) -> FloatTensor:              # (N, P, C_out)
         """
         Given a set of representative points, a point cloud, and its
         corresponding features, return a new set of representative points with
         features projected from the point cloud.
-        :type rep_pts: (N, *, dims)
-        :type pts: (N, K, dims)
-        :type fts: (N, K, C_in)
-        :rtype: (N, P, dims)
-        :param rep_pts: Representative points
-        :param pts: Regional point cloud such that fts[:,p_idx,:] is the feature associated with pts[:,p_idx,:]
-        :param fts: Regional features such that pts[:,p_idx,:] is the feature associated with fts[:,p_idx,:]
-        :return:
+        :param x: (rep_pts, pts, fts) where
+          - rep_pts: Representative points.
+          - pts: Regional point cloud such that fts[:,p_idx,:] is the
+          feature associated with pts[:,p_idx,:].
+          - fts: Regional features such that pts[:,p_idx,:] is the feature
+          associated with fts[:,p_idx,:].
+        :return: Features aggregated to rep_pts.
         """
         rep_pts, pts, fts = x
         fts = self.dense(fts) if fts is not None else fts
 
         # This step takes ~97% of the time. Prime target for optimization: KNN on GPU.
-        pts_idx = self.r_indices_func(rep_pts.cpu(), pts.cpu(), self.x_conv.K, self.D).cuda()
+        pts_idx = self.r_indices_func(rep_pts.cpu(), pts.cpu()).cuda()
         # -------------------------------------------------------------------------- #
 
         pts_regional = self.select_region(pts, pts_idx)
         fts_regional = self.select_region(fts, pts_idx) if fts is not None else fts
         fts_p = self.x_conv((rep_pts, pts_regional, fts_regional))
 
-        if False:
-            # Draw neighborhood points, for debugging.
-            t = 10
-            n = 0
-            test_point = rep_pts[n,t,:].cpu().data.numpy()
-            neighborhood = pts_regional[n,t,:,:].cpu().data.numpy()
-            plt.scatter(pts[n][:,0], pts[n][:,1])
-            plt.scatter(test_point[0], test_point[1], s = 100, c = 'green')
-            plt.scatter(neighborhood[:,0], neighborhood[:,1], s = 100, c = 'red')
-            plt.show()
         return fts_p
 
 class RandPointCNN(nn.Module):
     """ PointCNN with randomly subsampled representative points. """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int,
+                 r_indices_func : Callable[[UFloatTensor,  # (N, P, dims)
+                                            UFloatTensor,  # (N, x, dims)
+                                            int, int],
+                                           ULongTensor]    # (N, P, K)
+                ) -> None:
+        """ See documentation for PointCNN. """
         super(RandPointCNN, self).__init__()
-        self.pointcnn = PointCNN(*args, **kwargs)
+        self.pointcnn = PointCNN(C_in, C_out, dims, K, D, P, r_indices_func)
+        self.P = P
 
-        # This is safe because PointCNN requires P.
-        self.P = args[5] if len(args) > 5 else kwargs['P']
-
-    def forward(self, x :Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
+    def forward(self, x : Tuple[UFloatTensor,  # (N, x, dims)
+                                UFloatTensor]  # (N, x, dims)
+               ) -> Tuple[UFloatTensor,        # (N, P, dims)
+                          UFloatTensor]:       # (N, P, C_out)
+        """
+        Given a point cloud, and its corresponding features, return a new set
+        of randomly-sampled representative points with features projected from
+        the point cloud.
+        :param x: (pts, fts) where
+         - pts: Regional point cloud such that fts[:,p_idx,:] is the
+        feature associated with pts[:,p_idx,:].
+         - fts: Regional features such that pts[:,p_idx,:] is the feature
+        associated with fts[:,p_idx,:].
+        :return: Randomly subsampled points and their features.
+        """
         pts, fts = x
         if 0 < self.P < pts.size()[1]:
             # Select random set of indices of subsampled points.
@@ -263,39 +241,3 @@ class RandPointCNN(nn.Module):
             rep_pts = pts
         rep_pts_fts = self.pointcnn((rep_pts, pts, fts))
         return rep_pts, rep_pts_fts
-
-if __name__ == "__main__":
-    np.random.seed(0)
-
-    N = 1
-    num_points = 1000
-    dims = 2
-    C_in = 4
-    K = 10
-    D = 1
-
-    layer1 = RandPointCNN(C_in,   8, dims, K, D, 1000, knn_indices_func).cuda()
-    layer2 = RandPointCNN(   8,  16, dims, K, D,  500, knn_indices_func).cuda()
-    layer3 = RandPointCNN(  16,  32, dims, K, D,  250, knn_indices_func).cuda()
-    layer4 = RandPointCNN(  32,  64, dims, K, D,  125, knn_indices_func).cuda()
-    layer5 = RandPointCNN(  64, 128, dims, K, D,   50, knn_indices_func).cuda()
-
-    pts  = np.random.rand(N,num_points,dims).astype(np.float32)
-    fts  = np.random.rand(N,num_points,C_in).astype(np.float32)
-    pts = Variable(torch.from_numpy(pts)).cuda()
-    fts = Variable(torch.from_numpy(fts)).cuda()
-
-    if True:
-        pts, fts = layer1((pts, fts))
-    else:
-        plot(pts, fts)
-        pts, fts = layer1((pts, fts))
-        plot(pts, fts)
-        pts, fts = layer2((pts, fts))
-        plot(pts, fts)
-        pts, fts = layer3((pts, fts))
-        plot(pts, fts)
-        pts, fts = layer4((pts, fts))
-        plot(pts, fts)
-        pts, fts = layer5((pts, fts))
-        plot(pts, fts)
diff --git a/core/util_funcs.py b/core/util_funcs.py
new file mode 100644
index 0000000..2183670
--- /dev/null
+++ b/core/util_funcs.py
@@ -0,0 +1,66 @@
+# External Modules
+import torch
+from torch import cuda, FloatTensor, LongTensor
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.neighbors import NearestNeighbors
+from typing import Union
+
+# Types to allow for both CPU and GPU models.
+UFloatTensor = Union[FloatTensor, cuda.FloatTensor]
+ULongTensor = Union[LongTensor, cuda.LongTensor]
+
+def knn_indices_func_cpu(rep_pts : FloatTensor,  # (N, pts, dim)
+                         pts : FloatTensor,      # (N, x, dim)
+                         K : int, D : int
+                        ) -> LongTensor:         # (N, pts, K)
+    """
+    CPU-based Indexing function based on K-Nearest Neighbors search.
+    :param rep_pts: Representative points.
+    :param pts: Point cloud to get indices from.
+    :param K: Number of nearest neighbors to collect.
+    :param D: "Spread" of neighboring points.
+    :return: Array of indices, P_idx, into pts such that pts[n][P_idx[n],:]
+    is the set k-nearest neighbors for the representative points in pts[n].
+    """
+    rep_pts = rep_pts.data.numpy()
+    pts = pts.data.numpy()
+    region_idx = []
+
+    for n, p in enumerate(rep_pts):
+        P_particular = pts[n]
+        nbrs = NearestNeighbors(D*K + 1, algorithm = "ball_tree").fit(P_particular)
+        indices = nbrs.kneighbors(p)[1]
+        region_idx.append(indices[:,1::D])
+
+    region_idx = torch.from_numpy(np.stack(region_idx, axis = 0))
+    return region_idx
+
+def knn_indices_func_gpu(rep_pts : cuda.FloatTensor,  # (N, pts, dim)
+                         pts : cuda.FloatTensor,      # (N, x, dim)
+                         k : int, d : int
+                        ) -> cuda.LongTensor:         # (N, pts, K)
+    """
+    GPU-based Indexing function based on K-Nearest Neighbors search.
+    Very memory intensive, and thus unoptimal for large numbers of points.
+    :param rep_pts: Representative points.
+    :param pts: Point cloud to get indices from.
+    :param K: Number of nearest neighbors to collect.
+    :param D: "Spread" of neighboring points.
+    :return: Array of indices, P_idx, into pts such that pts[n][P_idx[n],:]
+    is the set k-nearest neighbors for the representative points in pts[n].
+    """
+    region_idx = []
+
+    for n, qry in enumerate(rep_pts):
+        ref = pts[n]
+        n, d = ref.size()
+        m, d = qry.size()
+        mref = ref.expand(m, n, d)
+        mqry = qry.expand(n, m, d).transpose(0, 1)
+        dist2 = torch.sum((mqry - mref)**2, 2).squeeze()
+        _, inds = torch.topk(dist2, k*d + 1, dim = 1, largest = False)
+        region_idx.append(inds[:,1::d])
+
+    region_idx = torch.stack(region_idx, dim = 0)
+    return region_idx
diff --git a/core/util_layers.py b/core/util_layers.py
new file mode 100644
index 0000000..290fddb
--- /dev/null
+++ b/core/util_layers.py
@@ -0,0 +1,148 @@
+import torch.nn as nn
+from typing import Callable, Union, Tuple
+
+from PointCNN.core.util_funcs import UFloatTensor
+
+def EndChannels(f, make_contiguous = False):
+    """ Class decorator to apply 2D convolution along end channels. """
+
+    class WrappedLayer(nn.Module):
+
+        def __init__(self):
+            super(WrappedLayer, self).__init__()
+            self.forward = lambda x: f(x.permute(0,3,1,2)).permute(0,2,3,1)
+
+    return WrappedLayer()
+
+class Dense(nn.Module):
+    """
+    Single layer perceptron with optional activation, batch normalization, and dropout.
+    """
+
+    def __init__(self, in_features : int, out_features : int,
+                 drop_rate : int = 0, with_bn : bool = True,
+                 activation : Callable[[UFloatTensor], UFloatTensor] = nn.ReLU()
+                ) -> None:
+        """
+        :param in_features: Length of input featuers (last dimension).
+        :param out_features: Length of output features (last dimension).
+        :param drop_rate: Drop rate to be applied after activation.
+        :param with_bn: Whether or not to apply batch normalization.
+        :param activation: Activation function.
+        """
+        super(Dense, self).__init__()
+
+        self.linear = nn.Linear(in_features, out_features)
+        self.activation = activation
+        # self.bn = LayerNorm(out_channels) if with_bn else None
+        self.drop = nn.Dropout(drop_rate) if drop_rate > 0 else None
+
+    def forward(self, x : UFloatTensor) -> UFloatTensor:
+        """
+        :param x: Any input tensor that can be input into nn.Linear.
+        :return: Tensor with linear layer and optional activation, batchnorm,
+        and dropout applied.
+        """
+        x = self.linear(x)
+        if self.activation:
+            x = self.activation(x)
+        # if self.bn:
+        #     x = self.bn(x)
+        if self.drop:
+            x = self.drop(x)
+        return x
+
+class Conv(nn.Module):
+    """
+    2D convolutional layer with optional activation and batch normalization.
+    """
+
+    def __init__(self, in_channels : int, out_channels : int,
+                 kernel_size : Union[int, Tuple[int, int]], with_bn : bool = True,
+                 activation : Callable[[UFloatTensor], UFloatTensor] = nn.relu()
+                ) -> None:
+        """
+        :param in_channels: Length of input featuers (first dimension).
+        :param out_channels: Length of output features (first dimension).
+        :param kernel_size: Size of convolutional kernel.
+        :param with_bn: Whether or not to apply batch normalization.
+        :param activation: Activation function.
+        """
+        super(Conv, self).__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias = not with_bn)
+        self.activation = activation
+        self.bn = nn.BatchNorm2d(out_channels, momentum = 0.9) if with_bn else None
+
+    def forward(self, x : UFloatTensor) -> UFloatTensor:
+        """
+        :param x: Any input tensor that can be input into nn.Conv2d.
+        :return: Tensor with convolutional layer and optional activation and batchnorm applied.
+        """
+        x = self.conv(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.bn:
+            x = self.bn(x)
+        return x
+
+class SepConv(nn.Module):
+    """ Depthwise separable convolution with optional activation and batch normalization"""
+
+    def __init__(self, in_channels : int, out_channels : int,
+                 kernel_size : Union[int, Tuple[int, int]],
+                 depth_multiplier : int = 1, with_bn : bool = True,
+                 activation : Callable[[UFloatTensor], UFloatTensor] = nn.ReLU()
+                 ) -> None:
+        """
+        :param in_channels: Length of input featuers (first dimension).
+        :param out_channels: Length of output features (first dimension).
+        :param kernel_size: Size of convolutional kernel.
+        :depth_multiplier: Depth multiplier for middle part of separable convolution.
+        :param with_bn: Whether or not to apply batch normalization.
+        :param activation: Activation function.
+        """
+        super(SepConv, self).__init__()
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels * depth_multiplier, kernel_size, groups = in_channels),
+            nn.Conv2d(in_channels * depth_multiplier, out_channels, 1, bias = not with_bn)
+        )
+
+        self.activation = activation
+        self.bn = nn.BatchNorm2d(out_channels, momentum = 0.9) if with_bn else None
+
+    def forward(self, x : UFloatTensor) -> UFloatTensor:
+        """
+        :param x: Any input tensor that can be input into nn.Conv2d.
+        :return: Tensor with depthwise separable convolutional layer and
+        optional activation and batchnorm applied.
+        """
+        x = self.conv(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.bn:
+            x = self.bn(x)
+        return x
+
+class LayerNorm(nn.Module):
+    """
+    Batch Normalization over ONLY the mini-batch layer (suitable for nn.Linear layers).
+    """
+
+    def __init__(self, N : int, dim : int, *args, **kwargs) -> None:
+        """
+        :param N: Batch size.
+        :param D: Dimensions.
+        """
+        super(LayerNorm, self).__init__()
+        if dim == 1:
+            self.bn = nn.BatchNorm1d(N, *args, **kwargs)
+        elif dim == 2:
+            self.bn = nn.BatchNorm2d(N, *args, **kwargs)
+        elif dim == 3:
+            self.bn = nn.BatchNorm3d(N, *args, **kwargs)
+        else:
+            raise ValueError("Dimensionality %i not supported" % dim)
+
+        self.forward = lambda x: self.bn(x.unsqueeze(0)).squeeze(0)
diff --git a/mnist/context.py b/mnist/context.py
deleted file mode 100644
index 126a413..0000000
--- a/mnist/context.py
+++ /dev/null
@@ -1,4 +0,0 @@
-import sys, os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
-from lib import pointcnn
diff --git a/mnist/model.py b/mnist/model.py
index 909f4ed..0c891f1 100644
--- a/mnist/model.py
+++ b/mnist/model.py
@@ -1,3 +1,8 @@
+"""
+I got tired of cleaning the code base, so this file will stay as is
+probably, unless I really want it to be cleaner.
+"""
+
 import os
 
 import math
@@ -11,27 +16,21 @@ from torch import nn
 from torch.autograd import Variable
 from torch.utils.data import Dataset, DataLoader
 import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D
 
-from pointcnn.core import rPointCNN
-from pointcnn.util import knn_indices_func_gpu
-from pointcnn.layers import Dense
-from visualize import *
+from PointCNN import RandPointCNN
+from PointCNN import knn_indices_func_gpu
+from PointCNN.core.util_layers import Dense
+
+from PointCNN.mnist.visualize import make_dot
 
 random.seed(0)
 
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 
-x = 2
-
-# N_neighbors, dilution, N_rep, C_out
-# 8          , 1,        all  , 16 * x
-# 8          , 2,        all  , 32 * x
-# 8          , 4,        all  , 48 * x
-# 12         , 4,        120  , 64 * x
-# 12         , 6,        120  , 80 * x
-
-# Data_dim = 3
+# C_in, C_out, D, N_neighbors, dilution, N_rep, r_indices_func, C_lifted = None, mlp_width = 2
+# (a, b, c, d, e) == (C_in, C_out, N_neighbors, dilution, N_rep)
+# Abbreviated PointCNN constructor.
+AbbPointCNN = lambda a,b,c,d,e: RandPointCNN(a, b, 3, c, d, e, knn_indices_func_gpu)
 
 class mnist_dataset(Dataset):
 
@@ -41,25 +40,21 @@ class mnist_dataset(Dataset):
 
     def __len__(self):
         return len(self.data)
-    
+
     def __getitem__(self, i):
         return self.data[i], self.labels[i]
 
-# C_in, C_out, D, N_neighbors, dilution, N_rep, r_indices_func, C_lifted = None, mlp_width = 2
-# (a, b, c, d, e) == (C_in, C_out, N_neighbors, dilution, N_rep)
-paPointCNN = lambda a,b,c,d,e: rPointCNN(a, b, 3, c, d, e, knn_indices_func_gpu)
-
 class Classifier(nn.Module):
-    
+
     def __init__(self):
         super(Classifier, self).__init__()
 
-        self.pcnn1 = paPointCNN(  1,  32,  8, 1,  -1)
+        self.pcnn1 = AbbPointCNN(  1,  32,  8, 1,  -1)
         self.pcnn2 = nn.Sequential(
-            paPointCNN( 32,  64,  8, 2,  -1),
-            paPointCNN( 64,  96,  8, 4,  -1),
-            paPointCNN( 96, 128, 12, 4, 120),
-            paPointCNN(128, 160, 12, 6, 120)
+            AbbPointCNN( 32,  64,  8, 2,  -1),
+            AbbPointCNN( 64,  96,  8, 4,  -1),
+            AbbPointCNN( 96, 128, 12, 4, 120),
+            AbbPointCNN(128, 160, 12, 6, 120)
         )
 
         self.fcn = nn.Sequential(
@@ -85,27 +80,6 @@ class Classifier(nn.Module):
         logits_mean = torch.mean(logits, dim = 1)
         return logits_mean
 
-"""
-def get_indices(batch_size, sample_num, point_num, random_sample = True):
-    indices = []
-    for i in range(batch_size):
-
-        if random_sample:
-            # point_num >= sample_num generally
-            choices = np.random.choice(point_num, sample_num, replace = (point_num < sample_num))
-        else:
-            # This modulo generally not used.
-            choices = np.arange(sample_num) % point_num
-
-        choices = np.expand_dims(choices, axis = 0)
-        b_idx_mat = np.full_like(choices, i)
-
-        # Each set of choices is paired with its batch index.
-        choices_2d = np.concatenate((b_idx_max, choices), axis = 0)
-        indices.append(choices_2d)
-    return np.stack(indices, axis = 1)  # (2, batch_size, 
-"""
-
 model = Classifier().cuda()
 
 num_class = 10
@@ -214,7 +188,7 @@ for e in range(1, num_epochs + 1):
         loss = loss_fn(out, Variable(label.long()).cuda())
         loss.backward()
         optimizer.step()
-        
+
         if global_step % 25 == 0:
             loss_v = loss.data[0]
             print("Loss:", loss_v)
diff --git a/mnist/train.py b/mnist/train.py
deleted file mode 100644
index e4cee7d..0000000
--- a/mnist/train.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import h5py
-
-f = h5py.File("./mnist/zips/train_0.h5", 'r')
-
-data = f["data"]
-label = f["label"]
-
-for r in data[0]:
-    print(r)
diff --git a/pointcnn/__init__.py b/pointcnn/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/pointcnn/layers.py b/pointcnn/layers.py
deleted file mode 100644
index d6dc0b4..0000000
--- a/pointcnn/layers.py
+++ /dev/null
@@ -1,136 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-
-try:
-    pass
-    # from .util import end_channels
-except:
-    # from util import end_channels
-    pass
-
-def end_channels(f, make_contiguous = False):
-    class wrapped_layer(nn.Module):
-        def __init__(self):
-            super(wrapped_layer, self).__init__()
-            self.f = f
-        def forward(self, x):
-            x = x.permute(0,3,1,2)
-            x = self.f(x)
-            x = x.permute(0,2,3,1)
-            return x
-    return wrapped_layer()
-
-class Dense(nn.Module):
-
-    def __init__(self, in_features, out_features, drop_rate = 0, with_bn = True,
-                 activation = nn.ReLU()):
-        super(Dense, self).__init__()
-
-        self.linear = nn.Linear(in_features, out_features)
-        self.activation = activation
-        # self.bn = LayerNorm(out_channels) if with_bn else None
-        self.drop = nn.Dropout(drop_rate) if drop_rate > 0 else None
-
-    def forward(self, x):
-        x = self.linear(x)
-        if self.activation:
-            x = self.activation(x)
-        # if self.bn:
-        #     x = self.bn(x)
-        if self.drop:
-            x = self.drop(x)
-        return x
-
-class Conv(nn.Module):
-
-    def __init__(self, in_channels, out_channels, kernel_size, with_bn = True,
-                 activation = nn.ReLU()):
-        super(Conv, self).__init__()
-
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias = not with_bn)
-        self.activation = activation
-        self.bn = nn.BatchNorm2d(out_channels, momentum = 0.9) if with_bn else None
-
-    def forward(self, x):
-        x = self.conv(x)
-        if self.activation:
-            x = self.activation(x)
-        if self.bn:
-            x = self.bn(x)
-        return x
-
-class SepConv(nn.Module):
-    """
-    Depthwise separable convolution
-    """
-
-    def __init__(self, in_channels, out_channels, kernel_size, depth_multiplier = 1,
-                 with_bn = True, activation = nn.ReLU()):
-        super(SepConv, self).__init__()
-
-        self.conv = nn.Sequential(
-            nn.Conv2d(in_channels, in_channels * depth_multiplier, kernel_size, groups = in_channels),
-            nn.Conv2d(in_channels * depth_multiplier, out_channels, 1, bias = not with_bn)
-        )
-
-        self.activation = activation
-        self.bn = nn.BatchNorm2d(out_channels, momentum = 0.9) if with_bn else None
-
-    def forward(self, x):
-        x = self.conv(x)
-        if self.activation:
-            x = self.activation(x)
-        if self.bn:
-            x = self.bn(x)
-        return x
-
-class LayerNorm(nn.Module):
-    """
-    Batch Normalization over ONLY the mini-batch layer
-    (suitable for nn.Linear layers).
-    """
-
-    def __init__(self, N, D, *args, **kwargs):
-        super(LayerNorm, self).__init__()
-        if D == 1:
-            self.bn = nn.BatchNorm1d(N, *args, **kwargs)
-        elif D == 2:
-            self.bn = nn.BatchNorm2d(N, *args, **kwargs)
-        elif D == 3:
-            self.bn = nn.BatchNorm3d(N, *args, **kwargs)
-        else:
-            raise ValueError("Dimensionality %i not supported" % D)
-
-        self.forward = lambda x: self.bn(x.unsqueeze(0)).squeeze(0)
-
-def MLP(layer_sizes, activation_layer = nn.ReLU(), batch_norm = True):
-    """
-    Creates a fully connected MLP of arbitrary depth.
-    :param layer_sizes: Sizes of MLP hidden layers.
-    :param activation_layer: Activation function to be applied in between layers.
-    :return: Multilayer perceptron module
-    """
-    if isinstance(layer_sizes, np.ndarray):
-        layer_sizes = layer_sizes.tolist()
-    if batch_norm:
-        return nn.Sequential(*[
-            nn.Sequential(nn.Linear(C_in, C_out),
-                          activation_layer,
-                          LayerNorm(D = 2, momentum = 0.9)
-            ) for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:])
-        ])
-    else:
-        return nn.Sequential(*[
-            nn.Sequential(nn.Linear(C_in, C_out),
-                          activation_layer,
-            ) for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:])
-        ])
-
-if __name__ == "__main__":
-    ftr_map = torch.autograd.Variable(torch.FloatTensor(2,8,100,100))
-    layer = SeparableConv2d(8, 16, 2)
-    out = layer(ftr_map)
-    print(out)
-
-    test = nn.SpatialConvolutionLocal(8, 16, 100, 100, 100, 100)
diff --git a/pointcnn/util.py b/pointcnn/util.py
deleted file mode 100644
index cb48ded..0000000
--- a/pointcnn/util.py
+++ /dev/null
@@ -1,118 +0,0 @@
-import sys, os
-
-import torch
-
-import numpy as np
-from sklearn.neighbors import NearestNeighbors
-
-torch.CUDA_LAUNCH_BLOCKING = 1
-
-CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) 
-
-sys.path.append(os.path.join(CURRENT_DIR, "..", "lib"))
-
-import pytorch_knn_cuda
-
-try:
-    # from .context import pytorch_knn_cuda
-    pass
-except SystemError:
-    # from context import pytorch_knn_cuda
-    pass
-
-def apply_along_dim(xs, f, dim):
-    """
-    PyTorch analog to np.apply_along_axis.
-    :param xs:
-    :param dim:
-    """
-    return torch.stack([f(x) for x in torch.unbind(xs, dim)], dim)
-
-def zipwith_matmul(xs, ys):
-    """
-    Given two lists of 2D matrices of appropriate size, zips them
-    together with matrix multiplication.
-    :param xs:
-    :param ys:
-    """
-    # xs of shape [N, n, m]
-    # ys of shape [N, m, p]
-    # return shape [N, n, p]
-    N = len(xs)
-    return torch.stack([torch.mm(xs[i], ys[i]) for i in range(N)], dim = 0)
-
-def knn_indices_func(ps, P, k, d):
-    """
-    Indexing function based on K-Nearest Neighbors search.
-    :type ps: FloatTensor (N, N_rep, D)
-    :type P: FloatTensor (N, *, D)
-    :type k: int
-    :rtype: FloatTensor (N, N_rep, N_neighbors)
-    :param ps: Representative point
-    :param P: Point cloud to get indices from
-    :param k: Number of nearest neighbors to collect.
-    :return: Array of indices, P_idx, into P such that P[n][P_idx[n],:]
-    is the set k-nearest neighbors for the representative points in P[n].
-    """
-    ps = ps.data.numpy()
-    P = P.data.numpy()
-
-    def single_batch_knn(p, P_particular):
-        nbrs = NearestNeighbors(d*k + 1, algorithm = "ball_tree").fit(P_particular)
-        indices = nbrs.kneighbors(p)[1]
-        return indices[:,1::d]
-
-    region_idx = np.stack([
-        single_batch_knn(p, P[n]) for n, p in enumerate(ps)
-    ], axis = 0)
-    return torch.from_numpy(region_idx)
-
-# def knn_indices_func_gpu(ps, P, k, d):
-# 
-#     def single_batch_knn(p, P_particular):
-#         nbrs_f = pytorch_knn_cuda.KNearestNeighbor(d*k + 1)
-#         indices = nbrs_f(P_particular, p)[0]
-#         return indices[:,1::d]
-# 
-#     region_idx = torch.stack([
-#         single_batch_knn(p, P[n]) for n, p in enumerate(ps)
-#     ], dim = 0)
-#     return region_idx
-
-def knn_indices_func_gpu(ps, P, k, d):
-
-    def single_batch_knn(qry, ref):
-        n, d = ref.size()
-        m, d = qry.size()
-        mref = ref.expand(m, n, d)
-        mqry = qry.expand(n, m, d).transpose(0, 1)
-        dist2 = torch.sum((mqry - mref)**2, 2).squeeze()
-        _, inds = torch.topk(dist2, k*d + 1, dim = 1, largest = False)
-        return inds[:,1::d]
-
-    region_idx = torch.stack([
-        single_batch_knn(p, P[n]) for n, p in enumerate(ps)
-    ], dim = 0)
-    return region_idx
-
-def plot(pts, fts):
-    num_F = fts.size()[2]
-    pts = pts[0].data.cpu().numpy()
-    plt.scatter(pts[:,0], pts[:,1], s = num_F, c = "k")
-    plt.savefig("./%i.png" % num_F)
-    plt.cla()
-
-if __name__ == "__main__":
-    from torch.autograd import Variable
-    N_rep = 1000
-    N = 2
-    num_points = 10000
-    D = 3
-    test_P  = np.random.rand(N,num_points,D).astype(np.float32)
-    idx = np.random.choice(test_P.shape[1], N_rep, replace = False)
-    test_ps = test_P[:,idx,:]
-    test_P = Variable(torch.from_numpy(test_P)).cuda()
-    test_ps = Variable(torch.from_numpy(test_ps)).cuda()
-    out = knn_indices_func_gpu(test_ps, test_P, 5, 3)
-
-    print(out)
diff --git a/tests/context.py b/tests/context.py
deleted file mode 100644
index fd3aeb0..0000000
--- a/tests/context.py
+++ /dev/null
@@ -1,4 +0,0 @@
-import sys, os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
-import pointcnn
diff --git a/tests/test.py b/tests/test.py
new file mode 100644
index 0000000..69c451b
--- /dev/null
+++ b/tests/test.py
@@ -0,0 +1,82 @@
+import unittest
+
+import torch
+from torch.autograd import Variable
+import numpy as np
+
+from PointCNN import XConv, RandPointCNN, knn_indices_func_cpu
+from PointCNN.tests.util_funcs import plot_pts_and_fts
+
+np.random.seed(0)
+
+class BasicTests(unittest.TestCase):
+    """ Basic test cases """
+
+    def test_xconv_shape(self):
+        self.assertTrue(True)
+
+        N = 4
+        D = 3
+        C_in = 8
+        C_out = 32
+        N_neighbors = 100
+
+        model = XConv(C_in, C_out, D, N_neighbors).cuda()
+        p = Variable(torch.from_numpy(np.random.rand(N,D).astype(np.float32))).cuda()
+        P = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,D).astype(np.float32))).cuda()
+        F = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,C_in).astype(np.float32))).cuda()
+        out = model(p, P, F)
+        self.assertEqual(out.size(), (N, C_out))
+
+    def test_knn(self):
+        P = np.array([[[0,0],
+                       [0,0.95],
+                       [1,0],
+                       [1,1]]])
+        ps = P[:,[0,3],:]
+
+        P = Variable(torch.from_numpy(P))
+        ps = Variable(torch.from_numpy(ps))
+
+        out = knn_indices_func_cpu(ps, P, 2).numpy()
+        target = np.array([[[1,2],
+                            [2,1]]])
+
+        self.assertTrue(np.array_equal(target, out))
+
+    def test_pointcnn_shape(self):
+        N = 1
+        num_points = 1000
+        dims = 2
+        C_in = 4
+        K = 10
+        D = 1
+
+        layer1 = RandPointCNN(C_in,   8, dims, K, D, 1000, knn_indices_func_cpu).cuda()
+        layer2 = RandPointCNN(   8,  16, dims, K, D,  500, knn_indices_func_cpu).cuda()
+        layer3 = RandPointCNN(  16,  32, dims, K, D,  250, knn_indices_func_cpu).cuda()
+        layer4 = RandPointCNN(  32,  64, dims, K, D,  125, knn_indices_func_cpu).cuda()
+        layer5 = RandPointCNN(  64, 128, dims, K, D,   50, knn_indices_func_cpu).cuda()
+
+        pts  = np.random.rand(N,num_points,dims).astype(np.float32)
+        fts  = np.random.rand(N,num_points,C_in).astype(np.float32)
+        pts = Variable(torch.from_numpy(pts)).cuda()
+        fts = Variable(torch.from_numpy(fts)).cuda()
+
+        if True:
+            pts, fts = layer1((pts, fts))
+        else:
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer1((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer2((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer3((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer4((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer5((pts, fts))
+            plot_pts_and_fts(pts, fts)
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_advanced.py b/tests/test_advanced.py
deleted file mode 100644
index 3f6c382..0000000
--- a/tests/test_advanced.py
+++ /dev/null
@@ -1,10 +0,0 @@
-import unittest
-
-class AdvancedTests(unittest.TestCase):
-    """ Basic test cases """
-
-    def test_example(self):
-        self.assertTrue(True)
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 2970dae..3c1ba80 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -4,14 +4,16 @@ import torch
 from torch.autograd import Variable
 import numpy as np
 
-from pointcnn.core import XConv, knn_indices_func
+from PointCNN import XConv, RandPointCNN, knn_indices_func_cpu
+from PointCNN.tests import plot_pts_and_fts
+
+np.random.seed(0)
 
 class BasicTests(unittest.TestCase):
     """ Basic test cases """
 
     def test_xconv_shape(self):
         self.assertTrue(True)
-        np.random.seed(0)
 
         N = 4
         D = 3
@@ -36,12 +38,45 @@ class BasicTests(unittest.TestCase):
         P = Variable(torch.from_numpy(P))
         ps = Variable(torch.from_numpy(ps))
 
-        out = knn_indices_func(ps, P, 2).numpy()
+        out = knn_indices_func_cpu(ps, P, 2).numpy()
         target = np.array([[[1,2],
                             [2,1]]])
 
         self.assertTrue(np.array_equal(target, out))
 
+    def test_pointcnn_shape(self):
+        N = 1
+        num_points = 1000
+        dims = 2
+        C_in = 4
+        K = 10
+        D = 1
+
+        layer1 = RandPointCNN(C_in,   8, dims, K, D, 1000, knn_indices_func_cpu).cuda()
+        layer2 = RandPointCNN(   8,  16, dims, K, D,  500, knn_indices_func_cpu).cuda()
+        layer3 = RandPointCNN(  16,  32, dims, K, D,  250, knn_indices_func_cpu).cuda()
+        layer4 = RandPointCNN(  32,  64, dims, K, D,  125, knn_indices_func_cpu).cuda()
+        layer5 = RandPointCNN(  64, 128, dims, K, D,   50, knn_indices_func_cpu).cuda()
+
+        pts  = np.random.rand(N,num_points,dims).astype(np.float32)
+        fts  = np.random.rand(N,num_points,C_in).astype(np.float32)
+        pts = Variable(torch.from_numpy(pts)).cuda()
+        fts = Variable(torch.from_numpy(fts)).cuda()
+
+        if True:
+            pts, fts = layer1((pts, fts))
+        else:
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer1((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer2((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer3((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer4((pts, fts))
+            plot_pts_and_fts(pts, fts)
+            pts, fts = layer5((pts, fts))
+            plot_pts_and_fts(pts, fts)
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/util_funcs.py b/tests/util_funcs.py
new file mode 100644
index 0000000..5924aa2
--- /dev/null
+++ b/tests/util_funcs.py
@@ -0,0 +1,49 @@
+# External Modules
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Internal Modules
+from PointCNN.core import UFloatTensor
+
+def plot_pts_and_fts(pts : UFloatTensor,  # (N, x, dims)
+                     fts : UFloatTensor   # (N, x, y)
+                    ) -> None:
+    """
+    Visualization function. Shows points and number of features, represented by
+    the size of the point.
+    :param pts: Point cloud such that fts[:,p_idx,:] is the feature associated
+    with pts[:,p_idx,:].
+    :param fts: Features such that pts[:,p_idx,:] is the feature associated
+    with fts[:,p_idx,:].
+    """
+    if pts.is_cuda:
+        pts = pts.cpu()
+    num_F = fts.size()[2]
+    pts = pts[0].data.numpy()
+    plt.scatter(pts[:,0], pts[:,1], s = num_F, c = "k")
+    plt.show()
+    plt.cla()
+
+def plot_neighborhood(pts : UFloatTensor,          # (N, x, dims)
+                      rep_pts : UFloatTensor,      # (N, P, dims)
+                      pts_regional : UFloatTensor  # (N, P, dims)
+                     ) -> None:
+    """
+    Visualization function. Shows neighborhood points around a randomly
+    selected representative.
+    :param pts: Point cloud.
+    :param rep_pts: Representative points.
+    :param pts_regional: Regional neighborhoods around representative points.
+    """
+    if rep_pts.is_cuda:
+        rep_pts = rep_pts.cpu()
+        pts_regional = pts_regional.cpu()
+    n = np.randint(0, rep_pts.shape[0])
+    t = np.randint(0, rep_pts.shape[1])
+    test_point = rep_pts[n,t,:].data.numpy()
+    neighborhood = pts_regional[n,t,:,:].data.numpy()
+    plt.scatter(pts[n][:,0], pts[n][:,1])
+    plt.scatter(test_point[0], test_point[1], s = 100, c = 'green')
+    plt.scatter(neighborhood[:,0], neighborhood[:,1], s = 100, c = 'red')
+    plt.show()
+    plt.cla()