b0a3849960
Using numpy methods that are intended for vectorized operations is substantially slower than using the math module, so when feasible the math module methods should be used.
560 lines
17 KiB
Python
560 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
from math import atan2, asin, sin, cos, sqrt, pi
|
||
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 = 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.setPos(0, 0)
|
||
tr = QtGui.QTransform()
|
||
self.setTransform(tr.translate(Point(self['pos'])).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.setRotation(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.setRotation(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 = atan2(rd[1], rd[0])
|
||
try:
|
||
ar = a1 - ai + asin((sin(ai) * ray['ior'] / ior))
|
||
except ValueError:
|
||
ar = np.nan
|
||
ray.setEnd(p1)
|
||
dp = Point(cos(ar), 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 = atan2(rd[1], rd[0])
|
||
ar = a1 + pi - 2*ai
|
||
ray.setEnd(p1)
|
||
dp = Point(cos(ar), 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 = asin(h2/r) * 180. / 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), atan2(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
|
||
|
||
norm = atan2(pt[1], pt[0])
|
||
if r < 0:
|
||
norm += pi
|
||
dp = p - pt
|
||
ang = atan2(dp[1], dp[0])
|
||
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)
|
||
|