Merge pull request statsmodels#8783 from aglebov/plot-acf-funcs

ENH: Plot cross-correlations and auto/cross-correlation matrix

Author: josef-pkt

statsmodels/graphics/tests/test_tsaplots.py

@@ -12,7 +12,9 @@
 from statsmodels.datasets import elnino, macrodata
 from statsmodels.graphics.tsaplots import (
     month_plot,
+    plot_accf_grid,
     plot_acf,
+    plot_ccf,
     plot_pacf,
     plot_predict,
     quarter_plot,
@@ -190,6 +192,55 @@ def test_plot_pacf_irregular(close_figures):
     plot_pacf(pacf, ax=ax, alpha=None, zero=False)
 
 
+@pytest.mark.matplotlib
+def test_plot_ccf(close_figures):
+    # Just test that it runs.
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    ar = np.r_[1.0, -0.9]
+    ma = np.r_[1.0, 0.9]
+    armaprocess = tsp.ArmaProcess(ar, ma)
+    rs = np.random.RandomState(1234)
+    x1 = armaprocess.generate_sample(100, distrvs=rs.standard_normal)
+    x2 = armaprocess.generate_sample(100, distrvs=rs.standard_normal)
+    plot_ccf(x1, x2)
+    plot_ccf(x1, x2, ax=ax, lags=10)
+    plot_ccf(x1, x2, ax=ax)
+    plot_ccf(x1, x2, ax=ax, alpha=None)
+    plot_ccf(x1, x2, ax=ax, negative_lags=True)
+    plot_ccf(x1, x2, ax=ax, adjusted=True)
+    plot_ccf(x1, x2, ax=ax, fft=True)
+    plot_ccf(x1, x2, ax=ax, title='CCF')
+    plot_ccf(x1, x2, ax=ax, auto_ylims=True)
+    plot_ccf(x1, x2, ax=ax, use_vlines=False)
+
+
+@pytest.mark.matplotlib
+def test_plot_accf_grid(close_figures):
+    # Just test that it runs.
+    fig = plt.figure()
+
+    ar = np.r_[1.0, -0.9]
+    ma = np.r_[1.0, 0.9]
+    armaprocess = tsp.ArmaProcess(ar, ma)
+    rs = np.random.RandomState(1234)
+    x = np.vstack([
+        armaprocess.generate_sample(100, distrvs=rs.standard_normal),
+        armaprocess.generate_sample(100, distrvs=rs.standard_normal),
+    ]).T
+    plot_accf_grid(x)
+    plot_accf_grid(pd.DataFrame({'x': x[:, 0], 'y': x[:, 1]}))
+    plot_accf_grid(x, fig=fig, lags=10)
+    plot_accf_grid(x, fig=fig)
+    plot_accf_grid(x, fig=fig, negative_lags=False)
+    plot_accf_grid(x, fig=fig, alpha=None)
+    plot_accf_grid(x, fig=fig, adjusted=True)
+    plot_accf_grid(x, fig=fig, fft=True)
+    plot_accf_grid(x, fig=fig, auto_ylims=True)
+    plot_accf_grid(x, fig=fig, use_vlines=False)
+
+
 @pytest.mark.matplotlib
 def test_plot_month(close_figures):
     dta = elnino.load_pandas().data

statsmodels/graphics/tsaplots.py

@@ -7,7 +7,8 @@
 import pandas as pd
 
 from statsmodels.graphics import utils
-from statsmodels.tsa.stattools import acf, pacf
+from statsmodels.tools.validation import array_like
+from statsmodels.tsa.stattools import acf, pacf, ccf
 
 
 def _prepare_data_corr_plot(x, lags, zero):
@@ -38,6 +39,7 @@ def _plot_corr(
     use_vlines,
     vlines_kwargs,
     auto_ylims=False,
+    skip_lag0_confint=True,
     **kwargs,
 ):
     if irregular:
@@ -66,13 +68,13 @@ def _plot_corr(
         )
 
     if confint is not None:
-        if lags[0] == 0:
+        if skip_lag0_confint and lags[0] == 0:
             lags = lags[1:]
             confint = confint[1:]
             acf_x = acf_x[1:]
         lags = lags.astype(float)
-        lags[0] -= 0.5
-        lags[-1] += 0.5
+        lags[np.argmin(lags)] -= 0.5
+        lags[np.argmax(lags)] += 0.5
         ax.fill_between(
             lags, confint[:, 0] - acf_x, confint[:, 1] - acf_x, alpha=0.25
         )
@@ -367,6 +369,279 @@ def plot_pacf(
     return fig
 
 
+def plot_ccf(
+        x,
+        y,
+        *,
+        ax=None,
+        lags=None,
+        negative_lags=False,
+        alpha=0.05,
+        use_vlines=True,
+        adjusted=False,
+        fft=False,
+        title="Cross-correlation",
+        auto_ylims=False,
+        vlines_kwargs=None,
+        **kwargs,
+):
+    """
+    Plot the cross-correlation function
+
+    Correlations between ``x`` and the lags of ``y`` are calculated.
+
+    The lags are shown on the horizontal axis and the correlations
+    on the vertical axis.
+
+    Parameters
+    ----------
+    x, y : array_like
+        Arrays of time-series values.
+    ax : AxesSubplot, optional
+        If given, this subplot is used to plot in, otherwise a new figure with
+        one subplot is created.
+    lags : {int, array_like}, optional
+        An int or array of lag values, used on the horizontal axis. Uses
+        ``np.arange(lags)`` when lags is an int.  If not provided,
+        ``lags=np.arange(len(corr))`` is used.
+    negative_lags: bool, optional
+        If True, negative lags are shown on the horizontal axis.
+    alpha : scalar, optional
+        If a number is given, the confidence intervals for the given level are
+        plotted, e.g. if alpha=.05, 95 % confidence intervals are shown.
+        If None, confidence intervals are not shown on the plot.
+    use_vlines : bool, optional
+        If True, shows vertical lines and markers for the correlation values.
+        If False, only shows markers.  The default marker is 'o'; it can
+        be overridden with a ``marker`` kwarg.
+    adjusted : bool
+        If True, then denominators for cross-correlations are n-k, otherwise n.
+    fft : bool, optional
+        If True, computes the CCF via FFT.
+    title : str, optional
+        Title to place on plot. Default is 'Cross-correlation'.
+    auto_ylims : bool, optional
+        If True, adjusts automatically the vertical axis limits to CCF values.
+    vlines_kwargs : dict, optional
+        Optional dictionary of keyword arguments that are passed to vlines.
+    **kwargs : kwargs, optional
+        Optional keyword arguments that are directly passed on to the
+        Matplotlib ``plot`` and ``axhline`` functions.
+
+    Returns
+    -------
+    Figure
+        The figure where the plot is drawn. This is either an existing figure
+        if the `ax` argument is provided, or a newly created figure
+        if `ax` is None.
+
+    See Also
+    --------
+    See notes and references for statsmodels.graphics.tsaplots.plot_acf
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> import matplotlib.pyplot as plt
+    >>> import statsmodels.api as sm
+
+    >>> dta = sm.datasets.macrodata.load_pandas().data
+    >>> diffed = dta.diff().dropna()
+    >>> sm.graphics.tsa.plot_ccf(diffed["unemp"], diffed["infl"])
+    >>> plt.show()
+    """
+    fig, ax = utils.create_mpl_ax(ax)
+
+    lags, nlags, irregular = _prepare_data_corr_plot(x, lags, True)
+    vlines_kwargs = {} if vlines_kwargs is None else vlines_kwargs
+
+    if negative_lags:
+        lags = -lags
+
+    ccf_res = ccf(
+        x, y, adjusted=adjusted, fft=fft, alpha=alpha, nlags=nlags + 1
+    )
+    if alpha is not None:
+        ccf_xy, confint = ccf_res
+    else:
+        ccf_xy = ccf_res
+        confint = None
+
+    _plot_corr(
+        ax,
+        title,
+        ccf_xy,
+        confint,
+        lags,
+        irregular,
+        use_vlines,
+        vlines_kwargs,
+        auto_ylims=auto_ylims,
+        skip_lag0_confint=False,
+        **kwargs,
+    )
+
+    return fig
+
+
+def plot_accf_grid(
+        x,
+        *,
+        varnames=None,
+        fig=None,
+        lags=None,
+        negative_lags=True,
+        alpha=0.05,
+        use_vlines=True,
+        adjusted=False,
+        fft=False,
+        missing="none",
+        zero=True,
+        auto_ylims=False,
+        bartlett_confint=False,
+        vlines_kwargs=None,
+        **kwargs,
+):
+    """
+    Plot auto/cross-correlation grid
+
+    Plots lags on the horizontal axis and the correlations
+    on the vertical axis of each graph.
+
+    Parameters
+    ----------
+    x : array_like
+        2D array of time-series values: rows are observations,
+        columns are variables.
+    varnames: sequence of str, optional
+        Variable names to use in plot titles. If ``x`` is a pandas dataframe
+        and ``varnames`` is provided, it overrides the column names
+        of the dataframe. If ``varnames`` is not provided and ``x`` is not
+        a dataframe, variable names ``x[0]``, ``x[1]``, etc. are generated.
+    fig : Matplotlib figure instance, optional
+        If given, this figure is used to plot in, otherwise a new figure
+        is created.
+    lags : {int, array_like}, optional
+        An int or array of lag values, used on horizontal axes. Uses
+        ``np.arange(lags)`` when lags is an int.  If not provided,
+        ``lags=np.arange(len(corr))`` is used.
+    negative_lags: bool, optional
+        If True, negative lags are shown on the horizontal axes of plots
+        below the main diagonal.
+    alpha : scalar, optional
+        If a number is given, the confidence intervals for the given level are
+        plotted, e.g. if alpha=.05, 95 % confidence intervals are shown.
+        If None, confidence intervals are not shown on the plot.
+    use_vlines : bool, optional
+        If True, shows vertical lines and markers for the correlation values.
+        If False, only shows markers.  The default marker is 'o'; it can
+        be overridden with a ``marker`` kwarg.
+    adjusted : bool
+        If True, then denominators for correlations are n-k, otherwise n.
+    fft : bool, optional
+        If True, computes the ACF via FFT.
+    missing : str, optional
+        A string in ['none', 'raise', 'conservative', 'drop'] specifying how
+        NaNs are to be treated.
+    zero : bool, optional
+        Flag indicating whether to include the 0-lag autocorrelations
+        (which are always equal to 1). Default is True.
+    auto_ylims : bool, optional
+        If True, adjusts automatically the vertical axis limits
+        to correlation values.
+    bartlett_confint : bool, default False
+        If True, use Bartlett's formula to calculate confidence intervals
+        in auto-correlation plots. See the description of ``plot_acf`` for
+        details. This argument does not affect cross-correlation plots.
+    vlines_kwargs : dict, optional
+        Optional dictionary of keyword arguments that are passed to vlines.
+    **kwargs : kwargs, optional
+        Optional keyword arguments that are directly passed on to the
+        Matplotlib ``plot`` and ``axhline`` functions.
+
+    Returns
+    -------
+    Figure
+        If `fig` is None, the created figure.  Otherwise, `fig` is returned.
+        Plots on the grid show the cross-correlation of the row variable
+        with the lags of the column variable.
+
+    See Also
+    --------
+    See notes and references for statsmodels.graphics.tsaplots
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> import matplotlib.pyplot as plt
+    >>> import statsmodels.api as sm
+
+    >>> dta = sm.datasets.macrodata.load_pandas().data
+    >>> diffed = dta.diff().dropna()
+    >>> sm.graphics.tsa.plot_accf_grid(diffed[["unemp", "infl"]])
+    >>> plt.show()
+    """
+    from statsmodels.tools.data import _is_using_pandas
+
+    array_like(x, "x", ndim=2)
+    m = x.shape[1]
+
+    fig = utils.create_mpl_fig(fig)
+    gs = fig.add_gridspec(m, m)
+
+    if _is_using_pandas(x, None):
+        varnames = varnames or list(x.columns)
+
+        def get_var(i):
+            return x.iloc[:, i]
+    else:
+        varnames = varnames or [f'x[{i}]' for i in range(m)]
+
+        x = np.asarray(x)
+
+        def get_var(i):
+            return x[:, i]
+
+    for i in range(m):
+        for j in range(m):
+            ax = fig.add_subplot(gs[i, j])
+            if i == j:
+                plot_acf(
+                    get_var(i),
+                    ax=ax,
+                    title=f'ACF({varnames[i]})',
+                    lags=lags,
+                    alpha=alpha,
+                    use_vlines=use_vlines,
+                    adjusted=adjusted,
+                    fft=fft,
+                    missing=missing,
+                    zero=zero,
+                    auto_ylims=auto_ylims,
+                    bartlett_confint=bartlett_confint,
+                    vlines_kwargs=vlines_kwargs,
+                    **kwargs,
+                )
+            else:
+                plot_ccf(
+                    get_var(i),
+                    get_var(j),
+                    ax=ax,
+                    title=f'CCF({varnames[i]}, {varnames[j]})',
+                    lags=lags,
+                    negative_lags=negative_lags and i > j,
+                    alpha=alpha,
+                    use_vlines=use_vlines,
+                    adjusted=adjusted,
+                    fft=fft,
+                    auto_ylims=auto_ylims,
+                    vlines_kwargs=vlines_kwargs,
+                    **kwargs,
+                )
+
+    return fig
+
+
 def seasonal_plot(grouped_x, xticklabels, ylabel=None, ax=None):
     """
     Consider using one of month_plot or quarter_plot unless you need