Coverage for python/gsfit/database_readers/st40_astra_mdsplus/setup_plasma.py: 0%
61 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
9if TYPE_CHECKING:
10 from . import DatabaseReader
13def setup_plasma(
14 self: "DatabaseReader",
15 pulseNo: int,
16 settings: dict[str, typing.Any],
17 times_to_reconstruct: npt.NDArray[np.float64],
18) -> Plasma:
19 """
20 This method initialises the Rust `Plasma` class.
22 :param pulseNo: Pulse number, used to read from the database
23 :param settings: Dictionary containing the JSON settings read from the `settings` directory
24 :param times_to_reconstruct: Times the equilibrium will be solved at [second]
26 **This method is specific to ST40's experimental MDSplus database.**
28 See `python/gsfit/database_readers/interface.py` for more details on how a new database_reader should be implemented.
29 """
31 # Initial plasma conditions
32 initial_guess_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"]
33 initial_guess_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_r"]
34 initial_guess_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_z"]
35 initial_guess_minor_radius = settings["GSFIT_code_settings.json"]["initial_guess"]["minor_radius"]
36 initial_guess_elongation = settings["GSFIT_code_settings.json"]["initial_guess"]["elongation"]
38 # Numerical settings the Grad-Shafranov solve is run with
39 n_iter_max = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_max"]
40 n_iter_min = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_min"]
41 n_iter_no_vertical_feedback = settings["GSFIT_code_settings.json"]["numerics"]["n_iter_no_vertical_feedback"]
42 grad_shafranov_deviation_tolerance = settings["GSFIT_code_settings.json"]["numerics"]["grad_shafranov_deviation_tolerance"]
43 use_anderson_mixing = settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["use"]
44 anderson_mixing_from_previous_iter = settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["mixing_from_previous_iter"]
45 # The pressure at the plasma boundary. Integrating `p'` leaves the pressure defined up to a constant, which
46 # GSFit historically fixed by demanding zero pressure at the boundary. Setting `free` makes it a degree of
47 # freedom of the fit instead. `.get` is used so that older settings directories keep working
48 pressure_edge_settings = settings["source_function_p_prime.json"].get(
49 "pressure_edge", settings["GSFIT_code_settings.json"]["numerics"].get("pressure_edge", {})
50 )
51 pressure_edge_free = pressure_edge_settings.get("free", False)
52 pressure_edge_regularisation_weight = pressure_edge_settings.get("regularisation_weight", 1.0e-4)
54 # Set the source functions types
55 p_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
56 ff_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
58 if settings["source_function_p_prime.json"]["method"] == "efit_polynomial":
59 n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"]
60 regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"])
61 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
62 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
63 if regularisations.shape == (1, 0):
64 regularisations = np.zeros((0, n_dof), dtype=np.float64)
65 p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
66 elif settings["source_function_p_prime.json"]["method"] == "tensioned_cubic_b_spline":
67 regularisations = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["regularizations"])
68 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
69 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
70 interior_knots = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interior_knots"])
71 n_dof = len(interior_knots) + 4
72 if regularisations.shape == (1, 0):
73 regularisations = np.zeros((0, n_dof), dtype=np.float64)
74 interval_tensions = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interval_tensions"])
75 p_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
76 else:
77 raise ValueError(f"Unknown method for p_prime source function: {settings['source_function_p_prime.json']['method']}")
79 if settings["source_function_ff_prime.json"]["method"] == "efit_polynomial":
80 n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"]
81 regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["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 ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
87 elif settings["source_function_ff_prime.json"]["method"] == "tensioned_cubic_b_spline":
88 regularisations = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["regularizations"])
89 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
90 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
91 interior_knots = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["interior_knots"])
92 n_dof = len(interior_knots) + 4
93 if regularisations.shape == (1, 0):
94 regularisations = np.zeros((0, n_dof), dtype=np.float64)
95 interval_tensions = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["interval_tensions"])
96 ff_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
97 else:
98 raise ValueError(f"Unknown method for ff_prime source function: {settings['source_function_ff_prime.json']['method']}")
100 # Grid size and shape
101 n_r = settings["GSFIT_code_settings.json"]["grid"]["n_r"]
102 n_z = settings["GSFIT_code_settings.json"]["grid"]["n_z"]
103 r_min = settings["GSFIT_code_settings.json"]["grid"]["r_min"]
104 r_max = settings["GSFIT_code_settings.json"]["grid"]["r_max"]
105 z_min = settings["GSFIT_code_settings.json"]["grid"]["z_min"]
106 z_max = settings["GSFIT_code_settings.json"]["grid"]["z_max"]
108 # Normalised poloidal flux grid
109 n_psi_n = settings["GSFIT_code_settings.json"]["n_psi_n"]
110 psi_n = np.linspace(0.0, 1.0, n_psi_n).astype(np.float64)
112 # Initialise the Plasma Rust class
113 plasma = Plasma(
114 n_r,
115 n_z,
116 r_min,
117 r_max,
118 z_min,
119 z_max,
120 psi_n, # BUXTON: perhaps better to send in `n_psi_n`
121 p_prime_source_function,
122 ff_prime_source_function,
123 initial_guess_ip,
124 initial_guess_cur_r,
125 initial_guess_cur_z,
126 initial_guess_minor_radius,
127 initial_guess_elongation,
128 n_iter_max,
129 n_iter_min,
130 n_iter_no_vertical_feedback,
131 grad_shafranov_deviation_tolerance,
132 use_anderson_mixing,
133 anderson_mixing_from_previous_iter,
134 times_to_reconstruct,
135 pressure_edge_free,
136 pressure_edge_regularisation_weight,
137 )
139 return plasma
142# import typing
143# from typing import TYPE_CHECKING
145# import gsfit_rs
146# import numpy as np
147# import numpy.typing as npt
148# from gsfit_rs import Plasma
149# from st40_database import GetData # type: ignore[import-not-found]
151# if TYPE_CHECKING:
152# from . import DatabaseReader
155# def setup_plasma(
156# self: "DatabaseReader",
157# pulseNo: int,
158# settings: dict[str, typing.Any],
159# ) -> Plasma:
160# """
161# This method initialises the Rust `Plasma` class.
163# :param pulseNo: Pulse number, used to read from the database
164# :param settings: Dictionary containing the JSON settings read from the `settings` directory
166# **This method is specific to ST40's ASTRA stored on MDSplus.**
168# See `python/gsfit/database_readers/interface.py` for more details on how a new database_reader should be implemented.
169# """
171# # Initial plasma conditions
172# initial_guess_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"]
173# initial_guess_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_r"]
174# initial_guess_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["cur_z"]
176# # Set the source functions types
177# p_prime_source_function: gsfit_rs.EfitPolynomial
178# ff_prime_source_function: gsfit_rs.EfitPolynomial
180# if settings["source_function_p_prime.json"]["method"] == "efit_polynomial":
181# n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"]
182# regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"])
183# p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
184# elif settings["source_function_p_prime.json"]["method"] == "efit_polynomial":
185# n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"]
186# regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"])
187# p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
188# else:
189# raise ValueError(f"Unknown method for p_prime source function: {settings['source_function_p_prime.json']['method']}")
191# if settings["source_function_ff_prime.json"]["method"] == "efit_polynomial":
192# n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"]
193# regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["regularizations"])
194# ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
195# elif settings["source_function_ff_prime.json"]["method"] == "efit_polynomial":
196# n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"]
197# regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["regularizations"])
198# ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
200# # Grid size and shape
201# n_r = settings["GSFIT_code_settings.json"]["grid"]["n_r"]
202# n_z = settings["GSFIT_code_settings.json"]["grid"]["n_z"]
203# r_min = settings["GSFIT_code_settings.json"]["grid"]["r_min"]
204# r_max = settings["GSFIT_code_settings.json"]["grid"]["r_max"]
205# z_min = settings["GSFIT_code_settings.json"]["grid"]["z_min"]
206# z_max = settings["GSFIT_code_settings.json"]["grid"]["z_max"]
208# # Normalised poloidal flux grid
209# n_psi_n = settings["GSFIT_code_settings.json"]["n_psi_n"]
210# psi_n = np.linspace(0.0, 1.0, n_psi_n).astype(np.float64)
212# # Limiter
213# elmag_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_astra_mdsplus"]["elmag_run_name"]
214# elmag = GetData(12050, f"ELMAG#{elmag_run_name}", is_fail_quiet=False)
215# limit_pts_r = typing.cast(npt.NDArray[np.float64], elmag.get("LIMITER.LIMIT_PTS.R"))
216# limit_pts_z = typing.cast(npt.NDArray[np.float64], elmag.get("LIMITER.LIMIT_PTS.Z"))
218# # Vacuum vessel where the plasma is allowed to be
219# vessel_r = limit_pts_r
220# vessel_z = limit_pts_z
222# # Add lower MC tiles
223# limit_pts_r = np.append(limit_pts_r, 0.7103)
224# limit_pts_z = np.append(limit_pts_z, -0.3131)
225# # Add upper MC tiles
226# limit_pts_r = np.append(limit_pts_r, 0.7103)
227# limit_pts_z = np.append(limit_pts_z, 0.3031)
229# # Initialise the Plasma Rust class
230# plasma = Plasma(
231# n_r,
232# n_z,
233# r_min,
234# r_max,
235# z_min,
236# z_max,
237# psi_n, # BUXTON: perhaps better to send in `n_psi_n`
238# limit_pts_r,
239# limit_pts_z,
240# vessel_r,
241# vessel_z,
242# p_prime_source_function,
243# ff_prime_source_function,
244# initial_guess_ip,
245# initial_guess_cur_r,
246# initial_guess_cur_z,
247# )
249# return plasma