qlib/qlib/contrib/estimator/estimator.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

# coding=utf-8

import pandas as pd

import os
import copy
import json
import yaml
import pickle

import qlib
from ..evaluate import risk_analysis
from ..evaluate import backtest as normal_backtest
from ..evaluate import long_short_backtest
from .config import ExperimentConfig
from .fetcher import create_fetcher_with_config

from ...log import get_module_logger, TimeInspector
from ...utils import get_module_by_module_path, compare_dict_value


class Estimator(object):
    def __init__(self, config_manager, sacred_ex):

        # Set logger.
        self.logger = get_module_logger("Estimator")

        # 1. Set config manager.
        self.config_manager = config_manager

        # 2. Set configs.
        self.ex_config = config_manager.ex_config
        self.data_config = config_manager.data_config
        self.model_config = config_manager.model_config
        self.trainer_config = config_manager.trainer_config
        self.strategy_config = config_manager.strategy_config
        self.backtest_config = config_manager.backtest_config

        # If experiment.mode is test or experiment.finetune is True, load the experimental results in the loader
        if self.ex_config.mode == self.ex_config.TEST_MODE or self.ex_config.finetune:
            self.compare_config_with_config_manger(self.config_manager)

        # 3. Set sacred_experiment.
        self.ex = sacred_ex

        # 4. Init data handler.
        self.data_handler = None
        self._init_data_handler()

        # 5. Init trainer.
        self.trainer = None
        self._init_trainer()

        # 6. Init strategy.
        self.strategy = None
        self._init_strategy()

    def _init_data_handler(self):
        handler_module = get_module_by_module_path(self.data_config.handler_module_path)

        # Set market
        market = self.data_config.handler_filter.get("market", None)
        if market is None:
            if "market" in self.data_config.handler_parameters:
                self.logger.warning(
                    "Warning: The market in data.args section is deprecated. "
                    "It only works when market is not set in data.filter section. "
                    "It will be overridden by market in the data.filter section."
                )
                market = self.data_config.handler_parameters["market"]
            else:
                market = "csi500"

        self.data_config.handler_parameters["market"] = market

        data_filter_list = []
        handler_filters = self.data_config.handler_filter.get("filter_pipeline", list())
        for h_filter in handler_filters:
            filter_module_path = h_filter.get("module_path", "qlib.data.filter")
            filter_class_name = h_filter.get("class", "")
            filter_parameters = h_filter.get("args", {})
            filter_module = get_module_by_module_path(filter_module_path)
            filter_class = getattr(filter_module, filter_class_name)
            data_filter = filter_class(**filter_parameters)
            data_filter_list.append(data_filter)

        self.data_config.handler_parameters["data_filter_list"] = data_filter_list
        handler_class = getattr(handler_module, self.data_config.handler_class)
        self.data_handler = handler_class(**self.data_config.handler_parameters)

    def _init_trainer(self):

        model_module = get_module_by_module_path(self.model_config.model_module_path)
        trainer_module = get_module_by_module_path(self.trainer_config.trainer_module_path)
        model_class = getattr(model_module, self.model_config.model_class)
        trainer_class = getattr(trainer_module, self.trainer_config.trainer_class)

        self.trainer = trainer_class(
            model_class,
            self.model_config.save_path,
            self.model_config.parameters,
            self.data_handler,
            self.ex,
            **self.trainer_config.parameters
        )

    def _init_strategy(self):

        module = get_module_by_module_path(self.strategy_config.strategy_module_path)
        strategy_class = getattr(module, self.strategy_config.strategy_class)
        self.strategy = strategy_class(**self.strategy_config.parameters)

    def run(self):
        if self.ex_config.mode == ExperimentConfig.TRAIN_MODE:
            self.trainer.train()
        elif self.ex_config.mode == ExperimentConfig.TEST_MODE:
            self.trainer.load()
        else:
            raise ValueError("unexpected mode: %s" % self.ex_config.mode)
        analysis = self.backtest()
        print(analysis)
        self.logger.info(
            "experiment id: {}, experiment name: {}".format(self.ex.experiment.current_run._id, self.ex_config.name)
        )

        # Remove temp dir
        # shutil.rmtree(self.ex_config.tmp_run_dir)

    def backtest(self):
        TimeInspector.set_time_mark()
        # 1. Get pred and prediction score of model(s).
        pred = self.trainer.get_test_score()
        try:
            performance = self.trainer.get_test_performance()
        except NotImplementedError:
            performance = None
        # 2. Normal Backtest.
        report_normal, positions_normal = self._normal_backtest(pred)
        # 3. Long-Short Backtest.
        # Deprecated
        # long_short_reports = self._long_short_backtest(pred)
        # 4. Analyze
        analysis_df = self._analyze(report_normal)
        # 5. Save.
        self._save_backtest_result(
            pred,
            analysis_df,
            positions_normal,
            report_normal,
            # long_short_reports,
            performance,
        )
        return analysis_df

    def _normal_backtest(self, pred):
        TimeInspector.set_time_mark()
        if "account" not in self.backtest_config.normal_backtest_parameters:
            if "account" in self.strategy_config.parameters:
                self.logger.warning(
                    "Warning: The account in strategy section is deprecated. "
                    "It only works when account is not set in backtest section. "
                    "It will be overridden by account in the backtest section."
                )
                self.backtest_config.normal_backtest_parameters["account"] = self.strategy_config.parameters["account"]
        report_normal, positions_normal = normal_backtest(
            pred, strategy=self.strategy, **self.backtest_config.normal_backtest_parameters
        )
        TimeInspector.log_cost_time("Finished normal backtest.")
        return report_normal, positions_normal

    def _long_short_backtest(self, pred):
        TimeInspector.set_time_mark()
        long_short_reports = long_short_backtest(pred, **self.backtest_config.long_short_backtest_parameters)
        TimeInspector.log_cost_time("Finished long-short backtest.")
        return long_short_reports

    @staticmethod
    def _analyze(report_normal):
        TimeInspector.set_time_mark()

        analysis = dict()
        # analysis["pred_long"] = risk_analysis(long_short_reports["long"])
        # analysis["pred_short"] = risk_analysis(long_short_reports["short"])
        # analysis["pred_long_short"] = risk_analysis(long_short_reports["long_short"])
        analysis["excess_return_without_cost"] = risk_analysis(report_normal["return"] - report_normal["bench"])
        analysis["excess_return_with_cost"] = risk_analysis(
            report_normal["return"] - report_normal["bench"] - report_normal["cost"]
        )
        analysis_df = pd.concat(analysis)  # type: pd.DataFrame
        TimeInspector.log_cost_time(
            "Finished generating analysis," " average turnover is: {0:.4f}.".format(report_normal["turnover"].mean())
        )
        return analysis_df

    def _save_backtest_result(self, pred, analysis, positions, report_normal, performance):
        # 1. Result dir.
        result_dir = os.path.join(self.config_manager.ex_config.tmp_run_dir, "result")
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)

        self.ex.add_info(
            "task_config",
            json.loads(json.dumps(self.config_manager.config, default=str)),
        )

        # 2. Pred.
        TimeInspector.set_time_mark()
        pred_pkl_path = os.path.join(result_dir, "pred.pkl")
        pred.to_pickle(pred_pkl_path)
        self.ex.add_artifact(pred_pkl_path)
        TimeInspector.log_cost_time("Finished saving pred.pkl to: {}".format(pred_pkl_path))

        # 3. Ana.
        TimeInspector.set_time_mark()
        analysis_pkl_path = os.path.join(result_dir, "analysis.pkl")
        analysis.to_pickle(analysis_pkl_path)
        self.ex.add_artifact(analysis_pkl_path)
        TimeInspector.log_cost_time("Finished saving analysis.pkl to: {}".format(analysis_pkl_path))

        # 4. Pos.
        TimeInspector.set_time_mark()
        positions_pkl_path = os.path.join(result_dir, "positions.pkl")
        with open(positions_pkl_path, "wb") as fp:
            pickle.dump(positions, fp)
        self.ex.add_artifact(positions_pkl_path)
        TimeInspector.log_cost_time("Finished saving positions.pkl to: {}".format(positions_pkl_path))

        # 5. Report normal.
        TimeInspector.set_time_mark()
        report_normal_pkl_path = os.path.join(result_dir, "report_normal.pkl")
        report_normal.to_pickle(report_normal_pkl_path)
        self.ex.add_artifact(report_normal_pkl_path)
        TimeInspector.log_cost_time("Finished saving report_normal.pkl to: {}".format(report_normal_pkl_path))

        # 6. Report long short.
        # Deprecated
        # for k, name in zip(
        #     ["long", "short", "long_short"],
        #     ["report_long.pkl", "report_short.pkl", "report_long_short.pkl"],
        # ):
        #     TimeInspector.set_time_mark()
        #     pkl_path = os.path.join(result_dir, name)
        #     long_short_reports[k].to_pickle(pkl_path)
        #     self.ex.add_artifact(pkl_path)
        #     TimeInspector.log_cost_time("Finished saving {} to: {}".format(name, pkl_path))

        # 7. Origin test label.
        TimeInspector.set_time_mark()
        label_pkl_path = os.path.join(result_dir, "label.pkl")
        self.data_handler.get_origin_test_label_with_date(
            self.trainer_config.parameters["test_start_date"],
            self.trainer_config.parameters["test_end_date"],
        ).to_pickle(label_pkl_path)
        self.ex.add_artifact(label_pkl_path)
        TimeInspector.log_cost_time("Finished saving label.pkl to: {}".format(label_pkl_path))

        # 8. Experiment info, save the model(s) performance here.
        TimeInspector.set_time_mark()
        cur_ex_id = self.ex.experiment.current_run._id
        exp_info = {
            "id": cur_ex_id,
            "name": self.ex_config.name,
            "performance": performance,
            "observer_type": self.ex_config.observer_type,
        }

        if self.ex_config.observer_type == ExperimentConfig.OBSERVER_MONGO:
            exp_info.update(
                {
                    "mongo_url": self.ex_config.mongo_url,
                    "db_name": self.ex_config.db_name,
                }
            )
        else:
            exp_info.update({"dir": self.ex_config.global_dir})

        with open(self.ex_config.exp_info_path, "w") as fp:
            json.dump(exp_info, fp, indent=4, sort_keys=True)
        self.ex.add_artifact(self.ex_config.exp_info_path)
        TimeInspector.log_cost_time("Finished saving ex_info to: {}".format(self.ex_config.exp_info_path))

    @staticmethod
    def compare_config_with_config_manger(config_manager):
        """Compare loader model args and current config with ConfigManage

        :param config_manager: ConfigManager
        :return:
        """
        fetcher = create_fetcher_with_config(config_manager, load_form_loader=True)
        loader_mode_config = fetcher.get_experiment(
            exp_name=config_manager.ex_config.loader_name,
            exp_id=config_manager.ex_config.loader_id,
            fields=["task_config"],
        )["task_config"]
        with open(config_manager.config_path) as fp:
            current_config = yaml.load(fp.read())
            current_config = json.loads(json.dumps(current_config, default=str))

        logger = get_module_logger("Estimator")

        loader_mode_config = copy.deepcopy(loader_mode_config)
        current_config = copy.deepcopy(current_config)

        # Require test_mode_config.test_start_date <= current_config.test_start_date
        loader_trainer_args = loader_mode_config.get("trainer", {}).get("args", {})
        cur_trainer_args = current_config.get("trainer", {}).get("args", {})
        loader_start_date = loader_trainer_args.pop("test_start_date")
        cur_test_start_date = cur_trainer_args.pop("test_start_date")
        assert (
            loader_start_date <= cur_test_start_date
        ), "Require: loader_mode_config.test_start_date <= current_config.test_start_date"

        # TODO: For the user's own extended `Trainer`, the support is not very good
        if "RollingTrainer" == current_config.get("trainer", {}).get("class", None):
            loader_period = loader_trainer_args.pop("rolling_period")
            cur_period = cur_trainer_args.pop("rolling_period")
            assert (
                loader_period == cur_period
            ), "Require: loader_mode_config.rolling_period == current_config.rolling_period"

        compare_section = ["trainer", "model", "data"]
        for section in compare_section:
            changes = compare_dict_value(loader_mode_config.get(section, {}), current_config.get(section, {}))
            if changes:
                logger.warning("Warning: Loader mode config and current config, `{}` are different:\n".format(section))