Updated tests and made Screen printing prettier

__main__.py

@@ -3,34 +3,27 @@ import matplotlib as mpl
 import matplotlib.pyplot as plt
 import elements
 import graphics
+import matrices
 
 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)
+BS1 = elements.BeamSplitter(previous = {"left": (source, "left"), "right": None}, matrix=matrices.HALF_MIRROR_LEFT_ACTIVE)
 
-X_left_1 = elements.Plate(previous = (BS1, "left"), jones_matrix=X)
-Z_left_2 = elements.Plate(previous = (X_left_1, "left"), jones_matrix = Z)
+X_left_1 = elements.Plate(previous = (BS1, "left"), jones_matrix=matrices.X)
+Z_left_2 = elements.Plate(previous = (X_left_1, "left"), jones_matrix = matrices.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)
+Z_right_1 = elements.Plate(previous = (BS1, "right"), jones_matrix = matrices.Z)
+X_right_2 = elements.Plate(previous = (Z_right_1, "left"), jones_matrix = matrices.X)
 
-BS2 = elements.BeamSplitter(previous = {"left": (Z_left_2, "left"), "right": (X_right_2, "left")}, matrix = B)
+BS2 = elements.BeamSplitter(previous = {"left": (Z_left_2, "left"), "right": (X_right_2, "left")}, matrix = matrices.HALF_MIRROR_RIGHT_ACTIVE)
 
-screen = elements.Screen(previous = (BS2, "left"))
+screen_right = elements.Screen(previous = (BS2, "right"))
+screen_left = elements.Screen(previous = (BS2, "left"))
 
-print(screen.out())
+screen_right.show("right")
+screen_left.show("left")

elements.py

@@ -14,9 +14,12 @@ class Screen:
     def __init__(self, previous=(None, None)):
         self.previous = previous
 
-    def show(self):
+    def show(self, name = 'ψ'):
         previous_component, previous_component_output_side = self.previous
-        print(previous_component.out(previous_component_output_side))
+        ket = previous_component.out(previous_component_output_side)
+        alpha = round(ket[0])
+        beta = round(ket[1])
+        print(f"|{name}⟩ = {alpha}|0⟩ + {beta}|1⟩")
 
     def out(self):
         previous_component, previous_component_output_side = self.previous

matrices.py

@@ -67,7 +67,7 @@ HALF_MIRROR_LEFT_ACTIVE = (1/np.sqrt(2))*np.array([[-1,  0,   1,   0],
                                                    [1,   0,   1,   0],
                                                    [0,   1,   0,  -1]], dtype='complex')
 
-HALF_MIRROR_LEFT_ACTIVE = (1/np.sqrt(2))*np.array([[1,   0,   1,   0],
+HALF_MIRROR_RIGHT_ACTIVE = (1/np.sqrt(2))*np.array([[1,   0,   1,   0],
                                                    [0,  -1,   0,   1],
                                                    [1,   0,  -1,   0],
                                                    [0,   1,   0,   1]], dtype='complex')