Upload afres_v2.py with huggingface_hub

afres_v2.py

@@ -0,0 +1,838 @@
+"""
+AFRES v2: Agentic Factor Revision and Evaluation System
+
+A faithful adaptation of APRES's rubric-discovery pipeline to quantitative
+factor generation.  Key design choices (per user request):
+
+  1. LLM agent generates factor expressions from a discovered rubric.
+  2. A regression model (factor-value → future-return) provides the fitness
+     signal.  Fitness = –MAE (or IC) on a held-out test set.
+  3. NO LLM-as-judge.  Factor quality is determined entirely by the data.
+  4. NO QD / MAP-Elites.  The search budget is spent on rubric discovery.
+  5. The rubric-discovery loop mirrors APRES Section 3.1:
+        Propose → Generate Factors → Evaluate via Regression → Select & Refine
+     with MultiAIDE-style tree search (branching, debug-and-retry).
+
+Architecture
+────────────
+┌─────────────┐     ┌─────────────────┐     ┌──────────────────┐
+│  Rubric     │────→│  LLM Factor     │────→│  Regression      │
+│  Proposer   │     │  Generator      │     │  Evaluator       │
+│  (LLM)      │     │  (prompt → expr)│     │  (sklearn)       │
+└─────────────┘     └─────────────────┘     └──────────────────┘
+       ↑                                              │
+       │           Select & Refine (MultiAIDE tree)    │
+       └──────────────────────────────────────────────┘
+
+References
+  • APRES   – arXiv:2603.03142 (Sec. 3.1 rubric search)
+  • MultiAIDE – Zhao et al. 2025 (tree-search scaffold)
+"""
+
+import json
+import re
+import time
+import copy
+import warnings
+from dataclasses import dataclass, field
+from typing import List, Dict, Optional, Tuple, Callable
+from enum import Enum
+from collections import defaultdict
+
+import numpy as np
+import pandas as pd
+from sklearn.linear_model import Ridge
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.metrics import mean_absolute_error, r2_score
+
+warnings.filterwarnings("ignore")
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 1.  DATA STRUCTURES
+# ═════════════════════════════════════════════════════════════════════════════
+
+class SignalType(Enum):
+    PRICE_BASED = "price_based"
+    VOLUME_BASED = "volume_based"
+    FUNDAMENTAL = "fundamental"
+    TECHNICAL = "technical"
+    CROSS_SECTIONAL = "cross_sectional"
+    TIME_SERIES = "time_series"
+
+@dataclass
+class RubricItem:
+    """One actionable design principle for factor generation."""
+    id: str
+    description: str
+    # Actionable constraints passed to the generator
+    feature_focus: List[str] = field(default_factory=list)
+    preferred_ops: List[str] = field(default_factory=list)
+    time_horizon_hint: str = "any"          # short / medium / long / any
+    complexity_hint: str = "any"            # low / medium / high / any
+    weight: float = 1.0
+
+    def to_dict(self) -> Dict:
+        return {
+            "id": self.id, "description": self.description,
+            "feature_focus": self.feature_focus,
+            "preferred_ops": self.preferred_ops,
+            "time_horizon_hint": self.time_horizon_hint,
+            "complexity_hint": self.complexity_hint,
+            "weight": self.weight,
+        }
+
+@dataclass
+class FactorRubric:
+    items: List[RubricItem]
+    def to_dict(self):
+        return {"items": [i.to_dict() for i in self.items]}
+
+@dataclass
+class Factor:
+    id: str
+    expression: str
+    rubric_id: str          # which rubric variant produced this factor
+    ic: float = 0.0
+    mae: float = 0.0        # regression MAE (lower = better)
+    r2: float = 0.0         # regression R²
+    sharpe: float = 0.0
+    returns: float = 0.0
+    generation: int = 0
+    valid: bool = True      # False if expression failed to evaluate
+
+    def to_dict(self):
+        return {
+            "id": self.id, "expression": self.expression,
+            "rubric_id": self.rubric_id, "ic": self.ic,
+            "mae": self.mae, "r2": self.r2, "sharpe": self.sharpe,
+            "returns": self.returns, "generation": self.generation,
+            "valid": self.valid,
+        }
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 2.  MARKET DATA  (panel: days × stocks, pre-computed features)
+# ═════════════════════════════════════════════════════════════════════════════
+
+class MarketData:
+    """
+    Synthetic market panel with baked-in predictive structure so that
+    well-designed factors can genuinely outperform random ones.
+    """
+
+    def __init__(self, n_stocks: int = 80, n_days: int = 600, seed: int = 42):
+        rng = np.random.RandomState(seed)
+        self.n_stocks = n_stocks
+        self.n_days = n_days
+        self.dates = pd.date_range("2020-01-01", periods=n_days, freq="B")
+        self.symbols = [f"S{i:03d}" for i in range(n_stocks)]
+
+        # ── raw price / volume panels ──
+        self.close = np.zeros((n_days, n_stocks))
+        self.open_ = np.zeros((n_days, n_stocks))
+        self.high = np.zeros((n_days, n_stocks))
+        self.low = np.zeros((n_days, n_stocks))
+        self.volume = np.zeros((n_days, n_stocks))
+
+        for i in range(n_stocks):
+            # base returns
+            ret = rng.normal(0.0002, 0.018, n_days)
+
+            # ── predictable structure (weak but real) ──
+            # momentum: past 5-day ret → next day (+)
+            mom = np.zeros(n_days)
+            mom[5:] = 0.25 * ret[:-5]
+            # mean-reversion: deviation from 10-day mean → next day (–)
+            rev = np.zeros(n_days)
+            for t in range(10, n_days):
+                rev[t] = -0.15 * (ret[t - 1] - ret[t - 10:t].mean())
+            # volume-return interaction
+            vol_signal = np.zeros(n_days)
+            vol_signal[5:] = 0.10 * np.abs(ret[:-5]) * rng.lognormal(0, 0.3, n_days - 5)
+
+            ret[10:] += mom[10:] + rev[10:] + vol_signal[10:]
+
+            prices = 100 * np.exp(np.cumsum(ret))
+            self.close[:, i] = prices
+            self.open_[:, i] = prices * (1 + rng.normal(0, 0.001, n_days))
+            self.high[:, i] = prices * (1 + np.abs(rng.normal(0, 0.008, n_days)))
+            self.low[:, i] = prices * (1 - np.abs(rng.normal(0, 0.008, n_days)))
+            self.volume[:, i] = rng.lognormal(15, 0.3, n_days)
+
+        # ── pre-compute feature panels ──
+        self.ret_1d = np.diff(self.close, axis=0, prepend=self.close[:1]) / (self.close + 1e-10)
+        self.ret_5d = np.zeros_like(self.close)
+        self.ret_5d[5:] = (self.close[5:] - self.close[:-5]) / (self.close[:-5] + 1e-10)
+        self.ret_20d = np.zeros_like(self.close)
+        self.ret_20d[20:] = (self.close[20:] - self.close[:-20]) / (self.close[:-20] + 1e-10)
+
+        df_close = pd.DataFrame(self.close)
+        self.sma_5 = df_close.rolling(5, min_periods=1).mean().values
+        self.sma_10 = df_close.rolling(10, min_periods=1).mean().values
+        self.sma_20 = df_close.rolling(20, min_periods=1).mean().values
+        self.vol_20d = pd.DataFrame(self.ret_1d).rolling(20, min_periods=1).std().values
+
+        df_volume = pd.DataFrame(self.volume)
+        self.vol_sma_20 = df_volume.rolling(20, min_periods=1).mean().values
+        self.high_20d = pd.DataFrame(self.high).rolling(20, min_periods=1).max().values
+        self.low_20d = pd.DataFrame(self.low).rolling(20, min_periods=1).min().values
+        self.vwap = self.close * (1 + rng.normal(0, 0.0003, (n_days, n_stocks)))
+
+        # ── target: next-day return ──
+        self.future_ret = np.zeros_like(self.close)
+        self.future_ret[:-1] = np.diff(self.close, axis=0) / (self.close[:-1] + 1e-10)
+
+        # ── train / test split ──
+        self.train_idx = np.arange(30, 450)   # skip first 30 for feature warmup
+        self.test_idx = np.arange(450, min(580, n_days - 1))
+
+    def eval_expr(self, expr: str) -> Optional[np.ndarray]:
+        """Evaluate a factor expression → (days × stocks) array."""
+        ns = {
+            'close': self.close, 'open': self.open_, 'high': self.high,
+            'low': self.low, 'volume': self.volume, 'vwap': self.vwap,
+            'returns_1d': self.ret_1d, 'returns_5d': self.ret_5d,
+            'returns_20d': self.ret_20d, 'sma_5': self.sma_5,
+            'sma_10': self.sma_10, 'sma_20': self.sma_20,
+            'volatility_20d': self.vol_20d, 'volume_sma_20': self.vol_sma_20,
+            'high_20d': self.high_20d, 'low_20d': self.low_20d,
+            'np': np, 'abs': np.abs, 'log': np.log, 'sqrt': np.sqrt,
+            'sign': np.sign,
+            'rank': lambda x: self._rank(x),
+            'ts_mean': lambda x, w: self._ts(x, w, 'mean'),
+            'ts_std': lambda x, w: self._ts(x, w, 'std'),
+            'ts_max': lambda x, w: self._ts(x, w, 'max'),
+            'ts_min': lambda x, w: self._ts(x, w, 'min'),
+            'ts_zscore': lambda x, w: (x - self._ts(x, w, 'mean')) /
+                                        (self._ts(x, w, 'std') + 1e-10),
+            'ts_delta': lambda x, w: self._delta(x, w),
+            'ts_corr': lambda x, y, w: self._corr(x, y, w),
+            'ts_cov': lambda x, y, w: self._cov(x, y, w),
+            'ts_rank': lambda x, w: self._tsrank(x, w),
+        }
+        try:
+            result = eval(expr, {"__builtins__": {}}, ns)
+            if isinstance(result, np.ndarray) and result.shape == (self.n_days, self.n_stocks):
+                return result
+            return None
+        except Exception:
+            return None
+
+    # ---- helpers ----
+    def _rank(self, x):
+        r = np.zeros_like(x)
+        for t in range(x.shape[0]):
+            valid = np.isfinite(x[t])
+            if valid.sum() > 0:
+                r[t, valid] = pd.Series(x[t, valid]).rank(pct=True).values
+        return r
+
+    def _ts(self, x, w, method):
+        df = pd.DataFrame(x)
+        if method == 'mean':   return df.rolling(w, min_periods=1).mean().values
+        if method == 'std':    return df.rolling(w, min_periods=1).std().values
+        if method == 'max':    return df.rolling(w, min_periods=1).max().values
+        if method == 'min':    return df.rolling(w, min_periods=1).min().values
+        return x
+
+    def _delta(self, x, w):
+        out = np.zeros_like(x)
+        out[w:] = x[w:] - x[:-w]
+        return out
+
+    def _corr(self, x, y, w):
+        r = np.zeros_like(x)
+        for t in range(x.shape[0]):
+            a = x[max(0, t - w + 1):t + 1].flatten()
+            b = y[max(0, t - w + 1):t + 1].flatten()
+            if len(a) > 1 and np.std(a) > 0 and np.std(b) > 0:
+                r[t] = np.corrcoef(a, b)[0, 1]
+        return r
+
+    def _cov(self, x, y, w):
+        c = np.zeros_like(x)
+        for t in range(x.shape[0]):
+            a = x[max(0, t - w + 1):t + 1].flatten()
+            b = y[max(0, t - w + 1):t + 1].flatten()
+            c[t] = np.cov(a, b)[0, 1] if len(a) > 1 else 0.0
+        return c
+
+    def _tsrank(self, x, w):
+        r = np.zeros_like(x)
+        for t in range(x.shape[0]):
+            vals = x[max(0, t - w + 1):t + 1].flatten()
+            if len(vals) >= w:
+                r[t] = pd.Series(vals).rank(pct=True).values[-1]
+            else:
+                r[t] = np.nan
+        return r
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 3.  REGRESSION EVALUATOR  (factor value → future return)
+# ═════════════════════════════════════════════════════════════════════════════
+
+class RegressionEvaluator:
+    """
+    Trains a regression model: factor_value → next-day return.
+    Fitness is reported as:
+        • MAE  (Mean Absolute Error)  – primary metric, lower = better
+        • IC   (Spearman rank correlation) – cross-sectional predictive power
+        • R²   (coefficient of determination)
+        • Sharpe & returns from a simple long-top-quintile strategy
+    """
+
+    def __init__(self, data: MarketData, model_type: str = "ridge"):
+        self.data = data
+        self.model_type = model_type
+
+    def evaluate(self, factor: Factor) -> Dict[str, float]:
+        vals = self.data.eval_expr(factor.expression)
+        if vals is None:
+            factor.valid = False
+            return {"mae": 1e6, "ic": 0.0, "r2": -1.0, "sharpe": 0.0, "returns": 0.0}
+
+        # ── flatten panel to (sample, feature) for regression ──
+        X_train, y_train = self._flatten(vals, self.data.train_idx)
+        X_test,  y_test  = self._flatten(vals, self.data.test_idx)
+
+        if len(X_train) < 100 or len(X_test) < 50:
+            factor.valid = False
+            return {"mae": 1e6, "ic": 0.0, "r2": -1.0, "sharpe": 0.0, "returns": 0.0}
+
+        # ── train regression ──
+        model = Ridge(alpha=1.0) if self.model_type == "ridge" else \
+                RandomForestRegressor(n_estimators=50, max_depth=6, random_state=42, n_jobs=-1)
+        model.fit(X_train, y_train)
+        pred_test = model.predict(X_test)
+
+        mae = mean_absolute_error(y_test, pred_test)
+        r2 = r2_score(y_test, pred_test)
+
+        # ── IC (cross-sectional rank correlation per day) ──
+        ics = []
+        for t in self.data.test_idx:
+            f_t = vals[t]
+            r_t = self.data.future_ret[t]
+            valid = np.isfinite(f_t) & np.isfinite(r_t)
+            if valid.sum() >= 10:
+                ic = np.corrcoef(pd.Series(f_t[valid]).rank().values,
+                                  pd.Series(r_t[valid]).rank().values)[0, 1]
+                if np.isfinite(ic):
+                    ics.append(ic)
+        ic = float(np.mean(ics)) if ics else 0.0
+
+        # ── simple portfolio Sharpe ──
+        port_rets = []
+        for t in self.data.test_idx:
+            f_t = vals[t]
+            r_t = self.data.future_ret[t]
+            valid = np.isfinite(f_t) & np.isfinite(r_t)
+            if valid.sum() >= 10:
+                q80 = np.percentile(f_t[valid], 80)
+                mask = (f_t >= q80) & valid
+                if mask.sum() > 0:
+                    port_rets.append(float(np.mean(r_t[mask])))
+
+        if len(port_rets) > 2:
+            rets_arr = np.array(port_rets)
+            ann_ret = float(np.mean(rets_arr) * 252)
+            sharpe = float((np.mean(rets_arr) / (np.std(rets_arr) + 1e-10)) * np.sqrt(252))
+        else:
+            ann_ret = 0.0
+            sharpe = 0.0
+
+        factor.mae = mae
+        factor.ic = ic
+        factor.r2 = r2
+        factor.sharpe = sharpe
+        factor.returns = ann_ret
+        factor.valid = True
+        return {"mae": mae, "ic": ic, "r2": r2, "sharpe": sharpe, "returns": ann_ret}
+
+    def _flatten(self, vals: np.ndarray, idx: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """Flatten panel slices to (n_samples, 1) feature matrix + target vector."""
+        X = vals[idx].flatten().reshape(-1, 1)
+        y = self.data.future_ret[idx].flatten()
+        mask = np.isfinite(X[:, 0]) & np.isfinite(y)
+        return X[mask], y[mask]
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 4.  LLM FACTOR GENERATOR  (interface + simulated implementation)
+# ═════════════════════════════════════════════════════════════════════════════
+
+class LLMFactorGenerator:
+    """
+    Generates factor expressions from a rubric.
+    In production this would call an LLM API (GPT-4, Claude, etc.).
+    The prototype uses a template-based simulator that respects the
+    actionable constraints in each RubricItem.
+    """
+
+    # operator library
+    OPS = ["ts_mean", "ts_std", "ts_max", "ts_min", "ts_zscore",
+           "ts_delta", "ts_corr", "ts_cov", "ts_rank", "rank",
+           "abs", "sign", "log", "sqrt"]
+
+    FEATURES = ["close", "open", "high", "low", "volume", "vwap",
+                "returns_1d", "returns_5d", "returns_20d",
+                "sma_5", "sma_10", "sma_20",
+                "volatility_20d", "volume_sma_20",
+                "high_20d", "low_20d"]
+
+    WINDOWS = [3, 5, 10, 20]
+
+    def __init__(self, seed: int = 99):
+        self.rng = np.random.RandomState(seed)
+        self._counter = 0
+
+    def generate(self, rubric: FactorRubric, n: int = 5) -> List[Factor]:
+        """Generate n factor expressions constrained by the rubric."""
+        # Aggregate constraints from rubric items
+        features = self._extract_features(rubric)
+        ops = self._extract_ops(rubric)
+        complexity_target = self._extract_complexity(rubric)
+
+        factors = []
+        for _ in range(n):
+            self._counter += 1
+            expr = self._build_expression(features, ops, complexity_target)
+            factors.append(Factor(
+                id=f"f_{self._counter}",
+                expression=expr,
+                rubric_id=id(rubric),  # temporary ID, will be overwritten
+                generation=0,
+            ))
+        return factors
+
+    # ── constraint extraction ──
+    def _extract_features(self, rubric: FactorRubric) -> List[str]:
+        feats = []
+        for item in rubric.items:
+            feats.extend(item.feature_focus)
+        return list(set(feats)) if feats else self.FEATURES
+
+    def _extract_ops(self, rubric: FactorRubric) -> List[str]:
+        ops = []
+        for item in rubric.items:
+            ops.extend(item.preferred_ops)
+        return list(set(ops)) if ops else self.OPS
+
+    def _extract_complexity(self, rubric: FactorRubric) -> int:
+        hints = [item.complexity_hint for item in rubric.items]
+        if "low" in hints:
+            return 2
+        if "high" in hints:
+            return 5
+        return 3  # default medium
+
+    # ── expression builder ──
+    def _build_expression(self, features, ops, complexity_target: int) -> str:
+        """Build a random expression of roughly target complexity."""
+        n_terms = max(1, complexity_target - 1 + self.rng.randint(-1, 2))
+        terms = []
+        for _ in range(n_terms):
+            terms.append(self._random_term(features, ops))
+        if len(terms) == 1:
+            return terms[0]
+        # Combine with binary ops
+        expr = terms[0]
+        for t in terms[1:]:
+            op = self.rng.choice([" + ", " - ", " * "])
+            expr = f"({expr}){op}({t})"
+        return expr
+
+    def _random_term(self, features, ops) -> str:
+        feat = self.rng.choice(features)
+        # Decide: raw feature, windowed operator, or cross-feature interaction
+        choice = self.rng.choice(["raw", "op", "interaction"], p=[0.25, 0.50, 0.25])
+        if choice == "raw":
+            return feat
+        if choice == "op":
+            op = self.rng.choice([o for o in ops if o.startswith("ts_") or o in ("rank", "abs", "sign", "log", "sqrt")])
+            if op.startswith("ts_") and op not in ("ts_corr", "ts_cov"):
+                w = self.rng.choice(self.WINDOWS)
+                return f"{op}({feat}, {w})"
+            if op in ("rank", "abs", "sign", "log", "sqrt"):
+                return f"{op}({feat})"
+            # fallback
+            return feat
+        # interaction
+        feat2 = self.rng.choice([f for f in features if f != feat])
+        op = self.rng.choice([" * ", " + ", " - "])
+        return f"{feat}{op}{feat2}"
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 5.  RUBRIC PROPOSER  (MultiAIDE-style tree search)
+# ═════════════════════════════════════════════════════════════════════════════
+
+class RubricProposer:
+    """
+    Implements the APRES rubric-discovery loop:
+        Propose → Generate Factors → Evaluate via Regression → Select & Refine
+    With MultiAIDE-style tree search:
+        • N0 initial rubric branches
+        • Each step: branch N new variants from best or buggy (prob p_debug)
+        • Max debug depth D_max per node
+    """
+
+    # Seed rubrics to start the search
+    SEED_RUBRICS = [
+        FactorRubric([
+            RubricItem("price_momentum", "Prefer price-based momentum signals",
+                       feature_focus=["close", "returns_1d", "returns_5d"],
+                       preferred_ops=["ts_mean", "ts_delta", "rank"],
+                       time_horizon_hint="short", complexity_hint="low"),
+        ]),
+        FactorRubric([
+            RubricItem("volume_confirm", "Use volume to confirm price signals",
+                       feature_focus=["volume", "volume_sma_20", "returns_1d"],
+                       preferred_ops=["ts_corr", "rank", "ts_mean"],
+                       time_horizon_hint="medium", complexity_hint="medium"),
+        ]),
+        FactorRubric([
+            RubricItem("mean_reversion", "Capture mean-reversion patterns",
+                       feature_focus=["close", "sma_20", "sma_10"],
+                       preferred_ops=["ts_zscore", "ts_delta", "abs"],
+                       time_horizon_hint="short", complexity_hint="medium"),
+        ]),
+    ]
+
+    def __init__(self,
+                 N0: int = 3,
+                 N: int = 3,
+                 p_debug: float = 0.3,
+                 D_max: int = 5,
+                 seed: int = 77):
+        self.N0 = N0
+        self.N = N
+        self.p_debug = p_debug
+        self.D_max = D_max
+        self.rng = np.random.RandomState(seed)
+        self._debug_counts = defaultdict(int)   # node_id → debug attempts
+
+    def propose_initial(self) -> List[FactorRubric]:
+        """N0 diverse seed rubrics."""
+        rubrics = [copy.deepcopy(r) for r in self.SEED_RUBRICS[:self.N0]]
+        # Pad with randomised variants if fewer seeds than N0
+        while len(rubrics) < self.N0:
+            base = copy.deepcopy(self.rng.choice(self.SEED_RUBRICS))
+            rubrics.append(self._mutate_rubric(base))
+        return rubrics
+
+    def propose_from_parent(self, parent: FactorRubric, is_buggy: bool = False) -> List[FactorRubric]:
+        """Branch N new rubric variants from a parent."""
+        children = []
+        for _ in range(self.N):
+            child = self._mutate_rubric(copy.deepcopy(parent))
+            if is_buggy:
+                # More aggressive mutation for debug branches
+                child = self._mutate_rubric(child)
+            children.append(child)
+        return children
+
+    # ── mutation operators ──
+    def _mutate_rubric(self, rubric: FactorRubric) -> FactorRubric:
+        """Apply one structural mutation to a rubric."""
+        items = rubric.items
+        mutation = self.rng.choice(["add", "remove", "replace", "tweak"])
+
+        if mutation == "add" or len(items) == 0:
+            items.append(self._random_item())
+        elif mutation == "remove" and len(items) > 1:
+            items.pop(self.rng.randint(0, len(items)))
+        elif mutation == "replace" and len(items) > 0:
+            idx = self.rng.randint(0, len(items))
+            items[idx] = self._random_item()
+        elif mutation == "tweak" and len(items) > 0:
+            idx = self.rng.randint(0, len(items))
+            items[idx] = self._tweak_item(items[idx])
+
+        return FactorRubric(items)
+
+    def _random_item(self) -> RubricItem:
+        templates = [
+            ("momentum", "Focus on momentum signals",
+             ["close", "returns_5d", "returns_20d"], ["ts_mean", "ts_delta", "rank"]),
+            ("volatility", "Exploit volatility patterns",
+             ["volatility_20d", "returns_1d", "close"], ["ts_zscore", "ts_std", "abs"]),
+            ("volume_price", "Volume-price interaction",
+             ["volume", "close", "returns_1d"], ["ts_corr", "rank", "ts_mean"]),
+            ("cross_sectional", "Cross-sectional ranking",
+             ["close", "volume", "returns_1d"], ["rank", "ts_zscore", "ts_rank"]),
+            ("mean_reversion", "Mean-reversion",
+             ["close", "sma_10", "sma_20"], ["ts_zscore", "ts_delta", "sign"]),
+            ("breakout", "Breakout patterns",
+             ["high_20d", "low_20d", "close"], ["ts_max", "ts_min", "ts_delta"]),
+            ("vwap", "VWAP deviation",
+             ["vwap", "close", "volume"], ["ts_zscore", "ts_mean", "ts_corr"]),
+        ]
+        t = templates[self.rng.randint(0, len(templates))]
+        return RubricItem(
+            id=t[0], description=t[1],
+            feature_focus=t[2], preferred_ops=t[3],
+            time_horizon_hint=self.rng.choice(["short", "medium", "long"]),
+            complexity_hint=self.rng.choice(["low", "medium", "high"]),
+            weight=1.0,
+        )
+
+    def _tweak_item(self, item: RubricItem) -> RubricItem:
+        """Small perturbation of one rubric item."""
+        tweak = self.rng.choice(["feature", "op", "horizon", "complexity"])
+        if tweak == "feature" and item.feature_focus:
+            all_feats = ["close", "volume", "returns_1d", "returns_5d", "vwap",
+                         "sma_10", "sma_20", "volatility_20d", "high_20d", "low_20d"]
+            item.feature_focus = list(set(
+                item.feature_focus + [self.rng.choice(all_feats)]
+            ))[:3]
+        elif tweak == "op" and item.preferred_ops:
+            all_ops = ["ts_mean", "ts_std", "ts_zscore", "ts_delta", "rank",
+                       "ts_corr", "ts_rank", "abs", "sign", "log"]
+            item.preferred_ops = list(set(
+                item.preferred_ops + [self.rng.choice(all_ops)]
+            ))[:3]
+        elif tweak == "horizon":
+            item.time_horizon_hint = self.rng.choice(["short", "medium", "long"])
+        elif tweak == "complexity":
+            item.complexity_hint = self.rng.choice(["low", "medium", "high"])
+        return item
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 6.  MAIN AFRES SYSTEM  (Rubric Discovery Loop)
+# ═════════════════════════════════════════════════════════════════════════════
+
+class AFRES:
+    """
+    Agentic Factor Revision and Evaluation System.
+
+    Phase 1 – Rubric Discovery (faithful to APRES §3.1):
+        1. Propose:      RubricProposer creates rubric variants
+        2. Generate:     LLMFactorGenerator produces factors per rubric
+        3. Evaluate:     RegressionEvaluator trains model, reports MAE
+        4. Select&Refine: MultiAIDE tree search keeps / branches best rubrics
+
+    Phase 2 – Best-Rubric Factor Generation:
+        Use the discovered rubric to generate a final pool of factors.
+    """
+
+    def __init__(self,
+                 data: MarketData,
+                 generator: LLMFactorGenerator,
+                 evaluator: RegressionEvaluator,
+                 proposer: RubricProposer,
+                 factors_per_rubric: int = 4):
+        self.data = data
+        self.generator = generator
+        self.evaluator = evaluator
+        self.proposer = proposer
+        self.k = factors_per_rubric
+
+        # Search state
+        self.best_rubric: Optional[FactorRubric] = None
+        self.best_fitness: float = -1e9
+        self.all_factors: List[Factor] = []
+        self.history: List[Dict] = []
+
+    def discover_rubric(self, max_iterations: int = 20) -> Tuple[FactorRubric, List[Factor]]:
+        """
+        Run the APRES-style rubric-discovery loop.
+
+        Returns the best discovered rubric and the best factors found under it.
+        """
+        print("=" * 70)
+        print("  PHASE 1: RUBRIC DISCOVERY (APRES-style MultiAIDE search)")
+        print("=" * 70)
+
+        # ── initial population ──
+        population = self.proposer.propose_initial()
+        scores = []
+        for rubric in population:
+            fitness, factors = self._evaluate_rubric(rubric)
+            scores.append((fitness, rubric, factors))
+            self.history.append({
+                "iteration": 0, "rubric_id": id(rubric),
+                "fitness": fitness, "n_factors": len(factors),
+                "mean_ic": np.mean([f.ic for f in factors]) if factors else 0.0,
+                "mean_mae": np.mean([f.mae for f in factors]) if factors else 1e6,
+            })
+
+        # Track best
+        best = max(scores, key=lambda x: x[0])
+        self.best_fitness, self.best_rubric, best_factors = best
+        print(f"\nInitial best fitness = {self.best_fitness:.4f}  "
+              f"(IC={np.mean([f.ic for f in best_factors]):.4f}, "
+              f"MAE={np.mean([f.mae for f in best_factors]):.4f})")
+
+        # ── iterative tree search ──
+        for it in range(1, max_iterations + 1):
+            t0 = time.time()
+
+            # Pick parent: best with prob (1-p_debug), else random buggy
+            if self.proposer.rng.random() > self.proposer.p_debug:
+                parent = self.best_rubric
+                is_buggy = False
+            else:
+                # Pick a random previously-evaluated rubric (could be bad)
+                parent = self.proposer.rng.choice([s[1] for s in scores])
+                is_buggy = True
+
+            children = self.proposer.propose_from_parent(parent, is_buggy)
+
+            child_scores = []
+            for child in children:
+                fitness, factors = self._evaluate_rubric(child)
+                child_scores.append((fitness, child, factors))
+                self.history.append({
+                    "iteration": it, "rubric_id": id(child),
+                    "fitness": fitness, "n_factors": len(factors),
+                    "mean_ic": np.mean([f.ic for f in factors]) if factors else 0.0,
+                    "mean_mae": np.mean([f.mae for f in factors]) if factors else 1e6,
+                    "buggy_branch": is_buggy,
+                })
+
+            # Update global best
+            local_best = max(child_scores, key=lambda x: x[0])
+            if local_best[0] > self.best_fitness:
+                self.best_fitness = local_best[0]
+                self.best_rubric = local_best[1]
+                best_factors = local_best[2]
+                print(f"  Iter {it:02d}: NEW BEST fitness={self.best_fitness:.4f}  "
+                      f"IC={np.mean([f.ic for f in best_factors]):.4f}  "
+                      f"MAE={np.mean([f.mae for f in best_factors]):.4f}  "
+                      f"({time.time()-t0:.1f}s)")
+            else:
+                print(f"  Iter {it:02d}: no improvement  "
+                      f"best={self.best_fitness:.4f}  "
+                      f"({time.time()-t0:.1f}s)")
+
+            scores.extend(child_scores)
+
+        print(f"\n{'='*70}")
+        print(f"  RUBRIC DISCOVERY COMPLETE  –  best fitness = {self.best_fitness:.4f}")
+        print(f"{'='*70}")
+        self._print_rubric(self.best_rubric)
+        return self.best_rubric, best_factors
+
+    def _evaluate_rubric(self, rubric: FactorRubric) -> Tuple[float, List[Factor]]:
+        """
+        Evaluate a rubric:
+            1. Generate k factor expressions
+            2. Evaluate each on market data via regression
+            3. Aggregate fitness (mean of valid factors)
+        """
+        rubric_id = f"rubric_{id(rubric)}"
+        factors = self.generator.generate(rubric, n=self.k)
+        for f in factors:
+            f.rubric_id = rubric_id
+
+        valid_factors = []
+        for f in factors:
+            self.evaluator.evaluate(f)
+            self.all_factors.append(f)
+            if f.valid:
+                valid_factors.append(f)
+
+        if not valid_factors:
+            return -1e6, factors  # severe penalty for rubric that generates no valid factors
+
+        # Fitness = mean IC (higher = better)
+        # Could also use -MAE or a blend; IC is cleaner for ranking
+        fitness = float(np.mean([f.ic for f in valid_factors]))
+        return fitness, factors
+
+    def generate_final_pool(self, n: int = 20) -> List[Factor]:
+        """Generate a large factor pool from the best discovered rubric."""
+        if self.best_rubric is None:
+            raise RuntimeError("Run discover_rubric() first")
+        print(f"\n{'='*70}")
+        print(f"  PHASE 2: FINAL FACTOR POOL  (best rubric, n={n})")
+        print(f"{'='*70}")
+        factors = self.generator.generate(self.best_rubric, n=n)
+        for f in factors:
+            f.rubric_id = "best"
+            self.evaluator.evaluate(f)
+            self.all_factors.append(f)
+
+        valid = [f for f in factors if f.valid]
+        valid.sort(key=lambda f: f.ic, reverse=True)
+        return valid
+
+    @staticmethod
+    def _print_rubric(rubric: FactorRubric):
+        print("\nBest discovered rubric:")
+        for item in rubric.items:
+            print(f"  • {item.id}: {item.description}")
+            print(f"    features={item.feature_focus}  ops={item.preferred_ops}  "
+                  f"horizon={item.time_horizon_hint}  complexity={item.complexity_hint}")
+
+    def summary(self) -> Dict:
+        """Produce a JSON-serialisable summary."""
+        valid_factors = [f for f in self.all_factors if f.valid]
+        if not valid_factors:
+            return {}
+
+        top = max(valid_factors, key=lambda f: f.ic)
+        return {
+            "best_rubric": self.best_rubric.to_dict() if self.best_rubric else None,
+            "best_fitness": float(self.best_fitness),
+            "total_factors_generated": len(self.all_factors),
+            "valid_factors": len(valid_factors),
+            "top_factor": top.to_dict(),
+            "mean_ic": float(np.mean([f.ic for f in valid_factors])),
+            "mean_mae": float(np.mean([f.mae for f in valid_factors])),
+            "mean_sharpe": float(np.mean([f.sharpe for f in valid_factors])),
+            "search_history": self.history,
+        }
+
+# ═════════════════════════════════════════════════════════════════════════════
+# 7.  DRIVER
+# ═════════════════════════════════════════════════════════════════════════════
+
+def main():
+    print("\n" + "=" * 70)
+    print("  AFRES v2: Agentic Factor Revision and Evaluation System")
+    print("  (Faithful APRES rubric-discovery adapted to factor generation)")
+    print("=" * 70 + "\n")
+
+    t0 = time.time()
+
+    # 1. Market data
+    print("[1/4]  Loading market data …")
+    data = MarketData(n_stocks=80, n_days=600, seed=42)
+
+    # 2. Components
+    print("[2/4]  Initialising components …")
+    generator = LLMFactorGenerator(seed=99)
+    evaluator = RegressionEvaluator(data, model_type="ridge")
+    proposer = RubricProposer(N0=3, N=3, p_debug=0.3, D_max=5, seed=77)
+
+    # 3. AFRES system
+    afres = AFRES(data, generator, evaluator, proposer, factors_per_rubric=4)
+
+    # 4. Run rubric discovery
+    print("[3/4]  Running rubric-discovery loop …")
+    best_rubric, best_factors = afres.discover_rubric(max_iterations=15)
+
+    # 5. Final pool
+    print("\n[4/4]  Generating final factor pool …")
+    final_pool = afres.generate_final_pool(n=20)
+
+    # 6. Report
+    print("\n" + "=" * 70)
+    print("  FINAL REPORT")
+    print("=" * 70)
+    print(f"\nTotal factors evaluated: {len(afres.all_factors)}")
+    print(f"Valid factors: {len([f for f in afres.all_factors if f.valid])}")
+    print(f"\nTop 5 factors by IC:")
+    for f in final_pool[:5]:
+        print(f"  {f.id:8s}  IC={f.ic:+.4f}  MAE={f.mae:.4f}  "
+              f"Sharpe={f.sharpe:6.2f}  {f.expression}")
+
+    # 7. Save
+    results = afres.summary()
+    with open("/app/afres_v2_results.json", "w") as fp:
+        json.dump(results, fp, indent=2, default=str)
+    print(f"\nResults saved to /app/afres_v2_results.json")
+    print(f"Total wall time: {time.time()-t0:.1f}s")
+
+
+if __name__ == "__main__":
+    main()