cadquery to h5m with moab and surface sense

cadquery-to-dagmc-with-surfaces-sense.py

# this scrip makes a cadquery geometry, uses cadquery to imprint the surfaces, uses cadquery to mesh the faces of each solid
# then makes a moab core object by adding faces one at a time.
# if a face appears in two solids then the surface sense will be reversed

# work in progress
import cadquery as cq

import openmc
from typing import Tuple
import numpy as np
from pymoab import core, types
from OCP.TopLoc import TopLoc_Location
from OCP.BRep import BRep_Tool
from OCP.TopAbs import TopAbs_Orientation


def tessellate_solid_faces(parts, tolerance: float = 0.1, angularTolerance: float = 0.1):
    """Creates a mesh / faceting / tessellation of the surface"""

    parts.mesh(tolerance, angularTolerance)
    offset = 0
    vertices = [] # type List[Vector]
    triangles = {}
    for c, f in enumerate(parts.Faces()):
        loc = TopLoc_Location()
        poly = BRep_Tool.Triangulation_s(f.wrapped, loc)
        Trsf = loc.Transformation()
        reverse = (
            True
            if f.wrapped.Orientation() == TopAbs_Orientation.TopAbs_REVERSED
            else False
        )
        # add vertices
        nodes = [poly.Node(i + 1) for i in range(poly.NbNodes())]
        face_verticles = [
            (v.X(), v.Y(), v.Z()) for v in (v.Transformed(Trsf) for v in nodes)
        ]
        vertices += face_verticles

        face_triangles = [
            (
                t.Value(1) + offset - 1,
                t.Value(3) + offset - 1,
                t.Value(2) + offset - 1,
            )
            if reverse
            else (
                t.Value(1) + offset - 1,
                t.Value(2) + offset - 1,
                t.Value(3) + offset - 1,
            )
            for t in poly.Triangles()
        ]
        triangles[f.hashCode()] = face_triangles

        offset += poly.NbNodes()

    solids_faces_triangles = {}
    for s in parts.Solids():
        solids_faces_triangles[s.hashCode()]={}
        for f in s.Faces():
            solids_faces_triangles[s.hashCode()][f.hashCode()] = triangles[f.hashCode()]

    return vertices, solids_faces_triangles




def define_moab_core_and_tags() -> Tuple[core.Core, dict]:
    """Creates a MOAB Core instance which can be built up by adding sets of
    triangles to the instance

    Returns:
        (pymoab Core): A pymoab.core.Core() instance
        (pymoab tag_handle): A pymoab.core.tag_get_handle() instance
    """

    # create pymoab instance
    moab_core = core.Core()

    tags = dict()

    sense_tag_name = "GEOM_SENSE_2"
    sense_tag_size = 2
    tags["surf_sense"] = moab_core.tag_get_handle(
        sense_tag_name,
        sense_tag_size,
        types.MB_TYPE_HANDLE,
        types.MB_TAG_SPARSE,
        create_if_missing=True,
    )

    tags["category"] = moab_core.tag_get_handle(
        types.CATEGORY_TAG_NAME,
        types.CATEGORY_TAG_SIZE,
        types.MB_TYPE_OPAQUE,
        types.MB_TAG_SPARSE,
        create_if_missing=True,
    )

    tags["name"] = moab_core.tag_get_handle(
        types.NAME_TAG_NAME,
        types.NAME_TAG_SIZE,
        types.MB_TYPE_OPAQUE,
        types.MB_TAG_SPARSE,
        create_if_missing=True,
    )

    tags["geom_dimension"] = moab_core.tag_get_handle(
        types.GEOM_DIMENSION_TAG_NAME,
        1,
        types.MB_TYPE_INTEGER,
        types.MB_TAG_DENSE,
        create_if_missing=True,
    )

    # Global ID is a default tag, just need the name to retrieve
    tags["global_id"] = moab_core.tag_get_handle(types.GLOBAL_ID_TAG_NAME)

    return moab_core, tags


# test geometry 1 two touching cubes
# assembly = cq.Assembly(name="TwoTouchingCuboids")
# cuboid1 = cq.Workplane().box(10, 10, 10)
# assembly.add(cuboid1)
# cuboid2 = cq.Workplane().moveTo(0.5 * 10 + 0.5 * 4).box(4, 4, 4)
# assembly.add(cuboid2)

# test geometry 2, sphere in sphell
assembly = cq.Assembly(name="sphericalshell")
sphere1 = cq.Workplane().sphere(10)
sphere2 = cq.Workplane().sphere(10 + 5).cut(sphere1)
assembly.add(sphere1)
assembly.add(sphere2)


material_tags = ["mat1", "mat2"]

imprinted_assembly, imprinted_solids_with_original_id = cq.occ_impl.assembly.imprint(
    assembly
)

#assumes all normals are correct
vertices, triangles_by_solid_by_face = tessellate_solid_faces(imprinted_assembly, 0.1)


import trimesh
triangles_by_solid_by_face_fixed_normals = {}
for solid_id, triangles_on_each_face in triangles_by_solid_by_face.items():
    triangles_by_solid_by_face_fixed_normals[solid_id]={}
    for face_id, triangles_on_face in triangles_on_each_face.items():
        
        mesh = trimesh.Trimesh(vertices=vertices, faces=triangles_on_face, process=False)
        mesh.fix_normals()
        trimesh.repair.fix_normals(mesh) 
        print(mesh.faces)
        triangles_by_solid_by_face_fixed_normals[solid_id][face_id]=mesh.faces

        # just prints out the triangles that are now different
        for t, t2 in zip(mesh.faces,triangles_on_face):
            if t[0]!=t2[0] or t[1]!=t2[1]:
                print('different',t,t2)

# builds a dictionary of face ids as keys and solid ids as values
face_ids_with_solid_ids = {}
for solid_id, triangles_on_each_face in triangles_by_solid_by_face_fixed_normals.items():
    for face_id, triangles_on_face in triangles_on_each_face.items():
        if face_id in face_ids_with_solid_ids.keys(): 
            face_ids_with_solid_ids[face_id].append(solid_id)
        else:
            face_ids_with_solid_ids[face_id] = [solid_id]


moab_core, tags = define_moab_core_and_tags()
face_ids_in_moab_core = []
volume_sets_added_by_face_id={}
volume_sets ={}
# for solid_id, triangles_on_each_face in triangles_by_solid_by_face.items()
for material_tag, (solid_id, triangles_on_each_face) in zip(material_tags, triangles_by_solid_by_face_fixed_normals.items()):
    print('adding moab details for solid id', solid_id,'and material tag',material_tag)

    volume_set = moab_core.create_meshset()
    volume_sets[solid_id]=volume_set

    # for face_id, triangles_on_face in triangles_on_each_face.items():
    #     volume_sets_added_by_face_id[face_id]=volume_set

print('moab volume_sets',volume_sets,'\n')
print('face_ids_with_solid_ids',face_ids_with_solid_ids,'\n')

# passing through the solids and faces once to make the surface senses for the geometry
sense_data_for_each_solid_and_face = {}
for material_tag, (solid_id, triangles_on_each_face) in zip(material_tags, triangles_by_solid_by_face_fixed_normals.items()):
    volume_set=volume_sets[solid_id]
    sense_data_for_each_solid_and_face[solid_id]={}
    for face_id, triangles_on_face in triangles_on_each_face.items():

        # faces appears twice in geometry
        if len(face_ids_with_solid_ids[face_id])==2:
            # get the other volume_set
            if volume_sets[face_ids_with_solid_ids[face_id][0]] == volume_sets[solid_id]:
                print('2nd volume_set is the other one')
                other_solid_id = face_ids_with_solid_ids[face_id][1]
            else:
                print('1st volume_set is the other one')
                other_solid_id = face_ids_with_solid_ids[face_id][0]
            other_volume_set = volume_sets[other_solid_id]
            sense_data = np.array( [volume_set,other_volume_set], dtype='uint64')
        else:
            #face appears in just one solid so 2nd sense is 0
            sense_data = np.array( [volume_set, 0], dtype='uint64')
        sense_data_for_each_solid_and_face[solid_id][face_id] = sense_data

print('sense_data_for_each_solid_and_face',sense_data_for_each_solid_and_face,'\n')

# passing through the solids and faces to make the moab object for each volume
for material_tag, (solid_id, triangles_on_each_face) in zip(material_tags, triangles_by_solid_by_face_fixed_normals.items()):
    print('adding moab details for solid id', solid_id,'and material tag',material_tag)
    # recent versions of MOAB handle this automatically
    # but best to go ahead and do it manually
    volume_set = volume_sets[solid_id]
    moab_core.tag_set_data(tags["global_id"], volume_set, solid_id)
    moab_core.tag_set_data(tags["geom_dimension"], volume_set, 3)
    moab_core.tag_set_data(tags["category"], volume_set, "Volume")

    for face_id, triangles_on_face in triangles_on_each_face.items():
        print('    adding face',face_id,'to moab core')
        # if the faces has previously been added we know the surface sense should be rverse of the previous time
        sense_data = sense_data_for_each_solid_and_face[solid_id][face_id]# reversing sense data

        # just for printing info
        if sense_data[1]==0:
            print('        face is not shared by solids')
        else:
            print('        face is shared by solids')
    

        if face_id in face_ids_in_moab_core:
            print(f'        face {face_id} has already been added, this must be a face shared by two solids')
            face_set=moab_core.create_meshset()
            moab_core.add_parent_child(volume_set, face_set)

            print('        face is present in two solids, making surface sense accordingly')
            solid1, solid2 = face_ids_with_solid_ids[face_id]
            print(f'        face {face_id} is in solid {solid1} and {solid2}')

            moab_core.tag_set_data(tags["surf_sense"], face_set, sense_data)
        else:

            print(f'        face {face_id} has not previously been added')
            
            # checks if this face appears more than once in the geometry and makes surf sense data differently if it does

            face_set = moab_core.create_meshset()
            # volume_sets_added_by_face_id[face_id] = volume_set
            moab_core.tag_set_data(tags["global_id"], face_set, face_id)
            moab_core.tag_set_data(tags["geom_dimension"], face_set, 2)
            moab_core.tag_set_data(tags["category"], face_set, "Surface")

            moab_core.add_parent_child(volume_set, face_set)
            moab_core.tag_set_data(tags["surf_sense"], face_set, sense_data)

            moab_verts = moab_core.create_vertices(vertices)
            moab_core.add_entity(face_set, moab_verts)

            for triangle in triangles_on_face:
                tri = (
                    moab_verts[int(triangle[0])],
                    moab_verts[int(triangle[1])],
                    moab_verts[int(triangle[2])],
                )

                moab_triangle = moab_core.create_element(types.MBTRI, tri)
                moab_core.add_entity(face_set, moab_triangle)
            
            face_ids_in_moab_core.append(face_id)
        print('        sense_data', sense_data)

    group_set = moab_core.create_meshset()

    moab_core.tag_set_data(tags["category"], group_set, "Group")

    moab_core.tag_set_data(tags["name"], group_set, f"mat:{material_tag}")

    moab_core.tag_set_data(tags["geom_dimension"], group_set, 4)

    moab_core.add_entity(group_set, volume_set)


all_sets = moab_core.get_entities_by_handle(0)

file_set = moab_core.create_meshset()

moab_core.add_entities(file_set, all_sets)

moab_core.write_file("dagmc.h5m")

mat1 = openmc.Material(name='mat1')
mat1.add_nuclide('Fe56', 1)
mat1.set_density('g/cm3', 1)
mat2 = openmc.Material(name='mat2')

mat2.add_nuclide('Fe56', 1)
mat2.set_density('g/cm3', 1)
my_materials = openmc.Materials([mat1, mat2])


mat1_filter = openmc.MaterialFilter(mat1)
tally1 = openmc.Tally(name='mat1_flux_tally')
tally1.filters = [mat1_filter]
tally1.scores = ['flux']

mat2_filter = openmc.MaterialFilter(mat2)
tally2 = openmc.Tally(name='mat2_flux_tally')
tally2.filters = [mat2_filter]
tally2.scores = ['flux']
my_tallies = openmc.Tallies([tally1, tally2])


my_settings = openmc.Settings()
my_settings.batches = 2
my_settings.inactive = 0
my_settings.particles = 5000
my_settings.run_mode = 'fixed source'

# Create a DT point source
my_source = openmc.Source()
my_source.space = openmc.stats.Point((0, 0, 0))
my_source.angle = openmc.stats.Isotropic()
my_source.energy = openmc.stats.Discrete([14e6], [1])
my_settings.source = my_source

universe = openmc.DAGMCUniverse('dagmc.h5m').bounded_universe()
geometry = openmc.Geometry(universe)
model = openmc.Model(geometry,my_materials,my_settings,my_tallies)

output_file_from_cad = model.run()
with openmc.StatePoint(output_file_from_cad) as sp_from_cad:
    cad_result1 = sp_from_cad.get_tally(name="mat1_flux_tally")
    cad_result2 = sp_from_cad.get_tally(name="mat2_flux_tally")

print(f'CAD tally mat 1 mean {cad_result1.mean} std dev {cad_result1.std_dev}')
# result should not be 0
print(f'CAD tally mat 2 mean {cad_result2.mean} std dev {cad_result2.std_dev}')