Coverage for python/gsfit/database_readers/freegsnke/setup_plasma.py: 64%
56 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 gsfit_rs
5import numpy as np
6import numpy.typing as npt
7from freegsnke.equilibrium_update import Equilibrium as FreeGsnkeEquilibrium
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 freegsnke_eqs: list[FreeGsnkeEquilibrium],
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
27 **This method is specific to FreeGSNKE.**
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_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"]
34 initial_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["r_cur"]
35 initial_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["z_cur"]
37 # Set the source functions types
38 p_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
39 ff_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline
41 # p_prime source function
42 if settings["source_function_p_prime.json"]["method"] == "efit_polynomial":
43 n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"]
44 regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"])
45 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
46 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
47 if regularisations.shape == (1, 0):
48 regularisations = np.zeros((0, n_dof), dtype=np.float64)
49 p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations)
50 elif settings["source_function_p_prime.json"]["method"] == "tensioned_cubic_b_spline":
51 regularisations = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["regularizations"])
52 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0).
53 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error
54 interior_knots = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interior_knots"])
55 n_dof = len(interior_knots) + 4
56 if regularisations.shape == (1, 0):
57 regularisations = np.zeros((0, n_dof), dtype=np.float64)
58 interval_tensions = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interval_tensions"])
59 p_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions)
60 else:
61 raise ValueError(f"Unknown method for p_prime source function: {settings['source_function_p_prime.json']['method']}")
63 # ff_prime source function
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 freegsnke_tokamak = freegsnke_eqs[0].tokamak # assume static data does not change in time
99 limit_pts_r = np.array(freegsnke_tokamak.limiter.R)
100 limit_pts_z = np.array(freegsnke_tokamak.limiter.Z)
102 # Vacuum vessel where the plasma is allowed to be
103 vessel_r = limit_pts_r
104 vessel_z = limit_pts_z
106 # Initialise the Plasma Rust class
107 plasma = Plasma(
108 n_r,
109 n_z,
110 r_min,
111 r_max,
112 z_min,
113 z_max,
114 psi_n, # BUXTON: perhaps better to send in `n_psi_n`
115 limit_pts_r,
116 limit_pts_z,
117 vessel_r,
118 vessel_z,
119 p_prime_source_function,
120 ff_prime_source_function,
121 initial_ip,
122 initial_cur_r,
123 initial_cur_z,
124 )
126 return plasma