Merge pull request #58 from fusion-energy/develop

replacing propert-tea with params

.circleci/config.yml

@@ -3,7 +3,7 @@ jobs:
   # test:
   build:
     docker:
-      - image: openmc/openmc
+      - image: openmc/openmc:v0.12.2
     steps:
       - checkout
       - run:

.github/workflows/ci.yml

@@ -14,7 +14,7 @@ jobs:
   testing:
     runs-on: ubuntu-latest
     container:
-      image: openmc/openmc
+      image: openmc/openmc:v0.12.2
     steps:
       - name: Checkout repository
         uses: actions/checkout@v2

.gitignore

@@ -134,3 +134,6 @@ dmypy.json
 
 # vim swap files
 *.swp
+
+# image files created by tests
+tests/*.png

README.md

@@ -46,7 +46,7 @@ my_source = TokamakSource(
   ).make_openmc_sources()
 ```
 
-For a more complete example check out the [example script](https://github.com/fusion-energy/openmc-plasma-source/blob/better_readme/examples/tokamak_source_example.py).
+For a more complete example check out the [example script](https://github.com/fusion-energy/openmc-plasma-source/blob/main/examples/tokamak_source_example.py).
 
 ![out](https://user-images.githubusercontent.com/40028739/135100022-330aa51c-e2a2-401c-9738-90f3e99c84d4.png)
  ![out](https://user-images.githubusercontent.com/40028739/135098576-a94709ef-96b4-4b8d-8fa0-76a201b6c5d2.png)

openmc_plasma_source/fuel_types.py

@@ -0,0 +1,26 @@
+"""fuel_types.py
+
+Defines dictionary for determining mean energy and mass of reactants
+for a given fusion fuel type.
+"""
+
+from param import Parameterized, Number
+
+
+class Fuel(Parameterized):
+
+    # mean energy, eV
+    mean_energy = Number(None, bounds=(0, None), inclusive_bounds=(False, False))
+
+    # mass of the reactants, AMU
+    mass_of_reactants = Number(None, bounds=(0, None), inclusive_bounds=(False, False))
+
+    def __init__(self, mean_energy, mass_of_reactants):
+        self.mean_energy = mean_energy
+        self.mass_of_reactants = mass_of_reactants
+
+
+fuel_types = {
+    "DD": Fuel(mean_energy=2450000.0, mass_of_reactants=4),
+    "DT": Fuel(mean_energy=14080000.0, mass_of_reactants=5),
+}

openmc_plasma_source/plotting/plot_tokamak_source.py

@@ -3,81 +3,151 @@
 import numpy as np
 
 
-def scatter_tokamak_source(source, quantity=None, **kwargs):
+def scatter_tokamak_source(source, ax=None, quantity=None, aspect="equal", **kwargs):
     """Create a 2D scatter plot of the tokamak source.
     See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.scatter.html
     for more arguments.
 
     Args:
         source (ops.TokamakSource): the plasma source
-        quantity ("str", optional): value by which the lines should be
+        ax (maplotlib.Axes, optional): Axes object on which to plot. If not
+            provided by the user, the current working Axes is retrieved using
+            matplotlib.pyplot.gca().
+        aspect (str, optional): Set aspect ratio of the plot. May be set to 'auto',
+            'equal', or a number denoting the ratio of the plot height to the plot
+            width.
+        quantity (str, optional): value by which the lines should be
             coloured. Defaults to None.
+        **kwargs: Keyword arguments compatible with matplotlib.pyplot.scatter
 
     Raises:
         ValueError: If the quantity is unknown
     """
 
-    quantity_to_attribute = {
-        "ion_temperature": source.temperatures,
-        "neutron_source_density": source.neutron_source_density,
-        "strength": source.strengths,
-    }
-    if quantity in quantity_to_attribute:
-        colours = quantity_to_attribute[quantity]
-    elif quantity is None:
-        colours = None
-    else:
-        raise ValueError("Unknown quantity")
-    plt.gca().set_aspect("equal")
+    # Define possible quantities, and link to arrays within tokamak_source
+    # If given a non-TokamakSource, this step will fail with an AttributeError
+    try:
+        quantity_to_attribute = {
+            "ion_temperature": source.temperatures,
+            "neutron_source_density": source.neutron_source_density,
+            "strength": source.strengths,
+        }
+    except AttributeError as e:
+        raise ValueError(
+            f"openmc_plasma_source.scatter_tokamak_source: argument 'source' "
+            f"must be of type TokamakSource"
+        ) from e
+
+    # For a given quantity, determine colours to plot for each point
+    # If given an incorrect quantity name, this step will fail with a KeyError
+    colours = None
+    if quantity is not None:
+        try:
+            colours = quantity_to_attribute[quantity]
+        except KeyError as e:
+            raise ValueError(
+                f"openmc_plasma_source.scatter_tokamak_source: Unknown 'quantity' "
+                f"provided, options are {quantity_to_attribute.keys()}"
+            ) from e
+
+    # If not provided with an Axes instance, retrieve the current Axes in focus
+    if ax is None:
+        ax = plt.gca()
 
-    return plt.scatter(source.RZ[0], source.RZ[1], c=colours, **kwargs)
+    # Scatter the source R and Z positions, optionally colouring using the chosen
+    # quantity.
+    ax.scatter(source.RZ[0], source.RZ[1], c=colours, **kwargs)
+
+    # Set the aspect ratio on the axes.
+    # Defaults to 'equal', so 1m on the x-axis has the same width as 1m on the y-axis
+    ax.set_aspect(aspect)
+
+    return ax
 
 
 def plot_tokamak_source_3D(
-    source, quantity=None, angles=[0, 1 / 2 * np.pi], colorbar="viridis", **kwargs
+    tokamak_source,
+    ax=None,
+    quantity=None,
+    angles=[0, 0.5 * np.pi],
+    colorbar=None,
+    **kwargs,
 ):
     """Creates a 3D plot of the tokamak source.
     See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.plot.html#matplotlib.pyplot.plot
     for more arguments.
 
     Args:
-        source (ops.TokamakSource): the plasma source
+        tokamak_source (ops.TokamakSource): the plasma source
+        ax (maplotlib.Axes, optional): Axes object on which to plot. If not
+            provided by the user, a new Axes is created by calling
+            matplotlib.pyplot.axes(projection="3d").
         quantity ("str", optional): value by which the lines should be
             coloured. Defaults to None.
         angles (list, optional): iterable of two floats defining the coverage.
             Defaults to [0, 1/2*np.pi].
         colorbar (str, optional): colorbar used if quantity is not None.
-            Defaults to "viridis".
+            When None, matplotlib currently defaults to "viridis".
 
     Raises:
         ValueError: If the quantity is unknown
     """
 
-    quantity_to_attribute = {
-        "ion_temperature": source.temperatures,
-        "neutron_source_density": source.neutron_source_density,
-        "strength": source.strengths,
-    }
-    if quantity in quantity_to_attribute:
-        values = quantity_to_attribute[quantity]
-    elif quantity is None:
-        values = None
+    # Define possible quantities, and link to arrays within tokamak_source
+    # If given a non-TokamakSource, this step will fail with an AttributeError
+    try:
+        quantity_to_attribute = {
+            "ion_temperature": tokamak_source.temperatures,
+            "neutron_source_density": tokamak_source.neutron_source_density,
+            "strength": tokamak_source.strengths,
+        }
+    except AttributeError as e:
+        raise ValueError(
+            f"openmc_plasma_source.plot_tokamak_source_3D: argument 'source' "
+            f"must be of type TokamakSource"
+        ) from e
+
+    # For a given quantity, get the corresponding array from tokamak_source
+    # If given an incorrect quantity name, this step will fail with a KeyError
+    if quantity is not None:
+        try:
+            quantity_values = quantity_to_attribute[quantity]
+        except KeyError as e:
+            raise ValueError(
+                f"openmc_plasma_source.plot_tokamak_source_3D: Unknown 'quantity' "
+                f"provided, options are {quantity_to_attribute.keys()}"
+            ) from e
     else:
-        raise ValueError("Unknown quantity")
-
-    colorbar = cm.get_cmap(colorbar)
-    axes = plt.axes(projection="3d")
-    theta = np.linspace(*angles, 100)
-    for i in range(source.sample_size):
-        if values is not None:
-            colour = colorbar(values[i] / max(values))
-        else:
-            colour = None
-        x = source.RZ[0][i] * np.sin(theta)
-        y = source.RZ[0][i] * np.cos(theta)
-        z = source.RZ[1][i]
-        plt.plot(x, y, z, color=colour, **kwargs)
-
-    axes.set_xlim(-source.major_radius, source.major_radius)
-    axes.set_ylim(-source.major_radius, source.major_radius)
-    axes.set_zlim(-source.major_radius, source.major_radius)
+        quantity_values = np.ones(tokamak_source.sample_size)
+
+    # Get the colours used to plot each curve
+    # If 'quantity' == None, all have the same colour, selected from the middle
+    # of the colormap.
+    cmap = cm.get_cmap(colorbar)
+    if quantity is not None:
+        colors = cmap(quantity_values / np.max(quantity_values))
+    else:
+        colors = cmap(np.full(tokamak_source.sample_size, 0.5))
+
+    # If not provided with an Axes object, create a new one
+    if ax is None:
+        ax = plt.axes(projection="3d")
+
+    # Get curves to plot
+    n_theta = 100
+    theta = np.linspace(*angles, n_theta)
+    xs = np.outer(tokamak_source.RZ[0], np.sin(theta))
+    ys = np.outer(tokamak_source.RZ[0], np.cos(theta))
+    zs = tokamak_source.RZ[1]
+
+    # Plot each curve in turn
+    for x, y, z, color in zip(xs, ys, zs, colors):
+        ax.plot(x, y, z, color=color, **kwargs)
+
+    # Set plot bounds
+    major_radius = tokamak_source.major_radius
+    ax.set_xlim(-major_radius, major_radius)
+    ax.set_ylim(-major_radius, major_radius)
+    ax.set_zlim(-major_radius, major_radius)
+
+    return ax

openmc_plasma_source/point_source.py

@@ -1,35 +1,47 @@
+import openmc
 from typing import Tuple
+from param import Parameterized, Number, NumericTuple, ListSelector
 
-import openmc
+from .fuel_types import fuel_types
 
 
-class FusionPointSource(openmc.Source):
+class FusionPointSource(openmc.Source, Parameterized):
     """An openmc.Source object with some presets to make it more convenient
     for fusion simulations using a point source. All attributes can be changed
     after initialization if required. Default isotropic point source at the
     origin with a Muir energy distribution.
+
+    Args:
+        coordinate (tuple[float,float,float]): Location of the point source.
+            Each component is measured in metres.
+        temperature (float): Temperature of the source (eV).
+        fuel_type (str): The fusion fuel mix. Either 'DT' or 'DD'.
     """
 
+    coordinate = NumericTuple(None, length=3)
+    temperature = Number(None, bounds=(0, None))  # temperature in eV
+    fuel_type = ListSelector(fuel_types.keys())
+
     def __init__(
         self,
         coordinate: Tuple[float, float, float] = (0, 0, 0),
         temperature: float = 20000.0,
         fuel: str = "DT",
     ):
+        # Set local attributes
+        self.coordinate = coordinate
+        self.temperature = temperature
+        self.fuel_type = fuel
+        self.fuel = fuel_types[self.fuel_type]
 
+        # Call init for openmc.Source
         super().__init__()
 
         # performed after the super init as these are Source attributes
-        self.space = openmc.stats.Point(coordinate)
+        self.space = openmc.stats.Point(self.coordinate)
         self.angle = openmc.stats.Isotropic()
-        if fuel == "DT":
-            mean_energy = 14080000.0  # mean energy in eV
-            mass_of_reactants = 5  # mass of the reactants (D + T) AMU
-        elif fuel == "DD":
-            mean_energy = 2450000.0  # mean energy in eV
-            mass_of_reactants = 4  # mass of the reactants (D + D) AMU
-        else:
-            raise ValueError(f'fuel must be either "DT" or "DD", not {fuel}')
         self.energy = openmc.stats.Muir(
-            e0=mean_energy, m_rat=mass_of_reactants, kt=temperature
+            e0=self.fuel.mean_energy,
+            m_rat=self.fuel.mass_of_reactants,
+            kt=self.temperature,
         )