"""Quantile regression models for heterogeneous price effect analysis.""" import logging from pathlib import Path from typing import Dict, List, Optional, Sequence, Tuple import numpy as np import pandas as pd import statsmodels.api as sm from statsmodels.regression.quantile_regression import QuantReg logger = logging.getLogger(__name__) DEFAULT_TAUS: List[float] = [0.10, 0.25, 0.50, 0.75, 0.90] def _available_cols(df: pd.DataFrame, requested: List[str]) -> List[str]: """Return the subset of requested columns present in df.""" present = [c for c in requested if c in df.columns] missing = set(requested) - set(present) if missing: logger.warning("Missing columns (skipped): %s", sorted(missing)) return present def _prepare_data( df: pd.DataFrame, y_col: str, x_cols: List[str], ) -> Tuple[pd.Series, pd.DataFrame]: """Drop NaN rows, add constant, return (y, X).""" cols = _available_cols(df, x_cols) if not cols: raise ValueError("No valid feature columns found in the DataFrame.") subset = df[[y_col] + cols].dropna() if subset.empty: raise ValueError("No observations remaining after dropping NaN rows.") y = subset[y_col].astype(float) X = sm.add_constant(subset[cols].astype(float)) return y, X # --------------------------------------------------------------------------- # Single quantile estimation # --------------------------------------------------------------------------- def estimate_quantile( df: pd.DataFrame, y_col: str, x_cols: List[str], tau: float, max_iter: int = 10_000, ) -> sm.regression.quantile_regression.QuantRegResults: """Single quantile regression at quantile tau.""" if not 0 < tau < 1: raise ValueError(f"tau must be in (0, 1), got {tau}") y, X = _prepare_data(df, y_col, x_cols) model = QuantReg(y, X) results = model.fit(q=tau, max_iter=max_iter) logger.info( "Quantile regression (tau=%.2f): pseudo-R2=%.4f, n=%d", tau, results.prsquared, int(results.nobs), ) return results # --------------------------------------------------------------------------- # Quantile sequence # --------------------------------------------------------------------------- def estimate_quantile_sequence( df: pd.DataFrame, y_col: str, x_cols: List[str], taus: Optional[List[float]] = None, max_iter: int = 10_000, ) -> Dict[float, sm.regression.quantile_regression.QuantRegResults]: """Estimate quantile regressions at multiple quantiles.""" if taus is None: taus = DEFAULT_TAUS results: Dict[float, sm.regression.quantile_regression.QuantRegResults] = {} for tau in sorted(taus): try: res = estimate_quantile(df, y_col, x_cols, tau, max_iter=max_iter) results[tau] = res except Exception as exc: logger.error("Quantile tau=%.2f failed: %s", tau, exc) logger.info("Estimated %d / %d quantile regressions.", len(results), len(taus)) return results # --------------------------------------------------------------------------- # Coefficient comparison across quantiles # --------------------------------------------------------------------------- def compare_quantile_coefficients( results_dict: Dict[float, sm.regression.quantile_regression.QuantRegResults], ) -> pd.DataFrame: """Extract coefficients and CIs across quantiles into a long DataFrame.""" records = [] for tau, res in sorted(results_dict.items()): ci = res.conf_int() for var in res.params.index: records.append({ "tau": tau, "variable": var, "coefficient": res.params[var], "std_error": res.bse[var], "p_value": res.pvalues[var], "ci_lower": ci.loc[var, 0], "ci_upper": ci.loc[var, 1], }) coef_df = pd.DataFrame(records) logger.info( "Coefficient comparison: %d variables x %d quantiles.", coef_df["variable"].nunique(), coef_df["tau"].nunique(), ) return coef_df # --------------------------------------------------------------------------- # Plotting # --------------------------------------------------------------------------- def plot_quantile_coefficients( coef_df: pd.DataFrame, variables: List[str], save_path: Optional[str] = None, figsize: Tuple[int, int] = (14, 4), ols_results: Optional[object] = None, ) -> None: """Plot coefficient evolution across quantiles with confidence bands.""" try: import matplotlib.pyplot as plt except ImportError: logger.error("matplotlib is required for plotting.") return n_vars = len(variables) if n_vars == 0: logger.warning("No variables specified for plotting.") return n_cols = min(3, n_vars) n_rows = int(np.ceil(n_vars / n_cols)) fig, axes = plt.subplots(n_rows, n_cols, figsize=(figsize[0], figsize[1] * n_rows)) if n_vars == 1: axes = np.array([axes]) axes = np.atleast_1d(axes).flatten() for idx, var in enumerate(variables): ax = axes[idx] sub = coef_df[coef_df["variable"] == var].sort_values("tau") if sub.empty: ax.set_title(f"{var} (no data)") continue taus = sub["tau"].values coefs = sub["coefficient"].values ci_lo = sub["ci_lower"].values ci_hi = sub["ci_upper"].values ax.plot(taus, coefs, "o-", color="steelblue", linewidth=2, label="QR coef") ax.fill_between(taus, ci_lo, ci_hi, alpha=0.2, color="steelblue", label="95% CI") if ols_results is not None and var in ols_results.params.index: ols_b = ols_results.params[var] ax.axhline(ols_b, color="firebrick", linestyle="--", linewidth=1, label="OLS") ax.axhline(0, color="grey", linestyle=":", linewidth=0.5) ax.set_xlabel("Quantile (tau)") ax.set_ylabel("Coefficient") ax.set_title(var) ax.legend(fontsize=7) # Hide unused axes for idx in range(n_vars, len(axes)): axes[idx].set_visible(False) fig.tight_layout() if save_path: Path(save_path).parent.mkdir(parents=True, exist_ok=True) fig.savefig(save_path, dpi=200, bbox_inches="tight") logger.info("Quantile coefficient plot saved to %s", save_path) plt.close(fig) # --------------------------------------------------------------------------- # Segment effects summary # --------------------------------------------------------------------------- def get_segment_effects( results_dict: Dict[float, sm.regression.quantile_regression.QuantRegResults], ) -> pd.DataFrame: """Summarize how effects differ across price segments (budget vs. luxury).""" coef_df = compare_quantile_coefficients(results_dict) segments = { "budget (q10)": 0.10, "lower-mid (q25)": 0.25, "median (q50)": 0.50, "upper-mid (q75)": 0.75, "luxury (q90)": 0.90, } available_taus = set(coef_df["tau"].unique()) records = [] variables = [v for v in coef_df["variable"].unique() if v != "const"] for var in variables: var_data = coef_df[coef_df["variable"] == var] row: Dict[str, object] = {"variable": var} for seg_name, tau in segments.items(): if tau in available_taus: match = var_data[var_data["tau"] == tau] if not match.empty: b = match["coefficient"].values[0] p = match["p_value"].values[0] star = "***" if p < 0.001 else ("**" if p < 0.01 else ("*" if p < 0.05 else "")) row[seg_name] = f"{b:.4f}{star}" else: row[seg_name] = "" else: row[seg_name] = "" # Heterogeneity indicator: range of coefficients coefs = var_data["coefficient"].values if len(coefs) >= 2: row["coef_range"] = round(float(np.max(coefs) - np.min(coefs)), 4) row["heterogeneity"] = ( "high" if row["coef_range"] > 0.1 else ("moderate" if row["coef_range"] > 0.03 else "low") ) else: row["coef_range"] = np.nan row["heterogeneity"] = "n/a" records.append(row) summary = pd.DataFrame(records) logger.info("Segment effects summarised for %d variables.", len(summary)) return summary