%matplotlib inline
import openmc

"""
Explanation: This notebook is intended to demonstrate the basic features of the Python API for constructing input files and running OpenMC. In it, we will show how to create a basic reflective pin-cell model that is equivalent to modeling an infinite array of fuel pins. If you have never used OpenMC, this can serve as a good starting point to learn the Python API. We highly recommend having a copy of the Python API reference documentation open in another browser tab that you can refer to.
End of explanation
"""


uo2 = openmc.Material(1, "uo2")
print(uo2)

"""
Explanation: Defining Materials
Materials in OpenMC are defined as a set of nuclides with specified atom/weight fractions. To begin, we will create a material by making an instance of the Material class. In OpenMC, many objects, including materials, are identified by a "unique ID" that is simply just a positive integer. These IDs are used when exporting XML files that the solver reads in. They also appear in the output and can be used for identification. Since an integer ID is not very useful by itself, you can also give a material a name as well.
End of explanation
"""


mat = openmc.Material()
print(mat)

"""
Explanation: On the XML side, you have no choice but to supply an ID. However, in the Python API, if you don't give an ID, one will be automatically generated for you:
End of explanation
"""


help(uo2.add_nuclide)

"""
Explanation: We see that an ID of 2 was automatically assigned. Let's now move on to adding nuclides to our uo2 material. The Material object has a method add_nuclide() whose first argument is the name of the nuclide and second argument is the atom or weight fraction.
End of explanation
"""


# Add nuclides to uo2
uo2.add_nuclide('U235', 0.03)
uo2.add_nuclide('U238', 0.97)
uo2.add_nuclide('O16', 2.0)

"""
Explanation: We see that by default it assumes we want an atom fraction.
End of explanation
"""


uo2.set_density('g/cm3', 10.0)

"""
Explanation: Now we need to assign a total density to the material. We'll use the set_density for this.
End of explanation
"""


zirconium = openmc.Material(2, "zirconium")
zirconium.add_element('Zr', 1.0)
zirconium.set_density('g/cm3', 6.6)

water = openmc.Material(3, "h2o")
water.add_nuclide('H1', 2.0)
water.add_nuclide('O16', 1.0)
water.set_density('g/cm3', 1.0)

"""
Explanation: You may sometimes be given a material specification where all the nuclide densities are in units of atom/b-cm. In this case, you just want the density to be the sum of the constituents. In that case, you can simply run mat.set_density('sum').
With UO2 finished, let's now create materials for the clad and coolant. Note the use of add_element() for zirconium.
End of explanation
"""


water.add_s_alpha_beta('c_H_in_H2O')

"""
Explanation: An astute observer might now point out that this water material we just created will only use free-atom cross sections. We need to tell it to use an $S(\alpha,\beta)$ table so that the bound atom cross section is used at thermal energies. To do this, there's an add_s_alpha_beta() method. Note the use of the GND-style name "c_H_in_H2O".
End of explanation
"""


mats = openmc.Materials([uo2, zirconium, water])

"""
Explanation: When we go to run the transport solver in OpenMC, it is going to look for a materials.xml file. Thus far, we have only created objects in memory. To actually create a materials.xml file, we need to instantiate a Materials collection and export it to XML.
End of explanation
"""


mats = openmc.Materials()
mats.append(uo2)
mats += [zirconium, water]
isinstance(mats, list)

"""
Explanation: Note that Materials is actually a subclass of Python's built-in list, so we can use methods like append(), insert(), pop(), etc.
End of explanation
"""


mats.export_to_xml()
!cat materials.xml

"""
Explanation: Finally, we can create the XML file with the export_to_xml() method. In a Jupyter notebook, we can run a shell command by putting ! before it, so in this case we are going to display the materials.xml file that we created.
End of explanation
"""


water.remove_nuclide('O16')
water.add_element('O', 1.0)

mats.export_to_xml()
!cat materials.xml

"""
Explanation: Element Expansion
Did you notice something really cool that happened to our Zr element? OpenMC automatically turned it into a list of nuclides when it exported it! The way this feature works is as follows:

First, it checks whether Materials.cross_sections has been set, indicating the path to a cross_sections.xml file.
If Materials.cross_sections isn't set, it looks for the OPENMC_CROSS_SECTIONS environment variable.
If either of these are found, it scans the file to see what nuclides are actually available and will expand elements accordingly.

Let's see what happens if we change O16 in water to elemental O.
End of explanation
"""


!cat $OPENMC_CROSS_SECTIONS | head -n 10
print('    ...')
!cat $OPENMC_CROSS_SECTIONS | tail -n 10

"""
Explanation: We see that now O16 and O17 were automatically added. O18 is missing because our cross sections file (which is based on ENDF/B-VII.1) doesn't have O18. If OpenMC didn't know about the cross sections file, it would have assumed that all isotopes exist.
The cross_sections.xml file
The cross_sections.xml tells OpenMC where it can find nuclide cross sections and $S(\alpha,\beta)$ tables. It serves the same purpose as MCNP's xsdir file and Serpent's xsdata file. As we mentioned, this can be set either by the OPENMC_CROSS_SECTIONS environment variable or the Materials.cross_sections attribute.
Let's have a look at what's inside this file:
End of explanation
"""


uo2_three = openmc.Material()
uo2_three.add_element('U', 1.0, enrichment=3.0)
uo2_three.add_element('O', 2.0)
uo2_three.set_density('g/cc', 10.0)

"""
Explanation: Enrichment
Note that the add_element() method has a special argument enrichment that can be used for Uranium. For example, if we know that we want to create 3% enriched UO2, the following would work:
End of explanation
"""


# Create PuO2 material
puo2 = openmc.Material()
puo2.add_nuclide('Pu239', 0.94)
puo2.add_nuclide('Pu240', 0.06)
puo2.add_nuclide('O16', 2.0)
puo2.set_density('g/cm3', 11.5)

# Create the mixture
mox = openmc.Material.mix_materials([uo2, puo2], [0.97, 0.03], 'wo')

"""
Explanation: Mixtures
In OpenMC it is also possible to define materials by mixing existing materials. For example, if we wanted to create MOX fuel out of a mixture of UO2 (97 wt%) and PuO2 (3 wt%) we could do the following:
End of explanation
"""


sph = openmc.Sphere(r=1.0)

"""
Explanation: The 'wo' argument in the mix_materials() method specifies that the fractions are weight fractions. Materials can also be mixed by atomic and volume fractions with 'ao' and 'vo', respectively. For 'ao' and 'wo' the fractions must sum to one. For 'vo', if fractions do not sum to one, the remaining fraction is set as void.
Defining Geometry
At this point, we have three materials defined, exported to XML, and ready to be used in our model. To finish our model, we need to define the geometric arrangement of materials. OpenMC represents physical volumes using constructive solid geometry (CSG), also known as combinatorial geometry. The object that allows us to assign a material to a region of space is called a Cell (same concept in MCNP, for those familiar). In order to define a region that we can assign to a cell, we must first define surfaces which bound the region. A surface is a locus of zeros of a function of Cartesian coordinates $x$, $y$, and $z$, e.g.

A plane perpendicular to the x axis: $x - x_0 = 0$
A cylinder parallel to the z axis: $(x - x_0)^2 + (y - y_0)^2 - R^2 = 0$
A sphere: $(x - x_0)^2 + (y - y_0)^2 + (z - z_0)^2 - R^2 = 0$

Between those three classes of surfaces (planes, cylinders, spheres), one can construct a wide variety of models. It is also possible to define cones and general second-order surfaces (tori are not currently supported).
Note that defining a surface is not sufficient to specify a volume -- in order to define an actual volume, one must reference the half-space of a surface. A surface half-space is the region whose points satisfy a positive or negative inequality of the surface equation. For example, for a sphere of radius one centered at the origin, the surface equation is $f(x,y,z) = x^2 + y^2 + z^2 - 1 = 0$. Thus, we say that the negative half-space of the sphere, is defined as the collection of points satisfying $f(x,y,z) < 0$, which one can reason is the inside of the sphere. Conversely, the positive half-space of the sphere would correspond to all points outside of the sphere.
Let's go ahead and create a sphere and confirm that what we've told you is true.
End of explanation
"""


inside_sphere = -sph
outside_sphere = +sph

"""
Explanation: Note that by default the sphere is centered at the origin so we didn't have to supply x0, y0, or z0 arguments. Strictly speaking, we could have omitted R as well since it defaults to one. To get the negative or positive half-space, we simply need to apply the - or + unary operators, respectively.
(NOTE: Those unary operators are defined by special methods: __pos__ and __neg__ in this case).
End of explanation
"""


print((0,0,0) in inside_sphere, (0,0,2) in inside_sphere)
print((0,0,0) in outside_sphere, (0,0,2) in outside_sphere)

"""
Explanation: Now let's see if inside_sphere actually contains points inside the sphere:
End of explanation
"""


z_plane = openmc.ZPlane(z0=0)
northern_hemisphere = -sph & +z_plane

"""
Explanation: Everything works as expected! Now that we understand how to create half-spaces, we can create more complex volumes by combining half-spaces using Boolean operators: &amp; (intersection), | (union), and ~ (complement). For example, let's say we want to define a region that is the top part of the sphere (all points inside the sphere that have $z > 0$.
End of explanation
"""


northern_hemisphere.bounding_box

"""
Explanation: For many regions, OpenMC can automatically determine a bounding box. To get the bounding box, we use the bounding_box property of a region, which returns a tuple of the lower-left and upper-right Cartesian coordinates for the bounding box:
End of explanation
"""


cell = openmc.Cell()
cell.region = northern_hemisphere

# or...
cell = openmc.Cell(region=northern_hemisphere)

"""
Explanation: Now that we see how to create volumes, we can use them to create a cell.
End of explanation
"""


cell.fill = water

"""
Explanation: By default, the cell is not filled by any material (void). In order to assign a material, we set the fill property of a Cell.
End of explanation
"""


universe = openmc.Universe()
universe.add_cell(cell)

# this also works
universe = openmc.Universe(cells=[cell])

"""
Explanation: Universes and in-line plotting
A collection of cells is known as a universe (again, this will be familiar to MCNP/Serpent users) and can be used as a repeatable unit when creating a model. Although we don't need it yet, the benefit of creating a universe is that we can visualize our geometry while we're creating it.
End of explanation
"""


universe.plot(width=(2.0, 2.0))

"""
Explanation: The Universe object has a plot method that will display our the universe as current constructed:
End of explanation
"""


universe.plot(width=(2.0, 2.0), basis='xz')

"""
Explanation: By default, the plot will appear in the $x$-$y$ plane. We can change that with the basis argument.
End of explanation
"""


universe.plot(width=(2.0, 2.0), basis='xz',
              colors={cell: 'fuchsia'})

"""
Explanation: If we have particular fondness for, say, fuchsia, we can tell the plot() method to make our cell that color.
End of explanation
"""


fuel_or = openmc.ZCylinder(r=0.39)
clad_ir = openmc.ZCylinder(r=0.40)
clad_or = openmc.ZCylinder(r=0.46)

"""
Explanation: Pin cell geometry
We now have enough knowledge to create our pin-cell. We need three surfaces to define the fuel and clad:

The outer surface of the fuel -- a cylinder parallel to the z axis
The inner surface of the clad -- same as above
The outer surface of the clad -- same as above

These three surfaces will all be instances of openmc.ZCylinder, each with a different radius according to the specification.
End of explanation
"""


fuel_region = -fuel_or
gap_region = +fuel_or & -clad_ir
clad_region = +clad_ir & -clad_or

"""
Explanation: With the surfaces created, we can now take advantage of the built-in operators on surfaces to create regions for the fuel, the gap, and the clad:
End of explanation
"""


fuel = openmc.Cell(1, 'fuel')
fuel.fill = uo2
fuel.region = fuel_region

gap = openmc.Cell(2, 'air gap')
gap.region = gap_region

clad = openmc.Cell(3, 'clad')
clad.fill = zirconium
clad.region = clad_region

"""
Explanation: Now we can create corresponding cells that assign materials to these regions. As with materials, cells have unique IDs that are assigned either manually or automatically. Note that the gap cell doesn't have any material assigned (it is void by default).
End of explanation
"""


pitch = 1.26
left = openmc.XPlane(x0=-pitch/2, boundary_type='reflective')
right = openmc.XPlane(x0=pitch/2, boundary_type='reflective')
bottom = openmc.YPlane(y0=-pitch/2, boundary_type='reflective')
top = openmc.YPlane(y0=pitch/2, boundary_type='reflective')

"""
Explanation: Finally, we need to handle the coolant outside of our fuel pin. To do this, we create x- and y-planes that bound the geometry.
End of explanation
"""


water_region = +left & -right & +bottom & -top & +clad_or

moderator = openmc.Cell(4, 'moderator')
moderator.fill = water
moderator.region = water_region

"""
Explanation: The water region is going to be everything outside of the clad outer radius and within the box formed as the intersection of four half-spaces.
End of explanation
"""


box = openmc.rectangular_prism(width=pitch, height=pitch,
                               boundary_type='reflective')
type(box)

"""
Explanation: OpenMC also includes a factory function that generates a rectangular prism that could have made our lives easier.
End of explanation
"""


water_region = box & +clad_or

"""
Explanation: Pay attention here -- the object that was returned is NOT a surface. It is actually the intersection of four surface half-spaces, just like we created manually before. Thus, we don't need to apply the unary operator (-box). Instead, we can directly combine it with +clad_or.
End of explanation
"""


root = openmc.Universe(cells=(fuel, gap, clad, moderator))

geom = openmc.Geometry()
geom.root_universe = root

# or...
geom = openmc.Geometry(root)
geom.export_to_xml()
!cat geometry.xml

"""
Explanation: The final step is to assign the cells we created to a universe and tell OpenMC that this universe is the "root" universe in our geometry. The Geometry is the final object that is actually exported to XML.
End of explanation
"""


point = openmc.stats.Point((0, 0, 0))
src = openmc.Source(space=point)

"""
Explanation: Starting source and settings
The Python API has a module openmc.stats with various univariate and multivariate probability distributions. We can use these distributions to create a starting source using the openmc.Source object.
End of explanation
"""


settings = openmc.Settings()
settings.source = src
settings.batches = 100
settings.inactive = 10
settings.particles = 1000

settings.export_to_xml()
!cat settings.xml

"""
Explanation: Now let's create a Settings object and give it the source we created along with specifying how many batches and particles we want to run.
End of explanation
"""


cell_filter = openmc.CellFilter(fuel)

t = openmc.Tally(1)
t.filters = [cell_filter]

"""
Explanation: User-defined tallies
We actually have all the required files needed to run a simulation. Before we do that though, let's give a quick example of how to create tallies. We will show how one would tally the total, fission, absorption, and (n,$\gamma$) reaction rates for $^{235}$U in the cell containing fuel. Recall that filters allow us to specify where in phase-space we want events to be tallied and scores tell us what we want to tally:
$$X = \underbrace{\int d\mathbf{r} \int d\mathbf{\Omega} \int dE}{\text{filters}} \; \underbrace{f(\mathbf{r},\mathbf{\Omega},E)}{\text{scores}} \psi (\mathbf{r},\mathbf{\Omega},E)$$
In this case, the where is "the fuel cell". So, we will create a cell filter specifying the fuel cell.
End of explanation
"""


t.nuclides = ['U235']
t.scores = ['total', 'fission', 'absorption', '(n,gamma)']

"""
Explanation: The what is the total, fission, absorption, and (n,$\gamma$) reaction rates in $^{235}$U. By default, if we only specify what reactions, it will gives us tallies over all nuclides. We can use the nuclides attribute to name specific nuclides we're interested in.
End of explanation
"""


tallies = openmc.Tallies([t])
tallies.export_to_xml()
!cat tallies.xml

"""
Explanation: Similar to the other files, we need to create a Tallies collection and export it to XML.
End of explanation
"""


openmc.run()

"""
Explanation: Running OpenMC
Running OpenMC from Python can be done using the openmc.run() function. This function allows you to set the number of MPI processes and OpenMP threads, if need be.
End of explanation
"""


!cat tallies.out

"""
Explanation: Great! OpenMC already told us our k-effective. It also spit out a file called tallies.out that shows our tallies. This is a very basic method to look at tally data; for more sophisticated methods, see other example notebooks.
End of explanation
"""


p = openmc.Plot()
p.filename = 'pinplot'
p.width = (pitch, pitch)
p.pixels = (200, 200)
p.color_by = 'material'
p.colors = {uo2: 'yellow', water: 'blue'}

"""
Explanation: Geometry plotting
We saw before that we could call the Universe.plot() method to show a universe while we were creating our geometry. There is also a built-in plotter in the codebase that is much faster than the Python plotter and has more options. The interface looks somewhat similar to the Universe.plot() method. Instead though, we create Plot instances, assign them to a Plots collection, export it to XML, and then run OpenMC in geometry plotting mode. As an example, let's specify that we want the plot to be colored by material (rather than by cell) and we assign yellow to fuel and blue to water.
End of explanation
"""


plots = openmc.Plots([p])
plots.export_to_xml()
!cat plots.xml

"""
Explanation: With our plot created, we need to add it to a Plots collection which can be exported to XML.
End of explanation
"""


openmc.plot_geometry()

"""
Explanation: Now we can run OpenMC in plotting mode by calling the plot_geometry() function. Under the hood this is calling openmc --plot.
End of explanation
"""


!convert pinplot.ppm pinplot.png

"""
Explanation: OpenMC writes out a peculiar image with a .ppm extension. If you have ImageMagick installed, this can be converted into a more normal .png file.
End of explanation
"""


from IPython.display import Image
Image("pinplot.png")

"""
Explanation: We can use functionality from IPython to display the image inline in our notebook:
End of explanation
"""


p.to_ipython_image()

"""
Explanation: That was a little bit cumbersome. Thankfully, OpenMC provides us with a method on the Plot class that does all that "boilerplate" work.
End of explanation
"""


