Y-oHr-N / OptGBM

Optuna + LightGBM = OptGBM
MIT License
34 stars 7 forks source link

Add plot_feature_importances to _BaseOGBMModel #10

Closed Y-oHr-N closed 4 years ago

Y-oHr-N commented 5 years ago
Y-oHr-N commented 4 years ago
from typing import Any
from typing import List
from typing import Optional
from typing import Union

import lightgbm as lgb
import numpy as np
import pandas as pd

TWO_DIM_ARRAYLIKE_TYPE = Union[np.ndarray, pd.DataFrame]

class _VotingBooster(object):
    @property
    def feature_name(self) -> List[str]:
        return self.boosters[0].feature_name

    def __init__(
        self,
        boosters: List[lgb.Booster],
        weights: Optional[np.ndarray] = None
    ) -> None:
        self.boosters = boosters
        self.weights = weights

    def predict(
        self,
        X: TWO_DIM_ARRAYLIKE_TYPE,
        **kwargs: Any
    ) -> TWO_DIM_ARRAYLIKE_TYPE:
        results = []

        for b in self.boosters:
            result = b.predict(X, **kwargs)

            results.append(result)

        return np.average(results, axis=0, weights=self.weights)

    def feature_importance(self, **kwargs: Any) -> np.ndarray:
        results = []

        for b in self.boosters:
            result = b.feature_importance(**kwargs)

            results.append(result)

        return np.average(results, axis=0, weights=self.weights)