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