Function bodies 261 total

def plot_model_predictions(
    predictions,
    df,
    stan_data,
    model_name: str,
    output_path: str = None,
    show_observations: bool = True,
    show_ci: bool = True,
    zoom_to = None,
    zoom_reference_date = None,
):
    """Plot model predictions at unobserved time points.

    Args:
        predictions: Dictionary from extract_predictions() with keys
                     t_pred, f_pred_mean, f_pred_lower, f_pred_upper, etc.
        df: Original DataFrame with timestamps (for start date reference)
        stan_data: Stan data dictionary with _t_max for scaling
        model_name: Name of model for plot title
        output_path: Path to save plot (optional)
        show_observations: Whether to show data points
        show_ci: Whether to show 95% credible interval for predictions
        zoom_to: Optional zoom specification. Can be:
                - Tuple of (start_date, end_date) as pandas Timestamps
                - String preset: 'last_week', 'last_month', 'last

def plot_hourly_predictions(
    predictions,
    df,
    stan_data,
    model_name: str,
    output_path: str = None,
    day_index: int = 0,
    show_ci: bool = True,
    show_all_days: bool = False,
):
    """Plot hourly predictions for a specific day or aggregated across days.

    Args:
        predictions: Dictionary from extract_predictions() with keys
                     t_pred, hour_of_day_pred, f_pred_mean, f_pred_lower, f_pred_upper, etc.
        df: Original DataFrame with timestamps (for start date reference)
        stan_data: Stan data dictionary with _t_max for scaling
        model_name: Name of model for plot title
        output_path: Path to save plot (optional)
        day_index: Which day to visualize (0-indexed within prediction grid)
        show_ci: Whether to show 95% credible interval
        show_all_days: If True, show aggregated hourly means across all days
                      If False, show predictions for specific day_index

    Returns:
        matpl

def plot_weekly_zoomed_predictions(
    predictions,
    df,
    stan_data,
    model_name: str,
    output_path: str = None,
    target_date=None,
    week_start_date=None,
    show_ci: bool = True,
    show_observations: bool = True,
):
    """Plot hourly predictions zoomed into a specific week.

    Args:
        predictions: Dictionary from extract_predictions() with keys
                     t_pred, hour_of_day_pred, f_pred_mean, f_pred_lower, f_pred_upper, etc.
        df: Original DataFrame with timestamps (for start date reference)
        stan_data: Stan data dictionary with _t_max for scaling
        model_name: Name of model for plot title
        output_path: Path to save plot (optional)
        target_date: Target date (pandas Timestamp) to center week on.
                     If None, uses the middle of prediction range.
        week_start_date: Explicit start date of week (pandas Timestamp).
                         If provided, overrides target_date.
        show_ci: Wh

def plot_state_space_expectations(
    idata,
    df_weight,
    df_intensity,
    stan_data,
    model_name: str,
    output_path: str = None,
    show_ci: bool = True,
):
    """Plot state-space model expectations for fitness and weight with data overlays.

    Creates a 2-panel figure:
    1. Latent fitness state over time with workout intensity bars
    2. Weight predictions with weight observations

    Args:
        idata: ArviZ InferenceData from state-space model fit
        df_weight: DataFrame with weight observations (columns: timestamp, weight_lbs)
        df_intensity: DataFrame with daily intensity (columns: date, intensity)
        stan_data: Stan data dictionary with scaling parameters
        model_name: Name of model for plot title
        output_path: Path to save plot (optional)
        show_ci: Whether to show 95% credible intervals

    Returns:
        matplotlib Figure object
    """
    # Back-transform scaling parameters
    y_mean = stan_data.get("_y_mean", 0

def plot_state_space_component_details(
    idata,
    df_weight,
    df_intensity,
    stan_data,
    model_name: str,
    output_path: str = None,
    show_ci: bool = True,
):
    """Plot detailed component-by-component views for state-space spline model.

    Creates a 2x2 grid showing:
    1. Fitness effect over time (gamma * fitness)
    2. GP trend component over time
    3. Daily spline component over time
    4. Residuals (observed - predicted) over time

    Args:
        idata: ArviZ InferenceData from state-space spline model fit
        df_weight: DataFrame with weight observations (columns: timestamp, weight_lbs)
        df_intensity: DataFrame with daily intensity (columns: date, intensity)
        stan_data: Stan data dictionary with scaling parameters
        model_name: Name of model for plot title
        output_path: Path to save plot (optional)
        show_ci: Whether to show 95% credible intervals

    Returns:
        matplotlib Figure object
    """
    # Back-t

def plot_weight_80ci(idata, df, stan_data, output_path: Path | str = None):
    """Plot the fitted weight model with 80% credible interval.

    Args:
        idata: ArviZ InferenceData object
        df: Original DataFrame with weight data
        stan_data: Stan data dictionary with scaling parameters
        output_path: Optional path to save plot
    """
    fig, ax = plt.subplots(1, 1, figsize=(12, 6))

    # Back-transform predictions to original scale
    y_mean = stan_data["_y_mean"]
    y_sd = stan_data["_y_sd"]

    # Extract posterior samples for f (GP function)
    f_samples = idata.posterior["f"].values

    # Compute posterior mean and 80% credible interval (10% - 90%)
    f_mean = f_samples.mean(axis=(0, 1)) * y_sd + y_mean
    f_lower = np.percentile(f_samples, 10, axis=(0, 1)) * y_sd + y_mean
    f_upper = np.percentile(f_samples, 90, axis=(0, 1)) * y_sd + y_mean

    # Plot observations
    ax.scatter(df["date"], df["weight_lbs"], alpha=0.5, s=20,
               label

def main():
    """Run model and create 80% CI plot."""
    # Fit the model (or load existing results if available)
    print("Fitting weight model...")
    fit, idata, df, stan_data = fit_weight_model()

    # Create plot with 80% CI
    print("Creating plot with 80% credible interval...")
    fig, ax = plot_weight_80ci(idata, df, stan_data, "output/weight_fit_80ci.png")

    # Show summary statistics
    print("\n" + "=" * 60)
    print("MODEL SUMMARY WITH 80% CI")
    print("=" * 60)

    # Key parameters
    summary = az.summary(idata, var_names=["alpha", "rho", "sigma", "trend_change"])
    print("\nKey parameters:")
    print(summary)

    # Back-transform trend change
    y_sd = stan_data["_y_sd"]
    trend_change = idata.posterior["trend_change"].values.mean() * y_sd
    print(f"\nTrend change (original scale): {trend_change:.2f} lbs")

    print("\nPlot saved to output/weight_fit_80ci.png")
    plt.show()

def add_zoom_capabilities(
    fig_or_ax,
    zoom_to: Optional[Tuple[float, float]] = None,
):
    """Add zooming capabilities to a matplotlib figure or axes.

    Args:
        fig_or_ax: matplotlib Figure or Axes object
        zoom_to: Optional (xmin, xmax) tuple to zoom to specific x-range.
                 If None, no programmatic zoom is applied.

    Returns:
        The input figure or axes (for chaining).
    """
    if zoom_to is not None:
        # Handle both Figure and Axes inputs
        if isinstance(fig_or_ax, mpl_fig.Figure):
            # It's a Figure
            for ax in fig_or_ax.axes:
                ax.set_xlim(zoom_to)
        else:
            # Assume it's an Axes
            fig_or_ax.set_xlim(zoom_to)

    return fig_or_ax

def zoom_to_date_range(
    ax,
    start_date: pd.Timestamp,
    end_date: pd.Timestamp,
):
    """Zoom axes to a specific date range.

    Args:
        ax: matplotlib Axes object
        start_date: Start date (pd.Timestamp or datetime-like)
        end_date: End date (pd.Timestamp or datetime-like)

    Returns:
        The input axes (for chaining).
    """
    # Convert dates to matplotlib numeric format
    start_num = date2num(start_date)
    end_num = date2num(end_date)

    ax.set_xlim(start_num, end_num)
    return ax

def link_axes(axes: List[plt.Axes]):
    """Link x-axes of multiple subplots for synchronized zooming.

    Args:
        axes: List of Axes objects to link.

    Returns:
        List of linked axes (same as input).
    """
    if len(axes) < 2:
        return axes

    # Link x-axes: sharex only works during subplot creation, so we manually sync
    # For now, just use sharex which should work if axes are from same figure
    # This is a minimal implementation
    for ax in axes[1:]:
        ax.sharex(axes[0])

    return axes

def zoom_to_preset(
    ax: plt.Axes,
    preset: str,
    reference_date: Optional[float] = None,
):
    """Zoom axes to a preset range.

    Args:
        ax: matplotlib Axes object
        preset: One of 'last_week', 'last_month', 'last_year', 'all'
        reference_date: Reference point for relative presets (e.g., latest date).
                       If None, uses current x-axis maximum.

    Returns:
        The input axes (for chaining).
    """
    # Get current x-axis limits and data range
    xlim = ax.get_xlim()
    xmin, xmax = xlim

    if reference_date is None:
        reference_date = xmax

    preset = preset.lower()
    if preset == 'all':
        # Reset to autoscale
        ax.autoscale(axis='x')
    elif preset == 'last_week':
        ax.set_xlim(reference_date - 7, reference_date)
    elif preset == 'last_month':
        ax.set_xlim(reference_date - 30, reference_date)
    elif preset == 'last_year':
        ax.set_xlim(reference_date - 365, reference_date)
    el