import os.path import httpx import json from httpx import BasicAuth import time from random import uniform import threading from concurrent.futures import ThreadPoolExecutor, as_completed 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): """读取本地账号密码""" 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("登录成功!") return True else: print(f"登录失败: {response.json()}") return False def simulate_alpha(self, expression, settings=None): """模拟Alpha因子""" 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} result = sim_progress_resp.json()["alpha"] print(f"生成的Alpha ID: {result}") return {"status": "success", "alpha_id": result} def close(self): """关闭连接""" if self.client: self.client.close() class AlphaSimulationManager: def __init__(self, credentials_file='account.txt'): self.credentials_file = credentials_file self.results = [] def format_time(self, seconds): """将秒数格式化为 xx分xx秒 格式""" if seconds < 60: return f"{seconds:.2f}秒" else: minutes = int(seconds // 60) remaining_seconds = seconds % 60 return f"{minutes}分{remaining_seconds:.2f}秒" def simulate_single_alpha(self, api, expression, settings=None): """模拟单个Alpha因子（线程安全）""" alpha_start_time = time.time() try: # 模拟Alpha因子 simulation_result = api.simulate_alpha(expression, settings) alpha_end_time = time.time() time_consuming = alpha_end_time - alpha_start_time # 根据模拟结果类型处理 if simulation_result["status"] == "success": # 模拟成功的结果 result = { "expression": expression, "time_consuming": time_consuming, "formatted_time": self.format_time(time_consuming), "alpha_id": simulation_result["alpha_id"], "status": "success", "description": "/" } print(f"✓ 因子模拟成功: {expression}") print(f" 耗时: {self.format_time(time_consuming)}，Alpha ID: {simulation_result['alpha_id']}") else: # 模拟失败的结果（API返回的错误） result = { "expression": expression, "time_consuming": time_consuming, "formatted_time": self.format_time(time_consuming), "alpha_id": "/", "status": "error", "description": simulation_result["message"] } print(f"✗ 因子模拟失败: {expression}") print(f" 耗时: {self.format_time(time_consuming)}，错误: {simulation_result['message']}") except Exception as e: # 其他异常情况 alpha_end_time = time.time() time_consuming = alpha_end_time - alpha_start_time result = { "expression": expression, "time_consuming": time_consuming, "formatted_time": self.format_time(time_consuming), "alpha_id": "/", "status": "failed", "description": str(e) } print(f"✗ 因子模拟异常: {expression}") print(f" 耗时: {self.format_time(time_consuming)}，异常: {str(e)}") return result def simulate_alpha_batch(self, alpha_batch, batch_number): """模拟一批Alpha因子（3个一组）""" print(f"\n{'=' * 60}") print(f"开始第 {batch_number} 批因子模拟 (共 {len(alpha_batch)} 个因子)") print(f"因子列表: {alpha_batch}") print(f"{'=' * 60}") batch_start_time = time.time() batch_results = [] # 创建API客户端实例（每个线程独立的客户端） api = WorldQuantBrainSimulate(self.credentials_file) try: if api.login(): # 使用线程池执行3个因子的模拟 with ThreadPoolExecutor(max_workers=3) as executor: # 提交所有任务 future_to_alpha = {executor.submit(self.simulate_single_alpha, api, alpha): alpha for alpha in alpha_batch} # 等待所有任务完成 for future in as_completed(future_to_alpha): alpha = future_to_alpha[future] try: result = future.result() batch_results.append(result) except Exception as e: print(f"因子 {alpha} 执行异常: {e}") except Exception as e: print(f"第 {batch_number} 批模拟过程中出错: {e}") finally: api.close() batch_end_time = time.time() batch_total_time = batch_end_time - batch_start_time print(f"\n第 {batch_number} 批模拟完成!") print(f"本批总耗时: {self.format_time(batch_total_time)}") print(f"{'=' * 60}") return batch_results def run_simulation(self, alpha_list, batch_size=3): """运行批量模拟""" print("开始Alpha因子批量模拟...") total_start_time = time.time() # 将因子列表分成每批3个 batches = [alpha_list[i:i + batch_size] for i in range(0, len(alpha_list), batch_size)] all_results = [] for i, batch in enumerate(batches, 1): # 模拟当前批次 batch_results = self.simulate_alpha_batch(batch, i) all_results.extend(batch_results) # 如果不是最后一批，则等待3-5秒 if i < len(batches): sleep_time = uniform(3, 5) print(f"\n等待 {sleep_time:.2f} 秒后开始下一批...") time.sleep(sleep_time) total_end_time = time.time() total_time = total_end_time - total_start_time # 输出最终结果汇总 self.print_summary(all_results, total_time) # 保存结果到文件 self.save_results(all_results) return all_results def print_summary(self, results, total_time): """打印结果汇总""" print(f"\n{'=' * 60}") print("模拟结果汇总") print(f"{'=' * 60}") success_count = sum(1 for r in results if r['status'] == 'success') error_count = sum(1 for r in results if r['status'] == 'error') failed_count = sum(1 for r in results if r['status'] == 'failed') print(f"总模拟因子数: {len(results)}") print(f"成功: {success_count} 个") print(f"模拟错误: {error_count} 个") print(f"执行异常: {failed_count} 个") print(f"总耗时: {self.format_time(total_time)}") print(f"{'=' * 60}") for i, result in enumerate(results, 1): status_icon = "✓" if result['status'] == 'success' else "✗" print(f"{i}. {status_icon} {result['expression']}") print(f" 状态: {result['status']}") print(f" 耗时: {result['formatted_time']}") print(f" Alpha ID: {result['alpha_id']}") if result['status'] != 'success': print(f" 原因: {result['description']}") print() def save_results(self, results): """保存结果到文件""" # 转换为可序列化的格式 serializable_results = [] for result in results: serializable_result = result.copy() serializable_result['time_consuming'] = round(serializable_result['time_consuming'], 2) serializable_results.append(serializable_result) # 将日志文件, 保存到当前目录下, result 文件夹中 if not os.path.exists('./result'): os.makedirs('./result') result_name = f"result/simulation_results-{str(int(time.time()))}.json" with open(result_name, 'w', encoding='utf-8') as f: json.dump(serializable_results, f, ensure_ascii=False, indent=2) print(f"结果已保存到 {result_name}") if __name__ == "__main__": # 待模拟因子列表 with open('alpha.txt', 'r', encoding='utf-8') as file: alpha_list = [line.strip() for line in file] if not alpha_list: print("alpha.txt 文件不存在") with open('alpha.txt', 'w', encoding='utf-8') as file: file.write("") print("已创建 alpha.txt 文件, 请添加因子后重新运行, 一行一个因子") exit(1) # 创建模拟管理器并运行 manager = AlphaSimulationManager() results = manager.run_simulation(alpha_list, batch_size=3)