Handle the different cases for tallying nuclear heating in OpenMC

src/jade/helper/openmc.py

@@ -25,6 +25,7 @@
 logger = logging.getLogger(__name__)
 
 PAT_DIGITS = re.compile(r"\d+")
+CHARGED_SECONDARY_PARTICLES = ("electron", "positron")
 
 
 @dataclass
@@ -525,12 +526,111 @@ def _get_tally_data(self, tally: openmc.Tally):
         df = self._normalise_df(df)
         return df
 
+    @staticmethod
+    def _particle_bins(tally: openmc.Tally) -> set[str]:
+        """Get the particle bins of a tally
+
+        Parameters
+        ----------
+        tally : openmc.Tally
+            openmc tally
+
+        Returns
+        -------
+        set[str]
+            particle bins of the tally, empty if it has no particle filter
+        """
+        for tally_filter in tally.filters:
+            if isinstance(tally_filter, openmc.ParticleFilter):
+                return set(tally_filter.bins)
+        return set()
+
+    @staticmethod
+    def _same_heating_domain(tally1: openmc.Tally, tally2: openmc.Tally) -> bool:
+        """Check if two tallies score over the same phase space, ignoring the
+        particle filter, so that their results can be summed bin by bin
+
+        Parameters
+        ----------
+        tally1 : openmc.Tally
+            first openmc tally
+        tally2 : openmc.Tally
+            second openmc tally
+
+        Returns
+        -------
+        bool
+            True if the tallies have the same filters (apart from the particle
+            filter) and the same nuclides
+        """
+        filters1 = [
+            f for f in tally1.filters if not isinstance(f, openmc.ParticleFilter)
+        ]
+        filters2 = [
+            f for f in tally2.filters if not isinstance(f, openmc.ParticleFilter)
+        ]
+        if len(filters1) != len(filters2):
+            return False
+        if list(tally1.nuclides) != list(tally2.nuclides):
+            return False
+
+        def _equal(filter1, filter2) -> bool:
+            try:
+                return bool(filter1 == filter2)
+            except ValueError:
+                return False
+
+        return all(any(_equal(f1, f2) for f2 in filters2) for f1 in filters1)
+
+    @staticmethod
+    def _fold_charged_into_photon(df: pd.DataFrame) -> pd.DataFrame:
+        """Sum the electron and positron heating rows into the matching photon
+        heating rows. Before OpenMC 0.16.0 the heating from photon-induced
+        electrons and positrons is scored in the electron and positron bins,
+        from 0.16.0 it is attributed to the parent photon and those bins are
+        zero, so the sum is valid for both.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            tally dataframe, possibly containing electron and positron rows
+
+        Returns
+        -------
+        pd.DataFrame
+            tally dataframe where heating electron and positron rows have been
+            summed into the photon ones
+        """
+        if "particle" not in df.columns:
+            return df
+        to_fold = (df["score"] == "heating") & df["particle"].isin(
+            CHARGED_SECONDARY_PARTICLES
+        )
+        if not to_fold.any():
+            return df
+        columns = df.columns
+        df = df.copy()
+        df.loc[to_fold, "particle"] = "photon"
+        df["variance"] = df["std. dev."].pow(2)
+        keys = [col for col in columns if col not in ["mean", "std. dev."]]
+        # Errors are combined in quadrature
+        df = (
+            df.groupby(keys, sort=False, dropna=False)[["mean", "variance"]]
+            .sum()
+            .reset_index()
+        )
+        df["std. dev."] = df["variance"].pow(0.5)
+        return df[columns]
+
     def _combine_heating_tallies(self, heating_tallies: dict) -> dict:
-        """Extract tally data from statepoint file
+        """Combine photon, electron and positron heating tallies. Electron and
+        positron heating can be scored either in the same tally of the photon
+        heating (additional particle bins) or in separate tallies with the same
+        filters. In both cases it is summed into the photon heating.
 
         Parameters
         ----------
-        tallies : dict
+        heating_tallies : dict
             Dictionary containing all OpenMC heating tallies
 
         Returns
@@ -539,40 +639,33 @@ def _combine_heating_tallies(self, heating_tallies: dict) -> dict:
             dictionary of pandas dataframes containing combined tally data
         """
         photon_tallies = {}
+        charged_tallies = {}
         heating_tallies_df = {}
         for id, tally in heating_tallies.items():
-            try:
-                particle_filter = tally.find_filter(openmc.ParticleFilter)
-            except ValueError:
-                particle_filter = None
-            if particle_filter:
-                if "neutron" in particle_filter.bins:
-                    heating_tallies_df[id] = self._get_tally_data(tally)
-                if "photon" in particle_filter.bins:
-                    photon_tallies[id] = tally
-                    heating_tallies_df[id] = self._get_tally_data(tally)
+            particle_bins = self._particle_bins(tally)
+            if particle_bins and particle_bins.issubset(CHARGED_SECONDARY_PARTICLES):
+                charged_tallies[id] = tally
+            elif "photon" in particle_bins:
+                photon_tallies[id] = tally
             else:
                 heating_tallies_df[id] = self._get_tally_data(tally)
         for id, photon_tally in photon_tallies.items():
-            photon_cell_filter = photon_tally.find_filter(openmc.CellFilter)
-            for _, tally in heating_tallies.items():
-                try:
-                    particle_filter = tally.find_filter(openmc.ParticleFilter)
-                except ValueError:
-                    particle_filter = None
-                cell_filter = tally.find_filter(openmc.CellFilter)
-                if particle_filter and cell_filter:
-                    if photon_cell_filter == cell_filter:
-                        if (
-                            "electron" in particle_filter.bins
-                            or "positron" in particle_filter.bins
-                        ):
-                            tally_df = self._get_tally_data(tally)
-                            heating_tallies_df[id]["mean"] += tally_df["mean"]
-                            heating_tallies_df[id]["std. dev."] = (
-                                heating_tallies_df[id]["std. dev."].pow(2)
-                                + tally_df["std. dev."].pow(2)
-                            ).pow(0.5)
+            photon_bins = self._particle_bins(photon_tally)
+            dfs = [self._get_tally_data(photon_tally)]
+            for tally in charged_tallies.values():
+                if not self._same_heating_domain(photon_tally, tally):
+                    continue
+                tally_df = self._get_tally_data(tally)
+                # only add the particles not already scored in the photon tally
+                dfs.append(
+                    tally_df[
+                        (tally_df["score"] == "heating")
+                        & ~tally_df["particle"].isin(photon_bins)
+                    ]
+                )
+            heating_tallies_df[id] = self._fold_charged_into_photon(
+                pd.concat(dfs, ignore_index=True)
+            )
         return heating_tallies_df
 
     def tallies_to_dataframes(self) -> dict:

tests/helper/test_openmc.py

@@ -4,10 +4,14 @@
 
 import pytest, os
 
+import pandas as pd
+
 from jade.helper.__optionals__ import OMC_AVAIL
 from tests.post.resources import openmc as resources
 
 if OMC_AVAIL:
+    import openmc
+
     import jade.helper.openmc as omc
 
 
@@ -43,3 +47,155 @@ def test_tallies_to_dataframes(self):
             tallies[24]["std. dev."][5]
             / out.cell_data.cell_volumes[tallies[24]["cell"][5]]
         )
+        # photon heating includes the separate electron and positron tallies
+        assert 561 not in tallies and 562 not in tallies
+        raw_sum = sum(
+            out._get_tally_data(out.statepoint.tallies[i])["mean"]
+            for i in (56, 561, 562)
+        )
+        assert list(tallies[56]["mean"]) == pytest.approx(list(raw_sum))
+
+
+def _heating_df(cells, particles, means, stds, score="heating", nuclide="total"):
+    """Build a dataframe with the same layout as openmc get_pandas_dataframe"""
+    rows = []
+    i = 0
+    for cell in cells:
+        for particle in particles:
+            rows.append([cell, particle, nuclide, score, means[i], stds[i]])
+            i += 1
+    return pd.DataFrame(
+        rows, columns=["cell", "particle", "nuclide", "score", "mean", "std. dev."]
+    )
+
+
+def _tally(tally_id, particles, filters, scores=("heating",)):
+    tally = openmc.Tally(tally_id=tally_id)
+    tally.filters = [*filters, openmc.ParticleFilter(particles)]
+    tally.scores = list(scores)
+    return tally
+
+
+@pytest.mark.skipif(not OMC_AVAIL, reason="OpenMC is not available")
+class TestHeatingCombination:
+    def test_fold_single_tally(self):
+        particles = ["neutron", "photon", "electron", "positron"]
+        df = _heating_df(
+            [1, 2], particles, [1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8]
+        )
+        out = omc.OpenMCStatePoint._fold_charged_into_photon(df)
+        assert list(out.columns) == list(df.columns)
+        assert list(out["particle"]) == ["neutron", "photon"] * 2
+        assert list(out["cell"]) == [1, 1, 2, 2]
+        assert list(out["mean"]) == [1, 9, 5, 21]
+        assert out["std. dev."][1] == pytest.approx((4 + 9 + 16) ** 0.5)
+        assert out["std. dev."][3] == pytest.approx((36 + 49 + 64) ** 0.5)
+
+    def test_fold_zero_charged(self):
+        # OpenMC >= 0.16.0 attributes the charged particle heating to photons
+        df = _heating_df([1], ["photon", "electron", "positron"], [2, 0, 0], [1, 0, 0])
+        out = omc.OpenMCStatePoint._fold_charged_into_photon(df)
+        assert len(out) == 1
+        assert out["mean"][0] == 2
+        assert out["std. dev."][0] == pytest.approx(1)
+
+    def test_fold_other_scores_untouched(self):
+        heating = _heating_df([1], ["photon", "electron"], [1, 2], [1, 1])
+        flux = _heating_df([1], ["photon", "electron"], [3, 4], [1, 1], score="flux")
+        df = pd.concat([heating, flux], ignore_index=True)
+        out = omc.OpenMCStatePoint._fold_charged_into_photon(df)
+        assert list(out["score"]) == ["heating", "flux", "flux"]
+        assert list(out["particle"]) == ["photon", "photon", "electron"]
+        assert list(out["mean"]) == [3, 3, 4]
+
+    def test_fold_no_particle(self):
+        df = _heating_df([1], ["photon"], [1], [1]).drop(columns="particle")
+        out = omc.OpenMCStatePoint._fold_charged_into_photon(df)
+        assert out.equals(df)
+
+    def _combine(self, monkeypatch, tallies: dict, dfs: dict) -> dict:
+        sp = object.__new__(omc.OpenMCStatePoint)
+        monkeypatch.setattr(sp, "_get_tally_data", lambda tally: dfs[tally.id])
+        return sp._combine_heating_tallies(tallies)
+
+    def test_combine_separate_tallies(self, monkeypatch):
+        materials = openmc.MaterialFilter([1, 2])
+        tallies = {
+            1: _tally(1, ["neutron"], [materials]),
+            2: _tally(2, ["photon"], [materials]),
+            3: _tally(3, ["electron"], [materials]),
+            4: _tally(4, ["positron"], [materials]),
+        }
+        dfs = {
+            1: _heating_df([1, 2], ["neutron"], [10, 20], [1, 1]),
+            2: _heating_df([1, 2], ["photon"], [1, 2], [1, 1]),
+            3: _heating_df([1, 2], ["electron"], [3, 4], [1, 1]),
+            4: _heating_df([1, 2], ["positron"], [5, 6], [1, 1]),
+        }
+        out = self._combine(monkeypatch, tallies, dfs)
+        assert set(out.keys()) == {1, 2}
+        assert list(out[1]["mean"]) == [10, 20]
+        assert list(out[2]["particle"]) == ["photon", "photon"]
+        assert list(out[2]["mean"]) == [9, 12]
+        assert out[2]["std. dev."][0] == pytest.approx(3**0.5)
+
+    def test_combine_electron_positron_tally(self, monkeypatch):
+        cells = openmc.CellFilter([1, 2])
+        tallies = {
+            2: _tally(2, ["photon"], [cells]),
+            3: _tally(3, ["electron", "positron"], [cells]),
+        }
+        dfs = {
+            2: _heating_df([1, 2], ["photon"], [1, 2], [1, 1]),
+            3: _heating_df([1, 2], ["electron", "positron"], [3, 4, 5, 6], [1] * 4),
+        }
+        out = self._combine(monkeypatch, tallies, dfs)
+        assert set(out.keys()) == {2}
+        assert list(out[2]["mean"]) == [8, 13]
+
+    def test_combine_single_tally(self, monkeypatch):
+        cells = openmc.CellFilter([1])
+        particles = ["neutron", "photon", "electron", "positron"]
+        tallies = {5: _tally(5, particles, [cells])}
+        dfs = {5: _heating_df([1], particles, [1, 2, 3, 4], [1] * 4)}
+        out = self._combine(monkeypatch, tallies, dfs)
+        assert list(out[5]["particle"]) == ["neutron", "photon"]
+        assert list(out[5]["mean"]) == [1, 9]
+
+    def test_combine_no_double_counting(self, monkeypatch):
+        cells = openmc.CellFilter([1])
+        tallies = {
+            2: _tally(2, ["photon", "electron"], [cells]),
+            3: _tally(3, ["electron", "positron"], [cells]),
+        }
+        dfs = {
+            2: _heating_df([1], ["photon", "electron"], [1, 2], [1, 1]),
+            3: _heating_df([1], ["electron", "positron"], [2, 4], [1, 1]),
+        }
+        out = self._combine(monkeypatch, tallies, dfs)
+        assert list(out[2]["mean"]) == [7]
+
+    def test_combine_different_filters(self, monkeypatch):
+        cells = openmc.CellFilter([1])
+        energy = openmc.EnergyFilter([0.0, 1e6, 2e7])
+        other_cells = openmc.CellFilter([2])
+        tallies = {
+            2: _tally(2, ["photon"], [cells, energy]),
+            3: _tally(3, ["electron"], [cells]),
+            4: _tally(4, ["positron"], [other_cells]),
+        }
+        dfs = {
+            2: _heating_df([1], ["photon"], [1], [1]),
+            3: _heating_df([1], ["electron"], [2], [1]),
+            4: _heating_df([2], ["positron"], [4], [1]),
+        }
+        out = self._combine(monkeypatch, tallies, dfs)
+        assert list(out[2]["mean"]) == [1]
+
+    def test_combine_no_particle_filter(self, monkeypatch):
+        tally = openmc.Tally(tally_id=7)
+        tally.filters = [openmc.CellFilter([1])]
+        tally.scores = ["heating"]
+        df = _heating_df([1], ["photon"], [1], [1]).drop(columns="particle")
+        out = self._combine(monkeypatch, {7: tally}, {7: df})
+        assert out[7].equals(df)