Chi Apply perturbation (#343)

* docstring and minor changes

* ApplyPerturbation

* test fix

* test fix

* test fix

* test fix

* test fix

* Update endf6.py

* Update endf6.py

* Update samples.py

---------

Co-authored-by: Luca Fiorito <[email protected]>

sandy/core/endf6.py

@@ -2651,6 +2651,28 @@ def apply_perturbations_xs(self, smps, processes=1, pendf=None, njoy_kws={}, **k
         >>> assert outs_31[0]["pendf"].data == outs_31[1]["pendf"].data
         >>> assert outs_31[0]["endf6"].data != outs_31[1]["endf6"].data
 
+        Default use case for chi and xs, perturbed together.
+
+        >>> smps_ = tape.get_perturbations(2, njoy_kws=dict(err=1, nubar=False, mubar=False), smp_kws=dict(seed33=3, seed35=5))
+        >>> outs_33_35 = tape.apply_perturbations(smps_, njoy_kws=dict(err=1), processes=1)
+
+        Compare to individual xs and chi perturbations with same seed.
+
+        >>> outs_33_ = tape.apply_perturbations({33: smps_[33]}, njoy_kws=dict(err=1), processes=1)
+        >>> outs_35 = tape.apply_perturbations({35: smps_[35]}, njoy_kws=dict(err=1), processes=1)
+        
+        >>> for i in range(2):
+        ...    assert(outs_33_[i]["endf6"].data == tape.data)
+        ...    assert(outs_35[i]["endf6"].data != tape.data)
+        ...    assert(outs_35[i]["endf6"].data == outs_33_35[i]["endf6"].data)
+        ...    assert(outs_33_[i]["pendf"].data != outs_35[i]["pendf"].data)
+        ...    assert(outs_33_[i]["pendf"].data == outs_33_35[i]["pendf"].data)
+
+        >>> assert outs_33_[0]["pendf"].data != outs_33_[1]["pendf"].data
+        >>> assert outs_33_[0]["endf6"].data == outs_33_[1]["endf6"].data
+        >>> assert outs_35[0]["pendf"].data == outs_35[1]["pendf"].data
+        >>> assert outs_35[0]["endf6"].data != outs_35[1]["endf6"].data
+
         Check that redundant nubar is also perturbed.
 
         >>> nu0 = sandy.Xs.from_endf6(outs_31[0]["endf6"].filter_by(listmt=[452, 455, 456]))
@@ -2699,7 +2721,7 @@ def apply_perturbations_xs(self, smps, processes=1, pendf=None, njoy_kws={}, **k
         >>> assert outs1[0]["pendf"].write_string() == outs2[0]["pendf"].write_string()
         """
 
-        if 33 not in smps and 31 not in smps:
+        if 33 not in smps and 31 not in smps and 35 not in smps:
             logging.info("no perturbation coefficient was found.")
             return
 
@@ -2713,7 +2735,10 @@ def apply_perturbations_xs(self, smps, processes=1, pendf=None, njoy_kws={}, **k
             data["pnu"] = smps[31].iterate_xs_samples()
         if 33 in smps:
             data["pxs"] = smps[33].iterate_xs_samples()
-        # this I'd like to change
+        if 35 in smps:
+            # At this level, the samples have the same index
+            # as a covariance matrix and can be treated as xs
+            data["pchi"] = smps[35].iterate_xs_samples()
 
         if processes == 1:
             outs = {}
@@ -2897,7 +2922,7 @@ def endf6_perturb_worker(e6, pendf, n,
                          pxs=None,
                          pnu=None,
                          plpc=None,
-                         pedistr=None,
+                         pchi=None,
                          verbose=False,
                          to_ace=False,
                          to_file=False,
@@ -2970,8 +2995,23 @@ def endf6_perturb_worker(e6, pendf, n,
         pass
 
     # apply edistr perturbation
-    if pedistr is not None:
-        pass
+    if pchi is not None:
+        # Applies the same perturbation to all incident particle energies and K
+        edistr_pert = []
+        for (ein, k), df in sandy.Edistr.from_endf6(endf6_pert).data.groupby(['EIN', 'K']):
+            dummy_xs = sandy.Xs(
+                df.rename({"EOUT": "E"}, axis=1).set_index(["MAT","MT"])[["E","VALUE"]].pivot(columns="E").T.droplevel(level=0)
+            )
+            dummy_xs_pert = sandy.core.xs.xs_perturb_worker(dummy_xs, n, pchi, verbose=verbose)
+            edistr_pert.append(
+                dummy_xs_pert.data.stack().stack().        # multiple column index to columns
+                to_frame().reset_index().                  # Edistr.data has every info in columns
+                rename({"E": "EOUT", 0: "VALUE"}, axis=1). # rename columns to match Edistr.data
+                assign(K=k, EIN=ein)[["MAT", "MT", "K", "EIN", "EOUT", "VALUE"]]  # sort columns to match Edistr.data
+            )
+        endf6_pert = sandy.Edistr(
+                                pd.concat(edistr_pert, ignore_index=True)
+                                ).normalize().to_endf6(endf6_pert).update_intro()
 
     # apply xs perturbation
     if pxs is not None:

sandy/core/samples.py

@@ -73,8 +73,6 @@ def read_fy_samples(file='PERT_MF8_MT454.xlsx'):
 
     # same sorting structure as when it was produced in get_perturbations_fy
     smp = pd.concat(smp, ignore_index=True).sort_values(by=["ZAM", "E", "ZAP", "SMP"])
-
-    # for some reasons get_perturbations_fy gives me this
 
     return smp
 
@@ -259,17 +257,14 @@ def iterate_xs_samples(self):
 
 
         In this example the original covariance contains data for MT=1 and MT=51.
-
         >>> endf6 = sandy.get_endf6_file('jeff_33', 'xs', 942400)
         >>> smps = endf6.get_perturbations(1, njoy_kws=dict(err=1, chi=False, mubar=False, nubar=False, errorr33_kws=dict(mt=[1, 51])))[33]
 
         Then, since MT=1 is redundant, samples are passed to its partial components (MT=2 and MT=3).
-
         >>> expected = pd.MultiIndex.from_product([[9440], [1, 51] + list(sandy.redundant_xs[1])], names=["MAT", "MT"])
         >>> assert next(smps.iterate_xs_samples())[1].columns.equals(expected)
         
         If case one of the partial components already has samples, i.e., MT=2...
-
         >>> endf6 = sandy.get_endf6_file('jeff_33', 'xs', 942400)
         >>> smps = endf6.get_perturbations(1, njoy_kws=dict(err=1, chi=False, mubar=False, nubar=False, errorr33_kws=dict(mt=[1, 2, 51])))[33]
 
@@ -278,6 +273,15 @@ def iterate_xs_samples(self):
 
         >>> expected = pd.MultiIndex.from_product([[9440], [1, 2, 51]], names=["MAT", "MT"])
         >>> assert next(smps.iterate_xs_samples())[1].columns.equals(expected)
+
+        Default use case for MF35.
+
+        >>> tape = sandy.get_endf6_file("jeff_33", "xs", 942390)
+        >>> smps = tape.get_perturbations(2, njoy_kws=dict(err=1, nubar=False, mubar=False))
+
+        Check that output is not empty, and with correct shape.
+
+        >>> assert(next(smps[35].iterate_xs_samples())[1].shape == (240, 5))
         """
         levels = Xs._columnsnames
         df = self.data.unstack(level=levels)
@@ -322,7 +326,8 @@ def _mean_convergence(self):
         foo = lambda x: smp.loc[:x].mean()
         return pd.DataFrame(map(foo, rng), index=rng)
 
-    def from_excel(file, beg=None, end=None):   # this should be corrected as a class method
+    @classmethod
+    def from_excel(cls, file, beg=None, end=None):
         """
         Read perturbation coefficients (for nubar and xs) from excel file.
         The file format is compatible with what written in
@@ -355,5 +360,4 @@ def from_excel(file, beg=None, end=None):   # this should be corrected as a clas
 
         df = df.iloc[:, loc+2:].loc[:, beg:end].reset_index(drop=True)
         df.index = idx
-        smp = Samples(df)
-        return smp
+        return cls(df)