From 34359e51e40f11d1d06a919483e5d9dda6e46e1e Mon Sep 17 00:00:00 2001 From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com> Date: Sat, 31 May 2025 13:56:42 -0400 Subject: [PATCH] Add figure size and display options --- scripts/bitcoin_m2_chart.py | 22 ++++++++++++++--- scripts/lagged_oil_unrate_chart_styled.py | 29 +++++++++++++++++++---- tests/test_bitcoin_m2.py | 21 +++++++++++++++- tests/test_plot.py | 17 ++++++++++++- 4 files changed, 79 insertions(+), 10 deletions(-) diff --git a/scripts/bitcoin_m2_chart.py b/scripts/bitcoin_m2_chart.py index 59e6ee3..746ed90 100644 --- a/scripts/bitcoin_m2_chart.py +++ b/scripts/bitcoin_m2_chart.py @@ -59,6 +59,10 @@ def plot_bitcoin_m2( m2: pd.DataFrame, offset_days: int, extend_years: int, + *, + width: float = 12, + height: float = 6, + dpi: int | None = None, ) -> plt.Figure: """Plot Bitcoin price and shifted M2 on dual axes.""" # Determine overlapping range @@ -70,7 +74,7 @@ def plot_bitcoin_m2( m2_shifted = m2_common.shift(offset_days) - fig, ax1 = plt.subplots(figsize=(12, 6)) + fig, ax1 = plt.subplots(figsize=(width, height), dpi=dpi) ax1.plot( btc_common.index, btc_common["value"], label="Bitcoin (LHS)", color="#1f77b4" ) @@ -137,6 +141,10 @@ def main(argv: list[str] | None = None) -> plt.Figure: p.add_argument( "--output", type=str, default=None, help="optional path to save the figure" ) + p.add_argument("--width", type=float, default=12.0, help="figure width in inches") + p.add_argument("--height", type=float, default=6.0, help="figure height in inches") + p.add_argument("--dpi", type=int, default=None, help="figure DPI") + p.add_argument("--no-show", action="store_true", help="do not display the figure") args = p.parse_args(argv) start_dt = datetime.fromisoformat(args.start) @@ -144,10 +152,18 @@ def main(argv: list[str] | None = None) -> plt.Figure: btc, m2 = fetch_series(start_dt, end_dt, args.btc_series, args.m2_series, args.db) validate_series(btc, m2, args.btc_series, args.m2_series) - fig = plot_bitcoin_m2(btc, m2, args.offset, args.extend_years) + fig = plot_bitcoin_m2( + btc, + m2, + args.offset, + args.extend_years, + width=args.width, + height=args.height, + dpi=args.dpi, + ) save_figure(fig, args.output, __file__) - if argv is None: + if argv is None and not args.no_show: fig.show() return fig diff --git a/scripts/lagged_oil_unrate_chart_styled.py b/scripts/lagged_oil_unrate_chart_styled.py index e2bfc59..635b70f 100644 --- a/scripts/lagged_oil_unrate_chart_styled.py +++ b/scripts/lagged_oil_unrate_chart_styled.py @@ -82,6 +82,10 @@ def plot_lagged( _start_date: datetime, _end_date: datetime, extend_years: int, + *, + width: float = 14, + height: float = 7, + dpi: int | None = None, ) -> plt.Figure: """ Render the dual-axis chart with fixed scales and log ticks on oil. @@ -107,7 +111,7 @@ def plot_lagged( # Apply lag/lead oil_shifted = oil_common.shift(offset_months) - fig, ax1 = plt.subplots(figsize=(14, 7)) + fig, ax1 = plt.subplots(figsize=(width, height), dpi=dpi) plt.title("Unemployment Rate and US Oil Price", fontsize=20, weight="bold") plt.suptitle( "Civilian Unemployment Rate and WTI Crude Oil Price", @@ -193,7 +197,7 @@ def plot_lagged( # ────────────────────────────────────────────────────────────────────────── # CLI # ────────────────────────────────────────────────────────────────────────── -def main() -> plt.Figure: +def main(argv: list[str] | None = None) -> plt.Figure: p = argparse.ArgumentParser() p.add_argument( "-o", @@ -232,7 +236,11 @@ def main() -> plt.Figure: default=None, help="optional path to save the figure (PNG or PDF)", ) - args = p.parse_args() + p.add_argument("--width", type=float, default=14.0, help="figure width in inches") + p.add_argument("--height", type=float, default=7.0, help="figure height in inches") + p.add_argument("--dpi", type=int, default=None, help="figure DPI") + p.add_argument("--no-show", action="store_true", help="do not display the figure") + args = p.parse_args(argv) start_dt = datetime.fromisoformat(args.start) end_dt = datetime.fromisoformat(args.end) @@ -256,9 +264,20 @@ def main() -> plt.Figure: oil = oil.loc[oil.index.isin(unrate.index)] validate_series(unrate, oil) - fig = plot_lagged(unrate, oil, args.offset, start_dt, end_dt, args.extend_years) + fig = plot_lagged( + unrate, + oil, + args.offset, + start_dt, + end_dt, + args.extend_years, + width=args.width, + height=args.height, + dpi=args.dpi, + ) save_figure(fig, args.output, __file__) - fig.show() + if argv is None and not args.no_show: + fig.show() return fig diff --git a/tests/test_bitcoin_m2.py b/tests/test_bitcoin_m2.py index 7f83ee5..473807b 100644 --- a/tests/test_bitcoin_m2.py +++ b/tests/test_bitcoin_m2.py @@ -32,7 +32,19 @@ def test_btc_m2_main(monkeypatch): return btc, m2 monkeypatch.setattr("scripts.bitcoin_m2_chart.fetch_series", fake_fetch) - fig = btc_m2_main(["--btc-series", "BTC", "--m2-series", "M2"]) + fig = btc_m2_main([ + "--btc-series", + "BTC", + "--m2-series", + "M2", + "--width", + "8", + "--height", + "4", + "--dpi", + "150", + "--no-show", + ]) assert len(fig.axes) == 2 @@ -52,6 +64,13 @@ def test_btc_m2_main_saves_output(tmp_path, monkeypatch): "M2", "--output", str(out_file), + "--width", + "8", + "--height", + "4", + "--dpi", + "150", + "--no-show", ] ) assert out_file.exists() diff --git a/tests/test_plot.py b/tests/test_plot.py index e5246e3..26cfaf4 100644 --- a/tests/test_plot.py +++ b/tests/test_plot.py @@ -8,7 +8,7 @@ import pandas as pd sys.path.append(os.path.abspath(".")) -from scripts.lagged_oil_unrate_chart_styled import plot_lagged +from scripts.lagged_oil_unrate_chart_styled import plot_lagged, main as lagged_main from scripts.custom_chart import plot_series, main as custom_main @@ -64,3 +64,18 @@ def test_custom_main_saves_output(tmp_path): ] ) assert out_file.exists() + + +def test_lagged_main_options(monkeypatch): + dates = pd.date_range("2021-01-01", periods=3, freq="M") + unrate = pd.DataFrame({"value": [1, 2, 3]}, index=dates) + oil = pd.DataFrame({"value": [4, 5, 6]}, index=dates) + + monkeypatch.setattr( + "scripts.lagged_oil_unrate_chart_styled.fetch_series", lambda *_: (unrate, oil) + ) + + fig = lagged_main( + ["--offset", "1", "--width", "8", "--height", "4", "--dpi", "150", "--no-show"] + ) + assert len(fig.axes) == 2