PyMOAB script for assigning materials from lcad file

lcad_assign.py

from pymoab import core, types
import numpy as np

from argparse import ArgumentParser


# creates a UW2 formatted material name from
# a material id and density
# returns this as a string
def create_mat_key(mat_id, density):
    
    key_val_separator = ":"
    prop_separator = "/"

    mat_key = "mat"
    density_key = "rho"

    mat_name = mat_key + key_val_separator + str(mat_id)
    mat_name += prop_separator
    mat_name += density_key + key_val_separator + str(density)

    return mat_name

#parses the volume information from a line
# the information is expected to be in the
# following format:
#
# volume_id material_number density importance (optional)
#
# this information is assumed to be space-delimited
#
# returns; material key (string w/ UW2 format), (int) volume_id
def parse_vol_line(line):

    # split line on spaces
    elems = line.split()

    vol_id = int(elems[0])

    mat_id = int(elems[1])

    if not mat_id:
        mat_name = "mat:void"
    else:
        density = float(elems[2])
        mat_name = create_mat_key(mat_id,density)
        
    return mat_name, vol_id


# parses a surface line
# removes special character from surface if necessary
# returns (int) surface_id
def parse_surf_line(line):
    surface_id = int(line.strip().strip("*"))
    return surface_id
    
# parses through an lcad file and creates
# a dictionary of volumes with material assignments as keys
# as well as a Set of surface ids and a Set of volume ids
# returns: material_dict, surface_set, volume_set
def parse_lcad(lcad_filename):

    lcad_file = open(lcad_filename, 'r')

    # containers for output
    material_dict = dict()
    surface_set = set()
    volume_set = set()

    # let's assume the lcad file isn't crazy big
    # and cache the lines in a list
    lines = [ line.strip("\n") for line in lcad_file ]

    # remove any leading blank lines
    while not lines[0]:
        lines.pop(0)
    
    # we'll start by parsing the volumes
    parsing_vols = True
    
    while len(lines):

        # get this line from list
        line = lines.pop(0)

        # if we hit an empty line, this should
        # be the break between the volume list
        # and the surface list, switch modes
        if not line:
            parsing_vols = False
            continue

        # skip implicit complement
        if "implicit complement" in line:
            continue
        
        # if we're parsing volumes still, parse
        # the volume information from the line
        if parsing_vols:
            mat, vol_id = parse_vol_line(line)
            if mat in material_dict.keys():
                material_dict[mat] += [vol_id,]
            else:
                material_dict[mat] = [vol_id,]

            volume_set.add(vol_id)
            
        # if we're parsing surfaces, parse the
        # surface from the line
        else:
            surf_id = parse_surf_line(line)
            surface_set.add(surf_id)

    # when done, return accumulated objects
    return material_dict, volume_set, surface_set


def update_mat_assignments(dagmc_model, lcad_file):

    # parse the lcad file (assume argument 1 is the filename)
    print "Parsing lcad file..."
    material_dict, lcad_vol_set, lcad_surf_set = parse_lcad(lcad_file)

    # start a pymoab instance
    mb = core.Core()

    print "Loading DAGMC model..."
    # load the dagmc model
    mb.load_file(dagmc_model)

    print "Updating material assignments..."
    # gather tags
    cat_tag = mb.tag_get_handle("CATEGORY")
    gid_tag = mb.tag_get_handle("GLOBAL_ID")
    geom_dim_tag = mb.tag_get_handle("GEOM_DIMENSION")
    name_tag = mb.tag_get_handle("NAME")

    vols  = mb.get_entities_by_type_and_tag(0, types.MBENTITYSET, cat_tag, np.array(["Volume",]))
    surfs = mb.get_entities_by_type_and_tag(0, types.MBENTITYSET, cat_tag, np.array(["Surface",]))

    vol_id_set = set(mb.tag_get_data(gid_tag, vols, flat = True))
    assert vol_id_set == lcad_vol_set, "Volumes in the DAGMC model and lcad file do not match" 

    # would be nice to check that number of surfaces are the same, but
    # the implicit complement throws this off I think
    surf_id_set = set(mb.tag_get_data(gid_tag, surfs, flat = True))
    # best we can do for now is make sure that the MOAB surfaces are a subset
    # of the lcad surfs
    assert surf_id_set <= lcad_surf_set, "Surfaces in the DAGMC model are not a subset of those in the lcad file" 
    
    groups = mb.get_entities_by_type_and_tag(0, types.MBENTITYSET, cat_tag, np.array(["Group",]))

    mat_groups = [ group for group in groups if "mat" in mb.tag_get_data(name_tag, group, flat = True)[0] ]

    # get rid of the existing material groups
    mb.delete_entities(mat_groups)

    assert( len(set(material_dict.keys())) == len(material_dict.keys()) )
    
    # create new mat groups for each mat found in the lcad file
    for mat in material_dict.keys():

        mat_group = mb.create_meshset()

        # tag group as group
        mb.tag_set_data(cat_tag, mat_group, np.array(["Group",]))
        # tag group name as key
        mb.tag_set_data(name_tag, mat_group, np.array([mat,]))

        group_vols = [ vol for vol in vols if mb.tag_get_data(gid_tag, vol, flat = True)[0] in material_dict[mat] ]

        assert(len(group_vols))
        
        mb.add_entities(mat_group, group_vols)
    
    return mb
    
        
if __name__ == "__main__":


    ap = ArgumentParser("Assigns materials to a dagmc model using an lcad file")

    ap.add_argument("dagmc_model", type = str, help = "DagMC geometry file (.h5m)")
    ap.add_argument("lcad_file", type = str, help = "DagMC geometry file (.h5m)")
    ap.add_argument("-o", dest = "outfile", type = str, help = "Optional output filename")

    args = ap.parse_args()

    # update material assignments
    mb = update_mat_assignments(args.dagmc_model, args.lcad_file)

    # setup output filename
    output_file_base = args.dagmc_model.split(".")[0]
    suffix = "_lcad_mats.h5m"
    output_filename = args.outfile if args.outfile else output_file_base+suffix
    
    # write new DAGMC model to file
    print "Writing file " + output_filename + "..."
    mb.write_file(output_filename)
    print "done"