From 61efad0a8d63884f535fa4adce4e04b846f8c714 Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Fri, 11 Nov 2022 12:10:10 -0500 Subject: [PATCH] added test and expand_dims kwarg --- openmc/tallies.py | 39 +++++++++++++++------- tests/unit_tests/test_filter_mesh.py | 48 ++++++++++++++++++++++++++-- 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/openmc/tallies.py b/openmc/tallies.py index 33a14974c3..688369de04 100644 --- a/openmc/tallies.py +++ b/openmc/tallies.py @@ -1405,7 +1405,7 @@ class Tally(IDManagerMixin): return df - def get_reshaped_data(self, value='mean'): + def get_reshaped_data(self, value='mean', expand_dims=False): """Returns an array of tally data with one dimension per filter. The tally data in OpenMC is stored as a 3D array with the dimensions @@ -1417,17 +1417,24 @@ class Tally(IDManagerMixin): This builds and returns a reshaped version of the tally data array with unique dimensions corresponding to each tally filter. For example, - suppose this tally has arrays of data with shape (8,5,5) corresponding - to two filters (2 and 4 bins, respectively), five nuclides and five + suppose this tally has arrays of data with shape (30,5,5) corresponding + to two filters (2 and 15 bins, respectively), five nuclides and five scores. This method will return a version of the data array with the - with a new shape of (2,4,5,5) such that the first two dimensions - correspond directly to the two filters with two and four bins. + with a new shape of (2,15,5,5) such that the first two dimensions + correspond directly to the two filters with two and fifteen bins. If + expand_dims is True and our filter above with 15 bins is an instance of + :class:`openmc.MeshFilter` with a shape of (3,5,1). The resulting tally + data array will have a new shape of (2,3,5,1,5,5). Parameters ---------- value : str A string for the type of value to return - 'mean' (default), 'std_dev', 'rel_err', 'sum', or 'sum_sq' are accepted + expand_dims : bool, optional + Whether or not to expand the dimensions of filters with multiple + dimensions. This will result in more than one dimension per filter + for the returned data array. Returns ------- @@ -1439,24 +1446,32 @@ class Tally(IDManagerMixin): # Get the 3D array of data in filters, nuclides and scores data = self.get_values(value=value) - # Build a new array shape with one dimension per filter + # Build a new array shape with one dimension per filter or expand + # multidimensional filters if desired new_shape = tuple() idx0 = None for i, f in enumerate(self.filters): - # Mesh filter indices are backwards so we need to flip them - if isinstance(f, openmc.MeshFilter): - fshape = f.shape[::-1] - new_shape += fshape - idx0, idx1 = i, i + len(fshape) - 1 + if expand_dims: + # Mesh filter indices are backwards so we need to flip them + if isinstance(f, openmc.MeshFilter): + fshape = f.shape[::-1] + new_shape += fshape + idx0, idx1 = i, i + len(fshape) - 1 + else: + new_shape += f.shape else: - new_shape += f.shape + new_shape += (np.prod(f.shape),) new_shape += (self.num_nuclides, self.num_scores) # Reshape the data with one dimension for each filter data = np.reshape(data, new_shape) + + # If we had a MeshFilter we should swap the axes to have the same shape + # for the data and the filter if idx0 is not None: data = np.swapaxes(data, idx0, idx1) + return data def hybrid_product(self, other, binary_op, filter_product=None, diff --git a/tests/unit_tests/test_filter_mesh.py b/tests/unit_tests/test_filter_mesh.py index 265642360b..3bb836e469 100644 --- a/tests/unit_tests/test_filter_mesh.py +++ b/tests/unit_tests/test_filter_mesh.py @@ -1,7 +1,10 @@ +import math + import numpy as np -import openmc from uncertainties import unumpy +import openmc + def test_spherical_mesh_estimators(run_in_tmpdir): """Test that collision/tracklength estimators agree for SphericalMesh""" @@ -62,7 +65,8 @@ def test_cylindrical_mesh_estimators(run_in_tmpdir): mat.add_nuclide('U235', 1.0) mat.set_density('g/cm3', 10.0) - cyl = openmc.model.RightCircularCylinder((0., 0., -5.), 10., 10.0, boundary_type='vacuum') + cyl = openmc.model.RightCircularCylinder((0., 0., -5.), 10., 10.0, + boundary_type='vacuum') cell = openmc.Cell(fill=mat, region=-cyl) model = openmc.Model() model.geometry = openmc.Geometry([cell]) @@ -107,3 +111,43 @@ def test_cylindrical_mesh_estimators(run_in_tmpdir): diff = unumpy.nominal_values(delta) std_dev = unumpy.std_devs(delta) assert np.all(diff < 3*std_dev) + +def test_get_reshaped_data(run_in_tmpdir): + """Test that expanding MeshFilter dimensions works as expected""" + + mat = openmc.Material() + mat.add_nuclide('U235', 1.0) + mat.set_density('g/cm3', 10.0) + + sphere = openmc.Sphere(r=10.0, boundary_type='vacuum') + cell = openmc.Cell(fill=mat, region=-sphere) + model = openmc.Model() + model.geometry = openmc.Geometry([cell]) + model.settings.particles = 1_000 + model.settings.inactive = 10 + model.settings.batches = 20 + + sph_mesh = openmc.SphericalMesh() + sph_mesh.r_grid = np.linspace(0.0, 5.0**3, 20)**(1/3) + sph_mesh.theta_grid = np.linspace(0, math.pi, 4) + sph_mesh.phi_grid = np.linspace(0, 2*math.pi, 3) + tally1 = openmc.Tally() + efilter = openmc.EnergyFilter([0, 1e5, 1e8]) + meshfilter = openmc.MeshFilter(sph_mesh) + assert meshfilter.shape == (19, 3, 2) + tally1.filters = [efilter, meshfilter] + tally1.scores = ['flux'] + + model.tallies = openmc.Tallies([tally1]) + + # Run OpenMC + sp_filename = model.run() + + # Get flux tally as reshaped data + with openmc.StatePoint(sp_filename) as sp: + t1 = sp.tallies[tally1.id] + data1 = t1.get_reshaped_data() + data2 = t1.get_reshaped_data(expand_dims=True) + + assert data1.shape == (2, 19*3*2, 1, 1) + assert data2.shape == (2, 19, 3, 2, 1, 1)