1af176a6e0
also did some pep8/pyflakes cleanup
122 lines
4.0 KiB
Python
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
|