NexQuant/rdagent/scenarios/qlib/developer/factor_runner.py

from pathlib import Path

import pandas as pd
from pandarallel import pandarallel

from rdagent.core.conf import RD_AGENT_SETTINGS
from rdagent.core.utils import cache_with_pickle

pandarallel.initialize(verbose=1)

from rdagent.app.qlib_rd_loop.conf import FactorBasePropSetting
from rdagent.components.runner import CachedRunner
from rdagent.core.exception import FactorEmptyError
from rdagent.log import rdagent_logger as logger
from rdagent.scenarios.qlib.developer.utils import process_factor_data
from rdagent.scenarios.qlib.experiment.factor_experiment import QlibFactorExperiment
from rdagent.scenarios.qlib.experiment.model_experiment import QlibModelExperiment

DIRNAME = Path(__file__).absolute().resolve().parent
DIRNAME_local = Path.cwd()

# TODO: supporting multiprocessing and keep previous results


class QlibFactorRunner(CachedRunner[QlibFactorExperiment]):
    """
    Docker run
    Everything in a folder
    - config.yaml
    - price-volume data dumper
    - `data.py` + Adaptor to Factor implementation
    - results in `mlflow`
    """

    def calculate_information_coefficient(
        self, concat_feature: pd.DataFrame, SOTA_feature_column_size: int, new_feature_columns_size: int
    ) -> pd.DataFrame:
        res = pd.Series(index=range(SOTA_feature_column_size * new_feature_columns_size))
        for col1 in range(SOTA_feature_column_size):
            for col2 in range(SOTA_feature_column_size, SOTA_feature_column_size + new_feature_columns_size):
                res.loc[col1 * new_feature_columns_size + col2 - SOTA_feature_column_size] = concat_feature.iloc[
                    :, col1
                ].corr(concat_feature.iloc[:, col2])
        return res

    def deduplicate_new_factors(self, SOTA_feature: pd.DataFrame, new_feature: pd.DataFrame) -> pd.DataFrame:
        # calculate the IC between each column of SOTA_feature and new_feature
        # if the IC is larger than a threshold, remove the new_feature column
        # return the new_feature

        concat_feature = pd.concat([SOTA_feature, new_feature], axis=1)
        IC_max = (
            concat_feature.groupby("datetime")
            .parallel_apply(
                lambda x: self.calculate_information_coefficient(x, SOTA_feature.shape[1], new_feature.shape[1])
            )
            .mean()
        )
        IC_max.index = pd.MultiIndex.from_product([range(SOTA_feature.shape[1]), range(new_feature.shape[1])])
        IC_max = IC_max.unstack().max(axis=0)
        return new_feature.iloc[:, IC_max[IC_max < 0.99].index]

    @cache_with_pickle(CachedRunner.get_cache_key, CachedRunner.assign_cached_result)
    def develop(self, exp: QlibFactorExperiment) -> QlibFactorExperiment:
        """
        Generate the experiment by processing and combining factor data,
        then passing the combined data to Docker for backtest results.
        """
        if exp.based_experiments and exp.based_experiments[-1].result is None:
            logger.info(f"Baseline experiment execution ...")
            exp.based_experiments[-1] = self.develop(exp.based_experiments[-1])

        fbps = FactorBasePropSetting()
        env_to_use = {
            "PYTHONPATH": "./",
            "train_start": fbps.train_start,
            "train_end": fbps.train_end,
            "valid_start": fbps.valid_start,
            "valid_end": fbps.valid_end,
            "test_start": fbps.test_start,
            "feature_names": str(list(exp.base_features.keys())),
            "feature_expressions": str(list(exp.base_features.values())),
        }
        if fbps.test_end is not None:
            env_to_use.update({"test_end": fbps.test_end})

        if exp.based_experiments:
            SOTA_factor = None
            # Filter and retain only QlibFactorExperiment instances
            sota_factor_experiments_list = [
                base_exp for base_exp in exp.based_experiments if isinstance(base_exp, QlibFactorExperiment)
            ]
            if len(sota_factor_experiments_list) > 1:
                logger.info(f"SOTA factor processing ...")
                SOTA_factor = process_factor_data(sota_factor_experiments_list)

            # Process the new factors data
            logger.info(f"New factor processing ...")
            new_factors = process_factor_data(exp)

            if new_factors.empty:
                raise FactorEmptyError("Factors failed to run on the full sample, this round of experiment failed.")

            # Combine the SOTA factor and new factors if SOTA factor exists
            if SOTA_factor is not None and not SOTA_factor.empty:
                new_factors = self.deduplicate_new_factors(SOTA_factor, new_factors)
                if new_factors.empty:
                    raise FactorEmptyError(
                        "The factors generated in this round are highly similar to the previous factors. Please change the direction for creating new factors."
                    )
                combined_factors = pd.concat([SOTA_factor, new_factors], axis=1).dropna()
            else:
                combined_factors = new_factors

            # Sort and nest the combined factors under 'feature'
            combined_factors = combined_factors.sort_index()
            combined_factors = combined_factors.loc[:, ~combined_factors.columns.duplicated(keep="last")]
            new_columns = pd.MultiIndex.from_product([["feature"], combined_factors.columns])
            combined_factors.columns = new_columns
            logger.info(f"Factor data processing completed.")

            num_features = len(exp.base_features) + len(combined_factors.columns)

            # Due to the rdagent and qlib docker image in the numpy version of the difference,
            # the `combined_factors_df.pkl` file could not be loaded correctly in qlib dokcer,
            # so we changed the file type of `combined_factors_df` from pkl to parquet.
            target_path = exp.experiment_workspace.workspace_path / "combined_factors_df.parquet"

            # Save the combined factors to the workspace
            combined_factors.to_parquet(target_path, engine="pyarrow")

            # If model exp exists in the previous experiment
            exist_sota_model_exp = False
            for base_exp in reversed(exp.based_experiments):
                if isinstance(base_exp, QlibModelExperiment):
                    sota_model_exp = base_exp
                    exist_sota_model_exp = True
                    break
            logger.info(f"Experiment execution ...")
            if exist_sota_model_exp:
                exp.experiment_workspace.inject_files(
                    **{"model.py": sota_model_exp.sub_workspace_list[0].file_dict["model.py"]}
                )
                sota_training_hyperparameters = sota_model_exp.sub_tasks[0].training_hyperparameters
                if sota_training_hyperparameters:
                    env_to_use.update(
                        {
                            "n_epochs": str(sota_training_hyperparameters.get("n_epochs", "100")),
                            "lr": str(sota_training_hyperparameters.get("lr", "2e-4")),
                            "early_stop": str(sota_training_hyperparameters.get("early_stop", 10)),
                            "batch_size": str(sota_training_hyperparameters.get("batch_size", 256)),
                            "weight_decay": str(sota_training_hyperparameters.get("weight_decay", 0.0001)),
                        }
                    )
                sota_model_type = sota_model_exp.sub_tasks[0].model_type
                if sota_model_type == "TimeSeries":
                    env_to_use.update(
                        {"dataset_cls": "TSDatasetH", "num_features": num_features, "step_len": 20, "num_timesteps": 20}
                    )
                elif sota_model_type == "Tabular":
                    env_to_use.update({"dataset_cls": "DatasetH", "num_features": num_features})

                # model + combined factors
                result, stdout = exp.experiment_workspace.execute(
                    qlib_config_name="conf_combined_factors_sota_model.yaml", run_env=env_to_use
                )
            else:
                # LGBM + combined factors
                result, stdout = exp.experiment_workspace.execute(
                    qlib_config_name="conf_combined_factors.yaml",
                    run_env=env_to_use,
                )
        else:
            logger.info(f"Experiment execution ...")
            if exp.base_feature_codes:
                factors = process_factor_data(exp)
                factors = factors.sort_index()
                factors = factors.loc[:, ~factors.columns.duplicated(keep="last")]
                new_columns = pd.MultiIndex.from_product([["feature"], factors.columns])
                factors.columns = new_columns
                target_path = exp.experiment_workspace.workspace_path / "combined_factors_df.parquet"
                # Save the combined factors to the workspace
                factors.to_parquet(target_path, engine="pyarrow")
                logger.info(f"Factor data processing completed.")
                result, stdout = exp.experiment_workspace.execute(
                    qlib_config_name="conf_combined_factors.yaml",
                    run_env=env_to_use,
                )
            else:
                result, stdout = exp.experiment_workspace.execute(
                    qlib_config_name="conf_baseline.yaml",
                    run_env=env_to_use,
                )

        if result is None:
            logger.error(f"Failed to run this experiment, because {stdout}")
            raise FactorEmptyError(f"Failed to run this experiment, because {stdout}")

        exp.result = result
        exp.stdout = stdout

        # Protection Manager: Check if factor passes risk criteria
        try:
            self._run_protection_check(exp, result)
        except Exception as e:
            logger.warning(f"Protection check failed for factor {exp.hypothesis.hypothesis}: {e}")
            # Don't block the workflow, just log the warning

        # Save results to database immediately after Docker execution
        try:
            self._save_result_to_database(exp, result)
        except Exception as e:
            logger.warning(f"Failed to save results to database: {e}")

        return exp

    def _save_result_to_database(self, exp, result) -> None:
        """
        Save backtest results to the ResultsDatabase and write factor JSON summary.

        This method is called immediately after Docker execution to ensure
        results are persisted before any potential failures.

        Parameters
        ----------
        exp : QlibFactorExperiment
            The experiment with backtest results
        result : pd.Series
            Backtest metrics from Qlib (qlib_res.csv) - a pd.Series with index
            containing metric names like 'IC', '1day.excess_return_with_cost.shar', etc.
        """
        try:
            import json
            import pandas as pd
            from pathlib import Path
            from rdagent.components.backtesting import ResultsDatabase

            # Get factor name from hypothesis
            factor_name = "unknown"
            if hasattr(exp, 'hypothesis') and exp.hypothesis is not None:
                factor_name = getattr(exp.hypothesis, 'hypothesis', 'unknown')

            # Check if already rejected by protection
            if getattr(exp, 'rejected_by_protection', False):
                logger.info(
                    f"Factor rejected by protection, skipping DB save: "
                    f"{getattr(exp, 'protection_reason', 'unknown')}"
                )
                return

            # Log result structure for debugging
            logger.info(f"Saving result for factor: {factor_name}")
            logger.info(f"Result type: {type(result)}")
            if isinstance(result, pd.Series):
                logger.info(f"Result index (metric names): {list(result.index)}")
                logger.info(f"Result values:\n{result.to_string()}")
            elif isinstance(result, dict):
                logger.info(f"Result keys: {list(result.keys())}")

            # Extract metrics from result (pd.Series from qlib_res.csv)
            metrics = {}
            if isinstance(result, pd.Series):
                metrics['ic'] = self._safe_float(result.get('IC', None))
                metrics['sharpe_ratio'] = self._safe_float(
                    result.get('1day.excess_return_with_cost.shar',
                    result.get('1day.excess_return_with_cost.sharpe', None))
                )
                metrics['annualized_return'] = self._safe_float(
                    result.get('1day.excess_return_with_cost.annualized_return', None)
                )
                metrics['max_drawdown'] = self._safe_float(
                    result.get('1day.excess_return_with_cost.max_drawdown', None)
                )
                metrics['win_rate'] = self._safe_float(result.get('win_rate', None))
                metrics['information_ratio'] = self._safe_float(
                    result.get('1day.excess_return_with_cost.information_ratio', None)
                )
                metrics['volatility'] = self._safe_float(
                    result.get('1day.excess_return_with_cost.std',
                    result.get('1day.excess_return_with_cost.volatility', None))
                )
                # Store raw metrics for JSON export
                metrics['raw_metrics'] = result.to_dict()
            elif isinstance(result, dict):
                metrics['ic'] = self._safe_float(result.get('IC', result.get('ic', None)))
                metrics['sharpe_ratio'] = self._safe_float(
                    result.get('sharpe', result.get('sharpe_ratio', None))
                )
                metrics['annualized_return'] = self._safe_float(result.get('annualized_return', None))
                metrics['max_drawdown'] = self._safe_float(result.get('max_drawdown', None))
                metrics['win_rate'] = self._safe_float(result.get('win_rate', None))
                metrics['information_ratio'] = None
                metrics['volatility'] = None
                metrics['raw_metrics'] = result

            # Only save if we have at least IC or Sharpe
            if metrics.get('ic') is None and metrics.get('sharpe_ratio') is None:
                logger.warning(
                    f"No valid IC/Sharpe for factor '{factor_name}', skipping DB save. "
                    f"IC={metrics.get('ic')}, Sharpe={metrics.get('sharpe_ratio')}"
                )
                return

            # Ensure DB directory exists before creating database
            db_path = Path(__file__).parent.parent.parent.parent / "results" / "db"
            db_path.mkdir(parents=True, exist_ok=True)
            db_file = db_path / "backtest_results.db"

            # Save to database
            db = ResultsDatabase(db_path=str(db_file))
            run_id = db.add_backtest(factor_name=factor_name[:100], metrics=metrics)
            logger.info(
                f"Factor result saved to DB: {factor_name[:60]} "
                f"(IC={metrics.get('ic')}, Sharpe={metrics.get('sharpe_ratio')}, run_id={run_id})"
            )

            # Also write a JSON summary to results/factors/ for file-based access
            self._save_factor_json(factor_name, metrics, run_id)

            db.close()

        except Exception as e:
            import traceback
            logger.error(
                f"Database save failed for factor '{getattr(exp.hypothesis, 'hypothesis', 'unknown')}': {e}\n"
                f"Traceback: {traceback.format_exc()}"
            )

    def _save_factor_json(self, factor_name: str, metrics: dict, run_id: int) -> None:
        """
        Save factor metrics as a JSON file for easy file-based access.

        Parameters
        ----------
        factor_name : str
            Name of the factor
        metrics : dict
            Extracted metrics dictionary
        run_id : int
            Database run ID
        """
        import json
        from datetime import datetime
        from pathlib import Path

        try:
            # Ensure factors directory exists
            factors_dir = Path(__file__).parent.parent.parent.parent / "results" / "factors"
            factors_dir.mkdir(parents=True, exist_ok=True)

            # Sanitize factor name for filename
            safe_name = factor_name.replace("/", "_").replace("\\", "_").replace(" ", "_")
            # Truncate if too long (filesystem limits)
            if len(safe_name) > 150:
                safe_name = safe_name[:150]

            json_path = factors_dir / f"{safe_name}.json"

            # Build summary document
            factor_summary = {
                "factor_name": factor_name,
                "run_id": run_id,
                "saved_at": datetime.now().isoformat(),
                "metrics": {
                    "ic": metrics.get("ic"),
                    "sharpe_ratio": metrics.get("sharpe_ratio"),
                    "annualized_return": metrics.get("annualized_return"),
                    "max_drawdown": metrics.get("max_drawdown"),
                    "win_rate": metrics.get("win_rate"),
                    "information_ratio": metrics.get("information_ratio"),
                    "volatility": metrics.get("volatility"),
                },
                "raw_metrics": {
                    k: v for k, v in metrics.get("raw_metrics", {}).items()
                    if isinstance(v, (int, float, str, bool, type(None)))
                },
            }

            json_path.write_text(json.dumps(factor_summary, indent=2, default=str), encoding="utf-8")
            logger.info(f"Factor JSON saved: {json_path}")

        except Exception as e:
            logger.warning(f"Failed to save factor JSON for '{factor_name}': {e}")

    def _safe_float(self, value):
        """Safely convert value to float, returning None for invalid values."""
        import pandas as pd
        if value is None:
            return None
        try:
            f = float(value)
            if pd.isna(f) or f == float('inf') or f == float('-inf'):
                return None
            return f
        except (ValueError, TypeError):
            return None

    def _run_protection_check(self, exp, result: dict) -> None:
        """
        Run protection checks on backtest results.

        Parameters
        ----------
        exp : QlibFactorExperiment
            The experiment with backtest results
        result : dict
            Backtest metrics dictionary
        """
        from rdagent.components.backtesting.protections import ProtectionManager

        protection_manager = ProtectionManager()
        protection_manager.create_default_protections()

        # Extract returns and equity curve from backtest results
        returns = result.get("returns", [])
        timestamps = result.get("timestamps", [])
        current_equity = result.get("final_equity", 100000)
        peak_equity = result.get("peak_equity", current_equity)

        # Get factor name from hypothesis
        factor_name = "unknown"
        if hasattr(exp, "hypothesis") and exp.hypothesis is not None:
            factor_name = getattr(exp.hypothesis, "hypothesis", "unknown")

        protection_result = protection_manager.check_all(
            returns=returns,
            timestamps=timestamps,
            current_equity=current_equity,
            peak_equity=peak_equity,
            factor_name=factor_name,
        )

        if protection_result.should_block:
            logger.warning(
                f"Factor {factor_name} rejected by protection manager: {protection_result.reason}"
            )
            # Mark factor as rejected by protection
            exp.rejected_by_protection = True
            exp.protection_reason = protection_result.reason