A script for checking scaling of a DAGMC model.

threaded_scaling.py

import sys

from matplotlib import pyplot as plt

import openmc
import numpy as np

def threaded_scaling(threads, model, output_filename=None):
    """Run simulation with various threads and collect timing results

    Parameters
    ----------

    threads : Iterable of int
        The thread counts to run with
    model : openmc.Model
        The model to run
    output_filename: str or pathlib.Path
        Name of the output file to be written in CSV
    """

    filename = 'scaling.csv' if output_filename is None else output_filename

    output = open(filename, 'w')
    header = f'threads, particles, inactive batches, total batches, time inactive (s), time active (s), time transport (s)\n'
    output.write(header)

    for thread_count in threads:
        sp_file = model.run(threads=thread_count)

        # collect timing
        with openmc.StatePoint(sp_file) as sp:

            inactive_time = sp.runtime['inactive batches']
            active_time = sp.runtime['active batches']
            transport_time = sp.runtime['transport']

        particles = model.settings.particles
        inactive = model.settings.inactive
        batches = model.settings.batches

        output_line = f'{thread_count}, {particles}, {inactive},' \
                      f'{batches}, {inactive_time}, {active_time}, {transport_time}\n'
        output.write(output_line)
        output.flush()
        # open the statepoint file and get timing

    output.close()

def ax_plot(ax, filename, label=''):
    data = np.loadtxt(filename, skiprows=1, delimiter=',').T
    ax.set_xlabel('# threads')
    ax.set_ylabel('particles/s')
    ax.plot(data[0, :], data[1,:] * data[2,:] / data[4, :], label=label)


def fig_plot(title, filename, label):
    fig, ax = plt.subplots()
    fig.suptitle(title)
    ax_plot(ax, filename, label)
    plt.legend()
    plt.show()


def run():

    l_values = [1, 10, 20, 40]

    threads = [1, 5, 10, 15, 20]

    model = openmc.Model.from_xml()
    model.settings.particles = 50_000
    model.settings.inactive = 10
    model.settings.batches = 11

    for l in l_values:
        dagmc_file = f'moab_mesh_l_{l}.h5m'
        model.geometry.root_universe = openmc.DAGMCUniverse(dagmc_file)
        scaling_filename = f'scaling_l_{l}.csv'
        print(f'Threaded scaling run l={l}...')
        threaded_scaling(threads, model, scaling_filename)

if __name__ == '__main__':
    if len(sys.argv) == 1 or 'r' in sys.argv[1]:
        run()
    if len(sys.argv) > 1 and 'p' in sys.argv[1]:
        plot(sys.argv[2])