Coverage for python/gsfit/plotting/two_d_plasma.py: 0%

1import matplotlib.axes 

2import numpy as np 

3import numpy.typing as npt 

4 

5from ..gsfit import Gsfit 

6 

7 

8def plot( 

9 gsfit_controller: Gsfit, 

10 ax: matplotlib.axes.Axes, 

11 time_desired: float, 

12 color: str = "blue", 

13 linestyle: str = "dashed", 

14 psi_n_levels: npt.NDArray[np.float64] | None = None, 

15) -> None: 

16 plasma = gsfit_controller.plasma 

17 

18 time = plasma.get_array1(["time"]) 

19 i_time = np.argmin(np.abs(time - time_desired)) 

20 

21 gsfit_r = plasma.get_array1(["grid", "r"]) 

22 gsfit_z = plasma.get_array1(["grid", "z"]) 

23 gsfit_psi = plasma.get_array3(["two_d", "psi"])[i_time, :, :] 

24 

25 gsfit_nbnd = plasma.get_vec_usize(["p_boundary", "nbnd"])[i_time] 

26 gsfit_boundary_r = plasma.get_array2(["p_boundary", "rbnd"])[i_time, :gsfit_nbnd] 

27 gsfit_boundary_z = plasma.get_array2(["p_boundary", "zbnd"])[i_time, :gsfit_nbnd] 

28 

29 # Default to 35 levels if not provided 

30 if psi_n_levels is None: 

31 psi_n_levels = np.linspace(np.min(gsfit_psi), np.max(gsfit_psi), 35) 

32 

33 ax.contour(gsfit_r, gsfit_z, gsfit_psi, levels=psi_n_levels, colors=color, linewidths=0.6, linestyles=linestyle) 

34 ax.plot(gsfit_boundary_r, gsfit_boundary_z, color=color, linestyle=linestyle, label="GSFit", linewidth=1.0)