Python file designed to update DAGMC model group names from the old conventions to the new ones.

update_dagmc_mat_names.py

#!/usr/bin/env python

import numpy as np
import argparse

from pymoab import core, types


# some convenience functions
def get_category(mbi, category):
    """
    get the category tag and find all of the groups
    """
    catTag = mbi.tag_get_handle("CATEGORY")
    return mbi.get_entities_by_type_and_tag(0,
                                            types.MBENTITYSET,
                                            catTag,
                                            np.array([category]))


def get_groups(mbi):
    """returns all meshsets with a 'Group' category tag"""
    return get_category(mbi, "Group")


def get_volumes(mbi):
    """returns all meshsets with a 'Volume' category tag"""
    return get_category(mbi, "Volume")


def update_file(filename, output_file=None, print_only=False, strip_rho=False):
    """
    updates a dagmc file's naming scheme and validates it's volume to group
    contaimnet
    """
    # create a new moab instance
    mb = core.Core()
    # load the file
    mb.load_file(filename)
    validate_group_topology(mb)

    if print_only:
        print("Group names of " + filename + " are:")
        for groupName in get_group_names(mb):
            print(groupName)
        return

    update_names(mb, strip_rho)
    # write file with "new_names" appended
    if output_file is None or output_file is "":
        ext = filename[-4:]
        basename = filename[:-4]
        suffix = "new_names"
        new_filename = basename+"_new_names"+ext
    else:
        new_filename = str(output_file)

    mb.write_file(new_filename)


def get_group_names(mbi, groupSets):
    """
    returns a list of group names from the input PyMOAB instance
    """
    # get the name tag (material specification) for each group
    nameTag = mbi.tag_get_handle("NAME")
    groupNames = mbi.tag_get_data(nameTag, groupSets, flat=True)
    return groupNames


def update_names(mbi, strip_density=False):
    """
    updates the group names of a MOAB instance to the new UWUW
    naming scheme
    """

    # get the group EntitySets
    groupSets = get_groups(mbi)

    groupNames = get_group_names(mbi, groupSets)

    print("Old Group Names", groupNames)
    # update the group names to the new naming scheme
    newGroupNames = []
    for name in groupNames:
        # if this is a material specification, update name
        if 'mat' in name:
            # split on old naming delimiter
            temp = name.split('_')
            # join each key/value pair with new ":" delimiter
            temp = [":".join(temp[i:i+2]) for i in range(0, len(temp), 2)]
            # separate material properties with new "/" delimiter
            name = "/".join(temp)
            if strip_density and 'rho' in name:
                # split name on new separator
                split_name = name.split('/')
                rho_entries = [s for s in split_name if 'rho:' in s]
                if len(rho_entries) > 1:
                    print("Warning: multiple density entries "
                          "found in mat: {}".format(name))
                for s in rho_entries:
                    split_name.remove(s)
                name = '/'.join(split_name)
        # update reflecting boundary condition
        if 'reflect' in name.lower():
            name = "boundary:Reflecting"
        # add name (modified or not) to the set of newGroupNames
        if 'graveyard' in name.lower():
            name = "mat:Graveyard"
        newGroupNames.append(name)
    # make newGroupNames into a numpy array
    newGroupNames = np.array(newGroupNames)
    print("New Group Names", newGroupNames)

    # set this data in the moab instance
    nameTag = mbi.tag_get_handle("NAME")
    mbi.tag_set_data(nameTag, groupSets, newGroupNames)


def get_group_name(mbi, group_set):
    """
    retrieves the group name for a given group
    """
    nameTag = mbi.tag_get_handle("NAME")
    groupName = mbi.tag_get_data(nameTag, group_set, flat=True)
    return groupName[0]


def validate_group_topology(mbi):
    # get all of the volumes
    vols = list(get_volumes(mbi))
    # get all of the groups
    groups = list(get_groups(mbi))
    # keys used to check for material
    mat_keys = ['mat', 'graveyard']
    # check each group's volumes
    for group in groups:
        # make sure the group is a material group
        groupName = get_group_name(mbi, group)
        if any(mat_key in groupName.lower() for mat_key in mat_keys):
            # get the group volumes
            group_vols = mbi.get_entities_by_handle(group)
            #
            for group_vol in group_vols:
                if group_vol in vols:
                    try:
                        vols.remove(group_vol)
                    except ValueError:
                        raise Exception("A volume exists in multiple groups")

    if len(vols) != 0:
        print("Found "+str(len(vols))+" empty volumes. "
              "They will be assigned material vacuum.")
        vacuum_group = None
        # look for a group with the material name vacuum
        for group in groups:
            # get the group name
            groupName = get_group_name(mbi, group)
            # check the group name
            if "Vacuum" in groupName:
                vacuum_group = group
                break
        # if there is no pre-existing vacuum group,
        # create one
        if vacuum_group is None:
            vacuum_group = mbi.create_meshset()
            mbi.add_entities(vacuum_group, vols)
            catTag = mbi.tag_get_handle("CATEGORY")
            mbi.tag_set_data(catTag, vacuum_group, ["Group"])
            nameTag = mbi.tag_get_handle("NAME")
            mbi.tag_set_data(nameTag, vacuum_group, ["mat:Vacuum"])
            geomTag = mbi.tag_get_handle("GEOM_DIMENSION")
            mbi.tag_set_data(geomTag, vacuum_group, [4])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Update naming syntax for "
                                                 "groups in a DAGMC model.")
    parser.add_argument('filename', type=str,
                        help="Filename of the DAGMC Model to update.")
    parser.add_argument('-o', type=str, dest='output_file',
                        default="",
                        help="Output file name "
                             "(default is input file + '_new_names')")
    parser.add_argument('-p', dest='print_current_mats', action='store_true',
                        help="Print current group names and exit")
    parser.add_argument('-s', dest='strip_rho', action='store_true',
                        default=False,
                        help="Strip density values "
                             "from material assignements.")
    args = parser.parse_args()

    # call main function
    update_file(args.filename,
                args.output_file,
                args.print_current_mats,
                args.strip_rho)