Merge pull request #3 from PR0M3TH3AN/codex/implement-custom_chart.py-with-argparse-and-plotting

Add flexible custom chart script

Makefile

@@ -12,10 +12,10 @@ data:
 plot:
 	./startup.sh
 
-# Custom plot
-# Usage: make plot-custom ARGS="--offset 12 --end 2025-05-31 --extend-years 5"
+# Custom plot of arbitrary series
+# Usage: make plot-custom ARGS="--series UNRATE CPIAUCSL --start 2000-01-01"
 plot-custom:
-	./startup.sh $(ARGS)
+        ./startup.sh python scripts/custom_chart.py $(ARGS)
 
 # Run tests
 test:

README.md

@@ -80,6 +80,23 @@ requested period before rendering the chart.
 | `--end`          | End date (YYYY-MM-DD)                           |
 | `--extend-years` | Years to extend the x-axis beyond the end date  |
 
+## Custom Charts
+
+Plot any combination of FRED series stored in `data/fred.db` using
+`scripts/custom_chart.py`.
+
+```bash
+# Example: unemployment and oil
+./startup.sh python scripts/custom_chart.py \
+    --series UNRATE DCOILWTICO --start 2000-01-01 --end 2020-12-31
+```
+
+The same invocation is available via Make:
+
+```bash
+make plot-custom ARGS="--series UNRATE DCOILWTICO --start 2000-01-01 --end 2020-12-31"
+```
+
 ---
 
 ## 📝 Details

scripts/custom_chart.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""custom_chart.py
+------------------
+Plot arbitrary FRED series stored in ``data/fred.db``.
+
+Example:
+    python scripts/custom_chart.py --series UNRATE DCOILWTICO --start 2000-01-01 --end 2020-12-31
+"""
+
+from __future__ import annotations
+
+import argparse
+from datetime import datetime
+from pathlib import Path
+from typing import Iterable
+
+import matplotlib.dates as mdates
+import matplotlib.pyplot as plt
+import pandas as pd
+import sqlite3
+
+
+def fetch_series_multi(
+    series: Iterable[str],
+    start: datetime,
+    end: datetime,
+    db_path: str | Path = Path("data/fred.db"),
+) -> pd.DataFrame:
+    """Load one or more FRED series from a local SQLite DB."""
+    db_path = Path(db_path)
+    if not db_path.exists():
+        raise FileNotFoundError(
+            f"{db_path} not found. Run scripts/refresh_data.py on a machine with internet, commit the DB, then retry."
+        )
+
+    frames: list[pd.DataFrame] = []
+    with sqlite3.connect(db_path) as conn:
+        for name in series:
+            df = pd.read_sql(
+                f"SELECT date, {name} FROM {name}",
+                conn,
+                parse_dates=["date"],
+                index_col="date",
+            ).loc[start:end]
+            df.rename(columns={name: name}, inplace=True)
+            frames.append(df)
+
+    if not frames:
+        raise ValueError("No series loaded")
+
+    combined = pd.concat(frames, axis=1).sort_index()
+    return combined
+
+
+def plot_series(df: pd.DataFrame) -> plt.Figure:
+    """Plot each column of ``df`` on a shared axis."""
+    fig, ax = plt.subplots(figsize=(10, 5))
+    for col in df.columns:
+        ax.plot(df.index, df[col], label=col)
+
+    ax.set_xlabel("Date")
+    ax.set_ylabel("Value")
+    ax.legend(frameon=False)
+    ax.grid(axis="y", linestyle="--", alpha=0.4)
+    ax.xaxis.set_major_locator(mdates.YearLocator(base=5))
+    ax.xaxis.set_minor_locator(mdates.YearLocator())
+    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y"))
+    fig.autofmt_xdate()
+    fig.tight_layout()
+    return fig
+
+
+def main(argv: list[str] | None = None) -> plt.Figure:
+    p = argparse.ArgumentParser(description="Plot FRED series from a local DB")
+    p.add_argument("--series", nargs="+", required=True, help="FRED series names")
+    p.add_argument(
+        "--start", type=str, default="1970-01-01", help="start date YYYY-MM-DD"
+    )
+    p.add_argument(
+        "--end",
+        type=str,
+        default=datetime.today().strftime("%Y-%m-%d"),
+        help="end date YYYY-MM-DD",
+    )
+    p.add_argument("--db", type=str, default="data/fred.db", help="path to SQLite DB")
+    args = p.parse_args(argv)
+
+    start_dt = datetime.fromisoformat(args.start)
+    end_dt = datetime.fromisoformat(args.end)
+
+    data = fetch_series_multi(args.series, start_dt, end_dt, args.db)
+    fig = plot_series(data)
+
+    if argv is None:
+        fig.show()
+    return fig
+
+
+if __name__ == "__main__":
+    main()

tests/test_plot.py

@@ -9,6 +9,7 @@ import pandas as pd
 sys.path.append(os.path.abspath("."))
 
 from scripts.lagged_oil_unrate_chart_styled import plot_lagged
+from scripts.custom_chart import plot_series, main as custom_main
 
 
 def test_plot_lagged_returns_figure():
@@ -20,3 +21,27 @@ def test_plot_lagged_returns_figure():
 
     assert len(fig.axes) == 2
     assert any(ax.get_lines() for ax in fig.axes)
+
+
+def test_plot_series_and_main():
+    dates = pd.date_range("2020-01-01", periods=3, freq="D")
+    df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}, index=dates)
+    fig = plot_series(df)
+    assert len(fig.axes) == 1
+    assert any(fig.axes[0].get_lines())
+
+    fig2 = custom_main(
+        [
+            "--series",
+            "UNRATE",
+            "DCOILWTICO",
+            "--start",
+            "2020-01-01",
+            "--end",
+            "2020-03-01",
+            "--db",
+            "data/fred.db",
+        ]
+    )
+    assert len(fig2.axes) == 1
+    assert any(fig2.axes[0].get_lines())