checking for incorrect values is more concise

pyproject.toml

@@ -18,7 +18,7 @@ classifiers = [
 authors = [
   { name="Rémi Delaporte-Mathurin" },
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 keywords = ["python", "neutron", "fusion", "source", "openmc", "energy", "tokamak"]
 dependencies = [
     "numpy>=1.9",

src/openmc_plasma_source/point_source.py

@@ -10,10 +10,10 @@ def fusion_point_source(
     temperature: float = 20000.0,
     fuel: dict = {"D": 0.5, "T": 0.5},
 ):
-    """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.
+    """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.

src/openmc_plasma_source/ring_source.py

@@ -12,11 +12,12 @@ def fusion_ring_source(
     z_placement: float = 0,
     temperature: float = 20000.0,
     fuel: dict = {"D": 0.5, "T": 0.5},
-):
-    """An openmc.Source object with some presets to make it more convenient
-    for fusion simulations using a ring source. All attributes can be changed
-    after initialization if required. Default isotropic ring source with a
-    realistic energy distribution.
+)-> list[openmc.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

src/openmc_plasma_source/tokamak_source.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import openmc
+import openmc.checkvalue as cv
 
 from .fuel_types import get_neutron_energy_distribution
 
@@ -26,20 +27,20 @@ def tokamak_source(
     angles: Tuple[float, float] = (0, 2 * np.pi),
     sample_size: int = 1000,
     fuel: dict = {"D": 0.5, "T": 0.5},
-):
-    """Plasma neutron source sampling.
-    This class greatly relies on models described in [1]
+) -> list[openmc.IndependentSource]:
+    """Creates a list of openmc.IndependentSource objects representing a tokamak plasma.
+
+    Resulting sources will have an energy distribution according to the fuel
+    composition.This function greatly relies on models described in [1]
 
     [1] : Fausser et al, 'Tokamak D-T neutron source models for different
     plasma physics confinement modes', Fus. Eng. and Design,
     https://doi.org/10.1016/j.fusengdes.2012.02.025
 
     Usage:
-        my_plasma = Plasma(**plasma_prms)
-        my_plasma.sample_sources()
-        print(my_plasma.RZ)
-        print(my_plasma.temperatures)
-        openmc_sources = my_plasma.make_openmc_sources()
+        my_source = tokamak_source(**plasma_prms)
+        my_settings = openmc.Settings()
+        my_settings.source = my_source
 
     Args:
         major_radius (float): Plasma major radius (cm)
@@ -73,65 +74,26 @@ def tokamak_source(
     """
 
     # Perform sanity checks for inputs not caught by properties
-
-    if not isinstance(major_radius, (int, float)):
-        raise ValueError("Major radius must be a number")
-
-    if not isinstance(minor_radius, (int, float)):
-        raise ValueError("Minor radius must be a number")
-
-    if not isinstance(elongation, (int, float)):
-        raise ValueError("Elongation must be a number")
-
-    if not isinstance(triangularity, (int, float)):
-        raise ValueError("Triangularity must be a number")
-
-    if not isinstance(ion_density_centre, (int, float)):
-        raise ValueError("ion_density_centre must be a number")
-
-    if not isinstance(ion_density_peaking_factor, (int, float)):
-        raise ValueError("ion_density_peaking_factor must be a number")
-
-    if not isinstance(ion_density_pedestal, (int, float)):
-        raise ValueError("ion_density_pedestal must be a number")
-
-    if not isinstance(ion_density_separatrix, (int, float)):
-        raise ValueError("ion_density_separatrix must be a number")
-
-    if minor_radius >= major_radius:
-        raise ValueError("Minor radius must be smaller than major radius")
-
-    if pedestal_radius >= minor_radius:
-        raise ValueError("Pedestal radius must be smaller than minor radius")
-
-    if abs(shafranov_factor) >= 0.5 * minor_radius:
-        raise ValueError("Shafranov factor must be smaller than 0.5*minor radius")
-
-    if major_radius < 0:
-        raise ValueError("Major radius must greater than 0")
-
-    if minor_radius < 0:
-        raise ValueError("Minor radius must greater than 0")
-
-    if elongation <= 0:
-        raise ValueError("Elongation must greater than 0")
-
-    if not triangularity <= 1.0:
-        raise ValueError("Triangularity must less than or equal to 1.")
-    if not triangularity >= -1.0:
-        raise ValueError("Triangularity must greater than or equal to -1.")
-
-    if mode not in ["H", "L", "A"]:
-        raise ValueError("Mode must be one of the following: ['H', 'L', 'A']")
-
-    if ion_density_centre < 0:
-        raise ValueError("ion_density_centre must greater than 0")
-
-    if ion_density_pedestal < 0:
-        raise ValueError("ion_density_pedestal must greater than 0")
-
-    if ion_density_separatrix < 0:
-        raise ValueError("ion_density_separatrix must greater than 0")
+    cv.check_type('major_radius', major_radius, (int, float))
+    cv.check_type('minor_radius', minor_radius, (int, float))
+    cv.check_type('elongation', elongation, (int, float))
+    cv.check_type('triangularity', triangularity, (int, float))
+    cv.check_type('ion_density_centre', ion_density_centre, (int, float))
+    cv.check_type('ion_density_peaking_factor', ion_density_peaking_factor, (int, float))
+    cv.check_type('ion_density_pedestal', ion_density_pedestal, (int, float))
+    cv.check_type('ion_density_separatrix', ion_density_separatrix, (int, float))
+    cv.check_less_than('minor_radius', minor_radius, major_radius)
+    cv.check_less_than('pedestal_radius', pedestal_radius, minor_radius)
+    cv.check_less_than('shafranov_factor', abs(shafranov_factor), 0.5 * minor_radius)
+    cv.check_greater_than('major_radius', major_radius, 0)
+    cv.check_greater_than('minor_radius', minor_radius, 0)
+    cv.check_greater_than('elongation', elongation, 0)
+    cv.check_less_than('triangularity', triangularity, 1., True)
+    cv.check_greater_than('triangularity', triangularity, -1., True)
+    cv.check_value('mode', mode, ["H", "L", "A"])
+    cv.check_greater_than('ion_density_centre', ion_density_centre, 0)
+    cv.check_greater_than('ion_density_pedestal', ion_density_pedestal, 0)
+    cv.check_greater_than('ion_density_separatrix', ion_density_separatrix, 0)
 
     if (
         isinstance(angles, tuple)

tests/test_tokamak_source.py

@@ -65,7 +65,7 @@ def test_bad_major_radius(tokamak_args_dict, minor_radius, major_radius):
     """Checks that tokamak_source creation rejects invalid major radius"""
     tokamak_args_dict["minor_radius"] = minor_radius
     tokamak_args_dict["major_radius"] = major_radius
-    with pytest.raises(ValueError):
+    with pytest.raises((ValueError, TypeError)):
         tokamak_source(**tokamak_args_dict)
 
 
@@ -90,7 +90,7 @@ def test_bad_minor_radius(tokamak_args_dict, major_radius, minor_radius):
     """Checks that tokamak_source creation rejects invalid minor radius"""
     tokamak_args_dict["major_radius"] = major_radius
     tokamak_args_dict["minor_radius"] = minor_radius
-    with pytest.raises(ValueError):
+    with pytest.raises((ValueError, TypeError)):
         tokamak_source(**tokamak_args_dict)
 
 
@@ -105,7 +105,7 @@ def test_elongation(tokamak_args_dict, elongation):
 def test_bad_elongation(tokamak_args_dict, elongation):
     """Checks that tokamak_source creation rejects invalid elongation"""
     tokamak_args_dict["elongation"] = elongation
-    with pytest.raises(ValueError):
+    with pytest.raises((ValueError, TypeError)):
         tokamak_source(**tokamak_args_dict)
 
 
@@ -120,7 +120,7 @@ def test_triangularity(tokamak_args_dict, triangularity):
 def test_bad_triangularity(tokamak_args_dict, triangularity):
     """Checks that tokamak_source creation rejects invalid triangularity"""
     tokamak_args_dict["triangularity"] = triangularity
-    with pytest.raises(ValueError):
+    with pytest.raises((ValueError, TypeError)):
         tokamak_source(**tokamak_args_dict)
 
 
@@ -184,11 +184,11 @@ def test_bad_angles(tokamak_args_dict, angles):
     # It should fail when given something that isn't a 2-tuple or similar
     # Contents should convert to float
     tokamak_args_dict["angles"] = angles
-    with pytest.raises(ValueError) as excinfo:
+    with pytest.raises((ValueError, TypeError)):
         tokamak_source(**tokamak_args_dict)
 
 
-def test_ion_density(tokamak_args_dict, tokamak_source_example):
+def test_ion_density(tokamak_args_dict):
     # test with values of r that are within acceptable ranges.
     r = np.linspace(0.0, tokamak_args_dict["minor_radius"], 100)
     density = tokamak_ion_density(
@@ -206,7 +206,7 @@ def test_ion_density(tokamak_args_dict, tokamak_source_example):
     assert np.all(np.isfinite(density))
 
 
-def test_bad_ion_density(tokamak_args_dict, tokamak_source_example):
+def test_bad_ion_density(tokamak_args_dict):
     # It should fail if given a negative r
     with pytest.raises(ValueError) as excinfo:
         r = [0, 5, -6]