smbinterp/interp/grid/gmsh.py

91 lines
2.6 KiB
Python
Raw Normal View History

2010-10-23 15:57:44 -07:00
import pickle
from itertools import combinations
from collections import defaultdict
import numpy as np
from scipy.spatial import KDTree
from interp.grid import grid
from interp.grid.simplex import face
from interp.tools import exact_func
class gmsh_grid(grid):
THREE_NODE_TRIANGLE = 2
def __init__(self, filename):
"""
construct an interp.grid.grid-compliant grid
object out of a 2D gmsh file
"""
gmsh_file = open(filename, 'r')
gmsh_file.readline() # $MeshFormat
format = gmsh_file.readline()
gmsh_file.readline() # $EndMeshFormat
gmsh_file.readline() # $Nodes
node_count = int(gmsh_file.readline())
self.verts = np.empty((node_count, 3))
self.q = np.empty(node_count)
for i in xrange(node_count):
cur_line = gmsh_file.readline()
(index, x,y,z) = cur_line.split()
index = int(index) - 1
self.verts[i][0] = float(x)
self.verts[i][1] = float(y)
self.verts[i][2] = float(z)
self.q[i] = exact_func(self.verts[i])
grid.__init__(self)
self.tree = KDTree(self.verts)
2010-10-23 15:57:44 -07:00
gmsh_file.readline() # $EndNodes
gmsh_file.readline() # $Elements
# temporary dict used to compute face connectivity
neighbors = {}
element_count = int(gmsh_file.readline())
for i in xrange(element_count):
cur_line = gmsh_file.readline()
cur_line = cur_line.split()
cur_face_index, node_type, rest = (int(cur_line[0]),
int(cur_line[1]),
[int(j) for j in cur_line[2:]])
if(node_type == gmsh_grid.THREE_NODE_TRIANGLE):
points_for_cur_face = [i-1 for i in rest[rest[0]+1:]]
cur_face = face(cur_face_index)
for cur_point in points_for_cur_face:
self.faces_for_vert[cur_point].append(cur_face)
cur_face.verts = points_for_cur_face
self.faces[cur_face_index] = cur_face
edges = [tuple(sorted(i)) for i in combinations(points_for_cur_face, 2)]
# edge is two verts
for edge in edges:
if edge in neighbors:
neighbors[edge].append(cur_face_index)
else:
neighbors[edge] = [cur_face_index]
for k,v in neighbors.iteritems():
if len(v) > 1:
self.faces[v[0]].add_neighbor(self.faces[v[1]])
self.faces[v[1]].add_neighbor(self.faces[v[0]])
def dump_to_blender_files(self, pfile = '/tmp/points.p', ffile = '/tmp/faces.p'):
pickle.dump([(p[0], p[1], p[2]) for p in self.verts], open(pfile, 'w'))
pickle.dump([f.verts for f in self.faces.itervalues()], open(ffile, 'w'))