Added plot functionality

Author: ActurialCapital

examples/example_calendar_split.py

@@ -26,11 +26,11 @@
 
     # Example 1
     # ---------
-    
+
     model = CalendarSplit(
         data.index,
-        n_train = 3,
-        n_test = 1,
+        n_train=3,
+        n_test=1,
         every='YS',
         sample_labels=["IS", "OOS"]
     )
@@ -98,15 +98,15 @@
 
     # Example 2
     # ---------
-    
+
     model = CalendarSplit(
         data.index,
-        n_train = 4,
-        n_test = 2,
+        n_train=4,
+        n_test=2,
         every='QS',
         sample_labels=["IS", "OOS"]
     )
-    
+
     # Base
     print(model.splits_arr)
     print(model.splits_arr.dtype)

fold/model_selection/base.py

@@ -1,12 +1,17 @@
 import abc
+from dataclasses import asdict
+import typing as tp
 import numpy as np
 import pandas as pd
 from pandas.tseries.offsets import BaseOffset as Offset
-import typing as tp
+from plotly.basedatatypes import BaseFigure
+from plotly.graph_objects import FigureWidget
+import plotly.express as px
 
 from fold.model_selection.config import Config, UpdateConfig
 from fold.model_selection.duration import Duration
 from fold.model_selection.coverage import Coverage
+from fold.model_selection.plots import Layout, Trace, Heatmap
 from fold.tools import (
     substitute,
     BaseTool,
@@ -119,7 +124,7 @@ def __init__(
         self._ndim = 2 if len(self._sample_labels) > 1 else config.ndim
         self._index = index
         self._splits_arr = splits_arr
-        
+
         super().__init__(
             index=self.index,
             splits_arr=self.splits_arr,
@@ -128,7 +133,7 @@ def __init__(
             split_labels=self.split_labels,
             sample_labels=self.sample_labels
         )
-        
+
     @property
     def split_labels(self) -> pd.Index:
         """
@@ -587,6 +592,98 @@ def _get_range_meta(i, j):
 
         return pd.Series(range_objs, index=keys, dtype=object)
 
+    def plot(
+        self,
+        mask_kwargs: tp.Optional[tp.Dict[str, tp.Any]] = None,
+        trace_kwargs: tp.Optional[tp.Dict[str, tp.Any]] = None,
+        add_trace_kwargs: tp.Optional[tp.Dict[str, tp.Any]] = None,
+        fig: tp.Optional[tp.Any] = None,
+        figure_kwargs: tp.Optional[tp.Dict[str, tp.Any]] = None,
+        layout_kwargs: tp.Optional[Layout] = None,
+    ) -> BaseFigure:
+        """
+        Plot Flow graph of training vs test data.
+
+        Parameters
+        ----------
+        mask_kwargs :tp.Dict[str, tp.Any]
+            Keyword arguments passed to `Splitter.get_iter_set_masks`.
+        trace_kwargs : tp.Dict[str, tp.Any]
+            Keyword arguments passed to `plotly` Heatmap`.
+        add_trace_kwargs : tp.Dict[str, tp.Any] 
+            Keyword arguments passed to `add_trace`.
+        fig : FigureWidget 
+            Figure to add traces to.
+        **layout_kwargs 
+            Keyword arguments for layout.
+
+        Examples
+        --------
+        ```pycon
+        >>> from fold import SklearnFold
+        >>> import pandas as pd
+        >>> from sklearn.model_selection import TimeSeriesSplit
+
+        >>> index = pd.date_range("2010", "2024", freq="W")
+        >>> model = SklearnFold(index, TimeSeriesSplit())
+        >>> model.plot().show()
+        ```
+
+        """
+        trace_kwargs = trace_kwargs or {}
+        add_trace_kwargs = add_trace_kwargs or {}
+        figure_kwargs = figure_kwargs or {}
+        mask_kwargs = mask_kwargs or {}
+        layout_kwargs = asdict(Layout())
+
+        # Figure
+        if fig is None:
+            fig = FigureWidget(**figure_kwargs)
+
+        # Update layout
+        fig.update_layout(**layout_kwargs)
+
+        # Colors
+        colorway = (
+            fig.layout.colorway
+            if fig.layout.colorway is not None
+            else fig.layout.template.layout.colorway
+        )
+        if len(self.sample_labels) > len(colorway):
+            colorway = px.colors.qualitative.Alphabet
+
+        if self.n_splits > 0 and self.n_samples > 0:
+
+            # Loop through Coverage.get_iter_sample_masks
+            for i, mask in enumerate(self.get_iter_sample_masks()):
+                # Data
+                df = mask.ffill()
+                df[mask] = i
+                data = df.transpose().iloc[::-1]
+
+                # Settings
+                color = colorway[i % len(colorway)]
+                name = str(self.sample_labels[i])
+                default_trace = asdict(Trace()) | dict(
+                    x=data.columns,
+                    y=data.index,
+                    legendgroup=str(self.sample_labels[i]),
+                    name=name,
+                    colorscale=[color, color],
+                    hovertemplate="%{x}<br>Split: %{y}<br>Sample: " + name,
+                )
+
+                # Plot
+                fig = Heatmap(
+                    data=data,
+                    trace_kwargs=default_trace | trace_kwargs,
+                    add_trace_kwargs=add_trace_kwargs,
+                    is_y_category=True,
+                    fig=fig,
+                ).figure
+
+        return fig
+
 
 class BasePurgedCV:
     """
@@ -781,13 +878,13 @@ def split(
             pred_times = pd.Series(pred_times, index=X.index)
         else:
             checks.assert_instance_of(
-                pred_times, 
-                pd.Series, 
+                pred_times,
+                pd.Series,
                 arg_name="pred_times"
             )
             checks.assert_index_equal(
-                X.index, 
-                pred_times.index, 
+                X.index,
+                pred_times.index,
                 check_names=False
             )
         if eval_times is None:

fold/model_selection/plots.py

@@ -0,0 +1,148 @@
+import abc
+from dataclasses import dataclass, field
+import typing as tp
+
+import numpy as np
+
+import plotly.graph_objects as go
+from plotly.basedatatypes import BaseFigure, BaseTraceType
+
+
+class BasePlotter(abc.ABC):
+    def __init__(self, figure: BaseFigure, traces: tp.Tuple[BaseTraceType, ...]):
+        """Base trace updating class."""
+        self._figure = figure
+        self._traces = traces
+
+    @property
+    def figure(self) -> BaseFigure:
+        """Figure."""
+        return self._figure
+
+    @property
+    def traces(self) -> tp.Tuple[BaseTraceType, ...]:
+        """Traces to update."""
+        return self._traces
+
+    @classmethod
+    @abc.abstractmethod
+    def update_trace(
+        cls,
+        trace: BaseTraceType,
+        data: np.ndarray,
+        *args,
+        **kwargs
+    ):
+        """Update one trace."""
+        pass
+
+    @abc.abstractmethod
+    def update(self, *args, **kwargs):
+        """Update all traces using new data."""
+        pass
+
+
+@dataclass
+class Layout:
+    width: int = 700
+    height: int = 350
+    margin: dict = field(
+        default_factory=lambda: dict(t=30, b=30, l=30, r=30)
+    )
+    legend: dict = field(
+        default_factory=lambda: dict(
+            x=1,
+            y=1.02, orientation="h",
+            yanchor="bottom",
+            xanchor="right",
+            traceorder="normal"
+        )
+    )
+    # Scikit-learn colors
+    colorway: list = field(
+        default_factory=lambda: [
+            "#007AB8",
+            "#F7931E",
+            "#505050"
+        ]
+    )
+
+
+@dataclass
+class Trace:
+    hoverongaps: bool = False
+    showscale: bool = False
+    showlegend: bool = True
+
+
+class Heatmap(BasePlotter):
+    """
+    Create a heatmap plot.
+
+    Parameters
+    ----------
+    data : array_like, optional
+        Data in any format that can be converted to NumPy.
+    is_x_category : bool, optional
+        Whether X-axis is a categorical axis. Default is False.
+    is_y_category : bool, optional
+        Whether Y-axis is a categorical axis. Default is False.
+    trace_kwargs : dict, optional
+        Keyword arguments passed to `plotly.graph_objects.Heatmap`.
+    add_trace_kwargs : dict, optional
+        Keyword arguments passed to `add_trace`.
+    fig : Figure or FigureWidget, optional
+        Figure to add traces to.
+    **layout_kwargs : dict
+        Keyword arguments for layout.
+    """
+
+    def __init__(
+        self,
+        data: tp.Optional[np.ndarray],
+        is_x_category: tp.Optional[bool] = False,
+        is_y_category: tp.Optional[bool] = False,
+        trace_kwargs: tp.Dict[str, tp.Any] = None,
+        add_trace_kwargs: tp.Dict[str, tp.Any] = None,
+        figure_kwargs: tp.Dict[str, tp.Any] = None,
+        fig: tp.Optional[BaseFigure] = None,
+        **layout_kwargs
+    ):
+        trace = go.Heatmap(**trace_kwargs)
+
+        if data is not None:
+            self.update_trace(trace, data)
+            
+        fig.add_trace(trace, **add_trace_kwargs)
+
+        axis_kwargs = dict()
+        if is_x_category:
+            if fig.data[-1]["xaxis"] is not None:
+                axis_kwargs["xaxis" + fig.data[-1]["xaxis"][1:]] = dict(type="category")
+            else:
+                axis_kwargs["xaxis"] = dict(type="category")
+
+        if is_y_category:
+            if fig.data[-1]["yaxis"] is not None:
+                axis_kwargs["yaxis" + fig.data[-1]["yaxis"][1:]] = dict(type="category")
+            else:
+                axis_kwargs["yaxis"] = dict(type="category")
+
+        fig.update_layout(**axis_kwargs)
+        fig.update_layout(**layout_kwargs)
+
+        super().__init__(fig, (fig.data[-1],))
+
+    @classmethod
+    def update_trace(
+        cls,
+        trace: BaseTraceType,
+        data: np.ndarray,
+        *args,
+        **kwargs
+    ):
+        trace.z = data
+
+    def update(self, data: np.ndarray):
+        with self.fig.batch_update():
+            self.update_trace(self.traces[0], data)