toko tests working

src/openmc_plasma_source/tokamak_source.py

@@ -73,6 +73,31 @@ 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")
 
@@ -82,24 +107,15 @@ def tokamak_source(
     if abs(shafranov_factor) >= 0.5 * minor_radius:
         raise ValueError("Shafranov factor must be smaller than 0.5*minor radius")
 
-    if not isinstance(major_radius, (int, float)):
-        raise ValueError("Major radius must be a number")
-
     if major_radius < 0:
         raise ValueError("Major radius must greater than 0")
 
-    if not isinstance(minor_radius, (int, float)):
-        raise ValueError("Minor radius must be a number")
     if minor_radius < 0:
         raise ValueError("Minor radius must greater than 0")
 
-    if not isinstance(elongation, (int, float)):
-        raise ValueError("Elongation must be a number")
-    if elongation < 0:
+    if elongation <= 0:
         raise ValueError("Elongation must greater than 0")
 
-    if not isinstance(triangularity, (int, float)):
-        raise ValueError("Triangularity must be a number")
     if not triangularity <= 1.0:
         raise ValueError("Triangularity must less than or equal to 1.")
     if not triangularity >= -1.0:
@@ -108,21 +124,13 @@ def tokamak_source(
     if mode not in ["H", "L", "A"]:
         raise ValueError("Mode must be one of the following: ['H', 'L', 'A']")
 
-    if not isinstance(ion_density_centre, (int, float)):
-        raise ValueError("ion_density_centre must be a number")
     if ion_density_centre < 0:
         raise ValueError("ion_density_centre must greater than 0")
 
-    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 ion_density_pedestal < 0:
         raise ValueError("ion_density_pedestal must greater than 0")
 
-    if not isinstance(ion_density_separatrix, (int, float)):
-        raise ValueError("ion_density_separatrix must be a number")
     if ion_density_separatrix < 0:
         raise ValueError("ion_density_separatrix must greater than 0")
 
@@ -278,7 +286,7 @@ def tokamak_ion_temperature(
             (
                 ion_temperature_pedestal
                 + (ion_temperature_centre - ion_temperature_pedestal)
-                * (1 - (r / pedestal_radius) ** ion_temperature_beta)
+                * (1 - (np.abs(r / pedestal_radius)) ** ion_temperature_beta)
                 ** ion_temperature_peaking_factor
             ),
             (

tests/test_tokamak_source.py

@@ -1,9 +1,9 @@
 import numpy as np
 import pytest
-from hypothesis import given, settings
+from hypothesis import given, settings, HealthCheck
 from hypothesis import strategies as st
 import openmc
-from openmc_plasma_source import tokamak_source, tokamak_ion_density
+from openmc_plasma_source import tokamak_source, tokamak_ion_density, tokamak_ion_temperature, tokamak_convert_a_alpha_to_R_Z
 
 
 @pytest.fixture
@@ -167,7 +167,7 @@ def test_bad_shafranov_factor(tokamak_args_dict, major_radius, minor_radius, sha
 def test_angles(tokamak_args_dict, angles):
     """Checks that custom angles can be set"""
     # Note: should accept negative angles and angles in reverse order
-    angles = tokamak_args_dict["angles"]
+    tokamak_args_dict["angles"] = angles
     sources = tokamak_source(**tokamak_args_dict)
     for source in sources:
         assert np.array_equal((source.space.phi.a, source.space.phi.b), angles)
@@ -180,63 +180,119 @@ def test_bad_angles(tokamak_args_dict, angles):
     # Contents should convert to float
     tokamak_args_dict["angles"] = angles
     with pytest.raises(ValueError) as excinfo:
-        tokamak_source = tokamak_source(**tokamak_args_dict)
-
-
-# def test_ion_density(tokamak_source_example):
-#     # test with values of r that are within acceptable ranges.
-#     r = np.linspace(0.0, tokamak_source_example.minor_radius, 100)
-#     density = tokamak_source_example.ion_density(r)
-#     assert isinstance(r, np.ndarray)
-#     assert len(density) == len(r)
-#     assert np.all(np.isfinite(density))
-
-
-# def test_bad_ion_density(tokamak_source_example):
-#     # It should fail if given a negative r
-#     with pytest.raises(ValueError) as excinfo:
-#         density = tokamak_ion_density([0, 5, -6])
-#     assert "must not be negative" in str(excinfo.value)
-
-
-# def test_ion_temperature(tokamak_source_example):
-#     # test with values of r that are within acceptable ranges.
-#     r = np.linspace(0.0, tokamak_source_example.minor_radius, 100)
-#     temperature = tokamak_source_example.ion_temperature(r)
-#     assert isinstance(temperature, np.ndarray)
-#     assert len(temperature) == len(r)
-#     assert np.all(np.isfinite(temperature))
+        tokamak_source(**tokamak_args_dict)
 
 
-# def test_bad_ion_temperature(tokamak_source_example):
-#     # It should fail if given a negative r
-#     with pytest.raises(ValueError) as excinfo:
-#         temperature = tokamak_source_example.ion_temperature([0, 5, -6])
-#     assert "must not be negative" in str(excinfo.value)
+def test_ion_density(tokamak_args_dict, tokamak_source_example):
+    # 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(
+        r=r,
+        mode='L',
+        ion_density_centre=tokamak_args_dict['ion_density_centre'], 
+        ion_density_peaking_factor=tokamak_args_dict['ion_density_peaking_factor'], 
+        ion_density_pedestal=tokamak_args_dict['ion_density_pedestal'], 
+        major_radius=tokamak_args_dict['major_radius'], 
+        pedestal_radius=tokamak_args_dict['pedestal_radius'], 
+        ion_density_separatrix=tokamak_args_dict['ion_density_separatrix'],
+    )
+    assert isinstance(r, np.ndarray)
+    assert len(density) == len(r)
+    assert np.all(np.isfinite(density))
 
 
-# def test_convert_a_alpha_to_R_Z(tokamak_source_example):
-#     # Similar to  test_source_locations_are_within_correct_range
-#     # Rather than going in detail, simply tests validity of inputs and outputs
-#     # Test with suitable values for a and alpha
-#     a = np.linspace(0.0, tokamak_source_example.minor_radius, 100)
-#     alpha = np.linspace(0.0, 2 * np.pi, 100)
-#     R, Z = tokamak_source_example.convert_a_alpha_to_R_Z(a, alpha)
-#     assert isinstance(R, np.ndarray)
-#     assert isinstance(Z, np.ndarray)
-#     assert len(R) == len(a)
-#     assert len(Z) == len(a)
-#     assert np.all(np.isfinite(R))
-#     assert np.all(np.isfinite(Z))
+def test_bad_ion_density(tokamak_args_dict, tokamak_source_example):
+    # It should fail if given a negative r
+    with pytest.raises(ValueError) as excinfo:
+        r=[0, 5, -6]
+        tokamak_ion_density(
+            r=r,
+            mode='L',
+            ion_density_centre=tokamak_args_dict['ion_density_centre'], 
+            ion_density_peaking_factor=tokamak_args_dict['ion_density_peaking_factor'], 
+            ion_density_pedestal=tokamak_args_dict['ion_density_pedestal'], 
+            major_radius=tokamak_args_dict['major_radius'], 
+            pedestal_radius=tokamak_args_dict['pedestal_radius'], 
+            ion_density_separatrix=tokamak_args_dict['ion_density_separatrix'],
+        )
+    assert "must not be negative" in str(excinfo.value)
+
+
+def test_ion_temperature(tokamak_args_dict, tokamak_source_example):
+    # test with values of r that are within acceptable ranges.
+    r = np.linspace(0.0, 2.9, 100)
+    temperature = tokamak_ion_temperature(
+        r=r,
+        mode=tokamak_args_dict['mode'],
+        pedestal_radius=tokamak_args_dict['pedestal_radius'],
+        ion_temperature_pedestal=tokamak_args_dict['ion_temperature_pedestal'],
+        ion_temperature_centre=tokamak_args_dict['ion_temperature_centre'],
+        ion_temperature_beta=tokamak_args_dict['ion_temperature_beta'],
+        ion_temperature_peaking_factor=tokamak_args_dict['ion_temperature_peaking_factor'],
+        ion_temperature_separatrix=tokamak_args_dict['ion_temperature_separatrix'],
+        major_radius=tokamak_args_dict['major_radius'],
+    )
+    assert isinstance(temperature, np.ndarray)
+    assert len(temperature) == len(r)
+    assert np.all(np.isfinite(temperature))
 
 
-# def test_bad_convert_a_alpha_to_R_Z(tokamak_source_example):
-#     # Repeat test_convert_a_alpha_to_R_Z, but show that negative a breaks it
-#     a = np.linspace(0.0, tokamak_source_example.minor_radius, 100)
-#     alpha = np.linspace(0.0, 2 * np.pi, 100)
-#     with pytest.raises(ValueError) as excinfo:
-#         R, Z = tokamak_source_example.convert_a_alpha_to_R_Z(-a, alpha)
-#     assert "must not be negative" in str(excinfo.value)
+def test_bad_ion_temperature(tokamak_args_dict):
+    # It should fail if given a negative r
+    with pytest.raises(ValueError) as excinfo:
+        r=[0, 5, -6]
+        tokamak_ion_temperature(
+            r=r,
+            mode=tokamak_args_dict['mode'],
+            pedestal_radius=tokamak_args_dict['pedestal_radius'],
+            ion_temperature_pedestal=tokamak_args_dict['ion_temperature_pedestal'],
+            ion_temperature_centre=tokamak_args_dict['ion_temperature_centre'],
+            ion_temperature_beta=tokamak_args_dict['ion_temperature_beta'],
+            ion_temperature_peaking_factor=tokamak_args_dict['ion_temperature_peaking_factor'],
+            ion_temperature_separatrix=tokamak_args_dict['ion_temperature_separatrix'],
+            major_radius=tokamak_args_dict['major_radius'],
+        )
+    assert "must not be negative" in str(excinfo.value)
+
+
+def test_convert_a_alpha_to_R_Z(tokamak_args_dict):
+    # Similar to  test_source_locations_are_within_correct_range
+    # Rather than going in detail, simply tests validity of inputs and outputs
+    # Test with suitable values for a and alpha
+    a = np.linspace(0.0, 2.9, 100)
+    alpha = np.linspace(0.0, 2 * np.pi, 100)
+    R, Z = tokamak_convert_a_alpha_to_R_Z(
+        a=a,
+        alpha=alpha,
+        shafranov_factor=tokamak_args_dict['shafranov_factor'], 
+        minor_radius=tokamak_args_dict['minor_radius'], 
+        major_radius=tokamak_args_dict['major_radius'], 
+        triangularity=tokamak_args_dict['triangularity'], 
+        elongation=tokamak_args_dict['elongation'],
+    )
+    assert isinstance(R, np.ndarray)
+    assert isinstance(Z, np.ndarray)
+    assert len(R) == len(a)
+    assert len(Z) == len(a)
+    assert np.all(np.isfinite(R))
+    assert np.all(np.isfinite(Z))
+
+
+def test_bad_convert_a_alpha_to_R_Z(tokamak_args_dict):
+    # Repeat test_convert_a_alpha_to_R_Z, but show that negative a breaks it
+    a = np.linspace(0.0, 2.9, 100)
+    alpha = np.linspace(0.0, 2 * np.pi, 100)
+    with pytest.raises(ValueError) as excinfo:
+        tokamak_convert_a_alpha_to_R_Z(
+            a=-a,
+            alpha=alpha,
+            shafranov_factor=tokamak_args_dict['shafranov_factor'], 
+            minor_radius=tokamak_args_dict['minor_radius'], 
+            major_radius=tokamak_args_dict['major_radius'], 
+            triangularity=tokamak_args_dict['triangularity'], 
+            elongation=tokamak_args_dict['elongation'],
+        )
+    assert "must not be negative" in str(excinfo.value)
 
 
 @st.composite
@@ -303,54 +359,63 @@ def tokamak_source_strategy(draw):
         ion_temperature_separatrix=0.1,
         mode="H",
         ion_temperature_beta=6,
-    )
+    ), {
+        'major_radius':major_radius,
+        'minor_radius':minor_radius,
+        'elongation':elongation,
+        'triangularity':triangularity,
+    }
 
 
-# @given(tokamak_source=tokamak_source_strategy())
-# @settings(max_examples=50)
-# def test_strengths_are_normalised(tokamak_source):
-#     """Tests that the sum of the strengths attribute is equal to"""
-#     assert pytest.approx(sum(tokamak_source.strengths)) == 1
-
-
-# @given(tokamak_source=tokamak_source_strategy())
-# @settings(max_examples=50)
-# 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_source=tokamak_source_strategy())
+@settings(max_examples=30, suppress_health_check=(HealthCheck.too_slow,))
+def test_strengths_are_normalised(tokamak_source):
+    """Tests that the sum of the strengths attribute is equal to"""
+    local_strength = 0
+    all_sources = tokamak_source[0]
+    for source in all_sources:
+        local_strength = local_strength + source.strength
+    assert pytest.approx(local_strength) == 1
+
+
+@given(tokamak_source=tokamak_source_strategy())
+@settings(max_examples=2)
+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[1]['major_radius']
+    A = tokamak_source[1]['minor_radius']
+    El = tokamak_source[1]['elongation']
+    delta = tokamak_source[1]['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[0]:
+        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)