ASCEE/pyqtgraph
5
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Added new demos:

Browse Source 
- Relativity simulator
- Optics simulator
- Mechanical chain simulator
This commit is contained in:
Luke Campagnola 2014-05-22 01:22:12 -04:00
parent c3fdcc9ae5
commit 0524bfa6e8
14 changed files with 1078 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
examples/__main__.py
View File

@ -32,6 +32,11 @@ examples = OrderedDict([
('Auto-range', 'PlotAutoRange.py'), ('Auto-range', 'PlotAutoRange.py'),
('Remote Plotting', 'RemoteSpeedTest.py'), ('Remote Plotting', 'RemoteSpeedTest.py'),
('HDF5 big data', 'hdf5.py'), ('HDF5 big data', 'hdf5.py'),
('Demos', OrderedDict([
('Optics', 'optics_demos.py'),
('Special relativity', 'relativity_demo.py'),
('Verlet chain', 'verlet_chain_demo.py'),
])),
('GraphicsItems', OrderedDict([ ('GraphicsItems', OrderedDict([
('Scatter Plot', 'ScatterPlot.py'), ('Scatter Plot', 'ScatterPlot.py'),
#('PlotItem', 'PlotItem.py'), #('PlotItem', 'PlotItem.py'),

1
examples/optics/__init__.py Normal file
View File

@ -0,0 +1 @@
from pyoptic import *

582
examples/optics/pyoptic.py Normal file
View File

@ -0,0 +1,582 @@
# -*- coding: utf-8 -*-
from PyQt4 import QtGui, QtCore
import pyqtgraph as pg
#from pyqtgraph.canvas import Canvas, CanvasItem
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(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 = map(float, [info['B1'], info['B2'], info['B3']])
C = 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)

BIN
examples/optics/schott_glasses.csv.gz Normal file
View File

Binary file not shown.

170
examples/optics_demos.py Normal file
View File

@ -0,0 +1,170 @@
# -*- coding: utf-8 -*-
"""
Optical system design demo
"""
import initExample ## Add path to library (just for examples; you do not need this)
from optics import *
import pyqtgraph as pg
import numpy as np
from pyqtgraph import Point
app = pg.QtGui.QApplication([])
w = pg.GraphicsWindow(border=0.5)
w.resize(1000, 900)
w.show()
### Curved mirror demo
view = w.addViewBox()
view.setAspectLocked()
#grid = pg.GridItem()
#view.addItem(grid)
view.setRange(pg.QtCore.QRectF(-50, -30, 100, 100))
optics = []
rays = []
m1 = Mirror(r1=-100, pos=(5,0), d=5, angle=-15)
optics.append(m1)
m2 = Mirror(r1=-70, pos=(-40, 30), d=6, angle=180-15)
optics.append(m2)
allRays = []
for y in np.linspace(-10, 10, 21):
r = Ray(start=Point(-100, y))
view.addItem(r)
allRays.append(r)
for o in optics:
view.addItem(o)
t1 = Tracer(allRays, optics)
### Dispersion demo
optics = []
view = w.addViewBox()
view.setAspectLocked()
#grid = pg.GridItem()
#view.addItem(grid)
view.setRange(pg.QtCore.QRectF(-10, -50, 90, 60))
optics = []
rays = []
l1 = Lens(r1=20, r2=20, d=10, angle=8, glass='Corning7980')
optics.append(l1)
allRays = []
for wl in np.linspace(355,1040, 25):
for y in [10]:
r = Ray(start=Point(-100, y), wl=wl)
view.addItem(r)
allRays.append(r)
for o in optics:
view.addItem(o)
t2 = Tracer(allRays, optics)
### Scanning laser microscopy demo
w.nextRow()
view = w.addViewBox(colspan=2)
optics = []
#view.setAspectLocked()
view.setRange(QtCore.QRectF(200, -50, 500, 200))
## Scan mirrors
scanx = 250
scany = 20
m1 = Mirror(dia=4.2, d=0.001, pos=(scanx, 0), angle=315)
m2 = Mirror(dia=8.4, d=0.001, pos=(scanx, scany), angle=135)
## Scan lenses
l3 = Lens(r1=23.0, r2=0, d=5.8, pos=(scanx+50, scany), glass='Corning7980') ## 50mm UVFS (LA4148)
l4 = Lens(r1=0, r2=69.0, d=3.2, pos=(scanx+250, scany), glass='Corning7980') ## 150mm UVFS (LA4874)
## Objective
obj = Lens(r1=15, r2=15, d=10, dia=8, pos=(scanx+400, scany), glass='Corning7980')
IROptics = [m1, m2, l3, l4, obj]
## Scan mirrors
scanx = 250
scany = 30
m1a = Mirror(dia=4.2, d=0.001, pos=(scanx, 2*scany), angle=315)
m2a = Mirror(dia=8.4, d=0.001, pos=(scanx, 3*scany), angle=135)
## Scan lenses
l3a = Lens(r1=46, r2=0, d=3.8, pos=(scanx+50, 3*scany), glass='Corning7980') ## 100mm UVFS (LA4380)
l4a = Lens(r1=0, r2=46, d=3.8, pos=(scanx+250, 3*scany), glass='Corning7980') ## 100mm UVFS (LA4380)
## Objective
obja = Lens(r1=15, r2=15, d=10, dia=8, pos=(scanx+400, 3*scany), glass='Corning7980')
IROptics2 = [m1a, m2a, l3a, l4a, obja]
for o in set(IROptics+IROptics2):
view.addItem(o)
IRRays = []
IRRays2 = []
for dy in [-0.4, -0.15, 0, 0.15, 0.4]:
IRRays.append(Ray(start=Point(-50, dy), dir=(1, 0), wl=780))
IRRays2.append(Ray(start=Point(-50, dy+2*scany), dir=(1, 0), wl=780))
for r in set(IRRays+IRRays2):
view.addItem(r)
IRTracer = Tracer(IRRays, IROptics)
IRTracer2 = Tracer(IRRays2, IROptics2)
phase = 0.0
def update():
global phase
if phase % (8*np.pi) > 4*np.pi:
m1['angle'] = 315 + 1.5*np.sin(phase)
m1a['angle'] = 315 + 1.5*np.sin(phase)
else:
m2['angle'] = 135 + 1.5*np.sin(phase)
m2a['angle'] = 135 + 1.5*np.sin(phase)
phase += 0.2
timer = QtCore.QTimer()
timer.timeout.connect(update)
timer.start(40)
## Start Qt event loop unless running in interactive mode or using pyside.
if __name__ == '__main__':
import sys
if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
QtGui.QApplication.instance().exec_()

1
examples/relativity Submodule

@ -0,0 +1 @@
Subproject commit 876a0a80b705dad71e5b1addab9b859cfc292f20

23
examples/relativity_demo.py Normal file
View File

@ -0,0 +1,23 @@
# -*- coding: utf-8 -*-
"""
Special relativity simulation
"""
import initExample ## Add path to library (just for examples; you do not need this)
import pyqtgraph as pg
from relativity import RelativityGUI
pg.mkQApp()
win = RelativityGUI()
win.setWindowTitle("Relativity!")
win.resize(1100,700)
win.show()
win.loadPreset(None, 'Twin Paradox (grid)')
## Start Qt event loop unless running in interactive mode or using pyside.
if __name__ == '__main__':
import sys
if (sys.flags.interactive != 1) or not hasattr(pg.QtCore, 'PYQT_VERSION'):
pg.QtGui.QApplication.instance().exec_()

1
examples/verlet_chain/__init__.py Normal file
View File

@ -0,0 +1 @@
from chain import ChainSim

110
examples/verlet_chain/chain.py Normal file
View File

@ -0,0 +1,110 @@
import pyqtgraph as pg
import numpy as np
import time
from relax import relax
class ChainSim(pg.QtCore.QObject):
stepped = pg.QtCore.Signal()
relaxed = pg.QtCore.Signal()
def __init__(self):
pg.QtCore.QObject.__init__(self)
self.damping = 0.1 # 0=full damping, 1=no damping
self.relaxPerStep = 10
self.maxTimeStep = 0.01
self.pos = None # (Npts, 2) float
self.mass = None # (Npts) float
self.fixed = None # (Npts) bool
self.links = None # (Nlinks, 2), uint
self.lengths = None # (Nlinks), float
self.push = None # (Nlinks), bool
self.pull = None # (Nlinks), bool
self.initialized = False
self.lasttime = None
self.lastpos = None
def init(self):
if self.initialized:
return
assert None not in [self.pos, self.mass, self.links, self.lengths]
if self.fixed is None:
self.fixed = np.zeros(self.pos.shape[0], dtype=bool)
if self.push is None:
self.push = np.ones(self.links.shape[0], dtype=bool)
if self.pull is None:
self.pull = np.ones(self.links.shape[0], dtype=bool)
# precompute relative masses across links
l1 = self.links[:,0]
l2 = self.links[:,1]
m1 = self.mass[l1]
m2 = self.mass[l2]
self.mrel1 = (m1 / (m1+m2))[:,np.newaxis]
self.mrel1[self.fixed[l1]] = 1 # fixed point constraint
self.mrel1[self.fixed[l2]] = 0
self.mrel2 = 1.0 - self.mrel1
for i in range(100):
self.relax(n=10)
self.initialized = True
def makeGraph(self):
#g1 = pg.GraphItem(pos=self.pos, adj=self.links[self.rope], pen=0.2, symbol=None)
brushes = np.where(self.fixed, pg.mkBrush(0,0,0,255), pg.mkBrush(50,50,200,255))
g2 = pg.GraphItem(pos=self.pos, adj=self.links[self.push & self.pull], pen=0.5, brush=brushes, symbol='o', size=(self.mass**0.33), pxMode=False)
p = pg.ItemGroup()
#p.addItem(g1)
p.addItem(g2)
return p
def update(self):
# approximate physics with verlet integration
now = pg.ptime.time()
if self.lasttime is None:
dt = 0
else:
dt = now - self.lasttime
self.lasttime = now
if self.lastpos is None:
self.lastpos = self.pos
# remember fixed positions
fixedpos = self.pos[self.fixed]
while dt > 0:
dt1 = min(self.maxTimeStep, dt)
dt -= dt1
# compute motion since last timestep
dx = self.pos - self.lastpos
self.lastpos = self.pos
# update positions for gravity and inertia
acc = np.array([[0, -5]]) * dt1
inertia = dx * (self.damping**(dt1/self.mass))[:,np.newaxis] # with mass-dependent damping
self.pos = self.pos + inertia + acc
self.pos[self.fixed] = fixedpos # fixed point constraint
# correct for link constraints
self.relax(self.relaxPerStep)
self.stepped.emit()
def relax(self, n=50):
# speed up with C magic
relax(self.pos, self.links, self.mrel1, self.mrel2, self.lengths, self.push, self.pull, n)
self.relaxed.emit()

3
examples/verlet_chain/make Executable file
View File

@ -0,0 +1,3 @@
gcc -fPIC -c relax.c
gcc -shared -o maths.so relax.o

BIN
examples/verlet_chain/maths.so Executable file
View File

Binary file not shown.

48
examples/verlet_chain/relax.c Normal file
View File

@ -0,0 +1,48 @@
#include <math.h>
#include <stdio.h>
void relax(
double* pos,
long* links,
double* mrel1,
double* mrel2,
double* lengths,
char* push,
char* pull,
int nlinks,
int iters)
{
int i, l, p1, p2;
double x1, x2, y1, y2, dx, dy, dist, change;
// printf("%d, %d\n", iters, nlinks);
for( i=0; i<iters; i++ ) {
for( l=0; l<nlinks; l++ ) {
p1 = 2*links[l*2];
p2 = 2*links[l*2 + 1];
x1 = pos[p1];
y1 = pos[p1 + 1];
x2 = pos[p2];
y2 = pos[p2 + 1];
dx = x2 - x1;
dy = y2 - y1;
// dist = pow(dx*dx + dy*dy, 0.5);
dist = sqrt(dx*dx + dy*dy);
if( push[l]==0 && dist < lengths[l] )
dist = lengths[l];
if( pull[l]==0 && dist > lengths[l] )
dist = lengths[l];
change = (lengths[l]-dist) / dist;
dx *= change;
dy *= change;
pos[p1] -= mrel2[l] * dx;
pos[p1+1] -= mrel2[l] * dy;
pos[p2] += mrel1[l] * dx;
pos[p2+1] += mrel1[l] * dy;
}
}
}

23
examples/verlet_chain/relax.py Normal file
View File

@ -0,0 +1,23 @@
import ctypes
import os
so = os.path.join(os.path.dirname(__file__), 'maths.so')
lib = ctypes.CDLL(so)
lib.relax.argtypes = [
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_int,
ctypes.c_int,
]
def relax(pos, links, mrel1, mrel2, lengths, push, pull, iters):
nlinks = links.shape[0]
lib.relax(pos.ctypes, links.ctypes, mrel1.ctypes, mrel2.ctypes, lengths.ctypes, push.ctypes, pull.ctypes, nlinks, iters)

111
examples/verlet_chain_demo.py Normal file
View File

@ -0,0 +1,111 @@
"""
Mechanical simulation of a chain using verlet integration.
"""
import initExample ## Add path to library (just for examples; you do not need this)
import pyqtgraph as pg
from pyqtgraph.Qt import QtCore, QtGui
import numpy as np
from verlet_chain import ChainSim
sim = ChainSim()
chlen1 = 80
chlen2 = 60
npts = chlen1 + chlen2
sim.mass = np.ones(npts)
sim.mass[chlen1-15] = 100
sim.mass[chlen1-1] = 500
sim.mass[npts-1] = 200
sim.fixed = np.zeros(npts, dtype=bool)
sim.fixed[0] = True
sim.fixed[chlen1] = True
sim.pos = np.empty((npts, 2))
sim.pos[:chlen1, 0] = 0
sim.pos[chlen1:, 0] = 10
sim.pos[:chlen1, 1] = np.arange(chlen1)
sim.pos[chlen1:, 1] = np.arange(chlen2)
links1 = [(j, i+j+1) for i in range(chlen1) for j in range(chlen1-i-1)]
links2 = [(j, i+j+1) for i in range(chlen2) for j in range(chlen2-i-1)]
sim.links = np.concatenate([np.array(links1), np.array(links2)+chlen1, np.array([[chlen1-1, npts-1]])])
p1 = sim.pos[sim.links[:,0]]
p2 = sim.pos[sim.links[:,1]]
dif = p2-p1
sim.lengths = (dif**2).sum(axis=1) ** 0.5
sim.lengths[(chlen1-1):len(links1)] *= 1.05 # let auxiliary links stretch a little
sim.lengths[(len(links1)+chlen2-1):] *= 1.05
sim.lengths[-1] = 7
push1 = np.ones(len(links1), dtype=bool)
push1[chlen1:] = False
push2 = np.ones(len(links2), dtype=bool)
push2[chlen2:] = False
sim.push = np.concatenate([push1, push2, np.array([True], dtype=bool)])
sim.pull = np.ones(sim.links.shape[0], dtype=bool)
sim.pull[-1] = False
mousepos = sim.pos[0]
def display():
global view, sim
view.clear()
view.addItem(sim.makeGraph())
def relaxed():
global app
display()
app.processEvents()
def mouse(pos):
global mousepos
pos = view.mapSceneToView(pos)
mousepos = np.array([pos.x(), pos.y()])
def update():
global mousepos
#sim.pos[0] = sim.pos[0] * 0.9 + mousepos * 0.1
s = 0.9
sim.pos[0] = sim.pos[0] * s + mousepos * (1.0-s)
sim.update()
app = pg.mkQApp()
win = pg.GraphicsLayoutWidget()
win.show()
view = win.addViewBox()
view.setAspectLocked(True)
view.setXRange(-100, 100)
#view.autoRange()
view.scene().sigMouseMoved.connect(mouse)
#display()
#app.processEvents()
sim.relaxed.connect(relaxed)
sim.init()
sim.relaxed.disconnect(relaxed)
sim.stepped.connect(display)
timer = pg.QtCore.QTimer()
timer.timeout.connect(update)
timer.start(16)
## Start Qt event loop unless running in interactive mode or using pyside.
if __name__ == '__main__':
import sys
if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
QtGui.QApplication.instance().exec_()