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

# In[1]:


import openmc.deplete
from math import pi
import matplotlib.pyplot as plt


# First create a simple one material, one region problem and deplete it over three 1-day timesteps.

# In[2]:


material = openmc.Material()
material.add_nuclide('U235', 1.0)
material.set_density('g/cm3', 5.0)
material.volume = 4/3 * pi * 10.0**3

model = openmc.Model()
sph = openmc.Sphere(r=10.0, boundary_type='vacuum')
cell = openmc.Cell(fill=material, region=-sph)
model.geometry = openmc.Geometry([cell])

model.settings = openmc.Settings()
model.settings.particles = 1000
model.settings.batches = 50
model.settings.inactive = 10

chain_file = 'chain_simple.xml'
op = openmc.deplete.Operator(model, chain_file)
power = 40.0e3
time_steps = [1.0]*3
integrator = openmc.deplete.PredictorIntegrator(op, time_steps, timestep_units='d', power=power)
integrator.integrate()


# Now let's look at the density of $^{135}$Xe from the results:

# In[3]:


results = openmc.deplete.ResultsList.from_hdf5('depletion_results.h5')
time, n_Xe135 = results.get_atoms('1', 'Xe135')
days = 24*60*60
plt.plot(time/days, n_Xe135)
plt.xlabel('Time [d]')
plt.ylabel('Xe135 [atoms]')


# For the next depletion run, we need to get a new set of materials that is equivalent to the last depletion step. For this, we can use the `ResultsList.export_to_materials` method. The change we'll make to the material is to remove ${135}$Xe so that it starts with a near-zero concentration.

# In[4]:


# Get materials at the end of the last simulation
model.materials = results.export_to_materials(len(time_steps))

# Now change the material by getting rid of Xe135
model.materials[0].remove_nuclide('Xe135')


# Now run the depletion integrator again --- note that we don't supply previous results (which would override the material change we made above).

# In[5]:


# We need to create a new operator; otherwise, it will still use the number densities
# from the end of the previous simulation
op = openmc.deplete.Operator(model, chain_file)

integrator = openmc.deplete.PredictorIntegrator(op, time_steps, timestep_units='d', power=power)
integrator.integrate()


# And finally let's look at the $^{135}$Xe density again:

# In[6]:


results = openmc.deplete.ResultsList.from_hdf5('depletion_results.h5')
time, n_Xe135 = results.get_atoms('1', 'Xe135')
days = 24*60*60
plt.plot(time/days, n_Xe135)
plt.xlabel('Time [d]')
plt.ylabel('Xe135 [atoms]')


# In[ ]: