diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a475db9..89b2bec 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,3 +31,6 @@ jobs: python examples/point_source_example.py python examples/ring_source_example.py python examples/tokamak_source_example.py + python examples/plot_tokamak_ion_temperature.py + python examples/plot_tokamak_neutron_source_density.py + python examples/plot_tokamak_neutron_source_strengths.py diff --git a/.gitignore b/.gitignore index e16264b..4a64908 100644 --- a/.gitignore +++ b/.gitignore @@ -135,8 +135,9 @@ dmypy.json # vim swap files *.swp -# image files created by tests +# image files created by tests and examples tests/*.png +*.png # openmc output files *.h5 diff --git a/README.md b/README.md index 3d77f75..26d0903 100644 --- a/README.md +++ b/README.md @@ -8,9 +8,9 @@ This python-based package offers a collection of pre-built [OpenMC](https://gith OpenMC is required to use this package. -To install openmc-plasma-source, simply run: +To install openmc_plasma_source, simply run: ``` -pip install openmc-plasma-source +pip install openmc_plasma_source ``` ## Usage @@ -24,26 +24,27 @@ Each source has its own strength (or probability that a neutron spawns in this l The equations implemented here are taken from [this paper](https://doi.org/10.1016/j.fusengdes.2012.02.025). ```python -from openmc_plasma_source import TokamakSource +from openmc_plasma_source import tokamak_source -my_source = TokamakSource( +my_sources = tokamak_source( elongation=1.557, ion_density_centre=1.09e20, - ion_density_peaking_factor=1, ion_density_pedestal=1.09e20, + ion_density_peaking_factor=1, ion_density_separatrix=3e19, - ion_temperature_centre=45.9, + ion_temperature_centre=45.9e3, + ion_temperature_pedestal=6.09e3, + ion_temperature_separatrix=0.1e3, 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, + ion_temperature_beta=6, + major_radius=906, + minor_radius=292.258, + pedestal_radius=0.8 * 292.258, mode="H", shafranov_factor=0.44789, triangularity=0.270, - ion_temperature_beta=6 - ).make_openmc_sources() + fuel={"D": 0.5, "T": 0.5}, +) ``` For a more complete example check out the [example script](https://github.com/fusion-energy/openmc-plasma-source/blob/main/examples/tokamak_source_example.py). @@ -57,13 +58,13 @@ For a more complete example check out the [example script](https://github.com/fu Create a ring source with temperature distribution of a 2000 eV plasma. ```python -from openmc_plasma_source import FusionRingSource +from openmc_plasma_source import fusion_ring_source -my_plasma = FusionRingSource( - angles = (0., 6.28318530718), # input is in radians - radius = 400, # units in cm - temperature = 20000., # ion temperature in eV - fuel='DT' # or 'DD' +my_source = fusion_ring_source( + radius=700, + angles=(0.0, 2 * math.pi), # 360deg source + temperature=20000.0, + fuel={"D": 0.5, "T": 0.5}, ) ``` ### Point Source @@ -72,12 +73,12 @@ Create a point source with temperature distribution of a 2000 eV plasma. ```python -from openmc_plasma_source import FusionPointSource +from openmc_plasma_source import fusion_point_source -my_plasma = FusionPointSource( - coordinate = (0, 0, 0), - temperature = 20000., # ion temperature in eV - fuel = 'DT' # or 'DD' +my_source = fusion_point_source( + coordinate=(0, 0, 0), + temperature=20000.0, + fuel={"D": 0.09, "T": 0.91}, # note this is mainly tritium fuel so that TT reactions are more likely ) ``` @@ -86,5 +87,5 @@ my_plasma = FusionPointSource( To run the tests, simply run ``` -pytest tests/ +pytest tests ``` diff --git a/examples/plot_tokamak_ion_temperature.py b/examples/plot_tokamak_ion_temperature.py new file mode 100644 index 0000000..933a5f2 --- /dev/null +++ b/examples/plot_tokamak_ion_temperature.py @@ -0,0 +1,42 @@ +import matplotlib.pyplot as plt +import numpy as np +from openmc_plasma_source import tokamak_convert_a_alpha_to_R_Z, tokamak_ion_temperature + +sample_size = 20000 +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, +) + +RZ = tokamak_convert_a_alpha_to_R_Z( + a=a, + alpha=alpha, + shafranov_factor=0.44789, + minor_radius=minor_radius, + major_radius=major_radius, + triangularity=0.270, + elongation=1.557, +) + +plt.scatter(RZ[0], RZ[1], c=temperatures) +plt.gca().set_aspect("equal") +plt.xlabel("R [cm]") +plt.ylabel("Z [cm]") +plt.colorbar(label="Ion temperature [eV]") + +plt.savefig("tokamak_source_ion_temperature.png") +print("written tokamak_source_ion_temperature.png") diff --git a/examples/plot_tokamak_neutron_source_density.py b/examples/plot_tokamak_neutron_source_density.py new file mode 100644 index 0000000..34b0e89 --- /dev/null +++ b/examples/plot_tokamak_neutron_source_density.py @@ -0,0 +1,62 @@ +import matplotlib.pyplot as plt +import numpy as np +from openmc_plasma_source import ( + tokamak_ion_temperature, + tokamak_convert_a_alpha_to_R_Z, + tokamak_neutron_source_density, + tokamak_ion_density, +) + +sample_size = 20000 +minor_radius = 292.258 +major_radius = 906 +mode = "L" +ion_density_centre = 45.9 + +# 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=mode, + pedestal_radius=0.8 * minor_radius, + ion_temperature_pedestal=6.09, + ion_temperature_centre=ion_density_centre, + ion_temperature_beta=2, + ion_temperature_peaking_factor=8.06, + ion_temperature_separatrix=0.1, + major_radius=major_radius, +) + +densities = tokamak_ion_density( + mode=mode, + ion_density_centre=ion_density_centre, + ion_density_peaking_factor=1, + ion_density_pedestal=1.09e20, + major_radius=major_radius, + pedestal_radius=0.8 * minor_radius, + ion_density_separatrix=3e19, + r=a, +) + +neutron_source_density = tokamak_neutron_source_density(densities, temperatures, "DD") + +RZ = tokamak_convert_a_alpha_to_R_Z( + a=a, + alpha=alpha, + shafranov_factor=0.44789, + minor_radius=minor_radius, + major_radius=major_radius, + triangularity=0.270, + elongation=1.557, +) + +plt.scatter(RZ[0], RZ[1], c=neutron_source_density) +plt.gca().set_aspect("equal") +plt.xlabel("R [cm]") +plt.ylabel("Z [cm]") +plt.colorbar(label="neutron source density") + +plt.savefig("tokamak_source_neutron_source_density.png") +print("written tokamak_source_neutron_source_density.png") diff --git a/examples/plot_tokamak_neutron_source_strengths.py b/examples/plot_tokamak_neutron_source_strengths.py new file mode 100644 index 0000000..5e1858e --- /dev/null +++ b/examples/plot_tokamak_neutron_source_strengths.py @@ -0,0 +1,64 @@ +import matplotlib.pyplot as plt +import numpy as np +from openmc_plasma_source import ( + tokamak_ion_temperature, + tokamak_convert_a_alpha_to_R_Z, + tokamak_neutron_source_density, + tokamak_ion_density, +) + +sample_size = 20000 +minor_radius = 292.258 +major_radius = 906 +mode = "L" +ion_density_centre = 45.9 + +# 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=mode, + pedestal_radius=0.8 * minor_radius, + ion_temperature_pedestal=6.09, + ion_temperature_centre=ion_density_centre, + ion_temperature_beta=2, + ion_temperature_peaking_factor=8.06, + ion_temperature_separatrix=0.1, + major_radius=major_radius, +) + +densities = tokamak_ion_density( + mode=mode, + ion_density_centre=ion_density_centre, + ion_density_peaking_factor=1, + ion_density_pedestal=1.09e20, + major_radius=major_radius, + pedestal_radius=0.8 * minor_radius, + ion_density_separatrix=3e19, + r=a, +) + +neutron_source_density = tokamak_neutron_source_density(densities, temperatures, "DD") + +strengths = neutron_source_density / sum(neutron_source_density) + +RZ = tokamak_convert_a_alpha_to_R_Z( + a=a, + alpha=alpha, + shafranov_factor=0.44789, + minor_radius=minor_radius, + major_radius=major_radius, + triangularity=0.270, + elongation=1.557, +) + +plt.scatter(RZ[0], RZ[1], c=strengths) +plt.gca().set_aspect("equal") +plt.xlabel("R [cm]") +plt.ylabel("Z [cm]") +plt.colorbar(label="neutron emission strength") + +plt.savefig("tokamak_source_neutron_emission_strength.png") +print("written tokamak_source_neutron_emission_strength.png") diff --git a/examples/point_source_example.py b/examples/point_source_example.py index fb040a1..584ce32 100644 --- a/examples/point_source_example.py +++ b/examples/point_source_example.py @@ -1,6 +1,8 @@ -import openmc -from openmc_plasma_source import FusionPointSource from pathlib import Path +import numpy as np + +import openmc +from openmc_plasma_source import fusion_point_source # just making use of a local cross section xml file, replace with your own cross sections or comment out openmc.config["cross_sections"] = Path(__file__).parent.resolve() / "cross_sections.xml" @@ -11,7 +13,14 @@ geometry = openmc.Geometry([cell]) # define the source -my_source = FusionPointSource(coordinate=(0, 0, 0), temperature=20000.0, fuel="DT") +my_source = fusion_point_source( + coordinate=(0, 0, 0), + temperature=20000.0, + fuel={ + "D": 0.09, + "T": 0.91, + }, # note this is mainly tritium fuel so that TT reactions are more likely +) # Tell OpenMC we're going to use our custom source settings = openmc.Settings() @@ -20,7 +29,20 @@ settings.particles = 1000 settings.source = my_source - model = openmc.model.Model(materials=None, geometry=geometry, settings=settings) model.run() + + +# optionally if you would like to plot the energy of particles then another package can be used +# https://github.com/fusion-energy/openmc_source_plotter + +from openmc_source_plotter import plot_source_energy + +plot = plot_source_energy( + this=settings, + n_samples=2000000, # increase this value for a smoother plot + energy_bins=np.linspace(0, 16e6, 1000), + yaxis_type="log", +) +plot.show() diff --git a/examples/ring_source_example.py b/examples/ring_source_example.py index f09d570..c72e652 100644 --- a/examples/ring_source_example.py +++ b/examples/ring_source_example.py @@ -1,8 +1,9 @@ -import openmc -from openmc_plasma_source import FusionRingSource import math from pathlib import Path +import openmc +from openmc_plasma_source import fusion_ring_source + # just making use of a local cross section xml file, replace with your own cross sections or comment out openmc.config["cross_sections"] = Path(__file__).parent.resolve() / "cross_sections.xml" @@ -12,10 +13,11 @@ geometry = openmc.Geometry([cell]) # define the source -my_source = FusionRingSource( +my_source = fusion_ring_source( radius=700, angles=(0.0, 2 * math.pi), # 360deg source temperature=20000.0, + fuel={"D": 0.5, "T": 0.5}, ) settings = openmc.Settings() @@ -25,6 +27,19 @@ # tell OpenMC we're going to use our custom source settings.source = my_source -model = openmc.model.Model(materials=None, geometry=geometry, settings=settings) +model = openmc.Model(materials=None, geometry=geometry, settings=settings) model.run() + + +# optionally if you would like to plot the location of particles then another package can be used +# https://github.com/fusion-energy/openmc_source_plotter + +from openmc_source_plotter import plot_source_position + +plot = plot_source_position( + this=settings, + n_samples=2000, +) + +plot.show() diff --git a/examples/tokamak_source_example.py b/examples/tokamak_source_example.py index 9c98390..4878e24 100644 --- a/examples/tokamak_source_example.py +++ b/examples/tokamak_source_example.py @@ -1,34 +1,37 @@ from pathlib import Path + import openmc -from openmc_plasma_source import TokamakSource + +from openmc_plasma_source import tokamak_source # just making use of a local cross section xml file, replace with your own cross sections or comment out openmc.config["cross_sections"] = Path(__file__).parent.resolve() / "cross_sections.xml" # minimal geometry -sphere_surface = openmc.Sphere(r=1000.0, boundary_type="vacuum") +sphere_surface = openmc.Sphere(r=100000.0, boundary_type="vacuum") cell = openmc.Cell(region=-sphere_surface) geometry = openmc.Geometry([cell]) # create a plasma source -my_plasma = TokamakSource( +my_sources = tokamak_source( elongation=1.557, ion_density_centre=1.09e20, - ion_density_peaking_factor=1, ion_density_pedestal=1.09e20, + ion_density_peaking_factor=1, ion_density_separatrix=3e19, - ion_temperature_centre=45.9, + ion_temperature_centre=45.9e3, + ion_temperature_pedestal=6.09e3, + ion_temperature_separatrix=0.1e3, 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, + ion_temperature_beta=6, + major_radius=906, + minor_radius=292.258, + pedestal_radius=0.8 * 292.258, mode="H", shafranov_factor=0.44789, triangularity=0.270, - ion_temperature_beta=6, + fuel={"D": 0.5, "T": 0.5}, ) # Tell OpenMC we're going to use our custom source @@ -36,9 +39,19 @@ settings.run_mode = "fixed source" settings.batches = 10 settings.particles = 1000 -settings.source = my_plasma.sources +settings.source = my_sources model = openmc.model.Model(materials=None, geometry=geometry, settings=settings) model.run() + + +# optionally if you would like to plot the direction of particles then another package can be used +# https://github.com/fusion-energy/openmc_source_plotter + +from openmc_source_plotter import plot_source_position + +plot = plot_source_position(this=settings, n_samples=200) + +plot.show() diff --git a/pyproject.toml b/pyproject.toml index c421db0..3379f2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,11 +18,12 @@ 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", "matplotlib>=3.2.2", + "NeSST>=1.1.0" ] [project.urls] @@ -35,7 +36,8 @@ write_to = "src/_version.py" [project.optional-dependencies] tests = [ "pytest>=5.4.3", - "hypothesis" + "hypothesis", + "NeSST>=1.1.0" ] [tool.setuptools] diff --git a/src/openmc_plasma_source/__init__.py b/src/openmc_plasma_source/__init__.py index e7553a0..0dc1642 100644 --- a/src/openmc_plasma_source/__init__.py +++ b/src/openmc_plasma_source/__init__.py @@ -1,7 +1,7 @@ try: - from importlib.metadata import version, PackageNotFoundError + from importlib.metadata import PackageNotFoundError, version except (ModuleNotFoundError, ImportError): - from importlib_metadata import version, PackageNotFoundError + from importlib_metadata import PackageNotFoundError, version try: __version__ = version("openmc_plasma_source") except PackageNotFoundError: @@ -11,6 +11,7 @@ __all__ = ["__version__"] -from .tokamak_source import TokamakSource -from .ring_source import FusionRingSource -from .point_source import FusionPointSource +from .fuel_types import get_neutron_energy_distribution +from .point_source import fusion_point_source +from .ring_source import fusion_ring_source +from .tokamak_source import * diff --git a/src/openmc_plasma_source/fuel_types.py b/src/openmc_plasma_source/fuel_types.py index 27299b9..5e527a5 100644 --- a/src/openmc_plasma_source/fuel_types.py +++ b/src/openmc_plasma_source/fuel_types.py @@ -1,37 +1,219 @@ -"""fuel_types.py +import NeSST as nst +import numpy as np +import openmc -Defines dictionary for determining mean energy and mass of reactants -for a given fusion fuel type. -""" +def neutron_energy_mean(ion_temperature: float, reaction: str) -> float: + """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 -class Fuel: - def __init__(self, mean_energy, mass_of_reactants): - self.mean_energy = mean_energy - self.mass_of_reactants = mass_of_reactants + Args: + ion_temperature (float): the temperature of the ions in eV + reaction (str): the two isotope that fuse, can be either 'DD' or 'DT' - @property - def mean_energy(self): - return self._mean_energy + Raises: + ValueError: if the reaction is not 'DD' or 'DT' then a ValueError is raised - @mean_energy.setter - def mean_energy(self, value): - if value <= 0: - raise (ValueError("mean_energy needs to be strictly positive")) - self._mean_energy = value + Returns: + float: the mean neutron energy in eV + """ - @property - def mass_of_reactants(self): - return self._mass_of_reactants + # values from Ballabio paper + if reaction == "DD": + a_1 = 4.69515 + a_2 = -0.040729 + a_3 = 0.47 + a_4 = 0.81844 + mean = 2.4495e6 # in units of eV + elif reaction == "DT": + a_1 = 5.30509 + a_2 = 2.4736e-3 + a_3 = 1.84 + a_4 = 1.3818 + mean = 14.021e6 # in units of eV + else: + raise ValueError(f'reaction must be either "DD" or "DT" not {reaction}') - @mass_of_reactants.setter - def mass_of_reactants(self, value): - if value <= 0: - raise (ValueError("mass_of_reactants needs to be strictly positive")) - self._mass_of_reactants = value + ion_temperature_kev = ion_temperature / 1e3 # Ballabio equation accepts KeV units + mean_delta = ( + a_1 + * ion_temperature_kev ** (2.0 / 3.0) + / (1.0 + a_2 * ion_temperature_kev**a_3) + + a_4 * ion_temperature_kev + ) + mean_delta *= 1e3 # converting back to eV + return mean + mean_delta -fuel_types = { - "DD": Fuel(mean_energy=2450000.0, mass_of_reactants=4), - "DT": Fuel(mean_energy=14080000.0, mass_of_reactants=5), -} +def neutron_energy_std_dev(ion_temperature: float, reaction: str) -> float: + """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' + + Raises: + ValueError: if the reaction is not 'DD' or 'DT' then a ValueError is raised + + Returns: + float: the mean neutron energy in eV + """ + + # values from Ballabio paper + if reaction == "DD": + w_0 = 82.542 + a_1 = 1.7013e-3 + a_2 = 0.16888 + a_3 = 0.49 + a_4 = 7.9460e-4 + elif reaction == "DT": + w_0 = 177.259 + a_1 = 5.1068e-4 + a_2 = 7.6223e-3 + a_3 = 1.78 + a_4 = 8.7691e-5 + else: + raise ValueError(f'reaction must be either "DD" or "DT" not {reaction}') + + ion_temperature_kev = ion_temperature / 1e3 # Ballabio equation accepts KeV units + delta = ( + a_1 + * ion_temperature_kev ** (2.0 / 3.0) + / (1.0 + a_2 * ion_temperature_kev**a_3) + + a_4 * ion_temperature_kev + ) + + # 2.3548200450309493 on the line below comes from equation 2* math.sqrt(math.log(2)*2) + variance = ((w_0 * (1 + delta)) ** 2 * ion_temperature_kev) / 2.3548200450309493**2 + variance *= 1e6 # converting keV^2 back to eV^2 + std_dev = np.sqrt(variance) + return std_dev + + +def get_reactions_from_fuel(fuel): + if ["D", "T"] == sorted(set(fuel.keys())): + return ["DT", "DD", "TT"] + elif ["D"] == sorted(set(fuel.keys())): + return ["DD"] + elif ["T"] == sorted(set(fuel.keys())): + return ["TT"] + else: + msg = '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, +) -> 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 + + Returns + ------- + openmc.stats.Discrete + energy distribution + """ + + sum_fuel_isotopes = sum(fuel.values()) + if sum_fuel_isotopes > 1.0: + raise ValueError( + f"isotope fractions within the fuel must sum to be below 1. Not {sum_fuel_isotopes}" + ) + + if sum_fuel_isotopes < 0.0: + raise ValueError( + f"isotope fractions must sum to be above 0. Not {sum_fuel_isotopes}" + ) + + if sum_fuel_isotopes != 1.0: + raise ValueError(f"isotope fractions must sum to 1. Not {sum_fuel_isotopes}") + + for k, v in fuel.items(): + if k not in ["D", "T"]: + raise ValueError( + f'Fuel dictionary keys must be either "D" or "T" not "{k}".' + ) + if v < 0: + raise ValueError(f'Fuel dictionary values must be above 0 not "{k}".') + if v > 1: + raise ValueError(f'Fuel dictionary values must be below 1 not "{k}".') + + # 1.0 neutron yield, all reactions scaled by this value + num_of_vals = 100 + # single grid for TT neutrons + E_pspec = np.linspace(0, 12e6, num_of_vals) # E_pspec is exspected in MeV units + + DDmean = neutron_energy_mean(ion_temperature=ion_temperature, reaction="DD") + DTmean = neutron_energy_mean(ion_temperature=ion_temperature, reaction="DT") + DD_std_dev = neutron_energy_std_dev(ion_temperature=ion_temperature, reaction="DD") + DT_std_dev = neutron_energy_std_dev(ion_temperature=ion_temperature, reaction="DT") + + reactions = get_reactions_from_fuel(fuel) + + if reactions == ["TT"]: + strength_TT = 1.0 + dNdE_TT = strength_TT * nst.dNdE_TT(E_pspec, ion_temperature) + tt_source = openmc.stats.Tabular(E_pspec * 1e6, dNdE_TT) + return tt_source + + elif reactions == ["DD"]: + strength_DD = 1.0 + dd_source = openmc.stats.Normal(mean_value=DDmean, std_dev=DD_std_dev) + return dd_source + + # DT, DD and TT reaction + else: + strength_DD = nst.yield_from_dt_yield_ratio( + "dd", 1.0, ion_temperature, fuel["D"], fuel["T"] + ) + strength_TT = nst.yield_from_dt_yield_ratio( + "tt", 1.0, ion_temperature, fuel["D"], fuel["T"] + ) + + dd_source = openmc.stats.Normal(mean_value=DDmean, std_dev=DD_std_dev) + # normal could be done with Muir but in this case we have the mean and std dev from NeSST + # dd_source = openmc.stats.muir(e0=DDmean * 1e6, m_rat=4, kt=ion_temperature) + + dt_source = openmc.stats.Normal(mean_value=DTmean, std_dev=DT_std_dev) + # normal could be done with Muir but in this case we have the mean and std dev from NeSST + # dt_source = openmc.stats.muir(e0=DTmean * 1e6, m_rat=5, kt=ion_temperature) + + dNdE_TT = strength_TT * nst.dNdE_TT(E_pspec, ion_temperature) + + # removing any zeros from the end of the array + dNdE_TT = np.trim_zeros(dNdE_TT, "b") + # making array lengths match + E_pspec = E_pspec[: len(dNdE_TT)] + + openmc_univariate = [dd_source, dt_source] + + # sometimes bins are empty + if len(E_pspec) == 0: + total_strength = sum([strength_DD, 1.0]) + probabilities = [strength_DD / total_strength, 1.0 / total_strength] + else: + total_strength = sum([strength_TT, strength_DD, 1.0]) + tt_source = openmc.stats.Tabular(E_pspec[1:], dNdE_TT[1:]) + probabilities = [ + strength_DD / total_strength, + 1.0 / total_strength, + strength_TT / total_strength, + ] + openmc_univariate.append(tt_source) + + return openmc.stats.Mixture(probabilities, openmc_univariate) + # bug reported for combine_distributions #3105 on openmc + # return openmc.data.combine_distributions( + # dists=openmc_univariate, + # probs=probabilities + # ) diff --git a/src/openmc_plasma_source/plotting/__init__.py b/src/openmc_plasma_source/plotting/__init__.py deleted file mode 100644 index 9ba21fa..0000000 --- a/src/openmc_plasma_source/plotting/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .plot_tokamak_source import plot_tokamak_source_3D, scatter_tokamak_source diff --git a/src/openmc_plasma_source/plotting/plot_tokamak_source.py b/src/openmc_plasma_source/plotting/plot_tokamak_source.py deleted file mode 100644 index 74d25f5..0000000 --- a/src/openmc_plasma_source/plotting/plot_tokamak_source.py +++ /dev/null @@ -1,153 +0,0 @@ -import matplotlib.pyplot as plt -from matplotlib import cm -import numpy as np - - -def scatter_tokamak_source(source, ax=None, quantity=None, aspect="equal", **kwargs): - """Create a 2D scatter plot of the tokamak source. - See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.scatter.html - for more arguments. - - Args: - source (ops.TokamakSource): the plasma source - ax (maplotlib.Axes, optional): Axes object on which to plot. If not - provided by the user, the current working Axes is retrieved using - matplotlib.pyplot.gca(). - aspect (str, optional): Set aspect ratio of the plot. May be set to 'auto', - 'equal', or a number denoting the ratio of the plot height to the plot - width. - quantity (str, optional): value by which the lines should be - coloured. Defaults to None. - **kwargs: Keyword arguments compatible with matplotlib.pyplot.scatter - - Raises: - ValueError: If the quantity is unknown - """ - - # Define possible quantities, and link to arrays within tokamak_source - # If given a non-TokamakSource, this step will fail with an AttributeError - try: - quantity_to_attribute = { - "ion_temperature": source.temperatures, - "neutron_source_density": source.neutron_source_density, - "strength": source.strengths, - } - except AttributeError as e: - raise ValueError( - f"openmc_plasma_source.scatter_tokamak_source: argument 'source' " - f"must be of type TokamakSource" - ) from e - - # For a given quantity, determine colours to plot for each point - # If given an incorrect quantity name, this step will fail with a KeyError - colours = None - if quantity is not None: - try: - colours = quantity_to_attribute[quantity] - except KeyError as e: - raise ValueError( - f"openmc_plasma_source.scatter_tokamak_source: Unknown 'quantity' " - f"provided, options are {quantity_to_attribute.keys()}" - ) from e - - # If not provided with an Axes instance, retrieve the current Axes in focus - if ax is None: - ax = plt.gca() - - # Scatter the source R and Z positions, optionally colouring using the chosen - # quantity. - ax.scatter(source.RZ[0], source.RZ[1], c=colours, **kwargs) - - # Set the aspect ratio on the axes. - # Defaults to 'equal', so 1m on the x-axis has the same width as 1m on the y-axis - ax.set_aspect(aspect) - - return ax - - -def plot_tokamak_source_3D( - tokamak_source, - ax=None, - quantity=None, - angles=[0, 0.5 * np.pi], - colorbar=None, - **kwargs, -): - """Creates a 3D plot of the tokamak source. - See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.plot.html#matplotlib.pyplot.plot - for more arguments. - - Args: - tokamak_source (ops.TokamakSource): the plasma source - ax (maplotlib.Axes, optional): Axes object on which to plot. If not - provided by the user, a new Axes is created by calling - matplotlib.pyplot.axes(projection="3d"). - quantity ("str", optional): value by which the lines should be - coloured. Defaults to None. - angles (list, optional): iterable of two floats defining the coverage. - Defaults to [0, 1/2*np.pi]. - colorbar (str, optional): colorbar used if quantity is not None. - When None, matplotlib currently defaults to "viridis". - - Raises: - ValueError: If the quantity is unknown - """ - - # Define possible quantities, and link to arrays within tokamak_source - # If given a non-TokamakSource, this step will fail with an AttributeError - try: - quantity_to_attribute = { - "ion_temperature": tokamak_source.temperatures, - "neutron_source_density": tokamak_source.neutron_source_density, - "strength": tokamak_source.strengths, - } - except AttributeError as e: - raise ValueError( - f"openmc_plasma_source.plot_tokamak_source_3D: argument 'source' " - f"must be of type TokamakSource" - ) from e - - # For a given quantity, get the corresponding array from tokamak_source - # If given an incorrect quantity name, this step will fail with a KeyError - if quantity is not None: - try: - quantity_values = quantity_to_attribute[quantity] - except KeyError as e: - raise ValueError( - f"openmc_plasma_source.plot_tokamak_source_3D: Unknown 'quantity' " - f"provided, options are {quantity_to_attribute.keys()}" - ) from e - else: - quantity_values = np.ones(tokamak_source.sample_size) - - # Get the colours used to plot each curve - # If 'quantity' == None, all have the same colour, selected from the middle - # of the colormap. - cmap = cm.get_cmap(colorbar) - if quantity is not None: - colors = cmap(quantity_values / np.max(quantity_values)) - else: - colors = cmap(np.full(tokamak_source.sample_size, 0.5)) - - # If not provided with an Axes object, create a new one - if ax is None: - ax = plt.axes(projection="3d") - - # Get curves to plot - n_theta = 100 - theta = np.linspace(*angles, n_theta) - xs = np.outer(tokamak_source.RZ[0], np.sin(theta)) - ys = np.outer(tokamak_source.RZ[0], np.cos(theta)) - zs = tokamak_source.RZ[1] - - # Plot each curve in turn - for x, y, z, color in zip(xs, ys, zs, colors): - ax.plot(x, y, z, color=color, **kwargs) - - # Set plot bounds - major_radius = tokamak_source.major_radius - ax.set_xlim(-major_radius, major_radius) - ax.set_ylim(-major_radius, major_radius) - ax.set_zlim(-major_radius, major_radius) - - return ax diff --git a/src/openmc_plasma_source/point_source.py b/src/openmc_plasma_source/point_source.py index e123d1c..0ba064b 100644 --- a/src/openmc_plasma_source/point_source.py +++ b/src/openmc_plasma_source/point_source.py @@ -1,79 +1,51 @@ -import openmc from typing import Tuple -from .fuel_types import fuel_types +import openmc + +from .fuel_types import get_neutron_energy_distribution + +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]: + """Creates a list of openmc.IndependentSource objects representing an ICF source. -class FusionPointSource(openmc.IndependentSource): - """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. + 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. Each component is measured in metres. temperature (float): Temperature of the source (eV). - fuel_type (str): The fusion fuel mix. Either 'DT' or 'DD'. + fuel (dict): Isotopes as keys and atom fractions as values """ - def __init__( - self, - coordinate: Tuple[float, float, float] = (0.0, 0.0, 0.0), - temperature: float = 20000.0, - fuel: str = "DT", + if ( + isinstance(coordinate, tuple) + and len(coordinate) == 3 + and all(isinstance(x, (int, float)) for x in coordinate) ): - # Set local attributes - self.coordinate = coordinate - self.temperature = temperature - self.fuel_type = fuel - self.fuel = fuel_types[self.fuel_type] - - # Call init for openmc.Source - super().__init__() - - # performed after the super init as these are Source attributes - self.space = openmc.stats.Point(self.coordinate) - self.angle = openmc.stats.Isotropic() - self.energy = openmc.stats.muir( - e0=self.fuel.mean_energy, - m_rat=self.fuel.mass_of_reactants, - kt=self.temperature, - ) - - @property - def coordinate(self): - return self._coordinate + pass + else: + raise ValueError("coordinate must be a tuple of three floats.") - @coordinate.setter - def coordinate(self, value): - if ( - isinstance(value, tuple) - and len(value) == 3 - and all(isinstance(x, (int, float)) for x in value) - ): - self._coordinate = value - 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.") - @property - def temperature(self): - return self._temperature + sources = [] - @temperature.setter - def temperature(self, value): - if isinstance(value, (int, float)) and value > 0: - self._temperature = value - else: - raise ValueError("Temperature must be strictly positive float.") + energy_distribution = get_neutron_energy_distribution( + ion_temperature=temperature, fuel=fuel + ) - @property - def fuel_type(self): - return self._fuel_type + source = openmc.IndependentSource() + source.energy = energy_distribution + source.space = openmc.stats.Point(coordinate) + source.angle = openmc.stats.Isotropic() + sources.append(source) - @fuel_type.setter - def fuel_type(self, value): - if value in fuel_types: - self._fuel_type = value - else: - raise KeyError("Invalid fuel type") + return sources diff --git a/src/openmc_plasma_source/ring_source.py b/src/openmc_plasma_source/ring_source.py index 9a4c501..1bb8726 100644 --- a/src/openmc_plasma_source/ring_source.py +++ b/src/openmc_plasma_source/ring_source.py @@ -1,117 +1,76 @@ -import openmc -import numpy as np from typing import Tuple -from .fuel_types import fuel_types +import numpy as np +import openmc + +from .fuel_types import get_neutron_energy_distribution + +def fusion_ring_source( + radius: float, + 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]: + """Creates a list of openmc.IndependentSource objects in a ring shape. -class FusionRingSource(openmc.IndependentSource): - """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 Muir - energy distribution. + 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 radians z_placement (float): Location of the ring source (m). Defaults to 0. temperature (float): the temperature to use in the Muir distribution in eV, - fuel_type (str): The fusion fuel mix. Either 'DT' or 'DD'. + fuel (dict): Isotopes as keys and atom fractions as values """ - def __init__( - self, - radius: float, - angles: Tuple[float, float] = (0, 2 * np.pi), - z_placement: float = 0, - temperature: float = 20000.0, - fuel: str = "DT", - ): - # Set local attributes - self.radius = radius - self.angles = angles - self.z_placement = z_placement - self.temperature = temperature - self.fuel_type = fuel - self.fuel = fuel_types[self.fuel_type] - - # Call init for openmc.Source - super().__init__() - - # performed after the super init as these are Source attributes - self.space = openmc.stats.CylindricalIndependent( - r=openmc.stats.Discrete([self.radius], [1]), - phi=openmc.stats.Uniform(a=self.angles[0], b=self.angles[1]), - z=openmc.stats.Discrete([self.z_placement], [1]), - origin=(0.0, 0.0, 0.0), - ) - self.angle = openmc.stats.Isotropic() - self.energy = openmc.stats.muir( - e0=self.fuel.mean_energy, - m_rat=self.fuel.mass_of_reactants, - kt=self.temperature, + if isinstance(radius, (int, float)) and radius > 0: + pass + else: + raise ValueError("Radius must be a float strictly greater than 0.") + + if ( + 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: + raise TypeError("Z placement must be a float.") + + if isinstance(temperature, (int, float)) and temperature > 0: + pass + else: + 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.space = openmc.stats.CylindricalIndependent( + r=openmc.stats.Discrete([radius], [1]), + phi=openmc.stats.Uniform(a=angles[0], b=angles[1]), + z=openmc.stats.Discrete([z_placement], [1]), + origin=(0.0, 0.0, 0.0), + ) + + source.energy = energy_distributions + source.angle = openmc.stats.Isotropic() + sources.append(source) - @property - def radius(self): - return self._radius - - @radius.setter - def radius(self, value): - if isinstance(value, (int, float)) and value > 0: - self._radius = value - else: - raise ValueError("Radius must be a float strictly greater than 0.") - - @property - def angles(self): - return self._angles - - @angles.setter - def angles(self, value): - if ( - isinstance(value, tuple) - and len(value) == 2 - and all( - isinstance(angle, (int, float)) and -2 * np.pi <= angle <= 2 * np.pi - for angle in value - ) - ): - self._angles = value - else: - raise ValueError( - "Angles must be a tuple of floats between zero and 2 * np.pi" - ) - - @property - def z_placement(self): - return self._z_placement - - @z_placement.setter - def z_placement(self, value): - if isinstance(value, (int, float)): - self._z_placement = value - else: - raise TypeError("Z placement must be a float.") - - @property - def temperature(self): - return self._temperature - - @temperature.setter - def temperature(self, value): - if isinstance(value, (int, float)) and value > 0: - self._temperature = value - else: - raise ValueError("Temperature must be a float strictly greater than 0.") - - @property - def fuel_type(self): - return self._fuel_type - - @fuel_type.setter - def fuel_type(self, value): - if value in fuel_types.keys(): - self._fuel_type = value - else: - raise KeyError("Invalid fuel type.") + return sources diff --git a/src/openmc_plasma_source/tokamak_source.py b/src/openmc_plasma_source/tokamak_source.py index 904d5c0..b7ff6c9 100644 --- a/src/openmc_plasma_source/tokamak_source.py +++ b/src/openmc_plasma_source/tokamak_source.py @@ -1,22 +1,48 @@ -import openmc -import numpy as np from typing import Tuple - -class TokamakSource: - """Plasma neutron source sampling. - This class greatly relies on models described in [1] +import numpy as np +import openmc +import openmc.checkvalue as cv +import NeSST as nst +from .fuel_types import get_neutron_energy_distribution, get_reactions_from_fuel +from NeSST.spectral_model import reac_DD, reac_TT, reac_DT + + +def tokamak_source( + major_radius: float, + minor_radius: float, + elongation: float, + triangularity: float, + mode: str, + ion_density_centre: float, + ion_density_peaking_factor: float, + ion_density_pedestal: float, + ion_density_separatrix: float, + ion_temperature_centre: float, + ion_temperature_peaking_factor: float, + ion_temperature_beta: float, + ion_temperature_pedestal: float, + ion_temperature_separatrix: float, + pedestal_radius: float, + shafranov_factor: float, + angles: Tuple[float, float] = (0, 2 * np.pi), + sample_size: int = 1000, + fuel: dict = {"D": 0.5, "T": 0.5}, + sample_seed: int = 122807528840384100672342137672332424406, +) -> 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) @@ -30,387 +56,349 @@ class TokamakSource: ion_density_pedestal (float): Ion density at pedestal (m-3) ion_density_separatrix (float): Ion density at separatrix (m-3) ion_temperature_centre (float): Ion temperature at the plasma - centre (keV) + centre (eV) ion_temperature_peaking_factor (float): Ion temperature peaking factor (referred in [1] as ion temperature exponent alpha_T) ion_temperature_beta (float): Ion temperature beta exponent (referred in [1] as ion temperature exponent beta_T) - ion_temperature_pedestal (float): Ion temperature at pedestal (keV) + ion_temperature_pedestal (float): Ion temperature at pedestal (eV) ion_temperature_separatrix (float): Ion temperature at separatrix - (keV) + (eV) pedestal_radius (float): Minor radius at pedestal (cm) shafranov_factor (float): Shafranov factor (referred in [1] as esh) also known as outward radial displacement of magnetic surfaces (cm) angles (iterable of floats): the start and stop angles of the ring in radians - sample_size int: number of neutron sources. Defaults to 1000. sample_seed int: the seed passed to numpy.random when sampling source location. Numpy recommend a large int value. Defaults to 122807528840384100672342137672332424406 + sample_size (int, optional): number of neutron sources. Defaults + to 1000. + fuel (dict): Isotopes as keys and atom fractions as values """ - def __init__( - self, - major_radius: float, - minor_radius: float, - elongation: float, - triangularity: float, - mode: str, - ion_density_centre: float, - ion_density_peaking_factor: float, - ion_density_pedestal: float, - ion_density_separatrix: float, - ion_temperature_centre: float, - ion_temperature_peaking_factor: float, - ion_temperature_beta: float, - ion_temperature_pedestal: float, - ion_temperature_separatrix: float, - pedestal_radius: float, - shafranov_factor: float, - angles: Tuple[float, float] = (0, 2 * np.pi), - sample_size: int = 1000, - sample_seed: int = 122807528840384100672342137672332424406, - ) -> None: - # Assign attributes - self.major_radius = major_radius - self.minor_radius = minor_radius - self.elongation = elongation - self.triangularity = triangularity - self.ion_density_centre = ion_density_centre - self.ion_density_peaking_factor = ion_density_peaking_factor - self.mode = mode - self.pedestal_radius = pedestal_radius - self.ion_density_pedestal = ion_density_pedestal - self.ion_density_separatrix = ion_density_separatrix - self.ion_temperature_centre = ion_temperature_centre - self.ion_temperature_peaking_factor = ion_temperature_peaking_factor - self.ion_temperature_pedestal = ion_temperature_pedestal - self.ion_temperature_separatrix = ion_temperature_separatrix - self.ion_temperature_beta = ion_temperature_beta - self.shafranov_factor = shafranov_factor - self.angles = angles - self.sample_size = sample_size - self.sample_seed = sample_seed - - # Perform sanity checks for inputs not caught by properties - if self.minor_radius >= self.major_radius: - raise ValueError("Minor radius must be smaller than major radius") - - if self.pedestal_radius >= self.minor_radius: - raise ValueError("Pedestal radius must be smaller than minor radius") - - if abs(self.shafranov_factor) >= 0.5 * minor_radius: - raise ValueError("Shafranov factor must be smaller than 0.5*minor radius") - - # Create a list of souces - self.sample_sources() - self.sources = self.make_openmc_sources() - - @property - def major_radius(self): - return self._major_radius - - @major_radius.setter - def major_radius(self, value): - if isinstance(value, (int, float)) and value > 0: - self._major_radius = value - else: - raise ValueError( - "Major radius must be a number within the specified bounds" - ) - - @property - def minor_radius(self): - return self._minor_radius - - @minor_radius.setter - def minor_radius(self, value): - if isinstance(value, (int, float)) and value > 0: - self._minor_radius = value - else: - raise ValueError( - "Minor radius must be a number within the specified bounds" - ) + # Perform sanity checks for inputs not caught by properties + 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.0, True) + cv.check_greater_than("triangularity", triangularity, -1.0, 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) + 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") + + # Create a list of sources + """Samples sample_size neutrons and creates attributes .densities + (ion density), .temperatures (ion temperature), .strengths + (neutron source density) and .RZ (coordinates) + """ + # create a sample of (a, alpha) coordinates + rng = np.random.default_rng(sample_seed) + a = rng.random(sample_size) * minor_radius + alpha = rng.random(sample_size) * 2 * np.pi + + # compute densities, temperatures + densities = tokamak_ion_density( + mode=mode, + ion_density_centre=ion_density_centre, + ion_density_peaking_factor=ion_density_peaking_factor, + ion_density_pedestal=ion_density_pedestal, + major_radius=major_radius, + pedestal_radius=pedestal_radius, + ion_density_separatrix=ion_density_separatrix, + r=a, + ) - @property - def elongation(self): - return self._elongation - - @elongation.setter - def elongation(self, value): - if isinstance(value, (int, float)) and value > 0: - self._elongation = value - else: - raise ValueError("Elongation must be a number within the specified bounds") - - @property - def triangularity(self): - return self._triangularity - - @triangularity.setter - def triangularity(self, value): - if isinstance(value, (int, float)) and -1.0 <= value <= 1.0: - self._triangularity = value - else: - raise ValueError( - "Triangularity must be a number within the specified bounds" - ) + # compute temperatures + temperatures = tokamak_ion_temperature( + r=a, + mode=mode, + pedestal_radius=pedestal_radius, + ion_temperature_pedestal=ion_temperature_pedestal / 1e3, + ion_temperature_centre=ion_temperature_centre / 1e3, + ion_temperature_beta=ion_temperature_beta, + ion_temperature_peaking_factor=ion_temperature_peaking_factor, + ion_temperature_separatrix=ion_temperature_separatrix / 1e3, + major_radius=major_radius, + ) - @property - def mode(self): - return self._mode - - @mode.setter - def mode(self, value): - if value in ["H", "L", "A"]: - self._mode = value - else: - raise ValueError("Mode must be one of the following: ['H', 'L', 'A']") - - @property - def ion_density_centre(self): - return self._ion_density_centre - - @ion_density_centre.setter - def ion_density_centre(self, value): - if isinstance(value, (int, float)) and value > 0: - self._ion_density_centre = value - else: - raise ValueError("Ion density centre must be a number greater than 0") - - @property - def ion_density_peaking_factor(self): - return self._ion_density_peaking_factor - - @ion_density_peaking_factor.setter - def ion_density_peaking_factor(self, value): - if isinstance(value, (int, float)): - self._ion_density_peaking_factor = value - else: - raise ValueError("Ion density peaking factor must be a number") - - @property - def ion_density_pedestal(self): - return self._ion_density_pedestal - - @ion_density_pedestal.setter - def ion_density_pedestal(self, value): - if isinstance(value, (int, float)) and value > 0: - self._ion_density_pedestal = value - else: - raise ValueError("Ion density pedestal must be a number greater than 0") - - @property - def ion_density_separatrix(self): - return self._ion_density_separatrix - - @ion_density_separatrix.setter - def ion_density_separatrix(self, value): - if isinstance(value, (int, float)) and value > 0: - self._ion_density_separatrix = value - else: - raise ValueError("Ion density separatrix must be a number greater than 0") - - @property - def angles(self): - return self._angles - - @angles.setter - def angles(self, value): - if ( - isinstance(value, tuple) - and len(value) == 2 - and all( - isinstance(angle, (int, float)) and -2 * np.pi <= angle <= 2 * np.pi - for angle in value - ) - ): - self._angles = value - else: - raise ValueError( - "Angles must be a tuple of floats between zero and 2 * np.pi" - ) + # convert coordinates + RZ = tokamak_convert_a_alpha_to_R_Z( + a=a, + alpha=alpha, + shafranov_factor=shafranov_factor, + minor_radius=minor_radius, + major_radius=major_radius, + triangularity=triangularity, + elongation=elongation, + ) - # TODO setters and getters for the rest + fuel_densities = {} + for key, value in fuel.items(): + fuel_densities[key] = densities * value + reactions = get_reactions_from_fuel(fuel) - def _bounds_check(value, bounds): - return bounds[0] < value + neutron_source_density = {} + total_source_density = 0 + for reaction in reactions: - def ion_density(self, r): - """Computes the ion density at a given position. The ion density is - only dependent on the minor radius. + if reaction == "DD": + fuel_density = fuel_densities["D"] * 0.5 + elif reaction == "TT": + fuel_density = fuel_densities["T"] * 0.5 + elif reaction == "DT": + fuel_density = fuel_densities["T"] * fuel_densities["D"] - Args: - r (float, ndarray): the minor radius (cm) + neutron_source_density[reaction] = tokamak_neutron_source_density( + fuel_density, temperatures, reaction + ) + if reaction == "TT": + # TT reaction emits 2 neutrons + neutron_source_density[reaction] = neutron_source_density[reaction] * 2 + + total_source_density += sum(neutron_source_density[reaction]) + + all_sources = [] + for reaction in reactions: + strengths = neutron_source_density[reaction] / total_source_density + + sources = tokamak_make_openmc_sources( + strengths=strengths, + angles=angles, + temperatures=temperatures, + fuel=fuel, + RZ=RZ, + ) + all_sources = all_sources + sources + return all_sources + + +def tokamak_ion_density( + mode, + ion_density_centre, + ion_density_peaking_factor, + ion_density_pedestal, + major_radius, + pedestal_radius, + ion_density_separatrix, + r, +): + """Computes the ion density at a given position. The ion density is + only dependent on the minor radius. - Returns: - float, ndarray: ion density in m-3 - """ + Args: + r (float, ndarray): the minor radius (cm) - r = np.asarray(r) - if np.any(r < 0): - raise ValueError("Minor radius must not be negative") + Returns: + float, ndarray: ion density in m-3 + """ - if self.mode == "L": - density = ( - self.ion_density_centre - * (1 - (r / self.major_radius) ** 2) ** self.ion_density_peaking_factor - ) - elif self.mode in ["H", "A"]: - density = np.where( - r < self.pedestal_radius, - ( - (self.ion_density_centre - self.ion_density_pedestal) - * (1 - (r / self.pedestal_radius) ** 2) - ** self.ion_density_peaking_factor - + self.ion_density_pedestal - ), - ( - (self.ion_density_pedestal - self.ion_density_separatrix) - * (self.major_radius - r) - / (self.major_radius - self.pedestal_radius) - + self.ion_density_separatrix - ), - ) - return density + r = np.asarray(r) + if np.any(r < 0): + raise ValueError("Minor radius must not be negative") - def ion_temperature(self, r): - """Computes the ion temperature at a given position. The ion - temperature is only dependent on the minor radius. + if mode == "L": + density = ( + ion_density_centre + * (1 - (r / major_radius) ** 2) ** ion_density_peaking_factor + ) + elif mode in ["H", "A"]: + density = np.where( + r < pedestal_radius, + ( + (ion_density_centre - ion_density_pedestal) + * (1 - (r / pedestal_radius) ** 2) ** ion_density_peaking_factor + + ion_density_pedestal + ), + ( + (ion_density_pedestal - ion_density_separatrix) + * (major_radius - r) + / (major_radius - pedestal_radius) + + ion_density_separatrix + ), + ) + return density + + +def tokamak_ion_temperature( + r, + mode, + pedestal_radius, + ion_temperature_pedestal, + ion_temperature_centre, + ion_temperature_beta, + ion_temperature_peaking_factor, + ion_temperature_separatrix, + major_radius, +): + """Computes the ion temperature at a given position. The ion + temperature is only dependent on the minor radius. - Args: - r (float, ndarray): minor radius (cm) + Args: + r (float, ndarray): minor radius (cm) - Returns: - float, ndarray: ion temperature (keV) - """ + Returns: + float, ndarray: ion temperature (eV) + """ - r = np.asarray(r) - if np.any(r < 0): - raise ValueError("Minor radius must not be negative") + r = np.asarray(r) + if np.any(r < 0): + raise ValueError("Minor radius must not be negative") - if self.mode == "L": - temperature = ( - self.ion_temperature_centre - * (1 - (r / self.major_radius) ** 2) - ** self.ion_temperature_peaking_factor - ) - elif self.mode in ["H", "A"]: - temperature = np.where( - r < self.pedestal_radius, - ( - self.ion_temperature_pedestal - + (self.ion_temperature_centre - self.ion_temperature_pedestal) - * (1 - (r / self.pedestal_radius) ** self.ion_temperature_beta) - ** self.ion_temperature_peaking_factor - ), - ( - self.ion_temperature_separatrix - + (self.ion_temperature_pedestal - self.ion_temperature_separatrix) - * (self.major_radius - r) - / (self.major_radius - self.pedestal_radius) - ), - ) - return temperature - - def convert_a_alpha_to_R_Z(self, a, alpha): - """Converts (r, alpha) cylindrical coordinates to (R, Z) cartesian - coordinates. - - Args: - a (float, ndarray): minor radius (cm) - alpha (float, ndarray): angle (rad) - - Returns: - ((float, ndarray), (float, ndarray)): (R, Z) coordinates - """ - a = np.asarray(a) - alpha = np.asarray(alpha) - if np.any(a < 0): - raise ValueError("Radius 'a' must not be negative") - - shafranov_shift = self.shafranov_factor * (1.0 - (a / self.minor_radius) ** 2) - R = ( - self.major_radius - + a * np.cos(alpha + (self.triangularity * np.sin(alpha))) - + shafranov_shift + if mode == "L": + temperature = ( + ion_temperature_centre + * (1 - (r / major_radius) ** 2) ** ion_temperature_peaking_factor ) - Z = self.elongation * a * np.sin(alpha) - return (R, Z) - - def sample_sources(self): - """Samples self.sample_size neutrons and creates attributes .densities - (ion density), .temperatures (ion temperature), .strengths - (neutron source density) and .RZ (coordinates) - """ - # create a sample of (a, alpha) coordinates - rng = np.random.default_rng(self.sample_seed) - a = rng.random(self.sample_size) * self.minor_radius - alpha = rng.random(self.sample_size) * 2 * np.pi - - # compute densities, temperatures, neutron source densities and - # convert coordinates - self.densities = self.ion_density(a) - self.temperatures = self.ion_temperature(a) - self.neutron_source_density = neutron_source_density( - self.densities, self.temperatures + elif mode in ["H", "A"]: + temperature = np.where( + r < pedestal_radius, + ( + ion_temperature_pedestal + + (ion_temperature_centre - ion_temperature_pedestal) + * (1 - (np.abs(r / pedestal_radius)) ** ion_temperature_beta) + ** ion_temperature_peaking_factor + ), + ( + ion_temperature_separatrix + + (ion_temperature_pedestal - ion_temperature_separatrix) + * (major_radius - r) + / (major_radius - pedestal_radius) + ), ) - self.strengths = self.neutron_source_density / sum(self.neutron_source_density) - self.RZ = self.convert_a_alpha_to_R_Z(a, alpha) - - def make_openmc_sources(self): - """Creates a list of OpenMC Sources() objects. The created sources are - ring sources based on the .RZ coordinates between two angles. The - energy of the sources are Muir energy spectra with ion temperatures - based on .temperatures. The strength of the sources (their probability) - is based on .strengths. - - Args: - angles ((float, float), optional): rotation of the ring source. + return temperature * 1e3 + + +def tokamak_convert_a_alpha_to_R_Z( + a, + alpha, + shafranov_factor, + minor_radius, + major_radius, + triangularity, + elongation, +): + """Converts (r, alpha) cylindrical coordinates to (R, Z) cartesian + coordinates. + + Args: + a (float, ndarray): minor radius (cm) + alpha (float, ndarray): angle (rad) + shafranov_factor: + minor_radius: + major_radius: + + Returns: + ((float, ndarray), (float, ndarray)): (R, Z) coordinates + """ + a = np.asarray(a) + alpha = np.asarray(alpha) + if np.any(a < 0): + raise ValueError("Radius 'a' must not be negative") + + shafranov_shift = shafranov_factor * (1.0 - (a / minor_radius) ** 2) + R = ( + major_radius + + a * np.cos(alpha + (triangularity * np.sin(alpha))) + + shafranov_shift + ) + Z = elongation * a * np.sin(alpha) + return (R, Z) + + +def tokamak_make_openmc_sources( + strengths, + angles, + temperatures, + fuel, + RZ, +): + """Creates a list of OpenMC Sources() objects. The created sources are + ring sources based on the .RZ coordinates between two angles. The + energy of the sources are Muir energy spectra with ion temperatures + based on .temperatures. The strength of the sources (their probability) + is based on .strengths. + + Args: + strengths + angles ((float, float), optional): rotation of the ring source. Defaults to (0, 2*np.pi). + temperatures + fuel + RZ + + Returns: + list: list of openmc.IndependentSource() + """ - Returns: - list: list of openmc.IndependentSource() - """ + sources = [] + # create a ring source for each sample in the plasma source + R_vals = RZ[0] + Z_vals = RZ[1] + assert len(Z_vals) == len(R_vals) == len(temperatures) == len(strengths) + for R_val, Z_val, temperature, strength in zip( + R_vals, Z_vals, temperatures, strengths + ): - sources = [] - # create a ring source for each sample in the plasma source - for i in range(len(self.strengths)): - my_source = openmc.IndependentSource() + if strength > 0.0: + radius = openmc.stats.Discrete([R_val], [1]) + z_values = openmc.stats.Discrete([Z_val], [1]) + angle = openmc.stats.Uniform(a=angles[0], b=angles[1]) - # extract the RZ values accordingly - radius = openmc.stats.Discrete([self.RZ[0][i]], [1]) - z_values = openmc.stats.Discrete([self.RZ[1][i]], [1]) - angle = openmc.stats.Uniform(a=self.angles[0], b=self.angles[1]) + my_source = openmc.IndependentSource() + my_source.energy = get_neutron_energy_distribution( + ion_temperature=temperature, + fuel=fuel, + ) # create a ring source my_source.space = openmc.stats.CylindricalIndependent( r=radius, phi=angle, z=z_values, origin=(0.0, 0.0, 0.0) ) - my_source.angle = openmc.stats.Isotropic() - my_source.energy = openmc.stats.muir( - e0=14080000.0, m_rat=5.0, kt=self.temperatures[i] - ) - # the strength of the source (its probability) is given by - # self.strengths - my_source.strength = self.strengths[i] + my_source.strength = strength # append to the list of sources sources.append(my_source) - return sources + return sources -def neutron_source_density(ion_density, ion_temperature): +def tokamak_neutron_source_density(ion_density, ion_temperature, reaction): """Computes the neutron source density given ion density and ion temperature. Args: ion_density (float, ndarray): Ion density (m-3) - ion_temperature (float, ndarray): Ion temperature (keV) - + ion_temperature (float, ndarray): Ion temperature (eV) + reaction (str): The fusion reactions to consider e.g. 'DD' Returns: float, ndarray: Neutron source density (neutron/s/m3) """ @@ -418,39 +406,13 @@ def neutron_source_density(ion_density, ion_temperature): ion_density = np.asarray(ion_density) ion_temperature = np.asarray(ion_temperature) - return ion_density**2 * DT_xs(ion_temperature) - - -def DT_xs(ion_temperature): - """Sadler–Van Belle formula - Ref : https://doi.org/10.1016/j.fusengdes.2012.02.025 - - Args: - ion_temperature (float, ndarray): ion temperature in keV - - Returns: - float, ndarray: the DT cross section at the given temperature - """ - - ion_temperature = np.asarray(ion_temperature) - - c = [ - 2.5663271e-18, - 19.983026, - 2.5077133e-2, - 2.5773408e-3, - 6.1880463e-5, - 6.6024089e-2, - 8.1215505e-3, - ] - - U = 1 - ion_temperature * ( - c[2] + ion_temperature * (c[3] - c[4] * ion_temperature) - ) / (1.0 + ion_temperature * (c[5] + c[6] * ion_temperature)) - - val = ( - c[0] - * np.exp(-c[1] * (U / ion_temperature) ** (1 / 3)) - / (U ** (5 / 6) * ion_temperature ** (2 / 3)) - ) - return val + if reaction == "DD": + return ion_density * reac_DD(ion_temperature) + elif reaction == "TT": + return ion_density * reac_TT(ion_temperature) + elif reaction == "DT": + return ion_density * reac_DT(ion_temperature) # could use _DT_xs instead + else: + raise ValueError( + 'Reaction {reaction} not in available options ["DD", "DT", "TT"]' + ) diff --git a/tests/test_fuel_types.py b/tests/test_fuel_types.py index 2d68717..b0a2c8a 100644 --- a/tests/test_fuel_types.py +++ b/tests/test_fuel_types.py @@ -1,32 +1,47 @@ -from openmc_plasma_source.fuel_types import Fuel, fuel_types import pytest +from openmc_plasma_source import get_neutron_energy_distribution -@pytest.mark.parametrize("energy,mass", [(2.5e7, 5), (15, 30)]) -def test_fuel_with_correct_inputs(energy, mass): + +@pytest.mark.parametrize( + "temperature, fuel", + [ + (2e3, {"D": 1.0}), + (2e3, {"T": 1.0}), + (2e3, {"T": 0.5, "D": 0.5}), + (2e3, {"T": 0.2, "D": 0.8}), + ], +) +def test_fuel_with_correct_inputs(temperature, fuel): # Should accept any non-zero positive inputs to either variable - fuel = Fuel(energy, mass) - assert fuel.mean_energy == energy - assert fuel.mass_of_reactants == mass + get_neutron_energy_distribution(temperature, fuel) @pytest.mark.parametrize( - "energy,mass", [(2.5e7, -5), (-12, 30), (1e7, 0), (0, 4), (-12, -12)] + "temperature, fuel", + [ + (2e3, {"D": 1.1}), + (2e3, {"T": 0.9}), + (2e3, {"T": -0.5, "D": 0.5}), + (2e3, {"T": -0.2, "D": -0.8}), + ], ) -def test_fuel_with_bad_inputs(energy, mass): +def test_fuel_with_bad_inputs(temperature, fuel): # Should reject any negative numbers and zeros. with pytest.raises(ValueError): - fuel = Fuel(energy, mass) - + get_neutron_energy_distribution(temperature, fuel) -@pytest.mark.parametrize("fuel_type", ["DT", "DD"]) -def test_fuel_types(fuel_type): - # Should accept 'DD' and 'DT' - assert isinstance(fuel_types[fuel_type], Fuel) - -@pytest.mark.parametrize("fuel_type", ["dt", "dd", "Dt", "dD", "hello world", 5]) -def test_incorrect_fuel_types(fuel_type): - # Should reject everything except 'DT' and 'DD' - with pytest.raises(KeyError): - my_fuel = fuel_types[fuel_type] +@pytest.mark.parametrize( + "temperature, fuel", + [ + (2e3, {"DD": 1.1}), + (2e3, {"DT": 0.9}), + (2e3, {"He3": -0.5, "D": 0.5}), + (2e3, {1: -0.2, "D": -0.8}), + ], +) +def test_fuel_with_incorrect_isotopese(temperature, fuel): + # Should reject anything which is not 'D' or 'T'. + with pytest.raises(ValueError): + get_neutron_energy_distribution(temperature, fuel) diff --git a/tests/test_plotting.py b/tests/test_plotting.py deleted file mode 100644 index 7f1a5f7..0000000 --- a/tests/test_plotting.py +++ /dev/null @@ -1,262 +0,0 @@ -import matplotlib.pyplot as plt -import numpy as np -from openmc_plasma_source import ( - TokamakSource, - plotting as ops_plt, -) -import pytest - - -@pytest.fixture -def tokamak_source(): - return 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, - sample_size=500, - ) - - -def test_scatter_tokamak_source_defaults(tokamak_source): - """Ensure plotting is successful without providing additional args""" - plt.figure() - assert not plt.gca().collections # Check current ax is empty - ops_plt.scatter_tokamak_source(tokamak_source) - assert plt.gca().collections # Check current ax is not empty - # Save for viewing, clean up - plt.xlabel("R") - plt.ylabel("Z", rotation=0) - plt.savefig("tests/test_scatter_tokamak_source_defaults.png") - plt.close() - - -def test_scatter_tokamak_source_with_ax(tokamak_source): - """Ensure plotting is successful for user-provided ax""" - fig = plt.figure() - ax = fig.gca() - assert not ax.collections # Check ax is empty - ops_plt.scatter_tokamak_source(tokamak_source, ax=ax) - assert ax.collections # Check ax is not empty - # Save for viewing, clean up - ax.set_xlabel("R") - ax.set_ylabel("Z", rotation=0) - fig.savefig("tests/test_scatter_tokamak_source_with_ax.png") - plt.close(fig) - - -def test_scatter_tokamak_source_with_subplots(tokamak_source): - """Ensure plotting is successful for multiple user-provided ax""" - fig, (ax1, ax2) = plt.subplots(1, 2) - # Plot on the first axes - assert not ax1.collections # Check ax is empty - ops_plt.scatter_tokamak_source(tokamak_source, ax=ax1) - assert ax1.collections # Check ax is not empty - # Generate new data - tokamak_source.sample_sources() - tokamak_source.sources = tokamak_source.make_openmc_sources() - # Plot on the other axes - assert not ax2.collections # Check ax is empty - ops_plt.scatter_tokamak_source(tokamak_source, ax=ax2) - assert ax2.collections # Check ax is not empty - # Save for viewing, clean up - ax1.set_xlabel("R") - ax1.set_ylabel("Z", rotation=0) - ax2.set_xlabel("R") - fig.savefig("tests/test_scatter_tokamak_source_subplots.png") - plt.close(fig) - - -@pytest.mark.parametrize( - "quantity", ["ion_temperature", "neutron_source_density", "strength"] -) -def test_scatter_tokamak_source_quantities(tokamak_source, quantity): - """Plot with colours set by 'quantity'""" - fig = plt.figure() - ax = fig.gca() - assert not ax.collections # Check ax is empty - ops_plt.scatter_tokamak_source(tokamak_source, ax=ax, quantity=quantity) - assert ax.collections # Check ax is not empty - # Save for viewing, clean up - ax.set_xlabel("R") - ax.set_ylabel("Z", rotation=0) - fig.savefig(f"tests/test_scatter_tokamak_source_quantities_{quantity}.png") - plt.close(fig) - - -@pytest.mark.parametrize("aspect", ["equal", "auto", 2]) -def test_scatter_tokamak_source_aspect(tokamak_source, aspect): - """Plot with various aspect ratios""" - fig = plt.figure() - ax = fig.gca() - assert not ax.collections # Check ax is empty - ops_plt.scatter_tokamak_source(tokamak_source, ax=ax, aspect=aspect) - assert ax.collections # Check ax is not empty - # Save for viewing, clean up - ax.set_xlabel("R") - ax.set_ylabel("Z", rotation=0) - fig.savefig(f"tests/test_scatter_tokamak_source_aspect_{aspect}.png") - plt.close(fig) - - -@pytest.mark.parametrize("kwargs", [{"alpha": 0.2}, {"marker": "x"}]) -def test_scatter_tokamak_source_kwargs(tokamak_source, kwargs): - """Plot with a kwarg compatible with 'scatter'""" - fig = plt.figure() - ax = fig.gca() - assert not ax.collections # Check ax is empty - ops_plt.scatter_tokamak_source(tokamak_source, ax=ax, **kwargs) - assert ax.collections # Check ax is not empty - # Save for viewing, clean up - ax.set_xlabel("R") - ax.set_ylabel("Z", rotation=0) - fig.savefig( - f"tests/test_scatter_tokamak_source_kwargs_{list(kwargs.keys())[0]}.png" - ) - plt.close(fig) - - -def test_scatter_tokamak_not_source(): - """Ensure failure when given non-TokamakSource to plot""" - with pytest.raises(ValueError) as excinfo: - fig = plt.figure() - ax = fig.gca() - ops_plt.scatter_tokamak_source("hello world", ax=ax) - plt.close() - assert "TokamakSource" in str(excinfo.value) - - -@pytest.mark.parametrize("quantity", ["coucou", "ion_density", 17]) -def test_scatter_tokamak_wrong_quantity(tokamak_source, quantity): - """Ensure failure when incorrect quantity specified""" - with pytest.raises(ValueError) as excinfo: - fig = plt.figure() - ax = fig.gca() - ops_plt.scatter_tokamak_source(tokamak_source, ax=ax, quantity=quantity) - plt.close() - assert "quantity" in str(excinfo.value) - - -def test_plot_tokamak_source_3D_default(tokamak_source): - """Ensure plots correctly with default inputs""" - plt.figure() - ops_plt.plot_tokamak_source_3D(tokamak_source) - assert plt.gca().lines # Check current ax is not empty - # Save for viewing, clean up - plt.savefig("tests/test_plot_tokamak_source_3D_defaults.png") - plt.close() - - -def test_plot_tokamak_source_3D_with_ax(tokamak_source): - """Ensure plots correctly given ax instance""" - fig = plt.figure() - ax = fig.add_subplot(1, 1, 1, projection="3d") - assert not ax.lines # Check ax is empty - ops_plt.plot_tokamak_source_3D(tokamak_source, ax=ax) - assert ax.lines # Check ax is not empty - # Save for viewing, clean up - fig.savefig("tests/test_plot_tokamak_source_3D_with_ax.png") - plt.close(fig) - - -@pytest.mark.parametrize( - "quantity", ["ion_temperature", "neutron_source_density", "strength"] -) -def test_plot_tokamak_source_3D_quantities(tokamak_source, quantity): - """Ensure plots correctly for each quantity""" - fig = plt.figure() - ax = fig.add_subplot(1, 1, 1, projection="3d") - assert not ax.lines # Check ax is empty - ops_plt.plot_tokamak_source_3D(tokamak_source, ax=ax, quantity=quantity) - assert ax.lines # Check ax is not empty - # Save for viewing, clean up - fig.savefig(f"tests/test_plot_tokamak_source_3D_quantities_{quantity}.png") - plt.close(fig) - - -@pytest.mark.parametrize("colorbar", ["plasma", "rainbow"]) -def test_plot_tokamak_source_3D_colorbars(tokamak_source, colorbar): - """Ensure plots correctly given colorbar choice""" - fig = plt.figure() - ax = fig.add_subplot(1, 1, 1, projection="3d") - assert not ax.lines # Check ax is empty - ops_plt.plot_tokamak_source_3D( - tokamak_source, ax=ax, quantity="ion_temperature", colorbar=colorbar - ) - assert ax.lines # Check ax is not empty - # Save for viewing, clean up - fig.savefig(f"tests/test_plot_tokamak_source_3D_colorbar_{colorbar}.png") - plt.close(fig) - - -@pytest.mark.parametrize( - "angles,name", - [ - [(0, np.pi / 4), "eighth"], - [(0, np.pi), "half"], - [(np.pi, 2 * np.pi), "half_offset"], - [(0, 2 * np.pi), "full"], - ], -) -def test_plot_tokamak_source_3D_angles(tokamak_source, angles, name): - """Ensure plots correctly given angles range""" - fig = plt.figure() - ax = fig.add_subplot(1, 1, 1, projection="3d") - assert not ax.lines # Check ax is empty - ops_plt.plot_tokamak_source_3D( - tokamak_source, ax=ax, quantity="ion_temperature", angles=angles - ) - assert ax.lines # Check ax is not empty - # Save for viewing, clean up - fig.savefig(f"tests/test_plot_tokamak_source_3D_angles_{name}.png") - plt.close(fig) - - -@pytest.mark.parametrize("kwargs", [{"alpha": 0.2}, {"linestyle": "dotted"}]) -def test_plot_tokamak_source_3D_kwargs(tokamak_source, kwargs): - """Ensure plots correctly given additonal keyword arguments to pass on to plot""" - fig = plt.figure() - ax = fig.add_subplot(1, 1, 1, projection="3d") - assert not ax.lines # Check ax is empty - ops_plt.plot_tokamak_source_3D( - tokamak_source, ax=ax, quantity="ion_temperature", **kwargs - ) - assert ax.lines # Check ax is not empty - # Save for viewing, clean up - fig.savefig( - f"tests/test_plot_tokamak_source_3D_kwargs_{list(kwargs.keys())[0]}.png" - ) - plt.close(fig) - - -def test_plot_tokamak_source_3D_not_source(): - """Ensure failure when given non-TokamakSource to plot""" - with pytest.raises(ValueError) as excinfo: - fig = plt.figure() - ax = fig.add_subplot(1, 1, 1, projection="3d") - ops_plt.plot_tokamak_source_3D("hello world", ax=ax) - plt.close() - assert "TokamakSource" in str(excinfo.value) - - -@pytest.mark.parametrize("quantity", ["coucou", "ion_density", 17]) -def test_plot_tokamak_source_3D_wrong_quantity(tokamak_source, quantity): - """Ensure failure when incorrect quantity specified""" - with pytest.raises(ValueError) as excinfo: - fig = plt.figure() - ax = fig.add_subplot(1, 1, 1, projection="3d") - ops_plt.plot_tokamak_source_3D(tokamak_source, ax=ax, quantity=quantity) - plt.close() - assert "quantity" in str(excinfo.value) diff --git a/tests/test_point_source.py b/tests/test_point_source.py index 2e4ab13..74f0b8d 100644 --- a/tests/test_point_source.py +++ b/tests/test_point_source.py @@ -1,20 +1,25 @@ -from openmc_plasma_source import FusionPointSource - +import numpy as np import openmc import pytest -import numpy as np + +from openmc_plasma_source import fusion_point_source def test_creation(): - my_source = FusionPointSource() + my_source = fusion_point_source() # Ensure it is of type openmc.IndependentSource - assert isinstance(my_source, openmc.IndependentSource) + for source in my_source: + assert isinstance(source, openmc.IndependentSource) - # Ensure it has space, angle, and energy set - assert isinstance(my_source.space, openmc.stats.Point) - assert isinstance(my_source.angle, openmc.stats.Isotropic) - assert isinstance(my_source.energy, openmc.stats.univariate.Normal) + # Ensure it has space, angle, and energy set + assert isinstance(source.space, openmc.stats.Point) + assert isinstance(source.angle, openmc.stats.Isotropic) + assert ( + isinstance(source.energy, openmc.stats.univariate.Normal) + or isinstance(source.energy, openmc.stats.univariate.Tabular) + or isinstance(source.energy, openmc.stats.Mixture) + ) @pytest.mark.parametrize( @@ -22,8 +27,7 @@ def test_creation(): ) def test_coordinate(coordinate): # Should allow any tuple of length 3 containing numbers - my_source = FusionPointSource(coordinate=coordinate) - assert np.array_equal(my_source.coordinate, coordinate) + fusion_point_source(coordinate=coordinate) @pytest.mark.parametrize( @@ -33,32 +37,30 @@ def test_bad_coordinate(coordinate): # Should reject iterables of length != 3, anything non-tuple, and anything # that can't convert to float with pytest.raises(ValueError): - FusionPointSource(coordinate=coordinate) + fusion_point_source(coordinate=coordinate) @pytest.mark.parametrize("temperature", [20000, 1e4, 0.1, 25000.0]) def test_temperature(temperature): # Should accept any positive float - my_source = FusionPointSource(temperature=temperature) - assert my_source.temperature == temperature + fusion_point_source(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): - FusionPointSource(temperature=temperature) + fusion_point_source(temperature=temperature) -@pytest.mark.parametrize("fuel", ["DT", "DD"]) +@pytest.mark.parametrize("fuel", [{"D": 0.5, "T": 0.5}, {"D": 1.0}, {"T": 1.0}]) def test_fuel(fuel): # Should accept either 'DD' or 'DT' - my_source = FusionPointSource(fuel=fuel) - assert my_source.fuel_type == fuel + fusion_point_source(fuel=fuel) -@pytest.mark.parametrize("fuel", ["топливо", 5]) +@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((KeyError, TypeError)): - FusionPointSource(fuel=fuel) + with pytest.raises(ValueError): + fusion_point_source(fuel=fuel) diff --git a/tests/test_ring_source.py b/tests/test_ring_source.py index 35de338..57c8ba8 100644 --- a/tests/test_ring_source.py +++ b/tests/test_ring_source.py @@ -1,41 +1,44 @@ -from openmc_plasma_source import FusionRingSource - +import numpy as np import openmc import pytest -import numpy as np + +from openmc_plasma_source import fusion_ring_source def test_creation(): - my_source = FusionRingSource(radius=1.0, z_placement=1.0) + my_source = fusion_ring_source(radius=1.0, z_placement=1.0) # Ensure it is of type openmc.IndependentSource - assert isinstance(my_source, openmc.IndependentSource) + for source in my_source: + assert isinstance(source, openmc.IndependentSource) - # Ensure it has space, angle, and energy set - assert isinstance(my_source.space, openmc.stats.CylindricalIndependent) - assert isinstance(my_source.angle, openmc.stats.Isotropic) - assert isinstance(my_source.energy, openmc.stats.univariate.Normal) + # Ensure it has space, angle, and energy set + assert isinstance(source.space, openmc.stats.CylindricalIndependent) + assert isinstance(source.angle, openmc.stats.Isotropic) + assert ( + isinstance(source.energy, openmc.stats.univariate.Normal) + or isinstance(source.energy, openmc.stats.univariate.Tabular) + or isinstance(source.energy, openmc.stats.Mixture) + ) @pytest.mark.parametrize("radius", [1, 5.6, 1e5, 7.0]) def test_radius(radius): # should allow any positive float - my_source = FusionRingSource(radius=radius) - assert my_source.radius == radius + 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 with pytest.raises(ValueError): - my_source = FusionRingSource(radius=radius) + fusion_ring_source(radius=radius) @pytest.mark.parametrize("angles", [(1, 2), (0.0, np.pi), (np.pi, 0.0)]) def test_angles(angles): # Should allow any tuple of length 2 with contents convertible to float - my_source = FusionRingSource(radius=1.0, angles=angles) - assert np.array_equal(my_source.angles, angles) + fusion_ring_source(radius=1.0, angles=angles) @pytest.mark.parametrize("angles", [(1.0,), [1, 2], 5, "ab", ("a", "b")]) @@ -43,38 +46,36 @@ def test_bad_angles(angles): # Should reject iterables of length != 2, anything non tuple, and anything # that can't convert to float with pytest.raises(ValueError): - FusionRingSource(radius=1.0, angles=angles) + fusion_ring_source(radius=1.0, angles=angles) @pytest.mark.parametrize("temperature", [20000.0, 1e4, 0.1, 25000]) def test_temperature(temperature): # Should accept any positive float - my_source = FusionRingSource(radius=1.0, temperature=temperature) - assert my_source.temperature == temperature + 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): - FusionRingSource(radius=1.0, temperature=temperature) + fusion_ring_source(radius=1.0, temperature=temperature) -@pytest.mark.parametrize("fuel", ["DT", "DD"]) +@pytest.mark.parametrize("fuel", [{"D": 0.5, "T": 0.5}, {"D": 1.0}]) def test_fuel(fuel): # Should accept either 'DD' or 'DT' - my_source = FusionRingSource(radius=1.0, fuel=fuel) - assert my_source.fuel_type == fuel + fusion_ring_source(radius=1.0, fuel=fuel) -@pytest.mark.parametrize("fuel", ["топливо", 5]) +@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((KeyError, TypeError)): - FusionRingSource(radius=1.0, fuel=fuel) + with pytest.raises(ValueError): + fusion_ring_source(radius=1.0, fuel=fuel) @pytest.mark.parametrize("z", ["coucou", [5, 2.0]]) def test_wrong_z_placement(z): with pytest.raises((TypeError)): - FusionRingSource(radius=1.0, z_placement=z) + fusion_ring_source(radius=1.0, z_placement=z) diff --git a/tests/test_tokamak_source.py b/tests/test_tokamak_source.py index 27379aa..9c43724 100644 --- a/tests/test_tokamak_source.py +++ b/tests/test_tokamak_source.py @@ -1,14 +1,19 @@ -from openmc_plasma_source import TokamakSource -from openmc import IndependentSource import numpy as np - import pytest -from hypothesis import given, settings, assume, strategies as st +from hypothesis import given, settings, HealthCheck +from hypothesis import strategies as st +import openmc +from openmc_plasma_source import ( + tokamak_source, + tokamak_ion_density, + tokamak_ion_temperature, + tokamak_convert_a_alpha_to_R_Z, +) @pytest.fixture def tokamak_args_dict(): - """Returns a dict of realistic inputs for TokamakSource""" + """Returns a dict of realistic inputs for tokamak_source""" args_dict = { "elongation": 1.557, "triangularity": 0.270, @@ -32,51 +37,49 @@ def tokamak_args_dict(): @pytest.fixture def tokamak_source_example(tokamak_args_dict): - """Returns a TokamakSource with realistic inputs""" - return TokamakSource(**tokamak_args_dict) + """Returns a tokamak_source with realistic inputs""" + return tokamak_source(**tokamak_args_dict) def test_creation(tokamak_source_example): - """Tests that the sources generated by TokamakSource are of + """Tests that the sources generated by tokamak_source are of type openmc.Source""" - for source in tokamak_source_example.sources: - assert isinstance(source, IndependentSource) + for source in tokamak_source_example: + assert isinstance(source, openmc.IndependentSource) @pytest.mark.parametrize( "minor_radius,major_radius", [(3.0, 10.0), (3.0, 100), (3.0, 3.00001)] ) def test_major_radius(tokamak_args_dict, minor_radius, major_radius): - """Checks that TokamakSource creation accepts valid major radius""" + """Checks that tokamak_source creation accepts valid major radius""" tokamak_args_dict["minor_radius"] = minor_radius tokamak_args_dict["major_radius"] = major_radius - tokamak_source = TokamakSource(**tokamak_args_dict) - assert tokamak_source.major_radius == major_radius + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize( "minor_radius,major_radius", [(3, 3), (3, 1), (3, -5), (3, "hello world")] ) def test_bad_major_radius(tokamak_args_dict, minor_radius, major_radius): - """Checks that TokamakSource creation rejects invalid 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): - tokamak_source = TokamakSource(**tokamak_args_dict) + with pytest.raises((ValueError, TypeError)): + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize( "major_radius,minor_radius", [(10.0, 3.0), (10.0, 9.9), (10.0, 0.1)] ) def test_minor_radius(tokamak_args_dict, major_radius, minor_radius): - """Checks that TokamakSource creation accepts valid minor radius""" + """Checks that tokamak_source creation accepts valid minor radius""" tokamak_args_dict["major_radius"] = major_radius tokamak_args_dict["minor_radius"] = minor_radius # Set shafranov factor to 0 and pedestal factor to 0.8*minor_radius for safety tokamak_args_dict["pedestal_radius"] = 0.8 * minor_radius tokamak_args_dict["shafranov_factor"] = 0.0 - tokamak_source = TokamakSource(**tokamak_args_dict) - assert tokamak_source.minor_radius == minor_radius + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize( @@ -84,43 +87,41 @@ def test_minor_radius(tokamak_args_dict, major_radius, minor_radius): [(10.0, 10.0), (10.0, 20.0), (10.0, 0), (10.0, -6), (10.0, "hello world")], ) def test_bad_minor_radius(tokamak_args_dict, major_radius, minor_radius): - """Checks that TokamakSource creation rejects invalid 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): - tokamak_source = TokamakSource(**tokamak_args_dict) + with pytest.raises((ValueError, TypeError)): + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize("elongation", [1.0, 1.667, 0.5, 20, 0.001]) def test_elongation(tokamak_args_dict, elongation): - """Checks that TokamakSource creation accepts valid elongation""" + """Checks that tokamak_source creation accepts valid elongation""" tokamak_args_dict["elongation"] = elongation - tokamak_source = TokamakSource(**tokamak_args_dict) - assert tokamak_source.elongation == elongation + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize("elongation", [0, -5, "hello world"]) def test_bad_elongation(tokamak_args_dict, elongation): - """Checks that TokamakSource creation rejects invalid elongation""" + """Checks that tokamak_source creation rejects invalid elongation""" tokamak_args_dict["elongation"] = elongation - with pytest.raises(ValueError): - tokamak_source = TokamakSource(**tokamak_args_dict) + with pytest.raises((ValueError, TypeError)): + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize("triangularity", [0.0, 0.5, 0.9, 1.0, -0.5, -0.9, -1.0]) def test_triangularity(tokamak_args_dict, triangularity): - """Checks that TokamakSource creation accepts valid triangularity""" + """Checks that tokamak_source creation accepts valid triangularity""" tokamak_args_dict["triangularity"] = triangularity - tokamak_source = TokamakSource(**tokamak_args_dict) - assert tokamak_source.triangularity == triangularity + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize("triangularity", [1.1, -1.1, 10, -10, "hello world"]) def test_bad_triangularity(tokamak_args_dict, triangularity): - """Checks that TokamakSource creation rejects invalid triangularity""" + """Checks that tokamak_source creation rejects invalid triangularity""" tokamak_args_dict["triangularity"] = triangularity - with pytest.raises(ValueError): - tokamak_source = TokamakSource(**tokamak_args_dict) + with pytest.raises((ValueError, TypeError)): + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize( @@ -136,13 +137,12 @@ def test_bad_triangularity(tokamak_args_dict, triangularity): ], ) def test_shafranov_factor(tokamak_args_dict, major_radius, minor_radius, shaf): - """Checks that TokamakSource creation accepts valid Shafranov factor""" + """Checks that tokamak_source creation accepts valid Shafranov factor""" tokamak_args_dict["major_radius"] = major_radius tokamak_args_dict["minor_radius"] = minor_radius tokamak_args_dict["pedestal_radius"] = 0.8 * minor_radius tokamak_args_dict["shafranov_factor"] = shaf - tokamak_source = TokamakSource(**tokamak_args_dict) - assert tokamak_source.shafranov_factor == shaf + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize( @@ -157,13 +157,13 @@ def test_shafranov_factor(tokamak_args_dict, major_radius, minor_radius, shaf): ], ) def test_bad_shafranov_factor(tokamak_args_dict, major_radius, minor_radius, shaf): - """Checks that TokamakSource creation rejects invalid Shafranov factor""" + """Checks that tokamak_source creation rejects invalid Shafranov factor""" tokamak_args_dict["major_radius"] = major_radius tokamak_args_dict["minor_radius"] = minor_radius tokamak_args_dict["pedestal_radius"] = 0.8 * minor_radius tokamak_args_dict["shafranov_factor"] = shaf with pytest.raises((ValueError, TypeError)): - tokamak_source = TokamakSource(**tokamak_args_dict) + tokamak_source(**tokamak_args_dict) @pytest.mark.parametrize( @@ -173,9 +173,8 @@ def test_angles(tokamak_args_dict, angles): """Checks that custom angles can be set""" # Note: should accept negative angles and angles in reverse order tokamak_args_dict["angles"] = angles - tokamak_source = TokamakSource(**tokamak_args_dict) - assert np.array_equal(tokamak_source.angles, angles) - for source in tokamak_source.sources: + sources = tokamak_source(**tokamak_args_dict) + for source in sources: assert np.array_equal((source.space.phi.a, source.space.phi.b), angles) @@ -185,49 +184,101 @@ 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: - tokamak_source = TokamakSource(**tokamak_args_dict) + with pytest.raises((ValueError, TypeError)): + tokamak_source(**tokamak_args_dict) -def test_ion_density(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_source_example.minor_radius, 100) - density = tokamak_source_example.ion_density(r) + 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_bad_ion_density(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: - density = tokamak_source_example.ion_density([0, 5, -6]) + 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_source_example): +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, tokamak_source_example.minor_radius, 100) - temperature = tokamak_source_example.ion_temperature(r) + 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_ion_temperature(tokamak_source_example): +def test_bad_ion_temperature(tokamak_args_dict): # It should fail if given a negative r with pytest.raises(ValueError) as excinfo: - temperature = tokamak_source_example.ion_temperature([0, 5, -6]) + 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_source_example): +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, tokamak_source_example.minor_radius, 100) + a = np.linspace(0.0, 2.9, 100) alpha = np.linspace(0.0, 2 * np.pi, 100) - R, Z = tokamak_source_example.convert_a_alpha_to_R_Z(a, alpha) + 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) @@ -236,18 +287,26 @@ def test_convert_a_alpha_to_R_Z(tokamak_source_example): assert np.all(np.isfinite(Z)) -def test_bad_convert_a_alpha_to_R_Z(tokamak_source_example): +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, tokamak_source_example.minor_radius, 100) + a = np.linspace(0.0, 2.9, 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) + 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 def tokamak_source_strategy(draw): - """Defines a hypothesis strategy that automatically generates a TokamakSource. + """Defines a hypothesis strategy that automatically generates a tokamak_source. Geometry attributes are varied, while plasma attributes are fixed. """ # Used to avoid generation of inappropriate float values @@ -292,7 +351,7 @@ def tokamak_source_strategy(draw): ) ) - return TokamakSource( + return tokamak_source( elongation=elongation, triangularity=triangularity, major_radius=major_radius, @@ -309,18 +368,27 @@ 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) +@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""" - assert pytest.approx(sum(tokamak_source.strengths)) == 1 + 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=50) +@settings(max_examples=50, suppress_health_check=(HealthCheck.too_slow,)) def test_source_locations_are_within_correct_range(tokamak_source): """Tests that each source has RZ locations within the expected range. @@ -331,10 +399,10 @@ def test_source_locations_are_within_correct_range(tokamak_source): 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 + 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""" @@ -343,7 +411,7 @@ def get_R_on_LCMS(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: + 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)