feat: init the repo
This commit is contained in:
Pratik Bhadane
@@ -0,0 +1,351 @@
|
||||
"""
|
||||
ferro_ta.viz — Charting and visualisation API.
|
||||
|
||||
Generates charts (matplotlib and/or Plotly) with indicators overlaid on price.
|
||||
|
||||
API
|
||||
---
|
||||
plot(ohlcv, indicators=None, *, backend='matplotlib', title=None,
|
||||
figsize=None, savefig=None, show=False)
|
||||
Generate a chart from OHLCV data and optional indicator series.
|
||||
Returns a figure object for further customisation.
|
||||
|
||||
Backends
|
||||
--------
|
||||
- ``'matplotlib'`` — requires ``matplotlib`` (recommended for static charts)
|
||||
- ``'plotly'`` — requires ``plotly`` (recommended for interactive charts)
|
||||
|
||||
Install optional backends::
|
||||
|
||||
pip install ferro-ta[plot] # adds matplotlib + plotly
|
||||
pip install matplotlib # matplotlib only
|
||||
pip install plotly # plotly only
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import numpy as np
|
||||
>>> from ferro_ta import RSI, SMA
|
||||
>>> from ferro_ta.tools.viz import plot
|
||||
>>> rng = np.random.default_rng(0)
|
||||
>>> n = 60
|
||||
>>> close = np.cumprod(1 + rng.normal(0, 0.01, n)) * 100
|
||||
>>> ohlcv = {"close": close, "open": close, "high": close * 1.01,
|
||||
... "low": close * 0.99, "volume": np.ones(n) * 1000}
|
||||
>>> fig = plot(ohlcv, indicators={"RSI(14)": RSI(close, timeperiod=14),
|
||||
... "SMA(20)": SMA(close, timeperiod=20)},
|
||||
... backend='matplotlib', show=False)
|
||||
>>> fig is not None
|
||||
True
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
import numpy as np
|
||||
import warnings
|
||||
from numpy.typing import ArrayLike, NDArray
|
||||
|
||||
__all__ = [
|
||||
"plot",
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# plot
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def plot(
|
||||
ohlcv: Any,
|
||||
indicators: Optional[dict[str, ArrayLike]] = None,
|
||||
*,
|
||||
backend: str = "matplotlib",
|
||||
title: Optional[str] = None,
|
||||
figsize: Optional[tuple[float, float]] = None,
|
||||
savefig: Optional[str] = None,
|
||||
show: bool = True,
|
||||
volume: bool = True,
|
||||
close_col: str = "close",
|
||||
volume_col: str = "volume",
|
||||
) -> Any:
|
||||
"""Generate a chart from OHLCV data and optional indicator series.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
ohlcv : dict, pandas.DataFrame, or array-like
|
||||
OHLCV data. At minimum a ``close`` key/column is required.
|
||||
indicators : dict {label: array}, optional
|
||||
Additional indicator series to plot below the price panel.
|
||||
Each entry is plotted in its own subplot.
|
||||
backend : str
|
||||
``'matplotlib'`` (default) or ``'plotly'``.
|
||||
title : str, optional
|
||||
Chart title.
|
||||
figsize : (width, height), optional
|
||||
Figure size in inches (matplotlib) or pixels (plotly).
|
||||
savefig : str, optional
|
||||
Save figure to this file path (e.g. ``'chart.png'``, ``'chart.html'``).
|
||||
show : bool
|
||||
If ``True``, call ``plt.show()`` or ``fig.show()`` interactively.
|
||||
volume : bool
|
||||
If ``True`` and a volume series is present, add a volume subplot.
|
||||
close_col, volume_col : str
|
||||
Column names when *ohlcv* is a DataFrame.
|
||||
|
||||
Returns
|
||||
-------
|
||||
matplotlib.figure.Figure or plotly.graph_objects.Figure
|
||||
|
||||
Raises
|
||||
------
|
||||
ImportError
|
||||
If the requested backend is not installed.
|
||||
"""
|
||||
close_arr, volume_arr = _extract_close_volume(ohlcv, close_col, volume_col)
|
||||
|
||||
if backend == "matplotlib":
|
||||
return _plot_matplotlib(
|
||||
close_arr,
|
||||
volume_arr if volume else None,
|
||||
indicators,
|
||||
title=title,
|
||||
figsize=figsize,
|
||||
savefig=savefig,
|
||||
show=show,
|
||||
)
|
||||
elif backend == "plotly":
|
||||
return _plot_plotly(
|
||||
close_arr,
|
||||
volume_arr if volume else None,
|
||||
indicators,
|
||||
title=title,
|
||||
figsize=figsize,
|
||||
savefig=savefig,
|
||||
show=show,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown backend {backend!r}. Supported: 'matplotlib', 'plotly'."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_close_volume(
|
||||
ohlcv: Any,
|
||||
close_col: str,
|
||||
volume_col: str,
|
||||
) -> tuple[NDArray[np.float64], Optional[NDArray[np.float64]]]:
|
||||
"""Extract close and (optional) volume from various input formats."""
|
||||
try:
|
||||
import pandas as pd
|
||||
|
||||
if isinstance(ohlcv, pd.DataFrame):
|
||||
close = ohlcv[close_col].values.astype(np.float64)
|
||||
volume = (
|
||||
ohlcv[volume_col].values.astype(np.float64)
|
||||
if volume_col in ohlcv.columns
|
||||
else None
|
||||
)
|
||||
return close, volume
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if isinstance(ohlcv, dict):
|
||||
close = np.asarray(
|
||||
ohlcv.get(close_col, ohlcv.get("close", [])), dtype=np.float64
|
||||
)
|
||||
vol_key = volume_col if volume_col in ohlcv else "volume"
|
||||
volume = (
|
||||
np.asarray(ohlcv[vol_key], dtype=np.float64) if vol_key in ohlcv else None
|
||||
)
|
||||
return close, volume
|
||||
|
||||
# Plain array
|
||||
return np.asarray(ohlcv, dtype=np.float64), None
|
||||
|
||||
|
||||
def _n_subplots(indicators: Optional[dict], volume_arr: Optional[NDArray]) -> int:
|
||||
n = 1 # price
|
||||
if volume_arr is not None:
|
||||
n += 1
|
||||
if indicators:
|
||||
n += len(indicators)
|
||||
return n
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Matplotlib backend
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _plot_matplotlib(
|
||||
close: NDArray,
|
||||
volume: Optional[NDArray],
|
||||
indicators: Optional[dict[str, ArrayLike]],
|
||||
*,
|
||||
title: Optional[str],
|
||||
figsize: Optional[tuple],
|
||||
savefig: Optional[str],
|
||||
show: bool,
|
||||
) -> Any:
|
||||
try:
|
||||
import matplotlib.gridspec as gridspec
|
||||
import matplotlib.pyplot as plt
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"matplotlib is required for the 'matplotlib' backend. "
|
||||
"Install with: pip install matplotlib"
|
||||
) from exc
|
||||
|
||||
n_subplots = _n_subplots(indicators, volume)
|
||||
height_ratios = [3] + [1] * (n_subplots - 1)
|
||||
fig_h = figsize[1] if figsize else 2.5 * n_subplots + 1
|
||||
fig_w = figsize[0] if figsize else 12.0
|
||||
fig = plt.figure(figsize=(fig_w, fig_h))
|
||||
gs = gridspec.GridSpec(n_subplots, 1, height_ratios=height_ratios, hspace=0.35)
|
||||
|
||||
ax_price = fig.add_subplot(gs[0])
|
||||
ax_price.plot(close, color="#1f77b4", linewidth=1.2, label="close")
|
||||
ax_price.set_ylabel("Price")
|
||||
ax_price.legend(loc="upper left", fontsize=8)
|
||||
ax_price.grid(alpha=0.3)
|
||||
if title:
|
||||
ax_price.set_title(title)
|
||||
|
||||
row = 1
|
||||
if volume is not None:
|
||||
ax_vol = fig.add_subplot(gs[row], sharex=ax_price)
|
||||
ax_vol.bar(range(len(volume)), volume, color="#aec7e8", alpha=0.7, width=0.8)
|
||||
ax_vol.set_ylabel("Volume")
|
||||
ax_vol.grid(alpha=0.3)
|
||||
row += 1
|
||||
|
||||
if indicators:
|
||||
colors = ["#d62728", "#2ca02c", "#9467bd", "#8c564b", "#e377c2", "#17becf"]
|
||||
for idx, (label, arr) in enumerate(indicators.items()):
|
||||
ax_ind = fig.add_subplot(gs[row], sharex=ax_price)
|
||||
color = colors[idx % len(colors)]
|
||||
arr_np = np.asarray(arr, dtype=np.float64)
|
||||
ax_ind.plot(arr_np, color=color, linewidth=1.0, label=label)
|
||||
ax_ind.set_ylabel(label, fontsize=8)
|
||||
ax_ind.legend(loc="upper left", fontsize=8)
|
||||
ax_ind.grid(alpha=0.3)
|
||||
row += 1
|
||||
|
||||
# Use tight_layout when possible but suppress known benign UserWarning
|
||||
# about incompatible Axes configurations.
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings(
|
||||
"ignore",
|
||||
message="This figure includes Axes that are not compatible with tight_layout.*",
|
||||
category=UserWarning,
|
||||
)
|
||||
plt.tight_layout()
|
||||
|
||||
if savefig:
|
||||
fig.savefig(savefig, dpi=100, bbox_inches="tight")
|
||||
if show:
|
||||
plt.show()
|
||||
return fig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Plotly backend
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _plot_plotly(
|
||||
close: NDArray,
|
||||
volume: Optional[NDArray],
|
||||
indicators: Optional[dict[str, ArrayLike]],
|
||||
*,
|
||||
title: Optional[str],
|
||||
figsize: Optional[tuple],
|
||||
savefig: Optional[str],
|
||||
show: bool,
|
||||
) -> Any:
|
||||
try:
|
||||
import plotly.graph_objects as go
|
||||
from plotly.subplots import make_subplots
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"plotly is required for the 'plotly' backend. "
|
||||
"Install with: pip install plotly"
|
||||
) from exc
|
||||
|
||||
n_subplots = _n_subplots(indicators, volume)
|
||||
row_heights = [0.5] + [0.1] * (n_subplots - 1)
|
||||
total = sum(row_heights)
|
||||
row_heights = [r / total for r in row_heights]
|
||||
shared_xaxes = True
|
||||
subplot_titles = ["Price"]
|
||||
if volume is not None:
|
||||
subplot_titles.append("Volume")
|
||||
if indicators:
|
||||
subplot_titles.extend(list(indicators.keys()))
|
||||
|
||||
fig = make_subplots(
|
||||
rows=n_subplots,
|
||||
cols=1,
|
||||
shared_xaxes=shared_xaxes,
|
||||
row_heights=row_heights,
|
||||
subplot_titles=subplot_titles,
|
||||
vertical_spacing=0.05,
|
||||
)
|
||||
x = list(range(len(close)))
|
||||
fig.add_trace(
|
||||
go.Scatter(
|
||||
x=x, y=close.tolist(), mode="lines", name="close", line={"color": "#1f77b4"}
|
||||
),
|
||||
row=1,
|
||||
col=1,
|
||||
)
|
||||
|
||||
row = 2
|
||||
if volume is not None:
|
||||
fig.add_trace(
|
||||
go.Bar(x=x, y=volume.tolist(), name="volume", marker_color="#aec7e8"),
|
||||
row=row,
|
||||
col=1,
|
||||
)
|
||||
row += 1
|
||||
|
||||
if indicators:
|
||||
colors = ["#d62728", "#2ca02c", "#9467bd", "#8c564b", "#e377c2", "#17becf"]
|
||||
for idx, (label, arr) in enumerate(indicators.items()):
|
||||
arr_np = np.asarray(arr, dtype=np.float64)
|
||||
color = colors[idx % len(colors)]
|
||||
fig.add_trace(
|
||||
go.Scatter(
|
||||
x=x,
|
||||
y=arr_np.tolist(),
|
||||
mode="lines",
|
||||
name=label,
|
||||
line={"color": color},
|
||||
),
|
||||
row=row,
|
||||
col=1,
|
||||
)
|
||||
row += 1
|
||||
|
||||
fig_w = figsize[0] if figsize else 900
|
||||
fig_h = figsize[1] if figsize else 500
|
||||
fig.update_layout(
|
||||
title=title or "ferro_ta Chart",
|
||||
width=fig_w,
|
||||
height=fig_h,
|
||||
showlegend=True,
|
||||
)
|
||||
|
||||
if savefig:
|
||||
if savefig.endswith(".html"):
|
||||
fig.write_html(savefig)
|
||||
else:
|
||||
fig.write_image(savefig)
|
||||
if show:
|
||||
fig.show()
|
||||
return fig
|
||||
Reference in New Issue
Block a user