Coverage for python/gsfit/database_readers/mock_st40_mdsplus/setup_plasma.py: 0%
50 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-14 15:25 +0000
1import typing
2from typing import TYPE_CHECKING
4import numpy as np
5import numpy.typing as npt
6from gsfit_rs import EfitPolynomial
7from gsfit_rs import Plasma
8from gsfit_rs import TensionedCubicBSpline
10if TYPE_CHECKING:
11 from . import DatabaseReader
14def setup_plasma(
15 self: "DatabaseReader",
16 pulseNo: int,
17 settings: dict[str, typing.Any],
18 times_to_reconstruct: npt.NDArray[np.float64],
19) -> Plasma:
20 """
21 This method initialises the Rust `Plasma` class.
23 :param pulseNo: Pulse number, used to select which mocked MDSplus tree to read
24 :param settings: Dictionary containing the JSON settings read from the `settings` directory
25 :param times_to_reconstruct: Times the equilibrium will be solved at [second]
27 **This method is specific to a mock of ST40's experimental MDSplus database.**
29 See `python/gsfit/database_readers/interface.py` for more details on how a new database_reader should be implemented.
30 """
32 # Initial plasma conditions
33 initial_guess_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"]
34 initial_guess_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_r"]
35 initial_guess_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_z"]
36 initial_guess_minor_radius = settings["GSFIT_code_settings.json"]["initial_guess"]["minor_radius"]
37 initial_guess_elongation = settings["GSFIT_code_settings.json"]["initial_guess"]["elongation"]
39 # Numerical settings the Grad-Shafranov solve is run with
40 n_iter_max = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_max"]
41 n_iter_min = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_min"]
42 n_iter_no_vertical_feedback = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_no_vertical_feedback"]
43 grad_shafranov_deviation_tolerance = settings["GSFIT_code_settings.json"]["numerics"]["grad_shafranov_deviation_tolerance"]
44 use_anderson_mixing = settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["use"]
45 anderson_mixing_from_previous_iter = settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["mixing_from_previous_iter"]
46 # The pressure at the plasma boundary. Integrating `p'` leaves the pressure defined up to a constant, which
47 # GSFit historically fixed by demanding zero pressure at the boundary. Setting `free` makes it a degree of
48 # freedom of the fit instead. Fall back to the old location for existing settings directories.
49 pressure_edge_settings = settings["source_function_p_prime.json"].get(
50 "pressure_edge", settings["GSFIT_code_settings.json"]["numerics"].get("pressure_edge", {})
51 )
52 pressure_edge_free = pressure_edge_settings.get("free", False)
53 pressure_edge_regularisation_weight = pressure_edge_settings.get("regularisation_weight", 1.0e-4)
55 # Set the source functions types
56 p_prime_source_function = build_source_function(settings["source_function_p_prime.json"])
57 ff_prime_source_function = build_source_function(settings["source_function_ff_prime.json"])
59 # Grid size and shape
60 n_r = settings["GSFIT_code_settings.json"]["grid"]["n_r"]
61 n_z = settings["GSFIT_code_settings.json"]["grid"]["n_z"]
62 r_min = settings["GSFIT_code_settings.json"]["grid"]["r_min"]
63 r_max = settings["GSFIT_code_settings.json"]["grid"]["r_max"]
64 z_min = settings["GSFIT_code_settings.json"]["grid"]["z_min"]
65 z_max = settings["GSFIT_code_settings.json"]["grid"]["z_max"]
67 # Normalised poloidal flux grid
68 n_psi_n = settings["GSFIT_code_settings.json"]["n_psi_n"]
69 psi_n = np.linspace(0.0, 1.0, n_psi_n).astype(np.float64)
71 # Initialise the Plasma Rust class
72 plasma = Plasma(
73 n_r,
74 n_z,
75 r_min,
76 r_max,
77 z_min,
78 z_max,
79 psi_n,
80 p_prime_source_function,
81 ff_prime_source_function,
82 initial_guess_ip,
83 initial_guess_cur_r,
84 initial_guess_cur_z,
85 initial_guess_minor_radius,
86 initial_guess_elongation,
87 n_iter_max,
88 n_iter_min,
89 n_iter_no_vertical_feedback,
90 grad_shafranov_deviation_tolerance,
91 use_anderson_mixing,
92 anderson_mixing_from_previous_iter,
93 times_to_reconstruct,
94 pressure_edge_free,
95 pressure_edge_regularisation_weight,
96 )
98 return plasma
101def build_source_function(source_settings: dict[str, typing.Any]) -> EfitPolynomial | TensionedCubicBSpline:
102 """
103 Construct the Rust source function (`p_prime` or `ff_prime`) requested by the settings.
105 :param source_settings: Contents of either `source_function_p_prime.json` or `source_function_ff_prime.json`
106 """
108 if source_settings["method"] == "efit_polynomial":
109 n_dof = source_settings["efit_polynomial"]["n_dof"]
110 regularisations = np.array(source_settings["efit_polynomial"]["regularizations"])
111 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
112 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
113 if regularisations.shape == (1, 0):
114 regularisations = np.zeros((0, n_dof), dtype=np.float64)
115 return EfitPolynomial(n_dof, regularisations)
117 if source_settings["method"] == "tensioned_cubic_b_spline":
118 regularisations = np.array(source_settings["tensioned_cubic_b_spline"]["regularizations"])
119 interior_knots = np.array(source_settings["tensioned_cubic_b_spline"]["interior_knots"])
120 n_dof = len(interior_knots) + 4
121 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
122 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
123 if regularisations.shape == (1, 0):
124 regularisations = np.zeros((0, n_dof), dtype=np.float64)
125 interval_tensions = np.array(source_settings["tensioned_cubic_b_spline"]["interval_tensions"])
126 return TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
128 raise ValueError(f"Unknown source function method: {source_settings['method']}")