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

1import matplotlib.axes 

2import numpy as np 

3from gsfit_rs.imas import equilibrium_paths as ep 

4 

5from ..gsfit import Gsfit 

6 

7 

8def plot( 

9 gsfit_controller: Gsfit, 

10 ax: matplotlib.axes.Axes, 

11) -> None: 

12 plasma = gsfit_controller.plasma 

13 

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

15 r = plasma.get(ep.time_slice[0].profiles_2d[0].grid.dim1) 

16 z = plasma.get(ep.time_slice[0].profiles_2d[0].grid.dim2) 

17 n_r = len(r) 

18 n_z = len(z) 

19 

20 for i_z in range(n_z): 

21 ax.plot([r[0], r[-1]], [z[i_z], z[i_z]], color="black", linestyle="solid", linewidth=0.5) 

22 for i_r in range(n_r): 

23 ax.plot([r[i_r], r[i_r]], [z[0], z[-1]], color="black", linestyle="solid", linewidth=0.5) 

24 

25 # `profiles_2d/r` and `/z` are the (R, Z) mesh; flattening them row-major gives the flat grid 

26 flat_r = np.asarray(plasma.get(ep.time_slice[0].profiles_2d[0].r)).flatten() 

27 flat_z = np.asarray(plasma.get(ep.time_slice[0].profiles_2d[0].z)).flatten() 

28 ax.plot(flat_r, flat_z, color="black", linestyle="none", marker="o", markersize=0.5)