Mesh Material Raytracing Test Script

gistfile1.txt

import os
import sys

import numpy as np
from scipy.spatial import Delaunay
from scipy.stats import qmc
import openmc

cube_corners = np.array([
    [0, 0, 0],
    [0, 0, 1],
    [0, 1, 0],
    [0, 1, 1],
    [1, 0, 0],
    [1, 0, 1],
    [1, 1, 0],
    [1, 1, 1]
])

def make_model():
    # Generate interior points using a Sobol sequence
    seed = 12345
    sobol_engine = qmc.Sobol(3, seed=seed)
    interior_points = sobol_engine.random_base2(6)

    # For convenience, define a function that returns n Sobol samples in 2D
    def get_sobol_2d_samples(m, seed):
        sampler = qmc.Sobol(d=2, seed=seed)
        return sampler.random_base2(m)

    xy = get_sobol_2d_samples(3, seed)
    face_x0 = np.column_stack((np.zeros(2**3), xy[:, 0], xy[:, 1]))
    xy = get_sobol_2d_samples(3, seed + 1)
    face_x1 = np.column_stack((np.ones(2**3), xy[:, 0], xy[:, 1]))
    xy = get_sobol_2d_samples(3, seed + 2)
    face_y0 = np.column_stack((xy[:, 0], np.zeros(2**3), xy[:, 1]))
    xy = get_sobol_2d_samples(3, seed + 3)
    face_y1 = np.column_stack((xy[:, 0], np.ones(2**3), xy[:, 1]))
    xy = get_sobol_2d_samples(3, seed + 4)
    face_z0 = np.column_stack((xy[:, 0], xy[:, 1], np.zeros(2**3)))
    xy = get_sobol_2d_samples(3, seed + 5)
    face_z1 = np.column_stack((xy[:, 0], xy[:, 1], np.ones(2**3)))

    # Combine corners and interior points
    points = np.vstack([cube_corners, interior_points, face_x0, face_x1, face_y0, face_y1, face_z0, face_z1])

    # Scale points to cover [-10, 10]^3
    points -= 0.5
    points *= 20

    # Create the Delaunay tetrahedralization
    tri = Delaunay(points)

    # Create list to store each tetrahedron as an OpenMC Cell
    cells = []

    for simplex in tri.simplices:
        # 4 vertices of the tetrahedron
        p0 = points[simplex[0]]
        p1 = points[simplex[1]]
        p2 = points[simplex[2]]
        p3 = points[simplex[3]]

        # 3-point faces of the tetrahedron
        faces = [
            (p0, p1, p2),
            (p0, p1, p3),
            (p0, p2, p3),
            (p1, p2, p3),
        ]

        # We need the centroid to determine which side of each plane is "inside"
        centroid = (p0 + p1 + p2 + p3) / 4.0

        # Create planes and figure out orientation for each face
        region_parts = []
        for (fp1, fp2, fp3) in faces:
            plane = openmc.Plane.from_points(fp1, fp2, fp3)

            # Turn into canonical form
            plane.a, plane.b, plane.c, plane.d = plane.normalize()

            # Evaluate the plane equation at the tetrahedron's centroid
            # plane.evaluate(x) > 0 => the point x is on the +plane side
            region_part = +plane if plane.evaluate(centroid) > 0 else -plane
            region_parts.append(region_part)

        # Create material for this cell
        m = openmc.Material()
        m.add_nuclide('H1', 1.0)
        m.set_density('g/cm3', 1.0)

        # Create a cell corresponding to this tetrahedron
        cell = openmc.Cell(fill=m, region=openmc.Intersection(region_parts))
        cell.volume = np.abs(np.linalg.det(np.array([p1 - p0, p2 - p0, p3 - p0]))) / 6.0
        m.volume = cell.volume
        cells.append(cell)

    # place all cells in a single universe and use that universe to fill a bounding region
    univ = openmc.Universe(cells=cells)
    box = openmc.model.RectangularParallelepiped(
        -10, 10, -10, 10, -10, 10, boundary_type='vacuum'
    )
    root_cell = openmc.Cell(fill=univ, region=-box)

    geom = openmc.Geometry([root_cell], merge_surfaces=True)

    model = openmc.Model(geometry=geom)
    model.settings.run_mode = 'fixed source'
    model.settings.particles = 1000
    model.settings.batches = 10
    model.export_to_model_xml('random_geom.xml')
    return model


def get_mesh_volmes(model, mesh_type, mesh_dims, n_samples, n_threads=None):
    # Mesh volumes
    if mesh_type == 'Regular':
        mesh = openmc.RegularMesh.from_domain(model.geometry.root_universe, dimension=mesh_dims)
    elif mesh_type == 'Tetrahedral':
        mesh = openmc.UnstructuredMesh(filename='test_mesh_tets.exo', library='moab')
    else:
        raise ValueError(f'Unsupported mesh type: {mesh_type}')

    init_args = {}
    if n_threads is not None:
        init_args['args'] = ['-s', str(n_threads)]

    print(f'\n---------------------------------------------')
    if mesh_type == 'Regular':
        print(f'Mesh dims: {mesh.dimension}')
    mesh_volumes = mesh.material_volumes(model=model, n_samples=n_samples, max_materials=300, output=True, **init_args)
    print(f'---------------------------------------------\n')

    write_bad_only = True
    bad_cells = 0
    element_diffs = []
    material_diffs = []

    volume_sums = {m_id : sum(volumes) for m_id, volumes in mesh_volumes.items()}

    element_sums = {e_id : sum(v[1] for v in mesh_volumes.by_element(e_id)) for e_id in range(mesh_volumes.num_elements)}

    for cell_id, cell in model.geometry.get_all_material_cells().items():
        material = cell.fill
        material_volume = cell.fill.volume

        if material.id not in volume_sums:
            continue
        mesh_volume = volume_sums[material.id]
        material_diffs.append(np.abs(material_volume - mesh_volume))
        if write_bad_only and material_diffs[-1] < 1e-2:
            continue
        bad_cells += 1
        # print(f'Material {cell.fill.id}:')
        # print(f'  Material volume: {material_volume:.6f}, Mesh volume: {mesh_volume:.6f}')

    print(f'Total bad cells: {bad_cells} out of {len(model.geometry.get_all_material_cells())}')
    material_diffs = np.array(material_diffs)
    print(f'Mean difference: {np.mean(material_diffs):.6f}')
    print(f'Max difference: {np.max(material_diffs):.6f}')
    print(f'Standard deviation: {np.std(material_diffs):.6f}')

    bad_elements = 0
    element_volumes = mesh.volumes.flatten()
    for e_id, e_vol in element_sums.items():
        element_diffs.append(np.abs(e_vol - element_volumes[e_id]))
        if write_bad_only and element_diffs[-1] < 1e-2:
            continue
        bad_elements += 1

        # print(f'Element {e_id}:')
        # print(f'  Element volume: {e_vol:.6f}, Mesh volume: {mesh_vol:.6f}')

    print(f'Total bad elements: {bad_elements} out of {mesh_volumes.num_elements}')
    element_diffs = np.array(element_diffs)
    print(f'Mean element difference: {np.mean(element_diffs):.6f}')
    print(f'Max element difference: {np.max(element_diffs):.6f}')
    print(f'Standard deviation: {np.std(element_diffs):.6f}')

    return material_diffs, element_diffs


def main():
    # parameters
    n_samples = 3*[2048]
    n_threads = 30
    mesh_dims = [(1, 1, 1), (2, 2, 2), (4, 4, 4), (8, 8, 8), (16, 16, 16)] #, (32, 32, 32)]
    diff_list = []
    element_diffs = []
    mesh_type = 'Regular'
    model = make_model()

    material_volume_reference = 'material_vol_ref.npy'
    element_volume_reference = 'element_vol_ref.npy'

    # create reference histogram
    mat_diffs, element_diff = get_mesh_volmes(model, mesh_type=mesh_type, mesh_dims=(1, 1, 1), n_samples=n_samples, n_threads=1)

    material_vol_ref_hist = np.histogram(mat_diffs, bins=50, density=True)[0]
    np.save(material_volume_reference, material_vol_ref_hist, allow_pickle=True)

    if mesh_type == 'Tetrahedral':
        mesh_dims = [(1, 1, 1)]

    for dims in mesh_dims:
        mat_diffs, element_diff = get_mesh_volmes(model, mesh_type=mesh_type, mesh_dims=dims, n_samples=n_samples, n_threads=n_threads)
        diff_list.append(mat_diffs)
        element_diffs.append(element_diff)

    # plot a histogram of the differences
    import matplotlib.pyplot as plt
    import seaborn as sns
    sns.set(style='whitegrid')
    fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(18, 12))
    fig.suptitle(f'Material Volume Differences; Threads: {n_threads}; Rays: {n_samples}')
    axes = axes.flatten()

    mat_vol_histograms = [np.histogram(mat_diffs, bins=50, density=True)[0] for mat_diffs in diff_list]
    ref_hist = np.load(material_volume_reference, allow_pickle=True)
    # check each histogram against the reference
    mat_vol_histchecks = [np.allclose(hist, ref_hist, atol=1e-6) for hist in mat_vol_histograms]

    for i, mat_diffs in enumerate(diff_list):
        linestyle = ['--', '-.', ':'][i % 3]
        sns.histplot(mat_diffs, bins=50, kde=True, stat='probability', linestyle=linestyle, alpha=0.7, ax=axes[i])
        if not mat_vol_histchecks[i]:
            axes[i].annotate('Mismatch', xy=(0.7, 0.9), xycoords='axes fraction', fontsize=12, color='red',
                            bbox=dict(facecolor='white', alpha=0.8))
        axes[i].set_xlabel('Volume difference')
        axes[i].set_xlim(left=0)
        axes[i].set_ylim(0.0, 0.5)
        axes[i].set_ylabel('Density')
        axes[i].set_title(f'{mesh_type} Mesh ({mesh_dims[i][0]}x{mesh_dims[i][1]}x{mesh_dims[i][2]})')
        axes[i].legend()

    plt.tight_layout()
    plt.savefig(f'material_volume_difference_histograms_{n_samples[0]}.png')

    fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(18, 12))
    fig.suptitle(f'Element Volume Differences; Threads: {n_threads}; Rays: {n_samples}')
    axes = axes.flatten()
    # perform a similar check for the mesh element volumes
    for i, element_diff in enumerate(element_diffs):
        element_diff = np.array(element_diff).flatten()
        sns.histplot(element_diff, bins=50, kde=True, stat='probability', linestyle=linestyle, alpha=0.7, ax=axes[i])
        axes[i].set_xlabel('Volume difference')
        axes[i].set_xlim(left=0)
        axes[i].set_ylim(0.0, 0.5)
        axes[i].set_ylabel('Density')
        axes[i].set_title(f'{mesh_type} Mesh ({mesh_dims[i][0]}x{mesh_dims[i][1]}x{mesh_dims[i][2]})')
        axes[i].legend()

    plt.tight_layout()
    plt.savefig(f'element_volume_difference_histograms_{n_samples[0]}.png')

    if '-s' in sys.argv:
        plt.show()


if __name__ == '__main__':
    main()