Add support for probability table (URR)

[dodu94]

My bad if I missed it but is seems that the ACE parser do not support probability tables (URR block). Similarly to what I was saying in #15 I had a need for plotting them for V&V purposes.
Again, sorry for not contributing properly with a PR, I will try to reopen the conversation with IT on the C compiling business. In the meantime I will store here some code in case anyone else volunteers to implement it:

import pandas as pd
from matplotlib.axes import Axes
from matplotlib.figure import Figure


# CONSTANTS FOR TABLE PARSING
CUMULATIVE = "cumulative_probability"
TOTAL = "total"
ELASTIC = "elastic"
FISSION = "fission"
CAPTURE = "capture"  # (n, gamma)
HEATING = "heating_number"

J_TYPES = {
    2: TOTAL,
    3: ELASTIC,
    4: FISSION,
    5: CAPTURE,
    6: HEATING,
}

# Interpolation
LOGLOG = 5
LINLIN = 2
# Factors Flag
IS_XS = 0
IS_FACTOR = 1


class CDF:
    def __init__(
        self, energy: float, cdf: tuple[np.ndarray, np.ndarray], xs_type: str
    ) -> None:
        """Simple class to store CDF data at a specific energy point. Attributes
        are the same as input parameters.

        Parameters
        ----------
        energy : float
            incident neutron energy for which the data is valid
        cdf : tuple[np.ndarray, np.ndarray]
            probabilities and xs/factor values
        xs_type : str
            xs type. One of the J_TYPES values (i.e., TOTAL, ELASTIC, etc.)

        Raises
        ------
        ValueError
            If the xs_type is not supported (i.e. not among J_TYPES)
        """

        if xs_type not in J_TYPES.values():
            raise ValueError(f"Type {xs_type} not supported")

        self.energies = energy
        self.cdf = cdf
        self.xs_type = xs_type


class PTable:
    def __init__(
        self, table: pd.DataFrame, log_log: bool = False, pdf: bool = True
    ) -> None:
        """Simple class to store a probability table. The table is a pivoted dataframe
        where the index is the xs or multiplying factor values,
        the columns are the incident neutron energies and the values are the PDF or CDF.


        Parameters
        ----------
        table : pd.DataFrame
            pivoted datafram that is the result of all assembled CDFs or PDFs
        log_log : bool, optional
            interpolation method that should be used between energies. If True is log-log
            otherwise is lin-lin. by default False
        pdf : bool, optional
            if True values are PDFs, otherwise they are CDFs, by default True

        Attributes
        ----------
        table : pd.DataFrame
            pivoted datafram that is the result of all assembled CDFs or PDFs
        log_log : bool
            interpolation method that should be used between energies. If True is log-log
            otherwise is lin-lin.
        xs_type : str
            xs type. One of the J_TYPES values (i.e., TOTAL, ELASTIC, etc.)
        pdf : bool
            if True values are PDFs, otherwise they are CDFs
        """

        self.table = table
        self.log_log = log_log  # if false is lin-lin interpolation
        self.xs_type = table.index.name
        self.pdf = pdf

    # TODO: implement an interpolation method either log-log or lin-lin, 2D matrix

    def plot(self) -> tuple[Figure, Axes]:
        """Plot the PTable

        Returns
        -------
        tuple[Figure, Axes]
            classic matplotlib Figure and Axes
        """
        fig, ax = plt.subplots()

        if self.pdf:
            to_plot = "PDF"
        else:
            to_plot = "CDF"

        df = self.table
        X = df.columns.values
        Y = df.index.values

        ax.set_yscale("log")

        heatmap = ax.pcolormesh(X, Y, df.values, shading="auto", cmap="viridis")
        # Add a color bar
        fig.colorbar(heatmap, label=to_plot)
        # Add some labels
        ax.set_xlabel("Incident neutron energy [MeV]")
        ax.set_ylabel(self.xs_type)
        ax.set_title(f"{self.xs_type} {to_plot} table")

        return fig, ax


class PTableSection:
    def __init__(
        self,
        INT: int,
        ILF: int,
        IOA: int,
        IFF: int,
        CDFs: dict[tuple[float, str], CDF],
        pdf: bool = True,
    ) -> None:
        """Store the parsed probability tables from an ACE file [URR block].

        Parameters
        ----------
        INT : int
            Interpolation scheme for the probability table. 2 for lin-lin, 5 for log-log
        ILF : int
            Inelastic competition flag
        IOA : int
            Other absorption flag
        IFF : int
            Factors Flag
        CDFs : dict[tuple[float, str], CDF]
            dictionary of CDFs
        pdf : bool, optional
            if True, PDF Ptables will be assembled, eitherwise, CDF ones. by default True

        Attributes
        ----------
        INT : int
            Interpolation scheme for the probability table. 2 for lin-lin, 5 for log-log
        ILF : int
            Inelastic competition flag
        IOA : int
            Other absorption flag
        IFF : int
            Factors Flag
        CDFs : dict[tuple[float, str], CDF]
            dictionary of CDFs
        ptables : dict[str, PTable]
            dictionary of PTables

        Raises
        ------
        ValueError
            if the interpolation scheme is not valid
        ValueError
            if the Factors Flag is not valid
        """
        self.INT = INT  # Interpolation scheme for the probability table. 2 for lin-lin, 5 for log-log
        if INT not in [LOGLOG, LINLIN]:
            raise ValueError(f"Interpolation scheme {INT} not valid")
        self.ILF = ILF  # Inelastic competition flag.
        self.IOA = IOA  # Outher absorption flag.
        if IFF not in [IS_XS, IS_FACTOR]:
            raise ValueError(f"Factors Flag {IFF} not valid")
        self.IFF = IFF  # Factors Flag.
        self.CDFs = CDFs  # ordered CDFs
        self.ptables = self._compute_all_tables(pdf=pdf)

    def _compute_all_tables(self, pdf: bool = True) -> dict[str, PTable]:
        """Compute all the tables for the various cross section types

        Parameters
        ----------
        pdf : bool, optional
            if True assemble a PDF, eitherwise a CDF, by default True

        Returns
        -------
        dict[str, PTable]
            dictionary of PTables
        """
        tables = {}
        for xs_type in J_TYPES.values():
            tables[xs_type] = self._compute_table(xs_type, pdf=pdf)
        return tables

    def _compute_table(self, xs_type: str, pdf: bool = True) -> PTable:
        """Assemble a PTable from the CDFs parsed from the ACE file.

        Parameters
        ----------
        xs_type : str
            one of the J_TYPES values (i.e., TOTAL, ELASTIC, etc.)
        pdf : bool, optional
            if True assemble a PDF, eitherwise a CDF, by default True

        Returns
        -------
        PTable
            Assembled PTable

        Raises
        ------
        ValueError
            If the xs_type is not supported (i.e. not among J_TYPES)
        """
        dfs = []
        if xs_type not in J_TYPES.values():
            raise ValueError(f"Type {xs_type} not supported")

        if self.IFF == IS_XS:
            label_y = "XS [b]"
        elif self.IFF == IS_FACTOR:
            label_y = "XS multiplication factor"

        for (energy, xs), table in self.CDFs.items():
            if xs == xs_type:
                df = pd.DataFrame()
                df["CDF"] = table.cdf[0]
                df["PDF"] = table.cdf[0] - np.array([0] + table.cdf[0].tolist())[:-1]
                df[label_y] = table.cdf[1]
                df["Energy [MeV]"] = energy
                dfs.append(df)

        if pdf:
            to_plot = "PDF"
        else:
            to_plot = "CDF"

        df = pd.concat(dfs)
        table = df.pivot_table(index=label_y, columns="Energy [MeV]", values=to_plot).bfill()

        if self.INT == LOGLOG:
            log_log = True
        else:
            log_log = False

        return PTable(table, log_log=log_log, pdf=pdf)

    @classmethod
    def from_ace(cls, ace_table: ace.Table) -> PTableSection:
        """Parse the URR block from an ACE table and return a PTableSection object.

        Parameters
        ----------
        ace_table : ace.Table
            ACE table object

        Returns
        -------
        PTableSection
            PTableSection object
        """
        N = int(
            ace_table.xss[ace_table.jxs[23]]
        )  # Number of incident energies where there is a probability table.
        M = int(ace_table.xss[ace_table.jxs[23] + 1])  # Length of probability table.
        INT = int(
            ace_table.xss[ace_table.jxs[23] + 2]
        )  # Interpolation scheme for the probability table. 2 for lin-lin, 5 for log-log
        ILF = int(ace_table.xss[ace_table.jxs[23] + 3])  # Inelastic competition flag.
        IOA = int(ace_table.xss[ace_table.jxs[23] + 4])  # Outher absorption flag.
        IFF = int(ace_table.xss[ace_table.jxs[23] + 5])  # Factors Flag.
        START_E_INDEX = int(
            ace_table.jxs[23] + 6
        )  # Start of incident energies. (up to N)
        PTABLE = int(ace_table.jxs[23] + 6 + N)  # Start of probability table. (up to M)

        # --- parse the different tables ---
        # energies are the same for all tables
        energies = ace_table.xss[START_E_INDEX : START_E_INDEX + N]
        tables = {}
        # iterate on each energy
        for energy_idx, energy_val in enumerate(energies):
            start_p = PTABLE + 6 * M * (energy_idx)
            probabilities = ace_table.xss[start_p : start_p + M]  # same for all CDFs
            for j, xs_type in J_TYPES.items():
                start = start_p + M * (j - 1)
                values = ace_table.xss[start : start + M]
                ptable = CDF(energy_val, (probabilities, values), xs_type)
                tables[energy_val, xs_type] = ptable

        return cls(INT, ILF, IOA, IFF, tables)

# example of plotting
file = r"R:\AC_ResultsDB\Neutronics\01_DATABASE_MODELS\04_ACE_LIBRARIES\TESTING\JEFF4.0T1\54-xe-129g_293.6.ace"
ace_table = ace.get_table(file)
ptable_section = PTableSection.from_ace(ace_table)
ptable_section.ptables[HEATING].plot()

[MicahGale]

It is possible to install some tools locally to just your user: https://stackoverflow.com/questions/3212099/install-gcc-on-linux-with-no-root-privilege, though this may still be a violation of company policy, so proceed with caution.