pyqtgraph/examples/optics/pyoptic.py

# -*- 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:
    # 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):
        return self.getParam(item)

    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)
            #print "surface intersection:", p1, ai*180/3.14159
            #trans = self.sceneTransform().inverted()[0] * surface.sceneTransform()
            #p1 = trans.map(p1)
            if p1 is None:
                ray.setEnd(None)
                break
            p1 = surface.mapToItem(ray, p1)
            
            #print "adjusted position:", p1
            #ior = self.ior(ray['wl'])
            rd = ray['dir']
            a1 = np.arctan2(rd[1], rd[0])
            ar = a1 - ai + np.arcsin((np.sin(ai) * ray['ior'] / ior))
            #print [x for x in [a1, ai, (np.sin(ai) * ray['ior'] / ior), ar]]
            #print ai, np.sin(ai), ray['ior'],  ior
            ray.setEnd(p1)
            dp = Point(np.cos(ar), np.sin(ar))
            #p2 = p1+dp
            #p1p = self.mapToScene(p1)
            #p2p = self.mapToScene(p2)
            #dpp = Point(p2p-p1p)
            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))
            
            #dx = abs(r) - (abs(r)**2 - abs(h2)**2)**0.5
            #p.moveTo(-d*w/2.+ d*dx, d*h2)
            arc = QtCore.QRectF(0, -r, r*2, r*2)
            #self.surfaces.append((arc.center(), r, h2))
            a1 = np.arcsin(h2/r) * 180. / np.pi
            a2 = -2*a1
            a1 += 180.
            self.path.arcMoveTo(arc, a1)
            self.path.arcTo(arc, a1, a2)
            #if d == -1:
                #p1 = QtGui.QPainterPath()
                #p1.addRect(arc)
                #self.paths.append(p1)
        self.h2 = h2
        
    def boundingRect(self):
        return self.path.boundingRect()
        
    def paint(self, p, *args):
        return  ## usually we let the optic draw.
        #p.setPen(pg.mkPen('r'))
        #p.drawPath(self.path)
            
    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)