smbinterp/bin/parse_gmsh.py

#!/usr/bin/env python
import atexit
import os
import sys
import readline
import rlcompleter

historyPath = os.path.expanduser("~/.pyhistory")

def save_history(historyPath=historyPath):
  import readline
  readline.write_history_file(historyPath)

if os.path.exists(historyPath):
  readline.read_history_file(historyPath)

atexit.register(save_history)
del os, atexit, readline, rlcompleter, save_history, historyPath


import sys
import pickle
from itertools import combinations
from collections import defaultdict
from scipy.spatial import KDTree

from grid.simplex import face

import numpy as np

THREE_NODE_TRIANGLE = 2

if __name__ == '__main__':
  if len(sys.argv) != 2:
    print >> sys.stderr, "usage: %s <gmsh file>" % sys.argv[0]
    sys.exit(1)

  gmsh_file = open(sys.argv[1], 'r')


  gmsh_file.readline() # $MeshFormat
  format = gmsh_file.readline()
  gmsh_file.readline() # $EndMeshFormat

  gmsh_file.readline() # $Nodes

  node_count = int(gmsh_file.readline())

  verts = np.zeros((node_count, 3))
  for i in xrange(node_count):
    cur_line = gmsh_file.readline()
    (index, x,y,z) = cur_line.split()
    index = int(index) - 1
    x     = float(x)
    y     = float(y)
    z     = float(z)

    verts[i][0] = x
    verts[i][1] = y
    verts[i][2] = z

  tree = KDTree(verts)

  gmsh_file.readline() # $E   ndNodes
  gmsh_file.readline() # $E   lements

  element_count = int(gmsh_file.readline())

  faces     = {}
  neighbors = {}
  faces_for_vert = defaultdict(list)

  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 == 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:
        faces_for_vert[cur_point].append(cur_face_index)

      cur_face.verts       = points_for_cur_face

      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:
      faces[v[0]].add_neighbor(faces[v[1]])
      faces[v[1]].add_neighbor(faces[v[0]])

  pickle.dump([(p[0], p[1], p[2]) for p in verts], open('/tmp/points.p', 'w'))
  pickle.dump([f.verts for f in faces.itervalues()], open('/tmp/faces.p',  'w'))