Input intensity

1 files changed, 6 insertions, 1 deletions

diff --git a/src/phys.py b/src/phys.py
index cf30667..4e78fe2 100644
--- a/src/phys.py
+++ b/src/phys.py
@@ -94,6 +94,7 @@ class Polarizer:
 class System:
 	def __init__(self):
 		self.elements = list()
+		self.input_intensity = 1
 
 	def recalculate(system):
 		system.states = [None] * len(system.elements)
@@ -102,6 +103,10 @@ class System:
 		state = None
 		for i, pol in enumerate(system.elements):
 			if pol.enable:
-				state = pol.mul(state)
+				new_state = pol.mul(state)
+				if state is None and new_state is not None:
+					state = new_state * np.sqrt(system.input_intensity)
+				else:
+					state = new_state
 			system.states[i] = state
 			system.ellipses.append(Ellipse(state))