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

1import matplotlib.axes 

2import numpy as np 

3import numpy.typing as npt 

4from gsfit_rs.imas import equilibrium_paths as ep 

5 

6from ..gsfit import Gsfit 

7 

8 

9def plot( 

10 gsfit_controller: Gsfit, 

11 ax: matplotlib.axes.Axes, 

12 time_desired: float, 

13 color: str = "blue", 

14 linestyle: str = "dashed", 

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

16) -> None: 

17 time = gsfit_controller.plasma.get(ep.time_slice[:].time) 

18 i_time = int(np.argmin(np.abs(time - time_desired))) 

19 

20 # `profiles_2d(0)` because GSFit solves on a single rectangular (R, Z) grid 

21 gsfit_r = gsfit_controller.plasma.get(ep.time_slice[i_time].profiles_2d[0].grid.dim1) 

22 gsfit_z = gsfit_controller.plasma.get(ep.time_slice[i_time].profiles_2d[0].grid.dim2) 

23 gsfit_psi = gsfit_controller.plasma.get(ep.time_slice[i_time].profiles_2d[0].psi) 

24 

25 # The boundary is stored at its own length for each time-slice, so there is no padding to trim 

26 gsfit_boundary_r = gsfit_controller.plasma.get(ep.time_slice[i_time].boundary.outline.r) 

27 gsfit_boundary_z = gsfit_controller.plasma.get(ep.time_slice[i_time].boundary.outline.z) 

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)