Coverage for src / qsmile / core / plot.py: 96%

1"""Plotting utilities for option chain representations.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7if TYPE_CHECKING: 

8 import matplotlib.figure 

9 

10 

11def _require_matplotlib(): 

12 """Import matplotlib or raise a helpful error.""" 

13 try: 

14 import matplotlib as mpl # noqa: F401 

15 except ImportError: 

16 msg = "matplotlib is required for plotting. Install it with: pip install qsmile[plot]" 

17 raise ImportError(msg) from None 

18 

19 

20def plot_bid_ask( 

21 x, 

22 mid, 

23 lower, 

24 upper, 

25 *, 

26 xlabel: str = "", 

27 ylabel: str = "", 

28 title: str = "", 

29 label: str | None = None, 

30 color: str | None = None, 

31 fmt: str = "none", 

32 ax=None, 

33 **kwargs, 

34) -> matplotlib.figure.Figure: 

35 """Plot bid/ask as error bars around mid values. 

36 

37 Parameters 

38 ---------- 

39 x : array-like 

40 X-axis values (e.g., strikes or unitised k). 

41 mid : array-like 

42 Mid values. 

43 lower : array-like 

44 Lower bound (bid). 

45 upper : array-like 

46 Upper bound (ask). 

47 xlabel, ylabel, title : str 

48 Axis labels and title. 

49 label : str, optional 

50 Legend label. 

51 color : str, optional 

52 Color for the series. 

53 ax : matplotlib Axes, optional 

54 Axes to plot on. If None, creates a new figure. 

55 

56 Returns: 

57 ------- 

58 matplotlib.figure.Figure 

59 """ 

60 _require_matplotlib() 

61 import matplotlib.pyplot as plt 

62 import numpy as np 

63 

64 x = np.asarray(x) 

65 mid = np.asarray(mid) 

66 lower = np.asarray(lower) 

67 upper = np.asarray(upper) 

68 

69 if ax is None: 

70 fig, ax = plt.subplots() 

71 else: 

72 fig = ax.get_figure() 

73 

74 yerr_lower = mid - lower 

75 yerr_upper = upper - mid 

76 

77 ax.errorbar(x, mid, yerr=[yerr_lower, yerr_upper], fmt=fmt, capsize=3, label=label, color=color, **kwargs) 

78 

79 if xlabel: 

80 ax.set_xlabel(xlabel) 

81 if ylabel: 

82 ax.set_ylabel(ylabel) 

83 if title: 

84 ax.set_title(title) 

85 if label: 

86 ax.legend() 

87 ax.grid(True, alpha=0.3) 

88 

89 return fig 

90 

91 

92def plot_line( 

93 x, 

94 y, 

95 *, 

96 xlabel: str = "", 

97 ylabel: str = "", 

98 title: str = "", 

99 label: str | None = None, 

100 color: str | None = None, 

101 ax=None, 

102 **kwargs, 

103) -> matplotlib.figure.Figure: 

104 """Plot a single curve. 

105 

106 Parameters 

107 ---------- 

108 x : array-like 

109 X-axis values. 

110 y : array-like 

111 Y-axis values. 

112 xlabel, ylabel, title : str 

113 Axis labels and title. 

114 label : str, optional 

115 Legend label. 

116 color : str, optional 

117 Color for the line. 

118 ax : matplotlib Axes, optional 

119 Axes to plot on. If None, creates a new figure. 

120 

121 Returns: 

122 ------- 

123 matplotlib.figure.Figure 

124 """ 

125 _require_matplotlib() 

126 import matplotlib.pyplot as plt 

127 import numpy as np 

128 

129 x = np.asarray(x) 

130 y = np.asarray(y) 

131 

132 if ax is None: 

133 fig, ax = plt.subplots() 

134 else: 

135 fig = ax.get_figure() 

136 

137 ax.plot(x, y, label=label, color=color, **kwargs) 

138 

139 if xlabel: 

140 ax.set_xlabel(xlabel) 

141 if ylabel: 

142 ax.set_ylabel(ylabel) 

143 if title: 

144 ax.set_title(title) 

145 if label: 

146 ax.legend() 

147 ax.grid(True, alpha=0.3) 

148 

149 return fig