Coverage for src/pyhrp/plot.py: 100%

1"""Plotly visualization for hierarchical clustering dendrograms. 

2 

3Kept separate from the allocation core (``hrp.py``) so that the optional 

4plotting dependency (plotly) and the visualization logic do not couple the 

5algorithm modules. ``Dendrogram.plot`` delegates here. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING 

11 

12import plotly.graph_objects as go 

13import scipy.cluster.hierarchy as sch 

14 

15if TYPE_CHECKING: 

16 from .dendrogram import Dendrogram 

17 

18__all__ = ["plot_dendrogram"] 

19 

20 

21def plot_dendrogram(dendrogram: Dendrogram, **kwargs: object) -> go.Figure: 

22 """Build and return a plotly dendrogram figure for a Dendrogram. 

23 

24 Args: 

25 dendrogram (Dendrogram): Clustering result carrying the linkage matrix 

26 and asset labels to render. 

27 **kwargs (object): Extra keyword arguments forwarded to 

28 ``scipy.cluster.hierarchy.dendrogram`` (e.g. ``color_threshold``). 

29 

30 Returns: 

31 go.Figure: A plotly figure drawing the dendrogram as line traces. 

32 

33 Raises: 

34 ValueError: If the dendrogram has no linkage matrix to plot. 

35 

36 Examples: 

37 >>> import polars as pl 

38 >>> from pyhrp.hrp import build_tree 

39 >>> from pyhrp.plot import plot_dendrogram 

40 >>> cor = pl.DataFrame({"A": [1.0, 0.5], "B": [0.5, 1.0]}) 

41 >>> fig = plot_dendrogram(build_tree(cor, method="ward")) 

42 >>> type(fig).__name__ 

43 'Figure' 

44 """ 

45 if dendrogram.linkage is None: 

46 msg = "Dendrogram has no linkage matrix to plot." 

47 raise ValueError(msg) 

48 ddata = sch.dendrogram(dendrogram.linkage, labels=dendrogram.assets, no_plot=True, **kwargs) 

49 fig = go.Figure() 

50 for xs, ys in zip(ddata["icoord"], ddata["dcoord"], strict=False): 

51 fig.add_trace(go.Scatter(x=xs, y=ys, mode="lines", line={"color": "steelblue"}, showlegend=False)) 

52 n = len(dendrogram.assets) 

53 fig.update_layout( 

54 xaxis={ 

55 "tickmode": "array", 

56 "tickvals": [5 + 10 * i for i in range(n)], 

57 "ticktext": ddata["ivl"], 

58 "tickangle": -90, 

59 }, 

60 ) 

61 return fig