smbinterp/interp/grid/gmsh.py
2010-10-29 12:41:39 -06:00

91 lines
2.6 KiB
Python

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)
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'))