Merge pull request #107 from fusion-energy/rem/refactoring

Refactoring and tidy up

.github/workflows/ci.yml

@@ -34,7 +34,7 @@ jobs:
       - name: Run tests
         shell: bash -l {0}
         run: |
-          pytest tests --cov opennmc_plasma_source --cov-report xml --cov-report term
+          pytest tests --cov openmc_plasma_source --cov-report xml --cov-report term
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v4

src/openmc_plasma_source/fuel_types.py

@@ -1,22 +1,25 @@
 import NeSST as nst
 import numpy as np
-import openmc
+import openmc.stats
+
+from typing import Dict, List
 
 
 def neutron_energy_mean(ion_temperature: float, reaction: str) -> float:
-    """Calculates the mean energy of the neutron emitted during DD or DT
+    """
+    Calculates the mean energy of the neutron emitted during DD or DT
     fusion accounting for temperature of the incident ions. Based on Ballabio
     fits, see Table III of L. Ballabio et al 1998 Nucl. Fusion 38 1723
 
     Args:
-        ion_temperature (float): the temperature of the ions in eV
-        reaction (str): the two isotope that fuse, can be either 'DD' or 'DT'
+        ion_temperature: the temperature of the ions in eV.
+        reaction: the two isotope that fuse, can be either 'DD' or 'DT'.
 
     Raises:
-        ValueError: if the reaction is not 'DD' or 'DT' then a ValueError is raised
+        ValueError: if the reaction is not 'DD' or 'DT' then a ValueError is raised.
 
     Returns:
-        float: the mean neutron energy in eV
+        The mean neutron energy in eV.
     """
 
     # values from Ballabio paper
@@ -47,19 +50,20 @@ def neutron_energy_mean(ion_temperature: float, reaction: str) -> float:
 
 
 def neutron_energy_std_dev(ion_temperature: float, reaction: str) -> float:
-    """Calculates the standard deviation of the neutron energy emitted during DD
+    """
+    Calculates the standard deviation of the neutron energy emitted during DD
      or DT fusion accounting for temperature of the incident ions. Based on
     Ballabio fits, see Table III of L. Ballabio et al 1998 Nucl. Fusion 38 1723
 
     Args:
-        ion_temperature (float): the temperature of the ions in eV
-        reaction (str): the two isotope that fuse, can be either 'DD' or 'DT'
+        ion_temperature: the temperature of the ions in eV.
+        reaction: the two isotope that fuse, can be either 'DD' or 'DT'.
 
     Raises:
-        ValueError: if the reaction is not 'DD' or 'DT' then a ValueError is raised
+        ValueError: if the reaction is not 'DD' or 'DT' then a ValueError is raised.
 
     Returns:
-        float: the mean neutron energy in eV
+        The mean neutron energy in eV
     """
 
     # values from Ballabio paper
@@ -93,35 +97,33 @@ def neutron_energy_std_dev(ion_temperature: float, reaction: str) -> float:
     return std_dev
 
 
-def get_reactions_from_fuel(fuel):
-    if ["D", "T"] == sorted(set(fuel.keys())):
+def get_reactions_from_fuel(fuel: Dict[str, float]) -> List[str]:
+    unique_keys = sorted(set(fuel.keys()))
+    if ["D", "T"] == unique_keys:
         return ["DT", "DD", "TT"]
-    elif ["D"] == sorted(set(fuel.keys())):
+    elif ["D"] == unique_keys:
         return ["DD"]
-    elif ["T"] == sorted(set(fuel.keys())):
+    elif ["T"] == unique_keys:
         return ["TT"]
     else:
-        msg = 'reactions of fuel {fuel} could not be found. Supported fuel keys are "T" and "D"'
+        msg = f'reactions of fuel {fuel} could not be found. Supported fuel keys are "T" and "D"'
         raise ValueError(msg)
 
 
 def get_neutron_energy_distribution(
     ion_temperature: float,
-    fuel: dict,
+    fuel: Dict[str, float],
 ) -> openmc.stats.Discrete:
     """Finds the energy distribution and their relative strengths.
 
     Parameters
     ----------
-    ion_temperature : float
-        temperature of plasma ions in eV
-    fuel : dict
-        isotopes as keys and atom fractions as values
+    ion_temperature : temperature of plasma ions in eV
+    fuel : isotopes as keys and atom fractions as values
 
     Returns
     -------
-    openmc.stats.Discrete
-        energy distribution
+    energy distribution
     """
 
     sum_fuel_isotopes = sum(fuel.values())

src/openmc_plasma_source/point_source.py

@@ -1,51 +1,52 @@
 from typing import Tuple
 
-import openmc
+import openmc.stats
+from openmc import IndependentSource
 
 from .fuel_types import get_neutron_energy_distribution
 
+from typing import Dict, List
+
 
 def fusion_point_source(
     coordinate: Tuple[float, float, float] = (0.0, 0.0, 0.0),
     temperature: float = 20000.0,
-    fuel: dict = {"D": 0.5, "T": 0.5},
-) -> list[openmc.IndependentSource]:
+    fuel: Dict[str, float] = {"D": 0.5, "T": 0.5},
+) -> List[IndependentSource]:
     """Creates a list of openmc.IndependentSource objects representing an ICF source.
 
     Resulting ICF (Inertial Confinement Fusion) source will have an energy
     distribution according to the fuel composition.
 
     Args:
-        coordinate (tuple[float,float,float]): Location of the point source.
+        coordinate: Location of the point source.
             Each component is measured in metres.
-        temperature (float): Temperature of the source (eV).
-        fuel (dict): Isotopes as keys and atom fractions as values
+        temperature: Temperature of the source (eV).
+        fuel: Isotopes as keys and atom fractions as values
+
+    Returns:
+        A list of one openmc.IndependentSource instance.
     """
 
-    if (
+    if not (
         isinstance(coordinate, tuple)
         and len(coordinate) == 3
         and all(isinstance(x, (int, float)) for x in coordinate)
     ):
-        pass
-    else:
         raise ValueError("coordinate must be a tuple of three floats.")
 
     if not isinstance(temperature, (int, float)):
         raise ValueError("Temperature must be a float.")
     if temperature <= 0:
         raise ValueError("Temperature must be positive float.")
 
-    sources = []
+    source = openmc.IndependentSource()
 
     energy_distribution = get_neutron_energy_distribution(
         ion_temperature=temperature, fuel=fuel
     )
-
-    source = openmc.IndependentSource()
     source.energy = energy_distribution
     source.space = openmc.stats.Point(coordinate)
     source.angle = openmc.stats.Isotropic()
-    sources.append(source)
 
-    return sources
+    return [source]

src/openmc_plasma_source/ring_source.py

@@ -1,65 +1,57 @@
 import numpy as np
-import openmc
+import openmc.stats
+from openmc import IndependentSource
 
 from .fuel_types import get_neutron_energy_distribution
 
+from typing import Tuple, List, Dict
+
 
 def fusion_ring_source(
     radius: float,
-    angles: tuple[float, float] = (0, 2 * np.pi),
+    angles: Tuple[float, float] = (0, 2 * np.pi),
     z_placement: float = 0,
     temperature: float = 20000.0,
-    fuel: dict = {"D": 0.5, "T": 0.5},
-) -> list[openmc.IndependentSource]:
+    fuel: Dict = {"D": 0.5, "T": 0.5},
+) -> List[IndependentSource]:
     """Creates a list of openmc.IndependentSource objects in a ring shape.
 
     Useful for simulations where all the plasma parameters are not known and
     this simplified geometry will suffice. Resulting ring source will have an
     energy distribution according to the fuel composition.
 
     Args:
-        radius (float): the inner radius of the ring source, in metres
-        angles (iterable of floats): the start and stop angles of the ring in
+        radius: the inner radius of the ring source, in metres
+        angles: the start and stop angles of the ring in
             radians
-        z_placement (float): Location of the ring source (m). Defaults to 0.
-        temperature (float): Temperature of the source (eV).
-        fuel (dict): Isotopes as keys and atom fractions as values
+        z_placement: Location of the ring source (m). Defaults to 0.
+        temperature: Temperature of the source (eV).
+        fuel: Isotopes as keys and atom fractions as values
+
+    Returns:
+        A list of one openmc.IndependentSource instance.
     """
 
-    if isinstance(radius, (int, float)) and radius > 0:
-        pass
-    else:
+    if not isinstance(radius, (int, float)) or radius <= 0:
         raise ValueError("Radius must be a float strictly greater than 0.")
 
-    if (
+    if not (
         isinstance(angles, tuple)
         and len(angles) == 2
         and all(
             isinstance(angle, (int, float)) and -2 * np.pi <= angle <= 2 * np.pi
             for angle in angles
         )
     ):
-        pass
-    else:
         raise ValueError("Angles must be a tuple of floats between zero and 2 * np.pi")
 
-    if isinstance(z_placement, (int, float)):
-        pass
-    else:
+    if not isinstance(z_placement, (int, float)):
         raise TypeError("Z placement must be a float.")
 
-    if isinstance(temperature, (int, float)) and temperature > 0:
-        pass
-    else:
+    if not (isinstance(temperature, (int, float)) and temperature > 0):
         raise ValueError("Temperature must be a float strictly greater than 0.")
 
-    sources = []
-
-    energy_distributions = get_neutron_energy_distribution(
-        ion_temperature=temperature, fuel=fuel
-    )
-
-    source = openmc.IndependentSource()
+    source = IndependentSource()
 
     source.space = openmc.stats.CylindricalIndependent(
         r=openmc.stats.Discrete([radius], [1]),
@@ -68,8 +60,11 @@ def fusion_ring_source(
         origin=(0.0, 0.0, 0.0),
     )
 
+    energy_distributions = get_neutron_energy_distribution(
+        ion_temperature=temperature, fuel=fuel
+    )
+
     source.energy = energy_distributions
     source.angle = openmc.stats.Isotropic()
-    sources.append(source)
 
-    return sources
+    return [source]