Testing TokamakSource generates points correctly

.circleci/config.yml

@@ -9,6 +9,8 @@ jobs:
       - run:
           name: install
           command: |
+            python -m pip install --upgrade pip
+            python -m pip install pytest hypothesis
             python setup.py install
             apt-get update -y && apt-get install -y curl
       - run:

tests/test_tokamak_source.py

@@ -1,74 +1,117 @@
 from openmc_plasma_source import TokamakSource
-
 from openmc import Source
+import numpy as np
 
 import pytest
+from hypothesis import given, settings, assume, strategies as st
+
+@st.composite
+def tokamak_source_strategy(draw,return_dict=False):
+    """Defines a hypothesis strategy that automatically generates a TokamakSource.
+    When passed `return_dict=True`, only a dict of the args is passed.
+    Geometry attributes are varied, while plasma attributes are fixed.
+    """
+    # Used to avoid generation of inappropriate float values
+    finites = {"allow_nan": False, "allow_infinity": False, "allow_subnormal": False}
+
+    # Specify the base strategies for each geometry input
+    minor_radius = draw(st.floats(min_value=0.0, max_value=100.0, **finites))
+    major_radius = draw(st.floats(min_value=0.0, max_value=100.0, **finites))
+    pedestal_radius = draw(st.floats(min_value=0.0, max_value=100.0, **finites))
+    elongation = draw(st.floats(min_value=1.0, max_value=10.0, **finites))
+    triangularity = draw(st.floats(min_value=-1.0, max_value=1.0, **finites))
+    shafranov_factor = draw(st.floats(min_value=0.0, max_value=1.0, **finites))
+
+    # Specify requirements that must be satisfied for a valid tokamak
+    assume(major_radius > minor_radius)
+    assume(minor_radius > pedestal_radius)
+    assume(minor_radius > shafranov_factor)
 
+    args_dict = {
+        "elongation" : elongation,
+        "triangularity" : triangularity,
+        "major_radius" : major_radius,
+        "minor_radius" : minor_radius,
+        "pedestal_radius" : pedestal_radius,
+        "shafranov_factor" : shafranov_factor,
+        "ion_density_centre" : 1.09e20,
+        "ion_density_peaking_factor" : 1,
+        "ion_density_pedestal" : 1.09e20,
+        "ion_density_separatrix" : 3e19,
+        "ion_temperature_centre" : 45.9,
+        "ion_temperature_peaking_factor" : 8.06,
+        "ion_temperature_pedestal" : 6.09,
+        "ion_temperature_separatrix" : 0.1,
+        "mode" : "H",
+        "ion_temperature_beta" : 6,
+    }
 
-def test_creation():
-    my_source = TokamakSource(
-        elongation=1.557,
-        ion_density_centre=1.09e20,
-        ion_density_peaking_factor=1,
-        ion_density_pedestal=1.09e20,
-        ion_density_separatrix=3e19,
-        ion_temperature_centre=45.9,
-        ion_temperature_peaking_factor=8.06,
-        ion_temperature_pedestal=6.09,
-        ion_temperature_separatrix=0.1,
-        major_radius=9.06,
-        minor_radius=2.92258,
-        pedestal_radius=0.8 * 2.92258,
-        mode="H",
-        shafranov_factor=0.44789,
-        triangularity=0.270,
-        ion_temperature_beta=6,
-    )
-    for source in my_source.make_openmc_sources():
+    return args_dict if return_dict else TokamakSource(**args_dict)
+
+
+@given(tokamak_source=tokamak_source_strategy())
+@settings(max_examples=1, deadline=None)
+def test_creation(tokamak_source):
+    """Tests that the sources generated by TokamakSource are of type openmc.Source"""
+    for source in tokamak_source.sources:
         assert isinstance(source, Source)
 
 
-def test_strengths_are_normalised():
+@given(tokamak_source=tokamak_source_strategy())
+@settings(max_examples=100, deadline=None)
+def test_strengths_are_normalised(tokamak_source):
     """Tests that the sum of the strengths attribute is equal to"""
-    my_source = TokamakSource(
-        elongation=1.557,
-        ion_density_centre=1.09e20,
-        ion_density_peaking_factor=1,
-        ion_density_pedestal=1.09e20,
-        ion_density_separatrix=3e19,
-        ion_temperature_centre=45.9,
-        ion_temperature_peaking_factor=8.06,
-        ion_temperature_pedestal=6.09,
-        ion_temperature_separatrix=0.1,
-        major_radius=9.06,
-        minor_radius=2.92258,
-        pedestal_radius=0.8 * 2.92258,
-        mode="H",
-        shafranov_factor=0.44789,
-        triangularity=0.270,
-        ion_temperature_beta=6,
-    )
-    assert pytest.approx(sum(my_source.strengths), 1)
-
-
-def test_angles():
+    assert pytest.approx(sum(tokamak_source.strengths), 1)
+
+
+@given(tokamak_source=tokamak_source_strategy())
+@settings(max_examples=100, deadline=None)
+def test_source_locations_are_within_correct_range(tokamak_source):
+    """Tests that each source has RZ locations within the expected range.
+    As the function converting (a,alpha) coordinates to (R,Z) is not bijective, we
+    cannot convert back to validate each individual point. However, we can determine
+    whether the generated points are contained within the shell of the last closed
+    magnetic surface.
+    See "Tokamak D-T neutron source models for different plasma physics confinement
+    modes", C. Fausser et al., Fusion Engineering and Design, 2012 for more info.
+    """
+    R_0 = tokamak_source.major_radius
+    A = tokamak_source.minor_radius
+    El = tokamak_source.elongation
+    delta = tokamak_source.triangularity
+
+    def get_R_on_LCMS(alpha):
+        """Gets R on the last closed magnetic surface for a given alpha"""
+        return R_0 + A * np.cos(alpha + delta * np.sin(alpha))
+
+    approx_lt = lambda x,y :  x < y or np.isclose(x,y)
+    approx_gt = lambda x,y :  x > y or np.isclose(x,y)
+
+    for source in tokamak_source.sources:
+        R, Z = source.space.r.x[0], source.space.z.x[0]
+        # First test that the point is contained with a simple box with
+        # lower left (r_min,-z_max) and upper right (r_max,z_max)
+        assert approx_gt(R,R_0 - A)
+        assert approx_lt(R,R_0 + A)
+        assert approx_lt(abs(Z), A*El)
+        # For a given Z, we can determine the two values of alpha where
+        # where a = minor_radius, and from there determine the upper and
+        # lower bounds for R.
+        alpha_1 = np.arcsin(abs(Z) / (El * A))
+        alpha_2 = np.pi - alpha_1
+        R_max, R_min = get_R_on_LCMS(alpha_1), get_R_on_LCMS(alpha_2)
+        assert approx_lt(R_max,R_0+A)
+        assert approx_gt(R_min,R_0-A)
+        assert approx_lt(R, R_max)
+        assert approx_gt(R, R_min)
+
+
+@given(tokamak_args_dict=tokamak_source_strategy(return_dict=True))
+@settings(max_examples=1, deadline=None)
+def test_angles(tokamak_args_dict):
     """Checks that custom angles can be set"""
-    my_source = TokamakSource(
-        elongation=1.557,
-        ion_density_centre=1.09e20,
-        ion_density_peaking_factor=1,
-        ion_density_pedestal=1.09e20,
-        ion_density_separatrix=3e19,
-        ion_temperature_centre=45.9,
-        ion_temperature_peaking_factor=8.06,
-        ion_temperature_pedestal=6.09,
-        ion_temperature_separatrix=0.1,
-        major_radius=9.06,
-        minor_radius=2.92258,
-        pedestal_radius=0.8 * 2.92258,
-        mode="H",
-        shafranov_factor=0.44789,
-        triangularity=0.270,
-        ion_temperature_beta=6,
-        angles=(0, 1),
-    )
+    tokamak_args_dict["angles"] = (0, 1)
+    tokamak_source = TokamakSource(**tokamak_args_dict)
+    assert tokamak_source.angles == (0, 1)
+    for source in tokamak_source.sources:
+        assert (source.space.phi.a, source.space.phi.b) == (0, 1)