1 files changed, 70 insertions, 7 deletions
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)