diff --git a/__init__.py b/__init__.py
deleted file mode 100644
index 9664db059251e7c9e9a2e29007f3fe619d4c6b51..0000000000000000000000000000000000000000
--- a/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-import numpy as np
-import cmath
diff --git a/__main__.py b/__main__.py
index 86d59a274a0b8d734ecb3c3f16b446e8ec19ce63..bc5b490ae5905afbdfbaa55d9d0b1ccf9deb8443 100644
--- a/__main__.py
+++ b/__main__.py
@@ -1,4 +1,34 @@
+import numpy as np
 import elements
 import graphics
 
 print("Hi World!")
+
+## Tests
+
+H = (1/np.sqrt(2))*np.array([[1, 1],[1, -1]], dtype='complex')
+H_rev = (1/np.sqrt(2))*np.array([[-1, 1],[1, 1]], dtype='complex')
+
+A = np.kron(H_rev, np.eye(2, dtype='complex')) # first beam splitter matrix
+B = np.kron(H, np.eye(2, dtype='complex')) # second beam splitter matrix
+
+X = np.array([[0, 1], [1, 0]], dtype='complex')
+Z = np.array([[1, 0], [0, -1]], dtype='complex')
+
+
+source = elements.Source(initial_vector = np.array([1.+0j, 0.]))
+
+BS1 = elements.BeamSplitter(previous = {"left": (source, "left"), "right": None}, matrix=A)
+
+X_left_1 = elements.Plate(previous = (BS1, "left"), jones_matrix=X)
+Z_left_2 = elements.Plate(previous = (X_left_1, "left"), jones_matrix = Z)
+
+
+Z_right_1 = elements.Plate(previous = (BS1, "right"), jones_matrix = Z)
+X_right_2 = elements.Plate(previous = (Z_right_1, "left"), jones_matrix = X)
+
+BS2 = elements.BeamSplitter(previous = {"left": (Z_left_2, "left"), "right": (X_right_2, "left")}, matrix = B)
+
+screen = elements.Screen(previous = (BS2, "left"))
+
+print(screen.out())
diff --git a/elements.py b/elements.py
index 2026c2b4875b5ed75f202040b8b47496a1af75bd..a8cc8c445aaa7ef9dd5e6756dd946388eb76a1b1 100644
--- a/elements.py
+++ b/elements.py
@@ -58,7 +58,7 @@ class BeamSplitter(Element):
             left_input_state = left_input_element.out(left_input_element_output_side)
 
         if self.previous["right"] is None:
-            left_input_state = np.zeros((2,), dtype="complex") # if no input on the right, then |\psi> = 0
+            right_input_state = np.zeros((2,), dtype="complex") # if no input on the right, then |\psi> = 0
         else:
             right_input_element = self.previous["right"][0]
             right_input_element_output_side = self.previous["right"][1]