Neural_Net_Visualizer/Visual_NN.py

import vpython as vp
import numpy as np

neuron_color = vp.vector(0.00, 0.49, 0.73)
weight_color = vp.vector(0.80, 0.80, 0.80)
text_color = vp.vector(1.0, 1.0, 1.0)


class Network:
    def __init__(self, architecture, act_funcs, title,
                 x_compress=2.0, y_compress=1.0):

        neuron_storage = []
        weight_storage = []
        net_amp = (len(architecture) - 1) / 2
        max_amp = (max(architecture) - 1) / 2

        x_text = net_amp * x_compress
        y_title = max_amp * y_compress + 0.5
        if act_funcs:
            y_text = max_amp * y_compress + 1.0
            y_funcs_text = y_text - 0.5
        else:
            y_text = max_amp * y_compress + 0.5
        
        vp.text(text=title, pos=vp.vector(0, y_title, 0),
                align="center", color=text_color, height=0.2)
        
        # Descriptive Text Below the Network
        vp.text(text="inputs", pos=vp.vector(-x_text, -y_text, 0),
                align="right", color=text_color, height=0.2)
        vp.text(text="|", pos=vp.vector(-x_text + 0.5, -y_text, 0),
                align="center", color=text_color, height=0.2)
        vp.text(text="hidden\nlayers", pos=vp.vector(0, -y_text, 0),
                align="center", color=text_color, height=0.2)
        vp.text(text="|", pos=vp.vector(x_text - 0.5, -y_text, 0),
                align="center", color=text_color, height=0.2)
        vp.text(text="outputs", pos=vp.vector(x_text, -y_text, 0),
                align="left", color=text_color, height=0.2)
        
        # Initialize the Neuron Storage
        for n in architecture:
            layer = [0 for i in range(n)]
            neuron_storage.append(layer)

        # Create and Store the Neuron Spheres
        for layer_num, num_neurons in enumerate(architecture):
            layer_amp = (num_neurons - 1) / 2
            x = (layer_num - net_amp) * x_compress

            if act_funcs:
                vp.text(text=act_funcs[layer_num],
                        pos=vp.vector(x, -y_funcs_text, 0),
                        align="center", color=text_color, height=0.2)
            
            for neuron_num in range(num_neurons):
                y = (neuron_num - layer_amp) * y_compress
                neuron_pos = vp.vector(x, y, 0)
                neuron_storage[layer_num][neuron_num] = vp.sphere(
                    pos=neuron_pos, radius=0.1, color=neuron_color)

        # Create and Store the Weight Bonds
        for layer_num, current_n in enumerate(architecture[:-1]):
            next_n = architecture[layer_num+1]

            weight_storage.append([])

            for i in range(current_n):
                weight_storage[-1].append([])
                
                x_start = (layer_num - net_amp) * x_compress
                x_end =   x_start + 1 * x_compress
                
                for j in range(next_n):
                    y_start = (i - (current_n - 1) / 2) * y_compress
                    y_end = (j - (next_n - 1) / 2) * y_compress
                    
                    start = vp.vector(x_start, y_start, 0)
                    end = vp.vector(x_end, y_end, 0)
                    
                    weight_storage[-1][-1].append(vp.cylinder(
                        pos=start, axis=end-start, radius=0.01,
                        color=weight_color))

         # Make Neuron and Weight Bond Storage Class Attributes
        self.neuron_storage = neuron_storage
        self.weight_storage = weight_storage
                

arch = [1, 13, 1]
funcs = ['linear', 'tanh', 'linear']
network = Network(
    arch, funcs, "Neural Network", x_compress=2.0, y_compress=0.33)