Create common helpers and refactor charts

This commit is contained in:
thePR0M3TH3AN committed 2025-05-31 08:55:43 -04:00
1 parent de889caa3d
commit e74d35dce6
4 files changed
+93 -136

No files matched your search

+4 -34
View File
@@ -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
+74
View File
@@ -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")
+2 -38
View File
@@ -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:
+13 -64
View File
@@ -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)
# ──────────────────────────────────────────────────────────────────────────