├── .gitignore ├── README.md ├── alphagen ├── config.py ├── data │ ├── calculator.py │ ├── exception.py │ ├── expression.py │ ├── parser.py │ ├── pool_update.py │ ├── tokens.py │ └── tree.py ├── models │ ├── alpha_pool.py │ └── linear_alpha_pool.py ├── rl │ ├── env │ │ ├── core.py │ │ └── wrapper.py │ └── policy.py ├── trade │ ├── base.py │ └── strategy.py └── utils │ ├── __init__.py │ ├── correlation.py │ ├── logging.py │ ├── maybe.py │ ├── misc.py │ ├── pytorch_utils.py │ └── random.py ├── alphagen_generic ├── features.py └── operators.py ├── alphagen_llm ├── client │ ├── __init__.py │ ├── base.py │ ├── llama_cpp_client.py │ ├── openai_client.py │ └── repl.py └── prompts │ ├── common.py │ ├── interaction.py │ └── system_prompt.py ├── alphagen_qlib ├── calculator.py ├── stock_data.py ├── strategy.py └── utils.py ├── backtest.py ├── data_collection ├── baostock_utils.py ├── fetch_baostock_data.py └── qlib_dump_bin.py ├── dso.py ├── dso ├── __init__.py ├── checkpoint.py ├── config │ ├── __init__.py │ ├── config_common.json │ └── config_regression.json ├── const.py ├── core.py ├── cyfunc.pyx ├── execute.py ├── functions.py ├── library.py ├── logeval.py ├── memory.py ├── policy │ ├── __init__.py │ ├── policy.py │ └── rnn_policy.py ├── policy_optimizer │ ├── __init__.py │ ├── pg_policy_optimizer.py │ ├── policy_optimizer.py │ ├── ppo_policy_optimizer.py │ └── pqt_policy_optimizer.py ├── prior.py ├── program.py ├── run.py ├── setup.py ├── subroutines.py ├── task │ ├── __init__.py │ ├── regression │ │ ├── __init__.py │ │ ├── benchmarks.csv │ │ ├── dataset.py │ │ ├── function_sets.csv │ │ ├── mat_mult_benchmark.py │ │ ├── polyfit.py │ │ ├── regression.py │ │ ├── sklearn.py │ │ └── test_sklearn.py │ └── task.py ├── tf_state_manager.py ├── train.py ├── train_stats.py ├── utils.py └── variance.py ├── gp.py ├── gplearn ├── __init__.py ├── _program.py ├── fitness.py ├── functions.py ├── genetic.py ├── tests │ ├── __init__.py │ ├── test_estimator_checks.py │ ├── test_examples.py │ ├── test_fitness.py │ ├── test_functions.py │ ├── test_genetic.py │ └── test_utils.py └── utils.py ├── images └── logo.jpg ├── requirements.txt ├── scripts ├── llm_only.py ├── llm_test_validity.py └── rl.py └── trade_decision.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | .idea/ 3 | .vscode/ 4 | .DS_STORE 5 | playground*.ipynb 6 | playground*.py 7 | /data/ 8 | /out/ 9 | logs/ 10 | tb_logs/ 11 | get_baostock_data.ipynb 12 | get_baostock_data.py 13 | /utils/ 14 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AlphaGen 2 | 3 |

4 | 5 |

6 | 7 | Automatic formulaic alpha generation with reinforcement learning. 8 | 9 | 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. 10 | 11 | ## Repository Structure 12 | 13 | - `/alphagen` contains the basic data structures and the essential modules for starting an alpha mining pipeline; 14 | - `/alphagen_qlib` contains the qlib-specific APIs for data preparation; 15 | - `/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; 16 | - `/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. 17 | - `/gplearn` and `/dso` contains modified versions of our baselines; 18 | - `/scripts` contains several scripts for running the experiments. 19 | 20 | ## Result Reproduction 21 | 22 | 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). 23 | 24 | ### Choice 1: Stock data preparation 25 | 26 | Builtin pipeline requires Qlib library and local-storaged stock data. 27 | 28 | - READ THIS! We need some of the metadata (but not the actual stock price/volume data) given by Qlib, so follow the data preparing process in [Qlib](https://github.com/microsoft/qlib#data-preparation) first. 29 | - The actual stock data we use are retrieved from [baostock](http://baostock.com/baostock/index.php/%E9%A6%96%E9%A1%B5), due to concerns on the timeliness and truthfulness of the data source used by Qlib. 30 | - The data can be downloaded by running the script `data_collection/fetch_baostock_data.py`. The newly downloaded data is saved into `~/.qlib/qlib_data/cn_data_baostock_fwdadj` by default. This path can be customized to fit your specific needs, but make sure to use the correct path when loading the data (In `alphagen_qlib/stock_data.py`, function `StockData._init_qlib`, the path should be passed to qlib with `qlib.init(provider_uri=path)`). 31 | 32 | ### Choice 2: Adapt to external pipelines 33 | 34 | Maybe you have better implements of alpha calculation, you can implement an adapter of `alphagen.data.calculator.AlphaCalculator`. The interface is defined as follows: 35 | 36 | ```python 37 | class AlphaCalculator(metaclass=ABCMeta): 38 | @abstractmethod 39 | def calc_single_IC_ret(self, expr: Expression) -> float: 40 | 'Calculate IC between a single alpha and a predefined target.' 41 | 42 | @abstractmethod 43 | def calc_single_rIC_ret(self, expr: Expression) -> float: 44 | 'Calculate Rank IC between a single alpha and a predefined target.' 45 | 46 | @abstractmethod 47 | def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]: 48 | 'Calculate both IC and Rank IC between a single alpha and a predefined target.' 49 | 50 | @abstractmethod 51 | def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float: 52 | 'Calculate IC between two alphas.' 53 | 54 | @abstractmethod 55 | def calc_pool_IC_ret(self, exprs: List[Expression], weights: List[float]) -> float: 56 | 'First combine the alphas linearly,' 57 | 'then Calculate IC between the linear combination and a predefined target.' 58 | 59 | @abstractmethod 60 | def calc_pool_rIC_ret(self, exprs: List[Expression], weights: List[float]) -> float: 61 | 'First combine the alphas linearly,' 62 | 'then Calculate Rank IC between the linear combination and a predefined target.' 63 | 64 | @abstractmethod 65 | def calc_pool_all_ret(self, exprs: List[Expression], weights: List[float]) -> Tuple[float, float]: 66 | 'First combine the alphas linearly,' 67 | 'then Calculate both IC and Rank IC between the linear combination and a predefined target.' 68 | ``` 69 | 70 | Reminder: the values evaluated from different alphas may have drastically different scales, we recommend that you should normalize them before combination. 71 | 72 | ### Before running 73 | 74 | All principle components of our expriment are located in [train_maskable_ppo.py](train_maskable_ppo.py). 75 | 76 | These parameters may help you build an `AlphaCalculator`: 77 | 78 | - instruments (Set of instruments) 79 | - start_time & end_time (Data range for each dataset) 80 | - target (Target stock trend, e.g., 20d return rate) 81 | 82 | These parameters will define a RL run: 83 | 84 | - batch_size (PPO batch size) 85 | - features_extractor_kwargs (Arguments for LSTM shared net) 86 | - device (PyTorch device) 87 | - save_path (Path for checkpoints) 88 | - tensorboard_log (Path for TensorBoard) 89 | 90 | ### Run the experiments 91 | 92 | Please run the individual scripts at the root directory of this project as modules, i.e. `python -m scripts.NAME ARGS...`. 93 | Use `python -m scripts.NAME -h` for information on the arguments. 94 | 95 | - `scripts/rl.py`: Main experiments of AlphaGen/HARLA 96 | - `scripts/llm_only.py`: Alpha generator based solely on iterative interactions with an LLM. 97 | - `scripts/llm_test_validity.py`: Tests on how the system prompt affects the valid alpha rate of an LLM. 98 | 99 | ### After running 100 | 101 | - Model checkpoints and alpha pools are located in `save_path`; 102 | - The model is compatiable with [stable-baselines3](https://github.com/DLR-RM/stable-baselines3) 103 | - Alpha pools are formatted in human-readable JSON. 104 | - Tensorboard logs are located in `tensorboard_log`. 105 | 106 | ## Baselines 107 | 108 | ### GP-based methods 109 | 110 | [gplearn](https://github.com/trevorstephens/gplearn) implements Genetic Programming, a commonly used method for symbolic regression. We maintained a modified version of gplearn to make it compatiable with our task. The corresponding experiment scipt is [gp.py](gp.py) 111 | 112 | ### Deep Symbolic Regression 113 | 114 | [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) 115 | 116 | ## Trading (Experimental) 117 | 118 | We implemented some trading strategies based on Qlib. See [backtest.py](backtest.py) and [trade_decision.py](trade_decision.py) for demos. 119 | 120 | ## Citing our work 121 | 122 | ```bibtex 123 | @inproceedings{alphagen, 124 | author = {Yu, Shuo and Xue, Hongyan and Ao, Xiang and Pan, Feiyang and He, Jia and Tu, Dandan and He, Qing}, 125 | title = {Generating Synergistic Formulaic Alpha Collections via Reinforcement Learning}, 126 | year = {2023}, 127 | doi = {10.1145/3580305.3599831}, 128 | booktitle = {Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining}, 129 | } 130 | ``` 131 | 132 | ## Contributing 133 | 134 | Feel free to submit Issues or Pull requests. 135 | 136 | ## Contributors 137 | 138 | This work is maintained by the MLDM research group, [IIP, ICT, CAS](http://iip.ict.ac.cn/). 139 | 140 | Maintainers include: 141 | 142 | - [Hongyan Xue](https://github.com/xuehongyanL) 143 | - [Shuo Yu](https://github.com/Chlorie) 144 | 145 | Thanks to the following contributors: 146 | 147 | - [@yigaza](https://github.com/yigaza) 148 | 149 | Thanks to the following in-depth research on our project: 150 | 151 | - *因子选股系列之九十五: DFQ强化学习因子组合挖掘系统* 152 | -------------------------------------------------------------------------------- /alphagen/config.py: -------------------------------------------------------------------------------- 1 | from typing import Type 2 | from alphagen.data.expression import * 3 | 4 | 5 | MAX_EXPR_LENGTH = 15 6 | MAX_EPISODE_LENGTH = 256 7 | 8 | OPERATORS: List[Type[Operator]] = [ 9 | # Unary 10 | Abs, # Sign, 11 | Log, 12 | # Binary 13 | Add, Sub, Mul, Div, Greater, Less, 14 | # Rolling 15 | Ref, Mean, Sum, Std, Var, # Skew, Kurt, 16 | Max, Min, 17 | Med, Mad, # Rank, 18 | Delta, WMA, EMA, 19 | # Pair rolling 20 | Cov, Corr 21 | ] 22 | 23 | DELTA_TIMES = [1, 5, 10, 20, 40] 24 | 25 | CONSTANTS = [-30., -10., -5., -2., -1., -0.5, -0.01, 0.01, 0.5, 1., 2., 5., 10., 30.] 26 | 27 | REWARD_PER_STEP = 0. 28 | -------------------------------------------------------------------------------- /alphagen/data/calculator.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | from typing import Tuple, Optional, Sequence 3 | from torch import Tensor 4 | import torch 5 | 6 | from alphagen.data.expression import Expression 7 | from alphagen.utils.correlation import batch_pearsonr, batch_spearmanr 8 | 9 | 10 | class AlphaCalculator(metaclass=ABCMeta): 11 | @abstractmethod 12 | def calc_single_IC_ret(self, expr: Expression) -> float: 13 | 'Calculate IC between a single alpha and a predefined target.' 14 | 15 | @abstractmethod 16 | def calc_single_rIC_ret(self, expr: Expression) -> float: 17 | 'Calculate Rank IC between a single alpha and a predefined target.' 18 | 19 | def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]: 20 | return self.calc_single_IC_ret(expr), self.calc_single_rIC_ret(expr) 21 | 22 | @abstractmethod 23 | def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float: 24 | 'Calculate IC between two alphas.' 25 | 26 | @abstractmethod 27 | def calc_pool_IC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float: 28 | 'First combine the alphas linearly,' 29 | 'then Calculate IC between the linear combination and a predefined target.' 30 | 31 | @abstractmethod 32 | def calc_pool_rIC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float: 33 | 'First combine the alphas linearly,' 34 | 'then Calculate Rank IC between the linear combination and a predefined target.' 35 | 36 | @abstractmethod 37 | def calc_pool_all_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float]: 38 | 'First combine the alphas linearly,' 39 | 'then Calculate both IC and Rank IC between the linear combination and a predefined target.' 40 | 41 | 42 | class TensorAlphaCalculator(AlphaCalculator): 43 | def __init__(self, target: Optional[Tensor]) -> None: 44 | self._target = target 45 | 46 | @property 47 | @abstractmethod 48 | def n_days(self) -> int: ... 49 | 50 | @property 51 | def target(self) -> Tensor: 52 | if self._target is None: 53 | raise ValueError("A target must be set before calculating non-mutual IC.") 54 | return self._target 55 | 56 | @abstractmethod 57 | def evaluate_alpha(self, expr: Expression) -> Tensor: 58 | 'Evaluate an alpha into a `Tensor` of shape (days, stocks).' 59 | 60 | def make_ensemble_alpha(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tensor: 61 | n = len(exprs) 62 | factors = [self.evaluate_alpha(exprs[i]) * weights[i] for i in range(n)] 63 | return torch.sum(torch.stack(factors, dim=0), dim=0) 64 | 65 | def _calc_IC(self, value1: Tensor, value2: Tensor) -> float: 66 | return batch_pearsonr(value1, value2).mean().item() 67 | 68 | def _calc_rIC(self, value1: Tensor, value2: Tensor) -> float: 69 | return batch_spearmanr(value1, value2).mean().item() 70 | 71 | def _IR_from_batch(self, batch: Tensor) -> float: 72 | mean, std = batch.mean(), batch.std() 73 | return (mean / std).item() 74 | 75 | def _calc_ICIR(self, value1: Tensor, value2: Tensor) -> float: 76 | return self._IR_from_batch(batch_pearsonr(value1, value2)) 77 | 78 | def _calc_rICIR(self, value1: Tensor, value2: Tensor) -> float: 79 | return self._IR_from_batch(batch_spearmanr(value1, value2)) 80 | 81 | def calc_single_IC_ret(self, expr: Expression) -> float: 82 | return self._calc_IC(self.evaluate_alpha(expr), self.target) 83 | 84 | def calc_single_IC_ret_daily(self, expr: Expression) -> Tensor: 85 | return batch_pearsonr(self.evaluate_alpha(expr), self.target) 86 | 87 | def calc_single_rIC_ret(self, expr: Expression) -> float: 88 | return self._calc_rIC(self.evaluate_alpha(expr), self.target) 89 | 90 | def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]: 91 | value = self.evaluate_alpha(expr) 92 | target = self.target 93 | return self._calc_IC(value, target), self._calc_rIC(value, target) 94 | 95 | def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float: 96 | return self._calc_IC(self.evaluate_alpha(expr1), self.evaluate_alpha(expr2)) 97 | 98 | def calc_mutual_IC_daily(self, expr1: Expression, expr2: Expression) -> Tensor: 99 | return batch_pearsonr(self.evaluate_alpha(expr1), self.evaluate_alpha(expr2)) 100 | 101 | def calc_pool_IC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float: 102 | with torch.no_grad(): 103 | value = self.make_ensemble_alpha(exprs, weights) 104 | return self._calc_IC(value, self.target) 105 | 106 | def calc_pool_rIC_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> float: 107 | with torch.no_grad(): 108 | value = self.make_ensemble_alpha(exprs, weights) 109 | return self._calc_rIC(value, self.target) 110 | 111 | def calc_pool_all_ret(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float]: 112 | with torch.no_grad(): 113 | value = self.make_ensemble_alpha(exprs, weights) 114 | target = self.target 115 | return self._calc_IC(value, target), self._calc_rIC(value, target) 116 | 117 | def calc_pool_all_ret_with_ir(self, exprs: Sequence[Expression], weights: Sequence[float]) -> Tuple[float, float, float, float]: 118 | "Returns IC, ICIR, Rank IC, Rank ICIR" 119 | with torch.no_grad(): 120 | value = self.make_ensemble_alpha(exprs, weights) 121 | target = self.target 122 | ics = batch_pearsonr(value, target) 123 | rics = batch_spearmanr(value, target) 124 | ic_mean, ic_std = ics.mean().item(), ics.std().item() 125 | ric_mean, ric_std = rics.mean().item(), rics.std().item() 126 | return ic_mean, ic_mean / ic_std, ric_mean, ric_mean / ric_std 127 | -------------------------------------------------------------------------------- /alphagen/data/exception.py: -------------------------------------------------------------------------------- 1 | class InvalidExpressionException(ValueError): 2 | pass 3 | -------------------------------------------------------------------------------- /alphagen/data/parser.py: -------------------------------------------------------------------------------- 1 | import re 2 | from typing import Type, List, Dict, Set, Union, Optional, cast, overload, Literal 3 | from .expression import * 4 | from ..utils.misc import find_last_if 5 | 6 | 7 | _PATTERN = re.compile(r'([+-]?[\d.]+|\W|\w+)') 8 | _NUMERIC = re.compile(r'[+-]?[\d.]+') 9 | _StackItem = Union[List[Type[Operator]], Expression] 10 | _OpMap = Dict[str, List[Type[Operator]]] 11 | _DTLike = Union[float, Constant, DeltaTime] 12 | 13 | 14 | class ExpressionParsingError(Exception): 15 | pass 16 | 17 | 18 | class ExpressionParser: 19 | def __init__( 20 | self, 21 | operators: List[Type[Operator]], 22 | ignore_case: bool = False, 23 | time_deltas_need_suffix: bool = False, 24 | non_positive_time_deltas_allowed: bool = True, 25 | feature_need_dollar_sign: bool = False, 26 | additional_operator_mapping: Optional[_OpMap] = None 27 | ): 28 | self._ignore_case = ignore_case 29 | self._allow_np_dt = non_positive_time_deltas_allowed 30 | self._suffix_needed = time_deltas_need_suffix 31 | self._dollar_needed = feature_need_dollar_sign 32 | self._features = {f.name.lower(): f for f in FeatureType} 33 | self._operators: _OpMap = {op.__name__: [op] for op in operators} 34 | if additional_operator_mapping is not None: 35 | self._merge_op_mapping(additional_operator_mapping) 36 | if ignore_case: 37 | self._operators = {k.lower(): v for k, v in self._operators.items()} 38 | self._stack: List[_StackItem] = [] 39 | self._tokens: List[str] = [] 40 | 41 | def parse(self, expr: str) -> Expression: 42 | self._stack = [] 43 | self._tokens = [t for t in _PATTERN.findall(expr) if not t.isspace()] 44 | self._tokens.reverse() 45 | while len(self._tokens) > 0: 46 | self._stack.append(self._get_next_item()) 47 | self._process_punctuation() 48 | if len(self._stack) != 1: 49 | raise ExpressionParsingError("Multiple items remain in the stack") 50 | if len(self._stack) == 0: 51 | raise ExpressionParsingError("Nothing was parsed") 52 | if isinstance(self._stack[0], Expression): 53 | return self._stack[0] 54 | raise ExpressionParsingError(f"{self._stack[0]} is not a valid expression") 55 | 56 | def _merge_op_mapping(self, map: _OpMap) -> None: 57 | for name, ops in map.items(): 58 | if (old_ops := self._operators.get(name)) is not None: 59 | self._operators[name] = list(dict.fromkeys(old_ops + ops)) 60 | else: 61 | self._operators[name] = ops 62 | 63 | def _get_next_item(self) -> _StackItem: 64 | top = self._pop_token() 65 | if top == '$': # Feature next 66 | top = self._pop_token() 67 | if (feature := self._features.get(top)) is None: 68 | raise ExpressionParsingError(f"Can't find the feature {top}") 69 | return Feature(feature) 70 | elif self._tokens_eq(top, "Constant"): 71 | if self._pop_token() != '(': 72 | raise ExpressionParsingError("\"Constant\" should be followed by a left parenthesis") 73 | value = self._to_float(self._pop_token()) 74 | if self._pop_token() != ')': 75 | raise ExpressionParsingError("\"Constant\" should be closed by a right parenthesis") 76 | return Constant(value) 77 | elif _NUMERIC.fullmatch(top) is not None: 78 | value = self._to_float(top) 79 | if self._peek_token() == 'd': 80 | self._pop_token() 81 | return self._as_delta_time(value) 82 | else: 83 | return Constant(value) 84 | else: 85 | if not self._dollar_needed and (feature := self._features.get(top)) is not None: 86 | return Feature(feature) 87 | elif (ops := self._operators.get(top)) is not None: 88 | return ops 89 | else: 90 | raise ExpressionParsingError(f"Cannot find the operator/feature name {top}") 91 | 92 | def _process_punctuation(self) -> None: 93 | if len(self._tokens) == 0: 94 | return 95 | top = self._pop_token() 96 | stack_top_is_ops = len(self._stack) != 0 and not isinstance(self._stack[-1], Expression) 97 | if (top == '(') != stack_top_is_ops: 98 | raise ExpressionParsingError("A left parenthesis should follow an operator name") 99 | if top == '(' or top == ',': 100 | return 101 | elif top == ')': 102 | self._build_one_subexpr() # Pop an operator with its operands 103 | self._process_punctuation() # There might be consecutive right parens 104 | else: 105 | raise ExpressionParsingError(f"Unexpected token {top}") 106 | 107 | def _build_one_subexpr(self) -> None: 108 | if (op_idx := find_last_if(self._stack, lambda item: isinstance(item, list))) == -1: 109 | raise ExpressionParsingError("Unmatched right parenthesis") 110 | ops = cast(List[Type[Operator]], self._stack[op_idx]) 111 | operands = self._stack[op_idx + 1:] 112 | self._stack = self._stack[:op_idx] 113 | if any(not isinstance(item, Expression) for item in operands): 114 | raise ExpressionParsingError("An operator name cannot be used as an operand") 115 | operands = cast(List[Expression], operands) 116 | dt_operands = operands 117 | if (not self._suffix_needed and 118 | isinstance(operands[-1], Constant) and 119 | (dt := self._as_delta_time(operands[-1], noexcept=True)) is not None): 120 | dt_operands = operands.copy() 121 | dt_operands[-1] = dt 122 | msgs: Set[str] = set() 123 | for op in ops: 124 | used_operands = operands 125 | if issubclass(op, (RollingOperator, PairRollingOperator)): 126 | used_operands = dt_operands 127 | if (msg := op.validate_parameters(*used_operands)).is_none: 128 | self._stack.append(op(*used_operands)) # type: ignore 129 | return 130 | else: 131 | msgs.add(msg.value_or("")) 132 | raise ExpressionParsingError("; ".join(msgs)) 133 | 134 | def _tokens_eq(self, lhs: str, rhs: str) -> bool: 135 | if self._ignore_case: 136 | return lhs.lower() == rhs.lower() 137 | else: 138 | return lhs == rhs 139 | 140 | @classmethod 141 | def _to_float(cls, token: str) -> float: 142 | try: 143 | return float(token) 144 | except: 145 | raise ExpressionParsingError(f"{token} can't be converted to float") 146 | 147 | def _pop_token(self) -> str: 148 | if len(self._tokens) == 0: 149 | raise ExpressionParsingError("No more tokens left") 150 | top = self._tokens.pop() 151 | return top.lower() if self._ignore_case else top 152 | 153 | def _peek_token(self) -> Optional[str]: 154 | return self._tokens[-1] if len(self._tokens) != 0 else None 155 | 156 | @overload 157 | def _as_delta_time(self, value: _DTLike, noexcept: Literal[False] = False) -> DeltaTime: ... 158 | @overload 159 | def _as_delta_time(self, value: _DTLike, noexcept: Literal[True]) -> Optional[DeltaTime]: ... 160 | @overload 161 | def _as_delta_time(self, value: _DTLike, noexcept: bool) -> Optional[DeltaTime]: ... 162 | 163 | def _as_delta_time(self, value: _DTLike, noexcept: bool = False): 164 | def maybe_raise(message: str) -> None: 165 | if not noexcept: 166 | raise ExpressionParsingError(message) 167 | 168 | if isinstance(value, DeltaTime): 169 | return value 170 | if isinstance(value, Constant): 171 | value = value.value 172 | if not float(value).is_integer(): 173 | maybe_raise(f"A DeltaTime should be integral, but {value} is not") 174 | return 175 | if int(value) <= 0 and not self._allow_np_dt: 176 | maybe_raise(f"A DeltaTime should refer to a positive time difference, but got {int(value)}d") 177 | return 178 | return DeltaTime(int(value)) 179 | 180 | 181 | def parse_expression(expr: str) -> Expression: 182 | "Parse an expression using the default expression parser." 183 | return ExpressionParser(Operators).parse(expr) 184 | -------------------------------------------------------------------------------- /alphagen/data/pool_update.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | from typing import List, Optional, cast 3 | from dataclasses import dataclass, MISSING 4 | 5 | from .expression import Expression 6 | 7 | 8 | @dataclass 9 | class PoolUpdate(metaclass=ABCMeta): 10 | @property 11 | @abstractmethod 12 | def old_pool(self) -> List[Expression]: ... 13 | 14 | @property 15 | @abstractmethod 16 | def new_pool(self) -> List[Expression]: ... 17 | 18 | @property 19 | @abstractmethod 20 | def old_pool_ic(self) -> Optional[float]: ... 21 | 22 | @property 23 | @abstractmethod 24 | def new_pool_ic(self) -> float: ... 25 | 26 | @property 27 | def ic_increment(self) -> float: 28 | return self.new_pool_ic - (self.old_pool_ic or 0.) 29 | 30 | @abstractmethod 31 | def describe(self) -> str: ... 32 | 33 | def describe_verbose(self) -> str: return self.describe() 34 | 35 | def _describe_ic_diff(self) -> str: 36 | return ( 37 | f"{self.old_pool_ic:.4f} -> {self.new_pool_ic:.4f} " 38 | f"(increment of {self.ic_increment:.4f})" 39 | ) 40 | 41 | def _describe_pool(self, title: str, pool: List[Expression]) -> str: 42 | list_exprs = "\n".join([f" {expr}" for expr in pool]) 43 | return f"{title}\n{list_exprs}" 44 | 45 | 46 | class _PoolUpdateStub: 47 | old_pool: List[Expression] = cast(List[Expression], MISSING) 48 | new_pool: List[Expression] = cast(List[Expression], MISSING) 49 | old_pool_ic: Optional[float] = cast(Optional[float], MISSING) 50 | new_pool_ic: float = cast(float, MISSING) 51 | 52 | 53 | @dataclass 54 | class SetPool(_PoolUpdateStub, PoolUpdate): 55 | old_pool: List[Expression] 56 | new_pool: List[Expression] 57 | old_pool_ic: Optional[float] 58 | new_pool_ic: float 59 | 60 | def describe(self) -> str: 61 | pool = self._describe_pool("Alpha pool:", self.new_pool) 62 | return f"{pool}\nIC of the combination: {self.new_pool_ic:.4f}" 63 | 64 | def describe_verbose(self) -> str: 65 | if len(self.old_pool) == 0: 66 | return self.describe() 67 | old_pool = self._describe_pool("Old alpha pool:", self.old_pool) 68 | new_pool = self._describe_pool("New alpha pool:", self.new_pool) 69 | perf = f"IC of the pools: {self._describe_ic_diff()})" 70 | return f"{old_pool}\n{new_pool}\n{perf}" 71 | 72 | 73 | @dataclass 74 | class AddRemoveAlphas(_PoolUpdateStub, PoolUpdate): 75 | added_exprs: List[Expression] 76 | removed_idx: List[int] 77 | old_pool: List[Expression] 78 | old_pool_ic: float 79 | new_pool_ic: float 80 | 81 | @property 82 | def new_pool(self) -> List[Expression]: 83 | remain = [True] * len(self.old_pool) 84 | for i in self.removed_idx: 85 | remain[i] = False 86 | return [expr for i, expr in enumerate(self.old_pool) if remain[i]] + self.added_exprs 87 | 88 | def describe(self) -> str: 89 | def describe_exprs(title: str, exprs: List[Expression]) -> str: 90 | if len(exprs) == 0: 91 | return "" 92 | if len(exprs) == 1: 93 | return f"{title}: {exprs[0]}\n" 94 | exprs_str = "\n".join([f" {expr}" for expr in exprs]) 95 | return f"{title}s:\n{exprs_str}\n" 96 | 97 | added = describe_exprs("Added alpha", self.added_exprs) 98 | removed = describe_exprs("Removed alpha", [self.old_pool[i] for i in self.removed_idx]) 99 | perf = f"IC of the combination: {self._describe_ic_diff()}" 100 | return f"{added}{removed}{perf}" 101 | 102 | def describe_verbose(self) -> str: 103 | old = self._describe_pool("Old alpha pool:", self.old_pool) 104 | return f"{old}\n{self.describe()}" 105 | -------------------------------------------------------------------------------- /alphagen/data/tokens.py: -------------------------------------------------------------------------------- 1 | from enum import IntEnum 2 | from typing import Type 3 | from alphagen_qlib.stock_data import FeatureType 4 | from alphagen.data.expression import Operator, Expression 5 | 6 | 7 | class SequenceIndicatorType(IntEnum): 8 | BEG = 0 9 | SEP = 1 10 | 11 | 12 | class Token: 13 | def __repr__(self): 14 | return str(self) 15 | 16 | 17 | class ConstantToken(Token): 18 | def __init__(self, constant: float) -> None: 19 | self.constant = constant 20 | 21 | def __str__(self): return str(self.constant) 22 | 23 | 24 | class DeltaTimeToken(Token): 25 | def __init__(self, delta_time: int) -> None: 26 | self.delta_time = delta_time 27 | 28 | def __str__(self): return str(self.delta_time) 29 | 30 | 31 | class FeatureToken(Token): 32 | def __init__(self, feature: FeatureType) -> None: 33 | self.feature = feature 34 | 35 | def __str__(self): return '$' + self.feature.name.lower() 36 | 37 | 38 | class OperatorToken(Token): 39 | def __init__(self, operator: Type[Operator]) -> None: 40 | self.operator = operator 41 | 42 | def __str__(self): return self.operator.__name__ 43 | 44 | 45 | class SequenceIndicatorToken(Token): 46 | def __init__(self, indicator: SequenceIndicatorType) -> None: 47 | self.indicator = indicator 48 | 49 | def __str__(self): return self.indicator.name 50 | 51 | 52 | class ExpressionToken(Token): 53 | def __init__(self, expr: Expression) -> None: 54 | self.expression = expr 55 | 56 | def __str__(self): return str(self.expression) 57 | 58 | 59 | BEG_TOKEN = SequenceIndicatorToken(SequenceIndicatorType.BEG) 60 | SEP_TOKEN = SequenceIndicatorToken(SequenceIndicatorType.SEP) 61 | -------------------------------------------------------------------------------- /alphagen/data/tree.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | from alphagen.data.exception import InvalidExpressionException 3 | from alphagen.data.expression import BinaryOperator, Constant, DeltaTime, Expression, Feature, PairRollingOperator, RollingOperator, UnaryOperator 4 | from alphagen.data.tokens import * 5 | 6 | 7 | class ExpressionBuilder: 8 | stack: List[Expression] 9 | 10 | def __init__(self): 11 | self.stack = [] 12 | 13 | def get_tree(self) -> Expression: 14 | if len(self.stack) == 1: 15 | return self.stack[0] 16 | else: 17 | raise InvalidExpressionException(f"Expected only one tree, got {len(self.stack)}") 18 | 19 | def add_token(self, token: Token): 20 | if not self.validate(token): 21 | raise InvalidExpressionException(f"Token {token} not allowed here, stack: {self.stack}.") 22 | if isinstance(token, OperatorToken): 23 | n_args: int = token.operator.n_args() 24 | children = [] 25 | for _ in range(n_args): 26 | children.append(self.stack.pop()) 27 | self.stack.append(token.operator(*reversed(children))) # type: ignore 28 | elif isinstance(token, ConstantToken): 29 | self.stack.append(Constant(token.constant)) 30 | elif isinstance(token, DeltaTimeToken): 31 | self.stack.append(DeltaTime(token.delta_time)) 32 | elif isinstance(token, FeatureToken): 33 | self.stack.append(Feature(token.feature)) 34 | elif isinstance(token, ExpressionToken): 35 | self.stack.append(token.expression) 36 | else: 37 | assert False 38 | 39 | def is_valid(self) -> bool: 40 | return len(self.stack) == 1 and self.stack[0].is_featured 41 | 42 | def validate(self, token: Token) -> bool: 43 | if isinstance(token, OperatorToken): 44 | return self.validate_op(token.operator) 45 | elif isinstance(token, DeltaTimeToken): 46 | return self.validate_dt() 47 | elif isinstance(token, ConstantToken): 48 | return self.validate_const() 49 | elif isinstance(token, (FeatureToken, ExpressionToken)): 50 | return self.validate_featured_expr() 51 | else: 52 | assert False 53 | 54 | def validate_op(self, op: Type[Operator]) -> bool: 55 | if len(self.stack) < op.n_args(): 56 | return False 57 | 58 | if issubclass(op, UnaryOperator): 59 | if not self.stack[-1].is_featured: 60 | return False 61 | elif issubclass(op, BinaryOperator): 62 | if not self.stack[-1].is_featured and not self.stack[-2].is_featured: 63 | return False 64 | if (isinstance(self.stack[-1], DeltaTime) or 65 | isinstance(self.stack[-2], DeltaTime)): 66 | return False 67 | elif issubclass(op, RollingOperator): 68 | if not isinstance(self.stack[-1], DeltaTime): 69 | return False 70 | if not self.stack[-2].is_featured: 71 | return False 72 | elif issubclass(op, PairRollingOperator): 73 | if not isinstance(self.stack[-1], DeltaTime): 74 | return False 75 | if not self.stack[-2].is_featured or not self.stack[-3].is_featured: 76 | return False 77 | else: 78 | assert False 79 | return True 80 | 81 | def validate_dt(self) -> bool: 82 | return len(self.stack) > 0 and self.stack[-1].is_featured 83 | 84 | def validate_const(self) -> bool: 85 | return len(self.stack) == 0 or self.stack[-1].is_featured 86 | 87 | def validate_featured_expr(self) -> bool: 88 | return not (len(self.stack) >= 1 and isinstance(self.stack[-1], DeltaTime)) 89 | -------------------------------------------------------------------------------- /alphagen/models/alpha_pool.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Dict, Any, Callable 2 | from abc import ABCMeta, abstractmethod 3 | 4 | import torch 5 | from ..data.calculator import AlphaCalculator 6 | from ..data.expression import Expression 7 | 8 | 9 | class AlphaPoolBase(metaclass=ABCMeta): 10 | def __init__( 11 | self, 12 | capacity: int, 13 | calculator: AlphaCalculator, 14 | device: torch.device = torch.device("cpu") 15 | ): 16 | self.size = 0 17 | self.capacity = capacity 18 | self.calculator = calculator 19 | self.device = device 20 | self.eval_cnt = 0 21 | self.best_ic_ret: float = -1. 22 | 23 | @property 24 | def vacancy(self) -> int: 25 | return self.capacity - self.size 26 | 27 | @property 28 | @abstractmethod 29 | def state(self) -> Dict[str, Any]: 30 | "Get a dictionary representing the state of this pool." 31 | 32 | @abstractmethod 33 | def to_json_dict(self) -> Dict[str, Any]: 34 | """ 35 | Serialize this pool into a dictionary that can be dumped as json, 36 | i.e. no complex objects. 37 | """ 38 | 39 | @abstractmethod 40 | def try_new_expr(self, expr: Expression) -> float: ... 41 | 42 | @abstractmethod 43 | def test_ensemble(self, calculator: AlphaCalculator) -> Tuple[float, float]: ... 44 | -------------------------------------------------------------------------------- /alphagen/rl/env/core.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Optional 2 | import gymnasium as gym 3 | import math 4 | 5 | from alphagen.config import MAX_EXPR_LENGTH 6 | from alphagen.data.tokens import * 7 | from alphagen.data.expression import * 8 | from alphagen.data.tree import ExpressionBuilder 9 | from alphagen.models.alpha_pool import AlphaPoolBase 10 | from alphagen.models.linear_alpha_pool import LinearAlphaPool 11 | from alphagen.utils import reseed_everything 12 | 13 | 14 | class AlphaEnvCore(gym.Env): 15 | pool: AlphaPoolBase 16 | _tokens: List[Token] 17 | _builder: ExpressionBuilder 18 | _print_expr: bool 19 | 20 | def __init__( 21 | self, 22 | pool: AlphaPoolBase, 23 | device: torch.device = torch.device('cuda:0'), 24 | print_expr: bool = False 25 | ): 26 | super().__init__() 27 | 28 | self.pool = pool 29 | self._print_expr = print_expr 30 | self._device = device 31 | 32 | self.eval_cnt = 0 33 | 34 | self.render_mode = None 35 | self.reset() 36 | 37 | def reset( 38 | self, *, 39 | seed: Optional[int] = None, 40 | return_info: bool = False, 41 | options: Optional[dict] = None 42 | ) -> Tuple[List[Token], dict]: 43 | reseed_everything(seed) 44 | self._tokens = [BEG_TOKEN] 45 | self._builder = ExpressionBuilder() 46 | return self._tokens, self._valid_action_types() 47 | 48 | def step(self, action: Token) -> Tuple[List[Token], float, bool, bool, dict]: 49 | if (isinstance(action, SequenceIndicatorToken) and 50 | action.indicator == SequenceIndicatorType.SEP): 51 | reward = self._evaluate() 52 | done = True 53 | elif len(self._tokens) < MAX_EXPR_LENGTH: 54 | self._tokens.append(action) 55 | self._builder.add_token(action) 56 | done = False 57 | reward = 0.0 58 | else: 59 | done = True 60 | reward = self._evaluate() if self._builder.is_valid() else -1. 61 | 62 | if math.isnan(reward): 63 | reward = 0. 64 | 65 | return self._tokens, reward, done, False, self._valid_action_types() 66 | 67 | def _evaluate(self): 68 | expr: Expression = self._builder.get_tree() 69 | if self._print_expr: 70 | print(expr) 71 | try: 72 | ret = self.pool.try_new_expr(expr) 73 | self.eval_cnt += 1 74 | return ret 75 | except OutOfDataRangeError: 76 | return 0. 77 | 78 | def _valid_action_types(self) -> dict: 79 | valid_op_unary = self._builder.validate_op(UnaryOperator) 80 | valid_op_binary = self._builder.validate_op(BinaryOperator) 81 | valid_op_rolling = self._builder.validate_op(RollingOperator) 82 | valid_op_pair_rolling = self._builder.validate_op(PairRollingOperator) 83 | 84 | valid_op = valid_op_unary or valid_op_binary or valid_op_rolling or valid_op_pair_rolling 85 | valid_dt = self._builder.validate_dt() 86 | valid_const = self._builder.validate_const() 87 | valid_feature = self._builder.validate_featured_expr() 88 | valid_stop = self._builder.is_valid() 89 | 90 | ret = { 91 | 'select': [valid_op, valid_feature, valid_const, valid_dt, valid_stop], 92 | 'op': { 93 | UnaryOperator: valid_op_unary, 94 | BinaryOperator: valid_op_binary, 95 | RollingOperator: valid_op_rolling, 96 | PairRollingOperator: valid_op_pair_rolling 97 | } 98 | } 99 | return ret 100 | 101 | def valid_action_types(self) -> dict: 102 | return self._valid_action_types() 103 | 104 | def render(self, mode='human'): 105 | pass 106 | -------------------------------------------------------------------------------- /alphagen/rl/env/wrapper.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, List, Optional 2 | import gymnasium as gym 3 | import numpy as np 4 | 5 | from alphagen.config import * 6 | from alphagen.data.tokens import * 7 | from alphagen.models.alpha_pool import AlphaPoolBase 8 | from alphagen.rl.env.core import AlphaEnvCore 9 | 10 | SIZE_NULL = 1 11 | SIZE_OP = len(OPERATORS) 12 | SIZE_FEATURE = len(FeatureType) 13 | SIZE_DELTA_TIME = len(DELTA_TIMES) 14 | SIZE_CONSTANT = len(CONSTANTS) 15 | SIZE_SEP = 1 16 | SIZE_ACTION = SIZE_OP + SIZE_FEATURE + SIZE_DELTA_TIME + SIZE_CONSTANT + SIZE_SEP 17 | 18 | 19 | class AlphaEnvWrapper(gym.Wrapper): 20 | state: np.ndarray 21 | env: AlphaEnvCore 22 | action_space: gym.spaces.Discrete 23 | observation_space: gym.spaces.Box 24 | counter: int 25 | 26 | def __init__( 27 | self, 28 | env: AlphaEnvCore, 29 | subexprs: Optional[List[Expression]] = None 30 | ): 31 | super().__init__(env) 32 | self.subexprs = subexprs or [] 33 | self.size_action = SIZE_ACTION + len(self.subexprs) 34 | self.action_space = gym.spaces.Discrete(self.size_action) 35 | self.observation_space = gym.spaces.Box( 36 | low=0, high=self.size_action + SIZE_NULL - 1, 37 | shape=(MAX_EXPR_LENGTH, ), 38 | dtype=np.uint8 39 | ) 40 | 41 | def reset(self, **kwargs) -> Tuple[np.ndarray, dict]: 42 | self.counter = 0 43 | self.state = np.zeros(MAX_EXPR_LENGTH, dtype=np.uint8) 44 | self.env.reset() 45 | return self.state, {} 46 | 47 | def step(self, action: int): 48 | _, reward, done, truncated, info = self.env.step(self.action(action)) 49 | if not done: 50 | self.state[self.counter] = action 51 | self.counter += 1 52 | return self.state, self.reward(reward), done, truncated, info 53 | 54 | def action(self, action: int) -> Token: 55 | return self.action_to_token(action) 56 | 57 | def reward(self, reward: float) -> float: 58 | return reward + REWARD_PER_STEP 59 | 60 | def action_masks(self) -> np.ndarray: 61 | res = np.zeros(self.size_action, dtype=bool) 62 | valid = self.env.valid_action_types() 63 | 64 | offset = 0 # Operators 65 | for i in range(offset, offset + SIZE_OP): 66 | if valid['op'][OPERATORS[i - offset].category_type()]: 67 | res[i] = True 68 | offset += SIZE_OP 69 | if valid['select'][1]: # Features 70 | res[offset:offset + SIZE_FEATURE] = True 71 | offset += SIZE_FEATURE 72 | if valid['select'][2]: # Constants 73 | res[offset:offset + SIZE_CONSTANT] = True 74 | offset += SIZE_CONSTANT 75 | if valid['select'][3]: # Delta time 76 | res[offset:offset + SIZE_DELTA_TIME] = True 77 | offset += SIZE_DELTA_TIME 78 | if valid['select'][1]: # Sub-expressions 79 | res[offset:offset + len(self.subexprs)] = True 80 | offset += len(self.subexprs) 81 | if valid['select'][4]: # SEP 82 | res[offset] = True 83 | return res 84 | 85 | def action_to_token(self, action: int) -> Token: 86 | if action < 0: 87 | raise ValueError 88 | if action < SIZE_OP: 89 | return OperatorToken(OPERATORS[action]) 90 | action -= SIZE_OP 91 | if action < SIZE_FEATURE: 92 | return FeatureToken(FeatureType(action)) 93 | action -= SIZE_FEATURE 94 | if action < SIZE_CONSTANT: 95 | return ConstantToken(CONSTANTS[action]) 96 | action -= SIZE_CONSTANT 97 | if action < SIZE_DELTA_TIME: 98 | return DeltaTimeToken(DELTA_TIMES[action]) 99 | action -= SIZE_DELTA_TIME 100 | if action < len(self.subexprs): 101 | return ExpressionToken(self.subexprs[action]) 102 | action -= len(self.subexprs) 103 | if action == 0: 104 | return SequenceIndicatorToken(SequenceIndicatorType.SEP) 105 | assert False 106 | 107 | 108 | def AlphaEnv(pool: AlphaPoolBase, subexprs: Optional[List[Expression]] = None, **kwargs): 109 | return AlphaEnvWrapper(AlphaEnvCore(pool=pool, **kwargs), subexprs=subexprs) 110 | -------------------------------------------------------------------------------- /alphagen/rl/policy.py: -------------------------------------------------------------------------------- 1 | import gymnasium as gym 2 | import math 3 | import torch.nn.functional as F 4 | from stable_baselines3.common.torch_layers import BaseFeaturesExtractor 5 | from torch import nn 6 | 7 | from alphagen.data.expression import * 8 | 9 | 10 | class PositionalEncoding(nn.Module): 11 | def __init__(self, d_model: int, max_len: int = 5000): 12 | super().__init__() 13 | position = torch.arange(max_len).unsqueeze(1) 14 | div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) 15 | pe = torch.zeros(max_len, d_model) 16 | pe[:, 0::2] = torch.sin(position * div_term) 17 | pe[:, 1::2] = torch.cos(position * div_term) 18 | self.register_buffer('_pe', pe) 19 | 20 | def forward(self, x: Tensor) -> Tensor: 21 | "x: ([batch_size, ]seq_len, embedding_dim)" 22 | seq_len = x.size(0) if x.dim() == 2 else x.size(1) 23 | return x + self._pe[:seq_len] # type: ignore 24 | 25 | 26 | class TransformerSharedNet(BaseFeaturesExtractor): 27 | def __init__( 28 | self, 29 | observation_space: gym.Space, 30 | n_encoder_layers: int, 31 | d_model: int, 32 | n_head: int, 33 | d_ffn: int, 34 | dropout: float, 35 | device: torch.device 36 | ): 37 | super().__init__(observation_space, d_model) 38 | 39 | assert isinstance(observation_space, gym.spaces.Box) 40 | n_actions = observation_space.high[0] + 1 # type: ignore 41 | 42 | self._device = device 43 | self._d_model = d_model 44 | self._n_actions: float = n_actions 45 | 46 | self._token_emb = nn.Embedding(n_actions + 1, d_model, 0) # Last one is [BEG] 47 | self._pos_enc = PositionalEncoding(d_model).to(device) 48 | 49 | self._transformer = nn.TransformerEncoder( 50 | nn.TransformerEncoderLayer( 51 | d_model=d_model, nhead=n_head, 52 | dim_feedforward=d_ffn, dropout=dropout, 53 | activation=lambda x: F.leaky_relu(x), # type: ignore 54 | batch_first=True, device=device 55 | ), 56 | num_layers=n_encoder_layers, 57 | norm=nn.LayerNorm(d_model, eps=1e-5, device=device) 58 | ) 59 | 60 | def forward(self, obs: Tensor) -> Tensor: 61 | bs, seqlen = obs.shape 62 | beg = torch.full((bs, 1), fill_value=self._n_actions, dtype=torch.long, device=obs.device) 63 | obs = torch.cat((beg, obs.long()), dim=1) 64 | pad_mask = obs == 0 65 | src = self._pos_enc(self._token_emb(obs)) 66 | res = self._transformer(src, src_key_padding_mask=pad_mask) 67 | return res.mean(dim=1) 68 | 69 | 70 | class LSTMSharedNet(BaseFeaturesExtractor): 71 | def __init__( 72 | self, 73 | observation_space: gym.Space, 74 | n_layers: int, 75 | d_model: int, 76 | dropout: float, 77 | device: torch.device 78 | ): 79 | super().__init__(observation_space, d_model) 80 | 81 | assert isinstance(observation_space, gym.spaces.Box) 82 | n_actions = observation_space.high[0] + 1 # type: ignore 83 | 84 | self._device = device 85 | self._d_model = d_model 86 | self._n_actions: float = n_actions 87 | 88 | self._token_emb = nn.Embedding(n_actions + 1, d_model, 0) # Last one is [BEG] 89 | self._pos_enc = PositionalEncoding(d_model).to(device) 90 | 91 | self._lstm = nn.LSTM( 92 | input_size=d_model, 93 | hidden_size=d_model, 94 | num_layers=n_layers, 95 | batch_first=True, 96 | dropout=dropout 97 | ) 98 | 99 | def forward(self, obs: Tensor) -> Tensor: 100 | bs, seqlen = obs.shape 101 | beg = torch.full((bs, 1), fill_value=self._n_actions, dtype=torch.long, device=obs.device) 102 | obs = torch.cat((beg, obs.long()), dim=1) 103 | real_len = (obs != 0).sum(1).max() 104 | src = self._pos_enc(self._token_emb(obs)) 105 | res = self._lstm(src[:,:real_len])[0] 106 | return res.mean(dim=1) 107 | 108 | 109 | class Decoder(BaseFeaturesExtractor): 110 | def __init__( 111 | self, 112 | observation_space: gym.Space, 113 | n_layers: int, 114 | d_model: int, 115 | n_head: int, 116 | d_ffn: int, 117 | dropout: float, 118 | device: torch.device 119 | ): 120 | super().__init__(observation_space, d_model) 121 | 122 | assert isinstance(observation_space, gym.spaces.Box) 123 | n_actions = observation_space.high[0] + 1 # type: ignore 124 | 125 | self._device = device 126 | self._d_model = d_model 127 | self._n_actions: float = n_actions 128 | 129 | self._token_emb = nn.Embedding(n_actions + 1, d_model, 0) # Last one is [BEG] 130 | self._pos_enc = PositionalEncoding(d_model).to(device) 131 | 132 | # Actually an encoder for now 133 | self._decoder = nn.TransformerEncoder( 134 | nn.TransformerEncoderLayer( 135 | d_model=d_model, nhead=n_head, dim_feedforward=d_ffn, 136 | dropout=dropout, batch_first=True, device=device 137 | ), 138 | n_layers, 139 | norm=nn.LayerNorm(d_model, device=device) 140 | ) 141 | 142 | def forward(self, obs: Tensor) -> Tensor: 143 | batch_size = obs.size(0) 144 | begins = torch.full(size=(batch_size, 1), fill_value=self._n_actions, 145 | dtype=torch.long, device=obs.device) 146 | obs = torch.cat((begins, obs.type(torch.long)), dim=1) # (bs, len) 147 | pad_mask = obs == 0 148 | res = self._token_emb(obs) # (bs, len, d_model) 149 | res = self._pos_enc(res) # (bs, len, d_model) 150 | res = self._decoder(res, src_key_padding_mask=pad_mask) # (bs, len, d_model) 151 | return res.mean(dim=1) # (bs, d_model) 152 | -------------------------------------------------------------------------------- /alphagen/trade/base.py: -------------------------------------------------------------------------------- 1 | from enum import IntEnum 2 | from typing import Dict, NamedTuple, Optional, Type 3 | 4 | import pandera as pa 5 | 6 | 7 | Amount = float 8 | Price = float 9 | 10 | StockCode = str 11 | StockSignal = float 12 | 13 | StockPosition = pa.DataFrameSchema({ 14 | 'code': pa.Column(StockCode), 15 | 'amount': pa.Column(Amount), 16 | 'days_holded': pa.Column(int), 17 | }) 18 | 19 | StockStatus = pa.DataFrameSchema({ 20 | 'code': pa.Column(StockCode), 21 | 'buyable': pa.Column(bool), 22 | 'sellable': pa.Column(bool), 23 | 'signal': pa.Column(StockSignal, nullable=True), 24 | }) 25 | 26 | 27 | class StockOrderDirection(IntEnum): 28 | BUY = 1 29 | SELL = 2 30 | 31 | 32 | class StockOrder: 33 | code: StockCode 34 | amount: Amount 35 | direction: Optional[StockOrderDirection] 36 | 37 | def __init__(self, 38 | code: StockCode, 39 | amount: Amount): 40 | self.code = code 41 | self.amount = amount 42 | self.direction = None 43 | 44 | def to_buy(self): 45 | self.direction = StockOrderDirection.BUY 46 | 47 | def to_sell(self): 48 | self.direction = StockOrderDirection.SELL 49 | 50 | def set_direction(self, direction: StockOrderDirection): 51 | self.direction = direction 52 | -------------------------------------------------------------------------------- /alphagen/trade/strategy.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | from typing import List, Optional, Tuple 3 | 4 | import pandas as pd 5 | 6 | from alphagen.trade.base import StockCode 7 | 8 | 9 | class Strategy(metaclass=ABCMeta): 10 | @abstractmethod 11 | def step_decision(self, 12 | status_df: pd.DataFrame, 13 | position_df: Optional[pd.DataFrame] = None 14 | ) -> Tuple[List[StockCode], List[StockCode]]: 15 | pass 16 | -------------------------------------------------------------------------------- /alphagen/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .correlation import batch_spearmanr 2 | from .random import reseed_everything 3 | from .logging import get_logger 4 | -------------------------------------------------------------------------------- /alphagen/utils/correlation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | 4 | from alphagen.utils.pytorch_utils import masked_mean_std 5 | 6 | 7 | def _mask_either_nan(x: Tensor, y: Tensor, fill_with: float = torch.nan): 8 | x = x.clone() # [days, stocks] 9 | y = y.clone() # [days, stocks] 10 | nan_mask = x.isnan() | y.isnan() 11 | x[nan_mask] = fill_with 12 | y[nan_mask] = fill_with 13 | n = (~nan_mask).sum(dim=1) 14 | return x, y, n, nan_mask 15 | 16 | 17 | def _rank_data_1d(x: Tensor) -> Tensor: 18 | _, inv, counts = x.unique(return_inverse=True, return_counts=True) 19 | cs = counts.cumsum(dim=0) 20 | cs = torch.cat((torch.zeros(1, dtype=x.dtype, device=x.device), cs)) 21 | rmin = cs[:-1] 22 | rmax = cs[1:] - 1 23 | ranks = (rmin + rmax) / 2 24 | return ranks[inv] 25 | 26 | 27 | def _rank_data(x: Tensor, nan_mask: Tensor) -> Tensor: 28 | rank = torch.stack([_rank_data_1d(row) for row in x]) 29 | rank[nan_mask] = 0 30 | return rank # [d, s] 31 | 32 | 33 | def _batch_pearsonr_given_mask( 34 | x: Tensor, y: Tensor, 35 | n: Tensor, mask: Tensor 36 | ) -> Tensor: 37 | x_mean, x_std = masked_mean_std(x, n, mask) 38 | y_mean, y_std = masked_mean_std(y, n, mask) 39 | cov = (x * y).sum(dim=1) / n - x_mean * y_mean 40 | stdmul = x_std * y_std 41 | stdmul[(x_std < 1e-3) | (y_std < 1e-3)] = 1 42 | corrs = cov / stdmul 43 | return corrs 44 | 45 | 46 | def batch_spearmanr(x: Tensor, y: Tensor) -> Tensor: 47 | x, y, n, nan_mask = _mask_either_nan(x, y) 48 | rx = _rank_data(x, nan_mask) 49 | ry = _rank_data(y, nan_mask) 50 | return _batch_pearsonr_given_mask(rx, ry, n, nan_mask) 51 | 52 | 53 | def batch_pearsonr(x: Tensor, y: Tensor) -> Tensor: 54 | return _batch_pearsonr_given_mask(*_mask_either_nan(x, y, fill_with=0.)) 55 | -------------------------------------------------------------------------------- /alphagen/utils/logging.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | from typing import Optional 4 | 5 | 6 | def get_logger(name: str, file_path: Optional[str] = None) -> logging.Logger: 7 | if file_path is not None: 8 | os.makedirs(os.path.dirname(file_path), exist_ok=True) 9 | 10 | logger = logging.getLogger(name) 11 | while logger.hasHandlers(): 12 | handler = logger.handlers[0] 13 | handler.close() 14 | logger.removeHandler(handler) 15 | 16 | logger.setLevel(logging.DEBUG) 17 | formatter = logging.Formatter("%(asctime)s-%(levelname)s-%(message)s") 18 | 19 | if file_path is not None: 20 | file_handler = logging.FileHandler(file_path) 21 | file_handler.setLevel(logging.DEBUG) 22 | file_handler.setFormatter(formatter) 23 | logger.addHandler(file_handler) 24 | 25 | stream_handler = logging.StreamHandler() 26 | stream_handler.setLevel(logging.INFO) 27 | stream_handler.setFormatter(formatter) 28 | logger.addHandler(stream_handler) 29 | 30 | return logger 31 | 32 | 33 | def get_null_logger() -> logging.Logger: 34 | logger = logging.getLogger("null_logger") 35 | logger.addHandler(logging.NullHandler()) 36 | logger.propagate = False 37 | return logger 38 | -------------------------------------------------------------------------------- /alphagen/utils/maybe.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, TypeVar, Generic, Type, Callable, cast 2 | 3 | 4 | _T = TypeVar("_T") 5 | _TRes = TypeVar("_TRes") 6 | 7 | 8 | class Maybe(Generic[_T]): 9 | def __init__(self, value: Optional[_T]) -> None: 10 | self._value = value 11 | 12 | @property 13 | def is_some(self) -> bool: return self._value is not None 14 | 15 | @property 16 | def is_none(self) -> bool: return self._value is None 17 | 18 | @property 19 | def value(self) -> Optional[_T]: return self._value 20 | 21 | def value_or(self, other: _T) -> _T: 22 | return cast(_T, self.value) if self.is_some else other 23 | 24 | def and_then(self, func: Callable[[_T], "Maybe[_TRes]"]) -> "Maybe[_TRes]": 25 | return func(cast(_T, self._value)) if self.is_some else Maybe(None) 26 | 27 | def map(self, func: Callable[[_T], _TRes]) -> "Maybe[_TRes]": 28 | return some(func(cast(_T, self._value))) if self.is_some else Maybe(None) 29 | 30 | def or_else(self, func: Callable[[], "Maybe[_T]"]) -> "Maybe[_T]": 31 | return self if self.is_some else func() 32 | 33 | 34 | def some(value: _T) -> Maybe[_T]: return Maybe(value) 35 | def none(_: Type[_T]) -> Maybe[_T]: return Maybe(None) 36 | -------------------------------------------------------------------------------- /alphagen/utils/misc.py: -------------------------------------------------------------------------------- 1 | from typing import TypeVar, List, Iterable, Tuple, Callable, Optional 2 | from types import FrameType 3 | import inspect 4 | 5 | 6 | _T = TypeVar("_T") 7 | 8 | 9 | def reverse_enumerate(lst: List[_T]) -> Iterable[Tuple[int, _T]]: 10 | for i in range(len(lst) - 1, -1, -1): 11 | yield i, lst[i] 12 | 13 | 14 | def find_last_if(lst: List[_T], predicate: Callable[[_T], bool]) -> int: 15 | for i in range(len(lst) - 1, -1, -1): 16 | if predicate(lst[i]): 17 | return i 18 | return -1 19 | 20 | 21 | def get_arguments_as_dict(frame: Optional[FrameType] = None) -> dict: 22 | if frame is None: 23 | frame = inspect.currentframe().f_back # type: ignore 24 | keys, _, _, values = inspect.getargvalues(frame) # type: ignore 25 | res = {} 26 | for k in keys: 27 | if k != "self": 28 | res[k] = values[k] 29 | return res 30 | 31 | 32 | def pprint_arguments(frame: Optional[FrameType] = None) -> dict: 33 | if frame is None: 34 | frame = inspect.currentframe().f_back # type: ignore 35 | args = get_arguments_as_dict(frame) 36 | formatted_args = '\n'.join(f" {k}: {v}" for k, v in args.items()) 37 | print(f"[Parameters]\n{formatted_args}") 38 | return args 39 | -------------------------------------------------------------------------------- /alphagen/utils/pytorch_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Optional 2 | import torch 3 | from torch import Tensor 4 | 5 | 6 | def masked_mean_std( 7 | x: Tensor, 8 | n: Optional[Tensor] = None, 9 | mask: Optional[Tensor] = None 10 | ) -> Tuple[Tensor, Tensor]: 11 | """ 12 | `x`: [days, stocks], input data 13 | `n`: [days], should be `(~mask).sum(dim=1)`, provide this to avoid unnecessary computations 14 | `mask`: [days, stocks], data masked as `True` will not participate in the computation, \ 15 | defaults to `torch.isnan(x)` 16 | """ 17 | if mask is None: 18 | mask = torch.isnan(x) 19 | if n is None: 20 | n = (~mask).sum(dim=1) 21 | x = x.clone() 22 | x[mask] = 0. 23 | mean = x.sum(dim=1) / n 24 | std = ((((x - mean[:, None]) * ~mask) ** 2).sum(dim=1) / n).sqrt() 25 | return mean, std 26 | 27 | 28 | def normalize_by_day(value: Tensor) -> Tensor: 29 | "The shape of the input and the output is (days, stocks)" 30 | mean, std = masked_mean_std(value) 31 | value = (value - mean[:, None]) / std[:, None] 32 | nan_mask = torch.isnan(value) 33 | value[nan_mask] = 0. 34 | return value 35 | -------------------------------------------------------------------------------- /alphagen/utils/random.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | import random 3 | import os 4 | import numpy as np 5 | import torch 6 | from torch.backends import cudnn 7 | 8 | 9 | def reseed_everything(seed: Optional[int]): 10 | if seed is None: 11 | return 12 | 13 | random.seed(seed) 14 | os.environ["PYTHONHASHSEED"] = str(seed) 15 | np.random.seed(seed) 16 | torch.manual_seed(seed) 17 | torch.cuda.manual_seed(seed) 18 | cudnn.deterministic = True 19 | cudnn.benchmark = True 20 | -------------------------------------------------------------------------------- /alphagen_generic/features.py: -------------------------------------------------------------------------------- 1 | from alphagen.data.expression import Feature, Ref 2 | from alphagen_qlib.stock_data import FeatureType 3 | 4 | 5 | high = High = HIGH = Feature(FeatureType.HIGH) 6 | low = Low = LOW = Feature(FeatureType.LOW) 7 | volume = Volume = VOLUME = Feature(FeatureType.VOLUME) 8 | open_ = Open = OPEN = Feature(FeatureType.OPEN) 9 | close = Close = CLOSE = Feature(FeatureType.CLOSE) 10 | vwap = Vwap = VWAP = Feature(FeatureType.VWAP) 11 | target = Ref(close, -20) / close - 1 12 | -------------------------------------------------------------------------------- /alphagen_generic/operators.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | 3 | import numpy as np 4 | from alphagen.data.expression import * 5 | 6 | 7 | GenericOperator = namedtuple('GenericOperator', ['name', 'function', 'arity']) 8 | 9 | unary_ops = [Abs, Log] 10 | binary_ops = [Add, Sub, Mul, Div, Greater, Less] 11 | rolling_ops = [Ref, Mean, Sum, Std, Var, Max, Min, Med, Mad, Delta, WMA, EMA] 12 | rolling_binary_ops = [Cov, Corr] 13 | 14 | 15 | def unary(cls): 16 | def _calc(a): 17 | n = len(a) 18 | return np.array([f'{cls.__name__}({a[i]})' for i in range(n)]) 19 | 20 | return _calc 21 | 22 | 23 | def binary(cls): 24 | def _calc(a, b): 25 | n = len(a) 26 | a = a.astype(str) 27 | b = b.astype(str) 28 | return np.array([f'{cls.__name__}({a[i]},{b[i]})' for i in range(n)]) 29 | 30 | return _calc 31 | 32 | 33 | def rolling(cls, day): 34 | def _calc(a): 35 | n = len(a) 36 | return np.array([f'{cls.__name__}({a[i]},{day})' for i in range(n)]) 37 | 38 | return _calc 39 | 40 | 41 | def rolling_binary(cls, day): 42 | def _calc(a, b): 43 | n = len(a) 44 | a = a.astype(str) 45 | b = b.astype(str) 46 | return np.array([f'{cls.__name__}({a[i]},{b[i]},{day})' for i in range(n)]) 47 | 48 | return _calc 49 | 50 | 51 | funcs: List[GenericOperator] = [] 52 | for op in unary_ops: 53 | funcs.append(GenericOperator(function=unary(op), name=op.__name__, arity=1)) 54 | for op in binary_ops: 55 | funcs.append(GenericOperator(function=binary(op), name=op.__name__, arity=2)) 56 | for op in rolling_ops: 57 | for day in [10, 20, 30, 40, 50]: 58 | funcs.append(GenericOperator(function=rolling(op, day), name=op.__name__ + str(day), arity=1)) 59 | for op in rolling_binary_ops: 60 | for day in [10, 20, 30, 40, 50]: 61 | funcs.append(GenericOperator(function=rolling_binary(op, day), name=op.__name__ + str(day), arity=2)) 62 | -------------------------------------------------------------------------------- /alphagen_llm/client/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import ChatClient, ChatConfig 2 | from .repl import ReplChatClient 3 | from .llama_cpp_client import LlamaCppClient 4 | from .openai_client import OpenAIClient 5 | -------------------------------------------------------------------------------- /alphagen_llm/client/base.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | from typing import Literal, List, Optional, Callable, Tuple, Union, overload 3 | from dataclasses import dataclass 4 | from logging import Logger 5 | 6 | from alphagen.utils.logging import get_null_logger 7 | 8 | 9 | Role = Literal["system", "user", "assistant"] 10 | 11 | 12 | @dataclass 13 | class Message: 14 | role: Role 15 | content: str 16 | 17 | 18 | Dialog = List[Message] 19 | MessageFormatter = Callable[[str, str], str] 20 | def _default_msg_fmt(role: str, content: str) -> str: return f"[{role}] {content}" 21 | 22 | 23 | @dataclass 24 | class ChatConfig: 25 | system_prompt: Optional[str] = None 26 | logger: Optional[Logger] = None 27 | message_formatter: MessageFormatter = _default_msg_fmt 28 | 29 | 30 | class ChatClient(metaclass=ABCMeta): 31 | def __init__(self, config: Optional[ChatConfig] = None) -> None: 32 | config = config or ChatConfig() 33 | self._system_prompt = config.system_prompt 34 | self._logger = config.logger or get_null_logger() 35 | self._msg_fmt = config.message_formatter 36 | self.reset(self._system_prompt) 37 | 38 | @property 39 | def dialog(self) -> Dialog: return self._dialog 40 | 41 | @property 42 | def logger(self) -> Logger: return self._logger 43 | 44 | @property 45 | def message_formatter(self) -> MessageFormatter: return self._msg_fmt 46 | 47 | @overload 48 | def log_message(self, msg: Tuple[str, str]) -> None: ... 49 | 50 | @overload 51 | def log_message(self, msg: Message) -> None: ... 52 | 53 | def log_message(self, msg: Union[Tuple[str, str], Message]) -> None: 54 | if isinstance(msg, Message): 55 | self._logger.debug(self._msg_fmt(msg.role, msg.content)) 56 | else: 57 | self._logger.debug(self._msg_fmt(*msg)) 58 | 59 | def reset(self, system_prompt: Optional[str] = None) -> None: 60 | self.log_message(("script", "Dialog history is reset!")) 61 | self._dialog = [] 62 | if (sys := system_prompt or self._system_prompt) is not None: 63 | self._add_message("system", sys, write_log=system_prompt is not None) 64 | 65 | @abstractmethod 66 | def chat_complete(self, content: str) -> str: ... 67 | 68 | def _add_message(self, role: Role, content: str, write_log: bool = True) -> None: 69 | msg = Message(role, content) 70 | self._dialog.append(msg) 71 | if write_log: 72 | self.log_message(msg) 73 | 74 | _system_prompt: Optional[str] 75 | _dialog: Dialog 76 | _logger: Logger 77 | _msg_fmt: MessageFormatter 78 | -------------------------------------------------------------------------------- /alphagen_llm/client/llama_cpp_client.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List 2 | import requests as req 3 | from .base import ChatClient, ChatConfig, Message 4 | 5 | 6 | _B_INST, _E_INST = "[INST]", "[/INST]" 7 | _B_SYS, _E_SYS = "<>\n", "\n<>\n\n" 8 | _SPECIAL_TAGS = [_B_INST, _E_INST, _B_SYS, _E_SYS] 9 | _BOS_ID, _EOS_ID = 1, 2 # TODO: Hardcoded Llama 2 token ID 10 | 11 | 12 | class LlamaCppClient(ChatClient): 13 | def __init__( 14 | self, 15 | endpoint: str, 16 | config: Optional[ChatConfig] = None 17 | ) -> None: 18 | super().__init__(config) 19 | if self._system_prompt is not None: 20 | self._ensure_no_special_tags(self._system_prompt) 21 | self._endpoint = endpoint 22 | self._slot_id = -1 23 | 24 | def chat_complete(self, content: str) -> str: 25 | self._ensure_no_special_tags(content) 26 | self._add_message("user", content) 27 | res = req.post(f"{self._endpoint}/completion", json={ 28 | "prompt": self._prompt(), 29 | "slot_id": self._slot_id 30 | }) 31 | obj = res.json() 32 | self._slot_id = obj["slot_id"] 33 | answer: str = obj["content"].strip() 34 | self._add_message("assistant", answer) 35 | return answer 36 | 37 | def tokenize(self, prompt: str, bos: bool = False, eos: bool = False) -> List[int]: 38 | res = req.post(f"{self._endpoint}/tokenize", json={"content": prompt}) 39 | ids: list[int] = res.json()["tokens"] 40 | if bos: 41 | ids.insert(0, _BOS_ID) 42 | if eos: 43 | ids.append(_EOS_ID) 44 | return ids 45 | 46 | def decode(self, token_ids: List[int]) -> str: 47 | res = req.post(f"{self._endpoint}/detokenize", json={"tokens": token_ids}) 48 | return res.json()["content"] 49 | 50 | _endpoint: str 51 | _slot_id: int 52 | 53 | def _prompt(self) -> List[int]: 54 | dialog = self._dialog 55 | if dialog[0].role == "system": 56 | assert len(dialog) > 1, "No user prompt after the system prompt." 57 | dialog = [Message( 58 | role=dialog[1].role, 59 | content=f"{_B_SYS}{dialog[0].content}{_E_SYS}{dialog[1].content}" 60 | )] + dialog[2:] 61 | assert ( 62 | len(dialog) % 2 == 1 and 63 | all(msg.role == "user" for msg in dialog[0::2]) and 64 | all(msg.role == "assistant" for msg in dialog[1::2]) 65 | ), ( 66 | "The roles in the dialog must be user/assistant alternating, " 67 | "with an optional system prompt in the front." 68 | ) 69 | tokens = sum( 70 | (self._tokenize_qa_pair(q, a) for q, a in zip(dialog[0:-1:2], dialog[1::2])), 71 | [] 72 | ) 73 | tokens += self._tokenize_qa_pair(dialog[-1]) 74 | return tokens 75 | 76 | def _tokenize_qa_pair(self, prompt: Message, answer: Optional[Message] = None) -> List[int]: 77 | prompt_str = f"{_B_INST} {prompt.content.strip()} {_E_INST}" 78 | if answer is None: 79 | return self.tokenize(prompt_str, bos=True, eos=False) 80 | prompt_str += f" {answer.content.strip()} " 81 | return self.tokenize(prompt_str, bos=True, eos=True) 82 | 83 | @classmethod 84 | def _ensure_no_special_tags(cls, prompt: str) -> None: 85 | assert not any(tag in prompt for tag in _SPECIAL_TAGS), "The message contains special tag." 86 | -------------------------------------------------------------------------------- /alphagen_llm/client/openai_client.py: -------------------------------------------------------------------------------- 1 | from dataclasses import asdict 2 | from openai import OpenAI 3 | import tokentrim as tt 4 | from tokentrim.model_map import MODEL_MAX_TOKENS 5 | 6 | from .base import ChatClient, ChatConfig 7 | 8 | 9 | class OpenAIClient(ChatClient): 10 | def __init__( 11 | self, 12 | client: OpenAI, 13 | config: ChatConfig, 14 | model: str = "gpt-3.5-turbo-0125", 15 | trim_to_token_limit: bool = True 16 | ) -> None: 17 | super().__init__(config) 18 | _update_model_max_tokens() 19 | self._client = client 20 | self._model = model 21 | self._trim = trim_to_token_limit 22 | 23 | def chat_complete(self, content: str) -> str: 24 | self._add_message("user", content) 25 | idx = int(self._system_prompt is not None) 26 | messages = [asdict(msg) for msg in self._dialog[idx:]] 27 | response = self._client.chat.completions.create( 28 | model=self._model, 29 | messages=tt.trim(messages, self._model, self._system_prompt) # type: ignore 30 | ) 31 | result: str = response.choices[0].message.content # type: ignore 32 | self._add_message("assistant", result) 33 | return result 34 | 35 | def _on_reset(self) -> None: 36 | self._start_idx = 0 37 | 38 | _client: OpenAI 39 | _model: str 40 | 41 | 42 | _UPDATED = False 43 | 44 | 45 | def _update_model_max_tokens(): 46 | global _UPDATED 47 | if _UPDATED: 48 | return 49 | MODEL_MAX_TOKENS["gpt-3.5-turbo"] = 16385 50 | MODEL_MAX_TOKENS["gpt-3.5-turbo-1106"] = 16385 51 | MODEL_MAX_TOKENS["gpt-3.5-turbo-0125"] = 16385 52 | -------------------------------------------------------------------------------- /alphagen_llm/client/repl.py: -------------------------------------------------------------------------------- 1 | from logging import Logger 2 | import sys 3 | from typing import Optional 4 | 5 | from .base import ChatClient, ChatConfig 6 | 7 | 8 | class ReplChatClient(ChatClient): 9 | def __init__(self, logger: Optional[Logger] = None): 10 | super().__init__(ChatConfig(logger=logger)) 11 | 12 | def chat_complete(self, content: str) -> str: 13 | self._add_message("user", content) 14 | print(f'{"=" * 28}QUERY{"=" * 28}') 15 | print(content) 16 | print(f'{"=" * 20}INPUT LLM ANSWER HERE{"=" * 20}') 17 | answer = "".join(sys.stdin.readlines()).rstrip('\n') 18 | self._add_message("assistant", answer) 19 | return answer 20 | -------------------------------------------------------------------------------- /alphagen_llm/prompts/common.py: -------------------------------------------------------------------------------- 1 | from enum import IntEnum 2 | from typing import Optional, List, Tuple 3 | from num2words import num2words 4 | from alphagen.data.expression import Expression 5 | from alphagen.data.parser import ExpressionParser 6 | 7 | 8 | class MetricDescriptionMode(IntEnum): 9 | NOT_INCLUDED = 0 # Description of this metric is not included in the prompt. 10 | INCLUDED = 1 # Description of this metric is included in the prompt. 11 | SORTED_BY = 2 # Description is included, and the alphas will be sorted according to this metric. 12 | 13 | 14 | def alpha_word(n: int) -> str: return "alpha" if n == 1 else "alphas" 15 | 16 | 17 | def alpha_phrase(n: int, adjective: Optional[str] = None) -> str: 18 | n_word = str(n) if n > 10 else num2words(n) 19 | adjective = f" {adjective}" if adjective is not None else "" 20 | return f"{n_word}{adjective} {alpha_word(n)}" 21 | 22 | 23 | def safe_parse(parser: ExpressionParser, expr_str: str) -> Optional[Expression]: 24 | try: 25 | return parser.parse(expr_str) 26 | except: 27 | return None 28 | 29 | 30 | def safe_parse_list(lines: List[str], parser: ExpressionParser) -> Tuple[List[Expression], List[str]]: 31 | parsed, invalid = [], [] 32 | for line in lines: 33 | if line == "": 34 | continue 35 | if (e := safe_parse(parser, line)) is not None: 36 | parsed.append(e) 37 | else: 38 | invalid.append(line) 39 | return parsed, invalid 40 | -------------------------------------------------------------------------------- /alphagen_qlib/calculator.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | from torch import Tensor 3 | from alphagen.data.calculator import TensorAlphaCalculator 4 | from alphagen.data.expression import Expression 5 | from alphagen.utils.pytorch_utils import normalize_by_day 6 | from alphagen_qlib.stock_data import StockData 7 | 8 | 9 | class QLibStockDataCalculator(TensorAlphaCalculator): 10 | def __init__(self, data: StockData, target: Optional[Expression] = None): 11 | super().__init__(normalize_by_day(target.evaluate(data)) if target is not None else None) 12 | self.data = data 13 | 14 | def evaluate_alpha(self, expr: Expression) -> Tensor: 15 | return normalize_by_day(expr.evaluate(self.data)) 16 | 17 | @property 18 | def n_days(self) -> int: 19 | return self.data.n_days 20 | -------------------------------------------------------------------------------- /alphagen_qlib/stock_data.py: -------------------------------------------------------------------------------- 1 | from typing import List, Union, Optional, Tuple 2 | from enum import IntEnum 3 | import numpy as np 4 | import pandas as pd 5 | import torch 6 | 7 | 8 | class FeatureType(IntEnum): 9 | OPEN = 0 10 | CLOSE = 1 11 | HIGH = 2 12 | LOW = 3 13 | VOLUME = 4 14 | VWAP = 5 15 | 16 | 17 | _DEFAULT_QLIB_DATA_PATH = "~/.qlib/qlib_data/cn_data" 18 | _QLIB_INITIALIZED = False 19 | 20 | 21 | def initialize_qlib(qlib_data_path: str = _DEFAULT_QLIB_DATA_PATH) -> None: 22 | import qlib 23 | from qlib.config import REG_CN 24 | qlib.init(provider_uri=qlib_data_path, region=REG_CN) 25 | global _QLIB_INITIALIZED 26 | _QLIB_INITIALIZED = True 27 | 28 | 29 | class StockData: 30 | _qlib_initialized: bool = False 31 | 32 | def __init__( 33 | self, 34 | instrument: Union[str, List[str]], 35 | start_time: str, 36 | end_time: str, 37 | max_backtrack_days: int = 100, 38 | max_future_days: int = 30, 39 | features: Optional[List[FeatureType]] = None, 40 | device: torch.device = torch.device("cuda:0"), 41 | preloaded_data: Optional[Tuple[torch.Tensor, pd.Index, pd.Index]] = None 42 | ) -> None: 43 | self._init_qlib() 44 | 45 | self._instrument = instrument 46 | self.max_backtrack_days = max_backtrack_days 47 | self.max_future_days = max_future_days 48 | self._start_time = start_time 49 | self._end_time = end_time 50 | self._features = features if features is not None else list(FeatureType) 51 | self.device = device 52 | data_tup = preloaded_data if preloaded_data is not None else self._get_data() 53 | self.data, self._dates, self._stock_ids = data_tup 54 | 55 | @classmethod 56 | def _init_qlib(cls) -> None: 57 | global _QLIB_INITIALIZED 58 | if not _QLIB_INITIALIZED: 59 | initialize_qlib() 60 | 61 | def _load_exprs(self, exprs: Union[str, List[str]]) -> pd.DataFrame: 62 | # This evaluates an expression on the data and returns the dataframe 63 | # It might throw on illegal expressions like "Ref(constant, dtime)" 64 | from qlib.data.dataset.loader import QlibDataLoader 65 | from qlib.data import D 66 | if not isinstance(exprs, list): 67 | exprs = [exprs] 68 | cal: np.ndarray = D.calendar() 69 | start_index = cal.searchsorted(pd.Timestamp(self._start_time)) # type: ignore 70 | end_index = cal.searchsorted(pd.Timestamp(self._end_time)) # type: ignore 71 | real_start_time = cal[start_index - self.max_backtrack_days] 72 | if cal[end_index] != pd.Timestamp(self._end_time): 73 | end_index -= 1 74 | real_end_time = cal[end_index + self.max_future_days] 75 | return (QlibDataLoader(config=exprs) # type: ignore 76 | .load(self._instrument, real_start_time, real_end_time)) 77 | 78 | def _get_data(self) -> Tuple[torch.Tensor, pd.Index, pd.Index]: 79 | features = ['$' + f.name.lower() for f in self._features] 80 | df = self._load_exprs(features) 81 | df = df.stack().unstack(level=1) 82 | dates = df.index.levels[0] # type: ignore 83 | stock_ids = df.columns 84 | values = df.values 85 | values = values.reshape((-1, len(features), values.shape[-1])) # type: ignore 86 | return torch.tensor(values, dtype=torch.float, device=self.device), dates, stock_ids 87 | 88 | def __getitem__(self, slc: slice) -> "StockData": 89 | "Get a subview of the data given a date slice or an index slice." 90 | if slc.step is not None: 91 | raise ValueError("Only support slice with step=None") 92 | if isinstance(slc.start, str): 93 | return self[self.find_date_slice(slc.start, slc.stop)] 94 | start, stop = slc.start, slc.stop 95 | start = start if start is not None else 0 96 | stop = (stop if stop is not None else self.n_days) + self.max_future_days + self.max_backtrack_days 97 | start = max(0, start) 98 | stop = min(self.data.shape[0], stop) 99 | idx_range = slice(start, stop) 100 | data = self.data[idx_range] 101 | remaining = data.isnan().reshape(-1, data.shape[-1]).all(dim=0).logical_not().nonzero().flatten() 102 | data = data[:, :, remaining] 103 | return StockData( 104 | instrument=self._instrument, 105 | start_time=self._dates[start + self.max_backtrack_days].strftime("%Y-%m-%d"), 106 | end_time=self._dates[stop - 1 - + self.max_future_days].strftime("%Y-%m-%d"), 107 | max_backtrack_days=self.max_backtrack_days, 108 | max_future_days=self.max_future_days, 109 | features=self._features, 110 | device=self.device, 111 | preloaded_data=(data, self._dates[idx_range], self._stock_ids[remaining.tolist()]) 112 | ) 113 | 114 | def find_date_index(self, date: str, exclusive: bool = False) -> int: 115 | ts = pd.Timestamp(date) 116 | idx: int = self._dates.searchsorted(ts) # type: ignore 117 | if exclusive and self._dates[idx] == ts: 118 | idx += 1 119 | idx -= self.max_backtrack_days 120 | if idx < 0 or idx > self.n_days: 121 | raise ValueError(f"Date {date} is out of range: available [{self._start_time}, {self._end_time}]") 122 | return idx 123 | 124 | def find_date_slice(self, start_time: Optional[str] = None, end_time: Optional[str] = None) -> slice: 125 | """ 126 | Find a slice of indices corresponding to the given date range. 127 | For the input, both ends are inclusive. The output is a normal left-closed right-open slice. 128 | """ 129 | start = None if start_time is None else self.find_date_index(start_time) 130 | stop = None if end_time is None else self.find_date_index(end_time, exclusive=False) 131 | return slice(start, stop) 132 | 133 | @property 134 | def n_features(self) -> int: 135 | return len(self._features) 136 | 137 | @property 138 | def n_stocks(self) -> int: 139 | return self.data.shape[-1] 140 | 141 | @property 142 | def n_days(self) -> int: 143 | return self.data.shape[0] - self.max_backtrack_days - self.max_future_days 144 | 145 | @property 146 | def stock_ids(self) -> pd.Index: 147 | return self._stock_ids 148 | 149 | def make_dataframe( 150 | self, 151 | data: Union[torch.Tensor, List[torch.Tensor]], 152 | columns: Optional[List[str]] = None 153 | ) -> pd.DataFrame: 154 | """ 155 | Parameters: 156 | - `data`: a tensor of size `(n_days, n_stocks[, n_columns])`, or 157 | a list of tensors of size `(n_days, n_stocks)` 158 | - `columns`: an optional list of column names 159 | """ 160 | if isinstance(data, list): 161 | data = torch.stack(data, dim=2) 162 | if len(data.shape) == 2: 163 | data = data.unsqueeze(2) 164 | if columns is None: 165 | columns = [str(i) for i in range(data.shape[2])] 166 | n_days, n_stocks, n_columns = data.shape 167 | if self.n_days != n_days: 168 | raise ValueError(f"number of days in the provided tensor ({n_days}) doesn't " 169 | f"match that of the current StockData ({self.n_days})") 170 | if self.n_stocks != n_stocks: 171 | raise ValueError(f"number of stocks in the provided tensor ({n_stocks}) doesn't " 172 | f"match that of the current StockData ({self.n_stocks})") 173 | if len(columns) != n_columns: 174 | raise ValueError(f"size of columns ({len(columns)}) doesn't match with " 175 | f"tensor feature count ({data.shape[2]})") 176 | if self.max_future_days == 0: 177 | date_index = self._dates[self.max_backtrack_days:] 178 | else: 179 | date_index = self._dates[self.max_backtrack_days:-self.max_future_days] 180 | index = pd.MultiIndex.from_product([date_index, self._stock_ids]) 181 | data = data.reshape(-1, n_columns) 182 | return pd.DataFrame(data.detach().cpu().numpy(), index=index, columns=columns) 183 | -------------------------------------------------------------------------------- /alphagen_qlib/strategy.py: -------------------------------------------------------------------------------- 1 | import copy 2 | from math import isnan 3 | from typing import List, Optional, Tuple 4 | import numpy as np 5 | import pandas as pd 6 | from qlib.contrib.strategy.signal_strategy import BaseSignalStrategy 7 | from qlib.backtest.decision import Order, OrderDir, TradeDecisionWO 8 | from alphagen.trade.base import StockCode, StockPosition, StockSignal, StockStatus 9 | 10 | from alphagen.trade.strategy import Strategy 11 | 12 | 13 | class TopKSwapNStrategy(BaseSignalStrategy, Strategy): 14 | def __init__( 15 | self, 16 | K, n_swap, 17 | min_hold_days=1, 18 | only_tradable=False, 19 | **kwargs, 20 | ): 21 | super().__init__(**kwargs) 22 | self.K = K 23 | self.n_swap = n_swap 24 | self.min_hold_days = min_hold_days 25 | self.only_tradable = only_tradable 26 | 27 | def step_decision(self, 28 | status_df: pd.DataFrame, 29 | position_df: Optional[pd.DataFrame] = None 30 | ) -> Tuple[List[StockCode], List[StockCode]]: 31 | signal = dict(zip(status_df['code'], status_df['signal'])) 32 | unbuyable = set(record['code'] for record in status_df.to_dict('records') if not record['buyable']) 33 | unsellable = set(record['code'] for record in status_df.to_dict('records') if not record['sellable']) 34 | 35 | if position_df is None: 36 | days_holded = dict() 37 | else: 38 | days_holded = dict(zip(position_df['code'], position_df['days_holded'])) 39 | 40 | all_valid_stocks = set(k for k, v in signal.items() if not isnan(v)) 41 | all_holding_stocks = days_holded.keys() 42 | valid_holding_stocks = all_holding_stocks & all_valid_stocks 43 | 44 | n_to_open = self.K - len(days_holded) 45 | not_holding_stocks = all_valid_stocks - valid_holding_stocks 46 | 47 | to_buy, to_sell, to_open = [], [], [] 48 | 49 | holding_priority = [] # All sellable stocks in descending order 50 | for stock_id in sorted(valid_holding_stocks, key=signal.get, reverse=True): 51 | if stock_id in unsellable: 52 | continue 53 | if days_holded[stock_id] < self.min_hold_days: 54 | continue 55 | holding_priority.append(stock_id) 56 | 57 | for stock_id in sorted(not_holding_stocks, key=signal.get, reverse=True): 58 | if stock_id in unbuyable: 59 | continue 60 | 61 | can_swap = len(to_buy) >= self.n_swap and holding_priority and signal[stock_id] > signal[holding_priority[-1]] 62 | if can_swap: 63 | to_sell.append(holding_priority.pop()) 64 | to_buy.append(stock_id) 65 | elif len(to_open) < n_to_open: 66 | to_open.append(stock_id) 67 | else: 68 | break 69 | 70 | return to_buy + to_open, to_sell 71 | 72 | def generate_trade_decision(self, execute_result=None): 73 | trade_step = self.trade_calendar.get_trade_step() 74 | trade_start_time, trade_end_time = self.trade_calendar.get_step_time(trade_step) 75 | pred_start_time, pred_end_time = self.trade_calendar.get_step_time(trade_step, shift=1) 76 | pred_score = self.signal.get_signal(start_time=pred_start_time, end_time=pred_end_time) 77 | time_per_step = self.trade_calendar.get_freq() 78 | current_temp = copy.deepcopy(self.trade_position) 79 | cash = current_temp.get_cash() 80 | 81 | if isinstance(pred_score, pd.DataFrame): 82 | pred_score = pred_score.iloc[:, 0] 83 | if pred_score is None: 84 | return TradeDecisionWO([], self) 85 | 86 | stock_signal: StockSignal = pred_score.to_dict() 87 | 88 | def get_holding(stock_id: str): 89 | amount = None # Not required yet 90 | days_holded = current_temp.get_stock_count(stock_id, bar=time_per_step) 91 | return dict(code=stock_id, amount=amount, days_holded=days_holded) 92 | 93 | position = pd.DataFrame.from_records([ 94 | get_holding(stock_id) for stock_id in current_temp.get_stock_list() 95 | ], columns=['code', 'days_holded']) 96 | 97 | def get_status(stock_id: str): 98 | if isnan(stock_signal[stock_id]): 99 | return dict(code=stock_id, signal=np.nan, buyable=False, sellable=False) 100 | buyable = self.trade_exchange.is_stock_tradable( 101 | stock_id=stock_id, 102 | start_time=trade_start_time, 103 | end_time=trade_end_time, 104 | direction=Order.BUY 105 | ) 106 | sellable = self.trade_exchange.is_stock_tradable( 107 | stock_id=stock_id, 108 | start_time=trade_start_time, 109 | end_time=trade_end_time, 110 | direction=Order.SELL 111 | ) 112 | return dict(code=stock_id, signal=stock_signal[stock_id], buyable=buyable, sellable=sellable) 113 | 114 | status = pd.DataFrame.from_records([get_status(code) for code in stock_signal]) 115 | 116 | to_buy, to_sell = self.step_decision(status_df=status, position_df=position) 117 | 118 | buy_orders, sell_orders = [], [] 119 | 120 | for code in to_sell: 121 | sell_amount = current_temp.get_stock_amount(code=code) 122 | sell_order = Order( 123 | stock_id=code, 124 | amount=sell_amount, 125 | start_time=trade_start_time, 126 | end_time=trade_end_time, 127 | direction=Order.SELL, 128 | ) 129 | if self.trade_exchange.check_order(sell_order): 130 | sell_orders.append(sell_order) 131 | trade_val, trade_cost, trade_price = self.trade_exchange.deal_order( 132 | sell_order, position=current_temp 133 | ) 134 | cash += trade_val - trade_cost 135 | 136 | value = cash * self.risk_degree / len(to_buy) if len(to_buy) > 0 else 0 137 | 138 | for code in to_buy: 139 | buy_price = self.trade_exchange.get_deal_price( 140 | stock_id=code, start_time=trade_start_time, end_time=trade_end_time, direction=OrderDir.BUY 141 | ) 142 | buy_amount = value / buy_price 143 | factor = self.trade_exchange.get_factor(stock_id=code, start_time=trade_start_time, end_time=trade_end_time) 144 | buy_amount = self.trade_exchange.round_amount_by_trade_unit(buy_amount, factor) 145 | buy_order = Order( 146 | stock_id=code, 147 | amount=buy_amount, 148 | start_time=trade_start_time, 149 | end_time=trade_end_time, 150 | direction=Order.BUY, 151 | ) 152 | buy_orders.append(buy_order) 153 | 154 | return TradeDecisionWO(sell_orders + buy_orders, self) 155 | -------------------------------------------------------------------------------- /alphagen_qlib/utils.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import json 3 | from typing import List, Tuple 4 | from alphagen.data.expression import * 5 | from alphagen_generic.features import * 6 | 7 | from alphagen_qlib.stock_data import StockData 8 | from alphagen.data.parser import ExpressionParser 9 | 10 | 11 | def load_recent_data(instrument: str, 12 | window_size: int = 365, 13 | offset: int = 1, 14 | **kwargs) -> Tuple[StockData, str]: 15 | today = datetime.date.today() 16 | start_date = str(today - datetime.timedelta(days=window_size)) 17 | end_date = str(today - datetime.timedelta(days=offset)) 18 | 19 | return StockData(instrument=instrument, 20 | start_time=start_date, 21 | end_time=end_date, 22 | max_future_days=0, 23 | **kwargs), end_date 24 | 25 | 26 | def load_alpha_pool(raw) -> Tuple[List[Expression], List[float]]: 27 | parser = ExpressionParser(Operators) 28 | exprs = [parser.parse(e) for e in raw["exprs"]] 29 | weights = raw["weights"] 30 | return exprs, weights 31 | 32 | 33 | def load_alpha_pool_by_path(path: str) -> Tuple[List[Expression], List[float]]: 34 | with open(path, encoding="utf-8") as f: 35 | raw = json.load(f) 36 | return load_alpha_pool(raw) 37 | -------------------------------------------------------------------------------- /backtest.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, TypeVar, Callable, Optional, Tuple 2 | import os 3 | import pickle 4 | import warnings 5 | import pandas as pd 6 | import json 7 | from pathlib import Path 8 | from dataclasses import dataclass 9 | from dataclasses_json import DataClassJsonMixin 10 | 11 | from qlib.backtest import backtest, executor as exec 12 | from qlib.contrib.evaluate import risk_analysis 13 | from qlib.contrib.report.analysis_position import report_graph 14 | from qlib.contrib.strategy import TopkDropoutStrategy 15 | 16 | from alphagen.data.expression import * 17 | from alphagen.data.parser import parse_expression 18 | from alphagen_generic.features import * 19 | from alphagen_qlib.stock_data import StockData, initialize_qlib 20 | from alphagen_qlib.calculator import QLibStockDataCalculator 21 | from alphagen_qlib.utils import load_alpha_pool_by_path 22 | 23 | 24 | _T = TypeVar("_T") 25 | 26 | 27 | def _create_parents(path: str) -> None: 28 | dir = os.path.dirname(path) 29 | if dir != "": 30 | os.makedirs(dir, exist_ok=True) 31 | 32 | 33 | def write_all_text(path: str, text: str) -> None: 34 | _create_parents(path) 35 | with open(path, "w") as f: 36 | f.write(text) 37 | 38 | 39 | def dump_pickle(path: str, 40 | factory: Callable[[], _T], 41 | invalidate_cache: bool = False) -> Optional[_T]: 42 | if invalidate_cache or not os.path.exists(path): 43 | _create_parents(path) 44 | obj = factory() 45 | with open(path, "wb") as f: 46 | pickle.dump(obj, f) 47 | return obj 48 | 49 | 50 | @dataclass 51 | class BacktestResult(DataClassJsonMixin): 52 | sharpe: float 53 | annual_return: float 54 | max_drawdown: float 55 | information_ratio: float 56 | annual_excess_return: float 57 | excess_max_drawdown: float 58 | 59 | 60 | class QlibBacktest: 61 | def __init__( 62 | self, 63 | benchmark: str = "SH000300", 64 | top_k: int = 50, 65 | n_drop: Optional[int] = None, 66 | deal: str = "close", 67 | open_cost: float = 0.0015, 68 | close_cost: float = 0.0015, 69 | min_cost: float = 5, 70 | ): 71 | self._benchmark = benchmark 72 | self._top_k = top_k 73 | self._n_drop = n_drop if n_drop is not None else top_k 74 | self._deal_price = deal 75 | self._open_cost = open_cost 76 | self._close_cost = close_cost 77 | self._min_cost = min_cost 78 | 79 | def run( 80 | self, 81 | prediction: Union[pd.Series, pd.DataFrame], 82 | output_prefix: Optional[str] = None 83 | ) -> Tuple[pd.DataFrame, BacktestResult]: 84 | prediction = prediction.sort_index() 85 | index: pd.MultiIndex = prediction.index.remove_unused_levels() # type: ignore 86 | dates = index.levels[0] 87 | 88 | def backtest_impl(last: int = -1): 89 | with warnings.catch_warnings(): 90 | warnings.simplefilter("ignore") 91 | strategy = TopkDropoutStrategy( 92 | signal=prediction, 93 | topk=self._top_k, 94 | n_drop=self._n_drop, 95 | only_tradable=True, 96 | forbid_all_trade_at_limit=True 97 | ) 98 | executor = exec.SimulatorExecutor( 99 | time_per_step="day", 100 | generate_portfolio_metrics=True 101 | ) 102 | return backtest( 103 | strategy=strategy, 104 | executor=executor, 105 | start_time=dates[0], 106 | end_time=dates[last], 107 | account=100_000_000, 108 | benchmark=self._benchmark, 109 | exchange_kwargs={ 110 | "limit_threshold": 0.095, 111 | "deal_price": self._deal_price, 112 | "open_cost": self._open_cost, 113 | "close_cost": self._close_cost, 114 | "min_cost": self._min_cost, 115 | } 116 | )[0] 117 | 118 | try: 119 | portfolio_metric = backtest_impl() 120 | except IndexError: 121 | print("Cannot backtest till the last day, trying again with one less day") 122 | portfolio_metric = backtest_impl(-2) 123 | 124 | report, _ = portfolio_metric["1day"] # type: ignore 125 | result = self._analyze_report(report) 126 | graph = report_graph(report, show_notebook=False)[0] 127 | if output_prefix is not None: 128 | dump_pickle(output_prefix + "-report.pkl", lambda: report, True) 129 | dump_pickle(output_prefix + "-graph.pkl", lambda: graph, True) 130 | write_all_text(output_prefix + "-result.json", result.to_json()) 131 | return report, result 132 | 133 | def _analyze_report(self, report: pd.DataFrame) -> BacktestResult: 134 | excess = risk_analysis(report["return"] - report["bench"] - report["cost"])["risk"] 135 | returns = risk_analysis(report["return"] - report["cost"])["risk"] 136 | 137 | def loc(series: pd.Series, field: str) -> float: 138 | return series.loc[field] # type: ignore 139 | 140 | return BacktestResult( 141 | sharpe=loc(returns, "information_ratio"), 142 | annual_return=loc(returns, "annualized_return"), 143 | max_drawdown=loc(returns, "max_drawdown"), 144 | information_ratio=loc(excess, "information_ratio"), 145 | annual_excess_return=loc(excess, "annualized_return"), 146 | excess_max_drawdown=loc(excess, "max_drawdown"), 147 | ) 148 | 149 | 150 | if __name__ == "__main__": 151 | initialize_qlib("~/.qlib/qlib_data/cn_data") 152 | qlib_backtest = QlibBacktest(top_k=50, n_drop=5) 153 | data = StockData( 154 | instrument="csi300", 155 | start_time="2022-01-01", 156 | end_time="2023-06-30" 157 | ) 158 | calc = QLibStockDataCalculator(data, None) 159 | 160 | def run_backtest(prefix: str, seed: int, exprs: List[Expression], weights: List[float]): 161 | df = data.make_dataframe(calc.make_ensemble_alpha(exprs, weights)) 162 | qlib_backtest.run(df, output_prefix=f"out/backtests/50-5/{prefix}/{seed}") 163 | 164 | for p in Path("out/gp").iterdir(): 165 | seed = int(p.name) 166 | with open(p / "40.json") as f: 167 | report = json.load(f) 168 | state = report["res"]["res"]["pool_state"] 169 | run_backtest("gp", seed, [parse_expression(e) for e in state["exprs"]], state["weights"]) 170 | exit(0) 171 | for p in Path("out/results").iterdir(): 172 | inst, size, seed, time, ver = p.name.split('_', 4) 173 | size, seed = int(size), int(seed) 174 | if inst != "csi300" or size != 20 or time < "20240923" or ver == "llm_d5": 175 | continue 176 | exprs, weights = load_alpha_pool_by_path(str(p / "251904_steps_pool.json")) 177 | run_backtest(ver, seed, exprs, weights) 178 | for p in Path("out/llm-tests/interaction").iterdir(): 179 | if not p.name.startswith("v1"): 180 | continue 181 | run = int(p.name[3]) 182 | with open(p / "report.json") as f: 183 | report = json.load(f) 184 | state = report[-1]["pool_state"] 185 | run_backtest("pure_llm", run, [parse_expression(t[0]) for t in state], [t[1] for t in state]) 186 | -------------------------------------------------------------------------------- /data_collection/baostock_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | from contextlib import redirect_stdout, contextmanager 3 | import baostock as bs 4 | 5 | 6 | def baostock_login(): 7 | with open(os.devnull, "w") as devnull: 8 | with redirect_stdout(devnull): 9 | bs.login() 10 | 11 | 12 | def baostock_logout(): 13 | with open(os.devnull, "w") as devnull: 14 | with redirect_stdout(devnull): 15 | bs.logout() 16 | 17 | 18 | def baostock_relogin(): 19 | with open(os.devnull, "w") as devnull: 20 | with redirect_stdout(devnull): 21 | bs.logout() 22 | bs.login() 23 | 24 | 25 | @contextmanager 26 | def baostock_login_context(): 27 | baostock_login() 28 | try: 29 | yield None 30 | finally: 31 | baostock_logout() 32 | -------------------------------------------------------------------------------- /dso.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from alphagen.data.expression import * 4 | from alphagen_qlib.calculator import QLibStockDataCalculator 5 | from dso import DeepSymbolicRegressor 6 | from dso.library import Token, HardCodedConstant 7 | from dso import functions 8 | from alphagen.models.linear_alpha_pool import MseAlphaPool 9 | from alphagen.utils import reseed_everything 10 | from alphagen_generic.operators import funcs as generic_funcs 11 | from alphagen_generic.features import * 12 | 13 | 14 | 15 | funcs = {func.name: Token(complexity=1, **func._asdict()) for func in generic_funcs} 16 | for i, feature in enumerate(['open', 'close', 'high', 'low', 'volume', 'vwap']): 17 | funcs[f'x{i+1}'] = Token(name=feature, arity=0, complexity=1, function=None, input_var=i) 18 | for v in [-30., -10., -5., -2., -1., -0.5, -0.01, 0.01, 0.5, 1., 2., 5., 10., 30.]: 19 | funcs[f'Constant({v})'] = HardCodedConstant(name=f'Constant({v})', value=v) 20 | 21 | 22 | instruments = 'csi300' 23 | import sys 24 | seed = int(sys.argv[1]) 25 | reseed_everything(seed) 26 | 27 | cache = {} 28 | device = torch.device('cuda:0') 29 | data_train = StockData(instruments, '2009-01-01', '2018-12-31', device=device) 30 | data_valid = StockData(instruments, '2019-01-01', '2019-12-31', device=device) 31 | data_test = StockData(instruments, '2020-01-01', '2021-12-31', device=device) 32 | calculator_train = QLibStockDataCalculator(data_train, target) 33 | calculator_valid = QLibStockDataCalculator(data_valid, target) 34 | calculator_test = QLibStockDataCalculator(data_test, target) 35 | 36 | 37 | if __name__ == '__main__': 38 | X = np.array([['open_', 'close', 'high', 'low', 'volume', 'vwap']]) 39 | y = np.array([[1]]) 40 | functions.function_map = funcs 41 | 42 | pool = MseAlphaPool( 43 | capacity=10, 44 | calculator=calculator_train, 45 | ic_lower_bound=None 46 | ) 47 | 48 | class Ev: 49 | def __init__(self, pool): 50 | self.cnt = 0 51 | self.pool: MseAlphaPool = pool 52 | self.results = {} 53 | 54 | def alpha_ev_fn(self, key): 55 | expr = eval(key) 56 | try: 57 | ret = self.pool.try_new_expr(expr) 58 | except OutOfDataRangeError: 59 | ret = -1. 60 | finally: 61 | self.cnt += 1 62 | if self.cnt % 100 == 0: 63 | test_ic = pool.test_ensemble(calculator_test)[0] 64 | self.results[self.cnt] = test_ic 65 | print(self.cnt, test_ic) 66 | return ret 67 | 68 | ev = Ev(pool) 69 | 70 | config = dict( 71 | task=dict( 72 | task_type='regression', 73 | function_set=list(funcs.keys()), 74 | metric='alphagen', 75 | metric_params=[lambda key: ev.alpha_ev_fn(key)], 76 | ), 77 | training={'n_samples': 20000, 'batch_size': 128, 'epsilon': 0.05}, 78 | prior={'length': {'min_': 2, 'max_': 20, 'on': True}} 79 | ) 80 | 81 | # Create the model 82 | model = DeepSymbolicRegressor(config=config) 83 | model.fit(X, y) 84 | 85 | print(ev.results) 86 | -------------------------------------------------------------------------------- /dso/__init__.py: -------------------------------------------------------------------------------- 1 | from dso.core import DeepSymbolicOptimizer 2 | from dso.task.regression.sklearn import DeepSymbolicRegressor 3 | 4 | -------------------------------------------------------------------------------- /dso/checkpoint.py: -------------------------------------------------------------------------------- 1 | import os 2 | from datetime import datetime, timedelta 3 | 4 | import tensorflow as tf 5 | import numpy as np 6 | import pandas as pd 7 | 8 | from dso.program import Program, from_tokens 9 | 10 | 11 | class Checkpoint(): 12 | """ 13 | A helper class to checkpoint models. 14 | 15 | Methods 16 | ------- 17 | save 18 | Save a checkpoint. 19 | 20 | load 21 | Load from a given checkpoint. 22 | 23 | update 24 | Maybe save a checkpoint depending on frequency configuration. 25 | """ 26 | 27 | def __init__(self, model, load_path=None, save_freq=23, units="hours", 28 | save_on_done=False): 29 | """ 30 | model : dso.DeepSymbolicOptimizer 31 | The model to checkpoint. 32 | 33 | load_path : str or None 34 | Path to initial checkpoint directory to load. If None, do not start from 35 | checkpoint. 36 | 37 | save_freq : float or None 38 | The frequency at which to save a checkpoint. If None, non-final checkpoints 39 | will not be automatically saved. 40 | 41 | units : str 42 | The units of save_freq. Supports "hours", "minutes", "seconds", "iterations". 43 | 44 | save_on_done : bool 45 | Whether to save a final checkpoint upon reaching model.trainer.done. 46 | """ 47 | 48 | self.model = model 49 | if model.save_path is not None: 50 | self.checkpoint_dir = os.path.join(model.save_path, "checkpoint") 51 | os.makedirs(self.checkpoint_dir, exist_ok=True) 52 | else: 53 | self.checkpoint_dir = None 54 | 55 | # Create the Saver 56 | self.saver = tf.train.Saver() 57 | 58 | # Load from existing checkpoint, if given 59 | if load_path is not None: 60 | self.load(load_path) 61 | 62 | # Setup time-based checkpointing 63 | if save_freq is not None and units in ["hours", "minutes", "seconds"]: 64 | if units == "hours": 65 | self.dt = timedelta(hours=save_freq) 66 | elif units == "minutes": 67 | self.dt = timedelta(minutes=save_freq) 68 | elif units == "seconds": 69 | self.dt = timedelta(seconds=save_freq) 70 | self.next_save_time = datetime.now() + self.dt 71 | else: 72 | self.next_save_time = None 73 | self.dt = None 74 | 75 | # Setup iteration-based checkpointing 76 | if save_freq is not None and units == "iterations": 77 | self.save_freq_iters = save_freq 78 | else: 79 | self.save_freq_iters = None 80 | 81 | self.save_on_done = save_on_done 82 | 83 | def update(self): 84 | """ 85 | Maybe a save a checkpoint, depending on configuration. This should be called 86 | each iteration, i.e. after model.trainer.run_one_step(). 87 | """ 88 | 89 | # Save final checkpoint if done 90 | if self.save_on_done and self.model.trainer.done: 91 | self.save() 92 | 93 | # Save if time-based frequency is met 94 | elif self.next_save_time is not None and datetime.now() > self.next_save_time: 95 | self.save() 96 | self.next_save_time = datetime.now() + self.dt 97 | 98 | # Save if iteration-based frequency is met 99 | elif self.save_freq_iters is not None and (self.model.trainer.iteration % self.save_freq_iters) == 0: 100 | self.save() 101 | 102 | def save(self, save_path=None): 103 | """ 104 | Save a checkpoint. 105 | 106 | Parameters 107 | ---------- 108 | save_path : str or None 109 | Directory in which to save checkpoint. If None, save in: 110 | /checkpoint/checkpoint_. 111 | """ 112 | 113 | # Determine the save path 114 | if save_path is None: 115 | assert self.checkpoint_dir is not None, "Cannot support automated checkpointing with model.save_dir=None." 116 | timestamp = datetime.now().strftime("%Y-%m-%d-%H%M%S") 117 | save_path = os.path.join(self.checkpoint_dir, 118 | "checkpoint_{}".format(timestamp)) 119 | if os.path.exists(save_path): 120 | paths = os.listdir(os.path.dirname(save_path)) 121 | paths = [path for path in paths if path.startswith(os.path.basename(save_path))] 122 | save_path += "_{}".format(len(paths)) 123 | os.makedirs(save_path, exist_ok=False) 124 | 125 | # Save the TensorFlow graph 126 | # print("Saving TensorFlow graph...") 127 | tf_save_path = os.path.join(save_path, "tf") 128 | self.saver.save(self.model.sess, tf_save_path) 129 | 130 | # Save the Trainer 131 | # print("Saving Trainer...") 132 | trainer_save_path = os.path.join(save_path, "trainer.json") 133 | self.model.trainer.save(trainer_save_path) 134 | 135 | # Save the priority queue, if applicable 136 | # TBD: This should be in self.model.trainer.save or self.model.trainer.policy_optimizer.save after refactoring PolicyOptimizers to handle their own bookkeeping 137 | if self.model.trainer.priority_queue is not None: 138 | priority_queue_save_path = os.path.join(save_path, "priority_queue.npz") 139 | self.model.trainer.priority_queue.save(priority_queue_save_path) 140 | 141 | # Save the cache 142 | # print("Saving cache...") 143 | # TBD: Abstract into cache saving function 144 | cache_save_path = os.path.join(save_path, "cache.csv") 145 | cache_programs = Program.cache.values() 146 | cache_tokens = [",".join(map(str, p.tokens.tolist())) for p in cache_programs] 147 | cache_rewards = [p.r for p in cache_programs] 148 | cache_data = { "tokens" : cache_tokens, "rewards" : cache_rewards } 149 | cache_df = pd.DataFrame(cache_data) 150 | cache_df.to_csv(cache_save_path, index=False) 151 | 152 | # Save the extra samples that were produced while attempting to 153 | # generate a batch of new and unique samples 154 | if self.model.trainer.policy.valid_extended_batch: 155 | self.model.trainer.policy.valid_extended_batch = False 156 | batch_save_path = os.path.join(save_path, "batch.npz") 157 | with open(batch_save_path, 'wb') as f: 158 | np.savez(f, self.model.trainer.policy.extended_batch) 159 | 160 | def load(self, load_path): 161 | """ 162 | Load model state from checkpoint. 163 | 164 | Parameters 165 | ---------- 166 | load_path : str 167 | Checkpoint directory to load model state. 168 | """ 169 | 170 | # Load the TensorFlow graph 171 | if self.model.sess is None: 172 | self.model.setup() 173 | tf_load_path = os.path.join(load_path, "tf") 174 | self.saver.restore(self.model.sess, tf_load_path) 175 | 176 | # Load the Trainer 177 | # print("Loading Trainer...") 178 | trainer_load_path = os.path.join(load_path, "trainer.json") 179 | self.model.trainer.load(trainer_load_path) 180 | 181 | # Load the priority queue, if applicable 182 | # TBD: This should be in self.model.trainer.load or self.model.trainer.policy_optimizer.load after refactoring PolicyOptimizers to handle their own bookkeeping 183 | if self.model.trainer.priority_queue is not None: 184 | priority_queue_load_path = os.path.join(load_path, "priority_queue.npz") 185 | self.model.trainer.priority_queue.load(priority_queue_load_path) 186 | 187 | # Load the cache 188 | # print("Loading cache...") 189 | cache_load_path = os.path.join(load_path, "cache.csv") 190 | cache_df = pd.read_csv(cache_load_path) 191 | cache_df["tokens"] = cache_df["tokens"].str.split(",") 192 | programs = [from_tokens(np.array(tokens, dtype=np.int32)) for tokens in cache_df["tokens"]] 193 | for p, r in zip(programs, cache_df["rewards"]): 194 | p.r = r 195 | 196 | # Load the extra samples 197 | batch_save_path = os.path.join(load_path, "batch.npz") 198 | if os.path.isfile(batch_save_path): 199 | npzfile = np.load(batch_save_path, allow_pickle=True) 200 | self.model.trainer.policy.extended_batch = npzfile['arr_0'] 201 | self.model.trainer.policy.valid_extended_batch = True 202 | -------------------------------------------------------------------------------- /dso/config/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import commentjson as json 4 | 5 | from dso.utils import safe_merge_dicts 6 | 7 | 8 | def get_base_config(task, language_prior): 9 | # Load base config 10 | with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "config_common.json"), encoding='utf-8') as f: 11 | base_config = json.load(f) 12 | 13 | # Load task specific config 14 | task_config_file = None 15 | if task in ["regression", None]: 16 | task_config_file = "config_regression.json" 17 | elif task in ["control"]: 18 | task_config_file = "config_control.json" 19 | else: 20 | # Custom tasks use config_common.json. 21 | task_config_file = "config_common.json" 22 | with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), task_config_file), encoding='utf-8') as f: 23 | task_config = json.load(f) 24 | 25 | # Load language prior config 26 | if language_prior: 27 | with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "config_language.json"), encoding='utf-8') as f: 28 | language_config = json.load(f) 29 | task_config = safe_merge_dicts(task_config, language_config) 30 | 31 | return safe_merge_dicts(base_config, task_config) 32 | 33 | 34 | def load_config(config=None): 35 | # Load user config 36 | if isinstance(config, str): 37 | with open(config, encoding='utf-8') as f: 38 | user_config = json.load(f) 39 | elif isinstance(config, dict): 40 | user_config = config 41 | else: 42 | assert config is None, "Config must be None, str, or dict." 43 | user_config = {} 44 | 45 | # Determine the task and language prior 46 | try: 47 | task = user_config["task"]["task_type"] 48 | except KeyError: 49 | task = "regression" 50 | print("WARNING: Task type not specified. Falling back to default task type '{}' to load config.".format(task)) 51 | try: 52 | language_prior = user_config["prior"]["language_model"]["on"] 53 | except KeyError: 54 | language_prior = False 55 | 56 | # Load task-specific base config 57 | base_config = get_base_config(task, language_prior) 58 | 59 | # Return combined configs 60 | return safe_merge_dicts(base_config, user_config) 61 | -------------------------------------------------------------------------------- /dso/config/config_common.json: -------------------------------------------------------------------------------- 1 | { 2 | // Experiment configuration. 3 | "experiment" : { 4 | 5 | // Root directory to save results. 6 | "logdir" : "./log", 7 | 8 | // Save dir within logdir (will be set automatically to _ if null) 9 | "exp_name" : null, 10 | 11 | // Random number seed. Don't forget to change this for multiple runs! 12 | "seed" : 0 13 | }, 14 | 15 | // Task-specific hyperparameters. See task-specific configs (e.g. 16 | // config_regression.json) for more info. 17 | "task" : { 18 | 19 | // Choice of ["regression", "control"]. 20 | "task_type" : null 21 | }, 22 | 23 | // Hyperparameters related to the main training loop. 24 | "training" : { 25 | 26 | // These parameters control the length of the run. 27 | "n_samples" : 2000000, 28 | "batch_size" : 1000, 29 | 30 | // To use the risk-seeking policy gradient, set epsilon < 1.0 and 31 | // baseline="R_e" 32 | "epsilon" : 0.05, 33 | "baseline" : "R_e", 34 | 35 | // Control variate parameters for vanilla policy gradient. If risk-seeking 36 | // is used, these have no effect. 37 | "alpha" : 0.5, 38 | "b_jumpstart" : false, 39 | 40 | // Number of cores to use when evaluating a batch of rewards. For batch 41 | // runs using run.py and --runs > 1, this will be overridden to 1. For 42 | // single runs, recommended to set this to as many cores as you can use! 43 | "n_cores_batch" : 1, 44 | 45 | // The complexity measure is only used to compute a Pareto front. It does 46 | // not affect the optimization. 47 | "complexity" : "token", 48 | 49 | // The constant optimizer used to optimized each "const" token. 50 | "const_optimizer" : "scipy", 51 | "const_params" : { 52 | "method" : "L-BFGS-B", 53 | "options" : { 54 | "gtol" : 1e-3 55 | } 56 | }, 57 | "verbose" : true, 58 | 59 | // Debug level 60 | "debug" : 0, 61 | 62 | // Whether to stop early if success condition is met 63 | "early_stopping" : true, 64 | 65 | // EXPERIMENTAL: Hyperparameters related to utilizing a memory buffer. 66 | "use_memory" : false, 67 | "memory_capacity" : 1e3, 68 | "warm_start" : null, 69 | "memory_threshold" : null 70 | }, 71 | 72 | // Parameters to control what outputs to save. 73 | "logging" : { 74 | "save_all_iterations" : false, 75 | "save_summary" : false, 76 | "save_positional_entropy" : false, 77 | "save_pareto_front" : true, 78 | "save_cache" : false, 79 | "save_cache_r_min" : 0.9, 80 | "save_freq" : 1, 81 | "save_token_count" : false, 82 | // Size of the "hall of fame" (top performers during training) to save. 83 | "hof" : 100 84 | }, 85 | 86 | // The State Manager defines the inputs to the Controller 87 | "state_manager": { 88 | "type" : "hierarchical", 89 | // Observation hyperparameters 90 | "observe_action" : false, 91 | "observe_parent" : true, 92 | "observe_sibling" : true, 93 | "observe_dangling" : false, 94 | "embedding" : false, 95 | "embedding_size" : 8 96 | }, 97 | 98 | // Hyperparameters related to the policy, i.e. parameterized distribution over objects. 99 | "policy" : { 100 | 101 | // Type of policy 102 | // "rnn" is equivalent to "dso.policy.rnn_policy:RNNPolicy" 103 | "policy_type" : "rnn", 104 | 105 | // Maximum sequence length. 106 | "max_length" : 64, 107 | 108 | // Policy parameters 109 | "cell" : "lstm", 110 | "num_layers" : 1, 111 | "num_units" : 32, 112 | "initializer" : "zeros" 113 | }, 114 | 115 | // Hyperparameters related to the discrete distribution model over objects. 116 | "policy_optimizer" : { 117 | 118 | // Type of policy optimizer 119 | // "pg" is equivalent to "dso.policy_optimizer.pg_policy_optimizer:PGPolicyOptimizer" 120 | "policy_optimizer_type" : "pg", 121 | 122 | // Whether to compute TensorBoard summaries. 123 | "summary" : false, 124 | 125 | // Optimizer hyperparameters. 126 | "learning_rate" : 0.001, 127 | "optimizer" : "adam", 128 | 129 | // Entropy regularizer hyperparameters. 130 | "entropy_weight" : 0.005, 131 | "entropy_gamma" : 1.0 132 | }, 133 | 134 | // Hyperparameters related to genetic programming hybrid methods. 135 | "gp_meld" : { 136 | "run_gp_meld" : false, 137 | "verbose" : false, 138 | "generations" : 20, 139 | "p_crossover" : 0.5, 140 | "p_mutate" : 0.5, 141 | "tournament_size" : 5, 142 | "train_n" : 50, 143 | "mutate_tree_max" : 3, 144 | // Speeds up processing when doing expensive evaluations. 145 | "parallel_eval" : false 146 | }, 147 | 148 | 149 | // Hyperparameters related to including in situ priors and constraints. Each 150 | // prior must explicitly be turned "on" or it will not be used. 151 | "prior": { 152 | // Whether to count number of constraints (useful diagnostic but has some overhead) 153 | "count_constraints" : false, 154 | 155 | // Generic constraint that [targets] cannot be the [relationship] of 156 | // [effectors]. See prior.py for supported relationships. 157 | "relational" : { 158 | "targets" : [], 159 | "effectors" : [], 160 | "relationship" : null, 161 | "on" : false 162 | }, 163 | 164 | // Constrains expression length to fall between min_ and max_, inclusive. 165 | "length" : { 166 | "min_" : 4, 167 | "max_" : 30, 168 | "on" : false 169 | }, 170 | 171 | // Constraints "const" to appear at most max_ times. 172 | "repeat" : { 173 | "tokens" : "const", 174 | "min_" : null, 175 | "max_" : 3, 176 | "on" : false 177 | }, 178 | 179 | // Prevents consecutive inverse unary operators, e.g. log(exp(x)). 180 | "inverse" : { 181 | "on" : false 182 | }, 183 | 184 | // Prevents nested trig operators (e.g. sin(1 + cos(x))). 185 | "trig" : { 186 | "on" : false 187 | }, 188 | 189 | // Prevents "const" from being the only child, e.g. sin(const) or 190 | // add(const, const). 191 | "const" : { 192 | "on" : false 193 | }, 194 | 195 | // Prevents expressions with no input variables, e.g. sin(1.0 + const). 196 | "no_inputs" : { 197 | "on" : false 198 | }, 199 | 200 | // Uniform arity prior: the prior probabilities over arities is uniform. 201 | "uniform_arity" : { 202 | "on" : false 203 | }, 204 | 205 | // Soft length prior: expressions are discouraged to have length far from 206 | // loc. 207 | "soft_length" : { 208 | "loc" : 10, 209 | "scale" : 5, 210 | "on" : false 211 | }, 212 | 213 | // Prevents first token if its range does not contain X range 214 | // and prevents input variable X if its parent's domain does not contain X domain 215 | "domain_range" : { 216 | "on" : false 217 | }, 218 | 219 | // Constraint on selection of multi-discrete actions 220 | "multi_discrete" : { 221 | "dense" : false, // all action dims are filled 222 | "ordered" : false, // action dim in ascending order 223 | "on" : false 224 | } 225 | }, 226 | 227 | // Postprocessing hyperparameters. 228 | "postprocess" : { 229 | "show_count" : 5, 230 | "save_plots" : true 231 | } 232 | } 233 | -------------------------------------------------------------------------------- /dso/config/config_regression.json: -------------------------------------------------------------------------------- 1 | { 2 | "task" : { 3 | // Deep Symbolic Regression 4 | "task_type" : "regression", 5 | 6 | // This can either be (1) the name of the benchmark dataset (see 7 | // benchmarks.csv for a list of supported benchmarks) or (2) a path to a 8 | // CSV file containing the data. 9 | "dataset" : "Nguyen-1", 10 | 11 | // To customize a function set, edit this! See functions.py for a list of 12 | // supported functions. Note "const" will add placeholder constants that 13 | // will be optimized within the training loop. This will considerably 14 | // increase runtime. 15 | "function_set": ["add", "sub", "mul", "div", "sin", "cos", "exp", "log"], 16 | 17 | // Metric to be used for the reward function. See regression.py for 18 | // supported metrics. 19 | "metric" : "inv_nrmse", 20 | "metric_params" : [1.0], 21 | 22 | // Optional alternate metric to be used at evaluation time. 23 | "extra_metric_test" : null, 24 | "extra_metric_test_params" : [], 25 | 26 | // NRMSE threshold for early stopping. This is useful for noiseless 27 | // benchmark problems when DSO discovers the true solution. 28 | "threshold" : 1e-12, 29 | 30 | // With protected=false, floating-point errors (e.g. log of negative 31 | // number) will simply returns a minimal reward. With protected=true, 32 | // "protected" functions will prevent floating-point errors, but may 33 | // introduce discontinuities in the learned functions. 34 | "protected" : false, 35 | 36 | // You can add artificial reward noise directly to the reward function. 37 | // Note this does NOT add noise to the dataset. 38 | "reward_noise" : 0.0, 39 | "reward_noise_type" : "r", 40 | "normalize_variance" : false, 41 | 42 | // Set of thresholds (shared by all input variables) for building 43 | // decision trees. Note that no StateChecker will be added to Library 44 | // if decision_tree_threshold_set is an empty list or null. 45 | "decision_tree_threshold_set" : [], 46 | 47 | // Parameters for optimizing the "poly" token. 48 | // Note: poly_optimizer is turned on if and only if "poly" is in function_set. 49 | "poly_optimizer_params" : { 50 | // The (maximal) degree of the polynomials used to fit the data 51 | "degree": 3, 52 | // Cutoff value for the coefficients of polynomials. Coefficients 53 | // with magnitude less than this value will be regarded as 0. 54 | "coef_tol": 1e-6, 55 | // linear models from sklearn: linear_regression, lasso, 56 | // and ridge are currently supported, or our own implementation 57 | // of least squares regressor "dso_least_squares". 58 | "regressor": "dso_least_squares", 59 | "regressor_params": { 60 | // Cutoff value for p-value of coefficients. Coefficients with 61 | // larger p-values are forced to zero. 62 | "cutoff_p_value": 1.0, 63 | // Maximum number of terms in the polynomial. If more coefficients are nonzero, 64 | // coefficients with larger p-values will be forced to zero. 65 | "n_max_terms": null, 66 | // Cutoff value for the coefficients of polynomials. Coefficients 67 | // with magnitude less than this value will be regarded as 0. 68 | "coef_tol": 1e-6 69 | } 70 | } 71 | }, 72 | 73 | // Only the key training hyperparameters are listed here. See 74 | // config_common.json for the full list. 75 | "training" : { 76 | "n_samples" : 2000000, 77 | "batch_size" : 1000, 78 | "epsilon" : 0.05, 79 | 80 | // Recommended to set this to as many cores as you can use! Especially if 81 | // using the "const" token. 82 | "n_cores_batch" : 1 83 | }, 84 | 85 | // // Only the key Policy Optimizer hyperparameters are listed here. See 86 | // // config_common.json for the full list. 87 | "policy_optimizer" : { 88 | "learning_rate" : 0.0005, 89 | "entropy_weight" : 0.03, 90 | "entropy_gamma" : 0.7, 91 | 92 | // EXPERIMENTAL: Proximal policy optimization hyperparameters. 93 | // "policy_optimizer_type" : "ppo", 94 | // "ppo_clip_ratio" : 0.2, 95 | // "ppo_n_iters" : 10, 96 | // "ppo_n_mb" : 4, 97 | 98 | // EXPERIMENTAL: Priority queue training hyperparameters. 99 | // "policy_optimizer_type" : "pqt", 100 | // "pqt_k" : 10, 101 | // "pqt_batch_size" : 1, 102 | // "pqt_weight" : 200.0, 103 | // "pqt_use_pg" : false 104 | 105 | }, 106 | 107 | // Hyperparameters related to including in situ priors and constraints. Each 108 | // prior must explicitly be turned "on" or it will not be used. See 109 | // config_common.json for descriptions of each prior. 110 | "prior": { 111 | "length" : { 112 | "min_" : 4, 113 | "max_" : 64, 114 | "on" : true 115 | }, 116 | "repeat" : { 117 | "tokens" : "const", 118 | "min_" : null, 119 | "max_" : 3, 120 | "on" : true 121 | }, 122 | "inverse" : { 123 | "on" : true 124 | }, 125 | "trig" : { 126 | "on" : true 127 | }, 128 | "const" : { 129 | "on" : true 130 | }, 131 | "no_inputs" : { 132 | "on" : true 133 | }, 134 | "uniform_arity" : { 135 | "on" : true 136 | }, 137 | "soft_length" : { 138 | "loc" : 10, 139 | "scale" : 5, 140 | "on" : true 141 | }, 142 | "domain_range" : { 143 | "on" : false 144 | } 145 | } 146 | } 147 | -------------------------------------------------------------------------------- /dso/const.py: -------------------------------------------------------------------------------- 1 | """Constant optimizer used for deep symbolic optimization.""" 2 | 3 | from functools import partial 4 | 5 | import numpy as np 6 | from scipy.optimize import minimize 7 | 8 | 9 | def make_const_optimizer(name, **kwargs): 10 | """Returns a ConstOptimizer given a name and keyword arguments""" 11 | 12 | const_optimizers = { 13 | None : Dummy, 14 | "dummy" : Dummy, 15 | "scipy" : ScipyMinimize, 16 | } 17 | 18 | return const_optimizers[name](**kwargs) 19 | 20 | 21 | class ConstOptimizer(object): 22 | """Base class for constant optimizer""" 23 | 24 | def __init__(self, **kwargs): 25 | self.kwargs = kwargs 26 | 27 | 28 | def __call__(self, f, x0): 29 | """ 30 | Optimizes an objective function from an initial guess. 31 | 32 | The objective function is the negative of the base reward (reward 33 | without penalty) used for training. Optimization excludes any penalties 34 | because they are constant w.r.t. to the constants being optimized. 35 | 36 | Parameters 37 | ---------- 38 | f : function mapping np.ndarray to float 39 | Objective function (negative base reward). 40 | 41 | x0 : np.ndarray 42 | Initial guess for constant placeholders. 43 | 44 | Returns 45 | ------- 46 | x : np.ndarray 47 | Vector of optimized constants. 48 | """ 49 | raise NotImplementedError 50 | 51 | 52 | class Dummy(ConstOptimizer): 53 | """Dummy class that selects the initial guess for each constant""" 54 | 55 | def __init__(self, **kwargs): 56 | super(Dummy, self).__init__(**kwargs) 57 | 58 | 59 | def __call__(self, f, x0): 60 | return x0 61 | 62 | 63 | class ScipyMinimize(ConstOptimizer): 64 | """SciPy's non-linear optimizer""" 65 | 66 | def __init__(self, **kwargs): 67 | super(ScipyMinimize, self).__init__(**kwargs) 68 | 69 | 70 | def __call__(self, f, x0): 71 | with np.errstate(divide='ignore'): 72 | opt_result = partial(minimize, **self.kwargs)(f, x0) 73 | x = opt_result['x'] 74 | return x 75 | -------------------------------------------------------------------------------- /dso/cyfunc.pyx: -------------------------------------------------------------------------------- 1 | ''' 2 | # cython: linetrace=True 3 | # distutils: define_macros=CYTHON_TRACE_NOGIL=1 4 | ''' 5 | # Uncomment the above lines for cProfile 6 | 7 | import numpy as np 8 | import array 9 | 10 | from dso.library import StateChecker, Polynomial 11 | 12 | # Cython specific C imports 13 | cimport numpy as np 14 | from cpython cimport array 15 | cimport cython 16 | from libc.stdlib cimport malloc, free 17 | from cpython.ref cimport PyObject 18 | 19 | # Static inits 20 | cdef list apply_stack = [[None for i in range(25)] for i in range(1024)] 21 | cdef int *stack_count = malloc(1024 * sizeof(int)) 22 | 23 | @cython.boundscheck(False) # turn off bounds-checking for entire function 24 | @cython.wraparound(False) # turn off negative index wrapping for entire function 25 | def execute(np.ndarray X, int len_traversal, list traversal, int[:] is_input_var): 26 | 27 | """Executes the program according to X. 28 | 29 | Parameters 30 | ---------- 31 | X : array-like, shape = [n_samples, n_features] 32 | Training vectors, where n_samples is the number of samples and 33 | n_features is the number of features. 34 | 35 | Returns 36 | ------- 37 | y_hats : array-like, shape = [n_samples] 38 | The result of executing the program on X. 39 | """ 40 | #sp = 0 # allow a dummy first row, requires a none type function with arity of -1 41 | 42 | # Init some ints 43 | cdef int sp = -1 # Stack pointer 44 | cdef int Xs = X.shape[0] 45 | 46 | # Give cdef hints for object types 47 | cdef int i 48 | cdef int n 49 | cdef int arity 50 | cdef np.ndarray intermediate_result 51 | cdef list stack_end 52 | cdef object stack_end_function 53 | 54 | for i in range(len_traversal): 55 | 56 | if not is_input_var[i]: 57 | sp += 1 58 | # Move this to the front with a memset call 59 | stack_count[sp] = 0 60 | # Store the reference to stack_count[sp] rather than keep calling 61 | apply_stack[sp][stack_count[sp]] = traversal[i] 62 | stack_end = apply_stack[sp] 63 | # The first element is the function itself 64 | stack_end_function = stack_end[0] 65 | arity = stack_end_function.arity 66 | else: 67 | # Not a function, so lazily evaluate later 68 | stack_count[sp] += 1 69 | stack_end[stack_count[sp]] = X[:, traversal[i].input_var] 70 | 71 | # Keep on doing this so long as arity matches up, we can 72 | # add in numbers above and complete the arity later. 73 | while stack_count[sp] == arity: 74 | # If stack_end_function is a StateChecker (xi < tj), which is associated with 75 | # the i-th state variable xi and threshold tj, then stack_end_function needs to know 76 | # the value of xi (through set_state_value) in order to evaluate the returned value. 77 | # The index i is specified by the attribute state_index of a StateChecker. 78 | if isinstance(stack_end_function, StateChecker): 79 | stack_end_function.set_state_value(X[:, stack_end_function.state_index]) 80 | if isinstance(stack_end_function, Polynomial): 81 | intermediate_result = stack_end_function(X) 82 | else: 83 | intermediate_result = stack_end_function(*stack_end[1:(stack_count[sp] + 1)]) # 85% of overhead 84 | 85 | # I think we can get rid of this line, but will require a major rewrite. 86 | if sp == 0: 87 | return intermediate_result 88 | 89 | sp -= 1 90 | # Adjust pointer at the end of the stack 91 | stack_end = apply_stack[sp] 92 | stack_count[sp] += 1 93 | stack_end[stack_count[sp]] = intermediate_result 94 | 95 | # The first element is the function itself 96 | stack_end_function = stack_end[0] 97 | arity = stack_end_function.arity 98 | 99 | # We should never get here 100 | assert False, "Function should never get here!" 101 | return None 102 | -------------------------------------------------------------------------------- /dso/execute.py: -------------------------------------------------------------------------------- 1 | try: 2 | from dso import cyfunc 3 | except ImportError as e: 4 | cyfunc = None 5 | import array 6 | 7 | from dso.library import StateChecker, Polynomial 8 | 9 | 10 | def python_execute(traversal, X): 11 | """ 12 | Executes the program according to X using Python. 13 | 14 | Parameters 15 | ---------- 16 | X : array-like, shape = [n_samples, n_features] 17 | Training vectors, where n_samples is the number of samples and 18 | n_features is the number of features. 19 | 20 | Returns 21 | ------- 22 | y_hats : array-like, shape = [n_samples] 23 | The result of executing the program on X. 24 | """ 25 | 26 | apply_stack = [] 27 | 28 | for node in traversal: 29 | apply_stack.append([node]) 30 | 31 | while len(apply_stack[-1]) == apply_stack[-1][0].arity + 1: 32 | token = apply_stack[-1][0] 33 | terminals = apply_stack[-1][1:] 34 | 35 | if token.input_var is not None: 36 | intermediate_result = X[:, token.input_var] 37 | else: 38 | if isinstance(token, StateChecker): 39 | token.set_state_value(X[:, token.state_index]) 40 | if isinstance(token, Polynomial): 41 | intermediate_result = token(X) 42 | else: 43 | intermediate_result = token(*terminals) 44 | if len(apply_stack) != 1: 45 | apply_stack.pop() 46 | apply_stack[-1].append(intermediate_result) 47 | else: 48 | return intermediate_result 49 | 50 | assert False, "Function should never get here!" 51 | return None 52 | 53 | def cython_execute(traversal, X): 54 | """ 55 | Execute cython function using given traversal over input X. 56 | 57 | Parameters 58 | ---------- 59 | 60 | traversal : list 61 | A list of nodes representing the traversal over a Program. 62 | X : np.array 63 | The input values to execute the traversal over. 64 | 65 | Returns 66 | ------- 67 | 68 | result : float 69 | The result of executing the traversal. 70 | """ 71 | if len(traversal) > 1: 72 | is_input_var = array.array('i', [t.input_var is not None for t in traversal]) 73 | return cyfunc.execute(X, len(traversal), traversal, is_input_var) 74 | else: 75 | return python_execute(traversal, X) 76 | 77 | 78 | -------------------------------------------------------------------------------- /dso/policy/__init__.py: -------------------------------------------------------------------------------- 1 | from dso.policy.policy import make_policy, Policy -------------------------------------------------------------------------------- /dso/policy/policy.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | from typing import Tuple, TypeVar 4 | 5 | import tensorflow as tf 6 | import dso 7 | from dso.prior import LengthConstraint 8 | from dso.program import Program 9 | from dso.utils import import_custom_source 10 | from dso.prior import JointPrior 11 | from dso.tf_state_manager import StateManager 12 | from dso.memory import Batch 13 | 14 | # Used for function annotations using the type system 15 | actions = tf.TensorArray 16 | obs = tf.TensorArray 17 | priors = tf.TensorArray 18 | neglogp = tf.TensorArray 19 | entropy = tf.TensorArray 20 | 21 | def make_policy(sess, prior, state_manager, policy_type, **config_policy): 22 | """Factory function for Policy object.""" 23 | 24 | if policy_type == "rnn": 25 | from dso.policy.rnn_policy import RNNPolicy 26 | policy_class = RNNPolicy 27 | else: 28 | # Custom policy import 29 | policy_class = import_custom_source(policy_type) 30 | assert issubclass(policy_class, Policy), \ 31 | "Custom policy {} must subclass dso.policy.Policy.".format(policy_class) 32 | 33 | policy = policy_class(sess, 34 | prior, 35 | state_manager, 36 | **config_policy) 37 | 38 | return policy 39 | 40 | class Policy(ABC): 41 | """Abstract class for a policy. A policy is a parametrized probability 42 | distribution over discrete objects. DSO algorithms optimize the parameters 43 | of this distribution to generate discrete objects with high rewards. 44 | """ 45 | 46 | def __init__(self, 47 | sess : tf.Session, 48 | prior : JointPrior, 49 | state_manager : StateManager, 50 | debug : int = 0, 51 | max_length : int = 30) -> None: 52 | '''Parameters 53 | ---------- 54 | sess : tf.Session 55 | TenorFlow Session object. 56 | 57 | prior : dso.prior.JointPrior 58 | JointPrior object used to adjust probabilities during sampling. 59 | 60 | state_manager: dso.tf_state_manager.StateManager 61 | Object that handles the state features to be used 62 | 63 | debug : int 64 | Debug level, also used in learn(). 0: No debug. 1: Print shapes and 65 | number of parameters for each variable. 66 | 67 | max_length : int or None 68 | Maximum sequence length. This will be overridden if a LengthConstraint 69 | with a maximum length is part of the prior. 70 | ''' 71 | self.sess = sess 72 | self.prior = prior 73 | self.state_manager = state_manager 74 | self.debug = debug 75 | 76 | # Set self.max_length depending on the Prior 77 | self._set_max_length(max_length) 78 | 79 | # Samples produced during attempt to get novel samples. 80 | # Will be combined with checkpoint-loaded samples for next training step 81 | self.extended_batch = None 82 | self.valid_extended_batch = False 83 | 84 | def _set_max_length(self, max_length : int) -> None: 85 | """Set the max legnth depending on the Prior 86 | """ 87 | # Find max_length from the LengthConstraint prior, if it exists 88 | # For binding task, max_length is # of allowed mutations or master-seq length 89 | # Both priors will never happen in the same experiment 90 | prior_max_length = None 91 | for single_prior in self.prior.priors: 92 | if isinstance(single_prior, LengthConstraint): 93 | if single_prior.max is not None: 94 | prior_max_length = single_prior.max 95 | self.max_length = prior_max_length 96 | break 97 | 98 | if prior_max_length is None: 99 | assert max_length is not None, "max_length must be specified if "\ 100 | "there is no LengthConstraint." 101 | self.max_length = max_length 102 | print("WARNING: Maximum length not constrained. Sequences will " 103 | "stop at {} and complete by repeating the first input " 104 | "variable.".format(self.max_length)) 105 | elif max_length is not None and max_length != self.max_length: 106 | print("WARNING: max_length ({}) will be overridden by value from " 107 | "LengthConstraint ({}).".format(max_length, self.max_length)) 108 | 109 | @abstractmethod 110 | def _setup_tf_model(self, **kwargs) -> None: 111 | """"Setup the TensorFlow graph(s). 112 | 113 | Returns 114 | ------- 115 | None 116 | """ 117 | raise NotImplementedError 118 | 119 | @abstractmethod 120 | def make_neglogp_and_entropy(self, 121 | B : Batch, 122 | entropy_gamma : float 123 | ) -> Tuple[neglogp, entropy]: 124 | """Computes the negative log-probabilities for a given 125 | batch of actions, observations and priors 126 | under the current policy. 127 | 128 | Returns 129 | ------- 130 | neglogp, entropy : 131 | Tensorflow tensors 132 | """ 133 | raise NotImplementedError 134 | 135 | @abstractmethod 136 | def sample(self, n : int) -> Tuple[actions, obs, priors]: 137 | """Sample batch of n expressions. 138 | 139 | Returns 140 | ------- 141 | actions, obs, priors : 142 | Or a batch 143 | """ 144 | raise NotImplementedError 145 | 146 | @abstractmethod 147 | def compute_probs(self, memory_batch, log=False): 148 | """Compute the probabilities of a Batch. 149 | 150 | Returns 151 | ------- 152 | probs : 153 | Or a batch 154 | """ 155 | raise NotImplementedError 156 | 157 | 158 | -------------------------------------------------------------------------------- /dso/policy_optimizer/__init__.py: -------------------------------------------------------------------------------- 1 | from dso.policy_optimizer.policy_optimizer import make_policy_optimizer, PolicyOptimizer -------------------------------------------------------------------------------- /dso/policy_optimizer/pg_policy_optimizer.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | from dso.policy_optimizer import PolicyOptimizer 4 | from dso.policy import Policy 5 | 6 | class PGPolicyOptimizer(PolicyOptimizer): 7 | """Vanilla policy gradient policy optimizer. 8 | 9 | Parameters 10 | ---------- 11 | cell : str 12 | Recurrent cell to use. Supports 'lstm' and 'gru'. 13 | 14 | num_layers : int 15 | Number of RNN layers. 16 | 17 | num_units : int or list of ints 18 | Number of RNN cell units in each of the RNN's layers. If int, the value 19 | is repeated for each layer. 20 | 21 | initiailizer : str 22 | Initializer for the recurrent cell. Supports 'zeros' and 'var_scale'. 23 | 24 | """ 25 | def __init__(self, 26 | sess : tf.Session, 27 | policy : Policy, 28 | debug : int = 0, 29 | summary : bool = False, 30 | # Optimizer hyperparameters 31 | optimizer : str = 'adam', 32 | learning_rate : float = 0.001, 33 | # Loss hyperparameters 34 | entropy_weight : float = 0.005, 35 | entropy_gamma : float = 1.0) -> None: 36 | super()._setup_policy_optimizer(sess, policy, debug, summary, optimizer, learning_rate, entropy_weight, entropy_gamma) 37 | 38 | 39 | def _set_loss(self): 40 | with tf.name_scope("losses"): 41 | # Retrieve rewards from batch 42 | r = self.sampled_batch_ph.rewards 43 | # Baseline is the worst of the current samples r 44 | self.pg_loss = tf.reduce_mean((r - self.baseline) * self.neglogp, name="pg_loss") 45 | # Loss already is set to entropy loss 46 | self.loss += self.pg_loss 47 | 48 | 49 | def _preppend_to_summary(self): 50 | with tf.name_scope("summary"): 51 | tf.summary.scalar("pg_loss", self.pg_loss) 52 | 53 | 54 | def train_step(self, baseline, sampled_batch): 55 | """Computes loss, trains model, and returns summaries.""" 56 | feed_dict = { 57 | self.baseline : baseline, 58 | self.sampled_batch_ph : sampled_batch 59 | } 60 | 61 | summaries, _ = self.sess.run([self.summaries, self.train_op], feed_dict=feed_dict) 62 | 63 | return summaries -------------------------------------------------------------------------------- /dso/policy_optimizer/ppo_policy_optimizer.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | import numpy as np 4 | from dso.policy_optimizer import PolicyOptimizer 5 | from dso.policy import Policy 6 | from dso.memory import Batch 7 | 8 | class PPOPolicyOptimizer(PolicyOptimizer): 9 | """Proximal policy optimization policy optimizer. 10 | 11 | Parameters 12 | ---------- 13 | 14 | ppo_clip_ratio : float 15 | Clip ratio to use for PPO. 16 | 17 | ppo_n_iters : int 18 | Number of optimization iterations for PPO. 19 | 20 | ppo_n_mb : int 21 | Number of minibatches per optimization iteration for PPO. 22 | 23 | """ 24 | def __init__(self, 25 | sess : tf.Session, 26 | policy : Policy, 27 | debug : int = 0, 28 | summary : bool = False, 29 | # Optimizer hyperparameters 30 | optimizer : str = 'adam', 31 | learning_rate : float = 0.001, 32 | # Loss hyperparameters 33 | entropy_weight : float = 0.005, 34 | entropy_gamma : float = 1.0, 35 | # PPO hyperparameters 36 | ppo_clip_ratio : float = 0.2, 37 | ppo_n_iters : int = 10, 38 | ppo_n_mb : int = 4) -> None: 39 | self.ppo_clip_ratio = ppo_clip_ratio 40 | self.ppo_n_iters = ppo_n_iters 41 | self.ppo_n_mb = ppo_n_mb 42 | self.rng = np.random.RandomState(0) # Used for PPO minibatch sampling 43 | super()._setup_policy_optimizer(sess, policy, debug, summary, optimizer, learning_rate, entropy_weight, entropy_gamma) 44 | 45 | 46 | def _set_loss(self): 47 | with tf.name_scope("losses"): 48 | # Retrieve rewards from batch 49 | r = self.sampled_batch_ph.rewards 50 | 51 | self.old_neglogp_ph = tf.placeholder(dtype=tf.float32, 52 | shape=(None,), name="old_neglogp") 53 | ratio = tf.exp(self.old_neglogp_ph - self.neglogp) 54 | clipped_ratio = tf.clip_by_value(ratio, 1. - self.ppo_clip_ratio, 55 | 1. + self.ppo_clip_ratio) 56 | ppo_loss = -tf.reduce_mean((r - self.baseline) * 57 | tf.minimum(ratio, clipped_ratio)) 58 | # Loss already is set to entropy loss 59 | self.loss += ppo_loss 60 | 61 | # Define PPO diagnostics 62 | clipped = tf.logical_or(ratio < (1. - self.ppo_clip_ratio), 63 | ratio > 1. + self.ppo_clip_ratio) 64 | self.clip_fraction = tf.reduce_mean(tf.cast(clipped, tf.float32)) 65 | self.sample_kl = tf.reduce_mean(self.neglogp - self.old_neglogp_ph) 66 | 67 | 68 | def _preppend_to_summary(self): 69 | with tf.name_scope("summary"): 70 | tf.summary.scalar("ppo_loss", self.ppo_loss) 71 | 72 | 73 | def train_step(self, baseline, sampled_batch): 74 | feed_dict = { 75 | self.baseline : baseline, 76 | self.sampled_batch_ph : sampled_batch 77 | } 78 | n_samples = sampled_batch.rewards.shape[0] 79 | 80 | 81 | # Compute old_neglogp to be used for training 82 | old_neglogp = self.sess.run(self.neglogp, feed_dict=feed_dict) 83 | 84 | # Perform multiple steps of minibatch training 85 | # feed_dict[self.old_neglogp_ph] = old_neglogp 86 | indices = np.arange(n_samples) 87 | for ppo_iter in range(self.ppo_n_iters): 88 | self.rng.shuffle(indices) # in-place 89 | # list of [ppo_n_mb] arrays 90 | minibatches = np.array_split(indices, self.ppo_n_mb) 91 | for i, mb in enumerate(minibatches): 92 | sampled_batch_mb = Batch( 93 | **{name: array[mb] for name, array 94 | in sampled_batch._asdict().items()}) 95 | mb_feed_dict = { 96 | self.baseline: baseline, 97 | self.batch_size: len(mb), 98 | self.old_neglogp_ph: old_neglogp[mb], 99 | self.sampled_batch_ph: sampled_batch_mb 100 | } 101 | 102 | summaries, _ = self.sess.run([self.summaries, self.train_op], 103 | feed_dict=mb_feed_dict) 104 | 105 | # Diagnostics 106 | # kl, cf, _ = self.sess.run( 107 | # [self.sample_kl, self.clip_fraction, self.train_op], 108 | # feed_dict=mb_feed_dict) 109 | # print("ppo_iter", ppo_iter, "i", i, "KL", kl, "CF", cf) 110 | 111 | return summaries 112 | 113 | 114 | 115 | 116 | -------------------------------------------------------------------------------- /dso/policy_optimizer/pqt_policy_optimizer.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | from dso.policy_optimizer import PolicyOptimizer 4 | from dso.policy import Policy 5 | from dso.utils import make_batch_ph 6 | 7 | class PQTPolicyOptimizer(PolicyOptimizer): 8 | """Priority Queue Ttraining policy gradient policy optimizer. 9 | 10 | Parameters 11 | ---------- 12 | pqt_k : int 13 | Size of priority queue. 14 | 15 | pqt_batch_size : int 16 | Size of batch to sample (with replacement) from priority queue. 17 | 18 | pqt_weight : float 19 | Coefficient for PQT loss function. 20 | 21 | pqt_use_pg : bool 22 | Use policy gradient loss when using PQT? 23 | 24 | """ 25 | def __init__(self, 26 | sess : tf.Session, 27 | policy : Policy, 28 | debug : int = 0, 29 | summary : bool = False, 30 | # Optimizer hyperparameters 31 | optimizer : str = 'adam', 32 | learning_rate : float = 0.001, 33 | # Loss hyperparameters 34 | entropy_weight : float = 0.005, 35 | entropy_gamma : float = 1.0, 36 | # PQT hyperparameters 37 | pqt_k : int = 10, 38 | pqt_batch_size : int = 1, 39 | pqt_weight : float = 200.0, 40 | pqt_use_pg: bool = False) -> None: 41 | self.pqt_k = pqt_k 42 | self.pqt_batch_size = pqt_batch_size 43 | self.pqt_weight = pqt_weight 44 | self. pqt_use_pg = pqt_use_pg 45 | super()._setup_policy_optimizer(sess, policy, debug, summary, optimizer, learning_rate, entropy_weight, entropy_gamma) 46 | 47 | 48 | def _set_loss(self): 49 | with tf.name_scope("losses"): 50 | # Create placeholder for PQT batch 51 | self.pqt_batch_ph = make_batch_ph("pqt_batch", self.n_choices) #self.n_choices is defined in parent class 52 | 53 | pqt_neglogp, _ = self.policy.make_neglogp_and_entropy(self.pqt_batch_ph, self.entropy_gamma) 54 | self.pqt_loss = self.pqt_weight * tf.reduce_mean(pqt_neglogp, name="pqt_loss") 55 | 56 | # Loss already is set to entropy loss 57 | self.loss += self.pqt_loss 58 | 59 | 60 | def _preppend_to_summary(self): 61 | with tf.name_scope("summary"): 62 | tf.summary.scalar("pqt_loss", self.pqt_loss) 63 | 64 | 65 | def train_step(self, baseline, sampled_batch, pqt_batch): 66 | feed_dict = { 67 | self.baseline : baseline, 68 | self.sampled_batch_ph : sampled_batch 69 | } 70 | feed_dict.update({ 71 | self.pqt_batch_ph : pqt_batch 72 | }) 73 | 74 | summaries, _ = self.sess.run([self.summaries, self.train_op], feed_dict=feed_dict) 75 | 76 | return summaries -------------------------------------------------------------------------------- /dso/run.py: -------------------------------------------------------------------------------- 1 | """Parallelized, single-point launch script to run DSO on a set of benchmarks.""" 2 | 3 | import os 4 | import sys 5 | import time 6 | import multiprocessing 7 | from copy import deepcopy 8 | from datetime import datetime 9 | 10 | import click 11 | 12 | from dso import DeepSymbolicOptimizer 13 | from dso.logeval import LogEval 14 | from dso.config import load_config 15 | from dso.utils import safe_update_summary 16 | 17 | 18 | def train_dso(config): 19 | """Trains DSO and returns dict of reward, expression, and traversal""" 20 | 21 | print("\n== TRAINING SEED {} START ============".format(config["experiment"]["seed"])) 22 | 23 | # For some reason, for the control task, the environment needs to be instantiated 24 | # before creating the pool. Otherwise, gym.make() hangs during the pool initializer 25 | if config["task"]["task_type"] == "control" and config["training"]["n_cores_batch"] > 1: 26 | import gym 27 | import dso.task.control # Registers custom and third-party environments 28 | gym.make(config["task"]["env"]) 29 | 30 | # Train the model 31 | model = DeepSymbolicOptimizer(deepcopy(config)) 32 | start = time.time() 33 | result = model.train() 34 | result["t"] = time.time() - start 35 | result.pop("program") 36 | 37 | save_path = model.config_experiment["save_path"] 38 | summary_path = os.path.join(save_path, "summary.csv") 39 | 40 | print("== TRAINING SEED {} END ==============".format(config["experiment"]["seed"])) 41 | 42 | return result, summary_path 43 | 44 | 45 | def print_summary(config, runs, messages): 46 | text = '\n== EXPERIMENT SETUP START ===========\n' 47 | text += 'Task type : {}\n'.format(config["task"]["task_type"]) 48 | if config["task"]["task_type"] == "regression": 49 | text += 'Dataset : {}\n'.format(config["task"]["dataset"]) 50 | elif config["task"]["task_type"] == "control": 51 | text += 'Environment : {}\n'.format(config["task"]["env"]) 52 | text += 'Starting seed : {}\n'.format(config["experiment"]["seed"]) 53 | text += 'Runs : {}\n'.format(runs) 54 | if len(messages) > 0: 55 | text += 'Additional context :\n' 56 | for message in messages: 57 | text += " {}\n".format(message) 58 | text += '== EXPERIMENT SETUP END =============' 59 | print(text) 60 | 61 | 62 | @click.command() 63 | @click.argument('config_template', default="") 64 | @click.option('--runs', '--r', default=1, type=int, help="Number of independent runs with different seeds") 65 | @click.option('--n_cores_task', '--n', default=1, help="Number of cores to spread out across tasks") 66 | @click.option('--seed', '--s', default=None, type=int, help="Starting seed (overwrites seed in config), incremented for each independent run") 67 | @click.option('--benchmark', '--b', default=None, type=str, help="Name of benchmark") 68 | @click.option('--exp_name', default=None, type=str, help="Name of experiment to manually generate log path") 69 | def main(config_template, runs, n_cores_task, seed, benchmark, exp_name): 70 | """Runs DSO in parallel across multiple seeds using multiprocessing.""" 71 | 72 | messages = [] 73 | 74 | # Load the experiment config 75 | config_template = config_template if config_template != "" else None 76 | config = load_config(config_template) 77 | 78 | # Overwrite named benchmark (for tasks that support them) 79 | task_type = config["task"]["task_type"] 80 | if benchmark is not None: 81 | # For regression, --b overwrites config["task"]["dataset"] 82 | if task_type == "regression": 83 | config["task"]["dataset"] = benchmark 84 | # For control, --b overwrites config["task"]["env"] 85 | elif task_type == "control": 86 | config["task"]["env"] = benchmark 87 | else: 88 | raise ValueError("--b is not supported for task {}.".format(task_type)) 89 | 90 | # Update save dir if provided 91 | if exp_name is not None: 92 | config["experiment"]["exp_name"] = exp_name 93 | 94 | # Overwrite config seed, if specified 95 | if seed is not None: 96 | if config["experiment"]["seed"] is not None: 97 | messages.append( 98 | "INFO: Replacing config seed {} with command-line seed {}.".format( 99 | config["experiment"]["seed"], seed)) 100 | config["experiment"]["seed"] = seed 101 | 102 | # Save starting seed and run command 103 | config["experiment"]["starting_seed"] = config["experiment"]["seed"] 104 | config["experiment"]["cmd"] = " ".join(sys.argv) 105 | 106 | # Set timestamp once to be used by all workers 107 | timestamp = datetime.now().strftime("%Y-%m-%d-%H%M%S") 108 | config["experiment"]["timestamp"] = timestamp 109 | 110 | # Fix incompatible configurations 111 | if n_cores_task == -1: 112 | n_cores_task = multiprocessing.cpu_count() 113 | if n_cores_task > runs: 114 | messages.append( 115 | "INFO: Setting 'n_cores_task' to {} because there are only {} runs.".format( 116 | runs, runs)) 117 | n_cores_task = runs 118 | if config["training"]["verbose"] and n_cores_task > 1: 119 | messages.append( 120 | "INFO: Setting 'verbose' to False for parallelized run.") 121 | config["training"]["verbose"] = False 122 | if config["training"]["n_cores_batch"] != 1 and n_cores_task > 1: 123 | messages.append( 124 | "INFO: Setting 'n_cores_batch' to 1 to avoid nested child processes.") 125 | config["training"]["n_cores_batch"] = 1 126 | if config["gp_meld"]["run_gp_meld"] and n_cores_task > 1 and runs > 1: 127 | messages.append( 128 | "INFO: Setting 'parallel_eval' to 'False' as we are already parallelizing.") 129 | config["gp_meld"]["parallel_eval"] = False 130 | 131 | 132 | # Start training 133 | print_summary(config, runs, messages) 134 | 135 | # Generate configs (with incremented seeds) for each run 136 | configs = [deepcopy(config) for _ in range(runs)] 137 | for i, config in enumerate(configs): 138 | config["experiment"]["seed"] += i 139 | 140 | # Farm out the work 141 | if n_cores_task > 1: 142 | pool = multiprocessing.Pool(n_cores_task) 143 | for i, (result, summary_path) in enumerate(pool.imap_unordered(train_dso, configs)): 144 | if not safe_update_summary(summary_path, result): 145 | print("Warning: Could not update summary stats at {}".format(summary_path)) 146 | print("INFO: Completed run {} of {} in {:.0f} s".format(i + 1, runs, result["t"])) 147 | else: 148 | for i, config in enumerate(configs): 149 | result, summary_path = train_dso(config) 150 | if not safe_update_summary(summary_path, result): 151 | print("Warning: Could not update summary stats at {}".format(summary_path)) 152 | print("INFO: Completed run {} of {} in {:.0f} s".format(i + 1, runs, result["t"])) 153 | 154 | # Evaluate the log files 155 | print("\n== POST-PROCESS START =================") 156 | log = LogEval(config_path=os.path.dirname(summary_path)) 157 | log.analyze_log( 158 | show_count=config["postprocess"]["show_count"], 159 | show_hof=config["logging"]["hof"] is not None and config["logging"]["hof"] > 0, 160 | show_pf=config["logging"]["save_pareto_front"], 161 | save_plots=config["postprocess"]["save_plots"]) 162 | print("== POST-PROCESS END ===================") 163 | 164 | 165 | if __name__ == "__main__": 166 | main() 167 | -------------------------------------------------------------------------------- /dso/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | import os 3 | from setuptools import dist 4 | 5 | dist.Distribution().fetch_build_eggs(['Cython', 'numpy']) 6 | 7 | import numpy 8 | from Cython.Build import cythonize 9 | 10 | required = [ 11 | "pytest", 12 | "cython", 13 | "numpy<=1.19", 14 | "tensorflow==1.14", 15 | "numba==0.53.1", 16 | "sympy", 17 | "pandas", 18 | "scikit-learn", 19 | "click", 20 | "deap", 21 | "pathos", 22 | "seaborn", 23 | "progress", 24 | "tqdm", 25 | "commentjson", 26 | "PyYAML", 27 | "prettytable" 28 | ] 29 | 30 | extras = { 31 | "control": [ 32 | "mpi4py", 33 | "gym[box2d]==0.15.4", 34 | "pybullet", 35 | "stable-baselines[mpi]==2.10.0" 36 | ], 37 | "regression": [] 38 | } 39 | extras['all'] = list(set([item for group in extras.values() for item in group])) 40 | 41 | setup( name='dso', 42 | version='1.0dev', 43 | description='Deep symbolic optimization.', 44 | author='LLNL', 45 | packages=['dso'], 46 | setup_requires=["numpy", "Cython"], 47 | ext_modules=cythonize([os.path.join('dso','cyfunc.pyx')]), 48 | include_dirs=[numpy.get_include()], 49 | install_requires=required, 50 | extras_require=extras 51 | ) 52 | -------------------------------------------------------------------------------- /dso/task/__init__.py: -------------------------------------------------------------------------------- 1 | from dso.task.task import make_task, set_task, Task, HierarchicalTask, SequentialTask 2 | -------------------------------------------------------------------------------- /dso/task/regression/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RL-MLDM/alphagen/1c187544df587c6109e25948bf0d16e42fce24bf/dso/task/regression/__init__.py -------------------------------------------------------------------------------- /dso/task/regression/function_sets.csv: -------------------------------------------------------------------------------- 1 | name,function_set 2 | Koza,"add,sub,mul,div,sin,cos,exp,log" 3 | CKoza,"add,sub,mul,div,sin,cos,exp,log,const" 4 | KozaPlus1,"add,sub,mul,div,sin,cos,exp,log,1.0" 5 | Korns,"add,sub,mul,div,sin,cos,exp,log,n2,n3,sqrt,tan,tanh,const" 6 | Keijzer,"add,mul,inv,neg,sqrt,const" 7 | KeijzerPlus1,"add,mul,inv,neg,sqrt,1.0,const" 8 | Vladislavleva-A,"add,sub,mul,div,n2" 9 | Vladislavleva-B,"add,sub,mul,div,n2,exp,expneg" 10 | Vladislavleva-C,"add,sub,mul,div,n2,exp,expneg,sin,cos" 11 | Real,"add,sub,mul,div,sin,cos,exp,log" 12 | None,"add,sub,mul,div,sin,cos,exp,log" 13 | Jin,"add,sub,mul,div,sin,cos,exp,n2,n3,const" 14 | GrammarVAE,"add,mul,div,sin,exp,1.0,2.0,3.0" 15 | KozaUserConst,"add,sub,mul,div,sin,cos,exp,log,3.14159265358979323846,2.71828182845904523536" 16 | PKoza,"add,sub,mul,div,sin,cos,exp,log,poly" 17 | CPKoza,"add,sub,mul,div,sin,cos,exp,log,const,poly" 18 | PKozaPlusSqrt,"add,sub,mul,div,sin,cos,exp,log,sqrt,poly" 19 | Base,"add,sub,mul,div,sqrt,n2,log,exp,sin,cos,1" 20 | BaseC,"add,sub,mul,div,sqrt,n2,log,exp,sin,cos,1,const" 21 | BaseP,"add,sub,mul,div,sqrt,n2,log,exp,sin,cos,1,poly" 22 | BaseCP,"add,sub,mul,div,sqrt,n2,log,exp,sin,cos,1,const,poly" -------------------------------------------------------------------------------- /dso/task/regression/mat_mult_benchmark.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import scipy 3 | from scipy import linalg 4 | import time 5 | 6 | import sys 7 | args = sys.argv 8 | if len(args) == 4: 9 | nparams = int(args[1]) 10 | nsamples = int(args[2]) 11 | ntries = int(args[3]) 12 | else: 13 | nparams = 100 14 | nsamples = 1000 15 | ntries = 100 16 | 17 | def mat_mult_benchmark(nparams, nsamples): 18 | # create data 19 | X = 100.0 * np.random.rand(nsamples, nparams) - 50.0 20 | X[:,0] = 1.0 21 | X_pinv = scipy.linalg.pinv(X) 22 | X_pinv_32 = X_pinv.astype(np.float32, copy=True) 23 | y = 100.0 * np.random.rand(nsamples) - 50.0 24 | # do least squares 25 | tls = time.time() 26 | beta = scipy.linalg.lstsq(X, y) 27 | tls = time.time() - tls 28 | # use pinv with dot 29 | tdot = time.time() 30 | beta = np.dot(X_pinv, y) 31 | tdot = time.time() - tdot 32 | # use pinv with matmul 33 | tmatmul = time.time() 34 | beta = np.matmul(X_pinv, y) 35 | tmatmul = time.time() - tmatmul 36 | # use pinv with matmul and float32 37 | tmatmul32 = time.time() 38 | y_32 = y.astype(np.float32, copy=True) 39 | beta = np.matmul(X_pinv_32, y_32) 40 | tmatmul32 = time.time() - tmatmul32 41 | # print results 42 | print("pinv-dot: ", tls/tdot, "x; pinv-matmul: ", tls/tmatmul, 43 | "x; pinv-matmul32:", tls/tmatmul32, "x") 44 | 45 | for i in range(ntries): 46 | mat_mult_benchmark(nparams, nsamples) 47 | -------------------------------------------------------------------------------- /dso/task/regression/sklearn.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | 3 | from sklearn.base import BaseEstimator, RegressorMixin 4 | from sklearn.utils.validation import check_is_fitted 5 | 6 | from dso import DeepSymbolicOptimizer 7 | 8 | 9 | class DeepSymbolicRegressor(DeepSymbolicOptimizer, 10 | BaseEstimator, RegressorMixin): 11 | """ 12 | Sklearn interface for deep symbolic regression. 13 | """ 14 | 15 | def __init__(self, config=None): 16 | if config is None: 17 | config = { 18 | "task" : {"task_type" : "regression"} 19 | } 20 | DeepSymbolicOptimizer.__init__(self, config) 21 | 22 | def fit(self, X, y): 23 | 24 | # Update the Task 25 | config = deepcopy(self.config) 26 | config["task"]["dataset"] = (X, y) 27 | 28 | # Turn off file saving 29 | config["experiment"]["logdir"] = None 30 | 31 | # TBD: Add support for gp-meld and sklearn interface. Currently, gp-meld 32 | # relies on BenchmarkDataset objects, not (X, y) data. 33 | if config["gp_meld"].get("run_gp_meld"): 34 | print("WARNING: GP-meld not yet supported for sklearn interface.") 35 | config["gp_meld"]["run_gp_meld"] = False 36 | 37 | self.set_config(config) 38 | 39 | train_result = self.train() 40 | self.program_ = train_result["program"] 41 | 42 | return self 43 | 44 | def predict(self, X): 45 | 46 | check_is_fitted(self, "program_") 47 | 48 | return self.program_.execute(X) 49 | -------------------------------------------------------------------------------- /dso/task/regression/test_sklearn.py: -------------------------------------------------------------------------------- 1 | """Tests for sklearn interface.""" 2 | 3 | import pytest 4 | import numpy as np 5 | 6 | from dso import DeepSymbolicRegressor 7 | from dso.test.generate_test_data import CONFIG_TRAINING_OVERRIDE 8 | 9 | 10 | @pytest.fixture 11 | def model(): 12 | return DeepSymbolicRegressor() 13 | 14 | 15 | def test_task(model): 16 | """Test regression for various configs.""" 17 | 18 | # Generate some data 19 | np.random.seed(0) 20 | X = np.random.random(size=(10, 3)) 21 | y = np.random.random(size=(10,)) 22 | 23 | model.config_training.update(CONFIG_TRAINING_OVERRIDE) 24 | model.fit(X, y) 25 | -------------------------------------------------------------------------------- /dso/task/task.py: -------------------------------------------------------------------------------- 1 | """Factory functions for generating symbolic search tasks.""" 2 | 3 | from abc import ABC, abstractmethod 4 | import numpy as np 5 | 6 | from dso.program import Program 7 | from dso.utils import import_custom_source 8 | from dso.subroutines import parents_siblings 9 | 10 | 11 | class Task(ABC): 12 | """ 13 | Object specifying a symbolic search task. 14 | 15 | Attributes 16 | ---------- 17 | library : Library 18 | Library of Tokens. 19 | 20 | stochastic : bool 21 | Whether the reward function of the task is stochastic. 22 | 23 | task_type : str 24 | Task type: regression, control, or binding. 25 | 26 | name : str 27 | Unique name for instance of this task. 28 | """ 29 | 30 | task_type = None 31 | 32 | @abstractmethod 33 | def reward_function(self, program, optimizing=False): 34 | """ 35 | The reward function for this task. 36 | 37 | Parameters 38 | ---------- 39 | program : dso.program.Program 40 | 41 | The Program to compute reward of. 42 | 43 | optimizing : bool 44 | 45 | Whether the reward is computed for PlaceholderConstant optimization. 46 | 47 | Returns 48 | ------- 49 | reward : float 50 | 51 | Fitness/reward of the program. 52 | """ 53 | raise NotImplementedError 54 | 55 | @abstractmethod 56 | def evaluate(self, program): 57 | """ 58 | The evaluation metric for this task. 59 | 60 | Parameters 61 | ---------- 62 | program : dso.program.Program 63 | 64 | The Program to evaluate. 65 | 66 | Returns 67 | ------- 68 | 69 | info : dict 70 | 71 | Dictionary of evaluation metrics. Special key "success" is used to 72 | trigger early stopping. 73 | """ 74 | raise NotImplementedError 75 | 76 | @abstractmethod 77 | def get_next_obs(self, actions, obs, already_finished): 78 | """ 79 | Produce the next observation and prior from the current observation and 80 | list of actions so far. Observations must be 1-D np.float32 vectors. 81 | 82 | Parameters 83 | ---------- 84 | 85 | actions : np.ndarray (dtype=np.int32) 86 | Actions selected so far, shape (batch_size, current_length) 87 | 88 | obs : np.ndarray (dtype=np.float32) 89 | Previous observation, shape (batch_size, OBS_DIM). 90 | 91 | already_finished : np.ndarray (dtype=bool) 92 | Whether the object has *already* been completed. 93 | 94 | Returns 95 | ------- 96 | 97 | next_obs : np.ndarray (dtype=np.float32) 98 | The next observation, shape (batch_size, OBS_DIM). 99 | 100 | prior : np.ndarray (dtype=np.float32) 101 | Prior for selecting the next token, shape (batch_size, 102 | self.library.L). 103 | 104 | finished : np.ndarray (dtype=bool) 105 | Whether the object has *ever* been completed. 106 | """ 107 | pass 108 | 109 | @abstractmethod 110 | def reset_task(self): 111 | """ 112 | Create the starting observation. 113 | 114 | Returns 115 | ------- 116 | 117 | obs : np.ndarray (dtype=np.float32) 118 | Starting observation, shape (batch_size, OBS_DIM). 119 | """ 120 | pass 121 | 122 | 123 | class HierarchicalTask(Task): 124 | """ 125 | A Task in which the search space is a binary tree. Observations include 126 | the previous action, the parent, the sibling, and/or the number of dangling 127 | (unselected) nodes. 128 | """ 129 | 130 | OBS_DIM = 4 # action, parent, sibling, dangling 131 | 132 | def __init__(self): 133 | super(Task).__init__() 134 | 135 | def get_next_obs(self, actions, obs, already_finished): 136 | 137 | dangling = obs[:, 3] # Shape of obs: (?, 4) 138 | action = actions[:, -1] # Current action 139 | lib = self.library 140 | 141 | # Compute parents and siblings 142 | parent, sibling = parents_siblings(actions, 143 | arities=lib.arities, 144 | parent_adjust=lib.parent_adjust, 145 | empty_parent=lib.EMPTY_PARENT, 146 | empty_sibling=lib.EMPTY_SIBLING) 147 | 148 | # Compute dangling 149 | dangling += lib.arities[action] - 1 150 | 151 | # Compute finished 152 | just_finished = (dangling == 0) # Trees that completed _this_ time step 153 | # [batch_size] 154 | finished = np.logical_or(just_finished, 155 | already_finished) 156 | 157 | # Compute priors 158 | prior = self.prior(actions, parent, sibling, dangling, finished) # (?, n_choices) 159 | 160 | # Combine observation dimensions 161 | next_obs = np.stack([action, parent, sibling, dangling], axis=1) # (?, 4) 162 | next_obs = next_obs.astype(np.float32) 163 | 164 | return next_obs, prior, finished 165 | 166 | def reset_task(self, prior): 167 | """ 168 | Returns the initial observation: empty action, parent, and sibling, and 169 | dangling is 1. 170 | """ 171 | 172 | self.prior = prior 173 | 174 | # Order of observations: action, parent, sibling, dangling 175 | initial_obs = np.array([self.library.EMPTY_ACTION, 176 | self.library.EMPTY_PARENT, 177 | self.library.EMPTY_SIBLING, 178 | 1], 179 | dtype=np.float32) 180 | return initial_obs 181 | 182 | 183 | class SequentialTask(Task): 184 | """ 185 | A Task in which the search space is a (possibly variable-length) sequence. 186 | The observation is simply the previous action. 187 | """ 188 | 189 | pass 190 | 191 | 192 | def make_task(task_type, **config_task): 193 | """ 194 | Factory function for Task object. 195 | 196 | Parameters 197 | ---------- 198 | 199 | task_type : str 200 | Type of task: 201 | "regression" : Symbolic regression task. 202 | "control" : Episodic reinforcement learning task. 203 | "binding": AbAg binding affinity optimization task. 204 | 205 | config_task : kwargs 206 | Task-specific arguments. See specifications of task_dict. 207 | 208 | Returns 209 | ------- 210 | 211 | task : Task 212 | Task object. 213 | """ 214 | 215 | # Lazy import of task factory functions 216 | if task_type == 'binding': 217 | from dso.task.binding.binding import BindingTask 218 | task_class = BindingTask 219 | elif task_type == "regression": 220 | from dso.task.regression.regression import RegressionTask 221 | task_class = RegressionTask 222 | elif task_type == "control": 223 | from dso.task.control.control import ControlTask 224 | task_class = ControlTask 225 | else: 226 | # Custom task import 227 | task_class = import_custom_source(task_type) 228 | assert issubclass(task_class, Task), \ 229 | "Custom task {} must subclass dso.task.Task.".format(task_class) 230 | 231 | task = task_class(**config_task) 232 | return task 233 | 234 | 235 | def set_task(config_task): 236 | """Helper function to make set the Program class Task and execute function 237 | from task config.""" 238 | 239 | # Use of protected functions is the same for all tasks, so it's handled separately 240 | protected = config_task["protected"] if "protected" in config_task else False 241 | 242 | Program.set_execute(protected) 243 | task = make_task(**config_task) 244 | Program.set_task(task) 245 | -------------------------------------------------------------------------------- /dso/tf_state_manager.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import tensorflow as tf 4 | 5 | from dso.program import Program 6 | 7 | 8 | class StateManager(ABC): 9 | """ 10 | An interface for handling the tf.Tensor inputs to the Policy. 11 | """ 12 | 13 | def setup_manager(self, policy): 14 | """ 15 | Function called inside the policy to perform the needed initializations (e.g., if the tf context is needed) 16 | :param policy the policy class 17 | """ 18 | self.policy = policy 19 | self.max_length = policy.max_length 20 | 21 | @abstractmethod 22 | def get_tensor_input(self, obs): 23 | """ 24 | Convert an observation from a Task into a Tesnor input for the 25 | Policy, e.g. by performing one-hot encoding or embedding lookup. 26 | 27 | Parameters 28 | ---------- 29 | obs : np.ndarray (dtype=np.float32) 30 | Observation coming from the Task. 31 | 32 | Returns 33 | -------- 34 | input_ : tf.Tensor (dtype=tf.float32) 35 | Tensor to be used as input to the Policy. 36 | """ 37 | return 38 | 39 | def process_state(self, obs): 40 | """ 41 | Entry point for adding information to the state tuple. 42 | If not overwritten, this functions does nothing 43 | """ 44 | return obs 45 | 46 | 47 | def make_state_manager(config): 48 | """ 49 | Parameters 50 | ---------- 51 | config : dict 52 | Parameters for this StateManager. 53 | 54 | Returns 55 | ------- 56 | state_manager : StateManager 57 | The StateManager to be used by the policy. 58 | """ 59 | manager_dict = { 60 | "hierarchical": HierarchicalStateManager 61 | } 62 | 63 | if config is None: 64 | config = {} 65 | 66 | # Use HierarchicalStateManager by default 67 | manager_type = config.pop("type", "hierarchical") 68 | 69 | manager_class = manager_dict[manager_type] 70 | state_manager = manager_class(**config) 71 | 72 | return state_manager 73 | 74 | 75 | class HierarchicalStateManager(StateManager): 76 | """ 77 | Class that uses the previous action, parent, sibling, and/or dangling as 78 | observations. 79 | """ 80 | 81 | def __init__(self, observe_parent=True, observe_sibling=True, 82 | observe_action=False, observe_dangling=False, embedding=False, 83 | embedding_size=8): 84 | """ 85 | Parameters 86 | ---------- 87 | observe_parent : bool 88 | Observe the parent of the Token being selected? 89 | 90 | observe_sibling : bool 91 | Observe the sibling of the Token being selected? 92 | 93 | observe_action : bool 94 | Observe the previously selected Token? 95 | 96 | observe_dangling : bool 97 | Observe the number of dangling nodes? 98 | 99 | embedding : bool 100 | Use embeddings for categorical inputs? 101 | 102 | embedding_size : int 103 | Size of embeddings for each categorical input if embedding=True. 104 | """ 105 | self.observe_parent = observe_parent 106 | self.observe_sibling = observe_sibling 107 | self.observe_action = observe_action 108 | self.observe_dangling = observe_dangling 109 | self.library = Program.library 110 | 111 | # Parameter assertions/warnings 112 | assert self.observe_action + self.observe_parent + self.observe_sibling + self.observe_dangling > 0, \ 113 | "Must include at least one observation." 114 | 115 | self.embedding = embedding 116 | self.embedding_size = embedding_size 117 | 118 | def setup_manager(self, policy): 119 | super().setup_manager(policy) 120 | # Create embeddings if needed 121 | if self.embedding: 122 | initializer = tf.random_uniform_initializer(minval=-1.0, 123 | maxval=1.0, 124 | seed=0) 125 | with tf.variable_scope("embeddings", initializer=initializer): 126 | if self.observe_action: 127 | self.action_embeddings = tf.get_variable("action_embeddings", 128 | [self.library.n_action_inputs, self.embedding_size], 129 | trainable=True) 130 | if self.observe_parent: 131 | self.parent_embeddings = tf.get_variable("parent_embeddings", 132 | [self.library.n_parent_inputs, self.embedding_size], 133 | trainable=True) 134 | if self.observe_sibling: 135 | self.sibling_embeddings = tf.get_variable("sibling_embeddings", 136 | [self.library.n_sibling_inputs, self.embedding_size], 137 | trainable=True) 138 | 139 | def get_tensor_input(self, obs): 140 | observations = [] 141 | unstacked_obs = tf.unstack(obs, axis=1) 142 | action, parent, sibling, dangling = unstacked_obs[:4] 143 | 144 | # Cast action, parent, sibling to int for embedding_lookup or one_hot 145 | action = tf.cast(action, tf.int32) 146 | parent = tf.cast(parent, tf.int32) 147 | sibling = tf.cast(sibling, tf.int32) 148 | 149 | # Action, parent, and sibling inputs are either one-hot or embeddings 150 | if self.observe_action: 151 | if self.embedding: 152 | x = tf.nn.embedding_lookup(self.action_embeddings, action) 153 | else: 154 | x = tf.one_hot(action, depth=self.library.n_action_inputs) 155 | observations.append(x) 156 | if self.observe_parent: 157 | if self.embedding: 158 | x = tf.nn.embedding_lookup(self.parent_embeddings, parent) 159 | else: 160 | x = tf.one_hot(parent, depth=self.library.n_parent_inputs) 161 | observations.append(x) 162 | if self.observe_sibling: 163 | if self.embedding: 164 | x = tf.nn.embedding_lookup(self.sibling_embeddings, sibling) 165 | else: 166 | x = tf.one_hot(sibling, depth=self.library.n_sibling_inputs) 167 | observations.append(x) 168 | 169 | # Dangling input is just the value of dangling 170 | if self.observe_dangling: 171 | x = tf.expand_dims(dangling, axis=-1) 172 | observations.append(x) 173 | 174 | input_ = tf.concat(observations, -1) 175 | # possibly concatenates additional observations (e.g., bert embeddings) 176 | if len(unstacked_obs) > 4: 177 | input_ = tf.concat([input_, tf.stack(unstacked_obs[4:], axis=-1)], axis=-1) 178 | return input_ 179 | -------------------------------------------------------------------------------- /dso/variance.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from dso.program import from_tokens 4 | from dso.utils import weighted_quantile 5 | 6 | 7 | def quantile_variance(memory_queue, policy, batch_size, epsilon, step, 8 | n_experiments=1000, estimate_bias=True, 9 | n_samples_bias=1e6): 10 | 11 | print("Running quantile variance/bias experiments...") 12 | empirical_quantiles = [] 13 | memory_augmented_quantiles = [] 14 | 15 | if len(memory_queue) < memory_queue.capacity: 16 | print("WARNING: Memory queue not yet at capacity.") 17 | 18 | memory_r = memory_queue.get_rewards() 19 | memory_w = memory_queue.compute_probs() 20 | for exp in range(n_experiments): 21 | actions, obs, priors = policy.sample(batch_size) 22 | programs = [from_tokens(a) for a in actions] 23 | r = np.array([p.r for p in programs]) 24 | quantile = np.quantile(r, 1 - epsilon, interpolation="higher") 25 | empirical_quantiles.append(quantile) 26 | unique_programs = [p for p in programs if p.str not in memory_queue.unique_items] 27 | N = len(unique_programs) 28 | sample_r = [p.r for p in unique_programs] 29 | combined_r = np.concatenate([memory_r, sample_r]) 30 | if N == 0: 31 | print("WARNING: Found no unique samples in batch!") 32 | combined_w = memory_w / memory_w.sum() # Renormalize 33 | else: 34 | sample_w = np.repeat((1 - memory_w.sum()) / N, N) 35 | combined_w = np.concatenate([memory_w, sample_w]) 36 | 37 | # Compute the weighted quantile 38 | quantile = weighted_quantile(values=combined_r, weights=combined_w, q=1 - epsilon) 39 | memory_augmented_quantiles.append(quantile) 40 | 41 | empirical_quantiles = np.array(empirical_quantiles) 42 | memory_augmented_quantiles = np.array(memory_augmented_quantiles) 43 | print("Train step:", step) 44 | print("Memory weight:", memory_w.sum()) 45 | print("Mean(empirical quantile):", np.mean(empirical_quantiles)) 46 | print("Var(empirical quantile):", np.var(empirical_quantiles)) 47 | print("Mean(Memory augmented quantile):", np.mean(memory_augmented_quantiles)) 48 | print("Var(Memory augmented quantile):", np.var(memory_augmented_quantiles)) 49 | if estimate_bias: 50 | actions, obs, priors = policy.sample(int(n_samples_bias)) 51 | programs = [from_tokens(a) for a in actions] 52 | r = np.array([p.r for p in programs]) 53 | true_quantile = np.quantile(r, 1 - epsilon, interpolation="higher") 54 | print("'True' empirical quantile:", true_quantile) 55 | print("Empirical quantile bias:", np.mean(np.abs(empirical_quantiles - true_quantile))) 56 | print("Memory-augmented quantile bias:", np.mean(np.abs(memory_augmented_quantiles - true_quantile))) 57 | exit() 58 | -------------------------------------------------------------------------------- /gp.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from collections import Counter 4 | from typing import Optional 5 | 6 | import numpy as np 7 | 8 | from alphagen.data.expression import * 9 | from alphagen.models.linear_alpha_pool import MseAlphaPool 10 | from alphagen.utils.random import reseed_everything 11 | from alphagen_generic.operators import funcs as generic_funcs 12 | from alphagen_generic.features import * 13 | from alphagen_qlib.calculator import QLibStockDataCalculator 14 | from alphagen_qlib.stock_data import initialize_qlib 15 | from gplearn.fitness import make_fitness 16 | from gplearn.functions import make_function 17 | from gplearn.genetic import SymbolicRegressor 18 | 19 | 20 | funcs = [make_function(**func._asdict()) for func in generic_funcs] 21 | 22 | instruments = 'csi300' 23 | seed = 2 24 | reseed_everything(seed) 25 | 26 | cache = {} 27 | device = torch.device("cuda:0") 28 | initialize_qlib("~/.qlib/qlib_data/cn_data_2024h1") 29 | data_train = StockData(instruments, "2012-01-01", "2021-12-31", device=device) 30 | data_test = StockData(instruments, "2022-01-01", "2023-06-30", device=device) 31 | calculator_train = QLibStockDataCalculator(data_train, target) 32 | calculator_test = QLibStockDataCalculator(data_test, target) 33 | 34 | pool = MseAlphaPool( 35 | capacity=20, 36 | calculator=calculator_train, 37 | ic_lower_bound=None, 38 | l1_alpha=5e-3, 39 | device=device 40 | ) 41 | 42 | 43 | def _metric(x, y, w): 44 | key = y[0] 45 | 46 | if key in cache: 47 | return cache[key] 48 | token_len = key.count('(') + key.count(')') 49 | if token_len > 20: 50 | return -1. 51 | 52 | expr = eval(key) 53 | try: 54 | ic = calculator_train.calc_single_IC_ret(expr) 55 | except OutOfDataRangeError: 56 | ic = -1. 57 | if np.isnan(ic): 58 | ic = -1. 59 | cache[key] = ic 60 | return ic 61 | 62 | 63 | Metric = make_fitness(function=_metric, greater_is_better=True) 64 | 65 | 66 | def try_single(): 67 | top_key = Counter(cache).most_common(1)[0][0] 68 | expr = eval(top_key) 69 | ic_test, ric_test = calculator_test.calc_single_all_ret(expr) 70 | return { 71 | 'ic_test': ic_test, 72 | 'ric_test': ric_test 73 | } 74 | 75 | 76 | def try_pool(capacity: int, mutual_ic_thres: Optional[float] = None): 77 | pool = MseAlphaPool( 78 | capacity=capacity, 79 | calculator=calculator_train, 80 | ic_lower_bound=None 81 | ) 82 | exprs = [] 83 | 84 | def acceptable(expr: str) -> bool: 85 | if mutual_ic_thres is None: 86 | return True 87 | return all(abs(pool.calculator.calc_mutual_IC(e, eval(expr))) <= mutual_ic_thres 88 | for e in exprs) 89 | 90 | most_common = dict(Counter(cache).most_common(capacity if mutual_ic_thres is None else None)) 91 | for key in most_common: 92 | if acceptable(key): 93 | exprs.append(eval(key)) 94 | if len(exprs) >= capacity: 95 | break 96 | pool.force_load_exprs(exprs) 97 | 98 | ic_train, ric_train = pool.test_ensemble(calculator_train) 99 | ic_test, ric_test = pool.test_ensemble(calculator_test) 100 | return { 101 | "ic_train": ic_train, 102 | "ric_train": ric_train, 103 | "ic_test": ic_test, 104 | "ric_test": ric_test, 105 | "pool_state": pool.to_json_dict() 106 | } 107 | 108 | 109 | generation = 0 110 | 111 | def ev(): 112 | global generation 113 | generation += 1 114 | directory = f"out/gp/{seed}" 115 | os.makedirs(directory, exist_ok=True) 116 | if generation % 4 != 0: 117 | return 118 | capacity = 20 119 | res = {"pool": capacity, "res": try_pool(capacity, mutual_ic_thres=0.7)} 120 | with open(f'{directory}/{generation}.json', 'w') as f: 121 | json.dump({'res': res, 'cache': cache}, f, indent=4) 122 | 123 | 124 | if __name__ == '__main__': 125 | features = ['open_', 'close', 'high', 'low', 'volume', 'vwap'] 126 | constants = [f'Constant({v})' for v in [-30., -10., -5., -2., -1., -0.5, -0.01, 0.01, 0.5, 1., 2., 5., 10., 30.]] 127 | terminals = features + constants 128 | 129 | X_train = np.array([terminals]) 130 | y_train = np.array([[1]]) 131 | 132 | est_gp = SymbolicRegressor( 133 | population_size=1000, 134 | generations=40, 135 | init_depth=(2, 6), 136 | tournament_size=600, 137 | stopping_criteria=1., 138 | p_crossover=0.3, 139 | p_subtree_mutation=0.1, 140 | p_hoist_mutation=0.01, 141 | p_point_mutation=0.1, 142 | p_point_replace=0.6, 143 | max_samples=0.9, 144 | verbose=1, 145 | parsimony_coefficient=0., 146 | random_state=seed, 147 | function_set=funcs, 148 | metric=Metric, # type: ignore 149 | const_range=None, 150 | n_jobs=1 151 | ) 152 | est_gp.fit(X_train, y_train, callback=ev) 153 | print(est_gp._program.execute(X_train)) 154 | -------------------------------------------------------------------------------- /gplearn/__init__.py: -------------------------------------------------------------------------------- 1 | """Genetic Programming in Python, with a scikit-learn inspired API 2 | 3 | ``gplearn`` is a set of algorithms for learning genetic programming models. 4 | 5 | """ 6 | __version__ = '0.5.dev0' 7 | 8 | __all__ = ['genetic', 'functions', 'fitness'] 9 | -------------------------------------------------------------------------------- /gplearn/fitness.py: -------------------------------------------------------------------------------- 1 | """Metrics to evaluate the fitness of a program. 2 | 3 | The :mod:`gplearn.fitness` module contains some metric with which to evaluate 4 | the computer programs created by the :mod:`gplearn.genetic` module. 5 | """ 6 | 7 | # Author: Trevor Stephens 8 | # 9 | # License: BSD 3 clause 10 | 11 | import numbers 12 | 13 | import numpy as np 14 | from joblib import wrap_non_picklable_objects 15 | from scipy.stats import rankdata 16 | 17 | __all__ = ['make_fitness'] 18 | 19 | 20 | class _Fitness(object): 21 | 22 | """A metric to measure the fitness of a program. 23 | 24 | This object is able to be called with NumPy vectorized arguments and return 25 | a resulting floating point score quantifying the quality of the program's 26 | representation of the true relationship. 27 | 28 | Parameters 29 | ---------- 30 | function : callable 31 | A function with signature function(y, y_pred, sample_weight) that 32 | returns a floating point number. Where `y` is the input target y 33 | vector, `y_pred` is the predicted values from the genetic program, and 34 | sample_weight is the sample_weight vector. 35 | 36 | greater_is_better : bool 37 | Whether a higher value from `function` indicates a better fit. In 38 | general this would be False for metrics indicating the magnitude of 39 | the error, and True for metrics indicating the quality of fit. 40 | 41 | """ 42 | 43 | def __init__(self, function, greater_is_better): 44 | self.function = function 45 | self.greater_is_better = greater_is_better 46 | self.sign = 1 if greater_is_better else -1 47 | 48 | def __call__(self, *args): 49 | return self.function(*args) 50 | 51 | 52 | def make_fitness(*, function, greater_is_better, wrap=True): 53 | """Make a fitness measure, a metric scoring the quality of a program's fit. 54 | 55 | This factory function creates a fitness measure object which measures the 56 | quality of a program's fit and thus its likelihood to undergo genetic 57 | operations into the next generation. The resulting object is able to be 58 | called with NumPy vectorized arguments and return a resulting floating 59 | point score quantifying the quality of the program's representation of the 60 | true relationship. 61 | 62 | Parameters 63 | ---------- 64 | function : callable 65 | A function with signature function(y, y_pred, sample_weight) that 66 | returns a floating point number. Where `y` is the input target y 67 | vector, `y_pred` is the predicted values from the genetic program, and 68 | sample_weight is the sample_weight vector. 69 | 70 | greater_is_better : bool 71 | Whether a higher value from `function` indicates a better fit. In 72 | general this would be False for metrics indicating the magnitude of 73 | the error, and True for metrics indicating the quality of fit. 74 | 75 | wrap : bool, optional (default=True) 76 | When running in parallel, pickling of custom metrics is not supported 77 | by Python's default pickler. This option will wrap the function using 78 | cloudpickle allowing you to pickle your solution, but the evolution may 79 | run slightly more slowly. If you are running single-threaded in an 80 | interactive Python session or have no need to save the model, set to 81 | `False` for faster runs. 82 | 83 | """ 84 | if not isinstance(greater_is_better, bool): 85 | raise ValueError('greater_is_better must be bool, got %s' 86 | % type(greater_is_better)) 87 | if not isinstance(wrap, bool): 88 | raise ValueError('wrap must be an bool, got %s' % type(wrap)) 89 | if function.__code__.co_argcount != 3: 90 | raise ValueError('function requires 3 arguments (y, y_pred, w),' 91 | ' got %d.' % function.__code__.co_argcount) 92 | 93 | if wrap: 94 | return _Fitness(function=wrap_non_picklable_objects(function), 95 | greater_is_better=greater_is_better) 96 | return _Fitness(function=function, 97 | greater_is_better=greater_is_better) 98 | 99 | 100 | def _weighted_pearson(y, y_pred, w): 101 | """Calculate the weighted Pearson correlation coefficient.""" 102 | with np.errstate(divide='ignore', invalid='ignore'): 103 | y_pred_demean = y_pred - np.average(y_pred, weights=w) 104 | y_demean = y - np.average(y, weights=w) 105 | corr = ((np.sum(w * y_pred_demean * y_demean) / np.sum(w)) / 106 | np.sqrt((np.sum(w * y_pred_demean ** 2) * 107 | np.sum(w * y_demean ** 2)) / 108 | (np.sum(w) ** 2))) 109 | if np.isfinite(corr): 110 | return np.abs(corr) 111 | return 0. 112 | 113 | 114 | def _weighted_spearman(y, y_pred, w): 115 | """Calculate the weighted Spearman correlation coefficient.""" 116 | y_pred_ranked = np.apply_along_axis(rankdata, 0, y_pred) 117 | y_ranked = np.apply_along_axis(rankdata, 0, y) 118 | return _weighted_pearson(y_pred_ranked, y_ranked, w) 119 | 120 | 121 | def _mean_absolute_error(y, y_pred, w): 122 | """Calculate the mean absolute error.""" 123 | return np.average(np.abs(y_pred - y), weights=w) 124 | 125 | 126 | def _mean_square_error(y, y_pred, w): 127 | """Calculate the mean square error.""" 128 | return np.average(((y_pred - y) ** 2), weights=w) 129 | 130 | 131 | def _root_mean_square_error(y, y_pred, w): 132 | """Calculate the root mean square error.""" 133 | return np.sqrt(np.average(((y_pred - y) ** 2), weights=w)) 134 | 135 | 136 | def _log_loss(y, y_pred, w): 137 | """Calculate the log loss.""" 138 | eps = 1e-15 139 | inv_y_pred = np.clip(1 - y_pred, eps, 1 - eps) 140 | y_pred = np.clip(y_pred, eps, 1 - eps) 141 | score = y * np.log(y_pred) + (1 - y) * np.log(inv_y_pred) 142 | return np.average(-score, weights=w) 143 | 144 | 145 | weighted_pearson = _Fitness(function=_weighted_pearson, 146 | greater_is_better=True) 147 | weighted_spearman = _Fitness(function=_weighted_spearman, 148 | greater_is_better=True) 149 | mean_absolute_error = _Fitness(function=_mean_absolute_error, 150 | greater_is_better=False) 151 | mean_square_error = _Fitness(function=_mean_square_error, 152 | greater_is_better=False) 153 | root_mean_square_error = _Fitness(function=_root_mean_square_error, 154 | greater_is_better=False) 155 | log_loss = _Fitness(function=_log_loss, 156 | greater_is_better=False) 157 | 158 | _fitness_map = {'pearson': weighted_pearson, 159 | 'spearman': weighted_spearman, 160 | 'mean absolute error': mean_absolute_error, 161 | 'mse': mean_square_error, 162 | 'rmse': root_mean_square_error, 163 | 'log loss': log_loss} 164 | -------------------------------------------------------------------------------- /gplearn/functions.py: -------------------------------------------------------------------------------- 1 | """The functions used to create programs. 2 | 3 | The :mod:`gplearn.functions` module contains all of the functions used by 4 | gplearn programs. It also contains helper methods for a user to define their 5 | own custom functions. 6 | """ 7 | 8 | # Author: Trevor Stephens 9 | # 10 | # License: BSD 3 clause 11 | 12 | import numpy as np 13 | from joblib import wrap_non_picklable_objects 14 | 15 | __all__ = ['make_function'] 16 | 17 | 18 | class _Function(object): 19 | 20 | """A representation of a mathematical relationship, a node in a program. 21 | 22 | This object is able to be called with NumPy vectorized arguments and return 23 | a resulting vector based on a mathematical relationship. 24 | 25 | Parameters 26 | ---------- 27 | function : callable 28 | A function with signature function(x1, *args) that returns a Numpy 29 | array of the same shape as its arguments. 30 | 31 | name : str 32 | The name for the function as it should be represented in the program 33 | and its visualizations. 34 | 35 | arity : int 36 | The number of arguments that the ``function`` takes. 37 | 38 | """ 39 | 40 | def __init__(self, function, name, arity): 41 | self.function = function 42 | self.name = name 43 | self.arity = arity 44 | 45 | def __call__(self, *args): 46 | return self.function(*args) 47 | 48 | 49 | def make_function(*, function, name, arity, wrap=True): 50 | """Make a function node, a representation of a mathematical relationship. 51 | 52 | This factory function creates a function node, one of the core nodes in any 53 | program. The resulting object is able to be called with NumPy vectorized 54 | arguments and return a resulting vector based on a mathematical 55 | relationship. 56 | 57 | Parameters 58 | ---------- 59 | function : callable 60 | A function with signature `function(x1, *args)` that returns a Numpy 61 | array of the same shape as its arguments. 62 | 63 | name : str 64 | The name for the function as it should be represented in the program 65 | and its visualizations. 66 | 67 | arity : int 68 | The number of arguments that the `function` takes. 69 | 70 | wrap : bool, optional (default=True) 71 | When running in parallel, pickling of custom functions is not supported 72 | by Python's default pickler. This option will wrap the function using 73 | cloudpickle allowing you to pickle your solution, but the evolution may 74 | run slightly more slowly. If you are running single-threaded in an 75 | interactive Python session or have no need to save the model, set to 76 | `False` for faster runs. 77 | 78 | """ 79 | if not isinstance(arity, int): 80 | raise ValueError('arity must be an int, got %s' % type(arity)) 81 | if not isinstance(function, np.ufunc): 82 | if function.__code__.co_argcount != arity: 83 | raise ValueError('arity %d does not match required number of ' 84 | 'function arguments of %d.' 85 | % (arity, function.__code__.co_argcount)) 86 | if not isinstance(name, str): 87 | raise ValueError('name must be a string, got %s' % type(name)) 88 | if not isinstance(wrap, bool): 89 | raise ValueError('wrap must be an bool, got %s' % type(wrap)) 90 | 91 | # Check output shape 92 | args = [np.ones(10) for _ in range(arity)] 93 | try: 94 | function(*args) 95 | except (ValueError, TypeError): 96 | raise ValueError('supplied function %s does not support arity of %d.' 97 | % (name, arity)) 98 | if not hasattr(function(*args), 'shape'): 99 | raise ValueError('supplied function %s does not return a numpy array.' 100 | % name) 101 | if function(*args).shape != (10,): 102 | raise ValueError('supplied function %s does not return same shape as ' 103 | 'input vectors.' % name) 104 | 105 | if wrap: 106 | return _Function(function=wrap_non_picklable_objects(function), 107 | name=name, 108 | arity=arity) 109 | return _Function(function=function, 110 | name=name, 111 | arity=arity) 112 | 113 | 114 | def _protected_division(x1, x2): 115 | """Closure of division (x1/x2) for zero denominator.""" 116 | with np.errstate(divide='ignore', invalid='ignore'): 117 | return np.where(np.abs(x2) > 0.001, np.divide(x1, x2), 1.) 118 | 119 | 120 | def _protected_sqrt(x1): 121 | """Closure of square root for negative arguments.""" 122 | return np.sqrt(np.abs(x1)) 123 | 124 | 125 | def _protected_log(x1): 126 | """Closure of log for zero and negative arguments.""" 127 | with np.errstate(divide='ignore', invalid='ignore'): 128 | return np.where(np.abs(x1) > 0.001, np.log(np.abs(x1)), 0.) 129 | 130 | 131 | def _protected_inverse(x1): 132 | """Closure of inverse for zero arguments.""" 133 | with np.errstate(divide='ignore', invalid='ignore'): 134 | return np.where(np.abs(x1) > 0.001, 1. / x1, 0.) 135 | 136 | 137 | def _sigmoid(x1): 138 | """Special case of logistic function to transform to probabilities.""" 139 | with np.errstate(over='ignore', under='ignore'): 140 | return 1 / (1 + np.exp(-x1)) 141 | 142 | 143 | add2 = _Function(function=np.add, name='add', arity=2) 144 | sub2 = _Function(function=np.subtract, name='sub', arity=2) 145 | mul2 = _Function(function=np.multiply, name='mul', arity=2) 146 | div2 = _Function(function=_protected_division, name='div', arity=2) 147 | sqrt1 = _Function(function=_protected_sqrt, name='sqrt', arity=1) 148 | log1 = _Function(function=_protected_log, name='log', arity=1) 149 | neg1 = _Function(function=np.negative, name='neg', arity=1) 150 | inv1 = _Function(function=_protected_inverse, name='inv', arity=1) 151 | abs1 = _Function(function=np.abs, name='abs', arity=1) 152 | max2 = _Function(function=np.maximum, name='max', arity=2) 153 | min2 = _Function(function=np.minimum, name='min', arity=2) 154 | sin1 = _Function(function=np.sin, name='sin', arity=1) 155 | cos1 = _Function(function=np.cos, name='cos', arity=1) 156 | tan1 = _Function(function=np.tan, name='tan', arity=1) 157 | sig1 = _Function(function=_sigmoid, name='sig', arity=1) 158 | 159 | _function_map = {'add': add2, 160 | 'sub': sub2, 161 | 'mul': mul2, 162 | 'div': div2, 163 | 'sqrt': sqrt1, 164 | 'log': log1, 165 | 'abs': abs1, 166 | 'neg': neg1, 167 | 'inv': inv1, 168 | 'max': max2, 169 | 'min': min2, 170 | 'sin': sin1, 171 | 'cos': cos1, 172 | 'tan': tan1} 173 | -------------------------------------------------------------------------------- /gplearn/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RL-MLDM/alphagen/1c187544df587c6109e25948bf0d16e42fce24bf/gplearn/tests/__init__.py -------------------------------------------------------------------------------- /gplearn/tests/test_estimator_checks.py: -------------------------------------------------------------------------------- 1 | """Testing the Genetic Programming module's underlying datastructure 2 | (gplearn.genetic._Program) as well as the classes that use it, 3 | gplearn.genetic.SymbolicRegressor and gplearn.genetic.SymbolicTransformer.""" 4 | 5 | # Author: Trevor Stephens 6 | # 7 | # License: BSD 3 clause 8 | 9 | from sklearn.utils.estimator_checks import check_estimator 10 | 11 | from gplearn.genetic import SymbolicClassifier, SymbolicRegressor 12 | from gplearn.genetic import SymbolicTransformer 13 | 14 | 15 | def test_sklearn_regressor_checks(): 16 | """Run the sklearn estimator validation checks on SymbolicRegressor""" 17 | 18 | check_estimator(SymbolicRegressor(population_size=1000, 19 | generations=5)) 20 | 21 | 22 | def test_sklearn_classifier_checks(): 23 | """Run the sklearn estimator validation checks on SymbolicClassifier""" 24 | 25 | check_estimator(SymbolicClassifier(population_size=50, 26 | generations=5)) 27 | 28 | 29 | def test_sklearn_transformer_checks(): 30 | """Run the sklearn estimator validation checks on SymbolicTransformer""" 31 | 32 | check_estimator(SymbolicTransformer(population_size=50, 33 | hall_of_fame=10, 34 | generations=5)) 35 | -------------------------------------------------------------------------------- /gplearn/tests/test_functions.py: -------------------------------------------------------------------------------- 1 | """Testing the Genetic Programming functions module.""" 2 | 3 | # Author: Trevor Stephens 4 | # 5 | # License: BSD 3 clause 6 | 7 | import pickle 8 | 9 | import numpy as np 10 | from numpy import maximum 11 | from sklearn.datasets import load_diabetes, load_breast_cancer 12 | from sklearn.utils._testing import assert_raises 13 | from sklearn.utils.validation import check_random_state 14 | 15 | from gplearn.functions import _protected_sqrt, make_function 16 | from gplearn.genetic import SymbolicRegressor, SymbolicTransformer 17 | from gplearn.genetic import SymbolicClassifier 18 | 19 | # load the diabetes dataset and randomly permute it 20 | rng = check_random_state(0) 21 | diabetes = load_diabetes() 22 | perm = rng.permutation(diabetes.target.size) 23 | diabetes.data = diabetes.data[perm] 24 | diabetes.target = diabetes.target[perm] 25 | 26 | # load the breast cancer dataset and randomly permute it 27 | cancer = load_breast_cancer() 28 | perm = check_random_state(0).permutation(cancer.target.size) 29 | cancer.data = cancer.data[perm] 30 | cancer.target = cancer.target[perm] 31 | 32 | 33 | def test_validate_function(): 34 | """Check that valid functions are accepted & invalid ones raise error""" 35 | 36 | # Check arity tests 37 | _ = make_function(function=_protected_sqrt, name='sqrt', arity=1) 38 | # non-integer arity 39 | assert_raises(ValueError, 40 | make_function, 41 | function=_protected_sqrt, 42 | name='sqrt', 43 | arity='1') 44 | assert_raises(ValueError, 45 | make_function, 46 | function=_protected_sqrt, 47 | name='sqrt', 48 | arity=1.0) 49 | # non-bool wrap 50 | assert_raises(ValueError, 51 | make_function, 52 | function=_protected_sqrt, 53 | name='sqrt', 54 | arity=1, 55 | wrap='f') 56 | # non-matching arity 57 | assert_raises(ValueError, 58 | make_function, 59 | function=_protected_sqrt, 60 | name='sqrt', 61 | arity=2) 62 | assert_raises(ValueError, 63 | make_function, 64 | function=maximum, 65 | name='max', 66 | arity=1) 67 | 68 | # Check name test 69 | assert_raises(ValueError, 70 | make_function, 71 | function=_protected_sqrt, 72 | name=2, 73 | arity=1) 74 | 75 | # Check return type tests 76 | def bad_fun1(x1, x2): 77 | return 'ni' 78 | assert_raises(ValueError, 79 | make_function, 80 | function=bad_fun1, 81 | name='ni', 82 | arity=2) 83 | 84 | # Check return shape tests 85 | def bad_fun2(x1): 86 | return np.ones((2, 1)) 87 | assert_raises(ValueError, 88 | make_function, 89 | function=bad_fun2, 90 | name='ni', 91 | arity=1) 92 | 93 | # Check closure for negatives test 94 | def _unprotected_sqrt(x1): 95 | with np.errstate(divide='ignore', invalid='ignore'): 96 | return np.sqrt(x1) 97 | assert_raises(ValueError, 98 | make_function, 99 | function=_unprotected_sqrt, 100 | name='sqrt', 101 | arity=1) 102 | 103 | # Check closure for zeros test 104 | def _unprotected_div(x1, x2): 105 | with np.errstate(divide='ignore', invalid='ignore'): 106 | return np.divide(x1, x2) 107 | assert_raises(ValueError, 108 | make_function, 109 | function=_unprotected_div, 110 | name='div', 111 | arity=2) 112 | 113 | 114 | def test_function_in_program(): 115 | """Check that using a custom function in a program works""" 116 | 117 | def logic(x1, x2, x3, x4): 118 | return np.where(x1 > x2, x3, x4) 119 | 120 | logical = make_function(function=logic, 121 | name='logical', 122 | arity=4) 123 | function_set = ['add', 'sub', 'mul', 'div', logical] 124 | est = SymbolicTransformer(generations=2, population_size=2000, 125 | hall_of_fame=100, n_components=10, 126 | function_set=function_set, 127 | parsimony_coefficient=0.0005, 128 | max_samples=0.9, random_state=0) 129 | est.fit(diabetes.data[:300, :], diabetes.target[:300]) 130 | 131 | formula = est._programs[0][3].__str__() 132 | expected_formula = ('add(X3, logical(div(X5, sub(X5, X5)), ' 133 | 'add(X9, -0.621), X8, X4))') 134 | assert(expected_formula == formula) 135 | 136 | 137 | def test_parallel_custom_function(): 138 | """Regression test for running parallel training with custom functions""" 139 | 140 | def _logical(x1, x2, x3, x4): 141 | return np.where(x1 > x2, x3, x4) 142 | 143 | logical = make_function(function=_logical, 144 | name='logical', 145 | arity=4) 146 | est = SymbolicRegressor(generations=2, 147 | function_set=['add', 'sub', 'mul', 'div', logical], 148 | random_state=0, 149 | n_jobs=2) 150 | est.fit(diabetes.data, diabetes.target) 151 | _ = pickle.dumps(est) 152 | 153 | # Unwrapped functions should fail 154 | logical = make_function(function=_logical, 155 | name='logical', 156 | arity=4, 157 | wrap=False) 158 | est = SymbolicRegressor(generations=2, 159 | function_set=['add', 'sub', 'mul', 'div', logical], 160 | random_state=0, 161 | n_jobs=2) 162 | est.fit(diabetes.data, diabetes.target) 163 | assert_raises(AttributeError, pickle.dumps, est) 164 | 165 | # Single threaded will also fail in non-interactive sessions 166 | est = SymbolicRegressor(generations=2, 167 | function_set=['add', 'sub', 'mul', 'div', logical], 168 | random_state=0) 169 | est.fit(diabetes.data, diabetes.target) 170 | assert_raises(AttributeError, pickle.dumps, est) 171 | 172 | 173 | def test_parallel_custom_transformer(): 174 | """Regression test for running parallel training with custom transformer""" 175 | 176 | def _sigmoid(x1): 177 | with np.errstate(over='ignore', under='ignore'): 178 | return 1 / (1 + np.exp(-x1)) 179 | 180 | sigmoid = make_function(function=_sigmoid, 181 | name='sig', 182 | arity=1) 183 | est = SymbolicClassifier(generations=2, 184 | transformer=sigmoid, 185 | random_state=0, 186 | n_jobs=2) 187 | est.fit(cancer.data, cancer.target) 188 | _ = pickle.dumps(est) 189 | 190 | # Unwrapped functions should fail 191 | sigmoid = make_function(function=_sigmoid, 192 | name='sig', 193 | arity=1, 194 | wrap=False) 195 | est = SymbolicClassifier(generations=2, 196 | transformer=sigmoid, 197 | random_state=0, 198 | n_jobs=2) 199 | est.fit(cancer.data, cancer.target) 200 | assert_raises(AttributeError, pickle.dumps, est) 201 | 202 | # Single threaded will also fail in non-interactive sessions 203 | est = SymbolicClassifier(generations=2, 204 | transformer=sigmoid, 205 | random_state=0) 206 | est.fit(cancer.data, cancer.target) 207 | assert_raises(AttributeError, pickle.dumps, est) 208 | -------------------------------------------------------------------------------- /gplearn/tests/test_utils.py: -------------------------------------------------------------------------------- 1 | """Testing the utils module.""" 2 | 3 | # Author: Trevor Stephens 4 | # 5 | # License: BSD 3 clause 6 | 7 | import numpy as np 8 | from sklearn.utils._testing import assert_raises 9 | 10 | from gplearn.utils import _get_n_jobs, check_random_state, cpu_count 11 | 12 | 13 | def test_check_random_state(): 14 | """Check the check_random_state utility function behavior""" 15 | 16 | assert(check_random_state(None) is np.random.mtrand._rand) 17 | assert(check_random_state(np.random) is np.random.mtrand._rand) 18 | 19 | rng_42 = np.random.RandomState(42) 20 | assert(check_random_state(42).randint(100) == rng_42.randint(100)) 21 | 22 | rng_42 = np.random.RandomState(42) 23 | assert(check_random_state(rng_42) is rng_42) 24 | 25 | rng_42 = np.random.RandomState(42) 26 | assert(check_random_state(43).randint(100) != rng_42.randint(100)) 27 | 28 | assert_raises(ValueError, check_random_state, "some invalid seed") 29 | 30 | 31 | def test_get_n_jobs(): 32 | """Check that _get_n_jobs returns expected values""" 33 | 34 | jobs = _get_n_jobs(4) 35 | assert(jobs == 4) 36 | 37 | jobs = -2 38 | expected = cpu_count() + 1 + jobs 39 | jobs = _get_n_jobs(jobs) 40 | assert(jobs == expected) 41 | 42 | assert_raises(ValueError, _get_n_jobs, 0) 43 | -------------------------------------------------------------------------------- /gplearn/utils.py: -------------------------------------------------------------------------------- 1 | """Utilities that are required by gplearn. 2 | 3 | Most of these functions are slightly modified versions of some key utility 4 | functions from scikit-learn that gplearn depends upon. They reside here in 5 | order to maintain compatibility across different versions of scikit-learn. 6 | 7 | """ 8 | 9 | import numbers 10 | 11 | import numpy as np 12 | from joblib import cpu_count 13 | 14 | 15 | def check_random_state(seed): 16 | """Turn seed into a np.random.RandomState instance 17 | 18 | Parameters 19 | ---------- 20 | seed : None | int | instance of RandomState 21 | If seed is None, return the RandomState singleton used by np.random. 22 | If seed is an int, return a new RandomState instance seeded with seed. 23 | If seed is already a RandomState instance, return it. 24 | Otherwise raise ValueError. 25 | 26 | """ 27 | if seed is None or seed is np.random: 28 | return np.random.mtrand._rand 29 | if isinstance(seed, (numbers.Integral, np.integer)): 30 | return np.random.RandomState(seed) 31 | if isinstance(seed, np.random.RandomState): 32 | return seed 33 | raise ValueError('%r cannot be used to seed a numpy.random.RandomState' 34 | ' instance' % seed) 35 | 36 | 37 | def _get_n_jobs(n_jobs): 38 | """Get number of jobs for the computation. 39 | 40 | This function reimplements the logic of joblib to determine the actual 41 | number of jobs depending on the cpu count. If -1 all CPUs are used. 42 | If 1 is given, no parallel computing code is used at all, which is useful 43 | for debugging. For n_jobs below -1, (n_cpus + 1 + n_jobs) are used. 44 | Thus for n_jobs = -2, all CPUs but one are used. 45 | 46 | Parameters 47 | ---------- 48 | n_jobs : int 49 | Number of jobs stated in joblib convention. 50 | 51 | Returns 52 | ------- 53 | n_jobs : int 54 | The actual number of jobs as positive integer. 55 | 56 | """ 57 | if n_jobs < 0: 58 | return max(cpu_count() + 1 + n_jobs, 1) 59 | elif n_jobs == 0: 60 | raise ValueError('Parameter n_jobs == 0 has no meaning.') 61 | else: 62 | return n_jobs 63 | 64 | 65 | def _partition_estimators(n_estimators, n_jobs): 66 | """Private function used to partition estimators between jobs.""" 67 | # Compute the number of jobs 68 | n_jobs = min(_get_n_jobs(n_jobs), n_estimators) 69 | 70 | # Partition estimators between jobs 71 | n_estimators_per_job = (n_estimators // n_jobs) * np.ones(n_jobs, 72 | dtype=int) 73 | n_estimators_per_job[:n_estimators % n_jobs] += 1 74 | starts = np.cumsum(n_estimators_per_job) 75 | 76 | return n_jobs, n_estimators_per_job.tolist(), [0] + starts.tolist() 77 | -------------------------------------------------------------------------------- /images/logo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RL-MLDM/alphagen/1c187544df587c6109e25948bf0d16e42fce24bf/images/logo.jpg -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | baostock==0.8.8 2 | gym==0.26.2 3 | matplotlib==3.3.4 4 | numpy==1.20.1 5 | pandas==1.2.4 6 | qlib==0.0.2.dev20 7 | sb3_contrib==2.0.0 8 | stable_baselines3==2.0.0 9 | torch==2.0.1 10 | shimmy==1.1.0 11 | fire 12 | openai==1.2.3 13 | num2words 14 | -------------------------------------------------------------------------------- /scripts/llm_only.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List 2 | from logging import Logger 3 | from datetime import datetime 4 | import json 5 | from itertools import accumulate 6 | 7 | import fire 8 | import torch 9 | from openai import OpenAI 10 | 11 | from alphagen.data.expression import Expression 12 | from alphagen.data.parser import ExpressionParser 13 | from alphagen.data.expression import * 14 | from alphagen.models.linear_alpha_pool import MseAlphaPool 15 | from alphagen_qlib.calculator import QLibStockDataCalculator 16 | from alphagen_qlib.stock_data import StockData, initialize_qlib 17 | from alphagen_generic.features import target 18 | from alphagen_llm.client import OpenAIClient, ChatConfig 19 | from alphagen_llm.prompts.interaction import DefaultInteraction, DefaultReport 20 | from alphagen_llm.prompts.system_prompt import EXPLAIN_WITH_TEXT_DESC 21 | from alphagen.utils import get_logger 22 | from alphagen.utils.misc import pprint_arguments 23 | 24 | 25 | def build_chat(system_prompt: str, logger: Optional[Logger] = None): 26 | return OpenAIClient( 27 | OpenAI(base_url="https://api.ai.cs.ac.cn/v1"), 28 | ChatConfig( 29 | system_prompt=system_prompt, 30 | logger=logger 31 | ) 32 | ) 33 | 34 | 35 | def build_parser(use_additional_mapping: bool = False) -> ExpressionParser: 36 | mapping = { 37 | "Max": [Greater], 38 | "Min": [Less], 39 | "Delta": [Sub] 40 | } 41 | return ExpressionParser( 42 | Operators, 43 | ignore_case=True, 44 | additional_operator_mapping=mapping if use_additional_mapping else None, 45 | non_positive_time_deltas_allowed=False 46 | ) 47 | 48 | 49 | def build_test_data(instruments: str, device: torch.device, n_half_years: int) -> List[Tuple[str, StockData]]: 50 | halves = (("01-01", "06-30"), ("07-01", "12-31")) 51 | 52 | def get_dataset(i: int) -> Tuple[str, StockData]: 53 | year = 2022 + i // 2 54 | start, end = halves[i % 2] 55 | return ( 56 | f"{year}h{i % 2 + 1}", 57 | StockData( 58 | instrument=instruments, 59 | start_time=f"{year}-{start}", 60 | end_time=f"{year}-{end}", 61 | device=device 62 | ) 63 | ) 64 | 65 | return [get_dataset(i) for i in range(n_half_years)] 66 | 67 | 68 | def run_experiment( 69 | pool_size: int = 20, 70 | n_replace: int = 3, 71 | n_updates: int = 20, 72 | without_weights: bool = False, 73 | contextful: bool = False, 74 | prefix: Optional[str] = None, 75 | force_remove: bool = False, 76 | also_report_history: bool = False 77 | ): 78 | """ 79 | :param pool_size: Maximum alpha pool size 80 | :param n_replace: Replace n alphas on each iteration 81 | :param n_updates: Run n iterations 82 | :param without_weights: Do not report the weights of the alphas to the LLM 83 | :param contextful: Keep context in the conversation 84 | :param prefix: Output location prefix 85 | :param force_remove: Force remove worst old alphas 86 | :param also_report_history: Also report alpha pool update history to the LLM 87 | """ 88 | 89 | args = pprint_arguments() 90 | 91 | initialize_qlib(f"~/.qlib/qlib_data/cn_data") 92 | instruments = "csi300" 93 | device = torch.device("cuda:0") 94 | timestamp = datetime.now().strftime("%Y%m%d%H%M%S") 95 | prefix = str(prefix) + "-" if prefix is not None else "" 96 | out_path = f"./out/llm-tests/interaction/{prefix}{timestamp}" 97 | logger = get_logger(name="llm", file_path=f"{out_path}/llm.log") 98 | 99 | with open(f"{out_path}/config.json", "w") as f: 100 | json.dump(args, f) 101 | 102 | data_train = StockData( 103 | instrument=instruments, 104 | start_time="2012-01-01", 105 | end_time="2021-12-31", 106 | device=device 107 | ) 108 | data_test = build_test_data(instruments, device, n_half_years=3) 109 | calculator_train = QLibStockDataCalculator(data_train, target) 110 | calculator_test = [QLibStockDataCalculator(d, target) for _, d in data_test] 111 | 112 | def make_pool(exprs: List[Expression]) -> MseAlphaPool: 113 | pool = MseAlphaPool( 114 | capacity=max(pool_size, len(exprs)), 115 | calculator=calculator_train, 116 | device=device 117 | ) 118 | pool.force_load_exprs(exprs) 119 | return pool 120 | 121 | def show_iteration(_, iter: int): 122 | print(f"Iteration {iter} finished...") 123 | 124 | inter = DefaultInteraction( 125 | parser=build_parser(), 126 | client=build_chat(EXPLAIN_WITH_TEXT_DESC, logger=logger), 127 | pool_factory=make_pool, 128 | calculator_train=calculator_train, 129 | calculators_test=calculator_test, 130 | replace_k=n_replace, 131 | force_remove=force_remove, 132 | forgetful=not contextful, 133 | no_actual_weights=without_weights, 134 | also_report_history=also_report_history, 135 | on_pool_update=show_iteration 136 | ) 137 | inter.run(n_updates=n_updates) 138 | 139 | with open(f"{out_path}/report.json", "w") as f: 140 | json.dump([r.to_dict() for r in inter.reports], f) 141 | 142 | cum_days = list(accumulate(d.n_days for _, d in data_test)) 143 | mean_ic_results = {} 144 | mean_ics, mean_rics = [], [] 145 | 146 | def get_rolling_means(ics: List[float]) -> List[float]: 147 | cum_ics = accumulate(ic * tup[1].n_days for ic, tup in zip(ics, data_test)) 148 | return [s / n for s, n in zip(cum_ics, cum_days)] 149 | 150 | for report in inter.reports: 151 | mean_ics.append(get_rolling_means(report.test_ics)) 152 | mean_rics.append(get_rolling_means(report.test_rics)) 153 | 154 | for i, (name, _) in enumerate(data_test): 155 | mean_ic_results[name] = { 156 | "ics": [step[i] for step in mean_ics], 157 | "rics": [step[i] for step in mean_rics] 158 | } 159 | 160 | with open(f"{out_path}/rolling_mean_ic.json", "w") as f: 161 | json.dump(mean_ic_results, f) 162 | 163 | 164 | if __name__ == "__main__": 165 | fire.Fire(run_experiment) 166 | -------------------------------------------------------------------------------- /scripts/llm_test_validity.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Dict, Tuple 2 | from dataclasses import dataclass 3 | import json 4 | from logging import Logger 5 | from datetime import datetime 6 | from openai import OpenAI 7 | import fire 8 | 9 | from alphagen.data.expression import * 10 | from alphagen.data.parser import ExpressionParser 11 | from alphagen_llm.client import ChatClient, OpenAIClient, ChatConfig 12 | from alphagen_llm.prompts.common import safe_parse 13 | from alphagen.utils import get_logger 14 | from alphagen_llm.prompts.system_prompt import * 15 | 16 | 17 | _GENERATE_ALPHAS_DEFAULT_PROMPT = "Generate me ten alphas that you think would be indicative of future stock price trend. Each alpha should be on its own line without numbering. Please do not output anything else." 18 | 19 | 20 | @dataclass 21 | class ValidityTestResult: 22 | n_parsers: int 23 | "Number of parsers used in the test." 24 | lines: Dict[str, List[Optional[Expression]]] 25 | "Map of lines to the expressions parsed by each parser." 26 | n_duplicates: int 27 | "Number of duplicated lines found." 28 | 29 | @property 30 | def n_total_lines(self) -> int: 31 | "Total number of lines output by the client. Includes duplicated expressions." 32 | return len(self.lines) + self.n_duplicates 33 | 34 | @property 35 | def duplicate_rate(self) -> float: 36 | "Fraction of total lines that are duplicates." 37 | return self.n_duplicates / self.n_total_lines 38 | 39 | @property 40 | def validity_stats(self) -> List[Tuple[int, float]]: 41 | "Number and fraction of valid lines for each parser. Excluding the duplicates." 42 | counts: List[int] = [sum(1 for parsed in self.lines.values() if parsed[i] is not None) 43 | for i in range(self.n_parsers)] 44 | n_lines = len(self.lines) 45 | return [(c, c / n_lines) for c in counts] 46 | 47 | def __str__(self) -> str: 48 | generic = f"Lines: {self.n_total_lines} ({len(self.lines)} + {self.n_duplicates} duplicates) | Valid: " 49 | return generic + ", ".join(f"{n} ({r * 100:.1f}%)" for n, r in self.validity_stats) 50 | 51 | 52 | def test_validity( 53 | client: ChatClient, 54 | parsers: List[ExpressionParser], 55 | n_repeats: int, 56 | generate_alphas_prompt: str = _GENERATE_ALPHAS_DEFAULT_PROMPT, 57 | with_tqdm: bool = False 58 | ) -> ValidityTestResult: 59 | lines: Dict[str, List[Optional[Expression]]] = {} 60 | duplicates = 0 61 | range_ = range 62 | if with_tqdm: 63 | from tqdm import trange 64 | range_ = trange 65 | for _ in range_(n_repeats): 66 | client.reset() 67 | response = client.chat_complete(generate_alphas_prompt) 68 | output_lines = [stripped for line in response.splitlines() if (stripped := line.strip()) != ""] 69 | for line in output_lines: 70 | if line in lines: 71 | duplicates += 1 72 | continue 73 | lines[line] = [safe_parse(parser, line) for parser in parsers] 74 | return ValidityTestResult(len(parsers), lines, duplicates) 75 | 76 | 77 | def build_chat_client(system_prompt: str, logger: Optional[Logger] = None): 78 | return OpenAIClient( 79 | OpenAI(base_url="https://api.ai.cs.ac.cn/v1"), 80 | ChatConfig( 81 | system_prompt=system_prompt, 82 | logger=logger 83 | ) 84 | ) 85 | 86 | 87 | def build_parser(use_additional_mapping: bool = False) -> ExpressionParser: 88 | mapping = { 89 | "Max": [Greater], 90 | "Min": [Less], 91 | "Delta": [Sub] 92 | } 93 | return ExpressionParser( 94 | Operators, 95 | ignore_case=True, 96 | additional_operator_mapping=mapping if use_additional_mapping else None, 97 | non_positive_time_deltas_allowed=False 98 | ) 99 | 100 | 101 | def run_experiment(n_repeats: int = 10): 102 | timestamp = datetime.now().strftime("%Y%m%d%H%M%S") 103 | file_prefix = f"./out/llm-tests/test_validity/{timestamp}" 104 | logger: Logger = get_logger(name="llm", file_path=f"{file_prefix}.log") 105 | parsers = [build_parser(use_additional_mapping=use) for use in (False, True)] 106 | chat = build_chat_client(EXPLAIN_WITH_TEXT_DESC, logger) 107 | results = test_validity(chat, parsers, n_repeats=n_repeats, with_tqdm=True) 108 | print(results) 109 | with open(f"{file_prefix}.json", "w") as f: 110 | parsed, invalid = [], [] 111 | for line, output in results.lines.items(): 112 | if all(e is None for e in output): 113 | invalid.append(line) 114 | else: 115 | parsed.append(line) 116 | json.dump(dict(parsed=parsed, invalid=invalid), f, indent=4) 117 | 118 | 119 | if __name__ == "__main__": 120 | fire.Fire(run_experiment) 121 | -------------------------------------------------------------------------------- /trade_decision.py: -------------------------------------------------------------------------------- 1 | from math import isnan 2 | 3 | import pandas as pd 4 | from alphagen.trade.base import StockPosition, StockStatus 5 | from alphagen_qlib.calculator import QLibStockDataCalculator 6 | 7 | from alphagen_qlib.strategy import TopKSwapNStrategy 8 | from alphagen_qlib.utils import load_alpha_pool_by_path, load_recent_data 9 | 10 | 11 | POOL_PATH = '/DATA/xuehy/logs/kdd_csi300_20_4_20230410071036/301056_steps_pool.json' 12 | 13 | 14 | if __name__ == '__main__': 15 | data, latest_date = load_recent_data(instrument='csi300', window_size=365, offset=1) 16 | calculator = QLibStockDataCalculator(data=data, target=None) 17 | exprs, weights = load_alpha_pool_by_path(POOL_PATH) 18 | 19 | ensemble_alpha = calculator.make_ensemble_alpha(exprs, weights) 20 | df = data.make_dataframe(ensemble_alpha) 21 | 22 | strategy = TopKSwapNStrategy(K=20, 23 | n_swap=10, 24 | signal=df, # placeholder 25 | min_hold_days=1) 26 | 27 | signal = df.xs(latest_date).to_dict()['0'] 28 | status = StockStatus(pd.DataFrame.from_records([ 29 | (k, not isnan(v), not isnan(v), v) for k, v in signal.items() 30 | ], columns=['code', 'buyable', 'sellable', 'signal'])) 31 | position = pd.DataFrame(columns=StockPosition.dtypes.keys()).astype( 32 | {col: str(dtype) for col, dtype in StockPosition.dtypes.items()} 33 | ) 34 | 35 | to_buy, to_sell = strategy.step_decision(status_df=status, 36 | position_df=position) 37 | 38 | for i, code in enumerate(to_buy): 39 | if (i + 1) % 4 == 0: 40 | print(code) 41 | else: 42 | print(code, end=' ') 43 | --------------------------------------------------------------------------------