From 2b1d979b39ff99cc90f4cc70f362808244309e3b Mon Sep 17 00:00:00 2001
From: Elliott Biondo <elliott.biondo@gmail.com>
Date: Fri, 12 Feb 2016 12:16:27 -0600
Subject: [PATCH] Adding capabilitiy of printing photon source totals

---
 pyne/r2s.py       | 27 +++++++++++++++++++++++++++
 scripts/r2s.py    | 16 ++++++++++++++--
 tests/test_r2s.py | 13 ++++++++++++-
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/pyne/r2s.py b/pyne/r2s.py
index 5c89f9ee6f..04a6402427 100644
--- a/pyne/r2s.py
+++ b/pyne/r2s.py
@@ -1,6 +1,7 @@
 from os.path import isfile
 from warnings import warn
 from pyne.utils import QAWarning
+import numpy as np
 
 from pyne.mesh import Mesh
 from pyne.mcnp import Meshtal
@@ -138,3 +139,29 @@ def photon_sampling_setup(mesh, phtn_src, tags):
     photon_source_to_hdf5(phtn_src)
     h5_file = phtn_src + ".h5"
     photon_source_hdf5_to_mesh(mesh, h5_file, tags)
+
+
+def total_photon_source_intensity(m, tag_name):
+    """This function reads mesh tagged with photon source densities and returns
+    the total photon emission desinty.
+
+    Parameters
+    ----------
+    m : PyNE Mesh
+       The mesh-based photon emission density distribution in p/cm3/s.
+    tag_name : str
+       The name of the tag on the mesh with the photon emission density information.
+
+    Returns
+    -------
+    intensity : float
+        The total photon emission density across the entire mesh (p/s).
+    """
+
+    sd_tag = m.mesh.getTagHandle(tag_name)
+    intensity = 0.
+    for idx, _, ve in m:
+        vol = m.elem_volume(ve)
+        ve_data = sd_tag[ve]
+        intensity += vol*np.sum(ve_data)
+    return intensity
diff --git a/scripts/r2s.py b/scripts/r2s.py
index bfb0d387f2..cac1583597 100644
--- a/scripts/r2s.py
+++ b/scripts/r2s.py
@@ -51,6 +51,9 @@
 # The prefix of the .h5m files containing the source density distributations for
 # each decay time.
 output: source
+# The name of the output files containing the total photon source intensities for
+# each decay time
+tot_phtn_src_intensities : total_photon_source_intensites.txt
 """
 
 alara_params =\
@@ -66,7 +69,7 @@
 
 output zone
        integrate_energy
-       photon_source  fendl3bin  phtn_src1 24 1.00E4 2.00E4 5.00E4 1.00E5
+       photon_source  fendl3bin  phtn_src 24 1.00E4 2.00E4 5.00E4 1.00E5
        2.00E5 3.00E5 4.00E5 6.00E5 8.00E5 1.00E6 1.22E6 1.44E6 1.66E6
        2.00E6 2.50E6 3.00E6 4.00E6 5.00E6 6.50E6 8.00E6 1.00E7 1.20E7
        1.40E7 2.00E7
@@ -142,16 +145,25 @@ def step2():
     structured = config.getboolean('general', 'structured')
     decay_times = config.get('step2', 'decay_times').split(',')
     output = config.get('step2', 'output')
+    tot_phtn_src_intensities = config.get('step2', 'tot_phtn_src_intensities')
+    tag_name = "source_density"
 
     h5_file = 'phtn_src.h5'
     if not isfile(h5_file):
         photon_source_to_hdf5('phtn_src')
+    intensities = "Total photon source intensities (p/s)\n"
     for i, dc in enumerate(decay_times):
         print('Writing source for decay time: {0}'.format(dc))
         mesh = Mesh(structured=structured, mesh='blank_mesh.h5m')
-        tags = {('TOTAL', dc): 'source_density'}
+        tags = {('TOTAL', dc): tag_name}
         photon_source_hdf5_to_mesh(mesh, h5_file, tags)
         mesh.mesh.save('{0}_{1}.h5m'.format(output, i+1))
+        intensity = total_photon_source_intensity(m, tag_name)
+        intensities += "{0}: {1}".format(dc, intensity)
+
+    with open(tot_phtn_src_intensities, 'r') as f:
+        f.write(intensities)
+
     print('R2S step2 complete.')
 
 def main():
diff --git a/tests/test_r2s.py b/tests/test_r2s.py
index e53bf20aef..149858bf45 100644
--- a/tests/test_r2s.py
+++ b/tests/test_r2s.py
@@ -20,7 +20,7 @@
 
 from pyne.utils import QAWarning
 warnings.simplefilter("ignore", QAWarning)
-from pyne.r2s import irradiation_setup, photon_sampling_setup
+from pyne.r2s import irradiation_setup, photon_sampling_setup, total_photon_source_intensity
 from pyne.material import Material
 from pyne.mesh import Mesh, IMeshTag
 from pyne.mcnp import Meshtal
@@ -275,3 +275,14 @@ def test_photon_sampling_setup_unstructured():
     for i, mat, ve in m:
         assert_array_equal(m.tag1[i], exp_tag1[i])
         assert_array_equal(m.tag2[i], exp_tag2[i])
+
+
+def test_total_photon_source_intensity():
+
+    m = Mesh(structured = True, structured_coords=[[0, 1, 2],[0, 1, 3], [0, 1]])
+    m.source_density = IMeshTag(2, float)
+    m.source_density[:] = [[1., 2.], [3., 4.], [5., 6.], [7., 8.]]
+
+    intensity = total_photon_source_intensity(m, "source_density")
+    assert_equal(intensity, 58)
+