From e74d35dce6cdd32dae5744186849116167544848 Mon Sep 17 00:00:00 2001 From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com> Date: Sat, 31 May 2025 08:55:43 -0400 Subject: [PATCH] Create common helpers and refactor charts --- scripts/bitcoin_m2_chart.py | 38 ++--------- scripts/common.py | 74 ++++++++++++++++++++++ scripts/custom_chart.py | 40 +----------- scripts/lagged_oil_unrate_chart_styled.py | 77 ++++------------------- 4 files changed, 93 insertions(+), 136 deletions(-) create mode 100644 scripts/common.py diff --git a/scripts/bitcoin_m2_chart.py b/scripts/bitcoin_m2_chart.py index 51bee70..572e0f2 100644 --- a/scripts/bitcoin_m2_chart.py +++ b/scripts/bitcoin_m2_chart.py @@ -19,21 +19,7 @@ import matplotlib.dates as mdates import matplotlib.pyplot as plt import pandas as pd from dateutil.relativedelta import relativedelta -import sqlite3 - - -def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path: - """Save ``fig`` to ``output`` or a timestamped path under ``outputs/``.""" - out_path = ( - Path(output) - if output is not None - else Path("outputs") - / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" - ) - out_path.parent.mkdir(parents=True, exist_ok=True) - fig.savefig(out_path) - print(f"Saved figure to {out_path}") - return out_path +from scripts.common import fetch_series_db, save_figure def fetch_series( @@ -44,25 +30,9 @@ def fetch_series( db_path: str | Path = Path("data/fred.db"), ) -> Tuple[pd.DataFrame, pd.DataFrame]: """Load Bitcoin and M2 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." - ) - - with sqlite3.connect(db_path) as conn: - btc = pd.read_sql( - f"SELECT date, {btc_series} AS value FROM {btc_series}", - conn, - parse_dates=["date"], - index_col="date", - ).loc[start:end] - m2 = pd.read_sql( - f"SELECT date, {m2_series} AS value FROM {m2_series}", - conn, - parse_dates=["date"], - index_col="date", - ).loc[start:end] + df = fetch_series_db([btc_series, m2_series], start, end, db_path) + btc = df[[btc_series]].rename(columns={btc_series: "value"}) + m2 = df[[m2_series]].rename(columns={m2_series: "value"}) return btc, m2 diff --git a/scripts/common.py b/scripts/common.py new file mode 100644 index 0000000..f3dfe9d --- /dev/null +++ b/scripts/common.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from datetime import datetime +from pathlib import Path +from typing import Iterable + +import matplotlib.pyplot as plt +import pandas as pd +import sqlite3 + + +def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path: + """Save ``fig`` to ``output`` or ``outputs/`` with timestamp.""" + out_path = ( + Path(output) + if output is not None + else Path("outputs") + / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" + ) + out_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out_path) + print(f"Saved figure to {out_path}") + return out_path + + +def fetch_series_db( + 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") + + return pd.concat(frames, axis=1).sort_index() + + +def validate_dataframe(df: pd.DataFrame, name: str) -> None: + """Ensure dataframe has a ``value`` column with data and DatetimeIndex.""" + if df.empty: + raise ValueError(f"{name} dataframe is empty") + if "value" not in df.columns: + raise ValueError(f"{name} missing 'value' column") + if df["value"].isna().any(): + raise ValueError(f"{name} contains null values") + if not isinstance(df.index, pd.DatetimeIndex): + raise ValueError("Indexes must be DatetimeIndex") + + +def validate_overlap(df1: pd.DataFrame, df2: pd.DataFrame) -> None: + """Verify that two series share a non-empty overlapping date range.""" + start = max(df1.index.min(), df2.index.min()) + end = min(df1.index.max(), df2.index.max()) + if start >= end: + raise ValueError("Series have no overlapping date range") diff --git a/scripts/custom_chart.py b/scripts/custom_chart.py index 915f500..46dfdad 100644 --- a/scripts/custom_chart.py +++ b/scripts/custom_chart.py @@ -17,21 +17,7 @@ from typing import Iterable import matplotlib.dates as mdates import matplotlib.pyplot as plt import pandas as pd -import sqlite3 - - -def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path: - """Save ``fig`` to ``output`` or a timestamped path under ``outputs/``.""" - out_path = ( - Path(output) - if output is not None - else Path("outputs") - / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" - ) - out_path.parent.mkdir(parents=True, exist_ok=True) - fig.savefig(out_path) - print(f"Saved figure to {out_path}") - return out_path +from scripts.common import fetch_series_db, save_figure def fetch_series_multi( @@ -41,29 +27,7 @@ def fetch_series_multi( 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 + return fetch_series_db(series, start, end, db_path) def plot_series(df: pd.DataFrame) -> plt.Figure: diff --git a/scripts/lagged_oil_unrate_chart_styled.py b/scripts/lagged_oil_unrate_chart_styled.py index cb65e43..4fa991a 100644 --- a/scripts/lagged_oil_unrate_chart_styled.py +++ b/scripts/lagged_oil_unrate_chart_styled.py @@ -17,24 +17,14 @@ from pathlib import Path import matplotlib.dates as mdates import matplotlib.pyplot as plt import matplotlib.ticker as mticker -import numpy as np import pandas as pd from dateutil.relativedelta import relativedelta -import sqlite3 - - -def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path: - """Save ``fig`` to ``output`` or a timestamped path under ``outputs/``.""" - out_path = ( - Path(output) - if output is not None - else Path("outputs") - / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" - ) - out_path.parent.mkdir(parents=True, exist_ok=True) - fig.savefig(out_path) - print(f"Saved figure to {out_path}") - return out_path +from scripts.common import ( + fetch_series_db, + save_figure, + validate_dataframe, + validate_overlap, +) # ────────────────────────────────────────────────────────────────────────── @@ -51,26 +41,9 @@ def fetch_series( * convert UNRATE to month-end timestamps * average WTI daily prices to month-end """ - 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 new DB, then retry." - ) - - with sqlite3.connect(db_path) as conn: - unrate = pd.read_sql( - "SELECT date, UNRATE AS value FROM UNRATE", - conn, - parse_dates=["date"], - index_col="date", - ).loc[start:end] - oil = pd.read_sql( - "SELECT date, DCOILWTICO AS value FROM DCOILWTICO", - conn, - parse_dates=["date"], - index_col="date", - ).loc[start:end] + df = fetch_series_db(["UNRATE", "DCOILWTICO"], start, end, db_path) + unrate = df[["UNRATE"]].rename(columns={"UNRATE": "value"}) + oil = df[["DCOILWTICO"]].rename(columns={"DCOILWTICO": "value"}) # Resample WTI to month-end average oil_monthly = oil.resample("M").mean() @@ -80,34 +53,10 @@ def fetch_series( def validate_series(unrate: pd.DataFrame, oil: pd.DataFrame) -> None: - """ - Sanity-check the fetched series before plotting. - - Both dataframes must contain a ``value`` column with no missing values - and share a non-empty overlapping date range. - """ - if unrate.empty: - raise ValueError("UNRATE dataframe is empty") - if oil.empty: - raise ValueError("DCOILWTICO dataframe is empty") - - for name, df in [("UNRATE", unrate), ("DCOILWTICO", oil)]: - if "value" not in df.columns: - raise ValueError(f"{name} missing 'value' column") - if df["value"].isna().any(): - raise ValueError(f"{name} contains null values") - - if not isinstance(unrate.index, pd.DatetimeIndex) or not isinstance( - oil.index, pd.DatetimeIndex - ): - raise ValueError("Indexes must be DatetimeIndex") - - # Replace strict range check with overlap check - overlap_start = max(unrate.index.min(), oil.index.min()) - overlap_end = min(unrate.index.max(), oil.index.max()) - - if overlap_start >= overlap_end: - raise ValueError("Series have no overlapping date range") + """Run basic sanity checks on the two series.""" + validate_dataframe(unrate, "UNRATE") + validate_dataframe(oil, "DCOILWTICO") + validate_overlap(unrate, oil) # ──────────────────────────────────────────────────────────────────────────