Coverage for tcvx21/plotting/plot_array_as_transparency_m.py: 88%

1""" 

2A helpful function which allows us to plot an array with the transparency set by the value 

3""" 

4import numpy as np 

5import matplotlib.pyplot as plt 

6from matplotlib.colors import Normalize 

7 

8def plot_array_as_transparency(ax, x: np.ndarray, y: np.ndarray, alphas, cmap=plt.cm.Greys, intensity=0.3, clip=True, invert=False): 

9 """ 

10 Displays a 2d field given by 'alphas' as a semi-transparent field, where the min-values are transparent and the 

11 max-values are opaque. 

12 

13 Particularly useful for adding a grey region around contour plots, which can be done with 

14 limits={'xmin': grid.r_s.min(), 'xmax': grid.r_s.max(), 'ymin': grid.z_s.min(), 'ymax':grid.z_s.max()} 

15 plot_array_as_transparency(axes, limits, alphas=np.isnan(shaped_data)) 

16 where shaped_data is the z-values that you pass to contour 

17 

18 invert switches the max and min transparencies 

19 clip=True will mean that masked/NaN alphas are assumed to have a value of '1.0' 

20 """ 

21 

22 limits = dict(xmin=x.min(), xmax=x.max(), ymin=y.min(), ymax=y.max()) 

23 

24 alphas = np.array(alphas, dtype=float) 

25 

26 # Normalize to the range [0, 1] 

27 alphas -= alphas.min() 

28 alphas /= alphas.max() 

29 if invert: 

30 alphas *= -1.0 

31 alphas += 1.0 

32 

33 assert alphas.min() == 0.0 and alphas.max() == 1.0 

34 

35 on_grid = np.ones_like(alphas) * intensity 

36 

37 colors = Normalize(0.0, 1.0, clip=clip)(on_grid) 

38 colors = cmap(colors) 

39 

40 colors[..., -1] = alphas 

41 

42 return ax.imshow(colors, extent=(limits['xmin'], limits['xmax'], limits['ymin'], limits['ymax']), origin='lower')