Merge pull request #14 from fusion-energy/develop

refactoring, adding specifc dose and spectra convertors, added doc strings

openmc_post_processor/__init__.py

@@ -1,12 +1,13 @@
 from .utils import (
     find_fusion_energy_per_reaction,
     get_tally_units,
+    get_tally_units_dose,
+    get_tally_units_spectra,
     check_for_dimentionality_difference,
     find_source_strength,
     compute_volume_of_voxels,
     process_tally,
     process_dose_tally,
     process_spectra_tally,
-    convert_units,
     scale_tally
 )

openmc_post_processor/utils.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Tuple
 
 import numpy as np
 import openmc
@@ -7,24 +8,39 @@
 ureg = pint.UnitRegistry()
 ureg.load_definitions(str(Path(__file__).parent / "neutronics_units.txt"))
 
+
 def process_spectra_tally(
     tally,
-    required_units=None,
-    source_strength=None,
-    volume=None,
-):
-    """Processes a spectra tally converting the tally with default units obtained
-    during simulation into the user specified units. In some cases
-    additional user inputs will be required such as volume or source strength."""
+    required_units: Tuple[str, str] = None,
+    source_strength: float = None,
+    volume: float = None,
+) -> tuple:
+    """Processes a spectra tally converting the tally with default units
+    obtained during simulation into the user specified units.
+
+    Args:
+        tally: The openmc.Tally object which should be a spectra tally. With a
+            score of flux or current and an EnergyFilter
+        required_units: A tuple of the units to convert the energy and tally into
+        source_strength: In some cases the source_strength will be required
+            to convert the base units into the required units. This optional
+            argument allows the user to specify the source_strength when needed
+        volume: In some cases the volume will be required to convert the base
+            units into the required units. In the case of a regular mesh the
+            volume is automatically found. This optional argument allows the
+            user to specify the volume when needed or overwrite the
+            automatically calculated volume.
+
+    Returns:
+        Tuple of spectra energies and tally results
+    """
 
-    # user makes use of openmc.StatePoint.get_tally to find tally
-    # passes the tally into this function along with the required units
-    # the tally can be returned with base units or with converted units
+    check_for_unit_modifying_filters(tally)
 
     data_frame = tally.get_pandas_dataframe()
 
     # checks for user provided base units
-    base_units = get_tally_units(tally)
+    base_units = get_tally_units_spectra(tally)
     print(f"tally {tally.name} base units {base_units}")
 
     for filter in tally.filters:
@@ -44,39 +60,51 @@ def process_spectra_tally(
             source_strength,
             volume,
         )
-        tally_results = convert_units(
-            [energy_base, scaled_tally_result], required_units
-        )
-        return tally_results
+        tally_in_required_units = scaled_tally_result.to(required_units[1])
+        energy_in_required_units = energy_base.to(required_units[0])
+
+        return energy_in_required_units, tally_in_required_units
     else:
-        return [energy_base, tally_base]
+        return energy_base, tally_base
 
 
 def process_dose_tally(
     tally,
-    required_units=None,
-    source_strength=None,
-    volume=None,
+    required_units: str = None,
+    source_strength: float = None,
+    volume: float = None,
 ):
-    """Processes the tally converting the tally with default units obtained
-    during simulation into the user specified units. In some cases
-    additional user inputs will be required. Units with"""
+    """Processes a dose tally converting the tally with default units
+    obtained during simulation into the user specified units.
 
-    # user makes use of StatePoint.get_tally to find tally
-    # passes the tally into this function along with the required units
-    # the tally can be returned with base units or with converted units
+    Args:
+        tally: The openmc.Tally object which should be a spectra tally. With a
+            score of flux or current and an EnergyFunctionFilter
+        required_units: The units to convert the energy and tally into
+        source_strength: In some cases the source_strength will be required
+            to convert the base units into the required units. This optional
+            argument allows the user to specify the source_strength when needed
+        volume: In some cases the volume will be required to convert the base
+            units into the required units. In the case of a regular mesh the
+            volume is automatically found. This optional argument allows the
+            user to specify the volume when needed or overwrite the
+            automatically calculated volume.
+
+    Returns:
+        The dose tally result in the required units
+    """
 
     data_frame = tally.get_pandas_dataframe()
 
     # checks for user provided base units
-    base_units = get_tally_units(tally)
+    base_units = get_tally_units_dose(tally)
     print(f"tally {tally.name} base units {base_units}")
 
     if len(data_frame["mean"]) == 1:
         # just a single number in the tally result
         tally_result = data_frame["mean"].sum() * base_units[0]
     else:
-        # more than one number, a mesh tally
+        # more than one number, a mesh tally or several cell tallies
         tally_filter = tally.find_filter(filter_type=openmc.MeshFilter)
         shape = tally_filter.mesh.dimension.tolist()
         if 1 in shape:
@@ -86,35 +114,48 @@ def process_dose_tally(
         tally_result = tally_result.reshape(shape)
 
     if required_units:
-        tally_result = scale_tally(
+        scaled_tally_result = scale_tally(
             tally,
             tally_result,
             base_units[0],
             ureg[required_units],
             source_strength,
             volume,
         )
-        tally_result = convert_units([tally_result], [required_units])[0]
+        tally_in_required_units = scaled_tally_result.to(required_units)
+        return tally_in_required_units
 
     return tally_result
 
 
 def process_tally(
     tally,
-    required_units=None,
-    source_strength=None,
-    volume=None,
+    required_units: str = None,
+    source_strength: float = None,
+    volume: float = None,
 ):
-    """Processes the tally converting the tally with default units obtained
-    during simulation into the user specified units. In some cases
-    additional user inputs will be required. Units with"""
+    """Processes a tally converting the tally with default units obtained
+     during simulation into the user specified units.
 
-    # user makes use of StatePoint.get_tally to find tally
-    # passes the tally into this function along with the required units
-    # the tally can be returned with base units or with converted units
+    Args:
+        tally: The openmc.Tally object to convert the units of
+        required_units: The units to convert the energy and tally into
+        source_strength: In some cases the source_strength will be required
+            to convert the base units into the required units. This optional
+            argument allows the user to specify the source_strength when needed
+        volume: In some cases the volume will be required to convert the base
+            units into the required units. In the case of a regular mesh the
+            volume is automatically found. This optional argument allows the
+            user to specify the volume when needed or overwrite the
+            automatically calculated volume.
 
+    Returns:
+        The dose tally result in the required units
+    """
     data_frame = tally.get_pandas_dataframe()
 
+    check_for_unit_modifying_filters(tally)
+
     # checks for user provided base units
     base_units = get_tally_units(tally)
     print(f"tally {tally.name} base units {base_units}")
@@ -135,24 +176,19 @@ def process_tally(
             tally_result = tally_result.reshape(shape)
 
         if required_units:
-            tally_result = scale_tally(
+            scaled_tally_result = scale_tally(
                 tally,
                 tally_result,
                 base_units[0],
                 ureg[required_units],
                 source_strength,
                 volume,
             )
-            tally_result = convert_units([tally_result], [required_units])[0]
+            tally_in_required_units = scaled_tally_result.to(required_units)
+            return tally_in_required_units
 
     return tally_result
 
-def convert_units(value_to_convert, required_units):
-    converted_units = []
-    for value, required in zip(value_to_convert, required_units):
-        converted_units.append(value.to(required))
-
-    return converted_units
 
 def scale_tally(
     tally, tally_result, base_units, required_units, source_strength, volume
@@ -198,8 +234,8 @@ def scale_tally(
             volume = volume * ureg["1 / centimeter ** 3"]
         else:
             # volume required but not provided so it is found from the mesh
-            volume = compute_volume_of_voxels(tally, ureg)
-        
+            volume = compute_volume_of_voxels(tally) * ureg["1 / centimeter ** 3"]
+
         if length_diff == -3:
             tally_result = tally_result / volume
         if length_diff == 3:
@@ -208,18 +244,19 @@ def scale_tally(
     return tally_result
 
 
-def compute_volume_of_voxels(tally, ureg):
+def compute_volume_of_voxels(tally):
     tally_filter = tally.find_filter(filter_type=openmc.MeshFilter)
     if tally_filter:
         mesh = tally_filter.mesh
-        x = abs(mesh.lower_left[0] - mesh.upper_right[0])/mesh.dimension[0]
-        y = abs(mesh.lower_left[1] - mesh.upper_right[1])/mesh.dimension[1]
-        z = abs(mesh.lower_left[2] - mesh.upper_right[2])/mesh.dimension[2]
-        volume = x*y*z
-        return volume * ureg["1 / centimeter ** 3"]
+        x = abs(mesh.lower_left[0] - mesh.upper_right[0]) / mesh.dimension[0]
+        y = abs(mesh.lower_left[1] - mesh.upper_right[1]) / mesh.dimension[1]
+        z = abs(mesh.lower_left[2] - mesh.upper_right[2]) / mesh.dimension[2]
+        volume = x * y * z
+        return volume
     else:
         raise ValueError(f"volume could not be obtained from tally {tally}")
 
+
 def find_fusion_energy_per_reaction(reactants: str) -> float:
     """Finds the average fusion energy produced per fusion reaction in joules
     from the fuel type.
@@ -283,27 +320,74 @@ def get_cell_ids_from_tally_filters(tally):
     return cell_ids
 
 
+def get_tally_units_spectra(tally):
+
+    units = get_tally_units(tally)
+
+    # check it is a spectra tally by looking for a openmc.filter.EnergyFilter
+    for filter in tally.filters:
+        if isinstance(filter, openmc.filter.EnergyFilter):
+            # spectra tally has units for the energy as well as the flux
+            units = [ureg.electron_volt, units[0]]
+            return units
+
+    raise ValueError(
+        "units for spectra tally can't be found, an EnergyFilter was not present"
+    )
+
+
+def get_tally_units_dose(tally):
+
+    # An EnergyFunctionFilter is exspeted in dose tallies
+    units = get_tally_units(tally)
+
+    # check it is a dose tally by looking for a openmc.filter.EnergyFunctionFilter
+    for filter in tally.filters:
+        if isinstance(filter, openmc.filter.EnergyFunctionFilter):
+            print("filter is EnergyFunctionFilter")
+            # effective_dose
+            # dose coefficients have pico sievert cm **2
+            # flux has cm2 / simulated_particle units
+            # dose on a surface uses a current score (units of per simulated_particle) and is therefore * area to get pSv / source particle
+            # dose on a volume uses a flux score (units of cm2 per simulated particle) and therefore gives pSv cm**4 / simulated particle
+            units = [ureg.picosievert * ureg.centimeter * units[0]]
+            return units
+
+    raise ValueError(
+        "units for dose tally can't be found, an EnergyFunctionFilter was not present"
+    )
+
+
+def check_for_unit_modifying_filters(tally):
+    # check for EnergyFunctionFilter which modify the units of the tally
+    for filter in tally.filters:
+        if isinstance(filter, openmc.filter.EnergyFunctionFilter):
+            msg = (
+                "An EnergyFunctionFilter was found in the tally. This "
+                "modifies the tally units and the base units of the"
+                "EnergyFunctionFilter are not known to OpenMC. Therefore "
+                "the units of this tally can not be found. If you have "
+                "applied dose coefficients to an EnergyFunctionFilter "
+                "the units of these are known and yo can use the "
+                "get_tally_units_dose function instead of the "
+                "get_tally_units"
+            )
+            raise ValueError(msg)
+
+
 def get_tally_units(tally):
     """ """
 
+    if tally.scores == ["current"]:
+        units = get_particles_from_tally_filters(tally, ureg)
+        units = [units / (ureg.simulated_particle * ureg.centimeter ** 2)]
+
     if tally.scores == ["flux"]:
-        print('score is flux')
+        print("score is flux")
         # tally has units of particle-cm2 per simulated_particle
         # https://openmc.discourse.group/t/normalizing-tally-to-get-flux-value/99/4
         units = get_particles_from_tally_filters(tally, ureg)
         units = [units * ureg.centimeter / ureg.simulated_particle]
-        for filter in tally.filters:
-            if isinstance(filter, openmc.filter.EnergyFilter):
-                # spectra tally has units for the energy as well as the flux
-                units = [ureg.electron_volt, units[0]]
-            if isinstance(filter, openmc.filter.EnergyFunctionFilter):
-                print('filter is EnergyFunctionFilter')
-                # effective_dose
-                # dose coefficients have pico sievert cm **2
-                # flux has cm2 / simulated_particle units
-                # dose on a surface uses a current score (units of per simulated_particle) and is therefore * area to get pSv / source particle
-                # dose on a volume uses a flux score (units of cm2 per simulated particle) and therefore gives pSv cm**4 / simulated particle
-                units = [ureg.picosievert * ureg.centimeter * units[0]]
 
     elif tally.scores == ["heating"]:
         # heating units are eV / simulated_particle
@@ -322,8 +406,3 @@ def check_for_dimentionality_difference(units_1, units_2, unit_to_compare):
     units_1_time_power = units_1.dimensionality.get(unit_to_compare)
     units_2_time_power = units_2.dimensionality.get(unit_to_compare)
     return units_1_time_power - units_2_time_power
-
-
-# import pint
-# ureg = pint.UnitRegistry()
-# diff = check_for_dimentionality_difference_for_time(ureg['cm**2 / s'],ureg['m * s**2']

requirements-test.txt

@@ -1,5 +1,4 @@
 openmc_data_downloader
 openmc_dagmc_wrapper
 openmc_plasma_source
-# paramak
 pytest-cov>=2.12.1

tests/test_cell_effective_dose.py

@@ -13,7 +13,7 @@ def setUp(self):
         self.my_tally = statepoint.get_tally(name="2_neutron_effective_dose")
 
     def test_tally_base_units(self):
-        base_units = opp.get_tally_units(self.my_tally)
+        base_units = opp.get_tally_units_dose(self.my_tally)
         assert base_units[0].units == "centimeter ** 2 * neutron * picosievert / simulated_particle"
 
     def test_cell_tally_dose_no_processing(self):