diff --git a/sandy/core/xs.py b/sandy/core/xs.py index 2f979c5c..a11c8c49 100644 --- a/sandy/core/xs.py +++ b/sandy/core/xs.py @@ -465,30 +465,20 @@ def _reconstruct_sums(self, drop=True, inplace=False): """ df = self.data.copy() for mat in self.data.columns.get_level_values("MAT").unique(): - for parent, daughters in sorted(redundant_xs.items(), reverse=True): - daughters = [x for x in daughters if x in df[mat]] - if daughters: - df[mat,parent] = df[mat][daughters].sum(axis=1) + parent = pd.Index(self.__class__.redundant_xs.keys()) + mask = parent.isin(df.columns.get_level_values('MT')) + for father in parent[mask]: + mask = df.columns.get_level_values('MT')\ + .isin(self.__class__.redundant_xs[father]) + df[mat, father] = df.loc[:, mask].sum(axis=1) # keep only mts present in the original file if drop: - todrop = [x for x in df[mat].columns if x not in self.data[mat].columns] - cols_to_drop = pd.MultiIndex.from_product([[mat], todrop]) - df.drop(cols_to_drop, axis=1, inplace=True) + columns = df.columns.intersection(self.data.columns) + df = df.loc[:, columns] if inplace: self.data = df else: - return Xs(df) -# frame = self.copy() -# for mat in frame.columns.get_level_values("MAT").unique(): -# for parent, daughters in sorted(Xs.redundant_xs.items(), reverse=True): -# daughters = [ x for x in daughters if x in frame[mat]] -# if daughters: -# frame[mat,parent] = frame[mat][daughters].sum(axis=1) -# # keep only mts present in the original file -# if drop: -# todrop = [ x for x in frame[mat].columns if x not in self.columns.get_level_values("MT") ] -# frame.drop(pd.MultiIndex.from_product([[mat], todrop]), axis=1, inplace=True) -# return Xs(frame) + return self.__class__(df) def _perturb(self, pert, method=2, **kwargs): """Perturb cross sections/nubar given a set of perturbations.