297 lines
12 KiB
Python
297 lines
12 KiB
Python
# main.py
|
||
|
||
import asyncio
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.optim as optim
|
||
from data.live_data import LiveDataManager
|
||
from model.transformer import Transformer
|
||
from training.train import train
|
||
from data.data_utils import preprocess_data # Import preprocess_data
|
||
import ccxt.async_support as ccxt
|
||
import time
|
||
import os
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
from model.trading_model import TradingModel
|
||
from training.rl_agent import ContinuousRLAgent, ReplayBuffer
|
||
from training.train_historical import train_on_historical_data, load_best_checkpoint, save_candles_cache, CACHE_FILE, BEST_DIR
|
||
from data.data_utils import get_aligned_candle_with_index, get_features_for_tf
|
||
|
||
async def main():
|
||
symbol = 'BTC/USDT'
|
||
data_manager = LiveDataManager(symbol)
|
||
|
||
# Model parameters (adjust for ~1B parameters)
|
||
input_dim = 6 + len([5, 10, 20, 60, 120, 200]) # OHLCV + EMAs
|
||
d_model = 512
|
||
num_heads = 8
|
||
num_layers = 6
|
||
d_ff = 2048
|
||
dropout = 0.1
|
||
|
||
model = Transformer(input_dim, d_model, num_heads, num_layers, d_ff, dropout)
|
||
|
||
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
|
||
|
||
# Define loss functions
|
||
criterion_candle = nn.MSELoss()
|
||
criterion_volume = nn.MSELoss() # Consider a different loss for volume if needed
|
||
criterion_ticks = nn.MSELoss()
|
||
|
||
# Check for CUDA availability and set device
|
||
if torch.cuda.is_available():
|
||
device = torch.device('cuda')
|
||
print("Using CUDA")
|
||
else:
|
||
device = torch.device('cpu')
|
||
print("Using CPU")
|
||
try:
|
||
await train(model, data_manager, optimizer, criterion_candle, criterion_volume, criterion_ticks, num_epochs=10, device=device)
|
||
except KeyboardInterrupt:
|
||
print("Training stopped manually.")
|
||
finally:
|
||
await data_manager.close()
|
||
|
||
# -------------------------------------
|
||
# Main Asynchronous Function for Training & Charting
|
||
# -------------------------------------
|
||
async def main_backtest():
|
||
symbol = 'BTC/USDT'
|
||
# Define timeframes: we'll use 5 different ones.
|
||
timeframes = ["1m", "5m", "15m", "1h", "1d"]
|
||
now = int(time.time() * 1000)
|
||
# Use the base timeframe period of 1500 candles. For 1m, that is 1500 minutes.
|
||
period_ms = 1500 * 60 * 1000
|
||
since = now - period_ms
|
||
end_time = now
|
||
|
||
# Initialize exchange using MEXC (or your preferred exchange).
|
||
mexc_api_key = os.environ.get('MEXC_API_KEY', 'YOUR_API_KEY')
|
||
mexc_api_secret = os.environ.get('MEXC_API_SECRET', 'YOUR_SECRET_KEY')
|
||
exchange = ccxt.mexc({
|
||
'apiKey': mexc_api_key,
|
||
'secret': mexc_api_secret,
|
||
'enableRateLimit': True,
|
||
})
|
||
|
||
candles_dict = {}
|
||
for tf in timeframes:
|
||
print(f"Fetching historical data for timeframe {tf}...")
|
||
candles = await fetch_historical_data(exchange, symbol, tf, since, end_time, batch_size=500)
|
||
candles_dict[tf] = candles
|
||
|
||
# Optionally, save the multi-timeframe cache.
|
||
save_candles_cache(CACHE_FILE, candles_dict)
|
||
|
||
# Create the backtest environment using multi-timeframe data.
|
||
env = BacktestEnvironment(candles_dict, base_tf="1m", timeframes=timeframes)
|
||
|
||
# Neural Network dimensions: each timeframe produces 7 features.
|
||
input_dim = len(timeframes) * 7 # 7 features * 5 timeframes = 35.
|
||
hidden_dim = 128
|
||
output_dim = 3 # Actions: SELL, HOLD, BUY.
|
||
|
||
model = TradingModel(input_dim, hidden_dim, output_dim)
|
||
optimizer = optim.Adam(model.parameters(), lr=1e-4)
|
||
replay_buffer = ReplayBuffer(capacity=10000)
|
||
rl_agent = ContinuousRLAgent(model, optimizer, replay_buffer, batch_size=32, gamma=0.99)
|
||
|
||
# Load best checkpoint if available.
|
||
load_best_checkpoint(model, BEST_DIR)
|
||
|
||
# Train the agent over the historical period.
|
||
num_epochs = 10 # Adjust as needed.
|
||
train_on_historical_data(env, rl_agent, num_epochs=num_epochs, epsilon=0.1)
|
||
|
||
# Run a final simulation (without exploration) to record trade history.
|
||
state = env.reset(clear_trade_history=True)
|
||
done = False
|
||
cumulative_reward = 0.0
|
||
while not done:
|
||
action = rl_agent.act(state, epsilon=0.0)
|
||
state, reward, next_state, done = env.step(action)
|
||
cumulative_reward += reward
|
||
state = next_state
|
||
print("Final simulation cumulative profit:", cumulative_reward)
|
||
|
||
# Evaluate trade performance.
|
||
trades = env.trade_history
|
||
num_trades = len(trades)
|
||
num_wins = sum(1 for trade in trades if trade["pnl"] > 0)
|
||
win_rate = (num_wins / num_trades * 100) if num_trades > 0 else 0.0
|
||
total_profit = sum(trade["pnl"] for trade in trades)
|
||
print(f"Total trades: {num_trades}, Wins: {num_wins}, Win rate: {win_rate:.2f}%, Total Profit: {total_profit:.4f}")
|
||
|
||
# Plot chart with buy/sell markers on the base timeframe ("1m").
|
||
plot_trade_history(candles_dict["1m"], trades)
|
||
|
||
await exchange.close()
|
||
|
||
# -------------------------------------
|
||
# Historical Data Fetching Function (for a given timeframe)
|
||
# -------------------------------------
|
||
async def fetch_historical_data(exchange, symbol, timeframe, since, end_time, batch_size=500):
|
||
candles = []
|
||
since_ms = since
|
||
while True:
|
||
try:
|
||
batch = await exchange.fetch_ohlcv(symbol, timeframe=timeframe, since=since_ms, limit=batch_size)
|
||
except Exception as e:
|
||
print(f"Error fetching historical data for {timeframe}:", e)
|
||
break
|
||
if not batch:
|
||
break
|
||
for c in batch:
|
||
candle_dict = {
|
||
'timestamp': c[0],
|
||
'open': c[1],
|
||
'high': c[2],
|
||
'low': c[3],
|
||
'close': c[4],
|
||
'volume': c[5]
|
||
}
|
||
candles.append(candle_dict)
|
||
last_timestamp = batch[-1][0]
|
||
if last_timestamp >= end_time:
|
||
break
|
||
since_ms = last_timestamp + 1
|
||
print(f"Fetched {len(candles)} candles for timeframe {timeframe}.")
|
||
return candles
|
||
|
||
# -------------------------------------
|
||
# Backtest Environment with Multi-Timeframe State
|
||
# -------------------------------------
|
||
class BacktestEnvironment:
|
||
def __init__(self, candles_dict, base_tf="1m", timeframes=None):
|
||
self.candles_dict = candles_dict # dict of timeframe: candles_list
|
||
self.base_tf = base_tf
|
||
if timeframes is None:
|
||
self.timeframes = [base_tf] # fallback to single timeframe
|
||
else:
|
||
self.timeframes = timeframes
|
||
self.trade_history = [] # record of closed trades
|
||
self.current_index = 0 # index on base_tf candles
|
||
self.position = None # active position record
|
||
|
||
def reset(self, clear_trade_history=True):
|
||
self.current_index = 0
|
||
self.position = None
|
||
if clear_trade_history:
|
||
self.trade_history = []
|
||
return self.get_state(self.current_index)
|
||
|
||
def get_state(self, index):
|
||
"""Construct the state as the concatenated features of all timeframes.
|
||
For each timeframe, find the aligned candle for the base timeframe’s timestamp."""
|
||
state_features = []
|
||
base_candle = self.candles_dict[self.base_tf][index]
|
||
base_ts = base_candle["timestamp"]
|
||
for tf in self.timeframes:
|
||
candles_list = self.candles_dict[tf]
|
||
# Get the candle from this timeframe that is closest to (and <=) base_ts.
|
||
aligned_index, _ = get_aligned_candle_with_index(candles_list, base_ts)
|
||
features = get_features_for_tf(candles_list, aligned_index, period=10)
|
||
state_features.extend(features)
|
||
return np.array(state_features, dtype=np.float32)
|
||
|
||
def step(self, action):
|
||
"""
|
||
Simulate a trading step based on the base timeframe.
|
||
- If not in a position and action is BUY (2), record entry at next candle's open.
|
||
- If in a position and action is SELL (0), record exit at next candle's open, computing PnL.
|
||
Returns: (current_state, reward, next_state, done)
|
||
"""
|
||
base_candles = self.candles_dict[self.base_tf]
|
||
if self.current_index >= len(base_candles) - 1:
|
||
return self.get_state(self.current_index), 0.0, None, True
|
||
|
||
current_state = self.get_state(self.current_index)
|
||
next_index = self.current_index + 1
|
||
next_state = self.get_state(next_index)
|
||
current_candle = base_candles[self.current_index]
|
||
next_candle = base_candles[next_index]
|
||
reward = 0.0
|
||
|
||
# Action mapping: 0 -> SELL, 1 -> HOLD, 2 -> BUY.
|
||
if self.position is None:
|
||
if action == 2: # BUY signal: enter position at next candle's open.
|
||
entry_price = next_candle["open"]
|
||
self.position = {"entry_price": entry_price, "entry_index": self.current_index}
|
||
else:
|
||
if action == 0: # SELL signal: close position at next candle's open.
|
||
exit_price = next_candle["open"]
|
||
reward = exit_price - self.position["entry_price"]
|
||
trade = {
|
||
"entry_index": self.position["entry_index"],
|
||
"entry_price": self.position["entry_price"],
|
||
"exit_index": next_index,
|
||
"exit_price": exit_price,
|
||
"pnl": reward
|
||
}
|
||
self.trade_history.append(trade)
|
||
self.position = None
|
||
|
||
self.current_index = next_index
|
||
done = (self.current_index >= len(base_candles) - 1)
|
||
return current_state, reward, next_state, done
|
||
|
||
# -------------------------------------
|
||
# Chart Plotting: Trade History & PnL
|
||
# -------------------------------------
|
||
def plot_trade_history(candles, trade_history):
|
||
close_prices = [candle["close"] for candle in candles]
|
||
x = list(range(len(close_prices)))
|
||
plt.figure(figsize=(12, 6))
|
||
plt.plot(x, close_prices, label="Close Price", color="black", linewidth=1)
|
||
|
||
# Use these flags so that the label "BUY" or "SELL" is only shown once in the legend.
|
||
buy_label_added = False
|
||
sell_label_added = False
|
||
|
||
for trade in trade_history:
|
||
in_idx = trade["entry_index"]
|
||
out_idx = trade["exit_index"]
|
||
in_price = trade["entry_price"]
|
||
out_price = trade["exit_price"]
|
||
pnl = trade["pnl"]
|
||
|
||
# Plot BUY marker ("IN")
|
||
if not buy_label_added:
|
||
plt.plot(in_idx, in_price, marker="^", color="green", markersize=10, label="BUY (IN)")
|
||
buy_label_added = True
|
||
else:
|
||
plt.plot(in_idx, in_price, marker="^", color="green", markersize=10)
|
||
plt.text(in_idx, in_price, " IN", color="green", fontsize=8, verticalalignment="bottom")
|
||
|
||
# Plot SELL marker ("OUT")
|
||
if not sell_label_added:
|
||
plt.plot(out_idx, out_price, marker="v", color="red", markersize=10, label="SELL (OUT)")
|
||
sell_label_added = True
|
||
else:
|
||
plt.plot(out_idx, out_price, marker="v", color="red", markersize=10)
|
||
plt.text(out_idx, out_price, " OUT", color="red", fontsize=8, verticalalignment="top")
|
||
|
||
# Annotate the PnL near the SELL marker.
|
||
plt.text(out_idx, out_price, f" {pnl:+.2f}", color="blue", fontsize=8, verticalalignment="bottom")
|
||
|
||
# Choose line color based on profitability.
|
||
if pnl > 0:
|
||
line_color = "green"
|
||
elif pnl < 0:
|
||
line_color = "red"
|
||
else:
|
||
line_color = "gray"
|
||
# Draw a dotted line between the buy and sell points.
|
||
plt.plot([in_idx, out_idx], [in_price, out_price], linestyle="dotted", color=line_color)
|
||
|
||
plt.title("Trade History with PnL")
|
||
plt.xlabel("Base Candle Index (1m)")
|
||
plt.ylabel("Price")
|
||
plt.legend()
|
||
plt.grid(True)
|
||
plt.show()
|
||
|
||
if __name__ == '__main__':
|
||
asyncio.run(main_backtest())
|