563 lines
17 KiB
Python
563 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
import pyqtgraph as pg
|
||
from pyqtgraph.Qt import QtGui, QtCore
|
||
import numpy as np
|
||
import csv, gzip, os
|
||
from pyqtgraph import Point
|
||
|
||
class GlassDB:
|
||
"""
|
||
Database of dispersion coefficients for Schott glasses
|
||
+ Corning 7980
|
||
"""
|
||
def __init__(self, fileName='schott_glasses.csv'):
|
||
path = os.path.dirname(__file__)
|
||
fh = gzip.open(os.path.join(path, 'schott_glasses.csv.gz'), 'rb')
|
||
r = csv.reader(map(str, fh.readlines()))
|
||
lines = [x for x in r]
|
||
self.data = {}
|
||
header = lines[0]
|
||
for l in lines[1:]:
|
||
info = {}
|
||
for i in range(1, len(l)):
|
||
info[header[i]] = l[i]
|
||
self.data[l[0]] = info
|
||
self.data['Corning7980'] = { ## Thorlabs UV fused silica--not in schott catalog.
|
||
'B1': 0.68374049400,
|
||
'B2': 0.42032361300,
|
||
'B3': 0.58502748000,
|
||
'C1': 0.00460352869,
|
||
'C2': 0.01339688560,
|
||
'C3': 64.49327320000,
|
||
'TAUI25/250': 0.95, ## transmission data is fabricated, but close.
|
||
'TAUI25/1400': 0.98,
|
||
}
|
||
|
||
for k in self.data:
|
||
self.data[k]['ior_cache'] = {}
|
||
|
||
|
||
def ior(self, glass, wl):
|
||
"""
|
||
Return the index of refraction for *glass* at wavelength *wl*.
|
||
|
||
The *glass* argument must be a key in self.data.
|
||
"""
|
||
info = self.data[glass]
|
||
cache = info['ior_cache']
|
||
if wl not in cache:
|
||
B = list(map(float, [info['B1'], info['B2'], info['B3']]))
|
||
C = list(map(float, [info['C1'], info['C2'], info['C3']]))
|
||
w2 = (wl/1000.)**2
|
||
n = np.sqrt(1.0 + (B[0]*w2 / (w2-C[0])) + (B[1]*w2 / (w2-C[1])) + (B[2]*w2 / (w2-C[2])))
|
||
cache[wl] = n
|
||
return cache[wl]
|
||
|
||
def transmissionCurve(self, glass):
|
||
data = self.data[glass]
|
||
keys = [int(x[7:]) for x in data.keys() if 'TAUI25' in x]
|
||
keys.sort()
|
||
curve = np.empty((2,len(keys)))
|
||
for i in range(len(keys)):
|
||
curve[0][i] = keys[i]
|
||
key = 'TAUI25/%d' % keys[i]
|
||
val = data[key]
|
||
if val == '':
|
||
val = 0
|
||
else:
|
||
val = float(val)
|
||
curve[1][i] = val
|
||
return curve
|
||
|
||
|
||
GLASSDB = GlassDB()
|
||
|
||
|
||
def wlPen(wl):
|
||
"""Return a pen representing the given wavelength"""
|
||
l1 = 400
|
||
l2 = 700
|
||
hue = np.clip(((l2-l1) - (wl-l1)) * 0.8 / (l2-l1), 0, 0.8)
|
||
val = 1.0
|
||
if wl > 700:
|
||
val = 1.0 * (((700-wl)/700.) + 1)
|
||
elif wl < 400:
|
||
val = wl * 1.0/400.
|
||
#print hue, val
|
||
color = pg.hsvColor(hue, 1.0, val)
|
||
pen = pg.mkPen(color)
|
||
return pen
|
||
|
||
|
||
class ParamObj(object):
|
||
# Just a helper for tracking parameters and responding to changes
|
||
def __init__(self):
|
||
self.__params = {}
|
||
|
||
def __setitem__(self, item, val):
|
||
self.setParam(item, val)
|
||
|
||
def setParam(self, param, val):
|
||
self.setParams(**{param:val})
|
||
|
||
def setParams(self, **params):
|
||
"""Set parameters for this optic. This is a good function to override for subclasses."""
|
||
self.__params.update(params)
|
||
self.paramStateChanged()
|
||
|
||
def paramStateChanged(self):
|
||
pass
|
||
|
||
def __getitem__(self, item):
|
||
# bug in pyside 1.2.2 causes getitem to be called inside QGraphicsObject.parentItem:
|
||
return self.getParam(item) # PySide bug: https://bugreports.qt.io/browse/PYSIDE-671
|
||
|
||
def __len__(self):
|
||
# Workaround for PySide bug: https://bugreports.qt.io/browse/PYSIDE-671
|
||
return 0
|
||
|
||
def getParam(self, param):
|
||
return self.__params[param]
|
||
|
||
|
||
class Optic(pg.GraphicsObject, ParamObj):
|
||
|
||
sigStateChanged = QtCore.Signal()
|
||
|
||
|
||
def __init__(self, gitem, **params):
|
||
ParamObj.__init__(self)
|
||
pg.GraphicsObject.__init__(self) #, [0,0], [1,1])
|
||
|
||
self.gitem = gitem
|
||
self.surfaces = gitem.surfaces
|
||
gitem.setParentItem(self)
|
||
|
||
self.roi = pg.ROI([0,0], [1,1])
|
||
self.roi.addRotateHandle([1, 1], [0.5, 0.5])
|
||
self.roi.setParentItem(self)
|
||
|
||
defaults = {
|
||
'pos': Point(0,0),
|
||
'angle': 0,
|
||
}
|
||
defaults.update(params)
|
||
self._ior_cache = {}
|
||
self.roi.sigRegionChanged.connect(self.roiChanged)
|
||
self.setParams(**defaults)
|
||
|
||
def updateTransform(self):
|
||
self.resetTransform()
|
||
self.setPos(0, 0)
|
||
self.translate(Point(self['pos']))
|
||
self.rotate(self['angle'])
|
||
|
||
def setParam(self, param, val):
|
||
ParamObj.setParam(self, param, val)
|
||
|
||
def paramStateChanged(self):
|
||
"""Some parameters of the optic have changed."""
|
||
# Move graphics item
|
||
self.gitem.setPos(Point(self['pos']))
|
||
self.gitem.resetTransform()
|
||
self.gitem.rotate(self['angle'])
|
||
|
||
# Move ROI to match
|
||
try:
|
||
self.roi.sigRegionChanged.disconnect(self.roiChanged)
|
||
br = self.gitem.boundingRect()
|
||
o = self.gitem.mapToParent(br.topLeft())
|
||
self.roi.setAngle(self['angle'])
|
||
self.roi.setPos(o)
|
||
self.roi.setSize([br.width(), br.height()])
|
||
finally:
|
||
self.roi.sigRegionChanged.connect(self.roiChanged)
|
||
|
||
self.sigStateChanged.emit()
|
||
|
||
def roiChanged(self, *args):
|
||
pos = self.roi.pos()
|
||
# rotate gitem temporarily so we can decide where it will need to move
|
||
self.gitem.resetTransform()
|
||
self.gitem.rotate(self.roi.angle())
|
||
br = self.gitem.boundingRect()
|
||
o1 = self.gitem.mapToParent(br.topLeft())
|
||
self.setParams(angle=self.roi.angle(), pos=pos + (self.gitem.pos() - o1))
|
||
|
||
def boundingRect(self):
|
||
return QtCore.QRectF()
|
||
|
||
def paint(self, p, *args):
|
||
pass
|
||
|
||
def ior(self, wavelength):
|
||
return GLASSDB.ior(self['glass'], wavelength)
|
||
|
||
|
||
|
||
class Lens(Optic):
|
||
def __init__(self, **params):
|
||
defaults = {
|
||
'dia': 25.4, ## diameter of lens
|
||
'r1': 50., ## positive means convex, use 0 for planar
|
||
'r2': 0, ## negative means convex
|
||
'd': 4.0,
|
||
'glass': 'N-BK7',
|
||
'reflect': False,
|
||
}
|
||
defaults.update(params)
|
||
d = defaults.pop('d')
|
||
defaults['x1'] = -d/2.
|
||
defaults['x2'] = d/2.
|
||
|
||
gitem = CircularSolid(brush=(100, 100, 130, 100), **defaults)
|
||
Optic.__init__(self, gitem, **defaults)
|
||
|
||
def propagateRay(self, ray):
|
||
"""Refract, reflect, absorb, and/or scatter ray. This function may create and return new rays"""
|
||
|
||
"""
|
||
NOTE:: We can probably use this to compute refractions faster: (from GLSL 120 docs)
|
||
|
||
For the incident vector I and surface normal N, and the
|
||
ratio of indices of refraction eta, return the refraction
|
||
vector. The result is computed by
|
||
k = 1.0 - eta * eta * (1.0 - dot(N, I) * dot(N, I))
|
||
if (k < 0.0)
|
||
return genType(0.0)
|
||
else
|
||
return eta * I - (eta * dot(N, I) + sqrt(k)) * N
|
||
The input parameters for the incident vector I and the
|
||
surface normal N must already be normalized to get the
|
||
desired results. eta == ratio of IORs
|
||
|
||
|
||
For reflection:
|
||
For the incident vector I and surface orientation N,
|
||
returns the reflection direction:
|
||
I – 2 ∗ dot(N, I) ∗ N
|
||
N must already be normalized in order to achieve the
|
||
desired result.
|
||
"""
|
||
iors = [self.ior(ray['wl']), 1.0]
|
||
for i in [0,1]:
|
||
surface = self.surfaces[i]
|
||
ior = iors[i]
|
||
p1, ai = surface.intersectRay(ray)
|
||
if p1 is None:
|
||
ray.setEnd(None)
|
||
break
|
||
p1 = surface.mapToItem(ray, p1)
|
||
|
||
rd = ray['dir']
|
||
a1 = np.arctan2(rd[1], rd[0])
|
||
ar = a1 - ai + np.arcsin((np.sin(ai) * ray['ior'] / ior))
|
||
ray.setEnd(p1)
|
||
dp = Point(np.cos(ar), np.sin(ar))
|
||
ray = Ray(parent=ray, ior=ior, dir=dp)
|
||
return [ray]
|
||
|
||
|
||
class Mirror(Optic):
|
||
def __init__(self, **params):
|
||
defaults = {
|
||
'r1': 0,
|
||
'r2': 0,
|
||
'd': 0.01,
|
||
}
|
||
defaults.update(params)
|
||
d = defaults.pop('d')
|
||
defaults['x1'] = -d/2.
|
||
defaults['x2'] = d/2.
|
||
gitem = CircularSolid(brush=(100,100,100,255), **defaults)
|
||
Optic.__init__(self, gitem, **defaults)
|
||
|
||
def propagateRay(self, ray):
|
||
"""Refract, reflect, absorb, and/or scatter ray. This function may create and return new rays"""
|
||
|
||
surface = self.surfaces[0]
|
||
p1, ai = surface.intersectRay(ray)
|
||
if p1 is not None:
|
||
p1 = surface.mapToItem(ray, p1)
|
||
rd = ray['dir']
|
||
a1 = np.arctan2(rd[1], rd[0])
|
||
ar = a1 + np.pi - 2*ai
|
||
ray.setEnd(p1)
|
||
dp = Point(np.cos(ar), np.sin(ar))
|
||
ray = Ray(parent=ray, dir=dp)
|
||
else:
|
||
ray.setEnd(None)
|
||
return [ray]
|
||
|
||
|
||
class CircularSolid(pg.GraphicsObject, ParamObj):
|
||
"""GraphicsObject with two circular or flat surfaces."""
|
||
def __init__(self, pen=None, brush=None, **opts):
|
||
"""
|
||
Arguments for each surface are:
|
||
x1,x2 - position of center of _physical surface_
|
||
r1,r2 - radius of curvature
|
||
d1,d2 - diameter of optic
|
||
"""
|
||
defaults = dict(x1=-2, r1=100, d1=25.4, x2=2, r2=100, d2=25.4)
|
||
defaults.update(opts)
|
||
ParamObj.__init__(self)
|
||
self.surfaces = [CircleSurface(defaults['r1'], defaults['d1']), CircleSurface(-defaults['r2'], defaults['d2'])]
|
||
pg.GraphicsObject.__init__(self)
|
||
for s in self.surfaces:
|
||
s.setParentItem(self)
|
||
|
||
if pen is None:
|
||
self.pen = pg.mkPen((220,220,255,200), width=1, cosmetic=True)
|
||
else:
|
||
self.pen = pg.mkPen(pen)
|
||
|
||
if brush is None:
|
||
self.brush = pg.mkBrush((230, 230, 255, 30))
|
||
else:
|
||
self.brush = pg.mkBrush(brush)
|
||
|
||
self.setParams(**defaults)
|
||
|
||
def paramStateChanged(self):
|
||
self.updateSurfaces()
|
||
|
||
def updateSurfaces(self):
|
||
self.surfaces[0].setParams(self['r1'], self['d1'])
|
||
self.surfaces[1].setParams(-self['r2'], self['d2'])
|
||
self.surfaces[0].setPos(self['x1'], 0)
|
||
self.surfaces[1].setPos(self['x2'], 0)
|
||
|
||
self.path = QtGui.QPainterPath()
|
||
self.path.connectPath(self.surfaces[0].path.translated(self.surfaces[0].pos()))
|
||
self.path.connectPath(self.surfaces[1].path.translated(self.surfaces[1].pos()).toReversed())
|
||
self.path.closeSubpath()
|
||
|
||
def boundingRect(self):
|
||
return self.path.boundingRect()
|
||
|
||
def shape(self):
|
||
return self.path
|
||
|
||
def paint(self, p, *args):
|
||
p.setRenderHints(p.renderHints() | p.Antialiasing)
|
||
p.setPen(self.pen)
|
||
p.fillPath(self.path, self.brush)
|
||
p.drawPath(self.path)
|
||
|
||
|
||
class CircleSurface(pg.GraphicsObject):
|
||
def __init__(self, radius=None, diameter=None):
|
||
"""center of physical surface is at 0,0
|
||
radius is the radius of the surface. If radius is None, the surface is flat.
|
||
diameter is of the optic's edge."""
|
||
pg.GraphicsObject.__init__(self)
|
||
|
||
self.r = radius
|
||
self.d = diameter
|
||
self.mkPath()
|
||
|
||
def setParams(self, r, d):
|
||
self.r = r
|
||
self.d = d
|
||
self.mkPath()
|
||
|
||
def mkPath(self):
|
||
self.prepareGeometryChange()
|
||
r = self.r
|
||
d = self.d
|
||
h2 = d/2.
|
||
self.path = QtGui.QPainterPath()
|
||
if r == 0: ## flat surface
|
||
self.path.moveTo(0, h2)
|
||
self.path.lineTo(0, -h2)
|
||
else:
|
||
## half-height of surface can't be larger than radius
|
||
h2 = min(h2, abs(r))
|
||
arc = QtCore.QRectF(0, -r, r*2, r*2)
|
||
a1 = np.arcsin(h2/r) * 180. / np.pi
|
||
a2 = -2*a1
|
||
a1 += 180.
|
||
self.path.arcMoveTo(arc, a1)
|
||
self.path.arcTo(arc, a1, a2)
|
||
self.h2 = h2
|
||
|
||
def boundingRect(self):
|
||
return self.path.boundingRect()
|
||
|
||
def paint(self, p, *args):
|
||
return ## usually we let the optic draw.
|
||
|
||
def intersectRay(self, ray):
|
||
## return the point of intersection and the angle of incidence
|
||
#print "intersect ray"
|
||
h = self.h2
|
||
r = self.r
|
||
p, dir = ray.currentState(relativeTo=self) # position and angle of ray in local coords.
|
||
#print " ray: ", p, dir
|
||
p = p - Point(r, 0) ## move position so center of circle is at 0,0
|
||
#print " adj: ", p, r
|
||
|
||
if r == 0:
|
||
#print " flat"
|
||
if dir[0] == 0:
|
||
y = 0
|
||
else:
|
||
y = p[1] - p[0] * dir[1]/dir[0]
|
||
if abs(y) > h:
|
||
return None, None
|
||
else:
|
||
return (Point(0, y), np.arctan2(dir[1], dir[0]))
|
||
else:
|
||
#print " curve"
|
||
## find intersection of circle and line (quadratic formula)
|
||
dx = dir[0]
|
||
dy = dir[1]
|
||
dr = (dx**2 + dy**2) ** 0.5
|
||
D = p[0] * (p[1]+dy) - (p[0]+dx) * p[1]
|
||
idr2 = 1.0 / dr**2
|
||
disc = r**2 * dr**2 - D**2
|
||
if disc < 0:
|
||
return None, None
|
||
disc2 = disc**0.5
|
||
if dy < 0:
|
||
sgn = -1
|
||
else:
|
||
sgn = 1
|
||
|
||
|
||
br = self.path.boundingRect()
|
||
x1 = (D*dy + sgn*dx*disc2) * idr2
|
||
y1 = (-D*dx + abs(dy)*disc2) * idr2
|
||
if br.contains(x1+r, y1):
|
||
pt = Point(x1, y1)
|
||
else:
|
||
x2 = (D*dy - sgn*dx*disc2) * idr2
|
||
y2 = (-D*dx - abs(dy)*disc2) * idr2
|
||
pt = Point(x2, y2)
|
||
if not br.contains(x2+r, y2):
|
||
return None, None
|
||
raise Exception("No intersection!")
|
||
|
||
norm = np.arctan2(pt[1], pt[0])
|
||
if r < 0:
|
||
norm += np.pi
|
||
#print " norm:", norm*180/3.1415
|
||
dp = p - pt
|
||
#print " dp:", dp
|
||
ang = np.arctan2(dp[1], dp[0])
|
||
#print " ang:", ang*180/3.1415
|
||
#print " ai:", (ang-norm)*180/3.1415
|
||
|
||
#print " intersection:", pt
|
||
return pt + Point(r, 0), ang-norm
|
||
|
||
|
||
class Ray(pg.GraphicsObject, ParamObj):
|
||
"""Represents a single straight segment of a ray"""
|
||
|
||
sigStateChanged = QtCore.Signal()
|
||
|
||
def __init__(self, **params):
|
||
ParamObj.__init__(self)
|
||
defaults = {
|
||
'ior': 1.0,
|
||
'wl': 500,
|
||
'end': None,
|
||
'dir': Point(1,0),
|
||
}
|
||
self.params = {}
|
||
pg.GraphicsObject.__init__(self)
|
||
self.children = []
|
||
parent = params.get('parent', None)
|
||
if parent is not None:
|
||
defaults['start'] = parent['end']
|
||
defaults['wl'] = parent['wl']
|
||
self['ior'] = parent['ior']
|
||
self['dir'] = parent['dir']
|
||
parent.addChild(self)
|
||
|
||
defaults.update(params)
|
||
defaults['dir'] = Point(defaults['dir'])
|
||
self.setParams(**defaults)
|
||
self.mkPath()
|
||
|
||
def clearChildren(self):
|
||
for c in self.children:
|
||
c.clearChildren()
|
||
c.setParentItem(None)
|
||
self.scene().removeItem(c)
|
||
self.children = []
|
||
|
||
def paramStateChanged(self):
|
||
pass
|
||
|
||
def addChild(self, ch):
|
||
self.children.append(ch)
|
||
ch.setParentItem(self)
|
||
|
||
def currentState(self, relativeTo=None):
|
||
pos = self['start']
|
||
dir = self['dir']
|
||
if relativeTo is None:
|
||
return pos, dir
|
||
else:
|
||
trans = self.itemTransform(relativeTo)[0]
|
||
p1 = trans.map(pos)
|
||
p2 = trans.map(pos + dir)
|
||
return Point(p1), Point(p2-p1)
|
||
|
||
def setEnd(self, end):
|
||
self['end'] = end
|
||
self.mkPath()
|
||
|
||
def boundingRect(self):
|
||
return self.path.boundingRect()
|
||
|
||
def paint(self, p, *args):
|
||
#p.setPen(pg.mkPen((255,0,0, 150)))
|
||
p.setRenderHints(p.renderHints() | p.Antialiasing)
|
||
p.setCompositionMode(p.CompositionMode_Plus)
|
||
p.setPen(wlPen(self['wl']))
|
||
p.drawPath(self.path)
|
||
|
||
def mkPath(self):
|
||
self.prepareGeometryChange()
|
||
self.path = QtGui.QPainterPath()
|
||
self.path.moveTo(self['start'])
|
||
if self['end'] is not None:
|
||
self.path.lineTo(self['end'])
|
||
else:
|
||
self.path.lineTo(self['start']+500*self['dir'])
|
||
|
||
|
||
def trace(rays, optics):
|
||
if len(optics) < 1 or len(rays) < 1:
|
||
return
|
||
for r in rays:
|
||
r.clearChildren()
|
||
o = optics[0]
|
||
r2 = o.propagateRay(r)
|
||
trace(r2, optics[1:])
|
||
|
||
|
||
class Tracer(QtCore.QObject):
|
||
"""
|
||
Simple ray tracer.
|
||
|
||
Initialize with a list of rays and optics;
|
||
calling trace() will cause rays to be extended by propagating them through
|
||
each optic in sequence.
|
||
"""
|
||
def __init__(self, rays, optics):
|
||
QtCore.QObject.__init__(self)
|
||
self.optics = optics
|
||
self.rays = rays
|
||
for o in self.optics:
|
||
o.sigStateChanged.connect(self.trace)
|
||
self.trace()
|
||
|
||
def trace(self):
|
||
trace(self.rays, self.optics)
|
||
|