use defaultdict

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
This commit is contained in:
Olek 2022-11-21 15:33:54 -06:00 committed by yardasol
parent d45c37b206
commit 8163d7a759

View file

@ -3,6 +3,7 @@
from math import floor
import shutil
from pathlib import Path
from collections import defaultdict
from difflib import unified_diff
import numpy as np
@ -122,18 +123,11 @@ def test_full(run_in_tmpdir, problem, multiproc):
n_tallies = np.empty(N + 1, dtype=int)
# Get statepoint files for all BOS points and EOL
runtimes = {}
runtimes = defaultdict(list)
for n in range(N + 1):
statepoint = openmc.StatePoint(f"openmc_simulation_n{n}.h5")
runtime = statepoint.runtime
if n == 0:
for measure, time in runtime.items():
runtimes.update({measure: np.array([time])})
else:
for measure, time in runtime.items():
current = runtimes[measure]
updated = np.append(current, time)
runtimes.update({measure: updated})
for measure, time in statepoint.runtime.items():
runtimes[measure].append(time)
k_n = statepoint.keff
k_state[n] = [k_n.nominal_value, k_n.std_dev]
n_tallies[n] = len(statepoint.tallies)
@ -144,6 +138,9 @@ def test_full(run_in_tmpdir, problem, multiproc):
# Check that no additional tallies are loaded from the files
assert np.all(n_tallies == 0)
# Convert values in runtimes to arrays
runtimes = {k: np.array(v) for k, v in runtimes.items()}
# Check that runtimes are qualitatively correct
assert runtimes['reading cross sections'][0] != 0
assert runtimes['total initialization'][0] != 0