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 cell from interp.tools import exact_func, exact_func_3D THREE_NODE_TRIANGLE = 2 FOUR_NODE_TET = 4 EDGES_FOR_FACE_CONNECTIVITY = 2 EDGES_FOR_VOLUME_CONNECTIVITY = 3 class gmsh_grid(grid): 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 cell 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_cell_index, node_type, rest = (int(cur_line[0]), int(cur_line[1]), [int(j) for j in cur_line[2:]]) if(node_type == THREE_NODE_TRIANGLE): points_for_cur_cell = [i-1 for i in rest[rest[0]+1:]] cur_cell = cell(cur_cell_index) for cur_point in points_for_cur_cell: self.cells_for_vert[cur_point].append(cur_cell) cur_cell.verts = points_for_cur_cell self.cells[cur_cell_index] = cur_cell edges = [tuple(sorted(i)) for i in combinations(points_for_cur_cell, EDGES_FOR_FACE_CONNECTIVITY)] # edge is two verts for edge in edges: if edge in neighbors: neighbors[edge].append(cur_cell_index) else: neighbors[edge] = [cur_cell_index] for k,v in neighbors.iteritems(): if len(v) > 1: self.cells[v[0]].add_neighbor(self.cells[v[1]]) self.cells[v[1]].add_neighbor(self.cells[v[0]]) def dump_to_blender_files(self, pfile = '/tmp/points.p', ffile = '/tmp/cells.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.cells.itervalues()], open(ffile, 'w')) class gmsh_grid3D(grid): def __init__(self, filename): """ construct an interp.grid.grid-compliant grid object out of a 3D 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_3D(self.verts[i]) grid.__init__(self) self.tree = KDTree(self.verts) gmsh_file.readline() # $EndNodes gmsh_file.readline() # $Elements # temporary dict used to compute cell 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_cell_index, node_type, rest = (int(cur_line[0]), int(cur_line[1]), [int(j) for j in cur_line[2:]]) if(node_type == FOUR_NODE_TET): points_for_cur_cell = [i-1 for i in rest[rest[0]+1:]] cur_cell = cell(cur_cell_index) for cur_point in points_for_cur_cell: self.cells_for_vert[cur_point].append(cur_cell) cur_cell.verts = points_for_cur_cell self.cells[cur_cell_index] = cur_cell edges = [tuple(sorted(i)) for i in combinations(points_for_cur_cell, EDGES_FOR_VOLUME_CONNECTIVITY)] for edge in edges: if edge in neighbors: neighbors[edge].append(cur_cell_index) else: neighbors[edge] = [cur_cell_index] for k,v in neighbors.iteritems(): if len(v) > 1: self.cells[v[0]].add_neighbor(self.cells[v[1]]) self.cells[v[1]].add_neighbor(self.cells[v[0]]) def dump_to_blender_files(self, pfile = '/tmp/points.p', ffile = '/tmp/cells.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.cells.itervalues()], open(ffile, 'w'))