Upload alpha_factory/deterministic/proven_templates.py

alpha_factory/deterministic/proven_templates.py

@@ -34,7 +34,7 @@ def _get_bf(field: BrainField) -> str:
     return field.id
 
 
-def generate_alpha15_variant(field: BrainField, group_key: str = None) -> str:
+def generate_alpha15_variant(field: BrainField, group_key: str = None, decay: int = 5) -> str:
     """
     Alpha 15 archetype: intraday position proxy + value/momentum leg.
     
@@ -57,12 +57,12 @@ def generate_alpha15_variant(field: BrainField, group_key: str = None) -> str:
         f"rank("
         f"zscore(ts_rank((high + low) / 2 - close, 252)) "
         f"+ {sign}zscore(ts_rank({field_ref}, 252))"
-        f"), {group_key}), 5)"
+        f"), {group_key}), {decay})"
     )
     return expr
 
 
-def generate_alpha6_variant(field: BrainField, group_key: str = None) -> str:
+def generate_alpha6_variant(field: BrainField, group_key: str = None, decay: int = 5) -> str:
     """
     Alpha 6 archetype: microstructure score + value leg.
     
@@ -84,12 +84,12 @@ def generate_alpha6_variant(field: BrainField, group_key: str = None) -> str:
         f"rank("
         f"rank((vwap - close) / close) * rank(volume / ts_mean(volume, 20)) "
         f"+ {sign}zscore(ts_rank({field_ref}, 252))"
-        f"), {group_key}), 5)"
+        f"), {group_key}), {decay})"
     )
     return expr
 
 
-def generate_pure_field_rank(field: BrainField, horizon: int = 252, group_key: str = None) -> str:
+def generate_pure_field_rank(field: BrainField, horizon: int = 252, group_key: str = None, decay: int = 5) -> str:
     """
     Simplest possible expression — just rank the field cross-sectionally.
     Often works for high-quality pre-computed model fields (model77).
@@ -104,12 +104,12 @@ def generate_pure_field_rank(field: BrainField, horizon: int = 252, group_key: s
         f"ts_decay_linear("
         f"group_neutralize("
         f"{sign}rank(ts_rank({field_ref}, {horizon}))"
-        f", {group_key}), 5)"
+        f", {group_key}), {decay})"
     )
     return expr
 
 
-def generate_delta_momentum(field: BrainField, horizon: int = 21, group_key: str = None) -> str:
+def generate_delta_momentum(field: BrainField, horizon: int = 21, group_key: str = None, decay: int = 5) -> str:
     """
     Delta/change momentum — buy stocks where the field is improving.
     Works well for score-type fields (model16 derivatives, earnings revisions).
@@ -124,12 +124,12 @@ def generate_delta_momentum(field: BrainField, horizon: int = 21, group_key: str
         f"ts_decay_linear("
         f"group_neutralize("
         f"{sign}zscore(ts_delta({field_ref}, {horizon}))"
-        f", {group_key}), 5)"
+        f", {group_key}), {decay})"
     )
     return expr
 
 
-def generate_mean_reversion(field: BrainField, horizon: int = 20, group_key: str = None) -> str:
+def generate_mean_reversion(field: BrainField, horizon: int = 20, group_key: str = None, decay: int = 5) -> str:
     """
     Mean reversion on a field — short stocks where field is high, long where low.
     Effective for contrarian indicators (option PCR, social buzz).
@@ -151,16 +151,20 @@ def generate_mean_reversion(field: BrainField, horizon: int = 20, group_key: str
         f"ts_decay_linear("
         f"group_neutralize("
         f"{sign}zscore(ts_rank({field_ref}, {horizon}))"
-        f", {group_key}), 5)"
+        f", {group_key}), {decay})"
     )
     return expr
 
 
-def generate_batch_from_proven_templates(count: int = 5) -> list[dict]:
+def generate_batch_from_proven_templates(count: int = 5, decay: int = 5) -> list[dict]:
     """
     Generate a batch of alphas using PROVEN templates with novel fields.
     This is the PRIMARY generation method — guaranteed valid structure.
     
+    Args:
+        count: Number of alphas to generate.
+        decay: Decay parameter for ts_decay_linear (default 5).
+    
     Returns list of dicts with expression, field, template, group_key.
     """
     # Priority: goldmine (AC=0) first, then AC=1, then AC≤5
@@ -191,7 +195,7 @@ def generate_batch_from_proven_templates(count: int = 5) -> list[dict]:
         template_name, template_fn = templates[i % len(templates)]
         group_key = get_group_for_expression(prefer_novel=True)
         
-        expr = template_fn(field, group_key)
+        expr = template_fn(field, group_key, decay)
         
         results.append({
             "expression": expr,
@@ -202,7 +206,7 @@ def generate_batch_from_proven_templates(count: int = 5) -> list[dict]:
             "theme": "proven_template",
             "archetype": template_name,
             "sign": field.sign.value,
-            "decay": 5,
+            "decay": decay,
             "neutralization": group_key,
         })