# -*- coding: utf-8 -*- import os.path import httpx import time from httpx import BasicAuth from typing import Dict, Any, Optional, Tuple from .models import AlphaMetrics, TrainMetrics, TestMetrics, AlphaInfo class WorldQuantBrainSimulate: 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) -> Tuple[str, str]: 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) -> bool: 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("登录成功!") return True else: print(f"登录失败: {response.json()}") return False """模拟Alpha因子""" def simulate_alpha(self, expression: str, settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: if self.client is None: raise Exception("请先登录") default_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, } if settings: default_settings.update(settings) simulation_data = { 'type': 'REGULAR', 'settings': default_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'] print(f"进度URL: {sim_progress_url}") 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 print(sim_progress_resp.json()) print(f"等待 {retry_after_sec} 秒...") time.sleep(retry_after_sec) # 如果因子模拟不通过, 获取一下失败信息 if sim_progress_resp.json()["status"] == "ERROR": result = sim_progress_resp.json()["message"] print(f"因子模拟失败: {result}") # 返回一个特殊标识，表示模拟失败 return {"status": "error", "message": result} alpha_id = sim_progress_resp.json()["alpha"] print(f"生成的Alpha ID: {alpha_id}") # 获取详细的性能指标 metrics = self.get_alpha_metrics(alpha_id) return {"status": "success", "alpha_id": alpha_id, "metrics": metrics} """获取Alpha因子的详细指标""" def get_alpha_metrics(self, alpha_id: str) -> AlphaMetrics: if self.client is None: raise Exception("请先登录") try: # 获取Alpha的基本信息和指标 alpha_url = f'{self.brain_api_url}/alphas/{alpha_id}' alpha_resp = self.client.get(alpha_url) if alpha_resp.status_code in [200, 201]: alpha_data = alpha_resp.json() return self._parse_alpha_metrics(alpha_data) else: return AlphaMetrics( train_metrics=TrainMetrics(), is_metrics=TestMetrics(), test_metrics=TestMetrics(), alpha_info=AlphaInfo() ) except Exception as e: print(f"获取指标时出错: {str(e)}") return AlphaMetrics( train_metrics=TrainMetrics(), is_metrics=TestMetrics(), test_metrics=TestMetrics(), alpha_info=AlphaInfo() ) """解析Alpha数据，提取关键指标""" def _parse_alpha_metrics(self, alpha_data: Dict[str, Any]) -> AlphaMetrics: # 解析训练集数据 train_metrics = TrainMetrics() if 'train' in alpha_data and alpha_data['train']: train_data = alpha_data['train'] train_metrics = TrainMetrics( sharpe_ratio=train_data.get('sharpe'), annual_return=train_data.get('returns'), max_drawdown=train_data.get('drawdown'), turnover=train_data.get('turnover'), fitness=train_data.get('fitness'), pnl=train_data.get('pnl'), book_size=train_data.get('bookSize'), long_count=train_data.get('longCount'), short_count=train_data.get('shortCount'), margin=train_data.get('margin'), ) # 解析样本内测试数据 is_metrics = TestMetrics() if 'is' in alpha_data and alpha_data['is']: is_data = alpha_data['is'] is_metrics = TestMetrics( sharpe_ratio=is_data.get('sharpe'), annual_return=is_data.get('returns'), max_drawdown=is_data.get('drawdown'), turnover=is_data.get('turnover'), fitness=is_data.get('fitness'), pnl=is_data.get('pnl'), ) # 解析样本外测试数据 test_metrics = TestMetrics() if 'test' in alpha_data and alpha_data['test']: test_data = alpha_data['test'] test_metrics = TestMetrics( sharpe_ratio=test_data.get('sharpe'), annual_return=test_data.get('returns'), max_drawdown=test_data.get('drawdown'), turnover=test_data.get('turnover'), fitness=test_data.get('fitness'), pnl=test_data.get('pnl'), ) # 解析Alpha基本信息 alpha_info = AlphaInfo( grade=alpha_data.get('grade'), stage=alpha_data.get('stage'), status=alpha_data.get('status'), date_created=alpha_data.get('dateCreated'), ) # 解析检查结果 if 'is' in alpha_data and 'checks' in alpha_data['is']: checks = alpha_data['is']['checks'] check_results = {} for check in checks: check_name = check.get('name', '') result = check.get('result', '') value = check.get('value', None) check_results[check_name.lower()] = { 'result': result, 'value': value, 'limit': check.get('limit', None) } alpha_info.checks = check_results return AlphaMetrics( train_metrics=train_metrics, is_metrics=is_metrics, test_metrics=test_metrics, alpha_info=alpha_info, alpha_id=alpha_data.get('id') ) def close(self): """关闭连接""" if self.client: self.client.close()