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

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 ax=None, 

32) -> matplotlib.figure.Figure: 

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

34 

35 Parameters 

36 ---------- 

37 x : array-like 

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

39 mid : array-like 

40 Mid values. 

41 lower : array-like 

42 Lower bound (bid). 

43 upper : array-like 

44 Upper bound (ask). 

45 xlabel, ylabel, title : str 

46 Axis labels and title. 

47 label : str, optional 

48 Legend label. 

49 color : str, optional 

50 Color for the series. 

51 ax : matplotlib Axes, optional 

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

53 

54 Returns: 

55 ------- 

56 matplotlib.figure.Figure 

57 """ 

58 _require_matplotlib() 

59 import matplotlib.pyplot as plt 

60 import numpy as np 

61 

62 x = np.asarray(x) 

63 mid = np.asarray(mid) 

64 lower = np.asarray(lower) 

65 upper = np.asarray(upper) 

66 

67 if ax is None: 

68 fig, ax = plt.subplots() 

69 else: 

70 fig = ax.get_figure() 

71 

72 yerr_lower = mid - lower 

73 yerr_upper = upper - mid 

74 

75 ax.errorbar(x, mid, yerr=[yerr_lower, yerr_upper], fmt="o-", capsize=3, label=label, color=color) 

76 

77 if xlabel: 

78 ax.set_xlabel(xlabel) 

79 if ylabel: 

80 ax.set_ylabel(ylabel) 

81 if title: 

82 ax.set_title(title) 

83 if label: 

84 ax.legend() 

85 ax.grid(True, alpha=0.3) 

86 

87 return fig