Coverage for python/gsfit/database_readers/st40_mdsplus/setup_plasma.py: 0%
68 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 gsfit_rs
5import numpy as np
6import numpy.typing as npt
7from gsfit_rs import Plasma
9from ...tensioned_cubic_splines_regularisations import make_tensioned_cubic_b_spline_regularisations
11if TYPE_CHECKING:
12 from . import DatabaseReader
15def setup_plasma(
16 self: "DatabaseReader",
17 pulseNo: int,
18 settings: dict[str, typing.Any],
19 times_to_reconstruct: npt.NDArray[np.float64],
20) -> Plasma:
21 """
22 This method initialises the Rust `Plasma` class.
24 :param pulseNo: Pulse number, used to read from the database
25 :param settings: Dictionary containing the JSON settings read from the `settings` directory
26 :param times_to_reconstruct: Times the equilibrium will be solved at [second]
28 **This method is specific to ST40's experimental MDSplus database.**
30 See `python/gsfit/database_readers/interface.py` for more details on how a new database_reader should be implemented.
31 """
33 # Initial plasma conditions
34 initial_guess_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"]
35 initial_guess_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_r"]
36 initial_guess_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_z"]
37 initial_guess_minor_radius = settings["GSFIT_code_settings.json"]["initial_guess"]["minor_radius"]
38 initial_guess_elongation = settings["GSFIT_code_settings.json"]["initial_guess"]["elongation"]
40 # Numerical settings the Grad-Shafranov solve is run with
41 n_iter_max = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_max"]
42 n_iter_min = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_min"]
43 n_iter_no_vertical_feedback = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_no_vertical_feedback"]
44 grad_shafranov_deviation_tolerance = settings["GSFIT_code_settings.json"]["numerics"]["grad_shafranov_deviation_tolerance"]
45 use_anderson_mixing = settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["use"]
46 anderson_mixing_from_previous_iter = settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["mixing_from_previous_iter"]
47 # The pressure at the plasma boundary. Integrating `p'` leaves the pressure defined up to a constant, which
48 # GSFit historically fixed by demanding zero pressure at the boundary. Setting `free` makes it a degree of
49 # freedom of the fit instead. `.get` is used so that older settings directories keep working
50 pressure_edge_settings = settings["source_function_p_prime.json"].get(
51 "pressure_edge", settings["GSFIT_code_settings.json"]["numerics"].get("pressure_edge", {})
52 )
53 pressure_edge_free = pressure_edge_settings.get("free", False)
54 pressure_edge_regularisation_weight = pressure_edge_settings.get("regularisation_weight", 1.0e-4)
56 # Set the source functions types
57 p_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
58 ff_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
60 if settings["source_function_p_prime.json"]["method"] == "efit_polynomial":
61 n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"]
62 regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"])
63 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
64 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
65 if regularisations.shape == (1, 0):
66 regularisations = np.zeros((0, n_dof), dtype=np.float64)
67 p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
68 elif settings["source_function_p_prime.json"]["method"] == "tensioned_cubic_b_spline":
69 p_prime_tensioned_cubic_b_spline_settings = settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]
70 interior_knots = np.array(p_prime_tensioned_cubic_b_spline_settings["interior_knots"])
71 interval_tensions = np.array(p_prime_tensioned_cubic_b_spline_settings["interval_tensions"])
72 n_dof = len(interior_knots) + 4
73 if p_prime_tensioned_cubic_b_spline_settings.get("make_cubic_spline_regularisations", False):
74 # Build the regularisations from settings (see `tensioned_cubic_splines_regularisations.py`)
75 regularisations = make_tensioned_cubic_b_spline_regularisations(
76 interior_knots,
77 interval_tensions,
78 p_prime_tensioned_cubic_b_spline_settings["regularisation_builder"],
79 )
80 else:
81 regularisations = np.array(p_prime_tensioned_cubic_b_spline_settings["regularizations"])
82 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
83 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
84 if regularisations.shape == (1, 0):
85 regularisations = np.zeros((0, n_dof), dtype=np.float64)
86 p_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
87 else:
88 raise ValueError(f"Unknown method for p_prime source function: {settings['source_function_p_prime.json']['method']}")
90 if settings["source_function_ff_prime.json"]["method"] == "efit_polynomial":
91 n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"]
92 regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["regularizations"])
93 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
94 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
95 if regularisations.shape == (1, 0):
96 regularisations = np.zeros((0, n_dof), dtype=np.float64)
97 ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
98 elif settings["source_function_ff_prime.json"]["method"] == "tensioned_cubic_b_spline":
99 ff_prime_tensioned_cubic_b_spline_settings = settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]
100 interior_knots = np.array(ff_prime_tensioned_cubic_b_spline_settings["interior_knots"])
101 interval_tensions = np.array(ff_prime_tensioned_cubic_b_spline_settings["interval_tensions"])
102 n_dof = len(interior_knots) + 4
103 if ff_prime_tensioned_cubic_b_spline_settings.get("make_cubic_spline_regularisations", False):
104 # Build the regularisations from settings (see `tensioned_cubic_splines_regularisations.py`)
105 regularisations = make_tensioned_cubic_b_spline_regularisations(
106 interior_knots,
107 interval_tensions,
108 ff_prime_tensioned_cubic_b_spline_settings["regularisation_builder"],
109 )
110 else:
111 regularisations = np.array(ff_prime_tensioned_cubic_b_spline_settings["regularizations"])
112 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
113 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
114 if regularisations.shape == (1, 0):
115 regularisations = np.zeros((0, n_dof), dtype=np.float64)
116 ff_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
117 else:
118 raise ValueError(f"Unknown method for ff_prime source function: {settings['source_function_ff_prime.json']['method']}")
120 # Grid size and shape
121 n_r = settings["GSFIT_code_settings.json"]["grid"]["n_r"]
122 n_z = settings["GSFIT_code_settings.json"]["grid"]["n_z"]
123 r_min = settings["GSFIT_code_settings.json"]["grid"]["r_min"]
124 r_max = settings["GSFIT_code_settings.json"]["grid"]["r_max"]
125 z_min = settings["GSFIT_code_settings.json"]["grid"]["z_min"]
126 z_max = settings["GSFIT_code_settings.json"]["grid"]["z_max"]
128 # Normalised poloidal flux grid
129 n_psi_n = settings["GSFIT_code_settings.json"]["n_psi_n"]
130 psi_n = np.linspace(0.0, 1.0, n_psi_n).astype(np.float64)
132 # Initialise the Plasma Rust class
133 plasma = Plasma(
134 n_r,
135 n_z,
136 r_min,
137 r_max,
138 z_min,
139 z_max,
140 psi_n, # BUXTON: perhaps better to send in `n_psi_n`
141 p_prime_source_function,
142 ff_prime_source_function,
143 initial_guess_ip,
144 initial_guess_cur_r,
145 initial_guess_cur_z,
146 initial_guess_minor_radius,
147 initial_guess_elongation,
148 n_iter_max,
149 n_iter_min,
150 n_iter_no_vertical_feedback,
151 grad_shafranov_deviation_tolerance,
152 use_anderson_mixing,
153 anderson_mixing_from_previous_iter,
154 times_to_reconstruct,
155 pressure_edge_free,
156 pressure_edge_regularisation_weight,
157 )
159 return plasma