diff --git a/openbb_platform/extensions/econometrics/integration/test_econometrics_api.py b/openbb_platform/extensions/econometrics/integration/test_econometrics_api.py index 068c6d02e3a6..2048f4486dba 100644 --- a/openbb_platform/extensions/econometrics/integration/test_econometrics_api.py +++ b/openbb_platform/extensions/econometrics/integration/test_econometrics_api.py @@ -71,8 +71,8 @@ def get_data(menu: Literal["equity", "crypto"]): @pytest.mark.parametrize( "params, data_type", [ - ({"data": ""}, "equity"), - ({"data": ""}, "crypto"), + ({"data": "", "method": "pearson"}, "equity"), + ({"data": "", "method": "pearson"}, "crypto"), ], ) @pytest.mark.integration diff --git a/openbb_platform/extensions/econometrics/integration/test_econometrics_python.py b/openbb_platform/extensions/econometrics/integration/test_econometrics_python.py index 752156ed6df8..03f596cb24ab 100644 --- a/openbb_platform/extensions/econometrics/integration/test_econometrics_python.py +++ b/openbb_platform/extensions/econometrics/integration/test_econometrics_python.py @@ -66,8 +66,8 @@ def get_data(menu: Literal["equity", "crypto"]): @parametrize( "params, data_type", [ - ({"data": ""}, "equity"), - ({"data": ""}, "crypto"), + ({"data": "", "method": "pearson"}, "equity"), + ({"data": "", "method": "pearson"}, "crypto"), ], ) @pytest.mark.integration diff --git a/openbb_platform/extensions/econometrics/openbb_econometrics/econometrics_router.py b/openbb_platform/extensions/econometrics/openbb_econometrics/econometrics_router.py index 5bd700a41e46..25744bf4f68a 100644 --- a/openbb_platform/extensions/econometrics/openbb_econometrics/econometrics_router.py +++ b/openbb_platform/extensions/econometrics/openbb_econometrics/econometrics_router.py @@ -25,7 +25,9 @@ APIEx(parameters={"data": APIEx.mock_data("timeseries")}), ], ) -def correlation_matrix(data: List[Data]) -> OBBject[List[Data]]: +def correlation_matrix( + data: List[Data], method: Literal["pearson", "kendall", "spearman"] = "pearson" +) -> OBBject[List[Data]]: """Get the correlation matrix of an input dataset. The correlation matrix provides a view of how different variables in your dataset relate to one another. @@ -37,6 +39,11 @@ def correlation_matrix(data: List[Data]) -> OBBject[List[Data]]: ---------- data : List[Data] Input dataset. + method : Literal["pearson", "kendall", "spearman"] + Method to use for correlation calculation. Default is "pearson". + pearson : standard correlation coefficient + kendall : Kendall Tau correlation coefficient + spearman : Spearman rank correlation Returns ------- @@ -49,9 +56,14 @@ def correlation_matrix(data: List[Data]) -> OBBject[List[Data]]: df = basemodel_to_df(data) # remove non float columns from the dataframe to perform the correlation - df = df.select_dtypes(include=["float64"]) - corr = df.corr() + if "symbol" in df.columns and len(df.symbol.unique()) > 1 and "close" in df.columns: + df = df.pivot( + columns="symbol", + values="close", + ) + + corr = df.corr(method=method, numeric_only=True) # replace nan values with None to allow for json serialization corr = corr.replace(np.NaN, None) diff --git a/openbb_platform/extensions/econometrics/openbb_econometrics/econometrics_views.py b/openbb_platform/extensions/econometrics/openbb_econometrics/econometrics_views.py new file mode 100644 index 000000000000..61c935a56b83 --- /dev/null +++ b/openbb_platform/extensions/econometrics/openbb_econometrics/econometrics_views.py @@ -0,0 +1,39 @@ +"""Views for the Econometrics Extension.""" + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from openbb_charting.core.openbb_figure import ( + OpenBBFigure, + ) + + +class EconometricsViews: + """Econometrics Views.""" + + @staticmethod + def econometrics_correlation_matrix( # noqa: PLR0912 + **kwargs, + ) -> tuple["OpenBBFigure", dict[str, Any]]: + """Correlation Matrix Chart. + + Parameters + ---------- + data : Union[list[Data], DataFrame] + Input dataset. + method : Literal["pearson", "kendall", "spearman"] + Method to use for correlation calculation. Default is "pearson". + pearson : standard correlation coefficient + kendall : Kendall Tau correlation coefficient + spearman : Spearman rank correlation + colorscale : str + Plotly colorscale to use for the heatmap. Default is "RdBu". + title : str + Title of the chart. Default is "Asset Correlation Matrix". + layout_kwargs : Dict[str, Any] + Additional keyword arguments to apply with figure.update_layout(), by default None. + """ + # pylint: disable=import-outside-toplevel + from openbb_charting.charts.correlation_matrix import correlation_matrix + + return correlation_matrix(**kwargs) diff --git a/openbb_platform/extensions/econometrics/pyproject.toml b/openbb_platform/extensions/econometrics/pyproject.toml index a0d825908b4b..04de5f298356 100644 --- a/openbb_platform/extensions/econometrics/pyproject.toml +++ b/openbb_platform/extensions/econometrics/pyproject.toml @@ -21,3 +21,6 @@ build-backend = "poetry.core.masonry.api" [tool.poetry.plugins."openbb_core_extension"] econometrics = "openbb_econometrics.econometrics_router:router" + +[tool.poetry.plugins."openbb_charting_extension"] +econometrics = "openbb_econometrics.econometrics_views:EconometricsViews" \ No newline at end of file diff --git a/openbb_platform/obbject_extensions/charting/integration/test_charting_api.py b/openbb_platform/obbject_extensions/charting/integration/test_charting_api.py index e03b76f233dd..fca564fc9d35 100644 --- a/openbb_platform/obbject_extensions/charting/integration/test_charting_api.py +++ b/openbb_platform/obbject_extensions/charting/integration/test_charting_api.py @@ -897,3 +897,44 @@ def test_charting_economy_survey_bls_series(params, headers): assert chart assert not fig assert list(chart.keys()) == ["content", "format"] + + +@parametrize( + "params", + [ + ( + { + "data": "", + "method": "pearson", + "chart": True, + } + ) + ], +) +@pytest.mark.integration +def test_charting_econometrics_correlation_matrix(params, headers): + """Test chart econometrics correlation matrix.""" + # pylint:disable=import-outside-toplevel + from pandas import DataFrame + + url = "http://0.0.0.0:8000/api/v1/equity/price/historical?symbol=AAPL,MSFT,GOOG&provider=yfinance" + result = requests.get(url, headers=headers, timeout=10) + df = DataFrame(result.json()["results"]) + df = df.pivot(index="date", columns="symbol", values="close").reset_index() + body = df.to_dict(orient="records") + + params = {p: v for p, v in params.items() if v} + + query_str = get_querystring(params, []) + url = f"http://0.0.0.0:8000/api/v1/econometrics/correlation_matrix?{query_str}" + result = requests.post(url, headers=headers, timeout=10, data=json.dumps(body)) + + assert isinstance(result, requests.Response) + assert result.status_code == 200 + + chart = result.json()["chart"] + fig = chart.pop("fig", {}) + + assert chart + assert not fig + assert list(chart.keys()) == ["content", "format"] diff --git a/openbb_platform/obbject_extensions/charting/integration/test_charting_python.py b/openbb_platform/obbject_extensions/charting/integration/test_charting_python.py index 2f440606f913..24bc4029bc54 100644 --- a/openbb_platform/obbject_extensions/charting/integration/test_charting_python.py +++ b/openbb_platform/obbject_extensions/charting/integration/test_charting_python.py @@ -728,3 +728,34 @@ def test_charting_economy_survey_bls_series(params, obb): assert len(result.results) > 0 assert result.chart.content assert isinstance(result.chart.fig, OpenBBFigure) + + +@parametrize( + "params", + [ + ( + { + "data": "", + "method": "pearson", + "chart": True, + } + ) + ], +) +@pytest.mark.integration +def test_charting_econometrics_correlation_matrix(params, obb): + """Test chart econometrics correlation matrix.""" + + symbols = "XRT,XLB,XLI,XLH,XLC,XLY,XLU,XLK".split(",") + params["data"] = ( + obb.equity.price.historical(symbol=symbols, provider="yfinance") + .to_df() + .pivot(columns="symbol", values="close") + .filter(items=symbols, axis=1) + ) + result = obb.econometrics.correlation_matrix(**params) + assert result + assert isinstance(result, OBBject) + assert len(result.results) > 0 + assert result.chart.content + assert isinstance(result.chart.fig, OpenBBFigure) diff --git a/openbb_platform/obbject_extensions/charting/openbb_charting/charting.py b/openbb_platform/obbject_extensions/charting/openbb_charting/charting.py index edf6f2caf92d..4f400ebd8c70 100644 --- a/openbb_platform/obbject_extensions/charting/openbb_charting/charting.py +++ b/openbb_platform/obbject_extensions/charting/openbb_charting/charting.py @@ -1,6 +1,6 @@ """Charting Class implementation.""" -# pylint: disable=too-many-arguments,unused-argument +# pylint: disable=too-many-arguments,unused-argument,too-many-positional-arguments from typing import ( TYPE_CHECKING, @@ -57,6 +57,8 @@ class Charting: Create a line chart from external data. create_bar_chart Create a bar chart, on a single x-axis with one or more values for the y-axis, from external data. + create_correlation_matrix + Create a correlation matrix from external data. toggle_chart_style Toggle the chart style, of an existing chart, between light and dark mode. """ @@ -367,6 +369,54 @@ def create_bar_chart( return fig + def create_correlation_matrix( + self, + data: Union[ + list[Data], + "DataFrame", + ], + method: Literal["pearson", "kendall", "spearman"] = "pearson", + colorscale: str = "RdBu", + title: str = "Asset Correlation Matrix", + layout_kwargs: Optional[Dict[str, Any]] = None, + ): + """Create a correlation matrix from external data. + + Parameters + ---------- + data : Union[list[Data], DataFrame] + Input dataset. + method : Literal["pearson", "kendall", "spearman"] + Method to use for correlation calculation. Default is "pearson". + pearson : standard correlation coefficient + kendall : Kendall Tau correlation coefficient + spearman : Spearman rank correlation + colorscale : str + Plotly colorscale to use for the heatmap. Default is "RdBu". + title : str + Title of the chart. Default is "Asset Correlation Matrix". + layout_kwargs : Dict[str, Any] + Additional keyword arguments to apply with figure.update_layout(), by default None. + + Returns + ------- + OpenBBFigure + The OpenBBFigure object. + """ + # pylint: disable=import-outside-toplevel + from openbb_charting.charts.correlation_matrix import correlation_matrix + + kwargs = { + "data": data, + "method": method, + "colorscale": colorscale, + "title": title, + "layout_kwargs": layout_kwargs, + } + fig, _ = correlation_matrix(**kwargs) + fig = self._set_chart_style(fig) + return fig + # pylint: disable=inconsistent-return-statements def show(self, render: bool = True, **kwargs): """Display chart and save it to the OBBject.""" diff --git a/openbb_platform/obbject_extensions/charting/openbb_charting/charts/correlation_matrix.py b/openbb_platform/obbject_extensions/charting/openbb_charting/charts/correlation_matrix.py new file mode 100644 index 000000000000..10c3c49f45ed --- /dev/null +++ b/openbb_platform/obbject_extensions/charting/openbb_charting/charts/correlation_matrix.py @@ -0,0 +1,117 @@ +"""Correlation Matrix Chart.""" + +from typing import TYPE_CHECKING, Any, Union + +if TYPE_CHECKING: + from plotly.graph_objs import Figure # noqa + from openbb_charting.core.openbb_figure import OpenBBFigure # noqa + + +def correlation_matrix( # noqa: PLR0912 + **kwargs, +) -> tuple[Union["OpenBBFigure", "Figure"], dict[str, Any]]: + """Correlation Matrix Chart.""" + # pylint: disable=import-outside-toplevel + from numpy import ones_like, triu # noqa + from openbb_core.app.utils import basemodel_to_df # noqa + from openbb_charting.core.openbb_figure import OpenBBFigure + from openbb_charting.core.chart_style import ChartStyle + from plotly.graph_objs import Figure, Heatmap, Layout + from pandas import DataFrame + + if "data" in kwargs and isinstance(kwargs["data"], DataFrame): + corr = kwargs["data"] + elif "data" in kwargs and isinstance(kwargs["data"], list): + corr = basemodel_to_df(kwargs["data"], index=kwargs.get("index", "date")) # type: ignore + else: + corr = basemodel_to_df( + kwargs["obbject_item"], index=kwargs.get("index", "date") # type: ignore + ) + if ( + "symbol" in corr.columns + and len(corr.symbol.unique()) > 1 + and "close" in corr.columns + ): + corr = corr.pivot( + columns="symbol", + values="close", + ) + + method = kwargs.get("method") or "pearson" + corr = corr.corr(method=method, numeric_only=True) + + X = corr.columns.to_list() + x_replace = X[-1] + Y = X.copy() + y_replace = Y[0] + X = [x if x != x_replace else "" for x in X] + Y = [y if y != y_replace else "" for y in Y] + mask = triu(ones_like(corr, dtype=bool)) + df = corr.mask(mask) + title = kwargs.get("title") or "Asset Correlation Matrix" + text_color = "white" if ChartStyle().plt_style == "dark" else "black" + colorscale = kwargs.get("colorscale") or "RdBu" + heatmap = Heatmap( + z=df, + x=X, + y=Y, + xgap=1, + ygap=1, + colorscale=colorscale, + colorbar=dict( + orientation="v", + x=0.9, + y=0.45, + xanchor="left", + yanchor="middle", + len=0.75, + bgcolor="rgba(0,0,0,0)" if text_color == "white" else "rgba(255,255,255,0)", + ), + text=df.fillna(""), + texttemplate="%{text:.4f}", + hoverinfo="skip", + ) + layout = Layout( + title=title, + title_x=0.5, + title_y=0.95, + xaxis=dict( + showgrid=False, + showline=False, + ticklen=0, + tickfont=dict(size=16), + ticklabelstandoff=10, + domain=[0.05, 1], + ), + yaxis=dict( + showgrid=False, + side="left", + autorange="reversed", + showline=False, + ticklen=0, + tickfont=dict(size=16), + ticklabelstandoff=15, + domain=[0.05, 1], + ), + margin=dict(r=20, t=0, b=50), + dragmode=False, + ) + fig = Figure(data=[heatmap], layout=layout) + figure = OpenBBFigure(fig=fig) + figure.update_layout( + font=dict(color=text_color), + paper_bgcolor=( + "rgba(0,0,0,0)" if text_color == "white" else "rgba(255,255,255,0)" + ), + plot_bgcolor=( + "rgba(0,0,0,0)" if text_color == "white" else "rgba(255,255,255,0)" + ), + ) + layout_kwargs = kwargs.get("layout_kwargs", {}) + + if layout_kwargs: + figure.update_layout(**layout_kwargs) + + content = figure.show(external=True).to_plotly_json() # type: ignore + + return figure, content