smbinterp/interp/grid/gmsh.py

122 lines
4.0 KiB
Python

from itertools import combinations
import numpy as np
from scipy.spatial import KDTree
from interp.grid import grid, contains
THREE_NODE_TRIANGLE = 2
FOUR_NODE_TET = 4
MAX_SEARCH_COUNT = 256
class ggrid(grid):
def __init__(self, filename, values=None, dimension=3):
"""
construct an interp.grid.grid-compliant grid
object out of a {2,3}D gmsh file
"""
self.dim = dimension
self.values = None
gmsh_file = open(filename, 'r')
gmsh_file.readline() # $MeshFormat
fmt = gmsh_file.readline()
self.version, self.file_type, self.data_size = fmt.split()
gmsh_file.readline() # $EndMeshFormat
gmsh_file.readline() # $Nodes
node_count = int(gmsh_file.readline())
self.points = np.empty((node_count, dimension), dtype=np.float64)
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.points[i][0] = float(x)
self.points[i][1] = float(y)
if self.dim == 3:
self.points[i][2] = float(z)
self.tree = KDTree(self.points)
gmsh_file.readline() # $EndNodes
gmsh_file.readline() # $Elements
neighbors = {}
simplices = []
self.point_to_simplex = [[] for i in xrange(len(self.points))]
simplex_counter = 0
element_count = int(gmsh_file.readline())
for simplex_id in xrange(element_count):
cur_line = gmsh_file.readline()
cur_line = cur_line.split()
simplex_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 and self.dim == 2) \
or (node_type == FOUR_NODE_TET and self.dim == 3):
points_for_simplex = [i - 1 for i in rest[rest[0] + 1:]]
for point in points_for_simplex:
self.point_to_simplex[point].append(simplex_counter)
simplices.append(points_for_simplex)
simplex_counter += 1
self.simplex_to_simplex = [[] for i in xrange(len(simplices))]
for simplex in xrange(len(simplices)):
edges = [tuple(sorted(i)) for i in \
combinations(simplices[simplex], self.dim)]
for edge in edges:
if edge in neighbors:
neighbors[edge].append(simplex)
else:
neighbors[edge] = [simplex]
for k, v in neighbors.iteritems():
if len(v) > 1:
self.simplex_to_simplex[v[0]].append(v[1])
self.simplex_to_simplex[v[1]].append(v[0])
self.simplices = np.array(simplices, dtype=np.int32)
def find_simplex(self, X, max_search_count=MAX_SEARCH_COUNT):
# get closest point
(dist, indicies) = self.tree.query(X, 2)
closest_point = indicies[0]
simplex = None
checked_cells = []
cells_to_check = list(self.point_to_simplex[closest_point])
attempts = 0
while not simplex and cells_to_check:
cur_cell = cells_to_check.pop(0)
checked_cells.append(cur_cell)
R = self.points[self.simplices[cur_cell]]
if contains(X, R):
simplex = cur_cell
continue
attempts += 1
if attempts >= max_search_count:
raise Exception("Is the search becoming exhaustive?"
" (%dth attempt)" % attempts)
for neighbor in self.simplex_to_simplex[cur_cell]:
if (neighbor not in checked_cells) \
and (neighbor not in cells_to_check):
cells_to_check.append(neighbor)
if not simplex:
raise Exception('no containing simplex found')
return simplex