Coverage for python/gsfit/database_readers/freegs/setup_plasma.py: 0%
55 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 13:12 +0000
1import typing
2from typing import TYPE_CHECKING
4import freegs
5import gsfit_rs
6import numpy as np
7import numpy.typing as npt
8from gsfit_rs import Plasma
10if TYPE_CHECKING:
11 from . import DatabaseReader
14def setup_plasma(
15 self: "DatabaseReader",
16 pulseNo: int,
17 settings: dict[str, typing.Any],
18 time: npt.NDArray[np.float64],
19 freegs_eqs: list[freegs.equilibrium.Equilibrium],
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 time: Measured time vector
27 :param freegs_eqs: List of FreeGS equilibrium objects, one for each time-slice
29 **This method is specific to FreeGS.**
31 See `python/gsfit/database_readers/interface.py` for more details on how a new database_reader should be implemented.
32 """
34 # Initial plasma conditions
35 initial_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"]
36 initial_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["r_cur"]
37 initial_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["z_cur"]
39 # Set the source functions types
40 p_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
41 ff_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
43 if settings["source_function_p_prime.json"]["method"] == "efit_polynomial":
44 n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"]
45 regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"])
46 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
47 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
48 if regularisations.shape == (1, 0):
49 regularisations = np.zeros((0, n_dof), dtype=np.float64)
50 p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
51 elif settings["source_function_p_prime.json"]["method"] == "tensioned_cubic_b_spline":
52 regularisations = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["regularizations"])
53 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
54 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
55 interior_knots = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interior_knots"])
56 n_dof = len(interior_knots) + 4
57 if regularisations.shape == (1, 0):
58 regularisations = np.zeros((0, n_dof), dtype=np.float64)
59 interval_tensions = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interval_tensions"])
60 p_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
61 else:
62 raise ValueError(f"Unknown method for p_prime source function: {settings['source_function_p_prime.json']['method']}")
64 if settings["source_function_ff_prime.json"]["method"] == "efit_polynomial":
65 n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"]
66 regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["regularizations"])
67 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
68 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
69 if regularisations.shape == (1, 0):
70 regularisations = np.zeros((0, n_dof), dtype=np.float64)
71 ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
72 elif settings["source_function_ff_prime.json"]["method"] == "tensioned_cubic_b_spline":
73 regularisations = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["regularizations"])
74 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
75 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
76 interior_knots = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["interior_knots"])
77 n_dof = len(interior_knots) + 4
78 if regularisations.shape == (1, 0):
79 regularisations = np.zeros((0, n_dof), dtype=np.float64)
80 interval_tensions = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["interval_tensions"])
81 ff_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
82 else:
83 raise ValueError(f"Unknown method for ff_prime source function: {settings['source_function_ff_prime.json']['method']}")
85 # Grid size and shape
86 n_r = settings["GSFIT_code_settings.json"]["grid"]["n_r"]
87 n_z = settings["GSFIT_code_settings.json"]["grid"]["n_z"]
88 r_min = settings["GSFIT_code_settings.json"]["grid"]["r_min"]
89 r_max = settings["GSFIT_code_settings.json"]["grid"]["r_max"]
90 z_min = settings["GSFIT_code_settings.json"]["grid"]["z_min"]
91 z_max = settings["GSFIT_code_settings.json"]["grid"]["z_max"]
93 # Normalised poloidal flux grid
94 n_psi_n = settings["GSFIT_code_settings.json"]["n_psi_n"]
95 psi_n = np.linspace(0.0, 1.0, n_psi_n).astype(np.float64)
97 # Limiter
98 limit_pts_r = freegs_eqs[0].tokamak.wall.R
99 limit_pts_z = freegs_eqs[0].tokamak.wall.Z
101 # Vacuum vessel where the plasma is allowed to be
102 vessel_r = limit_pts_r
103 vessel_z = limit_pts_z
105 # Initialise the Plasma Rust class
106 plasma = Plasma(
107 n_r,
108 n_z,
109 r_min,
110 r_max,
111 z_min,
112 z_max,
113 psi_n, # BUXTON: perhaps better to send in `n_psi_n`
114 limit_pts_r,
115 limit_pts_z,
116 vessel_r,
117 vessel_z,
118 p_prime_source_function,
119 ff_prime_source_function,
120 initial_ip,
121 initial_cur_r,
122 initial_cur_z,
123 )
125 return plasma