wqb/FactorSimulator.py

# -*- coding: utf-8 -*-
import os
import time
import json
import httpx
from httpx import BasicAuth


class AlphaSimulator:
    def __init__(self, credentials_file='account.txt'):
        self.credentials_file = credentials_file
        self.client = None
        self.brain_api_url = 'https://api.worldquantbrain.com'

    def load_credentials(self):
        if not os.path.exists(self.credentials_file):
            print("未找到 account.txt 文件")
            with open(self.credentials_file, 'w') as f:
                f.write("")
            print("account.txt 文件已创建，请填写账号密码, 格式: ['username', 'password]")
            exit(1)

        with open(self.credentials_file) as f:
            credentials = eval(f.read())
        return credentials[0], credentials[1]

    def login(self):
        username, password = self.load_credentials()
        self.client = httpx.Client(auth=BasicAuth(username, password))

        response = self.client.post(f'{self.brain_api_url}/authentication')
        print(f"登录状态: {response.status_code}")

        if response.status_code == 201:
            print(f"登录成功!:{response.json()}")
            return True
        else:
            print(f"登录失败: {response.json()}")
            return False

    def load_alpha_list(self, file_path='alpha.txt'):
        if not os.path.exists(file_path):
            print(f"{file_path} 文件不存在")
            with open(file_path, 'w', encoding='utf-8') as file:
                file.write("")
            print(f"已创建 {file_path} 文件, 请添加因子后重新运行, 一行一个因子")
            return []

        with open(file_path, 'r', encoding='utf-8') as file:
            alpha_list = [line.strip() for line in file if line.strip()]

        return alpha_list

    def simulate_alpha(self, expression):
        if self.client is None:
            raise Exception("请先登录")

        settings = {
            'instrumentType': 'EQUITY',
            'region': 'USA',
            'universe': 'TOP3000',
            'delay': 1,
            'decay': 0,
            'neutralization': 'INDUSTRY',
            'truncation': 0.08,
            'pasteurization': 'ON',
            'unitHandling': 'VERIFY',
            'nanHandling': 'OFF',
            'language': 'FASTEXPR',
            'visualization': False,
        }

        simulation_data = {
            'type': 'REGULAR',
            'settings': settings,
            'regular': expression
        }

        sim_resp = self.client.post(f'{self.brain_api_url}/simulations', json=simulation_data)
        print(f"模拟提交状态: {sim_resp.status_code}")

        sim_progress_url = sim_resp.headers['location']

        while True:
            sim_progress_resp = self.client.get(sim_progress_url)
            retry_after_sec = float(sim_progress_resp.headers.get("Retry-After", 0))

            if retry_after_sec == 0:
                break

            if sim_progress_resp.json():
                result = sim_progress_resp.json()
                progress = result.get('progress', 0)
                print(f"模拟进度: {float(progress) * 100}%")

            print(f"等待 {retry_after_sec} 秒...")
            time.sleep(retry_after_sec)

        if sim_progress_resp.json().get("status") == "ERROR":
            result = sim_progress_resp.json().get("message", "未知错误")
            print(f"因子模拟失败: {result}")
            return {"status": "error", "message": result}

        alpha_id = sim_progress_resp.json().get("alpha")
        print(f"生成的Alpha ID: {alpha_id}")

        return {"status": "success", "alpha_id": alpha_id}

    def run_batch_simulation(self, alpha_list, batch_size=3):
        success_results = []
        fail_results = []

        for i in range(0, len(alpha_list), batch_size):
            batch = alpha_list[i:i + batch_size]
            print(f"\n开始处理第 {i // batch_size + 1} 批因子，共 {len(batch)} 个")

            for expression in batch:
                print(f"\n模拟因子: {expression}")
                start_time = time.time()

                try:
                    result = self.simulate_alpha(expression)
                    end_time = time.time()
                    time_consuming = end_time - start_time

                    result_data = {
                        "expression": expression,
                        "time_consuming": round(time_consuming, 2),
                        "status": result["status"],
                        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
                    }

                    if result["status"] == "success":
                        result_data["alpha_id"] = result["alpha_id"]
                        success_results.append(result_data)
                        print(f"✅ 模拟成功 - Alpha ID: {result['alpha_id']} - 耗时: {time_consuming:.2f}秒")
                    else:
                        result_data["message"] = result.get("message", "")
                        fail_results.append(result_data)
                        print(f"❌ 模拟失败 - {result.get('message', '未知错误')}")

                except Exception as e:
                    end_time = time.time()
                    time_consuming = end_time - start_time
                    error_result = {
                        "expression": expression,
                        "time_consuming": round(time_consuming, 2),
                        "status": "error",
                        "message": str(e),
                        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
                    }
                    fail_results.append(error_result)
                    print(f"❌ 模拟异常 - {str(e)}")

            print(f"第 {i // batch_size + 1} 批处理完成")

        self._save_results(success_results, fail_results)
        return success_results, fail_results

    def _save_results(self, success_results, fail_results):
        base_dir = 'simulator_result'
        success_dir = os.path.join(base_dir, 'success')
        fail_dir = os.path.join(base_dir, 'fail')

        os.makedirs(success_dir, exist_ok=True)
        os.makedirs(fail_dir, exist_ok=True)

        timestamp = str(int(time.time()))

        if success_results:
            success_file = os.path.join(success_dir, f"success_results_{timestamp}.json")
            with open(success_file, 'w', encoding='utf-8') as f:
                json.dump(success_results, f, ensure_ascii=False, indent=2)
            print(f"成功结果已保存到 {success_file}")

        if fail_results:
            fail_file = os.path.join(fail_dir, f"fail_results_{timestamp}.json")
            with open(fail_file, 'w', encoding='utf-8') as f:
                json.dump(fail_results, f, ensure_ascii=False, indent=2)
            print(f"失败结果已保存到 {fail_file}")

        print(f"\n总计: 成功 {len(success_results)} 个, 失败 {len(fail_results)} 个")


def main():
    simulator = AlphaSimulator()

    if not simulator.login():
        return

    alpha_list = simulator.load_alpha_list('alpha.txt')
    if not alpha_list:
        print("未找到有效的因子表达式，请检查 alpha.txt 文件")
        return

    print(f"共加载 {len(alpha_list)} 个因子表达式")

    simulator.run_batch_simulation(alpha_list, batch_size=3)


if __name__ == "__main__":
    main()