432 lines
12 KiB
Go
432 lines
12 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"fmt"
|
||
"net/http"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
// BrainClient 是WorldQuant Brain API的客户端
|
||
// 类似于Python中的WorldQuantBrainSimulate类
|
||
type BrainClient struct {
|
||
httpClient *http.Client // HTTP客户端，负责发送请求
|
||
baseURL string // API的基础URL
|
||
username string // 用户名
|
||
password string // 密码
|
||
isLoggedIn bool // 登录状态
|
||
}
|
||
|
||
// NewBrainClient 创建一个新的Brain API客户端
|
||
// 这是Go中的构造函数，类似于Python的__init__
|
||
func NewBrainClient(username, password string) *BrainClient {
|
||
return &BrainClient{
|
||
// 创建HTTP客户端，设置30秒超时
|
||
httpClient: &http.Client{
|
||
Timeout: 30 * time.Second,
|
||
},
|
||
baseURL: "https://api.worldquantbrain.com",
|
||
username: username,
|
||
password: password,
|
||
isLoggedIn: false,
|
||
}
|
||
}
|
||
|
||
// Login 登录Brain API
|
||
// 对应Python中的login方法
|
||
func (c *BrainClient) Login() error {
|
||
// 1. 创建POST请求
|
||
req, err := http.NewRequest("POST", c.baseURL+"/authentication", nil)
|
||
if err != nil {
|
||
return fmt.Errorf("创建登录请求失败: %v", err)
|
||
}
|
||
|
||
// 2. 设置Basic认证（用户名和密码）
|
||
req.SetBasicAuth(c.username, c.password)
|
||
|
||
// 3. 发送请求
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("发送登录请求失败: %v", err)
|
||
}
|
||
defer resp.Body.Close() // 重要：确保函数结束时关闭响应体
|
||
|
||
fmt.Printf("登录状态: %d\n", resp.StatusCode)
|
||
|
||
// 4. 检查响应状态码
|
||
if resp.StatusCode == 201 {
|
||
fmt.Println("登录成功!")
|
||
c.isLoggedIn = true
|
||
return nil
|
||
}
|
||
|
||
// 5. 登录失败，读取错误信息
|
||
var errorResp map[string]interface{}
|
||
if err := json.NewDecoder(resp.Body).Decode(&errorResp); err != nil {
|
||
return fmt.Errorf("登录失败，状态码: %d", resp.StatusCode)
|
||
}
|
||
return fmt.Errorf("登录失败: %v", errorResp)
|
||
}
|
||
|
||
// SimulateAlpha 模拟Alpha因子
|
||
// 对应Python中的simulate_alpha方法
|
||
func (c *BrainClient) SimulateAlpha(expression string, settings map[string]interface{}) (*SimulationResult, error) {
|
||
// 检查是否已登录
|
||
if !c.isLoggedIn {
|
||
return nil, fmt.Errorf("请先调用Login()登录")
|
||
}
|
||
|
||
startTime := time.Now() // 记录开始时间
|
||
|
||
// ========== 1. 准备模拟设置 ==========
|
||
// 创建默认设置（对应Python中的default_settings）
|
||
defaultSettings := SimulationSettings{
|
||
InstrumentType: "EQUITY",
|
||
Region: "USA",
|
||
Universe: "TOP3000",
|
||
Delay: 1,
|
||
Decay: 12,
|
||
Truncation: 0.05,
|
||
Neutralization: "FAST",
|
||
Pasteurization: "ON",
|
||
UnitHandling: "VERIFY",
|
||
NanHandling: "ON",
|
||
Language: "FASTEXPR",
|
||
Visualization: false,
|
||
}
|
||
|
||
// 用传入的设置覆盖默认设置
|
||
// 注意：这里需要安全地进行类型断言
|
||
if region, ok := settings["region"]; ok {
|
||
if regionStr, ok := region.(string); ok {
|
||
defaultSettings.Region = regionStr
|
||
}
|
||
}
|
||
if universe, ok := settings["universe"]; ok {
|
||
if universeStr, ok := universe.(string); ok {
|
||
defaultSettings.Universe = universeStr
|
||
}
|
||
}
|
||
if instrumentType, ok := settings["instrumentType"]; ok {
|
||
if instrumentTypeStr, ok := instrumentType.(string); ok {
|
||
defaultSettings.InstrumentType = instrumentTypeStr
|
||
}
|
||
}
|
||
if decay, ok := settings["decay"]; ok {
|
||
// 注意：JSON中的数字可能是float64类型
|
||
if decayFloat, ok := decay.(float64); ok {
|
||
defaultSettings.Decay = int(decayFloat)
|
||
} else if decayInt, ok := decay.(int); ok {
|
||
defaultSettings.Decay = decayInt
|
||
}
|
||
}
|
||
if truncation, ok := settings["truncation"]; ok {
|
||
if truncationFloat, ok := truncation.(float64); ok {
|
||
defaultSettings.Truncation = truncationFloat
|
||
}
|
||
}
|
||
if neutralization, ok := settings["neutralization"]; ok {
|
||
if neutralizationStr, ok := neutralization.(string); ok {
|
||
defaultSettings.Neutralization = neutralizationStr
|
||
}
|
||
}
|
||
// 可以根据需要添加更多字段的覆盖逻辑
|
||
|
||
// ========== 2. 构建请求数据 ==========
|
||
simRequest := SimulationRequest{
|
||
Type: "REGULAR",
|
||
Settings: defaultSettings,
|
||
Regular: expression,
|
||
}
|
||
|
||
// 将结构体转换为JSON字节数组
|
||
simData, err := json.Marshal(simRequest)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("构建模拟请求数据失败: %v", err)
|
||
}
|
||
|
||
// ========== 3. 发送模拟请求 ==========
|
||
req, err := http.NewRequest("POST", c.baseURL+"/simulations", bytes.NewBuffer(simData))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("创建模拟请求失败: %v", err)
|
||
}
|
||
|
||
// 设置请求头
|
||
req.SetBasicAuth(c.username, c.password)
|
||
req.Header.Set("Content-Type", "application/json")
|
||
|
||
// 发送请求
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("发送模拟请求失败: %v", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
fmt.Printf("模拟提交状态: %d\n", resp.StatusCode)
|
||
|
||
// 检查响应状态
|
||
if resp.StatusCode != 201 && resp.StatusCode != 200 {
|
||
var errorResp map[string]interface{}
|
||
json.NewDecoder(resp.Body).Decode(&errorResp)
|
||
return nil, fmt.Errorf("模拟请求失败: %v", errorResp)
|
||
}
|
||
|
||
// ========== 4. 获取轮询URL ==========
|
||
location := resp.Header.Get("Location")
|
||
if location == "" {
|
||
return nil, fmt.Errorf("未获取到模拟进度URL")
|
||
}
|
||
fmt.Printf("进度URL: %s\n", location)
|
||
|
||
// ========== 5. 轮询等待结果 ==========
|
||
// 对应Python中的while循环
|
||
var finalResp map[string]interface{}
|
||
for {
|
||
// 创建轮询请求
|
||
pollReq, err := http.NewRequest("GET", location, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("创建轮询请求失败: %v", err)
|
||
}
|
||
pollReq.SetBasicAuth(c.username, c.password)
|
||
|
||
// 发送轮询请求
|
||
pollResp, err := c.httpClient.Do(pollReq)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("轮询请求失败: %v", err)
|
||
}
|
||
|
||
// 检查Retry-After头（告诉我们需要等待多久）
|
||
retryAfter := pollResp.Header.Get("Retry-After")
|
||
|
||
if retryAfter != "" {
|
||
// 需要等待，解析等待时间
|
||
waitSeconds, err := strconv.ParseFloat(retryAfter, 64)
|
||
if err != nil {
|
||
waitSeconds = 1.0 // 解析失败则默认等待1秒
|
||
}
|
||
fmt.Printf("等待 %.2f 秒...\n", waitSeconds)
|
||
|
||
// 先关闭当前响应体
|
||
pollResp.Body.Close()
|
||
// 等待指定时间
|
||
time.Sleep(time.Duration(waitSeconds) * time.Second)
|
||
// 继续轮询
|
||
continue
|
||
}
|
||
|
||
// 解析响应JSON
|
||
err = json.NewDecoder(pollResp.Body).Decode(&finalResp)
|
||
pollResp.Body.Close()
|
||
|
||
if err != nil {
|
||
return nil, fmt.Errorf("解析轮询响应失败: %v", err)
|
||
}
|
||
|
||
// 检查状态是否为ERROR
|
||
if status, ok := finalResp["status"]; ok && status == "ERROR" {
|
||
message := "未知错误"
|
||
if msg, ok := finalResp["message"].(string); ok {
|
||
message = msg
|
||
}
|
||
return nil, fmt.Errorf("因子模拟失败: %s", message)
|
||
}
|
||
|
||
// 检查是否已经完成（包含alpha字段）
|
||
if alphaID, ok := finalResp["alpha"]; ok && alphaID != nil {
|
||
// 模拟完成！
|
||
alphaIDStr := fmt.Sprintf("%v", alphaID)
|
||
fmt.Printf("生成的Alpha ID: %s\n", alphaIDStr)
|
||
|
||
// 获取详细指标
|
||
metrics, err := c.getAlphaMetrics(alphaIDStr)
|
||
if err != nil {
|
||
fmt.Printf("警告: 获取Alpha指标失败: %v\n", err)
|
||
// 即使获取指标失败，也继续返回基础信息
|
||
}
|
||
|
||
// 计算耗时
|
||
elapsed := time.Since(startTime).Seconds()
|
||
|
||
// 返回成功结果
|
||
return &SimulationResult{
|
||
Status: "success",
|
||
Expression: expression,
|
||
AlphaID: alphaIDStr,
|
||
TimeCost: elapsed,
|
||
FormattedTime: formatTime(elapsed),
|
||
Timestamp: time.Now().Format("2006-01-02 15:04:05"),
|
||
Metrics: metrics,
|
||
Message: "",
|
||
}, nil
|
||
}
|
||
|
||
// 如果既没有Retry-After，也没有alpha字段，可能是其他状态
|
||
// 打印当前状态以便调试
|
||
fmt.Printf("当前状态: %v\n", finalResp)
|
||
|
||
// 为了避免无限循环，如果没有Retry-After也没有完成，就退出
|
||
// 但在实际应用中，可能需要继续等待或处理其他状态
|
||
break
|
||
}
|
||
|
||
// 如果走到这里，说明轮询异常退出
|
||
elapsed := time.Since(startTime).Seconds()
|
||
return &SimulationResult{
|
||
Status: "failed",
|
||
Expression: expression,
|
||
AlphaID: "",
|
||
TimeCost: elapsed,
|
||
FormattedTime: formatTime(elapsed),
|
||
Timestamp: time.Now().Format("2006-01-02 15:04:05"),
|
||
Metrics: nil,
|
||
Message: "轮询超时或状态异常",
|
||
}, nil
|
||
}
|
||
|
||
// getAlphaMetrics 获取Alpha因子的详细指标
|
||
// 对应Python中的get_alpha_metrics和_parse_alpha_metrics
|
||
func (c *BrainClient) getAlphaMetrics(alphaID string) (*PerformanceMetrics, error) {
|
||
// 构建请求URL
|
||
url := fmt.Sprintf("%s/alphas/%s", c.baseURL, alphaID)
|
||
|
||
// 创建GET请求
|
||
req, err := http.NewRequest("GET", url, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("创建请求失败: %v", err)
|
||
}
|
||
req.SetBasicAuth(c.username, c.password)
|
||
|
||
// 发送请求
|
||
resp, err := c.httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("请求失败: %v", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// 检查响应状态
|
||
if resp.StatusCode != 200 {
|
||
return nil, fmt.Errorf("HTTP状态码: %d", resp.StatusCode)
|
||
}
|
||
|
||
// 解析JSON
|
||
var alphaData map[string]interface{}
|
||
err = json.NewDecoder(resp.Body).Decode(&alphaData)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("解析JSON失败: %v", err)
|
||
}
|
||
|
||
// 创建指标对象
|
||
metrics := &PerformanceMetrics{}
|
||
|
||
// ========== 解析returns字段 ==========
|
||
if returns, ok := alphaData["returns"].(map[string]interface{}); ok {
|
||
if sharpe, ok := returns["sharpe"].(float64); ok {
|
||
metrics.SharpeRatio = &sharpe
|
||
}
|
||
if annualReturn, ok := returns["annualReturn"].(float64); ok {
|
||
metrics.AnnualReturn = &annualReturn
|
||
}
|
||
if annualVolatility, ok := returns["annualVolatility"].(float64); ok {
|
||
metrics.AnnualVolatility = &annualVolatility
|
||
}
|
||
if maxDrawdown, ok := returns["maxDrawdown"].(float64); ok {
|
||
metrics.MaxDrawdown = &maxDrawdown
|
||
}
|
||
if informationRatio, ok := returns["informationRatio"].(float64); ok {
|
||
metrics.InformationRatio = &informationRatio
|
||
}
|
||
if tailRatio, ok := returns["tailRatio"].(float64); ok {
|
||
metrics.TailRatio = &tailRatio
|
||
}
|
||
if commonRatio, ok := returns["commonRatio"].(float64); ok {
|
||
metrics.CommonRatio = &commonRatio
|
||
}
|
||
}
|
||
|
||
// ========== 解析riskAdjustment字段 ==========
|
||
if riskAdj, ok := alphaData["riskAdjustment"].(map[string]interface{}); ok {
|
||
if score, ok := riskAdj["score"].(float64); ok {
|
||
metrics.Score = &score
|
||
}
|
||
if turnover, ok := riskAdj["turnover"].(float64); ok {
|
||
metrics.Turnover = &turnover
|
||
}
|
||
if specificReturn, ok := riskAdj["specificReturn"].(float64); ok {
|
||
metrics.SpecificReturn = &specificReturn
|
||
}
|
||
if specificRisk, ok := riskAdj["specificRisk"].(float64); ok {
|
||
metrics.SpecificRisk = &specificRisk
|
||
}
|
||
}
|
||
|
||
// ========== 解析quantiles字段 ==========
|
||
if quantiles, ok := alphaData["quantiles"].(map[string]interface{}); ok {
|
||
if topMinusBottom, ok := quantiles["topMinusBottom"].(float64); ok {
|
||
metrics.TopMinusBottom = &topMinusBottom
|
||
}
|
||
if topDecileReturn, ok := quantiles["topDecileReturn"].(float64); ok {
|
||
metrics.TopDecileReturn = &topDecileReturn
|
||
}
|
||
if bottomDecileReturn, ok := quantiles["bottomDecileReturn"].(float64); ok {
|
||
metrics.BottomDecileReturn = &bottomDecileReturn
|
||
}
|
||
if ic, ok := quantiles["ic"].(float64); ok {
|
||
metrics.IC = &ic
|
||
}
|
||
if icDecay, ok := quantiles["icDecay"].(float64); ok {
|
||
metrics.ICDecay = &icDecay
|
||
}
|
||
}
|
||
|
||
// ========== 解析其他字段 ==========
|
||
if totalReturn, ok := alphaData["totalReturn"].(float64); ok {
|
||
metrics.TotalReturn = &totalReturn
|
||
}
|
||
if capacity, ok := alphaData["capacity"].(float64); ok {
|
||
metrics.Capacity = &capacity
|
||
}
|
||
if fitness, ok := alphaData["fitness"].(float64); ok {
|
||
metrics.Fitness = &fitness
|
||
}
|
||
if instrumentCount, ok := alphaData["instrumentCount"].(float64); ok {
|
||
metrics.InstrumentCount = &instrumentCount
|
||
}
|
||
if startDate, ok := alphaData["startDate"].(string); ok {
|
||
metrics.StartDate = &startDate
|
||
}
|
||
if endDate, ok := alphaData["endDate"].(string); ok {
|
||
metrics.EndDate = &endDate
|
||
}
|
||
|
||
return metrics, nil
|
||
}
|
||
|
||
// formatTime 将秒数格式化为"xx分xx秒"或"xx秒"的格式
|
||
// 对应Python中的format_time方法
|
||
func formatTime(seconds float64) string {
|
||
if seconds < 60 {
|
||
return fmt.Sprintf("%.2f秒", seconds)
|
||
}
|
||
minutes := int(seconds / 60)
|
||
remainingSeconds := seconds - float64(minutes*60)
|
||
return fmt.Sprintf("%d分%.2f秒", minutes, remainingSeconds)
|
||
}
|
||
|
||
// Close 关闭客户端，释放资源
|
||
// 对应Python中的close方法
|
||
func (c *BrainClient) Close() error {
|
||
// Go的http.Client不需要显式关闭
|
||
// 但如果需要清理其他资源，可以在这里添加
|
||
c.httpClient = nil
|
||
c.isLoggedIn = false
|
||
return nil
|
||
}
|
||
|
||
// 辅助函数：打印分隔线（方便使用）
|
||
func printSeparator() {
|
||
fmt.Println(strings.Repeat("=", 60))
|
||
}
|