Squash commits for FCS version

Author: Chlorie

.gitignore

@@ -2,11 +2,12 @@ __pycache__
 .idea/
 .vscode/
 .DS_STORE
-playground.ipynb
-playground.py
-^data/
-out/
+playground*.ipynb
+playground*.py
+/data/
+/out/
 logs/
 tb_logs/
 get_baostock_data.ipynb
-get_baostock_data.py
+get_baostock_data.py
+/utils/

README.md

@@ -6,13 +6,20 @@
 
 Automatic formulaic alpha generation with reinforcement learning.
 
-Paper *Generating Synergistic Formulaic Alpha Collections via Reinforcement Learning* accepted by [KDD 2023](https://kdd.org/kdd2023/), Applied Data Science (ADS) track.
+This repository contains the code for our paper *Generating Synergistic Formulaic Alpha Collections via Reinforcement Learning* accepted by [KDD 2023](https://kdd.org/kdd2023/), Applied Data Science (ADS) track, publically available on [ACM DL](https://dl.acm.org/doi/10.1145/3580305.3599831). Some extensions upon this work are also included in this repo.
 
-Paper available on [ACM DL](https://dl.acm.org/doi/10.1145/3580305.3599831) or [arXiv](https://arxiv.org/abs/2306.12964).
+## Repository Structure
+
+- `/alphagen` contains the basic data structures and the essential modules for starting an alpha mining pipeline;
+- `/alphagen_qlib` contains the qlib-specific APIs for data preparation;
+- `/alphagen_generic` contains data structures and utils designed for our baselines, which basically follow [gplearn](https://github.com/trevorstephens/gplearn) APIs, but with modifications for quant pipeline;
+- `/alphagen_llm` contains LLM client abstractions and a set of prompts useful for LLM-based alpha generation, and also provides some LLM-based automatic iterative alpha-generation routines.
+- `/gplearn` and `/dso` contains modified versions of our baselines;
+- `/scripts` contains several scripts for running the experiments.
 
-## How to reproduce?
+## Result Reproduction
 
-Note that you can either use our builtin alpha calculation pipeline(see Choice 1), or implement an adapter to your own pipeline(see Choice 2).
+Note that you can either use our builtin alpha calculation pipeline (see Choice 1), or implement an adapter to your own pipeline (see Choice 2).
 
 ### Choice 1: Stock data preparation
 
@@ -80,13 +87,14 @@ These parameters will define a RL run:
 - save_path (Path for checkpoints)
 - tensorboard_log (Path for TensorBoard)
 
-### Run!
+### Run the experiments
 
-```shell
-python train_maskable_ppo.py --seed=SEED --pool=POOL_CAPACITY --code=INSTRUMENTS --step=NUM_STEPS
-```
+Please run the individual scripts at the root directory of this project as modules, i.e. `python -m scripts.NAME ARGS...`.
+Use `python -m scripts.NAME -h` for information on the arguments.
 
-Where `SEED` is random seed, e.g., `1` or `1,2`, `POOL_CAPACITY` is the size of combination model and, `NUM_STEPS` is the limit of RL steps.
+- `scripts/rl.py`: Main experiments of AlphaGen/HARLA
+- `scripts/llm_only.py`: Alpha generator based solely on iterative interactions with an LLM.
+- `scripts/llm_test_validity.py`: Tests on how the system prompt affects the valid alpha rate of an LLM.
 
 ### After running
 
@@ -105,13 +113,6 @@ Where `SEED` is random seed, e.g., `1` or `1,2`, `POOL_CAPACITY` is the size of
 
 [DSO](https://github.com/brendenpetersen/deep-symbolic-optimization) is a mature deep learning framework for symbolic optimization tasks. We maintained a minimal version of DSO to make it compatiable with our task. The corresponding experiment scipt is [dso.py](dso.py)
 
-## Repository Structure
-
-- `/alphagen` contains the basic data structures and the essential modules for starting an alpha mining pipeline;
-- `/alphagen_qlib` contains the qlib-specific APIs for data preparation;
-- `/alphagen_generic` contains data structures and utils designed for our baselines, which basically follow [gplearn](https://github.com/trevorstephens/gplearn) APIs, but with modifications for quant pipeline;
-- `/gplearn` and `/dso` contains modified versions of our baselines.
-
 ## Trading (Experimental)
 
 We implemented some trading strategies based on Qlib. See [backtest.py](backtest.py) and [trade_decision.py](trade_decision.py) for demos.
@@ -147,4 +148,4 @@ Thanks to the following contributors:
 
 Thanks to the following in-depth research on our project:
 
-- *因子选股系列之九十五:DFQ强化学习因子组合挖掘系统*
+- *因子选股系列之九十五: DFQ强化学习因子组合挖掘系统*

alphagen/config.py

@@ -1,10 +1,11 @@
+from typing import Type
 from alphagen.data.expression import *
 
 
-MAX_EXPR_LENGTH = 20
+MAX_EXPR_LENGTH = 15
 MAX_EPISODE_LENGTH = 256
 
-OPERATORS = [
+OPERATORS: List[Type[Operator]] = [
     # Unary
     Abs,  # Sign,
     Log,
@@ -19,7 +20,7 @@
     Cov, Corr
 ]
 
-DELTA_TIMES = [10, 20, 30, 40, 50]
+DELTA_TIMES = [1, 5, 10, 20, 40]
 
 CONSTANTS = [-30., -10., -5., -2., -1., -0.5, -0.01, 0.01, 0.5, 1., 2., 5., 10., 30.]
 

alphagen/data/calculator.py

@@ -1,7 +1,10 @@
 from abc import ABCMeta, abstractmethod
-from typing import List, Tuple
+from typing import Tuple, Optional, Sequence
+from torch import Tensor
+import torch
 
 from alphagen.data.expression import Expression
+from alphagen.utils.correlation import batch_pearsonr, batch_spearmanr
 
 
 class AlphaCalculator(metaclass=ABCMeta):
@@ -13,25 +16,111 @@ def calc_single_IC_ret(self, expr: Expression) -> float:
     def calc_single_rIC_ret(self, expr: Expression) -> float:
         'Calculate Rank IC between a single alpha and a predefined target.'
 
-    @abstractmethod
     def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]:
-        'Calculate both IC and Rank IC between a single alpha and a predefined target.'
+        return self.calc_single_IC_ret(expr), self.calc_single_rIC_ret(expr)
 
     @abstractmethod
     def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float:
         'Calculate IC between two alphas.'
 
     @abstractmethod
-    def calc_pool_IC_ret(self, exprs: List[Expression], weights: List[float]) -> float:
+    def calc_pool_IC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
         'First combine the alphas linearly,'
         'then Calculate IC between the linear combination and a predefined target.'
 
     @abstractmethod
-    def calc_pool_rIC_ret(self, exprs: List[Expression], weights: List[float]) -> float:
+    def calc_pool_rIC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
         'First combine the alphas linearly,'
         'then Calculate Rank IC between the linear combination and a predefined target.'
 
     @abstractmethod
-    def calc_pool_all_ret(self, exprs: List[Expression], weights: List[float]) -> Tuple[float, float]:
+    def calc_pool_all_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float]:
         'First combine the alphas linearly,'
         'then Calculate both IC and Rank IC between the linear combination and a predefined target.'
+
+
+class TensorAlphaCalculator(AlphaCalculator):
+    def __init__(self, target: Optional[Tensor]) -> None:
+        self._target = target
+
+    @property
+    @abstractmethod
+    def n_days(self) -> int: ...
+
+    @property
+    def target(self) -> Tensor:
+        if self._target is None:
+            raise ValueError("A target must be set before calculating non-mutual IC.")
+        return self._target
+
+    @abstractmethod
+    def evaluate_alpha(self, expr: Expression) -> Tensor:
+        'Evaluate an alpha into a `Tensor` of shape (days, stocks).'
+
+    def make_ensemble_alpha(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tensor:
+        n = len(exprs)
+        factors = [self.evaluate_alpha(exprs[i]) * weights[i] for i in range(n)]
+        return torch.sum(torch.stack(factors, dim=0), dim=0)
+
+    def _calc_IC(self, value1: Tensor, value2: Tensor) -> float:
+        return batch_pearsonr(value1, value2).mean().item()
+    
+    def _calc_rIC(self, value1: Tensor, value2: Tensor) -> float:
+        return batch_spearmanr(value1, value2).mean().item()
+    
+    def _IR_from_batch(self, batch: Tensor) -> float:
+        mean, std = batch.mean(), batch.std()
+        return (mean / std).item()
+    
+    def _calc_ICIR(self, value1: Tensor, value2: Tensor) -> float:
+        return self._IR_from_batch(batch_pearsonr(value1, value2))
+    
+    def _calc_rICIR(self, value1: Tensor, value2: Tensor) -> float:
+        return self._IR_from_batch(batch_spearmanr(value1, value2))
+
+    def calc_single_IC_ret(self, expr: Expression) -> float:
+        return self._calc_IC(self.evaluate_alpha(expr), self.target)
+    
+    def calc_single_IC_ret_daily(self, expr: Expression) -> Tensor:
+        return batch_pearsonr(self.evaluate_alpha(expr), self.target)
+
+    def calc_single_rIC_ret(self, expr: Expression) -> float:
+        return self._calc_rIC(self.evaluate_alpha(expr), self.target)
+    
+    def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]:
+        value = self.evaluate_alpha(expr)
+        target = self.target
+        return self._calc_IC(value, target), self._calc_rIC(value, target)
+
+    def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float:
+        return self._calc_IC(self.evaluate_alpha(expr1), self.evaluate_alpha(expr2))
+
+    def calc_mutual_IC_daily(self, expr1: Expression, expr2: Expression) -> Tensor:
+        return batch_pearsonr(self.evaluate_alpha(expr1), self.evaluate_alpha(expr2))
+
+    def calc_pool_IC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
+        with torch.no_grad():
+            value = self.make_ensemble_alpha(exprs, weights)
+            return self._calc_IC(value, self.target)
+
+    def calc_pool_rIC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float:
+        with torch.no_grad():
+            value = self.make_ensemble_alpha(exprs, weights)
+            return self._calc_rIC(value, self.target)
+
+    def calc_pool_all_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float]:
+        with torch.no_grad():
+            value = self.make_ensemble_alpha(exprs, weights)
+            target = self.target
+            return self._calc_IC(value, target), self._calc_rIC(value, target)
+        
+    def calc_pool_all_ret_with_ir(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float, float, float]:
+        "Returns IC, ICIR, Rank IC, Rank ICIR"
+        with torch.no_grad():
+            value = self.make_ensemble_alpha(exprs, weights)
+            target = self.target
+            ics = batch_pearsonr(value, target)
+            rics = batch_spearmanr(value, target)
+            ic_mean, ic_std = ics.mean().item(), ics.std().item()
+            ric_mean, ric_std = rics.mean().item(), rics.std().item()
+            return ic_mean, ic_mean / ic_std, ric_mean, ric_mean / ric_std

alphagen/data/exception.py

@@ -0,0 +1,2 @@
+class InvalidExpressionException(ValueError):
+    pass