Merge branch 'improved_energy_distribution' of github.com:fusion-ener…

…gy/openmc-plasma-source into improved_energy_distribution

examples/plot_tokamak_ion_temperature.py

@@ -2,25 +2,25 @@
 import numpy as np
 from openmc_plasma_source import tokamak_convert_a_alpha_to_R_Z, tokamak_ion_temperature
 
-sample_size=1000
-minor_radius=292.258
-major_radius=906
+sample_size = 1000
+minor_radius = 292.258
+major_radius = 906
 
 # create a sample of (a, alpha) coordinates
 a = np.random.random(sample_size) * minor_radius
 alpha = np.random.random(sample_size) * 2 * np.pi
 
 temperatures = tokamak_ion_temperature(
-        r=a,
-        mode='L',
-        pedestal_radius=0.8 * minor_radius,
-        ion_temperature_pedestal=6.09,
-        ion_temperature_centre=45.9,
-        ion_temperature_beta=2,
-        ion_temperature_peaking_factor=8.06,
-        ion_temperature_separatrix=0.1,
-        major_radius=major_radius,
-    )
+    r=a,
+    mode="L",
+    pedestal_radius=0.8 * minor_radius,
+    ion_temperature_pedestal=6.09,
+    ion_temperature_centre=45.9,
+    ion_temperature_beta=2,
+    ion_temperature_peaking_factor=8.06,
+    ion_temperature_separatrix=0.1,
+    major_radius=major_radius,
+)
 
 RZ = tokamak_convert_a_alpha_to_R_Z(
     a=a,
@@ -33,10 +33,10 @@
 )
 
 plt.scatter(RZ[0], RZ[1], c=temperatures)
-plt.gca().set_aspect('equal')
+plt.gca().set_aspect("equal")
 plt.xlabel("R [cm]")
 plt.ylabel("Z [cm]")
-plt.colorbar(label='Ion temperature [eV]')
+plt.colorbar(label="Ion temperature [eV]")
 
-plt.savefig('tokamak_source_ion_temperature.png')
-print('written tokamak_source_ion_temperature.png')
+plt.savefig("tokamak_source_ion_temperature.png")
+print("written tokamak_source_ion_temperature.png")

tests/test_ring_source.py

@@ -25,6 +25,7 @@ def test_radius(radius):
     # should allow any positive float
     fusion_ring_source(radius=radius)
 
+
 @pytest.mark.parametrize("radius", [-1.0, "hello world", [1e5]])
 def test_bad_radius(radius):
     # should reject any negative float or anything not convertible to float
@@ -51,19 +52,21 @@ def test_temperature(temperature):
     # Should accept any positive float
     fusion_ring_source(radius=1.0, temperature=temperature)
 
+
 @pytest.mark.parametrize("temperature", [-20000.0, "hello world", [10000]])
 def test_bad_temperature(temperature):
     # Should reject negative floats and anything that isn't convertible to float
     with pytest.raises(ValueError):
         fusion_ring_source(radius=1.0, temperature=temperature)
 
 
-@pytest.mark.parametrize("fuel", [{'D':0.5,'T':0.5}, {"D":1.}])
+@pytest.mark.parametrize("fuel", [{"D": 0.5, "T": 0.5}, {"D": 1.0}])
 def test_fuel(fuel):
     # Should accept either 'DD' or 'DT'
     fusion_ring_source(radius=1.0, fuel=fuel)
 
-@pytest.mark.parametrize("fuel", [{"топливо": 1.}])
+
+@pytest.mark.parametrize("fuel", [{"топливо": 1.0}])
 def test_wrong_fuel(fuel):
     # Should reject fuel types besides those listed in fuel_types.py
     with pytest.raises(ValueError):