smbinterp/interp/grid/gmsh.py
2011-02-20 19:23:01 -07:00

180 lines
5.2 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 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())
# for dim = 2, see note in next for loop
self.verts = np.empty((node_count, 2))
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)
# for the general baker method to work, it must have 2
# components if it it a 2D mesh, so I removed:
# self.verts[i][2] = float(z)
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'))