├── .gitignore ├── README.md ├── alphagen ├── config.py ├── data │ ├── calculator.py │ ├── expression.py │ ├── tokens.py │ └── tree.py ├── models │ ├── alpha_pool.py │ └── model.py ├── rl │ ├── env │ │ ├── core.py │ │ └── wrapper.py │ └── policy.py └── utils │ ├── __init__.py │ ├── correlation.py │ ├── pytorch_utils.py │ └── random.py ├── alphagen_generic ├── features.py └── operators.py ├── alphagen_qlib ├── calculator.py └── stock_data.py ├── combine_AFF.py ├── data_collection ├── fetch_baostock_data.py └── qlib_dump_bin.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 │ │ ├── dataset.py │ │ ├── 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 ├── exp_AFF_calc_result.ipynb ├── exp_DSO_calc_result.ipynb ├── exp_GP_calc_result.ipynb ├── exp_ML_train_and_result.ipynb ├── exp_RL_calc_result.ipynb ├── gan ├── __init__.py ├── dataset │ ├── __init__.py │ └── collector.py ├── network │ ├── __init__.py │ ├── generater.py │ ├── loss.py │ ├── masker.py │ └── predictor.py └── utils │ ├── __init__.py │ ├── builder.py │ ├── data.py │ ├── pool.py │ └── qlib.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 ├── requirements.txt ├── train_AFF.py ├── train_DSO.py ├── train_GP.py └── train_RL.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | .idea/ 3 | .vscode/ 4 | .DS_STORE 5 | playground.ipynb 6 | playground.py 7 | out/ 8 | out_bak/ 9 | out_gp/ 10 | out_ppo/ 11 | out_ml/ 12 | out_dso/ 13 | ppo2/ 14 | pkl/ 15 | pkl_sr/ 16 | pkl_zip/ 17 | checkpoints/ 18 | .ipynb_checkpoints/ 19 | logs/ 20 | tb_logs/ 21 | code_bak/ 22 | get_baostock_data.ipynb 23 | get_baostock_data.py 24 | *.tar.gz 25 | *.zip 26 | *.csv 27 | # *.ipynb 28 | *.whl 29 | *.pt 30 | *.pkl 31 | *.png 32 | *.npy 33 | *.svg 34 | out22/ 35 | ppo2/ 36 | 240719_infer_data_norank.ipynb 37 | combine_AFF_GRF.py 38 | combine_GRF.ipynb -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AlphaForge(AFF) 2 | 3 | 4 | ### Data Preparation 5 | Similar to [AlphaGen](https://github.com/RL-MLDM/alphagen), We Use [Qlib](https://github.com/microsoft/qlib#data-preparation) as data save tool and download data from free & open-source data source [baostock](http://baostock.com/baostock/index.php/%E9%A6%96%E9%A1%B5). 6 | 7 | Please install Qlib [Qlib](https://github.com/microsoft/qlib) first 8 | 9 | Then download stock data through running `data_collection/fetch_baostock_data.py` 10 | 11 | The next, Modify the correspoding `/path/for/qlib_data` in `gan.utils.data.py` to the data you downloaded (the dafault setting is `~/.qlib/qlib_data/cn_data_rolling`) 12 | 13 | 14 | ### Run Our Model 15 | 16 | #### stage1: Minning alpha factors 17 | ```shell 18 | python train_AFF.py --instruments=csi300 --train_end_year=2020 --seeds=[0,1,2,3,4] --save_name=test --zoo_size=100 19 | ``` 20 | 21 | Here, 22 | - `instruments` is the dataset to use, e.g., `csi300`,`csi500`. 23 | - `seeds` is random seed list, e.g., `[0,1,2]` or `[0]`. 24 | - `train_end_year` is the last year of training set, when train_end_year is 2020,the train,valid and test set is seperately: `2010-01-01 to 2020-12-31`,`2021-01-01 to 2021-12-31`,`2022-01-01 to 2022-12-31` 25 | - `save_name` is the prefix when saving running results. `zoo_size` is the num of factors to save at stage 1 mining model. 26 | 27 | #### stage2: Combining alpha factors 28 | ```shell 29 | python combine_AFF.py --instruments=csi300 --train_end_year=2020 --seeds=[0,1,2,3,4] --save_name=test --n_factors=10 --window=inf 30 | ``` 31 | Here `instruments,train_end_year,seeds,save_name`,` must be the same as it in stage 1 32 | - `n_factors` is the num of factors used at each day, it should be less than or equal to `zoo_size` in stage 1. 33 | - `window` is the slicing window that is used to evaluate the alpha factors in order to dynamicly select and cobine. 34 | 35 | #### stage3: Show the results 36 | 37 | You could run the ipython notebook file 38 | 39 | ```shell 40 | exp_AFF_calc_result.ipynb 41 | ``` 42 | 43 | to generate and concat experiment result. 44 | 45 | 46 | ### Run baseline experiments 47 | 48 | The experiment process of other models is similar to running our AFF model, Except that none of the other models have a combine step. 49 | 50 | #### GP: 51 | 52 | train: `train_RL.py`, show result: `exp_RL_calc_result.ipynb` 53 | 54 | #### RL: 55 | 56 | train: `train_RL.py`, show result: `exp_RL_calc_result.ipynb` 57 | 58 | #### DSO: 59 | 60 | train: `train_RL.py`, show result: `exp_RL_calc_result.ipynb` 61 | 62 | #### ML models including XGBoost, LightGBM and MLP: 63 | 64 | train & show results: `exp_ML_train_and_result.ipynb` 65 | 66 | 67 | 68 | -------------------------------------------------------------------------------- /alphagen/config.py: -------------------------------------------------------------------------------- 1 | from alphagen.data.expression import * 2 | 3 | 4 | MAX_EXPR_LENGTH = 20 5 | MAX_EPISODE_LENGTH = 256 6 | 7 | OPERATORS = Operators 8 | OPERATORS = [ 9 | # Unary 10 | # Abs, 11 | # Sign, 12 | # Log, 13 | Inv, 14 | S_log1p, 15 | # CSRank, 16 | 17 | # Binary, 18 | Add, Sub, Mul, Div, 19 | Pow, 20 | # Greater, Less, 21 | 22 | # Rolling 23 | Ref, ts_mean, ts_sum, ts_std, ts_var, 24 | # ts_skew, 25 | # ts_kurt, 26 | ts_max, ts_min, 27 | ts_med, ts_mad, 28 | # ts_rank, 29 | 30 | ts_div, 31 | ts_pctchange, 32 | # ts_ir, 33 | # ts_min_max_diff, 34 | # ts_max_diff,ts_min_diff, 35 | ts_delta, ts_wma, ts_ema, 36 | 37 | # Pair rolling 38 | ts_cov, ts_corr 39 | ] 40 | 41 | DELTA_TIMES = [1,5,10, 20, 30, 40, 50] 42 | 43 | CONSTANTS = [-30., -10., -5., -2., -1., -0.5, -0.01, 0.01, 0.5, 1., 2., 5., 10., 30.] 44 | 45 | REWARD_PER_STEP = 0. 46 | -------------------------------------------------------------------------------- /alphagen/data/calculator.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | from typing import List, Tuple 3 | 4 | from alphagen.data.expression import Expression 5 | 6 | 7 | class AlphaCalculator(metaclass=ABCMeta): 8 | @abstractmethod 9 | def calc_single_IC_ret(self, expr: Expression) -> float: 10 | 'Calculate IC between a single alpha and a predefined target.' 11 | 12 | @abstractmethod 13 | def calc_single_rIC_ret(self, expr: Expression) -> float: 14 | 'Calculate Rank IC between a single alpha and a predefined target.' 15 | 16 | @abstractmethod 17 | def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]: 18 | 'Calculate both IC and Rank IC between a single alpha and a predefined target.' 19 | 20 | @abstractmethod 21 | def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float: 22 | 'Calculate IC between two alphas.' 23 | 24 | @abstractmethod 25 | def calc_pool_IC_ret(self, exprs: List[Expression], weights: List[float]) -> float: 26 | 'First combine the alphas linearly,' 27 | 'then Calculate IC between the linear combination and a predefined target.' 28 | 29 | @abstractmethod 30 | def calc_pool_rIC_ret(self, exprs: List[Expression], weights: List[float]) -> float: 31 | 'First combine the alphas linearly,' 32 | 'then Calculate Rank IC between the linear combination and a predefined target.' 33 | 34 | @abstractmethod 35 | def calc_pool_all_ret(self, exprs: List[Expression], weights: List[float]) -> Tuple[float, float]: 36 | 'First combine the alphas linearly,' 37 | 'then Calculate both IC and Rank IC between the linear combination and a predefined target.' 38 | -------------------------------------------------------------------------------- /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 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 | BEG_TOKEN = SequenceIndicatorToken(SequenceIndicatorType.BEG) 53 | SEP_TOKEN = SequenceIndicatorToken(SequenceIndicatorType.SEP) 54 | -------------------------------------------------------------------------------- /alphagen/data/tree.py: -------------------------------------------------------------------------------- 1 | from alphagen.data.expression import * 2 | from alphagen.data.tokens import * 3 | 4 | 5 | class ExpressionBuilder: 6 | stack: List[Expression] 7 | 8 | def __init__(self): 9 | self.stack = [] 10 | 11 | def get_tree(self) -> Expression: 12 | if len(self.stack) == 1: 13 | return self.stack[0] 14 | else: 15 | raise InvalidExpressionException(f"Expected only one tree, got {len(self.stack)}") 16 | 17 | def add_token(self, token: Token): 18 | if not self.validate(token): 19 | raise InvalidExpressionException(f"Token {token} not allowed here, stack: {self.stack}.") 20 | if isinstance(token, OperatorToken): 21 | n_args: int = token.operator.n_args() 22 | children = [] 23 | for _ in range(n_args): 24 | children.append(self.stack.pop()) 25 | self.stack.append(token.operator(*reversed(children))) # type: ignore 26 | elif isinstance(token, ConstantToken): 27 | self.stack.append(Constant(token.constant)) 28 | elif isinstance(token, DeltaTimeToken): 29 | self.stack.append(DeltaTime(token.delta_time)) 30 | elif isinstance(token, FeatureToken): 31 | self.stack.append(Feature(token.feature)) 32 | else: 33 | assert False 34 | 35 | def is_valid(self) -> bool: 36 | return len(self.stack) == 1 and self.stack[0].is_featured 37 | 38 | def validate(self, token: Token) -> bool: 39 | if isinstance(token, OperatorToken): 40 | return self.validate_op(token.operator) 41 | elif isinstance(token, DeltaTimeToken): 42 | return self.validate_dt() 43 | elif isinstance(token, ConstantToken): 44 | return self.validate_const() 45 | elif isinstance(token, FeatureToken): 46 | return self.validate_feature() 47 | else: 48 | assert False,token 49 | 50 | def validate_op(self, op: Type[Operator]) -> bool: 51 | if len(self.stack) < op.n_args(): 52 | return False 53 | 54 | if issubclass(op, UnaryOperator): 55 | if not self.stack[-1].is_featured: 56 | return False 57 | elif issubclass(op, BinaryOperator): 58 | if not self.stack[-1].is_featured and not self.stack[-2].is_featured: 59 | return False 60 | if (isinstance(self.stack[-1], DeltaTime) or 61 | isinstance(self.stack[-2], DeltaTime)): 62 | return False 63 | elif issubclass(op, RollingOperator): 64 | if not isinstance(self.stack[-1], DeltaTime): 65 | return False 66 | if not self.stack[-2].is_featured: 67 | return False 68 | elif issubclass(op, PairRollingOperator): 69 | if not isinstance(self.stack[-1], DeltaTime): 70 | return False 71 | if not self.stack[-2].is_featured or not self.stack[-3].is_featured: 72 | return False 73 | else: 74 | assert False 75 | return True 76 | 77 | def validate_dt(self) -> bool: 78 | return len(self.stack) > 0 and self.stack[-1].is_featured 79 | 80 | def validate_const(self) -> bool: 81 | return len(self.stack) == 0 or self.stack[-1].is_featured 82 | 83 | def validate_feature(self) -> bool: 84 | return not (len(self.stack) >= 1 and isinstance(self.stack[-1], DeltaTime)) 85 | 86 | 87 | class InvalidExpressionException(ValueError): 88 | pass 89 | 90 | 91 | if __name__ == '__main__': 92 | tokens = [ 93 | FeatureToken(FeatureType.LOW), 94 | OperatorToken(Abs), 95 | DeltaTimeToken(-10), 96 | OperatorToken(Ref), 97 | FeatureToken(FeatureType.HIGH), 98 | FeatureToken(FeatureType.CLOSE), 99 | OperatorToken(Div), 100 | OperatorToken(Add), 101 | ] 102 | 103 | builder = ExpressionBuilder() 104 | for token in tokens: 105 | builder.add_token(token) 106 | 107 | print(f'res: {str(builder.get_tree())}') 108 | print(f'ref: Add(Ref(Abs($low),-10),Div($high,$close))') 109 | -------------------------------------------------------------------------------- /alphagen/models/model.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, List 2 | import math 3 | 4 | import torch 5 | from torch import nn, Tensor 6 | from torch.distributions import Categorical, Normal 7 | from alphagen.data.expression import Operators 8 | from alphagen.data.tokens import * 9 | 10 | 11 | class TokenEmbedding(nn.Module): 12 | def __init__(self, 13 | d_model: int, 14 | operators: List[Type[Operator]], 15 | delta_time_range: Tuple[int, int], 16 | device: torch.device): 17 | super().__init__() 18 | self._d_model = d_model 19 | self._operators = operators 20 | self._delta_time_range = delta_time_range 21 | self._device = device 22 | 23 | self._const_linear = nn.Linear(1, d_model, device=device) 24 | dt_count = delta_time_range[1] - delta_time_range[0] 25 | total_emb = (len(SequenceIndicatorType) + len(FeatureType) + 26 | len(Operators) + dt_count) 27 | self._emb = nn.Embedding(total_emb, d_model, device=device) 28 | 29 | def forward(self, tokens: List[Token]) -> Tensor: 30 | const_idx: List[int] = [] 31 | consts: List[float] = [] 32 | emb_idx: List[int] = [] 33 | emb_type_idx: List[int] = [] 34 | 35 | feat_offset = len(SequenceIndicatorType) 36 | op_offset = feat_offset + len(FeatureType) 37 | dt_offset = op_offset + len(Operators) 38 | 39 | for i, tok in enumerate(tokens): 40 | if isinstance(tok, ConstantToken): 41 | const_idx.append(i) 42 | consts.append(tok.constant) 43 | continue 44 | emb_idx.append(i) 45 | if isinstance(tok, SequenceIndicatorToken): 46 | emb_type_idx.append(int(tok.indicator)) 47 | elif isinstance(tok, FeatureToken): 48 | emb_type_idx.append(int(tok.feature) + feat_offset) 49 | elif isinstance(tok, OperatorToken): 50 | emb_type_idx.append(Operators.index(tok.operator) + op_offset) 51 | elif isinstance(tok, DeltaTimeToken): 52 | emb_type_idx.append(int(tok.delta_time) - self._delta_time_range[0] + dt_offset) 53 | else: 54 | assert False, "NullToken is not allowed here" 55 | 56 | result = torch.zeros(len(tokens), self._d_model, 57 | dtype=torch.float, device=self._device) 58 | if len(const_idx) != 0: 59 | const_tensor = torch.tensor(consts, device=self._device).unsqueeze(1) 60 | result[const_idx] = self._const_linear(const_tensor) 61 | if len(emb_idx) != 0: 62 | result[emb_idx] = self._emb(torch.tensor(emb_type_idx, dtype=torch.long, device=self._device)) 63 | 64 | return result 65 | 66 | 67 | class PositionalEncoding(nn.Module): 68 | def __init__(self, d_model: int, max_len: int = 5000): 69 | super().__init__() 70 | position = torch.arange(max_len).unsqueeze(1) 71 | div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) 72 | pe = torch.zeros(max_len, d_model) 73 | pe[:, 0::2] = torch.sin(position * div_term) 74 | pe[:, 1::2] = torch.cos(position * div_term) 75 | self.register_buffer('_pe', pe) 76 | 77 | def forward(self, x: Tensor) -> Tensor: 78 | "x: ([batch_size, ]seq_len, embedding_dim)" 79 | seq_len = x.size(0) if x.dim() == 2 else x.size(1) 80 | return x + self._pe[:seq_len] # type: ignore 81 | 82 | 83 | class ExpressionGenerator(nn.Module): 84 | def __init__(self, 85 | n_encoder_layers: int, 86 | n_decoder_layers: int, 87 | d_model: int, 88 | n_head: int, 89 | d_ffn: int, 90 | dropout: float, 91 | operators: List[Type[Operator]], 92 | delta_time_range: Tuple[int, int], 93 | device: torch.device): 94 | super().__init__() 95 | self._device = device 96 | self._d_model = d_model 97 | self._operators = operators 98 | self._dt_range = delta_time_range 99 | 100 | self._token_emb = TokenEmbedding(d_model, operators, delta_time_range, device) 101 | self._pos_enc = PositionalEncoding(d_model).to(device) 102 | self._encoder = nn.TransformerEncoder( 103 | nn.TransformerEncoderLayer( 104 | d_model=d_model, nhead=n_head, dim_feedforward=d_ffn, 105 | dropout=dropout, device=device 106 | ), 107 | n_encoder_layers, 108 | norm=nn.LayerNorm(d_model, device=device) 109 | ) 110 | self._decoder = nn.TransformerDecoder( 111 | nn.TransformerDecoderLayer( 112 | d_model=d_model, nhead=n_head, dim_feedforward=d_ffn, 113 | dropout=dropout, device=device 114 | ), 115 | n_decoder_layers, 116 | norm=nn.LayerNorm(d_model, device=device) 117 | ) 118 | self._op_feat_const_linear = nn.Linear(d_model, 3, device=device) 119 | self._op_linear = nn.Linear(d_model, len(operators), device=device) 120 | self._feat_linear = nn.Linear(d_model, len(FeatureType), device=device) 121 | self._const_linear = nn.Linear(d_model, 2, device=device) 122 | self._dt_linear = nn.Linear(d_model, delta_time_range[1] - delta_time_range[0], 123 | device=device) 124 | 125 | def embed_expressions(self, expr: List[Token]) -> Tensor: 126 | res: Tensor = self._token_emb(expr) 127 | return self._pos_enc(res) 128 | 129 | def encode_expressions(self, expr: List[Token]) -> Tensor: 130 | res = self.embed_expressions(expr) 131 | res = self._encoder.forward(res.unsqueeze(1)) 132 | return res.squeeze(1) 133 | 134 | def forward( 135 | self, 136 | encoder_state: Tensor, 137 | decoder_tokens: List[Token], 138 | sample_delta_time: bool = False) -> Tuple[Token, Tensor]: 139 | """ 140 | Sample the next token. 141 | Returns the token and its corresponding log-prob (normalized) 142 | """ 143 | decoder_tokens = decoder_tokens.copy() 144 | decoder_tokens.insert(0, SEP_TOKEN) 145 | length = len(decoder_tokens) 146 | causal_mask = torch.triu( # [L, L] 147 | torch.ones(length, length, dtype=torch.bool, device=self._device), 148 | diagonal=1) 149 | decoder_state = self.embed_expressions(decoder_tokens) 150 | res: Tensor = self._decoder( # [L, 1, D] 151 | decoder_state.unsqueeze(1), 152 | encoder_state.unsqueeze(1), 153 | tgt_mask=causal_mask 154 | ) 155 | res = res.mean(dim=0).reshape(-1) # [D] 156 | 157 | if sample_delta_time: 158 | dist = Categorical(logits=self._dt_linear(res)) 159 | idx = dist.sample() 160 | log_prob = dist.log_prob(idx) 161 | dt = self._dt_range[0] + int(idx) 162 | return DeltaTimeToken(dt), log_prob 163 | 164 | select_dist = Categorical(logits=self._op_feat_const_linear(res)) 165 | idx = select_dist.sample() 166 | select_log_prob = select_dist.log_prob(idx) 167 | idx = int(idx) 168 | if idx == 0: # Operators 169 | dist = Categorical(logits=self._op_linear(res)) 170 | idx = dist.sample() 171 | log_prob = select_log_prob + dist.log_prob(idx) 172 | return OperatorToken(self._operators[int(idx)]), log_prob 173 | elif idx == 1: # Features 174 | dist = Categorical(logits=self._feat_linear(res)) 175 | idx = dist.sample() 176 | log_prob = select_log_prob + dist.log_prob(idx) 177 | return FeatureToken(FeatureType(int(idx))), log_prob 178 | else: # Constants 179 | affine: Tensor = self._const_linear(res) 180 | mu, sigma = affine[0], affine[1].exp() 181 | dist = Normal(mu, sigma) 182 | z = dist.sample() 183 | log_prob = select_log_prob + dist.log_prob(z) 184 | return ConstantToken(float(z)), log_prob 185 | -------------------------------------------------------------------------------- /alphagen/rl/env/core.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Optional 2 | import 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, AlphaPool 10 | from alphagen.utils import reseed_everything 11 | 12 | 13 | class AlphaEnvCore(gym.Env): 14 | pool: AlphaPoolBase 15 | _tokens: List[Token] 16 | _builder: ExpressionBuilder 17 | _print_expr: bool 18 | 19 | def __init__(self, 20 | pool: AlphaPoolBase, 21 | device: torch.device = torch.device('cuda:0'), 22 | print_expr: bool = False 23 | ): 24 | super().__init__() 25 | 26 | self.pool = pool 27 | self._print_expr = print_expr 28 | self._device = device 29 | 30 | self.eval_cnt = 0 31 | 32 | def reset( 33 | self, *, 34 | seed: Optional[int] = None, 35 | return_info: bool = False, 36 | options: Optional[dict] = None 37 | ) -> Tuple[List[Token], dict]: 38 | reseed_everything(seed) 39 | self._tokens = [BEG_TOKEN] 40 | self._builder = ExpressionBuilder() 41 | return self._tokens, self._valid_action_types() 42 | 43 | def step(self, action: Token) -> Tuple[List[Token], float, bool, dict]: 44 | if (isinstance(action, SequenceIndicatorToken) and 45 | action.indicator == SequenceIndicatorType.SEP): 46 | reward = self._evaluate() 47 | done = True 48 | elif len(self._tokens) < MAX_EXPR_LENGTH: 49 | self._tokens.append(action) 50 | self._builder.add_token(action) 51 | done = False 52 | reward = 0.0 53 | else: 54 | done = True 55 | reward = self._evaluate() if self._builder.is_valid() else -0. 56 | 57 | if math.isnan(reward): 58 | reward = 0. 59 | return self._tokens, reward, done, self._valid_action_types() 60 | 61 | def _evaluate(self): 62 | expr: Expression = self._builder.get_tree() 63 | if self._print_expr: 64 | print(expr) 65 | try: 66 | ret = self.pool.try_new_expr(expr) 67 | self.eval_cnt += 1 68 | return ret 69 | except OutOfDataRangeError: 70 | return 0. 71 | 72 | def _valid_action_types(self) -> dict: 73 | valid_op_unary = self._builder.validate_op(UnaryOperator) 74 | valid_op_binary = self._builder.validate_op(BinaryOperator) 75 | valid_op_rolling = self._builder.validate_op(RollingOperator) 76 | valid_op_pair_rolling = self._builder.validate_op(PairRollingOperator) 77 | 78 | valid_op = valid_op_unary or valid_op_binary or valid_op_rolling or valid_op_pair_rolling 79 | valid_dt = self._builder.validate_dt() 80 | valid_const = self._builder.validate_const() 81 | valid_feature = self._builder.validate_feature() 82 | valid_stop = self._builder.is_valid() 83 | 84 | ret = { 85 | 'select': [valid_op, valid_feature, valid_const, valid_dt, valid_stop], 86 | 'op': { 87 | UnaryOperator: valid_op_unary, 88 | BinaryOperator: valid_op_binary, 89 | RollingOperator: valid_op_rolling, 90 | PairRollingOperator: valid_op_pair_rolling 91 | } 92 | } 93 | return ret 94 | 95 | def valid_action_types(self) -> dict: 96 | return self._valid_action_types() 97 | 98 | def render(self, mode='human'): 99 | pass 100 | -------------------------------------------------------------------------------- /alphagen/rl/env/wrapper.py: -------------------------------------------------------------------------------- 1 | import gym 2 | import gym.spaces 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, AlphaPool 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 | 17 | SIZE_ALL = SIZE_NULL + SIZE_OP + SIZE_FEATURE + SIZE_DELTA_TIME + SIZE_CONSTANT + SIZE_SEP 18 | SIZE_ACTION = SIZE_ALL - SIZE_NULL 19 | 20 | OFFSET_OP = SIZE_NULL 21 | OFFSET_FEATURE = OFFSET_OP + SIZE_OP 22 | OFFSET_DELTA_TIME = OFFSET_FEATURE + SIZE_FEATURE 23 | OFFSET_CONSTANT = OFFSET_DELTA_TIME + SIZE_DELTA_TIME 24 | OFFSET_SEP = OFFSET_CONSTANT + SIZE_CONSTANT 25 | 26 | 27 | def action2token(action_raw: int) -> Token: 28 | action = action_raw + 1 29 | if action < OFFSET_OP: 30 | raise ValueError 31 | elif action < OFFSET_FEATURE: 32 | return OperatorToken(OPERATORS[action - OFFSET_OP]) 33 | elif action < OFFSET_DELTA_TIME: 34 | return FeatureToken(FeatureType(action - OFFSET_FEATURE)) 35 | elif action < OFFSET_CONSTANT: 36 | return DeltaTimeToken(DELTA_TIMES[action - OFFSET_DELTA_TIME]) 37 | elif action < OFFSET_SEP: 38 | return ConstantToken(CONSTANTS[action - OFFSET_CONSTANT]) 39 | elif action == OFFSET_SEP: 40 | return SequenceIndicatorToken(SequenceIndicatorType.SEP) 41 | else: 42 | assert False 43 | 44 | 45 | class AlphaEnvWrapper(gym.Wrapper): 46 | state: np.ndarray 47 | env: AlphaEnvCore 48 | action_space: gym.spaces.Discrete 49 | observation_space: gym.spaces.Box 50 | counter: int 51 | 52 | def __init__(self, env: AlphaEnvCore): 53 | super().__init__(env) 54 | self.action_space = gym.spaces.Discrete(SIZE_ACTION) 55 | self.observation_space = gym.spaces.Box(low=0, high=SIZE_ALL - 1, shape=(MAX_EXPR_LENGTH, ), dtype=np.uint8) 56 | 57 | def reset(self, **kwargs) -> np.ndarray: 58 | self.counter = 0 59 | self.state = np.zeros(MAX_EXPR_LENGTH, dtype=np.uint8) 60 | self.env.reset() 61 | return self.state 62 | 63 | def step(self, action: int): 64 | _, reward, done, info = self.env.step(self.action(action)) 65 | if not done: 66 | self.state[self.counter] = action 67 | self.counter += 1 68 | return self.state, self.reward(reward), done, info 69 | 70 | def action(self, action: int) -> Token: 71 | return action2token(action) 72 | 73 | def reward(self, reward: float) -> float: 74 | return reward + REWARD_PER_STEP 75 | 76 | def action_masks(self) -> np.ndarray: 77 | res = np.zeros(SIZE_ACTION, dtype=bool) 78 | valid = self.env.valid_action_types() 79 | for i in range(OFFSET_OP, OFFSET_OP + SIZE_OP): 80 | if valid['op'][OPERATORS[i - OFFSET_OP].category_type()]: 81 | res[i - 1] = True 82 | if valid['select'][1]: # FEATURE 83 | for i in range(OFFSET_FEATURE, OFFSET_FEATURE + SIZE_FEATURE): 84 | res[i - 1] = True 85 | if valid['select'][2]: # CONSTANT 86 | for i in range(OFFSET_CONSTANT, OFFSET_CONSTANT + SIZE_CONSTANT): 87 | res[i - 1] = True 88 | if valid['select'][3]: # DELTA_TIME 89 | for i in range(OFFSET_DELTA_TIME, OFFSET_DELTA_TIME + SIZE_DELTA_TIME): 90 | res[i - 1] = True 91 | if valid['select'][4]: # SEP 92 | res[OFFSET_SEP - 1] = True 93 | return res 94 | 95 | 96 | def AlphaEnv(pool: AlphaPoolBase, **kwargs): 97 | return AlphaEnvWrapper(AlphaEnvCore(pool=pool, **kwargs)) 98 | -------------------------------------------------------------------------------- /alphagen/rl/policy.py: -------------------------------------------------------------------------------- 1 | import gym 2 | import gym.spaces 3 | import math 4 | import torch.nn.functional as F 5 | from stable_baselines3.common.torch_layers import BaseFeaturesExtractor 6 | from torch import nn 7 | 8 | from alphagen.data.expression import * 9 | 10 | 11 | class PositionalEncoding(nn.Module): 12 | def __init__(self, d_model: int, max_len: int = 5000): 13 | super().__init__() 14 | position = torch.arange(max_len).unsqueeze(1) 15 | div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) 16 | pe = torch.zeros(max_len, d_model) 17 | pe[:, 0::2] = torch.sin(position * div_term) 18 | pe[:, 1::2] = torch.cos(position * div_term) 19 | self.register_buffer('_pe', pe) 20 | 21 | def forward(self, x: Tensor) -> Tensor: 22 | "x: ([batch_size, ]seq_len, embedding_dim)" 23 | seq_len = x.size(0) if x.dim() == 2 else x.size(1) 24 | return x + self._pe[:seq_len] # type: ignore 25 | 26 | 27 | class TransformerSharedNet(BaseFeaturesExtractor): 28 | def __init__( 29 | self, 30 | observation_space: gym.Space, 31 | n_encoder_layers: int, 32 | d_model: int, 33 | n_head: int, 34 | d_ffn: int, 35 | dropout: float, 36 | device: torch.device 37 | ): 38 | super().__init__(observation_space, d_model) 39 | 40 | assert isinstance(observation_space, gym.spaces.Box) 41 | n_actions = observation_space.high[0] + 1 # type: ignore 42 | 43 | self._device = device 44 | self._d_model = d_model 45 | self._n_actions: float = n_actions 46 | 47 | self._token_emb = nn.Embedding(n_actions + 1, d_model, 0) # Last one is [BEG] 48 | self._pos_enc = PositionalEncoding(d_model).to(device) 49 | 50 | self._transformer = nn.TransformerEncoder( 51 | nn.TransformerEncoderLayer( 52 | d_model=d_model, nhead=n_head, 53 | dim_feedforward=d_ffn, dropout=dropout, 54 | activation=lambda x: F.leaky_relu(x), # type: ignore 55 | batch_first=True, device=device 56 | ), 57 | num_layers=n_encoder_layers, 58 | norm=nn.LayerNorm(d_model, eps=1e-5, device=device) 59 | ) 60 | 61 | def forward(self, obs: Tensor) -> Tensor: 62 | bs, seqlen = obs.shape 63 | beg = torch.full((bs, 1), fill_value=self._n_actions, dtype=torch.long, device=obs.device) 64 | obs = torch.cat((beg, obs.long()), dim=1) 65 | pad_mask = obs == 0 66 | src = self._pos_enc(self._token_emb(obs)) 67 | res = self._transformer(src, src_key_padding_mask=pad_mask) 68 | return res.mean(dim=1) 69 | 70 | 71 | class LSTMSharedNet(BaseFeaturesExtractor): 72 | def __init__( 73 | self, 74 | observation_space: gym.Space, 75 | n_layers: int, 76 | d_model: int, 77 | dropout: float, 78 | device: torch.device 79 | ): 80 | super().__init__(observation_space, d_model) 81 | 82 | assert isinstance(observation_space, gym.spaces.Box) 83 | n_actions = observation_space.high[0] + 1 # type: ignore 84 | 85 | self._device = device 86 | self._d_model = d_model 87 | self._n_actions: float = n_actions 88 | 89 | self._token_emb = nn.Embedding(n_actions + 1, d_model, 0) # Last one is [BEG] 90 | self._pos_enc = PositionalEncoding(d_model).to(device) 91 | 92 | self._lstm = nn.LSTM( 93 | input_size=d_model, 94 | hidden_size=d_model, 95 | num_layers=n_layers, 96 | batch_first=True, 97 | dropout=dropout 98 | ) 99 | 100 | def forward(self, obs: Tensor) -> Tensor: 101 | bs, seqlen = obs.shape 102 | beg = torch.full((bs, 1), fill_value=self._n_actions, dtype=torch.long, device=obs.device) 103 | obs = torch.cat((beg, obs.long()), dim=1) 104 | real_len = (obs != 0).sum(1).max() 105 | src = self._pos_enc(self._token_emb(obs)) 106 | res = self._lstm(src[:,:real_len])[0] 107 | return res.mean(dim=1) 108 | 109 | 110 | class Decoder(BaseFeaturesExtractor): 111 | def __init__( 112 | self, 113 | observation_space: gym.Space, 114 | n_layers: int, 115 | d_model: int, 116 | n_head: int, 117 | d_ffn: int, 118 | dropout: float, 119 | device: torch.device 120 | ): 121 | super().__init__(observation_space, d_model) 122 | 123 | assert isinstance(observation_space, gym.spaces.Box) 124 | n_actions = observation_space.high[0] + 1 # type: ignore 125 | 126 | self._device = device 127 | self._d_model = d_model 128 | self._n_actions: float = n_actions 129 | 130 | self._token_emb = nn.Embedding(n_actions + 1, d_model, 0) # Last one is [BEG] 131 | self._pos_enc = PositionalEncoding(d_model).to(device) 132 | 133 | # Actually an encoder for now 134 | self._decoder = nn.TransformerEncoder( 135 | nn.TransformerEncoderLayer( 136 | d_model=d_model, nhead=n_head, dim_feedforward=d_ffn, 137 | dropout=dropout, batch_first=True, device=device 138 | ), 139 | n_layers, 140 | norm=nn.LayerNorm(d_model, device=device) 141 | ) 142 | 143 | def forward(self, obs: Tensor) -> Tensor: 144 | batch_size = obs.size(0) 145 | begins = torch.full(size=(batch_size, 1), fill_value=self._n_actions, 146 | dtype=torch.long, device=obs.device) 147 | obs = torch.cat((begins, obs.type(torch.long)), dim=1) # (bs, len) 148 | pad_mask = obs == 0 149 | res = self._token_emb(obs) # (bs, len, d_model) 150 | res = self._pos_enc(res) # (bs, len, d_model) 151 | res = self._decoder(res, src_key_padding_mask=pad_mask) # (bs, len, d_model) 152 | return res.mean(dim=1) # (bs, d_model) 153 | -------------------------------------------------------------------------------- /alphagen/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .correlation import batch_spearmanr 2 | from .random import reseed_everything 3 | -------------------------------------------------------------------------------- /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 | # nan_mask = (~torch.isfinite(x)) | (~torch.isfinite(y)) 12 | x[nan_mask] = fill_with 13 | y[nan_mask] = fill_with 14 | n = (~nan_mask).sum(dim=1) 15 | return x, y, n, nan_mask 16 | 17 | 18 | def _rank_data(x: Tensor, nan_mask: Tensor) -> Tensor: 19 | rank = x.argsort().argsort().float() # [d, s] 20 | eq = x[:, None] == x[:, :, None] # [d, s, s] 21 | eq = eq / eq.sum(dim=2, keepdim=True) # [d, s, s] 22 | rank = (eq @ rank[:, :, None]).squeeze(dim=2) 23 | rank[nan_mask] = 0 24 | return rank # [d, s] 25 | 26 | 27 | # def _rank_data(x: Tensor, nan_mask: Tensor) -> Tensor: 28 | # x = x.clone() 29 | # x[nan_mask] = float('nan') 30 | # x = x.argsort().float() 31 | # x[nan_mask] = float('nan') 32 | # rank = x.argsort().float() 33 | # # rank = x.argsort().argsort().float() # [d, s] 34 | 35 | 36 | # # eq = x[:, None] == x[:, :, None] # [d, s, s] 37 | # # eq = eq / eq.sum(dim=2, keepdim=True) # [d, s, s] 38 | 39 | # # rank = (eq @ rank[:, :, None]).squeeze(dim=2) 40 | # rank[nan_mask] = 0 41 | # return rank # [d, s] 42 | 43 | 44 | def _batch_pearsonr_given_mask( 45 | x: Tensor, y: Tensor, 46 | n: Tensor, mask: Tensor 47 | ) -> Tensor: 48 | x_mean, x_std = masked_mean_std(x, n, mask) 49 | y_mean, y_std = masked_mean_std(y, n, mask) 50 | cov = (x * y).sum(dim=1) / n - x_mean * y_mean 51 | stdmul = x_std * y_std 52 | stdmul[(x_std < 1e-3) | (y_std < 1e-3)] = 1 53 | corrs = cov / stdmul 54 | return corrs 55 | 56 | def _batch_ret_given_mask( 57 | x: Tensor, y: Tensor, 58 | n: Tensor, mask: Tensor 59 | ) -> Tensor: 60 | x_mean, x_std = masked_mean_std(x, n, mask) 61 | y_mean, y_std = masked_mean_std(y, n, mask) 62 | cov = (x * y).sum(dim=1) / n - x_mean * y_mean 63 | stdmul = x_std 64 | stdmul[(x_std < 1e-3) | (y_std < 1e-3)] = 1 65 | corrs = cov / (stdmul**2) 66 | return corrs 67 | 68 | def batch_spearmanr(x: Tensor, y: Tensor) -> Tensor: 69 | x, y, n, nan_mask = _mask_either_nan(x, y) 70 | rx = _rank_data(x, nan_mask) 71 | ry = _rank_data(y, nan_mask) 72 | return _batch_pearsonr_given_mask(rx, ry, n, nan_mask) 73 | 74 | 75 | def batch_pearsonr(x: Tensor, y: Tensor) -> Tensor: 76 | res = _batch_pearsonr_given_mask(*_mask_either_nan(x, y, fill_with=0.)) 77 | # fillna 78 | res[res.isnan()] = 0 79 | return res 80 | 81 | def batch_ret(x:Tensor,y:Tensor)->Tensor: 82 | return _batch_ret_given_mask(*_mask_either_nan(x, y, fill_with=0.)) 83 | 84 | def _mask_either_nan_y_only(x: Tensor, y: Tensor, fill_with: float = torch.nan): 85 | x = x.clone() # [days, stocks] 86 | y = y.clone() # [days, stocks] 87 | nan_mask = y.isnan() 88 | # nan_mask = ~torch.isfinite(y) 89 | x[nan_mask] = fill_with 90 | y[nan_mask] = fill_with 91 | n = (~nan_mask).sum(dim=1) 92 | return x, y, n, nan_mask 93 | def batch_pearsonr_full_y(x:Tensor,y:Tensor)->Tensor: 94 | x,y,n,nan_mask = _mask_either_nan_y_only(x,y,fill_with=0.) 95 | return _batch_pearsonr_given_mask(x,y,n,nan_mask) -------------------------------------------------------------------------------- /alphagen/utils/pytorch_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Optional 2 | import torch 3 | from torch import nn, 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 necessary 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 | # mask = ~torch.isfinite(x) 20 | if n is None: 21 | n = (~mask).sum(dim=1) 22 | x = x.clone() 23 | x[mask] = 0. 24 | mean = x.sum(dim=1) / n 25 | std = ((((x - mean[:, None]) * ~mask) ** 2).sum(dim=1) / n).sqrt() 26 | return mean, std 27 | 28 | 29 | def normalize_by_day(value: Tensor) -> Tensor: 30 | mean, std = masked_mean_std(value) 31 | value = (value - mean[:, None]) / std[:, None] 32 | nan_mask = torch.isnan(value) 33 | # nan_mask = ~torch.isfinite(value) 34 | value[nan_mask] = 0. 35 | return value -------------------------------------------------------------------------------- /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 = Feature(FeatureType.HIGH) 6 | low = Feature(FeatureType.LOW) 7 | volume = Feature(FeatureType.VOLUME) 8 | open_ = Feature(FeatureType.OPEN) 9 | close = Feature(FeatureType.CLOSE) 10 | vwap = Feature(FeatureType.VWAP) 11 | # target = Ref(close, -20) / close - 1 12 | target = Ref(vwap,-21)/Ref(vwap,-1)-1 -------------------------------------------------------------------------------- /alphagen_generic/operators.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | 3 | import numpy as np 4 | from alphagen.data.expression import * 5 | 6 | # OPERATORS = [ 7 | # # Unary 8 | # # Abs, 9 | # # Sign, 10 | # # Log, 11 | # Inv, 12 | # S_log1p, 13 | # # CSRank, 14 | 15 | # # Binary, 16 | # Add, Sub, Mul, Div, 17 | # Pow, 18 | # # Greater, Less, 19 | 20 | # # Rolling 21 | # Ref, ts_mean, ts_sum, ts_std, ts_var, 22 | # # ts_skew, 23 | # # ts_kurt, 24 | # ts_max, ts_min, 25 | # ts_med, ts_mad, 26 | # # ts_rank, 27 | 28 | # ts_div, 29 | # ts_pctchange, 30 | # # ts_ir, 31 | # # ts_min_max_diff, 32 | # # ts_max_diff,ts_min_diff, 33 | # ts_delta, ts_wma, ts_ema, 34 | 35 | # # Pair rolling 36 | # ts_cov, ts_corr 37 | # ] 38 | 39 | GenericOperator = namedtuple('GenericOperator', ['name', 'function', 'arity']) 40 | 41 | unary_ops = [Inv, S_log1p] 42 | binary_ops = [Add, Sub, Mul, Div, Pow,] 43 | rolling_ops = [Ref, ts_mean, ts_sum, ts_std, ts_var, ts_max, ts_min,ts_med, ts_mad,ts_div, ts_pctchange, ts_delta, ts_wma, ts_ema,] 44 | rolling_binary_ops = [ts_cov, ts_corr] 45 | 46 | def unary(cls): 47 | def _calc(a): 48 | n = len(a) 49 | return np.array([f'{cls.__name__}({a[i]})' for i in range(n)]) 50 | 51 | return _calc 52 | 53 | 54 | def binary(cls): 55 | def _calc(a, b): 56 | n = len(a) 57 | a = a.astype(str) 58 | b = b.astype(str) 59 | return np.array([f'{cls.__name__}({a[i]},{b[i]})' for i in range(n)]) 60 | 61 | return _calc 62 | 63 | def rolling(cls, day): 64 | def _calc(a): 65 | n = len(a) 66 | return np.array([f'{cls.__name__}({a[i]},{day})' for i in range(n)]) 67 | 68 | return _calc 69 | 70 | def rolling_binary(cls, day): 71 | def _calc(a, b): 72 | n = len(a) 73 | a = a.astype(str) 74 | b = b.astype(str) 75 | return np.array([f'{cls.__name__}({a[i]},{b[i]},{day})' for i in range(n)]) 76 | 77 | return _calc 78 | 79 | funcs: List[GenericOperator] = [] 80 | for op in unary_ops: 81 | funcs.append(GenericOperator(function=unary(op), name=op.__name__, arity=1)) 82 | for op in binary_ops: 83 | funcs.append(GenericOperator(function=binary(op), name=op.__name__, arity=2)) 84 | for op in rolling_ops: 85 | for day in [10, 20, 30, 40, 50]: 86 | funcs.append(GenericOperator(function=rolling(op, day), name=op.__name__ + str(day), arity=1)) 87 | for op in rolling_binary_ops: 88 | for day in [10, 20, 30, 40, 50]: 89 | funcs.append(GenericOperator(function=rolling_binary(op, day), name=op.__name__ + str(day), arity=2)) 90 | -------------------------------------------------------------------------------- /alphagen_qlib/calculator.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple 2 | from torch import Tensor 3 | import torch 4 | from alphagen.data.calculator import AlphaCalculator 5 | from alphagen.data.expression import Expression 6 | from alphagen.utils.correlation import batch_pearsonr, batch_spearmanr 7 | from alphagen.utils.pytorch_utils import normalize_by_day 8 | from alphagen_qlib.stock_data import StockData 9 | 10 | 11 | class QLibStockDataCalculator(AlphaCalculator): 12 | def __init__(self, data: StockData, target: Optional[Expression]): 13 | self.data = data 14 | 15 | if target is None: # Combination-only mode 16 | self.target_value = None 17 | else: 18 | self.target_value = normalize_by_day(target.evaluate(self.data)) 19 | 20 | def _calc_alpha(self, expr: Expression) -> Tensor: 21 | return normalize_by_day(expr.evaluate(self.data)) 22 | 23 | def _calc_IC(self, value1: Tensor, value2: Tensor) -> float: 24 | return batch_pearsonr(value1, value2).mean().item() 25 | 26 | def _calc_rIC(self, value1: Tensor, value2: Tensor) -> float: 27 | # return batch_pearsonr(value1, value2).mean().item() 28 | def chunk_batch_spearmanr(x,y,chunk_size=300): 29 | n_days = len(x) 30 | spearmanr_list= [] 31 | cur_fct = 0 32 | for i in range(0,n_days,chunk_size): 33 | spearmanr_list.append(batch_spearmanr(x[i:i+chunk_size],y[i:i+chunk_size])) 34 | spearmanr_list = torch.cat(spearmanr_list,dim=0) 35 | return spearmanr_list 36 | return chunk_batch_spearmanr(value1, value2).mean().item() 37 | 38 | def _calc_real_rIC(self, value1: Tensor, value2: Tensor) -> float: 39 | # return batch_pearsonr(value1, value2).mean().item() 40 | return batch_spearmanr(value1, value2).mean().item() 41 | 42 | def _calc_real_rICIR(self, value1: Tensor, value2: Tensor) -> float: 43 | # return batch_pearsonr(value1, value2).mean().item() 44 | aaa = batch_spearmanr(value1, value2) 45 | RIC = aaa.mean().item() 46 | _std = aaa.std().item() 47 | return RIC/_std 48 | 49 | def make_ensemble_alpha(self, exprs: List[Expression], weights: List[float]) -> Tensor: 50 | n = len(exprs) 51 | factors: List[Tensor] = [self._calc_alpha(exprs[i]) * weights[i] for i in range(n)] 52 | return sum(factors) # type: ignore 53 | 54 | def calc_single_IC_ret(self, expr: Expression) -> float: 55 | value = self._calc_alpha(expr) 56 | return self._calc_IC(value, self.target_value) 57 | 58 | def calc_single_rIC_ret(self, expr: Expression) -> float: 59 | value = self._calc_alpha(expr) 60 | return self._calc_rIC(value, self.target_value) 61 | 62 | def calc_single_all_ret(self, expr: Expression) -> Tuple[float, float]: 63 | value = self._calc_alpha(expr) 64 | return self._calc_IC(value, self.target_value), self._calc_rIC(value, self.target_value) 65 | 66 | def calc_mutual_IC(self, expr1: Expression, expr2: Expression) -> float: 67 | value1, value2 = self._calc_alpha(expr1), self._calc_alpha(expr2) 68 | return self._calc_IC(value1, value2) 69 | 70 | def calc_pool_IC_ret(self, exprs: List[Expression], weights: List[float]) -> float: 71 | with torch.no_grad(): 72 | ensemble_value = self.make_ensemble_alpha(exprs, weights) 73 | return self._calc_IC(ensemble_value, self.target_value) 74 | 75 | def calc_pool_rIC_ret(self, exprs: List[Expression], weights: List[float]) -> float: 76 | with torch.no_grad(): 77 | ensemble_value = self.make_ensemble_alpha(exprs, weights) 78 | return self._calc_real_rIC(ensemble_value, self.target_value) 79 | 80 | def calc_pool_rICIR_ret(self, exprs: List[Expression], weights: List[float]) -> float: 81 | with torch.no_grad(): 82 | ensemble_value = self.make_ensemble_alpha(exprs, weights) 83 | return self._calc_real_rICIR(ensemble_value, self.target_value) 84 | 85 | def calc_pool_all_ret(self, exprs: List[Expression], weights: List[float]) -> Tuple[float, float]: 86 | with torch.no_grad(): 87 | ensemble_value = self.make_ensemble_alpha(exprs, weights) 88 | return self._calc_IC(ensemble_value, self.target_value), self._calc_rIC(ensemble_value, self.target_value) 89 | 90 | 91 | def calc_pool_all_ret_raw(self, exprs: List[Expression], weights: List[float]) -> Tuple[float, float]: 92 | with torch.no_grad(): 93 | ensemble_value = self.make_ensemble_alpha(exprs, weights) 94 | 95 | ic = batch_pearsonr(ensemble_value, self.target_value) 96 | ric = batch_spearmanr(ensemble_value, self.target_value) 97 | return ic,ric -------------------------------------------------------------------------------- /alphagen_qlib/stock_data.py: -------------------------------------------------------------------------------- 1 | from typing import List, Union, Optional, Tuple, Dict 2 | from enum import IntEnum 3 | import numpy as np 4 | import pandas as pd 5 | import torch 6 | 7 | class FeatureType(IntEnum): 8 | OPEN = 0 9 | CLOSE = 1 10 | HIGH = 2 11 | LOW = 3 12 | VOLUME = 4 13 | VWAP = 5 14 | 15 | def change_to_raw_min(features): 16 | result = [] 17 | for feature in features: 18 | if feature in ['$vwap']: 19 | result.append(f"$money/$volume") 20 | elif feature in ['$volume']: 21 | result.append(f"{feature}/100000") 22 | # result.append('$close') 23 | else: 24 | result.append(feature) 25 | return result 26 | 27 | def change_to_raw(features): 28 | result = [] 29 | for feature in features: 30 | if feature in ['$open','$close','$high','$low','$vwap']: 31 | result.append(f"{feature}*$factor") 32 | elif feature in ['$volume']: 33 | result.append(f"{feature}/$factor/1000000") 34 | # result.append('$close') 35 | else: 36 | raise ValueError(f"feature {feature} not supported") 37 | return result 38 | 39 | class StockData: 40 | _qlib_initialized: bool = False 41 | 42 | def __init__(self, 43 | instrument: Union[str, List[str]], 44 | start_time: str, 45 | end_time: str, 46 | max_backtrack_days: int = 100, 47 | max_future_days: int = 30, 48 | features: Optional[List[FeatureType]] = None, 49 | device: torch.device = torch.device('cuda:0'), 50 | raw:bool = False, 51 | qlib_path:Union[str,Dict] = "", 52 | freq:str = 'day', 53 | ) -> None: 54 | self._init_qlib(qlib_path) 55 | self.df_bak = None 56 | self.raw = raw 57 | self._instrument = instrument 58 | self.max_backtrack_days = max_backtrack_days 59 | self.max_future_days = max_future_days 60 | self._start_time = start_time 61 | self._end_time = end_time 62 | self._features = features if features is not None else list(FeatureType) 63 | self.device = device 64 | self.freq = freq 65 | self.data, self._dates, self._stock_ids = self._get_data() 66 | 67 | 68 | @classmethod 69 | def _init_qlib(cls,qlib_path) -> None: 70 | if cls._qlib_initialized: 71 | return 72 | import qlib 73 | from qlib.config import REG_CN 74 | qlib.init(provider_uri=qlib_path, region=REG_CN) 75 | cls._qlib_initialized = True 76 | 77 | def _load_exprs(self, exprs: Union[str, List[str]]) -> pd.DataFrame: 78 | # This evaluates an expression on the data and returns the dataframe 79 | # It might throw on illegal expressions like "Ref(constant, dtime)" 80 | from qlib.data.dataset.loader import QlibDataLoader 81 | from qlib.data import D 82 | if not isinstance(exprs, list): 83 | exprs = [exprs] 84 | cal: np.ndarray = D.calendar(freq=self.freq) 85 | start_index = cal.searchsorted(pd.Timestamp(self._start_time)) # type: ignore 86 | end_index = cal.searchsorted(pd.Timestamp(self._end_time)) # type: ignore 87 | real_start_time = cal[start_index - self.max_backtrack_days] 88 | if cal[end_index] != pd.Timestamp(self._end_time): 89 | end_index -= 1 90 | # real_end_time = cal[min(end_index + self.max_future_days,len(cal)-1)] 91 | real_end_time = cal[end_index + self.max_future_days] 92 | result = (QlibDataLoader(config=exprs,freq=self.freq) # type: ignore 93 | .load(self._instrument, real_start_time, real_end_time)) 94 | return result 95 | 96 | def _get_data(self) -> Tuple[torch.Tensor, pd.Index, pd.Index]: 97 | features = ['$' + f.name.lower() for f in self._features] 98 | if self.raw and self.freq == 'day': 99 | features = change_to_raw(features) 100 | elif self.raw: 101 | features = change_to_raw_min(features) 102 | df = self._load_exprs(features) 103 | self.df_bak = df 104 | # print(df) 105 | df = df.stack().unstack(level=1) 106 | dates = df.index.levels[0] # type: ignore 107 | stock_ids = df.columns 108 | values = df.values 109 | values = values.reshape((-1, len(features), values.shape[-1])) # type: ignore 110 | return torch.tensor(values, dtype=torch.float, device=self.device), dates, stock_ids 111 | 112 | @property 113 | def n_features(self) -> int: 114 | return len(self._features) 115 | 116 | @property 117 | def n_stocks(self) -> int: 118 | return self.data.shape[-1] 119 | 120 | @property 121 | def n_days(self) -> int: 122 | return self.data.shape[0] - self.max_backtrack_days - self.max_future_days 123 | 124 | def add_data(self,data:torch.Tensor,dates:pd.Index): 125 | data = data.to(self.device) 126 | self.data = torch.cat([self.data,data],dim=0) 127 | self._dates = pd.Index(self._dates.append(dates)) 128 | 129 | 130 | def make_dataframe( 131 | self, 132 | data: Union[torch.Tensor, List[torch.Tensor]], 133 | columns: Optional[List[str]] = None 134 | ) -> pd.DataFrame: 135 | """ 136 | Parameters: 137 | - `data`: a tensor of size `(n_days, n_stocks[, n_columns])`, or 138 | a list of tensors of size `(n_days, n_stocks)` 139 | - `columns`: an optional list of column names 140 | """ 141 | if isinstance(data, list): 142 | data = torch.stack(data, dim=2) 143 | if len(data.shape) == 2: 144 | data = data.unsqueeze(2) 145 | if columns is None: 146 | columns = [str(i) for i in range(data.shape[2])] 147 | n_days, n_stocks, n_columns = data.shape 148 | if self.n_days != n_days: 149 | raise ValueError(f"number of days in the provided tensor ({n_days}) doesn't " 150 | f"match that of the current StockData ({self.n_days})") 151 | if self.n_stocks != n_stocks: 152 | raise ValueError(f"number of stocks in the provided tensor ({n_stocks}) doesn't " 153 | f"match that of the current StockData ({self.n_stocks})") 154 | if len(columns) != n_columns: 155 | raise ValueError(f"size of columns ({len(columns)}) doesn't match with " 156 | f"tensor feature count ({data.shape[2]})") 157 | if self.max_future_days == 0: 158 | date_index = self._dates[self.max_backtrack_days:] 159 | else: 160 | date_index = self._dates[self.max_backtrack_days:-self.max_future_days] 161 | index = pd.MultiIndex.from_product([date_index, self._stock_ids]) 162 | data = data.reshape(-1, n_columns) 163 | return pd.DataFrame(data.detach().cpu().numpy(), index=index, columns=columns) 164 | 165 | -------------------------------------------------------------------------------- /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/DulyHao/AlphaForge/d0cfc27df23c60f271bc885fd43027b86b787746/dso/task/regression/__init__.py -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /exp_AFF_calc_result.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import torch \n", 10 | "from torch import nn\n", 11 | "import os\n", 12 | "os.environ[\"CUDA_VISIBLE_DEVICES\"]='1'\n", 13 | "# from alphagen.config import *\n", 14 | "# from alphagen.data.tokens import *\n", 15 | "from alphagen.models.alpha_pool import AlphaPoolBase, AlphaPool\n", 16 | "from alphagen.rl.env.core import AlphaEnvCore\n", 17 | "import torch.nn.functional as F\n", 18 | "from gan.dataset import Collector\n", 19 | "from gan.network.generater import NetG_DCGAN\n", 20 | "from gan.network.masker import NetM\n", 21 | "from gan.network.predictor import NetP, train_regression_model,train_regression_model_with_weight\n", 22 | "from alphagen.rl.env.wrapper import SIZE_ACTION,action2token\n", 23 | "\n", 24 | "from alphagen_generic.features import open_\n", 25 | "from gan.utils import Builders\n", 26 | "from alphagen_generic.features import *\n", 27 | "from alphagen.data.expression import *\n", 28 | "\n", 29 | "from gan.utils.data import get_data_by_year\n" 30 | ] 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": null, 35 | "metadata": {}, 36 | "outputs": [], 37 | "source": [ 38 | "\n", 39 | "instruments: str = \"csi300\"\n", 40 | "freq = 'day'\n", 41 | "save_name = 'test'\n", 42 | "window = float('inf')" 43 | ] 44 | }, 45 | { 46 | "cell_type": "code", 47 | "execution_count": null, 48 | "metadata": {}, 49 | "outputs": [], 50 | "source": [ 51 | "from alphagen.utils.correlation import batch_pearsonr,batch_spearmanr\n", 52 | "device = 'cuda:0'\n", 53 | "result = []\n", 54 | "pred_dfs = {}\n", 55 | "for n_factors in [1,10,20,50,100]:\n", 56 | " for seed in range(5):\n", 57 | " cur_seed_ic = []\n", 58 | " cur_seed_ric = []\n", 59 | " all_pred_df_list = []\n", 60 | " for train_end in range(2020,2021):\n", 61 | " print(n_factors,seed,train_end)\n", 62 | " returned = get_data_by_year(\n", 63 | " train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n", 64 | " instruments=instruments, target=target,freq=freq,\n", 65 | " )\n", 66 | " data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,name = returned\n", 67 | " \n", 68 | "\n", 69 | " path = f'out/{save_name}_{instruments}_{train_end}_{seed}/z_bld_zoo_final.pkl'\n", 70 | " tensor_save_path = f'out/{save_name}_{instruments}_{train_end}_{seed}/pred_{train_end}_{n_factors}_{window}_{seed}.pt'\n", 71 | "\n", 72 | " pred = torch.load(tensor_save_path).to(device)\n", 73 | " tgt = target.evaluate(data_all)\n", 74 | " \n", 75 | " \n", 76 | " ones = torch.ones_like(tgt)\n", 77 | " ones = ones * torch.nan\n", 78 | " ones[-data_test.n_days:] = pred\n", 79 | " cur_df = data_all.make_dataframe(ones)\n", 80 | " all_pred_df_list.append(cur_df.unstack().iloc[-data_test.n_days:].stack())\n", 81 | " \n", 82 | " tgt = tgt[-data_test.n_days:].to(device)\n", 83 | " \n", 84 | " \n", 85 | " ic_s = torch.nan_to_num(batch_pearsonr(pred,tgt),nan=0)\n", 86 | " rank_ic_s = torch.nan_to_num(batch_spearmanr(pred,tgt),nan=0)\n", 87 | "\n", 88 | " cur_seed_ic.append(ic_s)\n", 89 | " cur_seed_ric.append(rank_ic_s)\n", 90 | " \n", 91 | " pred_dfs[f\"{n_factors}_{seed}\"] = pd.concat(all_pred_df_list,axis=0)\n", 92 | " ic = torch.cat(cur_seed_ic)\n", 93 | " rank_ic = torch.cat(cur_seed_ric)\n", 94 | "\n", 95 | " ic_mean = ic.mean().item()\n", 96 | " rank_ic_mean = rank_ic.mean().item()\n", 97 | " ic_std = ic.std().item()\n", 98 | " rank_ic_std = rank_ic.std().item()\n", 99 | " tmp = dict(\n", 100 | " seed = seed,\n", 101 | " num = n_factors,\n", 102 | " ic = ic_mean,\n", 103 | " ric = rank_ic_mean,\n", 104 | " icir = ic_mean/ic_std,\n", 105 | " ricir = rank_ic_mean/rank_ic_std,\n", 106 | " )\n", 107 | " result.append(tmp)" 108 | ] 109 | }, 110 | { 111 | "cell_type": "code", 112 | "execution_count": null, 113 | "metadata": {}, 114 | "outputs": [], 115 | "source": [ 116 | "import pandas as pd\n", 117 | "run_result = pd.DataFrame(result).groupby(['num','seed']).mean().groupby('num').agg(['mean','std'])\n", 118 | "print(run_result)" 119 | ] 120 | } 121 | ], 122 | "metadata": { 123 | "kernelspec": { 124 | "display_name": "py38n1", 125 | "language": "python", 126 | "name": "python3" 127 | }, 128 | "language_info": { 129 | "codemirror_mode": { 130 | "name": "ipython", 131 | "version": 3 132 | }, 133 | "file_extension": ".py", 134 | "mimetype": "text/x-python", 135 | "name": "python", 136 | "nbconvert_exporter": "python", 137 | "pygments_lexer": "ipython3", 138 | "version": "3.8.16" 139 | } 140 | }, 141 | "nbformat": 4, 142 | "nbformat_minor": 2 143 | } 144 | -------------------------------------------------------------------------------- /exp_DSO_calc_result.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 2, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import torch \n", 10 | "# from alphagen.config import *\n", 11 | "# from alphagen.data.tokens import *\n", 12 | "from alphagen_generic.features import *\n", 13 | "from alphagen.data.expression import *\n", 14 | "from gan.utils.data import get_data_by_year" 15 | ] 16 | }, 17 | { 18 | "cell_type": "code", 19 | "execution_count": 3, 20 | "metadata": {}, 21 | "outputs": [], 22 | "source": [ 23 | "\n", 24 | "\n", 25 | "instruments: str = \"csi300\"\n", 26 | "from typing import Tuple\n", 27 | "import json\n", 28 | "\n", 29 | "def load_alpha_pool(raw) -> Tuple[List[Expression], List[float]]:\n", 30 | " exprs_raw = raw['exprs']\n", 31 | " exprs = [eval(expr_raw.replace('open', 'open_').replace('$', '')) for expr_raw in exprs_raw]\n", 32 | " weights = raw['weights']\n", 33 | " return exprs, weights\n", 34 | "\n", 35 | "def load_alpha_pool_by_path(path: str) -> Tuple[List[Expression], List[float]]:\n", 36 | " with open(path, encoding='utf-8') as f:\n", 37 | " raw = json.load(f)\n", 38 | " return load_alpha_pool(raw)\n", 39 | " \n", 40 | "import os\n", 41 | "def load_ppo_path(path,name_prefix):\n", 42 | " \n", 43 | " files = os.listdir(path)\n", 44 | " folder = [i for i in files if name_prefix in i][-1]\n", 45 | " names = [i for i in os.listdir(f\"{path}/{folder}\") if '.json' in i]\n", 46 | " name = sorted(names,key = lambda x:int(x.split('_')[0]))[-1]\n", 47 | " return f\"{path}/{folder}/{name}\"" 48 | ] 49 | }, 50 | { 51 | "cell_type": "markdown", 52 | "metadata": {}, 53 | "source": [ 54 | "# Infer" 55 | ] 56 | }, 57 | { 58 | "cell_type": "code", 59 | "execution_count": null, 60 | "metadata": {}, 61 | "outputs": [], 62 | "source": [ 63 | "from alphagen_qlib.calculator import QLibStockDataCalculator\n", 64 | "result = []\n", 65 | "paths = os.listdir('out_dso')\n", 66 | "name = 'test1'\n", 67 | "freq = 'day'\n", 68 | "for instruments in ['csi300','csi500']:\n", 69 | " for num in [1,10,20,50,100]:\n", 70 | " for seed in range(5):\n", 71 | " for train_end in range(2016,2021):\n", 72 | " returned = get_data_by_year(\n", 73 | " train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n", 74 | " instruments=instruments, target=target,freq=freq,\n", 75 | " )\n", 76 | " data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,_ = returned\n", 77 | " \n", 78 | " path = f\"out_dso/{name}_{instruments}_{num}_{train_end}_{seed}/pool.json\"\n", 79 | " dirname = os.path.dirname(path)\n", 80 | " print(path)\n", 81 | " exprs,weights = load_alpha_pool_by_path(path)\n", 82 | " \n", 83 | " # calculator_test = QLibStockDataCalculator(data_test, target)\n", 84 | " calculator_test = QLibStockDataCalculator(data_all, target)\n", 85 | "\n", 86 | " ensemble_value = calculator_test.make_ensemble_alpha(exprs, weights)\n", 87 | " ensemble_value = ensemble_value[-data_test.n_days:]\n", 88 | " \n", 89 | " torch.save(ensemble_value.cpu(),f\"{dirname}/{train_end}_{num}_{seed}.pt\")" 90 | ] 91 | }, 92 | { 93 | "cell_type": "markdown", 94 | "metadata": {}, 95 | "source": [ 96 | "# Read Experiment Result and Calculate Metrics" 97 | ] 98 | }, 99 | { 100 | "cell_type": "code", 101 | "execution_count": null, 102 | "metadata": {}, 103 | "outputs": [], 104 | "source": [ 105 | "\n", 106 | "\n", 107 | "from alphagen.utils.correlation import batch_pearsonr,batch_spearmanr\n", 108 | "device = 'cuda:0'\n", 109 | "name = 'test1'\n", 110 | "instruments = 'csi300'\n", 111 | "result = []\n", 112 | "for seed in range(5):\n", 113 | " cur_seed_ic = []\n", 114 | " cur_seed_ric = []\n", 115 | " \n", 116 | " for num in [1,10,20,50,100]:\n", 117 | " for train_end in range(2016,2021):\n", 118 | " returned = get_data_by_year(\n", 119 | " train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n", 120 | " instruments=instruments, target=target,freq=freq,\n", 121 | " )\n", 122 | " data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,_ = returned\n", 123 | "\n", 124 | " dirname = f\"out_dso/{name}_{instruments}_{num}_{train_end}_{seed}\"\n", 125 | " \n", 126 | " pred = torch.load(f\"{dirname}/{train_end}_{num}_{seed}.pt\").to('cuda:0')\n", 127 | " tgt = target.evaluate(data_test)\n", 128 | " tgt = target.evaluate(data_all)[-data_test.n_days:,:]\n", 129 | "\n", 130 | " ic_s = torch.nan_to_num(batch_pearsonr(pred,tgt),nan=0)\n", 131 | " rank_ic_s = torch.nan_to_num(batch_spearmanr(pred,tgt),nan=0)\n", 132 | "\n", 133 | " cur_seed_ic.append(ic_s)\n", 134 | " cur_seed_ric.append(rank_ic_s)\n", 135 | " ic = torch.cat(cur_seed_ic)\n", 136 | " rank_ic = torch.cat(cur_seed_ric)\n", 137 | " \n", 138 | " ic_mean = ic.mean().item()\n", 139 | " rank_ic_mean = rank_ic.mean().item()\n", 140 | " ic_std = ic.std().item()\n", 141 | " rank_ic_std = rank_ic.std().item()\n", 142 | " tmp = dict(\n", 143 | " seed = seed,\n", 144 | " num = num,\n", 145 | " ic = ic_mean,\n", 146 | " ric = rank_ic_mean,\n", 147 | " icir = ic_mean/ic_std,\n", 148 | " ricir = rank_ic_mean/rank_ic_std,\n", 149 | " )\n", 150 | " result.append(tmp)\n", 151 | "\n", 152 | "import pandas as pd\n", 153 | "exp_result = pd.DataFrame(result).groupby(['num','seed']).mean().groupby('num').agg(['mean','std'])\n", 154 | "print(exp_result)\n", 155 | " " 156 | ] 157 | }, 158 | { 159 | "cell_type": "code", 160 | "execution_count": null, 161 | "metadata": {}, 162 | "outputs": [], 163 | "source": [] 164 | } 165 | ], 166 | "metadata": { 167 | "kernelspec": { 168 | "display_name": "py38n1", 169 | "language": "python", 170 | "name": "python3" 171 | }, 172 | "language_info": { 173 | "codemirror_mode": { 174 | "name": "ipython", 175 | "version": 3 176 | }, 177 | "file_extension": ".py", 178 | "mimetype": "text/x-python", 179 | "name": "python", 180 | "nbconvert_exporter": "python", 181 | "pygments_lexer": "ipython3", 182 | "version": "3.8.16" 183 | } 184 | }, 185 | "nbformat": 4, 186 | "nbformat_minor": 2 187 | } 188 | -------------------------------------------------------------------------------- /exp_GP_calc_result.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "\n", 10 | "import os\n", 11 | "import argparse\n", 12 | "\n", 13 | "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '0'\n", 14 | "instruments = 'csi300'\n" 15 | ] 16 | }, 17 | { 18 | "cell_type": "code", 19 | "execution_count": null, 20 | "metadata": {}, 21 | "outputs": [], 22 | "source": [ 23 | "import json\n", 24 | "from collections import Counter\n", 25 | "from alphagen.data.expression import *\n", 26 | "from alphagen.models.alpha_pool import AlphaPool\n", 27 | "from alphagen.utils.correlation import batch_pearsonr, batch_spearmanr\n", 28 | "from alphagen_generic.features import *\n", 29 | "from gan.utils.data import get_data_by_year\n", 30 | "\n", 31 | "\n", 32 | "def pred_pool(capacity,data):\n", 33 | " from alphagen_qlib.calculator import QLibStockDataCalculator\n", 34 | " pool = AlphaPool(capacity=capacity,\n", 35 | " stock_data=data,\n", 36 | " target=target,\n", 37 | " ic_lower_bound=None)\n", 38 | " exprs = []\n", 39 | " for key in dict(Counter(cache).most_common(capacity)):\n", 40 | " exprs.append(eval(key))\n", 41 | " pool.force_load_exprs(exprs)\n", 42 | " pool._optimize(alpha=5e-3, lr=5e-4, n_iter=2000)\n", 43 | "\n", 44 | " exprs = pool.exprs[:pool.size]\n", 45 | " weights = pool.weights[:pool.size]\n", 46 | " calculator_test = QLibStockDataCalculator(data, target)\n", 47 | " ensemble_value = calculator_test.make_ensemble_alpha(exprs, weights)\n", 48 | " return ensemble_value\n", 49 | "\n", 50 | "\n" 51 | ] 52 | }, 53 | { 54 | "cell_type": "markdown", 55 | "metadata": {}, 56 | "source": [ 57 | "# Infer" 58 | ] 59 | }, 60 | { 61 | "cell_type": "code", 62 | "execution_count": null, 63 | "metadata": {}, 64 | "outputs": [], 65 | "source": [ 66 | "for seed in range(5):\n", 67 | " for train_end in range(2016,2021):\n", 68 | " for num in [1,10,20,50]:\n", 69 | " save_dir = f'out_gp/{instruments}_{train_end}_day_{seed}' \n", 70 | " print(save_dir)\n", 71 | " \n", 72 | " returned = get_data_by_year(\n", 73 | " train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n", 74 | " instruments=instruments, target=target,freq='day',\n", 75 | " )\n", 76 | " data_all,data,data_valid,data_valid_withhead,data_test,data_test_withhead,name = returned\n", 77 | "\n", 78 | " cache = json.load(open(f'{save_dir}/40.json'))['cache']\n", 79 | "\n", 80 | " features = ['open_', 'close', 'high', 'low', 'volume', 'vwap']\n", 81 | " 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.]]\n", 82 | " terminals = features + constants\n", 83 | "\n", 84 | " pred = pred_pool(num,data)\n", 85 | " pred = pred[-data_test.n_days:]\n", 86 | " torch.save(pred.detach().cpu(),f\"{save_dir}/pred_{num}.pt\")\n", 87 | " \n" 88 | ] 89 | }, 90 | { 91 | "cell_type": "markdown", 92 | "metadata": {}, 93 | "source": [ 94 | "# Read and combine result to show" 95 | ] 96 | }, 97 | { 98 | "cell_type": "code", 99 | "execution_count": null, 100 | "metadata": {}, 101 | "outputs": [], 102 | "source": [ 103 | "result = []\n", 104 | "for num in [1]:\n", 105 | " for seed in range(5):\n", 106 | " \n", 107 | " cur_seed_ic = []\n", 108 | " cur_seed_ric = []\n", 109 | " for train_end in range(2016,2021):\n", 110 | " #'/path/to/save/results'\n", 111 | " save_dir = f'out_gp/{instruments}_{train_end}_day_{seed}' \n", 112 | "\n", 113 | " returned = get_data_by_year(\n", 114 | " train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n", 115 | " instruments=instruments, target=target,freq='day',\n", 116 | " )\n", 117 | " data_all,data,data_valid,data_valid_withhead,data_test,data_test_withhead,name = returned\n", 118 | "\n", 119 | " pred = torch.load(f\"{save_dir}/pred_{num}.pt\").to('cuda:0')\n", 120 | " \n", 121 | " tgt = target.evaluate(data_test)\n", 122 | " tgt = target.evaluate(data_all)[-data_test.n_days:,:]\n", 123 | "\n", 124 | " ic_s = torch.nan_to_num(batch_pearsonr(pred,tgt),nan=0)\n", 125 | " rank_ic_s = torch.nan_to_num(batch_spearmanr(pred,tgt),nan=0)\n", 126 | "\n", 127 | " cur_seed_ic.append(ic_s)\n", 128 | " cur_seed_ric.append(rank_ic_s)\n", 129 | " \n", 130 | " ic = torch.cat(cur_seed_ic)\n", 131 | " rank_ic = torch.cat(cur_seed_ric)\n", 132 | "\n", 133 | " ic_mean = ic.mean().item()\n", 134 | " rank_ic_mean = rank_ic.mean().item()\n", 135 | " ic_std = ic.std().item()\n", 136 | " rank_ic_std = rank_ic.std().item()\n", 137 | " tmp = dict(\n", 138 | " seed = seed,\n", 139 | " num = num,\n", 140 | " ic = ic_mean,\n", 141 | " ric = rank_ic_mean,\n", 142 | " icir = ic_mean/ic_std,\n", 143 | " ricir = rank_ic_mean/rank_ic_std,\n", 144 | " )\n", 145 | " result.append(tmp)\n", 146 | " \n", 147 | "import pandas as pd\n", 148 | "print(pd.DataFrame(result).groupby(['num','seed']).mean().groupby('num').agg(['mean','std']))" 149 | ] 150 | } 151 | ], 152 | "metadata": { 153 | "language_info": { 154 | "name": "python" 155 | } 156 | }, 157 | "nbformat": 4, 158 | "nbformat_minor": 2 159 | } 160 | -------------------------------------------------------------------------------- /exp_RL_calc_result.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import torch \n", 10 | "from torch import nn\n", 11 | "import os\n", 12 | "os.environ[\"CUDA_VISIBLE_DEVICES\"]='1'\n", 13 | "# from alphagen.config import *\n", 14 | "# from alphagen.data.tokens import *\n", 15 | "from alphagen.models.alpha_pool import AlphaPoolBase, AlphaPool\n", 16 | "from alphagen.rl.env.core import AlphaEnvCore\n", 17 | "import torch.nn.functional as F\n", 18 | "from gan.dataset import Collector\n", 19 | "from gan.network.generater import NetG_DCGAN\n", 20 | "from gan.network.masker import NetM\n", 21 | "from gan.network.predictor import NetP, train_regression_model,train_regression_model_with_weight\n", 22 | "from alphagen.rl.env.wrapper import SIZE_ACTION,action2token\n", 23 | "from alphagen_generic.features import open_\n", 24 | "from gan.utils import Builders\n", 25 | "from alphagen_generic.features import *\n", 26 | "from alphagen.data.expression import *\n", 27 | "from gan.utils.data import get_data_by_year\n", 28 | "import os\n", 29 | "\n", 30 | "\n", 31 | "\n", 32 | "\n", 33 | "instruments: str = \"csi300\"\n", 34 | "from typing import Tuple\n", 35 | "import json\n", 36 | "\n", 37 | "def load_alpha_pool(raw) -> Tuple[List[Expression], List[float]]:\n", 38 | " exprs_raw = raw['exprs']\n", 39 | " exprs = [eval(expr_raw.replace('open', 'open_').replace('$', '')) for expr_raw in exprs_raw]\n", 40 | " weights = raw['weights']\n", 41 | " return exprs, weights\n", 42 | "\n", 43 | "def load_alpha_pool_by_path(path: str) -> Tuple[List[Expression], List[float]]:\n", 44 | " with open(path, encoding='utf-8') as f:\n", 45 | " raw = json.load(f)\n", 46 | " return load_alpha_pool(raw)\n", 47 | " \n", 48 | "import os\n", 49 | "def load_ppo_path(path,name_prefix):\n", 50 | " \n", 51 | " files = os.listdir(path)\n", 52 | " folder = [i for i in files if name_prefix in i][-1]\n", 53 | " names = [i for i in os.listdir(f\"{path}/{folder}\") if '.json' in i]\n", 54 | " name = sorted(names,key = lambda x:int(x.split('_')[0]))[-1]\n", 55 | " return f\"{path}/{folder}/{name}\"" 56 | ] 57 | }, 58 | { 59 | "cell_type": "markdown", 60 | "metadata": {}, 61 | "source": [ 62 | "# infer" 63 | ] 64 | }, 65 | { 66 | "cell_type": "code", 67 | "execution_count": null, 68 | "metadata": {}, 69 | "outputs": [], 70 | "source": [ 71 | "\n", 72 | "from alphagen_qlib.calculator import QLibStockDataCalculator\n", 73 | "freq = 'day'\n", 74 | "chk_path = \"out_ppo/checkpoints\"\n", 75 | "result = []\n", 76 | "for train_end in range(2016,2021):\n", 77 | " returned = get_data_by_year(\n", 78 | " train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n", 79 | " instruments=instruments, target=target,freq='day',\n", 80 | " )\n", 81 | " data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,name = returned\n", 82 | " for seed in range(5):\n", 83 | " for num in [1,10,20,50,100]:\n", 84 | " name_prefix = f\"csi300_{train_end}_{num}_{seed}\"\n", 85 | " path = load_ppo_path(chk_path,name_prefix)\n", 86 | " \n", 87 | " exprs,weights = load_alpha_pool_by_path(path)\n", 88 | " \n", 89 | " # calculator_test = QLibStockDataCalculator(data_test, target)\n", 90 | " calculator_test = QLibStockDataCalculator(data_all, target)\n", 91 | "\n", 92 | " ensemble_value = calculator_test.make_ensemble_alpha(exprs, weights)\n", 93 | " ensemble_value = ensemble_value[-data_test.n_days:]\n", 94 | " dirname = os.path.dirname(path)\n", 95 | " \n", 96 | " torch.save(ensemble_value.cpu(),f\"{dirname}/{train_end}_{num}_{seed}.pkl\")" 97 | ] 98 | }, 99 | { 100 | "cell_type": "markdown", 101 | "metadata": {}, 102 | "source": [ 103 | "# read the infer result and evaluate" 104 | ] 105 | }, 106 | { 107 | "cell_type": "code", 108 | "execution_count": null, 109 | "metadata": {}, 110 | "outputs": [], 111 | "source": [ 112 | "\n", 113 | "\n", 114 | "from alphagen.utils.correlation import batch_pearsonr,batch_spearmanr\n", 115 | "device = 'cuda:0'\n", 116 | "result = []\n", 117 | "for seed in range(5):\n", 118 | " cur_seed_ic = []\n", 119 | " cur_seed_ric = []\n", 120 | " \n", 121 | " for num in [50,100]:\n", 122 | " for train_end in range(2016,2021):\n", 123 | " returned = get_data_by_year(\n", 124 | " train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2,\n", 125 | " instruments=instruments, target=target,freq=freq,\n", 126 | " )\n", 127 | " data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,name = returned\n", 128 | "\n", 129 | " \n", 130 | " name_prefix = f\"n1230day_csi500_{train_end}_{num}_{seed}\"\n", 131 | " path = load_ppo_path(chk_path,name_prefix)\n", 132 | " dirname = os.path.dirname(path)\n", 133 | " pred = torch.load(f\"{dirname}/{train_end}_{num}_{seed}.pkl\").to(device)\n", 134 | " tgt = target.evaluate(data_test)\n", 135 | " tgt = target.evaluate(data_all)[-data_test.n_days:,:]\n", 136 | "\n", 137 | " ic_s = torch.nan_to_num(batch_pearsonr(pred,tgt),nan=0)\n", 138 | " rank_ic_s = torch.nan_to_num(batch_spearmanr(pred,tgt),nan=0)\n", 139 | "\n", 140 | " cur_seed_ic.append(ic_s)\n", 141 | " cur_seed_ric.append(rank_ic_s)\n", 142 | " ic = torch.cat(cur_seed_ic)\n", 143 | " rank_ic = torch.cat(cur_seed_ric)\n", 144 | "\n", 145 | " ic_mean = ic.mean().item()\n", 146 | " rank_ic_mean = rank_ic.mean().item()\n", 147 | " ic_std = ic.std().item()\n", 148 | " rank_ic_std = rank_ic.std().item()\n", 149 | " tmp = dict(\n", 150 | " seed = seed,\n", 151 | " num = num,\n", 152 | " ic = ic_mean,\n", 153 | " ric = rank_ic_mean,\n", 154 | " icir = ic_mean/ic_std,\n", 155 | " ricir = rank_ic_mean/rank_ic_std,\n", 156 | " )\n", 157 | " result.append(tmp)\n", 158 | "\n", 159 | "import pandas as pd\n", 160 | "exp_result = pd.DataFrame(result).groupby(['num','seed']).mean().groupby('num').agg(['mean','std'])\n", 161 | "print(exp_result)\n", 162 | " " 163 | ] 164 | } 165 | ], 166 | "metadata": { 167 | "kernelspec": { 168 | "display_name": "py38n1", 169 | "language": "python", 170 | "name": "python3" 171 | }, 172 | "language_info": { 173 | "codemirror_mode": { 174 | "name": "ipython", 175 | "version": 3 176 | }, 177 | "file_extension": ".py", 178 | "mimetype": "text/x-python", 179 | "name": "python", 180 | "nbconvert_exporter": "python", 181 | "pygments_lexer": "ipython3", 182 | "version": "3.8.16" 183 | } 184 | }, 185 | "nbformat": 4, 186 | "nbformat_minor": 2 187 | } 188 | -------------------------------------------------------------------------------- /gan/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DulyHao/AlphaForge/d0cfc27df23c60f271bc885fd43027b86b787746/gan/__init__.py -------------------------------------------------------------------------------- /gan/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from .collector import * 2 | -------------------------------------------------------------------------------- /gan/dataset/collector.py: -------------------------------------------------------------------------------- 1 | from gan.utils import Builders 2 | import torch 3 | 4 | class Collector: 5 | def __init__(self,seq_len ,n_actions): 6 | super().__init__() 7 | self.blds = Builders(0,max_len=seq_len,n_actions=n_actions) 8 | self.blds_bak = Builders(0,max_len=seq_len,n_actions=n_actions) 9 | self.seq_len = seq_len 10 | self.n_actions = n_actions 11 | 12 | def reset(self,data,target,metric): 13 | self.blds_bak += self.blds 14 | print('Reset bak_len:',self.blds_bak.batch_size) 15 | self.blds_bak.evaluate(data,target,metric) 16 | self.blds = Builders(0,max_len=self.seq_len,n_actions=self.n_actions) 17 | 18 | def collect(self,netG,netM,z,reset_net = False,random_method=None): 19 | netG.eval() 20 | with torch.no_grad(): 21 | if reset_net: 22 | netG.initialize_parameters() 23 | netG.eval() 24 | z = random_method(z) 25 | logit_raw = netG(z) 26 | masked_x,masks,blds= netM(logit_raw) 27 | return blds 28 | 29 | def collect_randomly(self,z,netM): 30 | logit_raw = torch.randn([z.shape[0],self.seq_len,self.n_actions]) 31 | print('logit_raw',logit_raw.shape) 32 | masked_x,masks,blds= netM(logit_raw) 33 | return blds 34 | 35 | 36 | def collect_target_num(self,netG,netM,z, 37 | data,target,metric = None, 38 | target_num=1000,reset_net =False, 39 | drop_invalid=False, 40 | randomly = False, 41 | random_method = lambda x:x.normal_(), 42 | max_iter = 1000, 43 | ): 44 | 45 | cnt = 0 46 | iter_num = 0 47 | while self.blds.batch_size<=target_num: 48 | if randomly: 49 | builders = self.collect_randomly(z,netM) 50 | else: 51 | builders = self.collect(netG,netM,z,reset_net=reset_net,random_method=random_method) 52 | if drop_invalid: 53 | builders.drop_invalid() 54 | self.blds += builders 55 | self.blds.drop_duplicated() 56 | cnt += 1 57 | iter_num += 1 58 | if iter_num%10==0: 59 | print(f"cnt:{cnt} builders_len:{builders.batch_size},all_len:{self.blds.batch_size}") 60 | if iter_num>max_iter and max_iter>0: 61 | print(f'iter_num>max_iter:{max_iter}') 62 | break 63 | 64 | self.blds.drop_duplicated() 65 | print(self.blds.batch_size) 66 | self.blds.evaluate(data,target,metric) 67 | return 68 | @property 69 | def blds_list(self): 70 | return [self.blds_bak,self.blds] 71 | 72 | 73 | 74 | 75 | # class Collector_2: 76 | # def __init__(self): 77 | # super().__init__() 78 | # self.blds = Builders(0) 79 | # self.blds_bak = Builders(0) 80 | 81 | # def reset(self,data,target,metric): 82 | # self.blds_bak += self.blds 83 | # print('Reset bak_len:',self.blds_bak.batch_size) 84 | # self.blds_bak.evaluate(data,target,metric) 85 | # self.blds = Builders(0) 86 | 87 | # def collect(self,netG,netM,z,reset_net = False,random_method=None): 88 | # with torch.no_grad(): 89 | # if reset_net: 90 | # netG.initialize_parameters() 91 | # netG.eval() 92 | # random_method(z) 93 | # logit_raw = netG(z) 94 | # masked_x,masks,blds= netM(logit_raw) 95 | # return blds 96 | 97 | # def collect_randomly(self,z,netM): 98 | # logit_raw = torch.randn([z.shape[0],self.seq_len,self.n_actions]) 99 | # masked_x,masks,blds= netM(logit_raw) 100 | # return blds 101 | # def collect_target_num(self,netG,netM,z,data,target,target_num=1000,reset_net =False,drop_invalid=False, 102 | # random_method = lambda x:x.normal_(),metric = None,randomly=False): 103 | 104 | # cnt = 0 105 | 106 | # iter_num = 0 107 | # while self.blds.batch_size<=target_num: 108 | # if randomly: 109 | # builders = self.collect_randomly(z,netM) 110 | # else: 111 | # builders = self.collect(netG,netM,z,reset_net=reset_net,random_method=random_method) 112 | # if drop_invalid: 113 | # builders.drop_invalid() 114 | # self.blds += builders 115 | # self.blds.drop_duplicated() 116 | # cnt += 1 117 | # iter_num += 1 118 | # if iter_num%10==0: 119 | # print(f"cnt:{cnt} builders_len:{builders.batch_size},all_len:{self.blds.batch_size}") 120 | # if iter_num>100: 121 | # print('iter_num>100') 122 | # break 123 | 124 | # self.blds.drop_duplicated() 125 | # print(self.blds.batch_size) 126 | # self.blds.evaluate(data,target,metric) 127 | # return 128 | # @property 129 | # def blds_list(self): 130 | # return [self.blds_bak,self.blds] 131 | -------------------------------------------------------------------------------- /gan/network/__init__.py: -------------------------------------------------------------------------------- 1 | from .generater import * 2 | from .masker import * 3 | from .predictor import * 4 | -------------------------------------------------------------------------------- /gan/network/loss.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | 4 | def loss_simi(loss_inputs,cfg): 5 | onehot_tensor_1 = loss_inputs['onehot_tensor_1'] #（batch_size,MAX_EXPR_LENGTH,SIZE_ACTION） 6 | onehot_tensor_2 = loss_inputs['onehot_tensor_2'] 7 | 8 | simi = torch.sum(onehot_tensor_1*onehot_tensor_2,dim=-1).sum(dim = -1) # (batch_size,) 9 | simi = simi / onehot_tensor_1.shape[1] 10 | 11 | simi = simi - cfg.l_simi_thresh 12 | simi = torch.relu(simi) 13 | # simi = simi**2 14 | simi = simi.mean() 15 | return simi 16 | 17 | def loss_pred(loss_inputs,cfg): 18 | pred_1 = loss_inputs['pred_1'][:,0] #（batch_size） 19 | 20 | return - pred_1.mean() 21 | 22 | def loss_potential(loss_inputs,cfg): 23 | 24 | 25 | epsilong=cfg.l_potential_epsilon 26 | u1,u2=loss_inputs['latent_1'],loss_inputs['latent_2']#(batch_size*n_sample,latent_size_netP) 27 | u1=u1.clip(epsilong,1-epsilong) 28 | u2=u2.clip(epsilong,1-epsilong) 29 | similarity=(u1*u2).sum(axis=1)/ ( 30 | ((u1**2).sum(axis=1))**0.5 * ((u2**2).sum(axis=1))**0.5 31 | ) -cfg.l_potential_thresh 32 | 33 | similarity=similarity*(similarity>0)# 针对每个元素选择大于0的 34 | return similarity.mean() 35 | 36 | def loss_entropy(loss_inputs,cfg): 37 | onehot_tensor_1 = loss_inputs['onehot_tensor_1'] #（batch_size,MAX_EXPR_LENGTH,SIZE_ACTION） 38 | onehot_tensor_2 = loss_inputs['onehot_tensor_2'] 39 | 40 | entropy_1 = -torch.sum(onehot_tensor_1*torch.log(onehot_tensor_1),dim=-1).sum(dim = -1) # (batch_size,) 41 | entropy_1 = entropy_1 / onehot_tensor_1.shape[1] 42 | 43 | entropy_2 = -torch.sum(onehot_tensor_2*torch.log(onehot_tensor_2),dim=-1).sum(dim = -1) # (batch_size,) 44 | entropy_2 = entropy_2 / onehot_tensor_2.shape[1] 45 | 46 | entropy = entropy_1 + entropy_2 47 | entropy = entropy.mean() 48 | return entropy 49 | def get_losses(loss_inputs,cfg): 50 | 51 | loss = 0 52 | if cfg.l_simi != 0 : 53 | loss += cfg.l_simi * loss_simi(loss_inputs,cfg) 54 | if cfg.l_pred != 0 : 55 | loss += cfg.l_pred * loss_pred(loss_inputs,cfg) 56 | if cfg.l_potential !=0 : 57 | loss += cfg.l_potential * loss_potential(loss_inputs,cfg) 58 | if cfg.l_entropy !=0 : 59 | loss += cfg.l_entropy * loss_entropy(loss_inputs,cfg) 60 | 61 | return loss -------------------------------------------------------------------------------- /gan/network/masker.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | import numpy as np 6 | 7 | from gan.utils.builder import Builders 8 | 9 | 10 | class NetM(nn.Module): 11 | def __init__( 12 | self, 13 | max_len = 20, 14 | size_action = 48, 15 | 16 | ): 17 | super().__init__() 18 | self.max_len = max_len 19 | self.size_action = size_action 20 | 21 | def forward(self, x: torch.Tensor): 22 | # x: (bs, seq_len, 48) 23 | 24 | device = x.device 25 | bs,seq_len,n_actions = x.shape 26 | blds = Builders(bs,max_len=seq_len,n_actions=n_actions) 27 | 28 | masks = torch.zeros(bs,seq_len,n_actions).to(device) 29 | masked_x = torch.zeros(bs,seq_len,n_actions).to(device) 30 | # prev_select = None 31 | for i in range(seq_len): 32 | 33 | # get masks 34 | if i<=self.max_len: 35 | mask = blds.get_valid_op()# (bs, n_actions) 36 | else: 37 | mask = np.zeros([bs, n_actions],dtype=bool) 38 | mask[:,n_actions-1] = True # the last one is sep 39 | 40 | mask_tensor = torch.from_numpy(mask).to(device) 41 | masks[:,i,:] = mask_tensor 42 | 43 | # get onehot and push to builders 44 | logit = x[:,i,:] 45 | tmp = logit.detach().cpu().numpy()# (bs, n_actions) 46 | tmp[~mask] = -1e8 47 | select = tmp.argmax(axis=1)# (bs,) 48 | # prev_select = select 49 | assert (mask[:,select]*1.).mean() 50 | blds.add_token(select) 51 | 52 | # get masked_x 53 | masked_x[:,i,:] = x[:,i,:] 54 | masked_x[:,i,:][~mask] = -1e8 55 | 56 | 57 | return masked_x,masks,blds 58 | -------------------------------------------------------------------------------- /gan/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import * 2 | -------------------------------------------------------------------------------- /gan/utils/data.py: -------------------------------------------------------------------------------- 1 | from alphagen_generic.features import open_ 2 | from gan.utils import Builders 3 | from alphagen_generic.features import * 4 | from alphagen.data.expression import * 5 | 6 | import os 7 | def get_data_by_year( 8 | train_start = 2010,train_end=2019,valid_year=2020,test_year =2021, 9 | instruments=None, target=None,freq=None, 10 | ): 11 | QLIB_PATH = { 12 | 'day':'path/for/qlib', 13 | } 14 | 15 | from gan.utils import load_pickle,save_pickle 16 | # from gan.utils.qlib import get_data_my 17 | get_data_my = StockData 18 | 19 | train_dates=(f"{train_start}-01-01", f"{train_end}-12-31") 20 | val_dates=(f"{valid_year}-01-01", f"{valid_year}-12-31") 21 | test_dates=(f"{test_year}-01-01", f"{test_year}-12-31") 22 | 23 | train_start,train_end = train_dates 24 | valid_start,valid_end = val_dates 25 | valid_head_start = f"{valid_year-2}-01-01" 26 | test_start,test_end = test_dates 27 | test_head_start = f"{test_year-2}-01-01" 28 | 29 | name = instruments + '_pkl_' + str(target).replace('/','_').replace(' ','') + '_' + freq 30 | name = f"{name}_{train_start}_{train_end}_{valid_start}_{valid_end}_{test_start}_{test_end}" 31 | try: 32 | 33 | data = load_pickle(f'pkl/{name}/data.pkl') 34 | data_valid = load_pickle(f'pkl/{name}/data_valid.pkl') 35 | data_valid_withhead = load_pickle(f'pkl/{name}/data_valid_withhead.pkl') 36 | data_test = load_pickle(f'pkl/{name}/data_test.pkl') 37 | data_test_withhead = load_pickle(f'pkl/{name}/data_test_withhead.pkl') 38 | 39 | except: 40 | print('Data not exist, load from qlib') 41 | data = get_data_my(instruments, train_start, train_end,raw = True,qlib_path = QLIB_PATH,freq=freq) 42 | data_valid = get_data_my(instruments, valid_start, valid_end,raw = True,qlib_path = QLIB_PATH,freq=freq) 43 | data_valid_withhead = get_data_my(instruments,valid_head_start, valid_end,raw = True,qlib_path = QLIB_PATH,freq=freq) 44 | data_test = get_data_my(instruments, test_start, test_end,raw = True,qlib_path = QLIB_PATH,freq=freq) 45 | data_test_withhead = get_data_my(instruments, test_head_start, test_end,raw = True,qlib_path = QLIB_PATH,freq=freq) 46 | 47 | os.makedirs(f"pkl/{name}",exist_ok=True) 48 | save_pickle(data,f'pkl/{name}/data.pkl') 49 | save_pickle(data_valid,f'pkl/{name}/data_valid.pkl') 50 | save_pickle(data_valid_withhead,f'pkl/{name}/data_valid_withhead.pkl') 51 | save_pickle(data_test,f'pkl/{name}/data_test.pkl') 52 | save_pickle(data_test_withhead,f'pkl/{name}/data_test_withhead.pkl') 53 | 54 | try: 55 | data_all = load_pickle(f'pkl/{name}/data_all.pkl') 56 | except: 57 | data_all = get_data_my(instruments, train_start, test_end,raw = True,qlib_path = QLIB_PATH,freq=freq) 58 | save_pickle(data_all,f'pkl/{name}/data_all.pkl') 59 | return data_all,data,data_valid,data_valid_withhead,data_test,data_test_withhead,name 60 | -------------------------------------------------------------------------------- /gan/utils/pool.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | import numpy as np 4 | from typing import List, Optional, Tuple, Set 5 | from alphagen.data.expression import Expression 6 | from alphagen.utils.correlation import batch_pearsonr, batch_spearmanr 7 | from alphagen_qlib.stock_data import StockData 8 | 9 | from alphagen.utils.pytorch_utils import normalize_by_day 10 | from tqdm import tqdm 11 | class MyPool: 12 | def __init__(self,data:StockData,target:Expression ): 13 | super().__init__() 14 | self.data = data 15 | self.target = target 16 | self.tgt = target.evaluate(data) #(n_days,n_stocks) 17 | self.nan_mask = torch.isfinite(self.tgt).flatten() #(n_days,n_stocks) 18 | self.tgt = normalize_by_day(self.tgt) 19 | self.flatten_tgt = self.tgt.flatten()[self.nan_mask] #(n_days*n_stocks,) 20 | n_days,n_stocks = self.tgt.shape 21 | 22 | self.best_score = 0. 23 | 24 | self.exprs = [] 25 | 26 | self.expr_tensor = torch.ones((n_days,n_stocks,1)).to(self.device) 27 | self.size = 0 28 | @property 29 | def device(self): 30 | return self.data.data.device 31 | 32 | def add_expr(self,expr:Expression): 33 | print(f"[Pool +] {expr}") 34 | fct = expr.evaluate(self.data) 35 | fct = normalize_by_day(fct) 36 | 37 | self.expr_tensor = torch.cat([self.expr_tensor,fct[...,None]],dim=2) 38 | self.exprs.append(expr) 39 | self.size += 1 40 | assert self.size == len(self.exprs) == self.expr_tensor.shape[2] - 1 41 | 42 | 43 | cur_tensor = self.expr_tensor #(n_days,n_stocks,n_exprs+1) 44 | X = cur_tensor.reshape(-1,self.size+1)[self.nan_mask] 45 | y = self.flatten_tgt #(ndays*nstocks,) 46 | 47 | # Perform multiple linear regression 48 | coefficients, residuals, _, _ = torch.linalg.lstsq(X, y) 49 | 50 | coefficients#(n_exprs+1,1) 51 | # Get the predicted values 52 | predicted = cur_tensor @ coefficients 53 | predicted = predicted.reshape(*fct.shape) 54 | 55 | 56 | cur_score = batch_pearsonr(predicted,self.tgt).mean() 57 | increment = cur_score - self.best_score 58 | if increment > 0: 59 | prev_score = self.best_score 60 | self.best_score = cur_score 61 | print(f"[Best Score] {prev_score:.6f} + {increment:.6f} = {self.best_score:.6f}") 62 | # print(f"coefficients:{coefficients.shape} ,{coefficients.flatten()}") 63 | self.weights = coefficients 64 | 65 | def test_new_expr(self, expr: Expression) -> float: 66 | fct = expr.evaluate(self.data) 67 | fct = normalize_by_day(fct)#(n_days,n_stocks) 68 | 69 | 70 | y = self.flatten_tgt #(ndays*nstocks,) 71 | cur_tensor = torch.cat([self.expr_tensor,fct[...,None]],dim=2) #(n_days,n_stocks,n_exprs+1) 72 | X = cur_tensor.reshape(-1,self.size+2)[self.nan_mask] 73 | 74 | # Perform multiple linear regression 75 | coefficients, residuals, _, _ = torch.linalg.lstsq(X, y) 76 | 77 | coefficients#(n_exprs+1,1) 78 | # Get the predicted values 79 | predicted = cur_tensor @ coefficients 80 | predicted = predicted.reshape(*fct.shape) 81 | 82 | cur_score = batch_pearsonr(predicted,self.tgt).mean() 83 | 84 | increment = cur_score - self.best_score 85 | increment = max(increment.item(),0) 86 | if np.isnan(increment): 87 | increment = 0 88 | return {'score':increment,'ret':np.array([0.,0.,0.])} 89 | 90 | def test_expr_list(self,exprs:List[Expression],verbose:bool = False): 91 | to_iter = exprs if not verbose else tqdm(exprs) 92 | return [self.test_new_expr(expr)['score'] for expr in to_iter] 93 | 94 | def __call__(self,data:StockData): 95 | with torch.no_grad(): 96 | combined_factor: Tensor = 0 97 | for i in range(self.size): 98 | factor = normalize_by_day(self.exprs[i].evaluate(data)) # type: ignore 99 | weighted_factor = factor * self.weights[i+1] 100 | combined_factor += weighted_factor 101 | 102 | return combined_factor 103 | 104 | def evaluate_data(self,data:StockData,target:Expression=None): 105 | if target is None: 106 | target = self.target 107 | tgt = target.evaluate(data) 108 | tgt = normalize_by_day(tgt) 109 | 110 | fct = self(data) 111 | return batch_pearsonr(fct,tgt).mean().item()*100,batch_spearmanr(fct,tgt).mean().item()*100 112 | -------------------------------------------------------------------------------- /gan/utils/qlib.py: -------------------------------------------------------------------------------- 1 | 2 | from alphagen_qlib.stock_data import StockData 3 | def get_data_my(instru,start,end,raw=False,qlib_path = '',freq = 'day'): 4 | import qlib 5 | from qlib.data import D 6 | qlib.init(provider_uri=qlib_path, region='cn') 7 | def get_instruments(name,start,end): 8 | instru = D.instruments(name) 9 | return D.list_instruments( 10 | instruments=instru, 11 | start_time=start, 12 | end_time=end, 13 | as_list=True, 14 | freq=freq, 15 | 16 | ) 17 | instru = get_instruments(instru,start,end) 18 | return StockData(instru,start,end,raw = raw,qlib_path = qlib_path,freq = freq) -------------------------------------------------------------------------------- /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/DulyHao/AlphaForge/d0cfc27df23c60f271bc885fd43027b86b787746/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 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | baostock==00.8.90 2 | gym==0.21.0 3 | matplotlib==3.7.1 4 | numpy==1.24.3 5 | pandas==1.5.3 6 | qlib==0.9.0 7 | sb3_contrib==1.8.0 8 | stable_baselines3==1.8.0 9 | torch==1.13.0 10 | shimmy==1.1.0 11 | fire -------------------------------------------------------------------------------- /train_DSO.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sklearn 3 | import tensorflow as tf 4 | import numpy as np 5 | import os,json 6 | 7 | from alphagen.data.expression import * 8 | # from alphagen_qlib.calculator import QLibStockDataCalculator 9 | from dso import DeepSymbolicRegressor 10 | from dso.library import Token, HardCodedConstant 11 | from dso import functions 12 | from alphagen.models.alpha_pool import AlphaPool 13 | from alphagen.utils import reseed_everything 14 | from alphagen_generic.operators import funcs as generic_funcs 15 | from alphagen_generic.features import * 16 | from gan.utils.data import get_data_by_year 17 | 18 | funcs = {func.name: Token(complexity=1, **func._asdict()) for func in generic_funcs} 19 | for i, feature in enumerate(['open', 'close', 'high', 'low', 'volume', 'vwap']): 20 | funcs[f'x{i+1}'] = Token(name=feature, arity=0, complexity=1, function=None, input_var=i) 21 | for v in [-30., -10., -5., -2., -1., -0.5, -0.01, 0.01, 0.5, 1., 2., 5., 10., 30.]: 22 | funcs[f'Constant({v})'] = HardCodedConstant(name=f'Constant({v})', value=v) 23 | 24 | def main( 25 | instruments:str='csi300', 26 | train_end:int=2018, 27 | seeds:list=[0], 28 | capacity:int=100, 29 | cuda:int=0, 30 | name:str='test', 31 | ): 32 | import os 33 | os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda) 34 | if isinstance(seeds,str): 35 | seeds = eval(seeds) 36 | for seed in seeds: 37 | tf.set_random_seed(seed) 38 | reseed_everything(seed) 39 | returned = get_data_by_year( 40 | train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2, 41 | instruments=instruments, target=target,freq='day',) 42 | data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,_ = returned 43 | 44 | cache = {} 45 | device = torch.device('cuda:0') 46 | 47 | X = np.array([['open_', 'close', 'high', 'low', 'volume', 'vwap']]) 48 | y = np.array([[1]]) 49 | functions.function_map = funcs 50 | 51 | pool = AlphaPool(capacity=capacity, 52 | stock_data=data, 53 | target=target, 54 | ic_lower_bound=None) 55 | save_path = f'out_dso/{name}_{instruments}_{capacity}_{train_end}_{seed}/' 56 | os.makedirs(save_path,exist_ok=True) 57 | 58 | class Ev: 59 | def __init__(self, pool): 60 | self.cnt = 0 61 | self.pool = pool 62 | self.results = {} 63 | 64 | def alpha_ev_fn(self, key): 65 | expr = eval(key) 66 | try: 67 | ret = self.pool.try_new_expr(expr) 68 | except OutOfDataRangeError: 69 | ret = -1. 70 | else: 71 | ret = -1. 72 | finally: 73 | self.cnt += 1 74 | if self.cnt % 100 == 0: 75 | test_ic = pool.test_ensemble(data_test,target)[0] 76 | self.results[self.cnt] = test_ic 77 | print(self.cnt, test_ic) 78 | return ret 79 | 80 | ev = Ev(pool) 81 | 82 | config = dict( 83 | task=dict( 84 | task_type='regression', 85 | function_set=list(funcs.keys()), 86 | metric='alphagen', 87 | metric_params=[lambda key: ev.alpha_ev_fn(key)], 88 | ), 89 | training={'n_samples': 5000, 'batch_size': 128, 'epsilon': 0.05}, 90 | prior={'length': {'min_': 2, 'max_': 20, 'on': True}}, 91 | experiment={'seed':seed}, 92 | ) 93 | 94 | # Create the model 95 | model = DeepSymbolicRegressor(config=config) 96 | model.fit(X, y) 97 | with open(f'{save_path}/pool.json', 'w') as f: 98 | json.dump(pool.to_dict(), f) 99 | print(ev.results) 100 | 101 | if __name__ == '__main__': 102 | import fire 103 | fire.Fire(main) -------------------------------------------------------------------------------- /train_GP.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import argparse 4 | parser = argparse.ArgumentParser() 5 | parser.add_argument('--instrument',type=str,default='csi300') 6 | parser.add_argument('--seed',type=str,default='[0,1,2,3,4]') 7 | parser.add_argument('--years',type=str,default='[2016]') 8 | parser.add_argument('--freq',type=str,default='day') 9 | parser.add_argument('--cuda',type=str,default='0') 10 | 11 | 12 | args = parser.parse_args() 13 | instruments = args.instrument 14 | args.seed = eval(args.seed) 15 | args.years = eval(args.years) 16 | 17 | os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda 18 | print('instruments',instruments) 19 | print('seed',args.seed) 20 | print('years',args.years) 21 | print('cuda',args.cuda) 22 | 23 | 24 | import json 25 | from collections import Counter 26 | 27 | import numpy as np 28 | 29 | from alphagen.data.expression import * 30 | from alphagen.models.alpha_pool import AlphaPool 31 | from alphagen.utils.correlation import batch_pearsonr, batch_spearmanr 32 | from alphagen.utils.pytorch_utils import normalize_by_day 33 | from alphagen.utils.random import reseed_everything 34 | from alphagen_generic.operators import funcs as generic_funcs 35 | from alphagen_generic.features import * 36 | from gplearn.fitness import make_fitness 37 | from gplearn.functions import make_function 38 | from gplearn.genetic import SymbolicRegressor 39 | from gan.utils.data import get_data_by_year 40 | 41 | 42 | def _metric(x, y, w): 43 | key = y[0] 44 | 45 | if key in cache: 46 | return cache[key] 47 | token_len = key.count('(') + key.count(')') 48 | if token_len > 20: 49 | return -1. 50 | 51 | expr = eval(key) 52 | try: 53 | factor = expr.evaluate(data) 54 | factor = normalize_by_day(factor) 55 | ic = batch_pearsonr(factor, target_factor) 56 | ic = torch.nan_to_num(ic).mean().item() 57 | except OutOfDataRangeError: 58 | ic = -1. 59 | if np.isnan(ic): 60 | ic = -1. 61 | cache[key] = ic 62 | return ic 63 | 64 | 65 | 66 | 67 | def try_single(): 68 | top_key = Counter(cache).most_common(1)[0][0] 69 | try: 70 | v_valid = eval(top_key).evaluate(data_valid) 71 | v_test = eval(top_key).evaluate(data_test) 72 | ic_test = batch_pearsonr(v_test, target_factor_test) 73 | ic_test = torch.nan_to_num(ic_test,nan=0,posinf=0,neginf=0).mean().item() 74 | ic_valid = batch_pearsonr(v_valid, target_factor_valid) 75 | ic_valid = torch.nan_to_num(ic_valid,nan=0,posinf=0,neginf=0).mean().item() 76 | ric_test = batch_spearmanr(v_test, target_factor_test) 77 | ric_test = torch.nan_to_num(ric_test,nan=0,posinf=0,neginf=0).mean().item() 78 | ric_valid = batch_spearmanr(v_valid, target_factor_valid) 79 | ric_valid = torch.nan_to_num(ric_valid,nan=0,posinf=0,neginf=0).mean().item() 80 | return {'ic_test': ic_test, 'ic_valid': ic_valid, 'ric_test': ric_test, 'ric_valid': ric_valid} 81 | except OutOfDataRangeError: 82 | print ('Out of data range') 83 | print(top_key) 84 | exit() 85 | return {'ic_test': -1., 'ic_valid': -1., 'ric_test': -1., 'ric_valid': -1.} 86 | 87 | 88 | def try_pool(capacity): 89 | pool = AlphaPool(capacity=capacity, 90 | stock_data=data, 91 | target=target, 92 | ic_lower_bound=None) 93 | 94 | exprs = [] 95 | for key in dict(Counter(cache).most_common(capacity)): 96 | exprs.append(eval(key)) 97 | pool.force_load_exprs(exprs) 98 | pool._optimize(alpha=5e-3, lr=5e-4, n_iter=2000) 99 | 100 | ic_test, ric_test = pool.test_ensemble(data_test, target) 101 | ic_valid, ric_valid = pool.test_ensemble(data_valid, target) 102 | return {'ic_test': ic_test, 'ic_valid': ic_valid, 'ric_test': ric_test, 'ric_valid': ric_valid} 103 | 104 | 105 | 106 | 107 | def ev(): 108 | global generation 109 | generation += 1 110 | res = ( 111 | [{'pool': 0, 'res': try_single()}] + 112 | [{'pool': cap, 'res': try_pool(cap)} for cap in (10, 20, 50, 100)] 113 | ) 114 | print(res) 115 | global save_dir 116 | dir_ = save_dir 117 | #'/path/to/save/results' 118 | os.makedirs(dir_, exist_ok=True) 119 | if generation % 2 == 0: 120 | with open(f'{dir_}/{generation}.json', 'w') as f: 121 | json.dump({'cache': cache, 'res': res}, f) 122 | 123 | 124 | 125 | 126 | 127 | for seed in args.seed: 128 | for train_end in args.years: 129 | #'/path/to/save/results' 130 | save_dir = f'out_gp/{instruments}_{train_end}_{args.freq}_{seed}' 131 | 132 | Metric = make_fitness(function=_metric, greater_is_better=True) 133 | funcs = [make_function(**func._asdict()) for func in generic_funcs] 134 | 135 | generation = 0 136 | cache = {} 137 | 138 | reseed_everything(seed) 139 | 140 | 141 | returned = get_data_by_year( 142 | train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2, 143 | instruments=instruments, target=target,freq=args.freq, 144 | ) 145 | data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,name = returned 146 | 147 | pool = AlphaPool(capacity=10, 148 | stock_data=data, 149 | target=target, 150 | ic_lower_bound=None) 151 | 152 | target_factor = target.evaluate(data) 153 | target_factor_valid = target.evaluate(data_valid) 154 | target_factor_test = target.evaluate(data_test) 155 | 156 | 157 | features = ['open_', 'close', 'high', 'low', 'volume', 'vwap'] 158 | 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.]] 159 | terminals = features + constants 160 | 161 | X_train = np.array([terminals]) 162 | y_train = np.array([[1]]) 163 | 164 | est_gp = SymbolicRegressor(population_size=1000, 165 | generations=40, 166 | init_depth=(2, 6), 167 | tournament_size=600, 168 | stopping_criteria=1., 169 | p_crossover=0.3, 170 | p_subtree_mutation=0.1, 171 | p_hoist_mutation=0.01, 172 | p_point_mutation=0.1, 173 | p_point_replace=0.6, 174 | max_samples=0.9, 175 | verbose=1, 176 | parsimony_coefficient=0., 177 | random_state=seed, 178 | function_set=funcs, 179 | metric=Metric, 180 | const_range=None, 181 | n_jobs=1) 182 | est_gp.fit(X_train, y_train, callback=ev) 183 | print(est_gp._program.execute(X_train)) 184 | -------------------------------------------------------------------------------- /train_RL.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | os.environ["CUDA_VISIBLE_DEVICES"]='0' 4 | import json 5 | from typing import Optional 6 | from datetime import datetime 7 | 8 | import numpy as np 9 | from sb3_contrib.ppo_mask import MaskablePPO 10 | from stable_baselines3.common.callbacks import BaseCallback 11 | 12 | from alphagen_generic.features import * 13 | from alphagen.data.expression import * 14 | from alphagen.models.alpha_pool import AlphaPool, SingleAlphaPool, AlphaPoolBase 15 | from alphagen.rl.env.wrapper import AlphaEnv 16 | from alphagen.rl.policy import LSTMSharedNet 17 | from alphagen.utils.random import reseed_everything 18 | from alphagen.rl.env.core import AlphaEnvCore 19 | 20 | import pickle 21 | def save_pickle(data,path): 22 | with open(path,'wb') as f: 23 | pickle.dump(data,f) 24 | def load_pickle(path): 25 | with open(path,'rb') as f: 26 | return pickle.load(f) 27 | 28 | class CustomCallback(BaseCallback): 29 | def __init__(self, 30 | save_freq: int, 31 | show_freq: int, 32 | save_path: str, 33 | train_data: StockData, 34 | train_target: Expression, 35 | valid_data: StockData, 36 | valid_target: Expression, 37 | test_data: StockData, 38 | test_target: Expression, 39 | name_prefix: str = 'rl_model', 40 | timestamp: Optional[str] = None, 41 | verbose: int = 0): 42 | super().__init__(verbose) 43 | self.save_freq = save_freq 44 | self.show_freq = show_freq 45 | self.save_path = save_path 46 | self.name_prefix = name_prefix 47 | 48 | self.train_data = train_data 49 | self.train_target = train_target 50 | self.valid_data = valid_data 51 | self.valid_target = valid_target 52 | self.test_data = test_data 53 | self.test_target = test_target 54 | 55 | if timestamp is None: 56 | self.timestamp = datetime.now().strftime('%Y%m%d%H%M%S') 57 | else: 58 | self.timestamp = timestamp 59 | 60 | def _init_callback(self) -> None: 61 | if self.save_path is not None: 62 | os.makedirs(self.save_path, exist_ok=True) 63 | 64 | def _on_step(self) -> bool: 65 | return True 66 | 67 | def _on_rollout_end(self) -> None: 68 | assert self.logger is not None 69 | self.logger.record('pool/size', self.pool.size) 70 | self.logger.record('pool/significant', (np.abs(self.pool.weights[:self.pool.size]) > 1e-4).sum()) 71 | self.logger.record('pool/best_ic_ret', self.pool.best_ic_ret) 72 | self.logger.record('pool/eval_cnt', self.pool.eval_cnt) 73 | 74 | ic_train, rank_ic_train = self.pool.test_ensemble(self.train_data, self.train_target) 75 | self.logger.record('_train/ic', ic_train) 76 | self.logger.record('_train/rank_ic', rank_ic_train) 77 | 78 | ic_test, rank_ic_test = self.pool.test_ensemble(self.test_data, self.test_target) 79 | self.logger.record('_test/ic', ic_test) 80 | self.logger.record('_test/rank_ic', rank_ic_test) 81 | 82 | ic_valid, rank_ic_valid = self.pool.test_ensemble(self.valid_data, self.valid_target) 83 | self.logger.record('_valid/ic', ic_valid) 84 | self.logger.record('_valid/rank_ic', rank_ic_valid) 85 | self.save_checkpoint() 86 | 87 | def save_checkpoint(self): 88 | path = os.path.join(self.save_path, f'{self.name_prefix}_{self.timestamp}', f'{self.num_timesteps}_steps') 89 | self.model.save(path) # type: ignore 90 | if self.verbose > 1: 91 | print(f'Saving model checkpoint to {path}') 92 | save_pickle(self.pool,path+'_pool.pkl') 93 | with open(f'{path}_pool.json', 'w') as f: 94 | json.dump(self.pool.to_dict(), f) 95 | 96 | def show_pool_state(self): 97 | state = self.pool.state 98 | n = len(state['exprs']) 99 | print('---------------------------------------------') 100 | for i in range(n): 101 | weight = state['weights'][i] 102 | expr_str = str(state['exprs'][i]) 103 | ic_ret = state['ics_ret'][i] 104 | print(f'> Alpha #{i}: {weight}, {expr_str}, {ic_ret}') 105 | print(f'>> Ensemble ic_ret: {state["best_ic_ret"]}') 106 | print('---------------------------------------------') 107 | 108 | @property 109 | def pool(self) -> AlphaPoolBase: 110 | return self.env_core.pool 111 | 112 | @property 113 | def env_core(self) -> AlphaEnvCore: 114 | return self.training_env.envs[0].unwrapped # type: ignore 115 | 116 | 117 | def main( 118 | seed: int = 0, 119 | instruments: str = "csi300", 120 | pool_capacity: int = 10, 121 | steps: int = 200_000, 122 | raw: bool = False, 123 | train_end: int = 2019, 124 | freq: str = 'day', 125 | ): 126 | reseed_everything(seed) 127 | 128 | device = torch.device('cuda:0') 129 | close = Feature(FeatureType.CLOSE) 130 | 131 | from alphagen_generic.features import open_ 132 | from gan.utils import Builders 133 | from gan.utils.data import get_data_by_year 134 | import os 135 | 136 | 137 | reseed_everything(seed) 138 | returned = get_data_by_year( 139 | train_start = 2010,train_end=train_end,valid_year=train_end+1,test_year =train_end+2, 140 | instruments=instruments, target=target,freq=freq, 141 | ) 142 | data_all, data,data_valid,data_valid_withhead,data_test,data_test_withhead,name = returned 143 | 144 | 145 | pool = AlphaPool( 146 | capacity=pool_capacity, 147 | stock_data=data, 148 | target=target, 149 | ic_lower_bound=None 150 | ) 151 | env = AlphaEnv(pool=pool, device=device, print_expr=True) 152 | 153 | name_prefix = f"n1227day_{instruments}_{train_end}_{pool_capacity}_{seed}" ## new_time 154 | timestamp = datetime.now().strftime('%Y%m%d%H%M%S') 155 | 156 | checkpoint_callback = CustomCallback( 157 | save_freq=10000, 158 | show_freq=10000, 159 | save_path='out_ppo/checkpoints', 160 | train_data=data, 161 | train_target=target, 162 | valid_data=data_valid, 163 | valid_target=target, 164 | test_data=data_test, 165 | test_target=target, 166 | name_prefix=name_prefix, 167 | timestamp=timestamp, 168 | verbose=1, 169 | ) 170 | 171 | model = MaskablePPO( 172 | 'MlpPolicy', 173 | env, 174 | policy_kwargs=dict( 175 | features_extractor_class=LSTMSharedNet, 176 | features_extractor_kwargs=dict( 177 | n_layers=2, 178 | d_model=128, 179 | dropout=0.1, 180 | device=device, 181 | ), 182 | ), 183 | gamma=1., 184 | ent_coef=0.01, 185 | batch_size=128, 186 | tensorboard_log='out_ppo/log2', 187 | device=device, 188 | verbose=1, 189 | ) 190 | model.learn( 191 | total_timesteps=steps, 192 | callback=checkpoint_callback, 193 | tb_log_name=f'{name_prefix}_{timestamp}', 194 | ) 195 | 196 | 197 | from gan.utils.qlib import get_data_my 198 | if __name__ == '__main__': 199 | steps = { 200 | 10: 250_000, 201 | 20: 300_000, 202 | 50: 350_000, 203 | 100: 400_000 204 | } 205 | train_end = 2020 206 | for capacity in [1,10,20,50,100]: 207 | for seed in range(5): 208 | for instruments in ["csi300"]: 209 | main( 210 | seed=seed, instruments=instruments, pool_capacity=capacity, 211 | steps=steps[capacity], raw = True, 212 | train_end=train_end, 213 | ) 214 | 215 | 216 | --------------------------------------------------------------------------------