Create perturbed FY files with ad-hoc covariance matrix (all energies)¶

sensitivity = rdd.get_chain_yield_sensitivity()
chfy = sandy.fy.get_chain_yields()
set([(row.ZAM, row.E) for i, row in chfy.iterrows()])

nsmp = 1000  # sample size
smp = {}
acovs = {}
rcovs = {}
acovs_post = {}
rcovs_post = {}
for e, etext in zip(energies, ["thermal", "fast", "high_energy"]):  # change this for the nuclide you choose

    # Run GLS update algorithm to create covariance matrix
    mask = f"E=={e} & MT==454"
    ify = fy.data.query(mask).set_index("ZAP")
    absolute_cov = sandy.CategoryCov.from_stdev(ify.DFY)
    relative_cov = sandy.CategoryCov.from_stdev((ify.DFY / ify.FY).fillna(0))

    acovs[e] = absolute_cov
    rcovs[e] = relative_cov

    
    extra = chfy.query(f"E=='{etext}' & ZAM=={za*10}").set_index("A")
    y_extra = extra.CHY
    Vy_extra = sandy.CategoryCov.from_stdev(extra.DCHY)

    fy_post, cov_post = gls_update(fy, za*10, e, sensitivity, y_extra, Vy_extra)

    ify_post = fy_post.data.query(mask).set_index("ZAP")
    absolute_cov_post = cov_post
    relative_cov_post = cov_post.data / (ify_post.FY.values.reshape(-1, 1) @ ify_post.FY.values.reshape(1, -1))
    # Symmetrize the resulting matrix
    relative_cov_post = sandy.CategoryCov((relative_cov_post + relative_cov_post.T).fillna(0) / 2)

    acovs_post[e] = absolute_cov_post
    rcovs_post[e] = relative_cov_post


    # Plot the generated covariance data and compare to original data 
    fig, axs = plt.subplots(1, 3, figsize=(12, 4))

    ax = axs[0]
    ify.reset_index().plot(x="ZAP", y="FY", ax=ax, kind="scatter")
    ify_post.reset_index().plot(x="ZAP", y="FY",ax=ax, kind="scatter", color="tomato", marker="x")
    ax.tick_params(labelbottom=False)
    ax.set(yscale="log")
    
    ax = axs[1]
    relative_cov.get_std().reset_index().plot(x="ZAP", y="STD", kind="scatter", ax=ax)
    relative_cov_post.get_std().reset_index().plot(x="ZAP", y="STD", kind="scatter", color="tomato", marker="x", ax=ax)
    ax.tick_params(labelbottom=False)
    ax.set(yscale="log")

    ax = axs[2]
    sns.heatmap(relative_cov_post.get_corr().data, vmin=-1, vmax=1., cmap="bwr", ax=ax)
    ax.tick_params(labelbottom=False, labelleft=False)
    
    fig.suptitle(f"ZAM={za*10} - E={e:.3e} eV")
    fig.tight_layout()


    # Draw random sample from the new covariance data
    seed = random.randrange(2**32 - 1)    # create a seed, need to change it for the different energies
    print(f"sampling IFY for energy {e:.3e} eV...")
    smp[e] = relative_cov_post.sampling(nsmp, seed=seed)   # sample from posterior covariance matrix

fig, ax = plt.subplots()
(smp[energies[0]].std(axis=1) / rcovs_post[energies[0]].get_std() - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="dodgerblue", label=f"E={energies[0]:.2e}")
(smp[energies[1]].std(axis=1) / rcovs_post[energies[1]].get_std() - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="tomato", alpha=.6, label=f"E={energies[1]:.2e}")
(smp[energies[2]].std(axis=1) / rcovs_post[energies[2]].get_std() - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="k", alpha=.3, label=f"E={energies[2]:.2e}")
ax.set(ylim=[-100, 100], ylabel="relative difference / $\\%$", xlabel="ZAP")
ax.tick_params(labelbottom=False)
fig.tight_layout()

smp_min = 0   # write ENDF-6 file only in the sample range [smp_min, smp_max]
smp_max = 99
file_template = "u235_fy_cov_{}.jeff33"
for ismp in range(smp_min, smp_max+1):
    file = file_template.format(ismp)
    f = sandy.Fy(nfpy.data.copy())
    for e, s in smp.items():
        idx_ify = nfpy.data.query(f"E=={float(e)} & MT==454").index
        idx_cfy = nfpy.data.query(f"E=={float(e)} & MT==459").index
        f.data.loc[idx_ify, "DFY"] = f.data.loc[idx_ify, "FY"]            # just for me, i copy the original IFYs where uncertainties should be, so i can compare them to the perturbed ones (anyways I don't use uncertainties)
        f.data.loc[idx_cfy, "DFY"] = f.data.loc[idx_cfy, "FY"]            # same but for CFYs
        f.data.loc[idx_ify, "FY"] *= s[ismp].values                # IMPORTANT, this does not update the CFYs, which in random ENDF-6 file are inconsistent with the perturbed IFYs
        #f = f.apply_qmatrix(922350, e, rdd, keep_fy_index=True)      # Run this if you want to update the CFYs (slower), or else comment it out
    print(f"writing file '{file}'...")
    f.to_endf6(tape).to_file(file)

fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharex=True, sharey=True)

ax = axs[0]
e = energies[0]
mask = f"E=={e} and MT==454"
d = pd.DataFrame({k: v.data.query(mask).set_index("ZAP").FY for k, v in sfy.items()})
(d.mean(axis=1) / nfpy.data.query(mask).set_index("ZAP").FY - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="dodgerblue", label="$\\frac{m- \\mu}{\\mu} \\times 100$")
((d.std(axis=1) / d.mean(axis=1)).fillna(0) / rcovs_post[energies[0]].get_std() - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="tomato", alpha=.6, label="$\\frac{s_{post}- \\sigma_{post}}{\\sigma_{post}} \\times 100$")
ax.set(ylim=[-100, 100], ylabel="relative difference / $\\%$", xlabel="ZAP", title=f"E={e:.2e} eV")
ax.tick_params(labelbottom=False)

ax = axs[1]
e = energies[1]
mask = f"E=={e} and MT==454"
d = pd.DataFrame({k: v.data.query(mask).set_index("ZAP").FY for k, v in sfy.items()})
(d.mean(axis=1) / nfpy.data.query(mask).set_index("ZAP").FY - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="dodgerblue", label="$\\frac{m- \\mu}{\\mu} \\times 100$")
((d.std(axis=1) / d.mean(axis=1)).fillna(0) / rcovs_post[energies[0]].get_std() - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="tomato", alpha=.6, label="$\\frac{s_{post}- \\sigma_{post}}{\\sigma_{post}} \\times 100$")
ax.set(ylim=[-100, 100], ylabel="relative difference / $\\%$", xlabel="ZAP", title=f"E={e:.2e} eV")
ax.tick_params(labelbottom=False)
ax.get_legend().remove()

ax = axs[2]
e = energies[2]
mask = f"E=={e} and MT==454"
d = pd.DataFrame({k: v.data.query(mask).set_index("ZAP").FY for k, v in sfy.items()})
(d.mean(axis=1) / nfpy.data.query(mask).set_index("ZAP").FY - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="dodgerblue", label="$\\frac{m- \\mu}{\\mu} \\times 100$")
((d.std(axis=1) / d.mean(axis=1)).fillna(0) / rcovs_post[energies[0]].get_std() - 1).fillna(0).multiply(100).rename("DIFF").reset_index().plot(x="ZAP", y="DIFF", kind="scatter", ax=ax, color="tomato", alpha=.6, label="$\\frac{s_{post}- \\sigma_{post}}{\\sigma_{post}} \\times 100$")
ax.set(ylim=[-100, 100], ylabel="relative difference / $\\%$", xlabel="ZAP", title=f"E={e:.2e} eV")
ax.tick_params(labelbottom=False)
ax.get_legend().remove()

fig.tight_layout()