#!/usr/bin/env python
# coding: utf-8

# # How to apply a custom perturbation to a given cross section

# In[1]:


import os

import numpy as np
import pandas as pd


# In[2]:


import sandy


# In[3]:


import matplotlib.pyplot as plt
import seaborn as sns

sns.set_style("whitegrid")


# In this notebook we apply a constant relative perturbation, say 5%, to an energy-dependent cross section.
# The perturbation is applied over an energy bin of choice, e.g. between 10 eV and 100 eV.

# As an example we use the JEFF-3.3 evaluation for Pu-239  downloaded from the [OECD/NEA website](https://www.oecd-nea.org/dbdata/jeff/jeff33/index.html) and processed with NJOY-2016 into PENDF format using the following input file:

# In[4]:


tape = sandy.get_endf6_file("jeff_33", "xs", 942390).get_pendf(
    temperature=900,
    purr=False,
    heatr=False,
    gaspr=False,
    verbose=True,
)


# In[5]:


xs = sandy.Xs.from_endf6(tape)


# Now, we apply perturbations according to the WIMS-69 energy grid, available on the [Serpent Wiki](http://serpent.vtt.fi/mediawiki/index.php/EPRI-CPM_69_group_structure).
# Still, any other energy grids would work just fine.

# In[6]:


egrid = sandy.energy_grids.WIMS69
egrid


# The energy grid was converted in eV to be consistent with the cross section data.

# ### Use a `sandy.Pert` object

# SANDY uses a dedicated class to store (relative) perturbation coefficients.

# In[7]:


sandy.Pert([1, 1.05], index=[10, 100])


# In this example the `Pert` instance defines a multiplicative perturbation factor 1 (0%) for all xs values between 0 and 10 eV, and 1.05 (5%) for all xs between 10 and 100 eV.

# ### Apply perturbations to cross sections

# Below we apply perturbations to the fission cross section (MT=18) of Pu-239 (MAT=9437).
# For each energy bin in the WIMS grid, we increase all xs points in that energy range by 30%.
# The process is repreated iteratively and 69 perturbed xs objects are created.

# In[8]:


mat = 9437
mt = 18
pert_coeff = 30 / 100  # large perturbationfor better visualization


# In[9]:


perturbed_xs = []
for i in range(1, egrid.size):
    e_start = egrid[i-1]
    e_stop = egrid[i]
    index = egrid[i-1: i+1]
    pert = sandy.Pert([1, 1 + pert_coeff], index=index)
    print(f"perturbed xs in energy bin #{i} [{e_start:.5e}, {e_stop:.5e}]")
    xspert = xs.custom_perturbation(mat, mt, pert)
    perturbed_xs.append(xspert)


# ### Plot the results

# Entire energy range

# In[10]:


fig, ax = plt.subplots(figsize=(7, 3.5), dpi=100)
ax = xs.data[(mat,mt)].plot(logx=True, logy=True, color="dodgerblue", linewidth=2, ax=ax)
for xspert in perturbed_xs:
    ax = xspert.data[(mat,mt)].plot(logx=True, logy=True, alpha=0.8, linewidth=0.9, ax=ax)

ax.set_ylabel("cross section / b")
ax.set_xlabel("energy / eV")
ax.set_xlim([1e-5, 2e7])
ax.set_ylim([1e-1, 1e4])
fig.tight_layout();


# Thermal part

# In[11]:


fig, ax = plt.subplots(figsize=(7, 3.5), dpi=100)
ax = xs.data[(mat,mt)].plot(logx=True, logy=True, color="dodgerblue", linewidth=2, ax=ax)
for xspert in perturbed_xs:
    ax = xspert.data[(mat,mt)].plot(logx=True, logy=True, alpha=0.8, linewidth=0.9, ax=ax)

ax.set_ylabel("cross section / b")
ax.set_xlabel("energy / eV")
ax.set_xlim([1e-4, 1])
ax.set_ylim([1e1, 1e4])
fig.tight_layout();


# Fast part

# In[12]:


fig, ax = plt.subplots(figsize=(7, 3.5), dpi=100)
ax = xs.data[(mat,mt)].plot(logx=True, logy=False, color="dodgerblue", linewidth=2, ax=ax)
for xspert in perturbed_xs:
    ax = xspert.data[(mat,mt)].plot(logx=True, logy=False, alpha=0.8, linewidth=0.9, ax=ax)

ax.set_ylabel("cross section / b")
ax.set_xlabel("energy / eV")
ax.set_xlim([1e3, 1e6])
ax.set_ylim([1e0, 1e1])
fig.tight_layout();


# ### Write perturbed xs to file 

# Finally, we can write each perturbed cross section into a copy of the original PENDF file for further use.

# In[13]:


ipert = 0
perturbed_xs[ipert].to_endf6(tape).to_file(f"copy_pert{ipert}.pendf")