Coverage for src/mafw_tools/plotting_tools.py: 100%

1# Copyright 2025 European Union

2# Author: Bulgheroni Antonio (antonio.bulgheroni@ec.europa.eu)

3# SPDX-License-Identifier: EUPL-1.2

4"""

5 Plotting tools module

6

7 This module provides utility functions for creating plots and visualizations

8 using matplotlib. It includes functionality for plotting images with various

9 customization options such as colorbars, axis control, and title management.

10

11.. versionadded:: 1.0.0

12"""

13

14 import matplotlib

15 from matplotlib import pyplot as plt

16 from mpl_toolkits.axes_grid1 import make_axes_locatable

17

18 def plot_image(

19 img,

20 ax: plt.Axes,

21 title: str = '',

22 preserve_axis_limits: bool = False,

23 axis_off: bool = False,

24 attach_colorbar: bool = False,

25 cbar_ticks: list = None,

26 **kwargs,

27) -> matplotlib.image.AxesImage:

28 """

29 Plot an image on a matplotlib axis with various customization options.

30

31 This function provides a flexible way to display images with options for

32 controlling axis appearance, adding colorbars, and preserving axis limits.

33

34 :param img: The image data to be plotted

35 :type img: numpy.ndarray

36 :param ax: The matplotlib axis object where the image will be displayed

37 :type ax: matplotlib.pyplot.Axes

38 :param title: Title to be displayed on the axis. Defaults to empty string

39 :type title: str

40 :param preserve_axis_limits: If True, preserves the current axis limits. Defaults to False

41 :type preserve_axis_limits: bool

42 :param axis_off: If True, turns off the axis ticks and labels. Defaults to False

43 :type axis_off: bool

44 :param attach_colorbar: If True, attaches a colorbar to the plot. Defaults to False

45 :type attach_colorbar: bool

46 :param cbar_ticks: List of tick locations for the colorbar. Defaults to None

47 :type cbar_ticks: list

48 :param kwargs: Additional keyword arguments passed to matplotlib's imshow function

49 :return: The matplotlib image object created

50 :rtype: matplotlib.image.AxesImage

51 :example:

52

53 >>> import matplotlib.pyplot as plt

54 >>> import numpy as np

55 >>> fig, ax = plt.subplots()

56 >>> data = np.random.rand(10, 10)

57 >>> plot_image(data, ax, title="Sample Image", attach_colorbar=True)

58 >>> plt.show()

59 """

60 if preserve_axis_limits:

61 x_lim, y_lim = ax.get_xlim(), ax.get_ylim()

62 ax.cla()

63 else:

64 x_lim = y_lim = (0, 1)

65 ax.set_xticks([])

66 ax.set_yticks([])

67 if axis_off:

68 ax.axis('off')

69 ax.set_title(title)

70 map = ax.imshow(img, **kwargs)

71 if preserve_axis_limits and x_lim != (0, 1) and y_lim != (0, 1):

72 ax.set_xlim(*x_lim)

73 ax.set_ylim(*y_lim)

74

75 if attach_colorbar:

76 # Use the make_axes_locatable to create a new axes for the colorbar

77 divider = make_axes_locatable(ax)

78 cax = divider.append_axes('right', size='5%', pad=0.05)

79

80 # Add the colorbar to the new axes

81 cbar = plt.colorbar(map, cax=cax)

82 if cbar_ticks is not None:

83 cbar.set_ticks(cbar_ticks)

84 cbar.set_ticklabels([f'{i}' for i in cbar_ticks])

85

86 return map