File: GlyphExtractor.py

package info (click to toggle)
esys-particle 2.1-4
  • links: PTS, VCS
  • area: main
  • in suites: wheezy
  • size: 7,284 kB
  • sloc: cpp: 77,304; python: 5,647; makefile: 1,176; sh: 10
file content (79 lines) | stat: -rw-r--r-- 2,861 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#############################################################
##                                                         ##
## Copyright (c) 2003-2011 by The University of Queensland ##
## Earth Systems Science Computational Centre (ESSCC)      ##
## http://www.uq.edu.au/esscc                              ##
##                                                         ##
## Primary Business: Brisbane, Queensland, Australia       ##
## Licensed under the Open Software License version 3.0    ##
## http://www.opensource.org/licenses/osl-3.0.php          ##
##                                                         ##
#############################################################

"""
"""
import types
import vtk as kwvtk

from esys.lsm.vis import core

typeVtkDataArrayDict = dict()
typeVtkDataArrayDict[types.IntType]   = kwvtk.vtkIntArray
typeVtkDataArrayDict[types.FloatType] = kwvtk.vtkDoubleArray
typeVtkDataArrayDict[types.LongType]  = kwvtk.vtkLongArray

def getVtkDataArrayClass(thing):
    if (hasattr(thing, "__len__")):
        return typeVtkDataArrayDict[type(thing[0])]
    return typeVtkDataArrayDict[type(thing)]

class FieldMap:
    def __init__(self, name, map):
        self.name = name
        self.map  = map
        self.vtkDataArray = None
        self.componentIdxList = []

    def initialise(self, record):
        val = self.map(record)
        self.vtkDataArray = getVtkDataArrayClass(val)()
        self.vtkDataArray.SetName(self.name)
        self.vtkDataArray.SetNumberOfComponents(getattr(val, "__len__", lambda:1)())
        self.componentIdxList = range(0,self.vtkDataArray.GetNumberOfComponents())

    def insertNext(self, record):
        tupleIdx = self.vtkDataArray.GetNumberOfTuples()
        tupyl = self.map(record)
        for j in self.componentIdxList:
            self.vtkDataArray.InsertComponent(tupleIdx, j, tupyl[j])

    def getVtkDataArray(self):
        return self.vtkDataArray

class GlyphExtractor:
    def __init__(self, pointsMap, fieldMapList):
        self.pointsMap = pointsMap
        self.fieldMapList = fieldMapList

    def addRecordToVtkData(self, record, vtkPoints):
        vtkPoints.InsertNextPoint(self.pointsMap(record))
        for fieldMap in self.fieldMapList:
            fieldMap.insertNext(record)

    def getVtkData(self, data):
        rIt = iter(data)
        record = rIt.next()
        for fieldMap in self.fieldMapList:
            fieldMap.initialise(record)

        vtkPoints = kwvtk.vtkPoints()
        self.addRecordToVtkData(record, vtkPoints)
        for record in rIt:
            self.addRecordToVtkData(record, vtkPoints)

        vtkUnstrGrid = kwvtk.vtkUnstructuredGrid()
        vtkUnstrGrid.SetPoints(vtkPoints)
        for fieldMap in self.fieldMapList:
            vtkUnstrGrid.GetPointData().AddArray(fieldMap.getVtkDataArray())
        return vtkUnstrGrid