_get_edistr_cov method + bug solve

notebooks/notebook_Edistr_cov_from_enf6.ipynb

@@ -0,0 +1,112 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sandy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import seaborn as sns\n",
+    "import matplotlib.pyplot as plt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "endf6 = sandy.get_endf6_file('jeff_33','xs', 922380)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_corr(cov_obj, ax, **kwargs):\n",
+    "    cov_obj_ = sandy.CategoryCov(cov_obj) if not isinstance(cov_obj, sandy.CategoryCov) else cov_obj\n",
+    "    add = {\"cbar\": True, \"vmin\": -1, \"vmax\": 1, \"cmap\": \"RdBu\"}\n",
+    "    for k, v in kwargs.items():\n",
+    "        add[k] = v\n",
+    "    ax = sns.heatmap(cov_obj_.get_corr().data, ax=ax, **add)\n",
+    "    return ax"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot energy distribution correlation matrix from Endf6 files (MF=35)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cov = endf6._get_edistr_cov()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=(12, 12), dpi=100)\n",
+    "plot_corr(cov, ax=ax)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot the energy distribution of the incoming and outgoing energy of 1.0e-5:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=(6, 6), dpi=100)\n",
+    "plot_corr(cov.data.loc[(9237, 18, 1.0e-5, 1.0e-5), (9237, 18, 1.0e-5, 1.0e-5)], ax=ax)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python3 (sandy-devel)",
+   "language": "python",
+   "name": "sandy-devel"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.10"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

sandy/core/cov.py

@@ -2395,9 +2395,8 @@ def segmented_pivot_table(data_stack, index, columns, values, rows=10000000):
             index=index,
             columns=columns,
             values=values,
-            fill_value=0,
             aggfunc=np.sum,
-            )
+            ).fillna(0)
         pivot_matrix.append(partial_pivot)
     pivot_matrix = pd.concat(pivot_matrix).fillna(0)
     # Because the default axis to concatenate is the 0, some duplicate

sandy/core/endf6.py

@@ -2475,3 +2475,66 @@ def covariance_info(self, nubar=True, xs=True, mubar=True, chi=True):
             'chi': run_chi if chi else False,
             }
         return cov_info
+
+    def _get_edistr_cov(self, mt=range(1, 10000), mat=range(1, 10000)):
+        """
+        Extract energy distribution coefficients covariance matrix
+        from `sandy.Endf6` instance.
+
+        Parameters
+        ----------
+        mat : `int`, optional
+            MAT number. The default is None.
+        mt : `int`, optional
+            MT number. The default is None.
+
+        Returns
+        -------
+        `sandy.CategoryCov`
+            Covariance matrix.
+
+        """
+        listmt_ = [mt] if isinstance(mt, int) else mt
+        listmat_ = [mat] if isinstance(mat, int) else mat
+        tape = self.filter_by(listmf=[35],
+                              listmt=listmt_,
+                              listmat=listmat_)
+        data = []
+        for mat, mf, mt in tape.data:
+            sec = tape.read_section(mat, mf, mt)
+            for sub, sub_info in sec["SUB"].items():
+                Ek = np.array(sub_info["EK"])
+                Fkk = np.array(sub_info["FKK"])
+                NE = sub_info["NE"]
+                # Covariance
+                cov_upper = np.zeros([NE-1, NE-1])
+                indices = np.triu_indices(NE-1)
+                cov_upper[indices] = Fkk
+                cov = sandy.cov.triu_matrix(cov_upper).data.values
+                # Normalize covariance matrix dividing by the energy bin.
+                dE = 1./(Ek[1:]-Ek[:-1])
+                cov = sandy.cov.corr2cov(cov, dE)
+                # Add zero row and column at the end of the matrix
+                cov = np.insert(cov, cov.shape[0], [0]*cov.shape[1], axis=0)
+                cov = np.insert(cov, cov.shape[1], [0]*cov.shape[0], axis=1)
+                # multiindex:
+                index = pd.MultiIndex.from_product(
+                    [[mat], [mt], [sub_info["ELO"]], [sub_info["EHI"]], Ek],
+                    names=["MAT", "MT", "ELO", "EHI", "E"],
+                )
+                columns = pd.MultiIndex.from_product(
+                    [[mat], [mt], [sub_info["ELO"]], [sub_info["EHI"]], Ek],
+                    names=["MAT1", "MT1", "ELO1", "EHI1", "E1"],
+                )
+                df = pd.DataFrame(cov,
+                                  index=index,
+                                  columns=columns)\
+                       .stack(level=["MAT1", "MT1", "ELO1", "EHI1", "E1"])\
+                       .rename("VAL").reset_index()
+                data.append(df)
+        data = pd.concat(data)
+        return sandy.CategoryCov.from_stack(data,
+                                            index=["MAT", "MT", "ELO", "EHI", "E"],
+                                            columns=["MAT1", "MT1", "ELO1", "EHI1", "E1"],
+                                            values='VAL',
+                                            kind='all')