Script to be used with Trelis and/or Cubit to export to VTK file formats, understandable by deal.II

exportVTK.py

#!python
# ---------------------------------------------------------------------
#
# Copyright (C) 2018 by the deal.II authors
#
# This file is part of the deal.II library.
#
# The deal.II library is free software; you can use it, redistribute
# it, and/or modify it under the terms of the GNU Lesser General
# Public License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
# The full text of the license can be found in the file LICENSE.md at
# the top level directory of deal.II.
#
# ---------------------------------------------------------------------
import cubit


def exportVTK(vtkfile):
    """
    This script will output whatever mesh you have currently in CUBIT
    in the VTK ascii format (as an UNSTRUCTURED_GRID type)
    In this script:
    - block id   => material id
    - sideset id => boundary id
    - nodeset id => manifold id
    All ids are assumed to be assigned to Geometrical objects, i.e., 
    Volumes, Surfaces, and Curves. If you assign block, sideset, or 
    nodeset ids to mesh elements, for example, your milage may vary.
    Different versions of Cubit/Trelis will have different behaviour.
    @author: Luca Heltai <luca.heltai@sissa.it>
    """
    outfile = open(vtkfile, "w")
    # ============================================================
    # Global dictionaries
    # ============================================================
    node_ids = {}
    #
    hex_material_ids = {}
    hex_manifold_ids = {}
    #
    quad_material_ids = {}
    quad_manifold_ids = {}
    #
    edge_material_ids = {}
    edge_manifold_ids = {}
    #
    n_nodes = 0
    n_hexes = 0
    n_quads = 0
    n_edges = 0
    #
    block_ids = cubit.get_block_id_list()
    nodeset_ids = cubit.get_nodeset_id_list()
    sideset_ids = cubit.get_sideset_id_list()
    #
    # ============================================================
    # Collect nodes
    # ============================================================
    group_id = cubit.get_id_from_name("temp_nodes")
    if group_id != 0:
        cubit.silent_cmd("delete group " + str(group_id))
    cubit.silent_cmd("group 'temp_nodes' add node all")
    group_id = cubit.get_id_from_name("temp_nodes")
    node_list = cubit.get_group_nodes(group_id)
    cubit.silent_cmd("delete group " + str(group_id))
    n_nodes = len(node_list)
    #
    i = 0
    for node in node_list:
        node_ids[node] = i
        i += 1
    #
    # ============================================================
    # Collect all hexes for 3d grids
    # set default material and manifold ids
    # ============================================================
    group_id = cubit.get_id_from_name("temp_hexes")
    if group_id != 0:
        cubit.silent_cmd("delete group " + str(group_id))
    cubit.silent_cmd("group 'temp_hexes' add hex all")
    group_id = cubit.get_id_from_name("temp_hexes")
    hex_list = cubit.get_group_hexes(group_id)
    cubit.silent_cmd("delete group " + str(group_id))
    n_hexes = len(hex_list)
    #
    for hex in hex_list:
        hex_material_ids[hex] = 0
        hex_manifold_ids[hex] = -1
    #
    # ============================================================
    # Collect hexes by block id ( => Material ID)
    # ============================================================
    for block_id in block_ids:
        volumes = cubit.get_block_volumes(block_id)
        for vol in volumes:
            hexes = cubit.get_volume_hexes(vol)
            print len(hexes), 'hexes in volume', vol, 'with material id', block_id
            for hex in hexes:
                hex_material_ids[hex] = block_id
    #
    # ============================================================
    # Collect hexes by nodeset id ( => Manifold ID)
    # ============================================================
    for nodeset_id in nodeset_ids:
        volumes = cubit.get_nodeset_volumes(nodeset_id)
        for vol in volumes:
            hexes = cubit.get_volume_hexes(vol)
            print len(hexes), 'hexes in volume', vol, 'with manifold id', nodeset_id
            for hex in hexes:
                hex_manifold_ids[hex] = nodeset_id
    #
    # ============================================================
    # Collect all quads (for 2d grids)
    # set default material and manifold ids
    # ============================================================
    if n_hexes == 0:
        surface_list = ()
        group_id = cubit.get_id_from_name("temp_surfs")
        if group_id != 0:
            cubit.silent_cmd("delete group " + str(group_id))
        cubit.silent_cmd("group 'temp_surfs' add surf all")
        group_id = cubit.get_id_from_name("temp_surfs")
        surface_list = cubit.get_group_surfaces(group_id)
        cubit.silent_cmd("delete group " + str(group_id))
        for surface_id in surface_list:
            quads = cubit.get_surface_quads(surface_id)
            for quad in quads:
                quad_material_ids[quad] = 0
                quad_manifold_ids[quad] = -1
    #
    n_quads = len(quad_material_ids)
    #
    # ============================================================
    # Collect quads by block id for 2D grids ( => Material ID)
    # and by sideset_id for 3d grids ( => Material ID == Boundary ID)
    # ============================================================
    if n_hexes == 0:
        for block_id in block_ids:
            surfaces = cubit.get_block_surfaces(block_id)
            for surface_id in surfaces:
                quads = cubit.get_surface_quads(surface_id)
                print len(quads), 'quads in surface', surface_id, 'with material id', block_id
                for quad in quads:
                    quad_material_ids[quad] = block_id
    else:
        for sideset_id in sideset_ids:
            surfaces = cubit.get_sideset_surfaces(sideset_id)
            for surface_id in surfaces:
                quads = cubit.get_surface_quads(surface_id)
                print len(quads), 'quads in surface', surface_id, 'with boundary id', sideset_id
                for quad in quads:
                    quad_material_ids[quad] = sideset_id
                    quad_manifold_ids[quad] = -1
    #
    # ============================================================
    # Collect quads by nodeset id ( => Manifold ID)
    # ============================================================
    for nodeset_id in nodeset_ids:
        surfaces = cubit.get_nodeset_surfaces(nodeset_id)
        for surface_id in surfaces:
            quads = cubit.get_surface_quads(surface_id)
            print len(quads), 'quads in surface', surface_id, 'with manifold id', nodeset_id
            for quad in quads:
                quad_manifold_ids[quad] = nodeset_id
                if quad not in quad_material_ids:
                    # this can only happen in 3d, so this must be an internal
                    # face
                    quad_material_ids[quad] = -1
    #
    n_quads = len(quad_material_ids)
    #
    # ============================================================
    # Collect edges by sideset id ( => Material ID)
    # ============================================================
    for sideset_id in sideset_ids:
        curves = cubit.get_sideset_curves(sideset_id)
        surfaces = cubit.get_sideset_surfaces(sideset_id)
        group_id = cubit.get_id_from_name("temp_edges")
        if group_id != 0:
            cubit.silent_cmd("delete group " + str(group_id))
        for curve in curves:
            cubit.silent_cmd(
                "group 'temp_edges' add edge all in curve " + str(curve))
        for surface in surfaces:
            cubit.silent_cmd(
                "group 'temp_edges' add edge all in surface " + str(surface))
        group_id = cubit.get_id_from_name("temp_edges")
        edges = cubit.get_group_edges(group_id)
        cubit.silent_cmd("delete group " + str(group_id))
        print len(edges), 'edges with boundary id', sideset_id
        for edge in edges:
            edge_material_ids[edge] = sideset_id
            edge_manifold_ids[edge] = -1
    #
    # ============================================================
    # Collect edges by nodeset id ( => Manifold ID)
    # ============================================================
    for nodeset_id in nodeset_ids:
        curves = cubit.get_nodeset_curves(nodeset_id)
        surfaces = cubit.get_nodeset_surfaces(nodeset_id)
        volumes = cubit.get_nodeset_volumes(nodeset_id)
        group_id = cubit.get_id_from_name("temp_edges")
        if group_id != 0:
            cubit.silent_cmd("delete group " + str(group_id))
        for curve in curves:
            cubit.silent_cmd(
                "group 'temp_edges' add edge all in curve " +
                str(curve))
        for surface in surfaces:
            cubit.silent_cmd(
                "group 'temp_edges' add edge all in surface " +
                str(surface))
        for volume in volumes:
            cubit.silent_cmd(
                "group 'temp_edges' add edge all in volume " +
                str(volume))
        group_id = cubit.get_id_from_name("temp_edges")
        edges = cubit.get_group_edges(group_id)
        cubit.silent_cmd("delete group " + str(group_id))
        print len(edges), 'edges with manifold', nodeset_id
        for edge in edges:
            edge_manifold_ids[edge] = nodeset_id
            if edge not in edge_material_ids:
                edge_material_ids[edge] = -1
    #
    n_edges = len(edge_material_ids)
    #
    # ============================================================
    # Now we write the header.
    # ============================================================
    outfile.write("# vtk DataFile Version 3.0\n")
    outfile.write("File generated with Trelis\nASCII\n")
    outfile.write("DATASET UNSTRUCTURED_GRID\nPOINTS " +
                  str(n_nodes) + " double\n")
    #
    n_elements = n_hexes + n_quads + n_edges
    #
    # ============================================================
    # The node list.
    # ============================================================
    for node_num in node_list:
        node_coord = cubit.get_nodal_coordinates(node_num)
        if abs(node_coord[0]) < 1e-15:
            outfile.write("0 ")
        else:
            outfile.write(str(node_coord[0]) + " ")
        if abs(node_coord[1]) < 1e-15:
            outfile.write("0 ")
        else:
            outfile.write(str(node_coord[1]) + " ")
        if abs(node_coord[2]) < 1e-15:
            outfile.write("0")
        else:
            outfile.write(str(node_coord[2]))
        outfile.write("\n")
    #
    # ============================================================
    # Now we write the for the CELLS.
    # ============================================================
    storage = n_hexes*9 + n_quads*5 + n_edges*3
    outfile.write("CELLS " + str(n_elements) + " " + str(storage) + "\n")
    #
    # ============================================================
    # The hex list
    # ============================================================
    for hex in hex_manifold_ids:
        outfile.write("8")
        n = cubit.get_connectivity("Hex", hex)
        for i in range(8):
            outfile.write(" " + str(node_ids[n[i]]))
        outfile.write("\n")
    #
    #
    # ============================================================
    # The quad list
    # ============================================================
    for quad in quad_manifold_ids:
        outfile.write("4")
        n = cubit.get_connectivity("Quad", quad)
        for i in range(4):
            outfile.write(" " + str(node_ids[n[i]]))
        outfile.write("\n")
    #
    # ============================================================
    # The edge list
    # ============================================================
    for edge in edge_manifold_ids:
        outfile.write("2")
        n = cubit.get_connectivity("Edge", edge)
        for i in range(2):
            outfile.write(" " + str(node_ids[n[i]]))
        outfile.write("\n")
    #
    # ============================================================
    # The CELL_TYPES section
    # ============================================================
    cell_types = "12 " * n_hexes + "9 "*n_quads + "3 "*n_edges+"\n"
    outfile.write("CELL_TYPES " + str(n_elements) + "\n" + cell_types)
    #
    # ============================================================
    # The DATA_SET set
    # ============================================================
    outfile.write("CELL_DATA " + str(n_elements) + "\n")
    outfile.write("SCALARS MaterialID int 1\nLOOKUP_TABLE default\n")
    for hex in hex_material_ids:
        outfile.write(str(hex_material_ids[hex]) + " ")
    for quad in quad_material_ids:
        outfile.write(str(quad_material_ids[quad]) + " ")
    for edge in edge_material_ids:
        outfile.write(str(edge_material_ids[edge]) + " ")
    outfile.write("\n")
    #
    outfile.write("SCALARS ManifoldID int 1\nLOOKUP_TABLE default\n")
    for hex in hex_manifold_ids:
        outfile.write(str(hex_manifold_ids[hex]) + " ")
    for quad in quad_manifold_ids:
        outfile.write(str(quad_manifold_ids[quad]) + " ")
    for edge in edge_manifold_ids:
        outfile.write(str(edge_manifold_ids[edge]) + " ")
    outfile.write("\n")
    #
    print n_nodes, "nodes"
    print n_hexes, "hexes"
    print n_quads, "quads"
    print n_edges, "edges"


print 'use as exportVTK("/tmp/output.vtk")'