1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-11 14:56:55 +08:00

Add segment args for pred and refine MultiSegRecord

This commit is contained in:
D-X-Y
2021-03-28 08:30:16 +00:00
parent 9d04ae4676
commit 8a2e7b62af
5 changed files with 47 additions and 21 deletions

View File

@@ -61,10 +61,10 @@ class LGBModel(ModelFT):
evals_result["train"] = list(evals_result["train"].values())[0]
evals_result["valid"] = list(evals_result["valid"].values())[0]
def predict(self, dataset):
def predict(self, dataset, segment="test"):
if self.model is None:
raise ValueError("model is not fitted yet!")
x_test = dataset.prepare("test", col_set="feature", data_key=DataHandlerLP.DK_I)
x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
return pd.Series(self.model.predict(x_test.values), index=x_test.index)
def finetune(self, dataset: DatasetH, num_boost_round=10, verbose_eval=20):

View File

@@ -84,8 +84,8 @@ class LinearModel(Model):
self.coef_ = coef
self.intercept_ = 0.0
def predict(self, dataset):
def predict(self, dataset, segment="test"):
if self.coef_ is None:
raise ValueError("model is not fitted yet!")
x_test = dataset.prepare("test", col_set="feature", data_key=DataHandlerLP.DK_I)
x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
return pd.Series(x_test.values @ self.coef_ + self.intercept_, index=x_test.index)

View File

@@ -57,8 +57,8 @@ class XGBModel(Model):
evals_result["train"] = list(evals_result["train"].values())[0]
evals_result["valid"] = list(evals_result["valid"].values())[0]
def predict(self, dataset):
def predict(self, dataset, segment="test"):
if self.model is None:
raise ValueError("model is not fitted yet!")
x_test = dataset.prepare("test", col_set="feature")
x_test = dataset.prepare(segment, col_set="feature")
return pd.Series(self.model.predict(xgb.DMatrix(x_test.values)), index=x_test.index)