cfspopcon 4.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. cfspopcon-4.0.0/LICENSE +19 -0
  2. cfspopcon-4.0.0/PKG-INFO +38 -0
  3. cfspopcon-4.0.0/README.md +8 -0
  4. cfspopcon-4.0.0/cfspopcon/__init__.py +49 -0
  5. cfspopcon-4.0.0/cfspopcon/algorithms/__init__.py +74 -0
  6. cfspopcon-4.0.0/cfspopcon/algorithms/algorithm_class.py +413 -0
  7. cfspopcon-4.0.0/cfspopcon/algorithms/beta.py +50 -0
  8. cfspopcon-4.0.0/cfspopcon/algorithms/composite_algorithm.py +67 -0
  9. cfspopcon-4.0.0/cfspopcon/algorithms/core_radiated_power.py +90 -0
  10. cfspopcon-4.0.0/cfspopcon/algorithms/extrinsic_core_radiator.py +99 -0
  11. cfspopcon-4.0.0/cfspopcon/algorithms/fusion_gain.py +64 -0
  12. cfspopcon-4.0.0/cfspopcon/algorithms/geometry.py +54 -0
  13. cfspopcon-4.0.0/cfspopcon/algorithms/heat_exhaust.py +78 -0
  14. cfspopcon-4.0.0/cfspopcon/algorithms/ohmic_power.py +56 -0
  15. cfspopcon-4.0.0/cfspopcon/algorithms/peaked_profiles.py +91 -0
  16. cfspopcon-4.0.0/cfspopcon/algorithms/plasma_current_from_q_star.py +44 -0
  17. cfspopcon-4.0.0/cfspopcon/algorithms/power_balance_from_tau_e.py +74 -0
  18. cfspopcon-4.0.0/cfspopcon/algorithms/q_star_from_plasma_current.py +43 -0
  19. cfspopcon-4.0.0/cfspopcon/algorithms/single_functions.py +75 -0
  20. cfspopcon-4.0.0/cfspopcon/algorithms/two_point_model_fixed_fpow.py +70 -0
  21. cfspopcon-4.0.0/cfspopcon/algorithms/two_point_model_fixed_qpart.py +71 -0
  22. cfspopcon-4.0.0/cfspopcon/algorithms/two_point_model_fixed_tet.py +67 -0
  23. cfspopcon-4.0.0/cfspopcon/algorithms/use_LOC_tau_e_below_threshold.py +87 -0
  24. cfspopcon-4.0.0/cfspopcon/algorithms/zeff_and_dilution_from_impurities.py +57 -0
  25. cfspopcon-4.0.0/cfspopcon/atomic_data/__init__.py +4 -0
  26. cfspopcon-4.0.0/cfspopcon/atomic_data/read_radas_data.py +80 -0
  27. cfspopcon-4.0.0/cfspopcon/cli.py +79 -0
  28. cfspopcon-4.0.0/cfspopcon/file_io.py +104 -0
  29. cfspopcon-4.0.0/cfspopcon/formulas/Q_thermal_gain_factor.py +39 -0
  30. cfspopcon-4.0.0/cfspopcon/formulas/__init__.py +94 -0
  31. cfspopcon-4.0.0/cfspopcon/formulas/average_fuel_ion_mass.py +30 -0
  32. cfspopcon-4.0.0/cfspopcon/formulas/beta.py +135 -0
  33. cfspopcon-4.0.0/cfspopcon/formulas/confinement_regime_threshold_powers.py +153 -0
  34. cfspopcon-4.0.0/cfspopcon/formulas/current_drive.py +254 -0
  35. cfspopcon-4.0.0/cfspopcon/formulas/density_peaking.py +52 -0
  36. cfspopcon-4.0.0/cfspopcon/formulas/divertor_metrics.py +40 -0
  37. cfspopcon-4.0.0/cfspopcon/formulas/energy_confinement_time_scalings/__init__.py +9 -0
  38. cfspopcon-4.0.0/cfspopcon/formulas/energy_confinement_time_scalings/tau_e_from_Wp.py +169 -0
  39. cfspopcon-4.0.0/cfspopcon/formulas/energy_confinement_time_scalings/tau_e_scalings.yaml +342 -0
  40. cfspopcon-4.0.0/cfspopcon/formulas/figures_of_merit.py +108 -0
  41. cfspopcon-4.0.0/cfspopcon/formulas/fusion_rates.py +241 -0
  42. cfspopcon-4.0.0/cfspopcon/formulas/fusion_reaction_data/__init__.py +53 -0
  43. cfspopcon-4.0.0/cfspopcon/formulas/fusion_reaction_data/reaction_energies.py +140 -0
  44. cfspopcon-4.0.0/cfspopcon/formulas/fusion_reaction_data/reaction_rate_coefficients.py +298 -0
  45. cfspopcon-4.0.0/cfspopcon/formulas/geometry.py +41 -0
  46. cfspopcon-4.0.0/cfspopcon/formulas/helpers.py +23 -0
  47. cfspopcon-4.0.0/cfspopcon/formulas/impurity_effects.py +65 -0
  48. cfspopcon-4.0.0/cfspopcon/formulas/operational_limits.py +61 -0
  49. cfspopcon-4.0.0/cfspopcon/formulas/plasma_profile_data/PRF/aLT.csv +12 -0
  50. cfspopcon-4.0.0/cfspopcon/formulas/plasma_profile_data/PRF/metadata.yaml +2 -0
  51. cfspopcon-4.0.0/cfspopcon/formulas/plasma_profile_data/PRF/width.csv +12 -0
  52. cfspopcon-4.0.0/cfspopcon/formulas/plasma_profile_data/__init__.py +1 -0
  53. cfspopcon-4.0.0/cfspopcon/formulas/plasma_profile_data/density_and_temperature_profile_fits.py +216 -0
  54. cfspopcon-4.0.0/cfspopcon/formulas/plasma_profiles.py +186 -0
  55. cfspopcon-4.0.0/cfspopcon/formulas/radiated_power/__init__.py +17 -0
  56. cfspopcon-4.0.0/cfspopcon/formulas/radiated_power/inherent.py +141 -0
  57. cfspopcon-4.0.0/cfspopcon/formulas/radiated_power/mavrin_coronal.py +192 -0
  58. cfspopcon-4.0.0/cfspopcon/formulas/radiated_power/mavrin_noncoronal.py +251 -0
  59. cfspopcon-4.0.0/cfspopcon/formulas/radiated_power/post_and_jensen.py +158 -0
  60. cfspopcon-4.0.0/cfspopcon/formulas/radiated_power/radas.py +67 -0
  61. cfspopcon-4.0.0/cfspopcon/formulas/radiated_power/radiated_power.py +55 -0
  62. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/__init__.py +15 -0
  63. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/lambda_q.py +92 -0
  64. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/momentum_loss_functions.py +62 -0
  65. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/parallel_heat_flux_density.py +35 -0
  66. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/required_power_loss_fraction.py +40 -0
  67. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/solve_target_first_two_point_model.py +121 -0
  68. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/solve_two_point_model.py +209 -0
  69. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/target_electron_density.py +93 -0
  70. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/target_electron_flux.py +87 -0
  71. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/target_electron_temp.py +88 -0
  72. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/total_pressure.py +65 -0
  73. cfspopcon-4.0.0/cfspopcon/formulas/scrape_off_layer_model/upstream_electron_temp.py +33 -0
  74. cfspopcon-4.0.0/cfspopcon/helpers.py +105 -0
  75. cfspopcon-4.0.0/cfspopcon/input_file_handling.py +64 -0
  76. cfspopcon-4.0.0/cfspopcon/named_options.py +127 -0
  77. cfspopcon-4.0.0/cfspopcon/plotting/__init__.py +12 -0
  78. cfspopcon-4.0.0/cfspopcon/plotting/coordinate_formatter.py +23 -0
  79. cfspopcon-4.0.0/cfspopcon/plotting/make_plot.py +156 -0
  80. cfspopcon-4.0.0/cfspopcon/plotting/plot_style_handling.py +21 -0
  81. cfspopcon-4.0.0/cfspopcon/point_selection.py +58 -0
  82. cfspopcon-4.0.0/cfspopcon/transform.py +173 -0
  83. cfspopcon-4.0.0/cfspopcon/unit_handling/__init__.py +28 -0
  84. cfspopcon-4.0.0/cfspopcon/unit_handling/decorator.py +294 -0
  85. cfspopcon-4.0.0/cfspopcon/unit_handling/default_units.py +260 -0
  86. cfspopcon-4.0.0/cfspopcon/unit_handling/setup_unit_handling.py +98 -0
  87. cfspopcon-4.0.0/pyproject.toml +115 -0
@@ -0,0 +1,19 @@
1
+ Copyright (c) 2023 Commonwealth Fusion Systems
2
+
3
+ Permission is hereby granted, free of charge, to any person obtaining a copy of
4
+ this software and associated documentation files (the “Software”), to deal in
5
+ the Software without restriction, including without limitation the rights to
6
+ use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
7
+ of the Software, and to permit persons to whom the Software is furnished to do
8
+ so, subject to the following conditions:
9
+
10
+ The above copyright notice and this permission notice shall be included in all
11
+ copies or substantial portions of the Software.
12
+
13
+ THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19
+ SOFTWARE.
@@ -0,0 +1,38 @@
1
+ Metadata-Version: 2.1
2
+ Name: cfspopcon
3
+ Version: 4.0.0
4
+ Summary: Empirically-derived scoping of tokamak operational space.
5
+ Author: Commonwealth Fusion Systems
6
+ Requires-Python: >=3.9,<3.12
7
+ Classifier: Development Status :: 5 - Production/Stable
8
+ Classifier: Intended Audience :: Science/Research
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.9
12
+ Classifier: Programming Language :: Python :: 3.10
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3 :: Only
15
+ Classifier: Topic :: Scientific/Engineering :: Physics
16
+ Requires-Dist: click (>=8.1.0,<9.0.0)
17
+ Requires-Dist: ipdb (>=0.13.13,<0.14.0)
18
+ Requires-Dist: numpy (>=1.22.4,<2.0.0)
19
+ Requires-Dist: pandas (>=1.4,<2.0)
20
+ Requires-Dist: pint (>=0.22,<0.23)
21
+ Requires-Dist: pint-xarray (>=0.3,<0.4)
22
+ Requires-Dist: pyyaml (>=6.0,<7.0)
23
+ Requires-Dist: scipy (>=1.8,<2.0)
24
+ Requires-Dist: seaborn (>=0.12,<0.13)
25
+ Requires-Dist: toml (>=0.10.2,<0.11.0)
26
+ Requires-Dist: typing-extensions (>=4.0.1,<5.0.0)
27
+ Requires-Dist: xarray (>=2023.4.1,<2024.0.0)
28
+ Description-Content-Type: text/markdown
29
+
30
+ cfspopcon: 0D Plasma Calculations & Plasma OPerating CONtours
31
+ --------------------------------------------------------------
32
+
33
+ [![Build Status](https://github.com/cfs-energy/cfspopcon/actions/workflows/workflow_actions.yml/badge.svg)](https://github.com/cfs-energy/cfspopcon/actions)
34
+ [![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/)
35
+ [![Documentation Status](https://readthedocs.org/projects/cfspopcon/badge/?version=latest)](https://cfspopcon.readthedocs.io/en/latest/?badge=latest)
36
+
37
+ For more information please have a look at our [documentation](https://cfspopcon.readthedocs.io/en/latest/).
38
+
@@ -0,0 +1,8 @@
1
+ cfspopcon: 0D Plasma Calculations & Plasma OPerating CONtours
2
+ --------------------------------------------------------------
3
+
4
+ [![Build Status](https://github.com/cfs-energy/cfspopcon/actions/workflows/workflow_actions.yml/badge.svg)](https://github.com/cfs-energy/cfspopcon/actions)
5
+ [![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/)
6
+ [![Documentation Status](https://readthedocs.org/projects/cfspopcon/badge/?version=latest)](https://cfspopcon.readthedocs.io/en/latest/?badge=latest)
7
+
8
+ For more information please have a look at our [documentation](https://cfspopcon.readthedocs.io/en/latest/).
@@ -0,0 +1,49 @@
1
+ """Physics calculations & lumped-parameter models."""
2
+ from importlib.metadata import metadata
3
+
4
+ __version__ = metadata(__package__)["Version"]
5
+ __author__ = metadata(__package__)["Author"]
6
+
7
+ from . import algorithms, file_io, formulas, helpers, named_options, unit_handling
8
+ from .algorithms.algorithm_class import Algorithm, CompositeAlgorithm
9
+ from .input_file_handling import read_case
10
+ from .plotting import read_plot_style
11
+ from .point_selection import find_coords_of_maximum, find_coords_of_minimum
12
+ from .unit_handling import (
13
+ Quantity,
14
+ Unit,
15
+ convert_to_default_units,
16
+ convert_units,
17
+ default_unit,
18
+ magnitude,
19
+ magnitude_in_default_units,
20
+ set_default_units,
21
+ ureg,
22
+ wraps_ufunc,
23
+ )
24
+
25
+ # export main classes users should need as well as the option enums
26
+ __all__ = [
27
+ "helpers",
28
+ "named_options",
29
+ "algorithms",
30
+ "formulas",
31
+ "unit_handling",
32
+ "ureg",
33
+ "Quantity",
34
+ "Unit",
35
+ "wraps_ufunc",
36
+ "magnitude_in_default_units",
37
+ "set_default_units",
38
+ "default_unit",
39
+ "convert_to_default_units",
40
+ "convert_units",
41
+ "magnitude",
42
+ "read_case",
43
+ "read_plot_style",
44
+ "find_coords_of_maximum",
45
+ "find_coords_of_minimum",
46
+ "Algorithm",
47
+ "CompositeAlgorithm",
48
+ "file_io",
49
+ ]
@@ -0,0 +1,74 @@
1
+ """POPCON algorithms."""
2
+ from typing import Union
3
+
4
+ from ..named_options import Algorithms
5
+ from .algorithm_class import Algorithm, CompositeAlgorithm
6
+ from .beta import calc_beta
7
+ from .composite_algorithm import predictive_popcon
8
+ from .core_radiated_power import calc_core_radiated_power
9
+ from .extrinsic_core_radiator import calc_extrinsic_core_radiator
10
+ from .fusion_gain import calc_fusion_gain
11
+ from .geometry import calc_geometry
12
+ from .heat_exhaust import calc_heat_exhaust
13
+ from .ohmic_power import calc_ohmic_power
14
+ from .peaked_profiles import calc_peaked_profiles
15
+ from .plasma_current_from_q_star import calc_plasma_current_from_q_star
16
+ from .power_balance_from_tau_e import calc_power_balance_from_tau_e
17
+ from .q_star_from_plasma_current import calc_q_star_from_plasma_current
18
+ from .single_functions import SINGLE_FUNCTIONS
19
+ from .two_point_model_fixed_fpow import two_point_model_fixed_fpow
20
+ from .two_point_model_fixed_qpart import two_point_model_fixed_qpart
21
+ from .two_point_model_fixed_tet import two_point_model_fixed_tet
22
+ from .use_LOC_tau_e_below_threshold import use_LOC_tau_e_below_threshold
23
+ from .zeff_and_dilution_from_impurities import calc_zeff_and_dilution_from_impurities
24
+
25
+ ALGORITHMS: dict[Algorithms, Union[Algorithm, CompositeAlgorithm]] = {
26
+ Algorithms["calc_beta"]: calc_beta,
27
+ Algorithms["calc_core_radiated_power"]: calc_core_radiated_power,
28
+ Algorithms["calc_extrinsic_core_radiator"]: calc_extrinsic_core_radiator,
29
+ Algorithms["calc_fusion_gain"]: calc_fusion_gain,
30
+ Algorithms["calc_geometry"]: calc_geometry,
31
+ Algorithms["calc_heat_exhaust"]: calc_heat_exhaust,
32
+ Algorithms["calc_ohmic_power"]: calc_ohmic_power,
33
+ Algorithms["calc_peaked_profiles"]: calc_peaked_profiles,
34
+ Algorithms["calc_plasma_current_from_q_star"]: calc_plasma_current_from_q_star,
35
+ Algorithms["calc_power_balance_from_tau_e"]: calc_power_balance_from_tau_e,
36
+ Algorithms["predictive_popcon"]: predictive_popcon,
37
+ Algorithms["calc_q_star_from_plasma_current"]: calc_q_star_from_plasma_current,
38
+ Algorithms["two_point_model_fixed_fpow"]: two_point_model_fixed_fpow,
39
+ Algorithms["two_point_model_fixed_qpart"]: two_point_model_fixed_qpart,
40
+ Algorithms["two_point_model_fixed_tet"]: two_point_model_fixed_tet,
41
+ Algorithms["calc_zeff_and_dilution_from_impurities"]: calc_zeff_and_dilution_from_impurities,
42
+ Algorithms["use_LOC_tau_e_below_threshold"]: use_LOC_tau_e_below_threshold,
43
+ **SINGLE_FUNCTIONS,
44
+ }
45
+
46
+
47
+ def get_algorithm(algorithm: Union[Algorithms, str]) -> Union[Algorithm, CompositeAlgorithm]:
48
+ """Accessor for algorithms."""
49
+ if isinstance(algorithm, str):
50
+ algorithm = Algorithms[algorithm]
51
+
52
+ return ALGORITHMS[algorithm]
53
+
54
+
55
+ __all__ = [
56
+ "calc_beta",
57
+ "calc_core_radiated_power",
58
+ "calc_extrinsic_core_radiator",
59
+ "calc_fusion_gain",
60
+ "calc_geometry",
61
+ "calc_heat_exhaust",
62
+ "calc_ohmic_power",
63
+ "calc_peaked_profiles",
64
+ "calc_plasma_current_from_q_star",
65
+ "calc_power_balance_from_tau_e",
66
+ "predictive_popcon",
67
+ "calc_q_star_from_plasma_current",
68
+ "two_point_model_fixed_fpow",
69
+ "two_point_model_fixed_qpart",
70
+ "two_point_model_fixed_tet",
71
+ "calc_zeff_and_dilution_from_impurities",
72
+ "ALGORITHMS",
73
+ "get_algorithm",
74
+ ]
@@ -0,0 +1,413 @@
1
+ """Defines a class for different POPCON algorithms."""
2
+ from __future__ import annotations
3
+
4
+ import inspect
5
+ from collections.abc import Callable, Sequence
6
+ from functools import wraps
7
+ from typing import Any, Optional, Union
8
+ from warnings import warn
9
+
10
+ import xarray as xr
11
+
12
+ from ..unit_handling import convert_to_default_units
13
+
14
+ FunctionType = Callable[..., dict[str, Any]]
15
+
16
+
17
+ class Algorithm:
18
+ """A class which handles the input and output of POPCON algorithms."""
19
+
20
+ def __init__(self, function: FunctionType, return_keys: list[str], name: Optional[str] = None):
21
+ """Initialise an Algorithm.
22
+
23
+ Args:
24
+ function: a callable function
25
+ return_keys: the arguments which are returned from the function
26
+ name: Descriptive name for algorithm
27
+ """
28
+ self._function = function
29
+ self._name = self._function.__name__ if name is None else name
30
+
31
+ self._signature = inspect.signature(function)
32
+ for p in self._signature.parameters.values():
33
+ if p.kind not in (
34
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
35
+ inspect.Parameter.KEYWORD_ONLY,
36
+ inspect.Parameter.VAR_KEYWORD,
37
+ ):
38
+ raise ValueError(
39
+ f"Algorithm only supports functions with keyword arguments, but {function}, has {p.kind} parameter {p.name}"
40
+ )
41
+ self.input_keys = list(self._signature.parameters.keys())
42
+ self.return_keys = return_keys
43
+
44
+ self.default_values = {
45
+ key: val.default for key, val in self._signature.parameters.items() if val.default is not inspect.Parameter.empty
46
+ }
47
+ self.default_keys = list(self.default_values.keys())
48
+
49
+ self.required_input_keys = [key for key in self.input_keys if key not in self.default_keys]
50
+
51
+ self.__doc__ = self._make_docstring()
52
+
53
+ self.run = self._make_run(self._function)
54
+
55
+ def _make_docstring(self) -> str:
56
+ """Makes a doc-string detailing the function inputs and outputs."""
57
+ return_string = (
58
+ f"Algorithm: {self._name}\n" + "Inputs:\n" + ", ".join(self.input_keys) + "\n" + "Outputs:\n" + ", ".join(self.return_keys)
59
+ )
60
+ return return_string
61
+
62
+ def __repr__(self) -> str:
63
+ """Return a simple string description of the Algorithm."""
64
+ return f"Algorithm: {self._name}"
65
+
66
+ @classmethod
67
+ def _make_run(cls, func: FunctionType) -> Callable[..., xr.Dataset]:
68
+ """Helper to create the `run()` function with correct doc string.
69
+
70
+ Args:
71
+ func: function to be wrapped
72
+
73
+ Returns: a xarray DataSet of the result
74
+ """
75
+
76
+ @wraps(func)
77
+ def run(**kwargs: Any) -> xr.Dataset:
78
+ result = func(**kwargs)
79
+ dataset = xr.Dataset(result)
80
+ return dataset
81
+
82
+ return run
83
+
84
+ def update_dataset(self, dataset: xr.Dataset, in_place: bool = False) -> Optional[xr.Dataset]:
85
+ """Retrieve inputs from passed dataset and return a new dataset combining input and output quantities.
86
+
87
+ Specifying in_place=True modifies the dataset in place (changing the input), whereas in_place=False will
88
+ return a copy of the dataset with the outputs appended.
89
+
90
+ Args:
91
+ dataset: input dataset
92
+ in_place: modify the dataset in place, otherwise return a modified dataset keeping the input unchanged.
93
+
94
+ Returns: modified dataset
95
+ """
96
+ if not in_place:
97
+ dataset = dataset.copy(deep=True)
98
+
99
+ input_values = {}
100
+ for key in self.input_keys:
101
+ if key in dataset.keys():
102
+ input_values[key] = dataset[key]
103
+ elif key in self.default_keys:
104
+ input_values[key] = self.default_values[key]
105
+ else:
106
+ sorted_dataset_keys = ", ".join(sorted(dataset.keys())) # type:ignore[arg-type]
107
+ sorted_default_keys = ", ".join(sorted(self.default_keys))
108
+ raise KeyError(f"Key '{key}' not in dataset keys [{sorted_dataset_keys}] or default values [{sorted_default_keys}]")
109
+
110
+ result = self._function(**input_values)
111
+
112
+ for key, val in result.items():
113
+ dataset[key] = val
114
+
115
+ if not in_place:
116
+ return dataset
117
+ else:
118
+ return None
119
+
120
+ def __add__(self, other: Union[Algorithm, CompositeAlgorithm]) -> CompositeAlgorithm:
121
+ """Build a CompositeAlgorithm composed of this Algorithm and another Algorithm or CompositeAlgorithm."""
122
+ if isinstance(other, CompositeAlgorithm):
123
+ return CompositeAlgorithm(algorithms=[self, *other.algorithms])
124
+ else:
125
+ return CompositeAlgorithm(algorithms=[self, other])
126
+
127
+ @classmethod
128
+ def from_single_function(
129
+ cls, func: Callable, return_keys: list[str], name: Optional[str] = None, skip_unit_conversion: bool = False
130
+ ) -> Algorithm:
131
+ """Build an Algorithm which wraps a single function."""
132
+
133
+ @wraps(func)
134
+ def wrapped_function(**kwargs: Any) -> dict:
135
+ result = func(**kwargs)
136
+
137
+ if not isinstance(result, tuple):
138
+ result = (result,)
139
+
140
+ result_dict = {}
141
+ for i, key in enumerate(return_keys):
142
+ if skip_unit_conversion:
143
+ result_dict[key] = result[i]
144
+ else:
145
+ result_dict[key] = convert_to_default_units(result[i], key)
146
+
147
+ return result_dict
148
+
149
+ return cls(wrapped_function, return_keys, name=name)
150
+
151
+ def validate_inputs(
152
+ self, configuration: Union[dict, xr.Dataset], quiet: bool = False, raise_error_on_missing_inputs: bool = False
153
+ ) -> bool:
154
+ """Check that all required inputs are defined, and warn if inputs are unused."""
155
+ return _validate_inputs(self, configuration, quiet=quiet, raise_error_on_missing_inputs=raise_error_on_missing_inputs)
156
+
157
+
158
+ class CompositeAlgorithm:
159
+ """A class which combined multiple Algorithms into a single object which behaves like an Algorithm."""
160
+
161
+ def __init__(self, algorithms: Sequence[Union[Algorithm, CompositeAlgorithm]], name: Optional[str] = None):
162
+ """Initialise a CompositeAlgorithm, combining several other Algorithms.
163
+
164
+ Args:
165
+ algorithms: a list of Algorithms, in the order that they should be executed.
166
+ name: a name used to refer to the composite algorithm.
167
+ """
168
+ if not (isinstance(algorithms, Sequence) and all(isinstance(alg, (Algorithm, CompositeAlgorithm)) for alg in algorithms)):
169
+ raise TypeError("Should pass a list of algorithms or composites to CompositeAlgorithm.")
170
+
171
+ self.algorithms: list[Algorithm] = []
172
+
173
+ # flattens composite algorithms into their respective list of plain Algorithms
174
+ for alg in algorithms:
175
+ if isinstance(alg, Algorithm):
176
+ self.algorithms.append(alg)
177
+ else:
178
+ self.algorithms.extend(alg.algorithms)
179
+
180
+ self.input_keys: list[str] = []
181
+ self.required_input_keys: list[str] = []
182
+ self.return_keys: list[str] = []
183
+ pars: list[inspect.Parameter] = []
184
+
185
+ # traverse list of algorithms in order.
186
+ # If an ouput from the set of previous algorithms provides an input to a following algorithm
187
+ # the input is not turned into an input to the CompositeAlgorithm
188
+ for alg in self.algorithms:
189
+ alg_sig = inspect.signature(alg.run)
190
+ for key in alg.default_keys:
191
+ if key not in self.return_keys:
192
+ self.input_keys.append(key)
193
+ pars.append(alg_sig.parameters[key])
194
+ for key in alg.required_input_keys:
195
+ if key not in self.return_keys:
196
+ self.input_keys.append(key)
197
+ self.required_input_keys.append(key)
198
+ pars.append(alg_sig.parameters[key])
199
+
200
+ for key in alg.return_keys:
201
+ if key not in self.return_keys:
202
+ self.return_keys.append(key)
203
+
204
+ # create a signature for the run() function
205
+ # This is a purely aesthetic change, that ensures the run() function
206
+ # has a helpful tooltip in editors and in the documentation
207
+
208
+ # 1. make sure the list of pars doesn't have any duplicates, if there are duplicates
209
+ # we pick the first one. We don't assert that the types of two parameters are compatible
210
+ # that's not easy to do.
211
+ seen_pars: dict[str, int] = {}
212
+ pars = [p for i, p in enumerate(pars) if seen_pars.setdefault(p.name, i) == i]
213
+
214
+ # ensure POSITIONAL_OR_KEYWORD are before kw only
215
+ pars = sorted(pars, key=lambda p: p.kind)
216
+
217
+ def_pars = [p for p in pars if p.default != inspect.Parameter.empty]
218
+ non_def_pars = [p for p in pars if p.default == inspect.Parameter.empty]
219
+
220
+ # methods are immutable and we don't want to set a signature on the class' run() method
221
+ # thus we wrap the original run method and then assign the __signature__ to the wrapped
222
+ # wrapper function
223
+ def _wrap(f: Callable[..., xr.Dataset]) -> Callable[..., xr.Dataset]:
224
+ def wrapper(**kwargs: Any) -> xr.Dataset:
225
+ return f(**kwargs)
226
+
227
+ wrapper.__doc__ = f.__doc__
228
+
229
+ return wrapper
230
+
231
+ self.run = _wrap(self._run)
232
+ # ignore due to mypy bug/missing feature https://github.com/python/mypy/issues/3482
233
+ self.run.__signature__ = inspect.Signature( # type:ignore[attr-defined]
234
+ non_def_pars + def_pars, return_annotation=xr.Dataset
235
+ )
236
+ self._name = name
237
+ self.__doc__ = self._make_docstring()
238
+
239
+ def _make_docstring(self) -> str:
240
+ """Makes a doc-string detailing the function inputs and outputs."""
241
+ components = f"[{', '.join(alg._name for alg in self.algorithms)}]"
242
+
243
+ return_string = (
244
+ f"CompositeAlgorithm: {self._name}\n"
245
+ if self._name is not None
246
+ else "CompositeAlgorithm\n"
247
+ f"Composed of {components}\n"
248
+ f"Inputs:\n{', '.join(self.input_keys)}\n"
249
+ f"Outputs:\n{', '.join(self.return_keys)}"
250
+ )
251
+ return return_string
252
+
253
+ def __repr__(self) -> str:
254
+ """Return a simple string description of the CompositeAlgorithm."""
255
+ return f"CompositeAlgorithm: {self._name}"
256
+
257
+ def _run(self, **kwargs: Any) -> xr.Dataset:
258
+ """Run the sub-Algorithms, one after the other and return a xarray.Dataset of the results.
259
+
260
+ Will throw a warning if parameters are not used by any sub-Algorithm.
261
+ """
262
+ result = kwargs
263
+
264
+ parameters_extra = set(kwargs) - set(self.required_input_keys)
265
+ parameters_missing = set(self.required_input_keys) - set(kwargs)
266
+ if parameters_missing:
267
+ raise TypeError(f"CompositeAlgorithm.run() missing arguments: {', '.join(parameters_missing)}")
268
+ if parameters_extra:
269
+ warn(f"Not all input parameters were used. Unused parameters: [{', '.join(parameters_extra)}]", stacklevel=3)
270
+
271
+ for alg in self.algorithms:
272
+
273
+ alg_kwargs = {key: result[key] for key in result.keys() if key in alg.input_keys}
274
+
275
+ alg_result = alg.run(**alg_kwargs)
276
+ result.update(alg_result) # type:ignore[arg-type] # dict.update() doesn't like KeysView[Hashable]
277
+
278
+ return xr.Dataset(result)
279
+
280
+ def update_dataset(self, dataset: xr.Dataset, in_place: bool = False) -> Optional[xr.Dataset]:
281
+ """Retrieve inputs from passed dataset and return a new dataset combining input and output quantities.
282
+
283
+ Specifying in_place=True modifies the dataset in place (changing the input), whereas in_place=False will
284
+ return a copy of the dataset with the outputs appended.
285
+
286
+ N.b. will not throw a warning if the dataset contains unused elements.
287
+
288
+ Args:
289
+ dataset: input dataset
290
+ in_place: modify the dataset in place, otherwise return a modified dataset keeping the input unchanged.
291
+
292
+ Returns: modified dataset
293
+ """
294
+ if not in_place:
295
+ dataset = dataset.copy(deep=True)
296
+
297
+ for alg in self.algorithms:
298
+ # We've already used copy on the dataset, so can now call update_dataset with
299
+ # in_place = True for each of the algorithms.
300
+ alg.update_dataset(dataset, in_place=True)
301
+
302
+ if not in_place:
303
+ return dataset
304
+ else:
305
+ return None
306
+
307
+ def __add__(self, other: Union[Algorithm, CompositeAlgorithm]) -> CompositeAlgorithm:
308
+ """Build a CompositeAlgorithm composed of this CompositeAlgorithm and another Algorithm or CompositeAlgorithm."""
309
+ if isinstance(other, Algorithm):
310
+ return CompositeAlgorithm(algorithms=[*self.algorithms, other])
311
+ else:
312
+ return CompositeAlgorithm(algorithms=[*self.algorithms, *other.algorithms])
313
+
314
+ def validate_inputs( # noqa: PLR0912
315
+ self,
316
+ configuration: Union[dict, xr.Dataset],
317
+ quiet: bool = False,
318
+ raise_error_on_missing_inputs: bool = True,
319
+ warn_for_overridden_variables: bool = False,
320
+ ) -> bool:
321
+ """Check that all required inputs are defined, and warn if inputs are unused."""
322
+ # Check if variables are being silently internally overwritten
323
+ config_keys = list(configuration.keys())
324
+ key_setter = {key: ["INPUT"] for key in config_keys}
325
+
326
+ for algorithm in self.algorithms:
327
+ for key in algorithm.return_keys:
328
+ if key not in key_setter.keys():
329
+ key_setter[key] = [algorithm._name]
330
+ else:
331
+ key_setter[key].append(algorithm._name)
332
+
333
+ overridden_variables = []
334
+ for variable, algs in key_setter.items():
335
+ if len(algs) > 1:
336
+ overridden_variables.append(f"{variable}: ({', '.join(algs)})")
337
+
338
+ if warn_for_overridden_variables and len(overridden_variables) > 0:
339
+ warn(
340
+ f"The following variables were overridden internally (given as variable: (list of algorithms setting variable)): {', '.join(overridden_variables)}",
341
+ stacklevel=3,
342
+ )
343
+
344
+ # Check that algorithms are ordered such that dependent algorithms follow those setting their required input keys
345
+ available_parameters = config_keys.copy()
346
+ out_of_order_parameters = {}
347
+ for algorithm in self.algorithms:
348
+ for key in algorithm.required_input_keys:
349
+ if key not in available_parameters:
350
+ out_of_order_parameters[key] = algorithm
351
+ for key in algorithm.return_keys:
352
+ available_parameters.append(key)
353
+
354
+ if len(out_of_order_parameters) > 0:
355
+ message = ""
356
+ for key, algorithm in out_of_order_parameters.items():
357
+ if key in key_setter and len(key_setter.get(key, [])) > 0:
358
+ message += f"{key} needed by {algorithm} defined by output of {key_setter[key]}."
359
+ if len(message) > 0:
360
+ message = f"Algorithms out of order. {message}. Rearrange the list of algorithms so that dependent algorithm are after algorithms setting their inputs."
361
+ if raise_error_on_missing_inputs:
362
+ raise RuntimeError(message)
363
+ if not quiet:
364
+ warn(message, stacklevel=3)
365
+
366
+ _validate_inputs(self, configuration, quiet=quiet, raise_error_on_missing_inputs=raise_error_on_missing_inputs)
367
+
368
+ return False
369
+ else:
370
+ return _validate_inputs(self, configuration, quiet=quiet, raise_error_on_missing_inputs=raise_error_on_missing_inputs)
371
+
372
+
373
+ def _validate_inputs(
374
+ algorithm: Union[Algorithm, CompositeAlgorithm],
375
+ configuration: Union[dict, xr.Dataset],
376
+ quiet: bool = False,
377
+ raise_error_on_missing_inputs: bool = False,
378
+ ) -> bool:
379
+ """Check that all required inputs are defined, and warn if inputs are unused."""
380
+ config_keys = list(configuration.keys())
381
+
382
+ unused_config_keys = config_keys.copy()
383
+ missing_input_keys = set(algorithm.required_input_keys)
384
+
385
+ for key in config_keys:
386
+ if key in missing_input_keys:
387
+ missing_input_keys.remove(key)
388
+
389
+ if key in algorithm.input_keys:
390
+ # required_input_keys gives the list of keys which must
391
+ # be provided, while input_puts gives the list of keys
392
+ # which can be provided (but which might have default values).
393
+ unused_config_keys.remove(key)
394
+
395
+ if len(missing_input_keys) == 0 and len(unused_config_keys) == 0:
396
+ return True
397
+
398
+ elif len(missing_input_keys) > 0 and len(unused_config_keys) > 0:
399
+ message = f"Missing input parameters [{', '.join(missing_input_keys)}]. Also had unused input parameters [{', '.join(unused_config_keys)}]."
400
+ if raise_error_on_missing_inputs:
401
+ raise RuntimeError(message)
402
+
403
+ elif len(missing_input_keys) > 0:
404
+ message = f"Missing input parameters [{', '.join(missing_input_keys)}]."
405
+ if raise_error_on_missing_inputs:
406
+ raise RuntimeError(message)
407
+
408
+ else:
409
+ message = f"Unused input parameters [{', '.join(unused_config_keys)}]."
410
+
411
+ if not quiet:
412
+ warn(message, stacklevel=3)
413
+ return False
@@ -0,0 +1,50 @@
1
+ """Calculate toroidal, poloidal, total and normalized beta."""
2
+ from .. import formulas
3
+ from ..unit_handling import Unitfull, convert_to_default_units
4
+ from .algorithm_class import Algorithm
5
+
6
+ RETURN_KEYS = [
7
+ "beta_toroidal",
8
+ "beta_poloidal",
9
+ "beta",
10
+ "normalized_beta",
11
+ ]
12
+
13
+
14
+ def run_calc_beta(
15
+ average_electron_density: Unitfull,
16
+ average_electron_temp: Unitfull,
17
+ average_ion_temp: Unitfull,
18
+ magnetic_field_on_axis: Unitfull,
19
+ plasma_current: Unitfull,
20
+ minor_radius: Unitfull,
21
+ ) -> dict[str, Unitfull]:
22
+ """Calculate toroidal, poloidal, total and normalized beta.
23
+
24
+ Args:
25
+ average_electron_density: :term:`glossary link<average_electron_density>`
26
+ average_electron_temp: :term:`glossary link<average_electron_temp>`
27
+ average_ion_temp: :term:`glossary link<average_ion_temp>`
28
+ magnetic_field_on_axis: :term:`glossary link<magnetic_field_on_axis>`
29
+ plasma_current: :term:`glossary link<plasma_current>`
30
+ minor_radius: :term:`glossary link<minor_radius>`
31
+
32
+ Returns:
33
+ :term:`beta_toroidal`, :term:`beta_poloidal`, :term:`beta_total`, :term:`beta_N`
34
+ """
35
+ beta_toroidal = formulas.calc_beta_toroidal(average_electron_density, average_electron_temp, average_ion_temp, magnetic_field_on_axis)
36
+ beta_poloidal = formulas.calc_beta_poloidal(
37
+ average_electron_density, average_electron_temp, average_ion_temp, plasma_current, minor_radius
38
+ )
39
+
40
+ beta = formulas.calc_beta_total(beta_toroidal=beta_toroidal, beta_poloidal=beta_poloidal)
41
+ normalized_beta = formulas.calc_beta_normalised(beta, minor_radius, magnetic_field_on_axis, plasma_current)
42
+
43
+ local_vars = locals()
44
+ return {key: convert_to_default_units(local_vars[key], key) for key in RETURN_KEYS}
45
+
46
+
47
+ calc_beta = Algorithm(
48
+ function=run_calc_beta,
49
+ return_keys=RETURN_KEYS,
50
+ )