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