Adding projection plot

examples/example_plot_projected_position_from_initial_source.py

@@ -0,0 +1,35 @@
+import openmc_source_plotter as osp
+import openmc
+import numpy as np
+
+# initialises a new source object
+my_source = openmc.Source()
+
+# the distribution of radius is just a single value
+radius = openmc.stats.Discrete([10], [1])
+
+# the distribution of source z values is just a single value
+z_values = openmc.stats.Discrete([0], [1])
+
+# the distribution of source azimuthal angles values is a uniform distribution between 0 and 2 Pi
+angle = openmc.stats.Uniform(a=0.0, b=2 * 3.14159265359)
+
+# this makes the ring source using the three distributions and a radius
+my_source.space = openmc.stats.CylindricalIndependent(
+    r=radius, phi=angle, z=z_values, origin=(0.0, 0.0, 0.0)
+)
+
+
+# makes an initial_source.h5 file with details of the particles
+initial_source_filename = osp.create_initial_particles(
+    source=my_source,
+    number_of_particles=10,
+    openmc_exec="/home/jshim/miniconda3/envs/openmc_0_11_0/bin/openmc",
+)
+
+# plots the particle energy distribution
+plot = osp.plot_position_projected_from_initial_source(
+    input_filename=initial_source_filename, axis="XY"
+)
+
+plot.show()

openmc_source_plotter/__init__.py

@@ -1,5 +1,6 @@
 from .core import plot_direction_from_initial_source
 from .core import plot_position_from_initial_source
+from .core import plot_position_projected_from_initial_source
 from .core import plot_energy_from_initial_source
 from .core import get_particle_data
 from .core import create_initial_particles

openmc_source_plotter/core.py

@@ -128,6 +128,43 @@ def plot_energy_from_initial_source(
     return fig
 
 
+def plot_projected_position_from_initial_source(
+    input_filename="initial_source.h5", axis="XY"
+):
+    """makes a plot of the initial creation locations of the particle source"""
+
+    data = get_particle_data(input_filename)
+
+    text = ["Energy = " + str(i) + " eV" for i in data["e_values"]]
+
+    fig = go.Figure()
+
+    if axis == "XY":
+        x = data["x_values"]
+        y = data["y_values"]
+    if axis == "XZ":
+        x = data["x_values"]
+        y = data["z_values"]
+
+    fig.add_trace(
+        go.Scatter(
+            x=x,
+            y=y,
+            hovertext=text,
+            text=text,
+            mode="markers",
+            marker={
+                "size": 2,
+                "color": data["e_values"],
+            },
+        )
+    )
+
+    fig.update_layout(title="Particle production coordinates - coloured by energy")
+
+    return fig
+
+
 def plot_position_from_initial_source(input_filename="initial_source.h5"):
     """makes a plot of the initial creation locations of the particle source"""