量化交易学习-week2

1
import plotly.express as px
2
import plotly.graph_objects as go
3
import polars as pl
4

5
df = pl.read_parquet("1.parquet")
6
status = []
7
money = 100_000
8
volumns = 0
9
total = money
10
is_holding = False
11

12
for each_line in df.iter_rows():
13
    action = None
14
    date, *_, close, _, MA5, MA20 = each_line
15
    if MA20 is None or MA5 is None:
16
        continue
17
    if MA5 > MA20 and not is_holding:
18
        is_holding = True
19
        action = "Buy"
20
        volumns = money // (close * 100) * 100
21
        money = money - volumns * close
22

23
    elif MA5 < MA20 and is_holding:
24
        is_holding = False
25
        action = "Sell"
26
        money = money + volumns * close
27
        volumns = 0
28

29
    total = money + volumns * close
30
    status.append(
31
        {
32
            "date": date,
33
            "volumns": volumns,
34
            "money": money,
35
            "total": total,
36
            "action": action,
37
        }
38
    )
39

40
status_df = pl.DataFrame(status).sort(by="date")
41

42
initial_total = status_df["total"][0]
43
final_total = status_df["total"][-1]
44
total_return = (final_total / initial_total) - 1
45
status_df = status_df.with_columns(pl.col("total").cum_max().alias("peak"))
46
status_df = status_df.with_columns(
47
    ((pl.col("total") - pl.col("peak")) / pl.col("peak")).alias("drawdown_pct")
48
)
49
max_drawdown = status_df["drawdown_pct"].min()
50
total_days = len(status_df)
51
TRADING_DAYS_PER_YEAR = 252
52
annualized_return = (1 + total_return) ** (TRADING_DAYS_PER_YEAR / total_days) - 1
53

54
print("\n=== 回测结果 ===")
55
print(f"初始资金: {initial_total:,.2f}")
56
print(f"最终资金: {final_total:,.2f}")
57
print(f"总收益率: {total_return:.2%}")
58
print(f"最大回测: {max_drawdown:.2%}")
59
print(f"年化收益: {annualized_return:.2%}")
60

61
fig = px.line(status_df, x="date", y="total")
62
fig.show()