wehub-resource-sync
140 lines
3.8 KiB
Python
140 lines
3.8 KiB
Python
import numpy as np
|
|
import polars as pl
|
|
from sklearn.linear_model import Lasso # type: ignore
|
|
|
|
from vnpy.alpha import (
|
|
AlphaDataset,
|
|
AlphaModel,
|
|
Segment,
|
|
logger
|
|
)
|
|
|
|
|
|
class LassoModel(AlphaModel):
|
|
"""LASSO regression learning algorithm"""
|
|
|
|
def __init__(
|
|
self,
|
|
alpha: float = 0.0005,
|
|
max_iter: int = 1000,
|
|
random_state: int | None = None,
|
|
) -> None:
|
|
"""
|
|
Parameters
|
|
----------
|
|
alpha : float
|
|
Regularization parameter
|
|
max_iter : int
|
|
Maximum number of iterations
|
|
random_state : int
|
|
Random seed
|
|
"""
|
|
self.alpha: float = alpha
|
|
self.max_iter: int = max_iter
|
|
self.random_state: int | None = random_state
|
|
|
|
self.model: Lasso = None
|
|
|
|
self.feature_names: list[str] = []
|
|
|
|
def fit(self, dataset: AlphaDataset) -> None:
|
|
"""
|
|
Fit the model with dataset
|
|
|
|
Parameters
|
|
----------
|
|
dataset : AlphaDataset
|
|
The dataset used for training
|
|
"""
|
|
# Get training data
|
|
df_train: pl.DataFrame = dataset.fetch_learn(Segment.TRAIN)
|
|
df_valid: pl.DataFrame = dataset.fetch_learn(Segment.VALID)
|
|
|
|
# Merge data, remove duplicates and sort
|
|
df_train = pl.concat([df_train, df_valid])
|
|
df_train = df_train.unique(subset=["datetime", "vt_symbol"])
|
|
df_train = df_train.sort(["datetime", "vt_symbol"])
|
|
|
|
# Extract feature names
|
|
self.feature_names = df_train.columns[2:-1]
|
|
|
|
# Convert to numpy arrays
|
|
X: np.ndarray = df_train.select(self.feature_names).to_numpy()
|
|
y: np.ndarray = np.array(df_train["label"])
|
|
|
|
# Create and train the model
|
|
self.model = Lasso(
|
|
alpha=self.alpha,
|
|
max_iter=self.max_iter,
|
|
random_state=self.random_state,
|
|
fit_intercept=False,
|
|
copy_X=False
|
|
)
|
|
self.model.fit(X, y)
|
|
|
|
def predict(self, dataset: AlphaDataset, segment: Segment) -> np.ndarray:
|
|
"""
|
|
Make predictions using the model
|
|
|
|
Parameters
|
|
----------
|
|
dataset : AlphaDataset
|
|
The dataset used for prediction
|
|
segment : Segment
|
|
The segment of data to use for prediction
|
|
|
|
Returns
|
|
-------
|
|
np.ndarray
|
|
Prediction results
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
If the model has not been fitted yet
|
|
"""
|
|
# Check if model exists
|
|
if self.model is None:
|
|
raise ValueError("model is not fitted yet!")
|
|
|
|
# Get data for prediction
|
|
df: pl.DataFrame = dataset.fetch_infer(segment)
|
|
df = df.sort(["datetime", "vt_symbol"])
|
|
|
|
# Convert to numpy array
|
|
data: np.ndarray = df.select(df.columns[2: -1]).to_numpy()
|
|
|
|
# Return prediction results
|
|
result: np.ndarray = self.model.predict(data)
|
|
|
|
return result
|
|
|
|
def detail(self) -> None:
|
|
"""
|
|
Output detailed information about the model
|
|
|
|
Displays feature importance based on the coefficients
|
|
of the LASSO model, showing only non-zero features
|
|
sorted by absolute value.
|
|
"""
|
|
# Get feature coefficients
|
|
coef: np.ndarray = self.model.coef_
|
|
|
|
# Extract feature coefficients
|
|
data: list[tuple[str, float]] = list(zip(self.feature_names, coef, strict=False))
|
|
|
|
# Filter non-zero features
|
|
data = [x for x in data if x[1]]
|
|
|
|
# Sort by absolute value
|
|
data.sort(key=lambda x: abs(x[1]), reverse=True)
|
|
|
|
# Filter out features with very small coefficients
|
|
data = [x for x in data if round(x[1], 6) != 0]
|
|
|
|
# Print feature importance
|
|
logger.info(f"LASSO模型特征总数量: {len(data)}")
|
|
|
|
for name, importance in data:
|
|
logger.info(f"{name}: {importance:.6f}")
|