"""Test all three chart types from the PIMCO reference PDF.""" import sys sys.path.insert(0, ".") import pandas as pd import numpy as np from app.models.chart_spec import ( ChartSpec, PanelSpec, AxisSpec, SeriesSpec, ShadedFillSpec, AnnotationSpec, ) from app.renderer.engine import render_chart def test_chart1_bond_yields(): """Chart 1: Multi-line time series (bond yields across 5 countries).""" np.random.seed(42) dates = pd.date_range("2020-08-01", "2025-08-31", freq="B") n = len(dates) def walk(start, drift, vol): r = np.random.normal(drift, vol, n) v = start + np.cumsum(r) return pd.Series(v).rolling(10, min_periods=1).mean().values df = pd.DataFrame({ "Date": dates, "U.S.": walk(0.7, 0.0025, 0.03), "Australia": walk(0.9, 0.0025, 0.03), "U.K.": walk(0.3, 0.0028, 0.035), "Germany": walk(-0.4, 0.0020, 0.03), "Japan": walk(0.0, 0.0008, 0.015), }) spec = ChartSpec( layout="single", panels=[PanelSpec( title="10-year government bond yields", x_axis=AxisSpec(date_format="%b %Y"), y_axis=AxisSpec(suffix="%", min_val=-1, max_val=6, tick_interval=1), series=[ SeriesSpec(label="U.S.", data_column="U.S.", color_index=0), SeriesSpec(label="Australia", data_column="Australia", color_index=1), SeriesSpec(label="U.K.", data_column="U.K.", color_index=2), SeriesSpec(label="Germany", data_column="Germany", color_index=3), SeriesSpec(label="Japan", data_column="Japan", color_index=4), ], )], ) svg = render_chart(spec, {"_default": df}) with open("output/chart1_bond_yields.svg", "w") as f: f.write(svg) print("Chart 1 (Bond Yields) saved.") def test_chart2_gdp_investment(): """Chart 2: Two-line chart with grey ellipse annotation.""" np.random.seed(99) dates = pd.date_range("1994-01-01", "2024-12-31", freq="QS") n = len(dates) tech = np.random.normal(0.5, 0.35, n) tech = pd.Series(tech).rolling(2, min_periods=1).mean().values.copy() nontech = np.random.normal(0.4, 0.4, n) nontech = pd.Series(nontech).rolling(2, min_periods=1).mean().values.copy() # Add recession dips for start_idx in [24, 32, 56]: # ~2000, 2002, 2008 tech[start_idx:start_idx+4] -= 0.7 nontech[start_idx:start_idx+4] -= 0.8 df = pd.DataFrame({"Date": dates, "Tech investment": tech, "Non-tech investment": nontech}) spec = ChartSpec( layout="single", panels=[PanelSpec( title="Quarterly contribution to real GDP growth", x_axis=AxisSpec(date_format="%Y"), y_axis=AxisSpec(label="Percentage points", min_val=-1.0, max_val=1.5, tick_interval=0.5), series=[ SeriesSpec(label="Tech investment", data_column="Tech investment", color_index=0), SeriesSpec(label="Non-tech investment", data_column="Non-tech investment", color_index=2), ], annotations=[ AnnotationSpec( type="ellipse", x_start="2023-06-01", x_end="2024-12-31", y_start=-0.3, y_end=1.2, ), ], )], ) svg = render_chart(spec, {"_default": df}) with open("output/chart2_gdp_investment.svg", "w") as f: f.write(svg) print("Chart 2 (GDP Investment) saved.") def test_chart3_dual_panel(): """Chart 3: Dual-panel with trend lines and shaded fills.""" np.random.seed(77) dates = pd.date_range("2022-01-01", "2026-01-31", freq="MS") n = len(dates) # Left panel: PIMCO Global IP trend_left = np.linspace(93, 101, n) ip_data = trend_left + np.random.normal(0, 1.0, n) ip_data = pd.Series(ip_data).rolling(3, min_periods=1).mean().values.copy() # Post-election dip election_idx = 34 # ~Oct 2024 ip_data[election_idx+2:election_idx+5] -= 2.5 election_level_left = [100.0] * n # Right panel: World Exports trend_right = np.linspace(93, 106, n) exports = trend_right + np.random.normal(0, 2.0, n) exports = pd.Series(exports).rolling(3, min_periods=1).mean().values.copy() # Tariff spike exports[election_idx+3:election_idx+6] += 15 exports[election_idx+6:election_idx+8] -= 8 election_level_right = [100.0] * n df_left = pd.DataFrame({ "Date": dates, "PIMCO Global IP": ip_data, "Trend": trend_left, "Election Level": election_level_left, }) df_right = pd.DataFrame({ "Date": dates, "World Exports to U.S.": exports, "Trend": trend_right, "Election Level": election_level_right, }) spec = ChartSpec( layout="dual_panel", panels=[ PanelSpec( title="PIMCO global industrial production", x_axis=AxisSpec(date_format="%b '%y"), y_axis=AxisSpec(min_val=90, max_val=106, tick_interval=2), series=[ SeriesSpec( label="PIMCO global IP (Oct. 2024 = 100)", data_column="PIMCO Global IP", color_index=0, shaded_fill=ShadedFillSpec( reference_series="Two-year pre-U.S. election trend", above_color="blue", below_color="pink", ), ), SeriesSpec( label="Two-year pre-U.S. election trend", data_column="Trend", color_index=1, line_style="dashed", ), SeriesSpec( label="Pre-U.S. election level", data_column="Election Level", color_index=2, line_style="dotted", ), ], ), PanelSpec( title="World exports to U.S.", x_axis=AxisSpec(date_format="%b '%y"), y_axis=AxisSpec(min_val=90, max_val=130, tick_interval=5), series=[ SeriesSpec( label="World exports to U.S. (Oct. 2024 = 100)", data_column="World Exports to U.S.", color_index=0, shaded_fill=ShadedFillSpec( reference_series="Two-year pre-U.S. election trend", above_color="blue", below_color="pink", ), ), SeriesSpec( label="Two-year pre-U.S. election trend", data_column="Trend", color_index=1, line_style="dashed", ), SeriesSpec( label="Pre-U.S. election level", data_column="Election Level", color_index=2, line_style="dotted", ), ], ), ], ) svg = render_chart(spec, {"_panel_0": df_left, "_panel_1": df_right}) with open("output/chart3_dual_panel.svg", "w") as f: f.write(svg) print("Chart 3 (Dual Panel) saved.") if __name__ == "__main__": test_chart1_bond_yields() test_chart2_gdp_investment() test_chart3_dual_panel() print("\nAll test charts saved to output/")