Shitty Poincare sphere (to be improved later)

3 files changed, 78 insertions, 15 deletions

diff --git a/src/phys.py b/src/phys.py
index 8065d14..7bf20e6 100644
--- a/src/phys.py
+++ b/src/phys.py
@@ -73,6 +73,12 @@ class Polarizer:
 		self.t2 = 0
 		self.enable = True
 
+	def matrix(self):
+		A = np.sqrt(np.array([[self.t1, 0], [0, self.t2]])) # FIXME: half-assed again
+		M =  np.array([[1, 0], [0, np.exp(1j * self.phase_retardation),]])
+		return np.matmul(R(-self.angle - self.delta),
+			np.matmul(np.matmul(M, A), R(self.angle + self.delta)))
+
 	def mul(self, state):
 		# unpolarized light
 		if state is None:
@@ -81,13 +87,7 @@ class Polarizer:
 				              np.array([[1], [0]])) * np.sqrt(self.t1)
 			else:
 				return None
-
-		# FIXME: half-assed again
-		A = np.sqrt(np.array([[self.t1, 0], [0, self.t2]]))
-		M =  np.array([[1, 0], [0, np.exp(1j * self.phase_retardation),]])
-		MR = np.matmul(R(-self.angle - self.delta), \
-		    np.matmul(np.matmul(M, A), R(self.angle + self.delta)))
-		return np.dot(MR, state)
+		return np.dot(self.matrix(), state)
 
 class System:
 	def __init__(self):
diff --git a/src/ui.py b/src/ui.py
index 22ae436..fe9bb83 100644
--- a/src/ui.py
+++ b/src/ui.py
@@ -69,7 +69,7 @@ def update():
 		row.stokes[4].setText("%g°" % (ellipse.alpha * 180 / np.pi))
 		row.stokes[5].setText("%g°" % (ellipse.theta * 180 / np.pi))
 
-	GUI.poincare.repaint()
+	GUI.poincare.update_system(system)
 
 	#if GUI.auto_optimize.isChecked():
 	#	optimize()
diff --git a/src/ui_widgets.py b/src/ui_widgets.py
index 291324f..93096ff 100644
--- a/src/ui_widgets.py
+++ b/src/ui_widgets.py
@@ -2,8 +2,11 @@ import numpy as np
 from PyQt5.QtWidgets import *
 from PyQt5.QtGui import *
 from PyQt5.QtCore import *
+import pyqtgraph as pg, pyqtgraph.opengl as gl
+import scipy.linalg
 
 from ui import GUI
+import phys
 
 
 
@@ -227,12 +230,72 @@ class AngleSlider(QHBoxLayout):
 
 
 
-class PoincareWidget(QWidget):
-	def __init__(self,):
-		QWidget.__init__(self)
+class PoincareWidget(gl.GLViewWidget):
+	def __init__(self):
+		super().__init__()
 		self.setSizePolicy(QSizePolicy(QSizePolicy.Expanding, \
 		                               QSizePolicy.Expanding))
-
-	def paintEvent(self, event):
-		pass
-
+		self.setMinimumWidth(200)
+
+		zgrid = gl.GLGridItem()
+		zgrid.scale(0.1, 0.1, 0.1)
+		self.addItem(zgrid)
+
+		axis = gl.GLAxisItem()
+		axis.scale(3, 3, 3)
+		self.addItem(axis)
+
+		sphere = gl.GLMeshItem(
+			meshdata=gl.MeshData.sphere(200, 200),
+			color=(0.5, 0.5, 0.5, 0.5), glOptions="translucent")
+		sphere.scale(0.995, 0.995, 0.955) # depth testing fuckery
+		self.addItem(sphere)
+
+		self.arcs = []
+
+
+	def make_arc(state, element, N=50):
+		P = scipy.linalg.fractional_matrix_power(element.matrix(), 1 / N)
+
+		points = np.array([
+			1.5 * phys.jones_to_stokes(state)[1:, 0],
+			phys.jones_to_stokes(state)[1:, 0]])
+
+		for i in range(1, N + 1):
+			M = np.linalg.matrix_power(P, i)
+			new = np.dot(M, state)
+			points = np.append(points, \
+				np.array([phys.jones_to_stokes(new)[1:, 0]]), axis=0)
+
+		return points
+
+	colors = [
+		(1, 0, 0, 1),
+		(1, 0.5, 0, 1),
+		(1, 1, 0, 1),
+		(0.5, 1, 0, 1),
+		(0, 1, 0, 1),
+		(0, 1, 0.5, 1),
+		(0, 1, 1, 1),
+		(0, 0.5, 1, 1),
+		(0, 0, 1, 1)
+	]
+
+	def update_system(self, system):
+		for arc in self.arcs:
+			self.removeItem(arc)
+		self.arcs = []
+
+		for i, element in enumerate(system.elements[1:]):
+			state = system.states[i] # i - 1 actually
+			if state is None:
+				continue
+
+			color = PoincareWidget.colors[i % len(PoincareWidget.colors)]
+
+			points = PoincareWidget.make_arc(state, element)
+			arc = gl.GLLinePlotItem(pos=points, color=color,
+				width=2, antialias=True, mode="line_strip",
+				glOptions="opaque")
+			self.addItem(arc)
+			self.arcs.append(arc)